SURE-Tune

Reference implementation of SURE-Tune: A Principled Statistical Framework for Unsupervised Test-Time Adaptation of Diffusion Models.

SURE-Tune adapts a pretrained diffusion prior to a single test instance using Stein's unbiased risk estimate (SURE). The codebase includes restoration experiments, OOD detection experiments, solver/timestep/rank ablations, and paper-table exporters.

Repository Layout

algorithms/
  sure_tune.py          # SURE-Tune, effective_sigma, compute_tmin
  diffpir.py            # DiffPIR baseline
  daps.py               # DAPS baseline
  sgps.py               # SGPS baseline
physics/
  blur.py               # Motion and Gaussian blur operators
  tomography.py         # CT/Radon operator wrapper
utils/
  metrics.py            # PSNR, SSIM, LPIPS
  dataset.py            # Generic dataset helpers
experiments/
  validate_deblur_table.py       # Deblur paper-table runs
  validate_ct_table.py           # CT paper-table runs
  tune_reconstruction.py         # DPS/SURE/MC-N2S/SGPS/DAPS tuning
  tune_diffpir.py                # DiffPIR tuning
  run_benchmark_restoration.py   # Restoration ablations
  run_benchmark_ood.py           # OOD detection benchmark
files/
  paper.tex, paper.bib, README_STY.md
tests/
  Unit tests for theory/code identities

Environment

On this machine the tested path is WSL with the diffinv conda environment:

cd /mnt/c/Users/sadeg/MSS/diffusion-inverse-problems
source ~/miniconda3/etc/profile.d/conda.sh
conda activate diffinv

If creating a fresh environment:

pip install -r requirements.txt

For faster Hugging Face downloads, set HF_TOKEN in the environment.

Data Paths

Deblur and OOD experiments download public datasets automatically.

The CT paper-table and tuning scripts use the external LoDoPaB path:

C:\Users\sadeg\MSS\Multilevel_Bilevel_DFO\dataset\LoDoPaB-CT
/mnt/c/Users/sadeg/MSS/Multilevel_Bilevel_DFO/dataset/LoDoPaB-CT

The CT scripts also import dataset.datasets.LoDoPaBCustom from:

C:\Users\sadeg\MSS\Multilevel_Bilevel_DFO
/mnt/c/Users/sadeg/MSS/Multilevel_Bilevel_DFO

Final Paper-Table Runs

Run these after hyperparameters are settled. Outputs go to experiments/results_visuals/deblur_table_validation/ and experiments/results_visuals/ct_table_validation/.

python -m experiments.validate_deblur_table \
  --task motion --kernel-size 31 --n-samples 50 \
  --noise-sigma 0.05 --paper-table --seed 0

python -m experiments.validate_deblur_table \
  --task gaussian --gaussian-sigma 7.0 --n-samples 50 \
  --noise-sigma 0.05 --paper-table --seed 0

python -m experiments.validate_ct_table \
  --angles 32 --n-samples 50 \
  --noise-sigma 0.05 --paper-table --seed 0

python -m experiments.validate_ct_table \
  --angle-range 0 120 --angles 32 --n-samples 50 \
  --noise-sigma 0.05 --paper-table --seed 0

Each run writes:

metrics_*.csv
summary_*.csv
table_*.tex

The validation scripts checkpoint after each completed sample. If a run is interrupted after it has started writing fresh metrics_*.csv files, rerun the same command with --resume to skip completed samples:

python -m experiments.validate_deblur_table \
  --task motion --kernel-size 31 --n-samples 50 \
  --noise-sigma 0.05 --paper-table --seed 0 --resume

Use --resume only with metrics produced by the same code/configuration. If the output directory contains old n=3 smoke results, start the final run without --resume or write to a fresh --out-dir.

DiffPIR Tuning

Use these when changing DiffPIR integration, CT normalization, or blur severity.

python -m experiments.tune_diffpir \
  --task motion --kernel-size 31 --n-samples 3 \
  --noise-sigma 0.05 --seed 0

python -m experiments.tune_diffpir \
  --task gaussian --gaussian-sigma 7.0 --n-samples 3 \
  --noise-sigma 0.05 --seed 0

python -m experiments.tune_diffpir \
  --task ct --angles 32 --n-samples 3 \
  --noise-sigma 0.05 --seed 0

python -m experiments.tune_diffpir \
  --task ct --angle-range 0 120 --angles 32 --n-samples 3 \
  --noise-sigma 0.05 --seed 0

Outputs:

experiments/results_visuals/diffpir_tune/

Reconstruction Tuning

Use this to tune DPS/SURE-Tune/MC-N2S/SGPS/DAPS reconstruction-phase hyperparameters.

python -m experiments.tune_reconstruction \
  --task motion --method all --n-samples 3 \
  --noise-sigma 0.05 --seed 0

python -m experiments.tune_reconstruction \
  --task gaussian --gaussian-sigma 7.0 --method all --n-samples 3 \
  --noise-sigma 0.05 --seed 0

python -m experiments.tune_reconstruction \
  --task ct --angles 32 --method all --n-samples 3 \
  --noise-sigma 0.05 --seed 0

python -m experiments.tune_reconstruction \
  --task ct --angle-range 0 120 --angles 32 --method all --n-samples 3 \
  --noise-sigma 0.05 --seed 0

For SGPS alpha sweeps:

python -m experiments.tune_reconstruction \
  --task motion --method sgps --n-samples 3 \
  --noise-sigma 0.05 --seed 0 \
  --alpha-grid 0 0.025 0.05 0.1

python -m experiments.tune_reconstruction \
  --task ct --angles 32 --method sgps --n-samples 3 \
  --noise-sigma 0.05 --seed 0 \
  --alpha-grid 0 0.025 0.05 0.1

Outputs:

experiments/results_visuals/reconstruction_tune/

Restoration Ablations

Convergence ablation:

python -m experiments.run_benchmark_restoration --ablation

LoRA-rank ablation:

python -m experiments.run_benchmark_restoration --ablation_lora

Optional second-prior runs with the LSUN Church DDPM:

python -m experiments.run_benchmark_restoration \
  --ablation --model-id google/ddpm-church-256

python -m experiments.run_benchmark_restoration \
  --ablation_lora --model-id google/ddpm-church-256

Outputs:

experiments/results_visuals/ablation_convergence.png
experiments/results_visuals/ablation_convergence_summary.csv
experiments/results_visuals/ablation_lora_rank.png
experiments/results_visuals/ablation_lora_rank.csv

Model-specific runs append a sanitized model-id suffix to the output filenames.

OOD Detection

Smoke test:

python -m experiments.run_benchmark_ood \
  --experiment cifar10 --n-samples 20 --batch-size 10 \
  --seed 0 --sure-fast --bootstrap 100 \
  --out-dir experiments/results_visuals/ood_detection_smoke

Final OOD run:

python -m experiments.run_benchmark_ood \
  --n-samples 500 --batch-size 20 \
  --seed 0 --bootstrap 1000 \
  --include-sure-variance-ablation

PowerShell logging wrapper:

wsl bash -lc "cd /mnt/c/Users/sadeg/MSS/diffusion-inverse-problems && source ~/miniconda3/etc/profile.d/conda.sh && conda activate diffinv && python -m experiments.run_benchmark_ood --n-samples 500 --batch-size 20 --seed 0 --bootstrap 1000 --include-sure-variance-ablation" 2>&1 | Tee-Object experiments/logs/ood_final_variance_20260506.log

Outputs:

experiments/results_visuals/ood_detection/ood_summary.csv
experiments/results_visuals/ood_detection/ood_scores.csv
experiments/results_visuals/ood_detection/ood_table.tex

Optional slow Jacobian baseline:

python -m experiments.run_benchmark_ood \
  --n-samples 500 --batch-size 20 \
  --seed 0 --bootstrap 1000 \
  --include-jacobian \
  --out-dir experiments/results_visuals/ood_detection_jacobian

The default OOD run uses matched priors only. The SVHN-ID experiment uses a CIFAR prior proxy and is therefore excluded by default; enable it only for a clearly labeled appendix run:

python -m experiments.run_benchmark_ood \
  --experiment svhn --include-svhn-proxy \
  --n-samples 500 --batch-size 20 --seed 0

Quick Sanity Checks

Compile the active experiment scripts:

python -m py_compile \
  experiments/validate_deblur_table.py \
  experiments/validate_ct_table.py \
  experiments/tune_reconstruction.py \
  experiments/tune_diffpir.py \
  experiments/run_benchmark_restoration.py \
  experiments/run_benchmark_ood.py

Run unit tests:

pytest tests/ -v

Check the worktree before committing:

git status --short
git diff --check

Notes on Generated Files

Large generated outputs live under:

experiments/results_visuals/
experiments/logs/
data/

Do not commit smoke-test folders or Python bytecode caches. Final paper tables should be copied from the generated table_*.tex files into files/paper.tex only after checking the matching summary_*.csv.

Citation

@inproceedings{salehi2026suretune,
  title     = {SURE-Tune: A Principled Statistical Framework for Unsupervised Test-Time Adaptation of Diffusion Models},
  author    = {Salehi, Mohammad Sadegh and others},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}

License

BSD 3-Clause. See LICENSE.