Add PixText transformer and attention modules

vendor/pixtext/modeling/attention.py

@@ -0,0 +1,541 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+import torch
+from torch import nn
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
+
+from pixtext.modeling.conditioning import (
+    IMAGE_TOKEN_TYPE,
+    PROMPT_TOKEN_TYPE,
+    TIME_TOKEN_TYPE,
+    PackedGroupSequence,
+)
+from pixtext.modeling.modes import ImageAttentionRoutingMode, TimestepTokenRoutingMode
+from pixtext.modeling.patches import PackedImagePatches
+from pixtext.modeling.rope import (
+    DEFAULT_ROPE_AXES_DIMS,
+    DEFAULT_ROPE_THETA,
+    RotaryEmbedding3D,
+    _apply_rotary_factors,
+)
+
+_compiled_flex_attention = torch.compile(flex_attention, dynamic=True)
+_FORCE_STRICT_FP32_FLEX_ATTENTION = ContextVar("force_strict_fp32_flex_attention", default=False)
+_FLEX_ATTENTION_FORWARD_NUM_STAGES = ContextVar("flex_attention_forward_num_stages", default=None)
+_INTERREAD = ImageAttentionRoutingMode.ITEMS_INTERREAD_GLYPH_READS_OWN_ITEM
+_INTERREAD_BIDIRECTIONAL_MAIN = (
+    ImageAttentionRoutingMode.ITEMS_INTERREAD_GLYPH_READS_OWN_ITEM_BIDIRECTIONAL_MAIN
+)
+
+
+@contextmanager
+def force_strict_fp32_flex_attention():
+    """Keep FP32 FlexAttention strict inside a broader TF32 forward policy."""
+    token = _FORCE_STRICT_FP32_FLEX_ATTENTION.set(True)
+    try:
+        yield
+    finally:
+        _FORCE_STRICT_FP32_FLEX_ATTENTION.reset(token)
+
+
+@contextmanager
+def use_flex_attention_forward_num_stages(num_stages: int):
+    """Select a scoped Triton forward pipeline depth for FlexAttention."""
+    if isinstance(num_stages, bool) or not isinstance(num_stages, int) or num_stages <= 0:
+        raise ValueError("FlexAttention num_stages must be a positive integer")
+    token = _FLEX_ATTENTION_FORWARD_NUM_STAGES.set(num_stages)
+    try:
+        yield
+    finally:
+        _FLEX_ATTENTION_FORWARD_NUM_STAGES.reset(token)
+
+
+def _normalize_and_apply_rotary(
+    values: torch.Tensor, weight: torch.Tensor, eps: float, cosine: torch.Tensor, sine: torch.Tensor
+) -> torch.Tensor:
+    """Run Q/K RMSNorm and precomputed RoPE as one compilable subgraph."""
+    output_dtype = values.dtype
+    variance = values.float().square().mean(dim=-1, keepdim=True)
+    normalized = values.float() * torch.rsqrt(variance + eps)
+    normalized = (normalized * weight.float()).to(output_dtype)
+    return _apply_rotary_factors(normalized, cosine, sine)
+
+
+_compiled_normalize_and_apply_rotary = torch.compile(_normalize_and_apply_rotary, dynamic=True)
+
+
+class _RMSNorm(nn.Module):
+    """Last-dimension RMSNorm with FP32 reduction statistics."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-06) -> None:
+        super().__init__()
+        self.eps = float(eps)
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        output_dtype = values.dtype
+        variance = values.float().square().mean(dim=-1, keepdim=True)
+        normalized = values.float() * torch.rsqrt(variance + self.eps)
+        return (normalized * self.weight.float()).to(output_dtype)
+
+
+def _validate_routing(
+    *,
+    sample_ids: torch.Tensor,
+    image_ids: torch.Tensor,
+    token_types: torch.Tensor,
+    timestep_roles: torch.Tensor,
+    branch_kinds: torch.Tensor,
+    slot_ids: torch.Tensor,
+    token_count: int,
+    device: torch.device,
+) -> None:
+    for name, values in (
+        ("sample_ids", sample_ids),
+        ("image_ids", image_ids),
+        ("token_types", token_types),
+        ("timestep_roles", timestep_roles),
+        ("branch_kinds", branch_kinds),
+        ("slot_ids", slot_ids),
+    ):
+        if not isinstance(values, torch.Tensor):
+            raise TypeError(f"routing {name} must be a torch.Tensor")
+        if values.shape != (token_count,):
+            raise ValueError(
+                f"routing {name} must contain {token_count} entries, got shape {tuple(values.shape)}"
+            )
+        if values.dtype != torch.int64:
+            raise ValueError(f"routing {name} must use torch.int64, got {values.dtype}")
+        if values.device != device:
+            raise ValueError(f"routing {name} must be on {device}, got {values.device}")
+    valid_token_types = (
+        (token_types == IMAGE_TOKEN_TYPE)
+        | (token_types == PROMPT_TOKEN_TYPE)
+        | (token_types == TIME_TOKEN_TYPE)
+    )
+    if not bool(valid_token_types.all()):
+        raise ValueError("routing token_types may only contain image, prompt, or timestep IDs")
+    valid_timestep_roles = (timestep_roles == 0) | (timestep_roles == 1)
+    if not bool(valid_timestep_roles.all()):
+        raise ValueError("routing timestep_roles may only contain 0 or 1")
+    if not bool(((branch_kinds >= -1) & (branch_kinds <= 2)).all()):
+        raise ValueError("routing branch_kinds may only contain -1, 0, 1, or 2")
+    if not bool((slot_ids >= -1).all()):
+        raise ValueError("routing slot_ids must be at least -1")
+
+
+def _validated_routing(
+    packed: PackedGroupSequence | PackedImagePatches,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    int,
+    torch.device,
+    TimestepTokenRoutingMode | None,
+    torch.Tensor,
+    torch.Tensor,
+    bool,
+]:
+    if isinstance(packed, PackedGroupSequence):
+        sample_ids = packed.sample_ids
+        image_ids = packed.image_ids
+        token_types = packed.token_types
+        timestep_roles = packed.timestep_roles
+        token_count = packed.tokens.shape[0]
+        device = packed.tokens.device
+        timestep_token_routing = packed.timestep_token_routing
+        explicit_branches = packed.branch_kinds is not None
+        if explicit_branches:
+            assert packed.branch_kinds is not None
+            assert packed.slot_ids is not None
+            branch_kinds = packed.branch_kinds
+            slot_ids = packed.slot_ids
+        else:
+            branch_kinds = torch.where(
+                image_ids.eq(0), torch.zeros_like(image_ids), torch.ones_like(image_ids)
+            )
+            slot_ids = torch.where(image_ids.gt(0), image_ids - 1, -1)
+    elif isinstance(packed, PackedImagePatches):
+        sample_ids = packed.sample_ids
+        image_ids = packed.image_ids
+        token_types = torch.full_like(image_ids, IMAGE_TOKEN_TYPE)
+        timestep_roles = image_ids.gt(0).to(torch.int64)
+        token_count = packed.patches.shape[0]
+        device = packed.patches.device
+        timestep_token_routing = None
+        explicit_branches = False
+        branch_kinds = torch.where(
+            image_ids.eq(0), torch.zeros_like(image_ids), torch.ones_like(image_ids)
+        )
+        slot_ids = torch.where(image_ids.gt(0), image_ids - 1, -1)
+    else:
+        raise TypeError(
+            f"packed must be PackedGroupSequence or PackedImagePatches, got {type(packed).__name__}"
+        )
+    _validate_routing(
+        sample_ids=sample_ids,
+        image_ids=image_ids,
+        token_types=token_types,
+        timestep_roles=timestep_roles,
+        branch_kinds=branch_kinds,
+        slot_ids=slot_ids,
+        token_count=token_count,
+        device=device,
+    )
+    contains_time_tokens = bool((token_types == TIME_TOKEN_TYPE).any())
+    if contains_time_tokens != (timestep_token_routing is not None):
+        raise ValueError(
+            "timestep token routing must be present exactly when timestep tokens are packed"
+        )
+    if timestep_token_routing is not None:
+        TimestepTokenRoutingMode.require(timestep_token_routing)
+        if timestep_token_routing is TimestepTokenRoutingMode.PER_IMAGE_CONTEXTUAL and bool(
+            (image_ids[token_types == TIME_TOKEN_TYPE] < 0).any()
+        ):
+            raise ValueError("per-image contextual timestep tokens require image IDs")
+    return (
+        sample_ids,
+        image_ids,
+        token_types,
+        timestep_roles,
+        token_count,
+        device,
+        timestep_token_routing,
+        branch_kinds,
+        slot_ids,
+        explicit_branches,
+    )
+
+
+def _allowed_image_pair(
+    *,
+    query_image_ids: torch.Tensor,
+    key_image_ids: torch.Tensor,
+    query_token_types: torch.Tensor,
+    key_token_types: torch.Tensor,
+    image_attention_routing: ImageAttentionRoutingMode,
+) -> torch.Tensor:
+    allowed = (query_token_types == IMAGE_TOKEN_TYPE) & (key_token_types == IMAGE_TOKEN_TYPE)
+    if image_attention_routing is ImageAttentionRoutingMode.ALL_IMAGES:
+        return allowed
+    if image_attention_routing is ImageAttentionRoutingMode.MAIN_IMAGE_HUB:
+        return allowed & ((query_image_ids == 0) | (key_image_ids == 0))
+    if image_attention_routing is ImageAttentionRoutingMode.ITEMS_READ_MAIN:
+        return allowed & (query_image_ids > 0) & (key_image_ids == 0)
+    if image_attention_routing in (_INTERREAD, _INTERREAD_BIDIRECTIONAL_MAIN):
+        if image_attention_routing is _INTERREAD_BIDIRECTIONAL_MAIN:
+            return allowed & ((query_image_ids > 0) | (key_image_ids > 0))
+        return allowed & (query_image_ids > 0)
+    raise AssertionError(f"unhandled image attention routing: {image_attention_routing!r}")
+
+
+def _allowed_routing_pair(
+    *,
+    query_image_ids: torch.Tensor,
+    key_image_ids: torch.Tensor,
+    query_token_types: torch.Tensor,
+    key_token_types: torch.Tensor,
+    query_timestep_roles: torch.Tensor,
+    key_timestep_roles: torch.Tensor,
+    image_attention_routing: ImageAttentionRoutingMode,
+    per_image_contextual_time: bool,
+    query_branch_kinds: torch.Tensor,
+    key_branch_kinds: torch.Tensor,
+    query_slot_ids: torch.Tensor,
+    key_slot_ids: torch.Tensor,
+    explicit_branches: bool,
+) -> torch.Tensor:
+    query_is_time = query_token_types == TIME_TOKEN_TYPE
+    key_is_time = key_token_types == TIME_TOKEN_TYPE
+    content_pair = ~query_is_time & ~key_is_time
+    same_image = query_image_ids == key_image_ids
+    if explicit_branches:
+        query_is_image = query_token_types == IMAGE_TOKEN_TYPE
+        key_is_image = key_token_types == IMAGE_TOKEN_TYPE
+        image_pair = query_is_image & key_is_image
+        if image_attention_routing is ImageAttentionRoutingMode.ALL_IMAGES:
+            allowed_image_pair = image_pair
+        elif image_attention_routing is ImageAttentionRoutingMode.MAIN_IMAGE_HUB:
+            allowed_image_pair = image_pair & (query_branch_kinds.eq(0) | key_branch_kinds.eq(0))
+        elif image_attention_routing is ImageAttentionRoutingMode.ITEMS_READ_MAIN:
+            item_reads_main = image_pair & query_branch_kinds.eq(1) & key_branch_kinds.eq(0)
+            item_reads_own_glyph = (
+                image_pair
+                & query_branch_kinds.eq(1)
+                & key_branch_kinds.eq(2)
+                & query_slot_ids.eq(key_slot_ids)
+            )
+            allowed_image_pair = item_reads_main | item_reads_own_glyph
+        elif image_attention_routing in (_INTERREAD, _INTERREAD_BIDIRECTIONAL_MAIN):
+            item_reads_main = image_pair & query_branch_kinds.eq(1) & key_branch_kinds.eq(0)
+            item_reads_items = image_pair & query_branch_kinds.eq(1) & key_branch_kinds.eq(1)
+            item_reads_own_glyph = (
+                image_pair
+                & query_branch_kinds.eq(1)
+                & key_branch_kinds.eq(2)
+                & query_slot_ids.eq(key_slot_ids)
+            )
+            glyph_reads_own_item = (
+                image_pair
+                & query_branch_kinds.eq(2)
+                & key_branch_kinds.eq(1)
+                & query_slot_ids.eq(key_slot_ids)
+            )
+            allowed_image_pair = (
+                item_reads_main | item_reads_items | item_reads_own_glyph | glyph_reads_own_item
+            )
+            if image_attention_routing is _INTERREAD_BIDIRECTIONAL_MAIN:
+                allowed_image_pair = allowed_image_pair | image_pair & query_branch_kinds.eq(
+                    0
+                ) & key_branch_kinds.eq(1)
+        else:
+            raise AssertionError(f"unhandled image attention routing: {image_attention_routing!r}")
+    else:
+        allowed_image_pair = _allowed_image_pair(
+            query_image_ids=query_image_ids,
+            key_image_ids=key_image_ids,
+            query_token_types=query_token_types,
+            key_token_types=key_token_types,
+            image_attention_routing=image_attention_routing,
+        )
+    if per_image_contextual_time:
+        reads_own_time_bank = key_is_time & same_image
+        time_reads_own_image = query_is_time & (key_token_types == IMAGE_TOKEN_TYPE) & same_image
+    else:
+        same_timestep_role = query_timestep_roles == key_timestep_roles
+        reads_own_time_bank = key_is_time & same_timestep_role
+        time_reads_own_image = torch.zeros_like(reads_own_time_bank)
+    return (
+        content_pair & (same_image | allowed_image_pair)
+        | reads_own_time_bank
+        | time_reads_own_image
+    )
+
+
+def build_dense_attention_mask(
+    packed: PackedGroupSequence | PackedImagePatches,
+    *,
+    image_attention_routing: ImageAttentionRoutingMode = ImageAttentionRoutingMode.ITEMS_INTERREAD_GLYPH_READS_OWN_ITEM_BIDIRECTIONAL_MAIN,
+) -> torch.Tensor:
+    """Return a ``[query, key]`` bool mask where ``True`` means allowed."""
+    routing = ImageAttentionRoutingMode.require(image_attention_routing)
+    (
+        sample_ids,
+        image_ids,
+        token_types,
+        timestep_roles,
+        _,
+        _,
+        timestep_token_routing,
+        branch_kinds,
+        slot_ids,
+        explicit_branches,
+    ) = _validated_routing(packed)
+    same_group = sample_ids[:, None] == sample_ids[None, :]
+    allowed_pair = _allowed_routing_pair(
+        query_image_ids=image_ids[:, None],
+        key_image_ids=image_ids[None, :],
+        query_token_types=token_types[:, None],
+        key_token_types=token_types[None, :],
+        query_timestep_roles=timestep_roles[:, None],
+        key_timestep_roles=timestep_roles[None, :],
+        image_attention_routing=routing,
+        per_image_contextual_time=timestep_token_routing
+        is TimestepTokenRoutingMode.PER_IMAGE_CONTEXTUAL,
+        query_branch_kinds=branch_kinds[:, None],
+        key_branch_kinds=branch_kinds[None, :],
+        query_slot_ids=slot_ids[:, None],
+        key_slot_ids=slot_ids[None, :],
+        explicit_branches=explicit_branches,
+    )
+    return same_group & allowed_pair
+
+
+def build_flex_attention_block_mask(
+    packed: PackedGroupSequence | PackedImagePatches,
+    *,
+    image_attention_routing: ImageAttentionRoutingMode = ImageAttentionRoutingMode.ITEMS_INTERREAD_GLYPH_READS_OWN_ITEM_BIDIRECTIONAL_MAIN,
+) -> BlockMask:
+    """Build the packed routing mask consumed by PyTorch FlexAttention."""
+    routing = ImageAttentionRoutingMode.require(image_attention_routing)
+    (
+        sample_ids,
+        image_ids,
+        token_types,
+        timestep_roles,
+        token_count,
+        device,
+        timestep_token_routing,
+        branch_kinds,
+        slot_ids,
+        explicit_branches,
+    ) = _validated_routing(packed)
+    if token_count == 0:
+        raise ValueError("FlexAttention requires at least one packed token")
+    per_image_contextual_time = (
+        timestep_token_routing is TimestepTokenRoutingMode.PER_IMAGE_CONTEXTUAL
+    )
+
+    def mask_mod(
+        batch_index: torch.Tensor,
+        head_index: torch.Tensor,
+        query_index: torch.Tensor,
+        key_index: torch.Tensor,
+    ) -> torch.Tensor:
+        del batch_index, head_index
+        query_in_bounds = query_index < token_count
+        key_in_bounds = key_index < token_count
+        safe_query_index = query_index.clamp(max=token_count - 1)
+        safe_key_index = key_index.clamp(max=token_count - 1)
+        same_group = sample_ids[safe_query_index] == sample_ids[safe_key_index]
+        allowed_pair = _allowed_routing_pair(
+            query_image_ids=image_ids[safe_query_index],
+            key_image_ids=image_ids[safe_key_index],
+            query_token_types=token_types[safe_query_index],
+            key_token_types=token_types[safe_key_index],
+            query_timestep_roles=timestep_roles[safe_query_index],
+            key_timestep_roles=timestep_roles[safe_key_index],
+            image_attention_routing=routing,
+            per_image_contextual_time=per_image_contextual_time,
+            query_branch_kinds=branch_kinds[safe_query_index],
+            key_branch_kinds=branch_kinds[safe_key_index],
+            query_slot_ids=slot_ids[safe_query_index],
+            key_slot_ids=slot_ids[safe_key_index],
+            explicit_branches=explicit_branches,
+        )
+        return query_in_bounds & key_in_bounds & same_group & allowed_pair
+
+    return create_block_mask(
+        mask_mod, B=None, H=None, Q_LEN=token_count, KV_LEN=token_count, device=str(device)
+    )
+
+
+def run_flex_attention(
+    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *, block_mask: BlockMask
+) -> torch.Tensor:
+    """Run packed CUDA self-attention without changing the model head dimension."""
+    if query.ndim != 4 or key.ndim != 4 or value.ndim != 4:
+        raise ValueError("query, key, and value must use (B, H, T, D) layout")
+    if query.shape != key.shape or query.shape != value.shape:
+        raise ValueError("packed self-attention requires identical query, key, and value shapes")
+    if query.device.type != "cuda":
+        raise RuntimeError("FlexAttention requires CUDA tensors")
+    force_strict = query.dtype is torch.float32 and _FORCE_STRICT_FP32_FLEX_ATTENTION.get()
+    old_matmul_tf32 = None
+    if force_strict:
+        old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32
+        torch.backends.cuda.matmul.allow_tf32 = False
+    kernel_options = {"FORCE_USE_FLEX_ATTENTION": True}
+    forward_num_stages = _FLEX_ATTENTION_FORWARD_NUM_STAGES.get()
+    if forward_num_stages is not None:
+        kernel_options["num_stages"] = forward_num_stages
+    try:
+        return _compiled_flex_attention(
+            query, key, value, block_mask=block_mask, kernel_options=kernel_options
+        )
+    finally:
+        if force_strict:
+            assert old_matmul_tf32 is not None
+            torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32
+
+
+class GroupSelfAttention(nn.Module):
+    """QKV, QK normalization, 3D RoPE, and packed FlexAttention."""
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int = 1280,
+        num_heads: int = 16,
+        axes_dims: Sequence[int] = DEFAULT_ROPE_AXES_DIMS,
+        theta: float = DEFAULT_ROPE_THETA,
+        proj_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        for name, value in (("hidden_size", hidden_size), ("num_heads", num_heads)):
+            if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+                raise ValueError(f"{name} must be a positive integer, got {value!r}")
+        if hidden_size % num_heads != 0:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads: got {hidden_size} and {num_heads}"
+            )
+        if (
+            isinstance(proj_dropout, bool)
+            or not isinstance(proj_dropout, (int, float))
+            or (not math.isfinite(proj_dropout))
+            or (not 0 <= proj_dropout < 1)
+        ):
+            raise ValueError(f"proj_dropout must be finite and in [0, 1), got {proj_dropout!r}")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.rotary = RotaryEmbedding3D(axes_dims=axes_dims, theta=theta)
+        if self.rotary.head_dim != self.head_dim:
+            raise ValueError(
+                f"axes_dims must sum to the attention head dimension: got {self.rotary.head_dim}, expected {self.head_dim}"
+            )
+        self.q_norm = _RMSNorm(self.head_dim)
+        self.k_norm = _RMSNorm(self.head_dim)
+        self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=True)
+        self.proj = nn.Linear(hidden_size, hidden_size, bias=True)
+        self.proj_dropout = nn.Dropout(float(proj_dropout))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.ones_(self.q_norm.weight)
+        nn.init.ones_(self.k_norm.weight)
+        nn.init.xavier_uniform_(self.qkv.weight)
+        nn.init.zeros_(self.qkv.bias)
+        nn.init.xavier_uniform_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        position_ids: torch.Tensor,
+        block_mask: BlockMask,
+        rotary_factors: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Transform packed ``[tokens, hidden_size]`` values without padding."""
+        if not isinstance(tokens, torch.Tensor):
+            raise TypeError(f"tokens must be a torch.Tensor, got {type(tokens).__name__}")
+        if tokens.ndim != 2 or tokens.shape[1] != self.hidden_size:
+            raise ValueError(
+                f"tokens must use (tokens, hidden_size) layout with hidden_size {self.hidden_size}, got shape {tuple(tokens.shape)}"
+            )
+        if not tokens.is_floating_point():
+            raise ValueError(f"tokens must be floating point, got {tokens.dtype}")
+        if not isinstance(block_mask, BlockMask):
+            raise TypeError(f"block_mask must be a BlockMask, got {type(block_mask).__name__}")
+        token_count = tokens.shape[0]
+        if block_mask.shape[-2:] != (token_count, token_count):
+            raise ValueError(
+                f"BlockMask query/key lengths must match the packed token count: got {block_mask.shape[-2:]}, expected {(token_count, token_count)}"
+            )
+        qkv = self.qkv(tokens).reshape(token_count, 3, self.num_heads, self.head_dim)
+        (query, key, value) = qkv.permute(1, 2, 0, 3).unbind(dim=0)
+        if rotary_factors is None:
+            rotary_factors = self.rotary.prepare(position_ids)
+        normalize_and_rotate = (
+            _compiled_normalize_and_apply_rotary
+            if tokens.device.type == "cuda"
+            else _normalize_and_apply_rotary
+        )
+        (cosine, sine) = rotary_factors
+        query = normalize_and_rotate(
+            query, self.q_norm.weight, self.q_norm.eps, cosine, sine
+        ).unsqueeze(0)
+        key = normalize_and_rotate(
+            key, self.k_norm.weight, self.k_norm.eps, cosine, sine
+        ).unsqueeze(0)
+        value = value.unsqueeze(0)
+        attended = run_flex_attention(query, key, value, block_mask=block_mask)
+        merged = attended.transpose(1, 2).reshape(token_count, self.hidden_size)
+        return self.proj_dropout(self.proj(merged))

vendor/pixtext/modeling/branches.py

@@ -0,0 +1,294 @@
+"""Post-projection packing for mixed RGB and canonical-glyph branches."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from pixtext.modeling.patches import (
+    DirectPatchProjection,
+    PackedImagePatches,
+    PatchLayout,
+    patchify,
+)
+from pixtext.preprocessing.geometry import GroupGeometry
+from pixtext.preprocessing.glyphs import GlyphCondition
+
+MAIN_BRANCH = 0
+ITEM_BRANCH = 1
+GLYPH_BRANCH = 2
+
+
+@dataclass(frozen=True, slots=True)
+class BranchDescriptor:
+    kind: str
+    slot_id: int | None
+
+    def __post_init__(self) -> None:
+        if self.kind not in {"main", "item", "glyph"}:
+            raise ValueError(f"unknown branch kind: {self.kind!r}")
+        if self.kind == "main":
+            if self.slot_id is not None:
+                raise ValueError("main branch slot_id must be None")
+        elif (
+            isinstance(self.slot_id, bool) or not isinstance(self.slot_id, int) or self.slot_id < 0
+        ):
+            raise ValueError("item/glyph slot_id must be a non-negative integer")
+
+    @property
+    def kind_id(self) -> int:
+        return {"main": MAIN_BRANCH, "item": ITEM_BRANCH, "glyph": GLYPH_BRANCH}[self.kind]
+
+
+class DirectGlyphPatchProjection(nn.Module):
+    """Independent direct projection for one-channel [0,1] coverage patches."""
+
+    def __init__(self, *, patch_size: int = 16, hidden_size: int = 1280) -> None:
+        super().__init__()
+        if patch_size <= 0 or hidden_size <= 0:
+            raise ValueError("patch_size and hidden_size must be positive")
+        self.patch_size = patch_size
+        self.in_channels = 1
+        self.proj = nn.Linear(patch_size**2, hidden_size, bias=True)
+        nn.init.xavier_uniform_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, patches: torch.Tensor) -> torch.Tensor:
+        if patches.ndim != 2 or patches.shape[1] != self.proj.in_features:
+            raise ValueError("glyph patches must use (tokens, patch_size**2) layout")
+        return self.proj(patches)
+
+
+@dataclass(frozen=True, slots=True)
+class PackedBranchTokens:
+    tokens: torch.Tensor
+    descriptors: tuple[BranchDescriptor, ...]
+    layouts: tuple[PatchLayout, ...]
+    sample_ids: torch.Tensor
+    image_ids: torch.Tensor
+    branch_kinds: torch.Tensor
+    slot_ids: torch.Tensor
+    image_offsets: torch.Tensor
+    sample_offsets: torch.Tensor
+    branches_per_sample: tuple[int, ...]
+    rgb_token_indices: torch.Tensor
+    rgb_branch_indices: tuple[int, ...]
+    rgb_images_per_sample: tuple[int, ...]
+
+
+class PackedBranchTokenizer:
+    """Hide mixed raw widths and merge RGB/glyphs after projection."""
+
+    def __init__(
+        self, *, rgb_projection: DirectPatchProjection, glyph_projection: DirectGlyphPatchProjection
+    ) -> None:
+        if not isinstance(rgb_projection, DirectPatchProjection):
+            raise TypeError("rgb_projection must be DirectPatchProjection")
+        if not isinstance(glyph_projection, DirectGlyphPatchProjection):
+            raise TypeError("glyph_projection must be DirectGlyphPatchProjection")
+        if rgb_projection.patch_size != glyph_projection.patch_size:
+            raise ValueError("RGB and glyph patch sizes must match")
+        if rgb_projection.proj.out_features != glyph_projection.proj.out_features:
+            raise ValueError("RGB and glyph projections must share hidden size")
+        self.rgb_projection = rgb_projection
+        self.glyph_projection = glyph_projection
+        self.patch_size = rgb_projection.patch_size
+
+    def tokenize(
+        self,
+        *,
+        packed_rgb: PackedImagePatches,
+        glyphs: Sequence[Sequence[GlyphCondition | None]],
+        geometries: Sequence[GroupGeometry],
+    ) -> PackedBranchTokens:
+        group_count = len(packed_rgb.images_per_sample)
+        if len(glyphs) != group_count or len(geometries) != group_count:
+            raise ValueError("RGB, glyph, and geometry group counts must match")
+        if any(
+            len(group_glyphs) != len(geometry.items)
+            for (group_glyphs, geometry) in zip(glyphs, geometries, strict=True)
+        ):
+            raise ValueError("glyph slots must match geometry items")
+        torch._assert_async(
+            torch.isfinite(packed_rgb.patches).all(), "RGB branches must remain finite"
+        )
+        main_patches = packed_rgb.patches[packed_rgb.image_ids.eq(MAIN_BRANCH)]
+        torch._assert_async(
+            main_patches.ge(-1).all() & main_patches.le(1).all(),
+            "clean main RGB branches must remain within [-1, 1]",
+        )
+        rgb_tokens = self.rgb_projection(packed_rgb.patches)
+        token_parts: list[torch.Tensor] = []
+        descriptors: list[BranchDescriptor] = []
+        layouts: list[PatchLayout] = []
+        sample_values: list[int] = []
+        image_values: list[int] = []
+        kind_values: list[int] = []
+        slot_values: list[int] = []
+        image_offsets = [0]
+        sample_offsets = [0]
+        branches_per_sample: list[int] = []
+        rgb_token_indices: list[torch.Tensor] = []
+        rgb_branch_indices: list[int] = []
+        flat_rgb = 0
+        compact_rgb_offset = 0
+        device = packed_rgb.patches.device
+        dtype = rgb_tokens.dtype
+        for sample_id, (rgb_count, group_glyphs, geometry) in enumerate(
+            zip(packed_rgb.images_per_sample, glyphs, geometries, strict=True)
+        ):
+            if rgb_count != 1 + len(geometry.items):
+                raise ValueError("packed RGB topology does not match geometry")
+            group_branch_count = 0
+            for rgb_image_id in range(rgb_count):
+                layout = packed_rgb.layouts[flat_rgb]
+                token_count = layout.num_patches
+                rgb_part = rgb_tokens[compact_rgb_offset : compact_rgb_offset + token_count]
+                descriptor = BranchDescriptor(
+                    "main" if rgb_image_id == 0 else "item",
+                    None if rgb_image_id == 0 else rgb_image_id - 1,
+                )
+                physical_id = group_branch_count
+                token_parts.append(rgb_part)
+                descriptors.append(descriptor)
+                layouts.append(layout)
+                sample_values.extend([sample_id] * token_count)
+                image_values.extend([physical_id] * token_count)
+                kind_values.extend([descriptor.kind_id] * token_count)
+                slot_values.extend(
+                    [-1 if descriptor.slot_id is None else descriptor.slot_id] * token_count
+                )
+                start = image_offsets[-1]
+                image_offsets.append(start + token_count)
+                rgb_token_indices.append(torch.arange(start, start + token_count, device=device))
+                rgb_branch_indices.append(len(descriptors) - 1)
+                compact_rgb_offset += token_count
+                flat_rgb += 1
+                group_branch_count += 1
+                if rgb_image_id == 0:
+                    continue
+                glyph = group_glyphs[rgb_image_id - 1]
+                if glyph is None:
+                    continue
+                if glyph.aligned_crop != geometry.items[rgb_image_id - 1].crop:
+                    raise ValueError("glyph crop provenance does not match geometry")
+                pixels = glyph.image.to(device=device, dtype=dtype, non_blocking=True)
+                torch._assert_async(
+                    torch.isfinite(pixels).all() & pixels.ge(0).all() & pixels.le(1).all(),
+                    "glyph coverage must remain finite and within [0, 1]",
+                )
+                patchified = patchify(pixels, self.patch_size)
+                if patchified.layout.num_patches != token_count:
+                    raise ValueError("glyph and paired item token counts must match")
+                glyph_tokens = self.glyph_projection(patchified.patches)
+                glyph_descriptor = BranchDescriptor("glyph", rgb_image_id - 1)
+                glyph_physical_id = group_branch_count
+                token_parts.append(glyph_tokens)
+                descriptors.append(glyph_descriptor)
+                layouts.append(patchified.layout)
+                sample_values.extend([sample_id] * token_count)
+                image_values.extend([glyph_physical_id] * token_count)
+                kind_values.extend([GLYPH_BRANCH] * token_count)
+                slot_values.extend([rgb_image_id - 1] * token_count)
+                image_offsets.append(image_offsets[-1] + token_count)
+                group_branch_count += 1
+            branches_per_sample.append(group_branch_count)
+            sample_offsets.append(image_offsets[-1])
+        tokens = torch.cat(token_parts)
+        return PackedBranchTokens(
+            tokens=tokens,
+            descriptors=tuple(descriptors),
+            layouts=tuple(layouts),
+            sample_ids=torch.tensor(sample_values, dtype=torch.int64, device=device),
+            image_ids=torch.tensor(image_values, dtype=torch.int64, device=device),
+            branch_kinds=torch.tensor(kind_values, dtype=torch.int64, device=device),
+            slot_ids=torch.tensor(slot_values, dtype=torch.int64, device=device),
+            image_offsets=torch.tensor(image_offsets, dtype=torch.int64, device=device),
+            sample_offsets=torch.tensor(sample_offsets, dtype=torch.int64, device=device),
+            branches_per_sample=tuple(branches_per_sample),
+            rgb_token_indices=torch.cat(rgb_token_indices),
+            rgb_branch_indices=tuple(rgb_branch_indices),
+            rgb_images_per_sample=packed_rgb.images_per_sample,
+        )
+
+
+def _branch_position_ids(
+    packed: PackedBranchTokens, geometries: Sequence[GroupGeometry]
+) -> torch.Tensor:
+    """Build main/item/glyph role positions from the shared descriptor order."""
+    parts: list[torch.Tensor] = []
+    branch_index = 0
+    device = packed.tokens.device
+    for sample_id, (geometry, branch_count) in enumerate(
+        zip(geometries, packed.branches_per_sample, strict=True)
+    ):
+        reference = geometry.reference.size
+        center_h = reference.height / (2 * packed.layouts[branch_index].patch_size)
+        center_w = reference.width / (2 * packed.layouts[branch_index].patch_size)
+        for image_id in range(branch_count):
+            descriptor = packed.descriptors[branch_index]
+            layout = packed.layouts[branch_index]
+            grid = layout.grid_size
+            rows = torch.arange(grid.height, dtype=torch.float32, device=device)
+            columns = torch.arange(grid.width, dtype=torch.float32, device=device)
+            (rows, columns) = torch.meshgrid(rows, columns, indexing="ij")
+            if descriptor.kind == "main":
+                h = (rows + 0.5) / geometry.main.scale_h - center_h
+                w = (columns + 0.5) / geometry.main.scale_w - center_w
+                n = torch.zeros_like(h)
+            else:
+                assert descriptor.slot_id is not None
+                item = geometry.items[descriptor.slot_id]
+                if layout.image_size != item.crop.size:
+                    raise ValueError(
+                        f"group {sample_id} branch {image_id} layout does not match slot"
+                    )
+                h = rows + 0.5 + item.crop.top / layout.patch_size - center_h
+                w = columns + 0.5 + item.crop.left / layout.patch_size - center_w
+                n = torch.full_like(h, 1.0 if descriptor.kind == "item" else 2.0)
+            parts.append(torch.stack((n, h, w), dim=-1).reshape(-1, 3))
+            branch_index += 1
+    positions = torch.cat(parts)
+    if positions.shape != (packed.tokens.shape[0], 3):
+        raise RuntimeError("branch position count does not match token count")
+    return positions
+
+
+def pack_branch_group_sequences(
+    *, branches: PackedBranchTokens, prompts, geometries: Sequence[GroupGeometry]
+):
+    from pixtext.modeling.conditioning import IMAGE_TOKEN_TYPE, PackedGroupSequence
+
+    "Assemble optional prompts with physical RGB/glyph branches.\n\n    ``prompts=None`` is the real no-text adapter: it emits no placeholder\n    tokens or prompt-owned topology and derives every offset from branches.\n    "
+    if prompts is None:
+        if len(geometries) != len(branches.branches_per_sample):
+            raise ValueError("branch and geometry group counts must match")
+        positions = _branch_position_ids(branches, geometries)
+        device = branches.tokens.device
+        token_count = branches.tokens.shape[0]
+        empty = torch.empty(0, dtype=torch.int64, device=device)
+        return PackedGroupSequence(
+            tokens=branches.tokens,
+            position_ids=positions,
+            sample_ids=branches.sample_ids,
+            image_ids=branches.image_ids,
+            token_types=torch.full(
+                (token_count,), IMAGE_TOKEN_TYPE, dtype=torch.int64, device=device
+            ),
+            timestep_roles=branches.branch_kinds.ne(MAIN_BRANCH).to(torch.int64),
+            sample_offsets=branches.sample_offsets,
+            image_token_indices=torch.arange(token_count, dtype=torch.int64, device=device),
+            prompt_token_indices=empty,
+            image_offsets=branches.image_offsets,
+            prompt_offsets=torch.zeros(1, dtype=torch.int64, device=device),
+            layouts=branches.layouts,
+            images_per_sample=branches.branches_per_sample,
+            timestep_token_routing=None,
+            branch_kinds=branches.branch_kinds,
+            slot_ids=branches.slot_ids,
+            branch_descriptors=branches.descriptors,
+        )
+    raise ValueError("PixText does not use prompt tokens")

vendor/pixtext/modeling/conditioning.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+from pixtext.modeling.modes import TimestepTokenRoutingMode
+from pixtext.modeling.patches import PatchLayout
+from pixtext.preprocessing.geometry import ImageSize
+
+IMAGE_TOKEN_TYPE = 0
+PROMPT_TOKEN_TYPE = 1
+TIME_TOKEN_TYPE = 2
+
+
+@dataclass(frozen=True, slots=True)
+class PackedGroupSequence:
+    """Group-contiguous image and prompt tokens ready for single-stream attention."""
+
+    tokens: torch.Tensor
+    position_ids: torch.Tensor
+    sample_ids: torch.Tensor
+    image_ids: torch.Tensor
+    token_types: torch.Tensor
+    timestep_roles: torch.Tensor
+    sample_offsets: torch.Tensor
+    image_token_indices: torch.Tensor
+    prompt_token_indices: torch.Tensor
+    image_offsets: torch.Tensor
+    prompt_offsets: torch.Tensor
+    layouts: tuple[PatchLayout, ...]
+    images_per_sample: tuple[int, ...]
+    timestep_token_routing: TimestepTokenRoutingMode | None
+    branch_kinds: torch.Tensor | None = None
+    slot_ids: torch.Tensor | None = None
+    branch_descriptors: tuple[object, ...] = ()
+    sample_token_counts: tuple[int, ...] = ()
+
+    @property
+    def image_shapes(self) -> tuple[ImageSize, ...]:
+        return tuple(layout.image_size for layout in self.layouts)
+
+
+def extract_image_tokens(
+    sequence_tokens: torch.Tensor, packed: PackedGroupSequence
+) -> torch.Tensor:
+    """Restore compact image-token order from a transformed group sequence."""
+    if not isinstance(sequence_tokens, torch.Tensor):
+        raise TypeError(
+            f"sequence_tokens must be a torch.Tensor, got {type(sequence_tokens).__name__}"
+        )
+    if sequence_tokens.ndim != 2:
+        raise ValueError(
+            f"sequence_tokens must use (tokens, hidden_size) layout, got shape {tuple(sequence_tokens.shape)}"
+        )
+    if sequence_tokens.shape[0] != packed.tokens.shape[0]:
+        raise ValueError(
+            f"sequence token count does not match packed sequence: got {sequence_tokens.shape[0]}, expected {packed.tokens.shape[0]}"
+        )
+    if sequence_tokens.device != packed.image_token_indices.device:
+        raise ValueError("sequence tokens and packed indices must share a device")
+    return sequence_tokens.index_select(0, packed.image_token_indices)

vendor/pixtext/modeling/modes.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class ImageAttentionRoutingMode(str, Enum):
+    """Control direct attention edges between images in the same group."""
+
+    ALL_IMAGES = "all_images"
+    MAIN_IMAGE_HUB = "main_image_hub"
+    ITEMS_READ_MAIN = "items_read_main"
+    ITEMS_INTERREAD_GLYPH_READS_OWN_ITEM = "items_interread_glyph_reads_own_item"
+    ITEMS_INTERREAD_GLYPH_READS_OWN_ITEM_BIDIRECTIONAL_MAIN = (
+        "items_interread_glyph_reads_own_item_bidirectional_main"
+    )
+
+    @classmethod
+    def require(
+        cls, value: object, *, name: str = "image_attention_routing"
+    ) -> ImageAttentionRoutingMode:
+        if not isinstance(value, cls):
+            raise TypeError(f"{name} must be an ImageAttentionRoutingMode, got {value!r}")
+        return value
+
+
+class TextConditioningMode(str, Enum):
+    """Choose whether the packed Transformer contains prompt branches."""
+
+    PACKED_PROMPT = "packed_prompt"
+    NONE = "none"
+
+    @classmethod
+    def require(cls, value: object, *, name: str = "text_conditioning") -> TextConditioningMode:
+        if not isinstance(value, cls):
+            raise TypeError(f"{name} must be a TextConditioningMode, got {value!r}")
+        return value
+
+
+class TimestepTokenRoutingMode(str, Enum):
+    """Choose the topology and directional routing of iMF timestep banks."""
+
+    SHARED_ROLE_ONE_WAY = "shared_role_one_way"
+    PER_IMAGE_CONTEXTUAL = "per_image_contextual"
+
+    @classmethod
+    def require(
+        cls, value: object, *, name: str = "timestep_token_routing"
+    ) -> TimestepTokenRoutingMode:
+        if not isinstance(value, cls):
+            raise TypeError(f"{name} must be a TimestepTokenRoutingMode, got {value!r}")
+        return value
+
+
+class NAxisMode(str, Enum):
+    """Assign the RoPE n coordinate from image role or group-local identity."""
+
+    ROLE = "role"
+    IMAGE_ID = "image_id"
+
+    @classmethod
+    def require(cls, value: object, *, name: str = "n_axis_mode") -> NAxisMode:
+        if not isinstance(value, cls):
+            raise TypeError(f"{name} must be an NAxisMode, got {value!r}")
+        return value
+
+    def position_for(self, image_id: int) -> float:
+        if self is NAxisMode.ROLE:
+            return float(image_id > 0)
+        return float(image_id)

vendor/pixtext/modeling/norm.py

@@ -0,0 +1,19 @@
+import torch
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-06):
+        """
+        LlamaRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)

vendor/pixtext/modeling/patches.py

@@ -0,0 +1,298 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from pixtext.preprocessing.geometry import ImageSize
+
+
+def _validate_positive_integer(name: str, value: int) -> None:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer, got {value!r}")
+
+
+@dataclass(frozen=True, slots=True)
+class PatchLayout:
+    """Explicit metadata required to reconstruct one rectangular image."""
+
+    image_size: ImageSize
+    patch_size: int
+    channels: int
+
+    def __post_init__(self) -> None:
+        _validate_positive_integer("patch_size", self.patch_size)
+        _validate_positive_integer("channels", self.channels)
+        if (
+            self.image_size.height % self.patch_size != 0
+            or self.image_size.width % self.patch_size != 0
+        ):
+            raise ValueError(
+                f"image height and width must be divisible by patch_size: got {self.image_size!r} and patch_size={self.patch_size}"
+            )
+
+    @property
+    def grid_size(self) -> ImageSize:
+        return ImageSize(
+            self.image_size.height // self.patch_size, self.image_size.width // self.patch_size
+        )
+
+    @property
+    def num_patches(self) -> int:
+        grid = self.grid_size
+        return grid.height * grid.width
+
+
+@dataclass(frozen=True, slots=True)
+class PatchifiedImage:
+    """Patch vectors together with their explicit rectangular layout."""
+
+    patches: torch.Tensor
+    layout: PatchLayout
+
+
+@dataclass(frozen=True, slots=True)
+class PackedImagePatches:
+    """Raw image patches and routing metadata for a rank-local group batch."""
+
+    patches: torch.Tensor
+    sample_ids: torch.Tensor
+    image_ids: torch.Tensor
+    sample_offsets: torch.Tensor
+    image_offsets: torch.Tensor
+    layouts: tuple[PatchLayout, ...]
+    images_per_sample: tuple[int, ...]
+
+    @property
+    def image_shapes(self) -> tuple[ImageSize, ...]:
+        return tuple(layout.image_size for layout in self.layouts)
+
+    def sample_token_range(self, sample_start: int, sample_end: int) -> tuple[int, int]:
+        if (
+            isinstance(sample_start, bool)
+            or not isinstance(sample_start, int)
+            or isinstance(sample_end, bool)
+            or (not isinstance(sample_end, int))
+            or (not 0 <= sample_start < sample_end <= len(self.images_per_sample))
+        ):
+            raise IndexError("packed image sample range is out of bounds")
+        image_start = sum(self.images_per_sample[:sample_start])
+        image_end = image_start + sum(self.images_per_sample[sample_start:sample_end])
+        return (
+            sum(layout.num_patches for layout in self.layouts[:image_start]),
+            sum(layout.num_patches for layout in self.layouts[:image_end]),
+        )
+
+    def select_samples(self, sample_start: int, sample_end: int) -> PackedImagePatches:
+        """Return a contiguous sample range with local packed routing."""
+        (token_start, token_end) = self.sample_token_range(sample_start, sample_end)
+        images_per_sample = self.images_per_sample[sample_start:sample_end]
+        image_start = sum(self.images_per_sample[:sample_start])
+        image_end = image_start + sum(images_per_sample)
+        layouts = self.layouts[image_start:image_end]
+        image_offsets = [0]
+        for layout in layouts:
+            image_offsets.append(image_offsets[-1] + layout.num_patches)
+        sample_offsets = [0]
+        layout_index = 0
+        for image_count in images_per_sample:
+            sample_offsets.append(
+                sample_offsets[-1]
+                + sum(
+                    layout.num_patches
+                    for layout in layouts[layout_index : layout_index + image_count]
+                )
+            )
+            layout_index += image_count
+        device = self.patches.device
+        return PackedImagePatches(
+            patches=self.patches[token_start:token_end],
+            sample_ids=self.sample_ids[token_start:token_end] - sample_start,
+            image_ids=self.image_ids[token_start:token_end],
+            sample_offsets=torch.tensor(sample_offsets, dtype=torch.int64, device=device),
+            image_offsets=torch.tensor(image_offsets, dtype=torch.int64, device=device),
+            layouts=layouts,
+            images_per_sample=images_per_sample,
+        )
+
+
+class DirectPatchProjection(nn.Module):
+    """Project flattened RGB patches directly into JiT token features."""
+
+    def __init__(
+        self, *, patch_size: int = 16, in_channels: int = 3, hidden_size: int = 1280
+    ) -> None:
+        super().__init__()
+        for name, value in (
+            ("patch_size", patch_size),
+            ("in_channels", in_channels),
+            ("hidden_size", hidden_size),
+        ):
+            _validate_positive_integer(name, value)
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.proj = nn.Linear(in_channels * patch_size**2, hidden_size, bias=True)
+        nn.init.xavier_uniform_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, patches: torch.Tensor) -> torch.Tensor:
+        if not isinstance(patches, torch.Tensor):
+            raise TypeError(f"patches must be a torch.Tensor, got {type(patches).__name__}")
+        if patches.ndim != 2:
+            raise ValueError(
+                f"patches must use (tokens, patch_width) layout, got shape {tuple(patches.shape)}"
+            )
+        if patches.shape[1] != self.proj.in_features:
+            raise ValueError(
+                f"patch width does not match projection: got {patches.shape[1]}, expected {self.proj.in_features}"
+            )
+        return self.proj(patches)
+
+
+def patchify(image: torch.Tensor, patch_size: int = 16) -> PatchifiedImage:
+    """Convert one ``[C,H,W]`` image into row-major flattened patches."""
+    _validate_positive_integer("patch_size", patch_size)
+    if not isinstance(image, torch.Tensor):
+        raise TypeError(f"image must be a torch.Tensor, got {type(image).__name__}")
+    if image.ndim != 3:
+        raise ValueError(f"image must use (C, H, W) layout, got shape {tuple(image.shape)}")
+    (channels, height, width) = image.shape
+    if height % patch_size != 0 or width % patch_size != 0:
+        raise ValueError(
+            f"image height and width must be divisible by patch_size: got {(height, width)} and patch_size={patch_size}"
+        )
+    grid_height = height // patch_size
+    grid_width = width // patch_size
+    patches = (
+        image.reshape(channels, grid_height, patch_size, grid_width, patch_size)
+        .permute(1, 3, 0, 2, 4)
+        .reshape(grid_height * grid_width, channels * patch_size**2)
+    )
+    return PatchifiedImage(
+        patches=patches,
+        layout=PatchLayout(
+            image_size=ImageSize(height, width), patch_size=patch_size, channels=channels
+        ),
+    )
+
+
+def unpatchify(patches: torch.Tensor, layout: PatchLayout) -> torch.Tensor:
+    """Reconstruct one image using explicit rectangular layout metadata."""
+    if not isinstance(patches, torch.Tensor):
+        raise TypeError(f"patches must be a torch.Tensor, got {type(patches).__name__}")
+    if patches.ndim != 2:
+        raise ValueError(
+            f"patches must use (tokens, patch_width) layout, got shape {tuple(patches.shape)}"
+        )
+    grid = layout.grid_size
+    if patches.shape[0] != layout.num_patches:
+        raise ValueError(
+            f"patch token count does not match layout: got {patches.shape[0]}, expected {layout.num_patches}"
+        )
+    expected_patch_width = layout.channels * layout.patch_size**2
+    if patches.shape[1] != expected_patch_width:
+        raise ValueError(
+            f"patch width does not match layout: got {patches.shape[1]}, expected {expected_patch_width}"
+        )
+    image = (
+        patches.reshape(
+            grid.height, grid.width, layout.channels, layout.patch_size, layout.patch_size
+        )
+        .permute(2, 0, 3, 1, 4)
+        .reshape(layout.channels, layout.image_size.height, layout.image_size.width)
+    )
+    return image
+
+
+def pack_image_groups(
+    groups: Sequence[Sequence[torch.Tensor]], patch_size: int = 16
+) -> PackedImagePatches:
+    """Patchify and concatenate ``main + ordered items`` for every group."""
+    _validate_positive_integer("patch_size", patch_size)
+    if len(groups) == 0:
+        raise ValueError("packed image batch must contain at least one group")
+    patch_parts = []
+    sample_id_values = []
+    image_id_values = []
+    layouts = []
+    images_per_sample = []
+    sample_offsets = [0]
+    image_offsets = [0]
+    expected_dtype = None
+    expected_device = None
+    expected_channels = None
+    for sample_id, group in enumerate(groups):
+        if len(group) == 0:
+            raise ValueError(f"group {sample_id} must contain at least one image (the main image)")
+        images_per_sample.append(len(group))
+        for image_id, image in enumerate(group):
+            patchified = patchify(image, patch_size)
+            if expected_dtype is None:
+                expected_dtype = patchified.patches.dtype
+                expected_device = patchified.patches.device
+                expected_channels = patchified.layout.channels
+            else:
+                location = f"group {sample_id} image {image_id}"
+                if patchified.patches.dtype != expected_dtype:
+                    raise ValueError(
+                        f"{location} dtype {patchified.patches.dtype} does not match packed dtype {expected_dtype}"
+                    )
+                if patchified.patches.device != expected_device:
+                    raise ValueError(
+                        f"{location} device {patchified.patches.device} does not match packed device {expected_device}"
+                    )
+                if patchified.layout.channels != expected_channels:
+                    raise ValueError(
+                        f"{location} channels {patchified.layout.channels} do not match packed channels {expected_channels}"
+                    )
+            patch_parts.append(patchified.patches)
+            layouts.append(patchified.layout)
+            token_count = patchified.layout.num_patches
+            sample_id_values.extend([sample_id] * token_count)
+            image_id_values.extend([image_id] * token_count)
+            image_offsets.append(image_offsets[-1] + token_count)
+        sample_offsets.append(image_offsets[-1])
+    patches = torch.cat(patch_parts)
+    device = patches.device
+    return PackedImagePatches(
+        patches=patches,
+        sample_ids=torch.tensor(sample_id_values, dtype=torch.int64, device=device),
+        image_ids=torch.tensor(image_id_values, dtype=torch.int64, device=device),
+        sample_offsets=torch.tensor(sample_offsets, dtype=torch.int64, device=device),
+        image_offsets=torch.tensor(image_offsets, dtype=torch.int64, device=device),
+        layouts=tuple(layouts),
+        images_per_sample=tuple(images_per_sample),
+    )
+
+
+def unpack_image_patches(
+    patches: torch.Tensor, packed: PackedImagePatches
+) -> tuple[tuple[torch.Tensor, ...], ...]:
+    """Restore packed patch pixels to the original group/image structure."""
+    if not isinstance(patches, torch.Tensor):
+        raise TypeError(f"patches must be a torch.Tensor, got {type(patches).__name__}")
+    if patches.ndim != 2:
+        raise ValueError(
+            f"patches must use (tokens, patch_width) layout, got shape {tuple(patches.shape)}"
+        )
+    expected_token_count = sum(layout.num_patches for layout in packed.layouts)
+    if patches.shape[0] != expected_token_count:
+        raise ValueError(
+            f"packed total token count mismatch: got {patches.shape[0]}, expected {expected_token_count}"
+        )
+    if packed.image_offsets.shape != (len(packed.layouts) + 1,):
+        raise ValueError("packed image offsets do not match the image layouts")
+    flat_images = []
+    start = 0
+    for layout in packed.layouts:
+        end = start + layout.num_patches
+        flat_images.append(unpatchify(patches[start:end], layout))
+        start = end
+    groups = []
+    image_index = 0
+    for image_count in packed.images_per_sample:
+        groups.append(tuple(flat_images[image_index : image_index + image_count]))
+        image_index += image_count
+    return tuple(groups)

vendor/pixtext/modeling/rope.py

@@ -0,0 +1,350 @@
+"""Geometry-aware 3D rotary positions for variable-length image groups."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+from pixtext.modeling.modes import NAxisMode
+from pixtext.modeling.patches import PackedImagePatches, PatchLayout
+from pixtext.preprocessing.geometry import GroupGeometry, ImageSize
+
+DEFAULT_ROPE_AXES_DIMS = (12, 34, 34)
+DEFAULT_ROPE_THETA = 10000.0
+
+
+def _apply_rotary_factors(
+    values: torch.Tensor, cosine: torch.Tensor, sine: torch.Tensor
+) -> torch.Tensor:
+    """Apply already-materialized pairwise rotations."""
+    output_dtype = values.dtype
+    pair_count = values.shape[-1] // 2
+    paired = values.float().reshape(*values.shape[:-1], pair_count, 2)
+    (first, second) = paired.unbind(dim=-1)
+    factor_shape = (*[1] * (values.ndim - 2), values.shape[-2], pair_count)
+    cosine = cosine.reshape(factor_shape)
+    sine = sine.reshape(factor_shape)
+    return (
+        torch.stack((first * cosine - second * sine, first * sine + second * cosine), dim=-1)
+        .flatten(start_dim=-2)
+        .to(output_dtype)
+    )
+
+
+_compiled_apply_rotary_factors = torch.compile(_apply_rotary_factors, dynamic=True)
+
+
+class RotaryEmbedding3D(nn.Module):
+    """Apply dynamic 3D RoPE to the last dimension of Q or K tensors."""
+
+    def __init__(
+        self, axes_dims: Sequence[int] = DEFAULT_ROPE_AXES_DIMS, theta: float = DEFAULT_ROPE_THETA
+    ) -> None:
+        super().__init__()
+        if len(axes_dims) != 3:
+            raise ValueError(f"axes_dims must contain exactly three dimensions, got {axes_dims!r}")
+        if any(
+            isinstance(axis_dim, bool)
+            or not isinstance(axis_dim, int)
+            or axis_dim <= 0
+            or (axis_dim % 2 != 0)
+            for axis_dim in axes_dims
+        ):
+            raise ValueError(
+                f"every axes_dims entry must be a positive even integer, got {axes_dims!r}"
+            )
+        if (
+            isinstance(theta, bool)
+            or not isinstance(theta, (int, float))
+            or (not math.isfinite(theta))
+            or (theta <= 0)
+        ):
+            raise ValueError(f"theta must be positive and finite, got {theta!r}")
+        self.axes_dims = tuple(axes_dims)
+        self.theta = float(theta)
+        self.head_dim = sum(self.axes_dims)
+        for axis_name, axis_dim in zip(("n", "h", "w"), self.axes_dims, strict=True):
+            inverse_frequency = 1.0 / self.theta ** (
+                torch.arange(0, axis_dim, 2, dtype=torch.float32, device="cpu") / axis_dim
+            )
+            self.register_buffer(f"inv_freq_{axis_name}", inverse_frequency, persistent=False)
+
+    def _apply(self, fn, recurse: bool = True):
+        """Move buffers with the module while retaining FP32 frequencies."""
+        super()._apply(fn, recurse)
+        self.inv_freq_n = self.inv_freq_n.float()
+        self.inv_freq_h = self.inv_freq_h.float()
+        self.inv_freq_w = self.inv_freq_w.float()
+        return self
+
+    def forward(self, values: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor:
+        """Rotate ``values[..., tokens, head_dim]`` using ``[tokens, 3]`` IDs."""
+        if not isinstance(values, torch.Tensor):
+            raise TypeError(f"values must be a torch.Tensor, got {type(values).__name__}")
+        if values.ndim < 2:
+            raise ValueError(
+                f"values must have at least (tokens, head_dim) dimensions, got shape {tuple(values.shape)}"
+            )
+        if not values.is_floating_point():
+            raise ValueError(f"values must be floating point, got {values.dtype}")
+        factors = self.prepare(position_ids)
+        return self.apply_precomputed(values, factors)
+
+    def prepare(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Materialize FP32 cosine/sine factors once for a packed sequence."""
+        if not isinstance(position_ids, torch.Tensor):
+            raise TypeError(
+                f"position_ids must be a torch.Tensor, got {type(position_ids).__name__}"
+            )
+        if position_ids.ndim != 2 or position_ids.shape[1] != 3:
+            raise ValueError(
+                f"position_ids must use (tokens, 3) layout, got shape {tuple(position_ids.shape)}"
+            )
+        float_positions = position_ids.float()
+        angles = []
+        for axis, inverse_frequency in enumerate(
+            (self.inv_freq_n, self.inv_freq_h, self.inv_freq_w)
+        ):
+            if inverse_frequency.device != position_ids.device:
+                raise ValueError(
+                    f"RotaryEmbedding3D buffers and position_ids must be on the same device: got {inverse_frequency.device} and {position_ids.device}"
+                )
+            angles.append(float_positions[:, axis, None] * inverse_frequency[None, :])
+        packed_angles = torch.cat(angles, dim=-1)
+        return (packed_angles.cos(), packed_angles.sin())
+
+    def apply_precomputed(
+        self, values: torch.Tensor, factors: tuple[torch.Tensor, torch.Tensor]
+    ) -> torch.Tensor:
+        """Rotate values using factors returned by :meth:`prepare`."""
+        if not isinstance(values, torch.Tensor):
+            raise TypeError(f"values must be a torch.Tensor, got {type(values).__name__}")
+        if values.ndim < 2:
+            raise ValueError(
+                f"values must have at least (tokens, head_dim) dimensions, got shape {tuple(values.shape)}"
+            )
+        if not values.is_floating_point():
+            raise ValueError(f"values must be floating point, got {values.dtype}")
+        if (
+            not isinstance(factors, tuple)
+            or len(factors) != 2
+            or (not all(isinstance(factor, torch.Tensor) for factor in factors))
+        ):
+            raise TypeError("factors must be the (cosine, sine) tensor pair")
+        (cosine, sine) = factors
+        if values.shape[-1] != self.head_dim:
+            raise ValueError(
+                f"values head dimension does not match axes_dims: got {values.shape[-1]}, expected {self.head_dim}"
+            )
+        if cosine.ndim != 2 or sine.ndim != 2:
+            raise ValueError("rotary factors must use (tokens, head_dim / 2) layout")
+        if cosine.shape[0] != values.shape[-2] or sine.shape[0] != values.shape[-2]:
+            raise ValueError(
+                f"position token count does not match values: got {cosine.shape[0]} and {sine.shape[0]}, expected {values.shape[-2]}"
+            )
+        expected_factor_shape = (values.shape[-2], self.head_dim // 2)
+        if cosine.shape != expected_factor_shape or sine.shape != expected_factor_shape:
+            raise ValueError(
+                f"rotary factors must use (tokens, head_dim / 2) layout: got {tuple(cosine.shape)} and {tuple(sine.shape)}, expected {expected_factor_shape}"
+            )
+        if cosine.device != values.device or sine.device != values.device:
+            raise ValueError("rotary factors and values must be on the same device")
+        apply = (
+            _compiled_apply_rotary_factors
+            if values.device.type == "cuda"
+            else _apply_rotary_factors
+        )
+        return apply(values, cosine, sine)
+
+
+def _validate_nonnegative_integer(name: str, value: int) -> None:
+    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
+        raise ValueError(f"{name} must be a non-negative integer, got {value!r}")
+
+
+def _validate_positive_integer(name: str, value: int) -> None:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer, got {value!r}")
+
+
+def _validate_layout(
+    layout: PatchLayout,
+    expected_size: ImageSize,
+    expected_patch_size: int,
+    *,
+    sample_id: int,
+    image_id: int,
+) -> None:
+    if layout.patch_size != expected_patch_size:
+        raise ValueError(
+            f"group {sample_id} image {image_id} patch_size {layout.patch_size} does not match position patch_size {expected_patch_size}"
+        )
+    if layout.image_size != expected_size:
+        raise ValueError(
+            f"group {sample_id} image {image_id} layout size {layout.image_size!r} does not match geometry size {expected_size!r}"
+        )
+
+
+def _spatial_grid(
+    layout: PatchLayout,
+    *,
+    offset_h: float,
+    offset_w: float,
+    scale_h: float = 1.0,
+    scale_w: float = 1.0,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    grid = layout.grid_size
+    rows = torch.arange(grid.height, dtype=torch.float32, device=device)
+    columns = torch.arange(grid.width, dtype=torch.float32, device=device)
+    (rows, columns) = torch.meshgrid(rows, columns, indexing="ij")
+    return ((rows + 0.5) / scale_h + offset_h, (columns + 0.5) / scale_w + offset_w)
+
+
+def build_image_position_ids(
+    packed: PackedImagePatches,
+    geometries: Sequence[GroupGeometry],
+    *,
+    patch_size: int = 16,
+    n_axis_mode: NAxisMode = NAxisMode.ROLE,
+) -> torch.Tensor:
+    """Build packed ``[n, h, w]`` coordinates on each reference canvas."""
+    _validate_positive_integer("patch_size", patch_size)
+    n_axis_mode = NAxisMode.require(n_axis_mode)
+    if len(geometries) != len(packed.images_per_sample):
+        raise ValueError(
+            f"geometry count does not match packed group count: got {len(geometries)}, expected {len(packed.images_per_sample)}"
+        )
+    device = packed.patches.device
+    position_parts: list[torch.Tensor] = []
+    flat_image_id = 0
+    for sample_id, (geometry, image_count) in enumerate(
+        zip(geometries, packed.images_per_sample, strict=True)
+    ):
+        expected_image_count = 1 + len(geometry.items)
+        if image_count != expected_image_count:
+            raise ValueError(
+                f"group {sample_id} has {image_count} packed images but geometry describes {expected_image_count}"
+            )
+        main_layout = packed.layouts[flat_image_id]
+        _validate_layout(
+            main_layout, geometry.main.size, patch_size, sample_id=sample_id, image_id=0
+        )
+        reference = geometry.reference.size
+        center_h = reference.height / (2 * patch_size)
+        center_w = reference.width / (2 * patch_size)
+        (h, w) = _spatial_grid(
+            main_layout,
+            offset_h=-center_h,
+            offset_w=-center_w,
+            scale_h=geometry.main.scale_h,
+            scale_w=geometry.main.scale_w,
+            device=device,
+        )
+        position_parts.append(torch.stack((torch.zeros_like(h), h, w), dim=-1).reshape(-1, 3))
+        flat_image_id += 1
+        for image_id, item in enumerate(geometry.items, start=1):
+            layout = packed.layouts[flat_image_id]
+            _validate_layout(
+                layout, item.crop.size, patch_size, sample_id=sample_id, image_id=image_id
+            )
+            (h, w) = _spatial_grid(
+                layout,
+                offset_h=item.crop.top / patch_size - center_h,
+                offset_w=item.crop.left / patch_size - center_w,
+                device=device,
+            )
+            n = torch.full_like(h, n_axis_mode.position_for(image_id))
+            position_parts.append(torch.stack((n, h, w), dim=-1).reshape(-1, 3))
+            flat_image_id += 1
+    positions = torch.cat(position_parts)
+    if positions.shape[0] != packed.patches.shape[0]:
+        raise ValueError(
+            f"generated position count does not match packed token count: got {positions.shape[0]}, expected {packed.patches.shape[0]}"
+        )
+    return positions
+
+
+def build_prompt_position_ids(
+    prompt_length: int,
+    image_id: int,
+    geometry: GroupGeometry,
+    *,
+    patch_size: int = 16,
+    n_axis_mode: NAxisMode = NAxisMode.ROLE,
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """Build one prompt's independent sequential ``[n, h, w]`` positions."""
+    _validate_nonnegative_integer("prompt_length", prompt_length)
+    _validate_nonnegative_integer("image_id", image_id)
+    _validate_positive_integer("patch_size", patch_size)
+    n_axis_mode = NAxisMode.require(n_axis_mode)
+    if image_id > len(geometry.items):
+        raise ValueError(
+            f"image_id {image_id} is outside geometry image range [0, {len(geometry.items)}]"
+        )
+    reference = geometry.reference.size
+    if reference.height % patch_size != 0 or reference.width % patch_size != 0:
+        raise ValueError(
+            f"reference size must be divisible by patch_size: got {reference!r} and patch_size={patch_size}"
+        )
+    text_base = math.ceil(
+        max(reference.height / (2 * patch_size), reference.width / (2 * patch_size))
+    )
+    sequence = torch.arange(prompt_length, dtype=torch.float32, device=device).add_(text_base)
+    image_axis = torch.full_like(sequence, n_axis_mode.position_for(image_id))
+    return torch.stack((image_axis, sequence, sequence), dim=-1)
+
+
+def build_timestep_position_ids(
+    geometry: GroupGeometry,
+    *,
+    bank_image_ids: Sequence[int],
+    tokens_per_bank: int = 4,
+    patch_size: int = 16,
+    device: torch.device | str | None = None,
+) -> torch.Tensor:
+    """Place ordered timestep banks before the spatial coordinate range.
+
+    Multiple item banks intentionally share the item-role ``n`` coordinate;
+    their exact image identity is enforced by attention routing rather than by
+    adding an arbitrary item-index position.
+    """
+    if not isinstance(geometry, GroupGeometry):
+        raise TypeError(f"geometry must be a GroupGeometry, got {type(geometry).__name__}")
+    if not isinstance(bank_image_ids, Sequence) or isinstance(bank_image_ids, (str, bytes)):
+        raise TypeError("bank_image_ids must be a sequence of image IDs")
+    bank_image_ids = tuple(bank_image_ids)
+    if not bank_image_ids:
+        raise ValueError("bank_image_ids must contain at least the main image")
+    if any(
+        isinstance(image_id, bool)
+        or not isinstance(image_id, int)
+        or (not 0 <= image_id <= len(geometry.items))
+        for image_id in bank_image_ids
+    ):
+        raise ValueError("bank_image_ids must lie within the geometry image range")
+    if bank_image_ids[0] != 0:
+        raise ValueError("the first timestep bank must target the main image")
+    _validate_positive_integer("tokens_per_bank", tokens_per_bank)
+    _validate_positive_integer("patch_size", patch_size)
+    reference = geometry.reference.size
+    if reference.height % patch_size != 0 or reference.width % patch_size != 0:
+        raise ValueError(
+            f"reference size must be divisible by patch_size: got {reference!r} and patch_size={patch_size}"
+        )
+    sequence_base = math.ceil(
+        max(reference.height / (2 * patch_size), reference.width / (2 * patch_size))
+    )
+    sequence = torch.arange(tokens_per_bank, dtype=torch.float32, device=device).add_(
+        -sequence_base - tokens_per_bank
+    )
+    banks = []
+    for image_id in bank_image_ids:
+        role = float(image_id > 0)
+        role_axis = torch.full_like(sequence, role)
+        banks.append(torch.stack((role_axis, sequence, sequence), dim=-1))
+    return torch.cat(banks)

vendor/pixtext/modeling/transformer.py

@@ -0,0 +1,524 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.attention.flex_attention import BlockMask
+
+from pixtext.modeling.attention import GroupSelfAttention, build_flex_attention_block_mask
+from pixtext.modeling.branches import (
+    GLYPH_BRANCH,
+    ITEM_BRANCH,
+    DirectGlyphPatchProjection,
+    PackedBranchTokenizer,
+    pack_branch_group_sequences,
+)
+from pixtext.modeling.conditioning import extract_image_tokens
+from pixtext.modeling.modes import ImageAttentionRoutingMode, NAxisMode, TextConditioningMode
+from pixtext.modeling.norm import RMSNorm
+from pixtext.modeling.patches import DirectPatchProjection, PackedImagePatches, unpack_image_patches
+from pixtext.modeling.rope import DEFAULT_ROPE_AXES_DIMS, DEFAULT_ROPE_THETA, RotaryEmbedding3D
+from pixtext.preprocessing.geometry import GroupGeometry
+from pixtext.preprocessing.glyphs import GlyphCondition
+from pixtext.sampling.timesteps import GroupTimesteps
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def _validate_positive_integer(name: str, value: int) -> None:
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError(f"{name} must be a positive integer, got {value!r}")
+
+
+def _validate_probability(name: str, value: float) -> float:
+    if (
+        isinstance(value, bool)
+        or not isinstance(value, (int, float))
+        or (not math.isfinite(value))
+        or (not 0 <= value < 1)
+    ):
+        raise ValueError(f"{name} must be finite and in [0, 1), got {value!r}")
+    return float(value)
+
+
+def _validate_packed_timestep_routing(
+    tokens: torch.Tensor,
+    timestep_embeddings: torch.Tensor,
+    conditioning_ids: torch.Tensor,
+    *,
+    hidden_size: int,
+) -> None:
+    if not isinstance(tokens, torch.Tensor):
+        raise TypeError(f"tokens must be a torch.Tensor, got {type(tokens).__name__}")
+    if tokens.ndim != 2 or tokens.shape[1] != hidden_size:
+        raise ValueError(
+            f"tokens must use (tokens, hidden_size) layout with hidden_size {hidden_size}, got shape {tuple(tokens.shape)}"
+        )
+    if not tokens.is_floating_point():
+        raise ValueError(f"tokens must be floating point, got {tokens.dtype}")
+    if not isinstance(timestep_embeddings, torch.Tensor):
+        raise TypeError(
+            f"timestep_embeddings must be a torch.Tensor, got {type(timestep_embeddings).__name__}"
+        )
+    if (
+        timestep_embeddings.ndim != 2
+        or timestep_embeddings.shape[0] == 0
+        or timestep_embeddings.shape[1] != hidden_size
+    ):
+        raise ValueError(
+            f"timestep_embeddings must use (images, hidden_size) layout with at least one image and hidden_size {hidden_size}, got shape {tuple(timestep_embeddings.shape)}"
+        )
+    if not timestep_embeddings.is_floating_point():
+        raise ValueError(
+            f"timestep_embeddings must be floating point, got {timestep_embeddings.dtype}"
+        )
+    if not isinstance(conditioning_ids, torch.Tensor):
+        raise TypeError(
+            f"conditioning_ids must be a torch.Tensor, got {type(conditioning_ids).__name__}"
+        )
+    if conditioning_ids.shape != (tokens.shape[0],):
+        raise ValueError(
+            f"conditioning_ids must contain {tokens.shape[0]} entries, got shape {tuple(conditioning_ids.shape)}"
+        )
+    if conditioning_ids.dtype != torch.int64:
+        raise ValueError(f"conditioning_ids must use torch.int64, got {conditioning_ids.dtype}")
+    if tokens.dtype != timestep_embeddings.dtype:
+        raise ValueError(
+            f"tokens and timestep_embeddings must share a dtype, got {tokens.dtype} and {timestep_embeddings.dtype}"
+        )
+    if tokens.device != timestep_embeddings.device or tokens.device != conditioning_ids.device:
+        raise ValueError("tokens, timestep_embeddings, and conditioning_ids must share a device")
+
+
+def _route_packed_timestep_embeddings(
+    tokens: torch.Tensor,
+    timestep_embeddings: torch.Tensor,
+    conditioning_ids: torch.Tensor,
+    *,
+    hidden_size: int,
+) -> torch.Tensor:
+    _validate_packed_timestep_routing(
+        tokens, timestep_embeddings, conditioning_ids, hidden_size=hidden_size
+    )
+    return timestep_embeddings.index_select(0, conditioning_ids)
+
+
+def _route_packed_adaln_parameters(
+    tokens: torch.Tensor,
+    image_conditions: torch.Tensor,
+    conditioning_ids: torch.Tensor,
+    *,
+    modulation: nn.Module,
+    hidden_size: int,
+    parameter_count: int,
+) -> tuple[torch.Tensor, ...]:
+    """Compute AdaLN once per image, then route its parameters to tokens."""
+    _validate_packed_timestep_routing(
+        tokens, image_conditions, conditioning_ids, hidden_size=hidden_size
+    )
+    image_parameters = modulation(image_conditions)
+    expected_shape = (image_conditions.shape[0], parameter_count * hidden_size)
+    if image_parameters.shape != expected_shape:
+        raise ValueError(
+            f"AdaLN modulation must return {expected_shape}, got {tuple(image_parameters.shape)}"
+        )
+    return image_parameters.index_select(0, conditioning_ids).chunk(parameter_count, dim=-1)
+
+
+def _flat_image_conditioning_ids(
+    sample_ids: torch.Tensor, image_ids: torch.Tensor, images_per_sample: tuple[int, ...]
+) -> torch.Tensor:
+    """Map group-local image IDs to the stable flat image order."""
+    group_offsets = []
+    next_offset = 0
+    for image_count in images_per_sample:
+        group_offsets.append(next_offset)
+        next_offset += image_count
+    offset_tensor = torch.tensor(group_offsets, device=sample_ids.device, dtype=torch.int64)
+    return offset_tensor.index_select(0, sample_ids) + image_ids
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000.0):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_freq = t_freq.to(self.mlp[0].weight.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class IntervalEmbedder(TimestepEmbedder):
+    """Zero-initialized additive embedding for MeanFlow interval lengths."""
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__(hidden_size, frequency_embedding_size)
+        nn.init.normal_(self.mlp[0].weight, std=0.02)
+        nn.init.zeros_(self.mlp[0].bias)
+        nn.init.zeros_(self.mlp[2].weight)
+        nn.init.zeros_(self.mlp[2].bias)
+
+    def forward(self, interval):
+        raw = super().forward(interval)
+        zeros = torch.zeros_like(interval)
+        return raw - super().forward(zeros)
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, drop=0.0, bias=True) -> None:
+        super().__init__()
+        hidden_dim = int(hidden_dim * 2 / 3)
+        self.w12 = nn.Linear(dim, 2 * hidden_dim, bias=bias)
+        self.w3 = nn.Linear(hidden_dim, dim, bias=bias)
+        self.ffn_dropout = nn.Dropout(drop)
+
+    def forward(self, x):
+        x12 = self.w12(x)
+        (x1, x2) = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(self.ffn_dropout(hidden))
+
+
+class PackedJiTBlock(nn.Module):
+    """Timestep-conditioned JiT block for one rank-local packed sequence."""
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int = 1280,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        axes_dims: Sequence[int] = DEFAULT_ROPE_AXES_DIMS,
+        theta: float = DEFAULT_ROPE_THETA,
+        proj_dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        if (
+            isinstance(mlp_ratio, bool)
+            or not isinstance(mlp_ratio, (int, float))
+            or (not math.isfinite(mlp_ratio))
+            or (mlp_ratio <= 0)
+        ):
+            raise ValueError(f"mlp_ratio must be positive and finite, got {mlp_ratio!r}")
+        self.hidden_size = hidden_size
+        self.proj_dropout = _validate_probability("proj_dropout", proj_dropout)
+        self.norm1 = RMSNorm(hidden_size, eps=1e-06)
+        self.attn = GroupSelfAttention(
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            axes_dims=axes_dims,
+            theta=theta,
+            proj_dropout=self.proj_dropout,
+        )
+        self.norm2 = RMSNorm(hidden_size, eps=1e-06)
+        self.mlp = SwiGLUFFN(hidden_size, int(hidden_size * mlp_ratio), drop=self.proj_dropout)
+        for linear in (self.mlp.w12, self.mlp.w3):
+            nn.init.xavier_uniform_(linear.weight)
+            nn.init.zeros_(linear.bias)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+        nn.init.zeros_(self.adaLN_modulation[-1].weight)
+        nn.init.zeros_(self.adaLN_modulation[-1].bias)
+
+    def _timestep_modulations(
+        self,
+        tokens: torch.Tensor,
+        timestep_embeddings: torch.Tensor,
+        conditioning_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, ...]:
+        return _route_packed_adaln_parameters(
+            tokens,
+            timestep_embeddings,
+            conditioning_ids,
+            modulation=self.adaLN_modulation,
+            hidden_size=self.hidden_size,
+            parameter_count=6,
+        )
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        timestep_embeddings: torch.Tensor,
+        conditioning_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        block_mask: BlockMask,
+        rotary_factors: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Transform packed tokens using per-image timesteps and fixed routing."""
+        (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = (
+            self._timestep_modulations(tokens, timestep_embeddings, conditioning_ids)
+        )
+        attention_input = self.norm1(tokens) * (1 + scale_msa) + shift_msa
+        tokens = tokens + gate_msa * self.attn(
+            attention_input, position_ids, block_mask, rotary_factors
+        )
+        mlp_input = self.norm2(tokens) * (1 + scale_mlp) + shift_mlp
+        return tokens + gate_mlp * self.mlp(mlp_input)
+
+
+class PackedFinalLayer(nn.Module):
+    """Predict flattened pixel patches from packed image tokens."""
+
+    def __init__(
+        self, *, hidden_size: int = 1280, patch_size: int = 16, out_channels: int = 3
+    ) -> None:
+        super().__init__()
+        for name, value in (
+            ("hidden_size", hidden_size),
+            ("patch_size", patch_size),
+            ("out_channels", out_channels),
+        ):
+            _validate_positive_integer(name, value)
+        self.hidden_size = hidden_size
+        self.patch_size = patch_size
+        self.out_channels = out_channels
+        self.norm_final = RMSNorm(hidden_size, eps=1e-06)
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)
+        )
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        nn.init.zeros_(self.adaLN_modulation[-1].weight)
+        nn.init.zeros_(self.adaLN_modulation[-1].bias)
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        timestep_embeddings: torch.Tensor,
+        conditioning_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Predict one flattened pixel patch per packed image token."""
+        (shift, scale) = _route_packed_adaln_parameters(
+            tokens,
+            timestep_embeddings,
+            conditioning_ids,
+            modulation=self.adaLN_modulation,
+            hidden_size=self.hidden_size,
+            parameter_count=2,
+        )
+        normalized = self.norm_final(tokens) * (1 + scale) + shift
+        return self.linear(normalized)
+
+
+class GroupJiT(nn.Module):
+    """The glyph-conditioned Group JiT architecture used by both releases."""
+
+    def __init__(
+        self,
+        *,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=1280,
+        depth=32,
+        num_heads=16,
+        mlp_ratio=4.0,
+        axes_dims=(12, 34, 34),
+        theta=10000.0,
+        interval_conditioning=False,
+        image_attention_routing="items_interread_glyph_reads_own_item_bidirectional_main",
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.depth = depth
+        self.text_conditioning = TextConditioningMode.NONE
+        self.n_axis_mode = NAxisMode.ROLE
+        self.image_attention_routing = ImageAttentionRoutingMode(image_attention_routing)
+        self.interval_conditioning = interval_conditioning
+        self.cfg_conditioning = False
+        self.cfg_embedder = None
+        self.rotary = RotaryEmbedding3D(axes_dims=axes_dims, theta=theta)
+        self.patch_projection = DirectPatchProjection(
+            patch_size=patch_size, in_channels=in_channels, hidden_size=hidden_size
+        )
+        self.glyph_projection = DirectGlyphPatchProjection(
+            patch_size=patch_size, hidden_size=hidden_size
+        )
+        self.branch_tokenizer = PackedBranchTokenizer(
+            rgb_projection=self.patch_projection, glyph_projection=self.glyph_projection
+        )
+        self.timestep_embedder = TimestepEmbedder(hidden_size)
+        self.interval_embedder = IntervalEmbedder(hidden_size) if interval_conditioning else None
+        self.blocks = nn.ModuleList(
+            (
+                PackedJiTBlock(
+                    hidden_size=hidden_size,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    axes_dims=axes_dims,
+                    theta=theta,
+                )
+                for _ in range(depth)
+            )
+        )
+        self.final_layer = PackedFinalLayer(
+            hidden_size=hidden_size, patch_size=patch_size, out_channels=in_channels
+        )
+
+    def forward(
+        self,
+        *,
+        packed_images,
+        prompts,
+        geometries,
+        timesteps,
+        glyph_conditions,
+        image_intervals=None,
+        return_packed_patches=False,
+    ):
+        if prompts is not None:
+            raise ValueError("PixText target text is specified by glyph images")
+        if glyph_conditions is None:
+            raise ValueError("glyph_conditions are required")
+        if self.interval_conditioning:
+            if image_intervals is None or image_intervals.shape != timesteps.images.shape:
+                raise ValueError("MeanFlow requires one interval per RGB image")
+        elif image_intervals is not None:
+            raise ValueError("Teacher does not accept interval conditioning")
+        return self._forward_canonical_glyph(
+            packed_images=packed_images,
+            prompts=None,
+            geometries=geometries,
+            timesteps=timesteps,
+            glyph_conditions=glyph_conditions,
+            return_packed_patches=return_packed_patches,
+            image_intervals=image_intervals,
+            cfg_scales=None,
+        )
+
+    def _forward_canonical_glyph(
+        self,
+        *,
+        packed_images: PackedImagePatches,
+        prompts: None,
+        geometries: Sequence[GroupGeometry],
+        timesteps: GroupTimesteps,
+        glyph_conditions: Sequence[Sequence[GlyphCondition | None]],
+        return_packed_patches: bool,
+        image_intervals: torch.Tensor | None,
+        cfg_scales: torch.Tensor | None,
+    ):
+        """Run mixed RGB/glyph branches while returning only RGB predictions."""
+        if self.branch_tokenizer is None or self.glyph_projection is None:
+            raise RuntimeError("canonical glyph tokenizer is not configured")
+        branches = self.branch_tokenizer.tokenize(
+            packed_rgb=packed_images, glyphs=glyph_conditions, geometries=geometries
+        )
+        projected_prompts = None
+        sequence = pack_branch_group_sequences(
+            branches=branches, prompts=projected_prompts, geometries=geometries
+        )
+        assert sequence.branch_kinds is not None
+        assert sequence.slot_ids is not None
+        sequence_timesteps = torch.ones(
+            sequence.tokens.shape[0], dtype=timesteps.images.dtype, device=sequence.tokens.device
+        )
+        flat_image_offset = 0
+        for sample_id, image_count in enumerate(packed_images.images_per_sample):
+            sample_mask = sequence.sample_ids.eq(sample_id)
+            for slot_id in range(image_count - 1):
+                item_mask = (
+                    sample_mask
+                    & sequence.branch_kinds.eq(ITEM_BRANCH)
+                    & sequence.slot_ids.eq(slot_id)
+                )
+                sequence_timesteps[item_mask] = timesteps.images[flat_image_offset + slot_id + 1]
+            flat_image_offset += image_count
+        if not bool(sequence_timesteps[sequence.branch_kinds.eq(GLYPH_BRANCH)].eq(1).all()):
+            raise RuntimeError("glyph timestep must remain fixed at 1")
+        timestep_embeddings = self.timestep_embedder(sequence_timesteps)
+        if self.interval_embedder is not None:
+            if image_intervals is None:
+                raise RuntimeError("interval-aware glyph model requires intervals")
+            sequence_intervals = torch.zeros_like(sequence_timesteps)
+            flat_image_offset = 0
+            for sample_id, image_count in enumerate(packed_images.images_per_sample):
+                sample_mask = sequence.sample_ids.eq(sample_id)
+                for slot_id in range(image_count - 1):
+                    item_mask = (
+                        sample_mask
+                        & sequence.branch_kinds.eq(ITEM_BRANCH)
+                        & sequence.slot_ids.eq(slot_id)
+                    )
+                    sequence_intervals[item_mask] = image_intervals[flat_image_offset + slot_id + 1]
+                flat_image_offset += image_count
+            timestep_embeddings = timestep_embeddings + self.interval_embedder(sequence_intervals)
+        if self.cfg_embedder is not None:
+            timestep_embeddings = timestep_embeddings + self.cfg_embedder(
+                cfg_scales[sequence.sample_ids]
+            )
+        conditioning_ids = torch.arange(
+            sequence.tokens.shape[0], dtype=torch.int64, device=sequence.tokens.device
+        )
+        sequence_mask = build_flex_attention_block_mask(
+            sequence, image_attention_routing=self.image_attention_routing
+        )
+        rotary_factors = self.rotary.prepare(sequence.position_ids)
+        sequence_tokens = sequence.tokens
+        for block in self.blocks:
+            sequence_tokens = block(
+                sequence_tokens,
+                timestep_embeddings,
+                conditioning_ids,
+                sequence.position_ids,
+                sequence_mask,
+                rotary_factors,
+            )
+        branch_image_tokens = extract_image_tokens(sequence_tokens, sequence)
+        rgb_tokens = branch_image_tokens.index_select(0, branches.rgb_token_indices)
+        rgb_timestep_embeddings = self.timestep_embedder(timesteps.images)
+        if self.interval_embedder is not None:
+            assert image_intervals is not None
+            rgb_timestep_embeddings = rgb_timestep_embeddings + self.interval_embedder(
+                image_intervals
+            )
+        rgb_conditioning_ids = _flat_image_conditioning_ids(
+            packed_images.sample_ids, packed_images.image_ids, packed_images.images_per_sample
+        )
+        if self.final_layer is None:
+            raise RuntimeError("glyph conditioning requires the standard final layer")
+        predicted_patches = self.final_layer(
+            rgb_tokens, rgb_timestep_embeddings, rgb_conditioning_ids
+        )
+        if return_packed_patches:
+            return predicted_patches
+        return unpack_image_patches(predicted_patches, packed_images)