SANA-SR · NeurIPS 2026

Efficient One-Step Diffusion Restoration Model with Compact Token Compression and Linear Attention

Official implementation, accepted at NeurIPS 2026.

Paper · Model

SANA-SR combines a 32× deep-compression autoencoder, a pretrained linear-attention diffusion transformer, and LoRA adaptation for one-step real-world image super-resolution. Training uses reconstruction supervision, frozen-prior alignment, and adapter consistency. Prompt-aware structured pruning reduces deployment cost.

• Training, inference, evaluation, and structured pruning code are available.

git clone https://github.com/Seiun-Sky-debug/SANASR.git
cd SANASR
pip install -r requirements.txt

Prepare a SANA-600M-512 Diffusers checkpoint, the RAM checkpoint, and the DAPE checkpoint. Install the RAM implementation supplied by SeeSR, or point to an existing checkout:

export SEESR_ROOT=/path/to/SeeSR
# Optional: a local BERT tokenizer for offline RAM loading.
export RAM_BERT_PATH=/path/to/bert-base-uncased

Use shared datasets and pretrained models when available. Set experiment output paths to your result storage; datasets and checkpoints are supplied separately.

Set the paths in configs/sanasr_hqprompt.yaml, including the pretrained SANA model, HQ training images, RAM, and DAPE. The optional realsr_train_hq_dir can be an empty string. HQ images are used to extract the training prompts.

CUDA_VISIBLE_DEVICES=0 python core/train.py \
  --config configs/sanasr_hqprompt.yaml --gpu 0

The default configuration uses 512×512 crops, 4× SR, batch size 4, AdamW at 5e-5, 100K updates, rank-64 LoRA, and EMA decay 0.999. The reconstruction, perceptual, alignment, and consistency weights are 1, 2, 1, 1. The image loader supports paired LQ/HQ folders and online stochastic degradation.

With a trained LoRA adapter:

python core/inference_osediff_sana.py \
  --sana_path /path/to/Sana_600M_512px_diffusers \
  --lora_path /path/to/checkpoint/lora \
  --input /path/to/LQ --output /path/to/results/SR \
  --ram_path /path/to/ram_swin_large_14m.pth \
  --ram_ft_path /path/to/DAPE.pth \
  --prompt_suffix "clean, extremely detailed, best quality, sharp, high-resolution" \
  --gen_timestep 900 --keep_original_size --upscale_factor 4 --align_multiple 32

Standalone inference extracts prompts from LQ images. For the paper's HQ-prompt benchmark protocol, use the evaluator below with paired HQ references.

python core/evaluate_osediff_sana.py \
  --sana_path /path/to/Sana_600M_512px_diffusers \
  --lora_path /path/to/checkpoint/lora \
  --lq_dir /path/to/LQ --hq_dir /path/to/HQ \
  --output_dir /path/to/results/benchmark \
  --ram_path /path/to/ram_swin_large_14m.pth \
  --ram_ft_path /path/to/DAPE.pth --prompt_source hq \
  --prompt_suffix "clean, extremely detailed, best quality, sharp, high-resolution" \
  --gen_timestep 900 --keep_original_size --upscale_factor 4 --align_multiple 32 \
  --eval_mode full --save_prompts \
  --save_csv /path/to/results/benchmark/metrics.csv

Evaluate DIV2K-Val (3,000 pairs), RealSR (100 pairs), and DRealSR (93 pairs) separately with the same checkpoint and settings. Full mode scores each supplied image at its complete size. PSNR and SSIM use the Y channel; LPIPS, DISTS, MANIQA, MUSIQ, CLIP-IQA, and NIQE use RGB. HQ prompting uses reference images and should be reported explicitly.

For previously generated images, supply --sr_dir /path/to/SR --hq_dir /path/to/HQ --eval_mode full instead of the model and inference arguments.

Set the calibration-data and pretrained-model paths in configs/sanasr_prune.yaml:

python core/prune_obsdiff_structured_osediff_sana.py \
  --config configs/sanasr_prune.yaml \
  --init_lora_path /path/to/checkpoint/lora \
  --output_path /path/to/results/sanasr_pruned \
  --target_params_billion 0.35 --calib_steps 400 --keep_first_last

The saved pipeline includes the merged adapter and retained transformer blocks. Use that directory as --sana_path and omit --lora_path for inference or evaluation.

@inproceedings{qiao2026sanasr,
  title={Efficient One-Step Diffusion Restoration Model with Compact Token Compression and Linear Attention},
  author={Qiao, Bingtian and Shi, Yue and Zhou, Yingjie and Guo, Yong and Zhai, Guangtao and Cao, Jiezhang},
  booktitle={Advances in Neural Information Processing Systems},
  year={2026}
}

Built on SANA, OSEDiff, SeeSR, Diffusers, PEFT, and PyIQA. External models and datasets retain their respective licenses. Repository code is available under the MIT license.

Model weights: Coming soon.