Medical Image VAE Models (Brain MRI)

This repository contains Variational Autoencoder (VAE) models specifically adapted and fine-tuned for medical imaging tasks, particularly for brain MRI. Two different VAE models with different compression factors are provided (klf8 and klf4). Models are trained with KL loss, LPIPS, and MSE loss.

These VAE models are crucial components utilized in advanced brain anomaly detection frameworks. Notably, they are used by:

REFLECT: Rectified Flows for Efficient Brain Anomaly Correction Transport. This work is described in the paper REFLECT: Rectified Flows for Efficient Brain Anomaly Correction Transport. The code for the REFLECT framework can be found at https://github.com/farzad-bz/REFLECT.
MAD-AD: Masked Diffusion for Unsupervised Brain Anomaly Detection. This model is described in the paper MAD-AD: Masked Diffusion for Unsupervised Brain Anomaly Detection. The code for the MAD-AD model can be found at https://github.com/farzad-bz/MAD-AD.

Original Model Source

These models are adapted and fine-tuned based on the latent diffusion VAE architecture originally developed by CompVis:

CompVis Latent Diffusion GitHub Repository

Usage

Here's a simple example of how you can load and use the pretrained VAE models. Please note that scaling with 0.18215 is devised to encourage standard normal distribution for latent distribution.

from huggingface_hub import hf_hub_download
import torch

# Download model
model_path = hf_hub_download(repo_id="farzadbz/Medical-VAE", filename="VAE-Medical-klf8.pt")

# Load the model
model = torch.load(model_path)
model.eval()

# Run inference
encoded = model.encode(x).sample().mul_(0.18215)
reconstruction = model.decode(encoded / 0.18215)

Replace "VAE-Medical-klf8.pt" with the desired model filename (VAE-Medical-klf8.pt or VAE-Medical-klf4.pt).

Citation & Reference

If you use these VAE models in your research, please cite the works that utilize them. For REFLECT, cite:

@article{beizaee2025reflect,
  title={REFLECT: Rectified Flows for Efficient Brain Anomaly Correction Transport},
  author={Beizaee, Farzad and Hajimiri, Sina and Ben Ayed, Ismail and Lodygensky, Gregory and Desrosiers, Christian and Dolz, Jose},
  journal={arXiv preprint arXiv:2508.02889},
  year={2025}
}

For MAD-AD, cite:

@article{beizaee2025mad,
  title={MAD-AD: Masked Diffusion for Unsupervised Brain Anomaly Detection},
  author={Beizaee, Farzad and Lodygensky, Gregory and Desrosiers, Christian and Dolz, Jose},
  journal={arXiv preprint arXiv:2502.16943},
  year={2025}
}