"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.

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 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.