DiffDoctor: Diagnosing Image Diffusion Models Before Treating

TL;DR

DiffDoctor presents a novel diagnose-then-treat framework that uses pixel-level artifact maps to actively steer image diffusion models away from generating artifacts. A robust artifact detector is trained on 1M+ samples through class-balanced data collection and human-in-the-loop labeling, producing per-pixel confidence maps $C(\cdot)$. These maps back-propagate into the diffusion process via a pixel-level loss $\mathcal{L}_{\text{pixel}}$, optionally complemented by an offline regularization term $\mathcal{L}_{\text{offline}}$, to yield improved artifact suppression while maintaining overall image quality on unseen prompts. The approach demonstrates reduced artifact frequencies across multiple backbones (e.g., FLUX.1, SDXL, Kolors) and outperforms baselines in both objective metrics and human studies, with demonstrated applicability to DreamBooth. Limitations include challenges with semantically complex artifacts, motivating future work on richer artifact reasoning within the detector.

Abstract

In spite of recent progress, image diffusion models still produce artifacts. A common solution is to leverage the feedback provided by quality assessment systems or human annotators to optimize the model, where images are generally rated in their entirety. In this work, we believe problem-solving starts with identification, yielding the request that the model should be aware of not just the presence of defects in an image, but their specific locations. Motivated by this, we propose DiffDoctor, a two-stage pipeline to assist image diffusion models in generating fewer artifacts. Concretely, the first stage targets developing a robust artifact detector, for which we collect a dataset of over 1M flawed synthesized images and set up an efficient human-in-the-loop annotation process, incorporating a carefully designed class-balance strategy. The learned artifact detector is then involved in the second stage to optimize the diffusion model by providing pixel-level feedback. Extensive experiments on text-to-image diffusion models demonstrate the effectiveness of our artifact detector as well as the soundness of our diagnose-then-treat design.

Figures (8)

• Figure 1: Illustrations of DiffDoctor. We train a robust detector to localize the artifacts (diagnosing) and provide pixel-level feedback to optimize the diffusion model (treating). After tuning on limited samples, the diffusion model generates significantly fewer artifacts on unseen prompts while maintaining the quality.
• Figure 2: Pipeline of DiffDoctor. The first part shows the training of an artifact detector -- the doctor. Starting with the initial dataset, the artifact detector is trained in a humans-in-a-loop manner. The second part shows our diagnose-then-treat design, where the patient -- a trainable diffusion model, is prompted to synthesize images. Then the frozen artifact detector diagnoses its result by predicting the artifact maps, on which it treats the patient by minimizing the per-pixel artifact confidence to back-propagate to the diffusion model.
• Figure 3: Qualitative ablation study of artifact detectors. We visualize the artifact maps predicted by the artifact detector on our hard benchmark by headmaps.
• Figure 4: Mode collapse due to naive artifact detector. The treating process collapses into blurriness. We further visualize the training and evaluation time artifact confidence curve.
• Figure 5: Qualitative ablation on offline regularization. These images are synthesized by FLUX.1 during treating with or without offline regularization in evaluation time. We frame the synthesized image in the model collapse stage with a red box.