File size: 7,760 Bytes
fb0dbae 2bdc504 4087b82 fb0dbae 2bdc504 fb0dbae 2bdc504 fb0dbae 2bdc504 fb0dbae 2bdc504 fb0dbae 2bdc504 fb0dbae a521d20 fb0dbae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """
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 "<age> 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
|