add: full app

.gitattributes

@@ -33,4 +33,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-model/*.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,14 @@
 ---
-title: Mnist Flow Demo
-emoji: 👁
-colorFrom: gray
-colorTo: pink
+title: MNIST Flow (TransformerLM)
+emoji: 🌊
+colorFrom: blue
+colorTo: cyan
 sdk: gradio
-sdk_version: 6.9.0
+sdk_version: 6.5.1
 python_version: '3.12'
 app_file: app.py
 pinned: false
+short_description: Flow-matching MNIST digit generation
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,49 @@
+import gradio as gr
+import spaces
+import torch
+
+from model import generate_grid_image, load_model
+
+MODEL_READY = False
+
+
+def ensure_model_loaded():
+    global MODEL_READY
+    if not MODEL_READY:
+        load_model()
+        MODEL_READY = True
+
+
+@spaces.GPU
+@torch.inference_mode()
+def predict(label: int, steps: int, num_samples: int):
+    ensure_model_loaded()
+    return generate_grid_image(label=label, steps=steps, num_samples=num_samples)
+
+
+with gr.Blocks(title="MNIST Flow") as demo:
+    gr.Markdown("# MNIST Flow")
+    gr.Markdown(
+        "Flow-matching DiT model for MNIST digits. "
+        "Sampling uses fixed CFG=2.0 and caps ODE steps at 128."
+    )
+
+    grid = gr.Image(label="Samples", show_label=True)
+
+    with gr.Row():
+        label = gr.Dropdown([str(i) for i in range(10)], value="1", label="Label")
+        steps = gr.Slider(1, 128, value=128, step=1, label="Steps")
+        num_samples = gr.Slider(1, 36, value=16, step=1, label="Samples")
+
+    generate_btn = gr.Button("Generate")
+
+    generate_btn.click(
+        fn=predict,
+        inputs=[label, steps, num_samples],
+        outputs=grid,
+        scroll_to_output=True,
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

model.py

@@ -0,0 +1,656 @@
+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,
+}
+
+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,
+) -> 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)
+
+    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
+    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
+
+
+@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,
+    )
+
+    h = int(MODEL_CONFIG["image_height"])
+    w = int(MODEL_CONFIG["image_width"])
+    images: List[Image.Image] = []
+    scale = 10
+    for i in range(num_samples):
+        arr = out[i].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)
+        images.append(img)
+    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

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+spaces
+torch
+einops
+safetensors
+numpy
+pillow