from __future__ import annotations

import os
from typing import List, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import einsum, rearrange
from PIL import Image
from safetensors.torch import load_file


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CHECKPOINT_PATH = os.path.join(BASE_DIR, "model", "model.safetensors")

MODEL_CONFIG = {
    "model_type": "image_dit",
    "label_vocab_size": 11,
    "vocab_size": 257,
    "pixel_bins": 256,
    "context_length": 784,
    "d_model": 256,
    "num_layers": 8,
    "num_heads": 16,
    "d_ff": 1024,
    "rope_theta": 10000.0,
    "attention_backend": "torch_sdpa",
    "attention_sdp_backend": "auto",
    "device": "cuda",
    "dtype": "float16",
    "null_label_id": 10,
    "use_rope_2d": True,
    "image_height": 28,
    "image_width": 28,
}

INFER_CONFIG = {
    "steps": 128,
    "cfg_scale": 2.0,
    "trajectory_checkpoints": 32,
}

DTYPES = {
    "float16": torch.float16,
    "float32": torch.float32,
    "bfloat16": torch.bfloat16,
}

ALLOWED_ATTENTION_BACKENDS = {"custom", "torch_sdpa"}
ALLOWED_SDP_BACKENDS = {"auto", "flash", "mem_efficient", "math"}


def _resolve_device_dtype(device: str, dtype_name: str) -> Tuple[str, torch.dtype]:
    resolved_device = device
    if device == "cuda" and not torch.cuda.is_available():
        resolved_device = "cpu"

    resolved_dtype = DTYPES[dtype_name]
    if resolved_device == "cpu" and resolved_dtype == torch.float16:
        resolved_dtype = torch.float32

    return resolved_device, resolved_dtype


def set_sdp_backend(backend: str) -> None:
    backend = backend.lower()
    if backend not in ALLOWED_SDP_BACKENDS:
        raise ValueError(f"attention_sdp_backend must be one of {sorted(ALLOWED_SDP_BACKENDS)}")
    if not torch.cuda.is_available():
        return
    if backend == "auto":
        torch.backends.cuda.enable_flash_sdp(True)
        torch.backends.cuda.enable_mem_efficient_sdp(True)
        torch.backends.cuda.enable_math_sdp(True)
        return
    torch.backends.cuda.enable_flash_sdp(backend == "flash")
    torch.backends.cuda.enable_mem_efficient_sdp(backend == "mem_efficient")
    torch.backends.cuda.enable_math_sdp(backend == "math")


def softmax(x: torch.Tensor, dim: int):
    x_max = x.max(dim=dim, keepdim=True).values
    x_stable = x - x_max
    exp_x = torch.exp(x_stable)
    sum_exp_x = exp_x.sum(dim=dim, keepdim=True)
    return exp_x / sum_exp_x


class Linear(nn.Module):
    def __init__(self, in_features, out_features, device=None, dtype=None):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(out_features, in_features, device=device, dtype=dtype))
        mean = 0.0
        std = 2 / (in_features + out_features)
        a = mean - 3 * std
        b = mean + 3 * std
        nn.init.trunc_normal_(self.weight, mean=mean, std=std, a=a, b=b)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return einsum(self.weight, x, "out_features in_features, ... in_features -> ... out_features")


class Embedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim, device=None, dtype=None):
        super().__init__()
        self.weight = nn.Parameter(torch.empty(num_embeddings, embedding_dim, device=device, dtype=dtype))
        nn.init.trunc_normal_(self.weight, mean=0, std=1, a=-3, b=3)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        return self.weight[token_ids]


class RMSNorm(nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.empty(d_model, device=device, dtype=dtype))
        nn.init.ones_(self.weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        in_dtype = x.dtype
        x = x.to(torch.float32)
        rms = torch.sqrt(torch.mean(x**2, dim=-1) + self.eps).unsqueeze(-1)
        x = (1.0 / rms) * (x * self.weight)
        return x.to(in_dtype)


class SwiGLU(nn.Module):
    def __init__(self, d_model: int, d_ff: int, device=None, dtype=None):
        super().__init__()
        self.w1 = Linear(d_model, d_ff, device=device, dtype=dtype)
        self.w2 = Linear(d_ff, d_model, device=device, dtype=dtype)
        self.w3 = Linear(d_model, d_ff, device=device, dtype=dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w1x = self.w1(x)
        w3x = self.w3(x)
        silu = w1x * torch.sigmoid(w1x)
        return self.w2(silu * w3x)


class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None):
        super().__init__()
        theta_i = theta ** (torch.arange(0, d_k, 2).float() / d_k)
        position = torch.arange(max_seq_len)
        phases = position.unsqueeze(1) / theta_i.unsqueeze(0)
        phases_combined = torch.stack([torch.cos(phases), torch.sin(phases)], dim=-1).to(device=device)
        self.register_buffer("phases", phases_combined, persistent=False)

    def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
        x = rearrange(x, "... (d_k p) -> ... d_k p", p=2)
        x1 = x[..., 0]
        x2 = x[..., 1]
        phases_cos = self.phases[..., 0][token_positions].to(dtype=x.dtype)
        phases_sin = self.phases[..., 1][token_positions].to(dtype=x.dtype)
        x_rotated = torch.stack(
            [
                x1 * phases_cos - x2 * phases_sin,
                x1 * phases_sin + x2 * phases_cos,
            ],
            dim=-1,
        )
        return x_rotated.flatten(-2)


def _prepare_attention_mask(attention_mask: torch.Tensor, ref_tensor: torch.Tensor) -> torch.Tensor:
    mask = attention_mask.to(device=ref_tensor.device, dtype=torch.bool)
    if mask.dim() == 2:
        mask = mask[:, None, None, :]
    elif mask.dim() == 3:
        mask = mask[:, None, :, :]
    elif mask.dim() != 4:
        raise ValueError("attention_mask must be 2D, 3D, or 4D")
    return mask


def scaled_dot_product_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
):
    scale = torch.tensor(q.shape[-1], device=q.device, dtype=q.dtype).sqrt()
    qk_score = einsum(q, k, "batch ... n d, batch ... m d -> batch ... n m") / scale
    if attention_mask is not None:
        mask = _prepare_attention_mask(attention_mask, qk_score)
        qk_score = qk_score.masked_fill(~mask, float("-inf"))
    return einsum(softmax(qk_score, dim=-1), v, "batch ... n m, batch ... m d -> batch ... n d")


def torch_scaled_dot_product_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
):
    mask = None
    if attention_mask is not None:
        mask = _prepare_attention_mask(attention_mask, q)
    return F.scaled_dot_product_attention(
        q.contiguous(),
        k.contiguous(),
        v.contiguous(),
        attn_mask=mask,
        dropout_p=0.0,
        is_causal=False,
    )


class MultiheadSelfAttentionRoPE2D(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        max_height: int,
        max_width: int,
        theta: float,
        attention_backend: str = "custom",
        device=None,
        dtype=None,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = self.d_model // self.num_heads
        if self.d_k % 4 != 0:
            raise ValueError("per-head dimension must be divisible by 4 for 2D RoPE")
        self.d_v = self.d_k
        if attention_backend not in ALLOWED_ATTENTION_BACKENDS:
            raise ValueError(f"attention_backend must be one of {sorted(ALLOWED_ATTENTION_BACKENDS)}")
        self.attention_backend = attention_backend
        self.q_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.k_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.v_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.output_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.d_k_half = self.d_k // 2
        self.row_rope = RotaryPositionalEmbedding(theta, self.d_k_half, int(max_height), device)
        self.col_rope = RotaryPositionalEmbedding(theta, self.d_k_half, int(max_width), device)

    def _apply_2d_rope(self, x: torch.Tensor, row_positions: torch.Tensor, col_positions: torch.Tensor) -> torch.Tensor:
        row_part = x[..., : self.d_k_half]
        col_part = x[..., self.d_k_half :]
        return torch.cat(
            [
                self.row_rope(row_part, row_positions),
                self.col_rope(col_part, col_positions),
            ],
            dim=-1,
        )

    def forward(
        self,
        x: torch.Tensor,
        row_positions: torch.Tensor,
        col_positions: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        wqx = rearrange(self.q_proj(x), "... seq (heads d) -> ... heads seq d", heads=self.num_heads, d=self.d_k)
        wkx = rearrange(self.k_proj(x), "... seq (heads d) -> ... heads seq d", heads=self.num_heads, d=self.d_k)
        wvx = rearrange(self.v_proj(x), "... seq (heads d) -> ... heads seq d", heads=self.num_heads, d=self.d_v)
        q = self._apply_2d_rope(wqx, row_positions, col_positions)
        k = self._apply_2d_rope(wkx, row_positions, col_positions)
        if self.attention_backend == "torch_sdpa":
            attn = torch_scaled_dot_product_attention(q, k, wvx, attention_mask=attention_mask)
        else:
            attn = scaled_dot_product_attention(q, k, wvx, attention_mask=attention_mask)
        out = rearrange(attn, "... heads seq d -> ... seq (heads d)", heads=self.num_heads, d=self.d_v)
        return self.output_proj(out)


class MultiheadCrossAttentionRoPE2D(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        max_height: int,
        max_width: int,
        theta: float,
        attention_backend: str = "custom",
        device=None,
        dtype=None,
    ):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = self.d_model // self.num_heads
        if self.d_k % 4 != 0:
            raise ValueError("per-head dimension must be divisible by 4 for 2D RoPE")
        self.d_v = self.d_k
        if attention_backend not in ALLOWED_ATTENTION_BACKENDS:
            raise ValueError(f"attention_backend must be one of {sorted(ALLOWED_ATTENTION_BACKENDS)}")
        self.attention_backend = attention_backend
        self.q_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.k_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.v_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.output_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.d_k_half = self.d_k // 2
        self.row_rope = RotaryPositionalEmbedding(theta, self.d_k_half, int(max_height), device)
        self.col_rope = RotaryPositionalEmbedding(theta, self.d_k_half, int(max_width), device)

    def _apply_2d_rope(self, x: torch.Tensor, row_positions: torch.Tensor, col_positions: torch.Tensor) -> torch.Tensor:
        row_part = x[..., : self.d_k_half]
        col_part = x[..., self.d_k_half :]
        return torch.cat(
            [
                self.row_rope(row_part, row_positions),
                self.col_rope(col_part, col_positions),
            ],
            dim=-1,
        )

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor,
        row_positions: torch.Tensor,
        col_positions: torch.Tensor,
        context_row_positions: torch.Tensor,
        context_col_positions: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        wqx = rearrange(self.q_proj(x), "... seq (heads d) -> ... heads seq d", heads=self.num_heads, d=self.d_k)
        wkx = rearrange(self.k_proj(context), "... seq (heads d) -> ... heads seq d", heads=self.num_heads, d=self.d_k)
        wvx = rearrange(self.v_proj(context), "... seq (heads d) -> ... heads seq d", heads=self.num_heads, d=self.d_v)
        q = self._apply_2d_rope(wqx, row_positions, col_positions)
        k = self._apply_2d_rope(wkx, context_row_positions, context_col_positions)
        if self.attention_backend == "torch_sdpa":
            attn = torch_scaled_dot_product_attention(q, k, wvx, attention_mask=attention_mask)
        else:
            attn = scaled_dot_product_attention(q, k, wvx, attention_mask=attention_mask)
        out = rearrange(attn, "... heads seq d -> ... seq (heads d)", heads=self.num_heads, d=self.d_v)
        return self.output_proj(out)


class TransformerImageBlock(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        max_seq_len: int,
        max_height: int | None,
        max_width: int | None,
        theta: float,
        d_ff: int,
        attention_backend: str = "custom",
        use_rope_2d: bool = False,
        device=None,
        dtype=None,
    ):
        super().__init__()
        self.ffn = SwiGLU(d_model, d_ff, device, dtype)
        self.use_rope_2d = bool(use_rope_2d)
        if not self.use_rope_2d:
            raise ValueError("This demo vendors only the 2D RoPE image path")
        if max_height is None or max_width is None:
            raise ValueError("max_height/max_width must be provided when use_rope_2d is True")
        self.self_attn = MultiheadSelfAttentionRoPE2D(
            d_model,
            num_heads,
            max_height,
            max_width,
            theta,
            attention_backend=attention_backend,
            device=device,
            dtype=dtype,
        )
        self.cross_attn = MultiheadCrossAttentionRoPE2D(
            d_model,
            num_heads,
            max_height,
            max_width,
            theta,
            attention_backend=attention_backend,
            device=device,
            dtype=dtype,
        )
        self.ln1 = RMSNorm(d_model, device=device, dtype=dtype)
        self.ln2 = RMSNorm(d_model, device=device, dtype=dtype)
        self.ln3 = RMSNorm(d_model, device=device, dtype=dtype)

    def forward(
        self,
        x: torch.Tensor,
        context: torch.Tensor,
        row_positions: torch.Tensor,
        col_positions: torch.Tensor,
        context_row_positions: torch.Tensor,
        context_col_positions: torch.Tensor,
    ) -> torch.Tensor:
        x = x + self.self_attn(self.ln1(x), row_positions, col_positions, attention_mask=None)
        x = x + self.cross_attn(
            self.ln2(x),
            context,
            row_positions,
            col_positions,
            context_row_positions,
            context_col_positions,
            attention_mask=None,
        )
        x = x + self.ffn(self.ln3(x))
        return x


def _timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0) -> torch.Tensor:
    if t.dim() != 1:
        raise ValueError("t must be 1D with shape (batch,)")
    half = dim // 2
    if half == 0:
        return t[:, None]
    freqs = torch.exp(
        -torch.log(torch.tensor(max_period, device=t.device, dtype=torch.float32))
        * torch.arange(half, device=t.device, dtype=torch.float32)
        / max(half - 1, 1)
    )
    args = t.to(torch.float32)[:, None] * freqs[None, :]
    emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
    if dim % 2 == 1:
        emb = torch.cat([emb, torch.zeros((t.shape[0], 1), device=t.device, dtype=emb.dtype)], dim=-1)
    return emb


class DiTImage(nn.Module):
    def __init__(
        self,
        context_length: int,
        d_model: int,
        num_layers: int,
        num_heads: int,
        d_ff: int,
        rope_theta: float,
        label_vocab_size: int,
        attention_backend: str = "custom",
        image_height: int | None = None,
        image_width: int | None = None,
        use_rope_2d: bool = False,
        device=None,
        dtype=None,
    ):
        super().__init__()
        self.context_length = int(context_length)
        self.use_rope_2d = bool(use_rope_2d)
        if not self.use_rope_2d:
            raise ValueError("This demo expects use_rope_2d=True")
        if image_height is None or image_width is None:
            raise ValueError("image_height/image_width must be set for the flow demo")
        self.image_height = int(image_height)
        self.image_width = int(image_width)
        self.input_proj = Linear(1, d_model, device, dtype)
        self.time_proj = Linear(d_model, d_model, device, dtype)
        self.label_embeddings = Embedding(label_vocab_size, d_model, device, dtype)
        self.layers = nn.ModuleList(
            [
                TransformerImageBlock(
                    d_model,
                    num_heads,
                    context_length,
                    self.image_height,
                    self.image_width,
                    rope_theta,
                    d_ff,
                    attention_backend=attention_backend,
                    use_rope_2d=True,
                    device=device,
                    dtype=dtype,
                )
                for _ in range(num_layers)
            ]
        )
        self.ln_final = RMSNorm(d_model, device=device, dtype=dtype)
        self.output_proj = Linear(d_model, 1, device, dtype)

    def forward(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor:
        if x.dim() != 2:
            raise ValueError("x must be 2D with shape (batch, seq)")
        if context is None or context.dim() != 1 or context.shape[0] != x.shape[0]:
            raise ValueError("context must be 1D with matching batch size")
        if t.dim() == 2 and t.shape[1] == 1:
            t = t[:, 0]
        if t.dim() != 1 or t.shape[0] != x.shape[0]:
            raise ValueError("t must be 1D with matching batch size")

        model_dtype = self.input_proj.weight.dtype
        output_seq = self.input_proj(x.to(dtype=model_dtype).unsqueeze(-1))
        t_emb = _timestep_embedding(t, output_seq.shape[-1]).to(dtype=model_dtype)
        context_emb = (self.time_proj(t_emb) + self.label_embeddings(context)).unsqueeze(-2)

        seq_len = output_seq.shape[-2]
        expected = self.image_height * self.image_width
        if seq_len != expected:
            raise ValueError(f"sequence length {seq_len} does not match image_height*image_width {expected}")
        row_positions = torch.arange(self.image_height, device=output_seq.device, dtype=torch.long).repeat_interleave(
            self.image_width
        )
        col_positions = torch.arange(self.image_width, device=output_seq.device, dtype=torch.long).repeat(
            self.image_height
        )
        context_row_positions = torch.zeros(context_emb.shape[-2], device=output_seq.device, dtype=torch.long)
        context_col_positions = torch.zeros(context_emb.shape[-2], device=output_seq.device, dtype=torch.long)

        for layer in self.layers:
            output_seq = layer(
                output_seq,
                context_emb,
                row_positions,
                col_positions,
                context_row_positions,
                context_col_positions,
            )
        return self.output_proj(self.ln_final(output_seq)).squeeze(-1)


@torch.no_grad()
def flow_image_generate(
    model,
    prompt_indices: torch.Tensor,
    *,
    context: torch.Tensor,
    steps: int,
    cfg_scale: float = 0.0,
    uncond_context: torch.Tensor | None = None,
    generator: torch.Generator | None = None,
    return_history: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, list[tuple[int, torch.Tensor]]]:
    if prompt_indices.dim() != 2:
        raise ValueError("prompt_indices must be 2D (batch, seq)")
    if context.dim() != 1 or prompt_indices.shape[0] != context.shape[0]:
        raise ValueError("context must be 1D with matching batch size")
    if prompt_indices.shape[1] != 0:
        raise ValueError("flow_image_generate expects empty prompt_indices for full-image generation")
    steps = max(1, min(int(steps), 128))

    batch_size = context.shape[0]
    gen_length = int(model.context_length)
    x = torch.randn((batch_size, gen_length), device=prompt_indices.device, dtype=torch.float32, generator=generator)
    dt = 1.0 / float(steps)

    if uncond_context is not None:
        if uncond_context.dim() != 1 or uncond_context.shape[0] != batch_size:
            raise ValueError("uncond_context must be 1D with matching batch size")
        uncond_context = uncond_context.to(device=context.device, dtype=context.dtype)

    history: list[tuple[int, torch.Tensor]] = []
    checkpoint_count = max(32, int(INFER_CONFIG["trajectory_checkpoints"]))
    checkpoint_indices = np.linspace(1, steps, num=checkpoint_count, dtype=int).tolist()
    checkpoint_indices = sorted(set(max(1, min(steps, idx)) for idx in checkpoint_indices))

    for k in range(steps):
        t = torch.full((batch_size,), float(k) / float(steps), device=x.device, dtype=x.dtype)
        if cfg_scale > 0.0:
            if uncond_context is None:
                raise ValueError("uncond_context must be set when cfg_scale > 0 for flow_image_generate")
            v_cond = model(x, t, context=context)
            v_uncond = model(x, t, context=uncond_context)
            v = v_uncond + (cfg_scale + 1.0) * (v_cond - v_uncond)
        else:
            v = model(x, t, context=context)
        x = x + dt * v
        step_idx = k + 1
        if return_history and step_idx in checkpoint_indices:
            history.append((step_idx, x.detach().clone()))
    if return_history:
        if not history or history[-1][0] != steps:
            history.append((steps, x.detach().clone()))
        return x, history
    return x


def flow_pixels_to_uint8(values: np.ndarray) -> np.ndarray:
    clipped = np.clip(values.astype(np.float32), -1.0, 1.0)
    restored = np.round((clipped + 1.0) * 127.5)
    return np.clip(restored, 0, 255).astype(np.uint8)


MODEL = None
DEVICE = None
DTYPE = None


def load_model():
    global MODEL, DEVICE, DTYPE
    if MODEL is not None:
        return MODEL, DEVICE, DTYPE
    if not os.path.exists(CHECKPOINT_PATH):
        raise FileNotFoundError(f"Missing checkpoint at {CHECKPOINT_PATH}")

    device, dtype = _resolve_device_dtype(MODEL_CONFIG["device"], MODEL_CONFIG["dtype"])
    set_sdp_backend(MODEL_CONFIG["attention_sdp_backend"])

    model = DiTImage(
        context_length=MODEL_CONFIG["context_length"],
        d_model=MODEL_CONFIG["d_model"],
        num_layers=MODEL_CONFIG["num_layers"],
        num_heads=MODEL_CONFIG["num_heads"],
        d_ff=MODEL_CONFIG["d_ff"],
        rope_theta=MODEL_CONFIG["rope_theta"],
        label_vocab_size=MODEL_CONFIG["label_vocab_size"],
        attention_backend=MODEL_CONFIG["attention_backend"],
        image_height=MODEL_CONFIG["image_height"],
        image_width=MODEL_CONFIG["image_width"],
        use_rope_2d=MODEL_CONFIG["use_rope_2d"],
        device=device,
        dtype=dtype,
    )
    model.load_state_dict(load_file(CHECKPOINT_PATH))
    model.eval().to(device)

    MODEL = model
    DEVICE = device
    DTYPE = dtype
    return MODEL, DEVICE, DTYPE


def _to_image(sample: torch.Tensor) -> Image.Image:
    h = int(MODEL_CONFIG["image_height"])
    w = int(MODEL_CONFIG["image_width"])
    scale = 10
    arr = sample.detach().cpu().to(torch.float32).numpy().reshape(h, w)
    img = Image.fromarray(flow_pixels_to_uint8(arr), mode="L")
    if scale > 1:
        img = img.resize((w * scale, h * scale), resample=Image.NEAREST)
    return img


@torch.inference_mode()
def generate_images(label: int, steps: int, num_samples: int) -> List[Image.Image]:
    model, device, _ = load_model()
    num_samples = int(num_samples)
    label = int(label)
    steps = max(1, min(int(steps), 128))

    context = torch.full((num_samples,), label, device=device, dtype=torch.long)
    prompt = torch.empty((num_samples, 0), device=device, dtype=torch.long)
    cfg_scale = float(INFER_CONFIG["cfg_scale"])
    null_label_id = int(MODEL_CONFIG["null_label_id"])
    uncond_context = torch.full((num_samples,), null_label_id, device=device, dtype=torch.long)

    out = flow_image_generate(
        model,
        prompt,
        context=context,
        steps=steps,
        cfg_scale=cfg_scale,
        uncond_context=uncond_context,
        generator=None,
    )

    images: List[Image.Image] = []
    for i in range(num_samples):
        images.append(_to_image(out[i]))
    return images


def _grid_dims(num_samples: int) -> Tuple[int, int]:
    cols = int(np.ceil(np.sqrt(num_samples)))
    rows = int(np.ceil(num_samples / cols))
    return rows, cols


@torch.inference_mode()
def generate_grid_image(label: int, steps: int, num_samples: int) -> Image.Image:
    images = generate_images(label=label, steps=steps, num_samples=num_samples)
    if not images:
        return Image.new("L", (1, 1), color=0)
    rows, cols = _grid_dims(len(images))
    w, h = images[0].size
    grid = Image.new("L", (cols * w, rows * h))
    for idx, img in enumerate(images):
        r = idx // cols
        c = idx % cols
        grid.paste(img, (c * w, r * h))
    return grid


@torch.inference_mode()
def iter_trajectory_frames(label: int, steps: int):
    model, device, _ = load_model()
    steps = max(32, min(int(steps), 128))
    context = torch.full((1,), int(label), device=device, dtype=torch.long)
    prompt = torch.empty((1, 0), device=device, dtype=torch.long)
    null_label_id = int(MODEL_CONFIG["null_label_id"])
    uncond_context = torch.full((1,), null_label_id, device=device, dtype=torch.long)

    batch_size = context.shape[0]
    gen_length = int(model.context_length)
    x = torch.randn((batch_size, gen_length), device=prompt.device, dtype=torch.float32)
    dt = 1.0 / float(steps)

    checkpoint_count = max(16, int(INFER_CONFIG["trajectory_checkpoints"]))
    checkpoint_indices = np.linspace(1, steps, num=checkpoint_count, dtype=int).tolist()
    checkpoint_indices = sorted(set(max(1, min(steps, idx)) for idx in checkpoint_indices))

    cfg_scale = float(INFER_CONFIG["cfg_scale"])
    for k in range(steps):
        t = torch.full((batch_size,), float(k) / float(steps), device=x.device, dtype=x.dtype)
        if cfg_scale > 0.0:
            v_cond = model(x, t, context=context)
            v_uncond = model(x, t, context=uncond_context)
            v = v_uncond + (cfg_scale + 1.0) * (v_cond - v_uncond)
        else:
            v = model(x, t, context=context)
        x = x + dt * v
        step_idx = k + 1
        if step_idx in checkpoint_indices:
            yield _to_image(x[0]), step_idx, steps, len(checkpoint_indices)