LEAF is a medical image segmentation latent diffusion model. It replaces the conventional noise-prediction objective with a direct prediction of the segmentation map during fine-tuning, which reduces variance in segmentation outputs. The model also applies a feature-distillation technique to align convolutional hidden states with features produced by a transformer-based vision encoder.

Model Architecture

Diffusers Usage

Note: To use this model with diffusers, replace the file diffusers/models/unets/unet_2d.py with the unet_2d.py provided in this repository.

python

import torch
from pipeline import LeafPipeline
from PIL import Image

device = torch.device("cuda")

# load pipeline
pipeline = LeafPipeline.from_pretrained("pearisli/LEAF-QaTa-COV19").to(device)

# open image
image = Image.open("image.png").convert("RGB")

# run pipeline in inference
mask = pipeline(image).prediction[0]

# save segmentation mask
mask.save("mask.png")

Citation

We kindly encourage citation of our work if you find it useful.

bibtex

@InProceedings{HuaQil_LEAF_MICCAI2025,
        author = { Huang, Qilin AND Lin, Tianyu AND Chen, Zhiguang AND Zheng, Fudan},
        title = { { LEAF: Latent Diffusion with Efficient Encoder Distillation for Aligned Features in Medical Image Segmentation } },
        booktitle = {proceedings of Medical Image Computing and Computer Assisted Intervention -- MICCAI 2025},
        year = {2025},
        publisher = {Springer Nature Switzerland},
        volume = {LNCS 15965},
        month = {September},
        page = {384 -- 393}
}

bibtex

@misc{huang2025leaflatentdiffusionefficient,
      title={LEAF: Latent Diffusion with Efficient Encoder Distillation for Aligned Features in Medical Image Segmentation}, 
      author={Qilin Huang and Tianyu Lin and Zhiguang Chen and Fudan Zheng},
      year={2025},
      eprint={2507.18214},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2507.18214}, 
}