ShirinShouhstari/FIRM-weights

FIRM — pretrained weights

Pretrained models for FIRM: Flow-based Imaging via Regularized Minimization.

Code: https://github.com/wustl-cig/FIRM_code

FIRM parameterizes the conditional velocity through the measurement-conditional posterior mean, v = (E[x1 | x_t, y] - x_t) / (1 - t), characterizes that mean as the minimizer of a variational objective with an explicit data-consistency term, and unrolls the minimization with half-quadratic splitting — so the forward operator enters the learned velocity at every iteration, and no guidance is needed at sampling time. Reconstruction takes two Euler steps (NFE = N x K = 2 x 5 = 10).

Files

Eleven models, one per (dataset, problem) pair. Each file holds the EMA weights, the epoch they were taken from, and a copy of the matching config section. The optimizer state is stripped: these are for inference and fine-tuning.

File	Dataset	Problem	Res	Params
celeba_inpaint_random.pt	CelebA	inpainting, 70% pixels dropped	128	34.5 M
celeba_inpaint_box.pt	CelebA	inpainting, centered 40x40 hole	128	34.5 M
celeba_denoise.pt	CelebA	denoising, sigma 0.2	128	34.5 M
celeba_deblur.pt	CelebA	Gaussian blur, std 1.0	128	34.5 M
celeba_sr.pt	CelebA	super-resolution x2	128	34.5 M
afhq_inpaint_random.pt	AFHQ-Cat	inpainting, 70% pixels dropped	256	59.9 M
afhq_inpaint_box.pt	AFHQ-Cat	inpainting, centered 80x80 hole	256	59.9 M
afhq_denoise.pt	AFHQ-Cat	denoising, sigma 0.2	256	59.9 M
afhq_deblur.pt	AFHQ-Cat	Gaussian blur, std 3.0	256	239.4 M
afhq_sr.pt	AFHQ-Cat	super-resolution x4	256	59.9 M
afhq_cs_mri.pt	AFHQ-Cat	Fourier CS, Cartesian 21.88%	256	34.5 M

The AFHQ degradations are harder than the CelebA ones (x4 instead of x2 super-resolution, blur std 3.0 instead of 1.0, an 80x80 instead of 40x40 hole), so the two datasets are not comparable row-by-row. afhq_deblur is a wider model (ch=64); afhq_cs_mri is fine-tuned from the AFHQ random-inpainting denoiser rather than trained from scratch.

Usage

git clone https://github.com/wustl-cig/FIRM_code.git && cd FIRM_code
pip install -r requirements.txt

python -c "
from huggingface_hub import hf_hub_download
hf_hub_download('ShirinShouhstari/FIRM-weights', 'celeba_inpaint_random.pt', local_dir='weights')
"

python main.py --config configs/celeba_inpaint_random.yaml \
  --set experiment.phase=test model.resume_path=weights/celeba_inpaint_random.pt \
        data.root=/path/to/celeba distributed.gpus=[0] distributed.use_ddp=false

The config filename matches the weight filename in every case.

Reference numbers

Two Euler steps, EMA weights, sampler integrated fully to t = 1 (sample.eps: 0.0, sample.t_max: 1.0, sample.alpha_eps: 0.0) — the settings already in each config. CelebA: 100 images at 128x128. AFHQ-Cat: first 100 validation images at 256x256. The CelebA rows are the paper's Table 1, Ours (N = 2).

Model	PSNR	SSIM	LPIPS	batch
celeba_inpaint_random	35.03	0.9629	0.0188	32
celeba_inpaint_box	33.25	0.9603	0.0225	32
celeba_denoise	33.67	0.9320	0.0344	32
celeba_deblur	36.05	0.9572	0.0295	32
celeba_sr	34.40	0.9492	0.0285	32
afhq_inpaint_random	34.72	0.9362	0.0436	4
afhq_inpaint_box	29.27	0.9243	0.0601	2
afhq_denoise	33.31	0.9039	0.0810	1
afhq_deblur	29.90	0.8091	0.2294	4
afhq_sr	29.11	0.8169	0.1675	2
afhq_cs_mri	32.07	0.8929	0.0743	5

The batch size is part of the protocol: the mask and the observation noise are drawn per batch, so re-batching the same images moves PSNR by a few hundredths of a dB.

Training data and intended use

Trained on CelebA (faces) and AFHQ-Cat (animal faces). Both datasets are released for non-commercial research, and these weights inherit that intent: they are published to reproduce and build on the paper's results. The MIT license above covers this repository's own contents.

Each model is trained for one forward operator at one noise level and is not expected to transfer to a different degradation without retraining. Face models inherit the demographic skew of CelebA and should not be relied on for anything identity-related.

Citation

@article{shoushtari2026firm,
  title   = {{FIRM}: Flow-based Imaging via Regularized Minimization},
  author  = {Shoushtari, Shirin and Chandler, Edward P. and Shi, Xiao and Kamilov, Ulugbek S.},
  journal = {arXiv preprint},
  year    = {2026},
  url     = {https://arxiv.org/pdf/2609.12953}
}