from __future__ import annotations

from typing import Any

import torch
from diffusers import DiffusionPipeline

from .base import TextToImageGenerator


class DiffusersTextToImageAdapter(TextToImageGenerator):
    def __init__(self, model_id: str, device: str = "cuda", torch_dtype: Any | None = None, pipeline: Any | None = None, **kwargs) -> None:
        if pipeline is None:
            dtype = torch_dtype
            if dtype is None and str(device).startswith("cuda"):
                dtype = torch.float16
            pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype=dtype, **kwargs)
        self.pipe = pipeline
        self.device = device
        if hasattr(self.pipe, "to"):
            self.pipe.to(device)
        if hasattr(self.pipe, "set_progress_bar_config"):
            self.pipe.set_progress_bar_config(disable=True)

    def generate(self, prompts: list[str], *, generator: Any | None = None, **kwargs) -> list[Any]:
        out = self.pipe(prompt=prompts, generator=generator, **kwargs)
        if hasattr(out, "images"):
            return list(out.images)
        if isinstance(out, list):
            return out
        raise TypeError("Text-to-image pipeline output does not expose `.images`.")

    def generate_batch(self, prompts, *, generators=None, **kwargs):
        out = self.pipe(prompt=list(prompts), generator=generators, **kwargs)
        if hasattr(out, "images"):
            return list(out.images)
        if isinstance(out, list):
            return out
        raise TypeError("Text-to-image pipeline output does not expose `.images`.")