lance_mlx/lance.py

"""Lance MLX wrapper — unified multimodal model built on mlx-vlm's Qwen2.5-VL.

Architecture (mirrors bytedance/Lance):
                          ┌────────────────────────────────────┐
    VAE latents ─► vae2llm ─►                                  ─► llm2vae ─► VAE latents
                          │      Qwen2.5-VL LLM (mlx-vlm)      │
    text tokens ─────────►│   shared hidden space (2048 dim)   │
                          │                                    │
    image patches ──► ViT ►│                                    │
                          └────────────────────────────────────┘
                                       ▲
                                       │
                                  TimestepEmbedder
                                  PositionEmbedding3D

Reuses mlx-vlm.models.qwen2_5_vl for the LLM + ViT backbone (saves ~2000 lines
of port work). Lance-specific additions live here:

  - vae2llm / llm2vae linear adapters
  - TimestepEmbedder + PositionEmbedding3D (already in modeling_utils.py)
  - Flow-matching sampler (validation_gen)         [TODO — Phase 1.5]
  - Classifier-free guidance                       [TODO — Phase 1.5]
  - KV-cache-aware generation loop                 [TODO — Phase 2]
  - NaViT variable-resolution image packing        [TODO — Phase 2]

Status (2026-05-19): primitives + config + adapter wiring + smoke test only.
Sampler is a stub.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional, Tuple
from typing import Optional as _Optional  # silence linter — already imported

import mlx.core as mx
import mlx.nn as nn

from mlx_vlm.models.qwen2_5_vl.config import ModelConfig as Qwen25VLConfig
from mlx_vlm.models.qwen2_5_vl.vision import VisionModel as Qwen25VLVisionModel

from .modeling_utils import TimestepEmbedder, PositionEmbedding3D
from .qwen2_navit_mlx import LanceLanguageModel, KVCache


# ---------------------------------------------------------------------------
# LanceConfig — additions on top of Qwen2.5-VL config
# ---------------------------------------------------------------------------

@dataclass
class LanceConfig:
    qwen_config: Qwen25VLConfig
    visual_gen: bool = True
    visual_und: bool = True
    latent_patch_size: Tuple[int, int, int] = (1, 1, 1)   # (pt, ph, pw) — Lance prod default
    max_latent_size: int = 64                              # max H,W per axis in latents
    max_num_frames: int = 0                                # max video frames before VAE (image variant: 0)
    max_num_latent_frames_override: Optional[int] = None   # if set, used instead of derived
    latent_channel: int = 48                               # Wan 2.2 VAE z_dim
    vae_downsample_spatial: int = 16
    vae_downsample_temporal: int = 4
    vit_max_num_patch_per_side: int = 70
    connector_act: str = "gelu_pytorch_tanh"
    timestep_shift: float = 1.0
    pos_safety: int = 1024                                 # text-before-video pos shift

    @property
    def hidden_size(self) -> int:
        return self.qwen_config.text_config.hidden_size

    @property
    def patch_latent_dim(self) -> int:
        pt, ph, pw = self.latent_patch_size
        return pt * ph * pw * self.latent_channel

    @property
    def max_num_latent_frames(self) -> int:
        if self.max_num_latent_frames_override is not None:
            return self.max_num_latent_frames_override
        return self.max_num_frames // self.vae_downsample_temporal + 1


# ---------------------------------------------------------------------------
# Lance — main model
# ---------------------------------------------------------------------------

class Lance(nn.Module):
    """Unified multimodal model.

    Modes (selected by the forward path used):
      t2i        — text -> image: noise latent + text -> denoised latent
      t2v        — text -> video: noise latent stack + text -> denoised stack
      image_edit — text + image latent -> edited latent
      video_edit — text + video latent -> edited latent
      x2t_image  — image -> text (autoregressive understanding)
      x2t_video  — video -> text
    """

    def __init__(self, config: LanceConfig):
        super().__init__()
        self.config = config

        # Language backbone: Lance's modified Qwen2 with MoE-gen experts.
        # ViT vision tower lives in a separate object (loaded from vit.safetensors)
        # to keep the model.safetensors load path clean.
        self.language_model = LanceLanguageModel(config.qwen_config.text_config)

        if config.visual_gen:
            h = config.hidden_size
            # vae2llm projects flattened VAE latent patches into LLM hidden space.
            self.vae2llm = nn.Linear(config.patch_latent_dim, h)
            # llm2vae projects LLM hidden back to VAE patch space (per-token).
            self.llm2vae = nn.Linear(h, config.patch_latent_dim)
            # DiT-style scalar timestep -> hidden vector, broadcast to latent tokens.
            self.time_embedder = TimestepEmbedder(h)
            # 3D position lookup over (max_t, max_h, max_w) latent positions.
            self.latent_pos_embed = PositionEmbedding3D(
                config.max_num_latent_frames,
                config.max_latent_size,
                h,
            )
            # Safety shift so video latent positions don't collide with text positions.
            self.pos_shift = (
                config.max_latent_size * config.max_latent_size * config.max_num_latent_frames
                + config.pos_safety
            )

        # ViT for understanding is already inside self.backbone.vision_tower.
        # No extra ViT-side parameters here.

    # -----------------------------------------------------------------------
    # Latent patching helpers — match PT data_utils get_flattened_position_ids
    # -----------------------------------------------------------------------

    def patchify_latent(self, latent: mx.array) -> mx.array:
        """Reshape (B, T, H, W, C) -> (B, N, patch_latent_dim) where
        N = T/pt * H/ph * W/pw and each token concatenates the patch volume.
        """
        pt, ph, pw = self.config.latent_patch_size
        B, T, H, W, C = latent.shape
        assert T % pt == 0 and H % ph == 0 and W % pw == 0, (
            f"latent shape ({T},{H},{W}) not divisible by patch {(pt, ph, pw)}"
        )
        # (B, T/pt, pt, H/ph, ph, W/pw, pw, C) -> (B, T/pt, H/ph, W/pw, pt*ph*pw*C)
        x = latent.reshape(B, T // pt, pt, H // ph, ph, W // pw, pw, C)
        x = mx.transpose(x, (0, 1, 3, 5, 2, 4, 6, 7))
        x = x.reshape(B, (T // pt) * (H // ph) * (W // pw), pt * ph * pw * C)
        return x

    def unpatchify_latent(self, tokens: mx.array, t: int, h: int, w: int) -> mx.array:
        """Inverse of patchify_latent.
        tokens: (B, N, patch_latent_dim) with N = t/pt * h/ph * w/pw
        returns: (B, t, h, w, C)
        """
        pt, ph, pw = self.config.latent_patch_size
        C = self.config.latent_channel
        B = tokens.shape[0]
        x = tokens.reshape(B, t // pt, h // ph, w // pw, pt, ph, pw, C)
        x = mx.transpose(x, (0, 1, 4, 2, 5, 3, 6, 7))
        x = x.reshape(B, t, h, w, C)
        return x

    # -----------------------------------------------------------------------
    # Forward (stub) — full diffusion + AR forward is multi-page port. The
    # skeleton below shows the building blocks; the sampler glue is TODO.
    # -----------------------------------------------------------------------

    def embed_latent_tokens(
        self,
        noisy_latent: mx.array,         # (B, T, H, W, C)
        timesteps: mx.array,            # (B,) flow-matching scalar in [0, 1]
        position_ids: Optional[mx.array] = None,
    ) -> mx.array:
        """Convert (noisy) VAE latents into LLM-hidden-space tokens.

        Steps mirror PT Lance:
          1. Patchify latents into (B, N, patch_latent_dim).
          2. Project through vae2llm -> (B, N, hidden).
          3. Add 3D position embedding looked up by `position_ids` (or default
             flattened ordering).
          4. Add broadcast TimestepEmbedder output so every latent token sees
             the current denoising step.
        """
        B, T, H, W, _ = noisy_latent.shape
        x = self.patchify_latent(noisy_latent)                  # (B, N, plat)
        x = self.vae2llm(x)                                     # (B, N, hidden)

        if position_ids is None:
            # Compute 3D (t, h, w) coords for each token in raster order, then
            # flatten via pe_index = t * (max_h * max_w) + h * max_w + w.
            # This indexes the full (max_t × max_h × max_w) sinusoid table so
            # the 3D structure of the latent grid is preserved (otherwise the
            # decoder gets only contiguous flat positions covering the first
            # few rows of the full grid → horizontal striping artefacts).
            pt, ph, pw = self.config.latent_patch_size
            tH, tW = H // ph, W // pw
            tT = T // pt
            max_h = self.latent_pos_embed.max_latent_size
            max_w = max_h
            t_coords = mx.repeat(mx.arange(tT, dtype=mx.int32), tH * tW)
            h_coords = mx.repeat(mx.tile(mx.arange(tH, dtype=mx.int32), (tT,)), tW)
            w_coords = mx.tile(mx.arange(tW, dtype=mx.int32), (tT * tH,))
            flat_ids = t_coords * (max_h * max_w) + h_coords * max_w + w_coords
            position_ids = mx.broadcast_to(flat_ids[None], (B, flat_ids.shape[0]))
        pos = self.latent_pos_embed(position_ids)              # (B, N, hidden)
        x = x + pos

        t_emb = self.time_embedder(timesteps)                  # (B, hidden)
        x = x + t_emb[:, None, :]                              # broadcast over N
        return x

    def project_latent_out(self, hidden: mx.array, t: int, h: int, w: int) -> mx.array:
        """Inverse of embed_latent_tokens (post-LLM): project hidden -> patches -> latent.
        hidden: (B, N, hidden)
        returns: (B, t, h, w, C)
        """
        patches = self.llm2vae(hidden)                          # (B, N, plat)
        return self.unpatchify_latent(patches, t, h, w)

    # -----------------------------------------------------------------------
    # T2I sampler — minimum viable port of validation_gen for text-to-image.
    # Single sample, single resolution, no CFG (cfg_scale=1.0), no KV cache.
    # Flow-matching denoising loop following the PT timestep_shift convention.
    # -----------------------------------------------------------------------

    def _build_sequence(
        self,
        prompt_token_ids: mx.array,      # (P,) int32 — raw user prompt tokens (no specials)
        latent_shape: Tuple[int, int, int],  # (T_lat, H_lat, W_lat)
        timestep: mx.array,              # (1,) in [0, 1]
        noisy_latent: mx.array,          # (1, T, H, W, C)
        special_token_ids: dict,         # {"bos": int, "eos": int, "start_of_image": int, "end_of_image": int, "image_token_id": int}
    ) -> Tuple[mx.array, mx.array, int, int]:
        """Build the packed sequence matching PT Lance T2I layout:
            <|im_start|> [prompt] <|im_end|> <|vision_start|> [N image_token placeholders] <|vision_end|>
        At the placeholder positions the embedding is REPLACED with the
        VAE-derived embedding (vae2llm + 3D pos + time).

        Returns:
          inputs_embeds   (1, S, H)
          attn_mask       (S, S)   additive
          latent_start    int      first latent-token index in S
          n_lat           int      number of latent tokens
        """
        t_lat, h_lat, w_lat = latent_shape
        pt, ph, pw = self.config.latent_patch_size
        n_lat = (t_lat // pt) * (h_lat // ph) * (w_lat // pw)

        bos = special_token_ids["bos"]
        eos = special_token_ids["eos"]
        soi = special_token_ids["start_of_image"]
        eoi = special_token_ids["end_of_image"]
        img = special_token_ids["image_token_id"]

        # Build the full token id sequence with vision wrappers + image placeholders.
        prompt_list = prompt_token_ids.tolist()
        full_ids = [bos] + prompt_list + [eos, soi] + [img] * n_lat + [eoi]
        full_arr = mx.array(full_ids, dtype=mx.int32)
        full_emb = self.language_model.model.embed_tokens(full_arr[None])  # (1, S, H)

        # Where the image_token placeholders sit in the sequence:
        latent_start = 1 + len(prompt_list) + 2     # bos + prompt + eos + soi
        latent_end = latent_start + n_lat            # exclusive
        S = full_emb.shape[1]

        # Compute the latent embedding (vae2llm + 3D pos + time) and splice it in.
        latent_tokens = self.embed_latent_tokens(noisy_latent, timestep)  # (1, n_lat, H)

        before = full_emb[:, :latent_start]                   # bos..soi inclusive
        after = full_emb[:, latent_end:]                       # eoi
        inputs_embeds = mx.concatenate([before, latent_tokens, after], axis=1)

        # Attention mask:
        #   text positions ∈ [0, latent_start) ∪ [latent_end, S): causal within themselves
        #   latent positions ∈ [latent_start, latent_end): bidirectional within block;
        #     latent can attend back to all earlier text; text cannot peek forward to latent.
        #   end_of_image at latent_end attends causally over everything before it.
        neg_inf = -1e9
        idx = mx.arange(S)
        i_grid, j_grid = mx.meshgrid(idx, idx, indexing="ij")
        # 1) Strict causal: j > i is blocked …
        causal_block = j_grid > i_grid
        # 2) … EXCEPT within the latent block, where latent-to-latent is bidirectional
        i_in_lat = (i_grid >= latent_start) & (i_grid < latent_end)
        j_in_lat = (j_grid >= latent_start) & (j_grid < latent_end)
        both_in_lat = i_in_lat & j_in_lat
        bad = causal_block & ~both_in_lat
        mask = mx.where(bad, neg_inf, 0.0)

        return inputs_embeds, mask, latent_start, n_lat

    def _forward_backbone(
        self,
        inputs_embeds: mx.array,         # (1, S, H)
        attn_mask: mx.array,             # (S, S) additive
        text_len: int,                   # number of leading text tokens (rest = VAE gen tokens)
        position_ids: mx.array,          # (3, 1, S) mrope position ids
    ) -> mx.array:
        """Run Lance's modified Qwen2 language model. Returns hidden (1, S, H)."""
        return self.language_model(
            inputs_embeds=inputs_embeds,
            text_len=text_len,
            position_ids=position_ids,
            mask=attn_mask,
        )

    def _build_position_ids(
        self,
        latent_start: int,
        latent_grid: Tuple[int, int, int],   # (t_lat, h_lat, w_lat)
        S: int,                              # full sequence length
    ) -> mx.array:
        """Build mrope position_ids of shape (3, 1, S).

        Layout: text positions 0..latent_start-1 (bos+prompt+eos+soi),
                then n_lat latent positions with 3D coords (t,h,w) offset
                by latent_start so they don't collide with text positions,
                then trailing text positions (eoi etc.) continue past max(latent).

        For mrope, axis 0=T positions, 1=H positions, 2=W positions.
        Text tokens use the same scalar across all 3 axes.
        Latent tokens use distinct (T, H, W) coords per the 3D grid.
        """
        t_lat, h_lat, w_lat = latent_grid
        n_lat = t_lat * h_lat * w_lat

        # Text positions before latent
        text_pre = mx.arange(latent_start, dtype=mx.int32)         # 0..latent_start-1
        # Latent 3D coords (broadcast to a flat list of (T, H, W) per token)
        t_coords = mx.repeat(mx.arange(t_lat, dtype=mx.int32), h_lat * w_lat)        # t * h * w
        h_coords = mx.repeat(mx.tile(mx.arange(h_lat, dtype=mx.int32), (t_lat,)), w_lat)  # ...
        w_coords = mx.tile(mx.arange(w_lat, dtype=mx.int32), (t_lat * h_lat,))
        # Shift latent positions past text
        t_coords = t_coords + latent_start
        h_coords = h_coords + latent_start
        w_coords = w_coords + latent_start

        # Text positions after latent (eoi and any trailing). They continue from
        # max(latent_coord) + 1 so they're strictly after the latent block.
        n_post = S - latent_start - n_lat
        post_start = latent_start + max(t_lat, h_lat, w_lat)
        text_post = mx.arange(post_start, post_start + n_post, dtype=mx.int32)

        # Assemble per-axis position id sequences (length S each)
        # Pre-text and post-text use the same scalar on all axes.
        axis_t = mx.concatenate([text_pre, t_coords, text_post])
        axis_h = mx.concatenate([text_pre, h_coords, text_post])
        axis_w = mx.concatenate([text_pre, w_coords, text_post])

        # Stack into (3, 1, S)
        ids = mx.stack([axis_t, axis_h, axis_w], axis=0)           # (3, S)
        return ids[:, None, :]                                     # (3, 1, S)

    def denoise_step(
        self,
        prompt_token_ids: mx.array,
        noisy_latent: mx.array,
        timestep: mx.array,
        special_token_ids: dict,
    ) -> mx.array:
        """One denoising step. Returns predicted velocity in latent space (1, T, H, W, C)."""
        B, T, H, W, C = noisy_latent.shape
        assert B == 1, "denoise_step is single-sample"
        inputs_embeds, attn_mask, latent_start, n_lat = self._build_sequence(
            prompt_token_ids, (T, H, W), timestep, noisy_latent, special_token_ids,
        )
        S = inputs_embeds.shape[1]
        position_ids = self._build_position_ids(latent_start, (T, H, W), S)
        # text_len for MoE-gen routing: tokens before the latent block use normal weights.
        # The latent block uses moe_gen. The trailing end_of_image uses normal weights —
        # but at inference the eoi is only 1 token at the end. For our route, we treat
        # everything outside [latent_start, latent_start + n_lat) as text. To match the
        # PT routing (gen tokens are ONLY the latent block), we slice the latent block
        # explicitly in the attention/MLP routing inside qwen2_navit_mlx.
        hidden = self._forward_backbone(inputs_embeds, attn_mask, latent_start, position_ids)
        latent_hidden = hidden[:, latent_start:latent_start + n_lat, :]
        predicted_v = self.project_latent_out(latent_hidden, T, H, W)
        return predicted_v

    def sample_t2i(
        self,
        prompt_token_ids: mx.array,      # (P,) int32 — raw prompt tokens
        latent_shape: Tuple[int, int, int],  # (T_lat=1, H_lat, W_lat)
        special_token_ids: dict,         # bos/eos/start_of_image/end_of_image/image_token_id
        num_steps: int = 30,
        timestep_shift: float = 3.0,
        seed: Optional[int] = None,
        cfg_scale: float = 1.0,
        uncond_token_ids: Optional[mx.array] = None,  # for CFG; if None and cfg>1, uses just [bos, eos]
    ) -> mx.array:
        """Flow-matching denoising loop. Mirrors PT lance.py:598 exactly.

        Returns final latent (1, T, H, W, C).
        """
        t_lat, h_lat, w_lat = latent_shape
        C = self.config.latent_channel

        if seed is not None:
            mx.random.seed(seed)
        x = mx.random.normal(shape=(1, t_lat, h_lat, w_lat, C))

        # PT schedule:
        #   ts = linspace(1, 0, num_steps + 1)
        #   ts = shift * ts / (1 + (shift - 1) * ts)
        #   dts = ts[:-1] - ts[1:]
        #   loop over (timesteps[:-1], dts)
        ts = mx.linspace(1.0, 0.0, num_steps + 1)
        ts = timestep_shift * ts / (1.0 + (timestep_shift - 1.0) * ts)
        dts = ts[:-1] - ts[1:]
        ts = ts[:-1]

        for i in range(num_steps):
            t_arr = ts[i:i+1]                 # shape (1,)
            v = self.denoise_step(prompt_token_ids, x, t_arr, special_token_ids)
            if cfg_scale != 1.0:
                if uncond_token_ids is None:
                    uncond_token_ids = mx.array(
                        [special_token_ids["bos"], special_token_ids["eos"]], dtype=mx.int32
                    )[:0]   # empty prompt
                v_uncond = self.denoise_step(uncond_token_ids, x, t_arr, special_token_ids)
                v = v_uncond + cfg_scale * (v - v_uncond)
            x = x - dts[i].item() * v
            mx.eval(x)
        return x

    # =======================================================================
    # TI2I (text + image → image) editing.
    # Sequence: <|im_start|>[prompt]<|im_end|><|vision_start|>[ViT tokens]<|vision_end|>
    #           <|vision_start|>[N target-noise placeholders]<|vision_end|>
    # Routing:  ViT tokens are UND (normal weights); target-noise is GEN (moe_gen).
    # =======================================================================

    def _build_edit_sequence(
        self,
        prompt_token_ids: mx.array,
        visual_embeds: mx.array,             # (1, N_vit, hidden) input-image features from ViT
        image_grid_hw: Tuple[int, int],      # post-merge (h, w) of the input image
        cond_latent: Optional[mx.array],     # (1, T_c, H_c, W_c, C) input image's VAE latent (or None)
        latent_shape: Tuple[int, int, int],  # target (T_lat, H_lat, W_lat)
        timestep: mx.array,                  # (1,)
        noisy_latent: mx.array,              # (1, T, H, W, C)
        special_token_ids: dict,
    ) -> Tuple[mx.array, mx.array, int, int, Tuple[int, int], Optional[Tuple[int, int, Tuple[int, int, int]]]]:
        """Returns (inputs_embeds, attn_mask, latent_start, n_lat, (vit_start, vit_end), cond_block_info).

        cond_block_info is (cond_start, cond_end, (t_c, h_c, w_c)) if cond_latent
        is provided, else None.
        """
        t_lat, h_lat, w_lat = latent_shape
        pt, ph, pw = self.config.latent_patch_size
        n_lat = (t_lat // pt) * (h_lat // ph) * (w_lat // pw)

        bos = special_token_ids["bos"]
        eos = special_token_ids["eos"]
        soi = special_token_ids["start_of_image"]
        eoi = special_token_ids["end_of_image"]
        img = special_token_ids["image_token_id"]

        # Compute condition-latent embedding (no noise, timestep=0)
        cond_emb = None
        cond_grid = None
        n_cond = 0
        if cond_latent is not None:
            t_c, h_c, w_c = cond_latent.shape[1], cond_latent.shape[2], cond_latent.shape[3]
            cond_grid = (t_c, h_c, w_c)
            n_cond = (t_c // pt) * (h_c // ph) * (w_c // pw)
            zero_t = mx.zeros((1,), dtype=mx.float32)
            cond_emb = self.embed_latent_tokens(cond_latent, zero_t)  # (1, n_cond, H)

        prompt_list = prompt_token_ids.tolist()
        N_vit = visual_embeds.shape[1]
        # Lance chat template for edit (matches PT render_qwenvl_prompt + edit system prompt):
        #   <|im_start|>system\n<Lance edit system prompt><|im_end|>\n
        #   <|im_start|>user\n<|vision_start|>[N_vit IPADs]<|vision_end|><instruction><|im_end|>\n
        #   <|im_start|>assistant\n<|vision_start|>[n_cond IPADs]<|vision_end|><|vision_start|>[n_lat IPADs]<|vision_end|>
        # The caller must include the system+user header tokens in `prompt_token_ids` already
        # (we expect prompt_token_ids = pre-tokenized header).  The structure here only adds
        # the special-token wrappers around the visual blocks.

        # `prompt_token_ids` from the caller IS just the user instruction.  We build the
        # template wrappers here.  Caller passes special_token_ids with extras:
        #   "sys_prefix_ids": list[int] for "<|im_start|>system\n<Lance edit sys>\n<|im_end|>\n<|im_start|>user\n"
        #   "user_to_assistant_ids": list[int] for "<|im_end|>\n<|im_start|>assistant\n"
        sys_prefix = special_token_ids.get("sys_prefix_ids") or [bos]
        u2a       = special_token_ids.get("user_to_assistant_ids") or [eos]

        ids_pre_vit_block  = list(sys_prefix) + [soi]
        ids_vit_pad        = [img] * N_vit
        ids_vit_close      = [eoi]
        ids_instr          = list(prompt_list) + list(u2a)            # instruction then "<|im_end|>\n<|im_start|>assistant\n"
        ids_cond_block     = ([soi] + [img] * n_cond + [eoi]) if cond_emb is not None else []
        ids_lat_open       = [soi]
        ids_lat_pad        = [img] * n_lat
        ids_lat_close      = [eoi]
        full_ids = (ids_pre_vit_block + ids_vit_pad + ids_vit_close + ids_instr
                    + ids_cond_block + ids_lat_open + ids_lat_pad + ids_lat_close)
        full_arr = mx.array(full_ids, dtype=mx.int32)
        full_emb = self.language_model.model.embed_tokens(full_arr[None])  # (1, S, H)

        # Splice indices
        vit_start = len(ids_pre_vit_block)
        vit_end   = vit_start + N_vit
        instr_end = vit_end + len(ids_vit_close) + len(ids_instr)
        if cond_emb is not None:
            cond_start = instr_end + 1                                  # +1 for the cond soi
            cond_end   = cond_start + n_cond
            lat_start  = cond_end + 1 + 1                               # eoi + lat soi
        else:
            cond_start = cond_end = None
            lat_start  = instr_end + 1                                  # +1 for the lat soi
        lat_end = lat_start + n_lat
        S = full_emb.shape[1]

        # Compute target latent embedding (vae2llm + 3D pos + time)
        latent_tokens = self.embed_latent_tokens(noisy_latent, timestep)

        # Splice ViT (+ optional cond) + target latent embeddings into the full embedding
        chunks = [full_emb[:, :vit_start], visual_embeds]
        if cond_emb is not None:
            chunks += [full_emb[:, vit_end:cond_start], cond_emb,
                       full_emb[:, cond_end:lat_start], latent_tokens,
                       full_emb[:, lat_end:]]
        else:
            chunks += [full_emb[:, vit_end:lat_start], latent_tokens,
                       full_emb[:, lat_end:]]
        inputs_embeds = mx.concatenate(chunks, axis=1)
        assert inputs_embeds.shape[1] == S, f"shape mismatch: {inputs_embeds.shape[1]} vs {S}"

        # Attention mask: causal everywhere EXCEPT inside the target latent block,
        # which is bidirectional (matches the T2I behavior).
        neg_inf = -1e9
        idx = mx.arange(S)
        i_grid, j_grid = mx.meshgrid(idx, idx, indexing="ij")
        causal_block = j_grid > i_grid
        i_in_lat = (i_grid >= lat_start) & (i_grid < lat_end)
        j_in_lat = (j_grid >= lat_start) & (j_grid < lat_end)
        both_in_lat = i_in_lat & j_in_lat
        bad = causal_block & ~both_in_lat
        mask = mx.where(bad, neg_inf, 0.0)

        cond_info = None
        if cond_emb is not None:
            cond_info = (cond_start, cond_end, cond_grid)
        return inputs_embeds, mask, lat_start, n_lat, (vit_start, vit_end), cond_info

    def _build_edit_position_ids(
        self,
        S: int,
        vit_block: Tuple[int, int],       # (start, end) of ViT-token positions in sequence
        vit_grid_hw: Tuple[int, int],     # post-merge (h, w) for the input image
        lat_start: int,
        lat_grid: Tuple[int, int, int],   # target (t_lat, h_lat, w_lat)
        cond_info: Optional[Tuple[int, int, Tuple[int, int, int]]] = None,  # (cond_start, cond_end, (t,h,w))
    ) -> mx.array:
        """Build (3, 1, S) mrope position ids with proper 3D coords inside BOTH
        the ViT block (input image) and the latent block (target image).
        """
        vit_s, vit_e = vit_block
        n_vit = vit_e - vit_s
        h_vit, w_vit = vit_grid_hw
        t_lat, h_lat, w_lat = lat_grid
        n_lat = t_lat * h_lat * w_lat

        # Pre-vit text positions: 0..vit_s-1
        pre = mx.arange(vit_s, dtype=mx.int32)
        # ViT block 3D positions (offset by vit_s, T=0)
        v_t = mx.zeros((n_vit,), dtype=mx.int32) + vit_s
        v_h = mx.repeat(mx.arange(h_vit, dtype=mx.int32), w_vit) + vit_s
        v_w = mx.tile(mx.arange(w_vit, dtype=mx.int32), (h_vit,)) + vit_s
        if cond_info is None:
            # Mid block (eoi + soi between the two image blocks): continue past max(vit pos)
            mid_start = vit_s + max(h_vit, w_vit)
            n_mid = lat_start - vit_e
            mid_t = mx.arange(mid_start, mid_start + n_mid, dtype=mx.int32)
            mid_h = mid_t
            mid_w = mid_t
            cond_t = cond_h = cond_w = mx.array([], dtype=mx.int32)
            post_pre_start = lat_start
        else:
            cond_start, cond_end, cond_grid = cond_info
            t_c, h_c, w_c = cond_grid
            # Mid1: between vit_end and cond_start (eoi + soi tokens)
            mid1_start = vit_s + max(h_vit, w_vit)
            n_mid1 = cond_start - vit_e
            mid1 = mx.arange(mid1_start, mid1_start + n_mid1, dtype=mx.int32)
            # Cond latent block 3D positions
            c_t = mx.repeat(mx.arange(t_c, dtype=mx.int32), h_c * w_c) + cond_start
            c_h = mx.repeat(mx.tile(mx.arange(h_c, dtype=mx.int32), (t_c,)), w_c) + cond_start
            c_w = mx.tile(mx.arange(w_c, dtype=mx.int32), (t_c * h_c,)) + cond_start
            # Mid2: between cond_end and lat_start
            mid2_start = cond_start + max(t_c, h_c, w_c)
            n_mid2 = lat_start - cond_end
            mid2 = mx.arange(mid2_start, mid2_start + n_mid2, dtype=mx.int32)
            mid_t = mx.concatenate([mid1, c_t, mid2])
            mid_h = mx.concatenate([mid1, c_h, mid2])
            mid_w = mx.concatenate([mid1, c_w, mid2])
            cond_t = cond_h = cond_w = mx.array([], dtype=mx.int32)  # absorbed into mid_*
            post_pre_start = lat_start

        # Latent block 3D positions (offset by lat_start)
        l_t = mx.repeat(mx.arange(t_lat, dtype=mx.int32), h_lat * w_lat) + lat_start
        l_h = mx.repeat(mx.tile(mx.arange(h_lat, dtype=mx.int32), (t_lat,)), w_lat) + lat_start
        l_w = mx.tile(mx.arange(w_lat, dtype=mx.int32), (t_lat * h_lat,)) + lat_start
        # Post-latent (eoi): continue past max(latent)
        post_start = lat_start + max(t_lat, h_lat, w_lat)
        n_post = S - (lat_start + n_lat)
        post = mx.arange(post_start, post_start + n_post, dtype=mx.int32)

        axis_t = mx.concatenate([pre, v_t, mid_t, l_t, post])
        axis_h = mx.concatenate([pre, v_h, mid_h, l_h, post])
        axis_w = mx.concatenate([pre, v_w, mid_w, l_w, post])
        return mx.stack([axis_t, axis_h, axis_w], axis=0)[:, None, :]

    def denoise_step_edit(
        self,
        prompt_token_ids: mx.array,
        visual_embeds: mx.array,
        vit_grid_hw: Tuple[int, int],
        cond_latent: Optional[mx.array],     # input image's VAE latent for conditioning (or None)
        noisy_latent: mx.array,
        timestep: mx.array,
        special_token_ids: dict,
    ) -> mx.array:
        B, T, H, W, C = noisy_latent.shape
        assert B == 1
        inputs_embeds, attn_mask, lat_start, n_lat, vit_block, cond_info = self._build_edit_sequence(
            prompt_token_ids, visual_embeds, vit_grid_hw, cond_latent, (T, H, W),
            timestep, noisy_latent, special_token_ids,
        )
        S = inputs_embeds.shape[1]
        position_ids = self._build_edit_position_ids(
            S, vit_block, vit_grid_hw, lat_start, (T, H, W), cond_info=cond_info,
        )
        # text_len = lat_start so the latent block uses moe_gen, everything before uses normal.
        hidden = self._forward_backbone(inputs_embeds, attn_mask, lat_start, position_ids)
        latent_hidden = hidden[:, lat_start:lat_start + n_lat, :]
        return self.project_latent_out(latent_hidden, T, H, W)

    def sample_edit(
        self,
        prompt_token_ids: mx.array,
        visual_embeds: mx.array,            # (1, N_vit, hidden) from ViT
        vit_grid_hw: Tuple[int, int],
        latent_shape: Tuple[int, int, int],
        special_token_ids: dict,
        cond_latent: Optional[mx.array] = None,   # (1, T, H, W, 48) VAE-encoded input image
        num_steps: int = 30,
        timestep_shift: float = 3.0,
        seed: Optional[int] = None,
        cfg_text_scale: float = 4.0,
        cfg_vit_scale: float = 1.5,
        uncond_prompt_ids: Optional[mx.array] = None,   # empty instruction for text-uncond
    ) -> mx.array:
        """Three-component CFG flow-matching for TI2I, mirroring PT Lance:
            v_full   = denoise(text, vit, cond)                   # text + ViT + VAE cond
            v_t_un   = denoise(empty, vit, cond)                   # text uncond, keep visuals
            v_tv_un  = denoise(empty, zero_vit, zero_cond)         # text + ViT + VAE uncond
            v_final  = v_tv_un + cfg_text * (v_full - v_t_un) + cfg_vit * (v_t_un - v_tv_un)
        """
        t_lat, h_lat, w_lat = latent_shape
        C = self.config.latent_channel
        if seed is not None:
            mx.random.seed(seed)
        x = mx.random.normal(shape=(1, t_lat, h_lat, w_lat, C))

        ts = mx.linspace(1.0, 0.0, num_steps + 1)
        ts = timestep_shift * ts / (1.0 + (timestep_shift - 1.0) * ts)
        dts = ts[:-1] - ts[1:]
        ts = ts[:-1]

        if uncond_prompt_ids is None:
            uncond_prompt_ids = mx.array([], dtype=mx.int32)
        zero_visual = mx.zeros_like(visual_embeds)
        zero_cond = mx.zeros_like(cond_latent) if cond_latent is not None else None

        for i in range(num_steps):
            t_arr = ts[i:i+1]
            v_full = self.denoise_step_edit(
                prompt_token_ids, visual_embeds, vit_grid_hw, cond_latent,
                x, t_arr, special_token_ids,
            )
            if cfg_text_scale != 1.0 or cfg_vit_scale != 1.0:
                v_t_un = self.denoise_step_edit(
                    uncond_prompt_ids, visual_embeds, vit_grid_hw, cond_latent,
                    x, t_arr, special_token_ids,
                )
                v_tv_un = self.denoise_step_edit(
                    uncond_prompt_ids, zero_visual, vit_grid_hw, zero_cond,
                    x, t_arr, special_token_ids,
                )
                v = (v_tv_un
                     + cfg_text_scale * (v_full - v_t_un)
                     + cfg_vit_scale  * (v_t_un - v_tv_un))
            else:
                v = v_full
            x = x - dts[i].item() * v
            mx.eval(x)
        return x

    # =======================================================================
    # X → T (understanding) — autoregressive sampling with KV cache.
    # All tokens use normal (non-moe_gen) weights, so we route everything as
    # "und" (text_len = full sequence length).
    # =======================================================================

    def attach_vit(self, vit_model) -> None:
        """Attach an mlx_vlm-style VisionModel for X→T modes.

        Loaded as a sibling so its weights live separately from the LLM (the
        ViT is published in its own safetensors shard).
        """
        self.vit_model = vit_model

    def _build_x2t_prefill(
        self,
        prompt_text_ids: mx.array,    # (P,) — full prompt token ids including bos/eos/specials
        visual_embeds: Optional[mx.array],   # (1, N_v, hidden) or None
        image_token_positions: Optional[mx.array],   # (N_v,) positions in prompt where to splice
        image_grid_hw: Optional[Tuple[int, int]] = None,   # (h_patches/merge, w_patches/merge)
    ) -> Tuple[mx.array, mx.array]:
        """Build the prefill embeddings + 3-axis mrope position ids for X→T.

        Position ids per axis:
          - Text tokens get the same scalar (running position) on all 3 axes.
          - Image tokens inside the <|vision_start|>..<|vision_end|> block get
            (t, h, w) coords offset by the running position at the start of the
            block. After the block, position continues from max + 1.
        """
        text_emb = self.language_model.model.embed_tokens(prompt_text_ids[None])  # (1, P, H)

        # First splice in ViT embeddings (no positional change yet).
        if visual_embeds is not None and image_token_positions is not None and image_token_positions.size > 0:
            P = prompt_text_ids.shape[0]
            pos_list = sorted(int(p) for p in image_token_positions.tolist())
            parts = []
            last = 0
            v_idx = 0
            for p in pos_list:
                if p > last:
                    parts.append(text_emb[:, last:p])
                parts.append(visual_embeds[:, v_idx:v_idx + 1])
                v_idx += 1
                last = p + 1
            if last < P:
                parts.append(text_emb[:, last:])
            text_emb = mx.concatenate(parts, axis=1)

        P = text_emb.shape[1]

        # Build per-axis position ids.
        if image_grid_hw is None or image_token_positions is None or image_token_positions.size == 0:
            pos = mx.arange(P, dtype=mx.int32)
            position_ids = mx.broadcast_to(pos[None, None, :], (3, 1, P))
            return text_emb, position_ids

        h_m, w_m = image_grid_hw           # grid dims after spatial-merge
        n_img = h_m * w_m
        img_positions = sorted(int(p) for p in image_token_positions.tolist())
        img_start = img_positions[0]
        img_end_exclusive = img_positions[-1] + 1
        assert img_end_exclusive - img_start == n_img, \
            f"image_token_positions ({len(img_positions)}) don't match grid ({n_img})"

        # Pre-image text positions: 0..img_start-1
        pre = mx.arange(img_start, dtype=mx.int32)
        # Image positions: T=img_start, H=img_start+row, W=img_start+col
        rows = mx.repeat(mx.arange(h_m, dtype=mx.int32), w_m) + img_start
        cols = mx.tile(mx.arange(w_m, dtype=mx.int32), (h_m,)) + img_start
        ts   = mx.zeros((n_img,), dtype=mx.int32) + img_start
        # Post-image text positions: continue from img_start + max(h_m, w_m)
        post_start = img_start + max(h_m, w_m)
        n_post = P - img_end_exclusive
        post = mx.arange(post_start, post_start + n_post, dtype=mx.int32)

        axis_t = mx.concatenate([pre, ts,   post])
        axis_h = mx.concatenate([pre, rows, post])
        axis_w = mx.concatenate([pre, cols, post])
        position_ids = mx.stack([axis_t, axis_h, axis_w], axis=0)[:, None, :]
        return text_emb, position_ids

    def x2t_generate(
        self,
        prompt_text_ids: mx.array,
        visual_embeds: Optional[mx.array],
        image_token_positions: Optional[mx.array],
        image_grid_hw: Optional[Tuple[int, int]] = None,
        max_new_tokens: int = 64,
        eos_token_id: Optional[int] = None,
        temperature: float = 0.0,
    ) -> list:
        """AR generation with KV cache.

        image_grid_hw: (h_after_merger, w_after_merger) used for 3D mrope inside the vision block.
        Returns list[int] of newly generated token ids (no echo of prompt).
        """
        caches = self.language_model.make_caches()
        embeds, pos_ids = self._build_x2t_prefill(
            prompt_text_ids, visual_embeds, image_token_positions, image_grid_hw,
        )
        S = embeds.shape[1]
        neg_inf = -1e9
        idx = mx.arange(S)
        i_grid, j_grid = mx.meshgrid(idx, idx, indexing="ij")
        mask = mx.where(j_grid > i_grid, neg_inf, 0.0)

        # Prefill (all tokens use normal weights via text_len=S)
        hidden = self.language_model(embeds, text_len=S, position_ids=pos_ids,
                                     mask=mask, caches=caches)
        mx.eval(hidden)
        last_hidden = hidden[:, -1:, :]
        logits = self.language_model.lm_head(last_hidden)
        if temperature <= 0:
            next_id = int(mx.argmax(logits[0, 0]).item())
        else:
            probs = mx.softmax(logits[0, 0] / temperature, axis=-1)
            next_id = int(mx.random.categorical(probs[None])[0].item())

        out_ids = [next_id]
        # Running position for AR decoding: continue past where the prefill left off
        if image_grid_hw is not None and image_token_positions is not None and image_token_positions.size > 0:
            img_positions = sorted(int(p) for p in image_token_positions.tolist())
            img_start = img_positions[0]
            img_end_excl = img_positions[-1] + 1
            h_m, w_m = image_grid_hw
            post_start = img_start + max(h_m, w_m)
            n_post = S - img_end_excl
            cur_pos = post_start + n_post
        else:
            cur_pos = S

        # AR decode
        for _ in range(max_new_tokens - 1):
            if eos_token_id is not None and next_id == eos_token_id:
                break
            tok = mx.array([next_id], dtype=mx.int32)
            emb = self.language_model.model.embed_tokens(tok[None])
            pos = mx.array([cur_pos], dtype=mx.int32)
            position_ids = mx.broadcast_to(pos[None, None, :], (3, 1, 1))
            hidden = self.language_model(emb, text_len=1, position_ids=position_ids,
                                          mask=None, caches=caches)
            mx.eval(hidden)
            logits = self.language_model.lm_head(hidden[:, -1:, :])
            if temperature <= 0:
                next_id = int(mx.argmax(logits[0, 0]).item())
            else:
                probs = mx.softmax(logits[0, 0] / temperature, axis=-1)
                next_id = int(mx.random.categorical(probs[None])[0].item())
            out_ids.append(next_id)
            cur_pos += 1

        return out_ids