| --- |
| license: cc-by-nc-4.0 |
| library_name: pytorch |
| pipeline_tag: unconditional-image-generation |
| tags: |
| - diffusion |
| - medical-imaging |
| - brats |
| extra_gated_prompt: >- |
| These models were trained on the BraTS 2024 adult glioma dataset (BraTS-GLI) and are released for non-commercial research only |
| (CC BY-NC 4.0). If you use the models or images generated with them, you must cite |
| our paper and the papers in the Citation section. The models are not for clinical use. |
| extra_gated_fields: |
| Name: text |
| Affiliation: text |
| I will use the models for non-commercial research only: checkbox |
| I will cite the required papers in all work that uses the models: checkbox |
| I will not use the models or their images for clinical decisions: checkbox |
| --- |
| |
| # rad-dit: LightningDiT-XL/1 on BraTS |
|
|
| A latent diffusion transformer ([LightningDiT](https://github.com/hustvl/LightningDiT)) with the VA-VAE tokeniser, initialised with ImageNet weights and trained on BraTS (256×256). |
|
|
| **Data:** "BraTS" here means the BraTS 2024 adult glioma dataset (BraTS-GLI, post-treatment MRI), not all datasets of the BraTS 2024 challenge. |
|
|
| - **Paper:** [From Natural Images to Medical Scans: Transferability of Diffusion Models for Medical Image Generation](https://arxiv.org/abs/2611.05192) (WACV 2027) |
| - **Code:** https://github.com/CAIROLab-Bern/rad-dit |
| - **Project page:** https://rad-dit.github.io |
| - **Other models:** [JiT on RadImageNet](https://huggingface.co/rad-dit/jit-radimagenet), [LightningDiT-XL/1 on RadImageNet](https://huggingface.co/rad-dit/ldit-radimagenet) |
|
|
| ## Available models |
|
|
| | File | Model | gFID ↓ | |
| |---|---|---| |
| | [`ldit-xl1-vavae.pt`](https://huggingface.co/rad-dit/ldit-brats/blob/main/ldit-xl1-vavae.pt) | LightningDiT-XL/1 (675M) | 7.66 | |
| | [`latents_stats.pt`](https://huggingface.co/rad-dit/ldit-brats/blob/main/latents_stats.pt) | ImageNet VA-VAE latent statistics, used in training and sampling (from [hustvl/vavae-imagenet256-f16d32-dinov2](https://huggingface.co/hustvl/vavae-imagenet256-f16d32-dinov2), MIT) | – | |
|
|
| gFID: 256×256, cfg 1.0, Euler with 250 steps (see the paper). `ldit-xl1-vavae.pt` holds `model`, `ema` and `config`, as the official LightningDiT checkpoints. Sampling uses `ema`. The optimizer state is not included. |
|
|
| **Architecture:** `inference.py` builds the model from the inference config in the code repository (`ldit/configs/inference_*.yaml`, see the usage below). The checkpoint also holds a copy of the training config (`config`), for reference only. The inference config reads `latents_stats.pt` from the VA-VAE folder; the copy here is the same file. |
|
|
| ## Classes |
|
|
| | Label | Class | |
| |---|---| |
| | 1001 | No tumour in the slice | |
| | 1002 | Tumour in the slice | |
| | 1003 | No class (classifier-free guidance row) | |
|
|
| The inference config sets `num_classes: 1003`. Each image is one axial MRI slice of one modality (T1n, T1c, T2w or T2f). The modality is not a condition. For example, `class_probs: [0, 1]` gives tumour slices only. |
|
|
| ## Usage |
|
|
| **1. Get the code and log in** |
| ```bash |
| git clone https://github.com/CAIROLab-Bern/rad-dit.git && cd rad-dit |
| # Install the environment: see environment/README.md |
| hf auth login # accept the terms on this page first |
| ``` |
|
|
| **2. Download the model and the VA-VAE tokeniser** |
| ```bash |
| hf download rad-dit/ldit-brats --local-dir ldit/checkpoints/brats |
| # or in Python: from huggingface_hub import snapshot_download |
| # snapshot_download("rad-dit/ldit-brats", local_dir="ldit/checkpoints/brats") |
| hf download hustvl/vavae-imagenet256-f16d32-dinov2 \ |
| --local-dir ldit/pretrained/vavae-imagenet256-f16d32-dinov2 |
| ``` |
|
|
| **3. Generate images** (no dataset needed) |
| ```bash |
| cd ldit |
| accelerate launch --num_processes=1 inference.py \ |
| --config configs/inference_brats_hastumor.yaml \ |
| --ckpt checkpoints/brats/ldit-xl1-vavae.pt \ |
| --output-dir samples --num-samples 256 |
| ``` |
| Images are made in batches of `per_proc_batch_size` (256 in the config). Lower it in the config for fewer images or for a smaller GPU. |
| The images are sampled uniformly over the classes. For selected classes, add `class_probs` to the `data` section of the config, with one value per class. |
|
|
| **4. gFID:** see `ldit/README.md` in the code repository. |
|
|
| ## Terms of use |
|
|
| - Use the models for **non-commercial research only** (CC BY-NC 4.0). |
| - Do **not** use the models or their images for diagnosis or other clinical decisions. |
| - Generated images are synthetic. They are not real patient data. |
|
|
| ## Citation |
|
|
| If you use these models or images generated with them in your research, please cite our paper and consider giving the [code repository](https://github.com/CAIROLab-Bern/rad-dit) a ⭐: |
|
|
| ```bibtex |
| @inproceedings{georgiou2027natural, |
| title = {From Natural Images to Medical Scans: Transferability of Diffusion Models for Medical Image Generation}, |
| author = {Georgiou, Efthymios and Kakogeorgiou, Ioannis and Br{\"u}ningk, Sarah}, |
| booktitle = {2027 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)}, |
| year = {2027} |
| } |
| ``` |
|
|
| If you use our pretrained models, you must also cite the BraTS papers, as the data terms require: |
|
|
| - de Verdier et al., *The 2024 Brain Tumor Segmentation (BraTS) Challenge: Glioma Segmentation on Post-treatment MRI*, 2024. [Paper](https://doi.org/10.48550/arXiv.2405.18368) |
| - Menze et al., *The Multimodal Brain Tumor Image Segmentation Benchmark (BRATS)*, IEEE Transactions on Medical Imaging, 2015. [Paper](https://doi.org/10.1109/TMI.2014.2377694) |
|
|
