# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Hub custom pipeline: JiTPipeline for FD-Loss post-trained JiT (legacy time, velocity Euler)."""

from __future__ import annotations

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

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

EXAMPLE_DOC_STRING = """
    Examples:
        ```py
        >>> from pathlib import Path
        >>> import torch
        >>> from diffusers import DiffusionPipeline

        >>> model_dir = Path("BiliSakura/FD-Loss-diffusers/JiT-B-16-SIM")
        >>> pipe = DiffusionPipeline.from_pretrained(
        ...     str(model_dir),
        ...     custom_pipeline=str(model_dir / "pipeline.py"),
        ...     trust_remote_code=True,
        ... )
        >>> pipe.to("cuda")
        >>> image = pipe(class_labels="golden retriever", num_inference_steps=1, guidance_scale=3.0).images[0]
        ```
"""

RECOMMENDED_NOISE_BY_SIZE = {256: 1.0, 512: 2.0}


class JiTPipeline(DiffusionPipeline):
    r"""
    Class-conditional JiT pipeline for FD-Loss checkpoints (1-NFE velocity Euler, legacy time convention).

    Args:
        transformer ([`JiTTransformer2DModel`]): JiT backbone.
        scheduler ([`FDLossFlowMatchScheduler`], *optional*): Bundled flow-matching scheduler (`t=1→0`).
            Accepts `generator` in `step()` for reproducible stochastic sampling.
        id2label (`dict[int, str]`, *optional*): ImageNet id → label strings from `model_index.json`.
    """

    model_cpu_offload_seq = "transformer"
    _optional_components = ["scheduler"]

    @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 _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)

    @staticmethod
    def _coerce_scheduler(scheduler, transformer):
        """Always load the bundled FD-Loss flow scheduler (ignore model_index diffusers schedulers)."""
        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_flow_match_fd.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_flow_match_fd", 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)
                    scheduler_cls = getattr(module, "FDLossFlowMatchScheduler")
                    return scheduler_cls.from_pretrained(str(scheduler_dir))
        scheduler_dir = Path(__file__).resolve().parent / "scheduler"
        module_path = scheduler_dir / "scheduling_flow_match_fd.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_flow_match_fd", 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)
                scheduler_cls = getattr(module, "FDLossFlowMatchScheduler")
                return scheduler_cls.from_pretrained(str(scheduler_dir))
        raise ValueError("FD-Loss JiT requires a bundled FDLossFlowMatchScheduler under `scheduler/`.")

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


    @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]:
        r"""Map ImageNet English label strings to class ids."""
        self._ensure_labels_loaded()
        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}. Examples: {preview}, ...")
        return [self.labels[item] for item in label]

    def _backbone_timestep(self, t: torch.Tensor) -> torch.Tensor:
        if self.legacy_time_convention:
            return 1.0 - t
        return t

    def _cfg_interval(self, interval_min: float, interval_max: float, device: torch.device) -> Tuple[float, float]:
        low = float(self._backbone_timestep(torch.tensor(interval_min, device=device)))
        high = float(self._backbone_timestep(torch.tensor(interval_max, device=device)))
        return min(low, high), max(low, high)

    def _flow_time_tensors(self, t: torch.Tensor, sample: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Broadcast scalar flow time and per-sample backbone timesteps for `sample` (B, C, H, W)."""
        flow_t = t.reshape(()).to(device=sample.device, dtype=sample.dtype)
        t_view = flow_t.reshape(*([1] * sample.ndim))
        t_bb = self._backbone_timestep(flow_t).expand(sample.shape[0])
        return t_view, t_bb

    def _predict_velocity(
        self,
        sample: torch.Tensor,
        t: torch.Tensor,
        class_labels: torch.Tensor,
        guidance_scale: float,
        cfg_interval: Tuple[float, float],
        t_eps: float,
        interpolate_pos_encoding: bool,
    ) -> torch.Tensor:
        t_view, t_bb = self._flow_time_tensors(t, sample)
        num_classes = int(getattr(self.transformer.config, "num_classes", 1000))

        x_cond = self.transformer(
            sample,
            timestep=t_bb,
            class_labels=class_labels,
            interpolate_pos_encoding=interpolate_pos_encoding,
        ).sample
        v_cond = (sample - x_cond) / t_view.clamp_min(t_eps)

        if guidance_scale <= 1.0:
            return v_cond

        class_null = torch.full_like(class_labels, num_classes)
        x_uncond = self.transformer(
            sample,
            timestep=t_bb,
            class_labels=class_null,
            interpolate_pos_encoding=interpolate_pos_encoding,
        ).sample
        v_uncond = (sample - x_uncond) / t_view.clamp_min(t_eps)

        low, high = cfg_interval
        flow_t = t.reshape(()).to(device=sample.device, dtype=sample.dtype)
        mask = (flow_t < high) & ((low == 0) | (flow_t > low))
        scale = torch.where(
            mask,
            torch.tensor(guidance_scale, device=sample.device, dtype=sample.dtype),
            torch.tensor(1.0, device=sample.device, dtype=sample.dtype),
        )
        while scale.ndim < v_cond.ndim:
            scale = scale.unsqueeze(-1)
        return v_uncond + scale * (v_cond - v_uncond)

    @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 _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()))

    def _ensure_labels_loaded(self) -> None:
        if self._labels_loaded_from_model_index:
            return
        path = getattr(self.config, "_name_or_path", None)
        if path:
            model_index_path = Path(path).resolve() / "model_index.json"
            if model_index_path.exists():
                raw = json.loads(model_index_path.read_text(encoding="utf-8"))
                id2label = raw.get("id2label")
                if isinstance(id2label, dict):
                    self._id2label = {int(k): v for k, v in id2label.items()}
                    self.labels = self._build_label2id(self._id2label)
        self._labels_loaded_from_model_index = True

    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)

    @torch.inference_mode()
    def __call__(
        self,
        class_labels: Union[int, str, List[Union[int, str]]],
        guidance_scale: float = 3.0,
        guidance_interval_min: float = 0.1,
        guidance_interval_max: float = 1.0,
        noise_scale: Optional[float] = None,
        t_eps: float = 5e-2,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        num_inference_steps: int = 1,
        height: Optional[int] = None,
        width: Optional[int] = None,
        interpolate_pos_encoding: bool = True,
        output_type: Optional[str] = "pil",
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
        r"""
        Generate images conditioned on ImageNet class ids or English labels.

        Args:
            class_labels: Class id(s), label string(s), or mixed list.
            guidance_scale: CFG scale (FD-Loss JiT-B default: 3.0).
            guidance_interval_min: Lower flow-time bound for CFG (default 0.1).
            guidance_interval_max: Upper flow-time bound for CFG (default 1.0).
            noise_scale: Initial Gaussian scale (default 1.0 at 256px).
            t_eps: Clamp for flow time in velocity denominator.
            generator: RNG for latent init.
            num_inference_steps: FD-Loss sampling steps (1 = single NFE).
            height: Output height (defaults to native `sample_size`).
            width: Output width (defaults to native `sample_size`).
            interpolate_pos_encoding: Interpolate positional embeddings for non-native sizes.
            output_type: `pil`, `np`, or `pt`.
            return_dict: Return [`ImagePipelineOutput`] if True.
        """
        # 1. Check inputs
        if num_inference_steps < 1:
            raise ValueError("num_inference_steps must be >= 1.")
        if output_type not in {"pil", "np", "pt"}:
            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")

        # 2. Call parameters
        class_label_ids = self._normalize_class_labels(class_labels)
        batch_size = len(class_label_ids)
        generator = self._resolve_inference_generator(self._execution_device, generator)
        image_size = int(self.transformer.config.sample_size)
        patch_size = int(self.transformer.config.patch_size)
        height = int(height or image_size)
        width = int(width or image_size)
        if height % patch_size != 0 or width % patch_size != 0:
            raise ValueError(f"height/width must be divisible by patch_size={patch_size}.")
        channels = int(self.transformer.config.in_channels)
        null_class_val = int(getattr(self.transformer.config, "num_classes", 1000))
        if noise_scale is None:
            noise_scale = RECOMMENDED_NOISE_BY_SIZE.get(max(height, width), 1.0)

        # 3. Class labels
        labels = torch.tensor(class_label_ids, device=self._execution_device, dtype=torch.long).reshape(-1)
        labels = labels.clamp(0, null_class_val - 1)
        cfg_interval = self._cfg_interval(guidance_interval_min, guidance_interval_max, labels.device)

        # 4. Latents
        latents = randn_tensor(
            (batch_size, channels, height, width),
            generator=generator,
            device=self._execution_device,
            dtype=self.transformer.dtype,
        ) * noise_scale

        # 5–6. Denoising loop (FD-Loss flow t=1 → 0; generator forwarded to scheduler.step)
        self.scheduler.set_timesteps(num_inference_steps, device=self._execution_device)
        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator)
        timesteps = self.scheduler.timesteps

        for i in self.progress_bar(range(num_inference_steps)):
            t_cur = timesteps[i]
            velocity = self._predict_velocity(
                latents, t_cur, labels, guidance_scale, cfg_interval, t_eps, interpolate_pos_encoding
            )
            latents = self.scheduler.step(velocity, t_cur, latents, **extra_step_kwargs).prev_sample

        # 8. Output
        images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
        if output_type == "pt":
            images = images_pt
        elif output_type == "np":
            images = images_pt.permute(0, 2, 3, 1).numpy()
        else:
            images = self.numpy_to_pil(images_pt.permute(0, 2, 3, 1).numpy())

        self.maybe_free_model_hooks()
        if not return_dict:
            return (images,)
        return ImagePipelineOutput(images=images)