mahmoudibra98/compdiff-chest-xray

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 (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 (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

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

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.

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).