SNaP -- pretrained checkpoints

Weights for SNaP, a one-step posterior sampler for linear inverse problems. The code, configs and download script live in the GitHub repository:

These files are meant to be fetched by the repository's downloader, which verifies every file against its sha256:

python scripts/download_checkpoints.py --all

Each checkpoint is only valid together with its frozen config in configs/pretrained/<model>.yaml in the code repository.

file problem one-step PSNR / SSIM / LPIPS size
celeba_inpaint.pt CelebA 128 - random inpainting, 70% pixels missing, sigma=0.01 31.98 / 0.928 / 0.018 35 MB
celeba_box.pt CelebA 128 - centered 40x40 box inpainting, sigma=0.05 30.73 / 0.934 / 0.021 35 MB
celeba_deblur.pt CelebA 128 - Gaussian deblurring, sigma_b=1.0, k=61, sigma=0.05 32.92 / 0.919 / 0.019 70 MB
celeba_sr.pt CelebA 128 - 2x super-resolution, sigma=0.05 31.27 / 0.903 / 0.023 35 MB
afhq_inpaint.pt AFHQ-Cat 256 - random inpainting, 70% pixels missing, sigma=0.01 30.49 / 0.854 / 0.066 422 MB
afhq_box.pt AFHQ-Cat 256 - centered 80x80 box inpainting, sigma=0.05 26.47 / 0.892 / 0.054 422 MB
afhq_deblur.pt AFHQ-Cat 256 - Gaussian deblurring, sigma_b=3.0, k=61, sigma=0.05 26.25 / 0.674 / 0.160 422 MB
afhq_sr.pt AFHQ-Cat 256 - 4x super-resolution, sigma=0.05 26.07 / 0.704 / 0.131 422 MB
brain_r4_20db.pt fastMRI brain - multi-coil CS-MRI, R=4, 20 dB 32.07 / 0.889 156 MB
brain_r4_30db.pt fastMRI brain - multi-coil CS-MRI, R=4, 30 dB 32.89 / 0.899 156 MB
brain_r8_20db.pt fastMRI brain - multi-coil CS-MRI, R=8, 20 dB 28.11 / 0.819 156 MB
brain_r8_30db.pt fastMRI brain - multi-coil CS-MRI, R=8, 30 dB 29.12 / 0.845 156 MB

Metrics: 100 held-out test images, one step (k=1), a single posterior draw. MRI uses the magnitude, per-slice dynamic-range convention and has no LPIPS. SHA256SUMS lists the checksum of every file.

License

The weights are trained on third-party datasets and inherit their terms: CelebA (non-commercial research only), AFHQ (CC BY-NC 4.0) and fastMRI (its data-use agreement). Use them for non-commercial research. The code is MIT-licensed separately.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support