Publish the FLIP tutorial models

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figures/samples_4x4.png filter=lfs diff=lfs merge=lfs -text

README.md

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

code/config.json

@@ -0,0 +1,43 @@
+{
+    "job_type": "standard",
+    "GLOBAL_ROUNDS": 2,
+    "LOCAL_ROUNDS": 300,
+    "VALIDATE_EVERY": 25,
+    "LR_START": 0.001,
+    "LR_END": 5e-05,
+    "LR_DECAY_EPOCHS": 200,
+    "GRAD_CLIP_NORM": 1.0,
+    "VAL_SPLIT": 0.2,
+    "BATCH_SIZE": 6,
+    "SAVE_DEBUG_SAMPLES": true,
+    "spatial_shape": [
+        128,
+        128
+    ],
+    "net_config": {
+        "spatial_dims": 2,
+        "diffusion_model": {
+            "in_channels": 1,
+            "out_channels": 1,
+            "channels": [
+                64,
+                128,
+                192,
+                256
+            ],
+            "num_res_blocks": [
+                2,
+                2,
+                2,
+                2
+            ],
+            "attention_levels": [
+                false,
+                false,
+                true,
+                true
+            ]
+        }
+    },
+    "DATALOADER_WORKERS": 8
+}

code/models.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2026 Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Diffusion model network for the FLIP pixel-space diffusion tutorial.
+
+A single ``DiffusionModelUNet`` that denoises **images directly**. There is no autoencoder and no
+discriminator here: the diffusion process runs at full image resolution, so nothing needs to be
+compressed first and nothing is adversarially trained.
+
+Because the UNet sees images rather than latents, its ``in_channels``/``out_channels`` are the
+**image** channel count (1 for the single-channel chest X-rays this tutorial uses), not an
+autoencoder's ``latent_channels``. Contrast the latent diffusion tutorial, where the same network
+operates on a 3-channel latent.
+"""
+
+import json
+from pathlib import Path
+
+from monai.networks.nets import DiffusionModelUNet
+from torch import nn
+
+
+def load_net_config():
+    """Read the ``net_config`` block from the ``config.json`` next to this file."""
+    config_path = Path(__file__).parent / "config.json"
+    with open(config_path, "r") as f:
+        config = json.load(f)
+    net_config = config.get("net_config", {})
+    print(f"Loaded network config: {net_config}")
+    return net_config
+
+
+class DiffusionModelNetwork(nn.Module):
+    """Creates a pixel-space diffusion model: a single UNet denoising images directly."""
+
+    def __init__(self):
+        super().__init__()
+        net_config = load_net_config()
+        self.diffusion_model = DiffusionModelUNet(
+            spatial_dims=net_config["spatial_dims"],
+            in_channels=net_config["diffusion_model"]["in_channels"],
+            out_channels=net_config["diffusion_model"]["out_channels"],
+            channels=net_config["diffusion_model"]["channels"],
+            num_res_blocks=net_config["diffusion_model"]["num_res_blocks"],
+            attention_levels=net_config["diffusion_model"]["attention_levels"],
+            # Fused scaled-dot-product attention: the same result, without materialising each head's map.
+            use_flash_attention=True,
+        )
+
+
+_net = DiffusionModelNetwork()
+
+
+def get_model() -> nn.Module:
+    """
+    Returns the model defined in this file.
+    NOTE: This function needs to exist and cannot take any input arguments. If you would like to parameterize the
+    configuration of your model, for example loaded from a config file, do it when instantiating the model above.
+    """
+    return _net

code/plot_utils.py

@@ -0,0 +1,605 @@
+# Copyright (c) 2026 Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Figures for the image-synthesis tutorials, for local development only.
+
+``save_grid`` and ``save_triplanar`` write debug images during a simulator run; run as a script, this
+module plots each site's losses from a simulator workspace. It sits in ``image_synthesis/dev/``,
+which is never uploaded: each tutorial's ``job.py`` stages it into simulator runs, and on a trust the
+trainers' import of it fails and debug output is simply off.
+
+Images are written only when ``config.json`` sets ``SAVE_DEBUG_SAMPLES`` **and** this is a
+Makefile-launched simulator run (see :func:`samples_enabled`), into ``DEBUG_SAMPLES_DIR``. Nothing
+here sends an image anywhere.
+
+Usage, from a tutorial's directory after ``make run``::
+
+    python ../dev/plot_utils.py                          # reads /tmp/nvflare/<this tutorial>/flip_fedavg
+    python ../dev/plot_utils.py --workspace <dir> --out-dir <dir>
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import re
+from pathlib import Path
+
+import numpy as np
+import torch
+from flip.constants import FlipConstants
+from torchvision.utils import make_grid, save_image
+
+logger = logging.getLogger(__name__)
+
+#: Set by the tutorial Makefiles' ``sim`` recipe only: where the images go, and half of the simulator
+#: check in :func:`in_simulation`. Never a ``config.json`` key.
+DEBUG_DIR_ENV = "DEBUG_SAMPLES_DIR"
+
+#: How many images to write per grid when ``DEBUG_SAMPLES_MAX`` is absent from the config.
+DEFAULT_MAX_IMAGES = 8
+
+#: Training iterations between tri-planar figures when ``DEBUG_PLOT_EVERY`` is absent.
+DEFAULT_PLOT_EVERY = 100
+
+#: The three orthogonal views: display name, the array axis each one holds fixed, and whether the
+#: slice needs a left-right flip after the shared rotation.
+#:
+#: Volumes reach here as ``(C, X, Y, Z)`` with the chain's ``Orientationd(axcodes="RAS")``
+#: guaranteeing what the axes mean: +X is right, +Y anterior, +Z superior. So a sagittal view fixes
+#: X, a coronal view fixes Y, and an axial view fixes Z. Without that guarantee these labels — and
+#: the flip below — would be guesses.
+#:
+#: One ``np.rot90`` puts the vertical axis the right way up in all three (superior up for sagittal
+#: and coronal, anterior up for axial), but it does not settle the horizontal one, and the sagittal
+#: view is the odd one out. Coronal and axial both have **X** across the page, so rot90 leaves them
+#: with the patient's left on the viewer's left — neurological convention. Sagittal has **Y** across
+#: the page instead, which leaves anterior on the *right*: the head faces backwards compared with how
+#: every other view here, and every standard neuroimaging viewer, shows it. Hence the flip.
+_PLANES = (("sagittal", 0, True), ("coronal", 1, False), ("axial", 2, False))
+
+
+def in_simulation() -> bool:
+    """Whether this process is a local simulator run launched by a tutorial Makefile.
+
+    Needs **both** signals, so that neither alone can open the gate on a trust:
+
+    - ``FlipConstants.LOCAL_DEV`` is true. Every trust deployment (compose and Helm) sets it to
+      ``false``, and the ``flip`` library already reads its data from the trust APIs on that value,
+      so a client with it true cannot be running a real cohort. It is not enough on its own because
+      it *defaults* to true when the variable is unset.
+    - ``DEBUG_SAMPLES_DIR`` is set and non-blank. Only the tutorial Makefiles' ``sim`` recipe sets
+      it; it is not a ``config.json`` key, so nothing uploaded with the app can supply it.
+
+    Returns:
+        bool: True only inside a Makefile-launched simulator run.
+    """
+    if not os.environ.get(DEBUG_DIR_ENV, "").strip():
+        return False
+    try:
+        return FlipConstants.LOCAL_DEV is True
+    except Exception:  # unreadable settings: fail closed
+        return False
+
+
+def samples_enabled(config: dict) -> bool:
+    """Whether this job writes client-local debug images.
+
+    True only when ``config.json`` sets ``SAVE_DEBUG_SAMPLES`` **and** :func:`in_simulation` holds.
+    The config is the user's switch — false means no images, simulator included. It is never what
+    separates simulation from production, because ``config.json`` travels with the app to the trust:
+    in production (``LOCAL_DEV=false``, no ``DEBUG_SAMPLES_DIR``) this returns False whatever the
+    config says, so patient-derived pixels are never written at a trust.
+
+    This is the single gate for both writers. ``save_grid`` and ``save_triplanar`` consult it and
+    no-op when it is false, so a caller must never add a second condition of its own: a sampling
+    block that runs on a looser test than this one pays for its output and then discards it.
+
+    Args:
+        config (dict): The parsed ``config.json``.
+
+    Returns:
+        bool: True when ``SAVE_DEBUG_SAMPLES`` is truthy and this is a simulator run.
+    """
+    return bool(config.get("SAVE_DEBUG_SAMPLES", False)) and in_simulation()
+
+
+def plot_every(config: dict) -> int:
+    """Iterations between tri-planar figures, from ``DEBUG_PLOT_EVERY`` (default 100).
+
+    Args:
+        config (dict): The parsed ``config.json``.
+
+    Returns:
+        int: The interval. Zero or negative disables the figures while leaving the grids on.
+    """
+    return int(config.get("DEBUG_PLOT_EVERY", DEFAULT_PLOT_EVERY))
+
+
+def due_for_plot(iteration: int, config: dict) -> bool:
+    """Whether ``iteration`` (0-based) is one of the every-N steps that gets a tri-planar figure.
+
+    True on iteration 0 as well as every Nth after it, so a run that crashes early still leaves one
+    figure behind — the first is also the most diagnostic, being the only look at an untrained model.
+
+    Args:
+        iteration (int): Zero-based training step, counted across rounds and epochs (not within one
+            epoch — an epoch shorter than the interval would otherwise fire once per epoch).
+        config (dict): The parsed ``config.json``.
+
+    Returns:
+        bool: True when a figure is due and debug output is enabled at all.
+    """
+    if not samples_enabled(config):
+        return False
+    interval = plot_every(config)
+    return interval > 0 and iteration % interval == 0
+
+
+def debug_dir() -> Path:
+    """``DEBUG_SAMPLES_DIR``, created if needed; the tutorial Makefiles point it under fl-tutorials/data/."""
+    target = Path(os.environ[DEBUG_DIR_ENV].strip()).expanduser().resolve()
+    target.mkdir(parents=True, exist_ok=True)
+    return target
+
+
+def _to_2d(batch: torch.Tensor) -> torch.Tensor:
+    """Reduce a batch to something ``make_grid`` can tile, taking a mid-slice of any volume.
+
+    Args:
+        batch (torch.Tensor): ``(B, C, H, W)`` or ``(B, C, H, W, D)``.
+
+    Returns:
+        torch.Tensor: ``(B, C, H, W)`` — volumes sliced at the centre of the last axis, which for the
+        RAS-oriented volumetric chain is the axial mid-slice.
+    """
+    if batch.ndim == 5:
+        return batch[..., batch.shape[-1] // 2]
+    return batch
+
+
+def _name_parts(name: str, site_name: str, step: int | None) -> list[str]:
+    """Filename parts shared by both writers: what, where, when — in that order."""
+    parts = [name]
+    if site_name:
+        parts.append(site_name)
+    if step is not None:
+        parts.append(f"step{step:04d}")
+    return parts
+
+
+def _prepare(batch: torch.Tensor, max_images: int) -> torch.Tensor:
+    """Detach, truncate and flatten a batch to greyscale 2-D images ready for tiling."""
+    prepared = _to_2d(batch.detach().float().cpu()[:max_images])
+    if prepared.shape[1] > 1:
+        # Multi-channel here means a latent, not a colour image; show the first channel rather than
+        # letting save_image interpret 3 channels as RGB.
+        prepared = prepared[:, :1]
+    return prepared
+
+
+def save_grid(
+    tensors: dict[str, torch.Tensor],
+    name: str,
+    config: dict,
+    *,
+    site_name: str = "",
+    step: int | None = None,
+) -> Path | None:
+    """Write one PNG tiling the given batches, one row per entry. No-op unless enabled.
+
+    Rows are normalised **jointly**, not per image: independent scaling would rescale a washed-out
+    reconstruction back onto its input's range and hide precisely the intensity drift the comparison
+    is meant to expose.
+
+    Never raises. Debug output must not be able to fail a training round, so any error (a read-only
+    workspace, a tensor of an unexpected rank) is logged and swallowed.
+
+    Args:
+        tensors (dict[str, torch.Tensor]): Batches to tile, one row each, in insertion order — e.g.
+            ``{"input": images, "reconstruction": recon}``. Keys appear in the filename.
+        name (str): Short label for what this grid shows, e.g. ``"reconstruction"``.
+        config (dict): The parsed ``config.json``; read for ``SAVE_DEBUG_SAMPLES``/``DEBUG_SAMPLES_MAX``.
+        site_name (str): Client name, included in the filename so a shared workspace cannot collide.
+        step (int | None): Round or epoch counter, included in the filename to order the images.
+
+    Returns:
+        Path | None: The file written, or None when disabled or on any failure.
+    """
+    if not samples_enabled(config) or not tensors:
+        return None
+
+    try:
+        max_images = int(config.get("DEBUG_SAMPLES_MAX", DEFAULT_MAX_IMAGES))
+        rows = [_prepare(batch, max_images) for batch in tensors.values()]
+        columns = min(row.shape[0] for row in rows)
+        grid = make_grid(
+            torch.cat([row[:columns] for row in rows]),
+            nrow=columns,
+            normalize=True,
+            scale_each=False,
+        )
+
+        label = "_".join([name, *tensors.keys()]) if len(tensors) > 1 else name
+        path = debug_dir() / ("_".join(_name_parts(label, site_name, step)) + ".png")
+
+        save_image(grid, str(path))
+        logger.info(f"[DEBUG] Wrote {columns} sample(s) to {path}")
+        return path
+    except Exception as err:
+        logger.warning(f"[DEBUG] Could not write debug samples {name!r}: {err}")
+        return None
+
+
+def _central_slice(volume: torch.Tensor, axis: int, flip_lr: bool = False) -> np.ndarray:
+    """The central slice of one 3-D volume along ``axis``, oriented for display.
+
+    Args:
+        volume (torch.Tensor): A single volume, ``(X, Y, Z)``, channel already removed.
+        axis (int): Spatial axis to hold fixed — 0 sagittal, 1 coronal, 2 axial.
+        flip_lr (bool): Mirror the slice horizontally after the rotation. True for the sagittal
+            view only; see :data:`_PLANES` for why it alone needs it.
+
+    Returns:
+        np.ndarray: The 2-D slice, rotated so the axis that runs superior (or anterior, for the
+        axial view) points up the page, and mirrored if this view needs it.
+    """
+    plane = volume.index_select(axis, torch.tensor([volume.shape[axis] // 2], device=volume.device)).squeeze(axis)
+    rotated = np.rot90(plane.numpy())
+    return np.fliplr(rotated) if flip_lr else rotated
+
+
+def save_triplanar(
+    volumes: dict[str, torch.Tensor],
+    name: str,
+    config: dict,
+    *,
+    site_name: str = "",
+    step: int | None = None,
+    sample: int = 0,
+) -> Path | None:
+    """Write a 3x2 tri-planar figure: one row per view, one column per batch. No-op unless enabled.
+
+    Rows are sagittal / coronal / axial central cuts; columns are the entries of ``volumes`` in
+    insertion order — so ``{"input": images, "reconstruction": recon}`` puts the input beside what
+    the model made of it, in the same three planes.
+
+    This is the figure that answers questions a mid-slice grid cannot. An autoencoder can look
+    plausible on one axial slice while having learnt nothing about through-plane structure, and that
+    shows up immediately in the sagittal and coronal views.
+
+    **Each panel is scaled to its own range**, unlike :func:`save_grid`, which normalises jointly.
+    The trade-off is deliberate and goes the other way here: a joint scale is the honest one for
+    spotting intensity drift, but an early reconstruction spans a far wider range than the input
+    (noise, including negatives), so a shared scale compresses the input to flat grey and neither
+    column can be read. These figures exist to show *structure* — is there a brain, are the
+    ventricles in the right place, does the tumour survive — and structure needs per-panel contrast.
+    Read absolute intensity off the grids, not off these.
+
+    Never raises: debug output must not be able to fail a training round.
+
+    Args:
+        volumes (dict[str, torch.Tensor]): Batches to compare, ``(B, C, X, Y, Z)``, one column each.
+        name (str): Short label for what this figure shows, used in the filename.
+        config (dict): The parsed ``config.json``.
+        site_name (str): Client name, included in the filename so a shared workspace cannot collide.
+        step (int | None): Iteration or epoch counter, included in the filename to order the figures.
+        sample (int): Which item of the batch to show.
+
+    Returns:
+        Path | None: The file written, or None when disabled, when the tensors are not volumetric,
+        or on any failure.
+    """
+    if not samples_enabled(config) or not volumes:
+        return None
+
+    try:
+        import matplotlib
+
+        # Agg before pyplot: an FL client has no display, and the default backend would try to pick
+        # an interactive one and fail at import inside a training round.
+        matplotlib.use("Agg")
+        import matplotlib.pyplot as plt
+
+        picked = {}
+        for label, batch in volumes.items():
+            if batch.ndim != 5:
+                logger.info(f"[DEBUG] {name!r}: {label} is not volumetric ({tuple(batch.shape)}); skipping figure.")
+                return None
+            picked[label] = batch.detach().float().cpu()[min(sample, batch.shape[0] - 1), 0]
+
+        figure, axes = plt.subplots(len(_PLANES), len(picked), figsize=(3.2 * len(picked), 3.2 * len(_PLANES)))
+        axes = np.atleast_2d(axes).reshape(len(_PLANES), len(picked))
+        figure.patch.set_facecolor("black")
+
+        for row, (view, axis, flip_lr) in enumerate(_PLANES):
+            for column, (label, volume) in enumerate(picked.items()):
+                axis_handle = axes[row][column]
+                # No vmin/vmax: matplotlib scales each panel to its own min/max, which is the point
+                # (see the note on contrast above). A flat panel is left to matplotlib, which
+                # renders a constant array as mid-grey rather than dividing by a zero range.
+                axis_handle.imshow(_central_slice(volume, axis, flip_lr), cmap="gray")
+                axis_handle.set_facecolor("black")
+                axis_handle.axis("off")
+                if row == 0:
+                    axis_handle.set_title(label, color="0.85", fontsize=11)
+                if column == 0:
+                    # A y-label survives axis("off"), which is why the view name rides here rather
+                    # than in a title that would collide with the column headings.
+                    axis_handle.set_ylabel(view, color="0.85", fontsize=10)
+                    axis_handle.axis("on")
+                    axis_handle.set_xticks([])
+                    axis_handle.set_yticks([])
+                    for spine in axis_handle.spines.values():
+                        spine.set_visible(False)
+
+        figure.tight_layout()
+        path = debug_dir() / ("_".join(_name_parts(name, site_name, step)) + ".png")
+        figure.savefig(path, facecolor=figure.get_facecolor(), dpi=110)
+        plt.close(figure)
+        logger.info(f"[DEBUG] Wrote tri-planar figure to {path}")
+        return path
+    except Exception as err:
+        logger.warning(f"[DEBUG] Could not write tri-planar figure {name!r}: {err}")
+        return None
+
+
+# --------------------------------------------------------------------------------------------
+# Analysis half: loss curves from a simulator log. Host-side only — see the module docstring.
+# --------------------------------------------------------------------------------------------
+
+#: Log lines arrive as ``<timestamp> - <logger> - INFO - <message>``. Splitting the message off is
+#: what stops the timestamp's own colons ("20:10:04") being read as data. A line that does *not*
+#: match is a continuation of the message above it — the per-epoch summary is logged as one
+#: multi-line string, so its numbers live on unprefixed lines — and is parsed whole.
+_MESSAGE_RE = re.compile(r" - (?:DEBUG|INFO|WARNING|ERROR|CRITICAL) - (?P<message>.*)$")
+
+#: One ``name: number`` pair inside a message. Deliberately generic: every tutorial here logs its
+#: losses in that shape but with different names, so hard-coding one tutorial's names would leave
+#: the other two with empty plots. The name charset excludes quotes,
+#: commas and hyphens, which is what keeps a name from bridging across ``a: 1, b: 2`` or reaching
+#: back into the logger prefix; ``nan``/``inf`` are accepted as values so a run that broke still
+#: plots up to the point it broke.
+_PAIR_RE = re.compile(r"(?P<name>[A-Za-z][\w ()]*?)\s*:\s*(?P<value>[-+]?(?:\d+\.?\d*(?:[eE][-+]?\d+)?|nan|inf))\b")
+
+#: Fixed colours for the names this tree logs today, so two sites' figures can be read side by side
+#: and successive runs stay comparable. Anything unrecognised is assigned from _PALETTE in sorted
+#: order, which is deterministic for a given set of names.
+COLOURS = {
+    "l1": "#4ea1ff",
+    "perceptual": "#ffd166",
+    "gan": "#ff6b6b",
+    "KL": "#9b7bff",
+    "Total loss D": "#7bffb0",
+    "Total loss DM": "#4ea1ff",
+    "Validation loss (L1)": "#4ea1ff",
+    "Validation SSIM (foreground)": "#ff9f4e",
+    "Validation DM": "#ff9f4e",
+    # The per-epoch aggregates, coloured to match the per-iteration term each one totals.
+    "L1": "#4ea1ff",
+    "Perceptual": "#ffd166",
+    "GAN": "#ff6b6b",
+    "KLD": "#9b7bff",
+    "Total loss G": "#eeeeee",
+}
+_PALETTE = ("#4ea1ff", "#ffd166", "#ff6b6b", "#9b7bff", "#7bffb0", "#ff9f4e", "#59d4d4", "#d98cc4")
+
+#: Smoothing for the per-iteration panel, which is unreadable raw at BATCH_SIZE 2.
+EMA_ALPHA = 0.02
+
+#: A series needs at least this many points to be worth an axis; it also filters out the stray
+#: ``name: number`` pairs that appear once in a startup banner.
+MIN_POINTS = 2
+
+#: Series logged at the same cadence share a length. Anything within this fraction of the longest
+#: series is treated as per-iteration; the rest is per-epoch. A tutorial that logs only per-epoch
+#: (the two diffusion ones) therefore yields a single panel rather than an empty one.
+#:
+#: The same ratio is then applied *within* the per-epoch band, which is what keeps incidental
+#: ``name: number`` pairs out of the figure. A log is full of them — ``localhost: 34953``,
+#: ``PID: 1005143``, ``train size: 30`` — and they are not distinguishable from a metric by name or
+#: by value. They are distinguishable by cadence: a metric is logged once per iteration or once per
+#: epoch, so it appears as many times as there were iterations or epochs, while a startup line
+#: appears once or twice however long the run goes on.
+DENSE_RATIO = 0.9
+
+
+def _ema(values: list[float], alpha: float = EMA_ALPHA) -> list[float]:
+    """Exponential moving average, skipping non-finite points rather than propagating them."""
+    out: list[float] = []
+    acc = next((v for v in values if v == v and abs(v) != float("inf")), 0.0)
+    for value in values:
+        if value == value and abs(value) != float("inf"):
+            acc = alpha * value + (1 - alpha) * acc
+        out.append(acc)
+    return out
+
+
+def parse_log(path: Path) -> dict[str, list[float]]:
+    """Return every ``name: number`` series found in one site's ``log.txt``, in order of appearance.
+
+    Args:
+        path (Path): The site's ``log.txt``.
+
+    Returns:
+        dict[str, list[float]]: Series name to its values, longest series first.
+    """
+    series: dict[str, list[float]] = {}
+    with path.open(errors="replace") as handle:
+        for line in handle:
+            prefixed = _MESSAGE_RE.search(line)
+            message = prefixed.group("message") if prefixed else line
+            for match in _PAIR_RE.finditer(message):
+                series.setdefault(match.group("name"), []).append(float(match.group("value")))
+    return dict(sorted(series.items(), key=lambda item: -len(item[1])))
+
+
+def split_by_cadence(series: dict[str, list[float]]) -> tuple[dict[str, list[float]], dict[str, list[float]]]:
+    """Split series into per-iteration and per-epoch groups by how often each was logged.
+
+    Args:
+        series (dict[str, list[float]]): Output of :func:`parse_log`.
+
+    Returns:
+        tuple: ``(dense, sparse)``; ``sparse`` is empty when everything shares one cadence.
+    """
+    usable = {name: values for name, values in series.items() if len(values) >= MIN_POINTS}
+    if not usable:
+        return {}, {}
+    longest = max(len(values) for values in usable.values())
+    dense = {n: v for n, v in usable.items() if len(v) >= longest * DENSE_RATIO}
+
+    remainder = {n: v for n, v in usable.items() if len(v) < longest * DENSE_RATIO}
+    if not remainder:
+        return dense, {}
+    epoch_cadence = max(len(values) for values in remainder.values())
+    sparse = {n: v for n, v in remainder.items() if len(v) >= epoch_cadence * DENSE_RATIO}
+    return dense, sparse
+
+
+def _colour_for(name: str, index: int) -> str:
+    return COLOURS.get(name, _PALETTE[index % len(_PALETTE)])
+
+
+def _style(axis, title: str) -> None:
+    """Apply the shared dark styling to one axis."""
+    axis.set_facecolor("#111111")
+    axis.set_title(title, color="#eeeeee")
+    axis.tick_params(colors="#bbbbbb")
+    axis.grid(True, color="#333333", linewidth=0.5)
+    for spine in axis.spines.values():
+        spine.set_color("#444444")
+    axis.xaxis.label.set_color("#bbbbbb")
+    axis.yaxis.label.set_color("#bbbbbb")
+
+
+def _legend(axis, handles=None, labels=None) -> None:
+    """Legend outside the axes on the right, so it cannot sit on top of the curves.
+
+    The generic parser finds however many series a tutorial logs — eight for the autoencoder's
+    per-epoch panel — which is more than fits inside the plot area.
+    """
+    axis.legend(
+        *([handles, labels] if handles is not None else []),
+        loc="upper left",
+        bbox_to_anchor=(1.12, 1.0),
+        fontsize=8,
+        facecolor="#1c1c1c",
+        edgecolor="#444444",
+        labelcolor="#eeeeee",
+    )
+
+
+def plot_site(site: str, series: dict[str, list[float]], out_dir: Path) -> Path:
+    """Write the loss figure for one site and return its path.
+
+    The per-iteration panel is log-scaled when every value is positive: the terms routinely span two
+    decades, and on a linear axis the smallest of them — usually the one that is misweighted, which
+    is exactly what you are looking for — is a flat line on the floor.
+
+    Args:
+        site (str): Site name, used in the title and filename.
+        series (dict[str, list[float]]): Output of :func:`parse_log`.
+        out_dir (Path): Directory to write into; created if absent.
+
+    Returns:
+        Path: The figure written.
+    """
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    dense, sparse = split_by_cadence(series)
+    panels = [p for p in (dense, sparse) if p]
+    figure, axes = plt.subplots(len(panels), 1, figsize=(13, 4 * len(panels)), facecolor="#111111", squeeze=False)
+
+    for axis, panel in zip(axes.flatten(), panels, strict=True):
+        is_dense = panel is dense
+        # SSIM is bounded [0, 1] where the losses are not, so it gets its own axis or it flattens
+        # every other series against the top of the plot.
+        twin_names = [n for n in panel if "SSIM" in n] if not is_dense else []
+        twin = axis.twinx() if twin_names else None
+
+        for index, (name, values) in enumerate(panel.items()):
+            colour = _colour_for(name, index)
+            target = twin if name in twin_names else axis
+            if is_dense:
+                target.plot(values, color=colour, alpha=0.15, linewidth=0.6)
+                target.plot(_ema(values), color=colour, linewidth=1.8, label=name)
+            else:
+                steps = range(1, len(values) + 1)
+                target.plot(steps, values, color=colour, marker="o", markersize=3, label=name)
+
+        # Log y whenever everything is positive. The terms routinely span two or three decades, and
+        # on a linear axis the smallest of them — usually the misweighted one, which is exactly what
+        # you are looking for — is a flat line on the floor. On the per-epoch panel the first epoch
+        # is several times any later one, so a linear axis spends most of its height on it.
+        plotted = [values for name, values in panel.items() if name not in twin_names]
+        finite = [v for values in plotted for v in values if v == v and abs(v) != float("inf")]
+        if finite and min(finite) > 0:
+            axis.set_yscale("log")
+        axis.set_xlabel("training iteration" if is_dense else "local epoch")
+        axis.set_ylabel("weighted contribution to $\\mathcal{L}_G$" if is_dense else "loss")
+        _style(axis, f"{site} — {'per-iteration loss terms' if is_dense else 'per-epoch summary'}")
+
+        if twin is not None:
+            twin.set_ylabel(" / ".join(twin_names), color=COLOURS.get(twin_names[0], _PALETTE[5]))
+            twin.tick_params(axis="y", colors=COLOURS.get(twin_names[0], _PALETTE[5]))
+            for spine in twin.spines.values():
+                spine.set_color("#444444")
+            handles, labels = axis.get_legend_handles_labels()
+            extra_handles, extra_labels = twin.get_legend_handles_labels()
+            _legend(axis, handles + extra_handles, labels + extra_labels)
+        else:
+            _legend(axis)
+
+    figure.tight_layout()
+    out_dir.mkdir(parents=True, exist_ok=True)
+    destination = out_dir / f"losses_{site}.png"
+    figure.savefig(destination, dpi=110, facecolor=figure.get_facecolor())
+    plt.close(figure)
+    return destination
+
+
+def main() -> None:
+    """Plot every site's losses from a simulator workspace. Host-side entry point."""
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        "--workspace",
+        type=Path,
+        default=Path("/tmp/nvflare") / Path.cwd().name / "flip_fedavg",
+        help="simulator workspace root (default: this tutorial's, from the current directory)",
+    )
+    parser.add_argument("--out-dir", type=Path, default=None, help="defaults to $DEBUG_SAMPLES_DIR, else ./")
+    args = parser.parse_args()
+
+    out_dir = args.out_dir or Path(os.environ.get(DEBUG_DIR_ENV, "."))
+
+    logs = sorted(args.workspace.glob("site-*/log.txt"))
+    if not logs:
+        raise SystemExit(f"no site logs under {args.workspace} — pass --workspace")
+
+    for log in logs:
+        site = log.parent.name
+        series = parse_log(log)
+        dense, sparse = split_by_cadence(series)
+        if not dense and not sparse:
+            print(f"{site}: no scalar lines yet, skipping")
+            continue
+        destination = plot_site(site, series, out_dir)
+        counts = ", ".join(f"{name} x{len(values)}" for name, values in list(dense.items())[:4])
+        print(f"{site}: {counts} -> {destination}")
+
+
+if __name__ == "__main__":
+    main()

code/trainer.py

@@ -0,0 +1,429 @@
+# Copyright (c) 2026 Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""FLARE Client API script for the FLIP pixel-space diffusion tutorial.
+
+Trains a denoising diffusion model (MONAI's ``DiffusionModelUNet``) to generate 2-D chest X-rays, as
+an ordinary FedAvg job. On a ``train`` task the client trains the received model for ``LOCAL_ROUNDS``
+epochs and sends back the weight difference; on an ``evaluate`` task it scores the aggregated model
+on its held-out split.
+
+Training noises each image at a random timestep and teaches the UNet to predict that noise (MSE).
+Every ``VALIDATE_EVERY`` epochs the same loss is scored on the held-out split and, in local runs,
+sample images are drawn. Training runs in full precision.
+"""
+
+import argparse
+import json
+import logging
+import math
+import os
+from pathlib import Path
+
+import numpy as np
+import nvflare.client as flare
+import pandas as pd
+import pydicom
+import torch
+from flip import FLIP
+from flip.constants import ResourceType
+from models import get_model
+from monai.data import DataLoader, Dataset
+from monai.inferers import DiffusionInferer
+from monai.networks.schedulers import DDPMScheduler
+from nvflare.client.tracking import SummaryWriter
+from transforms import get_xray_transforms
+
+# plot_utils.py is a local-development helper (../dev/): job.py stages it into simulator runs, and it is
+# never uploaded, so on a trust this import fails and debug output is simply off.
+try:
+    from plot_utils import samples_enabled, save_grid
+except ModuleNotFoundError as err:
+    if err.name != "plot_utils":
+        raise
+
+    def samples_enabled(config: dict) -> bool:
+        return False
+
+    def save_grid(*args: object, **kwargs: object) -> None:
+        return None
+
+
+logger = logging.getLogger(__name__)
+
+#: Epochs between in-loop validation and sampling when ``VALIDATE_EVERY`` is absent from the config.
+DEFAULT_VALIDATE_EVERY = 10
+
+#: Images drawn per sample grid.
+NUM_DEBUG_SAMPLES = 4
+
+#: Max gradient norm when ``GRAD_CLIP_NORM`` is absent from the config.
+DEFAULT_GRAD_CLIP_NORM = 1.0
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--project_id", type=str, default="")
+    return parser.parse_args()
+
+
+def load_query() -> str:
+    """Read the cohort query from the client app config.
+
+    FLARE's TaskScriptRunner does a naive whitespace split on task_script_args, so the SQL query
+    (which can contain spaces) is plumbed via the top-level ``query`` key in
+    ``config/config_fed_client.json`` rather than as a CLI flag. In dev/simulator mode this is
+    ignored by ``flip.get_dataframe``.
+    """
+    client_cfg = Path(__file__).parent.parent / "config" / "config_fed_client.json"
+    if client_cfg.exists():
+        try:
+            return json.loads(client_cfg.read_text()).get("query", "")
+        except Exception:
+            return ""
+    return ""
+
+
+def load_config() -> dict:
+    """Load the user-supplied config.json that sits next to this script."""
+    config_path = Path(__file__).parent.resolve() / "config.json"
+    with open(config_path) as f:
+        return json.load(f)
+
+
+def validate_every(config: dict) -> int:
+    """Epochs between in-loop validation and sampling, from ``VALIDATE_EVERY``; zero or less turns it off.
+
+    The default is above 1 because the pass also samples, which runs the full reverse diffusion.
+    """
+    return int(config.get("VALIDATE_EVERY", DEFAULT_VALIDATE_EVERY))
+
+
+def due_for_validation(epoch: int, epochs: int, config: dict) -> bool:
+    """Whether this 0-based epoch validates: every ``VALIDATE_EVERY`` epochs, and always the last."""
+    interval = validate_every(config)
+    if interval <= 0:
+        return False
+    return (epoch + 1) % interval == 0 or epoch == epochs - 1
+
+
+def grad_clip_norm(config: dict) -> float:
+    """Max global gradient norm, from ``GRAD_CLIP_NORM``; zero or less turns clipping off."""
+    return float(config.get("GRAD_CLIP_NORM", DEFAULT_GRAD_CLIP_NORM))
+
+
+def lr_at(epoch: int, config: dict) -> float:
+    """The learning rate at an epoch of the job: half a cosine from ``LR_START`` to ``LR_END``, then held.
+
+    The decay spans the job's first ``LR_DECAY_EPOCHS`` epochs, so ``epoch`` counts across rounds
+    (``global_round * LOCAL_ROUNDS + epoch``). Without ``LR_END`` or a positive ``LR_DECAY_EPOCHS`` the
+    rate stays at ``LR_START``.
+    """
+    start = float(config["LR_START"])
+    end = float(config.get("LR_END", start))
+    decay_epochs = int(config.get("LR_DECAY_EPOCHS", 0) or 0)
+    if decay_epochs <= 0:
+        return start
+    progress = min(epoch, decay_epochs) / decay_epochs
+    return end + (start - end) * 0.5 * (1.0 + math.cos(math.pi * progress))
+
+
+def set_learning_rate(optimizer: torch.optim.Optimizer, lr: float) -> None:
+    for group in optimizer.param_groups:
+        group["lr"] = lr
+
+
+def site_data_paths(site_name: str) -> tuple[Path, Path] | None:
+    """This simulated site's own ``(images_dir, dataframe)``, from ``SITE{N}_IMAGES_DIR`` / ``SITE{N}_DATAFRAME``.
+
+    Set only by the tutorial Makefile for a local run on per-site data ("site-1" reads ``SITE1_*``).
+    ``None`` when either is unset or blank: the site then reads the shared ``DEV_*`` dataset through
+    ``FLIP``, as it does on a trust.
+    """
+    prefix = site_name.replace("-", "").upper()
+    images_dir = os.environ.get(f"{prefix}_IMAGES_DIR", "").strip()
+    dataframe = os.environ.get(f"{prefix}_DATAFRAME", "").strip()
+    if not images_dir or not dataframe:
+        return None
+    return Path(images_dir), Path(dataframe)
+
+
+def split_for_site(items: list, site_name: str) -> list:
+    """In the simulator, give "site-1" and "site-2" half the shared dataset each; on a trust, all of it."""
+    half = len(items) // 2
+    site = site_name.replace("-", "")
+    if site == "site1":
+        return items[:half]
+    if site == "site2":
+        return items[half:]
+    return items
+
+
+def noise_prediction_loss(
+    model: torch.nn.Module, inferer: DiffusionInferer, scheduler: DDPMScheduler, images: torch.Tensor
+) -> torch.Tensor:
+    """The diffusion objective: noise each image at a random timestep and score the predicted noise."""
+    noise = torch.randn_like(images)
+    timesteps = torch.randint(0, scheduler.num_train_timesteps, (images.shape[0],), device=images.device)
+    noise_pred = inferer(inputs=images, diffusion_model=model.diffusion_model, noise=noise, timesteps=timesteps)
+    return torch.nn.functional.mse_loss(noise_pred.float(), noise.float())
+
+
+def score_split(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    inferer: DiffusionInferer,
+    scheduler: DDPMScheduler,
+    device: torch.device,
+) -> float:
+    """Mean noise-prediction loss over a split, in eval mode; restores the caller's mode afterwards."""
+    was_training = model.diffusion_model.training
+    model.diffusion_model.eval()
+    losses = []
+    try:
+        with torch.no_grad():
+            for batch in loader:
+                losses.append(noise_prediction_loss(model, inferer, scheduler, batch["image"].to(device)).item())
+    finally:
+        model.diffusion_model.train(was_training)
+    return float(np.mean(losses)) if losses else float("nan")
+
+
+def sample_and_save(
+    model: torch.nn.Module,
+    inferer: DiffusionInferer,
+    scheduler: DDPMScheduler,
+    config: dict,
+    device: torch.device,
+    *,
+    step: int | None = None,
+) -> None:
+    """Generate ``NUM_DEBUG_SAMPLES`` images and save them as one grid, in local runs only.
+
+    Samples are the honest read on a diffusion model: the loss barely moves between a model that
+    generates radiographs and one that generates texture.
+    """
+    if not samples_enabled(config):
+        return
+    shape = [NUM_DEBUG_SAMPLES, config["net_config"]["diffusion_model"]["in_channels"], *config["spatial_shape"]]
+    was_training = model.diffusion_model.training
+    model.diffusion_model.eval()
+    try:
+        with torch.no_grad():
+            samples = inferer.sample(
+                input_noise=torch.randn(shape, device=device),
+                diffusion_model=model.diffusion_model,
+                scheduler=scheduler,
+                save_intermediates=False,
+                # No progress bar: simulated sites share one terminal and would overwrite each other's.
+                verbose=False,
+            )
+    finally:
+        model.diffusion_model.train(was_training)
+    save_grid({"sample": samples}, "samples", config, site_name=flare.get_site_name(), step=step)
+
+
+class DiffusionTrainer:
+    """The model, optimiser, schedules and data for pixel-space diffusion training.
+
+    One instance lives for the whole ``flare.is_running()`` loop, so the optimiser state carries over
+    from round to round.
+    """
+
+    def __init__(self, config: dict, project_id: str, query: str):
+        self.config = config
+        self.project_id = project_id
+        self.epochs = config["LOCAL_ROUNDS"]
+
+        self.model = get_model()
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self.optimizer = torch.optim.Adam(
+            self.model.diffusion_model.parameters(), lr=config["LR_START"], weight_decay=1e-6, amsgrad=True
+        )
+        self.noise_scheduler = DDPMScheduler(
+            num_train_timesteps=1000,
+            schedule="scaled_linear_beta",
+            clip_sample=False,
+            prediction_type="epsilon",
+            beta_start=0.0015,
+            beta_end=0.0195,
+        )
+        self.inferer = DiffusionInferer(scheduler=self.noise_scheduler)
+
+        # Data
+        self.flip = FLIP()
+        site_name = flare.get_site_name()
+        site_paths = site_data_paths(site_name)
+        if site_paths is not None:
+            # Local run on per-site data: this site's own folder and CSV, used whole.
+            self.images_dir, dataframe_path = site_paths
+            dataframe = pd.read_csv(dataframe_path, sep=None, engine="python")
+            logger.info(f"Reading {dataframe_path} and {self.images_dir}")
+        else:
+            self.images_dir = None
+            dataframe = self.flip.get_dataframe(project_id=project_id, query=query)
+        train_items, val_items = self.build_datalist(dataframe)
+        if site_paths is None:
+            train_items = split_for_site(train_items, site_name)
+
+        workers = int(config.get("DATALOADER_WORKERS", 4))  # as in Ark+ (FLIP#1434)
+        self.train_loader = DataLoader(
+            Dataset(train_items, transform=get_xray_transforms()),
+            batch_size=config["BATCH_SIZE"],
+            shuffle=True,
+            num_workers=workers,
+        )
+        self.val_loader = DataLoader(
+            Dataset(val_items, transform=get_xray_transforms(is_validation=True)),
+            batch_size=config["BATCH_SIZE"],
+            shuffle=False,
+            num_workers=workers,
+        )
+
+    def build_datalist(self, dataframe: pd.DataFrame) -> tuple[list, list]:
+        """List each accession's readable DICOM images and split them into train and validation.
+
+        As ``xray_classification``'s ``build_datalist``, without the labels. Each header is parsed up
+        front, so an unreadable file is dropped here rather than failing a training step.
+        """
+        datalist: list[dict[str, str]] = []
+        for accession_id in dataframe["accession_id"]:
+            if self.images_dir is not None:
+                accession_folder_path = self.images_dir / str(accession_id)
+            else:
+                try:
+                    accession_folder_path = self.flip.get_by_accession_number(
+                        self.project_id, accession_id, resource_type=[ResourceType.DICOM]
+                    )
+                except Exception as err:
+                    logger.info(f"Could not get image data folder path for {accession_id}: {err}")
+                    continue
+
+            for image in sorted(accession_folder_path.rglob("*.dcm")):
+                try:
+                    pydicom.dcmread(str(image), stop_before_pixels=True)
+                except Exception as err:
+                    logger.warning(f"Skipping invalid DICOM {image.name}: {err}")
+                    continue
+                datalist.append({"image": str(image)})
+
+        logger.info(f"Found {len(datalist)} files in total.")
+        val_size = int(self.config["VAL_SPLIT"] * len(datalist))
+        return datalist[val_size:], datalist[:val_size]
+
+    def load_weights(self, weights: dict[str, torch.Tensor]) -> None:
+        self.model.load_state_dict(state_dict=weights)
+        self.model.diffusion_model.to(self.device)
+
+    def train_one_epoch(self) -> float:
+        """One pass over the training split; returns the mean loss."""
+        self.model.diffusion_model.train()
+        max_norm = grad_clip_norm(self.config)
+        total = 0.0
+        for batch in self.train_loader:
+            self.optimizer.zero_grad(set_to_none=True)
+            loss = noise_prediction_loss(self.model, self.inferer, self.noise_scheduler, batch["image"].to(self.device))
+            if not torch.isfinite(loss):
+                raise RuntimeError("Non-finite loss during training; stopping this site.")
+            loss.backward()
+            if max_norm > 0:
+                torch.nn.utils.clip_grad_norm_(self.model.diffusion_model.parameters(), max_norm)
+            self.optimizer.step()
+            total += loss.item()
+        return total / max(1, len(self.train_loader))
+
+    def train(self, writer: SummaryWriter, global_round: int) -> int:
+        """One ``train`` task: ``LOCAL_ROUNDS`` epochs. Returns the number of optimiser steps."""
+        for epoch in range(self.epochs):
+            step = global_round * self.epochs + epoch + 1
+            lr = lr_at(step - 1, self.config)
+            set_learning_rate(self.optimizer, lr)
+            epoch_loss = self.train_one_epoch()
+            logger.info(f"Epoch {epoch + 1} / {self.epochs}: loss {epoch_loss}")
+            writer.add_scalar("Total loss DM@epoch", epoch_loss, global_step=step)
+            writer.add_scalar("LR DM@epoch", lr, global_step=step)
+
+            if due_for_validation(epoch, self.epochs, self.config):
+                val_loss = score_split(self.model, self.val_loader, self.inferer, self.noise_scheduler, self.device)
+                logger.info(f"Validation loss: {val_loss}")
+                writer.add_scalar("Validation loss DM@epoch", val_loss, global_step=step)
+                sample_and_save(self.model, self.inferer, self.noise_scheduler, self.config, self.device, step=step)
+
+        return self.epochs * len(self.train_loader)
+
+    def validate(self, writer: SummaryWriter) -> float:
+        """One ``evaluate`` task: score the aggregated model on this site's held-out split, and sample."""
+        val_loss = score_split(self.model, self.val_loader, self.inferer, self.noise_scheduler, self.device)
+        sample_and_save(self.model, self.inferer, self.noise_scheduler, self.config, self.device)
+        writer.add_scalar("Total loss DM (val)", val_loss, global_step=0)
+        return val_loss
+
+
+def to_torch_weights(input_model: flare.FLModel) -> dict[str, torch.Tensor]:
+    return {k: torch.as_tensor(v) for k, v in input_model.params.items()}
+
+
+def send_weight_diff(original_params: dict, model: torch.nn.Module, n_iterations: int) -> None:
+    """Send the full-model weight diff for a completed training round.
+
+    Built one tensor at a time to avoid holding a second full copy of the model in RAM (mirrors
+    ``flip.utils.get_model_weights_diff``).
+    """
+    diff = {}
+    for k, v in model.state_dict().items():
+        new_arr = v.detach().cpu().numpy()
+        diff[k] = new_arr - np.asarray(original_params[k])
+        del new_arr
+
+    flare.send(
+        flare.FLModel(
+            params=diff,
+            params_type="DIFF",
+            meta={"NUM_STEPS_CURRENT_ROUND": n_iterations},
+        )
+    )
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    args = parse_args()
+    config = load_config()
+
+    flare.init()
+    writer = SummaryWriter()
+
+    trainer = DiffusionTrainer(config, project_id=args.project_id, query=load_query())
+
+    while flare.is_running():
+        input_model = flare.receive()
+        if input_model is None:
+            break
+
+        global_round = input_model.current_round or 0
+
+        if flare.is_train():
+            logger.info(f"[diffusion trainer] received train task (round {global_round})")
+            trainer.load_weights(to_torch_weights(input_model))
+            n_iterations = trainer.train(writer, global_round)
+            send_weight_diff(input_model.params, trainer.model, n_iterations)
+
+        elif flare.is_evaluate():
+            trainer.load_weights(to_torch_weights(input_model))
+            val_loss = trainer.validate(writer)
+            logger.info(f"Validating the aggregated diffusion model on {flare.get_site_name()}'s data: {val_loss}")
+            flare.send(flare.FLModel(metrics={"val_loss": val_loss}))
+
+        else:
+            logger.warning("Received unknown task; ignoring.")
+
+
+if __name__ == "__main__":
+    main()

code/transforms.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2026 Guy's and St Thomas' NHS Foundation Trust & King's College London
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MONAI transforms for chest-X-ray inputs."""
+
+import json
+from pathlib import Path
+
+import monai.transforms as mt
+
+
+def load_spatial_shape() -> list[int]:
+    """Read ``spatial_shape`` from the ``config.json`` next to this file."""
+    with open(Path(__file__).parent / "config.json", "r") as f:
+        return list(json.load(f)["spatial_shape"])
+
+
+SPATIAL_SHAPE = load_spatial_shape()
+
+
+def to_grayscale(image):
+    """Average an ``(H, W, 3)`` RGB-stored radiograph to ``(H, W)``; grayscale passes through."""
+    if image.ndim == 3 and image.shape[-1] == 3:
+        return image.float().mean(dim=-1)
+    return image
+
+
+def get_xray_transforms(is_validation: bool = False) -> mt.Compose:
+    """Return the MONAI transforms used for chest-X-ray training/validation.
+
+    Args:
+        is_validation (bool): When True, skip random augmentation.
+
+    Returns:
+        mt.Compose: Composed transform pipeline keyed on "image".
+    """
+    transforms = [
+        # A pinned reader: auto-detection picks ITKReader whenever itk is installed, which changes the
+        # axis order, and MONAI's default swap_ij=True transposes the radiograph, which no rotation or
+        # flip undoes (a "fix" mirrors it, silently swapping left and right). swap_ij=False keeps
+        # DICOM's (row, column) order; fl-tutorials/tests/ pins this for every app on this path.
+        mt.LoadImaged(keys=["image"], reader="PydicomReader", swap_ij=False),
+        mt.Lambdad(keys=["image"], func=to_grayscale),
+        mt.EnsureChannelFirstd(keys=["image"], channel_dim="no_channel"),
+        mt.Resized(keys=["image"], spatial_size=SPATIAL_SHAPE),
+        mt.ScaleIntensityd(keys=["image"]),
+    ]
+    if not is_validation:
+        transforms.append(mt.RandAffined(keys=["image"], rotate_range=[-0.05, 0.05], scale_range=[0.01, 0.05]))
+    return mt.Compose(transforms)

diffusion_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5118f8e8a33cd94fb76a4966bd8f8cc9f5bca539aafcb18d5273b8a072d4c081
+size 88359404

evaluation/sample_xray.py

@@ -0,0 +1,44 @@
+"""4x4 grid of chest X-rays from the final global model of the X-ray diffusion tutorial.
+
+The first n samples drawn from one seed, not picked, with the tutorial's own scheduler and sampler (the full
+1000-step DDPM), shown as the tutorial's debug grids are: one min-max over the whole grid. Runs on the CPU so it
+leaves the GPU to the queue. Also saves the samples as a plain tensor.
+
+python sample_xray.py <app custom dir> <FL_global_model.pt> <out.png> [n=16] [seed=0] [threads=8]
+"""
+
+import json
+import sys
+import time
+from pathlib import Path
+
+import torch
+from torchvision.utils import make_grid, save_image
+
+app, ckpt, out = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
+n, seed, threads = (int(sys.argv[i]) if len(sys.argv) > i else d for i, d in ((4, 16), (5, 0), (6, 8)))
+sys.path.insert(0, str(app))
+from models import get_model  # noqa: E402
+from monai.inferers import DiffusionInferer  # noqa: E402
+from monai.networks.schedulers import DDPMScheduler  # noqa: E402
+
+torch.set_num_threads(threads)
+config = json.loads((app / "config.json").read_text())
+model = get_model()
+blob = torch.load(ckpt, map_location="cpu", weights_only=False)  # our own run's output
+model.load_state_dict(blob["model"] if isinstance(blob.get("model"), dict) else blob, strict=True)
+model.eval()
+
+scheduler = DDPMScheduler(num_train_timesteps=1000, schedule="scaled_linear_beta", clip_sample=False,
+                          prediction_type="epsilon", beta_start=0.0015, beta_end=0.0195)  # as trainer.py
+inferer = DiffusionInferer(scheduler=scheduler)
+shape = [n, config["net_config"]["diffusion_model"]["in_channels"], *config["spatial_shape"]]
+generator = torch.Generator().manual_seed(seed)
+start = time.time()
+with torch.no_grad():
+    samples = inferer.sample(input_noise=torch.randn(shape, generator=generator), diffusion_model=model.diffusion_model,
+                             scheduler=scheduler, save_intermediates=False, verbose=False)
+print(f"sampled {n} in {time.time() - start:.0f} s; range {samples.min():.3f}..{samples.max():.3f}")
+torch.save(samples.as_subclass(torch.Tensor).clone(), out.with_suffix(".pt"))
+save_image(make_grid(samples, nrow=int(n**0.5), normalize=True, scale_each=False, padding=2), str(out))
+print("wrote", out)

provenance.json

@@ -0,0 +1,36 @@
+{
+ "flip_commit": "b3a42840ab5a2d73e45b6d596bdea25f1b383502",
+ "install": "pip install \"flip-utils @ git+https://github.com/londonaicentre/FLIP@b3a42840ab5a2d73e45b6d596bdea25f1b383502#subdirectory=flip-utils\"",
+ "launch": "cd fl-tutorials/nvflare/image_synthesis/diffusion_model && python job.py --n_clients 2 --num_rounds 2",
+ "note": "Runs predate the tutorials' move to fl-tutorials/nvflare/image_generation/ (a2c053ec4).",
+ "sites": 2,
+ "hardware": "1 x NVIDIA GeForce RTX 5090 (both simulated sites)",
+ "versions": {
+  "torch": "2.13.0+cu130",
+  "monai": "1.6.0",
+  "nvflare": "2.9.0"
+ },
+ "wall_time": {
+  "start": "2026-10-22 21:33:39",
+  "end": "2026-10-23 02:38:35",
+  "hours_minutes": "5 h 05 min"
+ },
+ "model": "E6",
+ "local_epochs": 300,
+ "code_files": [
+  "config.json",
+  "models.py",
+  "plot_utils.py",
+  "trainer.py",
+  "transforms.py"
+ ],
+ "weights": {
+  "file": "diffusion_model.safetensors",
+  "sha256": "5118f8e8a33cd94fb76a4966bd8f8cc9f5bca539aafcb18d5273b8a072d4c081",
+  "dtype": "float32"
+ },
+ "final_validation_loss": {
+  "site-2": 0.0036,
+  "site-1": 0.0039
+ }
+}