---
license: apache-2.0
library_name: monai
tags:
- medical-imaging
- chest-x-ray
- diffusion
- ddpm
- federated-learning
- nvflare
---

# Chest X-ray diffusion model: FLIP tutorial model

A 2-D denoising diffusion model (DDPM) that generates 128×128 chest X-rays, trained with the
chest X-ray diffusion tutorial of [FLIP](https://github.com/londonaicentre/FLIP), the Federated
Learning Interoperability Platform (`fl-tutorials/nvflare/image_generation/xray_diffusion/`).

It was trained federated, with NVIDIA FLARE simulating two sites on one GPU, entirely on **synthetic**
X-rays. It is a tutorial model, not for clinical use.

![16 chest X-rays sampled from this model, seed 0](figures/samples_4x4.png)

## Contents

```text
diffusion_model.safetensors   the weights, float32 as trained
provenance.json               commit, install line, launch command, versions, hardware, wall time, final losses
code/                         the app files exactly as the run executed them
evaluation/sample_xray.py     the script that drew the samples above
figures/
```

## Model and training

MONAI `DiffusionModelUNet` in 2-D, unconditional (its only input besides the image is the noise
step), 22.1 M parameters, trained from scratch: four levels of 64, 128, 192 and 256 channels, with
self-attention at the two coarsest. It predicts the noise added at a random step of a 1000-step DDPM
schedule, with an MSE loss.

| | |
| --- | --- |
| Data | Ark+ tutorial splits, about 1.9k images per site, two sites |
| Preprocessing | grayscale, resized to 128×128, scaled to [0, 1]; light random affine in training |
| Schedule | 2 rounds × 300 epochs, batch 6 |
| Learning rate | 1e-3 decaying to 5e-5 over the first 200 epochs |
| Final validation loss | 0.0039 (site 1), 0.0036 (site 2), on each site's held-out 20 % |
| Wall time | 5 h 05 min on one NVIDIA GeForce RTX 5090 |
| Software | PyTorch 2.13.0, MONAI 1.6.0, NVIDIA FLARE 2.9.0 |

Samples saved during training looked like chest X-rays from about epoch 150.

## Use

```python
import json
from monai.networks.nets import DiffusionModelUNet
from safetensors.torch import load_file

net_config = json.load(open("code/config.json"))["net_config"]
model = DiffusionModelUNet(spatial_dims=2, **net_config["diffusion_model"])
weights = load_file("diffusion_model.safetensors")
model.load_state_dict({k.removeprefix("diffusion_model."): v for k, v in weights.items()})
```

Sample with MONAI's `DDPMScheduler` and `DiffusionInferer` (1000 steps), as in
`evaluation/sample_xray.py`.

## Reproducing it

Install flip-utils from the commit in `provenance.json` with its install line, check out FLIP at
the same commit, copy `code/` over the tutorial's `app_files/`, download the data with
`make -C fl-tutorials download-arkplus-finetuning-data`, and run the launch command. The run predates
the tutorial's move to `image_generation/`, so `provenance.json` gives the directory it ran from.
Training on a GPU is not bit-for-bit repeatable, so expect samples of the same quality, not the same
samples.

## Data and licence

The Ark+ tutorial splits
([`aicentreflip/tutorials-arkplus-cxr-classification`](https://huggingface.co/datasets/aicentreflip/tutorials-arkplus-cxr-classification))
are fully synthetic chest X-rays (DeCaf `balanced_synthetic_split`), with no patient data, licensed
Apache-2.0. This model is licensed Apache-2.0 too.