mahmoudibra98/compdiff-chest-xray

CompDiff — Chest X-Ray

Demographically-conditioned latent diffusion model for synthetic chest radiograph generation, from the CompDiff project. It is a fine-tune of stabilityai/stable-diffusion-2-1-base (UNet + CLIP text encoder trained) with a lightweight Hierarchical Conditioner Network (HCN): a typed compositional conditioner that takes the three demographic attributes (sex, race, age) and emits four demographic tokens that the UNet reads next to the clinical-findings text tokens.

• Paper: arXiv:2603.16551
• Code: https://github.com/mahmoudibrahim98/CompDiff
• Base model: Stable Diffusion 2.1-base (512×512)
• Demographic attributes: sex (2 classes), race (4 classes), age (continuous, years)
• Companion model: compdiff-fundus

Model versions. This is the second release of the chest model (September 2026), the checkpoint used in the current version of the paper. It conditions on all three attributes through the HCN, with age encoded continuously inside the conditioner. The first release (July 2026; sex × race through the HCN, age written into the prompt) remains available unchanged under the v1 revision: snapshot_download(..., revision="v1"). The two releases have different conditioner code and a different pipeline interface, so do not mix files across revisions.

What the conditioner does

Attributes enter the conditioner in their native type: sex and race as nominal embeddings, age as a continuous value mapped through sinusoidal features and an MLP (so nearby ages get nearby representations). A pairwise-MLP hierarchy composes them (age×sex, age×race, sex×race, then all three), each attribute is re-contextualised against the composed state, and four tokens are projected into the UNet cross-attention space:

[ t_age, t_sex, t_race, t_cls ]  →  concatenated to the 77 CLIP text tokens  →  UNet cross-attention

During training each token was supervised by an auxiliary head (sex, race, age regression, joint cell), so the tokens the UNet reads carry the attribute information. The text encoder only ever saw clinical findings: age, sex and race were stripped from every prompt.

Contents

model_index.json            # diffusers StableDiffusionPipeline index
unet/ text_encoder/ vae/    # fine-tuned SD-2.1 UNet + text encoder (fp32); vae is the frozen base
tokenizer/ scheduler/ feature_extractor/
hcn/                        # conditioner: config.json + model.safetensors (6.3M params)
compdiff2.py                # self-contained conditioner class (CompDiff2Conditioner)
compdiff_pipeline.py        # turnkey CompDiffPipeline (demographic-conditioned generation)
training_config.yaml        # full training configuration of the released run

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 generation loop used for the paper's evaluation cohorts (classifier-free guidance 7.5, DDPM sampling, 75 steps, 512×512). 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")

All three attributes are required and go through the conditioner. Put only clinical findings in prompt; do not write age, sex or race into it (the pipeline warns if it sees demographic words). age is a number of years, not a bin. Each of prompt, sex, race, age may also be a list of length num_images to generate a mixed batch, and seed= makes a call reproducible.

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

Training prompts had the form "<AGE> year old <RACE> <SEX>. <IMPRESSION>"; CompDiff strips the demographic clause before the text encoder and routes the three attributes through the HCN instead. The text encoder therefore only ever saw the impression / findings text, for example:

"Cardiomegaly with small bilateral pleural effusions."
"No acute cardiopulmonary process."
"Right mid lung rounded opacity may represent a new mass or infection. Recommend CT for further evaluation."

An empty prompt is replaced by "Normal chest radiograph", the same fallback used in training.

Using the conditioner directly

from compdiff2 import CompDiff2Conditioner   # after sys.path.insert(0, path)
hcn = CompDiff2Conditioner.from_pretrained(f"{path}/hcn", device="cuda")
ctx, mu, logsigma, aux, _ = hcn(sex_idx=torch.tensor([1]).cuda(),
                                race_idx=torch.tensor([0]).cuda(),
                                age_continuous=torch.tensor([67.0]).cuda())
ctx.shape   # torch.Size([1, 4, 1024]) -> concatenate to the CLIP hidden states (dim=1)

For classifier-free guidance the unconditional branch is the empty prompt with four zero tokens in place of the demographic tokens, as in compdiff_pipeline.py.

Advanced: plain Stable Diffusion backbone

Loading the pipeline with standard diffusers gives the fine-tuned SD-2.1 backbone without demographic conditioning (the conditioner 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]

Note that this backbone was trained with the demographic tokens always present, so sampling it without them is out of distribution; use CompDiffPipeline for real use.

Training and evaluation summary

• Data: MIMIC-CXR postero-anterior views with complete (age, sex, race) metadata; folders p10–p18 for training (62,094 images) and validation (1,300), folder p19 held out as the test split (7,039 images).
• Recipe: SD 2.1-base initialisation, UNet + text encoder + conditioner trained jointly, 512×512, effective batch 48, learning rate 1e-5, bf16 mixed precision. The released weights are the step-10,000 checkpoint, selected on the validation split.
• Inference used for the numbers below: DDPM, 75 steps, guidance 7.5, one image per test-split prompt with the test-split demographics.

Metrics of this checkpoint on the MIMIC-CXR test split (7,039 generated images, one generation seed), computed with the project's evaluation code:

Metric	Value
FID (Inception)	58.7
FID (RadImageNet features)	5.78
Sex accuracy (XRV classifier)	0.999
Race accuracy (XRV classifier)	0.955
Age RMSE, years (XRV regressor)	8.44
Mean disease AUROC (XRV DenseNet-121)	0.823

The paper reports means and standard deviations across three independently trained runs on the validation split, which are lower-fidelity than this single test-split row; see the paper for the like-for-like comparison against RoentGen-v2 and FairDiffusion, the zero-shot held-out-intersection results, and the downstream classifier experiments.

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 behaviour is limited to the attribute categories and the adult age range the model was trained on (MIMIC-CXR patients, 18 years and older). Race labels follow the MIMIC-CXR self-reported categories; they are a conditioning signal, not a biological ground truth.
• The conditioner was trained with all three attributes always present; partial conditioning (leaving an attribute unspecified) is not supported by this checkpoint.

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