compdiff_pipeline.py

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