--- library_name: ilex tags: - jax - equinox - ilex - neuroimaging - 3d license: apache-2.0 license_link: https://www.apache.org/licenses/LICENSE-2.0 --- # 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 ```python 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](https://github.com/hypercoil/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.