DiffuMT

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

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}
}