iMF-B-SIM/pipeline.py

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

from __future__ import annotations

import importlib.util
import inspect

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

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

# imeanflow / FD-Loss sd-vae latent statistics (stabilityai/sd-vae-ft-mse)
LATENT_CHANNEL_MEAN = (0.86488, -0.27787343, 0.21616915, 0.3738409)
LATENT_CHANNEL_STD = (4.85503674, 5.31922414, 3.93725398, 3.9870003)


def _load_bundled_vae(transformer) -> Optional[Any]:
    transformer_path = getattr(transformer.config, "_name_or_path", None)
    if not transformer_path:
        return None
    vae_dir = Path(transformer_path).resolve().parent / "vae"
    if not vae_dir.is_dir() or not (vae_dir / "config.json").is_file():
        return None
    from diffusers import AutoencoderKL

    vae_dtype = getattr(transformer, "dtype", torch.float32)
    return AutoencoderKL.from_pretrained(str(vae_dir), torch_dtype=vae_dtype)


class IMFPipeline(DiffusionPipeline):

    @staticmethod
    def prepare_extra_step_kwargs(
        scheduler,
        generator=None,
        eta: float | None = None,
    ):
        kwargs = {}
        step_params = set(inspect.signature(scheduler.step).parameters.keys())
        if "generator" in step_params:
            kwargs["generator"] = generator
        if eta is not None and "eta" in step_params:
            kwargs["eta"] = eta
        return kwargs

    @staticmethod
    def _prepare_generator(
        generator: Optional[Union[torch.Generator, List[torch.Generator]]],
    ) -> Optional[Union[torch.Generator, List[torch.Generator]]]:
        if generator is None:
            return None
        if isinstance(generator, list):
            for gen in generator:
                if gen is not None:
                    gen.manual_seed(int(gen.initial_seed()))
            return generator
        generator.manual_seed(int(generator.initial_seed()))
        return generator

    @staticmethod
    def _coerce_scheduler(scheduler, transformer) -> Any:
        if scheduler is not None and not isinstance(scheduler, (list, tuple)):
            return scheduler
        variant_path = getattr(transformer.config, "_name_or_path", None)
        if variant_path:
            scheduler_dir = Path(variant_path).resolve().parent / "scheduler"
            module_path = scheduler_dir / "scheduling_imf.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_imf", 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.IMFScheduler.from_pretrained(str(scheduler_dir))
        raise ValueError(
            "IMFPipeline could not load IMFScheduler. Ensure the variant includes scheduler/scheduling_imf.py."
        )

    @staticmethod
    def _resolve_inference_generator(
        device: Union[str, torch.device],
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    ) -> Optional[Union[torch.Generator, List[torch.Generator]]]:
        if generator is None:
            return None
        if isinstance(device, str):
            device = torch.device(device)
        device_type = device.type

        def _relocate(gen: torch.Generator) -> torch.Generator:
            if gen.device.type == device_type:
                return gen
            return torch.Generator(device=device_type).manual_seed(gen.initial_seed())

        if isinstance(generator, list):
            return [_relocate(g) for g in generator]
        return _relocate(generator)

    model_cpu_offload_seq = "transformer->vae"

    def __init__(
        self,
        transformer,
        scheduler,
        vae=None,
        id2label: Optional[Dict[Union[int, str], str]] = None,
    ):
        super().__init__()
        scheduler = self._coerce_scheduler(scheduler, transformer)
        if scheduler is None:
            raise ValueError("IMFPipeline requires a scheduler loaded from the checkpoint.")
        if isinstance(vae, (list, tuple)):
            vae = None
        if vae is None:
            vae = _load_bundled_vae(transformer)
        if vae is not None:
            vae = vae.to(device=transformer.device, dtype=getattr(transformer, "dtype", torch.float32))
        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)
        self.latent_channel_mean = torch.tensor(LATENT_CHANNEL_MEAN, dtype=torch.float32).view(1, 4, 1, 1)
        self.latent_channel_std = torch.tensor(LATENT_CHANNEL_STD, dtype=torch.float32).view(1, 4, 1, 1)
        vae_scale_factor = 8
        if vae is not None:
            vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
        self.vae_scale_factor = vae_scale_factor
        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)

    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 {}
        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 id2label(self) -> Dict[int, str]:
        self._ensure_labels_loaded()
        return self._id2label

    def get_label_ids(self, label: Union[str, List[str]]) -> List[int]:
        self._ensure_labels_loaded()
        if not self.labels:
            raise ValueError("No 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 self.labels]
        if missing:
            preview = ", ".join(list(self.labels.keys())[:8])
            raise ValueError(f"Unknown label(s): {missing}. Example valid labels: {preview}, ...")
        return [self.labels[item] for item in label]

    def _normalize_class_labels(self, class_labels: Union[int, str, List[Union[int, str]]]) -> List[int]:
        if isinstance(class_labels, int):
            return [class_labels]
        if isinstance(class_labels, str):
            return self.get_label_ids(class_labels)
        if class_labels and isinstance(class_labels[0], str):
            return self.get_label_ids(class_labels)
        return list(class_labels)

    def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
        mean = self.latent_channel_mean.to(device=latents.device, dtype=latents.dtype)
        std = self.latent_channel_std.to(device=latents.device, dtype=latents.dtype)
        return latents * std + mean

    def decode_latents(self, latents: torch.Tensor, output_type: str = "pil"):
        if output_type == "latent":
            return latents
        if self.vae is None:
            raise ValueError(
                "Cannot decode latents without a VAE. Ensure model_index.json lists vae and the variant includes vae/."
            )
        vae_dtype = next(self.vae.parameters()).dtype
        if next(self.vae.parameters()).device != latents.device:
            self.vae.to(device=latents.device, dtype=vae_dtype)
        latents = self._denormalize_latents(latents).to(dtype=vae_dtype)
        image = self.vae.decode(latents).sample
        return self.image_processor.postprocess(image, output_type=output_type)

    def _predict_velocity_u(
        self,
        latents: torch.Tensor,
        timestep: torch.Tensor,
        time_gap: torch.Tensor,
        class_labels: torch.Tensor,
        class_null: torch.Tensor,
        guidance_scale: float,
        guidance_interval_start: float,
        guidance_interval_end: float,
        do_classifier_free_guidance: bool,
    ) -> torch.Tensor:
        if do_classifier_free_guidance:
            latents_in = torch.cat([latents, latents], dim=0)
            labels = torch.cat([class_labels, class_null], dim=0)
            omega = torch.tensor([guidance_scale, 1.0], device=latents.device, dtype=latents.dtype)
            t_min = torch.tensor([guidance_interval_start, 0.0], device=latents.device, dtype=latents.dtype)
            t_max = torch.tensor([guidance_interval_end, 1.0], device=latents.device, dtype=latents.dtype)
            batch = latents.shape[0]
            timestep_in = timestep.reshape(1).repeat(2 * batch).to(dtype=latents.dtype)
            time_gap_in = time_gap.reshape(1).repeat(2 * batch).to(dtype=latents.dtype)
            omega = omega.repeat(batch)
            t_min = t_min.repeat(batch)
            t_max = t_max.repeat(batch)
        else:
            latents_in = latents
            labels = class_labels
            batch = latents.shape[0]
            timestep_in = timestep.reshape(1).repeat(batch).to(dtype=latents.dtype)
            time_gap_in = time_gap.reshape(1).repeat(batch).to(dtype=latents.dtype)
            omega = torch.full((batch,), guidance_scale, device=latents.device, dtype=latents.dtype)
            t_min = torch.full((batch,), guidance_interval_start, device=latents.device, dtype=latents.dtype)
            t_max = torch.full((batch,), guidance_interval_end, device=latents.device, dtype=latents.dtype)

        outputs = self.transformer(
            sample=latents_in,
            timestep=timestep_in,
            class_labels=labels,
            time_gap=time_gap_in,
            guidance_scale=omega,
            guidance_interval_start=t_min,
            guidance_interval_end=t_max,
            return_dict=True,
        )
        velocity_u = outputs.velocity_u

        if not do_classifier_free_guidance:
            return velocity_u

        u_cond, u_uncond = velocity_u.chunk(2, dim=0)
        return u_uncond + guidance_scale * (u_cond - u_uncond)

    @torch.inference_mode()
    def __call__(
        self,
        class_labels: Union[int, str, List[Union[int, str]]],
        num_inference_steps: int = 1,
        guidance_scale: float = 2.7,
        guidance_interval_start: float = 0.1,
        guidance_interval_end: float = 0.9,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        latents: Optional[torch.Tensor] = None,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
        if output_type not in {"pil", "np", "pt", "latent"}:
            raise ValueError("output_type must be one of: 'pil', 'np', 'pt', 'latent'.")

        class_label_ids = self._normalize_class_labels(class_labels)
        do_classifier_free_guidance = guidance_scale > 1.0
        batch_size = len(class_label_ids)
        generator = self._resolve_inference_generator(self._execution_device, generator)

        image_size = int(self.transformer.config.sample_size)
        channels = int(self.transformer.config.in_channels)
        null_class_val = int(getattr(self.transformer.config, "num_classes", 1000))

        generator = self._prepare_generator(generator)

        if latents is None:
            latents = randn_tensor(
                shape=(batch_size, channels, image_size, image_size),
                generator=generator,
                device=self._execution_device,
                dtype=self.transformer.dtype,
            )

        class_labels_t = torch.tensor(class_label_ids, device=latents.device, dtype=torch.long).reshape(-1)
        class_labels_t = class_labels_t.clamp(0, null_class_val - 1)
        class_null = torch.full_like(class_labels_t, null_class_val)

        self.scheduler.set_timesteps(num_inference_steps, device=latents.device)
        timesteps = self.scheduler.timesteps

        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)

        for i in self.progress_bar(range(num_inference_steps)):
            t = timesteps[i]
            t_next = timesteps[i + 1]
            time_gap = t - t_next
            velocity_u = self._predict_velocity_u(
                latents,
                t,
                time_gap,
                class_labels_t,
                class_null,
                guidance_scale,
                guidance_interval_start,
                guidance_interval_end,
                do_classifier_free_guidance,
            )
            latents = self.scheduler.step(velocity_u, t, latents, **extra_step_kwargs).prev_sample

        images = self.decode_latents(latents, output_type=output_type)

        self.maybe_free_model_hooks()

        if not return_dict:
            return (images,)
        return ImagePipelineOutput(images=images)


IMFPipelineOutput = ImagePipelineOutput