""" 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-fundus", device="cuda") img = pipe.generate("glaucoma, severe vision loss, abnormal cup-disc ratio, myopia", 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 convention below). Check `pipe.num_race` and the paper for the exact label set of a given model. Fundus index convention: sex : 0=male, 1=female race: 0=White, 1=Black/African American, 2=Asian (3 classes; no 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 " years old. ..." (age is conditioned through text for the released checkpoints). """ sex_idx = _parse_sex(sex) race_idx = _parse_race(race) text = f"{int(age)} years 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