jiayangshi/synchrotron_latent_diffusion

Latent Diffusion Model – Synchrotron (DM4CT)

This repository contains the pretrained latent-space diffusion model used in the DM4CT: Benchmarking Diffusion Models for CT Reconstruction (ICLR 2026) benchmark.

🔗 Paper: DM4CT: Benchmarking Diffusion Models for Computed Tomography Reconstruction 🔗 Project Page: https://dm4ct.github.io/DM4CT/ 🔗 Codebase: https://github.com/DM4CT/DM4CT

🔬 Model Overview

This model learns a prior over CT reconstruction images in a compressed latent space using a denoising diffusion probabilistic model (DDPM).

Unlike pixel-based diffusion models, diffusion is performed in the latent space of a pretrained autoencoder.

• —Architecture:
• —VQ-VAE (image encoder/decoder)
• —2D UNet operating in latent space
• —Input resolution (image space): 768 × 768
• —Latent resolution: 192 × 192
• —Channels: 1 (grayscale CT slice)
• —Training objective: ε-prediction (standard DDPM formulation)
• —Noise schedule: Linear beta schedule
• —Training dataset: Synchrotron dataset of rocks (Synchrotron)
• —Intensity normalization: Rescaled to (-1, 1)

The diffusion model operates purely in latent space and relies on the autoencoder for encoding and decoding. This model is intended to be combined with data-consistency correction for CT reconstruction tasks.

📊 Dataset: Synchrotron

The model was trained on a real-world high-resolution CT dataset acquired at a high-energy synchrotron facility.

Source: https://zenodo.org/records/15420527

Preprocessing steps:

• —Train/test split
• —Rescale reconstructed slices to (-1, 1)
• —No geometry information is embedded in the model

🧠 Training Details

• —Optimizer: AdamW
• —Learning rate: 1e-4
• —Hardware: NVIDIA A100 GPU
• —Training script: train_latent.py

🚀 Usage

You can load and use this model using the diffusers library:

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained(
    "jiayangshi/synchrotron_latent_diffusion"
)
pipeline.to("cuda")

# Generate an unconditional sample from the CT prior
# Note: For reconstruction tasks, this model is typically used with 
# a custom solver incorporating CT data consistency.
output = pipeline()
image = output.images[0]
image.save("reconstruction_prior.png")

📝 Citation

@inproceedings{
shi2026dmct,
title={{DM}4{CT}: Benchmarking Diffusion Models for Computed Tomography Reconstruction},
author={Shi, Jiayang and Pelt, Dani{\in d}l M and Batenburg, K Joost},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=YE5scJekg5}
}