ProMoE-B-256/pipeline.py

"""Hub custom pipeline: ProMoEPipeline.
Load with native Hugging Face diffusers and trust_remote_code=True.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import torch
from PIL import Image

try:
    from diffusers.pipelines.pipeline_utils import DiffusionPipeline
except Exception:  # pragma: no cover
    class DiffusionPipeline:
        def __init__(self):
            self._execution_device = torch.device("cpu")

        def register_modules(self, **kwargs):
            for key, value in kwargs.items():
                setattr(self, key, value)

        def to(self, device):
            self._execution_device = torch.device(device)
            for module in (getattr(self, "transformer", None), getattr(self, "vae", None)):
                if module is not None and hasattr(module, "to"):
                    module.to(device)
            return self

        def progress_bar(self, iterable):
            return iterable

        def maybe_free_model_hooks(self):
            return None

@dataclass
class ProMoEPipelineOutput:
    images: Union[List[Image.Image], np.ndarray, torch.Tensor]

class ProMoEPipeline(DiffusionPipeline):
    r"""
    Pipeline for class-conditional image generation with ProMoE.

    Parameters:
        transformer ([`ProMoETransformer2DModel`]):
            Class-conditional ProMoE transformer for flow-matching in latent space.
        scheduler ([`ProMoEFlowMatchScheduler`]):
            Flow-matching scheduler used during denoising.
        vae ([`AutoencoderKL`], *optional*):
            Variational autoencoder used to decode latents to pixels.
        id2label (`dict[int, str]`, *optional*):
            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
    """

    model_cpu_offload_seq = "transformer->vae"
    _optional_components = ["vae"]

    def __init__(
        self,
        transformer,
        scheduler,
        vae=None,
        id2label: Optional[Dict[Union[int, str], str]] = None,
    ):
        super().__init__()
        self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
        self._id2label = self._normalize_id2label(id2label)
        self.labels = self._build_label2id(self._id2label)
        self._labels_loaded_from_model_index = bool(self._id2label)

    def _ensure_labels_loaded(self) -> None:
        if self._labels_loaded_from_model_index:
            return
        loaded = self._read_id2label_from_model_index(getattr(self.config, "_name_or_path", None))
        if loaded:
            self._id2label = loaded
            self.labels = self._build_label2id(self._id2label)
        self._labels_loaded_from_model_index = True

    @staticmethod
    def _normalize_id2label(id2label: Optional[Dict[Union[int, str], str]]) -> Dict[int, str]:
        if not id2label:
            return {}
        return {int(key): value for key, value in id2label.items()}

    @staticmethod
    def _read_id2label_from_model_index(variant_path: Optional[str]) -> Dict[int, str]:
        if not variant_path:
            return {}
        variant_dir = Path(variant_path).resolve()
        model_index_path = variant_dir / "model_index.json"
        if not model_index_path.exists():
            return {}
        raw = json.loads(model_index_path.read_text(encoding="utf-8"))
        id2label = raw.get("id2label")
        if not isinstance(id2label, dict):
            return {}
        return {int(key): value for key, value in id2label.items()}

    @staticmethod
    def _build_label2id(id2label: Dict[int, str]) -> Dict[str, int]:
        label2id: Dict[str, int] = {}
        for class_id, value in id2label.items():
            for synonym in value.split(","):
                synonym = synonym.strip()
                if synonym:
                    label2id[synonym] = int(class_id)
        return dict(sorted(label2id.items()))

    @property
    def id2label(self) -> Dict[int, str]:
        r"""ImageNet class id to English label string (comma-separated synonyms)."""
        self._ensure_labels_loaded()
        return self._id2label

    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
        r"""
        Map ImageNet label strings to class ids.

        Args:
            label (`str` or `list[str]`):
                One or more English label strings. Each string must match a synonym in `id2label`.
        """
        self._ensure_labels_loaded()
        label2id = self.labels
        if not label2id:
            raise ValueError("No English labels loaded. Ensure `id2label` exists in model_index.json.")

        if isinstance(label, str):
            label = [label]

        missing = [item for item in label if item not in label2id]
        if missing:
            preview = ", ".join(list(label2id.keys())[:8])
            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
        return [label2id[item] for item in label]

    def _get_vae_spatial_downsample(self) -> int:
        if self.vae is None:
            return 8
        block_out_channels = getattr(getattr(self.vae, "config", None), "block_out_channels", [0, 0, 0, 0])
        return 2 ** (len(block_out_channels) - 1)

    def _normalize_class_labels(
        self,
        class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor],
        device: torch.device,
    ) -> torch.LongTensor:
        if torch.is_tensor(class_labels):
            return class_labels.to(device=device, dtype=torch.long).reshape(-1)

        if isinstance(class_labels, int):
            class_label_ids = [class_labels]
        elif isinstance(class_labels, str):
            class_label_ids = self.get_label_ids(class_labels)
        elif class_labels and isinstance(class_labels[0], str):
            class_label_ids = self.get_label_ids(class_labels)
        else:
            class_label_ids = list(class_labels)

        return torch.tensor(class_label_ids, device=device, dtype=torch.long).reshape(-1)

    def _prepare_latents(
        self,
        batch_size: int,
        latent_height: int,
        latent_width: int,
        dtype: torch.dtype,
        device: torch.device,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
    ) -> torch.Tensor:
        shape = (batch_size, self.transformer.in_channels, latent_height, latent_width)
        if isinstance(generator, list):
            latents = [torch.randn((1, *shape[1:]), generator=g, device=device, dtype=dtype) for g in generator]
            return torch.cat(latents, dim=0)
        return torch.randn(shape, generator=generator, device=device, dtype=dtype)

    def _decode_latents(self, latents: torch.Tensor, output_type: str):
        if output_type == "latent":
            return latents
        if self.vae is not None:
            scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215)
            decode_dtype = next(self.vae.parameters()).dtype
            latents = (latents / scaling_factor).to(dtype=decode_dtype)
            image = self.vae.decode(latents, return_dict=False)[0]
        else:
            image = latents

        image = (image / 2 + 0.5).clamp(0, 1)
        if output_type == "pt":
            return image
        image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
        if output_type == "np":
            return image
        pil_images = [Image.fromarray((img * 255).round().astype("uint8")) for img in image]
        return pil_images

    @torch.no_grad()
    def __call__(
        self,
        class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor],
        height: int = 256,
        width: int = 256,
        num_inference_steps: int = 50,
        guidance_scale: float = 1.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        output_type: str = "pil",
        return_dict: bool = True,
    ) -> Union[ProMoEPipelineOutput, Tuple]:
        r"""
        Generate class-conditional images with ProMoE.

        Args:
            class_labels (`int`, `str`, `list[int]`, `list[str]`, or `torch.LongTensor`):
                ImageNet class indices or human-readable English label strings.
        """
        device = self._execution_device if hasattr(self, "_execution_device") else torch.device("cpu")
        model_dtype = next(self.transformer.parameters()).dtype
        class_labels = self._normalize_class_labels(class_labels, device)
        batch_size = class_labels.shape[0]

        vae_scale = self._get_vae_spatial_downsample()
        latent_height = height // vae_scale
        latent_width = width // vae_scale
        latents = self._prepare_latents(batch_size, latent_height, latent_width, model_dtype, device, generator)

        self.scheduler.set_timesteps(num_inference_steps, device=device)
        null_labels = torch.full_like(class_labels, getattr(self.transformer.backbone.y_embedder, "num_classes", 1000))

        for t in self.progress_bar(self.scheduler.timesteps):
            if guidance_scale > 1.0:
                latent_input = torch.cat([latents, latents], dim=0)
                labels = torch.cat([class_labels, null_labels], dim=0)
            else:
                latent_input = latents
                labels = class_labels
            timestep = torch.full((labels.shape[0],), t, device=device, dtype=model_dtype)
            model_output = self.transformer(
                hidden_states=latent_input,
                timestep=timestep,
                class_labels=labels,
                return_dict=True,
            ).sample
            if model_output.shape[1] != latents.shape[1]:
                model_output = model_output.chunk(2, dim=1)[0]
            if guidance_scale > 1.0:
                model_output_cond, model_output_uncond = model_output.chunk(2)
                model_output = model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond)
            latents = self.scheduler.step(model_output, t, latents, generator=generator).prev_sample

        images = self._decode_latents(latents, output_type)
        self.maybe_free_model_hooks()
        if not return_dict:
            return (images,)
        return ProMoEPipelineOutput(images=images)