"""
PixelDiT T2I Pipeline — thin orchestrator.

Usage:
    from pipeline import PixelDiTPipeline
    pipe = PixelDiTPipeline()
    images = pipe("a viking at sunset", height=512, width=512)
    images[0].save("out.jpg")
"""

import torch
from PIL import Image
from .modeling_pixeldit import load_pixeldit
from .modeling_pixeldit_hf import PixelDiTModel
from .text_encoder_gemma import GemmaEncoder
from .text_encoder_qwen import QwenEncoder
from .scheduling_flow import FlowScheduler


class PixelDiTPipeline:
    def __init__(
        self,
        text_encoder="gemma",   # "gemma" | "qwen"
        qwen_proj=None,
        device="cuda",
        dtype=torch.bfloat16,
        cfg=3.5,
        flow_shift=4.0,
        pretrained=None,        # HF dir or repo id — loads via from_pretrained instead of .pth
    ):
        self.device = torch.device(device)
        self.dtype  = dtype

        if text_encoder == "qwen":
            self.encoder = QwenEncoder(proj_path=qwen_proj, output_device=device, output_dtype=dtype)
        else:
            self.encoder = GemmaEncoder(output_device=device, output_dtype=dtype)

        if pretrained is not None:
            print(f"[pipeline] loading from HF: {pretrained}")
            self.model = (
                PixelDiTModel.from_pretrained(pretrained)
                .to(device).to(dtype).eval()
            )
        else:
            self.model = load_pixeldit(device=device, dtype=dtype)

        self.scheduler = FlowScheduler(self.model, cfg=cfg, flow_shift=flow_shift)

    @torch.no_grad()
    def __call__(
        self,
        prompt,
        negative_prompt="",
        height=512,
        width=512,
        steps=20,
        cfg=None,
        seed=None,
    ):
        if isinstance(prompt, str):
            prompts = [prompt]
        else:
            prompts = list(prompt)
        B = len(prompts)

        if cfg is not None:
            self.scheduler.cfg = cfg
        if seed is not None:
            torch.manual_seed(seed)

        cond   = self.encoder.encode(prompts)
        uncond = (self.encoder.encode_null(B) if not negative_prompt
                  else self.encoder.encode([negative_prompt] * B))

        noise = torch.randn(B, 3, height, width, device=self.device, dtype=self.dtype)
        imgs  = self.scheduler.sample(noise, cond, uncond, steps=steps)

        imgs = (imgs.clamp(-1, 1) + 1) / 2
        imgs = (imgs * 255).byte().permute(0, 2, 3, 1).cpu().numpy()
        return [Image.fromarray(img) for img in imgs]