Initial upload via tools/push_to_hf.py (architecture: ilex.models.brain_ldm.BrainLDM)

README.md

@@ -0,0 +1,65 @@
+---
+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.

config.json

@@ -0,0 +1,46 @@
+{
+  "_ilex": {
+    "architecture": "ilex.models.brain_ldm.model.BrainLDM",
+    "base_bundle": "ilex-hub/brain_ldm.vae.1",
+    "base_subpath": "vae",
+    "constructor_kwargs": {
+      "combined_attention": false
+    },
+    "format": "ilex",
+    "framework_version": {
+      "equinox": "0.13.8",
+      "ilex": "0.0.0.dev0",
+      "jax": "0.10.0",
+      "jaxlib": "0.10.0",
+      "numpy": "2.4.4",
+      "safetensors": "0.7.0"
+    },
+    "has_state": false,
+    "origin": "ilex-native",
+    "weights_status": "bundled"
+  },
+  "authors": "Walter H. L. Pinaya, Petru-Daniel Tudosiu, Jessica Dafflon, Pedro F. Da Costa, Virginia Fernandez, Parashkev Nachev, Sebastien Ourselin, M. Jorge Cardoso",
+  "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.",
+  "data_type": "nibabel",
+  "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.",
+  "equinox_version": "0.13.8",
+  "ilex_version": "0.0.0.dev0",
+  "image_classes": "None at inference time (the model generates samples unconditional on any input image). The conditioning is a 4-vector of normalised scalars.",
+  "intended_use": "Research tool / prototype. Generate synthetic T1-weighted brain MRI volumes conditional on a 4-dimensional demographic / morphometric vector. The 4 conditioning scalars must each be normalised into [0, 1] using the cohort percentiles described in the upstream preprocessing recipe (UK Biobank: age 44-82 linearly mapped; sex 0/1; vCSF and brain-volume percentiles). The output is a single 3D T1w volume at the same intensity normalisation as the upstream training data ([0, 1], 1 mm isotropic). NOT a clinical tool; generated samples should not be used in any patient-facing or diagnostic workflow.",
+  "jax_version": "0.10.0",
+  "network_data_format": {
+    "inputs": {},
+    "outputs": {}
+  },
+  "numpy_version": "2.4.4",
+  "pred_classes": "Single-channel 3D T1-weighted brain MRI, 1 mm isotropic, intensity in [0, 1]. The published training resolution is 160 x 224 x 160; each spatial dim must be a multiple of 8 (the VAE's three stride-2 downsamples).",
+  "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+)."
+  ],
+  "schema": "https://github.com/Project-MONAI/MONAI-extra-test-data/releases/download/0.8.1/meta_schema_20240725.json",
+  "task": "3D conditional latent diffusion for synthetic brain MRI",
+  "version": "0.0.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db8c0633b365d0eb9c72dc6490b4412718376d9ec60594e99bc51ef886057e58
+size 2212759243