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
- chest-xray
- compdiff
---

# CompDiff — Chest X-Ray (RoentGen-v2)

Demographically-conditioned latent diffusion model for synthetic **chest radiographs** 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 (4), 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-chest-xray")
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("Cardiomegaly with small bilateral pleural effusions",
                    sex="female", race="White", 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> year old. ..."`). Put only clinical findings in `prompt` — not sex/race.

Index convention (chest X-ray):
```
sex : 0 = male, 1 = female
race: 0 = White, 1 = Black/African American, 2 = Asian, 3 = Hispanic/Latino
```
`sex`/`race` accept an integer index (always safe) or a string (mapped with the convention above).

### 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> year old. <clinical findings>"
```

You only pass the `<clinical findings>` in `prompt`; `compdiff_pipeline.py` prepends the age
string for you when you pass `age=`. For example, `generate("Cardiomegaly ...", age=67)` is
encoded as `"67 year old. Cardiomegaly ..."`. Omit `age=` to drop the age clause entirely.
Do **not** put sex/race in the prompt — they are conditioned by the HCN.

## 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 chest radiograph", 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).