AnaDiffusion

Anatomically Compositional Latent Diffusion for Controllable 3D Brain MRI Generation.

AnaDiffusion composes whole-brain MRI generation from anatomical part models. The model first synthesizes high-fidelity regional parts and assembles them into a whole-brain scaffold. The scaffold is re-encoded into the whole-brain latent space and injected during denoising, enabling local anatomical detail and globally coherent whole-brain synthesis.

The repository includes data-preparation utilities, example metadata and split files, training scripts for autoencoders and latent diffusion models, and inference code for compositional whole-brain sampling.

The data examples in this repository use placeholder paths and de-identified example image IDs. Replace every <PATH-TO-...> value before running the pipeline on your own data.

• Repository Structure
• Environment
• Preparing Your Own Dataset
• Method
• Training
• Inference
• Evaluation
• Citation
• Acknowledgements

AnaDiffusion/
├── anadiffusion/             # Core method: templates, part experts, scaffold injection
│   ├── templates.py          # MNI152 template masks and binary morphology
│   ├── parts.py              # Part crops, frozen part experts (latent SDEdit), assembly
│   ├── compose.py            # Scaffold injection shared by training and inference
│   ├── models.py             # Autoencoder/UNet construction, schedules, checkpoints
│   ├── data.py               # Volume loading from the manifest CSV and split JSON
│   └── config.py             # JSON configs with command-line overrides
├── configs/
│   ├── train/                # AE and base LDM configs; ours/ holds the compositional config
│   ├── infer/                # Compositional generation and editing config
│   └── eval/                 # MedicalNet FID config
├── data/                     # Example metadata, splits, and template masks
│   └── template/dilated_r*/  # MNI152 brain and part masks dilated by r voxels
├── data_prep/                # Dataset conversion and part-volume preparation
├── scripts/
│   ├── train/                # AE, base LDM, and hemisphere LDM training
│   │   └── ours/             # Compositional (assemble-then-refine) training
│   ├── infer/                # Compositional generation and part editing
│   └── eval/                 # MedicalNet FID on template-defined supports
├── environment.yml
└── README.md

Create the conda environment from the provided YAML:

conda env create -f environment.yml
conda activate monai

Run commands from the AnaDiffusion/ directory unless noted otherwise.

Each sample is a brain MRI NIfTI volume indexed by a de-identified image ID such as I000000. Volumes should be preprocessed and registered to the MNI space before running this pipeline. See the Preprocessing section below.

The files below show the expected manifest and split-file formats:

{
  "I000000": {
    "Age": 76,
    "Sex": "F",
    "Group": "CN"
  }
}

{
  "train": ["I000000", "I000001"],
  "val": ["I000002"],
  "test": ["I000003"]
}

Age, sex, and diagnostic group are not explicitly used during training or inference. They are included to support balanced train, validation, and test splits across participant conditions.

AnaDiffusion uses Turboprep for preprocessing. If you use turboprep, its outputs are registered to the MNI152 template space and segmented using SynthSeg, organized as follows:

<PATH-TO-TURBOPREP-OUTPUT>/
├── I000000/
│   ├── normalized.nii.gz
│   ├── mask.nii.gz
│   └── segm.nii.gz
└── I000001/
    ├── normalized.nii.gz
    ├── mask.nii.gz
    └── segm.nii.gz

python data_prep/turboprep_postproc.py \
  --root <PATH-TO-TURBOPREP-OUTPUT> \
  --outdir <PATH-TO-PROCESSED-TURBOPREP-OUTPUT>

<PATH-TO-PROCESSED-TURBOPREP-OUTPUT>/
└── I000000/
    ├── normalized.nii.gz
    ├── mask.nii.gz
    └── segm.nii.gz

Then create whole-brain and part volumes plus training CSVs:

python data_prep/prep_data.py \
  --root <PATH-TO-PROCESSED-TURBOPREP-OUTPUT> \
  --metadata_json <PATH-TO-METADATA-JSON> \
  --outdir <PATH-TO-PROCESSED-PARTS-DIR> \
  --postfix processed

Generated CSV schemas:

Compositional training and inference use MNI152 template masks for part composition. Template mask folders are provided under data/template/dilated_r*.

Each folder contains masks such as:

• lhemi_mask.nii.gz
• rhemi_mask.nii.gz
• rhemi_mirror_mask.nii.gz
• sub_mask.nii.gz
• mask.nii.gz (Whole brain mask)

The suffix is the dilation radius in voxels. Compositional training and inference crop the parts with the dilated_r2 masks (crop_template_dir) and assemble them inside the dilated_r0 masks (paste_template_dir).

AnaDiffusion trains a whole-brain LDM and part LDMs (left hemisphere, right hemisphere, cerebellar-brainstem complex; the two hemispheres share one side-conditioned model in a left-canonical frame), then fine-tunes the whole-brain denoiser with an assemble-then-refine objective. The implementation lives in anadiffusion/:

1. Scaffold (compose.py): at timestep t the whole-brain denoiser's clean estimate x0_hat of the latent is decoded.
2. Part assets (parts.py): the decoded estimate is cropped with the fixed MNI152 template masks into the three part boxes, and each crop is refined by its frozen part expert with latent SDEdit (forward noise to t_part, DDIM back to 0).
3. Assembly (parts.py): the refined parts are pasted back and feather-blended over the decoded estimate inside their template supports within the head envelope; CSF, background, and the eroded cerebellar boundary keep the scaffold content and are left to the whole-brain refiner.
4. Injection (compose.py): the assembled image is re-encoded with the whole-brain autoencoder and injected into the latent inside the part supports while the noise component of the latent is kept, z_inj = z_t + sqrt(abar_t) * m * (E(x_scaf) - x0_hat).

During training (scripts/train/ours/train_compositional.py) a training latent is forward-noised to t_inj ~ U[100, 300], the scaffold is built from the denoiser's own estimate with parts regenerated from t_part = part_t_ratio * t_inj, and the denoiser is trained to predict the noise of the injected latent on the seam band around the part boundaries with the training scan as target (weighted by min(1, 1/SNR_t)), in addition to the standard denoising loss. At inference (scripts/infer/infer_compositional.py) the same injection is applied once, handoff_steps DDIM steps before the end of sampling, with part_t_ratio = 1.

Recommended order:

1. Prepare the dataset with data_prep/turboprep_postproc.py.
2. Generate part volumes and CSV manifests with data_prep/prep_data.py.
3. Create a split JSON with train, val, and test image IDs.
4. Train the whole-brain, hemisphere, and substructure AEs.
5. Train the base whole-brain, hemisphere, and substructure LDMs.
6. Train the compositional whole-brain model using the pretrained AEs and LDM checkpoints.
7. Run compositional whole-brain generation with the final model checkpoints.

Train the compositional whole-brain model after the base whole-brain, hemisphere, and substructure models are available. It fine-tunes the whole-brain UNet only; the autoencoders and part LDMs stay frozen. Multi-GPU training uses torchrun:

torchrun --standalone --nproc_per_node=4 scripts/train/ours/train_compositional.py \
  --config configs/train/ours/whole_brain_compositional_spacing1p5.json

Compositional training hyperparameters (configs/train/ours/whole_brain_compositional_spacing1p5.json):

Any key can be overridden on the command line, for example --part_t_ratio 5.0 --max_steps 5000. The checkpoints UNET_step*.pt store the latent scale factors that inference reads back.

python scripts/infer/infer_compositional.py --config configs/infer/infer_compositional_spacing1p5.json
torchrun --standalone --nproc_per_node=4 scripts/infer/infer_compositional.py --config configs/infer/infer_compositional_spacing1p5.json

Sample i is drawn from a generator seeded with seed + i, so the output does not depend on the number of GPUs. Inference settings (configs/infer/infer_compositional_spacing1p5.json):

Generated volumes are written to <outdir>/images/sample_XXXX.nii.gz.

Replace one part asset of a generated sample by the corresponding asset of another sample; the modified asset set is reassembled, re-encoded, injected, and refined by the compositional whole-brain model:

python scripts/infer/infer_compositional.py --config configs/infer/infer_compositional_spacing1p5.json \
  --edit true --recipient 0 --donor 1 --part lhemi --handoff_steps 5 --outdir samples/edits

The edited volume and the recipient's pre-injection estimate are saved to --outdir.

Knobs to tweak:

• --recipient, --donor: indices of the recipient and donor samples (sample i uses seed seed + i).
• --part: the part to replace, one of lhemi, rhemi, sub (cerebellar-brainstem complex).
• --handoff_steps: the refinement budget, which trades preservation of the inserted part against harmonisation. A small value (around 5) keeps the donor part, while the generation default (10) lets the whole-brain refiner reshape it.
• --cerebellum_erosion: for cerebellar edits, 0 injects the full template support instead of its eroded core.

scripts/eval/medicalnet_fid.py computes the Fréchet distance between generated cohorts and the real test split in the feature space of the pretrained MedicalNet ResNet-10 (downloaded from Torch Hub, warvito/MedicalNet-models). FID is reported on the whole brain, the left and right hemispheres, the cerebellar-brainstem complex, and the seam band around the part boundaries; the part supports are derived from the MNI152 template masks.

python scripts/eval/medicalnet_fid.py --config configs/eval/medicalnet_fid_spacing1p5.json \
  --cohort compositional=samples/compositional/images

Values are FID × 10⁴, the unit of the tables in the paper. Results are appended to the JSON file given by out in the config. Real-split features are cached at real_feature_cache.

If you find AnaDiffusion useful in your research, please cite:

@article{han2026anadiffusion,
  title={AnaDiffusion: Anatomically Compositional Latent Diffusion for Controllable 3D Brain MRI Generation},
  author={Han, Huiwen and Liu, Lulin and Liu, Bangya and Cai, Yuanhao and Chen, Nuo and Wang, Xiaoqing and Xie, Ziqian and You, Chenyu and Ji, Shuiwang and Zhi, Degui and others},
  journal={arXiv preprint arXiv:2608.23014},
  year={2026}
}

ADNI. Data used in the preparation of this work were obtained from the Alzheimer's Disease Neuroimaging Initiative (ADNI) database (adni.loni.usc.edu). We thank ADNI for sharing these data.

MONAI. We thank the MONAI team for making their code openly available. The autoencoders, latent diffusion models, and data pipelines build on MONAI and MONAI Generative Models, the open-source medical imaging framework initiated by NVIDIA and King's College London.

TurboPrep. We thank Lemuel Puglisi for TurboPrep, which we use for MRI preprocessing (skull stripping, bias-field correction, affine registration to MNI152, intensity normalization, and segmentation). TurboPrep builds on FreeSurfer (SynthSeg, SynthStrip), ANTs, and intensity-normalization. If you use it, please cite:

@software{puglisi2024turboprep,
  author = {Puglisi, Lemuel},
  title = {Turboprep},
  year = {2024},
  url = {https://github.com/LemuelPuglisi/turboprep}
}

MNI152 template. We thank the McConnell Brain Imaging Centre, Montreal Neurological Institute, for the MNI152 template. Volumes are registered to the MNI152 T1 1 mm brain template (MNI152_T1_1mm_brain.nii.gz, as distributed with FSL and the HCP Pipelines), and the template volume and part masks under data/template/ are derived from it.

MedicalNet. We thank the MedicalNet authors for releasing their pretrained models. Evaluation features come from the pretrained MedicalNet ResNet-10, loaded through warvito/MedicalNet-models.