Pixel Diffusion UNet – Real-world Synchrotron Dataset (DM4CT)
This repository contains the pretrained pixel-space diffusion UNet presented in the paper DM4CT: Benchmarking Diffusion Models for Computed Tomography Reconstruction.
🔗 Project Page: https://dm4ct.github.io/DM4CT/
🔗 Arxiv: https://arxiv.org/abs/2602.18589
🔗 Codebase: https://github.com/DM4CT/DM4CT
🔬 Model Overview
This model learns a prior over CT reconstruction images using a denoising diffusion probabilistic model (DDPM).
It operates directly in pixel space (not latent space).
- Architecture: 2D UNet (Diffusers
UNet2DModel) - Input resolution: 768 × 768
- Channels: 1 (grayscale CT slice)
- Training objective: ε-prediction (standard DDPM formulation)
- Noise schedule: Linear beta schedule
- Training dataset: Real-world Synchrotron Dataset of rocks
- Intensity normalization: Rescaled to (-1, 1)
This model is intended to be combined with data-consistency correction for CT reconstruction.
📊 Dataset: Real-world Synchrotron Dataset
Source: Zenodo
Preprocessing steps:
- Train/test split
- Rescale reconstructed slices to (-1, 1)
- No geometry information is embedded in the model
The model learns an unconditional image prior over CT slices.
🧠 Training Details
- Optimizer: AdamW
- Learning rate: 1e-4
- Hardware: NVIDIA A100 GPU
- Training script: train_pixel.py
🚀 Usage
You can use this model with the diffusers library as follows:
from diffusers import DDPMPipeline
# Load the pipeline
pipeline = DDPMPipeline.from_pretrained("jiayangshi/synchrotron_pixel_diffusion")
# Access the UNet model
model = pipeline.unet
model.eval()
Citation
@inproceedings{
shi2026dmct,
title={{DM}4{CT}: Benchmarking Diffusion Models for Computed Tomography Reconstruction},
author={Shi, Jiayang and Pelt, Dani{\"e}l M and Batenburg, K Joost},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=YE5scJekg5}
}