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