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

# 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.

from __future__ import annotations

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

import torch
from diffusers.image_processor import VaeImageProcessor
from diffusers.models import AutoencoderKL
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
from diffusers.schedulers import DDIMScheduler, KarrasDiffusionSchedulers
from diffusers.utils.torch_utils import randn_tensor


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

        >>> model_dir = Path("./DiCo-XL-256").resolve()
        >>> pipe = DiffusionPipeline.from_pretrained(
        ...     str(model_dir),
        ...     local_files_only=True,
        ...     custom_pipeline=str(model_dir / "pipeline.py"),
        ...     trust_remote_code=True,
        ...     torch_dtype=torch.bfloat16,
        ... )
        >>> pipe.to("cuda")

        >>> image = pipe(
        ...     class_labels="golden retriever",
        ...     num_inference_steps=250,
        ...     guidance_scale=1.4,
        ...     generator=torch.Generator("cuda").manual_seed(0),
        ... ).images[0]
        ```
"""


class DiCoPipeline(DiffusionPipeline):
    r"""
    Pipeline for class-conditional image generation with DiCo (Diffusion ConvNet).

    Parameters:
        transformer ([`DiCoTransformer2DModel`]):
            Class-conditional DiCo denoiser operating in VAE latent space.
        vae ([`AutoencoderKL`]):
            Variational autoencoder used to decode latents to pixels.
        scheduler ([`DDIMScheduler`]):
            Diffusion scheduler. Other [`KarrasDiffusionSchedulers`] can be swapped at inference time.
        id2label (`dict[int, str]`, *optional*):
            ImageNet class id to English label mapping. Values may contain comma-separated synonyms.
    """

    model_cpu_offload_seq = "transformer->vae"

    @staticmethod
    def prepare_extra_step_kwargs(
        scheduler: KarrasDiffusionSchedulers,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        eta: float = 0.0,
    ) -> Dict[str, object]:
        kwargs: Dict[str, object] = {}
        step_params = set(inspect.signature(scheduler.step).parameters.keys())
        if "generator" in step_params:
            kwargs["generator"] = generator
        if "eta" in step_params:
            kwargs["eta"] = eta
        return kwargs

    def __init__(
        self,
        transformer,
        vae: AutoencoderKL,
        scheduler: KarrasDiffusionSchedulers,
        id2label: Optional[Dict[Union[int, str], str]] = None,
    ):
        super().__init__()
        if scheduler is None:
            scheduler = DDIMScheduler(
                num_train_timesteps=1000,
                beta_start=0.0001,
                beta_end=0.02,
                beta_schedule="linear",
                clip_sample=False,
                set_alpha_to_one=True,
                steps_offset=0,
                prediction_type="epsilon",
            )
        self.register_modules(transformer=transformer, vae=vae, scheduler=scheduler)
        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
        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)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        model_kwargs = dict(kwargs)
        transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
        scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
        vae_subfolder = model_kwargs.pop("vae_subfolder", None)
        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
        base_path = Path(pretrained_model_name_or_path)

        if transformer_subfolder is None and (base_path / "transformer").exists():
            transformer_subfolder = "transformer"
        if scheduler_subfolder is None and (base_path / "scheduler").exists():
            scheduler_subfolder = "scheduler"
        if vae_subfolder is None and (base_path / "vae").exists():
            vae_subfolder = "vae"

        try:
            return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
        except Exception:
            transformer_path = str(base_path / transformer_subfolder) if transformer_subfolder else pretrained_model_name_or_path
            from transformer.transformer_dico import DiCoTransformer2DModel
            transformer = DiCoTransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
            try:
                scheduler = DDIMScheduler.from_pretrained(
                    pretrained_model_name_or_path,
                    subfolder=scheduler_subfolder,
                    **scheduler_kwargs,
                )
            except Exception:
                scheduler = DDIMScheduler(
                    num_train_timesteps=1000,
                    beta_start=0.0001,
                    beta_end=0.02,
                    beta_schedule="linear",
                    clip_sample=False,
                    set_alpha_to_one=True,
                    steps_offset=0,
                    prediction_type="epsilon",
                    **scheduler_kwargs,
                )
            try:
                vae = AutoencoderKL.from_pretrained(
                    pretrained_model_name_or_path,
                    subfolder=vae_subfolder,
                    **model_kwargs,
                )
            except Exception:
                vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema", **model_kwargs)
            id2label = cls._read_id2label_from_model_index(str(base_path))
            return cls(transformer=transformer, vae=vae, scheduler=scheduler, id2label=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 {}
        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 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 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 label]

    def _normalize_class_labels(
        self,
        class_labels: Union[int, str, List[Union[int, str]], torch.LongTensor],
    ) -> torch.LongTensor:
        if torch.is_tensor(class_labels):
            return class_labels.to(device=self._execution_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=self._execution_device, dtype=torch.long).reshape(-1)

    def _default_image_size(self) -> int:
        return int(self.transformer.config.input_size) * self.vae_scale_factor

    def check_inputs(
        self,
        height: int,
        width: int,
        num_inference_steps: int,
        output_type: 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 height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
            raise ValueError(
                f"height and width must be divisible by the VAE downsample factor {self.vae_scale_factor}."
            )
        latent_height = height // self.vae_scale_factor
        latent_width = width // self.vae_scale_factor
        expected_size = int(self.transformer.config.input_size)
        if latent_height != expected_size or latent_width != expected_size:
            raise ValueError(
                f"Requested latent size {(latent_height, latent_width)} does not match the pretrained "
                f"transformer input_size={expected_size}. Use height=width={self._default_image_size()}."
            )

    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]]],
    ) -> torch.Tensor:
        latent_height = height // self.vae_scale_factor
        latent_width = width // self.vae_scale_factor
        return randn_tensor(
            (batch_size, self.transformer.config.in_channels, latent_height, latent_width),
            generator=generator,
            device=device,
            dtype=dtype,
        )

    @staticmethod
    def _expand_timestep(timestep, batch: int, device: torch.device) -> torch.Tensor:
        if not torch.is_tensor(timestep):
            timestep = torch.tensor([timestep], dtype=torch.long, device=device)
        elif timestep.ndim == 0:
            timestep = timestep[None].to(device=device)
        return timestep.expand(batch)

    @staticmethod
    def _prepare_model_output_for_scheduler(
        model_output: torch.Tensor,
        latent_channels: int,
        scheduler: KarrasDiffusionSchedulers,
    ) -> torch.Tensor:
        if model_output.shape[1] != 2 * latent_channels:
            return model_output
        variance_type = getattr(scheduler.config, "variance_type", None)
        if scheduler.__class__.__name__ == "DDPMScheduler" and variance_type in ("learned", "learned_range"):
            return model_output
        model_output, _ = torch.split(model_output, latent_channels, dim=1)
        return model_output

    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)

    @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: int = 250,
        guidance_scale: float = 1.0,
        eta: float = 0.0,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
        output_type: str = "pil",
        return_dict: bool = True,
    ) -> Union[ImagePipelineOutput, Tuple]:
        default_size = self._default_image_size()
        height = int(height or default_size)
        width = int(width or default_size)
        self.check_inputs(height, width, num_inference_steps, output_type)

        device = self._execution_device
        model_dtype = next(self.transformer.parameters()).dtype
        class_labels_tensor = self._normalize_class_labels(class_labels)
        batch_size = class_labels_tensor.numel()
        latent_channels = int(self.transformer.config.in_channels)
        null_class_val = int(self.transformer.config.num_classes)
        do_cfg = guidance_scale > 1.0

        latents = self.prepare_latents(
            batch_size=batch_size,
            height=height,
            width=width,
            dtype=model_dtype,
            device=device,
            generator=generator,
        )
        latent_model_input = torch.cat([latents] * 2) if do_cfg else latents

        class_labels_input = class_labels_tensor
        if do_cfg:
            class_null = torch.full_like(class_labels_tensor, null_class_val)
            class_labels_input = torch.cat([class_labels_tensor, class_null], dim=0)

        self.scheduler.set_timesteps(num_inference_steps, device=device)
        extra_step_kwargs = self.prepare_extra_step_kwargs(self.scheduler, generator=generator, eta=eta)

        for t in self.progress_bar(self.scheduler.timesteps):
            if do_cfg:
                half = latent_model_input[: len(latent_model_input) // 2]
                latent_model_input = torch.cat([half, half], dim=0)
            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
            timesteps = self._expand_timestep(t, latent_model_input.shape[0], latent_model_input.device)

            noise_pred = self.transformer(
                hidden_states=latent_model_input,
                timestep=timesteps,
                class_labels=class_labels_input,
                return_dict=True,
            ).sample

            if do_cfg:
                eps, rest = noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
                cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
                half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
                eps = torch.cat([half_eps, half_eps], dim=0)
                noise_pred = torch.cat([eps, rest], dim=1)

            model_output = self._prepare_model_output_for_scheduler(noise_pred, latent_channels, self.scheduler)
            latent_model_input = self.scheduler.step(
                model_output, t, latent_model_input, return_dict=True, **extra_step_kwargs
            ).prev_sample

        if do_cfg:
            latents, _ = latent_model_input.chunk(2, dim=0)
        else:
            latents = latent_model_input

        image = self.decode_latents(latents, output_type=output_type)
        self.maybe_free_model_hooks()
        if not return_dict:
            return (image,)
        return ImagePipelineOutput(images=image)


DiCoPipelineOutput = ImagePipelineOutput