Diagnose-to-Generate: Eliminating Structural Distortions in Flow Models via Patch-Level Grounded Alignment

Abstract

Flow-based image generators achieve strong semantic synthesis but still produce localized structural distortions, such as malformed hands and limbs. Existing post-training methods typically optimize global scalar rewards, which provide weak spatial credit assignment and may alter well-formed regions while correcting local defects. We present \textit{Diagnose-to-Generate (D2G)}, a grounded alignment framework that separates visual diagnosis from generative optimization. We construct \textbf{Achilles-600K}, a region-annotated dataset for structural anomalies, and train \textbf{ARGUS} through supervised cold-start and Region-Aware Group Relative Policy Optimization to predict anomaly categories and bounding boxes. These diagnoses are converted into confidence-weighted spatial rewards for our \textbf{Patch-level Flow-GRPO}. We further identify that global stochastic exploration can shift semantic regions across group rollouts, making patch-wise relative advantages spatially inconsistent. Our \textbf{Trajectory-Preserving Stochastic Sampling (TPSS)} anchors early denoising with deterministic ODE trajectories and introduces stochasticity only at later steps, thereby limiting intra-group spatial drift while retaining local exploration. Experiments on FLUX.1-dev and SD3.5 show improved structural fidelity with preserved FID and CLIP scores. Consistent gains under HumanRefiner, zero-shot GPT-4.1, and blinded human evaluation further demonstrate that the improvements extend beyond the training reward model.