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

Contents

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.