DiffRGD: An Inference-Time Diffusion Guidance
Through Riemannian Gradient Descent

     

---

ECCV 2026

Jia-Wei Liao^1,4, Li-Xuan Peng², Mei-Heng Yueh³, Min Sun², Cheng-Fu Chou¹, Jun-Cheng Chen^4
1National Taiwan University, ²National Tsing Hua University,
³National Taiwan Normal University, ⁴Academia Sinica

💡 TL;DR

DiffRGD is a plug-and-play, inference-time guidance framework for pre-trained diffusion models. Existing guidance methods inject loss gradients directly into the sampling process, drifting the latent away from its per-step Gaussian distribution and degrading sample quality. DiffRGD instead formulates each DDIM step as a constrained optimization problem on a spherical manifold induced by the latent Gaussian distribution (via a polar decomposition of the isotropic Gaussian) and solves it with Riemannian Gradient Descent — projecting the guidance gradient onto the tangent space of the sphere and retracting the update back onto it. This preserves the latent distribution while steering the sample toward the condition, with a convergence guarantee to a stationary point.

---

🔧 Installation Guide

1. Clone this repository

git clone https://github.com/jwliao1209/DiffRGD
cd DiffRGD

2. Set up the environment with uv (Python ≥ 3.13)

pip install uv
uv sync

All experiments run on a single GPU (we used an NVIDIA RTX 4090).

Show Repository Structure

DiffRGD/
├── src/
│   ├── methods/               # Guidance methods: DiffRGD (ours), DPS, MPGD, DSG, ADMMDiff
│   ├── pipelines/
│   │   ├── pipeline_image_sampler.py      # Pixel-space DDIM pipeline with guidance
│   │   └── pipeline_sd_image_sampler.py   # Stable Diffusion pipeline with latent guidance
│   ├── losses/                # Guidance losses (inverse problem, face parsing/sketch/ID, style)
│   ├── operators.py           # Degradation operators A(.) for inverse problems
│   ├── metrics/               # PSNR, SSIM, LPIPS, FID, KID, mIoU, CLIP/style scores
│   ├── dataset.py             # Image and measurement datasets
│   └── model.py               # OpenAI guided-diffusion UNet (Diffusers-compatible)
├── configs/
│   ├── operators/             # Degradation operator configurations
│   └── prompts.json           # Text prompts for style-guided generation
├── evaluation/                # Metric computation for each experiment
├── scripts/                   # Downloads, checkpoint conversion, experiment runners, plots
├── style_images/              # WikiArt style references for style-guided generation
├── generate_dataset.py        # Build measurement datasets
├── run_inverse_problem.py     # Image restoration (Sec. 4.1, 4.2)
├── run_control.py             # Conditional generation (Sec. 4.3)
└── run_sd_style_transfer.py   # Style-guided generation with Stable Diffusion (Appendix E.2)

🔗 Pre-trained Models

Model	Used for	How to get it
FFHQ 256 DDPM (DPS)	FFHQ restoration	bash scripts/download_model_checkpoint.sh then convert (below)
ImageNet 256 DDPM (guided-diffusion)	ImageNet restoration / denoising	same as above
CelebA-HQ 256 DDPM	Conditional generation	auto-downloaded (google/ddpm-ema-celebahq-256)
Stable Diffusion v1.5	SD super-resolution / style transfer	auto-downloaded (runwayml/stable-diffusion-v1-5)
BiSeNet face parser	Segmentation guidance	place at src/losses/checkpoints/bisnet.pth (from face-parsing.PyTorch)
AODA sketch model	Sketch guidance	place at src/losses/checkpoints/sketchnet.pth (from FreeDoM's model zoo)
ArcFace IR-SE50	FaceID guidance	bash scripts/download_face_id.sh

Convert the raw OpenAI checkpoints into Diffusers pipelines (creates checkpoints/openai_ffhq and checkpoints/openai_imagenet):

bash scripts/download_model_checkpoint.sh
uv run scripts/convert_openai_model_to_diffuser.py -d ffhq
uv run scripts/convert_openai_model_to_diffuser.py -d imagenet

📦 Datasets

Place validation images under data/<dataset>/ (e.g. data/ffhq/, data/imagenet/, data/celeba_hq/). Download helpers:

uv run scripts/download_ffhq.py                  # FFHQ 256 (Kaggle)
uv run scripts/split_ffhq_valid.py               # split FFHQ into train/valid
uv run scripts/download_celeba_hq_dataset.py     # CelebA-HQ 256 (Kaggle)
bash scripts/download_imagenet_val_data.sh       # ImageNet validation set

🖼️ Image Restoration

Each task is a noisy linear inverse problem y = A(x) + n with n ~ N(0, 0.05²I). Tasks: random_inpainting (70% random mask), box_inpainting, super_resolution (4× bicubic), gaussian_blur (31×31, σ=3), motion_blur (61×61, intensity 0.5), and denoising_<sigma>.

Generate measurements, then run a method:

bash scripts/generate_dataset.sh 150 ffhq

uv run run_inverse_problem.py \
    -t random_inpainting \        # task
    -d ffhq \                     # ffhq | imagenet | celeba_hq
    -m diffrgd \                  # dps | mpgd | dsg | admmdiff | diffrgd
    -n 150 \                      # number of samples
    --num_inference_steps 100 \   # DDIM steps
    --step_size 5 \               # guidance strength eta_t
    -i 1 \                        # guidance interval
    --inner_max_iter 3 \          # RGD inner iterations K
    --noise gaussian

Reproduce the paper's tables with the tuned hyperparameters:

bash scripts/run_ffhq_t1000.sh 150      # Table 1: FFHQ, 1000 DDIM steps
bash scripts/run_ffhq_t100.sh 150       # Table 2: FFHQ, 100 DDIM steps
bash scripts/run_imagenet_t100.sh 150   # Table 3: ImageNet, 100 DDIM steps
bash scripts/run_denoising.sh 150       # Table 9: ImageNet denoising (σ = 0.1/0.2/0.3)

Evaluate (PSNR / SSIM / LPIPS / FID):

bash scripts/run_eval_ffhq.sh 150
bash scripts/run_eval_imagenet.sh 150

🔍 Super-Resolution with Stable Diffusion (8× / 12×)

FFHQ 1024×1024 images are downsampled to 64×64 and super-resolved to 512×512 (8×) or 768×768 (12×) with Stable Diffusion v1.5:

bash scripts/run_sd_super_resolution.sh 100
uv run evaluation/evaluate_sd_super_resolution.py -m diffrgd --resolution 512   # or 768

🎯 Conditional Generation

Face generation on CelebA-HQ guided by a segmentation map (BiSeNet), a sketch (AODA), or a FaceID embedding (ArcFace):

uv run run_control.py \
    -t segmentation \             # segmentation | sketch | face_id
    -m diffrgd \                  # dps (= FreeDoM) | mpgd | dsg | admmdiff | diffrgd
    -n 150 \
    --step_size 30 \
    --inner_max_iter 3

Reproduce Table 5 and evaluate (mIoU or condition-L2 / FID / KID):

bash scripts/run_control.sh 150
bash scripts/run_eval_control.sh 150

🎨 Style-Guided Generation

Style transfer with Stable Diffusion v1.5, using a CLIP Gram-matrix style loss over 5 WikiArt style images (style_images/) × 20 prompts (configs/prompts.json):

bash scripts/run_sd_style_transfer.sh
uv run evaluation/evaluate_sd_style_transfer.py -m diffrgd    # CLIP-score / Style-score (Table 10)

🙏 Acknowledgements

We are grateful for the foundational code provided by DPS, FreeDoM, MPGD, DSG, ADMMDiff, and 🤗 Diffusers. Utilizing their resources implies agreement with their respective licenses. This research is supported by the National Science and Technology Council, Taiwan, and Academia Sinica.

📚 Citation

If you use our work or our implementation in this repo or find them helpful, please consider giving a citation.

@inproceedings{liao2026diffrgd,
  title     = {DiffRGD: An Inference-Time Diffusion Guidance Through Riemannian Gradient Descent},
  author    = {Liao, Jia-Wei and Peng, Li-Xuan and Yueh, Mei-Heng and Sun, Min and Chou, Cheng-Fu and Chen, Jun-Cheng},
  booktitle = {Proceedings of the European Conference on Computer Vision (ECCV)},
  year      = {2026},
}