Brain LDM (conditional 3D latent diffusion for synthetic brain MRI) -- Brain LDM v1 (model.pt + DDIM sampler)

Description

Conditional 3D latent-diffusion model that generates synthetic T1-weighted brain MRI volumes from a 4-dimensional demographic / morphometric conditioning vector (age, sex, ventricular-CSF volume, brain volume, each in [0, 1]). Composite of three sub-networks: (1) a KL-regularised 3D VAE (3-channel latent at 8x spatial downsampling) shared with the brain_ldm.vae base bundle; (2) a 553M-param conditional diffusion U-Net with channels (256, 512, 768), cross-attention at the two deepest levels, ResBlock-up/downsampling, and dual conditioning (both spatial concatenation and cross-attention from the same 4-vector); (3) a DDIM scheduler at the published training schedule (1000 timesteps, linear-beta in [1e-4, 2e-2], set_alpha_to_one=True, epsilon prediction, eta=0 by default for deterministic sampling). Trained on 31,740 UK Biobank T1w MRIs at 1mm isotropic. The local bundle is a parsimonious delta of the VAE base bundle: only the U-Net weights and scheduler buffers are stored locally; the VAE is loaded from the base.

Intended use

Research tool / prototype. Sample synthetic T1w brain MRI volumes conditional on a 4-dim demographic vector (age, sex, vCSF, brain volume; each normalised to [0, 1] per the upstream's UK Biobank percentiles). Generates 1 mm iso volumes at any multiple-of-8 spatial size. NOT a clinical tool.

Usage

from ilex.models.brain_ldm import BrainLDM
model = BrainLDM.from_pretrained('ilex-hub/brain_ldm.1')

Authors

Walter H. L. Pinaya, Petru-Daniel Tudosiu, Jessica Dafflon, Pedro F. Da Costa, Virginia Fernandez, Parashkev Nachev, Sebastien Ourselin, M. Jorge Cardoso

Citation

Pinaya W. H. L., Tudosiu P.-D., Dafflon J., Da Costa P. F., Fernandez V., Nachev P., Ourselin S., Cardoso M. J. (2022). Brain imaging generation with latent diffusion models. MICCAI Workshop on Deep Generative Models, Springer, pp. 117-126.

References

  • Pinaya W. H. L., Tudosiu P.-D., Dafflon J., Da Costa P. F., Fernandez V., Nachev P., Ourselin S., Cardoso M. J. (2022). Brain imaging generation with latent diffusion models. MICCAI Workshop on Deep Generative Models, Springer, pp. 117-126.
  • Song J., Meng C., Ermon S. (2021). Denoising Diffusion Implicit Models. ICLR. https://arxiv.org/abs/2010.02502
  • Upstream bundle: huggingface.co/MONAI/brain_image_synthesis_latent_diffusion_model (model.pt 2.2 GB, ~553M params; autoencoder.pt ~13.77M params).
  • Architecture: monai.networks.nets.DiffusionModelUNet + monai.networks.schedulers.DDIMScheduler (MONAI 1.4+).

License

HF Hub license tag: apache-2.0

Upstream license reference: https://www.apache.org/licenses/LICENSE-2.0

Copyright

Network architecture and pretrained weights -- copyright (c) MONAI Consortium, released under the Apache License 2.0. JAX / Equinox port code -- copyright (c) the ilex authors, released under the Apache-2.0 / GPL-3.0 dual license used by ilex itself.

Upstream source

Original weights / reference implementation: https://huggingface.co/MONAI/brain_image_synthesis_latent_diffusion_model

Provenance

This artefact was produced by ilex's save/load pipeline. The architecture is implemented in ilex.models.brain_ldm.BrainLDM and the weights have been converted from their upstream format. See the upstream source above for the canonical reference.

Downloads last month
16
Safetensors
Model size
0.6B params
Tensor type
F32
·
I32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Paper for ilex-hub/brain_ldm.1