--- license: mit library_name: pytorch pipeline_tag: unconditional-image-generation tags: - medical-imaging - mri - brain - neuroimaging - 3d - flow-matching - wavelets - generative - rectified-flow arxiv: 2601.05212 --- # FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching FlowLet is a conditional generative framework that synthesizes age-conditioned 3D brain MRI volumes. It performs flow matching directly in an invertible 3D Haar wavelet domain, which gives multi-scale generation without any learned latent compression and avoids the reconstruction artifacts that latent diffusion models can introduce. Sampling is a deterministic Euler ODE, so high-fidelity volumes are produced in few steps. Age is injected through two complementary mechanisms (FiLM in the residual blocks for global modulation, and spatial cross-attention in the transformer blocks for spatially adaptive control). A motivating application is Brain Age Prediction (BAP): training BAP models with FlowLet-generated data improves performance for under-represented age groups, while region-based analysis confirms preservation of anatomical structure. > Status: the four checkpoints listed below are currently in training. ![FlowLet architecture](assets/FlowLet_Architecture.png) ## Links - Hugging Face paper page: https://huggingface.co/papers/2601.05212 - arXiv: https://arxiv.org/abs/2601.05212 - Code (GitHub): https://github.com/sisinflab/FlowLet - Project page: https://danesed.github.io/flowlet-page/ - Model repository (this page): https://huggingface.co/danesed/FlowLet ## Model description | Component | Value | | --- | --- | | Representation | Single-level 3D Haar DWT, producing 8 wavelet subbands (1 LLL approximation plus 7 detail), each at half spatial resolution | | Network I/O | Conditional 3D U-Net, 8 input and 8 output channels (one per subband), 3D convolutions throughout | | Backbone | 3D U-Net with `model_channels=128`, `num_res_blocks=2`, GroupNorm-32, and `SpatialTransformerConditional` attention blocks. Two configurations are released (see [Models](#models)). | | Conditioning | Age (a single scalar), via FiLM in the residual blocks plus cross-attention in the transformer blocks. Condition embedding dimension 512. | | Age normalization | Min-max to the [0, 1] interval using `condition_ranges.json`, then clamped to [0, 1] so values outside the training range do not extrapolate. | | Objective | Rectified Flow Matching (straight-line interpolation between noise and data, constant target velocity). | | Sampling | Euler ODE integration, deterministic given the seed. High quality in few steps (100 steps for the highest-fidelity results). | | Output | NIfTI (`.nii.gz`), intensities rescaled to [0, 1], identity affine. | The codebase also implements other flow formulations (`cfm`, `vp_diffusion`, `trigonometric`), but only the Rectified Flow Matching checkpoints are released here. ## Models Four checkpoints: two spatial resolutions, each in two U-Net configurations. All four use Rectified Flow Matching (`rfm`) and age conditioning. The "base" and "large" configurations differ in the U-Net channel multipliers and attention resolutions, and therefore in parameter count. | Model | Resolution (saved volume) | Config | U-Net params | Planned file | Status | | --- | --- | --- | --- | --- | --- | | FlowLet-RFM-91-base | 91 x 109 x 91 | base (channel_mult 1,2,3,4 / attn 16,8) | 356.4 M | `rfm-91-base/flowlet_rfm_91_base.pth` | In training, coming soon | | FlowLet-RFM-91-large | 91 x 109 x 91 | large (channel_mult 1,2,4,8 / attn 4,8) | 1.00 B | `rfm-91-large/flowlet_rfm_91_large.pth` | In training, coming soon | | FlowLet-RFM-182-base | 182 x 218 x 182 | base (channel_mult 1,2,3,4 / attn 16,8) | 356.4 M | `rfm-182-base/flowlet_rfm_182_base.pth` | In training, coming soon | | FlowLet-RFM-182-large | 182 x 218 x 182 | large (channel_mult 1,2,4,8 / attn 4,8) | 1.00 B | `rfm-182-large/flowlet_rfm_182_large.pth` | In training, coming soon | Each variant folder will also contain its `config.json` (the architecture the generation script rebuilds the model from) and its `condition_ranges.json` (the age range used for normalization). The 91 resolution uses a padded model input of 112 x 112 x 112, and the 182 resolution uses 224 x 224 x 224. ## How to use (ready for when the weights are released) FlowLet uses a custom 3D architecture, so it is loaded with the repository code plus the released `.pth`, not with `transformers` or `PyTorchModelHubMixin`. Once a checkpoint is available, download it with its sidecar JSON files, then run the repository generation script. ```bash # Code and environment git clone https://github.com/sisinflab/FlowLet && cd FlowLet conda create -n flowlet_env python=3.11 && conda activate flowlet_env pip install -r requirements.txt # torch==2.6.0, xformers optional ``` ```python # Download one variant (weights, config, age ranges). Available once Status shows released. from huggingface_hub import hf_hub_download repo_id = "danesed/FlowLet" variant = "rfm-91-base" # rfm-91-base | rfm-91-large | rfm-182-base | rfm-182-large fname = "flowlet_rfm_91_base.pth" ckpt = hf_hub_download(repo_id, f"{variant}/{fname}", revision="main") config = hf_hub_download(repo_id, f"{variant}/config.json", revision="main") ranges = hf_hub_download(repo_id, f"{variant}/condition_ranges.json", revision="main") print(ckpt, config, ranges) ``` ```bash # Generate. The script rebuilds the model from config.json and normalizes age with # condition_ranges.json. Arguments are a flat argparse (no subcommands), so flag order is free. PYTHONPATH=. python3 -u scripts/generate.py \ --checkpoint_path "$CKPT" \ --config_path "$CONFIG" \ --condition_ranges_path "$RANGES" \ --output_dir ./generated/rfm-91-base \ --generation_conditions "Age=45" "Age=70.5" \ --num_synthetic 5 \ --num_flow_steps 100 \ --save_size 91 109 91 ``` For the 182 resolution variants pass `--save_size 182 218 182` (the padded input size is read from the variant's `config.json`). Notes: - Attention uses `xformers` when available and falls back to native PyTorch attention automatically if it is not installed (a warning is logged). To force the fallback, set `"use_xformers": false` in the variant `config.json` before generating. - Loading: the released `.pth` files are slimmed (weights under `model_state_dict` plus a small config block). The generation script calls `torch.load(..., map_location=device)` without setting `weights_only`. On torch 2.6 (pinned here) the default is `weights_only=True`, and the slimmed files contain only tensors and JSON-serializable config, so they load under that default. ## Training data FlowLet was trained on preprocessed T1-weighted brain MRI from public research cohorts: - OpenBHB: https://baobablab.github.io/bhb/dataset - ADNI: https://adni.loni.usc.edu/ - OASIS-3: https://sites.wustl.edu/oasisbrains/ No imaging data is redistributed in this repository. Because of patient-privacy regulations and data-use agreements, the scans cannot be shared here. Access must be requested from the original providers under their respective agreements. Preprocessing (per the paper and the code repository): N4ITK bias-field correction (ANTs), affine registration to MNI152 (FSL FLIRT), skull stripping (FSL BET), resampling to 91 x 109 x 91, and z-score intensity normalization. The conditioning variable is the subject Age, and the released `condition_ranges.json` covers Age in [5.90, 95.46]. ## Intended use and limitations Intended use: research on generative modeling of brain MRI, data augmentation for downstream research (for example Brain Age Prediction), and benchmarking of flow-matching formulations. Limitations and out-of-scope use: - Not a medical device. No diagnostic, screening, or clinical use. - Synthetic volumes may contain anatomical artifacts and do not correspond to real individuals. - Outputs reflect the cohort bias of the training data (acquisition sites, scanners, demographics). - Age is clamped to the training range [5.90, 95.46]. Values outside it are silently clipped, so out-of-range ages do not produce reliable extrapolation. - Generation is conditioned on age only. Other clinical or morphological factors are not controlled. ## Citation ```bibtex @misc{danese2026flowletconditional3dbrain, title={FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching}, author={Danilo Danese and Angela Lombardi and Matteo Attimonelli and Giuseppe Fasano and Tommaso Di Noia}, year={2026}, eprint={2601.05212}, archivePrefix={arXiv}, primaryClass={cs.CV}, url={https://arxiv.org/abs/2601.05212}, } @article{danese2026flowlet, title = {FlowLet: Conditional 3D Brain MRI Synthesis using Wavelet Flow Matching}, author = {Danese, Danilo and Lombardi, Angela and Attimonelli, Matteo and Fasano, Giuseppe and Di Noia, Tommaso}, journal = {Medical Image Analysis}, year = {2026}, publisher = {Elsevier}, DOI = {TO_BE_ASSIGNED} } ``` ## License Released under the MIT License. See https://github.com/sisinflab/FlowLet/blob/main/LICENSE