[ICCV 2025] DiffDoctor: Diagnosing Image Diffusion Models Before Treating

Yiyang Wang · Xi Chen · Xiaogang Xu · Sihui Ji · Yu Liu · Yujun Shen · Hengshuang Zhao


The University of Hong Kong   |   Tongyi Lab |   Ant Financial Services Group |   The Chinese University of Hong Kong

TODOs

• [DONE✅] Release the artifact detector checkpoint (in checkpoints/).
• [DONE✅] Release the code for training the diffusion model using the artifact detector.

Installation

Install with pip:

pip install -r requirements.txt

Download Checkpoints

• Download the artifact detector checkpoint using git lfs and put it in checkpoints/.
• Download Flux.1-Schnell or Flux.1-Dev checkpoint from HuggingFace.

Diagnose: Inference on the Artifact Detector

After downloading and placing the checkpoint of the artifact detector, run the following command to use the demo code for the artifact detector.

python ad_inference.py

Treat: Training the Diffusion Model

Checklist before running the training code:

• Downloading and placing the checkpoint of the artifact detector (ad_pytorch_model.bin) into folder checkpoints/.
• Modify config/settings.py. You should change all the checkpoint paths to the backbone diffusion model. All these paths are annotated with "YOUR_PATH/xxx" (e.g., change 'YOUR_PATH/FLUX.1-dev' to 'black-forest-labs/flux.1-dev').
• [Optional] Further modify config/settings.py to change the training parameters, the used prompts, etc. For example, you can change config.prompt_fn or config.eval_prompt_fn to any txt file in src/prompt_files.

After preparation, run the following command to train the diffusion model using the artifact detector.

accelerate launch --config_file config/accelerate_single.yaml train_diffusion_model.py --config config/settings.py:ad 

We also include the training code by using feedback from HPS v2. After downloading the checkpoint HPS_v2_compressed.pt and modify config.hps_ckpt_path, you can run the training code by:

accelerate launch --config_file config/accelerate_single.yaml train_diffusion_model.py --config config/settings.py:hps

Acknowledgements

This project is developped on the codebase of AlignProp. We appreciate this great work!

Citation

If you find this codebase useful for your research, please use the following entry.

@article{wang2025diffdoctor,
  title={DiffDoctor: Diagnosing Image Diffusion Models Before Treating},
  author={Wang, Yiyang and Chen, Xi and Xu, Xiaogang and Ji, Sihui and Liu, Yu and Shen, Yujun and Zhao, Hengshuang},
  journal={ICCV},
  year={2025}
}