DIAMOND: Directed Inference for Artifact Mitigation in Flow Matching Models

TL;DR

DIAMOND tackles artifacts in text-to-image generation by introducing a training-free, inference-time trajectory correction mechanism for Rectified Flow and Diffusion models. It builds on on-the-fly estimation of the clean latent $\\hat{x}_{0,t}$ and a differentiable Artifact Detector to compute a gradient-based artifact loss $\\mathcal{L}_a$, guiding the sampling trajectory with updates $x_{t-\Delta t} = x_t - \Delta t\\,v_\theta(x_t,t) - \\delta_t$ where $\\delta_t$ is normalized by $\\nabla_{x_t}\\mathcal{L}_a$. The method generalizes to standard diffusion models, uses a schedule for the correction strength $\\lambda_t$, and supports corrections within a restricted time interval to preserve global structure. Empirically, DIAMOND substantially reduces artifact frequency and pixel-level artifacts across multiple datasets and models while maintaining prompt fidelity and semantic alignment, without any training or weight modifications. This offers a practical, real-time approach to artifact mitigation with broad applicability to modern generative architectures.

Abstract

Despite impressive results from recent text-to-image models like FLUX, visual and anatomical artifacts remain a significant hurdle for practical and professional use. Existing methods for artifact reduction, typically work in a post-hoc manner, consequently failing to intervene effectively during the core image formation process. Notably, current techniques require problematic and invasive modifications to the model weights, or depend on a computationally expensive and time-consuming process of regional refinement. To address these limitations, we propose DIAMOND, a training-free method that applies trajectory correction to mitigate artifacts during inference. By reconstructing an estimate of the clean sample at every step of the generative trajectory, DIAMOND actively steers the generation process away from latent states that lead to artifacts. Furthermore, we extend the proposed method to standard Diffusion Models, demonstrating that DIAMOND provides a robust, zero-shot path to high-fidelity, artifact-free image synthesis without the need for additional training or weight modifications in modern generative architectures. Code is available at https://gmum.github.io/DIAMOND/

Figures (24)

• Figure 1: We propose DIAMOND, an inference-time trajectory correction mechanism to mitigate artifacts in Rectified Flow and Diffusion Models. First, we generate an image $\mathcal{D}({x}_{1})$ from the initial probability distribution sample, which then undergoes the generative process. The Base Trajectory (red) leads to an image containing artifacts, whereas our corrected DIAMOND Trajectory (green) results in the artifact-free image. At timestep $t$, we estimate the final image $\hat{x}_{0,t}$ (gray dashed line) and use an Artifact Detector to apply a gradient-based trajectory correction (purple arrows), shifting it away from artifact region.
• Figure 2: Artifact generation in FLUX.2, FLUX.1 and SDXL, with corrections applied using DIAMOND. While these models produce visible artifacts, DIAMOND detects and corrects them automatically without additional training.
• Figure 3: Overview of DIAMOND. Our technique employs an inference-time pipeline specifically designed to mitigate artifacts without requiring additional training. Artifact suppression is achieved by correcting the trajectory using gradients derived from a pixel-wise segmentation loss. During inference, the latent representation is iteratively updated using controlled shifts along the trajectory, enabling progressive correction of artifact-prone regions.
• Figure 4: Qualitative comparisons on images from the people dataset for FLUX.1 [dev]. Our model uses 10 inference steps and here does not apply base-model identity preservation. Even without regularization $\mathcal{L}_{\text{rec}}$, the model diminishes artifacts without significantly altering the generated image, unlike HandsXL.
• Figure 5: Comparison of images generated using estimated clean data latents $\hat{x}_{0,t}$ and noisy latents $x_t$ during intermediate inference steps for FLUX.1 [dev]. Artifact masks are shown to highlight the effectiveness of using the estimated latents. Contrary to using $x_t$, the detector is able to find artifact regions even during early stages when using $\hat{x}_{0,t}$ estimation.