README.md

---
license: mit
tags:
  - diffusion
  - image-generation
  - microscopy
  - microtubule
  - mask-conditioned
  - ddim
  - biology
language:
  - en
pipeline_tag: image-to-image
---

# DiffuMT — Mask-Conditioned Diffusion Model for IRM Microtubule Images

This is the diagnostic-selected checkpoint (epoch 290) of a mask-conditioned DDPM
([Konz et al., 2024](https://arxiv.org/abs/2402.05210)) fine-tuned on IRM
(interference reflection microscopy) microtubule images. The model generates
realistic 256×256 synthetic microtubule images conditioned on a binary segmentation mask.

The checkpoint was selected using the three-axis diagnostic from the paper
*"Diagnosing Diversity Collapse and Validating Mask-Conditioned Diffusion for
Labeled Microtubule Microscopy"* (under review). Epoch 290 maximises
the geometric mean of DINOv2 inter-similarity (realism), intra-diversity (collapse
detector), and CIELAB color distribution match — before the model enters the late
mode-collapse phase visible at epoch 400.

## Quick start

```python
import torch
from diffusers import UNet2DModel, DDIMScheduler

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"

# Weights are cached after the first download (~455 MB)
unet = UNet2DModel.from_pretrained("HTW-KI-Werkstatt/DiffuMT", subfolder="unet").to(device)
scheduler = DDIMScheduler.from_pretrained("HTW-KI-Werkstatt/DiffuMT", subfolder="scheduler")
scheduler.set_timesteps(50)
```

See the full sampling walkthrough in
[`notebooks/01_sampling_demo.ipynb`](https://github.com/HTW-KI-Werkstatt/DiffuMT/blob/main/notebooks/01_sampling_demo.ipynb)
in the [DiffuMT source repo](https://github.com/HTW-KI-Werkstatt/DiffuMT).

## Critical: mask conditioning scale

The model was trained with segmentation masks scaled to `{0, 1/255}` — **not** `{0, 1}`.
Feeding a binary `{0, 1}` mask is a 255× stronger conditioning signal and pushes the
model out of distribution (dark, grainy, magenta-tinted output). Always apply `ScaleSeg`
from `utils.py`:

```python
from torchvision import transforms

def scale_seg(tensor):
    return (tensor > 0).float() / 255.0

seg_transform = transforms.Compose([
    transforms.ToTensor(),  # PNG {0,255} → float [0,1]
    scale_seg,              # → {0, 1/255}
])
seg = seg_transform(mask_pil).unsqueeze(0).to(device)  # (1,1,H,W)
```

## Sampling

DDIM with η=0 (deterministic reverse process, 50 steps). Diversity comes from
different random initial noise vectors `x_T`, not stochastic steps:

```python
import torch, numpy as np
from PIL import Image

def sample(unet, scheduler, seg, seed=42, steps=50):
    scheduler.set_timesteps(steps)
    gen = torch.Generator(device=device).manual_seed(seed)
    x = torch.randn((1, 3, 256, 256), generator=gen, device=device)
    with torch.no_grad():
        for t in scheduler.timesteps:
            noise_pred = unet(torch.cat([x, seg], dim=1), t).sample
            x = scheduler.step(noise_pred, t, x).prev_sample
    arr = (x / 2 + 0.5).clamp(0, 1)
    arr = (arr[0].cpu().permute(1, 2, 0).numpy() * 255).astype(np.uint8)
    return Image.fromarray(arr)

image = sample(unet, scheduler, seg)
```

## Model details

| Property | Value |
|---|---|
| Architecture | 2D U-Net (`UNet2DModel`, diffusers) |
| Input channels | 4 (3 RGB + 1 mask) |
| Image size | 256 × 256 |
| Scheduler | DDIMScheduler (linear β schedule) |
| Selected epoch | 290 |
| Training epochs | 1000 (early-stopped by diagnostic) |
| Offset noise | ✓ (matches IRM brightness distribution) |
| Parameters | ~113M |

## Dataset

The **DiffuMT** dataset (2800 mask/real/synthetic triplets) generated with this
checkpoint is at
[`HTW-KI-Werkstatt/DiffuMT`](https://huggingface.co/datasets/HTW-KI-Werkstatt/DiffuMT).

## Interactive demo

Draw a binary mask in the browser and watch DDIM sampling step by step:
[DiffuMT project page](https://huggingface.co/spaces/HTW-KI-Werkstatt/DiffuMT)

## Citation

```bibtex
@inproceedings{konz2024segguideddiffusion,
  title     = {Anatomically-Controllable Medical Image Generation with
               Segmentation-Guided Diffusion Models},
  author    = {Nicholas Konz and Yuwen Chen and Haoyu Dong and Maciej A. Mazurowski},
  booktitle = {MICCAI},
  year      = {2024}
}
```