"""Tiny text-to-image pipeline for dancing-chibi-figures (64px RGBA pixel diffusion, no VAE).
Ships inside the model repo as pipeline.py so this works anywhere:

    from diffusers import DiffusionPipeline
    pipe = DiffusionPipeline.from_pretrained("sprited/dancing-chibi-figures-t2i-64",
                                             custom_pipeline="sprited/dancing-chibi-figures-t2i-64",
                                             trust_remote_code=True).to("cuda")
    imgs = pipe(["a person jumps in place"] * 4, num_inference_steps=50, guidance_scale=3.0).images
"""
import torch
from diffusers import DiffusionPipeline, ImagePipelineOutput
from diffusers.utils.torch_utils import randn_tensor
from PIL import Image


class ChibiT2IPipeline(DiffusionPipeline):
    def __init__(self, unet, scheduler, text_encoder, tokenizer):
        super().__init__()
        self.register_modules(unet=unet, scheduler=scheduler, text_encoder=text_encoder, tokenizer=tokenizer)

    @torch.no_grad()
    def encode(self, texts):
        tok = self.tokenizer(list(texts), padding="max_length", max_length=32, truncation=True, return_tensors="pt")
        return self.text_encoder(tok.input_ids.to(self.device)).last_hidden_state

    @torch.no_grad()
    def __call__(self, prompt, num_inference_steps=50, guidance_scale=3.0, generator=None, output_type="pil"):
        if isinstance(prompt, str): prompt = [prompt]
        emb = self.encode(prompt)
        null = self.encode([""] * len(prompt))
        s = self.unet.config.sample_size
        x = randn_tensor((len(prompt), 4, s, s), generator=generator, device=self.device, dtype=emb.dtype)
        self.scheduler.set_timesteps(num_inference_steps)
        for t in self.progress_bar(self.scheduler.timesteps):
            eps_c = self.unet(x, t, encoder_hidden_states=emb).sample
            eps_u = self.unet(x, t, encoder_hidden_states=null).sample
            eps = eps_u + guidance_scale * (eps_c - eps_u)
            x = self.scheduler.step(eps, t, x).prev_sample
        x = ((x.clamp(-1, 1) + 1) / 2 * 255).round().byte().cpu()         # premultiplied RGBA
        if output_type == "pil":
            imgs = [Image.fromarray(v.permute(1, 2, 0).numpy(), "RGBA") for v in x]
        else:
            imgs = x.numpy()
        return ImagePipelineOutput(images=imgs)