README.md

---
license: openrail++
base_model: stabilityai/stable-diffusion-2-1-base
library_name: diffusers
pipeline_tag: text-to-image
tags:
- text-to-image
- diffusers
- stable-diffusion
- medical-imaging
- fundus
- compdiff
---

# 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`](https://huggingface.co/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](https://arxiv.org/abs/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

```bash
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](https://pytorch.org/get-started/locally/). Tested combo (A100, driver
550.x / CUDA 12.4):

```bash
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:

```python
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**](https://huggingface.co/datasets/harvardairobotics/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):

```python
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](https://github.com/mahmoudibrahim98/CompDiff).

## 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:

```bibtex
@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).