Add model card

README.md

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