""" CompDiffPipeline — turnkey inference for CompDiff medical image generators. CompDiff = Stable-Diffusion-2.1-base (fine-tuned UNet + CLIP text encoder) plus a Hierarchical Conditioner Network (HCN) that injects demographic attributes. For the released checkpoints, sex and race are conditioned through the HCN (a token appended to the text-encoder embeddings), while age is conditioned through the text prompt. This wrapper reproduces the exact generation logic used to produce the paper's synthetic datasets (classifier-free guidance, DDPM sampling, HCN token concat). Usage ----- from compdiff_pipeline import CompDiffPipeline pipe = CompDiffPipeline.from_pretrained("mahmoudibra98/compdiff-chest-xray", device="cuda") img = pipe.generate("Cardiomegaly with small bilateral pleural effusions", sex="female", race="White", age=67)[0] img.save("out.png") `sex` and `race` accept either an integer index (always safe — matches the training label set) or a string (mapped with the chest-X-ray convention below). Check `pipe.num_race` and the paper for the exact label set of a given model. Chest-X-ray index convention: sex : 0=male, 1=female race: 0=White, 1=Black/African American, 2=Asian, 3=Hispanic/Latino """ import os import sys import importlib.util from typing import List, Optional, Union import torch from PIL import Image from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel from transformers import CLIPTextModel, CLIPTokenizer _SEX_STR = {"male": 0, "man": 0, "m": 0, "female": 1, "woman": 1, "f": 1} def _parse_sex(v) -> int: if isinstance(v, int): return v return _SEX_STR[str(v).strip().lower()] def _parse_race(v) -> int: if isinstance(v, int): return v s = str(v).upper() if "BLACK" in s or "AFRICAN" in s: return 1 if "ASIAN" in s: return 2 if "HISPANIC" in s or "LATINO" in s: return 3 return 0 # White / default def _load_hcn_class(local_dir: str): """Import HierarchicalConditionerV8 from the bundled hcn_v7.py next to this file.""" hcn_path = os.path.join(local_dir, "hcn_v7.py") spec = importlib.util.spec_from_file_location("compdiff_hcn_v7", hcn_path) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod.HierarchicalConditionerV8 class CompDiffPipeline: def __init__(self, tokenizer, text_encoder, vae, unet, scheduler, hcn, device="cuda", dtype=torch.float16): self.tokenizer = tokenizer self.text_encoder = text_encoder.to(device, dtype).eval() self.vae = vae.to(device, torch.float32).eval() # VAE kept in fp32 self.unet = unet.to(device, dtype).eval() self.scheduler = scheduler self.hcn = hcn.to(device, dtype).eval() self.device = device self.dtype = dtype self.num_race = getattr(hcn, "num_race", None) self.num_sex = getattr(hcn, "num_sex", None) @classmethod def from_pretrained(cls, model_id_or_path: str, device: str = "cuda", dtype: torch.dtype = torch.float16, **snapshot_kwargs): if os.path.isdir(model_id_or_path): local_dir = model_id_or_path else: from huggingface_hub import snapshot_download local_dir = snapshot_download(model_id_or_path, **snapshot_kwargs) tokenizer = CLIPTokenizer.from_pretrained(local_dir, subfolder="tokenizer") text_encoder = CLIPTextModel.from_pretrained(local_dir, subfolder="text_encoder") vae = AutoencoderKL.from_pretrained(local_dir, subfolder="vae") unet = UNet2DConditionModel.from_pretrained(local_dir, subfolder="unet") scheduler = DDPMScheduler.from_pretrained(local_dir, subfolder="scheduler") HCN = _load_hcn_class(local_dir) hcn = HCN.from_pretrained(os.path.join(local_dir, "hcn")) return cls(tokenizer, text_encoder, vae, unet, scheduler, hcn, device=device, dtype=dtype) def _encode_text(self, prompts: List[str]) -> torch.Tensor: inp = self.tokenizer(prompts, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt").input_ids.to(self.device) return self.text_encoder(input_ids=inp, return_dict=False)[0] @torch.no_grad() def generate( self, prompt: str, sex: Union[int, str], race: Union[int, str], age: Optional[int] = None, num_images: int = 1, num_inference_steps: int = 75, guidance_scale: float = 7.5, negative_prompt: str = "", resolution: int = 512, seed: Optional[int] = None, ) -> List[Image.Image]: """Generate demographically-conditioned images. prompt : clinical findings text (do NOT include sex/race — they go through the HCN). sex, race : int index or string. age : optional age in years; prepended to the prompt as " year old. ..." (age is conditioned through text for the released checkpoints). """ sex_idx = _parse_sex(sex) race_idx = _parse_race(race) text = f"{int(age)} year old. {prompt}" if age is not None else prompt B = num_images text_embeddings = self._encode_text([text] * B) # [B, 77, D] sex_t = torch.full((B,), sex_idx, dtype=torch.long, device=self.device) race_t = torch.full((B,), race_idx, dtype=torch.long, device=self.device) # encode_age=False for the released HCN -> age_idx is not passed to the HCN age_arg = None if not getattr(self.hcn, "encode_age", False) else \ torch.zeros((B,), dtype=torch.long, device=self.device) ctx = self.hcn(sex_idx=sex_t, race_idx=race_t, age_idx=age_arg)[0] # [B, 1, D] text_embeddings = torch.cat([text_embeddings, ctx.to(text_embeddings.dtype)], dim=1) uncond = self._encode_text([negative_prompt] * B) # [B, 77, D] zero_ctx = torch.zeros((B, 1, uncond.shape[-1]), device=self.device, dtype=uncond.dtype) uncond = torch.cat([uncond, zero_ctx], dim=1) # [B, 78, D] encoder_hidden_states = torch.cat([uncond, text_embeddings], dim=0) lat_shape = (B, self.unet.config.in_channels, resolution // 8, resolution // 8) gen = None if seed is None else torch.Generator(device=self.device).manual_seed(seed) latents = torch.randn(lat_shape, generator=gen, device=self.device, dtype=self.dtype) latents = latents * getattr(self.scheduler, "init_noise_sigma", 1.0) self.scheduler.set_timesteps(num_inference_steps) # Keep `t` on the scheduler's device (CPU) for scheduler.step; only the # UNet call needs the timestep on the compute device. for t in self.scheduler.timesteps: t_dev = t.to(self.device) model_in = self.scheduler.scale_model_input(torch.cat([latents] * 2), t_dev) ts = t_dev.expand(model_in.shape[0]) if t_dev.ndim == 0 else t_dev.repeat(model_in.shape[0]) noise_pred = self.unet(model_in, ts, encoder_hidden_states=encoder_hidden_states).sample nu, nt = noise_pred.chunk(2) noise_pred = nu + guidance_scale * (nt - nu) latents = self.scheduler.step(noise_pred, t, latents).prev_sample latents = (1 / 0.18215) * latents.to(torch.float32) images = self.vae.decode(latents).sample images = (images / 2 + 0.5).clamp(0, 1).cpu().permute(0, 2, 3, 1).numpy() return [Image.fromarray((im * 255).round().astype("uint8")) for im in images] __call__ = generate