AFM-M-2-1NFE-guided/pipeline.py

"""Hub custom pipeline: AFMPipeline.

Load with native Hugging Face diffusers and trust_remote_code=True.
"""

from __future__ import annotations

import importlib.util
import json
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import torch
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.utils.torch_utils import randn_tensor


class AFMPipeline(DiffusionPipeline):
    model_cpu_offload_seq = "generator->vae"

    @staticmethod
    def _coerce_scheduler(scheduler, generator):
        if scheduler is not None and not isinstance(scheduler, (list, tuple)):
            return scheduler
        variant_path = getattr(generator.config, "_name_or_path", None)
        if variant_path:
            scheduler_dir = Path(variant_path).resolve().parent / "scheduler"
            module_path = scheduler_dir / "scheduling_continuous_flow.py"
            config_path = scheduler_dir / "scheduler_config.json"
            if module_path.is_file() and config_path.is_file():
                spec = importlib.util.spec_from_file_location("scheduling_continuous_flow", module_path)
                if spec is not None and spec.loader is not None:
                    module = importlib.util.module_from_spec(spec)
                    spec.loader.exec_module(module)
                    return module.ContinuousFlowMatchScheduler.from_pretrained(str(scheduler_dir))
        raise ValueError("AFMPipeline could not load ContinuousFlowMatchScheduler from the variant folder.")

    def __init__(
        self,
        generator,
        vae,
        scheduler=None,
        pred_type: str = "x",
        id2label: Optional[Dict[Union[int, str], str]] = None,
    ) -> None:
        super().__init__()
        scheduler = self._coerce_scheduler(scheduler, generator)
        self.register_modules(generator=generator, vae=vae, scheduler=scheduler)
        self.register_to_config(pred_type=pred_type)
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
        self._id2label = self._normalize_id2label(id2label)
        self.labels = self._build_label2id(self._id2label)
        self._labels_loaded_from_model_index = bool(self._id2label)

    @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 {}
        model_index_path = Path(variant_path).resolve() / "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 vae_scale_factor(self) -> int:
        block_out_channels = getattr(self.vae.config, "block_out_channels", None)
        if block_out_channels:
            return int(2 ** (len(block_out_channels) - 1))
        return 8

    @property
    def id2label(self) -> Dict[int, str]:
        self._ensure_labels_loaded()
        return 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

    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
        self._ensure_labels_loaded()
        labels = [label] if isinstance(label, str) else label
        if not self.labels:
            raise ValueError("No id2label mapping is available in this checkpoint.")
        missing = [item for item in labels if item not in self.labels]
        if missing:
            preview = ", ".join(list(self.labels.keys())[:8])
            raise ValueError(f"Unknown English label(s): {missing}. Example valid labels: {preview}, ...")
        return [self.labels[item] for item in labels]

    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 _default_image_size(self) -> int:
        transformer = getattr(self.generator, "transformer", None)
        if transformer is not None:
            return int(transformer.config.sample_size) * self.vae_scale_factor
        return 256

    def check_inputs(
        self,
        height: int,
        width: int,
        num_inference_steps: int,
        output_type: str,
        sampler: str,
    ) -> None:
        if num_inference_steps < 1:
            raise ValueError("`num_inference_steps` must be >= 1.")
        if output_type not in {"pil", "np", "pt", "latent"}:
            raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.")
        if sampler not in {"euler", "heun"}:
            raise ValueError("sampler must be one of: 'euler', 'heun'.")
        if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
            raise ValueError(
                f"`height` and `width` must be divisible by {self.vae_scale_factor}, got ({height}, {width})."
            )

    def prepare_latents(
        self,
        batch_size: int,
        height: int,
        width: int,
        dtype: torch.dtype,
        device: torch.device,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
        latents: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if latents is not None:
            return latents.to(device=device, dtype=dtype)
        latent_height = height // self.vae_scale_factor
        latent_width = width // self.vae_scale_factor
        in_channels = int(getattr(self.generator, "in_channels", 4))
        return randn_tensor(
            (batch_size, in_channels, latent_height, latent_width),
            generator=generator,
            device=device,
            dtype=dtype,
        )

    def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"):
        if output_type == "latent":
            return latents
        scaling_factor = getattr(self.vae.config, "scaling_factor", 0.18215)
        image = self.vae.decode(latents / scaling_factor).sample
        if output_type == "pt":
            return image
        return self.image_processor.postprocess(image, output_type=output_type)

    def _run_denoising_loop(
        self,
        latents: torch.Tensor,
        class_labels_tensor: torch.Tensor,
        batch_size: int,
        device: torch.device,
        num_inference_steps: int,
        sampler: str,
        pred_type: str,
    ) -> torch.Tensor:
        self.scheduler.set_timesteps(num_inference_steps, device=device, solver=sampler)
        timesteps = self.scheduler.timesteps
        self.generator.eval()
        for t_src, t_tgt in self.progress_bar(list(zip(timesteps[:-1], timesteps[1:]))):
            t_src_batch = t_src.expand(batch_size)
            t_tgt_batch = t_tgt.expand(batch_size)
            outputs = self.generator(latents, class_labels_tensor, t_src_batch, t_tgt_batch)
            if pred_type == "v":
                if sampler == "heun":
                    latents_next = self.scheduler.step(
                        outputs, t_src, t_tgt, latents, prediction_type="v"
                    ).prev_sample
                    outputs_next = self.generator(latents_next, class_labels_tensor, t_tgt_batch, t_tgt_batch)
                    latents = self.scheduler.step(
                        outputs,
                        t_src,
                        t_tgt,
                        latents,
                        model_output_next=outputs_next,
                        prediction_type="v",
                    ).prev_sample
                else:
                    latents = self.scheduler.step(
                        outputs, t_src, t_tgt, latents, prediction_type="v"
                    ).prev_sample
            else:
                latents = outputs
        return latents

    @torch.inference_mode()
    def __call__(
        self,
        class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor],
        height: Optional[int] = None,
        width: Optional[int] = None,
        num_inference_steps: Optional[int] = None,
        sampler: str = "euler",
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        output_type: str = "pil",
        return_dict: bool = True,
    ):
        default_size = self._default_image_size()
        height = int(height or default_size)
        width = int(width or default_size)
        if num_inference_steps is None:
            num_inference_steps = int(getattr(self.generator.config, "num_inference_steps", 1))
        self.check_inputs(height, width, num_inference_steps, output_type, sampler)

        device = getattr(self, "_execution_device", None) or next(self.generator.parameters()).device
        dtype = next(self.generator.parameters()).dtype
        pred_type = self.config.pred_type
        class_labels_tensor = self._normalize_class_labels(class_labels, device=device)
        batch_size = class_labels_tensor.shape[0]
        latents = self.prepare_latents(
            batch_size=batch_size,
            height=height,
            width=width,
            dtype=dtype,
            device=device,
            generator=generator,
            latents=latents,
        )
        latents = self._run_denoising_loop(
            latents,
            class_labels_tensor,
            batch_size,
            device,
            num_inference_steps,
            sampler,
            pred_type,
        )
        image = self.decode_latents(latents, output_type=output_type)
        self.maybe_free_model_hooks()
        if not return_dict:
            return (image,)
        return ImagePipelineOutput(images=image)