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) 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
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
in the DiffuMT source repo.
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:
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:
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.
Interactive demo
Draw a binary mask in the browser and watch DDIM sampling step by step: DiffuMT project page
Citation
@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}
}