Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/bridgev2_obs.jpg filter=lfs diff=lfs merge=lfs -text
+examples/umi_sock_obs.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,27 @@
 ---
-title: Rxbrain Embodied Cognition
-emoji: 📚
-colorFrom: pink
-colorTo: indigo
+title: RxBrain Embodied Cognition
+emoji: 🔮
+colorFrom: red
+colorTo: yellow
 sdk: gradio
 sdk_version: 6.20.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Embodied reasoning + visual imagination (Hy-RxBrain 1.0)
+python_version: "3.12"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# RxBrain — Embodied Cognition Foundation Model
+
+Interactive demo of **[tencent/Hy-Embodied-RxBrain-1.0](https://huggingface.co/tencent/Hy-Embodied-RxBrain-1.0)**,
+a unified Mixture-of-Transformers model that couples language reasoning with visual imagination.
+
+Three capabilities:
+
+- **Visual QA** — question answering over images.
+- **Text-to-Image** — imagine an image via the flow-matching head (decoded through a frozen FLUX VAE).
+- **Embodied Planning** — decompose a task into interleaved steps, emitting both the action (text)
+  and the imagined goal frame (image) per step.
+
+Runs on ZeroGPU. Imagination requires the external FLUX VAE (`ae.safetensors`), fetched from
+`black-forest-labs/FLUX.1-schnell`.

app.py

@@ -0,0 +1,340 @@
+"""RxBrain — Embodied Cognition Foundation Model demo (ZeroGPU).
+
+tencent/Hy-Embodied-RxBrain-1.0 : a unified Mixture-of-Transformers model that
+couples language reasoning with visual imagination. This Space exposes three
+capabilities, each mapped 1:1 onto the authors' official inference scripts:
+
+  * Visual Question Answering (vqa_inference.answer)
+  * Text-to-Image imagination (text2image_inference.generate_image)
+  * Interleaved embodied planning (interleave_inference.generate_step_joint)
+"""
+import os
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  — must precede torch / CUDA-touching imports
+import tempfile  # noqa: E402
+import traceback  # noqa: E402
+
+import torch  # noqa: E402
+import numpy as np  # noqa: E402
+import gradio as gr  # noqa: E402
+from PIL import Image  # noqa: E402
+from huggingface_hub import snapshot_download, hf_hub_download  # noqa: E402
+
+# ---------------------------------------------------------------------------
+# Model / VAE loading (module scope, eager .to("cuda") for ZeroGPU)
+# ---------------------------------------------------------------------------
+MODEL_ID = "tencent/Hy-Embodied-RxBrain-1.0"
+VAE_REPO = "black-forest-labs/FLUX.1-schnell"
+DTYPE = torch.bfloat16
+
+print("Downloading RxBrain checkpoint …")
+CKPT_DIR = snapshot_download(MODEL_ID)
+print(f"Checkpoint at {CKPT_DIR}")
+
+print("Downloading FLUX VAE (ae.safetensors) …")
+VAE_PATH = hf_hub_download(VAE_REPO, "ae.safetensors")
+print(f"VAE at {VAE_PATH}")
+
+from transformers.models.hunyuan_vl_mot import HunYuanVLMoTProcessor  # noqa: E402
+from model import UnifiedMoTForConditionalGeneration, maybe_init_generation_path  # noqa: E402
+from vae_model.autoencoder import load_ae  # noqa: E402
+
+# Inference helpers ported straight from the official scripts.
+from vqa_inference import answer as _vqa_answer  # noqa: E402
+from text2image_inference import generate_image as _t2i_generate  # noqa: E402
+from interleave_inference import generate_step_joint as _interleave_step  # noqa: E402
+
+print("Loading processor …")
+processor = HunYuanVLMoTProcessor.from_pretrained(CKPT_DIR, trust_remote_code=True)
+
+print("Loading model …")
+model = UnifiedMoTForConditionalGeneration.from_pretrained(CKPT_DIR, dtype=DTYPE)
+maybe_init_generation_path(model, model_load_path=CKPT_DIR)
+model.to("cuda").eval()
+
+print("Loading VAE …")
+vae, _ = load_ae(VAE_PATH)
+vae.requires_grad_(False)
+vae.eval()
+vae.to("cuda", dtype=DTYPE)
+
+DEVICE = torch.device("cuda")
+print("All components ready.")
+
+
+def _tensor_to_pil(img):
+    """(3, H, W) tensor in [-1, 1] -> PIL.Image."""
+    arr = ((img.cpu().permute(1, 2, 0).numpy() + 1.0) * 127.5).clip(0, 255).astype("uint8")
+    return Image.fromarray(arr)
+
+
+def _save_input_image(image):
+    """Persist a Gradio image (PIL or path) to a temp jpg, return the path."""
+    if image is None:
+        return None
+    if isinstance(image, str):
+        return image
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    f = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+    image.convert("RGB").save(f.name, quality=95)
+    return f.name
+
+
+# ---------------------------------------------------------------------------
+# 1. Visual Question Answering
+# ---------------------------------------------------------------------------
+@spaces.GPU(duration=120)
+def vqa(image, question, max_new_tokens=256):
+    """Answer a question about an image using RxBrain's understanding path.
+
+    Args:
+        image: the input image to reason over.
+        question: a natural-language question about the image.
+        max_new_tokens: maximum number of answer tokens to decode.
+    Returns:
+        The model's answer text.
+    """
+    if image is None:
+        return "Please provide an input image."
+    if not question or not question.strip():
+        return "Please enter a question."
+    path = _save_input_image(image)
+    text = _vqa_answer(
+        model, processor, image_paths=[path], question=question.strip(),
+        device=DEVICE, dtype=DTYPE, max_new_tokens=int(max_new_tokens),
+    )
+    return text.strip()
+
+
+# ---------------------------------------------------------------------------
+# 2. Text-to-Image imagination
+# ---------------------------------------------------------------------------
+def _t2i_duration(prompt, height=256, width=256, num_steps=25, cfg_scale=1.0, seed=0):
+    base = 40
+    per_step = 1.5 * (2 if cfg_scale != 1.0 else 1)
+    return int(min(180, base + num_steps * per_step))
+
+
+@spaces.GPU(duration=_t2i_duration)
+def text_to_image(prompt, height=256, width=256, num_steps=25, cfg_scale=1.0, seed=0):
+    """Imagine an image from a text prompt via the flow-matching head.
+
+    Args:
+        prompt: the text description to imagine.
+        height: output image height in pixels (multiple of 16).
+        width: output image width in pixels (multiple of 16).
+        num_steps: number of Euler-ODE flow-matching steps.
+        cfg_scale: classifier-free guidance scale (>1 enables CFG).
+        seed: RNG seed for reproducibility.
+    Returns:
+        The generated PIL image.
+    """
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    torch.manual_seed(int(seed))
+    h = int(height) // 16 * 16
+    w = int(width) // 16 * 16
+    img = _t2i_generate(
+        model, vae, processor, prompt.strip(),
+        height=h, width=w, num_steps=int(num_steps),
+        device=DEVICE, dtype=DTYPE, cfg_scale=float(cfg_scale),
+    )
+    return _tensor_to_pil(img)
+
+
+# ---------------------------------------------------------------------------
+# 3. Interleaved Embodied Planning (text plan + goal images, step by step)
+# ---------------------------------------------------------------------------
+def _plan_duration(image, task, num_frames=3, num_steps=50, height=144, width=256, seed=0):
+    return int(min(230, 30 + num_frames * (25 + num_steps * 1.3)))
+
+
+@spaces.GPU(duration=_plan_duration)
+def embodied_plan(image, task, num_frames=3, num_steps=50, height=144, width=256, seed=0):
+    """Decompose an embodied task into steps, emitting both the next action
+    (text) and the imagined goal frame (image) for each step.
+
+    Args:
+        image: an observation frame of the current scene.
+        task: the embodied task to plan for.
+        num_frames: number of interleaved plan steps to roll out.
+        num_steps: number of flow-matching ODE steps per imagined frame.
+        height: imagined-frame height in pixels (multiple of 16).
+        width: imagined-frame width in pixels (multiple of 16).
+        seed: RNG seed for reproducibility.
+    Yields:
+        (gallery_of_goal_frames, plan_markdown) as each step completes.
+    """
+    if image is None:
+        raise gr.Error("Please provide an observation image.")
+    if not task or not task.strip():
+        raise gr.Error("Please enter a task description.")
+    torch.manual_seed(int(seed))
+    obs_path = _save_input_image(image)
+    h = int(height) // 16 * 16
+    w = int(width) // 16 * 16
+
+    prior_steps = []          # (text, frame_path) accumulated across steps
+    gallery = []              # (image, caption) for the gr.Gallery
+    plan_md_lines = [f"**Task:** {task.strip()}", ""]
+
+    for k in range(int(num_frames)):
+        text, img = _interleave_step(
+            model, vae, processor,
+            input_image_paths=[obs_path], task_text=task.strip(),
+            height=h, width=w, num_steps=int(num_steps),
+            device=DEVICE, dtype=DTYPE, max_text_tokens=96,
+            prior_steps=list(prior_steps),
+        )
+        pil = _tensor_to_pil(img)
+        frame_f = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        pil.save(frame_f.name)
+        prior_steps.append((text, frame_f.name))
+        gallery.append((pil, f"Step {k + 1}"))
+        plan_md_lines.append(f"**Step {k + 1}:** {text.strip() or '(imagined goal frame only)'}")
+        yield gallery, "\n\n".join(plan_md_lines)
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+INTRO = """# 🔮 RxBrain — Embodied Cognition Foundation Model
+
+**[tencent/Hy-Embodied-RxBrain-1.0](https://huggingface.co/tencent/Hy-Embodied-RxBrain-1.0)** — a
+unified Mixture-of-Transformers model that couples **language reasoning** with **visual imagination**.
+It answers questions about scenes, imagines images from text, and decomposes embodied tasks into
+interleaved step-by-step *plans + goal images*.
+
+*Tencent Robotics X × Futian Laboratory × Tencent Hunyuan · imagination decoded through a frozen FLUX VAE.*
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="RxBrain") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(INTRO)
+
+        with gr.Tabs():
+            # ---------------- VQA ----------------
+            with gr.Tab("🧠 Visual QA"):
+                with gr.Row():
+                    with gr.Column():
+                        vqa_image = gr.Image(label="Image", type="pil", height=320)
+                        vqa_question = gr.Textbox(
+                            label="Question",
+                            placeholder="What objects are on the stovetop?",
+                        )
+                        with gr.Accordion("Advanced", open=False):
+                            vqa_max_tokens = gr.Slider(
+                                16, 512, value=256, step=16, label="Max new tokens")
+                        vqa_btn = gr.Button("Answer", variant="primary")
+                    with gr.Column():
+                        vqa_out = gr.Textbox(label="Answer", lines=10)
+
+                vqa_btn.click(
+                    vqa, inputs=[vqa_image, vqa_question, vqa_max_tokens],
+                    outputs=vqa_out, api_name="vqa",
+                )
+                gr.Examples(
+                    examples=[
+                        ["examples/bridgev2_obs.jpg",
+                         "What objects are on the stovetop, and where is the green toy?"],
+                        ["examples/umi_sock_obs.jpg",
+                         "Describe the scene and the objects the robot could manipulate."],
+                        ["examples/xtrainer_garment_obs.jpg",
+                         "What garment is on the surface and what is its current state?"],
+                    ],
+                    inputs=[vqa_image, vqa_question],
+                    outputs=vqa_out, fn=vqa,
+                    cache_examples=False, run_on_click=True,
+                )
+
+            # ---------------- T2I ----------------
+            with gr.Tab("🎨 Text-to-Image"):
+                with gr.Row():
+                    with gr.Column():
+                        t2i_prompt = gr.Textbox(
+                            label="Prompt",
+                            placeholder="a watercolor painting of a cat",
+                        )
+                        with gr.Accordion("Advanced", open=False):
+                            with gr.Row():
+                                t2i_h = gr.Slider(128, 384, value=256, step=16, label="Height")
+                                t2i_w = gr.Slider(128, 384, value=256, step=16, label="Width")
+                            t2i_steps = gr.Slider(10, 50, value=25, step=1, label="ODE steps")
+                            t2i_cfg = gr.Slider(1.0, 7.0, value=1.0, step=0.5,
+                                                label="CFG scale (>1 enables guidance)")
+                            t2i_seed = gr.Number(value=0, precision=0, label="Seed")
+                        t2i_btn = gr.Button("Imagine", variant="primary")
+                    with gr.Column():
+                        t2i_out = gr.Image(label="Imagined image", height=320)
+
+                t2i_btn.click(
+                    text_to_image,
+                    inputs=[t2i_prompt, t2i_h, t2i_w, t2i_steps, t2i_cfg, t2i_seed],
+                    outputs=t2i_out, api_name="text_to_image",
+                )
+                gr.Examples(
+                    examples=[
+                        ["a watercolor painting of a cat"],
+                        ["a red apple on a wooden table, photorealistic"],
+                        ["a robot arm gripping a colorful toy on a kitchen counter"],
+                    ],
+                    inputs=[t2i_prompt],
+                    outputs=t2i_out, fn=text_to_image,
+                    cache_examples=False, run_on_click=True,
+                )
+
+            # ---------------- Interleaved Planning ----------------
+            with gr.Tab("🧩 Embodied Planning"):
+                gr.Markdown(
+                    "Give an observation frame + a task; RxBrain rolls out an interleaved "
+                    "plan — each step emits an **action (text)** and the **imagined goal frame**."
+                )
+                with gr.Row():
+                    with gr.Column():
+                        plan_image = gr.Image(label="Observation", type="pil", height=280)
+                        plan_task = gr.Textbox(label="Task", lines=3,
+                                               placeholder="Move the green toy over the metal pot.")
+                        with gr.Accordion("Advanced", open=False):
+                            plan_nframes = gr.Slider(1, 5, value=3, step=1, label="Plan steps")
+                            plan_steps = gr.Slider(10, 50, value=50, step=1, label="ODE steps/frame")
+                            with gr.Row():
+                                plan_h = gr.Slider(128, 256, value=144, step=16, label="Frame height")
+                                plan_w = gr.Slider(128, 384, value=256, step=16, label="Frame width")
+                            plan_seed = gr.Number(value=0, precision=0, label="Seed")
+                        plan_btn = gr.Button("Plan", variant="primary")
+                    with gr.Column():
+                        plan_gallery = gr.Gallery(label="Imagined goal frames", columns=3, height=280)
+                        plan_text = gr.Markdown(label="Plan")
+
+                plan_btn.click(
+                    embodied_plan,
+                    inputs=[plan_image, plan_task, plan_nframes, plan_steps,
+                            plan_h, plan_w, plan_seed],
+                    outputs=[plan_gallery, plan_text], api_name="embodied_plan",
+                )
+                gr.Examples(
+                    examples=[
+                        ["examples/bridgev2_obs.jpg",
+                         "Move the green toy from the left side of the stovetop and hold it over "
+                         "the metal pot on the left burner. Please give the full plan as concrete "
+                         "physical actions, one step per action, and imagine the state once each "
+                         "step is completed. If only one step remains, provide only that final step."],
+                        ["examples/umi_sock_obs.jpg",
+                         "Fold the purple sock in half. Please lay out the high-level subgoal plan "
+                         "and imagine the resulting frame after each major stage. If only one "
+                         "subgoal remains, provide only that final subgoal."],
+                    ],
+                    inputs=[plan_image, plan_task],
+                    outputs=[plan_gallery, plan_text], fn=embodied_plan,
+                    cache_examples=False, run_on_click=True,
+                )
+
+if __name__ == "__main__":
+    demo.queue(max_size=12).launch(mcp_server=True)

inference_utils.py

@@ -0,0 +1,60 @@
+"""Inference-time utilities for HY-Unified.
+
+Self-contained slice of the pieces the inference scripts need — the image
+transform used to derive train/inference resolution parity, and the
+task-instruction prompts that cue each generation mode.
+
+The SAME task strings are used at train and inference (parity is load-bearing:
+``interleave_inference`` / ``multiframe_inference`` build the prompt the same
+way the model was trained). Single source of truth, imported by the scripts.
+"""
+
+import torch
+from PIL import Image
+import torchvision.transforms as T
+import torchvision.transforms.functional as TF
+from torchvision.transforms import InterpolationMode
+
+
+# Task-instruction prompts appended to the user turn so each generation MODE is
+# explicitly cued.
+TASK_INSTRUCTION_SINGLE_FRAME = "Generate future image of the goal"
+TASK_INSTRUCTION_MULTI_FRAME = "Generate future video of the goal"
+TASK_INSTRUCTION_INTERLEAVE = "Generate interleave goal planning"
+
+
+class VAEImageTransform:
+    """Aspect-preserving resize + normalize to [-1, 1] for VAE encoding.
+
+    Sides are clamped to be ``stride``-divisible and within ``[min_size,
+    max_size]``; extreme aspect ratios are re-scaled to stay inside ``max_size``
+    (otherwise latent position IDs could go out of bounds).
+    """
+
+    def __init__(self, max_size=512, min_size=256, stride=16):
+        self.max_size = max_size
+        self.min_size = min_size
+        self.stride = stride
+
+    def _make_divisible(self, val):
+        return max(self.stride, round(val / self.stride) * self.stride)
+
+    def __call__(self, img: Image.Image) -> torch.Tensor:
+        w, h = img.size
+        scale = min(self.max_size / max(w, h), 1.0)
+        scale = max(scale, self.min_size / min(w, h))
+        new_w = self._make_divisible(round(w * scale))
+        new_h = self._make_divisible(round(h * scale))
+        # Clamp to max_size to handle extreme aspect ratios where the min_size
+        # constraint overrides the max_size constraint (e.g. w=1000,h=50 would
+        # produce new_w=5120 which causes position ID out-of-bounds).
+        max_div = self._make_divisible(self.max_size)
+        if new_w > max_div or new_h > max_div:
+            # Re-scale to fit within max_size while preserving aspect ratio.
+            shrink = min(max_div / new_w, max_div / new_h)
+            new_w = self._make_divisible(round(new_w * shrink))
+            new_h = self._make_divisible(round(new_h * shrink))
+        img = TF.resize(img, (new_h, new_w), InterpolationMode.BICUBIC, antialias=True)
+        tensor = TF.to_tensor(img)  # [0, 1]
+        tensor = T.Normalize([0.5] * 3, [0.5] * 3)(tensor)  # [-1, 1]
+        return tensor

interleave_inference.py

@@ -0,0 +1,345 @@
+"""Interleaved multi-image inference with a SigLIP-context handoff.
+
+Generates a sequence of frames autoregressively:
+
+    for k in 1..Y:
+        seq_k = chat_template(user[obs + decoded f_1..f_{k-1}] + text) + <Image>+LAT*N+</Image>+EOS
+        x_0   = Euler-ODE(seq_k)            # prior frames condition via SigLIP
+        f_k   = VAE.decode(x_0)             # decode to pixels
+        # f_k is fed (as a SigLIP image) into seq_{k+1}
+
+The crucial property: a previously generated frame enters the next step's context
+as a **SigLIP-encoded image in a user turn** — NOT as its VAE latent — so the
+context representation is identical at train and test (clean pixels either way),
+avoiding the noised-VAE-latent train/test mismatch of single-sequence interleave.
+
+`build_conditioned_sequence` is the single sequence constructor; training builds
+the prompt the same way (apply_chat_template([user(images)+text, assistant:""])
+-> strip EOS -> <Image>+LAT*N+</Image>+EOS), so train and infer token layouts
+match exactly.
+
+Run:
+    python interleave_inference.py --ckpt <ckpt> --vae <vae> \
+        --frames obs.jpg --task "put the cup on the shelf" \
+        --max_frames 3 --num_steps 50 --out_dir out_interleave
+"""
+from __future__ import annotations
+
+import argparse
+import os
+from typing import List, Optional
+
+import torch
+from PIL import Image
+from transformers.models.hunyuan_vl_mot import HunYuanVLMoTProcessor
+
+# NOTE: heavy imports (model package -> flash_attn, vae_model) are done lazily
+# inside the functions that need them, so `build_conditioned_sequence` can be
+# imported with only a processor.
+
+# Input-image placeholder token ids the upstream HunYuanVL processor scatters
+# ViT features into.
+INPUT_IMAGE_PLACEHOLDER_IDS = (120687, 120688)
+
+
+def build_conditioned_sequence(processor, input_image_paths: List[str], prompt_text: str,
+                               prior_steps=None, task_instruction: Optional[str] = None):
+    """Build the prompt token ids + processor vision tensors for a
+    (multi-image + text) -> 1-generated-image sample.
+
+    apply_chat_template([user(...), assistant:""]) then strip trailing EOS.
+
+    prior_steps: optional list of (text, frame_path) for already-produced steps.
+        When given, the user turn is the interleaved context the JOINT converter
+        emits — [obs frames] + task + [text_1, frame_1, ..., text_m, frame_m] —
+        so prior plan texts stay in context and prior frames enter as SigLIP inputs.
+        input_image_paths are then the OBSERVATION frames only.
+        When None, the legacy [all images] + text layout is used.
+    task_instruction: optional mode instruction appended as the LAST user-text item,
+        e.g. "Generate interleave goal planning".
+
+    Returns (prompt_ids: list[int], proc_inputs: dict).
+    """
+    user_content = [{"type": "image", "image": p} for p in input_image_paths]
+    user_content.append({"type": "text", "text": prompt_text})
+    for t, f in (prior_steps or []):
+        if t:
+            user_content.append({"type": "text", "text": t})
+        user_content.append({"type": "image", "image": f})
+    if task_instruction:
+        user_content.append({"type": "text", "text": task_instruction})
+    prompt_messages = [
+        {"role": "user", "content": user_content},
+        {"role": "assistant", "content": ""},
+    ]
+    proc_inputs = processor.apply_chat_template(
+        prompt_messages, return_dict=True, tokenize=True, add_generation_prompt=False,
+    )
+    eos = processor.tokenizer.eos_token_id
+    prompt_ids = list(proc_inputs["input_ids"][0])
+    while prompt_ids and prompt_ids[-1] == eos:
+        prompt_ids.pop()
+    return prompt_ids, proc_inputs
+
+
+@torch.no_grad()
+def generate_step_joint(
+    model, vae, processor,
+    input_image_paths: List[str],
+    task_text: str,
+    height: int, width: int, num_steps: int,
+    device, dtype,
+    max_text_tokens: int = 64,
+    prior_steps=None,
+):
+    """v2 joint step: autoregressively decode the plan TEXT, then (on <Image>)
+    flow-match the FRAME. Returns (text_str, image_tensor).
+
+    prior_steps: list of (text, frame_path) already produced — threaded into the
+    context so the rollout matches the JOINT training layout (prior texts kept,
+    prior frames as SigLIP). When given, input_image_paths = observation frames only.
+
+    The prefix is built by the SAME `build_conditioned_sequence` used at train
+    time, so the text the model emits after `</answer>` and the <Image> trigger
+    match what training taught."""
+    from model.flow_matching_modules import unpatchify_latent
+    from text2image_inference import get_2d_position_ids
+
+    cfg = model.config
+    eos = cfg.eos_token_id
+    image_start = cfg.image_start_token_id
+    latent_ph = cfg.flow_latent_placeholder_id
+    inner = model.model
+
+    # Append the interleave mode instruction (train==infer parity).
+    from inference_utils import TASK_INSTRUCTION_INTERLEAVE
+    prompt_ids, proc = build_conditioned_sequence(processor, input_image_paths, task_text,
+                                                   prior_steps=prior_steps,
+                                                   task_instruction=TASK_INSTRUCTION_INTERLEAVE)
+    pixel_values = proc.get("pixel_values")
+    image_grid_thw = proc.get("image_grid_thw")
+    if pixel_values is not None:
+        pixel_values = pixel_values.to(device=device, dtype=dtype)
+    if image_grid_thw is not None:
+        image_grid_thw = image_grid_thw.to(device=device)
+
+    # ---- 1. autoregressive text decode until <Image> (KV-cached: exact, ~2.8x) ----
+    # prefill the prompt once (use_cache), then feed one token/step against the growing
+    # KV cache — the model's native past_key_values plumbing (Mode C handles the decode
+    # shape). Validated bit-identical to the non-cached per-token full-forward decode.
+    from transformers.cache_utils import DynamicCache
+
+    def _dmask(ids):
+        seq_len = len(ids)
+        inp = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+        mod = torch.zeros(1, seq_len, dtype=torch.long, device=device)
+        iim = torch.zeros(1, seq_len, dtype=torch.bool, device=device)
+        for pid in INPUT_IMAGE_PLACEHOLDER_IDS:
+            m = inp[0] == pid
+            mod[0, m] = 1
+            iim[0, m] = True
+        return inp, mod, iim
+
+    _empty_g = torch.zeros((0, 2), dtype=torch.int32, device=device)
+    pkv = DynamicCache()
+    prompt_len = len(prompt_ids)
+    inp, mod, iim = _dmask(list(prompt_ids))
+    out = inner(input_ids=inp, position_ids=torch.arange(prompt_len, device=device).unsqueeze(0),
+                pixel_values=pixel_values, image_grid_thw=image_grid_thw,
+                cu_seqlens=torch.tensor([0, prompt_len], dtype=torch.int32, device=device),
+                sample_ids=torch.zeros(1, prompt_len, dtype=torch.int32, device=device),
+                modality_mask=mod, input_image_mask=iim,
+                flow_embeds=None, flow_positions=None, g_seqlens=_empty_g,
+                use_cache=True, past_key_values=pkv, cache_position=torch.arange(prompt_len, device=device))
+    pkv = out.past_key_values
+    nxt = int(out.logits[0, -1].argmax().item())
+    cur = prompt_len
+    text_ids = []
+    for _ in range(max_text_tokens):
+        if nxt == image_start or nxt == eos:
+            break
+        text_ids.append(nxt)
+        out = inner(input_ids=torch.tensor([[nxt]], dtype=torch.long, device=device),
+                    position_ids=torch.tensor([[cur]], device=device),
+                    pixel_values=None, image_grid_thw=None,
+                    cu_seqlens=torch.tensor([0, cur + 1], dtype=torch.int32, device=device),
+                    sample_ids=torch.zeros(1, 1, dtype=torch.int32, device=device),
+                    modality_mask=torch.zeros(1, 1, dtype=torch.long, device=device),
+                    input_image_mask=torch.zeros(1, 1, dtype=torch.bool, device=device),
+                    flow_embeds=None, flow_positions=None, g_seqlens=_empty_g,
+                    use_cache=True, past_key_values=pkv, cache_position=torch.tensor([cur], device=device))
+        pkv = out.past_key_values
+        nxt = int(out.logits[0, -1].argmax().item())
+        cur += 1
+    seq = list(prompt_ids) + text_ids
+    text_str = processor.tokenizer.decode(text_ids, skip_special_tokens=True)
+
+    # ---- 2. flow-match the frame — diffusion PREFIX-CACHE: prefill (prompt+text+IMAGE_START)
+    # once, then each denoise step forwards ONLY the n latent tokens against the cached prefix
+    # (Mode C causal=False -> latents attend bidirectionally to prefix+latents, == the gen-block).
+    # Validated equivalent to the full-forward loop within ROCm non-determinism. ----
+    p = model.latent_patch_size
+    ds = cfg.vae_image_downsample
+    h_lat, w_lat = height // ds, width // ds
+    n_latent = h_lat * w_lat
+    prefix = seq + [image_start]
+    prefix_len = len(prefix)                      # = latent_start
+    inp_p, mod_p, iim_p = _dmask(prefix)
+    pkv_f = DynamicCache()
+    inner(input_ids=inp_p, position_ids=torch.arange(prefix_len, device=device).unsqueeze(0),
+          pixel_values=pixel_values, image_grid_thw=image_grid_thw,
+          cu_seqlens=torch.tensor([0, prefix_len], dtype=torch.int32, device=device),
+          sample_ids=torch.zeros(1, prefix_len, dtype=torch.int32, device=device),
+          modality_mask=mod_p, input_image_mask=iim_p,
+          flow_embeds=None, flow_positions=None, g_seqlens=_empty_g,
+          use_cache=True, past_key_values=pkv_f, cache_position=torch.arange(prefix_len, device=device))
+
+    latent_pos_ids = get_2d_position_ids(h_lat, w_lat, cfg.max_latent_size).to(device)
+    latent_pos_emb = model.latent_pos_embed(latent_pos_ids).to(dtype)
+    lat_ids = torch.tensor([[latent_ph] * n_latent], dtype=torch.long, device=device)
+    fp_rel = torch.tensor([[0, n_latent]], dtype=torch.int32, device=device)   # rel to the n-token input
+    mod2 = torch.full((1, n_latent), 2, dtype=torch.long, device=device)
+    iim2 = torch.zeros(1, n_latent, dtype=torch.bool, device=device)
+    pos_l = torch.arange(prefix_len, prefix_len + n_latent, device=device).unsqueeze(0)
+    cu_l = torch.tensor([0, prefix_len + n_latent], dtype=torch.int32, device=device)
+    sid_l = torch.zeros(1, n_latent, dtype=torch.int32, device=device)
+
+    x = torch.randn(n_latent, p * p * cfg.vae_z_channels, device=device, dtype=dtype)
+    ts = torch.linspace(1.0, 0.0, num_steps + 1, device=device, dtype=dtype)
+    for i in range(num_steps):
+        t, dt = ts[i], ts[i] - ts[i + 1]
+        fe = (model.vae2llm(x.to(model.vae2llm.weight.dtype)).to(dtype)
+              + model.time_embedder(t.expand(n_latent)).to(dtype) + latent_pos_emb)
+        out = inner(input_ids=lat_ids, inputs_embeds=None, attention_mask=None,
+                    position_ids=pos_l, pixel_values=None, image_grid_thw=None,
+                    cu_seqlens=cu_l, sample_ids=sid_l, modality_mask=mod2, input_image_mask=iim2,
+                    flow_embeds=fe, flow_positions=fp_rel, g_seqlens=_empty_g,
+                    use_cache=True, past_key_values=pkv_f, cache_position=pos_l[0])
+        v = model.llm2vae(out.hidden_states[0, 0:n_latent]).to(dtype)
+        x = x - dt * v
+        pkv_f.crop(prefix_len)                    # drop the latents; keep the fixed prefix for next step
+    x_lat = unpatchify_latent(x.float(), h_lat, w_lat, p, cfg.vae_z_channels)
+    x_lat = x_lat.unsqueeze(0).to(device=device, dtype=next(vae.parameters()).dtype)
+    img = vae.decode(x_lat)
+    if hasattr(img, "sample"):
+        img = img.sample
+    return text_str, img.squeeze(0).float().clamp(-1, 1)
+
+
+@torch.no_grad()
+def interleave_generate_joint(
+    model, vae, processor,
+    obs_frames: List[str], task_text: str, num_frames: int,
+    height: int, width: int, num_steps: int, device, dtype, out_dir: str,
+    max_text_tokens: int = 64,
+):
+    """v2 rollout: at each step the model emits plan text + a frame; the decoded
+    frame AND its plan text are fed back as context for the next step — prior
+    frames via SigLIP, prior texts kept — matching the JOINT training layout."""
+    import os
+    os.makedirs(out_dir, exist_ok=True)
+    gen_paths, texts = [], []
+    prior_steps = []  # (text, frame_path) accumulated across steps
+    for k in range(num_frames):
+        text, img = generate_step_joint(
+            model, vae, processor, list(obs_frames), task_text,
+            height, width, num_steps, device, dtype, max_text_tokens,
+            prior_steps=list(prior_steps),
+        )
+        arr = ((img.cpu().permute(1, 2, 0).numpy() + 1.0) * 127.5).clip(0, 255).astype("uint8")
+        path = os.path.join(out_dir, f"joint_step{k+1}.png")
+        Image.fromarray(arr).save(path)
+        gen_paths.append(path)
+        texts.append(text)
+        prior_steps.append((text, path))
+        print(f"  [joint step {k+1}/{num_frames}] prior={len(prior_steps)-1} | TEXT: {text[:80]!r} -> {path}")
+    return gen_paths, texts
+
+
+def _row(paths, w, h):
+    """Horizontal strip of images (resized to w x h, 4px gaps)."""
+    imgs = [Image.open(p).convert("RGB").resize((w, h)) for p in paths]
+    n = len(imgs)
+    strip = Image.new("RGB", (w * n + 4 * (n - 1), h), (20, 20, 20))
+    for i, im in enumerate(imgs):
+        strip.paste(im, (i * (w + 4), 0))
+    return strip
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Joint interleaved rollout: decode ALL plan text + ALL frames.")
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--vae", required=True)
+    ap.add_argument("--frames", nargs="+", required=True, help="observation frame path(s)")
+    ap.add_argument("--task", required=True, help="overall task text")
+    ap.add_argument("--max_frames", type=int, default=None, help="cap #frames to generate")
+    ap.add_argument("--out_dir", default="out_interleave")
+    ap.add_argument("--height", type=int, default=144)
+    ap.add_argument("--width", type=int, default=256)
+    ap.add_argument("--num_steps", type=int, default=50, help="ODE steps per frame")
+    ap.add_argument("--max_text_tokens", type=int, default=96)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    ap.add_argument("--understanding_max_pixels", type=int, default=524288,
+                    help="Cap obs/prior-frame ViT input pixels to MATCH training. "
+                         "Default 524288 (~494 tok/frame); the model default 4194304 (~4050 tok/frame) "
+                         "is an 8x train/infer resolution mismatch that degrades eval. Set 0 to disable.")
+    args = ap.parse_args()
+
+    from model import UnifiedMoTForConditionalGeneration, maybe_init_generation_path
+    from vae_model.autoencoder import load_ae
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    torch.manual_seed(args.seed)
+
+    # ---- resolve obs frames + task ----
+    obs, task, n_frames = args.frames, args.task, (args.max_frames or 1)
+
+    processor = HunYuanVLMoTProcessor.from_pretrained(args.ckpt, trust_remote_code=True)
+    # train==infer parity: cap obs/prior-frame ViT pixels to the SAME value training
+    # used. Prior frames re-enter via SigLIP at rollout, so this must match training too.
+    if args.understanding_max_pixels and args.understanding_max_pixels > 0:
+        ip = processor.image_processor
+        ip.max_pixels = args.understanding_max_pixels
+        if isinstance(getattr(ip, "size", None), dict) and "longest_edge" in ip.size:
+            ip.size["longest_edge"] = args.understanding_max_pixels
+    model = UnifiedMoTForConditionalGeneration.from_pretrained(args.ckpt, dtype=dtype)
+    maybe_init_generation_path(model, model_load_path=args.ckpt)
+    model.to(device).eval()
+    vae, _ = load_ae(args.vae)
+    vae.requires_grad_(False)
+    vae.eval()
+    vae.to(device, dtype=dtype)
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    print(f"JOINT interleaved rollout | obs={len(obs)} frame(s) | steps={n_frames}")
+    print(f"TASK: {task}\n")
+
+    gen_paths, gen_texts = interleave_generate_joint(
+        model, vae, processor, obs, task, n_frames,
+        args.height, args.width, args.num_steps, device, dtype, args.out_dir,
+        max_text_tokens=args.max_text_tokens,
+    )
+
+    # ---- write the FULL decoded output (text) + montages (image) ----
+    lines = [f"TASK: {task}", ""]
+    for k in range(len(gen_texts)):
+        lines.append(f"--- step {k+1} ---")
+        lines.append(f"GEN text: {gen_texts[k]}")
+        lines.append("")
+    txt_path = os.path.join(args.out_dir, "result.txt")
+    open(txt_path, "w").write("\n".join(lines))
+    print("\n".join(lines))
+
+    # image montage: GEN row
+    width, height = args.width, args.height
+    gen_row = _row(gen_paths, width, height)
+    gen_row.save(os.path.join(args.out_dir, "rollout_GEN.png"))
+    print(f"\nDecoded {len(gen_paths)} frames + {len(gen_texts)} texts -> {args.out_dir}")
+    print(f"  text  : {txt_path}")
+    print(f"  images: {args.out_dir}/joint_step*.png + rollout montage")
+
+
+if __name__ == "__main__":
+    main()

model/__init__.py

@@ -0,0 +1,64 @@
+"""HY-Unified model package.
+
+New entry points — built on transformers 4.57's upstream
+hunyuan_vl_mot base model, extended with the third (`_g`) generation MoT path
+and Flow Matching modules.
+
+Importing this package triggers the module-level replacement of upstream's
+`_flash_attention_forward_mot` with our packed/padded/decode three-stage
+version (see `attention_mot_packed`).
+
+Public surface:
+    UnifiedMoTConfig, UnifiedMoTTextConfig
+    UnifiedMoTForConditionalGeneration, UnifiedMoTModel, UnifiedMoTOutput
+    MoTDecoderLayer
+    mask_apply_3way
+    maybe_init_generation_path
+    TimestepEmbedder, PositionEmbedding
+    pack_latents, sample_timesteps, build_noisy_latent, patch_boundary_loss
+"""
+# Order matters: attention rebind must happen before any HunYuanVLMoT* instantiation.
+from . import attention_mot_packed  # noqa: F401  (side effect)
+
+from .config_unified_mot import UnifiedMoTConfig, UnifiedMoTTextConfig
+from .modeling_decoder_mot import MoTDecoderLayer, mask_apply_3way
+from .modeling_text_model_mot import MoTTextModel, MoTTextForCausalLM
+from .modeling_unified_mot import (
+    UnifiedMoTForConditionalGeneration,
+    UnifiedMoTModel,
+    UnifiedMoTOutput,
+    HY_VL_MOT_IMAGE_TOKEN_ID,
+    HY_VL_MOT_VIDEO_TOKEN_ID,
+)
+from .mot_init_utils import (
+    maybe_init_generation_path,
+    init_mlp_g_from_mlp_v_net2wider,
+    verify_mlp_g_equals_mlp_v,
+    checkpoint_has_g_keys,
+)
+from .flow_matching_modules import (
+    TimestepEmbedder,
+    PositionEmbedding,
+    pack_latents,
+    sample_timesteps,
+    build_noisy_latent,
+    patch_boundary_loss,
+)
+
+
+__all__ = [
+    # Config
+    "UnifiedMoTConfig", "UnifiedMoTTextConfig",
+    # Models
+    "UnifiedMoTForConditionalGeneration", "UnifiedMoTModel", "UnifiedMoTOutput",
+    "MoTDecoderLayer", "MoTTextModel", "MoTTextForCausalLM",
+    "mask_apply_3way",
+    # Token IDs
+    "HY_VL_MOT_IMAGE_TOKEN_ID", "HY_VL_MOT_VIDEO_TOKEN_ID",
+    # Initialization
+    "maybe_init_generation_path", "init_mlp_g_from_mlp_v_net2wider",
+    "verify_mlp_g_equals_mlp_v", "checkpoint_has_g_keys",
+    # Flow Matching utils
+    "TimestepEmbedder", "PositionEmbedding",
+    "pack_latents", "sample_timesteps", "build_noisy_latent", "patch_boundary_loss",
+]

model/attention_mot_packed.py

@@ -0,0 +1,435 @@
+"""Three-stage packed/padded MoT flash attention for UnifiedMoT.
+
+This module replaces the upstream `_flash_attention_forward_mot` (a module-level
+function in `transformers.models.hunyuan_vl_mot.modeling_hunyuan_vl_mot`) with a
+version that supports three modes:
+
+  A) Packed prefill (P-Pack training)
+     attention_mask is a dict carrying `cu_seqlens` (N+1,), `sample_ids` (1, T),
+     and `g_seqlens` (M, 2). Inputs are (1, T, H, D) — a single packed row of
+     concatenated samples. Three-pass attention:
+       1) Causal varlen — `cu_seqlens` enforces sample boundaries
+       2) Visual bidirectional override on each `v_seqlens` segment (absolute coords)
+       3) Generation-block override on each `g_seqlens` segment, with KV range
+          `[sample_start, g_e]` so KV NEVER crosses sample boundaries
+
+  B) Padded prefill (legacy / inference fallback)
+     attention_mask dict has `padding_mask` (B, S). Same as upstream:
+     unpad → varlen → repad. Visual override layered on top. Generation override
+     not used in this mode.
+
+  C) Decode (KV-cache active, S_q != S_k)
+     Simple (B, 1, H, D) path with no varlen.
+
+The replacement is **module-scope rebinding** of upstream's
+`_flash_attention_forward_mot` — surgical and reversible. Apply by importing
+this module before instantiating any HunYuanVLMoT* class:
+
+    from model import attention_mot_packed  # noqa: F401
+"""
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from flash_attn import flash_attn_varlen_func
+from transformers.models.hunyuan_vl_mot import modeling_hunyuan_vl_mot as _upstream
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _has_ranges(obj) -> bool:
+    if obj is None:
+        return False
+    if torch.is_tensor(obj):
+        return obj.numel() > 0 and obj.shape[-1] in (2, 3)
+    if isinstance(obj, list):
+        if not obj:
+            return False
+        if torch.is_tensor(obj[0]):
+            return any(t.numel() > 0 for t in obj)
+        return True
+    return False
+
+
+def _iter_g_ranges(obj):
+    """Yield (g_start, g_end, batch_idx) triples from `g_seqlens`.
+
+    Accepts:
+      - Tensor of shape (M, 2) — absolute packed coords (b_idx implicit = 0)
+      - Tensor of shape (M, 3) — (s, e, b_idx)
+      - List of tuples — same as above
+      - List of (N_i, 2) tensors (per-sample, padded mode)
+    """
+    if torch.is_tensor(obj):
+        for row in obj.tolist():
+            if len(row) == 2:
+                yield int(row[0]), int(row[1]), 0
+            else:
+                yield int(row[0]), int(row[1]), int(row[2])
+    elif isinstance(obj, list):
+        if obj and torch.is_tensor(obj[0]):
+            for b_idx, t in enumerate(obj):
+                if t.numel() == 0:
+                    continue
+                for row in t.tolist():
+                    yield int(row[0]), int(row[1]), b_idx
+        else:
+            for item in obj:
+                if torch.is_tensor(item):
+                    item = item.tolist()
+                if len(item) == 2:
+                    yield int(item[0]), int(item[1]), 0
+                else:
+                    yield int(item[0]), int(item[1]), int(item[2])
+
+
+def _find_sample_start(cu_seqlens: torch.Tensor, g_s: int) -> int:
+    """Largest cu_seqlens[i] <= g_s. Used when sample_ids is unavailable."""
+    cu = cu_seqlens.tolist() if torch.is_tensor(cu_seqlens) else cu_seqlens
+    lo, hi = 0, len(cu) - 1
+    while lo < hi:
+        mid = (lo + hi + 1) // 2
+        if cu[mid] <= g_s:
+            lo = mid
+        else:
+            hi = mid - 1
+    return int(cu[lo])
+
+
+# ---------------------------------------------------------------------------
+# Step 2 helper: visual bidirectional override on (B, S, H, D) tensors
+# ---------------------------------------------------------------------------
+
+def _apply_visual_bidirectional(
+    attn_output: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    v_seqlens,
+) -> torch.Tensor:
+    """Override input-image segments with non-causal flash attention.
+
+    Operates on (B, S, H, D). v_seqlens may be a (M, 2) tensor (B==1) or a list
+    of (N_i, 2) tensors (per-sample for B>1 padded mode).
+
+    PERF: convert each segs tensor to CPU once (`.tolist()`) instead of
+    issuing two `.item()` syncs per row inside the inner loop. Each layer
+    calls this once during forward and once during backward; without the
+    fix, 32 layers × ~25 segments × 2 .item()/seg × 2 (fwd+bwd) ≈ 3200
+    GPU→CPU syncs per step, each blocking on the previous attention kernel.
+    """
+    device = query.device
+    if isinstance(v_seqlens, list):
+        segs_list = v_seqlens
+    else:
+        segs_list = [v_seqlens]  # B==1 path
+
+    visual_q = []
+    visual_k = []
+    visual_v = []
+    cu_v = [0]
+    max_v = 0
+    write_back = []  # (b_idx, s, e)
+    has_any = False
+
+    for b_idx, segs in enumerate(segs_list):
+        if segs is None or not torch.is_tensor(segs) or segs.numel() == 0:
+            continue
+        # PERF: one bulk D2H copy instead of per-element .item()
+        segs_cpu = segs.tolist()
+        for row in segs_cpu:
+            s = int(row[0])
+            e = int(row[1])
+            if e <= s:
+                continue
+            has_any = True
+            visual_q.append(query[b_idx, s:e])
+            visual_k.append(key[b_idx, s:e])
+            visual_v.append(value[b_idx, s:e])
+            ln = e - s
+            cu_v.append(cu_v[-1] + ln)
+            max_v = max(max_v, ln)
+            write_back.append((b_idx, s, e))
+
+    if not has_any:
+        # Preserve autograd graph topology across ranks even when this rank has
+        # no visual segments (matches upstream's `fake_visual` trick).
+        dummy = query[:1, :1].sum() * 0
+        return attn_output + dummy
+
+    vq = torch.cat(visual_q, dim=0)
+    vk = torch.cat(visual_k, dim=0)
+    vv = torch.cat(visual_v, dim=0)
+    cu_v_t = torch.tensor(cu_v, device=device, dtype=torch.int32)
+
+    vis_out = flash_attn_varlen_func(
+        vq, vk, vv,
+        cu_seqlens_q=cu_v_t, cu_seqlens_k=cu_v_t,
+        max_seqlen_q=max_v, max_seqlen_k=max_v,
+        causal=False,
+    )
+
+    attn_output = attn_output.clone()
+    off = 0
+    for b_idx, s, e in write_back:
+        ln = e - s
+        attn_output[b_idx, s:e] = vis_out[off:off + ln]
+        off += ln
+    return attn_output
+
+
+# ---------------------------------------------------------------------------
+# Step 3 helper: generation-block override (P-Pack only)
+# ---------------------------------------------------------------------------
+
+def _apply_generation_block(
+    attn_output: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    g_seqlens,
+    cu_seqlens: torch.Tensor,
+    sample_ids: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Q over each generation block, K/V from sample-start..g_end (no cross-sample).
+
+    PERF: previously each iteration of the per-block loop did:
+      - `sample_ids[0, g_s].item()` — D2H sync
+      - `cu_seqlens[sid].item()` — D2H sync
+      - `torch.tensor([0, q_len], device=cuda, ...)` — small tensor alloc + H2D
+      - same for cu_k_g
+    For 25 gen blocks × 32 layers × 2 (fwd+bwd) ≈ 3200 ops/step, each blocking
+    on the prior attention kernel. Now we do **one** D2H copy of cu_seqlens
+    and sample_ids[0] up front, build all `cu_q_g` / `cu_k_g` pairs as a
+    single (M, 2) tensor with one H2D, and view rows in the loop.
+    """
+    device = query.device
+
+    # ---- One-shot metadata D2H ------------------------------------------
+    if sample_ids is not None and torch.is_tensor(sample_ids):
+        sample_ids_cpu = (
+            sample_ids[0].tolist() if sample_ids.dim() == 2 else sample_ids.tolist()
+        )
+    else:
+        sample_ids_cpu = None
+    if torch.is_tensor(cu_seqlens):
+        cu_seqlens_cpu = cu_seqlens.tolist()
+    else:
+        cu_seqlens_cpu = list(cu_seqlens)
+
+    # ---- Pre-compute per-block specs entirely on CPU --------------------
+    specs = []  # (g_s, g_e, b_idx, sample_start, q_len, kv_len)
+    for g_s, g_e, b_idx in _iter_g_ranges(g_seqlens):
+        if g_e <= g_s:
+            continue
+        if sample_ids_cpu is not None:
+            sid = sample_ids_cpu[g_s]
+            sample_start = cu_seqlens_cpu[sid]
+        else:
+            sample_start = _find_sample_start(cu_seqlens_cpu, g_s)
+        q_len = g_e - g_s
+        kv_len = g_e - sample_start
+        specs.append((g_s, g_e, b_idx, sample_start, q_len, kv_len))
+
+    if not specs:
+        # No generation blocks on this rank — preserve autograd topology.
+        dummy = query[:1, :1].sum() * 0
+        return attn_output + dummy
+
+    # ---- Single H2D for all cu pairs ------------------------------------
+    cu_q_all = torch.tensor(
+        [[0, s[4]] for s in specs], device=device, dtype=torch.int32
+    )
+    cu_k_all = torch.tensor(
+        [[0, s[5]] for s in specs], device=device, dtype=torch.int32
+    )
+
+    attn_output = attn_output.clone()
+    for i, (g_s, g_e, b_idx, sample_start, q_len, kv_len) in enumerate(specs):
+        gen_out = flash_attn_varlen_func(
+            query[b_idx, g_s:g_e].contiguous(),
+            key[b_idx, sample_start:g_e].contiguous(),
+            value[b_idx, sample_start:g_e].contiguous(),
+            cu_seqlens_q=cu_q_all[i], cu_seqlens_k=cu_k_all[i],
+            max_seqlen_q=q_len, max_seqlen_k=kv_len,
+            causal=False,
+        )
+        attn_output[b_idx, g_s:g_e] = gen_out
+    return attn_output
+
+
+# ---------------------------------------------------------------------------
+# Main entry point — drop-in replacement for upstream's _flash_attention_forward_mot
+# ---------------------------------------------------------------------------
+
+def flash_attention_forward_mot_packed(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask,
+    dropout: float = 0.0,
+    scaling: Optional[float] = None,
+    sliding_window: Optional[int] = None,
+    softcap: Optional[float] = None,
+    **kwargs,
+):
+    """Packed/padded/decode dispatcher. See module docstring."""
+    _upstream._check_flash_attn()
+
+    if kwargs.get("output_attentions", False):
+        logger.warning_once("`flash_attention_2` does not support `output_attentions=True`.")
+
+    # (B, heads, S, D) -> (B, S, heads, D)
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
+    value = value.transpose(1, 2)
+
+    if query.dtype == torch.float32:
+        if torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_gpu_dtype()
+        elif hasattr(module.config, "_pre_quantization_dtype"):
+            target_dtype = module.config._pre_quantization_dtype
+        else:
+            target_dtype = next(m for m in module.modules() if isinstance(m, nn.Linear)).weight.dtype
+        query, key, value = query.to(target_dtype), key.to(target_dtype), value.to(target_dtype)
+
+    bsz, s_q, n_heads, head_dim = query.shape
+    s_k = key.shape[1]
+
+    if isinstance(attention_mask, dict):
+        v_seqlens = attention_mask.get("v_seqlens", None)
+        g_seqlens = attention_mask.get("g_seqlens", None)
+        cu_seqlens_packed = attention_mask.get("cu_seqlens", None)
+        padding_mask = attention_mask.get("padding_mask", None)
+        sample_ids = attention_mask.get("sample_ids", None)
+        # PERF: pre-computed by the language_model wrapper so we avoid a D2H
+        # sync on every layer. None means we still have to compute it here.
+        max_seqlen_packed = attention_mask.get("max_seqlen_packed", None)
+    else:
+        # Dummy: allow odd callers (e.g., HF generation builders) that pass tensor masks.
+        v_seqlens = None
+        g_seqlens = None
+        cu_seqlens_packed = None
+        padding_mask = None
+        sample_ids = None
+        max_seqlen_packed = None
+
+    packed_mode = (cu_seqlens_packed is not None) and (s_q == s_k) and (bsz == 1)
+
+    # ---------------- Mode A: Packed prefill ---------------------------------
+    if packed_mode:
+        h_kv = key.shape[2]
+        q_flat = query.contiguous().view(s_q, n_heads, head_dim)
+        k_flat = key.contiguous().view(s_k, h_kv, head_dim)
+        v_flat = value.contiguous().view(s_k, h_kv, head_dim)
+
+        cu = cu_seqlens_packed.to(device=query.device, dtype=torch.int32)
+        # Use pre-computed max_seqlen if the wrapper supplied it (saves a
+        # D2H sync per layer × 32 layers).
+        if max_seqlen_packed is not None:
+            max_seqlen = max_seqlen_packed
+        else:
+            with torch.no_grad():
+                max_seqlen = int((cu[1:] - cu[:-1]).max().item())
+
+        attn_flat = flash_attn_varlen_func(
+            q_flat, k_flat, v_flat,
+            cu_seqlens_q=cu, cu_seqlens_k=cu,
+            max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
+            causal=True,
+        )
+        attn_output = attn_flat.reshape(1, s_q, n_heads, head_dim)
+
+        if v_seqlens is not None:
+            attn_output = _apply_visual_bidirectional(attn_output, query, key, value, v_seqlens)
+
+        if _has_ranges(g_seqlens):
+            attn_output = _apply_generation_block(
+                attn_output, query, key, value, g_seqlens, cu, sample_ids
+            )
+
+        return attn_output, None
+
+    # ---------------- Mode B: Padded prefill --------------------------------
+    if padding_mask is not None and bsz > 1 and s_q == s_k:
+        pad_bool = padding_mask.bool()
+        seqlens = padding_mask.sum(dim=-1).to(torch.int32)
+        cu_seqlens = torch.zeros(bsz + 1, device=query.device, dtype=torch.int32)
+        cu_seqlens[1:] = torch.cumsum(seqlens, dim=0)
+        max_seqlen = int(seqlens.max().item())
+
+        q_unpad = query[pad_bool]
+        k_unpad = key[pad_bool]
+        v_unpad = value[pad_bool]
+
+        out_unpad = flash_attn_varlen_func(
+            q_unpad, k_unpad, v_unpad,
+            cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
+            causal=True,
+        )
+        attn_output = query.new_zeros(bsz, s_q, n_heads, head_dim)
+        attn_output[pad_bool] = out_unpad
+
+        if v_seqlens is not None:
+            attn_output = _apply_visual_bidirectional(attn_output, query, key, value, v_seqlens)
+
+        # Padded path supports generation-block override too (B>1 with per-sample triples)
+        if _has_ranges(g_seqlens):
+            # Build a synthetic cu_seqlens from padding_mask for sample_start lookup
+            cu_pad = torch.zeros(bsz + 1, device=query.device, dtype=torch.int32)
+            cu_pad[1:] = torch.cumsum(seqlens, dim=0)
+            attn_output = _apply_generation_block(
+                attn_output, query, key, value, g_seqlens, cu_pad, sample_ids
+            )
+
+        return attn_output, None
+
+    # ---------------- Mode C: Single-sample / decode ------------------------
+    h_kv = key.shape[2]
+    q_flat = query.contiguous().view(bsz * s_q, n_heads, head_dim)
+    k_flat = key.contiguous().view(bsz * s_k, h_kv, head_dim)
+    v_flat = value.contiguous().view(bsz * s_k, h_kv, head_dim)
+    cu_q_t = torch.arange(0, bsz + 1, dtype=torch.int32, device=query.device) * s_q
+    cu_k_t = torch.arange(0, bsz + 1, dtype=torch.int32, device=query.device) * s_k
+
+    attn_flat = flash_attn_varlen_func(
+        q_flat, k_flat, v_flat,
+        cu_seqlens_q=cu_q_t, cu_seqlens_k=cu_k_t,
+        max_seqlen_q=s_q, max_seqlen_k=s_k,
+        causal=(s_q == s_k),
+    )
+    attn_output = attn_flat.reshape(bsz, s_q, n_heads, head_dim)
+
+    # Visual override is valid during prefill (S_q == S_k) only
+    if v_seqlens is not None and s_q == s_k:
+        attn_output = _apply_visual_bidirectional(attn_output, query, key, value, v_seqlens)
+
+    if _has_ranges(g_seqlens) and s_q == s_k and bsz == 1:
+        # B==1 inference: synthetic cu_seqlens covers the whole row
+        cu_one = torch.tensor([0, s_q], device=query.device, dtype=torch.int32)
+        attn_output = _apply_generation_block(
+            attn_output, query, key, value, g_seqlens, cu_one, sample_ids
+        )
+
+    return attn_output, None
+
+
+# ---------------------------------------------------------------------------
+# Activate: rebind upstream module-level reference
+# ---------------------------------------------------------------------------
+
+_upstream._flash_attention_forward_mot = flash_attention_forward_mot_packed
+logger.info("attention_mot_packed: replaced upstream _flash_attention_forward_mot")
+
+
+__all__ = ["flash_attention_forward_mot_packed"]

model/config_unified_mot.py

@@ -0,0 +1,103 @@
+"""Configuration for UnifiedMoT.
+
+Subclasses transformers.models.hunyuan_vl_mot config to add HY-Unified
+project-level fields used by the Flow Matching path:
+
+  - mlp_g_intermediate_size / mlp_g_init_noise_std  (Net2Wider on generation MLP)
+  - boundary_loss_weight                            (patch-boundary smoothness L1)
+  - latent_patch_size / max_latent_size             (FM latent grid)
+  - vae_image_downsample / vae_z_channels           (VAE encoder spec)
+  - timestep_shift                                  (t-distribution shift for FM)
+  - use_moe                                         (toggle moe_layers wrapper)
+  - flow_loss_weight / text_loss_weight             (loss aggregation knobs)
+
+The text-level fields (mlp_g_*) live on text_config because the proxy in
+HunYuanVLMoTConfig surfaces them at the top level too.
+"""
+from __future__ import annotations
+
+from typing import Optional
+
+from transformers.models.hunyuan_vl_mot.configuration_hunyuan_vl_mot import (
+    HunYuanVLMoTConfig,
+    HunYuanVLMoTTextConfig,
+    HunYuanVLMoTVisionConfig,
+)
+
+
+class UnifiedMoTTextConfig(HunYuanVLMoTTextConfig):
+    """Extends upstream text config with the three-path MoT (mlp_g) knobs."""
+
+    model_type = "unified_mot"
+
+    def __init__(
+        self,
+        mlp_g_intermediate_size: Optional[int] = None,
+        mlp_g_init_noise_std: float = 1e-4,
+        use_moe: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.mlp_g_intermediate_size = mlp_g_intermediate_size
+        self.mlp_g_init_noise_std = float(mlp_g_init_noise_std)
+        self.use_moe = bool(use_moe)
+
+
+class UnifiedMoTConfig(HunYuanVLMoTConfig):
+    """Top-level config for the unified flow-matching model.
+
+    Inherits HunYuanVLMoTConfig (text_config + vision_config + token ids) and
+    layers on Flow-Matching specific fields.
+    """
+
+    model_type = "unified_mot"
+    sub_configs = {
+        "vision_config": HunYuanVLMoTVisionConfig,
+        "text_config": UnifiedMoTTextConfig,
+    }
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        # Flow Matching geometry
+        latent_patch_size: int = 2,
+        max_latent_size: int = 32,
+        vae_image_downsample: int = 16,
+        vae_z_channels: int = 16,
+        timestep_shift: float = 1.0,
+        # Loss aggregation
+        boundary_loss_weight: float = 0.0,
+        flow_loss_weight: float = 1.0,
+        text_loss_weight: float = 1.0,
+        # Special token used as the FM latent placeholder in the input sequence.
+        # Defaults to upstream's `latent_token_id` (120690) so the same checkpoint
+        # works without changes.
+        flow_latent_placeholder_id: Optional[int] = None,
+        **kwargs,
+    ):
+        # If text_config wasn't pre-built, route the **kwargs through so caller-level
+        # fields like mlp_g_intermediate_size still land on text_config.
+        if text_config is None:
+            text_config = UnifiedMoTTextConfig(**kwargs)
+        elif isinstance(text_config, dict):
+            text_config = UnifiedMoTTextConfig(**text_config)
+
+        super().__init__(text_config=text_config, vision_config=vision_config, **kwargs)
+
+        self.latent_patch_size = int(latent_patch_size)
+        self.max_latent_size = int(max_latent_size)
+        self.vae_image_downsample = int(vae_image_downsample)
+        self.vae_z_channels = int(vae_z_channels)
+        self.timestep_shift = float(timestep_shift)
+        self.boundary_loss_weight = float(boundary_loss_weight)
+        self.flow_loss_weight = float(flow_loss_weight)
+        self.text_loss_weight = float(text_loss_weight)
+        self.flow_latent_placeholder_id = (
+            int(flow_latent_placeholder_id)
+            if flow_latent_placeholder_id is not None
+            else int(self.latent_token_id)
+        )
+
+
+__all__ = ["UnifiedMoTConfig", "UnifiedMoTTextConfig"]

model/flow_matching_modules.py

@@ -0,0 +1,135 @@
+"""Flow Matching helper modules for UnifiedMoT.
+
+Provides:
+  - TimestepEmbedder, PositionEmbedding (re-exported from model_utils)
+  - patchify_latent / unpatchify_latent — reshape between (C, H_lat, W_lat)
+    and (h*w, p*p*C) packed-patch token form
+  - sample_timesteps — sigmoid(N(0,1)) → [0,1] with timestep_shift applied
+  - build_noisy_latent — linear interp x_t = (1-t)*x_0 + t*noise
+  - patch_boundary_loss — L1 across patch boundaries on predicted x_0
+"""
+from __future__ import annotations
+
+from typing import List, Tuple
+
+import torch
+
+# Re-exported from model_utils for a stable import path (see __all__ below).
+from .model_utils import (  # noqa: F401  # pylint: disable=unused-import
+    TimestepEmbedder,
+    PositionEmbedding,
+)
+
+
+# ---------------------------------------------------------------------------
+# Patchify / Unpatchify
+# ---------------------------------------------------------------------------
+
+def patchify_latent(latent: torch.Tensor, h: int, w: int, patch_size: int, channels: int) -> torch.Tensor:
+    """(C, H_lat, W_lat) → (h*w, p*p*C) where H_lat=h*p, W_lat=w*p."""
+    p = patch_size
+    latent = latent[:, : h * p, : w * p].reshape(channels, h, p, w, p)
+    latent = torch.einsum("chpwq->hwpqc", latent).reshape(-1, p * p * channels)
+    return latent
+
+
+def unpatchify_latent(packed: torch.Tensor, h: int, w: int, patch_size: int, channels: int) -> torch.Tensor:
+    """(h*w, p*p*C) → (C, H_lat, W_lat). Inverse of patchify_latent."""
+    p = patch_size
+    x = packed.reshape(h, w, p, p, channels)
+    x = torch.einsum("hwpqc->chpwq", x).reshape(channels, h * p, w * p)
+    return x
+
+
+def pack_latents(
+    padded_latents: List[torch.Tensor],
+    shapes: List[Tuple[int, int]],
+    patch_size: int,
+    channels: int,
+) -> torch.Tensor:
+    """Concatenate per-image patchified latents into one (sum_h*w, p*p*C) tensor."""
+    parts = []
+    for latent, (h, w) in zip(padded_latents, shapes):
+        parts.append(patchify_latent(latent, h, w, patch_size, channels))
+    return torch.cat(parts, dim=0)
+
+
+# ---------------------------------------------------------------------------
+# Timestep sampling
+# ---------------------------------------------------------------------------
+
+def sample_timesteps(
+    shapes: List[Tuple[int, int]],
+    timestep_shift: float,
+    device,
+    dtype,
+) -> torch.Tensor:
+    """Per-image t ~ sigmoid(N(0,1)) broadcast over each image's tokens, then shifted.
+
+    Returns a flat tensor of length sum(h*w), one t per token.
+    """
+    raw_t = torch.randn(len(shapes), device=device, dtype=dtype)
+    pieces = []
+    for (h, w), t in zip(shapes, raw_t):
+        pieces.append(torch.sigmoid(t).expand(h * w).clone())
+    timesteps = torch.cat(pieces, dim=0)
+    timesteps = timestep_shift * timesteps / (1.0 + (timestep_shift - 1.0) * timesteps)
+    return timesteps
+
+
+def build_noisy_latent(
+    clean: torch.Tensor,
+    noise: torch.Tensor,
+    timesteps: torch.Tensor,
+) -> torch.Tensor:
+    """x_t = (1 - t) * x_0 + t * noise, t broadcast over feature dim."""
+    return (1.0 - timesteps[:, None]) * clean + timesteps[:, None] * noise
+
+
+# ---------------------------------------------------------------------------
+# Boundary loss
+# ---------------------------------------------------------------------------
+
+def patch_boundary_loss(
+    x0_pred: torch.Tensor,
+    shapes: List[Tuple[int, int]],
+    patch_size: int,
+    channels: int,
+) -> torch.Tensor:
+    """L1 across horizontal/vertical patch boundaries on predicted x_0.
+
+    x0_pred: (sum_h*w, p*p*C). For each image's hxw grid, compares the right
+    column of patch (i,j) to the left column of patch (i,j+1) and similarly
+    for vertical neighbours. Encourages spatial smoothness in the latent
+    that the VAE decoder will turn into pixels.
+    """
+    p = patch_size
+    losses = []
+    offset = 0
+    for (h, w) in shapes:
+        n = h * w
+        x = x0_pred[offset : offset + n].reshape(h, w, p, p, channels)
+        if w > 1:
+            right_edge = x[:, :-1, :, p - 1, :]   # (h, w-1, p, C)
+            left_edge = x[:, 1:, :, 0, :]          # (h, w-1, p, C)
+            losses.append((right_edge - left_edge).abs().mean())
+        if h > 1:
+            bottom_edge = x[:-1, :, p - 1, :, :]   # (h-1, w, p, C)
+            top_edge = x[1:, :, 0, :, :]           # (h-1, w, p, C)
+            losses.append((bottom_edge - top_edge).abs().mean())
+        offset += n
+    if not losses:
+        return x0_pred.sum() * 0
+    return torch.stack(losses).mean()
+
+
+__all__ = [
+    "TimestepEmbedder",
+    "PositionEmbedding",
+    "patchify_latent",
+    "unpatchify_latent",
+    "pack_latents",
+    "sample_timesteps",
+    "build_noisy_latent",
+    "patch_boundary_loss",
+]

model/model_utils.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# Independent reimplementation of standard sinusoidal position- and
+# timestep-embedding utilities, written from their public mathematical
+# definitions (see the per-symbol references below). This file is NOT derived
+# from the DiT source tree and carries no third-party (CC BY-NC) copyright; it
+# is released under Apache-2.0 like the rest of this package.
+#
+# References for the underlying formulas (math only, no code reuse):
+#   * Sinusoidal position encoding — Vaswani et al., "Attention Is All You
+#     Need" (2017), Section 3.5.
+#   * Sinusoidal timestep embedding — Ho et al., "Denoising Diffusion
+#     Probabilistic Models" (2020); the same closed form is reused across
+#     diffusion and flow-matching models.
+#
+# The module layout, parameter names, and numerical outputs are intentionally
+# identical to the previous version so that existing checkpoints load unchanged
+# and inference results are bit-for-bit reproducible.
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+
+# --------------------------------------------------------------------------- #
+# Sinusoidal (sin-cos) position embeddings
+#
+# For a position p and channel index i the embedding interleaves
+#   sin(p / 10000^(2i/d))  and  cos(p / 10000^(2i/d)),
+# i.e. the classic Transformer positional encoding. The 2D variant encodes the
+# height and width axes with half the channels each and concatenates them.
+# --------------------------------------------------------------------------- #
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """1D sin-cos embedding for a flat array of positions.
+
+    Args:
+        embed_dim: even output dimension per position.
+        pos: array of positions, any shape; flattened to (M,).
+    Returns:
+        (M, embed_dim) array, [sin | cos] halves concatenated.
+    """
+    assert embed_dim % 2 == 0
+    # Inverse frequencies 1 / 10000^(k/(embed_dim/2)) for k = 0 .. embed_dim/2-1.
+    inv_freq = np.arange(embed_dim // 2, dtype=np.float64)
+    inv_freq /= embed_dim / 2.0
+    inv_freq = 1.0 / 10000**inv_freq  # (embed_dim/2,)
+
+    angles = np.einsum("m,d->md", pos.reshape(-1), inv_freq)  # outer product (M, D/2)
+    return np.concatenate([np.sin(angles), np.cos(angles)], axis=1)  # (M, D)
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    """2D sin-cos embedding: half the channels encode each spatial axis."""
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+    return np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    """Sin-cos table for a ``grid_size x grid_size`` patch grid.
+
+    Returns an ``(grid_size**2 [+ extra_tokens], embed_dim)`` array. When
+    ``cls_token`` is set and ``extra_tokens > 0``, that many zero rows are
+    prepended.
+    """
+    axis_h = np.arange(grid_size, dtype=np.float32)
+    axis_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(axis_w, axis_h)  # width varies fastest
+    grid = np.stack(grid, axis=0).reshape([2, 1, grid_size, grid_size])
+
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+# --------------------------------------------------------------------------- #
+# Timestep embedding
+# --------------------------------------------------------------------------- #
+class TimestepEmbedder(nn.Module):
+    """Embed scalar (possibly fractional) timesteps into vectors.
+
+    Sinusoidal frequency features are fed through a two-layer MLP with a SiLU
+    non-linearity. This is the standard timestep conditioning used by diffusion
+    and flow-matching models.
+    """
+
+    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):
+        """Sinusoidal features for a 1D tensor of (possibly fractional) timesteps.
+
+        Args:
+            t: (N,) tensor of timestep values.
+            dim: output feature dimension.
+            max_period: lowest angular frequency (longest period).
+        Returns:
+            (N, dim) tensor of [cos | sin] features (zero-padded if ``dim`` is odd).
+        """
+        half = dim // 2
+        # Geometrically spaced frequencies over [1, 1/max_period].
+        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):
+        freq_features = self.timestep_embedding(t, self.frequency_embedding_size)
+        freq_features = freq_features.to(next(self.mlp.parameters()).dtype)
+        return self.mlp(freq_features)
+
+
+class PositionEmbedding(nn.Module):
+    """Fixed 2D sin-cos position table exposed as a lookup by position id."""
+
+    def __init__(self, max_num_patch_per_side, hidden_size):
+        super().__init__()
+        self.max_num_patch_per_side = max_num_patch_per_side
+        self.hidden_size = hidden_size
+        # Build the (non-trainable) table eagerly. If it were created lazily it
+        # could stay on the meta device through transformers'
+        # from_pretrained(dtype=...) path and later materialize as uninitialized
+        # memory when moved to the GPU.
+        table = get_2d_sincos_pos_embed(hidden_size, max_num_patch_per_side)
+        self.pos_embed = nn.Parameter(
+            torch.from_numpy(table).float(),
+            requires_grad=False,
+        )
+
+    def _reset_parameters(self):
+        """Recompute the table after a meta-init path (call post-from_pretrained)."""
+        table = get_2d_sincos_pos_embed(self.hidden_size, self.max_num_patch_per_side)
+        on_meta = self.pos_embed.is_meta or self.pos_embed.device.type == "meta"
+        materialized = torch.from_numpy(table).to(
+            device="cpu" if on_meta else self.pos_embed.device,
+            dtype=torch.float32 if on_meta else self.pos_embed.dtype,
+        )
+        if on_meta:
+            self.pos_embed = nn.Parameter(materialized.float(), requires_grad=False)
+        else:
+            self.pos_embed.data.copy_(materialized.to(self.pos_embed.dtype))
+
+    def forward(self, position_ids):
+        return self.pos_embed[position_ids]

model/modeling_decoder_mot.py

@@ -0,0 +1,263 @@
+"""MoT decoder layer with three-path routing (text / vision-input / vision-gen).
+
+Subclasses `HunYuanVLMoTDecoderLayer` from upstream to add the third (`_g`)
+path used by Flow Matching. Upstream provides `mlp / mlp_v` and matching
+LayerNorms; we add `mlp_g / input_layernorm_g / post_attention_layernorm_g`.
+
+`modality_mask` semantics:
+  * 0 = text token            → mlp / input_layernorm / post_attention_layernorm
+  * 1 = input vision token    → mlp_v / *_v
+  * 2 = generated vision token (FM latent) → mlp_g / *_g
+
+Attention QKVO is **not** extended to a third path — generation tokens reuse
+the `_v` projection (matches the existing HY-Unified design where
+`q_proj_v / k_proj_v / v_proj_v / o_proj_v` cover both input and generated
+vision tokens). We just need to convert the int modality mask to a bool
+"vision" mask before calling `self.self_attn`.
+
+mlp_g may optionally have a wider intermediate (Net2Wider scale-up). When
+`config.mlp_g_intermediate_size` is set and larger than `intermediate_size`,
+we build mlp_g with a cloned config that overrides `intermediate_size`.
+"""
+from __future__ import annotations
+
+import copy
+from typing import Optional
+
+import torch
+
+from transformers.cache_utils import Cache
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+from transformers.utils.deprecation import deprecate_kwarg
+
+from transformers.models.hunyuan_vl_mot.modeling_hunyuan_vl_mot import (
+    HunYuanVLMoTDecoderLayer,
+    HunYuanVLMoTMLP,
+    HunYuanVLMoTRMSNorm,
+)
+
+# Side-effect import: replaces upstream's _flash_attention_forward_mot
+from . import attention_mot_packed  # noqa: F401  # pylint: disable=unused-import
+
+
+# ---------------------------------------------------------------------------
+# Three-way mask_apply
+# ---------------------------------------------------------------------------
+
+def mask_apply_3way(
+    hidden_states: torch.Tensor,
+    modality_mask: Optional[torch.Tensor],
+    text_funcs,
+    vision_funcs,
+    gen_funcs,
+    out_dims=None,
+    padding_mask: Optional[torch.Tensor] = None,
+):
+    """Routes tokens to text/vision/generation function lists by modality_mask.
+
+    hidden_states: (B, S, D)
+    modality_mask: (B, S) int tensor with values in {0, 1, 2}, or None
+        0 = text, 1 = vision input, 2 = vision generation
+    padding_mask: (B, S) int tensor, 1 = valid token, 0 = padding
+
+    Returns a list of stacked (B, S, out_d) tensors — one per `text_funcs[i]`.
+    """
+    if modality_mask is None:
+        # All-text: skip routing entirely
+        return [text_funcs[i](hidden_states) for i in range(len(text_funcs))]
+
+    bsz, seq_len, hidden_dim = hidden_states.size()
+    flat = hidden_states.reshape(bsz * seq_len, hidden_dim)
+    mask_flat = modality_mask.reshape(bsz * seq_len)
+    if padding_mask is not None:
+        valid_flat = padding_mask.reshape(bsz * seq_len).bool()
+    else:
+        valid_flat = None
+
+    placeholder = hidden_states[0:1, 0:1, :]  # (1, 1, D)
+    zero_feature = 0
+
+    num_outputs = len(text_funcs)
+    if out_dims is None:
+        out_dims_resolved = [hidden_dim] * num_outputs
+    else:
+        out_dims_resolved = list(out_dims)
+
+    # Pre-allocate output buffers (empty, not zeros — we overwrite all valid
+    # positions and the rest are masked out by the caller / padding).
+    out_flat = [
+        torch.empty(bsz * seq_len, od, device=flat.device, dtype=flat.dtype)
+        for od in out_dims_resolved
+    ]
+    # Padding positions need to be zeroed — we won't touch them in the
+    # gather/scatter below. Cheaper than zeroing the whole tensor: only
+    # zero the rows that won't be hit by any of the three modalities.
+    if valid_flat is not None:
+        invalid_flat = ~valid_flat
+        if invalid_flat.any():
+            for buf in out_flat:
+                buf[invalid_flat] = 0
+    # else: all rows will be hit by exactly one of {text, vision, gen}.
+
+    def _dispatch(idx_mask, funcs):
+        """Run `funcs` on rows selected by `idx_mask`, scatter back. If no
+        rows are selected, multiply through a placeholder so the params still
+        receive grad (avoids "unused parameter" DDP errors)."""
+        nonlocal zero_feature
+        if idx_mask.any():
+            hs_sel = flat[idx_mask]
+            for i, fn in enumerate(funcs):
+                out_flat[i][idx_mask] = fn(hs_sel)
+        else:
+            for fn in funcs:
+                zero_feature = zero_feature + fn(placeholder).mean() * 0
+
+    # Text: mask == 0
+    text_idx = (mask_flat == 0)
+    if valid_flat is not None:
+        text_idx = text_idx & valid_flat
+    _dispatch(text_idx, text_funcs)
+
+    # Vision input: mask == 1
+    vis_idx = (mask_flat == 1)
+    if valid_flat is not None:
+        vis_idx = vis_idx & valid_flat
+    _dispatch(vis_idx, vision_funcs)
+
+    # Generation: mask == 2
+    gen_idx = (mask_flat == 2)
+    if valid_flat is not None:
+        gen_idx = gen_idx & valid_flat
+    _dispatch(gen_idx, gen_funcs)
+
+    result = [out.view(bsz, seq_len, -1) for out in out_flat]
+    result[0] = result[0] + zero_feature
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Decoder layer subclass
+# ---------------------------------------------------------------------------
+
+def _make_g_config(config, mlp_g_intermediate_size: int):
+    """Clone config with a wider intermediate_size for mlp_g (Net2Wider).
+
+    MUST be `copy.deepcopy`, NOT `copy.copy`. Reason:
+      HunYuanVLMoTConfig.__setattr__ is a proxy that re-routes any attribute
+      writes (when the key is in text_config.__dict__) into self.text_config.
+      Shallow-copying the outer config keeps `g_cfg.text_config is
+      config.text_config` — same object — so the subsequent
+      `g_cfg.intermediate_size = ...` setattr ends up mutating the SHARED
+      text_config used by every other decoder layer's mlp/mlp_v construction.
+      Layer 0 escapes (its mlp/mlp_v are built before _make_g_config runs),
+      but layer 1+ then sees text_config.intermediate_size = 12288 and builds
+      mlp/mlp_v at the wider size — at that point ckpt weights (which are 6144
+      for mlp/mlp_v) no longer match and `from_pretrained` fails with size
+      mismatch. Deepcopy the whole config tree to break this aliasing.
+    """
+    g_cfg = copy.deepcopy(config)
+    g_cfg.intermediate_size = int(mlp_g_intermediate_size)
+    return g_cfg
+
+
+class MoTDecoderLayer(HunYuanVLMoTDecoderLayer):
+    """Three-path decoder layer: text / vision-input / vision-gen.
+
+    Adds `mlp_g`, `input_layernorm_g`, `post_attention_layernorm_g` on top of
+    upstream's two-path layer. `modality_mask` is interpreted as int{0,1,2}
+    instead of bool.
+    """
+
+    def __init__(self, config, layer_idx: int):
+        super().__init__(config, layer_idx)
+
+        # Text-config-aware lookup (HunYuanVLMoTConfig proxies these to text_config)
+        intermediate_size = getattr(config, "intermediate_size", None)
+        mlp_g_inter = getattr(config, "mlp_g_intermediate_size", None)
+
+        if mlp_g_inter is None or intermediate_size is None or mlp_g_inter == intermediate_size:
+            g_cfg = config
+        else:
+            assert mlp_g_inter >= intermediate_size, (
+                f"mlp_g_intermediate_size ({mlp_g_inter}) must be >= "
+                f"intermediate_size ({intermediate_size}); shrinking is not supported."
+            )
+            g_cfg = _make_g_config(config, mlp_g_inter)
+
+        self.mlp_g = HunYuanVLMoTMLP(g_cfg)
+        self.input_layernorm_g = HunYuanVLMoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm_g = HunYuanVLMoTRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        modality_mask: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        # Pull padding_mask from attention_mask dict for GEMM exclusion
+        padding_mask = None
+        if isinstance(attention_mask, dict):
+            pm = attention_mask.get("padding_mask", None)
+            if pm is not None and hidden_states.shape[1] == pm.shape[1]:
+                padding_mask = pm
+
+        # Convert int modality_mask to bool for attention QKVO routing
+        # (attention has only 2 paths: text vs vision; gen tokens go through _v)
+        attn_modality_mask = None
+        if modality_mask is not None:
+            attn_modality_mask = (modality_mask > 0)
+
+        residual = hidden_states
+
+        # Pre-attention LayerNorm — three paths
+        hidden_states = mask_apply_3way(
+            hidden_states, modality_mask,
+            [self.input_layernorm],
+            [self.input_layernorm_v],
+            [self.input_layernorm_g],
+            padding_mask=padding_mask,
+        )[0]
+
+        # Self-attention (upstream's two-path attention with bool mask)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            modality_mask=attn_modality_mask,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # MLP — three paths (LayerNorm + MLP fused per path to match upstream layout)
+        residual = hidden_states
+        hidden_states = mask_apply_3way(
+            hidden_states, modality_mask,
+            [lambda x: self.mlp(self.post_attention_layernorm(x))],
+            [lambda x: self.mlp_v(self.post_attention_layernorm_v(x))],
+            [lambda x: self.mlp_g(self.post_attention_layernorm_g(x))],
+            padding_mask=padding_mask,
+        )[0]
+        hidden_states = residual + hidden_states
+
+        # Zero out padding positions to keep residuals clean (matches upstream behavior)
+        if padding_mask is not None and hidden_states.shape[1] == padding_mask.shape[1]:
+            hidden_states = hidden_states * padding_mask.unsqueeze(-1)
+
+        return hidden_states
+
+
+__all__ = ["MoTDecoderLayer", "mask_apply_3way"]

model/modeling_text_model_mot.py

@@ -0,0 +1,246 @@
+"""Inner text-decoder subclass that threads packed kwargs into the causal_mask dict.
+
+The upstream `_HunYuanVLMoTTextModel.forward` builds:
+
+    causal_mask = {"v_seqlens": visual_segs, "padding_mask": attention_mask}
+
+and passes it to each decoder layer. For P-Pack training we additionally need
+`cu_seqlens / sample_ids / g_seqlens / input_image_mask` in that dict so our
+patched flash-attention function can dispatch to the packed code path.
+
+Rather than monkey-patching upstream's forward, we subclass `_HunYuanVLMoTTextModel`
+and `_HunYuanVLMoTTextForCausalLM` cleanly and replace the decoder layers with
+`MoTDecoderLayer` (three-path mlp_t / mlp_v / mlp_g) inside __init__.
+"""
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs
+
+from transformers.models.hunyuan_vl_mot.modeling_hunyuan_vl_mot import (
+    _HunYuanVLMoTTextModel,
+    _HunYuanVLMoTTextForCausalLM,
+    _modality_mask_to_segments,
+)
+
+from .modeling_decoder_mot import MoTDecoderLayer
+
+
+class MoTTextModel(_HunYuanVLMoTTextModel):
+    """Pure text decoder using MoTDecoderLayer (3-path) and packed kwargs.
+
+    Drop-in replacement for `_HunYuanVLMoTTextModel`:
+      * Replaces all decoder layers with `MoTDecoderLayer`
+      * Forward accepts `cu_seqlens / sample_ids / g_seqlens / input_image_mask`
+        and threads them into the `attention_mask` dict that decoder layers see
+      * `modality_mask` may be int{0,1,2} (text/vision-input/vision-gen);
+        v_seqlens are derived from `input_image_mask` if provided, else from
+        `modality_mask > 0`
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Replace decoder layers with MoT (3-path). Init weights match upstream
+        # for the shared keys; mlp_g/_g get default init (overwritten later by
+        # Net2Wider in mot_init_utils.maybe_init_generation_path).
+        self.layers = nn.ModuleList(
+            [MoTDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        modality_mask: Optional[torch.Tensor] = None,
+        # Packed-mode signals
+        cu_seqlens: Optional[torch.Tensor] = None,
+        sample_ids: Optional[torch.Tensor] = None,
+        g_seqlens: Optional[torch.Tensor] = None,
+        input_image_mask: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        # Position IDs:
+        #   * Packed mode (cu_seqlens given): caller must pass position_ids (per-sample arange)
+        #   * Padded B>1: derive from attention_mask cumsum
+        #   * B==1 / decode: cache_position
+        if position_ids is None:
+            if cu_seqlens is not None:
+                raise ValueError("Packed mode requires position_ids from the collator.")
+            if attention_mask is not None and attention_mask.shape[0] > 1:
+                position_ids = attention_mask.long().cumsum(dim=-1) - 1
+                position_ids = position_ids.clamp(min=0)
+                seq_len = inputs_embeds.shape[1]
+                if position_ids.shape[1] > seq_len:
+                    position_ids = position_ids[:, -seq_len:]
+            else:
+                position_ids = cache_position.unsqueeze(0)
+        text_position_ids = position_ids
+
+        if modality_mask is None:
+            modality_mask = torch.zeros(
+                inputs_embeds.shape[:-1], dtype=torch.long, device=inputs_embeds.device
+            )
+
+        # v_seqlens source: input_image_mask (input images only) when training
+        # flow generation; otherwise modality_mask > 0.
+        if input_image_mask is not None:
+            vis_mask = input_image_mask.bool()
+        else:
+            vis_mask = modality_mask > 0
+        visual_segs = _modality_mask_to_segments(vis_mask)
+
+        # Truncate modality_mask if shape mismatch (decode KV cache)
+        seq_len = inputs_embeds.shape[1]
+        if modality_mask.shape[1] > seq_len:
+            modality_mask = modality_mask[:, -seq_len:]
+
+        # Build extended causal_mask dict with packed signals.
+        # When cu_seqlens is set, padding_mask MUST be None (we treat as packed).
+        # PERF: pre-compute max_seqlen ONCE here so the per-layer attention
+        # forward doesn't redo `.max().item()` (a D2H sync) 32 times.
+        max_seqlen_packed = None
+        if cu_seqlens is not None:
+            with torch.no_grad():
+                max_seqlen_packed = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item())
+
+        causal_mask = {
+            "v_seqlens": visual_segs,
+            "g_seqlens": g_seqlens,
+            "cu_seqlens": cu_seqlens,
+            "sample_ids": sample_ids,
+            "padding_mask": attention_mask if cu_seqlens is None else None,
+            "max_seqlen_packed": max_seqlen_packed,
+        }
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, text_position_ids)
+
+        for decoder_layer in self.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=text_position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                modality_mask=modality_mask,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class MoTTextForCausalLM(_HunYuanVLMoTTextForCausalLM):
+    """Inner text + lm_head wrapper using MoTTextModel.
+
+    Threads packed kwargs and `shift_labels` through. When `shift_labels` is
+    provided, computes split text/image loss aggregates so the outer wrapper
+    can report them separately.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Replace inner text model with MoT (3-path) version
+        self.model = MoTTextModel(config)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        modality_mask: Optional[torch.Tensor] = None,
+        # Packed kwargs threaded to inner MoTTextModel
+        cu_seqlens: Optional[torch.Tensor] = None,
+        sample_ids: Optional[torch.Tensor] = None,
+        g_seqlens: Optional[torch.Tensor] = None,
+        input_image_mask: Optional[torch.Tensor] = None,
+        # Optional pre-shifted labels (used when image regions need -100 masking)
+        shift_labels: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            modality_mask=modality_mask,
+            cu_seqlens=cu_seqlens,
+            sample_ids=sample_ids,
+            g_seqlens=g_seqlens,
+            input_image_mask=input_image_mask,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+
+        # Choose label source: explicit shift_labels (collator-prepared, with
+        # image-region -100 masking) takes precedence over `labels`.
+        train_labels = shift_labels if shift_labels is not None else labels
+
+        if train_labels is not None:
+            flat_hs = hidden_states.reshape(-1, hidden_states.size(-1))
+            flat_labels = train_labels.reshape(-1)
+            valid = flat_labels >= 0
+            if valid.sum() == 0:
+                flat_hs_v = flat_hs[:1]
+                flat_labels_v = flat_labels[:1]
+            else:
+                flat_hs_v = flat_hs[valid]
+                flat_labels_v = flat_labels[valid]
+            logits = self.lm_head(flat_hs_v)
+            loss = self.loss_function(
+                logits=logits, labels=flat_labels_v, vocab_size=self.config.vocab_size, **kwargs
+            )
+        else:
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :])
+            loss = None
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.last_hidden_state,
+        )
+
+
+__all__ = ["MoTTextModel", "MoTTextForCausalLM"]

model/modeling_unified_mot.py

@@ -0,0 +1,523 @@
+"""UnifiedMoT main model: HunYuanVLMoT base + Flow Matching for image generation.
+
+Architecture:
+  * Inherits `transformers.models.hunyuan_vl_mot.HunYuanVLMoTModel` and uses
+    its HYViT2 vision tower via `self.visual`. The inner language model is
+    swapped to `MoTTextForCausalLM` so all decoder layers are 3-path
+    (mlp_t / mlp_v / mlp_g).
+  * Adds Flow Matching modules on the `ForConditionalGeneration` class:
+    `time_embedder`, `vae2llm`, `llm2vae`, `latent_pos_embed`.
+  * Forward supports both packed (P-Pack training) and padded modes.
+  * Computes split text-CE + flow-MSE + optional patch-boundary L1 loss
+    and exposes loss aggregates on the output (image_loss_sum / text_loss_sum /
+    image_loss_count / text_loss_count) for trainer-side weighting.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import ModelOutput
+
+from transformers.models.hunyuan_vl_mot.modeling_hunyuan_vl_mot import (
+    HY_VL_MOT_IMAGE_TOKEN_ID,
+    HY_VL_MOT_VIDEO_TOKEN_ID,
+    HunYuanVLMoTForConditionalGeneration,
+    HunYuanVLMoTModel,
+    HunYuanVLMoTPreTrainedModel,
+    HunYuanVLMoTModelOutputWithPast,
+)
+
+# Side-effect imports: replace upstream attention + extend decoder/text-model
+from . import attention_mot_packed  # noqa: F401  # pylint: disable=unused-import
+
+from .config_unified_mot import UnifiedMoTConfig
+from .modeling_text_model_mot import MoTTextForCausalLM
+from .flow_matching_modules import (
+    TimestepEmbedder,
+    PositionEmbedding,
+    pack_latents,
+    sample_timesteps,
+    build_noisy_latent,
+    patch_boundary_loss,
+)
+
+
+# ---------------------------------------------------------------------------
+# Output dataclass
+# ---------------------------------------------------------------------------
+
+@dataclass
+class UnifiedMoTOutput(ModelOutput):
+    """Output of UnifiedMoTForConditionalGeneration.
+
+    Includes the standard CausalLMOutputWithPast fields plus split text/image
+    loss aggregates so the trainer can apply per-modality weights.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    # Split aggregates (sum-reduced; count of contributing tokens)
+    image_loss_sum: Optional[torch.FloatTensor] = None
+    image_loss_count: Optional[torch.LongTensor] = None
+    text_loss_sum: Optional[torch.FloatTensor] = None
+    text_loss_count: Optional[torch.LongTensor] = None
+    boundary_loss: Optional[torch.FloatTensor] = None
+
+
+# ---------------------------------------------------------------------------
+# UnifiedMoTModel — wraps HunYuanVLMoTModel, swaps language_model to MoT version
+# ---------------------------------------------------------------------------
+
+class UnifiedMoTModel(HunYuanVLMoTModel):
+    """HunYuanVLMoT wrapper with 3-path MoT decoder layers.
+
+    `__init__` first calls upstream's super-__init__ to set up `self.visual`,
+    `self.language_model`, and the misc plumbing, then **rebuilds**
+    `self.language_model` as `MoTTextForCausalLM` so its decoder layers are
+    `MoTDecoderLayer` (with `mlp_g`).
+
+    Forward injects optional `flow_embeds` at `flow_positions` (absolute packed
+    coords) before delegating to language_model. All packed kwargs pass
+    straight through to the inner text model.
+    """
+    config_class = UnifiedMoTConfig
+
+    def __init__(self, config: UnifiedMoTConfig):
+        super().__init__(config)
+        # Replace upstream's language_model with our MoT (3-path) version.
+        # Done via `_from_config` to honor PretrainedConfig conventions.
+        self.language_model = MoTTextForCausalLM._from_config(config)
+
+    # -----------------------------------------------------------------
+    # Forward
+    # -----------------------------------------------------------------
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        labels: Optional[torch.LongTensor] = None,
+        # Packed-mode signals
+        cu_seqlens: Optional[torch.Tensor] = None,
+        sample_ids: Optional[torch.Tensor] = None,
+        modality_mask: Optional[torch.Tensor] = None,
+        input_image_mask: Optional[torch.Tensor] = None,
+        # Flow Matching
+        flow_embeds: Optional[torch.FloatTensor] = None,    # (total_flow_tokens, D)
+        flow_positions: Optional[torch.Tensor] = None,       # (M, 2) absolute packed coords
+        g_seqlens: Optional[torch.Tensor] = None,            # (M, 2) absolute, used by attention
+        **kwargs,
+    ) -> HunYuanVLMoTModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("Specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+
+        # ----- Vision: encode input images/videos and scatter into embeds -----
+        union_mask = None
+        has_input_media = (pixel_values is not None) or (pixel_values_videos is not None)
+        if self.training or has_input_media:
+            if has_input_media:
+                image_embeds, video_embeds, zero_feature = self.get_image_video_features(
+                    pixel_values, pixel_values_videos,
+                    image_grid_thw, video_grid_thw,
+                    inputs_embeds.device, inputs_embeds.dtype,
+                )
+                if len(image_embeds) > 0:
+                    image_embeds_cat = torch.cat(image_embeds, dim=0).to(
+                        inputs_embeds.device, inputs_embeds.dtype
+                    )
+                    image_mask, _, union_mask = self.get_placeholder_mask(
+                        input_ids, inputs_embeds, image_features=image_embeds_cat
+                    )
+                    inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds_cat)
+                if len(video_embeds) > 0:
+                    video_embeds_cat = torch.cat(video_embeds, dim=0).to(
+                        inputs_embeds.device, inputs_embeds.dtype
+                    )
+                    _, video_mask, union_mask = self.get_placeholder_mask(
+                        input_ids, inputs_embeds, video_features=video_embeds_cat
+                    )
+                    inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds_cat)
+                inputs_embeds = inputs_embeds + zero_feature
+
+        # ----- Flow Matching: inject pre-built flow embeds at flow_positions ---
+        if flow_embeds is not None and flow_positions is not None and flow_positions.numel() > 0:
+            seq_len = inputs_embeds.shape[1]
+            replacement = inputs_embeds[0].clone()
+            position_mask = torch.zeros(seq_len, device=inputs_embeds.device, dtype=torch.bool)
+            # PERF: bulk D2H of flow_positions (small tensor) avoids 2 .item()
+            # calls per range × ~25 ranges per forward.
+            flow_positions_cpu = flow_positions.tolist()
+            idx = 0
+            for row in flow_positions_cpu:
+                s = int(row[0])
+                e = int(row[1])
+                n = e - s
+                replacement[s:e] = flow_embeds[idx : idx + n].to(inputs_embeds.dtype)
+                position_mask[s:e] = True
+                idx += n
+            mask_f = position_mask.to(inputs_embeds.dtype).unsqueeze(-1)
+            inputs_embeds = (
+                inputs_embeds[0] * (1.0 - mask_f) + replacement * mask_f
+            ).unsqueeze(0)
+
+        # ----- modality_mask sanity ------------------------------------------
+        if modality_mask is None:
+            if union_mask is not None:
+                # union_mask is bool (input image/video tokens). Promote to int.
+                modality_mask = union_mask.long()
+            else:
+                modality_mask = torch.zeros(
+                    inputs_embeds.shape[:-1], dtype=torch.long, device=inputs_embeds.device
+                )
+
+        # ----- Forward through inner language model -------------------------
+        outputs = self.language_model(
+            input_ids=None,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            labels=labels,
+            modality_mask=modality_mask,
+            cu_seqlens=cu_seqlens,
+            sample_ids=sample_ids,
+            g_seqlens=g_seqlens,
+            input_image_mask=input_image_mask,
+            **kwargs,
+        )
+
+        return HunYuanVLMoTModelOutputWithPast(
+            loss=outputs.loss,
+            logits=outputs.logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+        )
+
+
+# ---------------------------------------------------------------------------
+# UnifiedMoTForConditionalGeneration — main entry point
+# ---------------------------------------------------------------------------
+
+class UnifiedMoTForConditionalGeneration(HunYuanVLMoTForConditionalGeneration):
+    """Full unified model: MoT VL backbone + Flow Matching head."""
+    config_class = UnifiedMoTConfig
+
+    # mlp_g and FM modules don't exist in upstream HunYuanVLMoT base checkpoints;
+    # let `from_pretrained` skip them.
+    _keys_to_ignore_on_load_missing = [
+        r"time_embedder\..*",
+        r"vae2llm\..*",
+        r"llm2vae\..*",
+        r"latent_pos_embed\..*",
+        r".*\.mlp_g\..*",
+        r".*\.input_layernorm_g\..*",
+        r".*\.post_attention_layernorm_g\..*",
+    ]
+
+    def __init__(self, config: UnifiedMoTConfig):
+        # Bypass upstream's __init__ (which would build a regular HunYuanVLMoTModel
+        # and miss our 3-path layers). Init the PreTrainedModel base directly.
+        HunYuanVLMoTPreTrainedModel.__init__(self, config)
+        self.model = UnifiedMoTModel(config)
+        self.config = config
+
+        hidden_size = config.text_config.hidden_size
+
+        # Flow Matching geometry (cached for hot path)
+        self.latent_patch_size = config.latent_patch_size
+        self.timestep_shift = config.timestep_shift
+        self.max_latent_size = config.max_latent_size
+        self.latent_channel = config.vae_z_channels
+        # Effective downsample from pixel to packed-patch latent token grid
+        self.latent_downsample = config.vae_image_downsample  # already includes patch_size
+        self.patch_latent_dim = self.latent_patch_size ** 2 * self.latent_channel
+
+        # FM modules — initialized from scratch (zero-loaded by from_pretrained)
+        self.time_embedder = TimestepEmbedder(hidden_size)
+        self.vae2llm = nn.Linear(self.patch_latent_dim, hidden_size)
+        self.llm2vae = nn.Linear(hidden_size, self.patch_latent_dim)
+        self.latent_pos_embed = PositionEmbedding(self.max_latent_size, hidden_size)
+
+        # Loss weights
+        self.boundary_loss_weight = config.boundary_loss_weight
+
+        self.post_init()
+
+    # -----------------------------------------------------------------
+    # Helpers — delegate vision encoding/placeholder to upstream methods
+    # -----------------------------------------------------------------
+
+    def get_image_features(self, *args, **kwargs):
+        return self.model.get_image_features(*args, **kwargs)
+
+    def get_video_features(self, *args, **kwargs):
+        return self.model.get_video_features(*args, **kwargs)
+
+    def get_image_video_features(self, *args, **kwargs):
+        return self.model.get_image_video_features(*args, **kwargs)
+
+    def get_placeholder_mask(self, *args, **kwargs):
+        return self.model.get_placeholder_mask(*args, **kwargs)
+
+    def get_input_embeddings(self):
+        return self.model.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.model.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, value):
+        self.model.language_model.set_output_embeddings(value)
+
+    # -----------------------------------------------------------------
+    # FM input builder — turn padded VAE latents into flow_embeds
+    # -----------------------------------------------------------------
+
+    def _build_flow_embeds(
+        self,
+        padded_latents: List[torch.Tensor],
+        latent_shapes: List[Tuple[int, int]],
+        latent_positions: Optional[torch.Tensor],
+        device,
+        dtype,
+    ):
+        """Pack VAE latents → noisy x_t → vae2llm + time_embed + pos_embed.
+
+        Returns (flow_embeds, vae_latent_clean, noise, timesteps).
+        """
+        clean = pack_latents(padded_latents, latent_shapes, self.latent_patch_size, self.latent_channel)
+        clean = clean.to(device=device, dtype=dtype)
+        noise = torch.randn_like(clean)
+        timesteps = sample_timesteps(latent_shapes, self.timestep_shift, device, dtype)
+        noisy = build_noisy_latent(clean, noise, timesteps)
+
+        flow_embeds = self.vae2llm(noisy.to(self.vae2llm.weight.dtype)) + self.time_embedder(timesteps)
+        if latent_positions is not None:
+            flow_embeds = flow_embeds + self.latent_pos_embed(latent_positions)
+        return flow_embeds, clean, noise, timesteps
+
+    def _build_input_latent_embeds(
+        self,
+        padded_latents: List[torch.Tensor],
+        latent_shapes: List[Tuple[int, int]],
+        latent_positions: Optional[torch.Tensor],
+        device,
+        dtype,
+    ):
+        """For TI2I: clean (t=0) input image VAE latents projected into LLM space."""
+        clean = pack_latents(padded_latents, latent_shapes, self.latent_patch_size, self.latent_channel)
+        clean = clean.to(device=device, dtype=dtype)
+        t_zero = torch.zeros(clean.shape[0], device=device, dtype=dtype)
+        embeds = self.vae2llm(clean.to(self.vae2llm.weight.dtype)) + self.time_embedder(t_zero)
+        if latent_positions is not None:
+            embeds = embeds + self.latent_pos_embed(latent_positions)
+        return embeds
+
+    # -----------------------------------------------------------------
+    # Forward
+    # -----------------------------------------------------------------
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        labels: Optional[torch.LongTensor] = None,
+        shift_labels: Optional[torch.LongTensor] = None,
+        # Packed signals
+        cu_seqlens: Optional[torch.Tensor] = None,
+        sample_ids: Optional[torch.Tensor] = None,
+        modality_mask: Optional[torch.Tensor] = None,
+        input_image_mask: Optional[torch.Tensor] = None,
+        # FM training inputs
+        use_flow_matching: bool = False,
+        padded_latent: Optional[List[torch.Tensor]] = None,
+        patchified_vae_latent_shapes: Optional[List[Tuple[int, int]]] = None,
+        vae_latent_indexes: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        vae_latent_positions: Optional[torch.Tensor] = None,
+        flow_positions: Optional[torch.Tensor] = None,
+        g_seqlens: Optional[torch.Tensor] = None,
+        # TI2I (input image VAE latent conditioning)
+        input_padded_latent: Optional[List[torch.Tensor]] = None,
+        input_vae_latent_shapes: Optional[List[Tuple[int, int]]] = None,
+        input_vae_latent_indexes: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        input_vae_latent_positions: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> UnifiedMoTOutput:
+        device = (input_ids if input_ids is not None else inputs_embeds).device
+        flow_embeds = None
+        vae_latent_clean = None
+        noise = None
+        timesteps = None
+
+        # ---- TI2I clean input latents (t=0, gets injected at input_*_indexes) -
+        # When provided, prepend them into inputs_embeds via flow_embeds path
+        # using the same flow_positions mechanism. Many call sites only use one
+        # of (input, generation) latents, so we handle them additively.
+        input_flow_embeds = None
+        if input_padded_latent is not None and input_vae_latent_shapes is not None:
+            if inputs_embeds is None and input_ids is not None:
+                inputs_embeds = self.get_input_embeddings()(input_ids)
+                input_ids = None  # already materialized
+            input_flow_embeds = self._build_input_latent_embeds(
+                input_padded_latent, input_vae_latent_shapes,
+                input_vae_latent_positions, device, inputs_embeds.dtype,
+            )
+            # Splice into inputs_embeds at the input_vae_latent_indexes positions
+            if input_vae_latent_indexes is not None:
+                b_idx, s_idx = input_vae_latent_indexes
+                inputs_embeds = inputs_embeds.clone()
+                inputs_embeds[b_idx, s_idx] = input_flow_embeds.to(inputs_embeds.dtype)
+
+        # ---- Generation flow embeds (noisy x_t, used during FM training) ----
+        if use_flow_matching and padded_latent is not None and patchified_vae_latent_shapes:
+            if inputs_embeds is None and input_ids is not None:
+                inputs_embeds = self.get_input_embeddings()(input_ids)
+                input_ids = None
+            flow_embeds, vae_latent_clean, noise, timesteps = self._build_flow_embeds(
+                padded_latent, patchified_vae_latent_shapes,
+                vae_latent_positions, device, inputs_embeds.dtype,
+            )
+
+        # ---- Forward through model -----------------------------------------
+        # Prefer pre-built inputs_embeds when we touched them above.
+        outputs = self.model(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            cache_position=cache_position,
+            logits_to_keep=logits_to_keep,
+            labels=None,  # we compute loss ourselves
+            cu_seqlens=cu_seqlens,
+            sample_ids=sample_ids,
+            modality_mask=modality_mask,
+            input_image_mask=input_image_mask,
+            flow_embeds=flow_embeds,
+            flow_positions=flow_positions,
+            g_seqlens=g_seqlens,
+            **kwargs,
+        )
+
+        last_hidden_state = outputs.hidden_states  # actually carries last_hidden_state per upstream output type
+        logits = outputs.logits
+
+        # ---- Compute losses ------------------------------------------------
+        text_loss_sum = None
+        text_loss_count = None
+        image_loss_sum = None
+        image_loss_count = None
+        boundary_loss = None
+        loss = None
+
+        # ---- Text CE loss (on shift_labels mask)
+        train_labels = shift_labels if shift_labels is not None else labels
+        if train_labels is not None and last_hidden_state is not None:
+            flat_hs = last_hidden_state.reshape(-1, last_hidden_state.size(-1))
+            flat_labels = train_labels.reshape(-1)
+            valid = flat_labels >= 0
+            n_valid = int(valid.sum().item())
+            if n_valid > 0:
+                hs_v = flat_hs[valid]
+                lbl_v = flat_labels[valid]
+                text_logits = self.get_output_embeddings()(hs_v).float()
+                ce_per_tok = F.cross_entropy(text_logits, lbl_v, reduction="none")
+                text_loss_sum = ce_per_tok.sum()
+                text_loss_count = torch.tensor(n_valid, device=device, dtype=torch.long)
+
+        # ---- Flow MSE loss + optional boundary loss
+        if vae_latent_clean is not None and vae_latent_indexes is not None and last_hidden_state is not None:
+            b_idx, s_idx = vae_latent_indexes
+            mse_hidden = last_hidden_state[b_idx, s_idx]
+            mse_preds = self.llm2vae(mse_hidden)
+            target = noise - vae_latent_clean  # velocity target
+            has_mse = timesteps > 0
+            if has_mse.any():
+                mse_pred_f32 = mse_preds[has_mse].float()
+                target_f32 = target[has_mse].float()
+                # Per-token MSE summed for split aggregation; trainer divides by count
+                mse_per_tok = ((mse_pred_f32 - target_f32) ** 2).mean(dim=-1)
+                image_loss_sum = mse_per_tok.sum()
+                # PERF: avoid D2H .item() round-trip — has_mse.sum() is already
+                # a GPU tensor, just cast it to long.
+                image_loss_count = has_mse.sum().long()
+
+            if self.boundary_loss_weight > 0.0:
+                # Recover predicted x_0 = x_t - t * v_pred
+                noisy = build_noisy_latent(vae_latent_clean.detach(), noise.detach(), timesteps.detach())
+                x0_pred = noisy - timesteps[:, None].detach() * mse_preds
+                boundary_loss = patch_boundary_loss(
+                    x0_pred, patchified_vae_latent_shapes,
+                    self.latent_patch_size, self.latent_channel,
+                )
+
+        # ---- Aggregate scalar loss -----------------------------------------
+        # By default: mean text-CE + mean flow-MSE + boundary_weight * boundary
+        parts = []
+        if text_loss_sum is not None and text_loss_count is not None and text_loss_count.item() > 0:
+            parts.append(self.config.text_loss_weight * (text_loss_sum / text_loss_count.float()))
+        if image_loss_sum is not None and image_loss_count is not None and image_loss_count.item() > 0:
+            parts.append(self.config.flow_loss_weight * (image_loss_sum / image_loss_count.float()))
+        if boundary_loss is not None:
+            parts.append(self.boundary_loss_weight * boundary_loss)
+        if parts:
+            loss = parts[0]
+            for extra in parts[1:]:
+                loss = loss + extra
+
+        return UnifiedMoTOutput(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=last_hidden_state,
+            image_loss_sum=image_loss_sum,
+            image_loss_count=image_loss_count,
+            text_loss_sum=text_loss_sum,
+            text_loss_count=text_loss_count,
+            boundary_loss=boundary_loss,
+        )
+
+
+__all__ = [
+    "UnifiedMoTModel",
+    "UnifiedMoTForConditionalGeneration",
+    "UnifiedMoTOutput",
+    "HY_VL_MOT_IMAGE_TOKEN_ID",
+    "HY_VL_MOT_VIDEO_TOKEN_ID",
+]

model/mot_init_utils.py

@@ -0,0 +1,366 @@
+"""Initialization utilities for the MoT generation path (mlp_g, layernorm_g).
+
+Two regimes are handled by a single entry point `maybe_init_generation_path`:
+
+    (a) same-size copy     : mlp_g.intermediate_size == mlp_v.intermediate_size
+                             -> plain state_dict copy (legacy behavior).
+    (b) Net2Wider expansion: mlp_g.intermediate_size >  mlp_v.intermediate_size
+                             -> function-preserving neuron duplication (Chen et al.,
+                             "Net2Net") with small Gaussian noise to break symmetry.
+
+Only integer-multiple widening is exact; non-integer multiples use round-robin
+replicate-and-trim with reciprocal-count scaling (still function-preserving, but
+logged as a warning — integer multiples are strongly preferred).
+
+All layernorm_g state is always plain-copied from layernorm_v because layernorms
+are on hidden_size and unaffected by MLP widening.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+@torch.no_grad()
+def _net2wider_expand_mlp(
+    mlp_g: nn.Module,
+    mlp_v: nn.Module,
+    noise_std: float = 1e-4,
+    generator: Optional[torch.Generator] = None,
+) -> None:
+    """In-place initialize `mlp_g` to be function-equivalent to `mlp_v` via neuron
+    duplication. After this call, for any input x, `mlp_g(x) ~= mlp_v(x)` within
+    O(noise_std * ||x||).
+
+    Expects a SwiGLU-style MLP with linear modules `.gate_proj`, `.up_proj`,
+    `.down_proj`, all bias-free, with shapes:
+        gate_proj / up_proj : (D, H)
+        down_proj           : (H, D)
+    """
+    d_v = mlp_v.gate_proj.out_features
+    d_g = mlp_g.gate_proj.out_features
+    assert mlp_v.gate_proj.in_features == mlp_g.gate_proj.in_features, (
+        f"hidden_size mismatch: mlp_v={mlp_v.gate_proj.in_features}, "
+        f"mlp_g={mlp_g.gate_proj.in_features}"
+    )
+    assert mlp_v.down_proj.out_features == mlp_g.down_proj.out_features, (
+        f"hidden_size mismatch on down_proj: mlp_v={mlp_v.down_proj.out_features}, "
+        f"mlp_g={mlp_g.down_proj.out_features}"
+    )
+    assert d_g >= d_v, (
+        f"mlp_g width {d_g} < mlp_v width {d_v} (shrinking unsupported)"
+    )
+
+    # Same-size fast path: plain copy, ignore noise.
+    if d_g == d_v:
+        mlp_g.load_state_dict(mlp_v.state_dict())
+        return
+
+    if d_g % d_v != 0:
+        logger.warning(
+            f"Net2Wider: d_g={d_g} not an integer multiple of d_v={d_v}; "
+            f"using replicate-and-trim. Integer multiples (e.g. 2x) are strongly "
+            f"preferred for cleanest function preservation."
+        )
+
+    device = mlp_v.gate_proj.weight.device
+    # indices[j] = source neuron index for destination neuron j
+    # Round-robin: for D_g = k*D_v this produces [0,1,...,D_v-1, 0,1,...,D_v-1, ...]
+    # which makes twins "adjacent in source" rather than "adjacent in destination".
+    # Either layout works mathematically; this choice keeps each source's twins
+    # spread out along the destination axis, which is marginally friendlier to
+    # downstream gather/scatter operations if any later code happens to assume
+    # contiguity within a source group.
+    indices = torch.arange(d_g, device=device) % d_v
+    counts = torch.bincount(indices, minlength=d_v)  # (D_v,)
+
+    # --- gate_proj and up_proj: row-wise replication ---------------------------
+    for name in ("gate_proj", "up_proj"):
+        src = getattr(mlp_v, name).weight.data          # (D_v, H)
+        dst_param = getattr(mlp_g, name).weight
+        dst = src[indices].clone().to(dtype=dst_param.dtype)  # (D_g, H)
+        if noise_std > 0:
+            noise = torch.empty(
+                dst.shape, dtype=dst.dtype, device=dst.device,
+            ).normal_(mean=0.0, std=noise_std, generator=generator)
+            dst.add_(noise)
+        dst_param.data.copy_(dst)
+
+    # --- down_proj: column-wise replication with reciprocal-count scaling -----
+    # down_v: (H, D_v), down_g: (H, D_g).
+    # For each destination column j with source s = indices[j]:
+    #     down_g[:, j] = down_v[:, s] / counts[s]
+    # so that   Sum_{j: indices[j]=s}  down_g[:, j] * h_s
+    #         = down_v[:, s] * h_s          (function preserved exactly).
+    down_v = mlp_v.down_proj.weight.data                # (H, D_v)
+    dst_param = mlp_g.down_proj.weight
+    scale = (1.0 / counts.to(down_v.dtype))             # (D_v,)
+    dst = (down_v[:, indices] * scale[indices]).clone().to(dtype=dst_param.dtype)  # (H, D_g)
+    if noise_std > 0:
+        noise = torch.empty(
+            dst.shape, dtype=dst.dtype, device=dst.device,
+        ).normal_(mean=0.0, std=noise_std, generator=generator)
+        dst.add_(noise)
+    dst_param.data.copy_(dst)
+
+
+@torch.no_grad()
+def init_mlp_g_from_mlp_v_net2wider(
+    layer,
+    noise_std: float = 1e-4,
+    generator: Optional[torch.Generator] = None,
+) -> dict:
+    """Initialize a single decoder layer's generation path (mlp_g + layernorm_g)
+    from its vision path (mlp_v + layernorm_v).
+
+    Auto-detects same-size vs expansion case by comparing MLP widths.
+    Layernorms are always plain-copied (shape is hidden_size, unchanged).
+
+    Defensive against meta-device parameters: when transformers'
+    `from_pretrained(dtype=...)` leaves `_g` params on the meta device (because
+    they aren't in the upstream state_dict), a plain `load_state_dict` on
+    those params doesn't materialize them — they stay meta and the `.to(...)`
+    that follows produces garbage GPU memory. To avoid that we rebuild the
+    `_g` parameters as fresh tensors on the same device/dtype as their `_v`
+    counterparts.
+    """
+    d_v = layer.mlp_v.gate_proj.out_features
+    d_g = layer.mlp_g.gate_proj.out_features
+
+    # Materialize meta-device params before in-place copy: rebuild _g linear
+    # weights on the same device/dtype as _v.
+    def _materialize_linear(dst_lin: nn.Linear, src_lin: nn.Linear):
+        if dst_lin.weight.is_meta or dst_lin.weight.device.type == "meta":
+            new_w = torch.empty_like(src_lin.weight)
+            dst_lin.weight = nn.Parameter(new_w, requires_grad=dst_lin.weight.requires_grad)
+            if dst_lin.bias is not None:
+                new_b = torch.empty_like(src_lin.bias) if src_lin.bias is not None else torch.zeros(
+                    dst_lin.out_features, dtype=src_lin.weight.dtype, device=src_lin.weight.device,
+                )
+                dst_lin.bias = nn.Parameter(new_b, requires_grad=dst_lin.bias.requires_grad)
+
+    def _materialize_norm(dst_norm, src_norm):
+        if dst_norm.weight.is_meta or dst_norm.weight.device.type == "meta":
+            new_w = torch.empty_like(src_norm.weight)
+            dst_norm.weight = nn.Parameter(new_w, requires_grad=dst_norm.weight.requires_grad)
+
+    for name in ("gate_proj", "up_proj", "down_proj"):
+        if d_g == d_v:
+            _materialize_linear(getattr(layer.mlp_g, name), getattr(layer.mlp_v, name))
+        # else: widening path will rebuild mlp_g entirely in maybe_init_generation_path
+    _materialize_norm(layer.input_layernorm_g, layer.input_layernorm_v)
+    _materialize_norm(layer.post_attention_layernorm_g, layer.post_attention_layernorm_v)
+
+    _net2wider_expand_mlp(
+        layer.mlp_g, layer.mlp_v,
+        noise_std=noise_std if d_g > d_v else 0.0,
+        generator=generator,
+    )
+    layer.input_layernorm_g.load_state_dict(layer.input_layernorm_v.state_dict())
+    layer.post_attention_layernorm_g.load_state_dict(
+        layer.post_attention_layernorm_v.state_dict()
+    )
+    return {"D_v": d_v, "D_g": d_g, "expanded": d_g > d_v, "noise_std": noise_std}
+
+
+@torch.no_grad()
+def verify_mlp_g_equals_mlp_v(
+    layer,
+    num_samples: int = 4,
+    seq_len: int = 8,
+) -> float:
+    """Diagnostic: forward random activations through both paths (under their
+    respective post-attention layernorms) and return max absolute difference.
+
+    Called from unit tests and logged right after init on layer 0 as a sanity
+    check. Uses fp32 computation when possible (casts output back if needed).
+    """
+    hidden = layer.hidden_size
+    dev = layer.mlp_v.gate_proj.weight.device
+    dt = layer.mlp_v.gate_proj.weight.dtype
+    x = torch.randn(num_samples, seq_len, hidden, device=dev, dtype=dt)
+    y_v = layer.mlp_v(layer.post_attention_layernorm_v(x))
+    y_g = layer.mlp_g(layer.post_attention_layernorm_g(x))
+    return (y_v.float() - y_g.float()).abs().max().item()
+
+
+def checkpoint_has_g_keys(model_load_path: Optional[str]) -> bool:
+    """True iff the checkpoint on disk already contains `_g.` weights.
+
+    This is the robust way to tell "generation path was previously saved; load
+    it as-is" from "generation path is absent; we need to initialize it from
+    mlp_v".
+
+    Returns False if the path does not exist or no safetensors file is found
+    (caller may fall back to a tensor-level check if needed).
+    """
+    if model_load_path is None:
+        return False
+    import os
+    import json as _json
+    index_path = os.path.join(model_load_path, "model.safetensors.index.json")
+    single_path = os.path.join(model_load_path, "model.safetensors")
+    if os.path.exists(index_path):
+        with open(index_path) as f:
+            keys = set(_json.load(f)["weight_map"].keys())
+        return any("_g." in k for k in keys)
+    if os.path.exists(single_path):
+        try:
+            from safetensors import safe_open
+        except ImportError:
+            return False
+        with safe_open(single_path, framework="pt") as f:
+            return any("_g." in k for k in f.keys())
+    return False
+
+
+def _resolve_decoder_layers(model):
+    """Walk common wrapper paths to find the list of transformer decoder layers.
+
+    Supports:
+        UnifiedMoTForConditionalGeneration -> .model.language_model.model.layers
+        bare HunYuan MoT model            -> .language_model.model.layers
+        even-barer HF model               -> .model.layers
+    """
+    for path in (
+        "model.language_model.model.layers",
+        "language_model.model.layers",
+        "model.layers",
+    ):
+        obj = model
+        ok = True
+        for attr in path.split("."):
+            if not hasattr(obj, attr):
+                ok = False
+                break
+            obj = getattr(obj, attr)
+        if ok:
+            return obj
+    raise AttributeError(
+        "Could not locate decoder layers on model; tried "
+        "model.language_model.model.layers / language_model.model.layers / model.layers"
+    )
+
+
+def maybe_init_generation_path(
+    model,
+    model_load_path: Optional[str] = None,
+    noise_std: float = 1e-4,
+    logger_: Optional[logging.Logger] = None,
+    seed: int = 0xA17C0D1F,
+    target_mlp_g_intermediate_size: Optional[int] = None,
+) -> bool:
+    """Top-level entry point used by train / inference / eval scripts.
+
+    Behavior:
+      * If `model_load_path` is provided and the checkpoint contains any `_g.`
+        key, do nothing and return False (generation path was already loaded
+        from the checkpoint — this is the T2I-resume / TI2I-continuation path).
+      * Otherwise iterate every decoder layer and call
+        `init_mlp_g_from_mlp_v_net2wider`:
+            - Same-size case  -> plain state_dict copy (noise ignored).
+            - Expansion case  -> Net2Wider neuron duplication + Gaussian noise
+              (std = `noise_std`).
+
+    `target_mlp_g_intermediate_size`: if provided and larger than the current
+    `mlp_g.intermediate_size`, REBUILD each layer's `mlp_g` to the wider size
+    before running Net2Wider init. Use this when the base checkpoint was
+    loaded with mlp_g at mlp_v's width (so from_pretrained accepted the
+    shapes) and you want to widen post-load.
+
+    Returns True if initialization was performed.
+    """
+    log = logger_ or logger
+
+    # ---- Re-materialize the latent_pos_embed (sin-cos table) ---------------
+    # PositionEmbedding's parameter is computed at __init__ time but
+    # `from_pretrained(dtype=...)` uses meta-init that bypasses the
+    # constructor's data assignment, leaving the param uninitialized.
+    pos_embed_module = None
+    for path in ("latent_pos_embed", "model.latent_pos_embed"):
+        obj = model
+        ok = True
+        for attr in path.split("."):
+            if not hasattr(obj, attr):
+                ok = False
+                break
+            obj = getattr(obj, attr)
+        if ok:
+            pos_embed_module = obj
+            break
+    if pos_embed_module is not None and hasattr(pos_embed_module, "_reset_parameters"):
+        pos_embed_module._reset_parameters()
+        log.info(
+            f"latent_pos_embed re-initialized from sin-cos table "
+            f"(shape={tuple(pos_embed_module.pos_embed.shape)})"
+        )
+
+    if checkpoint_has_g_keys(model_load_path):
+        log.info(
+            "Generation path present in checkpoint (found `_g.` keys); "
+            "skipping mlp_g initialization."
+        )
+        return False
+
+    layers = _resolve_decoder_layers(model)
+    d_v = layers[0].mlp_v.gate_proj.out_features
+    d_g_current = layers[0].mlp_g.gate_proj.out_features
+
+    # ---- Optional: rebuild mlp_g at a wider size before init ----
+    if target_mlp_g_intermediate_size is not None and target_mlp_g_intermediate_size > d_g_current:
+        from transformers.models.hunyuan_vl_mot.modeling_hunyuan_vl_mot import HunYuanVLMoTMLP
+        for layer in layers:
+            base_cfg = layer.mlp_g.gate_proj.weight  # to inherit device/dtype
+            device = base_cfg.device
+            dtype = base_cfg.dtype
+            # Use a duck-typed mini config matching upstream MLP's expected fields
+            class _MiniCfg:
+                pass
+            mini = _MiniCfg()
+            mini.hidden_size = layer.mlp_v.gate_proj.in_features
+            mini.intermediate_size = target_mlp_g_intermediate_size
+            mini.hidden_act = "silu"
+            mini.mlp_bias = (layer.mlp_v.gate_proj.bias is not None)
+            new_mlp_g = HunYuanVLMoTMLP(mini).to(device=device, dtype=dtype)
+            layer.mlp_g = new_mlp_g
+        d_g = target_mlp_g_intermediate_size
+        log.info(f"Rebuilt mlp_g for all layers at intermediate_size={d_g}.")
+    else:
+        d_g = d_g_current
+
+    expansion = d_g > d_v
+    mode = "copy" if not expansion else f"net2wider (x{d_g / d_v:.3f})"
+    log.info(
+        f"Initializing mlp_g/layernorm_g from mlp_v/layernorm_v: "
+        f"mode={mode} (d_v={d_v}, d_g={d_g}, "
+        f"noise_std={noise_std if expansion else 0.0:g})"
+    )
+
+    # Single deterministic generator for reproducibility across layers.
+    gen = torch.Generator(device=layers[0].mlp_v.gate_proj.weight.device)
+    gen.manual_seed(seed)
+
+    for layer in layers:
+        init_mlp_g_from_mlp_v_net2wider(
+            layer,
+            noise_std=noise_std if expansion else 0.0,
+            generator=gen,
+        )
+
+    # Spot-check layer 0: forward a random input through both paths. The diff
+    # should be ~O(noise_std * sqrt(D_v) * ||x||); much larger indicates a bug.
+    try:
+        diff = verify_mlp_g_equals_mlp_v(layers[0])
+        log.info(
+            f"  post-init forward diff (layer 0, random input): "
+            f"max |mlp_v - mlp_g| = {diff:.3e}"
+        )
+    except Exception as e:  # pylint: disable=broad-except  # pragma: no cover - diagnostic only
+        log.warning(f"verify_mlp_g_equals_mlp_v diagnostic failed: {e}")
+
+    return True

multiframe_inference.py

@@ -0,0 +1,391 @@
+"""Joint fixed-N multi-frame generation inference (world-model / x_to_N).
+
+Unlike `interleave_inference.py` (SigLIP-handoff autoregressive rollout — a prior
+frame re-enters the next step as a *clean SigLIP image*), multi-frame generation
+lays out **all N `<Image>` blocks in ONE sequence** and integrates them **jointly**
+with a single Euler ODE. A prior frame stays in context as its **VAE latent**
+(at the shared ODE timestep), exactly matching training, where the assistant turn
+held N latent blocks denoised together (per-frame independent timesteps + hybrid
+block attention: causal across frames, bidirectional within). So inference is:
+
+    seq = prompt(+ViT obs) + "open text" + N×(<Image>[LAT]*n</Image>) + EOS
+    x_k ~ N(0,1) for k in 1..N
+    for t in linspace(1,0,steps):                       # ONE shared schedule
+        flow_embed_k = vae2llm(x_k) + time(t) + pos     # k = 1..N, concatenated
+        hidden       = model(... flow_embeds @ N flow_positions ...)   # ONE forward
+        x_k         -= dt * llm2vae(hidden @ block_k)   # k = 1..N
+    frame_k = VAE.decode(x_k)
+
+The N frames cohere because frame k attends to frames 1..k (cross-frame causal),
+all in the single forward pass — no per-frame autoregressive loop.
+
+Alignment with training (this is the load-bearing correspondence)
+-----------------------------------------------------------------
+TRAINING (modeling_unified_mot._build_flow_embeds + flow_matching_modules.sample_timesteps):
+  the N target frames are packed into one sequence; `sample_timesteps` draws ONE
+  INDEPENDENT timestep per frame (`randn(len(shapes))`), each broadcast over that
+  frame's tokens; `x_t=(1-t)·clean+t·noise` per frame; a SINGLE forward predicts
+  the velocity for ALL N frames and the FM-MSE is summed over all N. So: every
+  frame independently noised, one forward → N frames.
+
+INFERENCE (here): one forward PER Euler step injects all N frames' current x_t
+  (each with its own `time_embedder(t_k)`) at the N flow_positions and reads the
+  velocity for ALL N frames — structurally identical to the training forward
+  (one forward, N frames, per-frame time embedding). We therefore do NOT loop
+  per frame and do NOT re-encode prior frames via SigLIP.
+
+The only inference-time choice is the per-frame timestep SCHEDULE `t_k(step)`:
+  - lockstep (default, validated ~28 dB): all frames share the same t each step,
+    swept 1→0 together. This is the equal-t diagonal of the joint t-space that the
+    independent-per-frame training already covers, so it is in-distribution.
+  - Because training randomized t per frame (incl. cases where earlier frames are
+    much cleaner than later ones), a staggered "diffusion-forcing" schedule
+    (frame k offset so it lags frame k-1) is also in-distribution and tends to
+    help long autoregressive rollouts — left as a future `--schedule` option.
+  Either way each step is ONE forward over all N frames; we never collapse the
+  model's per-frame timestep capability into N separate passes.
+
+Run:
+    python multiframe_inference.py --ckpt <ckpt> --vae <ae.safetensors> \
+        --frames obs0.jpg obs1.jpg --task "imagine the next frames" \
+        --num_frames 4 --num_steps 50 --out_dir multiframe_out
+"""
+from __future__ import annotations
+
+import argparse
+import os
+from typing import List, Optional
+
+import torch
+from PIL import Image
+from transformers.models.hunyuan_vl_mot import HunYuanVLMoTProcessor
+
+# Reuse the (obs frames + task) prompt constructor, ViT placeholder ids, and the
+# montage helper from the interleave inference module.
+from interleave_inference import (
+    build_conditioned_sequence,
+    INPUT_IMAGE_PLACEHOLDER_IDS,
+    _row,
+)
+
+
+@torch.no_grad()
+def _ar_decode_opening_text(
+    inner, processor, prompt_ids, pixel_values, image_grid_thw,
+    image_start_id, eos_id, device, max_text_tokens,
+):
+    """Greedily decode the assistant opening text until the model emits <Image>.
+
+    Mirrors the text loop in interleave_inference.generate_step_joint. Returns the
+    decoded token ids (without the terminating <Image>/EOS)."""
+    seq = list(prompt_ids)
+    text_ids: List[int] = []
+
+    def logits_last(ids):
+        seq_len = len(ids)
+        inp = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+        mod = torch.zeros(1, seq_len, dtype=torch.long, device=device)
+        iim = torch.zeros(1, seq_len, dtype=torch.bool, device=device)
+        for pid in INPUT_IMAGE_PLACEHOLDER_IDS:
+            m = inp[0] == pid
+            mod[0, m] = 1
+            iim[0, m] = True
+        out = inner(
+            input_ids=inp, inputs_embeds=None, attention_mask=None,
+            position_ids=torch.arange(seq_len, device=device).unsqueeze(0),
+            pixel_values=pixel_values, image_grid_thw=image_grid_thw,
+            cu_seqlens=torch.tensor([0, seq_len], dtype=torch.int32, device=device),
+            sample_ids=torch.zeros(1, seq_len, dtype=torch.int32, device=device),
+            modality_mask=mod, input_image_mask=iim,
+            flow_embeds=None, flow_positions=None,
+            g_seqlens=torch.zeros((0, 2), dtype=torch.int32, device=device),
+        )
+        return out.logits[0, -1]
+
+    for _ in range(max_text_tokens):
+        nxt = int(logits_last(seq).argmax().item())
+        if nxt == image_start_id or nxt == eos_id:
+            break
+        text_ids.append(nxt)
+        seq.append(nxt)
+    return text_ids
+
+@torch.no_grad()
+def generate_multiframe_joint(
+    model, vae, processor,
+    obs_frames: List[str], task_text: str, num_frames: int,
+    height: int, width: int, num_steps: int, device, dtype,
+    max_text_tokens: int = 96, opening_text: Optional[str] = None,
+):
+    """Jointly denoise N frames in one sequence. Returns (opening_text, [imgs])."""
+    from model.flow_matching_modules import unpatchify_latent
+    from text2image_inference import get_2d_position_ids
+
+    cfg = model.config
+    eos = cfg.eos_token_id
+    # Multi-frame uses a single <Video> ... </Video> wrapper instead of N×<Image>.
+    video_start = cfg.video_start_token_id  # 120122
+    video_end = cfg.video_end_token_id      # 120123
+    latent_ph = cfg.flow_latent_placeholder_id
+    inner = model.model
+
+    p = model.latent_patch_size
+    ds = cfg.vae_image_downsample
+    h_lat, w_lat = height // ds, width // ds
+    n_latent = h_lat * w_lat
+    patch_dim = p * p * cfg.vae_z_channels
+
+    # Append the multi-frame task instruction to the user turn (train==infer parity).
+    # build_conditioned_sequence folds task_text into the prompt.
+    from inference_utils import TASK_INSTRUCTION_MULTI_FRAME
+    instr_text = (task_text + "\n" + TASK_INSTRUCTION_MULTI_FRAME) if task_text else TASK_INSTRUCTION_MULTI_FRAME
+    prompt_ids, proc = build_conditioned_sequence(processor, obs_frames, instr_text)
+    pixel_values = proc.get("pixel_values")
+    image_grid_thw = proc.get("image_grid_thw")
+    if pixel_values is not None:
+        pixel_values = pixel_values.to(device=device, dtype=dtype)
+    if image_grid_thw is not None:
+        image_grid_thw = image_grid_thw.to(device=device)
+
+    # ---- opening text: teacher-force from GT or AR-decode until <Video> ----
+    if opening_text is not None:
+        text_ids = processor.tokenizer.encode(opening_text, add_special_tokens=False)
+        text_str = opening_text
+    else:
+        text_ids = _ar_decode_opening_text(
+            inner, processor, prompt_ids, pixel_values, image_grid_thw,
+            video_start, eos, device, max_text_tokens,
+        )
+        text_str = processor.tokenizer.decode(text_ids, skip_special_tokens=True)
+
+    # ---- lay out the full sequence: <Video> + N×[LAT] + </Video> (bare latents,
+    # matches the dataset's multi-frame layout; one flow_positions row per frame) ----
+    seq = list(prompt_ids) + list(text_ids)
+    seq.append(video_start)
+    frame_spans = []  # (latent_start, latent_end) per frame
+    for _ in range(num_frames):
+        ls = len(seq)
+        seq.extend([latent_ph] * n_latent)
+        le = len(seq)
+        frame_spans.append((ls, le))
+    seq.append(video_end)
+    seq.append(eos)
+    seq_len = len(seq)
+    input_ids = torch.tensor(seq, dtype=torch.long, device=device).unsqueeze(0)
+
+    mod = torch.zeros(1, seq_len, dtype=torch.long, device=device)
+    iim = torch.zeros(1, seq_len, dtype=torch.bool, device=device)
+    for pid in INPUT_IMAGE_PLACEHOLDER_IDS:
+        m = input_ids[0] == pid
+        mod[0, m] = 1
+        iim[0, m] = True
+    for ls, le in frame_spans:
+        mod[0, ls:le] = 2  # generation-latent route
+
+    # flow_positions / g_seqlens in frame order — must match flow_embeds concat order
+    # (model injects flow_embeds slices into flow_positions rows in order; see
+    #  modeling_unified_mot.py:177-183).
+    flow_positions = torch.tensor(
+        [[ls, le] for ls, le in frame_spans], dtype=torch.int32, device=device
+    )
+    g_seqlens = flow_positions.clone()
+    latent_pos_ids = get_2d_position_ids(h_lat, w_lat, cfg.max_latent_size).to(device)
+    cu = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
+    sid = torch.zeros(1, seq_len, dtype=torch.int32, device=device)
+    pos = torch.arange(seq_len, device=device).unsqueeze(0)
+
+    # ---- N latent buffers, joint Euler ODE (ONE forward/step over all N frames) ----
+    # Per-frame timestep schedule t_k(step): training noised each frame at its OWN t,
+    # so we keep a per-frame t here too. Default = lockstep (all frames share the
+    # step's t) — the validated in-distribution schedule. `frame_t_offset[k]` is the
+    # hook for a staggered "diffusion-forcing" schedule (kept 0 = lockstep).
+    xs = [torch.randn(n_latent, patch_dim, device=device, dtype=dtype) for _ in range(num_frames)]
+    ts = torch.linspace(1.0, 0.0, num_steps + 1, device=device, dtype=dtype)
+    frame_t_offset = [0.0] * num_frames  # all 0 -> lockstep; set per-frame for diffusion-forcing
+    pos_emb = model.latent_pos_embed(latent_pos_ids).to(dtype)  # same per-frame (ids restart at 0)
+    for i in range(num_steps):
+        dt = ts[i] - ts[i + 1]
+        fe_parts = []
+        for k in range(num_frames):
+            t_k = (ts[i] + frame_t_offset[k]).clamp(0.0, 1.0)        # this frame's timestep
+            time_emb_k = model.time_embedder(t_k.expand(n_latent)).to(dtype)
+            x_proj = model.vae2llm(xs[k].to(model.vae2llm.weight.dtype)).to(dtype)
+            fe_parts.append(x_proj + time_emb_k + pos_emb)
+        flow_embeds = torch.cat(fe_parts, dim=0)  # (N*n_latent, D), frame order == flow_positions
+        out = inner(
+            input_ids=input_ids, inputs_embeds=None, attention_mask=None,
+            position_ids=pos, pixel_values=pixel_values, image_grid_thw=image_grid_thw,
+            cu_seqlens=cu, sample_ids=sid, modality_mask=mod, input_image_mask=iim,
+            flow_embeds=flow_embeds, flow_positions=flow_positions, g_seqlens=g_seqlens,
+        )
+        hidden = out.hidden_states
+        for k, (ls, le) in enumerate(frame_spans):
+            v = model.llm2vae(hidden[0, ls:le]).to(dtype)
+            xs[k] = xs[k] - dt * v
+
+    # ---- decode each frame ----
+    imgs = []
+    vae_dtype = next(vae.parameters()).dtype
+    for k in range(num_frames):
+        x_lat = unpatchify_latent(xs[k].float(), h_lat, w_lat, p, cfg.vae_z_channels)
+        x_lat = x_lat.unsqueeze(0).to(device=device, dtype=vae_dtype)
+        img = vae.decode(x_lat)
+        if hasattr(img, "sample"):
+            img = img.sample
+        imgs.append(img.squeeze(0).float().clamp(-1, 1))
+    return text_str, imgs
+
+
+def vae_target_hw(img_path: str, max_size: int = 256, min_size: int = 96, stride: int = 16):
+    """(h_px, w_px) that training's VAEImageTransform produces for this image — so the
+    generated frame's latent grid matches what the model was trained to output. Runs the
+    REAL VAEImageTransform for exact parity; returns None on failure.
+
+    Why this matters: training resized each target frame aspect-preserving to <= max_size,
+    stride-divisible — jaka/umi 848x480 -> 256x144 (144 latent tok), xtrainer 640x480 ->
+    256x192 (192 latent tok). Generating every robot at a fixed 256x144 is wrong for
+    xtrainer (4:3): wrong token count + wrong aspect, which corrupts both its output and
+    its PSNR-vs-GT comparison (GT got squished 4:3 -> 16:9).
+    """
+    try:
+        from inference_utils import VAEImageTransform
+        t = VAEImageTransform(max_size=max_size, min_size=min_size, stride=stride)
+        tensor = t(Image.open(img_path).convert("RGB"))   # (3, H, W)
+        return int(tensor.shape[1]), int(tensor.shape[2])
+    except Exception:  # pylint: disable=broad-except  # any failure → caller falls back
+        return None
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Joint fixed-N multi-frame generation.")
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--vae", required=True)
+    ap.add_argument("--frames", nargs="+", required=True, help="observation frame path(s)")
+    ap.add_argument("--task", required=True, help="overall task text")
+    ap.add_argument("--num_frames", type=int, default=4,
+                    help="number of future frames to generate")
+    ap.add_argument("--out_dir", default="multiframe_out")
+    ap.add_argument("--height", type=int, default=None,
+                    help="override auto-derived generation height (px). Default: derive "
+                         "per-record from the GT/obs frame via the training VAEImageTransform "
+                         "(jaka/umi 848x480 -> 144, xtrainer 640x480 -> 192).")
+    ap.add_argument("--width", type=int, default=None,
+                    help="override auto-derived generation width (px). Default: derived (256).")
+    ap.add_argument("--max_image_size", type=int, default=256,
+                    help="VAEImageTransform max side — MUST match training (256).")
+    ap.add_argument("--min_image_size", type=int, default=96,
+                    help="VAEImageTransform min side — MUST match training (96).")
+    ap.add_argument("--image_stride", type=int, default=16,
+                    help="VAEImageTransform stride — MUST match training (16).")
+    ap.add_argument("--num_steps", type=int, default=50)
+    ap.add_argument("--max_text_tokens", type=int, default=96)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    ap.add_argument("--understanding_max_pixels", type=int, default=524288,
+                    help="Cap obs-frame ViT input pixels to MATCH training. "
+                         "Default 524288 (~494 tok/frame); the model default 4194304 (~4050 tok/frame) "
+                         "is an 8x train/infer resolution mismatch that degrades eval. Set 0 to disable.")
+    args = ap.parse_args()
+
+    from model import UnifiedMoTForConditionalGeneration, maybe_init_generation_path
+    from vae_model.autoencoder import load_ae
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    torch.manual_seed(args.seed)
+
+    # ---- resolve obs frames + task ----
+    obs, task = args.frames, args.task
+    n_frames = args.num_frames or 1
+    if not n_frames:
+        raise ValueError("num_frames resolved to 0 — pass --num_frames >= 1")
+
+    processor = HunYuanVLMoTProcessor.from_pretrained(args.ckpt, trust_remote_code=True)
+    # train==infer parity: cap obs-frame ViT pixels to the SAME value training used.
+    # Without this, inference obs frames are ~8x higher-res than what the model saw at train.
+    if args.understanding_max_pixels and args.understanding_max_pixels > 0:
+        ip = processor.image_processor
+        ip.max_pixels = args.understanding_max_pixels
+        if isinstance(getattr(ip, "size", None), dict) and "longest_edge" in ip.size:
+            ip.size["longest_edge"] = args.understanding_max_pixels
+    model = UnifiedMoTForConditionalGeneration.from_pretrained(args.ckpt, dtype=dtype)
+    maybe_init_generation_path(model, model_load_path=args.ckpt)
+    model.to(device).eval()
+    vae, _ = load_ae(args.vae)
+    vae.requires_grad_(False)
+    vae.eval()
+    vae.to(device, dtype=dtype)
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    print(f"JOINT multi-frame | obs={len(obs)} frame(s) | num_frames={n_frames}")
+    print(f"TASK: {task}\n")
+
+    # ---- generation resolution: match training's VAEImageTransform per record ----
+    # Training resized each TARGET frame aspect-preserving to <=max_image_size,
+    # stride-divisible (jaka/umi -> 256x144, xtrainer -> 256x192). Generating every
+    # robot at a fixed 256x144 was wrong for xtrainer (4:3). Derive from the data;
+    # explicit --height/--width override.
+    if args.height and args.width:
+        height, width = args.height, args.width
+    else:
+        ref = obs[0] if obs else None
+        hw = (vae_target_hw(ref, args.max_image_size, args.min_image_size, args.image_stride)
+              if ref else None)
+        if hw is None:
+            height, width = 144, 256
+            print("[warn] could not derive resolution from data; falling back to 144x256")
+        else:
+            height, width = hw
+    ds = model.config.vae_image_downsample
+    print(f"  gen resolution: {width}x{height} px  (latent {width // ds}x{height // ds} = {(width // ds) * (height // ds)} tok/frame)")
+
+    text_str, imgs = generate_multiframe_joint(
+        model, vae, processor, obs, task, n_frames,
+        height, width, args.num_steps, device, dtype,
+        max_text_tokens=args.max_text_tokens, opening_text=None,
+    )
+
+    # ---- save frames + montage + result.txt ----
+    gen_paths = []
+    for k, img in enumerate(imgs):
+        arr = ((img.cpu().permute(1, 2, 0).numpy() + 1.0) * 127.5).clip(0, 255).astype("uint8")
+        path = os.path.join(args.out_dir, f"frame{k + 1}.png")
+        Image.fromarray(arr).save(path)
+        gen_paths.append(path)
+
+    lines = [f"TASK: {task}", "",
+             f"OPEN text (decoded): {text_str}"]
+    lines.append("")
+    for k in range(len(imgs)):
+        lines.append(f"frame {k + 1}")
+    txt_path = os.path.join(args.out_dir, "result.txt")
+    open(txt_path, "w").write("\n".join(lines))
+    print("\n".join(lines))
+
+    from PIL import ImageDraw
+    # montage: INPUT (obs frames) / GEN (ours),
+    # each a horizontal strip of W×H cells (W,H = the derived generation resolution),
+    # with a left label column.
+    label_w = 64
+    gap = 6
+    rows = [("INPUT", list(obs)),
+            ("GEN", gen_paths)]
+    rows = [(lbl, paths) for lbl, paths in rows if paths]   # drop empty
+    strips = [(lbl, _row(paths, width, height)) for lbl, paths in rows]
+    montage_w = label_w + max(s.width for _, s in strips)
+    montage_h = sum(s.height for _, s in strips) + gap * (len(strips) - 1)
+    montage = Image.new("RGB", (montage_w, montage_h), (0, 0, 0))
+    draw = ImageDraw.Draw(montage)
+    y = 0
+    for lbl, strip in strips:
+        montage.paste(strip, (label_w, y))
+        draw.text((6, y + strip.height // 2 - 4), lbl, fill=(255, 255, 255))
+        y += strip.height + gap
+    out_png = os.path.join(args.out_dir, "multiframe_input_GEN.png")
+    montage.save(out_png)
+    print(f"\nGenerated {len(imgs)} frames -> {args.out_dir}")
+    print(f"  montage (INPUT/GEN): {out_png}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,9 @@
+--extra-index-url https://download.pytorch.org/whl/cu130
+torch==2.10.0
+torchvision==0.25.0
+safetensors
+einops
+numpy
+pillow
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt210-cu130-cp312/flash_attn-2.8.3-cp312-cp312-linux_x86_64.whl
+git+https://github.com/huggingface/transformers@9293856c419762ebf98fbe2bd9440f9ce7069f1a

text2image_inference.py

@@ -0,0 +1,267 @@
+"""Minimal text-to-image inference for the packed UnifiedMoT model.
+
+Performs Euler-ODE sampling in the LLM hidden space:
+  x_T ~ N(0, 1)
+  for t in linspace(1, 0, num_steps + 1)[:-1]:
+      flow_embed = vae2llm(x_t) + time_embedder(t) + latent_pos_embed
+      hidden    = model(... flow_embed injected at latent positions ...)
+      v_pred    = llm2vae(hidden_at_latent_positions)
+      x_{t-dt} = x_t - dt * v_pred
+
+Then VAE-decodes the final x_0 to pixels.
+
+This script is single-sample / single-image. Classifier-free guidance (CFG)
+is optional via ``--cfg_scale`` (>1 enables it): each ODE step runs a second
+forward with an empty prompt, matching the 10% ``text_cond_dropout`` used in
+training. ``--cfg_scale 1`` (default) disables CFG for the fastest path.
+
+Run:
+    # no CFG (fastest)
+    python text2image_inference.py --ckpt /path/to/checkpoint \\
+                               --prompt "a watercolor cat" \\
+                               --vae /path/to/vae.safetensors \\
+                               --out out.png \\
+                               --height 256 --width 256 --num_steps 25
+
+    # with CFG
+    python text2image_inference.py --ckpt /path/to/checkpoint \\
+                               --prompt "a watercolor cat" \\
+                               --vae /path/to/vae.safetensors \\
+                               --cfg_scale 5.0 --num_steps 50
+"""
+from __future__ import annotations
+
+import argparse
+
+import torch
+from PIL import Image
+from transformers.models.hunyuan_vl_mot import HunYuanVLMoTProcessor
+
+from model import (
+    UnifiedMoTConfig,
+    UnifiedMoTForConditionalGeneration,
+    maybe_init_generation_path,
+)
+from model.flow_matching_modules import (
+    unpatchify_latent,
+)
+from vae_model.autoencoder import load_ae
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", required=True, help="Model checkpoint directory")
+    p.add_argument("--vae", required=True, help="VAE safetensors path")
+    p.add_argument("--prompt", required=True)
+    p.add_argument("--out", default="out.png")
+    p.add_argument("--height", type=int, default=256)
+    p.add_argument("--width", type=int, default=256)
+    p.add_argument("--num_steps", type=int, default=25)
+    p.add_argument("--cfg_scale", type=float, default=1.0,
+                   help="Classifier-free guidance scale. >1 enables CFG "
+                        "(runs an extra empty-prompt forward per step; 2-5 typical). "
+                        "1.0 disables CFG.")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    return p.parse_args()
+
+
+def get_2d_position_ids(h: int, w: int, max_per_side: int) -> torch.Tensor:
+    """2D-flattened position ids matching model.PositionEmbedding lookup table."""
+    rows = torch.arange(h)[:, None] * max_per_side
+    cols = torch.arange(w)[None, :]
+    return (rows + cols).reshape(-1).long()
+
+
+@torch.no_grad()
+def _build_seq_meta(processor, prompt, cfg, n_latent_tokens, device):
+    """Build the packed (1, T) input sequence + routing tensors for one prompt.
+
+    Sequence:
+        chat_template(user/assistant) + <Image> + LATENT*N + </Image> + EOS
+    Returns a dict of everything the per-step forward needs (constant across
+    ODE steps except the latent region, which the caller patches each step).
+    """
+    eos = cfg.eos_token_id
+    prompt_messages = [
+        {"role": "user", "content": [{"type": "text", "text": prompt}]},
+        {"role": "assistant", "content": ""},
+    ]
+    prompt_inputs = processor.apply_chat_template(
+        prompt_messages, return_dict=True, tokenize=True, add_generation_prompt=False,
+    )
+    prompt_ids = list(prompt_inputs["input_ids"][0])
+    # Strip trailing EOS so the assistant turn flows into <Image>
+    while prompt_ids and prompt_ids[-1] == eos:
+        prompt_ids.pop()
+
+    latent_ph = cfg.flow_latent_placeholder_id  # upstream latent_token_id
+    image_start = cfg.image_start_token_id
+    image_end = cfg.image_end_token_id
+
+    seq = (
+        prompt_ids
+        + [image_start]
+        + [latent_ph] * n_latent_tokens
+        + [image_end]
+        + [eos]
+    )
+    input_ids = torch.tensor(seq, dtype=torch.long, device=device).unsqueeze(0)  # (1, T)
+    seq_len = input_ids.shape[1]
+    latent_start = len(prompt_ids) + 1  # right after <Image>
+    latent_end = latent_start + n_latent_tokens
+
+    modality_mask = torch.zeros(1, seq_len, dtype=torch.long, device=device)
+    modality_mask[0, latent_start:latent_end] = 2  # gen-latent route
+    flow_positions = torch.tensor([[latent_start, latent_end]], dtype=torch.int32, device=device)
+    g_seqlens = flow_positions.clone()
+    # Single sample → packed degenerate: still feed cu_seqlens/sample_ids so the
+    # same attention path runs as during training.
+    cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32, device=device)
+    sample_ids = torch.zeros(1, seq_len, dtype=torch.int32, device=device)
+    position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0)
+    return {
+        "input_ids": input_ids, "T": seq_len,
+        "latent_start": latent_start, "latent_end": latent_end,
+        "modality_mask": modality_mask, "g_seqlens": g_seqlens,
+        "cu_seqlens": cu_seqlens, "sample_ids": sample_ids,
+        "position_ids": position_ids,
+    }
+
+
+@torch.no_grad()
+def _forward_v(model, inner, base_embeds, flow_embed, meta, dtype):
+    """Run one forward with `flow_embed` injected at the latent span; return velocity."""
+    inputs_embeds = base_embeds.clone()
+    inputs_embeds[0, meta["latent_start"]:meta["latent_end"]] = flow_embed
+    out = inner(
+        input_ids=None,
+        inputs_embeds=inputs_embeds,
+        attention_mask=None,
+        position_ids=meta["position_ids"],
+        cu_seqlens=meta["cu_seqlens"],
+        sample_ids=meta["sample_ids"],
+        modality_mask=meta["modality_mask"],
+        input_image_mask=torch.zeros(1, meta["T"], dtype=torch.bool, device=inputs_embeds.device),
+        flow_embeds=None,  # we already pre-built inputs_embeds
+        flow_positions=None,
+        g_seqlens=meta["g_seqlens"],
+    )
+    # FM velocity target during training is v = noise - x_0
+    return model.llm2vae(out.hidden_states[0, meta["latent_start"]:meta["latent_end"]]).to(dtype)
+
+
+@torch.no_grad()
+def generate_image(
+    model: UnifiedMoTForConditionalGeneration,
+    vae,
+    processor: HunYuanVLMoTProcessor,
+    prompt: str,
+    height: int,
+    width: int,
+    num_steps: int,
+    device,
+    dtype,
+    cfg_scale: float = 1.0,
+):
+    """T2I single-sample sampling (optional CFG). Returns a (3, H, W) tensor in [-1, 1].
+
+    With ``cfg_scale > 1`` each ODE step runs cond + uncond forwards and combines
+      v = v_uncond + cfg_scale * (v_cond - v_uncond)
+    The uncond branch uses an empty prompt, matching text_cond_dropout=0.1 in
+    training where 10% of samples have the caption replaced with "".
+    """
+    cfg: UnifiedMoTConfig = model.config
+    p = model.latent_patch_size
+    downsample = cfg.vae_image_downsample  # pixel → patch-token (e.g. 16 = VAE(8) * patch(2))
+    h_lat = height // downsample
+    w_lat = width // downsample
+    n_latent_tokens = h_lat * w_lat
+
+    do_cfg = (cfg_scale != 1.0)
+    embed_layer = model.get_input_embeddings()
+
+    cond_meta = _build_seq_meta(processor, prompt, cfg, n_latent_tokens, device)
+    cond_base = embed_layer(cond_meta["input_ids"])
+    if do_cfg:
+        uncond_meta = _build_seq_meta(processor, "", cfg, n_latent_tokens, device)
+        uncond_base = embed_layer(uncond_meta["input_ids"])
+
+    # 2D position ids for latent_pos_embed lookup
+    latent_pos_ids = get_2d_position_ids(h_lat, w_lat, cfg.max_latent_size).to(device)
+
+    # Initial noise x_T
+    patch_latent_dim = p * p * cfg.vae_z_channels
+    x = torch.randn(n_latent_tokens, patch_latent_dim, device=device, dtype=dtype)
+
+    # Euler ODE: t from 1.0 → 0.0 in `num_steps`
+    ts = torch.linspace(1.0, 0.0, num_steps + 1, device=device, dtype=dtype)
+    inner = model.model  # UnifiedMoTModel
+
+    for i in range(num_steps):
+        t = ts[i]
+        dt = ts[i] - ts[i + 1]  # positive
+
+        # Build flow_embed for the current x_t (shared by cond & uncond branches)
+        time_emb = model.time_embedder(t.expand(n_latent_tokens)).to(dtype)
+        x_proj = model.vae2llm(x.to(model.vae2llm.weight.dtype)).to(dtype)
+        pos_emb = model.latent_pos_embed(latent_pos_ids).to(dtype)
+        flow_embed = x_proj + time_emb + pos_emb  # (n_latent_tokens, D)
+
+        v_cond = _forward_v(model, inner, cond_base, flow_embed, cond_meta, dtype)
+        if do_cfg:
+            v_uncond = _forward_v(model, inner, uncond_base, flow_embed, uncond_meta, dtype)
+            v = v_uncond + cfg_scale * (v_cond - v_uncond)
+        else:
+            v = v_cond
+        # Euler step: x_{t-dt} = x_t - dt * v
+        x = x - dt * v
+
+    # Unpatchify + VAE decode
+    x_lat = unpatchify_latent(x.float(), h_lat, w_lat, p, cfg.vae_z_channels)  # (C, H_lat, W_lat)
+    vae_dtype = next(vae.parameters()).dtype
+    x_lat = x_lat.unsqueeze(0).to(device=device, dtype=vae_dtype)
+    img = vae.decode(x_lat)
+    if hasattr(img, "sample"):
+        img = img.sample
+    return img.squeeze(0).float().clamp(-1, 1)
+
+
+def main():
+    args = parse_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[args.dtype]
+    torch.manual_seed(args.seed)
+
+    print(f"Loading processor from {args.ckpt}...")
+    processor = HunYuanVLMoTProcessor.from_pretrained(args.ckpt, trust_remote_code=True)
+
+    print(f"Loading model from {args.ckpt}...")
+    model = UnifiedMoTForConditionalGeneration.from_pretrained(args.ckpt, dtype=dtype)
+    
+    maybe_init_generation_path(model, model_load_path=args.ckpt)
+    model.to(device)
+    model.eval()
+
+    print(f"Loading VAE from {args.vae}...")
+    vae, _ = load_ae(args.vae)
+    vae.requires_grad_(False)
+    vae.eval()
+    vae.to(device, dtype=dtype)
+
+    cfg_note = f", CFG {args.cfg_scale}" if args.cfg_scale != 1.0 else " (no CFG)"
+    print(f"Generating: '{args.prompt}' @ {args.width}x{args.height}, {args.num_steps} ODE steps{cfg_note}")
+    img = generate_image(
+        model, vae, processor, args.prompt,
+        height=args.height, width=args.width, num_steps=args.num_steps,
+        device=device, dtype=dtype, cfg_scale=args.cfg_scale,
+    )
+
+    # Save: (C, H, W) in [-1, 1] → uint8 PNG
+    arr = ((img.cpu().permute(1, 2, 0).numpy() + 1.0) * 127.5).clip(0, 255).astype("uint8")
+    Image.fromarray(arr).save(args.out)
+    print(f"Saved → {args.out}")
+
+
+if __name__ == "__main__":
+    main()

vae_model/__init__.py

@@ -0,0 +1 @@
+from .autoencoder import AutoEncoder

vae_model/autoencoder.py

@@ -0,0 +1,360 @@
+# Copyright (c) 2024 Black Forest Labs.
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
+#
+# Original file was released under Apache-2.0, with the full license text
+# available at https://github.com/black-forest-labs/flux/blob/main/LICENSE.
+#
+# This modified file is released under the same license.
+
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+from safetensors.torch import load_file as load_sft
+
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    downsample: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        # no asymmetric padding in torch conv, must do it ourselves
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        # downsampling
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        ch: int,
+        out_ch: int,
+        ch_mult: list[int],
+        num_res_blocks: int,
+        in_channels: int,
+        resolution: int,
+        z_channels: int,
+    ):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+
+        # z to block_in
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean + std * torch.randn_like(mean)
+        else:
+            return mean
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams):
+        super().__init__()
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian()
+
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+    if len(missing) > 0 and len(unexpected) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+        print("\n" + "-" * 79 + "\n")
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+    elif len(missing) > 0:
+        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+    elif len(unexpected) > 0:
+        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def load_ae(local_path: str) -> AutoEncoder:
+    ae_params = AutoEncoderParams(
+            resolution=256,
+            in_channels=3,
+            downsample=8,
+            ch=128,
+            out_ch=3,
+            ch_mult=[1, 2, 4, 4],
+            num_res_blocks=2,
+            z_channels=16,
+            scale_factor=0.3611,
+            shift_factor=0.1159,
+    )
+
+    # Loading the autoencoder
+    ae = AutoEncoder(ae_params)
+
+    if local_path is not None:
+        sd = load_sft(local_path)
+        missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
+        print_load_warning(missing, unexpected)
+    return ae, ae_params

vqa_inference.py

@@ -0,0 +1,113 @@
+"""Minimal VQA / text-understanding inference for RxBrain (UnifiedMoT).
+
+Pure autoregressive text decoding: (image(s) + question) -> answer text.
+No VAE / flow-matching needed. Reuses the SAME prompt construction
+(`build_conditioned_sequence`) and KV-cached decode plumbing the interleaved
+planner uses for its text step, but decodes all the way to EOS instead of
+stopping at the <Image> modality-transition token.
+"""
+from __future__ import annotations
+
+import argparse
+from typing import List
+
+import torch
+from transformers.models.hunyuan_vl_mot import HunYuanVLMoTProcessor
+from transformers.cache_utils import DynamicCache
+
+from model import UnifiedMoTForConditionalGeneration, maybe_init_generation_path
+from interleave_inference import build_conditioned_sequence, INPUT_IMAGE_PLACEHOLDER_IDS
+
+
+@torch.no_grad()
+def answer(model, processor, image_paths: List[str], question: str,
+           device, dtype, max_new_tokens: int = 256) -> str:
+    cfg = model.config
+    eos = cfg.eos_token_id
+    inner = model.model
+
+    # No generation task-instruction: this is understanding, not image gen.
+    prompt_ids, proc = build_conditioned_sequence(processor, image_paths, question)
+
+    pixel_values = proc.get("pixel_values")
+    image_grid_thw = proc.get("image_grid_thw")
+    if pixel_values is not None:
+        pixel_values = pixel_values.to(device=device, dtype=dtype)
+    if image_grid_thw is not None:
+        image_grid_thw = image_grid_thw.to(device=device)
+
+    def _dmask(ids):
+        n = len(ids)
+        inp = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+        mod = torch.zeros(1, n, dtype=torch.long, device=device)
+        iim = torch.zeros(1, n, dtype=torch.bool, device=device)
+        for pid in INPUT_IMAGE_PLACEHOLDER_IDS:
+            m = inp[0] == pid
+            mod[0, m] = 1
+            iim[0, m] = True
+        return inp, mod, iim
+
+    _empty_g = torch.zeros((0, 2), dtype=torch.int32, device=device)
+    pkv = DynamicCache()
+    prompt_len = len(prompt_ids)
+    inp, mod, iim = _dmask(list(prompt_ids))
+    out = inner(input_ids=inp, position_ids=torch.arange(prompt_len, device=device).unsqueeze(0),
+                pixel_values=pixel_values, image_grid_thw=image_grid_thw,
+                cu_seqlens=torch.tensor([0, prompt_len], dtype=torch.int32, device=device),
+                sample_ids=torch.zeros(1, prompt_len, dtype=torch.int32, device=device),
+                modality_mask=mod, input_image_mask=iim,
+                flow_embeds=None, flow_positions=None, g_seqlens=_empty_g,
+                use_cache=True, past_key_values=pkv,
+                cache_position=torch.arange(prompt_len, device=device))
+    pkv = out.past_key_values
+    nxt = int(out.logits[0, -1].argmax().item())
+    cur = prompt_len
+    text_ids = []
+    for _ in range(max_new_tokens):
+        if nxt == eos:
+            break
+        text_ids.append(nxt)
+        out = inner(input_ids=torch.tensor([[nxt]], dtype=torch.long, device=device),
+                    position_ids=torch.tensor([[cur]], device=device),
+                    pixel_values=None, image_grid_thw=None,
+                    cu_seqlens=torch.tensor([0, cur + 1], dtype=torch.int32, device=device),
+                    sample_ids=torch.zeros(1, 1, dtype=torch.int32, device=device),
+                    modality_mask=torch.zeros(1, 1, dtype=torch.long, device=device),
+                    input_image_mask=torch.zeros(1, 1, dtype=torch.bool, device=device),
+                    flow_embeds=None, flow_positions=None, g_seqlens=_empty_g,
+                    use_cache=True, past_key_values=pkv,
+                    cache_position=torch.tensor([cur], device=device))
+        pkv = out.past_key_values
+        nxt = int(out.logits[0, -1].argmax().item())
+        cur += 1
+    return processor.tokenizer.decode(text_ids, skip_special_tokens=True)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", required=True)
+    ap.add_argument("--images", nargs="+", required=True)
+    ap.add_argument("--question", required=True)
+    ap.add_argument("--max_new_tokens", type=int, default=256)
+    ap.add_argument("--dtype", choices=["bfloat16", "float16", "float32"], default="bfloat16")
+    args = ap.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = getattr(torch, args.dtype)
+
+    print(f"Loading processor from {args.ckpt}...")
+    processor = HunYuanVLMoTProcessor.from_pretrained(args.ckpt, trust_remote_code=True)
+    print(f"Loading model from {args.ckpt}...")
+    model = UnifiedMoTForConditionalGeneration.from_pretrained(args.ckpt, dtype=dtype)
+    maybe_init_generation_path(model, model_load_path=args.ckpt)
+    model.to(device)
+    model.eval()
+
+    print(f"\nIMAGES: {args.images}")
+    print(f"Q: {args.question}")
+    ans = answer(model, processor, args.images, args.question, device, dtype, args.max_new_tokens)
+    print(f"A: {ans}")
+
+
+if __name__ == "__main__":
+    main()