AG-LDM — Anatomically Guided Latent Diffusion for Brain MRI

Pretrained weights for AG-LDM, a two-stage latent diffusion framework for conditional follow-up synthesis in longitudinal brain MRI. Given a baseline scan and clinical covariates (target age, sex, diagnosis), AG-LDM generates the follow-up scan, with segmentation-derived anatomical supervision applied in both training stages.

Code: https://github.com/JornyWan/AG-LDM

checkpoints/
  ag-ldm_stage1_autoencoder.pth   # Stage-1 fine-tuned KL autoencoder   (13.77 M params, 53 MB)
  ag-ldm_stage2_diffusion.pth     # Stage-2 channel-conditioned UNet    (475.44 M params, 1.8 GB)
code/                             # snapshot of the GitHub repository

These are the exact checkpoints used for the results reported in the paper. The Stage-2 file has been stripped of optimizer/scaler state (5.67 GB → 1.8 GB); it retains the model weights, the latent scale_factor, and the architecture/training configuration.

Usage

import torch
from huggingface_hub import hf_hub_download
from brlp import networks          # from https://github.com/JornyWan/AG-LDM

ae_path   = hf_hub_download("Jorny/AG-LDM", "checkpoints/ag-ldm_stage1_autoencoder.pth")
diff_path = hf_hub_download("Jorny/AG-LDM", "checkpoints/ag-ldm_stage2_diffusion.pth")

autoencoder = networks.init_autoencoder(ae_path).eval()
diffusion   = networks.init_latent_diffusion_channel_cond(diff_path).eval()

# latent scale factor used at sampling time
scale_factor = torch.load(diff_path, map_location="cpu")["scale_factor"]  # 0.960971

Or download everything:

hf download Jorny/AG-LDM --local-dir ./AG-LDM-weights

Model details

	Stage 1 (autoencoder)	Stage 2 (diffusion UNet)
Params	13.77 M	475.44 M
Input	1×120×144×120 MRI	11×16×20×16 latent
Latent	3×15×18×15 (1.5 mm, MNI)	—
Conditioning	—	channel concat: 3 noisy + 3 baseline + 5 covariates
Trained	10 epochs	20 epochs (71,600 steps)

Covariate channels: starting age, follow-up age, sex, starting diagnosis, follow-up diagnosis. Latents are DivisiblePad(k=4)-padded to 16×20×16 before the UNet.

Training data: ADNI longitudinal pairs (25,498 training pairs). Preprocessing follows BrLP: skull-strip, MNI registration, 1.5 mm resample to 122×146×122.

Segmentation teacher (not included)

Training uses a frozen WarpSeg tissue segmenter as the anatomical-guidance teacher. Its model definition and checkpoint are not redistributed here; obtain them from WarpSeg. The teacher is only required for training; the checkpoints above run inference without it.

Intended use and limitations

Research use only. Not a medical device and not for clinical decision-making.

The model performs conditional follow-up synthesis; it contains no explicit model of progression dynamics and no term coupling successive time points.
Anatomical accuracy was assessed with SynthSeg, an independent automated morphometric pipeline rather than anatomical ground truth.
Trained on ADNI and evaluated zero-shot on OASIS-3, both research cohorts with standardized acquisition; behaviour under population-scale or clinical-routine acquisition heterogeneity has not been measured.

Acknowledgments

Built on BrLP and MONAI / MONAI Generative Models. Segmentation teacher: WarpSeg.

Downloads last month: -; Downloads are not tracked for this model. How to track