X-Diffusion is a state-of-the-art deep learning model for generating detailed 3D MRI volumes from single 2D images using cross-sectional diffusion models. This project implements the methods described in our paper, providing a robust framework for medical image synthesis.

• Generate complete 3D MRI volumes from single 2D slices
• Optimized inference with mixed precision support
• Support for BRATS and UK Biobank datasets
• Modular architecture for easy customization
• Comprehensive evaluation metrics (PSNR, SSIM)
• Pre-trained models available
• Docker support for easy deployment
• Documentation

• Python 3.8 or higher
• CUDA 11.0+ (for GPU support)
• 32GB RAM recommended
• 30GB VRAM for training (GPU)

1. Clone the repository:

git clone https://github.com/EmmanuelleB985/X-Diffusion.git
cd X-Diffusion

2. Create a virtual environment:

conda create -n XDiffusion python=3.9
conda activate XDiffusion

3. Install dependencies:

pip install -r requirements.txt

4. Install development dependencies (optional):

pip install -r requirements-dev.txt

5. Install pre-commit hooks (for contributors):

pre-commit install

6. Download pre-trained models:

# Zero-123 checkpoint
wget https://cv.cs.columbia.edu/zero123/assets/300000.ckpt -P Zero123/

# X-Diffusion checkpoint 
tar -xzvf ckpt.tar.gz

python scripts/main.py \
    -t \
    --base configs/sd-brats-finetune-c_concat-256.yaml \
    --gpus 0 \
    --scale_lr False \
    --num_nodes 1 \
    --seed 42 \
    --check_val_every_n_epoch 10 \
    --finetune_from Zero123/300000.ckpt

# Single image inference
python scripts/inference.py \
    --model_path checkpoints/best_model.ckpt \
    --input_image path/to/image.png \
    --output_dir results/

# Batch inference
python scripts/inference.py \
    --model_path checkpoints/best_model.ckpt \
    --input_dir data/test_images/ \
    --output_dir results/ \
    --batch_size 4

Run all checks:

# Format code
black scripts/ tests/

# Sort imports
isort scripts/ tests/

# Run linting
flake8 scripts/ tests/

# Type checking
mypy scripts/

# Run pylint
pylint scripts/

# Check docstring coverage
interrogate scripts/ -v --fail-under=90

Or use pre-commit to run all checks automatically:

pre-commit run --all-files

Run the test suite:

# Run all tests
pytest

# Run with coverage
pytest --cov=scripts --cov-report=html

# Run specific test categories
pytest -m unit           # Unit tests only
pytest -m integration    # Integration tests only
pytest -m "not slow"     # Skip slow tests

Build the documentation:

cd docs
make html
# Open docs/_build/html/index.html in your browser

Download the BRATS2023 dataset:

1. Create an account on Synapse
2. Download the dataset
3. Extract to data/BRATS2023/

Access UK Biobank data:

1. Register on the UK Biobank platform
2. Follow instructions from UKBiobankDXAMRIPreprocessing

X-Diffusion uses a cross-sectional diffusion approach based on:

• Backbone: Modified Stable Diffusion architecture
• Conditioning: Cross-sectional image embeddings
• Denoising: Progressive refinement through diffusion steps
• 3D Reconstruction: Volumetric assembly from 2D slices

Our model achieves state-of-the-art performance on:

If you use this code in your research, please cite our paper:

@misc{bourigault2025xdiffusion,
    title={X-Diffusion: Generating Detailed 3D MRI Volumes From a Single Image Using Cross-Sectional Diffusion Models}, 
    author={Emmanuelle Bourigault and Abdullah Hamdi and Amir Jamaludin},
    year={2025},
    eprint={2404.19604},
    archivePrefix={arXiv},
    primaryClass={eess.IV}
}

This project builds upon:

• Zero-123
• Stable Diffusion
• PyTorch Lightning

We thank the authors for making their code publicly available.

This project is licensed under the MIT License - see the LICENSE file for details.

We welcome contributions!

1. Fork the repository
2. Create your feature branch (git checkout -b feature/my-feature)
3. Commit your changes (git commit -m 'Add my feature')
4. Push to the branch (git push origin feature/my-feature)
5. Open a Pull Request