"""AuroraIMG2-style text-to-image pipeline (load with trust_remote_code=True).

Based on the AuroraIMG2-50M-Preview pipeline (Apache-2.0), extended with a second text encoder type:
  * AuroraTextEncoder    -> the released from-scratch encoder (unchanged)
  * AuroraLMTextEncoder  -> a frozen pretrained causal LM (SlayerLab/Slayer149 GPT, Apache-2.0) + a small bidirectional
                            refiner; its LM weights are stored inside text_encoder/ so the repo is self-contained
Generated by trenowanie_generatora_obrazkow/export.py.

Components (all stored in the repo as separate sub-folders):
  * tokenizer     -> from-scratch byte-level BPE (8192 tokens), loaded with transformers
  * text_encoder  -> AuroraTextEncoder, ~12.7M params, trained from scratch (defined below)
  * transformer   -> AuroraDenoiser, ~50M params flow-matching DiT, trained from scratch (defined below)
  * vae           -> pretrained FLUX.2 small decoder (AutoencoderKLFlux2, Apache-2.0), frozen

Sampling is plain Euler on t = linspace(1, 0, steps + 1) with classifier-free guidance, i.e. exactly
what was used for the reported benchmark numbers.

NOTE: diffusers' custom-pipeline loader expects exactly ONE DiffusionPipeline subclass in this file.
"""
import contextlib
import math
from typing import List, Optional, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

from diffusers import ConfigMixin, DiffusionPipeline, ModelMixin
from diffusers.configuration_utils import register_to_config
from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
from diffusers.utils.torch_utils import randn_tensor


# =============================================================================
# Building blocks (identical to the training code)
# =============================================================================
class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps).to(x.dtype) * self.weight


def rotate_half_pairs(x: torch.Tensor) -> torch.Tensor:
    even = x[..., 0::2]
    odd = x[..., 1::2]
    return torch.stack((-odd, even), dim=-1).flatten(-2)


def _angles(positions: torch.Tensor, dim: int) -> torch.Tensor:
    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
    return torch.outer(positions.float(), inv_freq)


def build_rope_1d(n: int, d: int):
    ang = _angles(torch.arange(n), d)
    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]


def build_rope_2d(side: int, d: int):
    """First half of the head dim encodes the row, second half the column. Tokens are row-major."""
    coords = torch.arange(side)
    rows, cols = torch.meshgrid(coords, coords, indexing="ij")
    ang = torch.cat([_angles(rows.reshape(-1), d // 2), _angles(cols.reshape(-1), d // 2)], dim=-1)
    return ang.cos().repeat_interleave(2, dim=-1)[None, None], ang.sin().repeat_interleave(2, dim=-1)[None, None]


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    return x * cos.to(x.dtype) + rotate_half_pairs(x) * sin.to(x.dtype)


class SwiGLU(nn.Module):
    def __init__(self, dim: int, hidden: int):
        super().__init__()
        self.in_proj = nn.Linear(dim, hidden * 2, bias=False)
        self.out_proj = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        a, b = self.in_proj(x).chunk(2, dim=-1)
        return self.out_proj(F.silu(a) * b)


class TextAttention(nn.Module):
    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int):
        super().__init__()
        self.heads, self.kv_heads, self.head_dim = heads, kv_heads, head_dim
        self.qkv = nn.Linear(dim, (heads + 2 * kv_heads) * head_dim, bias=False)
        self.out = nn.Linear(heads * head_dim, dim, bias=False)

    def forward(self, x, valid_mask, cos, sin):
        b, t, _ = x.shape
        q, k, v = self.qkv(x).split(
            [self.heads * self.head_dim, self.kv_heads * self.head_dim, self.kv_heads * self.head_dim], dim=-1)
        q = q.view(b, t, self.heads, self.head_dim).transpose(1, 2)
        k = k.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
        v = v.view(b, t, self.kv_heads, self.head_dim).transpose(1, 2)
        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
        rep = self.heads // self.kv_heads
        k, v = k.repeat_interleave(rep, dim=1), v.repeat_interleave(rep, dim=1)
        h = F.scaled_dot_product_attention(q, k, v, attn_mask=valid_mask[:, None, None, :].bool())
        return self.out(h.transpose(1, 2).reshape(b, t, -1))


class TextBlock(nn.Module):
    def __init__(self, dim: int, heads: int, kv_heads: int, head_dim: int, mlp_hidden: int):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.attn = TextAttention(dim, heads, kv_heads, head_dim)
        self.norm2 = RMSNorm(dim)
        self.mlp = SwiGLU(dim, mlp_hidden)

    def forward(self, x, valid_mask, cos, sin):
        x = x + self.attn(self.norm1(x), valid_mask, cos, sin)
        return x + self.mlp(self.norm2(x))


class SelfAttention(nn.Module):
    """Global self-attention over the 16x16 latent tokens with 2D RoPE and QK-norm."""

    def __init__(self, dim: int, heads: int):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out = nn.Linear(dim, dim, bias=False)
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)

    def forward(self, x, cos, sin):
        b, t, c = x.shape
        q, k, v = self.qkv(x).view(b, t, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
        q, k = apply_rope(self.q_norm(q), cos, sin), apply_rope(self.k_norm(k), cos, sin)
        h = F.scaled_dot_product_attention(q, k, v)
        return self.out(h.transpose(1, 2).reshape(b, t, c))


class CrossAttention(nn.Module):
    def __init__(self, dim: int, heads: int, text_dim: int):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.q = nn.Linear(dim, dim, bias=False)
        self.kv = nn.Linear(text_dim, 2 * dim, bias=False)
        self.out = nn.Linear(dim, dim, bias=False)

    def forward(self, x, text, text_valid):
        b, t, c = x.shape
        n = text.shape[1]
        q = self.q(x).view(b, t, self.heads, self.head_dim).transpose(1, 2)
        k, v = self.kv(text).chunk(2, dim=-1)
        k = k.view(b, n, self.heads, self.head_dim).transpose(1, 2)
        v = v.view(b, n, self.heads, self.head_dim).transpose(1, 2)
        h = F.scaled_dot_product_attention(q, k, v, attn_mask=text_valid[:, None, None, :].bool())
        return self.out(h.transpose(1, 2).reshape(b, t, c))


class DenoiserBlock(nn.Module):
    """DiT block with PixArt-style adaLN-single: a shared time->(6*dim) projection plus a per-block table."""

    def __init__(self, dim: int, heads: int, text_dim: int, mlp_hidden: int):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.attn = SelfAttention(dim, heads)
        self.norm2 = RMSNorm(dim)
        self.cross = CrossAttention(dim, heads, text_dim)
        self.norm3 = RMSNorm(dim)
        self.mlp = SwiGLU(dim, mlp_hidden)
        self.mod_table = nn.Parameter(torch.zeros(6, dim))

    def forward(self, x, t6, text, text_valid, cos, sin):
        shift1, scale1, gate1, shift2, scale2, gate2 = (t6 + self.mod_table[None]).to(x.dtype).unbind(1)
        h = self.norm1(x) * (1 + scale1[:, None]) + shift1[:, None]
        x = x + gate1[:, None] * self.attn(h, cos, sin)
        x = x + self.cross(self.norm2(x), text, text_valid)
        h = self.norm3(x) * (1 + scale2[:, None]) + shift2[:, None]
        return x + gate2[:, None] * self.mlp(h)


def timestep_embedding(t: torch.Tensor, dim: int = 256) -> torch.Tensor:
    t = t.float() * 1000.0
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device).float() / max(half - 1, 1))
    angles = t[:, None] * freqs[None, :]
    return torch.cat([angles.sin(), angles.cos()], dim=-1)


# =============================================================================
# Frozen-LM text encoder (GoLLeM / Slayer GPT, inference architecture, Apache-2.0)
# =============================================================================
class LMRMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d))
        self.eps = eps

    def forward(self, x):
        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight


def lm_apply_rope(x, base=100000.0):
    """Parameter-free RoPE on [B, H, T, D], interleaved convention (even/odd channel pairs)."""
    _, _, T, dim = x.shape
    pos = torch.arange(T, device=x.device, dtype=torch.float32)
    freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
    ang = torch.outer(pos, freq)
    cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None]
    even, odd = x[..., ::2], x[..., 1::2]
    return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)


class LMSwiGLU(nn.Module):
    def __init__(self, d, hidden):
        super().__init__()
        self.gate = nn.Linear(d, hidden, bias=False)
        self.up = nn.Linear(d, hidden, bias=False)
        self.down = nn.Linear(hidden, d, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class LMBlock(nn.Module):
    def __init__(self, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, is_first=False):
        super().__init__()
        self.ln1, self.ln2 = LMRMSNorm(d, eps), LMRMSNorm(d, eps)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.mlp = LMSwiGLU(d, int(round(ffn_mult * d)))
        self.nh, self.d, self.rope_theta = nh, d, rope_theta
        self.value_residual, self.qk_norm, self.is_first = value_residual, qk_norm, is_first
        if value_residual and not is_first:
            self.vr_lambda = nn.Parameter(torch.zeros(1))
        if qk_norm:
            self.q_norm, self.k_norm = LMRMSNorm(d // nh, eps), LMRMSNorm(d // nh, eps)

    def forward(self, x, v0=None):
        B, T, D = x.size()
        q, k, v = self.qkv(self.ln1(x)).split(self.d, dim=2)
        hd = D // self.nh
        q, k, v = (t.view(B, T, self.nh, hd).transpose(1, 2) for t in (q, k, v))
        if self.qk_norm:
            q, k = self.q_norm(q), self.k_norm(k)
        q, k = lm_apply_rope(q, self.rope_theta), lm_apply_rope(k, self.rope_theta)
        if self.value_residual:
            if self.is_first:
                v0 = v
            else:
                v = v + self.vr_lambda * v0
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        x = x + self.proj(y.transpose(1, 2).contiguous().view(B, T, D))
        return x + self.mlp(self.ln2(x)), v0


class LMTrunk(nn.Module):
    """Token embedding + the first n blocks of the LM (no final norm, no head)."""

    def __init__(self, vocab, n_blocks, d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm):
        super().__init__()
        self.tok = nn.Embedding(vocab, d)
        self.blocks = nn.ModuleList([LMBlock(d, nh, ffn_mult, eps, rope_theta, value_residual, qk_norm, i == 0)
                                     for i in range(n_blocks)])

    def forward(self, idx):
        x, v0 = self.tok(idx), None
        for b in self.blocks:
            x, v0 = b(x, v0)
        return x


# =============================================================================
# FLUX.2 latent packing layouts (used to turn the model's (128,16,16) latents into the decoder's input)
# =============================================================================
def _unpatchify_flux2(z):  # official FLUX.2 pipeline layout: channel = c*4 + dh*2 + dw
    b, c4, h, w = z.shape
    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c4 // 4, h * 2, w * 2)


def _unpatchify_swap_subpixel(z):
    b, c4, h, w = z.shape
    return z.reshape(b, c4 // 4, 2, 2, h, w).permute(0, 1, 4, 3, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)


def _unpatchify_subpixel_major(z):
    b, c4, h, w = z.shape
    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, c4 // 4, h * 2, w * 2)


def _unpatchify_subpixel_major_swapped(z):
    b, c4, h, w = z.shape
    return z.reshape(b, 2, 2, c4 // 4, h, w).permute(0, 3, 4, 2, 5, 1).reshape(b, c4 // 4, h * 2, w * 2)


UNPATCHIFY_LAYOUTS = {
    "flux2_pipeline": _unpatchify_flux2,
    "swap_subpixel": _unpatchify_swap_subpixel,
    "subpixel_major": _unpatchify_subpixel_major,
    "subpixel_major_swapped": _unpatchify_subpixel_major_swapped,
}


# =============================================================================
# Diffusers-style models
# =============================================================================
class AuroraTextEncoder(ModelMixin, ConfigMixin):
    """Small from-scratch text encoder: GQA transformer with 1D RoPE, SwiGLU and RMSNorm."""

    @register_to_config
    def __init__(
        self,
        vocab_size: int = 8192,
        dim: int = 384,
        depth: int = 6,
        heads: int = 6,
        kv_heads: int = 2,
        head_dim: int = 64,
        mlp_hidden: int = 1024,
        max_tokens: int = 64,
        pad_id: int = 0,
        bos_id: int = 2,
        eos_id: int = 3,
        null_id: int = 4,
        mask_id: int = 5,
    ):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, dim)
        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
        self.norm = RMSNorm(dim)
        cos, sin = build_rope_1d(max_tokens, head_dim)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, ids: torch.Tensor):
        valid = ids.ne(self.config.pad_id)
        h = self.token_embedding(ids)
        for block in self.blocks:
            h = block(h, valid, self.rope_cos, self.rope_sin)
        return self.norm(h), valid


class AuroraLMTextEncoder(ModelMixin, ConfigMixin):
    """Frozen causal LM trunk -> RMSNorm (+ learned null prompt) -> bidirectional refiner blocks -> denoiser text tokens.
    Captions: [BOS] ids[:max_tokens-1] [PAD]...; the CFG null prompt is [NULL] [PAD]... (ids outside the LM vocab)."""

    @register_to_config
    def __init__(
        self,
        kind: str = "lm",
        vocab_size: int = 12288,
        lm_blocks: int = 18,  # LM blocks used (the trained run takes the stream before the last of 19)
        lm_dim: int = 768,
        lm_heads: int = 12,
        lm_ffn_mult: float = 2.8125,
        lm_norm_eps: float = 1e-6,
        lm_rope_theta: float = 100000.0,
        lm_value_residual: bool = True,
        lm_qk_norm: bool = True,
        dim: int = 768,
        depth: int = 2,
        heads: int = 12,
        kv_heads: int = 4,
        head_dim: int = 64,
        mlp_hidden: int = 2048,
        max_tokens: int = 96,
        bos_id: int = 12285,
        pad_id: int = 65535,
        null_id: int = 65534,
    ):
        super().__init__()
        self.lm = LMTrunk(vocab_size, lm_blocks, lm_dim, lm_heads, lm_ffn_mult, lm_norm_eps, lm_rope_theta,
                          lm_value_residual, lm_qk_norm)
        self.in_norm = RMSNorm(lm_dim)
        self.null = nn.Parameter(torch.zeros(1, 1, lm_dim))
        self.in_proj = nn.Linear(lm_dim, dim, bias=False) if dim != lm_dim else nn.Identity()
        self.blocks = nn.ModuleList([TextBlock(dim, heads, kv_heads, head_dim, mlp_hidden) for _ in range(depth)])
        self.norm = RMSNorm(dim)
        cos, sin = build_rope_1d(max_tokens, head_dim)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, ids: torch.Tensor):
        cfg = self.config
        valid = ids.ne(cfg.pad_id)
        is_null = ids[:, 0].eq(cfg.null_id)
        h = self.lm(ids.masked_fill(~valid | ids.eq(cfg.null_id), 0))  # causal + right padding: pads never seen
        h = self.in_norm(h)
        h = torch.where(is_null[:, None, None], self.null.to(h.dtype).expand_as(h), h)
        h = self.in_proj(h)
        for block in self.blocks:
            h = block(h, valid, self.rope_cos, self.rope_sin)
        return self.norm(h), valid


class AuroraDenoiser(ModelMixin, ConfigMixin):
    """Flow-matching DiT: velocity predictor on (B, 128, 16, 16) FLUX.2 latents, one latent pixel per token."""

    @register_to_config
    def __init__(
        self,
        latent_channels: int = 128,
        latent_size: int = 16,
        dim: int = 640,
        depth: int = 8,
        heads: int = 10,
        mlp_hidden: int = 1536,
        text_dim: int = 384,
        latent_layout: str = "flux2_pipeline",
    ):
        super().__init__()
        self.patch_in = nn.Linear(latent_channels, dim, bias=False)
        self.time_mlp = nn.Sequential(nn.Linear(256, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.ada = nn.Linear(dim, 6 * dim)
        self.blocks = nn.ModuleList([DenoiserBlock(dim, heads, text_dim, mlp_hidden) for _ in range(depth)])
        self.norm = RMSNorm(dim)
        self.patch_out = nn.Linear(dim, latent_channels, bias=False)
        cos, sin = build_rope_2d(latent_size, dim // heads)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, z_t: torch.Tensor, t: torch.Tensor, text: torch.Tensor, text_valid: torch.Tensor):
        b, c, h, w = z_t.shape
        dim = self.config.dim
        x = self.patch_in(z_t.flatten(2).transpose(1, 2))
        t_emb = self.time_mlp(timestep_embedding(t, 256))
        t6 = self.ada(F.silu(t_emb)).view(b, 6, dim)
        for block in self.blocks:
            x = block(x, t6, text, text_valid, self.rope_cos, self.rope_sin)
        return self.patch_out(self.norm(x)).transpose(1, 2).reshape(b, c, h, w)


# =============================================================================
# Pipeline
# =============================================================================
class AuroraIMG2Pipeline(DiffusionPipeline):
    """Text-to-image pipeline (256x256). Usage:

        pipe = DiffusionPipeline.from_pretrained(repo, custom_pipeline=repo, trust_remote_code=True)
        image = pipe("a red fox in the snow", num_inference_steps=50, guidance_scale=3.0).images[0]
    """

    def __init__(self, tokenizer, text_encoder, transformer, vae):
        super().__init__()
        self.register_modules(tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae)

    # ---- helpers -------------------------------------------------------------------------------
    def _autocast(self):
        device = self._execution_device
        if device.type == "cuda":
            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
            return torch.autocast(device_type="cuda", dtype=dtype)
        return contextlib.nullcontext()

    def _is_lm(self) -> bool:
        return getattr(self.text_encoder.config, "kind", "") == "lm"

    def _tokenize(self, texts: List[str]) -> torch.Tensor:
        """[BOS] + BPE ids (max 62) + [EOS], padded with [PAD] to max_tokens. Same as in training."""
        cfg = self.text_encoder.config
        if self._is_lm():  # [BOS] ids[:max-1] [PAD]..., no leading space (prepare_data.py --encoder slayer149)
            cleaned = [" ".join(str(t).split()) for t in texts]
            encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
            rows = torch.full((len(cleaned), int(cfg.max_tokens)), int(cfg.pad_id), dtype=torch.long)
            for i, ids in enumerate(encoded):
                ids = list(ids)[: int(cfg.max_tokens) - 1]
                rows[i, 0] = int(cfg.bos_id)
                if ids:
                    rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
            return rows
        max_len = int(cfg.max_tokens)
        cleaned = [" ".join(str(t).split()) for t in texts]
        # The training tokenizer uses ByteLevel(add_prefix_space=True), i.e. every caption was tokenised with a
        # leading space. The transformers wrapper can drop that setting when the tokenizer is saved/reloaded, so the
        # space is added explicitly here (a no-op if the backend already adds it).
        cleaned = [(" " + t) if t else t for t in cleaned]
        encoded = self.tokenizer(cleaned, add_special_tokens=False)["input_ids"]
        rows = torch.full((len(cleaned), max_len), int(cfg.pad_id), dtype=torch.long)
        for i, ids in enumerate(encoded):
            ids = list(ids)[: max_len - 2]
            rows[i, 0] = int(cfg.bos_id)
            if ids:
                rows[i, 1:1 + len(ids)] = torch.tensor(ids, dtype=torch.long)
            rows[i, 1 + len(ids)] = int(cfg.eos_id)
        return rows

    def _null_ids(self, batch: int) -> torch.Tensor:
        """The learned empty prompt used for classifier-free guidance: [BOS] [NULL] [EOS] [PAD]..."""
        cfg = self.text_encoder.config
        row = torch.full((int(cfg.max_tokens),), int(cfg.pad_id), dtype=torch.long)
        if self._is_lm():
            row[0] = int(cfg.null_id)
            return row[None, :].repeat(batch, 1)
        row[0], row[1], row[2] = int(cfg.bos_id), int(cfg.null_id), int(cfg.eos_id)
        return row[None, :].repeat(batch, 1)

    def encode_prompt(self, prompts: List[str], negative_prompts: Optional[List[str]], do_cfg: bool, device):
        ids = self._tokenize(prompts).to(device)
        with self._autocast():
            cond, cond_valid = self.text_encoder(ids)
            if not do_cfg:
                return cond, cond_valid
            if negative_prompts is None:
                uncond_ids = self._null_ids(len(prompts))
            else:
                uncond_ids = self._tokenize(negative_prompts)
            uncond, uncond_valid = self.text_encoder(uncond_ids.to(device))
        # unconditional first, conditional second (the order the model was evaluated with)
        return torch.cat([uncond, cond], dim=0), torch.cat([uncond_valid, cond_valid], dim=0)

    @torch.no_grad()
    def decode_latents(self, latents: torch.Tensor, chunk_size: int = 32) -> torch.Tensor:
        """Model latents (B,128,16,16) -> images in [-1, 1] via the FLUX.2 small decoder."""
        vae = self.vae
        eps = float(getattr(vae.config, "batch_norm_eps", 1e-4))
        mean = vae.bn.running_mean.detach().float().view(1, -1, 1, 1).to(latents.device)
        std = torch.sqrt(vae.bn.running_var.detach().float() + eps).view(1, -1, 1, 1).to(latents.device)
        unpatchify = UNPATCHIFY_LAYOUTS[self.transformer.config.latent_layout]
        outs = []
        for i in range(0, latents.shape[0], chunk_size):
            z = latents[i:i + chunk_size].float() * std + mean
            z = unpatchify(z)
            outs.append(vae.decode(z.to(vae.dtype)).sample.float().clamp(-1, 1))
        return torch.cat(outs, dim=0)

    # ---- main entry point ----------------------------------------------------------------------
    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 3.0,
        num_images_per_prompt: int = 1,
        generator: Optional[torch.Generator] = None,
        latents: Optional[torch.Tensor] = None,
        output_type: str = "pil",
        return_dict: bool = True,
    ):
        """
        Args:
            prompt: text prompt(s). English; prompts longer than 62 BPE tokens are truncated.
            negative_prompt: optional. None (default) uses the learned empty prompt the model was trained
                with for classifier-free guidance. Arbitrary negative prompts were not part of training.
            num_inference_steps: Euler steps. 50 = recommended (used for the benchmarks), 20 = fast.
            guidance_scale: classifier-free guidance scale. 3.0 recommended.
            output_type: "pil", "np", "pt" (tensor in [0,1]) or "latent" (raw model latents).
        """
        base = [prompt] if isinstance(prompt, str) else list(prompt)
        prompts = [p for p in base for _ in range(num_images_per_prompt)]
        negs = None
        if negative_prompt is not None:
            neg_base = [negative_prompt] * len(base) if isinstance(negative_prompt, str) else list(negative_prompt)
            if len(neg_base) != len(base):
                raise ValueError("negative_prompt must be a string or have the same length as prompt.")
            negs = [n for n in neg_base for _ in range(num_images_per_prompt)]

        device = self._execution_device
        b = len(prompts)
        do_cfg = guidance_scale != 1.0
        context, context_valid = self.encode_prompt(prompts, negs, do_cfg, device)

        tcfg = self.transformer.config
        shape = (b, tcfg.latent_channels, tcfg.latent_size, tcfg.latent_size)
        if latents is None:
            z = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
        else:
            z = latents.to(device=device, dtype=torch.float32)

        ts = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device)
        for i in self.progress_bar(range(num_inference_steps)):
            t0, t1 = ts[i], ts[i + 1]
            with self._autocast():
                if do_cfg:
                    v = self.transformer(torch.cat([z, z], dim=0), t0.expand(2 * b), context, context_valid)
                    v_uncond, v_cond = v.chunk(2, dim=0)
                    velocity = v_uncond + guidance_scale * (v_cond - v_uncond)
                else:
                    velocity = self.transformer(z, t0.expand(b), context, context_valid)
            z = z + (t1 - t0) * velocity.float()

        if output_type == "latent":
            images = z
        else:
            img = (self.decode_latents(z) + 1.0) / 2.0
            if output_type == "pt":
                images = img
            else:
                arr = img.cpu().permute(0, 2, 3, 1).float().numpy()
                images = self.numpy_to_pil(arr) if output_type == "pil" else arr

        self.maybe_free_model_hooks()
        if not return_dict:
            return (images,)
        return ImagePipelineOutput(images=images)