from __future__ import annotations

from typing import Sequence, Union

import torch

from diffusers import DiffusionPipeline
from diffusers.pipelines.pipeline_utils import ImagePipelineOutput


class BitDanceImageNetPipeline(DiffusionPipeline):
    model_cpu_offload_seq = "transformer"

    def __init__(self, transformer, autoencoder=None):
        super().__init__()
        self.register_modules(transformer=transformer, autoencoder=autoencoder)

    @torch.no_grad()
    def __call__(
        self,
        class_labels: Union[int, Sequence[int]] = 0,
        num_images_per_label: int = 1,
        sample_steps: int = 100,
        cfg_scale: float = 4.6,
        chunk_size: int = 0,
        output_type: str = "pil",
        return_dict: bool = True,
    ):
        device = self._execution_device

        if isinstance(class_labels, int):
            labels = [class_labels]
        else:
            labels = list(class_labels)

        class_ids = torch.tensor(labels, device=device, dtype=torch.long)
        if num_images_per_label > 1:
            class_ids = class_ids.repeat_interleave(num_images_per_label)

        images = self.transformer.sample(
            class_ids=class_ids,
            sample_steps=sample_steps,
            cfg_scale=cfg_scale,
            chunk_size=chunk_size,
        )

        images = (images / 2 + 0.5).clamp(0, 1)
        images = images.cpu().permute(0, 2, 3, 1).float().numpy()

        if output_type == "pil":
            images = self.numpy_to_pil(images)

        if not return_dict:
            return (images,)

        return ImagePipelineOutput(images=images)