README.md

---
license: mit
language:
- en
tags:
- medical-imaging
- super-resolution
- diffusion-model
- MRI
- chest-x-ray
- latent-diffusion
datasets:
- sebasmos/latent-sr-embeddings
pipeline_tag: image-to-image
---

# Latent-SR: Domain-Specific Diffusion Weights for Medical Image Super-Resolution

Trained diffusion model checkpoints from the paper:

> **"Domain-Specific Latent Representations Improve the Fidelity of Diffusion-Based Medical Image Super-Resolution"**  
> Sebastian Cajas, Ashaba Judith, Rahul Gorijavolu, Sahil Kapadia, Hillary Clinton Kasimbazi, Leo Kinyera, Emmanuel Paul Kwesiga, Sri Sri Jaithra Varma Manthena, Luis Filipe Nakayama, Ninsiima Doreen, Leo Anthony Celi.  
> arXiv:2604.12152 (2026) — under review at Nature Scientific Reports.
>
> 📄 Paper: https://arxiv.org/abs/2604.12152  
> 💻 Code: https://github.com/sebasmos/latent-sr

## Model Description

These are the trained x₀-prediction latent diffusion model (LDM) weights for 4× medical image super-resolution. Each checkpoint contains the diffusion **UNet only** — the encoder/decoder VAE is **not** bundled here and is loaded separately at inference time (MedVAE, SD-VAE, and KL-f4 base weights are obtained from their original sources). This repository therefore contains **only trained diffusion checkpoints**, not pretrained base VAEs.

All models were trained with:
- **Schedule**: cosine β, T=1000 train / T=100 inference
- **Prediction target**: x₀ (clean latent)
- **Pipeline**: frozen VAE encoder → diffusion UNet → frozen VAE decoder

## Checkpoints

This repository holds **21 trained diffusion checkpoints**: 7 latent geometries × 3 datasets
(BraTS 2023 brain MRI, MIMIC-CXR chest X-ray, MRNet knee MRI).

> **Note on metrics:** the paper's super-resolution results were recomputed against true
> high-resolution references. To avoid publishing stale or uncertain per-checkpoint PSNR
> values, they are **not** tabulated here — see the paper for the authoritative,
> reference-based numbers. Each checkpoint below is described by its VAE type, latent
> geometry, dataset, and experimental role.

| File | VAE (latent) | Dataset | Role |
|------|--------------|---------|------|
| `medvae-4-3_brats.ckpt` | MedVAE 4_3 (3×64×64) | BraTS 2023 brain MRI | **Paper main result** (domain-specific) |
| `medvae-4-3_cxr.ckpt`   | MedVAE 4_3 (3×64×64) | MIMIC-CXR chest X-ray | **Paper main result** (domain-specific) |
| `medvae-4-3_mrnet.ckpt` | MedVAE 4_3 (3×64×64) | MRNet knee MRI | **Paper main result** (domain-specific) |
| `sdvae_brats.ckpt`      | SD-VAE (4×32×32) | BraTS 2023 | Baseline (natural-image SD-VAE) |
| `sdvae_cxr.ckpt`        | SD-VAE (4×32×32) | MIMIC-CXR | Baseline (natural-image SD-VAE) |
| `sdvae_mrnet.ckpt`      | SD-VAE (4×32×32) | MRNet | Baseline (natural-image SD-VAE) |
| `klf4_brats.ckpt`       | KL-f4 (3×64×64) | BraTS 2023 | Natural-image control (matched geometry) |
| `klf4_cxr.ckpt`         | KL-f4 (3×64×64) | MIMIC-CXR | Natural-image control (matched geometry) |
| `klf4_mrnet.ckpt`       | KL-f4 (3×64×64) | MRNet | Natural-image control (matched geometry) |
| `medvae-4-1_brats.ckpt` | MedVAE 4_1 (1×64×64) | BraTS | Capacity control (1-channel) |
| `medvae-4-1_cxr.ckpt`   | MedVAE 4_1 (1×64×64) | MIMIC-CXR | Capacity control (1-channel) |
| `medvae-4-1_mrnet.ckpt` | MedVAE 4_1 (1×64×64) | MRNet | Capacity control (1-channel) |
| `medvae-4-4_brats.ckpt` | MedVAE 4_4 (4×64×64) | BraTS | Higher-capacity anchor (4-channel) |
| `medvae-4-4_cxr.ckpt`   | MedVAE 4_4 (4×64×64) | MIMIC-CXR | Higher-capacity anchor (4-channel) |
| `medvae-4-4_mrnet.ckpt` | MedVAE 4_4 (4×64×64) | MRNet | Higher-capacity anchor (4-channel) |
| `medvae-8-1_brats.ckpt` | MedVAE 8_1 (1×32×32) | BraTS | Capacity control (8× downsample) |
| `medvae-8-1_cxr.ckpt`   | MedVAE 8_1 (1×32×32) | MIMIC-CXR | Capacity control (8× downsample) |
| `medvae-8-1_mrnet.ckpt` | MedVAE 8_1 (1×32×32) | MRNet | Capacity control (8× downsample) |
| `medvae-8-4_brats.ckpt` | MedVAE 8_4 (4×32×32) | BraTS | Domain-matched control (8× downsample) |
| `medvae-8-4_cxr.ckpt`   | MedVAE 8_4 (4×32×32) | MIMIC-CXR | Domain-matched control (8× downsample) |
| `medvae-8-4_mrnet.ckpt` | MedVAE 8_4 (4×32×32) | MRNet | Domain-matched control (8× downsample) |

**Latent geometries.** `medvae-4-3` (3×64×64) is the domain-specific main result. `klf4`
(KL-f4, 3×64×64) is the natural-image VAE at the **same latent geometry** as the main result —
the matched-geometry control isolating domain vs. geometry. `sdvae` (SD-VAE, 4×32×32) is the
natural-image baseline. The remaining `medvae-4-1/4-4/8-1/8-4` checkpoints vary channel count
(1↔4) and downsampling factor (4×↔8×) as capacity/geometry controls, with `medvae-4-4`
(4×64×64) serving as the higher-capacity anchor.

The base VAE encoders/decoders (MedVAE, SD-VAE, KL-f4) are **not** stored in this repo and are
loaded separately at inference — see the [code repository](https://github.com/sebasmos/latent-sr).

## Usage

```python
from huggingface_hub import hf_hub_download
import torch

# Download the main MedVAE BraTS checkpoint
ckpt_path = hf_hub_download(
    repo_id="sebasmos/latent-sr-weights",
    filename="medvae-4-3_brats.ckpt"
)
state = torch.load(ckpt_path, map_location="cpu")
```

See the [code repository](https://github.com/sebasmos/latent-sr) for the full inference pipeline.

## Citation

```bibtex
@article{cajas2026domain,
  title   = {Domain-Specific Latent Representations Improve the Fidelity of Diffusion-Based Medical Image Super-Resolution},
  author  = {{Sebastian Cajas} and {Ashaba Judith} and {Rahul Gorijavolu} and {Sahil Kapadia} and {Hillary Clinton Kasimbazi} and {Leo Kinyera} and {Emmanuel Paul Kwesiga} and {Sri Sri Jaithra Varma Manthena} and {Luis Filipe Nakayama} and {Ninsiima Doreen} and {Leo Anthony Celi}},
  journal = {arXiv preprint arXiv:2604.12152},
  year    = {2026},
  url     = {https://arxiv.org/abs/2604.12152}
}
```

Code: https://github.com/sebasmos/latent-sr · Paper: https://arxiv.org/abs/2604.12152

## License

MIT License. The pretrained VAE weights (MedVAE, SD-VAE) are subject to their respective licenses.