Initial Space demo

.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/cats.jpg filter=lfs diff=lfs merge=lfs -text
+examples/street.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,13 +1,19 @@
 ---
 title: PixelUMM
-emoji: 🌖
-colorFrom: gray
-colorTo: blue
+emoji: 🎨
+colorFrom: green
+colorTo: gray
 sdk: gradio
-sdk_version: 6.29.1
-python_version: '3.13'
+sdk_version: "5.50.0"
 app_file: app.py
 pinned: false
+license: other
+short_description: Encoder-free unified image/video understanding + generation
+models:
+  - nvidia/PixelUMM
+  - Qwen/Qwen3-8B
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Gradio ZeroGPU demo for [nvidia/PixelUMM](https://huggingface.co/nvidia/PixelUMM): text-to-image, text-to-video, image understanding and video understanding.
+
+Model code vendored from [nv-tlabs/PixelUMM](https://github.com/nv-tlabs/PixelUMM) (Apache-2.0; `modeling/pixelumm/modeling_utils.py` keeps its CC BY-NC 4.0 notice), with FlashAttention varlen replaced by a PyTorch SDPA fallback. Model weights are under the NVIDIA One-Way Noncommercial License.

app.py

@@ -0,0 +1,431 @@
+import os
+import random
+import tempfile
+import time
+from pathlib import Path
+
+import spaces  # must be imported before torch on ZeroGPU
+import gradio as gr
+import numpy as np
+import torch
+import torchvision.transforms.functional as transforms_f
+from huggingface_hub import snapshot_download
+from PIL import Image
+
+from data.pixelumm_smart_resize import smart_resize_pixels
+from data.pixelumm_video_decoder import decode_pixelumm_video
+from eval.vlm.release_preprocess_contract import IMAGE_VLM_CONTRACT, VIDEO_VLM_CONTRACT
+from train.eval_utils import (
+    _generate_t2i_one,
+    _generate_t2v_one,
+    _generate_v2t_one,
+    generate_vlm_interleaved,
+    save_video_mp4,
+)
+from train.model_factory import build_pixelumm_model
+from train.release_checkpoint import (
+    DEFAULT_RELEASE_CONFIG,
+    load_checkpoint_weights,
+    parse_release_arguments,
+)
+
+# Inference code adapted from https://github.com/nv-tlabs/PixelUMM (inference.py).
+
+MODEL_ID = "nvidia/PixelUMM"
+MODEL_REVISION = "b621a91463d8a438f64c5d0116f42fb13cf26ce7"
+CHECKPOINT = "S8-F22-R05"  # default checkpoint; supports all four tasks
+QWEN_ID = "Qwen/Qwen3-8B"
+QWEN_REVISION = "b968826d9c46dd6066d109eabc6255188de91218"
+MAX_SEED = 2**31 - 1
+NEGATIVE_PROMPT = (
+    Path(__file__).parent / "experiments" / "s8_f22_r07" / "t2v_negative_prompt.txt"
+).read_text(encoding="utf-8").strip()
+
+# ---------------------------------------------------------------------------
+# Model loading (module level, as required by ZeroGPU)
+# ---------------------------------------------------------------------------
+model_root = snapshot_download(
+    MODEL_ID,
+    revision=MODEL_REVISION,
+    allow_patterns=[f"{CHECKPOINT}/__SAVE_COMPLETE", f"{CHECKPOINT}/model/*"],
+)
+qwen_dir = snapshot_download(
+    QWEN_ID,
+    revision=QWEN_REVISION,
+    allow_patterns=["config.json", "tokenizer.json", "tokenizer_config.json", "vocab.json", "merges.txt"],
+)
+
+model_args, data_args, training_args = parse_release_arguments(
+    DEFAULT_RELEASE_CONFIG, extra_args=["--llm_path", qwen_dir]
+)
+_prev_dtype = torch.get_default_dtype()
+torch.set_default_dtype(torch.bfloat16)
+try:
+    import logging
+
+    logger = logging.getLogger("pixelumm")
+    logger.addHandler(logging.StreamHandler())
+    logger.setLevel(logging.INFO)
+    model, tokenizer, new_token_ids = build_pixelumm_model(
+        model_args, data_args, training_args, logger
+    )
+finally:
+    torch.set_default_dtype(_prev_dtype)
+model = load_checkpoint_weights(os.path.join(model_root, CHECKPOINT), model, logger=logger)
+model = model.to(dtype=torch.bfloat16).eval()
+for p in model.parameters():
+    p.requires_grad = False
+model.to("cuda")
+PATCH = int(model_args.pixel_token_patch_size)
+print("PixelUMM loaded")
+
+
+def _seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def _resolve_seed(seed, randomize):
+    return random.randint(0, MAX_SEED) if randomize else int(seed)
+
+
+def _snap(v: int) -> int:
+    return max(PATCH, int(round(int(v) / PATCH)) * PATCH)
+
+
+# ---------------------------------------------------------------------------
+# Text to image
+# ---------------------------------------------------------------------------
+def _t2i_duration(prompt, height, width, steps, cfg, shift, seed):
+    return int(30 + steps * 0.6 * max(1.0, (height * width) / (512 * 512)))
+
+
+@spaces.GPU(duration=_t2i_duration)
+def _t2i_gpu(prompt, height, width, steps, cfg, shift, seed):
+    _seed_everything(seed)
+    device = torch.device("cuda")
+    with torch.inference_mode():
+        return _generate_t2i_one(
+            model, tokenizer, new_token_ids, prompt, (height, width), device,
+            steps, shift, "dpm-solver", cfg, 1.0, (0.0, 1.0), "none",
+        )
+
+
+def text_to_image(prompt, height, width, steps, cfg, shift, seed, randomize_seed):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    seed = _resolve_seed(seed, randomize_seed)
+    t0 = time.time()
+    image = _t2i_gpu(prompt.strip(), _snap(height), _snap(width), int(steps), float(cfg), float(shift), seed)
+    print(f"t2i {height}x{width} steps={steps} took {time.time() - t0:.1f}s")
+    return image, seed
+
+
+# ---------------------------------------------------------------------------
+# Text to video
+# ---------------------------------------------------------------------------
+VIDEO_SIZES = {
+    "320x176 (default)": (176, 320),
+    "176x320 (portrait)": (320, 176),
+    "480x272": (272, 480),
+    "640x352": (352, 640),
+    "832x464 (480p)": (464, 832),
+}
+XLARGE_TOKEN_THRESHOLD = 12_000
+
+
+def _video_tokens(height, width, frames):
+    return (frames // 4) * (height // PATCH) * (width // PATCH)
+
+
+def _t2v_duration(prompt, negative, height, width, frames, steps, cfg, shift, seed):
+    tokens = _video_tokens(height, width, frames)
+    passes = steps * (2 if cfg > 1.0 else 1)
+    est = passes * (tokens / 5000.0) * 0.35 * (1 + tokens / 40000.0)
+    return int(min(300, 40 + est))
+
+
+def _t2v_impl(prompt, negative, height, width, frames, steps, cfg, shift, seed):
+    _seed_everything(seed)
+    device = torch.device("cuda")
+    with torch.inference_mode():
+        return _generate_t2v_one(
+            model, tokenizer, new_token_ids, prompt, (frames, height, width), device,
+            steps, shift, "unipc", cfg, 1.0, (0.0, 1.0), "none",
+            video_fps=24.0, negative_prompt=negative,
+        )
+
+
+@spaces.GPU(duration=_t2v_duration)
+def _t2v_gpu(*args):
+    return _t2v_impl(*args)
+
+
+@spaces.GPU(duration=_t2v_duration, size="xlarge")
+def _t2v_gpu_xlarge(*args):
+    return _t2v_impl(*args)
+
+
+def text_to_video(prompt, negative, size, frames, steps, cfg, shift, seed, randomize_seed):
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+    height, width = VIDEO_SIZES[size]
+    frames = max(4, int(frames) // 4 * 4)
+    seed = _resolve_seed(seed, randomize_seed)
+    tokens = _video_tokens(height, width, frames)
+    fn = _t2v_gpu_xlarge if tokens > XLARGE_TOKEN_THRESHOLD else _t2v_gpu
+    t0 = time.time()
+    frames_out = fn(
+        prompt.strip(), (negative or "").strip(), height, width, frames,
+        int(steps), float(cfg), float(shift), seed,
+    )
+    print(f"t2v {width}x{height}x{frames} tokens={tokens} steps={steps} took {time.time() - t0:.1f}s")
+    out_dir = tempfile.mkdtemp()
+    path = save_video_mp4(frames_out, out_dir, fps=24.0, filename="pixelumm.mp4")
+    if path is None:
+        raise gr.Error("Generation produced no frames.")
+    return path, seed
+
+
+# ---------------------------------------------------------------------------
+# Image understanding
+# ---------------------------------------------------------------------------
+def _resize_image_vlm(image: Image.Image) -> Image.Image:
+    image = image.convert("RGB")
+    h, w = smart_resize_pixels(
+        image.height,
+        image.width,
+        factor=PATCH,
+        min_pixels=IMAGE_VLM_CONTRACT.min_image_pixels,
+        max_pixels=IMAGE_VLM_CONTRACT.max_image_tokens * PATCH**2,
+    )
+    return transforms_f.resize(
+        image, [h, w], interpolation=transforms_f.InterpolationMode.BICUBIC, antialias=True
+    )
+
+
+@spaces.GPU(duration=90)
+def _image_vlm_gpu(image, prompt, max_new_tokens, temperature):
+    device = torch.device("cuda")
+    with torch.inference_mode():
+        text = generate_vlm_interleaved(
+            model=model,
+            tokenizer=tokenizer,
+            new_token_ids=new_token_ids,
+            images=[image],
+            instruction=prompt,
+            device=device,
+            max_length=max_new_tokens,
+            temperature=temperature,
+            fixed_length=False,
+            reasoning_mode="standard_bare",
+            exact_max_new_tokens=True,
+        )
+    return text[0] if isinstance(text, tuple) else text
+
+
+def image_understanding(image, prompt, max_new_tokens, temperature):
+    if image is None:
+        raise gr.Error("Please upload an image.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a question.")
+    image = _resize_image_vlm(image)
+    return _image_vlm_gpu(image, prompt.strip(), int(max_new_tokens), float(temperature)).strip()
+
+
+# ---------------------------------------------------------------------------
+# Video understanding
+# ---------------------------------------------------------------------------
+def _decode_video_vlm(path: str):
+    video_path = Path(path)
+    c = VIDEO_VLM_CONTRACT
+    decoded = decode_pixelumm_video(
+        key=video_path.name,
+        data=video_path.read_bytes(),
+        min_fps=1.0,
+        max_fps=240.0,
+        target_fps=c.target_fps,
+        dense_target_fps=c.dense_target_fps,
+        sparse_target_fps=c.sparse_target_fps,
+        max_sampled_frames=c.max_sampled_frames,
+        sparse_max_sampled_frames=c.sparse_max_sampled_frames,
+        dense_representation_probability=c.dense_representation_probability,
+        max_raw_patch_tokens=c.max_raw_patch_tokens,
+        patch_size=PATCH,
+        resize_multiple=PATCH,
+        min_frame_pixels=c.min_frame_pixels,
+        max_frame_pixels=c.max_frame_pixels,
+        temporal_patch_size=c.temporal_patch_size,
+        drop_incomplete_dense_tube=c.drop_incomplete_dense_tube,
+        representation_selector_key=str(video_path),
+    )
+    if decoded is None:
+        raise gr.Error("Could not decode this video (unsupported format, frame rate or length).")
+    videos = decoded["videos"]
+    frames = [
+        Image.fromarray(videos[:, i].permute(1, 2, 0).cpu().numpy()).convert("RGB")
+        for i in range(int(videos.shape[1]))
+    ]
+    source_fps = float(decoded["source_fps"])
+    timestamps = [float(i) / source_fps for i in decoded["frame_indices"]]
+    return frames, timestamps, int(decoded["temporal_patch_size"])
+
+
+@spaces.GPU(duration=120)
+def _video_vlm_gpu(frames, timestamps, temporal_patch_size, prompt, max_new_tokens, temperature):
+    device = torch.device("cuda")
+    with torch.inference_mode():
+        return _generate_v2t_one(
+            model, tokenizer, new_token_ids, frames, timestamps, prompt, device,
+            max_new_tokens + 1, temperature,
+            temporal_patch_size=temporal_patch_size,
+            timestamp_mode=VIDEO_VLM_CONTRACT.timestamp_mode,
+        )
+
+
+def video_understanding(video, prompt, max_new_tokens, temperature):
+    if not video:
+        raise gr.Error("Please upload a video.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a question.")
+    frames, timestamps, tps = _decode_video_vlm(video)
+    return _video_vlm_gpu(
+        frames, timestamps, tps, prompt.strip(), int(max_new_tokens), float(temperature)
+    ).strip()
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+EX_DIR = Path(__file__).parent / "examples"
+
+with gr.Blocks(title="PixelUMM") as demo:
+    gr.Markdown("# PixelUMM")
+    gr.Markdown(
+        f"[Model on Hugging Face](https://huggingface.co/{MODEL_ID}) · "
+        "[Code](https://github.com/nv-tlabs/PixelUMM) · "
+        "[Paper](https://arxiv.org/abs/2609.38597)\n\n"
+        "Encoder-free unified multimodal model from NVIDIA: one decoder-only Transformer "
+        "(Qwen3-8B backbone, 15.2B params) that reads and writes raw pixels, with no VAE and no vision encoder. "
+        f"This demo runs the default `{CHECKPOINT}` checkpoint. "
+        "Weights are released under the NVIDIA One-Way Noncommercial License (research/evaluation only)."
+    )
+
+    with gr.Tabs():
+        with gr.Tab("Text to Image"):
+            with gr.Row():
+                with gr.Column():
+                    t2i_prompt = gr.Textbox(label="Prompt", lines=3, placeholder="A red panda reading beside a window")
+                    t2i_btn = gr.Button("Generate", variant="primary")
+                    with gr.Accordion("Advanced settings", open=False):
+                        with gr.Row():
+                            t2i_w = gr.Slider(256, 768, value=512, step=16, label="Width")
+                            t2i_h = gr.Slider(256, 768, value=512, step=16, label="Height")
+                        t2i_steps = gr.Slider(10, 100, value=50, step=1, label="Steps (DPM-Solver++)")
+                        t2i_cfg = gr.Slider(1.0, 10.0, value=3.5, step=0.1, label="CFG scale")
+                        t2i_shift = gr.Slider(1.0, 10.0, value=3.0, step=0.1, label="Timestep shift")
+                        with gr.Row():
+                            t2i_seed = gr.Number(value=4396, label="Seed", precision=0)
+                            t2i_rand = gr.Checkbox(value=True, label="Randomize seed")
+                with gr.Column():
+                    t2i_out = gr.Image(label="Image", type="pil", format="png")
+                    t2i_seed_out = gr.Number(label="Seed used", interactive=False)
+            gr.Examples(
+                examples=[
+                    ["A red panda reading beside a window"],
+                    ["A cozy cabin in a snowy pine forest at dusk, warm light glowing from the windows"],
+                    ["A bowl of ramen with a soft-boiled egg, studio food photography"],
+                    ["An astronaut riding a horse on the surface of Mars, cinematic lighting"],
+                ],
+                inputs=[t2i_prompt],
+            )
+            t2i_btn.click(
+                text_to_image,
+                inputs=[t2i_prompt, t2i_h, t2i_w, t2i_steps, t2i_cfg, t2i_shift, t2i_seed, t2i_rand],
+                outputs=[t2i_out, t2i_seed_out],
+            )
+
+        with gr.Tab("Text to Video"):
+            with gr.Row():
+                with gr.Column():
+                    t2v_prompt = gr.Textbox(
+                        label="Prompt",
+                        lines=3,
+                        placeholder="A golden retriever runs across a grassy field in bright daylight.",
+                    )
+                    t2v_btn = gr.Button("Generate", variant="primary")
+                    t2v_size = gr.Dropdown(list(VIDEO_SIZES), value="320x176 (default)", label="Resolution (W x H)")
+                    t2v_frames = gr.Slider(16, 96, value=96, step=4, label="Frames (24 fps; 96 = 4 s)")
+                    with gr.Accordion("Advanced settings", open=False):
+                        t2v_neg = gr.Textbox(label="Negative prompt", value=NEGATIVE_PROMPT, lines=3)
+                        t2v_steps = gr.Slider(10, 50, value=35, step=1, label="Steps (UniPC)")
+                        t2v_cfg = gr.Slider(1.0, 10.0, value=6.0, step=0.1, label="CFG scale")
+                        t2v_shift = gr.Slider(1.0, 15.0, value=10.0, step=0.5, label="Timestep shift")
+                        with gr.Row():
+                            t2v_seed = gr.Number(value=4396, label="Seed", precision=0)
+                            t2v_rand = gr.Checkbox(value=True, label="Randomize seed")
+                    gr.Markdown(
+                        "Larger resolutions with many frames run on a full-size GPU and use more quota. "
+                        "The official Cosmos guardrail checks are not run in this demo."
+                    )
+                with gr.Column():
+                    t2v_out = gr.Video(label="Video", autoplay=True)
+                    t2v_seed_out = gr.Number(label="Seed used", interactive=False)
+            gr.Examples(
+                examples=[
+                    ["A golden retriever runs across a grassy field in bright daylight. The camera tracks smoothly beside the dog."],
+                    ["Ocean waves crash against black volcanic rocks at sunset, slow motion, spray in the air."],
+                    ["A steaming cup of coffee on a wooden table, morning sunlight, steam slowly rising."],
+                    ["Aerial drone shot flying over a winding river through an autumn forest."],
+                ],
+                inputs=[t2v_prompt],
+            )
+            t2v_btn.click(
+                text_to_video,
+                inputs=[t2v_prompt, t2v_neg, t2v_size, t2v_frames, t2v_steps, t2v_cfg, t2v_shift, t2v_seed, t2v_rand],
+                outputs=[t2v_out, t2v_seed_out],
+            )
+
+        with gr.Tab("Image Understanding"):
+            with gr.Row():
+                with gr.Column():
+                    iu_image = gr.Image(label="Image", type="pil")
+                    iu_prompt = gr.Textbox(label="Question", value="Describe the image.", lines=2)
+                    iu_btn = gr.Button("Ask", variant="primary")
+                    with gr.Accordion("Advanced settings", open=False):
+                        iu_tokens = gr.Slider(16, 1024, value=256, step=16, label="Max new tokens")
+                        iu_temp = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)")
+                with gr.Column():
+                    iu_out = gr.Textbox(label="Answer", lines=12)
+            gr.Examples(
+                examples=[
+                    [str(EX_DIR / "cats.jpg"), "Describe the image in detail."],
+                    [str(EX_DIR / "cats.jpg"), "How many cats are there, and what are they lying on?"],
+                    [str(EX_DIR / "street.jpg"), "What objects are in this room?"],
+                ],
+                inputs=[iu_image, iu_prompt],
+            )
+            iu_btn.click(image_understanding, inputs=[iu_image, iu_prompt, iu_tokens, iu_temp], outputs=iu_out)
+
+        with gr.Tab("Video Understanding"):
+            with gr.Row():
+                with gr.Column():
+                    vu_video = gr.Video(label="Video")
+                    vu_prompt = gr.Textbox(label="Question", value="Describe the video.", lines=2)
+                    vu_btn = gr.Button("Ask", variant="primary")
+                    with gr.Accordion("Advanced settings", open=False):
+                        vu_tokens = gr.Slider(16, 1024, value=256, step=16, label="Max new tokens")
+                        vu_temp = gr.Slider(0.0, 1.5, value=0.0, step=0.05, label="Temperature (0 = greedy)")
+                    gr.Markdown("Frames are sampled at 1 fps (up to 96 frames, at most 448x448 each).")
+                with gr.Column():
+                    vu_out = gr.Textbox(label="Answer", lines=12)
+            vu_btn.click(video_understanding, inputs=[vu_video, vu_prompt, vu_tokens, vu_temp], outputs=vu_out)
+
+    gr.Markdown("[Twitter / X](https://x.com/realmrfakename)")
+
+
+if __name__ == "__main__":
+    demo.queue().launch()

data/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0

data/data_utils.py

@@ -0,0 +1,187 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+
+import math
+
+import torch
+from torch.nn.attention.flex_attention import or_masks, and_masks
+
+
+PIXELUMM_VIDEO_MROPE_BASE_FPS = 24.0
+
+
+def temporal_rope_units_per_second(fps, temporal_compression_factor):
+    """Return physical-time RoPE units per second for a temporal token grid."""
+    fps = float(fps)
+    temporal_compression_factor = float(temporal_compression_factor)
+    if not math.isfinite(fps) or fps <= 0:
+        raise ValueError(f"fps must be finite and positive, got {fps}")
+    if not math.isfinite(temporal_compression_factor) or temporal_compression_factor <= 0:
+        raise ValueError(
+            "temporal_compression_factor must be finite and positive, got "
+            f"{temporal_compression_factor}"
+        )
+    return fps / temporal_compression_factor
+
+
+def create_sparse_mask(document_lens, split_lens, attn_modes, device, total_length=None):
+    def causal_mask(b, h, q_idx, kv_idx):
+        return q_idx >= kv_idx
+
+    def full_and_noise_mask(b, h, q_idx, kv_idx):
+        return (full_and_noise_seq_id[q_idx] == full_and_noise_seq_id[kv_idx]) & (full_and_noise_seq_id[q_idx] >= 0)
+
+    def remove_noise_mask(b, h, q_idx, kv_idx):
+        return (~((noise_seq_id[kv_idx] >= 0) & (noise_seq_id[q_idx] != noise_seq_id[kv_idx])))
+
+    def sample_mask(b, h, q_idx, kv_idx):
+        return (document_id[q_idx] >= 0) & (document_id[q_idx] == document_id[kv_idx])
+
+    full_and_noise_tmp = []
+    noise_tmp = []
+
+    for i, (length, model) in enumerate(zip(split_lens, attn_modes)):
+        value = i if model in ['full', 'noise'] else -1
+        full_and_noise_tmp.extend([value] * length)
+        value_noise = i if model == 'noise' else -1
+        noise_tmp.extend([value_noise] * length)
+
+    logical_length = sum(document_lens)
+    total_length = logical_length if total_length is None else int(total_length)
+    if total_length < logical_length:
+        raise ValueError((total_length, logical_length))
+    trailing_padding = total_length - logical_length
+    full_and_noise_tmp.extend([-1] * trailing_padding)
+    noise_tmp.extend([-1] * trailing_padding)
+    full_and_noise_seq_id = torch.Tensor(full_and_noise_tmp).to(device)
+    noise_seq_id = torch.Tensor(noise_tmp).to(device)
+
+    document_id = torch.cat(
+        [torch.full((l,), i) for i, l in enumerate(document_lens, start=1)]
+        + ([torch.full((trailing_padding,), -1)] if trailing_padding else [])
+    ).to(device)
+
+    return and_masks(or_masks(causal_mask, full_and_noise_mask), remove_noise_mask, sample_mask)
+
+
+def tubeify(video, patch_size, temporal_patch_size):
+    """Patchify a fixed-length video clip into spatio-temporal tube patches.
+
+    Args:
+        video: (T, C, H, W) tensor.
+        patch_size: spatial patch size.
+        temporal_patch_size: number of frames per tube token.
+
+    Returns:
+        (num_tubes, temporal_patch_size * patch_size * patch_size * C)
+    """
+    t, c, h, w = video.shape
+    p = patch_size
+    assert t % temporal_patch_size == 0
+    assert h % p == 0 and w % p == 0
+    num_tubes_t = t // temporal_patch_size
+    video = video.reshape(num_tubes_t, temporal_patch_size, c, h // p, p, w // p, p)
+    # Canonical video token order is THW:
+    # temporal group first, then the full spatial grid for that group.
+    video = torch.einsum("utchpwq->uhwtpqc", video)
+    video = video.reshape(-1, temporal_patch_size * p ** 2 * c)
+    return video
+
+
+def build_pixelumm_vision_mrope_position_ids(
+    temporal_offset,
+    t_tokens,
+    h_tokens,
+    w_tokens,
+    device=None,
+    fps=None,
+    base_fps=PIXELUMM_VIDEO_MROPE_BASE_FPS,
+    temporal_compression_factor=4,
+):
+    """Build PixelUMM VFM unified_3d_mrope ids for one THW-flattened diffusion/pixel grid."""
+    if fps is not None:
+        tps = temporal_rope_units_per_second(fps, temporal_compression_factor)
+        base_tps = temporal_rope_units_per_second(base_fps, temporal_compression_factor)
+        t = torch.arange(t_tokens, dtype=torch.float32, device=device)
+        t = t / tps * base_tps + float(temporal_offset)
+        t = t.view(-1, 1).expand(-1, h_tokens * w_tokens)
+    else:
+        t = torch.arange(t_tokens, dtype=torch.long, device=device).view(-1, 1).expand(-1, h_tokens * w_tokens)
+        t = t + int(temporal_offset)
+    h = torch.arange(h_tokens, dtype=torch.long, device=device).view(1, -1, 1).expand(t_tokens, -1, w_tokens)
+    w = torch.arange(w_tokens, dtype=torch.long, device=device).view(1, 1, -1).expand(t_tokens, h_tokens, -1)
+    if fps is not None:
+        ids = torch.stack([t.flatten(), h.flatten().to(torch.float32), w.flatten().to(torch.float32)], dim=0)
+    else:
+        ids = torch.stack([t.flatten(), h.flatten(), w.flatten()], dim=0)
+    next_temporal_offset = math.ceil(ids.max().item()) + 1 if ids.numel() else int(temporal_offset)
+    return ids, next_temporal_offset
+
+
+def untubeify(video_tubes, image_size, num_frames, patch_size):
+    """Reconstruct video frames from THW-ordered spatio-temporal tube tokens."""
+    H, W = image_size
+    h_patches, w_patches = H // patch_size, W // patch_size
+    tube_frames = video_tubes.shape[-1] // (patch_size * patch_size * 3)
+    if num_frames % tube_frames != 0:
+        raise ValueError(
+            f"Expected num_frames to be divisible by tube_frames={tube_frames}, got {num_frames}"
+        )
+    temporal_groups = num_frames // tube_frames
+    frames = video_tubes.float().cpu().reshape(
+        temporal_groups, h_patches, w_patches, tube_frames, patch_size, patch_size, 3
+    )
+    frames = frames.permute(0, 3, 1, 4, 2, 5, 6).contiguous().reshape(num_frames, H, W, 3)
+    return frames
+
+
+def add_special_tokens(tokenizer):
+    all_special_tokens = []
+    for k, v in tokenizer.special_tokens_map.items():
+        if isinstance(v, str):
+            all_special_tokens.append(v)
+        elif isinstance(v, list):
+            all_special_tokens += v
+
+    new_tokens = []
+
+    if '<|im_start|>' not in all_special_tokens:
+        new_tokens.append('<|im_start|>')
+
+    if '<|im_end|>' not in all_special_tokens:
+        new_tokens.append('<|im_end|>')
+
+    if '<|vision_start|>' not in all_special_tokens:
+        new_tokens.append('<|vision_start|>')
+
+    if '<|vision_end|>' not in all_special_tokens:
+        new_tokens.append('<|vision_end|>')
+
+    num_new_tokens = tokenizer.add_tokens(new_tokens)
+    bos_token_id = tokenizer.convert_tokens_to_ids('<|im_start|>')
+    eos_token_id = tokenizer.convert_tokens_to_ids('<|im_end|>')
+    start_of_image = tokenizer.convert_tokens_to_ids('<|vision_start|>')
+    end_of_image = tokenizer.convert_tokens_to_ids('<|vision_end|>')
+
+    new_token_ids = dict(
+        bos_token_id=bos_token_id, 
+        eos_token_id=eos_token_id, 
+        start_of_image=start_of_image, 
+        end_of_image=end_of_image, 
+    )
+
+    return tokenizer, new_token_ids, num_new_tokens
+
+
+def len2weight(x, loss_reduction='square'):
+    if x == 0:
+        return x
+    if loss_reduction == 'token':
+        return 1
+    if loss_reduction == 'sample':
+        return 1 / x
+    if loss_reduction == 'square':
+        return 1 / (x ** 0.5)
+    raise NotImplementedError(loss_reduction)

data/pixelumm_image_resize.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""PixelUMM-owned image conversion and aspect-preserving resize helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from data.pixelumm_smart_resize import smart_resize_pixels
+
+
+def pil_to_chw_uint8(image: Any) -> torch.Tensor:
+    """Convert a PIL-compatible image to contiguous RGB ``uint8`` CHW."""
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    return (
+        torch.frombuffer(bytearray(image.tobytes()), dtype=torch.uint8)
+        .view(image.height, image.width, 3)
+        .permute(2, 0, 1)
+        .contiguous()
+    )
+
+
+def resize_image_to_exact_hw(
+    image: torch.Tensor,
+    target_hw: tuple[int, int],
+) -> torch.Tensor:
+    """Resize one CHW tensor with the released bicubic-antialias contract."""
+
+    target_h, target_w = (int(target_hw[0]), int(target_hw[1]))
+    if target_h <= 0 or target_w <= 0:
+        raise ValueError(f"target H/W must be positive, got {target_hw!r}")
+    if (int(image.shape[-2]), int(image.shape[-1])) == (target_h, target_w):
+        return image.contiguous()
+    resized = torch.nn.functional.interpolate(
+        image.float().unsqueeze(0),
+        size=(target_h, target_w),
+        mode="bicubic",
+        align_corners=False,
+        antialias=True,
+    ).squeeze(0)
+    if image.dtype == torch.uint8:
+        return resized.round().clamp_(0, 255).to(torch.uint8).contiguous()
+    return resized.to(dtype=image.dtype).contiguous()
+
+
+def resize_image_native_smart_resize(
+    image: torch.Tensor,
+    *,
+    patch_size: int,
+    max_tokens: int | None,
+    min_pixels: int,
+    enabled: bool,
+) -> torch.Tensor:
+    """Preserve native aspect ratio under explicit raw-pixel bounds."""
+
+    if not enabled:
+        return image.contiguous()
+    if max_tokens is None or int(max_tokens) <= 0:
+        raise ValueError("native_smart_resize requires a positive image-token budget")
+    max_pixels = int(max_tokens) * int(patch_size) * int(patch_size)
+    min_pixels = int(min_pixels)
+    if min_pixels <= 0:
+        raise ValueError("native_smart_resize requires positive min_image_pixels")
+    target_hw = smart_resize_pixels(
+        int(image.shape[-2]),
+        int(image.shape[-1]),
+        factor=int(patch_size),
+        min_pixels=min_pixels,
+        max_pixels=max_pixels,
+    )
+    return resize_image_to_exact_hw(image, target_hw)

data/pixelumm_smart_resize.py

@@ -0,0 +1,55 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared aspect-preserving smart-resize geometry for PixelUMM media."""
+
+from __future__ import annotations
+
+import math
+
+
+_MAX_ASPECT_RATIO = 200.0
+
+
+def smart_resize_pixels(
+    height: int,
+    width: int,
+    *,
+    factor: int,
+    min_pixels: int,
+    max_pixels: int,
+) -> tuple[int, int]:
+    """Resize H/W to a factor-aligned grid within explicit pixel bounds."""
+
+    height = int(height)
+    width = int(width)
+    factor = int(factor)
+    min_pixels = int(min_pixels)
+    max_pixels = int(max_pixels)
+    if height <= 0 or width <= 0 or factor <= 0:
+        raise ValueError(
+            f"Invalid resize geometry height={height} width={width} factor={factor}"
+        )
+    if min_pixels <= 0 or max_pixels < min_pixels:
+        raise ValueError(f"Invalid pixel bounds min={min_pixels} max={max_pixels}")
+    aspect_ratio = max(height, width) / min(height, width)
+    if aspect_ratio > _MAX_ASPECT_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be <= {_MAX_ASPECT_RATIO}, got {aspect_ratio}"
+        )
+
+    resized_height = max(factor, round(height / factor) * factor)
+    resized_width = max(factor, round(width / factor) * factor)
+    if resized_height * resized_width > max_pixels:
+        scale = math.sqrt((height * width) / max_pixels)
+        resized_height = max(
+            factor, math.floor(height / scale / factor) * factor
+        )
+        resized_width = max(
+            factor, math.floor(width / scale / factor) * factor
+        )
+    elif resized_height * resized_width < min_pixels:
+        scale = math.sqrt(min_pixels / (height * width))
+        resized_height = math.ceil(height * scale / factor) * factor
+        resized_width = math.ceil(width * scale / factor) * factor
+    return int(resized_height), int(resized_width)

data/pixelumm_video_decoder.py

@@ -0,0 +1,686 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""PixelUMM-native video sampling and smart-resize utilities."""
+
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor, wait
+from dataclasses import dataclass
+import hashlib
+import math
+import re
+from typing import TYPE_CHECKING
+
+import torch
+import torchvision.transforms.functional as transforms_F
+
+if TYPE_CHECKING:
+    from torchcodec.decoders import VideoDecoder
+
+from .pixelumm_smart_resize import smart_resize_pixels
+
+
+_VIDEO_EXTENSIONS = frozenset({"mp4", "avi", "webm", "mov"})
+# Keep the deterministic train-data split stable across the product rename.
+_VIDEO_REPRESENTATION_COMPATIBILITY_SALT = bytes.fromhex(
+    "706978656c626167656c2d766964656f2d726570726573656e746174696f6e2d763100"
+)
+
+
+def _open_video_decoder(*args: object, **kwargs: object) -> "VideoDecoder":
+    """Import TorchCodec only when a caller actually decodes video media."""
+
+    try:
+        from torchcodec.decoders import VideoDecoder
+    except (ImportError, RuntimeError) as error:
+        raise RuntimeError(
+            "Video decoding requires the torchcodec version pinned in "
+            "requirements.txt and system FFmpeg libraries"
+        ) from error
+    return VideoDecoder(*args, **kwargs)
+
+
+@dataclass(frozen=True)
+class PixelUMMVideoLayout:
+    """Deterministic model-input layout selected before decoding frames."""
+
+    frame_indices: tuple[int, ...]
+    conditioning_fps: float
+    resized_height: int
+    resized_width: int
+    raw_patch_tokens: int
+    packed_patch_tokens: int
+
+
+class PixelUMMVideoRejectedError(ValueError):
+    """Expected media rejection with a stable, machine-readable reason."""
+
+    def __init__(self, reason: str, message: str, **details: object) -> None:
+        super().__init__(message)
+        self.reason = str(reason)
+        self.details = dict(details)
+
+
+@dataclass(frozen=True)
+class PixelUMMVideoRepresentation:
+    """Per-video temporal representation selected before frame decoding."""
+
+    mode: str
+    nominal_frame_count: int
+    max_sampled_frames: int
+    temporal_patch_size: int
+    selection_reason: str
+    dense_probability: float
+    selection_draw: float | None
+    target_fps: float
+    dense_target_fps: float
+    sparse_target_fps: float
+    dense_nominal_frame_count: int
+    sparse_nominal_frame_count: int
+    dense_grid_supported: bool
+    sampling_mode: str
+
+
+def stable_video_representation_draw(identity: str) -> float:
+    """Map an immutable sample identity to a reproducible value in ``[0, 1)``."""
+
+    digest = hashlib.blake2b(
+        _VIDEO_REPRESENTATION_COMPATIBILITY_SALT + str(identity).encode("utf-8"),
+        digest_size=8,
+    ).digest()
+    return int.from_bytes(digest, "big") / float(1 << 64)
+
+
+def _decode_video_frames_with_timeout(
+    decoder: VideoDecoder,
+    frame_indices: tuple[int, ...],
+    *,
+    timeout_seconds: int,
+) -> torch.Tensor | None:
+    """Decode an explicit frame list without leaking uncancellable C++ work.
+
+    TorchCodec frame decode is already executing in C++ when ``Future.cancel``
+    returns false.  Returning at the timeout and calling ``shutdown(wait=False)``
+    therefore leaves a ghost decode consuming the DataLoader worker's CPU.  The
+    configured timeout is a slow-operation diagnostic threshold, not a false
+    promise that a Python thread can hard-cancel native decode.
+    """
+
+    executor = ThreadPoolExecutor(max_workers=1)
+    future = executor.submit(decoder.get_frames_at, indices=list(frame_indices))
+    try:
+        _, unfinished = wait((future,), timeout=int(timeout_seconds))
+        if unfinished:
+            wait((future,))
+        batch = future.result()
+    finally:
+        executor.shutdown(wait=True, cancel_futures=True)
+    frames = batch.data
+    if frames.ndim != 4 or int(frames.shape[0]) != len(frame_indices):
+        raise ValueError(
+            f"Unexpected torchcodec output shape {tuple(frames.shape)} for "
+            f"{len(frame_indices)} requested frames"
+        )
+    return frames
+
+
+def select_video_representation(
+    *,
+    total_frames: int,
+    source_fps: float,
+    target_fps: float,
+    dense_target_fps: float | None = None,
+    sparse_target_fps: float | None = None,
+    dense_max_sampled_frames: int,
+    dense_temporal_patch_size: int,
+    sparse_max_sampled_frames: int = 0,
+    dense_representation_probability: float = 1.0,
+    representation_draw: float | None = None,
+    start_frame: int = 0,
+    end_frame: int | None = None,
+) -> PixelUMMVideoRepresentation:
+    """Choose dense tubes or sparse frame-as-image encoding for one clip.
+
+    ``target_fps`` is used when explicit dense/sparse cadences are absent.
+    The release uses 24 FPS for generated 4-frame tubes and 1 FPS for sparse
+    frame-as-image understanding. A dense route is legal only
+    when the selected clip fits the dense cap and the source provides every
+    requested timestamp as a unique frame; low-FPS media is never duplicated.
+    A short clip that cannot provide 4 FPS but can provide 1 FPS falls back to
+    the strict sparse timestamp grid.  Full-clip uniform sampling is used when
+    the nominal dense grid exceeds its frame cap or the source cannot provide
+    even the sparse cadence. A zero sparse cap selects dense-only behavior.
+    """
+
+    dense_max_sampled_frames = int(dense_max_sampled_frames)
+    dense_temporal_patch_size = int(dense_temporal_patch_size)
+    sparse_max_sampled_frames = int(sparse_max_sampled_frames)
+    dense_representation_probability = float(dense_representation_probability)
+    dense_target_fps = float(target_fps if dense_target_fps is None else dense_target_fps)
+    sparse_target_fps = float(target_fps if sparse_target_fps is None else sparse_target_fps)
+    if dense_max_sampled_frames <= 0 or dense_temporal_patch_size <= 0:
+        raise ValueError(
+            "Dense video frame and temporal-patch values must be positive: "
+            f"frames={dense_max_sampled_frames} patch={dense_temporal_patch_size}"
+        )
+    if sparse_max_sampled_frames < 0:
+        raise ValueError(
+            "sparse_max_sampled_frames must be non-negative, got "
+            f"{sparse_max_sampled_frames}"
+        )
+    if not 0.0 <= dense_representation_probability <= 1.0:
+        raise ValueError(
+            "dense_representation_probability must be in [0, 1], got "
+            f"{dense_representation_probability}"
+        )
+    if representation_draw is not None and not 0.0 <= float(representation_draw) < 1.0:
+        raise ValueError(f"representation_draw must be in [0, 1), got {representation_draw}")
+    if dense_target_fps <= 0 or sparse_target_fps <= 0:
+        raise ValueError(
+            "dense_target_fps and sparse_target_fps must be positive: "
+            f"dense={dense_target_fps} sparse={sparse_target_fps}"
+        )
+
+    # total_frames is a strict upper bound on unique sampled source frames, so
+    # this returns the exact uncapped target-FPS grid without a second policy.
+    dense_nominal_indices, _ = sample_frame_indices(
+        total_frames=total_frames,
+        source_fps=source_fps,
+        target_fps=dense_target_fps,
+        max_sampled_frames=int(total_frames),
+        start_frame=start_frame,
+        end_frame=end_frame,
+    )
+    sparse_nominal_indices, _ = sample_frame_indices(
+        total_frames=total_frames,
+        source_fps=source_fps,
+        target_fps=sparse_target_fps,
+        max_sampled_frames=int(total_frames),
+        start_frame=start_frame,
+        end_frame=end_frame,
+    )
+    clip_frames = (int(total_frames) if end_frame is None else int(end_frame)) - int(start_frame)
+    clip_duration = clip_frames / float(source_fps)
+    requested_dense_count = max(1, int(math.ceil(clip_duration * dense_target_fps - 1e-12)))
+    requested_sparse_count = max(1, int(math.ceil(clip_duration * sparse_target_fps - 1e-12)))
+    dense_nominal_frame_count = len(dense_nominal_indices)
+    sparse_nominal_frame_count = len(sparse_nominal_indices)
+    dense_grid_supported = (
+        dense_nominal_frame_count == requested_dense_count
+        and dense_nominal_frame_count >= dense_temporal_patch_size
+    )
+    sparse_grid_supported = sparse_nominal_frame_count == requested_sparse_count
+
+    def sparse(reason: str) -> PixelUMMVideoRepresentation:
+        return PixelUMMVideoRepresentation(
+            mode="sparse_frame_image",
+            nominal_frame_count=sparse_nominal_frame_count,
+            max_sampled_frames=sparse_max_sampled_frames,
+            temporal_patch_size=1,
+            selection_reason=reason,
+            dense_probability=dense_representation_probability,
+            selection_draw=(
+                float(representation_draw) if representation_draw is not None else None
+            ),
+            target_fps=sparse_target_fps,
+            dense_target_fps=dense_target_fps,
+            sparse_target_fps=sparse_target_fps,
+            dense_nominal_frame_count=dense_nominal_frame_count,
+            sparse_nominal_frame_count=sparse_nominal_frame_count,
+            dense_grid_supported=dense_grid_supported,
+            sampling_mode=(
+                "uniform_clip"
+                if reason == "dense_grid_above_cap" or not sparse_grid_supported
+                else "timestamp_grid"
+            ),
+        )
+
+    # Decide the long-clip fallback from the requested dense cadence, not from
+    # the number of unique frames a low-FPS source happens to provide.
+    if sparse_max_sampled_frames > 0 and requested_dense_count > dense_max_sampled_frames:
+        return sparse("dense_grid_above_cap")
+    if sparse_max_sampled_frames > 0 and not dense_grid_supported:
+        return sparse("dense_grid_unavailable")
+    if sparse_max_sampled_frames > 0 and dense_representation_probability < 1.0:
+        if representation_draw is None and dense_representation_probability not in {0.0, 1.0}:
+            raise ValueError(
+                "A deterministic representation_draw is required when short-video "
+                "dense_representation_probability is strictly between 0 and 1"
+            )
+        draw = float(representation_draw or 0.0)
+        if draw >= dense_representation_probability:
+            return sparse("short_probability_sparse")
+    return PixelUMMVideoRepresentation(
+        mode=(
+            f"dense_tube_t{dense_temporal_patch_size}"
+            if dense_temporal_patch_size > 1
+            else "frame_image"
+        ),
+        nominal_frame_count=dense_nominal_frame_count,
+        max_sampled_frames=dense_max_sampled_frames,
+        temporal_patch_size=dense_temporal_patch_size,
+        selection_reason=(
+            "short_probability_dense"
+            if sparse_max_sampled_frames > 0 and dense_representation_probability < 1.0
+            else "dense_default"
+        ),
+        dense_probability=dense_representation_probability,
+        selection_draw=(
+            float(representation_draw) if representation_draw is not None else None
+        ),
+        target_fps=dense_target_fps,
+        dense_target_fps=dense_target_fps,
+        sparse_target_fps=sparse_target_fps,
+        dense_nominal_frame_count=dense_nominal_frame_count,
+        sparse_nominal_frame_count=sparse_nominal_frame_count,
+        dense_grid_supported=dense_grid_supported,
+        sampling_mode="timestamp_grid",
+    )
+
+
+def sample_frame_indices(
+    *,
+    total_frames: int,
+    source_fps: float,
+    target_fps: float,
+    max_sampled_frames: int,
+    start_frame: int = 0,
+    end_frame: int | None = None,
+    align_to_multiple: int = 1,
+    uniform_over_clip: bool = False,
+) -> tuple[tuple[int, ...], float]:
+    """Sample a strict timestamp grid, uniformly thinning only above the cap.
+
+    A nominal 1 FPS contract means timestamps ``0s, 1s, 2s, ...`` rather than
+    choosing ``round(duration)`` frames with an endpoint-inclusive linspace.
+    This preserves the requested cadence for short/medium clips.  Only when
+    the timestamp grid exceeds ``max_sampled_frames`` do we uniformly thin the
+    grid while retaining full-clip coverage. Temporal-tube divisibility is a
+    later padding concern and must never remove valid sampled frames here.
+    """
+
+    total_frames = int(total_frames)
+    source_fps = float(source_fps)
+    target_fps = float(target_fps)
+    max_sampled_frames = int(max_sampled_frames)
+    align_to_multiple = int(align_to_multiple)
+    start_frame = int(start_frame)
+    end_frame = total_frames if end_frame is None else int(end_frame)
+    if total_frames <= 0:
+        raise ValueError(f"total_frames must be positive, got {total_frames}")
+    if not math.isfinite(source_fps) or source_fps <= 0:
+        raise ValueError(f"source_fps must be positive and finite, got {source_fps}")
+    if not math.isfinite(target_fps) or target_fps <= 0:
+        raise ValueError(f"target_fps must be positive and finite, got {target_fps}")
+    if max_sampled_frames <= 0:
+        raise ValueError(f"max_sampled_frames must be positive, got {max_sampled_frames}")
+    if align_to_multiple <= 0:
+        raise ValueError(f"align_to_multiple must be positive, got {align_to_multiple}")
+    if start_frame < 0 or end_frame > total_frames or start_frame >= end_frame:
+        raise ValueError(
+            f"Invalid frame range [{start_frame}, {end_frame}) for total_frames={total_frames}"
+        )
+
+    available_frames = end_frame - start_frame
+    duration = available_frames / source_fps
+    if uniform_over_clip:
+        count = min(available_frames, max_sampled_frames)
+        indices_tensor = torch.linspace(
+            start_frame,
+            end_frame - 1,
+            count,
+            dtype=torch.float64,
+        ).round().to(torch.int64)
+    else:
+        sample_period = 1.0 / target_fps
+        requested_timestamps = torch.arange(
+            0.0,
+            duration,
+            sample_period,
+            dtype=torch.float64,
+        )
+        if requested_timestamps.numel() == 0:
+            requested_timestamps = torch.zeros(1, dtype=torch.float64)
+        # The timestamp grid is mathematically exact at equal source/target
+        # rates, but binary floating point can place values such as 7/24 * 24
+        # infinitesimally below 7.  Flooring those values silently duplicates
+        # the preceding frame (96 @ 24 FPS previously collapsed to 85 unique
+        # indices).  A tiny boundary tolerance preserves the intended integer
+        # grid without moving any real sub-frame timestamp across a boundary.
+        indices_tensor = torch.floor(
+            requested_timestamps * source_fps + 1e-9
+        ).to(torch.int64)
+        indices_tensor = (indices_tensor + start_frame).clamp(
+            min=start_frame,
+            max=end_frame - 1,
+        )
+    indices_tensor = torch.unique_consecutive(indices_tensor)
+    if indices_tensor.numel() > max_sampled_frames:
+        keep = torch.linspace(
+            0,
+            indices_tensor.numel() - 1,
+            max_sampled_frames,
+        ).round().to(torch.int64)
+        indices_tensor = indices_tensor[keep]
+    if align_to_multiple > 1 and indices_tensor.numel() >= align_to_multiple:
+        aligned = (indices_tensor.numel() // align_to_multiple) * align_to_multiple
+        indices_tensor = indices_tensor[:aligned]
+    indices = indices_tensor.tolist()
+    sampled_frames = len(indices)
+    conditioning_fps = sampled_frames / duration
+    return tuple(int(index) for index in indices), float(conditioning_fps)
+
+
+def compute_video_layout(
+    *,
+    total_frames: int,
+    source_fps: float,
+    source_height: int,
+    source_width: int,
+    target_fps: float,
+    max_sampled_frames: int,
+    max_raw_patch_tokens: int,
+    patch_size: int,
+    resize_multiple: int,
+    min_frame_pixels: int,
+    max_frame_pixels: int | None = None,
+    temporal_patch_size: int = 1,
+    start_frame: int = 0,
+    end_frame: int | None = None,
+    align_to_temporal_patch: bool = False,
+    uniform_over_clip: bool = False,
+) -> PixelUMMVideoLayout:
+    """Resolve sampled frames and H/W under a valid-frame raw-patch budget.
+
+    Temporal repeat padding changes the number of packed tube tokens, but must
+    not lower the spatial resolution selected for the valid sampled frames.
+    """
+
+    max_raw_patch_tokens = int(max_raw_patch_tokens)
+    patch_size = int(patch_size)
+    resize_multiple = int(resize_multiple)
+    min_frame_pixels = int(min_frame_pixels)
+    temporal_patch_size = int(temporal_patch_size)
+    if (
+        max_raw_patch_tokens <= 0
+        or patch_size <= 0
+        or resize_multiple <= 0
+        or temporal_patch_size <= 0
+    ):
+        raise ValueError(
+            "PixelUMM video token, patch, resize, and temporal values must be positive: "
+            f"raw_tokens={max_raw_patch_tokens} patch={patch_size} "
+            f"multiple={resize_multiple} "
+            f"temporal_patch_size={temporal_patch_size}"
+        )
+    if resize_multiple % patch_size != 0:
+        raise ValueError(
+            f"resize_multiple={resize_multiple} must be divisible by patch_size={patch_size}"
+        )
+    if min_frame_pixels <= 0:
+        raise ValueError(f"min_frame_pixels must be positive, got {min_frame_pixels}")
+
+    frame_indices, conditioning_fps = sample_frame_indices(
+        total_frames=total_frames,
+        source_fps=source_fps,
+        target_fps=target_fps,
+        max_sampled_frames=max_sampled_frames,
+        start_frame=start_frame,
+        end_frame=end_frame,
+        align_to_multiple=(temporal_patch_size if align_to_temporal_patch else 1),
+        uniform_over_clip=uniform_over_clip,
+    )
+    num_valid_frames = len(frame_indices)
+    temporal_groups = math.ceil(num_valid_frames / temporal_patch_size)
+    per_frame_max_pixels = max_raw_patch_tokens * patch_size**2 // num_valid_frames
+    if max_frame_pixels is not None:
+        if int(max_frame_pixels) <= 0:
+            raise ValueError(
+                f"max_frame_pixels must be positive when set, got {max_frame_pixels}"
+            )
+        per_frame_max_pixels = min(per_frame_max_pixels, int(max_frame_pixels))
+    if per_frame_max_pixels < min_frame_pixels:
+        raise ValueError(
+            "Video frame cap and minimum spatial budget exceed the total PixelUMM budget: "
+            f"frames={num_valid_frames} temporal_groups={temporal_groups} "
+            f"per_frame_max_pixels={per_frame_max_pixels} "
+            f"min_frame_pixels={min_frame_pixels} "
+            f"max_raw_tokens={max_raw_patch_tokens}"
+        )
+    resized_height, resized_width = smart_resize_pixels(
+        int(source_height),
+        int(source_width),
+        factor=resize_multiple,
+        min_pixels=min_frame_pixels,
+        max_pixels=per_frame_max_pixels,
+    )
+    spatial_patch_tokens = (
+        (int(resized_height) // patch_size) * (int(resized_width) // patch_size)
+    )
+    raw_patch_tokens = num_valid_frames * spatial_patch_tokens
+    packed_patch_tokens = temporal_groups * spatial_patch_tokens
+    if raw_patch_tokens > max_raw_patch_tokens:
+        raise RuntimeError(
+            "PixelUMM smart resize exceeded its raw-patch budget: "
+            f"actual={raw_patch_tokens} max={max_raw_patch_tokens} "
+            f"frames={num_valid_frames} temporal_groups={temporal_groups} "
+            f"size={resized_height}x{resized_width}"
+        )
+    if packed_patch_tokens > raw_patch_tokens:
+        raise RuntimeError(
+            "Packed tube tokens cannot exceed valid-frame raw patch tokens: "
+            f"packed={packed_patch_tokens} raw={raw_patch_tokens} "
+            f"frames={num_valid_frames} temporal_patch_size={temporal_patch_size}"
+        )
+    return PixelUMMVideoLayout(
+        frame_indices=frame_indices,
+        conditioning_fps=conditioning_fps,
+        resized_height=int(resized_height),
+        resized_width=int(resized_width),
+        raw_patch_tokens=int(raw_patch_tokens),
+        packed_patch_tokens=int(packed_patch_tokens),
+    )
+
+
+def decode_pixelumm_video(
+    *,
+    key: str,
+    data: bytes,
+    min_fps: float,
+    max_fps: float,
+    target_fps: float,
+    dense_target_fps: float | None = None,
+    sparse_target_fps: float | None = None,
+    max_sampled_frames: int,
+    max_raw_patch_tokens: int,
+    patch_size: int,
+    resize_multiple: int,
+    min_frame_pixels: int,
+    max_frame_pixels: int | None = None,
+    temporal_patch_size: int = 1,
+    sparse_max_sampled_frames: int = 0,
+    dense_representation_probability: float = 1.0,
+    drop_incomplete_dense_tube: bool = False,
+    representation_selector_key: str | None = None,
+    start_seconds: float | None = None,
+    end_seconds: float | None = None,
+    required_sampled_frames: int = 0,
+    num_threads: int = 1,
+    decoding_timeout: int = 60,
+) -> dict[str, object] | None:
+    """Decode a video directly into the PixelUMM V2T model-input layout."""
+
+    extension = re.sub(r".*[.]", "", str(key)).lower()
+    if extension not in _VIDEO_EXTENSIONS:
+        return None
+    video_reader = _open_video_decoder(data, num_ffmpeg_threads=int(num_threads))
+    total_frames = video_reader.metadata.num_frames
+    source_fps = video_reader.metadata.average_fps
+    if total_frames is None or source_fps is None:
+        raise ValueError(
+            "torchcodec missing video metadata "
+            f"(num_frames={total_frames}, average_fps={source_fps})"
+        )
+    total_frames = int(total_frames)
+    source_fps = float(source_fps)
+    required_sampled_frames = int(required_sampled_frames)
+    if required_sampled_frames < 0:
+        raise ValueError(
+            "required_sampled_frames must be non-negative, got "
+            f"{required_sampled_frames}"
+        )
+    if source_fps < float(min_fps) or source_fps > float(max_fps):
+        raise PixelUMMVideoRejectedError(
+            "source_fps",
+            f"Video source_fps={source_fps} is outside "
+            f"[{float(min_fps)}, {float(max_fps)}]",
+            source_fps=source_fps,
+            min_fps=float(min_fps),
+            max_fps=float(max_fps),
+        )
+
+    if (start_seconds is None) != (end_seconds is None):
+        raise ValueError("start_seconds and end_seconds must be provided together")
+    start_frame = 0
+    end_frame = total_frames
+    if start_seconds is not None:
+        start = max(0.0, float(start_seconds))
+        end = min(total_frames / source_fps, float(end_seconds))
+        if not start < end:
+            raise ValueError(
+                f"invalid video clip [{start}, {end}] for {total_frames} "
+                f"frames at {source_fps} FPS"
+            )
+        start_frame = min(total_frames - 1, int(math.floor(start * source_fps)))
+        end_frame = min(
+            total_frames,
+            max(start_frame + 1, int(math.ceil(end * source_fps))),
+        )
+
+    representation_draw = (
+        stable_video_representation_draw(representation_selector_key)
+        if 0.0 < float(dense_representation_probability) < 1.0
+        and representation_selector_key is not None
+        else None
+    )
+    representation = select_video_representation(
+        total_frames=total_frames,
+        source_fps=source_fps,
+        target_fps=target_fps,
+        dense_target_fps=dense_target_fps,
+        sparse_target_fps=sparse_target_fps,
+        dense_max_sampled_frames=max_sampled_frames,
+        dense_temporal_patch_size=temporal_patch_size,
+        sparse_max_sampled_frames=sparse_max_sampled_frames,
+        dense_representation_probability=dense_representation_probability,
+        representation_draw=representation_draw,
+        start_frame=start_frame,
+        end_frame=end_frame,
+    )
+
+    frame_indices, _ = sample_frame_indices(
+        total_frames=total_frames,
+        source_fps=source_fps,
+        target_fps=representation.target_fps,
+        max_sampled_frames=representation.max_sampled_frames,
+        start_frame=start_frame,
+        end_frame=end_frame,
+        align_to_multiple=(
+            representation.temporal_patch_size
+            if drop_incomplete_dense_tube and representation.temporal_patch_size > 1
+            else 1
+        ),
+        uniform_over_clip=representation.sampling_mode == "uniform_clip",
+    )
+
+    # This is the authoritative fixed-window filter.  It uses the same
+    # source metadata, clip bounds, representation policy, timestamp grid,
+    # and temporal alignment as the real decode, but runs before any frame
+    # payload is materialized.  Short dense samples therefore never pay the
+    # expensive get_frames_at/decode cost, while acceptance semantics remain
+    # exactly identical to the former post-decode shape check.
+    sampled_frame_count = len(frame_indices)
+    if (
+        required_sampled_frames
+        and sampled_frame_count != required_sampled_frames
+    ):
+        raise PixelUMMVideoRejectedError(
+            "sampled_frame_count",
+            "PixelUMM fixed-window sample has "
+            f"{sampled_frame_count} sampled frames; required "
+            f"{required_sampled_frames}",
+            sampled_frame_count=sampled_frame_count,
+            required_sampled_frames=required_sampled_frames,
+            total_frames=total_frames,
+            source_fps=source_fps,
+            clip_start_frame=start_frame,
+            clip_end_frame=end_frame,
+        )
+
+    frames = _decode_video_frames_with_timeout(
+        video_reader,
+        frame_indices,
+        timeout_seconds=int(decoding_timeout),
+    )
+    if frames is None:
+        return None
+    source_height = int(frames.shape[-2])
+    source_width = int(frames.shape[-1])
+    layout = compute_video_layout(
+        total_frames=total_frames,
+        source_fps=source_fps,
+        source_height=source_height,
+        source_width=source_width,
+        target_fps=representation.target_fps,
+        max_sampled_frames=representation.max_sampled_frames,
+        max_raw_patch_tokens=max_raw_patch_tokens,
+        patch_size=patch_size,
+        resize_multiple=resize_multiple,
+        min_frame_pixels=min_frame_pixels,
+        max_frame_pixels=max_frame_pixels,
+        temporal_patch_size=representation.temporal_patch_size,
+        start_frame=start_frame,
+        end_frame=end_frame,
+        align_to_temporal_patch=bool(drop_incomplete_dense_tube),
+        uniform_over_clip=representation.sampling_mode == "uniform_clip",
+    )
+    if tuple(frames.shape[-2:]) != (layout.resized_height, layout.resized_width):
+        frames = transforms_F.resize(
+            frames,
+            [layout.resized_height, layout.resized_width],
+            interpolation=transforms_F.InterpolationMode.BICUBIC,
+            antialias=True,
+        )
+    return {
+        "videos": frames.permute(1, 0, 2, 3).contiguous(),
+        "source_fps": source_fps,
+        "conditioning_fps": layout.conditioning_fps,
+        "frame_indices": layout.frame_indices,
+        "original_num_frames": total_frames,
+        "source_height": source_height,
+        "source_width": source_width,
+        "source_duration_seconds": total_frames / source_fps,
+        "clip_start_frame": start_frame,
+        "clip_end_frame": end_frame,
+        "clip_start_seconds": start_frame / source_fps,
+        "clip_end_seconds": end_frame / source_fps,
+        "sampling_mode": representation.sampling_mode,
+        "video_representation": representation.mode,
+        "nominal_target_fps_frame_count": representation.nominal_frame_count,
+        "dense_nominal_frame_count": representation.dense_nominal_frame_count,
+        "sparse_nominal_frame_count": representation.sparse_nominal_frame_count,
+        "dense_grid_supported": representation.dense_grid_supported,
+        "dense_max_sampled_frames": int(max_sampled_frames),
+        "sparse_max_sampled_frames": int(sparse_max_sampled_frames),
+        "dense_representation_probability": representation.dense_probability,
+        "representation_selection_draw": representation.selection_draw,
+        "representation_selection_reason": representation.selection_reason,
+        "temporal_patch_size": representation.temporal_patch_size,
+        "drop_incomplete_dense_tube": bool(drop_incomplete_dense_tube),
+        "target_fps": representation.target_fps,
+        "dense_target_fps": representation.dense_target_fps,
+        "sparse_target_fps": representation.sparse_target_fps,
+        "raw_patch_tokens": layout.raw_patch_tokens,
+        "packed_patch_tokens": layout.packed_patch_tokens,
+    }

eval/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0

eval/vlm/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0

eval/vlm/release_preprocess_contract.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""First-party preprocessing constants for the released R07 VLM suites."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ImagePreprocessContract:
+    resize_mode: str = "native_smart_resize"
+    min_image_pixels: int = 3_136
+    max_image_tokens: int = 2_560_000 // (16**2)
+    contract_name: str = "pixelumm-r07-image-vlm-v1"
+
+    @property
+    def recipe_names(self) -> tuple[str, ...]:
+        """Compatibility label consumed by existing eval receipts."""
+
+        return (self.contract_name,)
+
+
+@dataclass(frozen=True)
+class VideoPreprocessContract:
+    target_fps: float = 1.0
+    dense_target_fps: float = 4.0
+    sparse_target_fps: float = 1.0
+    max_sampled_frames: int = 96
+    sparse_max_sampled_frames: int = 96
+    dense_representation_probability: float = 0.0
+    drop_incomplete_dense_tube: bool = False
+    max_raw_patch_tokens: int = 96 * (448 // 16) ** 2
+    max_frame_pixels: int = 448**2
+    min_frame_pixels: int = 16**2
+    temporal_patch_size: int = 4
+    timestamp_mode: str = "tube_start"
+    mrope_mode: str = "sequence_hw"
+    contract_name: str = "pixelumm-r07-video4-sparse-v1"
+
+
+IMAGE_VLM_CONTRACT = ImagePreprocessContract()
+VIDEO_VLM_CONTRACT = VideoPreprocessContract()

experiments/s8_f22_r07/release.yaml

@@ -0,0 +1,54 @@
+schema: pixelumm_scientific_config_v1
+
+# Model configuration for the default PixelUMM checkpoint.
+model:
+  llm_backend: qwen3
+  qwen_mlp_sequence_chunk_size: 32768
+  qwen3_mrope_type: sensenova
+  qwen3_rope_theta: 1000000.0
+  llm_path: Qwen/Qwen3-8B
+  layer_module: Qwen3MoTDecoderLayer
+  enable_pixel_video: true
+  pixel_token_patch_size: 16
+  pixel_embedder_type: image_raw_patch_linear
+  pixel_separate_und_gen_embedder: true
+  pixel_video_embedder_type: video_raw_tube_linear
+  pixel_video_separate_und_gen_embedder: true
+  pixel_video_temporal_patch_size: 4
+  pixel_image_additive_pos_embed_type: none
+  pixel_video_additive_pos_embed_type: none
+  pixel_head_type: minit2i_linear
+  pixel_video_head_type: jit_style
+  sensenova_vision_rope_theta: 10000.0
+  sensenova_vision_max_position_embeddings: 10000
+  add_timestep_embedding: false
+  add_noise_scale_embedding: false
+  noise_scale: 1.0
+  noise_scale_mode: constant
+  packed_attention_impl: flex
+  inference_attention_backend: flash_attn_varlen
+  packed_expert_routing: pixelumm_mot
+  ce_loss_checkpoint_chunk_size: 4096
+  text_cond_dropout_prob: 0.1
+
+# Only the sequence layout is part of model construction. Dataset identity and
+# distributed data-position semantics are deliberately absent from the release path.
+data:
+  packed_sequence_layout: pixelumm_mot
+
+# Objective fields are retained because they are serialized into PixelUMMConfig
+# and are also reused by the local toy-training example.
+training:
+  gen_space: pixel
+  und_space: pixel
+  prediction_type: x
+  loss_type: v_loss
+  x_pred_t_min: 0.05
+  visual_gen: true
+  visual_und: true
+  freeze_und: false
+  conditional_branch_gradient_contract: pixelumm_zero_v1
+  noise_time_logit_mean: 0.0
+  noise_time_logit_std: 1.0
+  time_schedule: pixelumm_resolution
+  mse_loss_reduction: sample

experiments/s8_f22_r07/t2v_negative_prompt.txt

@@ -0,0 +1 @@
+The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality.

modeling/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+from . import pixelumm, qwen3
+
+__all__ = ["pixelumm", "qwen3"]

modeling/pixelumm/__init__.py

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+
+from modeling.qwen3.configuration_qwen3 import Qwen3Config
+
+from .pixelumm import PixelUMM, PixelUMMConfig
+from .qwen3_navit import Qwen3ForCausalLM
+
+
+__all__ = [
+    "PixelUMMConfig",
+    "PixelUMM",
+    "Qwen3Config",
+    "Qwen3ForCausalLM",
+]

modeling/pixelumm/dpm_solver.py

@@ -0,0 +1,253 @@
+# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Minimal flow-matching DPM-Solver++ used by PixelUMM image/video inference.
+
+This is the exact release path: flow prediction, shifted uniform-flow time
+steps, and second-order multistep DPM-Solver++. Generic VP schedules,
+classifier/PAG guidance, inverse solving, adaptive stepping, and third-order
+solvers are intentionally not part of the public runtime.
+"""
+
+from collections.abc import Callable
+
+import torch
+from tqdm import tqdm
+
+
+def sample_flow_dpm(
+    velocity_fn: Callable[[torch.Tensor, float], torch.Tensor],
+    x: torch.Tensor,
+    *,
+    num_steps: int,
+    shift: float,
+) -> torch.Tensor:
+    """Sample image or video flow; the caller owns conditioning and CFG."""
+    noise_schedule = NoiseScheduleFlow(schedule="discrete_flow")
+
+    def model(x: torch.Tensor, t_input: torch.Tensor) -> torch.Tensor:
+        t_value = t_input.flatten()[0].item() / noise_schedule.total_N
+        return velocity_fn(x, t_value)
+
+    model_fn = model_wrapper(
+        model, noise_schedule, model_type="flow", guidance_type="uncond"
+    )
+    solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+    return solver.sample(
+        x,
+        steps=num_steps,
+        order=2,
+        skip_type="time_uniform_flow",
+        method="multistep",
+        flow_shift=shift,
+        t_start=1.0,
+        t_end=0.001,
+    )
+
+
+class NoiseScheduleFlow:
+    """Rectified-flow schedule used by the released PixelUMM checkpoint."""
+
+    def __init__(self, schedule: str = "discrete_flow") -> None:
+        if schedule != "discrete_flow":
+            raise ValueError(f"PixelUMM requires schedule='discrete_flow', got {schedule!r}")
+        self.T = 1.0
+        self.t0 = 0.001
+        self.schedule = schedule
+        self.total_N = 1000
+
+    @staticmethod
+    def marginal_alpha(t: torch.Tensor) -> torch.Tensor:
+        return 1 - t
+
+    def marginal_log_mean_coeff(self, t: torch.Tensor) -> torch.Tensor:
+        return torch.log(self.marginal_alpha(t))
+
+    @staticmethod
+    def marginal_std(t: torch.Tensor) -> torch.Tensor:
+        return t
+
+    def marginal_lambda(self, t: torch.Tensor) -> torch.Tensor:
+        return self.marginal_log_mean_coeff(t) - torch.log(self.marginal_std(t))
+
+
+def _expand_dims(value: torch.Tensor, dims: int) -> torch.Tensor:
+    return value[(...,) + (None,) * (dims - 1)]
+
+
+def model_wrapper(
+    model: Callable[..., torch.Tensor],
+    noise_schedule: NoiseScheduleFlow,
+    *,
+    model_type: str = "flow",
+    guidance_type: str = "uncond",
+) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
+    """Convert PixelUMM's flow velocity into DPM-Solver noise prediction."""
+
+    if model_type != "flow":
+        raise ValueError(f"PixelUMM DPM inference requires model_type='flow', got {model_type!r}")
+    if guidance_type != "uncond":
+        raise ValueError(f"PixelUMM applies CFG inside the model; got guidance_type={guidance_type!r}")
+    if noise_schedule.schedule != "discrete_flow":
+        raise ValueError(f"Unsupported noise schedule {noise_schedule.schedule!r}")
+
+    def model_fn(x: torch.Tensor, t_continuous: torch.Tensor) -> torch.Tensor:
+        t_input = t_continuous * noise_schedule.total_N
+        velocity = model(x, t_input)
+        if not isinstance(velocity, torch.Tensor):
+            raise TypeError(f"PixelUMM flow model must return a Tensor, got {type(velocity).__name__}")
+        sigma_t = noise_schedule.marginal_std(t_continuous)
+        return (1 - _expand_dims(sigma_t, x.dim()).to(x)) * velocity + x
+
+    return model_fn
+
+
+class DPM_Solver:
+    """Second-order multistep DPM-Solver++ for the PixelUMM flow schedule."""
+
+    def __init__(
+        self,
+        model_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        noise_schedule: NoiseScheduleFlow,
+        algorithm_type: str = "dpmsolver++",
+    ) -> None:
+        if algorithm_type != "dpmsolver++":
+            raise ValueError(f"PixelUMM requires algorithm_type='dpmsolver++', got {algorithm_type!r}")
+        if noise_schedule.schedule != "discrete_flow":
+            raise ValueError(f"Unsupported noise schedule {noise_schedule.schedule!r}")
+        self.model = lambda x, t: model_fn(x, t.expand(x.shape[0]))
+        self.noise_schedule = noise_schedule
+
+    def _data_prediction(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        noise = self.model(x, t)
+        alpha_t = self.noise_schedule.marginal_alpha(t)
+        sigma_t = self.noise_schedule.marginal_std(t)
+        return (x - sigma_t * noise) / alpha_t
+
+    @staticmethod
+    def _time_steps(
+        *,
+        t_start: float,
+        t_end: float,
+        steps: int,
+        device: torch.device,
+        flow_shift: float,
+    ) -> torch.Tensor:
+        betas = torch.linspace(t_start, t_end, steps + 1, device=device)
+        sigmas = 1.0 - betas
+        return (flow_shift * sigmas / (1 + (flow_shift - 1) * sigmas)).flip(dims=[0])
+
+    def _first_order_update(
+        self,
+        x: torch.Tensor,
+        s: torch.Tensor,
+        t: torch.Tensor,
+        model_s: torch.Tensor,
+    ) -> torch.Tensor:
+        schedule = self.noise_schedule
+        h = schedule.marginal_lambda(t) - schedule.marginal_lambda(s)
+        sigma_s = schedule.marginal_std(s)
+        sigma_t = schedule.marginal_std(t)
+        alpha_t = torch.exp(schedule.marginal_log_mean_coeff(t))
+        phi_1 = torch.expm1(-h)
+        return sigma_t / sigma_s * x - alpha_t * phi_1 * model_s
+
+    def _second_order_update(
+        self,
+        x: torch.Tensor,
+        model_prev: list[torch.Tensor],
+        t_prev: list[torch.Tensor],
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        schedule = self.noise_schedule
+        model_prev_1, model_prev_0 = model_prev[-2], model_prev[-1]
+        t_prev_1, t_prev_0 = t_prev[-2], t_prev[-1]
+        lambda_prev_1 = schedule.marginal_lambda(t_prev_1)
+        lambda_prev_0 = schedule.marginal_lambda(t_prev_0)
+        lambda_t = schedule.marginal_lambda(t)
+        sigma_prev_0 = schedule.marginal_std(t_prev_0)
+        sigma_t = schedule.marginal_std(t)
+        alpha_t = torch.exp(schedule.marginal_log_mean_coeff(t))
+
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        d1 = (h / h_0) * (model_prev_0 - model_prev_1)
+        phi_1 = torch.expm1(-h)
+        return (
+            (sigma_t / sigma_prev_0) * x
+            - (alpha_t * phi_1) * model_prev_0
+            - 0.5 * (alpha_t * phi_1) * d1
+        )
+
+    @torch.no_grad()
+    def sample(
+        self,
+        x: torch.Tensor,
+        *,
+        steps: int = 20,
+        t_start: float = 1.0,
+        t_end: float = 0.001,
+        order: int = 2,
+        skip_type: str = "time_uniform_flow",
+        method: str = "multistep",
+        lower_order_final: bool = True,
+        solver_type: str = "dpmsolver",
+        flow_shift: float = 1.0,
+    ) -> torch.Tensor:
+        """Sample with the one solver configuration shipped by PixelUMM."""
+
+        if steps < 2:
+            raise ValueError(f"DPM-Solver++ requires at least two steps, got {steps}")
+        if order != 2:
+            raise ValueError(f"PixelUMM DPM inference requires order=2, got {order}")
+        if skip_type != "time_uniform_flow":
+            raise ValueError(f"PixelUMM requires skip_type='time_uniform_flow', got {skip_type!r}")
+        if method != "multistep":
+            raise ValueError(f"PixelUMM requires method='multistep', got {method!r}")
+        if solver_type != "dpmsolver":
+            raise ValueError(f"PixelUMM requires solver_type='dpmsolver', got {solver_type!r}")
+        if not lower_order_final:
+            raise ValueError("PixelUMM requires lower_order_final=True")
+        if t_start <= 0 or t_end <= 0:
+            raise ValueError("DPM-Solver time range must be positive")
+        if flow_shift <= 0:
+            raise ValueError(f"flow_shift must be positive, got {flow_shift}")
+
+        timesteps = self._time_steps(
+            t_start=t_start,
+            t_end=t_end,
+            steps=steps,
+            device=x.device,
+            flow_shift=flow_shift,
+        )
+        t_prev = [timesteps[0]]
+        model_prev = [self._data_prediction(x, timesteps[0])]
+
+        x = self._first_order_update(x, t_prev[-1], timesteps[1], model_prev[-1])
+        t_prev.append(timesteps[1])
+        model_prev.append(self._data_prediction(x, timesteps[1]))
+
+        for step in tqdm(range(2, steps + 1), desc="DPM-Solver++"):
+            t = timesteps[step]
+            if step == steps:
+                x = self._first_order_update(x, t_prev[-1], t, model_prev[-1])
+            else:
+                x = self._second_order_update(x, model_prev, t_prev, t)
+            t_prev[0], t_prev[1] = t_prev[1], t
+            model_prev[0] = model_prev[1]
+            if step < steps:
+                model_prev[1] = self._data_prediction(x, t)
+        return x

modeling/pixelumm/flow_unipc.py

@@ -0,0 +1,69 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""PixelUMM flow-UniPC inference loop with no PixelUMM runtime dependency."""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+
+from .flow_unipc_scheduler import FlowUniPCScheduler
+
+
+_NUM_TRAIN_TIMESTEPS = 1000
+
+
+@torch.no_grad()
+def sample_flow_unipc(
+    velocity_fn: Callable[[torch.Tensor, float], torch.Tensor],
+    initial_noise: torch.Tensor,
+    *,
+    num_steps: int = 35,
+    shift: float = 1.0,
+) -> torch.Tensor:
+    """Integrate a PixelUMM flow-prediction trajectory with UniPC.
+
+    The public callback receives normalized timesteps in ``[0, 1]``. The
+    local scheduler retains the exact R07 discrete ``[0, 1000]`` convention
+    internally and keeps the comparatively small solver state in FP32.
+    """
+    if not torch.is_floating_point(initial_noise):
+        raise TypeError("UniPC initial_noise must be floating point")
+    if num_steps < 1:
+        raise ValueError(f"UniPC num_steps must be positive, got {num_steps}")
+    if not torch.isfinite(torch.tensor(shift)) or shift <= 0:
+        raise ValueError(f"UniPC shift must be finite and positive, got {shift}")
+
+    solver_state = (
+        initial_noise.float()
+        if initial_noise.dtype in (torch.float16, torch.bfloat16)
+        else initial_noise
+    )
+    scheduler = FlowUniPCScheduler(num_train_timesteps=_NUM_TRAIN_TIMESTEPS)
+    scheduler.set_timesteps(
+        int(num_steps),
+        device=initial_noise.device,
+        shift=float(shift),
+    )
+
+    latent = solver_state
+    for timestep in scheduler.timesteps:
+        normalized_timestep = float(timestep.item()) / _NUM_TRAIN_TIMESTEPS
+        velocity = velocity_fn(latent, normalized_timestep)
+        if velocity.shape != latent.shape:
+            raise ValueError(
+                "UniPC velocity_fn must preserve sample shape: "
+                f"expected {tuple(latent.shape)}, got {tuple(velocity.shape)}"
+            )
+        latent = scheduler.step(
+            model_output=velocity,
+            timestep=timestep,
+            sample=latent.unsqueeze(0),
+        ).squeeze(0)
+
+    return latent
+
+
+__all__ = ["sample_flow_unipc"]

modeling/pixelumm/flow_unipc_scheduler.py

@@ -0,0 +1,249 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Minimal flow-prediction UniPC scheduler used by PixelUMM inference.
+
+Adapted from the Apache-2.0 licensed Wan2.2 flow-UniPC implementation at
+``wan/utils/fm_solvers_unipc.py`` (Wan-Video/Wan2.2 commit
+42bf4cfaa384bc21833865abc2f9e6c0e67233dc), itself adapted from Diffusers
+v0.31.0. This modified file intentionally retains only the fixed execution
+path used by the released PixelUMM R07 inference contract: second-order UniPC,
+flow prediction, B(h) solver type ``bh2``, no dynamic thresholding, and a
+zero terminal sigma.
+
+Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+Licensed under the Apache License, Version 2.0.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+
+
+class FlowUniPCScheduler:
+    """Fixed-contract flow-UniPC scheduler for released PixelUMM models."""
+
+    solver_order = 2
+
+    def __init__(self, *, num_train_timesteps: int = 1000) -> None:
+        if num_train_timesteps < 1:
+            raise ValueError("num_train_timesteps must be positive")
+        self.num_train_timesteps = int(num_train_timesteps)
+
+        alphas = np.linspace(
+            1,
+            1 / self.num_train_timesteps,
+            self.num_train_timesteps,
+        )[::-1].copy()
+        sigmas = 1.0 - alphas
+        self.sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device="cpu")
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        self.model_outputs: list[torch.Tensor | None] = [None] * self.solver_order
+        self.timestep_list: list[torch.Tensor | None] = [None] * self.solver_order
+        self.lower_order_nums = 0
+        self.last_sample: torch.Tensor | None = None
+        self.this_order = 0
+        self.step_index: int | None = None
+        self.num_inference_steps: int | None = None
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int,
+        *,
+        device: torch.device,
+        shift: float,
+    ) -> None:
+        if num_inference_steps < 1:
+            raise ValueError("num_inference_steps must be positive")
+        if not math.isfinite(shift) or shift <= 0:
+            raise ValueError("shift must be finite and positive")
+
+        sigmas = np.linspace(
+            self.sigma_max,
+            self.sigma_min,
+            num_inference_steps + 1,
+        ).copy()[:-1]
+        sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+        timesteps = sigmas * self.num_train_timesteps
+        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
+
+        self.sigmas = torch.from_numpy(sigmas).to("cpu")
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+        self.num_inference_steps = len(timesteps)
+        self.model_outputs = [None] * self.solver_order
+        self.timestep_list = [None] * self.solver_order
+        self.lower_order_nums = 0
+        self.last_sample = None
+        self.this_order = 0
+        self.step_index = None
+
+    @staticmethod
+    def _alpha_sigma(sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return 1 - sigma, sigma
+
+    def _convert_velocity(
+        self,
+        velocity: torch.Tensor,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self.step_index is not None
+        sigma = self.sigmas[self.step_index]
+        return sample - sigma * velocity
+
+    def _predictor_update(
+        self,
+        sample: torch.Tensor,
+        *,
+        order: int,
+    ) -> torch.Tensor:
+        assert self.step_index is not None
+        if order not in (1, 2):
+            raise ValueError(f"unsupported UniPC predictor order: {order}")
+
+        model_output_list = self.model_outputs
+        s0 = self.timestep_list[-1]
+        m0 = model_output_list[-1]
+        if s0 is None or m0 is None:
+            raise RuntimeError("UniPC predictor state is incomplete")
+
+        sigma_t_raw = self.sigmas[self.step_index + 1]
+        sigma_s0_raw = self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._alpha_sigma(sigma_t_raw)
+        alpha_s0, sigma_s0 = self._alpha_sigma(sigma_s0_raw)
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        h = lambda_t - lambda_s0
+
+        hh = -h
+        h_phi_1 = torch.expm1(hh)
+        b_h = torch.expm1(hh)
+
+        correction: torch.Tensor | int = 0
+        if order == 2:
+            previous = model_output_list[-2]
+            if previous is None:
+                raise RuntimeError("UniPC second-order predictor state is incomplete")
+            previous_sigma = self.sigmas[self.step_index - 1]
+            previous_alpha, previous_sigma = self._alpha_sigma(previous_sigma)
+            previous_lambda = torch.log(previous_alpha) - torch.log(previous_sigma)
+            rk = (previous_lambda - lambda_s0) / h
+            d1 = (previous - m0) / rk
+            correction = 0.5 * d1
+
+        predicted = sigma_t / sigma_s0 * sample - alpha_t * h_phi_1 * m0
+        predicted = predicted - alpha_t * b_h * correction
+        return predicted.to(sample.dtype)
+
+    def _corrector_update(
+        self,
+        this_model_output: torch.Tensor,
+        *,
+        last_sample: torch.Tensor,
+        this_sample: torch.Tensor,
+        order: int,
+    ) -> torch.Tensor:
+        assert self.step_index is not None
+        if order not in (1, 2):
+            raise ValueError(f"unsupported UniPC corrector order: {order}")
+
+        model_output_list = self.model_outputs
+        m0 = model_output_list[-1]
+        if m0 is None:
+            raise RuntimeError("UniPC corrector state is incomplete")
+
+        sigma_t_raw = self.sigmas[self.step_index]
+        sigma_s0_raw = self.sigmas[self.step_index - 1]
+        alpha_t, sigma_t = self._alpha_sigma(sigma_t_raw)
+        alpha_s0, sigma_s0 = self._alpha_sigma(sigma_s0_raw)
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        h = lambda_t - lambda_s0
+
+        hh = -h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_2 = h_phi_1 / hh - 1
+        b_h = torch.expm1(hh)
+
+        if order == 1:
+            rhos = torch.tensor([0.5], dtype=last_sample.dtype, device=last_sample.device)
+            correction: torch.Tensor | int = 0
+        else:
+            previous = model_output_list[-2]
+            if previous is None:
+                raise RuntimeError("UniPC second-order corrector state is incomplete")
+            previous_sigma = self.sigmas[self.step_index - 2]
+            previous_alpha, previous_sigma = self._alpha_sigma(previous_sigma)
+            previous_lambda = torch.log(previous_alpha) - torch.log(previous_sigma)
+            rk = (previous_lambda - lambda_s0) / h
+            d1 = (previous - m0) / rk
+
+            rks = torch.tensor([rk, 1.0], device=last_sample.device)
+            matrix = torch.stack([torch.ones_like(rks), rks])
+            h_phi_3 = h_phi_2 / hh - 0.5
+            rhs = torch.tensor(
+                [h_phi_2 / b_h, h_phi_3 * 2 / b_h],
+                device=last_sample.device,
+            )
+            rhos = torch.linalg.solve(matrix, rhs).to(last_sample.dtype)
+            correction = rhos[0] * d1
+
+        d1_t = this_model_output - m0
+        corrected = sigma_t / sigma_s0 * last_sample - alpha_t * h_phi_1 * m0
+        corrected = corrected - alpha_t * b_h * (correction + rhos[-1] * d1_t)
+        return corrected.to(last_sample.dtype)
+
+    def _index_for_timestep(self, timestep: torch.Tensor) -> int:
+        indices = (self.timesteps == timestep.to(self.timesteps.device)).nonzero()
+        if len(indices) == 0:
+            raise ValueError(f"timestep {timestep.item()} is not in the inference schedule")
+        position = 1 if len(indices) > 1 else 0
+        return indices[position].item()
+
+    def step(
+        self,
+        *,
+        model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.num_inference_steps is None:
+            raise RuntimeError("set_timesteps must be called before step")
+        if self.step_index is None:
+            self.step_index = self._index_for_timestep(timestep)
+
+        use_corrector = self.step_index > 0 and self.last_sample is not None
+        converted = self._convert_velocity(model_output, sample)
+        if use_corrector:
+            sample = self._corrector_update(
+                converted,
+                last_sample=self.last_sample,
+                this_sample=sample,
+                order=self.this_order,
+            )
+
+        self.model_outputs[0] = self.model_outputs[1]
+        self.timestep_list[0] = self.timestep_list[1]
+        self.model_outputs[1] = converted
+        self.timestep_list[1] = timestep
+
+        remaining_steps = len(self.timesteps) - self.step_index
+        order = min(self.solver_order, remaining_steps)
+        self.this_order = min(order, self.lower_order_nums + 1)
+        if self.this_order <= 0:
+            raise RuntimeError("invalid UniPC solver order")
+
+        self.last_sample = sample
+        previous = self._predictor_update(sample, order=self.this_order)
+        if self.lower_order_nums < self.solver_order:
+            self.lower_order_nums += 1
+        self.step_index += 1
+        return previous
+
+
+__all__ = ["FlowUniPCScheduler"]

modeling/pixelumm/loss_reduction.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Loss weights used by released PixelUMM training."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def sample_mean_token_weights(
+    group_lens: torch.Tensor | Sequence[int],
+    *,
+    output_size: int,
+    device: torch.device,
+) -> torch.Tensor:
+    """Expand per-group ``1/N`` weights to aligned token weights.
+
+    Summing one group's weighted token losses yields that media item's mean
+    loss. Dividing by the global weight denominator therefore computes an
+    exact global mean over generated media items, even when ranks pack
+    different numbers of examples.
+    """
+    if not torch.is_tensor(group_lens):
+        group_lens = torch.tensor(group_lens, dtype=torch.long, device=device)
+    else:
+        group_lens = group_lens.to(device=device, dtype=torch.long)
+    group_lens = group_lens.reshape(-1)
+    if group_lens.numel() == 0 and output_size != 0:
+        raise ValueError("Non-empty token output requires at least one loss group")
+    if group_lens.device.type == "cpu" and bool((group_lens <= 0).any()):
+        raise ValueError(f"Loss group lengths must be positive, got {group_lens.tolist()}")
+    token_weights = torch.repeat_interleave(
+        torch.reciprocal(group_lens.to(dtype=torch.float32)),
+        group_lens,
+        output_size=int(output_size),
+    )
+    if token_weights.shape[0] != int(output_size):
+        raise ValueError(
+            "Loss group lengths do not match output size: "
+            f"{token_weights.shape[0]} vs {output_size}"
+        )
+    return token_weights
+
+
+def loss_weights_for_reduction(
+    token_losses: torch.Tensor,
+    reduction: str,
+    *,
+    non_token_weights: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Return aligned weights for an explicit loss-reduction contract."""
+    if reduction == "token":
+        return torch.ones(
+            token_losses.numel(),
+            dtype=torch.float32,
+            device=token_losses.device,
+        )
+    if non_token_weights is None:
+        raise ValueError(f"{reduction} reduction requires aligned loss weights")
+    non_token_weights = non_token_weights.reshape(-1).to(
+        device=token_losses.device,
+        dtype=torch.float32,
+    )
+    if non_token_weights.numel() != token_losses.numel():
+        raise ValueError(
+            f"{reduction} weight length mismatch: "
+            f"{non_token_weights.numel()} vs {token_losses.numel()}"
+        )
+    return non_token_weights

modeling/pixelumm/modeling_utils.py

@@ -0,0 +1,165 @@
+# Copyright (c) 2022 Facebook, Inc. and its affiliates.
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: CC BY-NC 4.0
+#
+# This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
+#
+# Original file was released under CC BY-NC 4.0, with the full license text
+# available at https://github.com/facebookresearch/DiT/blob/main/LICENSE.txt.
+#
+# This modified file is released under the same license.
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+
+# --------------------------------------------------------
+# TimestepEmbedder
+# Reference:
+# DiT: https://github.com/facebookresearch/DiT/blob/main/models.py
+# --------------------------------------------------------
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+# ======================== Pixel-space modules ========================
+
+def patchify_image(image, patch_size: int, max_size: int):
+    """Patchify a PIL image into flat pixel patches in [-1, 1].
+
+    Args:
+        image: PIL.Image (any size)
+        patch_size: patch size (e.g. 16)
+        max_size: longest edge resize target
+    Returns:
+        patches: (num_patches, patch_size*patch_size*3)
+        shape: (h_patches, w_patches)
+    """
+    from PIL import Image
+    img = image.convert("RGB")
+    orig_w, orig_h = img.size
+    scale = max_size / max(orig_w, orig_h)
+    new_w = max(patch_size, round(orig_w * scale / patch_size) * patch_size)
+    new_h = max(patch_size, round(orig_h * scale / patch_size) * patch_size)
+    img = img.resize((new_w, new_h), Image.BICUBIC)
+    img_tensor = torch.from_numpy(np.array(img)).float() / 255.0 * 2.0 - 1.0
+    H, W, C = img_tensor.shape
+    h_patches = H // patch_size
+    w_patches = W // patch_size
+    patches = img_tensor.reshape(h_patches, patch_size, w_patches, patch_size, C)
+    patches = patches.permute(0, 2, 1, 3, 4).reshape(-1, patch_size * patch_size * C)
+    return patches, (h_patches, w_patches)
+
+
+
+
+class RawPixelPatchLinearEmbed(nn.Module):
+    """Single-linear image patch embedding: raw pixel patch -> hidden_size."""
+
+    def __init__(self, patch_size=32, in_channels=3, hidden_size=3584):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.patch_dim = patch_size * patch_size * in_channels
+        self.proj = nn.Linear(self.patch_dim, hidden_size, bias=True)
+
+    def forward(self, pixel_patches):
+        """pixel_patches: (num_patches, patch_size*patch_size*3) -> (num_patches, hidden_size)"""
+        if pixel_patches.shape[-1] != self.patch_dim:
+            raise ValueError(
+                f"RawPixelPatchLinearEmbed expected last dim {self.patch_dim} "
+                f"for {self.patch_size}x{self.patch_size}x{self.in_channels} patches, "
+                f"got {pixel_patches.shape[-1]}"
+            )
+        return self.proj(pixel_patches)
+
+
+class RawPixelVideoTubeLinearEmbed(nn.Module):
+    """Single-linear video tube embedding: raw pixel tube -> hidden_size."""
+
+    def __init__(self, patch_size=32, in_channels=3, temporal_patch_size=4, hidden_size=3584):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.temporal_patch_size = temporal_patch_size
+        self.patch_dim = patch_size * patch_size * in_channels
+        self.tube_dim = temporal_patch_size * self.patch_dim
+        self.proj = nn.Linear(self.tube_dim, hidden_size, bias=True)
+
+    def forward(self, pixel_video_tubes):
+        """pixel_video_tubes: (num_tubes, temporal_patch_size*patch_size*patch_size*3) -> hidden."""
+        if pixel_video_tubes.shape[-1] != self.tube_dim:
+            raise ValueError(
+                f"RawPixelVideoTubeLinearEmbed expected last dim {self.tube_dim} "
+                f"for {self.temporal_patch_size}x{self.patch_size}x{self.patch_size}x{self.in_channels} tubes, "
+                f"got {pixel_video_tubes.shape[-1]}"
+            )
+        return self.proj(pixel_video_tubes)
+
+
+
+class UnpatchifyHead(nn.Module):
+    """Predict pixel patches from hidden states. Zero-initialized output."""
+    def __init__(self, hidden_size=3584, patch_size=16, out_channels=3):
+        super().__init__()
+        self.patch_dim = patch_size * patch_size * out_channels
+        self.norm = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, self.patch_dim, bias=True)
+        nn.init.constant_(self.linear.weight, 0)
+        nn.init.constant_(self.linear.bias, 0)
+
+    def forward(self, hidden_states):
+        """Project hidden states to raw RGB patches."""
+        return self.linear(self.norm(hidden_states))
+
+
+class VideoUnpatchifyHead(nn.Module):
+    """Predict raw video tube patches from hidden states."""
+
+    def __init__(self, hidden_size=3584, patch_size=16, temporal_patch_size=4, out_channels=3):
+        super().__init__()
+        self.patch_dim = temporal_patch_size * patch_size * patch_size * out_channels
+        self.norm = nn.RMSNorm(hidden_size, eps=1e-6)
+        self.linear = nn.Linear(hidden_size, self.patch_dim, bias=True)
+        nn.init.constant_(self.linear.weight, 0)
+        nn.init.constant_(self.linear.bias, 0)
+
+    def forward(self, hidden_states):
+        return self.linear(self.norm(hidden_states))

modeling/pixelumm/pixelumm.py

@@ -0,0 +1,1043 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Tuple, Optional
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.attention.flex_attention import create_block_mask
+from torch.utils.checkpoint import checkpoint
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+from data.data_utils import PIXELUMM_VIDEO_MROPE_BASE_FPS, build_pixelumm_vision_mrope_position_ids, create_sparse_mask, tubeify
+from .qwen3_navit import NaiveCache
+from .loss_reduction import sample_mean_token_weights
+from .modeling_utils import TimestepEmbedder, RawPixelPatchLinearEmbed, RawPixelVideoTubeLinearEmbed, UnpatchifyHead, VideoUnpatchifyHead, patchify_image
+from .dpm_solver import sample_flow_dpm
+from .flow_unipc import sample_flow_unipc
+from .time_schedule import flow_shift_noise_timestep, pixelumm_resolution_shift, sample_logit_normal_noise_time
+_CHATML_USER_START = '<|im_start|>user\n'
+_CHATML_ASSISTANT_START = '<|im_start|>assistant\n'
+_CHATML_SYSTEM_START = '<|im_start|>system\n'
+_CHATML_END = '<|im_end|>'
+_CHATML_SEP = '\n'
+_DEFAULT_SYSTEM_PROMPT = 'You are a helpful assistant.'
+_PIXELUMM_TEXT_TOKEN_LIMIT = 4096
+
+def _checkpointed_chunked_lm_head_cross_entropy(lm_head: nn.Module, hidden_states: torch.Tensor, labels: torch.Tensor, chunk_size: int) -> torch.Tensor:
+    """Compute token CE without retaining one full ``tokens x vocab`` tensor.
+
+    Each chunk is activation-checkpointed, so its LM-head logits are discarded
+    after the forward pass and recomputed during backward.  This preserves the
+    exact per-token CE contract used by task logging/reweighting while bounding
+    peak logits/workspace memory by ``chunk_size x vocab_size``.
+    """
+    chunk_size = int(chunk_size)
+    if chunk_size <= 0 or hidden_states.shape[0] <= chunk_size:
+        logits = lm_head(hidden_states)
+        return F.cross_entropy(logits, labels, reduction='none')
+
+    def chunk_loss(chunk_hidden: torch.Tensor, chunk_labels: torch.Tensor) -> torch.Tensor:
+        logits = lm_head(chunk_hidden)
+        return F.cross_entropy(logits, chunk_labels, reduction='none')
+    losses = []
+    for start in range(0, hidden_states.shape[0], chunk_size):
+        end = min(start + chunk_size, hidden_states.shape[0])
+        chunk_hidden = hidden_states[start:end]
+        chunk_labels = labels[start:end]
+        if torch.is_grad_enabled():
+            loss = checkpoint(chunk_loss, chunk_hidden, chunk_labels, use_reentrant=False)
+        else:
+            loss = chunk_loss(chunk_hidden, chunk_labels)
+        losses.append(loss)
+    return torch.cat(losses, dim=0)
+
+class PixelUMMConfig(PretrainedConfig):
+
+    def __init__(
+        self,
+        visual_gen=True,
+        visual_und=True,
+        llm_config=None,
+        llm_backend='qwen3',
+        gen_space='pixel',
+        und_space='pixel',
+        patch_size=16,
+        pixel_embedder_type='image_raw_patch_linear',
+        pixel_image_additive_pos_embed_type='none',
+        pixel_video_additive_pos_embed_type='none',
+        add_timestep_embedding=False,
+        add_noise_scale_embedding=False,
+        noise_scale=1.0,
+        noise_scale_mode='constant',
+        enable_pixel_video=True,
+        pixel_video_embedder_type='video_raw_tube_linear',
+        pixel_separate_und_gen_embedder=True,
+        pixel_video_separate_und_gen_embedder=True,
+        pixel_video_temporal_patch_size=4,
+        noise_time_logit_mean=0.0,
+        noise_time_logit_std=1.0,
+        time_schedule='pixelumm_resolution',
+        prediction_type='x',
+        loss_type='v_loss',
+        x_pred_t_min=0.05,
+        pixel_head_type='minit2i_linear',
+        pixel_video_head_type='jit_style',
+        ce_loss_checkpoint_chunk_size=0,
+        mse_loss_reduction='sample',
+        packed_attention_impl='flex',
+        packed_expert_routing='pixelumm_mot',
+        packed_sequence_layout='pixelumm_mot',
+        conditional_branch_gradient_contract='pixelumm_zero_v1',
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.visual_gen = visual_gen
+        self.visual_und = visual_und
+        if not self.visual_gen or not self.visual_und:
+            raise ValueError('PixelUMM-release requires both R07 visual GEN and visual UND branches')
+        self.llm_config = llm_config
+        self.llm_backend = llm_backend
+        if self.llm_backend != 'qwen3':
+            raise ValueError(f'PixelUMM-release requires the R07 Qwen3 backend; got llm_backend={self.llm_backend!r}')
+        self.gen_space = gen_space
+        self.und_space = und_space
+        if self.gen_space != 'pixel' or self.und_space != 'pixel':
+            raise ValueError(f'PixelUMM-release requires the R07 raw-pixel GEN/UND spaces; got gen_space={self.gen_space!r}, und_space={self.und_space!r}')
+        self.conditional_branch_gradient_contract = conditional_branch_gradient_contract
+        if self.conditional_branch_gradient_contract != 'pixelumm_zero_v1':
+            raise ValueError(f"PixelUMM-release requires the R07 conditional gradient contract 'pixelumm_zero_v1', got {self.conditional_branch_gradient_contract!r}")
+        self.patch_size = patch_size
+        self.pixel_embedder_type = pixel_embedder_type
+        self.pixel_image_additive_pos_embed_type = pixel_image_additive_pos_embed_type
+        self.pixel_video_additive_pos_embed_type = pixel_video_additive_pos_embed_type
+        if self.pixel_embedder_type != 'image_raw_patch_linear':
+            raise ValueError(f"PixelUMM-release only supports the R07 image embedder pixel_embedder_type='image_raw_patch_linear', got {self.pixel_embedder_type!r}")
+        if self.pixel_image_additive_pos_embed_type != 'none':
+            raise ValueError("PixelUMM-release requires the R07 image positional policy pixel_image_additive_pos_embed_type='none'")
+        if self.pixel_video_additive_pos_embed_type != 'none':
+            raise ValueError("PixelUMM-release requires the R07 video positional policy pixel_video_additive_pos_embed_type='none'")
+        if add_noise_scale_embedding and (not add_timestep_embedding):
+            raise ValueError('add_noise_scale_embedding=True requires add_timestep_embedding=True')
+        if getattr(llm_config, 'qwen3_mrope_type', None) != 'sensenova':
+            raise ValueError("PixelUMM-release requires qwen3_mrope_type='sensenova'")
+        self.add_timestep_embedding = add_timestep_embedding
+        self.add_noise_scale_embedding = add_noise_scale_embedding
+        if self.add_timestep_embedding or self.add_noise_scale_embedding:
+            raise ValueError('PixelUMM-release requires the R07 embedding policy: add_timestep_embedding=False and add_noise_scale_embedding=False')
+        self.noise_scale = noise_scale
+        self.noise_scale_mode = noise_scale_mode
+        if self.noise_scale_mode != 'constant' or float(self.noise_scale) != 1.0:
+            raise ValueError("PixelUMM-release requires the R07 unit Gaussian noise scale: noise_scale_mode='constant' and noise_scale=1.0")
+        self.enable_pixel_video = enable_pixel_video
+        if not self.enable_pixel_video:
+            raise ValueError('PixelUMM-release requires the R07 pixel-video path')
+        self.pixel_video_embedder_type = pixel_video_embedder_type
+        self.pixel_separate_und_gen_embedder = bool(pixel_separate_und_gen_embedder)
+        if not isinstance(pixel_video_separate_und_gen_embedder, bool):
+            raise ValueError('pixel_video_separate_und_gen_embedder must be a boolean')
+        self.pixel_video_separate_und_gen_embedder = pixel_video_separate_und_gen_embedder
+        if not self.pixel_separate_und_gen_embedder:
+            raise ValueError('PixelUMM-release requires separate UND/GEN image embedders')
+        if self.pixel_video_embedder_type != 'video_raw_tube_linear':
+            raise ValueError(f"PixelUMM-release only supports the R07 video embedder pixel_video_embedder_type='video_raw_tube_linear', got {self.pixel_video_embedder_type!r}")
+        self.pixel_video_temporal_patch_size = pixel_video_temporal_patch_size
+        self.noise_time_logit_mean = noise_time_logit_mean
+        self.noise_time_logit_std = noise_time_logit_std
+        self.time_schedule = time_schedule
+        if self.time_schedule != 'pixelumm_resolution':
+            raise ValueError("PixelUMM-release requires time_schedule='pixelumm_resolution'")
+        self.prediction_type = prediction_type
+        self.loss_type = loss_type
+        if self.prediction_type != 'x' or self.loss_type != 'v_loss':
+            raise ValueError("PixelUMM-release requires prediction_type='x' and loss_type='v_loss'")
+        self.x_pred_t_min = x_pred_t_min
+        self.pixel_head_type = pixel_head_type
+        self.pixel_video_head_type = pixel_video_head_type
+        if self.pixel_head_type != 'minit2i_linear':
+            raise ValueError("PixelUMM-release requires pixel_head_type='minit2i_linear'")
+        if self.pixel_video_head_type != 'jit_style':
+            raise ValueError("PixelUMM-release requires pixel_video_head_type='jit_style'")
+        self.ce_loss_checkpoint_chunk_size = int(ce_loss_checkpoint_chunk_size)
+        if self.ce_loss_checkpoint_chunk_size < 0:
+            raise ValueError(f'ce_loss_checkpoint_chunk_size must be non-negative; got {self.ce_loss_checkpoint_chunk_size}.')
+        self.mse_loss_reduction = mse_loss_reduction
+        if self.mse_loss_reduction != 'sample':
+            raise ValueError('PixelUMM-release requires the R07 sample-mean MSE reduction')
+        self.packed_attention_impl = packed_attention_impl
+        if self.packed_attention_impl != 'flex':
+            raise ValueError('PixelUMM-release requires the R07 FlexAttention backend')
+        self.packed_expert_routing = packed_expert_routing
+        self.packed_sequence_layout = packed_sequence_layout
+        if self.packed_expert_routing != 'pixelumm_mot':
+            raise ValueError("PixelUMM-release requires packed_expert_routing='pixelumm_mot'")
+        if self.packed_sequence_layout != 'pixelumm_mot':
+            raise ValueError("PixelUMM-release requires packed_sequence_layout='pixelumm_mot'")
+
+class PixelUMM(PreTrainedModel):
+    config_class = PixelUMMConfig
+    base_model_prefix = 'pixelumm'
+
+    def __init__(self, language_model, config: PixelUMMConfig):
+        super().__init__(config)
+        self.language_model = language_model
+        self.hidden_size = config.llm_config.hidden_size
+        if getattr(config.llm_config, 'layer_module', None) != 'Qwen3MoTDecoderLayer':
+            raise ValueError("PixelUMM-release requires layer_module='Qwen3MoTDecoderLayer'")
+        self.num_heads = config.llm_config.num_attention_heads
+        self.time_embedder = TimestepEmbedder(self.hidden_size)
+        self.time_embedder.requires_grad_(False)
+
+        def build_pixel_image_embedder():
+            return RawPixelPatchLinearEmbed(patch_size=config.patch_size, in_channels=3, hidden_size=self.hidden_size)
+        self.pixel_patch_embed = build_pixel_image_embedder()
+        self.pixel_patch_embed_mot_gen = build_pixel_image_embedder()
+        self.pixel_video_patch_embed_mot_gen = RawPixelVideoTubeLinearEmbed(patch_size=config.patch_size, in_channels=3, temporal_patch_size=config.pixel_video_temporal_patch_size, hidden_size=self.hidden_size)
+        if config.pixel_video_separate_und_gen_embedder:
+            self.pixel_video_patch_embed = RawPixelVideoTubeLinearEmbed(patch_size=config.patch_size, in_channels=3, temporal_patch_size=config.pixel_video_temporal_patch_size, hidden_size=self.hidden_size)
+        self.unpatchify_head = UnpatchifyHead(hidden_size=self.hidden_size, patch_size=config.patch_size, out_channels=3)
+        self.video_unpatchify_head = VideoUnpatchifyHead(hidden_size=self.hidden_size, patch_size=config.patch_size, temporal_patch_size=config.pixel_video_temporal_patch_size, out_channels=3)
+        self.config = config
+
+    def _image_patch_embedder(self, gen_model=False):
+        return self.pixel_patch_embed_mot_gen if gen_model else self.pixel_patch_embed
+
+    def embed_pixel_video_tubes(self, pixel_video_tubes: torch.Tensor, gen_model: bool=True) -> torch.Tensor:
+        if not gen_model:
+            raise RuntimeError('Pixel video embedding is generation-only; video understanding uses image UND frames.')
+        return self.pixel_video_patch_embed_mot_gen(pixel_video_tubes)
+
+    def _sample_train_timesteps(self, raw_timesteps: torch.Tensor, modality: Optional[str]=None) -> torch.Tensor:
+        if modality not in {None, 'image', 'video'}:
+            raise ValueError(f'Unsupported timestep modality={modality!r}')
+        return sample_logit_normal_noise_time(
+            raw_timesteps,
+            logit_mean=float(self.config.noise_time_logit_mean),
+            logit_std=float(self.config.noise_time_logit_std),
+        )
+
+    def _apply_time_schedule(self, t: torch.Tensor, spatial_hw: Optional[Tuple[int, int]]=None) -> torch.Tensor:
+        """Shift PixelUMM noise-time ``t`` toward noise when shift > 1.
+
+        Unlike the upstream SenseNova implementation, PixelUMM defines
+        ``z_t = (1 - t) * clean + t * noise``.  ``t`` is therefore already
+        the noise fraction and must not be complemented before shifting.
+        """
+        if spatial_hw is None:
+            raise ValueError('time_schedule=pixelumm_resolution requires spatial_hw')
+        return flow_shift_noise_timestep(t, pixelumm_resolution_shift(*spatial_hw))
+
+    def _apply_time_schedule_per_sample(self, timesteps: torch.Tensor, token_lens, grid_hw, *, patch_size: int) -> torch.Tensor:
+        """Apply a spatially selected schedule without mixing packed samples."""
+        token_lens = self._tensor_to_int_list(token_lens)
+        if torch.is_tensor(grid_hw):
+            grid_hw = [(int(h), int(w)) for (h, w) in grid_hw.tolist()]
+        else:
+            grid_hw = [(int(h), int(w)) for (h, w) in grid_hw]
+        if len(token_lens) != len(grid_hw):
+            raise ValueError(f'Per-sample timestep metadata length mismatch: token_lens={len(token_lens)} grid_hw={len(grid_hw)}')
+        if sum(token_lens) != int(timesteps.numel()):
+            raise ValueError(f'Per-sample timestep token lengths do not match packed timesteps: {sum(token_lens)} vs {int(timesteps.numel())}')
+        chunks = []
+        cursor = 0
+        for (token_len, (grid_h, grid_w)) in zip(token_lens, grid_hw):
+            chunk = timesteps[cursor:cursor + token_len]
+            chunks.append(self._apply_time_schedule(chunk, spatial_hw=(grid_h * int(patch_size), grid_w * int(patch_size))))
+            cursor += token_len
+        return torch.cat(chunks, dim=0) if chunks else timesteps
+
+    def _text_mrope_position_ids(self, start_position, length, device=None):
+        pos = torch.arange(start_position, start_position + length, dtype=torch.long, device=device)
+        zeros = torch.zeros_like(pos)
+        return torch.stack([pos, zeros, zeros], dim=0)
+
+    @staticmethod
+    def _pixelumm_vfm_vision_mrope_position_ids(temporal_offset, t_tokens, h_tokens, w_tokens, device=None, fps=None, base_fps=PIXELUMM_VIDEO_MROPE_BASE_FPS, temporal_compression_factor=4):
+        return build_pixelumm_vision_mrope_position_ids(temporal_offset, t_tokens, h_tokens, w_tokens, device=device, fps=fps, base_fps=base_fps, temporal_compression_factor=temporal_compression_factor)
+
+    def _image_mrope_position_ids(self, start_position, h_tokens, w_tokens, device=None, include_start=True, include_end=True):
+        grid_t = start_position + int(include_start)
+        rows = torch.arange(h_tokens, dtype=torch.long, device=device).repeat_interleave(w_tokens)
+        cols = torch.arange(w_tokens, dtype=torch.long, device=device).repeat(h_tokens)
+        t = torch.full((h_tokens * w_tokens,), grid_t, dtype=torch.long, device=device)
+        (t_chunks, h_chunks, w_chunks) = ([], [], [])
+        if include_start:
+            start = torch.full((1,), start_position, dtype=torch.long, device=device)
+            t_chunks.append(start)
+            h_chunks.append(start.new_zeros(1))
+            w_chunks.append(start.new_zeros(1))
+        t_chunks.append(t)
+        h_chunks.append(rows)
+        w_chunks.append(cols)
+        next_position = grid_t + 1
+        if include_end:
+            end = torch.full((1,), next_position, dtype=torch.long, device=device)
+            t_chunks.append(end)
+            h_chunks.append(end.new_zeros(1))
+            w_chunks.append(end.new_zeros(1))
+            next_position += 1
+        legacy_t = torch.cat(t_chunks)
+        axes = [legacy_t, torch.cat(h_chunks), torch.cat(w_chunks)]
+        return (torch.stack(axes, dim=0), next_position)
+
+    def _video_mrope_position_ids(self, start_position, temporal_groups, h_tokens, w_tokens, device=None, fps=None, temporal_compression_factor=None, include_start=True, include_end=True):
+        grid_start = start_position + int(include_start)
+        if fps is not None:
+            (vision_ids, next_position) = self._pixelumm_vfm_vision_mrope_position_ids(grid_start, temporal_groups, h_tokens, w_tokens, device=device, fps=fps, base_fps=PIXELUMM_VIDEO_MROPE_BASE_FPS, temporal_compression_factor=temporal_compression_factor if temporal_compression_factor is not None else self.config.pixel_video_temporal_patch_size)
+            t = vision_ids[0]
+            h = vision_ids[1]
+            w = vision_ids[2]
+        else:
+            t = torch.arange(temporal_groups, dtype=torch.long, device=device)[:, None, None]
+            h = torch.arange(h_tokens, dtype=torch.long, device=device)[None, :, None]
+            w = torch.arange(w_tokens, dtype=torch.long, device=device)[None, None, :]
+            t = t.expand(temporal_groups, h_tokens, w_tokens).reshape(-1) + grid_start
+            h = h.expand(temporal_groups, h_tokens, w_tokens).reshape(-1)
+            w = w.expand(temporal_groups, h_tokens, w_tokens).reshape(-1)
+            next_position = grid_start + max(1, temporal_groups)
+        (t_chunks, h_chunks, w_chunks) = ([], [], [])
+        if include_start:
+            start = torch.full((1,), start_position, dtype=torch.long, device=device)
+            t_chunks.append(start.to(dtype=t.dtype))
+            h_chunks.append(torch.zeros(1, dtype=h.dtype, device=device))
+            w_chunks.append(torch.zeros(1, dtype=w.dtype, device=device))
+        t_chunks.append(t)
+        h_chunks.append(h)
+        w_chunks.append(w)
+        if include_end:
+            end = torch.full((1,), next_position, dtype=t.dtype, device=device)
+            t_chunks.append(end)
+            h_chunks.append(torch.zeros(1, dtype=h.dtype, device=device))
+            w_chunks.append(torch.zeros(1, dtype=w.dtype, device=device))
+            next_position += 1
+        legacy_t = torch.cat(t_chunks)
+        axes = [legacy_t, torch.cat(h_chunks), torch.cat(w_chunks)]
+        return (torch.stack(axes, dim=0), next_position)
+
+    def _advance_text_query_position_ids(self, packed_query_position_ids):
+        next_position_ids = packed_query_position_ids.clone()
+        next_position_ids[0] = next_position_ids[0] + 1
+        return next_position_ids
+
+    def _pixel_video_token_lens(self, packed_pixel_video_patches, pixel_video_patch_seqlens=None):
+        if pixel_video_patch_seqlens is not None:
+            return [int(x) for x in pixel_video_patch_seqlens.tolist()]
+        return [int(packed_pixel_video_patches.shape[0])]
+
+    @staticmethod
+    def _tensor_to_int_list(values):
+        if values is None:
+            return None
+        if torch.is_tensor(values):
+            values = values.detach().cpu().flatten().tolist()
+        return [int(x) for x in values]
+
+    @staticmethod
+    def _has_trainable_parameters(module: Optional[nn.Module]) -> bool:
+        return module is not None and any((parameter.requires_grad for parameter in module.parameters()))
+
+    @staticmethod
+    def _zero_graph_anchor(value: torch.Tensor) -> torch.Tensor:
+        return value.sum() * 0.0
+
+    def _dummy_pixel_image_embedder_anchor(self, module: nn.Module, reference: torch.Tensor) -> torch.Tensor:
+        patch_dim = getattr(module, 'patch_dim', None)
+        if patch_dim is None:
+            raise RuntimeError(f'pixelumm_zero_v1 cannot build an image-embedder dummy for {type(module).__name__}: missing patch_dim')
+        dummy = reference.new_zeros((1, int(patch_dim)))
+        output = module(dummy)
+        return self._zero_graph_anchor(output)
+
+    def _dummy_pixel_video_embedder_anchor(self, module: nn.Module, reference: torch.Tensor, *, gen_model: bool) -> torch.Tensor:
+        tube_dim = getattr(module, 'tube_dim', None)
+        if tube_dim is None:
+            tube_dim = getattr(module, 'video_tube_dim', None)
+        if tube_dim is None:
+            tube_dim = getattr(module, 'patch_dim', None)
+        if tube_dim is None:
+            raise RuntimeError(f'pixelumm_zero_v1 cannot build a video-embedder dummy for {type(module).__name__}: missing tube_dim/video_tube_dim/patch_dim')
+        dummy = reference.new_zeros((1, int(tube_dim)))
+        position_ids = torch.zeros((1, 3), dtype=torch.long, device=reference.device)
+        if gen_model:
+            output = self.embed_pixel_video_tubes(dummy, gen_model=True)
+        else:
+            try:
+                output = module(dummy, position_ids=position_ids)
+            except TypeError:
+                output = module(dummy)
+        return self._zero_graph_anchor(output)
+
+    def _dummy_pixel_image_head_anchor(self, reference: torch.Tensor) -> torch.Tensor:
+        head = self.unpatchify_head
+        dummy_hidden = reference.new_zeros((1, self.hidden_size))
+        output = head(dummy_hidden)
+        return self._zero_graph_anchor(output)
+
+    def _dummy_pixel_video_head_anchor(self, reference: torch.Tensor) -> torch.Tensor:
+        dummy_hidden = reference.new_zeros((1, self.hidden_size))
+        output = self.video_unpatchify_head(dummy_hidden)
+        return self._zero_graph_anchor(output)
+
+    def _pixelumm_zero_root_graph_anchor(self, last_hidden_state: torch.Tensor, *, has_und_image: bool, has_und_video: bool, has_gen_image: bool, has_gen_video: bool, has_text_output: bool) -> torch.Tensor:
+        """Keep every inactive *trainable* root branch in the autograd graph.
+
+        This mirrors PixelUMM VFM's dummy-forward contract.  Frozen coarse
+        branches remain excluded because their parameters have
+        ``requires_grad=False``; the helper therefore cannot create optimizer
+        state or parameter updates for F18/F19's disabled branch.
+        """
+        anchor = last_hidden_state.new_zeros(())
+        if not has_text_output and self._has_trainable_parameters(self.language_model.lm_head):
+            anchor = anchor + self._zero_graph_anchor(self.language_model.lm_head(last_hidden_state[:0]))
+        image_und_embedder = getattr(self, 'pixel_patch_embed', None)
+        if not has_und_image and self._has_trainable_parameters(image_und_embedder):
+            anchor = anchor + self._dummy_pixel_image_embedder_anchor(image_und_embedder, last_hidden_state)
+        video_und_embedder = getattr(self, 'pixel_video_patch_embed', None)
+        if not has_und_video and self._has_trainable_parameters(video_und_embedder):
+            anchor = anchor + self._dummy_pixel_video_embedder_anchor(video_und_embedder, last_hidden_state, gen_model=False)
+        image_gen_embedder = self._image_patch_embedder(gen_model=True)
+        if not has_gen_image and self._has_trainable_parameters(image_gen_embedder):
+            anchor = anchor + self._dummy_pixel_image_embedder_anchor(image_gen_embedder, last_hidden_state)
+        if not has_gen_image and self._has_trainable_parameters(getattr(self, 'unpatchify_head', None)):
+            anchor = anchor + self._dummy_pixel_image_head_anchor(last_hidden_state)
+        video_gen_embedder = getattr(self, 'pixel_video_patch_embed_mot_gen', None)
+        if not has_gen_video and self._has_trainable_parameters(video_gen_embedder):
+            anchor = anchor + self._dummy_pixel_video_embedder_anchor(video_gen_embedder, last_hidden_state, gen_model=True)
+        video_head = getattr(self, 'video_unpatchify_head', None)
+        if not has_gen_video and self._has_trainable_parameters(video_head):
+            anchor = anchor + self._dummy_pixel_video_head_anchor(last_hidden_state)
+        return anchor
+
+    def forward(
+        self,
+        sequence_length: int,
+        packed_text_ids: torch.LongTensor,
+        packed_text_indexes: torch.LongTensor,
+        sample_lens: List[int],
+        packed_position_ids: torch.LongTensor,
+        packed_mrope_position_ids: Optional[torch.LongTensor] = None,
+        split_lens: List[int] = None,
+        attn_modes: List[str] = None,
+        ce_loss_indexes: Optional[torch.BoolTensor] = None,
+        packed_label_ids: Optional[torch.LongTensor] = None,
+        packed_pixel_patches_und: Optional[torch.Tensor] = None,
+        packed_pixel_patch_indexes_und: Optional[torch.LongTensor] = None,
+        packed_pixel_video_patches_und: Optional[torch.Tensor] = None,
+        packed_pixel_video_patch_indexes_und: Optional[torch.LongTensor] = None,
+        packed_pixel_patches_gen: Optional[torch.Tensor] = None,
+        packed_pixel_patch_indexes_gen: Optional[torch.LongTensor] = None,
+        packed_pixel_grid_hw_gen: Optional[torch.LongTensor] = None,
+        packed_gen_image_lens: Optional[torch.LongTensor] = None,
+        packed_pixel_video_patches_gen: Optional[torch.Tensor] = None,
+        packed_pixel_video_patch_indexes_gen: Optional[torch.LongTensor] = None,
+        pixel_video_patch_seqlens_gen: Optional[torch.IntTensor] = None,
+        pixel_video_grid_hw_gen: Optional[torch.LongTensor] = None,
+        packed_gen_video_lens: Optional[torch.LongTensor] = None,
+        packed_timesteps: Optional[torch.LongTensor] = None,
+    ) -> dict:
+        """
+        Args:
+            sequence_length: length of sequence.
+            packed_text_ids: 1-D int tensor, packed text token ids.
+            packed_text_indexes: 1-D int tensor, packed text token indexes in sequence.
+            sample_lens: A list of N ints, length of each sample in packed_sequence.
+            packed_position_ids: packed 1-D positions, an image has only one global position shared
+                by all latent tokens.
+
+            packed_label_ids: 1-D int tensor, packed label token ids.
+            ce_loss_indexes: 1-D bool tensor, where to compute ce loss.
+            packed_timesteps: 1-D float tensor, flow timesteps. 0 indicates use clean image.
+        """
+        packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
+        packed_sequence = packed_text_embedding.new_zeros(size=(sequence_length, self.hidden_size))
+        packed_sequence[packed_text_indexes] = packed_text_embedding
+        logical_seqlen = sum(sample_lens)
+        flex_seqlen = (int(sequence_length) + 127) // 128 * 128
+        if flex_seqlen < logical_seqlen:
+            raise ValueError((logical_seqlen, flex_seqlen))
+        sparse_mask = create_sparse_mask(
+            sample_lens,
+            split_lens,
+            attn_modes,
+            packed_text_embedding.device,
+            total_length=flex_seqlen,
+        )
+        attention_mask = create_block_mask(
+            sparse_mask,
+            B=1,
+            H=1,
+            Q_LEN=flex_seqlen,
+            KV_LEN=flex_seqlen,
+            device=packed_text_embedding.device,
+            BLOCK_SIZE=128,
+            _compile=True,
+        )
+        packed_und_image_indexes = None
+        if packed_pixel_patches_und is not None:
+            und_embed = self.pixel_patch_embed(packed_pixel_patches_und)
+            packed_sequence[packed_pixel_patch_indexes_und] = und_embed
+            packed_und_image_indexes = packed_pixel_patch_indexes_und
+        if packed_pixel_video_patches_und is not None:
+            if not hasattr(self, 'pixel_video_patch_embed'):
+                raise ValueError('R07 video UND patches require the dedicated video embedder')
+            video_und_embed = self.pixel_video_patch_embed(packed_pixel_video_patches_und)
+            packed_sequence[packed_pixel_video_patch_indexes_und] = video_und_embed
+            if packed_und_image_indexes is None:
+                packed_und_image_indexes = packed_pixel_video_patch_indexes_und
+            else:
+                packed_und_image_indexes = torch.cat([packed_und_image_indexes, packed_pixel_video_patch_indexes_und], dim=0).sort().values
+        packed_gen_token_indexes = None
+        image_clean_target = None
+        image_noisy_patches = None
+        image_timesteps = None
+        video_clean_target = None
+        video_noisy_patches = None
+        video_timesteps = None
+        has_gen_tokens = packed_pixel_patches_gen is not None or packed_pixel_video_patches_gen is not None
+        if has_gen_tokens and packed_timesteps is None:
+            raise ValueError('Generation tokens are present but packed_timesteps is missing')
+        raw_timesteps = packed_timesteps if has_gen_tokens else None
+        if not has_gen_tokens and raw_timesteps is not None:
+            raise ValueError('PixelUMM-release received timesteps without R07 pixel GEN tokens')
+        all_gen_indexes = []
+        if packed_pixel_patches_gen is not None:
+            all_gen_indexes.append(packed_pixel_patch_indexes_gen)
+        if packed_pixel_video_patches_gen is not None:
+            all_gen_indexes.append(packed_pixel_video_patch_indexes_gen)
+        if all_gen_indexes:
+            sorted_gen_indexes = torch.cat(all_gen_indexes, dim=0).sort().values
+            if sorted_gen_indexes.numel() != raw_timesteps.numel():
+                raise ValueError(f'packed_timesteps length does not match image/video generation token count: {raw_timesteps.numel()} vs {sorted_gen_indexes.numel()}')
+            timestep_by_seq = raw_timesteps.new_zeros(sequence_length)
+            timestep_by_seq[sorted_gen_indexes] = raw_timesteps
+        if packed_pixel_video_patches_gen is not None:
+            clean_patches = packed_pixel_video_patches_gen
+            video_timesteps = timestep_by_seq[packed_pixel_video_patch_indexes_gen]
+            video_timesteps = self._sample_train_timesteps(video_timesteps, modality='video')
+            video_token_lens = self._pixel_video_token_lens(packed_pixel_video_patches_gen, pixel_video_patch_seqlens_gen)
+            if sum(video_token_lens) != int(video_timesteps.shape[0]):
+                raise ValueError(f'pixel_video_patch_seqlens_gen does not match packed video gen token count: {sum(video_token_lens)} vs {int(video_timesteps.shape[0])}')
+            if pixel_video_grid_hw_gen is None:
+                raise ValueError('time_schedule=pixelumm_resolution requires pixel_video_grid_hw_gen')
+            video_timesteps = self._apply_time_schedule_per_sample(video_timesteps, video_token_lens, pixel_video_grid_hw_gen, patch_size=self.config.patch_size)
+            noise = torch.randn_like(clean_patches)
+            noisy_patches = (1 - video_timesteps[:, None]) * clean_patches + video_timesteps[:, None] * noise
+            gen_embed = self.embed_pixel_video_tubes(noisy_patches)
+            packed_sequence[packed_pixel_video_patch_indexes_gen] = gen_embed
+            packed_gen_token_indexes = packed_pixel_video_patch_indexes_gen
+            video_clean_target = clean_patches
+            video_noisy_patches = noisy_patches
+        if packed_pixel_patches_gen is not None:
+            clean_patches = packed_pixel_patches_gen
+            image_timesteps = timestep_by_seq[packed_pixel_patch_indexes_gen]
+            image_timesteps = self._sample_train_timesteps(image_timesteps, modality='image')
+            if packed_pixel_grid_hw_gen is None:
+                raise ValueError('time_schedule=pixelumm_resolution requires packed_pixel_grid_hw_gen')
+            schedule_token_lens = [int(h) * int(w) for (h, w) in packed_pixel_grid_hw_gen.tolist()]
+            image_timesteps = self._apply_time_schedule_per_sample(image_timesteps, schedule_token_lens, packed_pixel_grid_hw_gen, patch_size=self.config.patch_size)
+            noise = torch.randn_like(clean_patches)
+            noisy_patches = (1 - image_timesteps[:, None]) * clean_patches + image_timesteps[:, None] * noise
+            gen_embed = self._image_patch_embedder(gen_model=True)(noisy_patches)
+            packed_sequence[packed_pixel_patch_indexes_gen] = gen_embed
+            if packed_gen_token_indexes is None:
+                packed_gen_token_indexes = packed_pixel_patch_indexes_gen
+            else:
+                packed_gen_token_indexes = torch.cat([packed_gen_token_indexes, packed_pixel_patch_indexes_gen], dim=0).sort().values
+            image_clean_target = clean_patches
+            image_noisy_patches = noisy_patches
+        packed_und_token_indexes = packed_text_indexes
+        if packed_und_image_indexes is not None:
+            packed_und_token_indexes = torch.cat([packed_und_token_indexes, packed_und_image_indexes], dim=0)
+        extra_inputs = {'packed_und_token_indexes': packed_und_token_indexes, 'packed_gen_token_indexes': packed_gen_token_indexes}
+        llm_position_ids = packed_mrope_position_ids if packed_mrope_position_ids is not None else packed_position_ids
+        if llm_position_ids is not None:
+            llm_position_ids = llm_position_ids.contiguous()
+        last_hidden_state = self.language_model(packed_sequence=packed_sequence, sample_lens=sample_lens, attention_mask=attention_mask, packed_position_ids=llm_position_ids, **extra_inputs)
+        mse = None
+        mse_loss_reduction_weights = None
+        if image_clean_target is not None or video_clean_target is not None:
+            mse_parts = []
+            mse_weight_parts = []
+
+            def per_token_mse(pred, target, loss_mask=None):
+                sq = (pred.float() - target.float()) ** 2
+                if loss_mask is None:
+                    return sq.mean(dim=-1, keepdim=True)
+                mask = loss_mask.to(device=sq.device, dtype=sq.dtype)
+                denom = mask.sum(dim=-1, keepdim=True).clamp_min(1.0)
+                return (sq * mask).sum(dim=-1, keepdim=True) / denom
+            if image_clean_target is not None:
+                image_has_mse = image_timesteps > 0
+                image_hidden = last_hidden_state[packed_pixel_patch_indexes_gen[image_has_mse]]
+                image_model_out = self.unpatchify_head(image_hidden)
+                t = image_timesteps[:, None]
+                t_f32 = t[image_has_mse].float().clamp(min=self.config.x_pred_t_min)
+                v_pred = (image_noisy_patches[image_has_mse].float() - image_model_out.float()) / t_f32
+                v_target = (image_noisy_patches[image_has_mse].float() - image_clean_target[image_has_mse].float()) / t_f32
+                image_mse = per_token_mse(v_pred, v_target)
+                mse_parts.append(image_mse)
+                if packed_gen_image_lens is None:
+                    raise ValueError('R07 sample-mean MSE requires packed_gen_image_lens')
+                mse_weight_parts.append(sample_mean_token_weights(packed_gen_image_lens, output_size=image_mse.shape[0], device=image_mse.device))
+            if video_clean_target is not None:
+                video_has_mse = video_timesteps > 0
+                video_hidden = last_hidden_state[packed_pixel_video_patch_indexes_gen[video_has_mse]]
+                video_model_out = self.video_unpatchify_head(video_hidden)
+                t = video_timesteps[:, None]
+                t_f32 = t[video_has_mse].float().clamp(min=self.config.x_pred_t_min)
+                v_pred = (video_noisy_patches[video_has_mse].float() - video_model_out.float()) / t_f32
+                v_target = (video_noisy_patches[video_has_mse].float() - video_clean_target[video_has_mse].float()) / t_f32
+                video_mse = per_token_mse(v_pred, v_target)
+                mse_parts.append(video_mse)
+                if packed_gen_video_lens is None:
+                    raise ValueError('R07 sample-mean MSE requires packed_gen_video_lens')
+                mse_weight_parts.append(sample_mean_token_weights(packed_gen_video_lens, output_size=video_mse.shape[0], device=video_mse.device))
+            mse = torch.cat(mse_parts, dim=0)
+            if mse_weight_parts:
+                mse_loss_reduction_weights = torch.cat(mse_weight_parts)
+        ce = None
+        if ce_loss_indexes is not None:
+            ce = _checkpointed_chunked_lm_head_cross_entropy(self.language_model.lm_head, last_hidden_state[ce_loss_indexes], packed_label_ids, self.config.ce_loss_checkpoint_chunk_size)
+
+        def _has_elements(value: Optional[torch.Tensor]) -> bool:
+            return value is not None and int(value.numel()) > 0
+        conditional_graph_anchor = self._pixelumm_zero_root_graph_anchor(last_hidden_state, has_und_image=_has_elements(packed_pixel_patches_und), has_und_video=_has_elements(packed_pixel_video_patches_und), has_gen_image=_has_elements(packed_pixel_patches_gen), has_gen_video=_has_elements(packed_pixel_video_patches_gen), has_text_output=ce is not None and int(ce.numel()) > 0)
+        return dict(mse=mse, ce=ce, mse_loss_reduction_weights=mse_loss_reduction_weights, conditional_graph_anchor=conditional_graph_anchor)
+
+    def _encode_prompt_text(self, tokenizer, prompt: str) -> list[int]:
+        try:
+            return tokenizer.encode(prompt, add_special_tokens=False)
+        except TypeError:
+            return tokenizer.encode(prompt)
+
+    def _text_to_visual_chatml_prompt_ids(self, tokenizer, prompt: str) -> list[int]:
+        caption_ids = self._encode_prompt_text(tokenizer, prompt)[:_PIXELUMM_TEXT_TOKEN_LIMIT]
+        return self._encode_prompt_text(tokenizer, _CHATML_USER_START) + caption_ids + self._encode_prompt_text(tokenizer, f'{_CHATML_END}{_CHATML_SEP}{_CHATML_ASSISTANT_START}')
+
+    def _image_to_text_chatml_prefix_ids(self, tokenizer) -> list[int]:
+        return self._encode_prompt_text(tokenizer, f'{_CHATML_SYSTEM_START}{_DEFAULT_SYSTEM_PROMPT}{_CHATML_END}{_CHATML_SEP}{_CHATML_USER_START}')
+
+    def _image_to_text_chatml_suffix_ids(self, tokenizer, instruction: str) -> list[int]:
+        return self._encode_prompt_text(tokenizer, f'{instruction}{_CHATML_END}{_CHATML_SEP}{_CHATML_ASSISTANT_START}')
+
+    def _prepare_token_ids_as_text_cache_input(self, curr_kvlens, curr_rope, prompt_ids_list):
+        packed_text_ids = list()
+        packed_text_position_ids = list()
+        text_token_lens = list()
+        packed_text_indexes = list()
+        packed_key_value_indexes = list()
+        curr = 0
+        (newlens, new_rope) = (list(), list())
+        for (text_ids, curr_kvlen, curr_position_id) in zip(prompt_ids_list, curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+            text_token_lens.append(len(text_ids))
+            packed_text_ids.extend(text_ids)
+            packed_text_position_ids.append(self._text_mrope_position_ids(curr_position_id, len(text_ids)))
+            packed_text_indexes.extend(range(curr, curr + len(text_ids)))
+            newlens.append(curr_kvlen + len(text_ids))
+            new_rope.append(curr_position_id + len(text_ids))
+            curr += len(text_ids)
+        packed_text_position_ids = torch.cat(packed_text_position_ids, dim=1)
+        generation_input = {'text_token_lens': torch.tensor(text_token_lens, dtype=torch.int), 'packed_text_ids': torch.tensor(packed_text_ids, dtype=torch.long), 'packed_text_position_ids': packed_text_position_ids, 'packed_text_indexes': torch.tensor(packed_text_indexes, dtype=torch.long), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long), 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int)}
+        return (generation_input, newlens, new_rope)
+
+    def prepare_text_to_visual_chatml_prompts(self, curr_kvlens, curr_rope, prompts, tokenizer, new_token_ids):
+        """T2I/T2V inference prefill matching the ChatML training prompt."""
+        prompt_ids_list = [self._text_to_visual_chatml_prompt_ids(tokenizer, prompt) + [new_token_ids['eos_token_id']] for prompt in prompts]
+        return self._prepare_token_ids_as_text_cache_input(curr_kvlens, curr_rope, prompt_ids_list)
+
+    def prepare_image_to_text_chatml_prefixes(self, curr_kvlens, curr_rope, tokenizer):
+        """I2T inference prefill before the image tokens."""
+        prompt_ids = self._image_to_text_chatml_prefix_ids(tokenizer)
+        prompt_ids_list = [prompt_ids for _ in curr_kvlens]
+        return self._prepare_token_ids_as_text_cache_input(curr_kvlens, curr_rope, prompt_ids_list)
+
+    def prepare_image_to_text_chatml_suffix_contexts_for_decode(self, curr_kvlens, curr_rope, instructions, tokenizer):
+        """I2T inference prefill after image tokens, aligned with train-time CE shift."""
+        suffix_ids_list = [self._image_to_text_chatml_suffix_ids(tokenizer, instruction) for instruction in instructions]
+        if any((len(ids) == 0 for ids in suffix_ids_list)):
+            raise ValueError('image-to-text ChatML suffix cannot be empty')
+        context_ids_list = [ids[:-1] for ids in suffix_ids_list]
+        start_token_ids = [ids[-1] for ids in suffix_ids_list]
+        (generation_input, newlens, new_rope) = self._prepare_token_ids_as_text_cache_input(curr_kvlens, curr_rope, context_ids_list)
+        return (generation_input, newlens, new_rope, start_token_ids)
+
+    @torch.no_grad
+    def forward_cache_update_text(self, past_key_values: NaiveCache, packed_text_ids: torch.IntTensor, packed_text_position_ids: torch.LongTensor, text_token_lens: torch.LongTensor, packed_text_indexes: torch.LongTensor, packed_key_value_indexes: torch.LongTensor, key_values_lens: torch.IntTensor):
+        packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
+        output = self.language_model.forward_inference(packed_query_sequence=packed_text_embedding, query_lens=text_token_lens, packed_query_position_ids=packed_text_position_ids, packed_query_indexes=packed_text_indexes, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, key_values_lens=key_values_lens, update_past_key_values=True, is_causal=True, mode='und')
+        past_key_values = output.past_key_values
+        return past_key_values
+
+    def _validate_pixel_generation_prefix_contract(self, visual_start_prefilled: bool) -> None:
+        if not visual_start_prefilled:
+            raise ValueError('pixelumm_mot pixel generation requires one causal <vision_start> in the KV cache')
+
+    def prepare_pixel_noise(self, curr_kvlens, curr_rope, image_sizes, *, visual_start_prefilled: bool=False):
+        """Generate image noise and construct the denoising query.
+
+        For ``pixelumm_mot``, callers must first cache exactly one causal
+        ``<vision_start>`` with :meth:`prepare_visual_delimiter_tokens`.  This
+        method then returns a patch-only, bidirectional noise query, matching
+        the train-time ``causal start -> full/noise patches`` mask exactly.
+        """
+        self._validate_pixel_generation_prefix_contract(visual_start_prefilled)
+        (packed_text_ids, packed_text_indexes) = (list(), list())
+        (packed_pixel_token_indexes, packed_init_noises) = (list(), list())
+        (packed_position_ids, packed_seqlens, packed_indexes) = (list(), list(), list())
+        packed_key_value_indexes = list()
+        query_curr = curr = 0
+        for ((H, W), curr_kvlen, curr_position_id) in zip(image_sizes, curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+            (h, w) = (H // self.config.patch_size, W // self.config.patch_size)
+            num_image_tokens = h * w
+            patch_dim = self.config.patch_size * self.config.patch_size * 3
+            packed_init_noises.append(torch.randn(num_image_tokens, patch_dim))
+            packed_pixel_token_indexes.extend(range(query_curr, query_curr + num_image_tokens))
+            packed_indexes.extend(range(curr, curr + num_image_tokens))
+            curr += num_image_tokens
+            query_curr += num_image_tokens
+            (mrope_ids, _) = self._image_mrope_position_ids(curr_position_id, h, w, include_start=False, include_end=False)
+            packed_position_ids.append(mrope_ids)
+            packed_seqlens.append(num_image_tokens + 0)
+        packed_position_ids = torch.cat(packed_position_ids, dim=1)
+        generation_input = {'packed_text_ids': torch.tensor(packed_text_ids, dtype=torch.long), 'packed_text_indexes': torch.tensor(packed_text_indexes, dtype=torch.long), 'packed_init_noises': torch.cat(packed_init_noises, dim=0), 'packed_pixel_token_indexes': torch.tensor(packed_pixel_token_indexes, dtype=torch.long), 'packed_seqlens': torch.tensor(packed_seqlens, dtype=torch.int), 'packed_position_ids': packed_position_ids, 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int), 'packed_indexes': torch.tensor(packed_indexes, dtype=torch.long), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long)}
+        return generation_input
+
+    def prepare_pixel_video_noise(self, curr_kvlens, curr_rope, video_sizes, fps=None, *, visual_start_prefilled: bool=False):
+        """Pixel-video version of :meth:`prepare_pixel_noise`."""
+        self._validate_pixel_generation_prefix_contract(visual_start_prefilled)
+        (packed_text_ids, packed_text_indexes) = (list(), list())
+        (packed_pixel_video_token_indexes, packed_init_noises) = (list(), list())
+        (packed_position_ids, packed_seqlens, packed_indexes) = (list(), list(), list())
+        packed_key_value_indexes = list()
+        query_curr = curr = 0
+        tube_frames = self.config.pixel_video_temporal_patch_size
+        patch_dim = tube_frames * self.config.patch_size * self.config.patch_size * 3
+        for ((num_frames, H, W), curr_kvlen, curr_position_id) in zip(video_sizes, curr_kvlens, curr_rope):
+            if num_frames % tube_frames != 0:
+                raise ValueError(f'Expected num_frames to be divisible by tube_frames={tube_frames}, got {num_frames}')
+            else:
+                temporal_groups = num_frames // tube_frames
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+            (h, w) = (H // self.config.patch_size, W // self.config.patch_size)
+            num_video_tokens = h * w * temporal_groups
+            packed_init_noises.append(torch.randn(num_video_tokens, patch_dim))
+            packed_pixel_video_token_indexes.extend(range(query_curr, query_curr + num_video_tokens))
+            packed_indexes.extend(range(curr, curr + num_video_tokens))
+            curr += num_video_tokens
+            query_curr += num_video_tokens
+            (mrope_ids, _) = self._video_mrope_position_ids(curr_position_id, temporal_groups, h, w, fps=fps, temporal_compression_factor=tube_frames, include_start=False, include_end=False)
+            packed_position_ids.append(mrope_ids)
+            packed_seqlens.append(num_video_tokens + 0)
+        packed_position_ids = torch.cat(packed_position_ids, dim=1)
+        generation_input = {'packed_text_ids': torch.tensor(packed_text_ids, dtype=torch.long), 'packed_text_indexes': torch.tensor(packed_text_indexes, dtype=torch.long), 'packed_init_noises': torch.cat(packed_init_noises, dim=0), 'packed_pixel_token_indexes': torch.tensor(packed_pixel_video_token_indexes, dtype=torch.long), 'packed_seqlens': torch.tensor(packed_seqlens, dtype=torch.int), 'packed_position_ids': packed_position_ids, 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int), 'packed_indexes': torch.tensor(packed_indexes, dtype=torch.long), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long)}
+        return generation_input
+
+    def prepare_pixel_images(self, curr_kvlens, curr_rope, images):
+        """Prepare clean pixel patches for understanding (KV cache prefill).
+
+        Pixel-space version of prepare_vit_images. No ViT — PIL image → pixel patchify.
+        Uses 1D RoPE (same as rest of unified PixelUMM).
+        """
+        packed_pixel_token_indexes = list()
+        packed_pixel_patches = []
+        (packed_text_ids, packed_text_indexes) = (list(), list())
+        (packed_seqlens, packed_position_ids, packed_indexes) = (list(), list(), list())
+        packed_key_value_indexes = list()
+        _curr = curr = 0
+        (newlens, new_rope) = (list(), list())
+        for (image, curr_kvlen, curr_position_id) in zip(images, curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+            ps = self.config.patch_size
+            (patches, (h_patches, w_patches)) = patchify_image(image, ps, max_size=max(image.size))
+            num_img_tokens = patches.shape[0]
+            (mrope_h, mrope_w) = (h_patches, w_patches)
+            packed_pixel_patches.append(patches)
+            packed_pixel_token_indexes.extend(range(_curr, _curr + num_img_tokens))
+            packed_indexes.extend(range(curr, curr + num_img_tokens))
+            curr += num_img_tokens
+            _curr += num_img_tokens
+            (mrope_ids, next_position_id) = self._image_mrope_position_ids(curr_position_id, mrope_h, mrope_w, include_start=False, include_end=False)
+            packed_position_ids.append(mrope_ids)
+            new_rope.append(next_position_id)
+            packed_seqlens.append(num_img_tokens)
+            newlens.append(curr_kvlen + num_img_tokens)
+        packed_position_ids = torch.cat(packed_position_ids, dim=1)
+        generation_input = {'packed_text_ids': torch.tensor(packed_text_ids, dtype=torch.long), 'packed_text_indexes': torch.tensor(packed_text_indexes, dtype=torch.long), 'packed_pixel_patches': torch.cat(packed_pixel_patches, dim=0), 'packed_pixel_token_indexes': torch.tensor(packed_pixel_token_indexes, dtype=torch.long), 'packed_position_ids': packed_position_ids, 'packed_seqlens': torch.tensor(packed_seqlens, dtype=torch.int), 'packed_indexes': torch.tensor(packed_indexes, dtype=torch.long), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long), 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int)}
+        return (generation_input, newlens, new_rope)
+
+    def prepare_text_spans_for_cache(self, curr_kvlens, curr_rope, texts, tokenizer):
+        """Causal mid-sequence text spans (e.g. V2T tube timestamp markers)."""
+        ids_list = [self._encode_prompt_text(tokenizer, text) for text in texts]
+        if any((len(ids) == 0 for ids in ids_list)):
+            raise ValueError('Cache text span cannot be empty')
+        return self._prepare_token_ids_as_text_cache_input(curr_kvlens, curr_rope, ids_list)
+
+    def prepare_pixel_video_und_tubes_for_text(self, curr_kvlens, curr_rope, tube_tensors):
+        """Pack video-UND tubes with the released 3D sequence/H/W coordinates."""
+        if not self.config.pixel_video_separate_und_gen_embedder:
+            raise ValueError('This checkpoint profile has no video UND embedder; video-VLM tubes are unsupported')
+        if not len(curr_kvlens) == len(curr_rope) == len(tube_tensors):
+            raise ValueError('V2T tube batches must match cache batch size')
+        temporal_patch_size = int(self.config.pixel_video_temporal_patch_size)
+        ps = self.config.patch_size
+        packed_pixel_video_patches = []
+        packed_pixel_video_patch_indexes = []
+        packed_position_ids = []
+        packed_seqlens = []
+        packed_indexes = []
+        packed_key_value_indexes = []
+        newlens = []
+        new_rope = []
+        query_curr = 0
+        combined_curr = 0
+        for (curr_kvlen, curr_position_id, tube) in zip(curr_kvlens, curr_rope, tube_tensors):
+            if tube.dim() != 4 or tube.shape[0] != temporal_patch_size:
+                raise ValueError(f'V2T tube must be (temporal_patch_size, C, H, W) with temporal_patch_size={temporal_patch_size}, got {tuple(tube.shape)}')
+            (_, _, H, W) = tube.shape
+            if H % ps or W % ps:
+                raise ValueError(f'V2T tube spatial size {(H, W)} not divisible by patch={ps}')
+            h = H // ps
+            w = W // ps
+            num_tokens = h * w
+            packed_key_value_indexes.extend(range(combined_curr, combined_curr + curr_kvlen))
+            combined_curr += curr_kvlen
+            tokens = tubeify(tube, ps, temporal_patch_size)
+            packed_pixel_video_patches.append(tokens)
+            packed_pixel_video_patch_indexes.extend(range(query_curr, query_curr + num_tokens))
+            packed_indexes.extend(range(combined_curr, combined_curr + num_tokens))
+            query_curr += num_tokens
+            combined_curr += num_tokens
+            grid_s = float(curr_position_id)
+            position_axes = [torch.full((num_tokens,), grid_s, dtype=torch.float32)]
+            position_axes.extend([torch.arange(h, dtype=torch.float32).repeat_interleave(w), torch.arange(w, dtype=torch.float32).repeat(h)])
+            packed_position_ids.append(torch.stack(position_axes, dim=0))
+            packed_seqlens.append(num_tokens)
+            newlens.append(curr_kvlen + num_tokens)
+            new_rope.append(int(curr_position_id) + 1)
+        generation_input = {'packed_text_ids': torch.zeros(0, dtype=torch.long), 'packed_text_indexes': torch.zeros(0, dtype=torch.long), 'packed_pixel_video_patches': torch.cat(packed_pixel_video_patches, dim=0), 'packed_pixel_video_patch_indexes': torch.tensor(packed_pixel_video_patch_indexes, dtype=torch.long), 'packed_position_ids': torch.cat(packed_position_ids, dim=1), 'packed_seqlens': torch.tensor(packed_seqlens, dtype=torch.int), 'packed_indexes': torch.tensor(packed_indexes, dtype=torch.long), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long), 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int)}
+        return (generation_input, newlens, new_rope)
+
+    @torch.no_grad
+    def forward_cache_update_pixel(self, past_key_values: NaiveCache, packed_text_ids: torch.LongTensor, packed_text_indexes: torch.LongTensor, packed_pixel_patches: Optional[torch.Tensor]=None, packed_pixel_token_indexes: Optional[torch.LongTensor]=None, packed_position_ids: Optional[torch.LongTensor]=None, packed_seqlens: Optional[torch.IntTensor]=None, packed_indexes: Optional[torch.LongTensor]=None, packed_key_value_indexes: Optional[torch.LongTensor]=None, key_values_lens: Optional[torch.IntTensor]=None, packed_pixel_video_patches: Optional[torch.Tensor]=None, packed_pixel_video_patch_indexes: Optional[torch.LongTensor]=None):
+        """Prefill clean pixel patches into KV cache (understanding)."""
+        if packed_pixel_video_patches is not None and not self.config.pixel_video_separate_und_gen_embedder:
+            raise ValueError('This checkpoint profile has no video UND embedder; video-VLM tubes are unsupported')
+        packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
+        packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
+        packed_sequence[packed_text_indexes] = packed_text_embedding
+        model_dtype = packed_text_embedding.dtype
+        if packed_pixel_video_patches is not None:
+            if packed_pixel_video_patch_indexes is None:
+                raise ValueError('Video pixel prefill requires video patch indexes')
+            pixel_embed = self.pixel_video_patch_embed(packed_pixel_video_patches.to(model_dtype))
+            pixel_token_indexes = packed_pixel_video_patch_indexes
+        else:
+            if packed_pixel_patches is None or packed_pixel_token_indexes is None:
+                raise ValueError('Pixel prefill requires image or video pixel patches')
+            pixel_embed = self.pixel_patch_embed(packed_pixel_patches.to(model_dtype))
+            pixel_token_indexes = packed_pixel_token_indexes
+        if pixel_embed.dtype != packed_sequence.dtype:
+            pixel_embed = pixel_embed.to(packed_sequence.dtype)
+        packed_sequence[pixel_token_indexes] = pixel_embed
+        output = self.language_model.forward_inference(packed_query_sequence=packed_sequence, query_lens=packed_seqlens, packed_query_position_ids=packed_position_ids, packed_query_indexes=packed_indexes, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, key_values_lens=key_values_lens, update_past_key_values=True, is_causal=False, mode='und')
+        return output.past_key_values
+
+    def generate_image(self, packed_text_ids: torch.LongTensor, packed_text_indexes: torch.LongTensor, packed_init_noises: torch.Tensor, packed_pixel_token_indexes: torch.LongTensor, packed_seqlens: torch.IntTensor, packed_position_ids: torch.LongTensor, packed_indexes: torch.LongTensor, past_key_values: NaiveCache, key_values_lens: torch.IntTensor, packed_key_value_indexes: torch.LongTensor, num_timesteps: int=24, timestep_shift: float=1.0, sampler: str='dpm-solver', cfg_renorm_min: float=0.0, cfg_renorm_type: str='global', cfg_interval: Optional[Tuple[float, float]]=[0, 1], cfg_text_scale: float=1.0, cfg_text_packed_query_indexes: Optional[torch.LongTensor]=None, cfg_text_packed_position_ids: Optional[torch.LongTensor]=None, cfg_text_past_key_values: Optional[NaiveCache]=None, cfg_text_key_values_lens: Optional[torch.IntTensor]=None, cfg_text_packed_key_value_indexes: Optional[torch.LongTensor]=None, cfg_img_scale: float=1.0, cfg_img_packed_query_indexes: Optional[torch.LongTensor]=None, cfg_img_packed_position_ids: Optional[torch.LongTensor]=None, cfg_img_past_key_values: Optional[NaiveCache]=None, cfg_img_key_values_lens: Optional[torch.IntTensor]=None, cfg_img_packed_key_value_indexes: Optional[torch.LongTensor]=None):
+        x_t = packed_init_noises
+
+        def _v_at(x, t_scalar):
+            """Compute velocity at state x and timestep t_scalar (with CFG gating)."""
+            ts = torch.tensor([t_scalar] * x.shape[0], device=x.device)
+            if t_scalar > cfg_interval[0] and t_scalar <= cfg_interval[1]:
+                (cts, cis) = (cfg_text_scale, cfg_img_scale)
+            else:
+                (cts, cis) = (1.0, 1.0)
+            return self._forward_flow(x_t=x, timestep=ts, packed_pixel_token_indexes=packed_pixel_token_indexes, packed_text_ids=packed_text_ids, packed_text_indexes=packed_text_indexes, packed_position_ids=packed_position_ids, packed_indexes=packed_indexes, packed_seqlens=packed_seqlens, key_values_lens=key_values_lens, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, cfg_renorm_min=cfg_renorm_min, cfg_renorm_type=cfg_renorm_type, cfg_text_scale=cts, cfg_text_packed_position_ids=cfg_text_packed_position_ids, cfg_text_packed_query_indexes=cfg_text_packed_query_indexes, cfg_text_key_values_lens=cfg_text_key_values_lens, cfg_text_past_key_values=cfg_text_past_key_values, cfg_text_packed_key_value_indexes=cfg_text_packed_key_value_indexes, cfg_img_scale=cis, cfg_img_packed_position_ids=cfg_img_packed_position_ids, cfg_img_packed_query_indexes=cfg_img_packed_query_indexes, cfg_img_key_values_lens=cfg_img_key_values_lens, cfg_img_past_key_values=cfg_img_past_key_values, cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes)
+        if sampler != 'dpm-solver':
+            raise ValueError(f"PixelUMM-release T2I requires sampler='dpm-solver', got {sampler!r}")
+        x_t = sample_flow_dpm(_v_at, x_t, num_steps=num_timesteps, shift=timestep_shift)
+        delimiter_tokens = 0
+        generated_lens = packed_seqlens - delimiter_tokens
+        unpacked_latent = x_t.split(generated_lens.tolist())
+        return unpacked_latent
+
+    @torch.no_grad
+    def generate_video(self, packed_text_ids: torch.LongTensor, packed_text_indexes: torch.LongTensor, packed_init_noises: torch.Tensor, packed_pixel_token_indexes: torch.LongTensor, packed_seqlens: torch.IntTensor, packed_position_ids: torch.LongTensor, packed_indexes: torch.LongTensor, past_key_values: NaiveCache, key_values_lens: torch.IntTensor, packed_key_value_indexes: torch.LongTensor, num_timesteps: int=24, timestep_shift: float=1.0, sampler: str='unipc', cfg_renorm_min: float=0.0, cfg_renorm_type: str='global', cfg_interval: Optional[Tuple[float, float]]=[0, 1], cfg_text_scale: float=1.0, cfg_text_packed_query_indexes: Optional[torch.LongTensor]=None, cfg_text_packed_position_ids: Optional[torch.LongTensor]=None, cfg_text_past_key_values: Optional[NaiveCache]=None, cfg_text_key_values_lens: Optional[torch.IntTensor]=None, cfg_text_packed_key_value_indexes: Optional[torch.LongTensor]=None, cfg_img_scale: float=1.0, cfg_img_packed_query_indexes: Optional[torch.LongTensor]=None, cfg_img_packed_position_ids: Optional[torch.LongTensor]=None, cfg_img_past_key_values: Optional[NaiveCache]=None, cfg_img_key_values_lens: Optional[torch.IntTensor]=None, cfg_img_packed_key_value_indexes: Optional[torch.LongTensor]=None):
+        x_t = packed_init_noises
+
+        def _v_at(x, t_scalar):
+            ts = torch.tensor([t_scalar] * x.shape[0], device=x.device)
+            if t_scalar > cfg_interval[0] and t_scalar <= cfg_interval[1]:
+                (cts, cis) = (cfg_text_scale, cfg_img_scale)
+            else:
+                (cts, cis) = (1.0, 1.0)
+            return self._forward_flow(x_t=x, timestep=ts, packed_pixel_token_indexes=packed_pixel_token_indexes, packed_text_ids=packed_text_ids, packed_text_indexes=packed_text_indexes, packed_position_ids=packed_position_ids, packed_indexes=packed_indexes, packed_seqlens=packed_seqlens, key_values_lens=key_values_lens, past_key_values=past_key_values, packed_key_value_indexes=packed_key_value_indexes, cfg_renorm_min=cfg_renorm_min, cfg_renorm_type=cfg_renorm_type, cfg_text_scale=cts, cfg_text_packed_position_ids=cfg_text_packed_position_ids, cfg_text_packed_query_indexes=cfg_text_packed_query_indexes, cfg_text_key_values_lens=cfg_text_key_values_lens, cfg_text_past_key_values=cfg_text_past_key_values, cfg_text_packed_key_value_indexes=cfg_text_packed_key_value_indexes, cfg_img_scale=cis, cfg_img_packed_position_ids=cfg_img_packed_position_ids, cfg_img_packed_query_indexes=cfg_img_packed_query_indexes, cfg_img_key_values_lens=cfg_img_key_values_lens, cfg_img_past_key_values=cfg_img_past_key_values, cfg_img_packed_key_value_indexes=cfg_img_packed_key_value_indexes, use_pixel_video=True)
+        if sampler == 'unipc':
+            x_t = sample_flow_unipc(_v_at, x_t, num_steps=num_timesteps, shift=timestep_shift)
+        elif sampler == 'dpm-solver':
+            if num_timesteps < 3:
+                raise ValueError('T2V DPM-Solver requires at least 3 steps')
+            x_t = sample_flow_dpm(_v_at, x_t, num_steps=num_timesteps, shift=timestep_shift)
+        else:
+            raise ValueError(f"PixelUMM T2V supports 'unipc' or 'dpm-solver', got {sampler!r}")
+        delimiter_tokens = 0
+        generated_lens = packed_seqlens - delimiter_tokens
+        return x_t.split(generated_lens.tolist())
+
+    @torch.no_grad
+    def _forward_flow(self, x_t: torch.Tensor, timestep: torch.LongTensor, packed_pixel_token_indexes: torch.LongTensor, packed_text_ids: torch.LongTensor, packed_text_indexes: torch.LongTensor, packed_indexes: torch.LongTensor, packed_position_ids: torch.LongTensor, packed_seqlens: torch.IntTensor, key_values_lens: torch.IntTensor, past_key_values: NaiveCache, packed_key_value_indexes: torch.LongTensor, cfg_renorm_min: float=0.0, cfg_renorm_type: str='global', cfg_text_scale: float=1.0, cfg_text_packed_position_ids: Optional[torch.LongTensor]=None, cfg_text_packed_query_indexes: Optional[torch.LongTensor]=None, cfg_text_key_values_lens: Optional[torch.Tensor]=None, cfg_text_past_key_values: Optional[NaiveCache]=None, cfg_text_packed_key_value_indexes: Optional[torch.LongTensor]=None, cfg_img_scale: float=1.0, cfg_img_packed_position_ids: Optional[torch.LongTensor]=None, cfg_img_packed_query_indexes: Optional[torch.LongTensor]=None, cfg_img_key_values_lens: Optional[torch.Tensor]=None, cfg_img_past_key_values: Optional[NaiveCache]=None, cfg_img_packed_key_value_indexes: Optional[torch.LongTensor]=None, use_pixel_video: bool=False):
+        packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
+        packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
+        packed_sequence[packed_text_indexes] = packed_text_embedding
+        assert timestep.unique().shape[0] == 1
+        t_val = timestep[0].item()
+        x_t_raw = x_t
+        if use_pixel_video:
+            x_t = self.embed_pixel_video_tubes(x_t)
+        else:
+            image_embedder = self._image_patch_embedder(gen_model=True)
+            x_t = image_embedder(x_t)
+        if x_t.dtype != packed_sequence.dtype:
+            x_t = x_t.to(packed_sequence.dtype)
+        packed_generation_token_indexes = packed_pixel_token_indexes
+        packed_sequence[packed_generation_token_indexes] = x_t
+        extra_inputs = {'mode': 'gen', 'packed_gen_token_indexes': packed_generation_token_indexes, 'packed_und_token_indexes': packed_text_indexes}
+        output = self.language_model.forward_inference(packed_query_sequence=packed_sequence, query_lens=packed_seqlens, packed_query_position_ids=packed_position_ids, packed_query_indexes=packed_indexes, past_key_values=past_key_values, key_values_lens=key_values_lens, packed_key_value_indexes=packed_key_value_indexes, update_past_key_values=False, is_causal=False, **extra_inputs)
+
+        def _pixel_head(hidden_seq):
+            hidden_sub = hidden_seq[packed_pixel_token_indexes]
+            if use_pixel_video:
+                return self.video_unpatchify_head(hidden_sub)
+            return self.unpatchify_head(hidden_sub)
+        model_out = _pixel_head(output.packed_query_sequence)
+        v_t = (x_t_raw.float() - model_out.float()) / max(t_val, self.config.x_pred_t_min)
+        if cfg_text_scale > 1.0:
+            cfg_text_output = self.language_model.forward_inference(packed_query_sequence=packed_sequence, query_lens=packed_seqlens, packed_query_position_ids=cfg_text_packed_position_ids, packed_query_indexes=cfg_text_packed_query_indexes, past_key_values=cfg_text_past_key_values, key_values_lens=cfg_text_key_values_lens, packed_key_value_indexes=cfg_text_packed_key_value_indexes, update_past_key_values=False, is_causal=False, **extra_inputs)
+            cfg_text_out = _pixel_head(cfg_text_output.packed_query_sequence)
+            cfg_text_v_t = (x_t_raw.float() - cfg_text_out.float()) / max(t_val, self.config.x_pred_t_min)
+        if cfg_img_scale > 1.0:
+            cfg_img_output = self.language_model.forward_inference(packed_query_sequence=packed_sequence, query_lens=packed_seqlens, packed_query_position_ids=cfg_img_packed_position_ids, packed_query_indexes=cfg_img_packed_query_indexes, past_key_values=cfg_img_past_key_values, key_values_lens=cfg_img_key_values_lens, packed_key_value_indexes=cfg_img_packed_key_value_indexes, update_past_key_values=False, is_causal=False, **extra_inputs)
+            cfg_img_out = _pixel_head(cfg_img_output.packed_query_sequence)
+            cfg_img_v_t = (x_t_raw.float() - cfg_img_out.float()) / max(t_val, self.config.x_pred_t_min)
+        if cfg_text_scale > 1.0:
+            if cfg_renorm_type == 'text_channel':
+                v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t)
+                norm_v_t = torch.norm(v_t, dim=-1, keepdim=True)
+                norm_v_t_text_ = torch.norm(v_t_text_, dim=-1, keepdim=True)
+                scale = (norm_v_t / (norm_v_t_text_ + 1e-08)).clamp(min=cfg_renorm_min, max=1.0)
+                v_t_text = v_t_text_ * scale
+                if cfg_img_scale > 1.0:
+                    v_t = cfg_img_v_t + cfg_img_scale * (v_t_text - cfg_img_v_t)
+                else:
+                    v_t = v_t_text
+            else:
+                v_t_text_ = cfg_text_v_t + cfg_text_scale * (v_t - cfg_text_v_t)
+                if cfg_img_scale > 1.0:
+                    v_t_ = cfg_img_v_t + cfg_img_scale * (v_t_text_ - cfg_img_v_t)
+                else:
+                    v_t_ = v_t_text_
+                if cfg_renorm_type == 'none':
+                    v_t = v_t_
+                else:
+                    if cfg_renorm_type == 'global':
+                        norm_v_t = torch.norm(v_t)
+                        norm_v_t_ = torch.norm(v_t_)
+                    elif cfg_renorm_type == 'channel':
+                        norm_v_t = torch.norm(v_t, dim=-1, keepdim=True)
+                        norm_v_t_ = torch.norm(v_t_, dim=-1, keepdim=True)
+                    else:
+                        raise NotImplementedError(f'{cfg_renorm_type} is not supported')
+                    scale = (norm_v_t / (norm_v_t_ + 1e-08)).clamp(min=cfg_renorm_min, max=1.0)
+                    v_t = v_t_ * scale
+        return v_t
+
+    def prepare_text_start_tokens(self, curr_kvlens, curr_rope, start_token_ids):
+        (packed_start_tokens, packed_key_value_indexes) = (list(), list())
+        packed_query_position_ids = list()
+        curr = 0
+        for (curr_kvlen, curr_position_id, token_id) in zip(curr_kvlens, curr_rope, start_token_ids):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            packed_start_tokens.append(int(token_id))
+            packed_query_position_ids.append(self._text_mrope_position_ids(curr_position_id, 1))
+            curr += curr_kvlen
+        packed_query_position_ids = torch.cat(packed_query_position_ids, dim=1)
+        generation_input = {'packed_start_tokens': torch.tensor(packed_start_tokens, dtype=torch.long), 'packed_query_position_ids': packed_query_position_ids, 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long)}
+        return generation_input
+
+    def prepare_visual_delimiter_tokens(self, curr_kvlens, curr_rope, new_token_ids, delimiter: str):
+        """Prepare a single causal visual delimiter per sample for pixelumm_mot inference.
+
+        ``<soi>`` and ``<eoi>`` are text embeddings routed through the und expert,
+        but their MRoPE coordinates are the visual marker coordinates. For both
+        SenseNova and Qwen-VL MRoPE, those marker coordinates match
+        ``_text_mrope_position_ids(curr_position_id, 1)``.
+        """
+        if delimiter == 'start':
+            token_id = new_token_ids['start_of_image']
+        elif delimiter == 'end':
+            token_id = new_token_ids['end_of_image']
+        else:
+            raise ValueError(f"Unknown visual delimiter {delimiter!r}; expected 'start' or 'end'")
+        (packed_text_ids, packed_text_indexes) = (list(), list())
+        (packed_text_position_ids, packed_key_value_indexes) = (list(), list())
+        (newlens, new_rope) = (list(), list())
+        curr = 0
+        for (curr_kvlen, curr_position_id) in zip(curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+            packed_text_ids.append(token_id)
+            packed_text_indexes.append(curr)
+            packed_text_position_ids.append(self._text_mrope_position_ids(curr_position_id, 1))
+            curr += 1
+            newlens.append(curr_kvlen + 1)
+            new_rope.append(curr_position_id + 1)
+        packed_text_position_ids = torch.cat(packed_text_position_ids, dim=1)
+        generation_input = {'text_token_lens': torch.ones(len(curr_kvlens), dtype=torch.int), 'packed_text_ids': torch.tensor(packed_text_ids, dtype=torch.long), 'packed_text_position_ids': packed_text_position_ids, 'packed_text_indexes': torch.tensor(packed_text_indexes, dtype=torch.long), 'packed_key_value_indexes': torch.tensor(packed_key_value_indexes, dtype=torch.long), 'key_values_lens': torch.tensor(curr_kvlens, dtype=torch.int)}
+        return (generation_input, newlens, new_rope)
+
+    @torch.no_grad
+    def generate_text(self, past_key_values: NaiveCache, packed_key_value_indexes: torch.LongTensor, key_values_lens: torch.IntTensor, packed_start_tokens: torch.LongTensor, packed_query_position_ids: torch.LongTensor, max_length: int, do_sample: bool=False, temperature: float=1.0, end_token_id: int=None, return_generation_metadata: bool=False):
+        step = 0
+        generated_sequence = []
+        curr_tokens = packed_start_tokens
+        hit_end_token = False
+        while step < max_length:
+            generated_sequence.append(curr_tokens)
+            packed_text_embedding = self.language_model.model.embed_tokens(curr_tokens)
+            query_lens = torch.ones_like(curr_tokens)
+            packed_query_indexes = torch.cumsum(key_values_lens, dim=0) + torch.arange(0, len(key_values_lens), device=key_values_lens.device, dtype=key_values_lens.dtype)
+            uppacked = list(packed_key_value_indexes.split(key_values_lens.tolist(), dim=0))
+            for i in range(len(uppacked)):
+                uppacked[i] += i
+            packed_key_value_indexes = torch.cat(uppacked, dim=0)
+            output = self.language_model.forward_inference(packed_query_sequence=packed_text_embedding, query_lens=query_lens, packed_query_position_ids=packed_query_position_ids, packed_query_indexes=packed_query_indexes, past_key_values=past_key_values, key_values_lens=key_values_lens, packed_key_value_indexes=packed_key_value_indexes, update_past_key_values=True, is_causal=True, mode='und')
+            past_key_values = output.past_key_values
+            packed_query_sequence = output.packed_query_sequence
+            pred_logits = self.language_model.lm_head(packed_query_sequence)
+            if do_sample:
+                probs = nn.functional.softmax(pred_logits / temperature, dim=-1)
+                curr_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+            else:
+                curr_tokens = torch.argmax(pred_logits, dim=-1)
+            uppacked = list(packed_key_value_indexes.split(key_values_lens.tolist(), dim=0))
+            for i in range(len(uppacked)):
+                uppacked[i] = torch.cat([uppacked[i], torch.tensor([uppacked[i][-1] + 1], device=uppacked[i].device)], dim=0)
+            packed_key_value_indexes = torch.cat(uppacked, dim=0)
+            key_values_lens = key_values_lens + 1
+            packed_query_position_ids = self._advance_text_query_position_ids(packed_query_position_ids)
+            step += 1
+            if end_token_id is not None and curr_tokens[0] == end_token_id:
+                hit_end_token = True
+                break
+        output_device = generated_sequence[0].device
+        generated_sequence = torch.stack([i.to(output_device) for i in generated_sequence], dim=0)
+        if return_generation_metadata:
+            return (generated_sequence, {'hit_end_token': hit_end_token})
+        return generated_sequence

modeling/pixelumm/qwen3_navit.py

@@ -0,0 +1,1056 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024 The Qwen Team and The HuggingFace Inc. team.
+# SPDX-FileCopyrightText: Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0
+
+# Qwen3 MoT + NaViT runtime for PixelUMM
+# Based on Bagel/modeling/bagel/qwen2_navit.py (ByteDance, Apache-2.0).
+# PixelUMM keeps only the released SenseNova T/H/W MRoPE topology:
+# upstream FlexAttention for packed training and public FlashAttention varlen
+# for inference.
+
+import copy
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+from torch import nn
+from torch.nn.attention.flex_attention import flex_attention
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.qwen3.modeling_qwen3 import (
+    Qwen3Attention,
+    Qwen3MLP,
+    Qwen3RMSNorm,
+)
+from transformers.utils import ModelOutput
+
+from modeling.qwen3.configuration_qwen3 import Qwen3Config
+
+try:
+    from flash_attn import flash_attn_varlen_func
+except ImportError:
+    flash_attn_varlen_func = None
+
+
+torch._dynamo.config.cache_size_limit = 512
+torch._dynamo.config.accumulated_cache_size_limit = 4096
+from torch._inductor.runtime.hints import TRITON_MAX_BLOCK
+
+TRITON_MAX_BLOCK["X"] = 65536
+# Space patch: flex_attention is only used by the training path, and ZeroGPU
+# does not support torch.compile, so it is left uncompiled here.
+
+def _apply_sequence_chunked(
+    fn,
+    seq: torch.Tensor,
+    chunk_size: int,
+) -> torch.Tensor:
+    """Apply a token-wise module in bounded sequence chunks.
+
+    Qwen3MLP is independent along the sequence dimension, so splitting that
+    dimension preserves its mathematical contract while bounding the transient
+    SwiGLU intermediate from ``sequence_length * intermediate_size`` to
+    ``chunk_size * intermediate_size``. A non-positive chunk size keeps the
+    original path. Empty tensors still call ``fn`` so pixelumm-zero branches keep
+    their zero-gradient autograd edges.
+    """
+    if chunk_size <= 0 or seq.shape[0] <= chunk_size:
+        return fn(seq)
+    return torch.cat(
+        tuple(fn(chunk) for chunk in seq.split(chunk_size, dim=0)),
+        dim=0,
+    )
+
+
+def _sdpa_varlen_attention(
+    query_states, key_states, value_states, query_lengths, key_lengths, *, causal,
+):
+    """Space patch: PyTorch SDPA replacement for flash_attn_varlen_func.
+
+    Matches FlashAttention semantics: GQA by head repetition and bottom-right
+    aligned causal masking when key_length > query_length.
+    """
+    import torch.nn.functional as F
+
+    n_rep = query_states.shape[1] // key_states.shape[1]
+    outputs = []
+    q_start = k_start = 0
+    for q_len, k_len in zip(query_lengths, key_lengths):
+        q = query_states[q_start:q_start + q_len].transpose(0, 1).unsqueeze(0)
+        k = key_states[k_start:k_start + k_len].transpose(0, 1)
+        v = value_states[k_start:k_start + k_len].transpose(0, 1)
+        if n_rep > 1:
+            k = k.repeat_interleave(n_rep, dim=0)
+            v = v.repeat_interleave(n_rep, dim=0)
+        k = k.unsqueeze(0)
+        v = v.unsqueeze(0)
+        mask = None
+        is_causal = False
+        if causal and q_len > 1:
+            if q_len == k_len:
+                is_causal = True
+            else:
+                offset = k_len - q_len
+                rows = torch.arange(q_len, device=q.device).unsqueeze(1)
+                cols = torch.arange(k_len, device=q.device).unsqueeze(0)
+                mask = cols <= rows + offset
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=is_causal)
+        outputs.append(out.squeeze(0).transpose(0, 1))
+        q_start += q_len
+        k_start += k_len
+    return torch.cat(outputs, dim=0)
+
+def _flash_varlen_inference_attention(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    query_lens,
+    key_lens,
+    *,
+    causal: bool,
+) -> torch.Tensor:
+    """Run the public FlashAttention varlen API with packed Q/K/V tensors."""
+
+
+    def _as_int_list(lengths) -> list[int]:
+        if isinstance(lengths, torch.Tensor):
+            lengths = lengths.detach().cpu().tolist()
+        return [int(length) for length in lengths]
+
+    query_lengths = _as_int_list(query_lens)
+    key_lengths = _as_int_list(key_lens)
+    if len(query_lengths) != len(key_lengths):
+        raise ValueError("query_lens and key_lens must contain the same number of samples")
+    if not query_lengths or any(length <= 0 for length in query_lengths + key_lengths):
+        raise ValueError("inference attention segments must be non-empty")
+    if sum(query_lengths) != query_states.shape[0]:
+        raise ValueError("query_lens do not match the flattened query tensor")
+    if sum(key_lengths) != key_states.shape[0] or key_states.shape[0] != value_states.shape[0]:
+        raise ValueError("key_lens do not match the flattened key/value tensors")
+    if causal and any(
+        key_length < query_length
+        for query_length, key_length in zip(query_lengths, key_lengths)
+    ):
+        raise ValueError("causal inference requires key_length >= query_length")
+    if query_states.shape[1] % key_states.shape[1] != 0:
+        raise ValueError(
+            "FlashAttention GQA requires query heads divisible by KV heads: "
+            f"query_heads={query_states.shape[1]} kv_heads={key_states.shape[1]}"
+        )
+
+    def _cumulative_lengths(lengths: list[int]) -> torch.Tensor:
+        cumulative = torch.zeros(
+            len(lengths) + 1,
+            dtype=torch.int32,
+            device=query_states.device,
+        )
+        cumulative[1:] = torch.tensor(
+            lengths,
+            dtype=torch.int32,
+            device=query_states.device,
+        ).cumsum(dim=0)
+        return cumulative
+
+    if flash_attn_varlen_func is None:
+        return _sdpa_varlen_attention(
+            query_states, key_states, value_states,
+            query_lengths, key_lengths, causal=causal,
+        )
+    result = flash_attn_varlen_func(
+        q=query_states.contiguous(),
+        k=key_states.contiguous(),
+        v=value_states.contiguous(),
+        cu_seqlens_q=_cumulative_lengths(query_lengths),
+        cu_seqlens_k=_cumulative_lengths(key_lengths),
+        max_seqlen_q=max(query_lengths),
+        max_seqlen_k=max(key_lengths),
+        causal=causal,
+    )
+    return result[0] if isinstance(result, tuple) else result
+
+class NaiveCache:
+    def __init__(self, num_layers):
+        self.key_cache = {k: None for k in range(num_layers)}
+        self.value_cache = {k: None for k in range(num_layers)}
+
+    @property
+    def num_layers(self):
+        return len(self.key_cache)
+
+@dataclass
+class BaseNavitOutputWithPast(ModelOutput):
+    packed_query_sequence: torch.FloatTensor = None
+    past_key_values: Optional[NaiveCache] = None
+
+def _compute_default_rope_parameters(config, device=None, **_kwargs):
+    """Default RoPE frequencies, copied from SenseNova's Qwen3 fork."""
+    base = config.rope_theta
+    partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0)
+    head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    dim = int(head_dim * partial_rotary_factor)
+    attention_factor = 1.0
+    inv_freq = 1.0 / (
+        base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)
+    )
+    return inv_freq, attention_factor
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input, copied from SenseNova's Qwen3 fork."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Apply RoPE exactly like SenseNova's Qwen3 fork.
+
+    ``position_ids`` is kept for API parity with the fork; it is unused.
+    """
+    del position_ids
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+class SenseNovaQwen3RotaryEmbedding(nn.Module):
+    """Qwen3 rotary embedding copied from SenseNova's fork.
+
+    This mirrors ``references/SenseNova-U1/.../modeling_qwen3.py``. SenseNova
+    computes RoPE frequencies for a 2x wider axis and keeps every other
+    frequency, instead of recomputing a fresh reduced-dim frequency range.
+    """
+
+    inv_freq: torch.Tensor
+
+    @staticmethod
+    def compute_default_rope_parameters(config, device=None, seq_len=None):
+        del seq_len
+        inv_freq, attention_scaling = _compute_default_rope_parameters(config, device)
+
+        cfg2 = copy.deepcopy(config)
+        head_dim = getattr(cfg2, "head_dim", None)
+        if head_dim is None:
+            head_dim = cfg2.hidden_size // cfg2.num_attention_heads
+            setattr(cfg2, "head_dim", head_dim)
+        cfg2.head_dim = int(head_dim) * 2
+
+        inv_freq_full, _ = _compute_default_rope_parameters(cfg2, device)
+        return inv_freq_full[::2], attention_scaling
+
+    def __init__(self, config, device=None):
+        super().__init__()
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.config = config
+
+        if self.rope_type == "default" or self.rope_type is None:
+            base_rope_init_fn = self.compute_default_rope_parameters
+        else:
+            raise ValueError(
+                f"SenseNovaQwen3RotaryEmbedding only supports the default release RoPE, got {self.rope_type}"
+            )
+
+        def _rope_init_fn_keep_freq_range(cfg, dev=None):
+            return base_rope_init_fn(cfg, dev)
+
+        self.rope_init_fn = _rope_init_fn_keep_freq_range
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    @torch.no_grad()
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+def pad_sequence(tensor, pad_size):
+    H, L, D = tensor.shape
+    pad_tensor = tensor.new_zeros((H, pad_size, D))
+    return torch.cat([tensor, pad_tensor], dim=1)
+
+class SenseNovaMRoPEMixin:
+    def _init_sensenova_mrope(self, config):
+        if config.qwen3_mrope_type != "sensenova":
+            raise ValueError("PixelUMM requires qwen3_mrope_type='sensenova'")
+        if self.head_dim % 4 != 0:
+            raise ValueError(
+                "SenseNova split Q/K norm requires head_dim divisible by 4, "
+                f"got {self.head_dim}"
+            )
+        self.qwen3_mrope_type = "sensenova"
+
+        t_config = copy.deepcopy(config)
+        t_config.head_dim = self.head_dim // 2
+        t_config.rope_parameters = copy.deepcopy(config.rope_parameters)
+        self.rotary_emb_t = SenseNovaQwen3RotaryEmbedding(config=t_config)
+
+        hw_config = copy.deepcopy(config)
+        hw_config.head_dim = self.head_dim // 4
+        hw_config.rope_theta = getattr(config, "rope_theta_hw", 10000.0)
+        hw_config.max_position_embeddings = getattr(
+            config, "max_position_embeddings_hw", 10000
+        )
+        hw_config.rope_parameters = copy.deepcopy(config.rope_parameters)
+        hw_config.rope_parameters["rope_theta"] = hw_config.rope_theta
+        self.rotary_emb_hw = SenseNovaQwen3RotaryEmbedding(config=hw_config)
+
+        norm_dim = self.head_dim // 2
+        self.q_norm_t = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+        self.q_norm_hw = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+        self.k_norm_t = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+        self.k_norm_hw = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+
+    def _remove_unused_base_qk_norms(self):
+        num_removed = int(self.q_norm is not None) + int(self.k_norm is not None)
+        self.q_norm = None
+        self.k_norm = None
+        return num_removed
+
+    def _init_sensenova_mot_norms(self, config):
+        norm_dim = self.head_dim // 2
+        self.q_norm_t_moe_gen = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+        self.q_norm_hw_moe_gen = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+        self.k_norm_t_moe_gen = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+        self.k_norm_hw_moe_gen = Qwen3RMSNorm(norm_dim, eps=config.rms_norm_eps)
+
+    @staticmethod
+    def _apply_sensenova_qk_norm(
+        query_states, key_states, q_norm_t, q_norm_hw, k_norm_t, k_norm_hw
+    ):
+        query_states_t, query_states_hw = query_states.chunk(2, dim=-1)
+        key_states_t, key_states_hw = key_states.chunk(2, dim=-1)
+        return (
+            torch.cat(
+                [q_norm_t(query_states_t), q_norm_hw(query_states_hw)], dim=-1
+            ),
+            torch.cat(
+                [k_norm_t(key_states_t), k_norm_hw(key_states_hw)], dim=-1
+            ),
+        )
+
+    def _apply_axis_rope(
+        self, query_states, key_states, rotary_emb, position_ids
+    ):
+        query_states = query_states.transpose(0, 1).unsqueeze(0)
+        key_states = key_states.transpose(0, 1).unsqueeze(0)
+        cos, sin = rotary_emb(query_states, position_ids.unsqueeze(0))
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+        return (
+            query_states.squeeze(0).transpose(0, 1),
+            key_states.squeeze(0).transpose(0, 1),
+        )
+
+    def _apply_sensenova_mrope(
+        self, query_states, key_states, packed_position_ids
+    ):
+        packed_position_ids = packed_position_ids[:, : query_states.shape[0]]
+        if packed_position_ids.ndim != 2 or packed_position_ids.shape[0] != 3:
+            raise ValueError(
+                "SenseNova MRoPE expects T/H/W position IDs shaped (3, seq), "
+                f"got {tuple(packed_position_ids.shape)}"
+            )
+
+        query_t, query_hw = query_states.chunk(2, dim=-1)
+        query_h, query_w = query_hw.chunk(2, dim=-1)
+        key_t, key_hw = key_states.chunk(2, dim=-1)
+        key_h, key_w = key_hw.chunk(2, dim=-1)
+
+        query_t, key_t = self._apply_axis_rope(
+            query_t, key_t, self.rotary_emb_t, packed_position_ids[0]
+        )
+        query_h, key_h = self._apply_axis_rope(
+            query_h, key_h, self.rotary_emb_hw, packed_position_ids[1]
+        )
+        query_w, key_w = self._apply_axis_rope(
+            query_w, key_w, self.rotary_emb_hw, packed_position_ids[2]
+        )
+        return (
+            torch.cat([query_t, query_h, query_w], dim=-1),
+            torch.cat([key_t, key_h, key_w], dim=-1),
+        )
+
+class PackedAttentionMoT(SenseNovaMRoPEMixin, Qwen3Attention):
+    def __init__(self, config, layer_idx: Optional[int] = None):
+        super().__init__(config, layer_idx)
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self._init_sensenova_mrope(config)
+        self._init_sensenova_mot_norms(config)
+
+        self.q_proj_moe_gen = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj_moe_gen = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj_moe_gen = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj_moe_gen = nn.Linear(
+            self.num_heads * self.head_dim,
+            self.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, *args, **kwargs):
+        if self.training:
+            return self.forward_train(*args, **kwargs)
+        return self.forward_inference(*args, **kwargs)
+
+    def forward_train(
+        self,
+        packed_sequence: torch.Tensor,
+        sample_lens: Optional[List[int]],
+        attention_mask,
+        packed_position_embeddings: torch.Tensor,
+        packed_und_token_indexes: torch.LongTensor,
+        packed_gen_token_indexes: torch.LongTensor,
+    ) -> torch.Tensor:
+        del sample_lens
+        if packed_und_token_indexes is None or packed_gen_token_indexes is None:
+            raise ValueError("MoT training requires UND and GEN token indexes")
+
+        query = packed_sequence.new_zeros(
+            packed_sequence.shape[0], self.num_heads * self.head_dim
+        )
+        key = packed_sequence.new_zeros(
+            packed_sequence.shape[0],
+            self.num_key_value_heads * self.head_dim,
+        )
+        value = packed_sequence.new_zeros(
+            packed_sequence.shape[0],
+            self.num_key_value_heads * self.head_dim,
+        )
+        und = packed_sequence[packed_und_token_indexes]
+        gen = packed_sequence[packed_gen_token_indexes]
+        query[packed_und_token_indexes] = self.q_proj(und)
+        query[packed_gen_token_indexes] = self.q_proj_moe_gen(gen)
+        key[packed_und_token_indexes] = self.k_proj(und)
+        key[packed_gen_token_indexes] = self.k_proj_moe_gen(gen)
+        value[packed_und_token_indexes] = self.v_proj(und)
+        value[packed_gen_token_indexes] = self.v_proj_moe_gen(gen)
+
+        query = query.view(-1, self.num_heads, self.head_dim)
+        key = key.view(-1, self.num_key_value_heads, self.head_dim)
+        value = value.view(-1, self.num_key_value_heads, self.head_dim)
+        normalized_query = torch.zeros_like(query)
+        normalized_key = torch.zeros_like(key)
+        normalized_query[packed_und_token_indexes], normalized_key[
+            packed_und_token_indexes
+        ] = self._apply_sensenova_qk_norm(
+            query[packed_und_token_indexes],
+            key[packed_und_token_indexes],
+            self.q_norm_t,
+            self.q_norm_hw,
+            self.k_norm_t,
+            self.k_norm_hw,
+        )
+        normalized_query[packed_gen_token_indexes], normalized_key[
+            packed_gen_token_indexes
+        ] = self._apply_sensenova_qk_norm(
+            query[packed_gen_token_indexes],
+            key[packed_gen_token_indexes],
+            self.q_norm_t_moe_gen,
+            self.q_norm_hw_moe_gen,
+            self.k_norm_t_moe_gen,
+            self.k_norm_hw_moe_gen,
+        )
+        query, key = self._apply_sensenova_mrope(
+            normalized_query,
+            normalized_key,
+            packed_position_embeddings,
+        )
+
+        if not hasattr(attention_mask, "shape"):
+            raise TypeError(
+                "PixelUMM training requires a FlexAttention BlockMask"
+            )
+        pad_size = attention_mask.shape[-2] - query.shape[0]
+        if pad_size < 0:
+            raise ValueError("attention mask is shorter than the packed sequence")
+        query = pad_sequence(query.permute(1, 0, 2), pad_size)
+        key = pad_sequence(key.permute(1, 0, 2), pad_size)
+        value = pad_sequence(value.permute(1, 0, 2), pad_size)
+        output = flex_attention(
+            query.unsqueeze(0),
+            key.unsqueeze(0),
+            value.unsqueeze(0),
+            enable_gqa=True,
+            block_mask=attention_mask,
+        )
+        end_index = output.shape[2] - pad_size if pad_size else output.shape[2]
+        output = output[0, :, :end_index, :]
+        output = output.transpose(0, 1).reshape(
+            -1, self.num_heads * self.head_dim
+        )
+
+        projected = output.new_zeros(output.shape[0], self.hidden_size)
+        projected[packed_und_token_indexes] = self.o_proj(
+            output[packed_und_token_indexes]
+        )
+        projected[packed_gen_token_indexes] = self.o_proj_moe_gen(
+            output[packed_gen_token_indexes]
+        )
+        return projected
+
+    def forward_inference(
+        self,
+        packed_query_sequence: torch.Tensor,
+        query_lens: torch.Tensor,
+        packed_query_position_embeddings: torch.Tensor,
+        packed_query_indexes: torch.Tensor,
+        past_key_values: Optional[NaiveCache] = None,
+        key_values_lens: Optional[torch.Tensor] = None,
+        packed_key_value_indexes: Optional[torch.Tensor] = None,
+        update_past_key_values: bool = True,
+        is_causal: bool = True,
+        mode: str = "und",
+        packed_gen_token_indexes=None,
+        packed_und_token_indexes=None,
+    ):
+        if mode not in {"und", "gen"}:
+            raise ValueError(f"Unsupported inference mode: {mode!r}")
+
+        if mode == "und":
+            query = self.q_proj(packed_query_sequence).view(
+                -1, self.num_heads, self.head_dim
+            )
+            key = self.k_proj(packed_query_sequence).view(
+                -1, self.num_key_value_heads, self.head_dim
+            )
+            value = self.v_proj(packed_query_sequence).view(
+                -1, self.num_key_value_heads, self.head_dim
+            )
+            query, key = self._apply_sensenova_qk_norm(
+                query,
+                key,
+                self.q_norm_t,
+                self.q_norm_hw,
+                self.k_norm_t,
+                self.k_norm_hw,
+            )
+        else:
+            if packed_gen_token_indexes is None or packed_und_token_indexes is None:
+                raise ValueError("GEN inference requires UND and GEN token indexes")
+            sequence = packed_query_sequence.to(torch.bfloat16)
+            query = sequence.new_zeros(
+                sequence.shape[0], self.num_heads * self.head_dim
+            )
+            key = sequence.new_zeros(
+                sequence.shape[0],
+                self.num_key_value_heads * self.head_dim,
+            )
+            value = sequence.new_zeros(
+                sequence.shape[0],
+                self.num_key_value_heads * self.head_dim,
+            )
+            und = sequence[packed_und_token_indexes]
+            gen = sequence[packed_gen_token_indexes]
+            query[packed_und_token_indexes] = self.q_proj(und)
+            query[packed_gen_token_indexes] = self.q_proj_moe_gen(gen)
+            key[packed_und_token_indexes] = self.k_proj(und)
+            key[packed_gen_token_indexes] = self.k_proj_moe_gen(gen)
+            value[packed_und_token_indexes] = self.v_proj(und)
+            value[packed_gen_token_indexes] = self.v_proj_moe_gen(gen)
+            query = query.view(-1, self.num_heads, self.head_dim).float()
+            key = key.view(-1, self.num_key_value_heads, self.head_dim).float()
+            value = value.view(-1, self.num_key_value_heads, self.head_dim)
+            query[packed_und_token_indexes], key[
+                packed_und_token_indexes
+            ] = self._apply_sensenova_qk_norm(
+                query[packed_und_token_indexes],
+                key[packed_und_token_indexes],
+                self.q_norm_t,
+                self.q_norm_hw,
+                self.k_norm_t,
+                self.k_norm_hw,
+            )
+            query[packed_gen_token_indexes], key[
+                packed_gen_token_indexes
+            ] = self._apply_sensenova_qk_norm(
+                query[packed_gen_token_indexes],
+                key[packed_gen_token_indexes],
+                self.q_norm_t_moe_gen,
+                self.q_norm_hw_moe_gen,
+                self.k_norm_t_moe_gen,
+                self.k_norm_hw_moe_gen,
+            )
+
+        query, key = self._apply_sensenova_mrope(
+            query,
+            key,
+            packed_query_position_embeddings,
+        )
+        query = query.to(torch.bfloat16)
+        key = key.to(torch.bfloat16)
+        value = value.to(torch.bfloat16)
+
+        if past_key_values is not None and past_key_values.key_cache[
+            self.layer_idx
+        ] is not None:
+            if key_values_lens is None or packed_key_value_indexes is None:
+                raise ValueError(
+                    "cached inference requires key lengths and KV indexes"
+                )
+            past_key = past_key_values.key_cache[self.layer_idx]
+            past_value = past_key_values.value_cache[self.layer_idx]
+            sequence_length = sum(query_lens) + sum(key_values_lens)
+            merged_key = past_key.new_zeros(
+                sequence_length,
+                self.num_key_value_heads,
+                self.head_dim,
+            )
+            merged_value = past_value.new_zeros(
+                sequence_length,
+                self.num_key_value_heads,
+                self.head_dim,
+            )
+            merged_key[packed_query_indexes] = key
+            merged_key[packed_key_value_indexes] = past_key
+            merged_value[packed_query_indexes] = value
+            merged_value[packed_key_value_indexes] = past_value
+            effective_key_lens = key_values_lens + query_lens
+        else:
+            merged_key = key
+            merged_value = value
+            effective_key_lens = query_lens
+
+        output = _flash_varlen_inference_attention(
+            query,
+            merged_key,
+            merged_value,
+            query_lens,
+            effective_key_lens,
+            causal=is_causal,
+        ).reshape(-1, self.num_heads * self.head_dim)
+        if mode == "und":
+            output = self.o_proj(output)
+        else:
+            projected = output.new_zeros(output.shape[0], self.hidden_size)
+            projected[packed_und_token_indexes] = self.o_proj(
+                output[packed_und_token_indexes]
+            )
+            projected[packed_gen_token_indexes] = self.o_proj_moe_gen(
+                output[packed_gen_token_indexes]
+            )
+            output = projected
+
+        if update_past_key_values:
+            if past_key_values is None:
+                raise ValueError("update_past_key_values requires a NaiveCache")
+            past_key_values.key_cache[self.layer_idx] = merged_key
+            past_key_values.value_cache[self.layer_idx] = merged_value
+        return output, past_key_values
+
+class Qwen3MoTDecoderLayer(nn.Module):
+    def __init__(self, config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.qwen_mlp_sequence_chunk_size = int(
+            getattr(config, "qwen_mlp_sequence_chunk_size", 0)
+        )
+        if self.qwen_mlp_sequence_chunk_size < 0:
+            raise ValueError("qwen_mlp_sequence_chunk_size must be >= 0")
+        self.self_attn = PackedAttentionMoT(config, layer_idx)
+        self.mlp = Qwen3MLP(config)
+        self.mlp_moe_gen = Qwen3MLP(config)
+        self.input_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.input_layernorm_moe_gen = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm_moe_gen = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(self, *args, **kwargs):
+        if self.training:
+            return self.forward_train(*args, **kwargs)
+        return self.forward_inference(*args, **kwargs)
+
+    def forward_train(
+        self,
+        packed_sequence: torch.Tensor,
+        sample_lens: Optional[List[int]],
+        attention_mask,
+        packed_position_embeddings: torch.Tensor,
+        packed_und_token_indexes: torch.LongTensor,
+        packed_gen_token_indexes: torch.LongTensor,
+    ) -> torch.Tensor:
+        if packed_und_token_indexes is None or packed_gen_token_indexes is None:
+            raise ValueError("MoT training requires UND and GEN token indexes")
+
+        residual = packed_sequence
+        normalized = torch.zeros_like(packed_sequence)
+        und_normalized = self.input_layernorm(
+            packed_sequence[packed_und_token_indexes]
+        )
+        gen_normalized = self.input_layernorm_moe_gen(
+            packed_sequence[packed_gen_token_indexes]
+        )
+        normalized[packed_und_token_indexes] = und_normalized
+        normalized[packed_gen_token_indexes] = gen_normalized
+        packed_sequence = residual + self.self_attn(
+            packed_sequence=normalized,
+            sample_lens=sample_lens,
+            attention_mask=attention_mask,
+            packed_position_embeddings=packed_position_embeddings,
+            packed_und_token_indexes=packed_und_token_indexes,
+            packed_gen_token_indexes=packed_gen_token_indexes,
+        )
+
+        residual = packed_sequence
+        mlp_output = torch.zeros_like(packed_sequence)
+        und_mlp = _apply_sequence_chunked(
+            self.mlp,
+            self.post_attention_layernorm(
+                packed_sequence[packed_und_token_indexes]
+            ),
+            self.qwen_mlp_sequence_chunk_size,
+        )
+        gen_mlp = _apply_sequence_chunked(
+            self.mlp_moe_gen,
+            self.post_attention_layernorm_moe_gen(
+                packed_sequence[packed_gen_token_indexes]
+            ),
+            self.qwen_mlp_sequence_chunk_size,
+        )
+        mlp_output[packed_und_token_indexes] = und_mlp
+        mlp_output[packed_gen_token_indexes] = gen_mlp
+        return residual + mlp_output
+
+    def forward_inference(
+        self,
+        packed_query_sequence: torch.Tensor,
+        query_lens: torch.Tensor,
+        packed_query_position_embeddings: torch.Tensor,
+        packed_query_indexes: torch.Tensor,
+        past_key_values: Optional[NaiveCache] = None,
+        key_values_lens: Optional[torch.Tensor] = None,
+        packed_key_value_indexes: Optional[torch.Tensor] = None,
+        update_past_key_values: bool = True,
+        is_causal: bool = True,
+        mode: str = "und",
+        packed_gen_token_indexes=None,
+        packed_und_token_indexes=None,
+    ):
+        if mode == "und":
+            normalized = self.input_layernorm(packed_query_sequence)
+        elif mode == "gen":
+            if packed_gen_token_indexes is None or packed_und_token_indexes is None:
+                raise ValueError("GEN inference requires UND and GEN token indexes")
+            normalized = torch.zeros_like(packed_query_sequence)
+            normalized[packed_und_token_indexes] = self.input_layernorm(
+                packed_query_sequence[packed_und_token_indexes]
+            )
+            normalized[packed_gen_token_indexes] = self.input_layernorm_moe_gen(
+                packed_query_sequence[packed_gen_token_indexes]
+            )
+        else:
+            raise ValueError(f"Unsupported inference mode: {mode!r}")
+
+        residual = packed_query_sequence
+        attention_output, past_key_values = self.self_attn(
+            packed_query_sequence=normalized,
+            query_lens=query_lens,
+            packed_query_position_embeddings=packed_query_position_embeddings,
+            packed_query_indexes=packed_query_indexes,
+            past_key_values=past_key_values,
+            key_values_lens=key_values_lens,
+            packed_key_value_indexes=packed_key_value_indexes,
+            update_past_key_values=update_past_key_values,
+            is_causal=is_causal,
+            mode=mode,
+            packed_gen_token_indexes=packed_gen_token_indexes,
+            packed_und_token_indexes=packed_und_token_indexes,
+        )
+        packed_query_sequence = residual + attention_output
+
+        residual = packed_query_sequence
+        if mode == "und":
+            mlp_output = self.mlp(
+                self.post_attention_layernorm(packed_query_sequence)
+            )
+        else:
+            mlp_output = torch.zeros_like(packed_query_sequence).to(
+                torch.bfloat16
+            )
+            mlp_output[packed_und_token_indexes] = self.mlp(
+                self.post_attention_layernorm(
+                    packed_query_sequence[packed_und_token_indexes]
+                ).to(torch.bfloat16)
+            )
+            mlp_output[packed_gen_token_indexes] = self.mlp_moe_gen(
+                self.post_attention_layernorm_moe_gen(
+                    packed_query_sequence[packed_gen_token_indexes]
+                ).to(torch.bfloat16)
+            )
+        return residual + mlp_output, past_key_values
+
+class Qwen3PreTrainedModel(PreTrainedModel):
+    config_class = Qwen3Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["Qwen3MoTDecoderLayer"]
+
+
+class Qwen3Model(Qwen3PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        if config.layer_module != "Qwen3MoTDecoderLayer":
+            raise ValueError("PixelUMM requires Qwen3MoTDecoderLayer")
+        if config.qwen3_mrope_type != "sensenova":
+            raise ValueError("PixelUMM requires qwen3_mrope_type='sensenova'")
+        if config.freeze_und:
+            raise ValueError("PixelUMM release does not support freeze_und")
+        if config.conditional_branch_gradient_contract != "pixelumm_zero_v1":
+            raise ValueError(
+                "PixelUMM requires conditional_branch_gradient_contract="
+                "'pixelumm_zero_v1'"
+            )
+
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                Qwen3MoTDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.norm_moe_gen = Qwen3RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_init()
+
+    def forward(self, *args, **kwargs):
+        if self.training:
+            return self.forward_train(*args, **kwargs)
+        return self.forward_inference(*args, **kwargs)
+
+    def forward_train(
+        self,
+        packed_sequence: torch.Tensor,
+        sample_lens: List[int],
+        attention_mask,
+        packed_position_ids: torch.Tensor,
+        packed_und_token_indexes: Optional[torch.LongTensor] = None,
+        packed_gen_token_indexes: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        if packed_position_ids.ndim != 2 or packed_position_ids.shape[0] != 3:
+            raise ValueError(
+                "PixelUMM requires T/H/W position IDs shaped (3, seq), "
+                f"got {tuple(packed_position_ids.shape)}"
+            )
+        if packed_und_token_indexes is None:
+            raise ValueError("packed_und_token_indexes is required")
+        if packed_gen_token_indexes is None:
+            packed_gen_token_indexes = packed_und_token_indexes.new_empty((0,))
+        position_ids = packed_position_ids[
+            :, : packed_sequence.shape[0]
+        ].contiguous()
+
+        for layer in self.layers:
+            packed_sequence = layer(
+                packed_sequence=packed_sequence,
+                sample_lens=None,
+                attention_mask=attention_mask,
+                packed_position_embeddings=position_ids,
+                packed_und_token_indexes=packed_und_token_indexes,
+                packed_gen_token_indexes=packed_gen_token_indexes,
+            )
+
+        output = torch.zeros_like(packed_sequence)
+        output[packed_und_token_indexes] = self.norm(
+            packed_sequence[packed_und_token_indexes]
+        )
+        output[packed_gen_token_indexes] = self.norm_moe_gen(
+            packed_sequence[packed_gen_token_indexes]
+        )
+        return output
+
+    def forward_inference(
+        self,
+        packed_query_sequence: torch.Tensor,
+        query_lens: torch.Tensor,
+        packed_query_position_ids: torch.Tensor,
+        packed_query_indexes: torch.Tensor,
+        past_key_values: Optional[NaiveCache] = None,
+        key_values_lens: Optional[torch.Tensor] = None,
+        packed_key_value_indexes: Optional[torch.Tensor] = None,
+        update_past_key_values: bool = True,
+        is_causal: bool = True,
+        mode: str = "und",
+        packed_gen_token_indexes=None,
+        packed_und_token_indexes=None,
+    ) -> BaseNavitOutputWithPast:
+        if (
+            packed_query_position_ids.ndim != 2
+            or packed_query_position_ids.shape[0] != 3
+        ):
+            raise ValueError(
+                "PixelUMM requires T/H/W query position IDs shaped (3, seq), "
+                f"got {tuple(packed_query_position_ids.shape)}"
+            )
+        position_ids = packed_query_position_ids[
+            :, : packed_query_sequence.shape[0]
+        ].contiguous()
+        extra_inputs = {"mode": mode}
+        if mode == "gen":
+            if packed_gen_token_indexes is None or packed_und_token_indexes is None:
+                raise ValueError("GEN inference requires UND and GEN token indexes")
+            extra_inputs.update(
+                packed_gen_token_indexes=packed_gen_token_indexes,
+                packed_und_token_indexes=packed_und_token_indexes,
+            )
+        elif mode != "und":
+            raise ValueError(f"Unsupported inference mode: {mode!r}")
+
+        for layer in self.layers:
+            packed_query_sequence, past_key_values = layer(
+                packed_query_sequence=packed_query_sequence,
+                query_lens=query_lens,
+                packed_query_position_embeddings=position_ids,
+                packed_query_indexes=packed_query_indexes,
+                past_key_values=past_key_values,
+                key_values_lens=key_values_lens,
+                packed_key_value_indexes=packed_key_value_indexes,
+                update_past_key_values=update_past_key_values,
+                is_causal=is_causal,
+                **extra_inputs,
+            )
+
+        if mode == "und":
+            packed_query_sequence = self.norm(packed_query_sequence)
+        else:
+            output = torch.zeros_like(packed_query_sequence)
+            output[packed_und_token_indexes] = self.norm(
+                packed_query_sequence[packed_und_token_indexes]
+            )
+            output[packed_gen_token_indexes] = self.norm_moe_gen(
+                packed_query_sequence[packed_gen_token_indexes]
+            )
+            packed_query_sequence = output
+        return BaseNavitOutputWithPast(
+            packed_query_sequence=packed_query_sequence,
+            past_key_values=past_key_values,
+        )
+
+
+class Qwen3ForCausalLM(Qwen3PreTrainedModel):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Qwen3Model(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(
+            config.hidden_size, config.vocab_size, bias=False
+        )
+        self.post_init()
+
+    def remove_unused_sensenova_base_qk_norms(self):
+        num_removed = 0
+        for layer in self.model.layers:
+            num_removed += layer.self_attn._remove_unused_base_qk_norms()
+        return num_removed
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(self, *args, **kwargs):
+        if self.training:
+            return self.forward_train(*args, **kwargs)
+        return self.forward_inference(*args, **kwargs)
+
+    def forward_train(
+        self,
+        packed_sequence: torch.Tensor,
+        sample_lens: List[int],
+        attention_mask,
+        packed_position_ids: torch.Tensor,
+        packed_und_token_indexes: Optional[torch.LongTensor] = None,
+        packed_gen_token_indexes: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        return self.model(
+            packed_sequence=packed_sequence,
+            sample_lens=sample_lens,
+            packed_position_ids=packed_position_ids,
+            attention_mask=attention_mask,
+            packed_und_token_indexes=packed_und_token_indexes,
+            packed_gen_token_indexes=packed_gen_token_indexes,
+        )
+
+    def forward_inference(
+        self,
+        packed_query_sequence: torch.Tensor,
+        query_lens: torch.Tensor,
+        packed_query_position_ids: torch.Tensor,
+        packed_query_indexes: torch.Tensor,
+        past_key_values: Optional[NaiveCache] = None,
+        key_values_lens: Optional[torch.Tensor] = None,
+        packed_key_value_indexes: Optional[torch.Tensor] = None,
+        update_past_key_values: bool = True,
+        is_causal: bool = True,
+        mode: str = "und",
+        packed_gen_token_indexes=None,
+        packed_und_token_indexes=None,
+    ) -> BaseNavitOutputWithPast:
+        return self.model(
+            packed_query_sequence=packed_query_sequence,
+            query_lens=query_lens,
+            packed_query_position_ids=packed_query_position_ids,
+            packed_query_indexes=packed_query_indexes,
+            past_key_values=past_key_values,
+            key_values_lens=key_values_lens,
+            packed_key_value_indexes=packed_key_value_indexes,
+            update_past_key_values=update_past_key_values,
+            is_causal=is_causal,
+            mode=mode,
+            packed_gen_token_indexes=packed_gen_token_indexes,
+            packed_und_token_indexes=packed_und_token_indexes,
+        )

modeling/pixelumm/time_schedule.py

@@ -0,0 +1,57 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Released PixelUMM noise-time sampling and spatial shift."""
+
+from __future__ import annotations
+
+import torch
+
+
+def sample_logit_normal_noise_time(
+    raw: torch.Tensor,
+    *,
+    logit_mean: float,
+    logit_std: float,
+) -> torch.Tensor:
+    """Map N(0,1) variates to PixelUMM noise time; -inf means clean."""
+
+    if logit_std <= 0:
+        raise ValueError(f"logit_std must be positive, got {logit_std}")
+    noise_time = torch.sigmoid(raw * float(logit_std) + float(logit_mean))
+    return torch.where(torch.isneginf(raw), torch.zeros_like(noise_time), noise_time)
+
+
+def pixelumm_resolution_shift(height: int, width: int) -> float:
+    """Select the released spatial shift from the target short edge."""
+
+    height = int(height)
+    width = int(width)
+    if height <= 0 or width <= 0:
+        raise ValueError(f"Spatial dimensions must be positive, got {(height, width)}")
+    short_edge = min(height, width)
+    if short_edge <= 256:
+        return 1.0
+    if short_edge <= 640:
+        return 2.0
+    if short_edge <= 960:
+        return 3.0
+    raise ValueError(
+        "PixelUMM resolution shift supports short edges up to 960 px; "
+        f"got shape {(height, width)}"
+    )
+
+
+def flow_shift_noise_timestep(t: torch.Tensor, shift: float) -> torch.Tensor:
+    """Shift PixelUMM noise time toward noise for shifts greater than one."""
+
+    if shift <= 0:
+        raise ValueError(f"shift must be positive, got {shift}")
+    return shift * t / (1 + (shift - 1) * t)
+
+
+__all__ = [
+    "flow_shift_noise_timestep",
+    "pixelumm_resolution_shift",
+    "sample_logit_normal_noise_time",
+]

modeling/qwen3/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from .configuration_qwen3 import Qwen3Config
+
+__all__ = ["Qwen3Config"]

modeling/qwen3/configuration_qwen3.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Qwen3 configuration for PixelUMM — standalone, no HF transformers dependency."""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Qwen3Config(PretrainedConfig):
+    """
+    Qwen3-0.6B config with MoT extensions.
+
+    Qwen3-0.6B reference values:
+        hidden_size=1024, intermediate_size=3072, num_hidden_layers=28,
+        num_attention_heads=16, num_key_value_heads=8, head_dim=128,
+        attention_bias=False, rope_theta=1e6, vocab_size=151936
+    """
+
+    model_type = "qwen3"
+
+    def __init__(
+        self,
+        vocab_size=151936,
+        hidden_size=1024,
+        intermediate_size=3072,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        head_dim=128,
+        hidden_act="silu",
+        max_position_embeddings=40960,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        tie_word_embeddings=True,
+        rope_theta=1000000.0,
+        rope_parameters=None,  # dict with rope_type etc; None → default
+        attention_bias=False,
+        attention_dropout=0.0,
+        # Hugging Face Qwen compatibility fields.
+        # use_sliding_window，sliding_window，max_window_layers，layer_types
+        # PixelUMM's qwen3_navit path currently keeps full attention and does not consume
+        # sliding-window attention, but newer transformer configs may define them.
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        layer_types=None,
+        pad_token_id=None,
+        bos_token_id=151643,
+        eos_token_id=151645,
+        # MoT extensions (same as Qwen2 version)
+        layer_module="Qwen3MoTDecoderLayer",
+        freeze_und=False,
+        qk_norm=True,  # always True for Qwen3
+        # PixelUMM's released SenseNova T/H/W MRoPE topology.
+        qwen3_mrope_type="sensenova",
+        mrope_section=None,
+        rope_theta_hw=10000.0,
+        max_position_embeddings_hw=10000,
+        inference_attention_backend="flash_attn_varlen",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.tie_word_embeddings = tie_word_embeddings
+        self.rope_theta = rope_theta
+        # MRoPE section (set before rope_parameters which needs it)
+        if mrope_section is None:
+            self.mrope_section = [24, 20, 20]  # default for head_dim=128
+        else:
+            self.mrope_section = mrope_section
+        self.qwen3_mrope_type = qwen3_mrope_type
+        self.rope_theta_hw = rope_theta_hw
+        self.max_position_embeddings_hw = max_position_embeddings_hw
+        # rope_parameters: ensure all fields needed by Qwen3 rotary and local interleaved MRoPE.
+        if rope_parameters is None:
+            rope_parameters = {"rope_type": "default"}
+        rope_parameters.setdefault("rope_theta", rope_theta)
+        rope_parameters.setdefault("mrope_section", self.mrope_section)
+        self.rope_parameters = rope_parameters
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if self.use_sliding_window else None
+        self.max_window_layers = max_window_layers
+        if layer_types is None:
+            layer_types = [
+                "sliding_attention"
+                if self.sliding_window is not None and i >= self.max_window_layers
+                else "full_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+        self.layer_types = layer_types
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.layer_module = layer_module
+        self.freeze_und = freeze_und
+        self.qk_norm = qk_norm
+        self.inference_attention_backend = str(inference_attention_backend)
+        if self.inference_attention_backend != "flash_attn_varlen":
+            raise ValueError(
+                "inference_attention_backend must be 'flash_attn_varlen'; "
+                f"got {self.inference_attention_backend!r}"
+            )

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -0,0 +1,12 @@
+torchvision==0.28.0
+torchcodec==0.16.0
+transformers==4.57.1
+huggingface_hub>=0.34.0,<1.0
+accelerate
+einops==0.8.1
+numpy>=2.0,<2.3
+Pillow
+PyYAML
+safetensors
+tqdm
+imageio-ffmpeg==0.6.0

train/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2025 Bytedance Ltd. and/or its affiliates.
+# SPDX-License-Identifier: Apache-2.0

train/checkpoint_export.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Integrity receipts for static DCP exports without a training completion marker.
+
+This verifies locally supplied model files, not upstream training completion,
+publisher authenticity, optimizer resumability, or model/config compatibility.
+Only prepare receipts for a trusted, fully downloaded, immutable export.
+"""
+
+import hashlib
+import json
+from pathlib import Path, PurePosixPath
+
+EXPORT_MANIFEST = "checkpoint_export.json"
+EXPORT_SCHEMA = "pixelumm_dcp_export_v1"
+
+
+def _model_file(root: Path, name: str) -> Path:
+    if not isinstance(name, str):
+        raise ValueError("Export file path must be a string")
+    parts = PurePosixPath(name).parts
+    if (
+        len(parts) != 2 or parts[0] != "model"
+        or name != "/".join(parts) or "\\" in name
+        or (parts[1] != ".metadata" and not parts[1].endswith(".distcp"))
+    ):
+        raise ValueError(f"Invalid export model path: {name!r}")
+    path = root / name
+    if path.resolve().parent != (root / "model").resolve() or path.is_symlink():
+        raise ValueError(f"Export model path must stay inside model/: {name!r}")
+    return path
+
+
+def _digest(path: Path) -> str:
+    before = path.stat()
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""):
+            digest.update(chunk)
+    after = path.stat()
+    if (before.st_size, before.st_mtime_ns) != (after.st_size, after.st_mtime_ns):
+        raise RuntimeError(f"Export changed during verification: {path}")
+    return digest.hexdigest()
+
+
+def validate_export_manifest(root: Path) -> set[str]:
+    root = Path(root)
+    payload = json.loads((root / EXPORT_MANIFEST).read_text(encoding="utf-8"))
+    if not isinstance(payload, dict) or payload.get("schema") != EXPORT_SCHEMA:
+        raise ValueError("Unsupported checkpoint export manifest schema")
+    entries = payload.get("files")
+    if not isinstance(entries, list) or not entries:
+        raise ValueError("Checkpoint export manifest has no files")
+    names = set()
+    for entry in entries:
+        if not isinstance(entry, dict) or set(entry) != {"path", "bytes", "sha256"}:
+            raise ValueError("Malformed checkpoint export file entry")
+        name = entry["path"]
+        path = _model_file(root, name)
+        if name in names:
+            raise ValueError(f"Duplicate export file: {name}")
+        names.add(name)
+        if type(entry["bytes"]) is not int or entry["bytes"] <= 0:
+            raise ValueError(f"Invalid export file size: {name}")
+        expected = entry["sha256"]
+        if not isinstance(expected, str) or len(expected) != 64 or any(c not in "0123456789abcdef" for c in expected):
+            raise ValueError(f"Invalid export SHA256: {name}")
+        if not path.is_file() or path.stat().st_size != entry["bytes"]:
+            raise RuntimeError(f"Missing or wrong-size export file: {name}")
+        if _digest(path) != expected:
+            raise RuntimeError(f"Export SHA256 mismatch: {name}")
+    if "model/.metadata" not in names:
+        raise ValueError("Export manifest must include model/.metadata")
+    return names
+
+
+def prepare_export_manifest(root: Path) -> Path:
+    from train.release_checkpoint import _validate_dcp_model
+
+    root = Path(root).expanduser().resolve()
+    output = root / EXPORT_MANIFEST
+    if output.exists():
+        raise FileExistsError(f"Preserving existing export manifest: {output}")
+    required = _validate_dcp_model(root)
+    if (root / "model.safetensors").exists():
+        raise ValueError("Prepare a DCP-only export, without model.safetensors")
+    files = []
+    for name in sorted(required):
+        path = _model_file(root, name)
+        files.append({"path": name, "bytes": path.stat().st_size, "sha256": _digest(path)})
+    payload = {
+        "schema": EXPORT_SCHEMA,
+        "verification": "local DCP file integrity; not a training completion assertion",
+        "files": files,
+    }
+    with output.open("x", encoding="utf-8") as stream:
+        stream.write(json.dumps(payload, indent=2) + "\n")
+    return output

train/config.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Typed configuration for the released F18 and F22 PixelUMM profiles."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, fields, replace
+from pathlib import Path
+from typing import Optional, TypeVar
+
+from train.scientific_config import load_scientific_config
+
+
+@dataclass(frozen=True)
+class ModelArguments:
+    llm_path: str = "Qwen/Qwen3-8B"
+    llm_qk_norm: bool = True
+    llm_backend: str = "qwen3"
+    qwen3_mrope_type: str = "sensenova"
+    qwen3_mrope_section: str = ""
+    qwen3_rope_theta: float = 1_000_000.0
+    tie_word_embeddings: bool = False
+    layer_module: str = "Qwen3MoTDecoderLayer"
+    pixel_token_patch_size: int = 16
+    pixel_embedder_type: str = "image_raw_patch_linear"
+    pixel_image_additive_pos_embed_type: str = "none"
+    pixel_video_additive_pos_embed_type: str = "none"
+    sensenova_vision_rope_theta: float = 10_000.0
+    sensenova_vision_max_position_embeddings: int = 10_000
+    add_timestep_embedding: bool = False
+    add_noise_scale_embedding: bool = False
+    noise_scale: float = 1.0
+    noise_scale_mode: str = "constant"
+    enable_pixel_video: bool = True
+    pixel_video_temporal_patch_size: int = 4
+    pixel_video_embedder_type: str = "video_raw_tube_linear"
+    pixel_separate_und_gen_embedder: bool = True
+    pixel_video_separate_und_gen_embedder: bool = True
+    packed_attention_impl: str = "flex"
+    inference_attention_backend: str = "flash_attn_varlen"
+    packed_expert_routing: str = "pixelumm_mot"
+    pixel_head_type: str = "minit2i_linear"
+    pixel_video_head_type: str = "jit_style"
+    qwen_mlp_sequence_chunk_size: int = 32_768
+    ce_loss_checkpoint_chunk_size: int = 4096
+    text_cond_dropout_prob: float = 0.1
+    pixel_gen_cond_dropout_prob: float = 0.1
+    pixel_und_cond_dropout_prob: float = 0.0
+
+
+@dataclass(frozen=True)
+class DataArguments:
+    packed_sequence_layout: str = "pixelumm_mot"
+
+
+@dataclass(frozen=True)
+class TrainingArguments:
+    gen_space: str = "pixel"
+    und_space: str = "pixel"
+    prediction_type: str = "x"
+    loss_type: str = "v_loss"
+    x_pred_t_min: float = 0.05
+    visual_gen: bool = True
+    visual_und: bool = True
+    freeze_und: bool = False
+    conditional_branch_gradient_contract: str = "pixelumm_zero_v1"
+    noise_time_logit_mean: float = 0.0
+    noise_time_logit_std: float = 1.0
+    time_schedule: str = "pixelumm_resolution"
+    mse_loss_reduction: str = "sample"
+    eval_num_timesteps: int = 50
+    eval_sampler: str = "dpm-solver"
+    eval_timestep_shift: float = 3.0
+    eval_cfg_text_scale: float = 3.5
+    eval_cfg_renorm_type: str = "none"
+    eval_video_sampler: str = "unipc"
+    eval_video_timestep_shift: float = 10.0
+    eval_video_cfg_text_scale: float = 6.0
+
+
+T = TypeVar("T")
+
+
+def _instantiate(cls: type[T], values: dict, *, section: str) -> T:
+    allowed = {item.name for item in fields(cls)}
+    unknown = sorted(set(values) - allowed)
+    if unknown:
+        raise RuntimeError(
+            f"release config has unsupported {section} fields: {unknown}"
+        )
+    return cls(**values)
+
+
+def load_release_config(
+    path: str | Path,
+    *,
+    llm_path: str | None = None,
+) -> tuple[ModelArguments, DataArguments, TrainingArguments]:
+    sections = load_scientific_config(Path(path))
+    model_args = _instantiate(ModelArguments, sections["model"], section="model")
+    data_args = _instantiate(DataArguments, sections["data"], section="data")
+    training_args = _instantiate(
+        TrainingArguments, sections["training"], section="training"
+    )
+    if llm_path is not None:
+        model_args = replace(model_args, llm_path=str(llm_path))
+    return model_args, data_args, training_args
+
+
+def parse_mrope_section(value: Optional[str]) -> Optional[list[int]]:
+    if value is None or not value.strip():
+        return None
+    section = [int(item.strip()) for item in value.strip("[]()").split(",") if item.strip()]
+    if len(section) != 3:
+        raise ValueError(f"Expected qwen3_mrope_section T,H,W, got {value!r}")
+    return section
+
+
+def configure_qwen3_mrope_config(llm_config, model_args: ModelArguments) -> None:
+    llm_config.qwen3_mrope_type = model_args.qwen3_mrope_type
+    if not getattr(llm_config, "rope_parameters", None):
+        llm_config.rope_parameters = {}
+    if model_args.qwen3_rope_theta > 0:
+        theta = float(model_args.qwen3_rope_theta)
+        llm_config.rope_theta = theta
+        llm_config.rope_parameters["rope_theta"] = theta
+    section = parse_mrope_section(model_args.qwen3_mrope_section)
+    if section is None:
+        section = getattr(llm_config, "mrope_section", None)
+    if section is None:
+        section = llm_config.rope_parameters.get("mrope_section")
+    if section is not None:
+        llm_config.mrope_section = section
+        llm_config.rope_parameters["mrope_section"] = section
+        head_dim = getattr(
+            llm_config,
+            "head_dim",
+            llm_config.hidden_size // llm_config.num_attention_heads,
+        )
+        rotary_dim = 2 * sum(section)
+        if rotary_dim > head_dim:
+            raise ValueError(
+                f"mrope section {section} gives rotary_dim={rotary_dim} > head_dim={head_dim}"
+            )
+        llm_config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim

train/eval_utils.py

@@ -0,0 +1,875 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""PixelUMM image, video, and multimodal inference helpers."""
+
+from __future__ import annotations
+
+import re
+from copy import deepcopy
+from pathlib import Path
+
+import imageio_ffmpeg
+import numpy as np
+import torch
+from PIL import Image
+
+from data.data_utils import untubeify
+
+
+def save_video_mp4(frames, output_dir, fps=16, filename="t2v.mp4"):
+    if not frames:
+        return None
+    path = str(Path(output_dir) / filename)
+    width, height = frames[0].size
+    fps = max(float(fps), 1.0)
+    writer = imageio_ffmpeg.write_frames(
+        path,
+        (width, height),
+        fps=fps,
+        codec="libx264",
+        pix_fmt_in="rgb24",
+        pix_fmt_out="yuv420p",
+        output_params=["-movflags", "+faststart"],
+        ffmpeg_log_level="error",
+    )
+    writer.send(None)
+    try:
+        for frame in frames:
+            writer.send(np.asarray(frame.convert("RGB")))
+    finally:
+        writer.close()
+    return path
+
+
+def _video_tubes_to_pil_frames(video_tubes, image_size, num_frames, patch_size):
+    if not torch.isfinite(video_tubes).all().item():
+        raise RuntimeError("Video generation produced non-finite pixels")
+    frames = untubeify(video_tubes, image_size, num_frames, patch_size)
+    frames = ((frames.clamp(-1, 1) + 1) / 2 * 255).to(torch.uint8).numpy()
+    return [Image.fromarray(frames[i]) for i in range(num_frames)]
+
+
+def _prefill_text_to_visual_chatml_context(
+    model, tokenizer, new_token_ids, prompt, device, kv, kv_lens, ropes,
+):
+    """Prefill the exact T2I/T2V training ChatML context into a KV cache."""
+    text_input, kv_lens, ropes = model.prepare_text_to_visual_chatml_prompts(
+        curr_kvlens=kv_lens,
+        curr_rope=ropes,
+        prompts=[prompt],
+        tokenizer=tokenizer,
+        new_token_ids=new_token_ids,
+    )
+    text_input = {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in text_input.items()
+    }
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        kv = model.forward_cache_update_text(kv, **text_input)
+    return kv, kv_lens, ropes
+
+
+def _prefill_generation_visual_start(
+    model, new_token_ids, device, kv, kv_lens, ropes,
+):
+    """Cache the causal ``<vision_start>`` used by PixelUMM generation."""
+    start_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+        kv_lens, ropes, new_token_ids, delimiter="start",
+    )
+    start_input = {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in start_input.items()
+    }
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        kv = model.forward_cache_update_text(kv, **start_input)
+    return kv, kv_lens, ropes, True
+
+
+def _generate_t2i_one(
+    model, tokenizer, new_token_ids, prompt, image_size, device,
+    num_timesteps, timestep_shift, sampler,
+    cfg_text_scale, cfg_img_scale, cfg_interval, cfg_renorm_type,
+):
+    """Run one T2I sample through the FSDP-wrapped model. Mirrors
+    test_inference_pixel.py:generate_t2i. Returns a PIL image."""
+    from modeling.pixelumm.qwen3_navit import NaiveCache
+
+    n_layers = model.config.llm_config.num_hidden_layers
+    kv = NaiveCache(n_layers)
+    kv_lens, ropes = [0], [0]
+
+    cfg_text_kv = deepcopy(kv) if cfg_text_scale > 1.0 else None
+    cfg_text_kv_lens = list(kv_lens) if cfg_text_scale > 1.0 else None
+    cfg_text_ropes = list(ropes) if cfg_text_scale > 1.0 else None
+    cfg_img_kv = deepcopy(kv) if cfg_img_scale > 1.0 else None
+    cfg_img_kv_lens = list(kv_lens) if cfg_img_scale > 1.0 else None
+    cfg_img_ropes = list(ropes) if cfg_img_scale > 1.0 else None
+
+    kv, kv_lens, ropes = _prefill_text_to_visual_chatml_context(
+        model, tokenizer, new_token_ids, prompt, device, kv, kv_lens, ropes,
+    )
+
+    kv, kv_lens, ropes, visual_start_prefilled = _prefill_generation_visual_start(
+        model, new_token_ids, device, kv, kv_lens, ropes,
+    )
+
+    generation_input = model.prepare_pixel_noise(
+        curr_kvlens=kv_lens, curr_rope=ropes,
+        image_sizes=[image_size],
+        visual_start_prefilled=visual_start_prefilled,
+    )
+    generation_input = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                        for k, v in generation_input.items()}
+    # Preserve the full-precision RNG sample. Model forward remains under BF16
+    # autocast, while every sampler starts from an FP32 numerical state.
+    generation_input['packed_init_noises'] = generation_input['packed_init_noises'].float()
+
+    cfg_kwargs = {}
+    if cfg_text_scale > 1.0 and cfg_text_kv is not None:
+        # CFG text-uncond reference must match training's ChatML shell dropout, NOT an
+        # empty prefill. Training drops the caption to cfg_dropout_input_ids =
+        # <|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n[EOS] (an empty caption
+        # through the same ChatML builder). Feed prompts=[""] so the uncond context the
+        # model denoises from equals what it saw as the dropout target.
+        cfg_text_kv, cfg_text_kv_lens, cfg_text_ropes = _prefill_text_to_visual_chatml_context(
+            model, tokenizer, new_token_ids, "", device,
+            cfg_text_kv, cfg_text_kv_lens, cfg_text_ropes,
+        )
+        (
+            cfg_text_kv,
+            cfg_text_kv_lens,
+            cfg_text_ropes,
+            cfg_text_start_prefilled,
+        ) = _prefill_generation_visual_start(
+            model, new_token_ids, device,
+            cfg_text_kv, cfg_text_kv_lens, cfg_text_ropes,
+        )
+        cfg_text_input = model.prepare_pixel_noise(
+            curr_kvlens=cfg_text_kv_lens, curr_rope=cfg_text_ropes,
+            image_sizes=[image_size],
+            visual_start_prefilled=cfg_text_start_prefilled,
+        )
+        cfg_text_input = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                          for k, v in cfg_text_input.items()}
+        cfg_kwargs.update(
+            cfg_text_past_key_values=cfg_text_kv,
+            cfg_text_packed_position_ids=cfg_text_input['packed_position_ids'],
+            cfg_text_packed_query_indexes=cfg_text_input['packed_indexes'],
+            cfg_text_key_values_lens=cfg_text_input['key_values_lens'],
+            cfg_text_packed_key_value_indexes=cfg_text_input['packed_key_value_indexes'],
+        )
+    if cfg_img_scale > 1.0 and cfg_img_kv is not None:
+        (
+            cfg_img_kv,
+            cfg_img_kv_lens,
+            cfg_img_ropes,
+            cfg_img_start_prefilled,
+        ) = _prefill_generation_visual_start(
+            model, new_token_ids, device,
+            cfg_img_kv, cfg_img_kv_lens, cfg_img_ropes,
+        )
+        cfg_img_input = model.prepare_pixel_noise(
+            curr_kvlens=cfg_img_kv_lens, curr_rope=cfg_img_ropes,
+            image_sizes=[image_size],
+            visual_start_prefilled=cfg_img_start_prefilled,
+        )
+        cfg_img_input = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                         for k, v in cfg_img_input.items()}
+        cfg_kwargs.update(
+            cfg_img_past_key_values=cfg_img_kv,
+            cfg_img_packed_position_ids=cfg_img_input['packed_position_ids'],
+            cfg_img_packed_query_indexes=cfg_img_input['packed_indexes'],
+            cfg_img_key_values_lens=cfg_img_input['key_values_lens'],
+            cfg_img_packed_key_value_indexes=cfg_img_input['packed_key_value_indexes'],
+        )
+
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        unpacked_pixels = model.generate_image(
+            past_key_values=kv,
+            num_timesteps=num_timesteps,
+            timestep_shift=timestep_shift,
+            sampler=sampler,
+            cfg_text_scale=cfg_text_scale,
+            cfg_img_scale=cfg_img_scale,
+            cfg_interval=cfg_interval,
+            cfg_renorm_min=0.0,
+            cfg_renorm_type=cfg_renorm_type,
+            **cfg_kwargs,
+            **generation_input,
+        )
+
+    H, W = image_size
+    ps = model.config.patch_size
+    h_patches, w_patches = H // ps, W // ps
+    pixels = unpacked_pixels[0]
+    img = pixels.float().cpu()
+    img = img.reshape(h_patches, w_patches, ps, ps, 3)
+    img = img.permute(0, 2, 1, 3, 4).reshape(H, W, 3)
+    img = ((img.clamp(-1, 1) + 1) / 2 * 255).to(torch.uint8).numpy()
+    return Image.fromarray(img)
+
+
+def _generate_t2v_one(
+    model, tokenizer, new_token_ids, prompt, video_size, device,
+    num_timesteps, timestep_shift, sampler,
+    cfg_text_scale, cfg_img_scale, cfg_interval, cfg_renorm_type,
+    video_fps=16,
+    negative_prompt="",
+):
+    from modeling.pixelumm.qwen3_navit import NaiveCache
+
+    n_layers = model.config.llm_config.num_hidden_layers
+    kv = NaiveCache(n_layers)
+    kv_lens, ropes = [0], [0]
+
+    cfg_text_kv = deepcopy(kv) if cfg_text_scale > 1.0 else None
+    cfg_text_kv_lens = list(kv_lens) if cfg_text_scale > 1.0 else None
+    cfg_text_ropes = list(ropes) if cfg_text_scale > 1.0 else None
+    cfg_img_kv = deepcopy(kv) if cfg_img_scale > 1.0 else None
+    cfg_img_kv_lens = list(kv_lens) if cfg_img_scale > 1.0 else None
+    cfg_img_ropes = list(ropes) if cfg_img_scale > 1.0 else None
+
+    kv, kv_lens, ropes = _prefill_text_to_visual_chatml_context(
+        model, tokenizer, new_token_ids, prompt, device, kv, kv_lens, ropes,
+    )
+
+    kv, kv_lens, ropes, visual_start_prefilled = _prefill_generation_visual_start(
+        model, new_token_ids, device, kv, kv_lens, ropes,
+    )
+
+    generation_input = model.prepare_pixel_video_noise(
+        curr_kvlens=kv_lens, curr_rope=ropes,
+        video_sizes=[video_size],
+        fps=video_fps,
+        visual_start_prefilled=visual_start_prefilled,
+    )
+    generation_input = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                        for k, v in generation_input.items()}
+    # Preserve the full-precision RNG sample. Model forward remains under BF16
+    # autocast, while every sampler starts from an FP32 numerical state.
+    generation_input['packed_init_noises'] = generation_input['packed_init_noises'].float()
+
+    cfg_kwargs = {}
+    if cfg_text_scale > 1.0 and cfg_text_kv is not None:
+        # Empty preserves the training-dropout shell. An explicit negative
+        # prompt reproduces the published video's CFG reference context.
+        cfg_text_kv, cfg_text_kv_lens, cfg_text_ropes = _prefill_text_to_visual_chatml_context(
+            model, tokenizer, new_token_ids, negative_prompt, device,
+            cfg_text_kv, cfg_text_kv_lens, cfg_text_ropes,
+        )
+        (
+            cfg_text_kv,
+            cfg_text_kv_lens,
+            cfg_text_ropes,
+            cfg_text_start_prefilled,
+        ) = _prefill_generation_visual_start(
+            model, new_token_ids, device,
+            cfg_text_kv, cfg_text_kv_lens, cfg_text_ropes,
+        )
+        cfg_text_input = model.prepare_pixel_video_noise(
+            curr_kvlens=cfg_text_kv_lens, curr_rope=cfg_text_ropes,
+            video_sizes=[video_size],
+            fps=video_fps,
+            visual_start_prefilled=cfg_text_start_prefilled,
+        )
+        cfg_text_input = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                          for k, v in cfg_text_input.items()}
+        cfg_kwargs.update(
+            cfg_text_past_key_values=cfg_text_kv,
+            cfg_text_packed_position_ids=cfg_text_input['packed_position_ids'],
+            cfg_text_packed_query_indexes=cfg_text_input['packed_indexes'],
+            cfg_text_key_values_lens=cfg_text_input['key_values_lens'],
+            cfg_text_packed_key_value_indexes=cfg_text_input['packed_key_value_indexes'],
+        )
+    if cfg_img_scale > 1.0 and cfg_img_kv is not None:
+        (
+            cfg_img_kv,
+            cfg_img_kv_lens,
+            cfg_img_ropes,
+            cfg_img_start_prefilled,
+        ) = _prefill_generation_visual_start(
+            model, new_token_ids, device,
+            cfg_img_kv, cfg_img_kv_lens, cfg_img_ropes,
+        )
+        cfg_img_input = model.prepare_pixel_video_noise(
+            curr_kvlens=cfg_img_kv_lens, curr_rope=cfg_img_ropes,
+            video_sizes=[video_size],
+            fps=video_fps,
+            visual_start_prefilled=cfg_img_start_prefilled,
+        )
+        cfg_img_input = {k: v.to(device) if isinstance(v, torch.Tensor) else v
+                         for k, v in cfg_img_input.items()}
+        cfg_kwargs.update(
+            cfg_img_past_key_values=cfg_img_kv,
+            cfg_img_packed_position_ids=cfg_img_input['packed_position_ids'],
+            cfg_img_packed_query_indexes=cfg_img_input['packed_indexes'],
+            cfg_img_key_values_lens=cfg_img_input['key_values_lens'],
+            cfg_img_packed_key_value_indexes=cfg_img_input['packed_key_value_indexes'],
+        )
+
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        unpacked_video = model.generate_video(
+            past_key_values=kv,
+            num_timesteps=num_timesteps,
+            timestep_shift=timestep_shift,
+            sampler=sampler,
+            cfg_text_scale=cfg_text_scale,
+            cfg_img_scale=cfg_img_scale,
+            cfg_interval=cfg_interval,
+            cfg_renorm_min=0.0,
+            cfg_renorm_type=cfg_renorm_type,
+            **cfg_kwargs,
+            **generation_input,
+        )
+
+    num_frames, H, W = video_size
+    return _video_tubes_to_pil_frames(
+        unpacked_video[0],
+        (H, W),
+        num_frames,
+        model.config.patch_size,
+    )
+
+
+def generate_vlm_interleaved(
+    model,
+    tokenizer,
+    new_token_ids,
+    images,
+    instruction,
+    device,
+    max_length,
+    temperature=0.0,
+    fixed_length=False,
+    image_placeholder_policy="strict",
+    reasoning_mode="standard_bare",
+    exact_max_new_tokens=False,
+    return_generation_metadata=False,
+):
+    """Generate an I2T answer with train-aligned interleaved media prefill.
+
+    Training represents every image as four consecutive attention segments:
+    causal user text, causal ``<vision_start>``, a full-attention patch block,
+    then causal ``<vision_end>``.  This helper applies the same sequence to one
+    or more images and preserves ``<image>`` / ``<image N>`` locations in the
+    benchmark question.  When a question has no placeholders, all images are
+    placed before the question, matching ordinary F14 I2T samples.
+
+    ``fixed_length=True`` is reserved for in-training FSDP eval, where every
+    rank must execute the same number of decode iterations.  Standalone
+    benchmark replicas can stop at EOS and should leave it false.
+
+    ``reasoning_mode="standard_bare"`` preserves the train-aligned bare
+    assistant prefix. Source responses may begin directly, with an empty
+    shell, or with a filled trace. ``exact_max_new_tokens=True`` makes the standalone benchmark's
+    output budget match Hugging Face/lmms-eval semantics without changing the
+    fixed-iteration contract used by in-training FSDP eval. The defaults
+    preserve all existing callers.
+    """
+    from modeling.pixelumm.qwen3_navit import NaiveCache
+
+    images = list(images)
+    if not images:
+        raise ValueError("VLM generation requires at least one image")
+
+    if image_placeholder_policy not in {"strict", "numbered_references"}:
+        raise ValueError(
+            f"Unsupported image placeholder policy: {image_placeholder_policy}"
+        )
+    if reasoning_mode != "standard_bare":
+        raise ValueError(f"Unsupported VLM reasoning mode: {reasoning_mode}")
+
+    placeholder = re.compile(r"<image(?:\s+(\d+))?>", flags=re.IGNORECASE)
+    matches = list(placeholder.finditer(instruction))
+
+    parts = []
+    if matches:
+        numbered = [match.group(1) is not None for match in matches]
+        if any(numbered) and not all(numbered):
+            raise ValueError(
+                "VLM prompts cannot mix <image> and numbered <image N> placeholders"
+            )
+
+        if image_placeholder_policy == "numbered_references" and not all(numbered):
+            raise ValueError(
+                "numbered_references requires numbered <image N> placeholders"
+            )
+
+        if image_placeholder_policy == "numbered_references":
+            labels = [int(match.group(1)) for match in matches]
+            unique_labels = sorted(set(labels))
+            if len(images) == 1 and len(unique_labels) > 1:
+                # Some benchmarks provide one pre-composed panel while their
+                # question uses <image N> as textual sub-panel references.
+                normalized = placeholder.sub(
+                    lambda match: f"Panel {int(match.group(1))}",
+                    instruction,
+                )
+                parts.extend((
+                    ("image", 0),
+                    ("text", normalized),
+                ))
+            else:
+                expected_labels = list(range(1, len(images) + 1))
+                if unique_labels != expected_labels:
+                    raise ValueError(
+                        "Numbered image references must cover each benchmark visual: "
+                        f"labels={unique_labels} expected={expected_labels}"
+                    )
+                cursor = 0
+                seen_labels = set()
+                for label, match in zip(labels, matches):
+                    parts.append(("text", instruction[cursor : match.start()]))
+                    if label not in seen_labels:
+                        parts.append(("image", label - 1))
+                        seen_labels.add(label)
+                    else:
+                        parts.append(("text", f"Picture {label}"))
+                    cursor = match.end()
+                parts.append(("text", instruction[cursor:]))
+        else:
+            if len(matches) != len(images):
+                raise ValueError(
+                    "Image placeholder count must match the number of benchmark visuals: "
+                    f"placeholders={len(matches)} images={len(images)}"
+                )
+            image_indices = (
+                [int(match.group(1)) - 1 for match in matches]
+                if all(numbered)
+                else list(range(len(matches)))
+            )
+            if sorted(image_indices) != list(range(len(images))):
+                raise ValueError(
+                    "Numbered image placeholders must reference every visual exactly once: "
+                    f"indices={image_indices} images={len(images)}"
+                )
+            cursor = 0
+            for image_index, match in zip(image_indices, matches):
+                parts.append(("text", instruction[cursor : match.start()]))
+                parts.append(("image", image_index))
+                cursor = match.end()
+            parts.append(("text", instruction[cursor:]))
+    else:
+        parts.extend(("image", image_index) for image_index in range(len(images)))
+        parts.append(("text", instruction))
+
+    n_layers = model.config.llm_config.num_hidden_layers
+    kv = NaiveCache(n_layers)
+    kv_lens, ropes = [0], [0]
+
+    def to_device(inputs):
+        return {
+            key: value.to(device) if isinstance(value, torch.Tensor) else value
+            for key, value in inputs.items()
+        }
+
+    def update_text_ids(token_ids):
+        nonlocal kv, kv_lens, ropes
+        if not token_ids:
+            return
+        text_input, kv_lens, ropes = model._prepare_token_ids_as_text_cache_input(
+            kv_lens,
+            ropes,
+            [list(token_ids)],
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(text_input))
+
+    prefix_input, kv_lens, ropes = model.prepare_image_to_text_chatml_prefixes(
+        curr_kvlens=kv_lens,
+        curr_rope=ropes,
+        tokenizer=tokenizer,
+    )
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        kv = model.forward_cache_update_text(kv, **to_device(prefix_input))
+
+    if not parts or parts[-1][0] != "text":
+        raise AssertionError("Interleaved VLM prompt must end in a text segment")
+    trailing_instruction = parts[-1][1]
+
+    for part_type, value in parts[:-1]:
+        if part_type == "text":
+            update_text_ids(model._encode_prompt_text(tokenizer, value))
+            continue
+
+        image_index = int(value)
+        if len(images) > 1:
+            update_text_ids(
+                model._encode_prompt_text(tokenizer, f"Picture {image_index + 1}: ")
+            )
+
+        start_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+            kv_lens,
+            ropes,
+            new_token_ids,
+            delimiter="start",
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(start_input))
+
+        image_input, kv_lens, ropes = model.prepare_pixel_images(
+            curr_kvlens=kv_lens,
+            curr_rope=ropes,
+            images=[images[image_index]],
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_pixel(kv, **to_device(image_input))
+
+        end_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+            kv_lens,
+            ropes,
+            new_token_ids,
+            delimiter="end",
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(end_input))
+
+    # Tokenize the trailing question text and ChatML turn boundary together,
+    # in one encode call to preserve the boundary before <|im_end|>.
+    suffix_ids = model._image_to_text_chatml_suffix_ids(
+        tokenizer,
+        trailing_instruction,
+    )
+    assistant_prefill = ""
+    if not suffix_ids:
+        raise ValueError("Image-to-text ChatML suffix cannot be empty")
+    update_text_ids(suffix_ids[:-1])
+    start_token_ids = [suffix_ids[-1]]
+    start_input = to_device(
+        model.prepare_text_start_tokens(kv_lens, ropes, start_token_ids)
+    )
+
+    # ``generate_text`` counts ``start_token_ids`` in its returned length.
+    # Add one iteration when exact_max_new_tokens is requested; the start
+    # token is removed below.
+    decode_max_length = int(max_length) + (1 if exact_max_new_tokens else 0)
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        token_ids, generation_metadata = model.generate_text(
+            past_key_values=kv,
+            max_length=decode_max_length,
+            do_sample=(temperature > 0.0),
+            temperature=max(temperature, 1e-5),
+            end_token_id=None if fixed_length else new_token_ids["eos_token_id"],
+            return_generation_metadata=True,
+            **start_input,
+        )
+
+    token_ids = token_ids.squeeze(-1).cpu().tolist()
+    eos_id = new_token_ids["eos_token_id"]
+    bos_id = new_token_ids["bos_token_id"]
+    text_ids = []
+    hit_eos = bool(generation_metadata["hit_end_token"])
+    for token_id in token_ids:
+        if token_id == eos_id:
+            break
+        text_ids.append(token_id)
+    if text_ids and text_ids[0] == start_token_ids[0]:
+        text_ids = text_ids[1:]
+    if text_ids and text_ids[0] == bos_id:
+        text_ids = text_ids[1:]
+    text = tokenizer.decode(text_ids, skip_special_tokens=True)
+    if not return_generation_metadata:
+        return text
+    return text, {
+        "assistant_prefill": assistant_prefill,
+        "generated_tokens": len(text_ids),
+        "hit_eos": hit_eos,
+        "max_new_tokens": int(max_length),
+        "reasoning_mode": reasoning_mode,
+    }
+
+
+def _v2t_frames_to_und_tubes(model, frames, timestamps, *, timestamp_mode="tube_start"):
+    """Group processed PIL frames into the released video-UND tube contract."""
+    import numpy as np
+
+    temporal_patch_size = int(model.config.pixel_video_temporal_patch_size)
+    mrope_type = str(model.config.llm_config.qwen3_mrope_type)
+    if mrope_type != "sensenova":
+        raise ValueError(f"PixelUMM V2T requires qwen3_mrope_type='sensenova', got {mrope_type!r}")
+    if timestamp_mode != "tube_start":
+        raise ValueError(f"PixelUMM V2T requires timestamp_mode='tube_start', got {timestamp_mode!r}")
+    if temporal_patch_size <= 0:
+        raise ValueError(
+            f"pixel_video_temporal_patch_size must be positive, got {temporal_patch_size}"
+        )
+    if not frames:
+        raise ValueError("V2T tube eval requires at least one processed frame")
+
+    timeline_origin = float(timestamps[0])
+    local_timestamps = [float(value) - timeline_origin for value in timestamps]
+    pad = (-len(frames)) % temporal_patch_size
+    if pad:
+        frames = list(frames) + [frames[-1]] * pad
+        local_timestamps = local_timestamps + [local_timestamps[-1]] * pad
+
+    def frame_to_tensor(frame):
+        array = np.asarray(frame.convert("RGB"), dtype=np.float32)
+        tensor = torch.from_numpy(array).permute(2, 0, 1)
+        return tensor / 255.0 * 2.0 - 1.0
+
+    tubes = []
+    for start in range(0, len(frames), temporal_patch_size):
+        tube_frames = frames[start : start + temporal_patch_size]
+        tube_start_seconds = float(local_timestamps[start])
+        tube_tensor = torch.stack([frame_to_tensor(frame) for frame in tube_frames], dim=0)
+        tubes.append(
+            {
+                "tensor": tube_tensor,
+                "timestamp_seconds": tube_start_seconds,
+            }
+        )
+    return tubes
+
+
+def _prefill_v2t_und_tubes(
+    model,
+    tokenizer,
+    new_token_ids,
+    kv,
+    kv_lens,
+    ropes,
+    frames,
+    timestamps,
+    device,
+    *,
+    timestamp_mode="tube_start",
+):
+    """Prefill the model-aligned tube V2T media context, split-by-split.
+
+    Training packs, per tube: a causal ``<t.t seconds>`` text span, a causal
+    ``<vision_start>``, one full-attention tube, then a causal ``<vision_end>``.
+    """
+
+    def to_device(prepared):
+        return {
+            key: value.to(device) if isinstance(value, torch.Tensor) else value
+            for key, value in prepared.items()
+        }
+
+    tubes = _v2t_frames_to_und_tubes(
+        model, frames, timestamps, timestamp_mode=timestamp_mode
+    )
+    for tube in tubes:
+        text_input, kv_lens, ropes = model.prepare_text_spans_for_cache(
+            kv_lens,
+            ropes,
+            [f"<{tube['timestamp_seconds']:.1f} seconds>"],
+            tokenizer,
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(text_input))
+
+        start_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+            kv_lens, ropes, new_token_ids, delimiter="start"
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(start_input))
+
+        tube_input, kv_lens, ropes = model.prepare_pixel_video_und_tubes_for_text(
+            kv_lens,
+            ropes,
+            [tube["tensor"]],
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_pixel(kv, **to_device(tube_input))
+
+        end_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+            kv_lens, ropes, new_token_ids, delimiter="end"
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(end_input))
+
+    return kv, kv_lens, ropes
+
+
+def _prefill_v2t_und_frames(
+    model,
+    tokenizer,
+    new_token_ids,
+    kv,
+    kv_lens,
+    ropes,
+    frames,
+    timestamps,
+    device,
+):
+    """Prefill sparse long-video frames through the image UND embedder.
+
+    This is the released Video4 sparse-frame contract: each sampled frame gets
+    its own lexical clip-relative
+    timestamp and visual island. No temporal tube projection or physical-time
+    MRoPE is used.
+    """
+
+    def to_device(prepared):
+        return {
+            key: value.to(device) if isinstance(value, torch.Tensor) else value
+            for key, value in prepared.items()
+        }
+
+    timeline_origin = float(timestamps[0])
+    local_timestamps = [float(value) - timeline_origin for value in timestamps]
+    for frame, timestamp_seconds in zip(frames, local_timestamps, strict=True):
+        text_input, kv_lens, ropes = model.prepare_text_spans_for_cache(
+            kv_lens,
+            ropes,
+            [f"<{timestamp_seconds:.1f} seconds>"],
+            tokenizer,
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(text_input))
+
+        start_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+            kv_lens, ropes, new_token_ids, delimiter="start"
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(start_input))
+
+        image_input, kv_lens, ropes = model.prepare_pixel_images(
+            curr_kvlens=kv_lens,
+            curr_rope=ropes,
+            images=[frame],
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_pixel(kv, **to_device(image_input))
+
+        end_input, kv_lens, ropes = model.prepare_visual_delimiter_tokens(
+            kv_lens, ropes, new_token_ids, delimiter="end"
+        )
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            kv = model.forward_cache_update_text(kv, **to_device(end_input))
+
+    return kv, kv_lens, ropes
+
+
+def _generate_v2t_one(
+    model,
+    tokenizer,
+    new_token_ids,
+    frames,
+    timestamps,
+    instruction,
+    device,
+    max_length,
+    temperature,
+    *,
+    temporal_patch_size=None,
+    timestamp_mode="tube_start",
+):
+    """Decode one V2T video context (frame-as-image or video-UND tubes)."""
+    from modeling.pixelumm.qwen3_navit import NaiveCache
+
+    if len(frames) != len(timestamps):
+        raise ValueError(
+            f"V2T frame/timestamp length mismatch: {len(frames)} != {len(timestamps)}"
+        )
+    n_layers = model.config.llm_config.num_hidden_layers
+    kv = NaiveCache(n_layers)
+    kv_lens, ropes = [0], [0]
+
+    prefix_input, kv_lens, ropes = model.prepare_image_to_text_chatml_prefixes(
+        curr_kvlens=kv_lens,
+        curr_rope=ropes,
+        tokenizer=tokenizer,
+    )
+    prefix_input = {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in prefix_input.items()
+    }
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        kv = model.forward_cache_update_text(kv, **prefix_input)
+
+    actual_temporal_patch_size = int(
+        temporal_patch_size
+        if temporal_patch_size is not None
+        else model.config.pixel_video_temporal_patch_size
+    )
+    if actual_temporal_patch_size == 1:
+        kv, kv_lens, ropes = _prefill_v2t_und_frames(
+            model, tokenizer, new_token_ids, kv, kv_lens, ropes,
+            frames, timestamps, device,
+        )
+    elif actual_temporal_patch_size == int(model.config.pixel_video_temporal_patch_size):
+        kv, kv_lens, ropes = _prefill_v2t_und_tubes(
+            model, tokenizer, new_token_ids, kv, kv_lens, ropes,
+            frames, timestamps, device, timestamp_mode=timestamp_mode,
+        )
+    else:
+        raise ValueError(
+            "Unsupported V2T temporal patch override: "
+            f"{actual_temporal_patch_size}; expected 1 or "
+            f"{model.config.pixel_video_temporal_patch_size}"
+        )
+
+    return _decode_v2t_suffix(
+        model,
+        tokenizer,
+        new_token_ids,
+        kv,
+        kv_lens,
+        ropes,
+        instruction,
+        device,
+        max_length,
+        temperature,
+    )
+
+
+def _decode_v2t_suffix(
+    model,
+    tokenizer,
+    new_token_ids,
+    kv,
+    kv_lens,
+    ropes,
+    instruction,
+    device,
+    max_length,
+    temperature,
+):
+    """Shared V2T ChatML suffix prefill and answer decoding."""
+    suffix_input, kv_lens, ropes, start_token_ids = (
+        model.prepare_image_to_text_chatml_suffix_contexts_for_decode(
+            curr_kvlens=kv_lens,
+            curr_rope=ropes,
+            instructions=[instruction],
+            tokenizer=tokenizer,
+        )
+    )
+    suffix_input = {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in suffix_input.items()
+    }
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        kv = model.forward_cache_update_text(kv, **suffix_input)
+
+    start_input = model.prepare_text_start_tokens(kv_lens, ropes, start_token_ids)
+    start_input = {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in start_input.items()
+    }
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        token_ids = model.generate_text(
+            past_key_values=kv,
+            max_length=max_length,
+            do_sample=(temperature > 0.0),
+            temperature=max(temperature, 1e-5),
+            end_token_id=None,
+            **start_input,
+        )
+
+    token_ids = token_ids.squeeze(-1).cpu().tolist()
+    eos_id = new_token_ids["eos_token_id"]
+    bos_id = new_token_ids["bos_token_id"]
+    text_ids = []
+    for token_id in token_ids:
+        if token_id == eos_id:
+            break
+        text_ids.append(token_id)
+    if text_ids and text_ids[0] == start_token_ids[0]:
+        text_ids = text_ids[1:]
+    if text_ids and text_ids[0] == bos_id:
+        text_ids = text_ids[1:]
+    return tokenizer.decode(text_ids, skip_special_tokens=True)

train/model_factory.py

@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Construct a PixelUMM model from the selected architecture profile."""
+
+from __future__ import annotations
+
+from transformers import AutoTokenizer
+from transformers.modeling_utils import no_init_weights
+
+from data.data_utils import add_special_tokens
+from modeling.pixelumm.pixelumm import PixelUMM, PixelUMMConfig
+from modeling.pixelumm.qwen3_navit import Qwen3ForCausalLM
+from modeling.qwen3.configuration_qwen3 import Qwen3Config
+from train.config import (
+    DataArguments,
+    ModelArguments,
+    TrainingArguments,
+    configure_qwen3_mrope_config,
+)
+
+
+def _count_parameters(module) -> int:
+    return sum(parameter.numel() for parameter in module.parameters())
+
+
+def build_pixelumm_model(
+    model_args: ModelArguments,
+    data_args: DataArguments,
+    training_args: TrainingArguments,
+    logger,
+):
+    """Build the selected topology and tokenizer without loading checkpoint weights."""
+    if model_args.llm_backend != "qwen3":
+        raise ValueError(
+            "PixelUMM-release supports only the R07 Qwen3 backend; "
+            f"got {model_args.llm_backend!r}"
+        )
+
+    llm_config = Qwen3Config.from_pretrained(model_args.llm_path)
+    llm_config.layer_module = model_args.layer_module
+    llm_config.qk_norm = model_args.llm_qk_norm
+    llm_config.tie_word_embeddings = model_args.tie_word_embeddings
+    llm_config.freeze_und = training_args.freeze_und
+    llm_config.qwen_mlp_sequence_chunk_size = int(
+        model_args.qwen_mlp_sequence_chunk_size
+    )
+    llm_config.inference_attention_backend = model_args.inference_attention_backend
+    llm_config.conditional_branch_gradient_contract = (
+        training_args.conditional_branch_gradient_contract
+    )
+
+    configure_qwen3_mrope_config(llm_config, model_args)
+    llm_config.rope_theta_hw = model_args.sensenova_vision_rope_theta
+    llm_config.max_position_embeddings_hw = (
+        model_args.sensenova_vision_max_position_embeddings
+    )
+
+    with no_init_weights():
+        language_model = Qwen3ForCausalLM(llm_config)
+    removed = language_model.remove_unused_sensenova_base_qk_norms()
+    logger.info(
+        "Removed %d unused inherited Q/K norm tensors.",
+        removed,
+    )
+
+    config = PixelUMMConfig(
+        visual_gen=training_args.visual_gen,
+        visual_und=training_args.visual_und,
+        llm_config=llm_config,
+        llm_backend=model_args.llm_backend,
+        gen_space=training_args.gen_space,
+        und_space=training_args.und_space,
+        patch_size=model_args.pixel_token_patch_size,
+        pixel_embedder_type=model_args.pixel_embedder_type,
+        pixel_image_additive_pos_embed_type=model_args.pixel_image_additive_pos_embed_type,
+        pixel_video_additive_pos_embed_type=model_args.pixel_video_additive_pos_embed_type,
+        add_timestep_embedding=model_args.add_timestep_embedding,
+        add_noise_scale_embedding=model_args.add_noise_scale_embedding,
+        noise_scale=model_args.noise_scale,
+        noise_scale_mode=model_args.noise_scale_mode,
+        enable_pixel_video=model_args.enable_pixel_video,
+        pixel_video_embedder_type=model_args.pixel_video_embedder_type,
+        pixel_separate_und_gen_embedder=model_args.pixel_separate_und_gen_embedder,
+        pixel_video_separate_und_gen_embedder=model_args.pixel_video_separate_und_gen_embedder,
+        pixel_video_temporal_patch_size=model_args.pixel_video_temporal_patch_size,
+        noise_time_logit_mean=training_args.noise_time_logit_mean,
+        noise_time_logit_std=training_args.noise_time_logit_std,
+        time_schedule=training_args.time_schedule,
+        prediction_type=training_args.prediction_type,
+        loss_type=training_args.loss_type,
+        x_pred_t_min=training_args.x_pred_t_min,
+        pixel_head_type=model_args.pixel_head_type,
+        pixel_video_head_type=model_args.pixel_video_head_type,
+        ce_loss_checkpoint_chunk_size=model_args.ce_loss_checkpoint_chunk_size,
+        mse_loss_reduction=training_args.mse_loss_reduction,
+        packed_attention_impl=model_args.packed_attention_impl,
+        packed_expert_routing=model_args.packed_expert_routing,
+        packed_sequence_layout=data_args.packed_sequence_layout,
+        conditional_branch_gradient_contract=(
+            training_args.conditional_branch_gradient_contract
+        ),
+    )
+    model = PixelUMM(language_model, config)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_args.llm_path)
+    tokenizer, new_token_ids, num_new_tokens = add_special_tokens(tokenizer)
+    current_vocab_size = model.language_model.model.embed_tokens.num_embeddings
+    if num_new_tokens > 0 or current_vocab_size != len(tokenizer):
+        logger.info(
+            "Resizing token embeddings: %d -> %d", current_vocab_size, len(tokenizer)
+        )
+        model.language_model.resize_token_embeddings(len(tokenizer))
+        model.config.llm_config.vocab_size = len(tokenizer)
+        model.language_model.config.vocab_size = len(tokenizer)
+
+    total_param_count = _count_parameters(model)
+    lm_param_count = _count_parameters(model.language_model)
+    logger.info(
+        "Model parameter count: %.2fB (LM-only: %.2fB)",
+        total_param_count / 1e9,
+        lm_param_count / 1e9,
+    )
+    return model, tokenizer, new_token_ids

train/release_checkpoint.py

@@ -0,0 +1,324 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Load complete PixelUMM checkpoints using a model configuration."""
+
+from __future__ import annotations
+
+import logging
+import os
+from dataclasses import dataclass
+from datetime import timedelta
+from pathlib import Path
+from typing import Sequence
+
+import torch
+import torch.distributed as dist
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+from train.config import DataArguments, ModelArguments, TrainingArguments, load_release_config
+from train.model_factory import build_pixelumm_model
+from train.checkpoint_export import EXPORT_MANIFEST, validate_export_manifest
+
+
+ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_RELEASE_CONFIG = ROOT / "experiments" / "s8_f22_r07" / "release.yaml"
+
+
+@dataclass(frozen=True)
+class ReleasedModel:
+    model: torch.nn.Module
+    tokenizer: object
+    new_token_ids: dict
+    model_args: ModelArguments
+    data_args: DataArguments
+    training_args: TrainingArguments
+    checkpoint_path: Path
+    config_path: Path
+
+
+def dtype_from_name(name: str) -> torch.dtype:
+    try:
+        return {
+            "bf16": torch.bfloat16,
+            "bfloat16": torch.bfloat16,
+            "fp16": torch.float16,
+            "float16": torch.float16,
+            "half": torch.float16,
+            "fp32": torch.float32,
+            "float32": torch.float32,
+            "float": torch.float32,
+        }[str(name).lower()]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported PixelUMM dtype: {name}") from exc
+
+
+def parse_release_arguments(
+    config_path: str | os.PathLike[str] = DEFAULT_RELEASE_CONFIG,
+    *,
+    extra_args: Sequence[str] = (),
+) -> tuple[ModelArguments, DataArguments, TrainingArguments]:
+    """Parse the checked-in release YAML without rebuilding a training CLI."""
+
+    config_path = Path(config_path).expanduser().resolve()
+    extra_args = tuple(map(str, extra_args))
+    if extra_args:
+        if len(extra_args) != 2 or extra_args[0] != "--llm_path":
+            raise ValueError(
+                "Released config accepts only the explicit --llm_path override"
+            )
+        llm_path = extra_args[1]
+    else:
+        llm_path = None
+    model_args, data_args, training_args = load_release_config(
+        config_path,
+        llm_path=llm_path,
+    )
+    _validate_release_contract(model_args, data_args, training_args)
+    return model_args, data_args, training_args
+
+
+def _validate_release_contract(model_args, data_args, training_args) -> None:
+    checks = {
+        "llm_backend=qwen3": model_args.llm_backend == "qwen3",
+        "mrope=sensenova": model_args.qwen3_mrope_type == "sensenova",
+        "image_patch=16": int(model_args.pixel_token_patch_size) == 16,
+        "video_enabled": bool(model_args.enable_pixel_video),
+        "video_und_embedder=boolean": isinstance(model_args.pixel_video_separate_und_gen_embedder, bool),
+        "video_temporal_patch=4": int(model_args.pixel_video_temporal_patch_size) == 4,
+        "image_embedder=image_raw_patch_linear": model_args.pixel_embedder_type
+        == "image_raw_patch_linear",
+        "video_embedder=video_raw_tube_linear": model_args.pixel_video_embedder_type
+        == "video_raw_tube_linear",
+        "inference_attention=flash_attn_varlen": model_args.inference_attention_backend
+        == "flash_attn_varlen",
+        "layout=pixelumm_mot": data_args.packed_sequence_layout == "pixelumm_mot",
+        "visual_gen": bool(training_args.visual_gen),
+        "visual_und": bool(training_args.visual_und),
+        "gen_space=pixel": training_args.gen_space == "pixel",
+        "und_space=pixel": training_args.und_space == "pixel",
+        "prediction=x": training_args.prediction_type == "x",
+        "loss=v_loss": training_args.loss_type == "v_loss",
+    }
+    failed = [name for name, passed in checks.items() if not passed]
+    if failed:
+        raise RuntimeError(
+            "Released checkpoint config is not a supported PixelUMM contract: "
+            + ", ".join(failed)
+        )
+
+
+def validate_released_checkpoint(checkpoint_path: str | os.PathLike[str]) -> Path:
+    checkpoint_path = Path(checkpoint_path).expanduser().resolve()
+    if not checkpoint_path.is_dir():
+        raise FileNotFoundError(f"Released checkpoint directory is missing: {checkpoint_path}")
+    export_manifest = checkpoint_path / EXPORT_MANIFEST
+    exported_files = validate_export_manifest(checkpoint_path) if export_manifest.is_file() else None
+    if not (checkpoint_path / "__SAVE_COMPLETE").is_file() and exported_files is None:
+        raise FileNotFoundError(
+            f"Released checkpoint lacks __SAVE_COMPLETE or verified {EXPORT_MANIFEST}: {checkpoint_path}"
+        )
+    if (checkpoint_path / "model" / ".metadata").is_file():
+        required_files = _validate_dcp_model(checkpoint_path)
+        if exported_files is not None and exported_files != required_files:
+            raise RuntimeError("Export manifest must cover exactly DCP metadata and all referenced shards")
+        if exported_files is not None and (checkpoint_path / "model.safetensors").exists():
+            raise RuntimeError("DCP export must not contain an unverified model.safetensors")
+    else:
+        safetensors_path = checkpoint_path / "model.safetensors"
+        if not safetensors_path.is_file() or safetensors_path.stat().st_size <= 0:
+            raise FileNotFoundError(
+                "Released checkpoint has neither complete DCP model/.metadata nor "
+                f"a non-empty model.safetensors: {checkpoint_path}"
+            )
+    return checkpoint_path
+
+
+def _validate_dcp_model(checkpoint_path: Path) -> set[str]:
+    """Fail closed on missing or truncated DCP model shards."""
+
+    from torch.distributed.checkpoint import FileSystemReader
+
+    model_dir = checkpoint_path / "model"
+    metadata_path = model_dir / ".metadata"
+    metadata = FileSystemReader(str(model_dir)).read_metadata()
+    if not metadata.storage_data:
+        raise RuntimeError(f"Released checkpoint DCP metadata is empty: {metadata_path}")
+
+    required_sizes: dict[str, int] = {}
+    for storage_info in metadata.storage_data.values():
+        relative_path = str(storage_info.relative_path)
+        required_sizes[relative_path] = max(
+            required_sizes.get(relative_path, 0),
+            int(storage_info.offset) + int(storage_info.length),
+        )
+    missing = []
+    truncated = []
+    for relative_path, required_size in sorted(required_sizes.items()):
+        shard = model_dir / relative_path
+        if shard.resolve().parent != model_dir.resolve() or not relative_path.endswith('.distcp'):
+            raise RuntimeError(f"Invalid DCP shard path: {relative_path}")
+        if not shard.is_file():
+            missing.append(relative_path)
+        elif shard.stat().st_size < required_size:
+            truncated.append((relative_path, shard.stat().st_size, required_size))
+    if missing or truncated:
+        raise RuntimeError(
+            f"Incomplete released DCP model at {checkpoint_path}: "
+            f"missing={missing[:8]} truncated={truncated[:8]}"
+        )
+    return {"model/.metadata", *(f"model/{name}" for name in required_sizes)}
+
+
+def validate_checkpoint_model_schema(
+    checkpoint_path: str | os.PathLike[str], model: torch.nn.Module
+) -> None:
+    """Compare stored keys/shapes, including DCP keys a load plan would ignore.
+
+    A model on the meta device is sufficient: this check reads only checkpoint
+    metadata, without allocating model weights or touching a CUDA device.
+    """
+    checkpoint_path = Path(checkpoint_path)
+    safetensors_path = checkpoint_path / "model.safetensors"
+    if safetensors_path.is_file():
+        with safe_open(str(safetensors_path), framework="pt", device="cpu") as reader:
+            shapes = {key: tuple(reader.get_slice(key).get_shape()) for key in reader.keys()}
+    else:
+        from torch.distributed.checkpoint import FileSystemReader
+        from torch.distributed.checkpoint.metadata import TensorStorageMetadata
+
+        metadata = FileSystemReader(str(checkpoint_path / "model")).read_metadata()
+        non_tensors = [
+            key for key, value in metadata.state_dict_metadata.items()
+            if not isinstance(value, TensorStorageMetadata)
+        ]
+        if non_tensors:
+            raise RuntimeError(f"Non-tensor entries in released model: {non_tensors[:8]}")
+        shapes = {key: tuple(value.size) for key, value in metadata.state_dict_metadata.items()}
+
+    state = model.state_dict()
+    missing = sorted(set(state) - set(shapes))
+    unexpected = sorted(set(shapes) - set(state))
+    mismatched = [
+        (key, tuple(state[key].shape), shapes[key])
+        for key in sorted(set(state) & set(shapes))
+        if tuple(state[key].shape) != shapes[key]
+    ]
+    if missing or unexpected or mismatched:
+        raise RuntimeError(
+            f"Checkpoint/model schema mismatch at {checkpoint_path}: "
+            f"missing={missing[:8]} unexpected={unexpected[:8]} "
+            f"shape_mismatches(model, checkpoint)={mismatched[:8]}"
+        )
+
+
+def load_checkpoint_weights(
+    checkpoint_path: str | os.PathLike[str],
+    model: torch.nn.Module,
+    *,
+    logger: logging.Logger | None = None,
+) -> torch.nn.Module:
+    """Strictly load only learned model tensors from a released checkpoint."""
+
+    checkpoint_path = validate_released_checkpoint(checkpoint_path)
+    validate_checkpoint_model_schema(checkpoint_path, model)
+    logger = logger or logging.getLogger("pixelumm.release")
+    safetensors_path = checkpoint_path / "model.safetensors"
+    if safetensors_path.is_file():
+        logger.info("Loading released safetensors weights from %s", safetensors_path)
+        state_dict = load_file(str(safetensors_path), device="cpu")
+    else:
+        from torch.distributed.checkpoint import FileSystemReader, load
+
+        model_dir = checkpoint_path / "model"
+        logger.info("Loading released DCP weights from %s", model_dir)
+        state_dict = model.state_dict()
+        load(state_dict, storage_reader=FileSystemReader(str(model_dir)))
+    incompatible = model.load_state_dict(state_dict, strict=True)
+    if incompatible.missing_keys or incompatible.unexpected_keys:
+        raise RuntimeError(
+            "Strict released checkpoint load returned incompatible keys: "
+            f"missing={incompatible.missing_keys[:8]} "
+            f"unexpected={incompatible.unexpected_keys[:8]}"
+        )
+    return model
+
+
+def initialize_single_process(device: str | torch.device = "cuda") -> torch.device:
+    """Initialize the world-size-one process group required by DCP loading."""
+
+    device = torch.device(device)
+    if device.type == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("Released PixelUMM inference requires a CUDA GPU")
+        index = 0 if device.index is None else int(device.index)
+        torch.cuda.set_device(index)
+        device = torch.device("cuda", index)
+    if not dist.is_initialized():
+        os.environ.setdefault("RANK", "0")
+        os.environ.setdefault("WORLD_SIZE", "1")
+        os.environ.setdefault("LOCAL_RANK", "0")
+        os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+        os.environ.setdefault("MASTER_PORT", "29501")
+        dist.init_process_group("gloo", timeout=timedelta(seconds=1800))
+    if dist.get_world_size() != 1:
+        raise RuntimeError(
+            "release inference.py is a one-process entrypoint; use the benchmark "
+            "launchers for distributed evaluation"
+        )
+    return device
+
+
+def load_released_model(
+    checkpoint_path: str | os.PathLike[str],
+    *,
+    config_path: str | os.PathLike[str] = DEFAULT_RELEASE_CONFIG,
+    device: str | torch.device = "cuda",
+    dtype: str = "bf16",
+    llm_path: str | None = None,
+    logger: logging.Logger | None = None,
+) -> ReleasedModel:
+    """Load a released DCP/safetensors checkpoint without training lineage state."""
+
+    checkpoint_path = validate_released_checkpoint(checkpoint_path)
+    config_path = Path(config_path).expanduser().resolve()
+    extra_args = [] if llm_path is None else ["--llm_path", str(llm_path)]
+    model_args, data_args, training_args = parse_release_arguments(
+        config_path, extra_args=extra_args
+    )
+    device = initialize_single_process(device)
+    parameter_dtype = dtype_from_name(dtype)
+    logger = logger or logging.getLogger("pixelumm.release")
+    if not logger.handlers:
+        logger.addHandler(logging.StreamHandler())
+    logger.setLevel(logging.INFO)
+
+    # The released checkpoint owns every learned tensor.  The Qwen path is a
+    # config/tokenizer dependency only; loading its five weight shards here
+    # would waste tens of GB and then be immediately overwritten by DCP.
+    previous_default_dtype = torch.get_default_dtype()
+    try:
+        torch.set_default_dtype(parameter_dtype)
+        model, tokenizer, new_token_ids = build_pixelumm_model(
+            model_args,
+            data_args,
+            training_args,
+            logger,
+        )
+    finally:
+        torch.set_default_dtype(previous_default_dtype)
+    model = load_checkpoint_weights(checkpoint_path, model, logger=logger)
+    model = model.to(device=device, dtype=parameter_dtype).eval()
+    for parameter in model.parameters():
+        parameter.requires_grad = False
+    return ReleasedModel(
+        model=model,
+        tokenizer=tokenizer,
+        new_token_ids=new_token_ids,
+        model_args=model_args,
+        data_args=data_args,
+        training_args=training_args,
+        checkpoint_path=checkpoint_path,
+        config_path=config_path,
+    )

train/scientific_config.py

@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Strict loader for immutable PixelUMM scientific YAML configurations.
+
+The release launcher owns runtime state such as checkpoint paths, output
+directories, topology, and tracking identity. This module deliberately loads
+only model/data/training science knobs into the typed release configuration.
+Keeping that boundary explicit prevents a portable
+scientific config from silently capturing one cluster's filesystem state.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+SCHEMA = "pixelumm_scientific_config_v1"
+SECTIONS = ("model", "data", "training")
+
+
+class _UniqueKeyLoader(yaml.SafeLoader):
+    """Safe YAML loader that rejects duplicate mapping keys."""
+
+
+def _construct_unique_mapping(
+    loader: _UniqueKeyLoader,
+    node: yaml.MappingNode,
+    deep: bool = False,
+) -> dict[Any, Any]:
+    # PyYAML normally keeps the last duplicate key.  A reviewed scientific
+    # contract must reject that ambiguity instead of silently changing it.
+    loader.flatten_mapping(node)
+    mapping: dict[Any, Any] = {}
+    for key_node, value_node in node.value:
+        key = loader.construct_object(key_node, deep=deep)
+        try:
+            duplicated = key in mapping
+        except TypeError as exc:
+            raise yaml.constructor.ConstructorError(
+                "while constructing a mapping",
+                node.start_mark,
+                "found an unhashable mapping key",
+                key_node.start_mark,
+            ) from exc
+        if duplicated:
+            raise yaml.constructor.ConstructorError(
+                "while constructing a mapping",
+                node.start_mark,
+                f"found duplicate key {key!r}",
+                key_node.start_mark,
+            )
+        mapping[key] = loader.construct_object(value_node, deep=deep)
+    return mapping
+
+
+_UniqueKeyLoader.add_constructor(
+    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
+    _construct_unique_mapping,
+)
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _cli_scalar(value: Any, *, location: str) -> str:
+    if isinstance(value, bool):
+        return "True" if value else "False"
+    if value is None:
+        return "None"
+    if isinstance(value, (str, int, float)):
+        return str(value)
+    raise RuntimeError(
+        f"scientific config value must be a scalar at {location}; "
+        f"got {type(value).__name__}"
+    )
+
+
+def load_scientific_config(
+    path: Path,
+    *,
+    expected_sha256: str | None = None,
+) -> dict[str, dict[str, Any]]:
+    """Load one fail-closed, sectioned scientific configuration."""
+
+    path = Path(path).resolve()
+    if not path.is_file() or path.stat().st_size == 0:
+        raise RuntimeError(f"scientific config is absent or empty: {path}")
+    if expected_sha256 is not None and sha256_file(path) != expected_sha256:
+        raise RuntimeError(f"scientific config identity mismatch: {path}")
+    try:
+        payload = yaml.load(
+            path.read_text(encoding="utf-8"),
+            Loader=_UniqueKeyLoader,
+        )
+    except (OSError, yaml.YAMLError) as exc:
+        raise RuntimeError(f"scientific config is invalid: {path}") from exc
+    if not isinstance(payload, dict):
+        raise RuntimeError("scientific config must be a YAML mapping")
+    expected_top_level = {"schema", *SECTIONS}
+    if set(payload) != expected_top_level:
+        raise RuntimeError(
+            "scientific config top-level keys changed: "
+            f"expected={sorted(expected_top_level)} actual={sorted(payload)}"
+        )
+    if payload["schema"] != SCHEMA:
+        raise RuntimeError(
+            f"unexpected scientific config schema: {payload['schema']!r}"
+        )
+
+    sections: dict[str, dict[str, Any]] = {}
+    owners: dict[str, str] = {}
+    for section in SECTIONS:
+        values = payload[section]
+        if not isinstance(values, dict) or not values:
+            raise RuntimeError(
+                f"scientific config section {section!r} must be a non-empty mapping"
+            )
+        normalized: dict[str, Any] = {}
+        for raw_name, value in values.items():
+            if not isinstance(raw_name, str) or not raw_name:
+                raise RuntimeError(
+                    f"scientific config contains an invalid key in {section!r}"
+                )
+            if raw_name.startswith("--") or not raw_name.replace("_", "a").isalnum():
+                raise RuntimeError(
+                    f"scientific config key must be a dataclass field name: {raw_name!r}"
+                )
+            if raw_name in owners:
+                raise RuntimeError(
+                    f"scientific config key {raw_name!r} appears in both "
+                    f"{owners[raw_name]!r} and {section!r}"
+                )
+            _cli_scalar(value, location=f"{section}.{raw_name}")
+            owners[raw_name] = section
+            normalized[raw_name] = value
+        sections[section] = normalized
+    return sections