mahmoudibra98/compdiff-fundus

CompDiff — Retinal Fundus (Fundus)

Demographically-conditioned latent diffusion model for synthetic retinal fundus images generation, from the CompDiff project. It is a fine-tune of stabilityai/stable-diffusion-2-1-base (UNet + CLIP text encoder trained) augmented with a lightweight Hierarchical Conditioner Network (HCN) that injects demographic attributes (sex, race, age) into the generation.

• Paper: arXiv:2603.16551
• Code: https://github.com/mahmoudibrahim98/CompDiff
• Base model: Stable Diffusion 2.1-base (512×512)
• Demographic attributes: sex (2), race (3), age (5 bins)

Contents

model_index.json           # diffusers StableDiffusionPipeline index
unet/ text_encoder/ vae/    # fine-tuned SD-2.1 components (vae is the frozen base)
tokenizer/ scheduler/ feature_extractor/
hcn/                        # HCN module: config.json + pytorch_model.bin
hcn_v7.py                   # self-contained HCN class (HierarchicalConditionerV8)
compdiff_pipeline.py        # turnkey CompDiffPipeline (demographic-conditioned generation)
training_config.yaml        # full training configuration

Requirements

pip install "diffusers>=0.35" transformers accelerate huggingface_hub safetensors pillow

Install torch to match your CUDA driver — a bare pip install torch may pull a build newer than your driver supports (e.g. a cu130 wheel on a CUDA 12.4 driver fails with "NVIDIA driver too old" / cuda available: False). Pick the wheel for your CUDA version from pytorch.org. Tested combo (A100, driver 550.x / CUDA 12.4):

pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124

Quickstart (demographic-conditioned)

The bundled compdiff_pipeline.py reproduces the exact generation used in the paper — it appends the HCN demographic token to the text embeddings and runs classifier-free guided DDPM sampling. This is the recommended entry point:

import torch
from huggingface_hub import snapshot_download

path = snapshot_download("mahmoudibra98/compdiff-fundus")
import sys; sys.path.insert(0, path)
from compdiff_pipeline import CompDiffPipeline

pipe = CompDiffPipeline.from_pretrained(path, device="cuda", dtype=torch.float16)
img = pipe.generate("glaucoma, severe vision loss, abnormal cup-disc ratio, myopia",
                    sex="female", race=0, age=67)[0]
img.save("out.png")

For the released checkpoints, sex and race are conditioned through the HCN, while age is conditioned through the prompt (pass age= and it is prepended as "<age> years old. ..."). Put only clinical findings in prompt — not sex/race.

Index convention:

sex : 0 = male, 1 = female
race: 0 = White, 1 = Black/African American, 2 = Asian   (this model uses 3 race classes)

String parsing works for these 3 classes (race="White", "Black", "Asian"), but note that unlike the chest model there is no Hispanic/Latino class here — index 3 is out of range. Pass integer indices if in doubt.

Prompt format

The model was trained with demographics stripped from the text (they enter through the HCN), so the text encoder only ever saw age + clinical findings. The effective template the encoder sees is:

"<age> years old. <clinical findings>"

You only pass the <clinical findings> in prompt; compdiff_pipeline.py prepends the age string for you when you pass age=. Omit age= to drop the age clause entirely. Do not put sex/race in the prompt — they are conditioned by the HCN.

Clinical findings vocabulary

Unlike the chest model, the fundus clinical findings are not free text — they are a fixed, comma-joined list of glaucoma-domain descriptors drawn from the Harvard-Ophthalmology FairGenMed label set (as used in the FairDiffusion work). For in-distribution results, build prompt as ", ".join(...) of the following slots, in this order:

Slot	Values
Glaucoma status	glaucoma / non-glaucoma
Vision loss (mean-deviation severity)	normal vision, or <severity> vision loss (e.g. mild vision loss, moderate vision loss, severe vision loss)
Cup-to-disc ratio (optional)	normal cup-disc ratio / borderline cup-disc ratio / abnormal cup-disc ratio
Refraction (optional)	hyperopia / emmetropia / myopia

Examples:

non-glaucoma, normal vision, normal cup-disc ratio, emmetropia
glaucoma, severe vision loss, abnormal cup-disc ratio, myopia

Free-form radiology-style prompts (e.g. "diabetic retinopathy with microaneurysms") are out of distribution for this model and will give unreliable results.

Advanced: plain Stable Diffusion backbone

Loading the pipeline with standard diffusers gives the fine-tuned SD-2.1 backbone without demographic conditioning (the HCN is not part of the diffusers pipeline):

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(path, dtype=torch.float16, safety_checker=None).to("cuda")
image = pipe("a retinal fundus image", num_inference_steps=75, guidance_scale=7.5).images[0]

To wire the HCN in manually, see compdiff_pipeline.py or generate_synthetic_dataset.py in the CompDiff repository.

Intended use & limitations

• Research use only. This is a generative model for studying demographic fairness of synthetic medical images. It is not a medical device and must not be used for diagnosis, screening, or any clinical decision-making.
• Synthetic images may contain artifacts and may not faithfully represent real pathology.
• Demographic behavior is limited to the attribute categories the model was trained on.

Citation

If you use this model, please cite:

@article{ibrahim2026compdiff,
  title   = {CompDiff: Hierarchical Compositional Diffusion for Fair and Zero-Shot Intersectional Medical Image Generation},
  author  = {Ibrahim, Mahmoud and Elen, Bart and Sun, Chang and Ertaylan, Gokhan and Dumontier, Michel},
  journal = {arXiv preprint arXiv:2603.16551},
  year    = {2026},
  url     = {https://arxiv.org/abs/2603.16551}
}

License

Model weights are released under the CreativeML OpenRAIL++-M license inherited from Stable Diffusion 2.1-base. Project code is MIT-licensed (see the CompDiff repository).