Spectral Alignment (SPA)

This repository contains the official implementation of our paper "Spectral Prior for Reducing Exposure Bias in Diffusion Models (ECCV2026)"

Inference-time correction of spectral mismatch in diffusion models. This directory contains the sampling code for Stable Diffusion 2.0 and SDXL.

Files

• spectral_alignment.py — SPA core: RAPS extraction, target spectrum evaluation from the spectral prior, and the guidance step.
• sample.py — text-to-image sampling script (SD2.0 / SDXL) with optional SPA.
• collect_spectrum.py — stage 1 of prior computation: collect the empirical power spectrum of x_{0|t} over a dataset.
• regression.py — stage 2 of prior computation: fit the parametric spectral prior to the collected statistics.
• laion.py — cached LAION-Art dataset used by collect_spectrum.py.
• scripts/sample_sd2.sh, scripts/sample_sdxl.sh — example sampling scripts.
• scripts/collect_spectrum.sh, scripts/regression.sh — prior computation (stage 1 / stage 2).

Requirements

• PyTorch, diffusers, transformers, accelerate, numpy, scipy, Pillow
• For collecting a spectral prior from scratch: datasets, requests, torchvision

Spectral prior

SPA requires a pre-computed spectral prior (--spectrum_ref): a pickle file containing, for each latent channel, cubic-spline interpolators of the parameters of the spectrum model S(t, f) = p_t * f^{-q_t} + r_t * f + s_t. See the paper (Sec. "Target Power Spectrum Modeling") for how it is fitted. Priors for SD2.0 and SDXL are provided under stats/.

Computing the prior from scratch

The prior is computed in two stages from a set of real images (LAION-Art):

1. Collection (collect_spectrum.py, needs a GPU): for each image and each timestep on a fixed grid, forward-diffuse the VAE latent, run one conditional UNet prediction, estimate x_{0|t} via Tweedie's formula, and extract its radially averaged power spectrum (RAPS). The per-channel RAPS are averaged over all images (the paper uses 10k). Outputs spectrum_1d_{sd20,sdxl}.pkl and frequencies_{sd20,sdxl}.pkl.

2. Regression (regression.py, CPU): fit S(t, f) to the collected spectrum per timestep and channel, and interpolate the parameters across timesteps with cubic splines. Outputs the spectral prior consumed by sample.py.

# Stage 1: collect statistics (SD2.0)
python collect_spectrum.py \
    --model_id stabilityai/stable-diffusion-2-base \
    --batch_size 25 --num_batches 400 \
    --output_dir stats

# Stage 2: fit the prior (channel 3 carries the linear term for SD2.0;
# use --linear-channels 1 for SDXL)
python regression.py \
    --spectrum stats/spectrum_1d_sd20.pkl \
    --frequencies stats/frequencies_sd20.pkl \
    --linear-channels 3 \
    --output stats/spline_sd20_laion.pkl

scripts/collect_spectrum.sh and scripts/regression.sh run these two stages for both SD2.0 and SDXL. Collecting from a gated dataset needs HuggingFace authentication (huggingface-cli login or HF_TOKEN).

Usage

# Single prompt, single GPU
python sample.py \
    --model_id stabilityai/stable-diffusion-2-base \
    --spectrum_ref stats/spline_sd20_laion.pkl \
    --prompt "a photograph of an astronaut riding a horse" \
    --cfg_scale 7.5 --steps 50 \
    --spa_scale 0.8 --penalty 3.0 \
    --output_dir outputs/sd2_spa

Setting --spa_scale 0 disables SPA (vanilla sampling baseline).

For large-scale generation (e.g., FID evaluation), use --prompt_file with one prompt per line (sample i uses line i mod #lines) and launch with accelerate for multi-GPU:

accelerate launch --num_processes 4 sample.py \
    --prompt_file captions.txt --n_samples 5000 --batch_size 16 ...

The initial noise of sample i is generated from seed (--seed + i), so samples are reproducible and directly comparable across runs (e.g., with and without SPA).

Main arguments

Argument	Description
--model_id	HuggingFace model ID (SD2.x or SDXL)
--spectrum_ref	Path to the spectral prior (.pkl)
--spa_scale	SPA guidance strength η (0 disables SPA)
--penalty	Intensity a of the asymmetric penalty (a ≥ 1)
--cfg_scale	Classifier-free guidance scale
--prompt / --prompt_file	Fixed prompt, or a file with one prompt per line