license: mit tags: - pytorch - diffusers - unconditional-image-generation - diffusion-models-class - medical-imaging - brain-mri - multiple-sclerosis

Brain MRI Synthesis with Latent Diffusion (from scratch)

This model is a diffusion-based model for unconditional image generation of latent representations of brain MRI FLAIR slices. The model is designed to synthesize high-resolution brain MRI images (256x256 pixels) through a Latent Diffusion process, leveraging a U-Net architecture with ResNet and Attention-based blocks.

Training Details

Architecture: Latent Diffusion Model (LDM)
Resolution: Latent resolution of 32x32 to generate 256x256 final images
Dataset: Lesion2D VH split (FLAIR MRI slices) (70% of the dataset)
Channels: 4 (latents are multi-channel representations of the original images)
Epochs: 100
Batch size: 16
Optimizer: AdamW with:
- Learning Rate: 1.0e-4
- Betas: (0.95, 0.999)
- Weight Decay: 1.0e-6
- Epsilon: 1.0e-8
Scheduler: Cosine with 500 warm-up steps
Gradient Accumulation: 1 step
Mixed Precision: No
Gradient Clipping: Max norm of 1.0
Noise Scheduler: Linear schedule with:
- num_train_timesteps: 1000
- beta_start: 0.0001
- beta_end: 0.02
Hardware: Trained on NVIDIA GPUs with a distributed dataloader using 12 workers.
Memory Consumption: Approx. 2.5 GB during training.

U-Net Architecture

Down Blocks: [DownBlock2D, DownBlock2D, DownBlock2D, DownBlock2D, AttnDownBlock2D, DownBlock2D]
Up Blocks: [UpBlock2D, AttnUpBlock2D, UpBlock2D, UpBlock2D, UpBlock2D, UpBlock2D]
Layers per Block: 2
Block Channels: [128, 128, 256, 256, 512, 512]

The model is designed to learn a compressed representation of the brain MRI images at a latent level, making the synthesis process more memory-efficient while maintaining high fidelity.

Usage

You can use the model directly with the diffusers library:

from diffusers import LatentDiffusionPipeline
import torch

# Load the model
pipeline = LatentDiffusionPipeline.from_pretrained("benetraco/latent_scratch")
pipeline.to("cuda")  # or "cpu"

# Generate an image
image = pipeline(batch_size=1).images[0]

# Display the image
image.show()