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 model with different compression factor are provided (klf8, and klf4). Models are trained with KL loss, LIPIPs, and MSE loss.

This model is used as part of the MAD-AD model, 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 not 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-klf4.pt")

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

# Run inference
encoded = vae.encode(x).sample().mul_(0.18215)
recontruction = x0 = vae.decode(encoded / 0.18215)

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

Citation & Reference

If you use this model in your research, please cite this work.

@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}
}