"""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<|im_end|>\n # <|im_start|>user\n<|vision_start|>[N_vit IPADs]<|vision_end|><|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\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