[New Model] Lance (ByteDance) — unified AR+diffusion multimodal (#3697)

Adds Lance (bytedance-research/Lance), a 3B unified AR+diffusion
multimodal model on a Qwen2.5-VL backbone. Lance is BAGEL-lineage:
it reuses vLLM-Omni's BAGEL Mixture-of-Transformers core verbatim and
specializes only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the
checkpoint layout.

All six single-stage Lance modalities verified end-to-end on 1x B300:
t2i, t2v, image_edit, video_edit, x2t_image (image understanding) and
x2t_video (video understanding). 3-D mRoPE is wired through
BagelRotaryEmbedding for the x2t / video / edit paths. Includes the
upstream gap-analysis quality fixes (task-specific system prompts, full
Qwen chat template, upstream inference defaults and aspect-ratio
buckets).

Ships an offline example (examples/offline_inference/lance/), an
OpenAI-compatible online-serving example
(examples/online_serving/lance/) and an e2e test
(tests/e2e/online_serving/test_lance.py).

See LANCE_PR.md for the architecture rationale, per-modality status and
the two remaining follow-ups (two-stage AR+DiT, video_edit polish).

Closes #3697.

Signed-off-by: lishunyang <lishunyang03@gmail.com>

LANCE_PR.md

@@ -0,0 +1,76 @@
+# [New Model] Lance (ByteDance) — #3697
+
+Adds Lance (`bytedance-research/Lance`), a 3B unified AR+diffusion
+multimodal model on a Qwen2.5-VL backbone.
+
+## Why this is BAGEL, not WAN
+
+The issue thread suggested remapping against `WanTransformer3DModel`. The
+released `Lance_3B` checkpoint is instead **BAGEL Mixture-of-Transformers**:
+every layer carries `mlp_moe_gen` / `q_proj_moe_gen` /
+`input_layernorm_moe_gen` twins plus `q_norm`/`k_norm`, with top-level
+`vae2llm` / `llm2vae` / `time_embedder` / `latent_pos_embed` connectors —
+identical naming to vLLM-Omni's BAGEL. The only WAN element is the VAE.
+Lance therefore reuses BAGEL's transformer core verbatim and specializes
+only the LLM (Qwen2.5-VL-MoT), ViT (Qwen2.5-VL vision), VAE (Wan2.2) and
+checkpoint layout.
+
+## Scope and status
+
+All six single-stage Lance modalities, verified end-to-end on 1× B300:
+
+| Modality | Status |
+|----------|--------|
+| `t2i` (text→image) | working — 1024×1024, 0 missing keys |
+| `t2v` (text→video) | working — photorealistic frames with motion; 768×768 trained bucket |
+| `image_edit` (img2img) | working — Lance-native VAE prefill scatters Wan2.2 latents into the LLM query sequence |
+| `video_edit` (video2video) | working — multi-frame VAE prefill + 3-D mRoPE; needs the `Lance_3B_Video` checkpoint |
+| `x2t_image` (image understanding) | working — Qwen2.5-VL ViT prefill + 3-D mRoPE |
+| `x2t_video` (video understanding) | working — Qwen2-VL video processor + 3-D `video_grid_thw` |
+
+`rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired
+through `BagelRotaryEmbedding`: t2i uses scalar position ids
+(BAGEL-equivalent); x2t / video / edit thread per-axis 3-D position ids.
+Geometry constants (`latent_patch_size_spatial = 1`, `max_latent_size = 64`,
+untied `lm_head`, Wan2.2 VAE `z=48 /16 spatial /4 temporal`) were locked in
+by inspecting the released checkpoint directly.
+
+The two-stage AR+DiT topology is the only deferred path (needs
+`LanceConfig` / `LanceProcessor` in the `vllm` package — a separate
+human-authored PR per vLLM's contribution policy).
+
+## Test plan
+
+Offline (`examples/offline_inference/lance/end2end.py`), e.g. t2i:
+
+```bash
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance \
+    --prompts "a corgi astronaut on the moon, cinematic" \
+    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 --output ./out
+```
+
+`--modality {text2video,img2img,video2video,img2text,video2text}` selects
+the other paths; video paths use `--model bytedance-research/Lance/Lance_3B_Video`.
+
+Online (OpenAI-compatible, `examples/online_serving/lance/`):
+
+```bash
+bash examples/online_serving/lance/run_server.sh        # vllm serve … --omni
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "A cute corgi astronaut on the moon, cinematic" \
+    --modality text2img --output corgi.png
+```
+
+E2E coverage: `tests/e2e/online_serving/test_lance.py`. Tested on 1×
+NVIDIA B300, `torch 2.11.0+cu130`, stock attention backend.
+
+## Follow-ups
+
+1. **Two-stage AR + DiT** — register `LanceConfig` / `LanceProcessor` in
+   the `vllm` package; separate human-authored PR per vLLM policy.
+2. **`video_edit` quality polish** — output is more abstract than `t2v`
+   at the same resolution (position-id offset between the VAE-prefill and
+   gen-latent blocks). Functionally correct end-to-end.
+
+Closes #3697.

examples/offline_inference/lance/README.md

@@ -0,0 +1,50 @@
+# Lance (ByteDance) — offline inference
+
+[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
+autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
+**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
+checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
+implements it by reusing the BAGEL transformer core and specializing only the
+ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
+
+This example covers all six Lance modalities from the upstream HF model card:
+`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
+`x2t_video` (video understanding).
+
+## Hardware
+
+Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
+
+## Run
+
+```bash
+# Text-to-image
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance \
+    --prompts "a corgi astronaut on the moon, cinematic" \
+    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
+    --output ./out
+
+# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
+# choices for all six task variants)
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance/Lance_3B_Video --modality text2video \
+    --num-frames 25 --video-height 480 --video-width 768 \
+    --prompts "a cat playing piano, cinematic" \
+    --steps 30 --fps 8 --output ./out
+```
+
+`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
+the 3-D `latent_pos_embed` table is loaded; the other paths can point at
+the top-level `bytedance-research/Lance` repo and resolve the right
+sub-checkpoint automatically.
+
+The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
+`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
+
+## Defaults
+
+Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
+text CFG 4.0, seed 42, 768×768.  For the understanding paths (`img2text` /
+`video2text`), sampling is enabled by default at `--text-temperature 0.8`
+because Lance's greedy decoder emits an immediate EOS for many prompts.

examples/offline_inference/lance/end2end.py

@@ -0,0 +1,268 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance-3B offline inference (single-stage).
+
+Lance is BAGEL-lineage; this example covers the single-stage paths wired in
+this PR:
+
+    text2img  — image generation from text
+    img2text  — image understanding (caption / VQA)
+
+Examples:
+
+    python examples/offline_inference/lance/end2end.py \
+        --model bytedance-research/Lance \
+        --prompts "a corgi astronaut on the moon, cinematic" \
+        --output ./out
+
+    python examples/offline_inference/lance/end2end.py \
+        --model bytedance-research/Lance --modality img2text \
+        --image-path /path/to/photo.jpg \
+        --prompts "Describe this image in detail."
+"""
+
+import argparse
+import os
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
+    p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
+    p.add_argument(
+        "--modality",
+        default="text2img",
+        choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video"],
+        help="Lance single-stage modality.",
+    )
+    p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
+    p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
+    p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
+    p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
+    p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
+    p.add_argument("--fps", type=int, default=8, help="Output video FPS when saving MP4.")
+    p.add_argument("--output", type=str, default=".", help="Output directory.")
+    p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
+    p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
+    p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
+    p.add_argument("--negative-prompt", type=str, default=None)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
+    p.add_argument(
+        "--do-sample",
+        action="store_true",
+        default=True,
+        help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
+    )
+    p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
+    p.add_argument(
+        "--text-temperature",
+        type=float,
+        default=0.8,
+        help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
+    )
+    p.add_argument(
+        "--deploy-config",
+        type=str,
+        default="vllm_omni/deploy/lance.yaml",
+        help="Deploy YAML (single-stage Lance).",
+    )
+
+    from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
+
+    nullify_stage_engine_defaults(p)
+    return p.parse_args()
+
+
+from vllm_omni.diffusion.models.lance.prompts import (  # noqa: E402
+    VIDEO_PAD,
+    VISION_END,
+    VISION_START,
+    render_lance_prompt,
+)
+
+# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
+# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
+_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
+
+
+def _format_text2img_prompts(prompts):
+    return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
+
+
+def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
+    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
+    return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
+
+
+def _format_img2img_prompts(prompts, image_path):
+    import os as _os
+
+    if not image_path or not _os.path.exists(image_path):
+        raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
+    from PIL import Image as _Image
+
+    img = _Image.open(image_path).convert("RGB")
+    return [
+        {
+            "prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"img2img": img},
+            "modalities": ["image"],
+        }
+        for p in prompts
+    ]
+
+
+def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
+    import os as _os
+
+    if not video_path or not _os.path.exists(video_path):
+        raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
+    import imageio.v3 as iio
+
+    video = iio.imread(video_path)
+    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
+    return [
+        {
+            "prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"video": video},
+            "modalities": ["video"],
+            "extra_args": extra,
+        }
+        for p in prompts
+    ]
+
+
+def _format_video2text_prompts(prompts, video_path):
+    import os as _os
+
+    if not video_path or not _os.path.exists(video_path):
+        raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
+    import imageio.v3 as iio
+
+    video = iio.imread(video_path)
+    return [
+        {
+            "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"video": video},
+            "modalities": ["text"],
+        }
+        for p in prompts
+    ]
+
+
+def _format_img2text_prompts(prompts, image_path):
+    from PIL import Image
+
+    if not image_path or not os.path.exists(image_path):
+        raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
+    img = Image.open(image_path).convert("RGB")
+    return [
+        {
+            "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"image": img},
+            "modalities": ["text"],
+        }
+        for p in prompts
+    ]
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output, exist_ok=True)
+
+    if args.txt_prompts:
+        with open(args.txt_prompts, encoding="utf-8") as f:
+            prompts = [ln.strip() for ln in f if ln.strip()]
+    elif args.modality == "img2text":
+        prompts = args.prompts or ["Describe this image in detail."]
+    else:
+        prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
+
+    from vllm_omni.entrypoints.omni import Omni
+
+    omni_kwargs = vars(args).copy()
+    omni_kwargs["deploy_config"] = args.deploy_config
+    omni_kwargs["model"] = args.model
+    omni = Omni(**omni_kwargs)
+
+    if args.modality == "img2text":
+        formatted = _format_img2text_prompts(prompts, args.image_path)
+    elif args.modality == "text2video":
+        formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
+    elif args.modality == "video2text":
+        formatted = _format_video2text_prompts(prompts, args.video_path)
+    elif args.modality == "img2img":
+        formatted = _format_img2img_prompts(prompts, args.image_path)
+    elif args.modality == "video2video":
+        formatted = _format_video2video_prompts(
+            prompts, args.video_path, args.num_frames, args.video_height, args.video_width
+        )
+    else:
+        formatted = _format_text2img_prompts(prompts)
+
+    params_list = omni.default_sampling_params_list
+    diffusion_params = params_list[0]  # single-stage: one param set
+    diffusion_params.num_inference_steps = args.steps  # type: ignore
+    if args.seed is not None:
+        diffusion_params.seed = args.seed  # type: ignore
+    extra = getattr(diffusion_params, "extra_args", {}) or {}
+    extra["cfg_text_scale"] = args.cfg_text_scale
+    extra["timestep_shift"] = args.timestep_shift
+    if args.modality in ("img2text", "video2text"):
+        extra["max_think_tokens"] = args.max_text_tokens
+        extra["do_sample"] = args.do_sample
+        extra["text_temperature"] = args.text_temperature
+    if args.negative_prompt is not None:
+        extra["negative_prompt"] = args.negative_prompt
+    diffusion_params.extra_args = extra  # type: ignore
+
+    outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
+
+    if args.modality in ("img2text", "video2text"):
+        for i, req_output in enumerate(outputs):
+            text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
+            if text is None and hasattr(req_output, "custom_output"):
+                text = (req_output.custom_output or {}).get("text_output")
+            print(f"[Output {i}] {text}")
+        return
+
+    if args.modality in ("text2video", "video2video"):
+        for i, req_output in enumerate(outputs):
+            custom = getattr(req_output, "custom_output", None) or {}
+            frames = custom.get("video_frames")
+            if not frames:
+                print(f"[Output {i}] no video frames returned")
+                continue
+            mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
+            try:
+                import imageio.v3 as iio
+                import numpy as np
+
+                arr = np.stack([np.asarray(f) for f in frames], axis=0)
+                iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
+                print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
+            except Exception as exc:
+                print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
+                frame_dir = os.path.join(args.output, f"lance_{i}_frames")
+                os.makedirs(frame_dir, exist_ok=True)
+                for j, frame in enumerate(frames):
+                    frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
+                print(f"  → {frame_dir} ({len(frames)} frames)")
+        return
+
+    idx = 0
+    for req_output in outputs:
+        images = getattr(req_output, "images", None)
+        if not images:
+            continue
+        for j, img in enumerate(images):
+            path = os.path.join(args.output, f"lance_{idx}_{j}.png")
+            img.save(path)
+            print(f"[Output] Saved image to {path}")
+        idx += 1
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/lance/README.md

@@ -0,0 +1,40 @@
+# Lance — online serving
+
+OpenAI-compatible chat completions API for Lance, served by `vllm-omni`.
+
+## Start the server
+
+```bash
+bash examples/online_serving/lance/run_server.sh
+# or, with overrides
+MODEL=bytedance-research/Lance \
+DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
+PORT=8091 \
+    bash examples/online_serving/lance/run_server.sh
+```
+
+The default deploy config is the single-stage `lance.yaml`; for `video_edit`
+or `text2video` workloads, point ``DEPLOY_CONFIG`` at a video-checkpoint
+config (or pass ``--model bytedance-research/Lance/Lance_3B_Video`` to
+`run_server.sh`).
+
+## Send a request
+
+```bash
+# text-to-image
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "A cute corgi astronaut on the moon, cinematic" \
+    --modality text2img \
+    --output corgi.png
+
+# image edit
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "Convert this into a vibrant cartoon-style illustration" \
+    --modality img2img \
+    --image-url path/to/photo.png \
+    --output edited.png
+```
+
+The client is shared with BAGEL — same OpenAI message format, same
+``modalities`` and ``num_inference_steps`` / ``seed`` / ``height`` /
+``width`` knobs.

examples/online_serving/lance/openai_chat_client.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""
+Lance OpenAI-compatible chat client for image generation and multimodal tasks.
+
+Usage:
+    python examples/online_serving/lance/openai_chat_client.py --prompt "A cute cat" --output output.png
+    python examples/online_serving/lance/openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
+"""
+
+import argparse
+import base64
+from pathlib import Path
+
+import requests
+
+
+def generate_image(
+    prompt: str,
+    server_url: str = "http://localhost:8091",
+    image_url: str | None = None,
+    height: int | None = None,
+    width: int | None = None,
+    steps: int | None = None,
+    seed: int | None = None,
+    negative_prompt: str | None = None,
+    modality: str = "text2img",  # "text2img" (default), "img2img", "img2text", "text2text"
+) -> bytes | str | None:
+    """Generate an image or text using the chat completions API.
+
+    Args:
+        prompt: Text description or prompt
+        server_url: Server URL
+        image_url: URL or path to input image (for img2img/img2text)
+        height: Image height in pixels
+        width: Image width in pixels
+        steps: Number of inference steps
+        seed: Random seed
+        negative_prompt: Negative prompt
+        modality: Task modality hint
+
+    Returns:
+        Image bytes (for image outputs) or Text string (for text outputs) or None if failed
+    """
+
+    # Construct Message Content
+    content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
+
+    if image_url:
+        # Check if local file
+        if Path(image_url).exists():
+            with open(image_url, "rb") as f:
+                b64_data = base64.b64encode(f.read()).decode("utf-8")
+                final_image_url = f"data:image/jpeg;base64,{b64_data}"
+        else:
+            final_image_url = image_url
+
+        content.append({"type": "image_url", "image_url": {"url": final_image_url}})
+
+    messages = [{"role": "user", "content": content}]
+
+    # Build request payload with all parameters at top level
+    # Note: vLLM ignores "extra_body", so we put parameters directly in the payload
+    payload = {"messages": messages}
+
+    # Set output modalities at top level
+    if modality == "text2img" or modality == "img2img":
+        payload["modalities"] = ["image"]
+    elif modality == "img2text" or modality == "text2text":
+        payload["modalities"] = ["text"]
+
+    # Add generation parameters directly to payload
+    if height is not None:
+        payload["height"] = height
+    if width is not None:
+        payload["width"] = width
+    if steps is not None:
+        payload["num_inference_steps"] = steps
+    if seed is not None:
+        payload["seed"] = seed
+    if negative_prompt:
+        payload["negative_prompt"] = negative_prompt
+
+    # Send request
+    try:
+        print(f"Sending request to {server_url} with modality {modality}...")
+        response = requests.post(
+            f"{server_url}/v1/chat/completions",
+            headers={"Content-Type": "application/json"},
+            json=payload,
+            timeout=300,
+        )
+        response.raise_for_status()
+        data = response.json()
+
+        # Extract content - check ALL choices since server may return multiple
+        # (e.g., text in choices[0], image in choices[1])
+        choices = data.get("choices", [])
+
+        # First pass: look for image output in any choice
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+
+            # Handle Image Output
+            if isinstance(choice_content, list) and len(choice_content) > 0:
+                first_item = choice_content[0]
+                if isinstance(first_item, dict) and "image_url" in first_item:
+                    img_url_str = first_item["image_url"].get("url", "")
+                    if img_url_str.startswith("data:image"):
+                        _, b64_data = img_url_str.split(",", 1)
+                        return base64.b64decode(b64_data)
+
+        # Second pass: look for text output if no image found
+        for choice in choices:
+            choice_content = choice.get("message", {}).get("content")
+            if isinstance(choice_content, str) and choice_content:
+                return choice_content
+
+        print(f"Unexpected response format: {choices}")
+        return None
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
+    parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
+    parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
+    parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
+
+    # Modality Control
+    parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
+    parser.add_argument(
+        "--modality",
+        "-m",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Task modality",
+    )
+
+    # Generation Params
+    parser.add_argument("--height", type=int, default=512, help="Image height")
+    parser.add_argument("--width", type=int, default=512, help="Image width")
+    parser.add_argument("--steps", type=int, default=25, help="Inference steps")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--negative", help="Negative prompt")
+
+    args = parser.parse_args()
+
+    print(f"Mode: {args.modality}")
+    if args.image_url:
+        print(f"Input Image: {args.image_url}")
+
+    result = generate_image(
+        prompt=args.prompt,
+        server_url=args.server,
+        image_url=args.image_url,
+        height=args.height,
+        width=args.width,
+        steps=args.steps,
+        seed=args.seed,
+        negative_prompt=args.negative,
+        modality=args.modality,
+    )
+
+    if result:
+        if isinstance(result, bytes):
+            # It's an image
+            output_path = Path(args.output)
+            output_path.write_bytes(result)
+            print(f"Image saved to: {output_path}")
+            print(f"Size: {len(result) / 1024:.1f} KB")
+        elif isinstance(result, str):
+            # It's text
+            print("Response:")
+            print(result)
+    else:
+        print("Failed to generate response")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/lance/run_server.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Lance online serving startup script
+
+MODEL="${MODEL:-bytedance-research/Lance}"
+PORT="${PORT:-8091}"
+DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/lance.yaml}"
+
+echo "Starting Lance server..."
+echo "Model:         $MODEL"
+echo "Deploy config: $DEPLOY_CONFIG"
+echo "Port:          $PORT"
+
+vllm serve "$MODEL" --omni \
+    --deploy-config "$DEPLOY_CONFIG" \
+    --port "$PORT"

tests/e2e/online_serving/test_lance.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end online serving test for Lance (single-stage).
+
+Verifies that the Lance pipeline can serve text-to-image and image-edit
+requests via the OpenAI-compatible chat completions API exposed by
+``vllm-omni serve``.
+
+Equivalent to running:
+
+    vllm-omni serve "bytedance-research/Lance" --omni \\
+        --deploy-config vllm_omni/deploy/lance.yaml --port 8091
+
+    # text2img
+    python3 examples/online_serving/lance/openai_chat_client.py \\
+        --prompt "A cute corgi astronaut" --modality text2img
+"""
+
+import base64
+import os
+from io import BytesIO
+
+import pytest
+from vllm.assets.image import ImageAsset
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.diffusion.models.lance.prompts import (
+    VIDEO_PAD,
+    VISION_END,
+    VISION_START,
+    render_lance_prompt,
+)
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+MODEL = "bytedance-research/Lance"
+STAGE_CONFIGS_PATH = get_deploy_config_path("ci/lance.yaml")
+
+TEXT2IMG_PROMPT = "A cute corgi astronaut on the moon, cinematic"
+IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration"
+
+# Lance vision block: <|vision_start|><|video_pad|><|vision_end|>
+_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=STAGE_CONFIGS_PATH,
+        stage_init_timeout=300,
+    ),
+]
+
+
+def _build_text2img_messages(prompt: str) -> list[dict]:
+    """Build OpenAI-format messages for text2img generation."""
+    rendered = render_lance_prompt("t2i", prompt)
+    return [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": rendered}],
+        }
+    ]
+
+
+def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
+    """Build OpenAI-format messages for img2img generation."""
+    rendered = render_lance_prompt("image_edit", prompt, vision_token=_VISION_BLOCK)
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": rendered},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
+                },
+            ],
+        }
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_lance_text2img_online(omni_server, openai_client) -> None:
+    """Lance text2img via the OpenAI-compatible chat completions API."""
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_text2img_messages(TEXT2IMG_PROMPT),
+        "modalities": ["image"],
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            "guidance_scale": 0.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_lance_img2img_online(omni_server, openai_client) -> None:
+    """Lance image_edit via the OpenAI-compatible chat completions API."""
+    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    buffer = BytesIO()
+    input_image.save(buffer, format="JPEG")
+    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64),
+        "modalities": ["image"],
+        "extra_body": {
+            "num_inference_steps": 2,
+            "guidance_scale": 0.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/helpers/stage_config.py

@@ -495,6 +495,15 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = {
             },
         ],
     },
+    "lance": {
+        "base_config": "lance.yaml",
+        "stages": [
+            {
+                "stage_id": 0,
+                "max_num_seqs": 1,
+            },
+        ],
+    },
     "bagel_mooncake": {
         "base_config": "bagel.yaml",
         "stages": [

vllm_omni/config/pipeline_registry.py

@@ -65,6 +65,17 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.bagel.pipeline",
         "BAGEL_SINGLE_STAGE_PIPELINE",
     ),
+    # Lance (ByteDance) — BAGEL-lineage unified AR+diffusion.
+    # Single-stage is the runnable/registered topology; lance_two_stage is a
+    # follow-up (needs LanceConfig/LanceProcessor in the vllm package).
+    "lance": (
+        "vllm_omni.model_executor.models.lance.pipeline",
+        "LANCE_SINGLE_STAGE_PIPELINE",
+    ),
+    "lance_two_stage": (
+        "vllm_omni.model_executor.models.lance.pipeline",
+        "LANCE_PIPELINE",
+    ),
     "glm_image": (
         "vllm_omni.model_executor.models.glm_image.pipeline",
         "GLM_IMAGE_PIPELINE",

vllm_omni/deploy/lance.yaml

@@ -0,0 +1,29 @@
+# Lance-3B (ByteDance) deploy — single-stage, CUDA. Target: 1x NVIDIA A100 80GB.
+#
+# Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion). The single-stage
+# DiT stage is self-contained: it owns the Qwen2-MoT LLM, the Qwen2.5-VL ViT,
+# the Wan2.2 VAE and the tokenizer (text2img / image_edit / x2t_image).
+#
+# Lance-3B BF16 footprint: ~7GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE — comfortably
+# within one A100 80GB. The video checkpoint (Lance_3B_Video) is a follow-up.
+#
+# Select via:
+#   vllm serve bytedance-research/Lance --omni \
+#       --deploy-config vllm_omni/deploy/lance.yaml
+# or programmatically:
+#   Omni(model="bytedance-research/Lance",
+#        deploy_config_path="vllm_omni/deploy/lance.yaml")
+
+pipeline: lance
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_batched_tokens: 32768
+    max_num_seqs: 1
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    devices: "0"
+    default_sampling_params:
+      seed: 42

vllm_omni/diffusion/data.py

@@ -790,6 +790,21 @@ class OmniDiffusionConfig:
         self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
         self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
 
+    @staticmethod
+    def _looks_like_lance_subfolder(model: str | None) -> bool:
+        """Return True when ``--model`` points at a Lance per-component subfolder.
+
+        Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and
+        ``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may
+        reasonably hand the AR-style sub-checkpoint path directly.  The
+        ``LancePipeline`` constructor knows to walk up to the repo root from
+        either subfolder name.
+        """
+        if not model:
+            return False
+        base = os.path.basename(str(model).rstrip("/"))
+        return base in {"Lance_3B", "Lance_3B_Video"}
+
     def enrich_config(self) -> None:
         """Load model metadata from HuggingFace and populate config fields.
 
@@ -843,6 +858,15 @@ class OmniDiffusionConfig:
             else:
                 cfg = get_hf_file_to_dict("config.json", self.model)
                 if cfg is None:
+                    # Lance ships its top-level config.json one directory above
+                    # the per-checkpoint subfolders (``Lance_3B/`` or
+                    # ``Lance_3B_Video/``).  Try to recover that case before
+                    # raising.
+                    if self._looks_like_lance_subfolder(self.model):
+                        self.model_class_name = "LancePipeline"
+                        self.set_tf_model_config(TransformerConfig())
+                        self.update_multimodal_support()
+                        return
                     raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
 
                 self.set_tf_model_config(TransformerConfig.from_dict(cfg))
@@ -853,6 +877,19 @@ class OmniDiffusionConfig:
                     self.model_class_name = "BagelPipeline"
                     self.set_tf_model_config(TransformerConfig())
                     self.update_multimodal_support()
+                elif (
+                    model_type == "lance"
+                    or "LancePipeline" in architectures
+                    or cfg.get("model_name") == "Lance"
+                    or self._looks_like_lance_subfolder(self.model)
+                ):
+                    # Lance ships a non-HF top-level config.json (model_name only)
+                    # plus per-component subfolders; resolve to the Lance pipeline.
+                    # Also accept --model pointing directly at the ``Lance_3B`` or
+                    # ``Lance_3B_Video`` subfolder by walking up to the repo root.
+                    self.model_class_name = "LancePipeline"
+                    self.set_tf_model_config(TransformerConfig())
+                    self.update_multimodal_support()
                 elif model_type == "neo_chat":
                     self.model_class_name = "SenseNovaU1Pipeline"
                     self.tf_model_config = TransformerConfig()

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -108,7 +108,16 @@ class BagelRotaryEmbedding(nn.Module):
     linear, dynamic-NTK, YaRN, …), we delegate the ``inv_freq`` /
     ``attention_scaling`` computation to HF's ``ROPE_INIT_FUNCTIONS`` so
     that the frequency basis and scaling factor are identical to the
-    original checkpoint.  This module has no learnable parameters.
+    original checkpoint.
+
+    For Qwen2.5-VL-style multimodal RoPE (``rope_scaling.rope_type == "mrope"``)
+    the ``inv_freq`` basis is the standard default-rope one; the difference
+    is that position ids are 3-D ``(t, h, w)`` per token and the
+    ``mrope_section`` describes how the head dimension is split across axes.
+    This module accepts either 2-D scalar position ids ``(B, S)`` or 3-D
+    multimodal position ids ``(B, 3, S)`` and dispatches accordingly so the
+    same module works for both BAGEL (1-D rope) and Lance (Qwen2.5-VL mrope).
+    This module has no learnable parameters.
     """
 
     def __init__(self, config):
@@ -122,10 +131,17 @@ class BagelRotaryEmbedding(nn.Module):
         rope_type = (
             rope_scaling.get("rope_type") or rope_scaling.get("type") or rope_parameters.get("rope_type") or "default"
         )
+        # Cache mrope_section for the forward-time section-split.
+        self._mrope_section: list[int] | None = None
+        if rope_type == "mrope":
+            section = rope_scaling.get("mrope_section") or rope_parameters.get("mrope_section")
+            if section is None:
+                raise ValueError("rope_scaling.rope_type == 'mrope' requires 'mrope_section'.")
+            self._mrope_section = list(section)
 
-        if rope_type == "default":
-            # transformers>=5.0 removed the 'default' entry from
-            # ROPE_INIT_FUNCTIONS; use the plain sinusoidal formula.
+        if rope_type in ("default", "mrope"):
+            # mrope shares the default sinusoidal frequency basis; the
+            # multimodal split happens at forward time, not at init.
             rope_theta = (
                 getattr(config, "rope_theta", None)
                 or rope_parameters.get("rope_theta")
@@ -149,11 +165,36 @@ class BagelRotaryEmbedding(nn.Module):
 
         Args:
             x: Input tensor (only used for dtype inference).
-            position_ids: Position indices, shape (batch_size, seq_len).
+            position_ids: Either 2-D scalar ``(batch_size, seq_len)`` for plain
+                1-D RoPE, or 3-D multimodal ``(batch_size, 3, seq_len)`` for
+                Qwen2.5-VL-style mRoPE.  The latter is auto-detected from
+                ``position_ids.ndim``.
 
         Returns:
             cos, sin: Rotary embeddings, each of shape (batch_size, seq_len, dim).
         """
+        if position_ids.ndim == 3 and self._mrope_section is not None:
+            # multimodal path: position_ids is (B, 3, S) with rows = (t, h, w).
+            # Compute per-axis frequencies, then assemble the per-section
+            # rotary basis matching Qwen2-VL's ``apply_multimodal_rotary_pos_emb``.
+            B = position_ids.shape[0]
+            inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(B, 3, -1, 1)
+            position_ids_expanded = position_ids[:, :, None, :].float()
+            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3)
+            # ``freqs`` is (B, 3, S, head_dim/2); double along the last axis
+            # to get the full ``head_dim`` rotary basis.
+            emb = torch.cat((freqs, freqs), dim=-1)  # (B, 3, S, head_dim)
+            cos_per_axis = emb.cos() * self.attention_scaling
+            sin_per_axis = emb.sin() * self.attention_scaling
+            # ``mrope_section`` (e.g. [16, 24, 24] for Qwen2.5-VL) sums to
+            # head_dim/2; doubled it sums to head_dim and cycles axis = i % 3.
+            sec_full = self._mrope_section * 2
+            cos_split = cos_per_axis.split(sec_full, dim=-1)
+            sin_split = sin_per_axis.split(sec_full, dim=-1)
+            cos = torch.cat([c[:, i % 3] for i, c in enumerate(cos_split)], dim=-1)
+            sin = torch.cat([s[:, i % 3] for i, s in enumerate(sin_split)], dim=-1)
+            return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
         inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
         position_ids_expanded = position_ids[:, None, :].float()
         freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
@@ -1982,7 +2023,9 @@ class Bagel(nn.Module):
                 "num_branches": num_branches,
                 "seq_len": seq_len,
                 "batched_query_lens": packed_seqlens.repeat(num_branches),
-                "batched_position_ids": torch.cat(branches_pid),
+                # Concatenate along the sequence (last) dim so this works for both
+                # scalar position ids ``(S,)`` and 3-D mRoPE ``(3, S)`` layouts.
+                "batched_position_ids": torch.cat(branches_pid, dim=-1),
                 "batched_kv_lens": torch.cat(branches_kvl),
                 "batched_query_indexes": torch.cat(
                     [qi + merged_offsets[b_idx] for b_idx, qi in enumerate(branches_qi)]

vllm_omni/diffusion/models/lance/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance (ByteDance) diffusion model components.
+
+Lance is a unified autoregressive + diffusion multimodal model on a
+Qwen2.5-VL-3B backbone.  Architecturally it is the BAGEL family (ByteDance
+Mixture-of-Transformers): the released ``Lance_3B`` checkpoint uses the exact
+same ``*_moe_gen`` MoT weight layout as BAGEL, plus ``vae2llm`` / ``llm2vae`` /
+``time_embedder`` / ``latent_pos_embed`` connectors.  The deltas vs BAGEL are:
+
+  * backbone is Qwen2.5-VL (mRoPE) instead of Qwen2,
+  * understanding ViT is Qwen2.5-VL vision (not SigLIP), loaded from the base
+    ``Qwen/Qwen2.5-VL-3B-Instruct`` rather than from the Lance checkpoint,
+  * VAE is Wan2.2 (reused from the vLLM-Omni WAN path) instead of the BAGEL AE,
+  * video path adds 3D latent position embeddings (follow-up; this module
+    implements the image path first).
+
+Because Lance is BAGEL-lineage, the transformer core is reused verbatim from
+``vllm_omni.diffusion.models.bagel.bagel_transformer`` and only the pipeline
+wiring (ViT / VAE / checkpoint layout) is specialized here.
+"""
+
+from vllm_omni.diffusion.models.lance.pipeline_lance import (
+    LancePipeline,
+    get_lance_post_process_func,
+)
+
+__all__ = [
+    "LancePipeline",
+    "get_lance_post_process_func",
+]

vllm_omni/diffusion/models/lance/lance_transformer.py

@@ -0,0 +1,944 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance transformer pieces.
+
+The Lance LLM is BAGEL's Qwen2-MoT transformer verbatim — the released
+``Lance_3B`` checkpoint uses the identical ``*_moe_gen`` / ``q_norm`` /
+``vae2llm`` / ``llm2vae`` / ``time_embedder`` / ``latent_pos_embed`` layout, so
+``Bagel`` / ``Qwen2MoTForCausalLM`` / ``Qwen2MoTConfig`` / ``NaiveCache`` are
+re-exported unchanged.  Only the understanding ViT differs (Qwen2.5-VL vision
+instead of SigLIP) and the video path adds a 3D latent position embedding.
+
+This module also provides :class:`LanceBagel`, a thin :class:`Bagel` subclass
+that overrides the two ViT entry points to consume Qwen2.5-VL's packed
+``pixel_values`` + ``image_grid_thw`` layout directly (BAGEL itself assumes
+SigLIP-style ``(C, H, W)`` tensors that get patchified inside the model).
+
+NOTE ON mRoPE: Lance's backbone is Qwen2.5-VL and its understanding /
+video paths use multimodal RoPE (``mrope_section=[16,24,24]``).  BAGEL's
+``BagelRotaryEmbedding`` is plain 1-D RoPE on scalar position ids.  For the
+text2img generation path Lance assigns scalar positions (the gen latent block
+shares a single rope position, same as BAGEL), so the reused rotary is
+correct there.  Full mRoPE for the x2t / video understanding path is a
+follow-up — see ``LancePositionEmbedding3D``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch import nn
+
+# Lance LLM == BAGEL Qwen2-MoT: reuse verbatim (checkpoint weight names match).
+from vllm_omni.diffusion.models.bagel.bagel_transformer import (  # noqa: F401
+    Bagel,
+    BaseNavitOutputWithPast,
+    MLPconnector,
+    NaiveCache,
+    PackedAttentionMoT,
+    PositionEmbedding,
+    Qwen2MoTConfig,
+    Qwen2MoTDecoderLayer,
+    Qwen2MoTForCausalLM,
+    Qwen2MoTModel,
+    TimestepEmbedder,
+    get_1d_sincos_pos_embed_from_grid,
+    patchify,
+)
+
+__all__ = [
+    "Bagel",
+    "BaseNavitOutputWithPast",
+    "LanceBagel",
+    "LanceIdentityConnector",
+    "LancePositionEmbedding3D",
+    "LanceQwen2_5_VLNaViTWrapper",
+    "LanceZeroVitPosEmbed",
+    "MLPconnector",
+    "NaiveCache",
+    "PackedAttentionMoT",
+    "PositionEmbedding",
+    "Qwen2MoTConfig",
+    "Qwen2MoTDecoderLayer",
+    "Qwen2MoTForCausalLM",
+    "Qwen2MoTModel",
+    "TimestepEmbedder",
+    "patchify",
+]
+
+# mRoPE temporal scaling constants for Lance (Qwen2.5-VL backbone).
+# Upstream Lance computes vision rope positions as
+# ``t_index = frame_idx * tokens_per_second * second_per_grid_t`` so adjacent
+# latent frames sit at well-separated rope coordinates.  Lance bundles a
+# ``Qwen2.5-VL-ViT/config.json`` with ``tokens_per_second: 2`` and uses
+# ``second_per_grid_ts = 1.0`` in its inference script.
+LANCE_TOKENS_PER_SECOND = 2
+LANCE_SECONDS_PER_GRID = 1.0
+
+
+def get_3d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray:
+    """3-D sin-cos positional embedding (t, h, w), matching the upstream Lance
+    ``modeling/lance/modeling_utils.py`` dimension split exactly."""
+    assert embed_dim % 2 == 0, "Embedding dimension must be even for 3D embeddings"
+    d = embed_dim // 3
+    d = d if d % 2 == 0 else d - 1
+    dim_t, dim_h = d, d
+    dim_w = embed_dim - 2 * d
+    assert dim_w % 2 == 0
+    emb_t = get_1d_sincos_pos_embed_from_grid(dim_t, grid[0])
+    emb_h = get_1d_sincos_pos_embed_from_grid(dim_h, grid[1])
+    emb_w = get_1d_sincos_pos_embed_from_grid(dim_w, grid[2])
+    return np.concatenate([emb_t, emb_h, emb_w], axis=1)
+
+
+def get_3d_sincos_pos_embed(embed_dim: int, t: int, h: int, w: int) -> np.ndarray:
+    grid_t = np.arange(t, dtype=np.float32)
+    grid_h = np.arange(h, dtype=np.float32)
+    grid_w = np.arange(w, dtype=np.float32)
+    tt, hh, ww = np.meshgrid(grid_t, grid_h, grid_w, indexing="ij")
+    grid = np.stack([tt, hh, ww], axis=0)
+    return get_3d_sincos_pos_embed_from_grid(embed_dim, grid)
+
+
+class LancePositionEmbedding3D(nn.Module):
+    """Frozen 3-D sin-cos latent position embedding for the video path.
+
+    BAGEL only ships a 2-D ``PositionEmbedding`` (image latents).  Lance's
+    ``Lance_3B_Video`` checkpoint adds a temporal axis; this mirrors upstream
+    ``modeling/lance/modeling_utils.py::PositionEmbedding3D``.  The image path
+    uses ``t=1`` and is numerically equivalent to the 2-D embedding.
+    """
+
+    def __init__(self, max_num_frames: int, max_num_patch_per_side: int, hidden_size: int):
+        super().__init__()
+        self.max_num_frames = max_num_frames
+        self.max_num_patch_per_side = max_num_patch_per_side
+        self.hidden_size = hidden_size
+        n = max_num_frames * max_num_patch_per_side * max_num_patch_per_side
+        self.pos_embed = nn.Parameter(torch.zeros(n, hidden_size), requires_grad=False)
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        pe = get_3d_sincos_pos_embed(
+            self.hidden_size,
+            self.max_num_frames,
+            self.max_num_patch_per_side,
+            self.max_num_patch_per_side,
+        )
+        self.pos_embed.data.copy_(torch.from_numpy(pe).float())
+
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
+        return self.pos_embed[position_ids]
+
+
+class LanceIdentityConnector(nn.Module):
+    """No-op connector for Lance.
+
+    BAGEL's ``connector`` projects the ViT hidden size to the LLM hidden size.
+    Qwen2.5-VL's vision tower (which Lance uses) already projects to the LLM
+    hidden size internally via ``merger`` (``out_hidden_size = hidden_size``),
+    and the released Lance safetensors carry no ``connector.*`` weights.  We
+    therefore plug in an ``Identity`` connector so ``forward_cache_update_vit``
+    keeps its existing call site without a separate code path.
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:  # noqa: D401
+        return x
+
+
+class LanceZeroVitPosEmbed(nn.Module):
+    """No-op positional embedding for Lance's ViT tokens.
+
+    BAGEL adds an extra 2-D sin-cos ``vit_pos_embed`` on top of the ViT output.
+    Qwen2.5-VL's vision tower already carries its own (rotary) positional
+    encoding, and the released Lance safetensors carry no
+    ``vit_pos_embed.*`` weights.  This module returns a broadcast-friendly
+    zero so the addition in ``forward_cache_update_vit`` is a no-op without
+    requiring a code-path branch.
+    """
+
+    def forward(self, position_ids: torch.Tensor) -> torch.Tensor:  # noqa: D401
+        return torch.zeros((), device=position_ids.device)
+
+
+class LanceQwen2_5_VLNaViTWrapper(nn.Module):
+    """Packed (NaViT-style) wrapper around the Qwen2.5-VL vision tower.
+
+    Bridges BAGEL's ``vit(packed_pixel_values, ...) -> [num_tokens, vit_hidden]``
+    surface to the HF ``Qwen2_5_VisionTransformerPretrainedModel`` which
+    consumes ``(hidden_states, grid_thw)``.  The packed call additionally needs
+    a per-image ``image_grid_thw`` so non-square images (and the
+    spatial-merge token count) line up — :class:`LanceBagel` stashes the grid
+    on the wrapper before invoking the ViT.
+    """
+
+    def __init__(self, vision_model: nn.Module, spatial_merge_size: int = 2):
+        super().__init__()
+        # Accept either the full Qwen2_5_VLForConditionalGeneration.visual or
+        # the bare vision transformer.
+        self.vision_model = getattr(vision_model, "visual", vision_model)
+        self.spatial_merge_size = spatial_merge_size
+        # Set by ``LanceBagel.forward_cache_update_vit`` before each call so
+        # the wrapper can pass true per-image (T, H, W) to the HF ViT.
+        self._pending_grid_thw: torch.Tensor | None = None
+
+    @property
+    def config(self):  # parity with SiglipNaViTWrapper.vision_model.config access
+        return self.vision_model.config
+
+    def set_pending_grid_thw(self, grid_thw: torch.Tensor) -> None:
+        self._pending_grid_thw = grid_thw
+
+    def forward(
+        self,
+        packed_pixel_values: torch.Tensor,
+        packed_flattened_position_ids: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        max_seqlen: int,
+    ) -> torch.Tensor:
+        if self._pending_grid_thw is not None:
+            grid_thw_t = self._pending_grid_thw.to(packed_pixel_values.device)
+            self._pending_grid_thw = None
+        else:
+            # Fallback for square images when grid_thw was not pre-stashed.
+            cu = cu_seqlens.tolist()
+            grid_thw = []
+            for i in range(len(cu) - 1):
+                n = cu[i + 1] - cu[i]
+                side = int(round(float(n) ** 0.5))
+                grid_thw.append([1, side, side])
+            grid_thw_t = torch.tensor(grid_thw, dtype=torch.long, device=packed_pixel_values.device)
+
+        hidden = packed_pixel_values
+        if hasattr(self.vision_model, "dtype"):
+            hidden = hidden.to(self.vision_model.dtype)
+        out = self.vision_model(hidden_states=hidden, grid_thw=grid_thw_t)
+        # Qwen2.5-VL returns ``BaseModelOutputWithPooling(last_hidden_state=...,
+        # pooler_output=merged)``; the merged sequence (post spatial-merger,
+        # shape ``[sum_tokens_after_merge, out_hidden_size]``) is what the LLM
+        # consumes.
+        if hasattr(out, "pooler_output") and out.pooler_output is not None:
+            return out.pooler_output
+        if isinstance(out, tuple):
+            return out[0]
+        return out
+
+
+class LanceBagel(Bagel):
+    """Bagel subclass with Lance-specific ViT handling.
+
+    The released Lance checkpoint pairs BAGEL's Qwen2-MoT trunk with the
+    Qwen2.5-VL vision tower (whose ``merger`` already projects to LLM
+    ``hidden_size`` and which carries its own rotary positional encoding).
+    Two BAGEL assumptions therefore break for Lance:
+
+    1.  ``prepare_vit_images`` calls ``transforms(image) -> (C, H, W)`` and
+        then ``patchify(...)``; Lance's image processor returns
+        ``(num_patches_flat, patch_features) + image_grid_thw`` already.
+    2.  ``forward_cache_update_vit`` adds a ``connector`` projection plus a
+        2-D ``vit_pos_embed`` that has no checkpoint weights and would
+        also double-count the ViT's own positional encoding.
+
+    This subclass overrides exactly those two methods.  Everything else —
+    LLM trunk, VAE flow, generation loop, ``forward_cache_update_text`` /
+    ``forward_cache_update_vae`` — is reused unchanged.
+    """
+
+    @staticmethod
+    def _qwen_vl_processor_call(processor, image):
+        """Invoke the Qwen2-VL image processor and return ``(pixel_values, grid_thw)``."""
+        proc_out = processor(images=image, return_tensors="pt")
+        pixel_values = proc_out["pixel_values"]
+        grid_thw = proc_out["image_grid_thw"]
+        return pixel_values, grid_thw
+
+    # ------------------------------------------------------------------ #
+    # 3-D mRoPE plumbing
+    # ------------------------------------------------------------------ #
+    @staticmethod
+    def _mrope_broadcast(position_ids: torch.Tensor) -> torch.Tensor:
+        """Convert scalar ``(S,)`` position ids to mrope ``(3, S)`` by
+        broadcasting the scalar value to all three (t, h, w) axes.
+
+        Lance/Qwen2.5-VL's mRoPE expects 3-D positions per token; for text
+        tokens (and image-edit context tokens) all three axes share the
+        same scalar position id.  Vision-generation paths override this with
+        per-token ``(t, h, w)`` indices.
+        """
+        if position_ids is None:
+            return position_ids
+        if position_ids.ndim == 1:
+            return position_ids.unsqueeze(0).expand(3, -1).contiguous()
+        return position_ids
+
+    def prepare_prompts(self, *args, **kwargs):
+        gen_input, newlens, new_rope = super().prepare_prompts(*args, **kwargs)
+        if "packed_text_position_ids" in gen_input:
+            gen_input["packed_text_position_ids"] = self._mrope_broadcast(gen_input["packed_text_position_ids"])
+        return gen_input, newlens, new_rope
+
+    def prepare_vae_latent(self, *args, **kwargs):
+        gen_input = super().prepare_vae_latent(*args, **kwargs)
+        if "packed_position_ids" in gen_input:
+            gen_input["packed_position_ids"] = self._mrope_broadcast(gen_input["packed_position_ids"])
+        return gen_input
+
+    def prepare_vae_latent_cfg(self, *args, **kwargs):
+        gen_input = super().prepare_vae_latent_cfg(*args, **kwargs)
+        if "cfg_packed_position_ids" in gen_input:
+            gen_input["cfg_packed_position_ids"] = self._mrope_broadcast(gen_input["cfg_packed_position_ids"])
+        return gen_input
+
+    def prepare_start_tokens(self, *args, **kwargs):
+        gen_input = super().prepare_start_tokens(*args, **kwargs)
+        if isinstance(gen_input, dict):
+            for k in (
+                "packed_query_position_ids",
+                "packed_text_position_ids",
+                "packed_position_ids",
+            ):
+                if k in gen_input and torch.is_tensor(gen_input[k]):
+                    gen_input[k] = self._mrope_broadcast(gen_input[k])
+        return gen_input
+
+    def prepare_vit_videos(self, curr_kvlens, curr_rope, videos, new_token_ids):
+        """Multi-frame ViT prefill for the ``x2t_video`` understanding path.
+
+        ``videos`` is a list of per-request video tensors / numpy arrays of
+        shape ``(T, H, W, 3)``.  We call the Qwen2-VL **video** processor
+        (stored on ``self._lance_video_processor``) which produces
+        ``pixel_values_videos`` already in packed
+        ``(num_patches_flat_3d, patch_features)`` layout plus a 3-D
+        ``video_grid_thw`` (``[T_lat, H_patches, W_patches]``).  We forward the
+        same pixel-values tensor to the Qwen2.5-VL ViT (wrapped by
+        :class:`LanceQwen2_5_VLNaViTWrapper`), set the wrapper's pending
+        grid, and place the post-merger tokens in the LLM sequence with
+        proper 3-D mRoPE positions ``(P + t, P + h, P + w)`` over the
+        ``T_lat × H_lat × W_lat`` post-merge grid.
+        """
+        processor = getattr(self, "_lance_video_processor", None)
+        if processor is None:
+            raise RuntimeError(
+                "LanceBagel.prepare_vit_videos requires the pipeline to set "
+                "``bagel._lance_video_processor`` (a Qwen2-VL-compatible video processor)."
+            )
+
+        packed_vit_token_indexes: list[int] = []
+        vit_token_seqlens: list[int] = []
+        packed_vit_tokens: list[torch.Tensor] = []
+        packed_vit_position_ids: list[torch.Tensor] = []
+        packed_text_ids: list[int] = []
+        packed_text_indexes: list[int] = []
+        packed_seqlens: list[int] = []
+        packed_indexes: list[int] = []
+        packed_key_value_indexes: list[int] = []
+        grid_thw_list: list[torch.Tensor] = []
+        per_axis_pos: list[tuple[int, int, int]] = []
+
+        merge_size = int(getattr(self.vit_model, "spatial_merge_size", 2))
+        merge_factor = merge_size * merge_size
+
+        _curr = curr = 0
+        newlens: list[int] = []
+        new_rope: list[int] = []
+        for video, curr_kvlen, curr_position_id in zip(videos, curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+
+            packed_text_ids.append(new_token_ids["start_of_image"])
+            packed_text_indexes.append(_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            _curr += 1
+            per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id))
+
+            proc_out = processor(videos=video, return_tensors="pt")
+            # Qwen2-VL video processor: pixel_values_videos is the same packed
+            # (num_patches_flat, patch_features) layout as the image processor
+            # but with a temporal axis ``T_lat`` in video_grid_thw.
+            pixel_values = proc_out["pixel_values_videos"]
+            grid_thw = proc_out["video_grid_thw"]
+            T_lat, H, W = (int(v) for v in grid_thw[0].tolist())
+            num_patches_pre_merge = T_lat * H * W
+            assert num_patches_pre_merge == pixel_values.shape[0], (
+                f"pixel_values rows ({pixel_values.shape[0]}) != T_lat*H*W ({num_patches_pre_merge}); "
+                f"video_grid_thw={grid_thw.tolist()}"
+            )
+            num_vit_tokens = num_patches_pre_merge // merge_factor
+
+            packed_vit_tokens.append(pixel_values)
+            vit_token_seqlens.append(num_patches_pre_merge)
+            packed_vit_position_ids.append(torch.arange(num_patches_pre_merge, dtype=torch.long))
+            grid_thw_list.append(grid_thw[0].to(torch.long))
+
+            packed_vit_token_indexes.extend(range(_curr, _curr + num_vit_tokens))
+            packed_indexes.extend(range(curr, curr + num_vit_tokens))
+            curr += num_vit_tokens
+            _curr += num_vit_tokens
+            h_merged = H // merge_size
+            w_merged = W // merge_size
+            t_scale = LANCE_TOKENS_PER_SECOND * LANCE_SECONDS_PER_GRID
+            for ti in range(T_lat):
+                for hi in range(h_merged):
+                    for wi in range(w_merged):
+                        per_axis_pos.append(
+                            (
+                                curr_position_id + int(ti * t_scale),
+                                curr_position_id + hi,
+                                curr_position_id + wi,
+                            )
+                        )
+
+            packed_text_ids.append(new_token_ids["end_of_image"])
+            packed_text_indexes.append(_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            _curr += 1
+            end_p = curr_position_id + max(int((T_lat - 1) * t_scale), h_merged - 1, w_merged - 1) + 1
+            per_axis_pos.append((end_p, end_p, end_p))
+
+            packed_seqlens.append(num_vit_tokens + 2)
+            newlens.append(curr_kvlen + num_vit_tokens + 2)
+            new_rope.append(end_p + 1)
+
+        pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long)
+        pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long)
+        pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long)
+        packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0)
+
+        generation_input = {
+            "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long),
+            "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long),
+            "vit_token_seqlens": torch.tensor(vit_token_seqlens, dtype=torch.int),
+            "packed_vit_tokens": torch.cat(packed_vit_tokens, dim=0),
+            "packed_vit_position_ids": torch.cat(packed_vit_position_ids, dim=0),
+            "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long),
+            "packed_position_ids": packed_position_ids_3d,
+            "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),
+            "_lance_grid_thw": torch.stack(grid_thw_list, dim=0),
+        }
+        return generation_input, newlens, new_rope
+
+    def prepare_vit_images(self, curr_kvlens, curr_rope, images, transforms, new_token_ids):
+        # ``transforms`` is set up by ``BagelPipeline.forward`` to be
+        # ``processor(images=img, return_tensors='pt').pixel_values[0]``, which
+        # discards the ``image_grid_thw``.  For Lance we need both, so we ignore
+        # the supplied lambda and re-call the processor here directly.  Lance's
+        # pipeline always sets ``self.image_processor`` to ``Qwen2VLImageProcessor``.
+        if not hasattr(self, "_lance_image_processor"):
+            raise RuntimeError(
+                "LanceBagel.prepare_vit_images requires the pipeline to set "
+                "``bagel._lance_image_processor`` (a Qwen2-VL-compatible processor)."
+            )
+        processor = self._lance_image_processor
+
+        packed_vit_token_indexes: list[int] = []
+        vit_token_seqlens: list[int] = []
+        packed_vit_tokens: list[torch.Tensor] = []
+        packed_vit_position_ids: list[torch.Tensor] = []
+        packed_text_ids: list[int] = []
+        packed_text_indexes: list[int] = []
+        packed_seqlens: list[int] = []
+        packed_indexes: list[int] = []
+        packed_key_value_indexes: list[int] = []
+        grid_thw_list: list[torch.Tensor] = []
+
+        merge_size = int(getattr(self.vit_model, "spatial_merge_size", 2))
+        merge_factor = merge_size * merge_size
+
+        _curr = curr = 0
+        newlens: list[int] = []
+        new_rope: list[int] = []
+        # 3-D position ids per query token (mRoPE) — text framing tokens use
+        # scalar broadcast, image patches use per-token ``(P + t, P + h_merged,
+        # P + w_merged)`` so the Qwen2.5-VL backbone sees the spatial layout
+        # along the (h, w) axes of mRoPE.
+        per_axis_pos: list[tuple[int, int, int]] = []
+        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
+
+            packed_text_ids.append(new_token_ids["start_of_image"])
+            packed_text_indexes.append(_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            _curr += 1
+            per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id))
+
+            pixel_values, grid_thw = self._qwen_vl_processor_call(processor, image)
+            # pixel_values: (num_patches_flat, patch_features). grid_thw: (1, 3) with (T, H, W) in patch units.
+            T, H, W = (int(v) for v in grid_thw[0].tolist())
+            num_patches_pre_merge = T * H * W
+            assert num_patches_pre_merge == pixel_values.shape[0], (
+                f"pixel_values rows ({pixel_values.shape[0]}) != T*H*W ({num_patches_pre_merge}); "
+                f"image_grid_thw={grid_thw.tolist()}"
+            )
+            num_img_tokens = num_patches_pre_merge // merge_factor
+
+            packed_vit_tokens.append(pixel_values)
+            vit_token_seqlens.append(num_patches_pre_merge)
+            packed_vit_position_ids.append(torch.arange(num_patches_pre_merge, dtype=torch.long))
+            grid_thw_list.append(grid_thw[0].to(torch.long))
+
+            packed_vit_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
+            # 3-D mRoPE positions for image patches: scan in row-major over the
+            # post-merge grid (H/merge × W/merge).  Temporal axis stays at the
+            # block's scalar position (single still image).
+            h_merged = H // merge_size
+            w_merged = W // merge_size
+            for hi in range(h_merged):
+                for wi in range(w_merged):
+                    per_axis_pos.append((curr_position_id, curr_position_id + hi, curr_position_id + wi))
+
+            packed_text_ids.append(new_token_ids["end_of_image"])
+            packed_text_indexes.append(_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            _curr += 1
+            end_p = curr_position_id + max(h_merged - 1, w_merged - 1) + 1
+            per_axis_pos.append((end_p, end_p, end_p))
+
+            packed_seqlens.append(num_img_tokens + 2)
+            newlens.append(curr_kvlen + num_img_tokens + 2)
+            new_rope.append(end_p + 1)
+
+        pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long)
+        pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long)
+        pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long)
+        packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0)
+
+        generation_input = {
+            "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long),
+            "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long),
+            "vit_token_seqlens": torch.tensor(vit_token_seqlens, dtype=torch.int),
+            "packed_vit_tokens": torch.cat(packed_vit_tokens, dim=0),
+            "packed_vit_position_ids": torch.cat(packed_vit_position_ids, dim=0),
+            "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long),
+            "packed_position_ids": packed_position_ids_3d,
+            "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),
+            # Lance-extra: per-image (T, H, W) so the ViT wrapper can pass the
+            # true grid to the HF tower (non-square images would otherwise be
+            # reconstructed incorrectly as ``sqrt(n)`` squares).
+            "_lance_grid_thw": torch.stack(grid_thw_list, dim=0),
+        }
+        return generation_input, newlens, new_rope
+
+    def forward_cache_update_vit(self, *args, **kwargs):
+        # Pluck the Lance-specific extra and stash on the wrapper for the
+        # upcoming ViT call.  Everything else delegates to the parent so we
+        # don't have to copy 60 lines of LLM-call boilerplate.
+        grid_thw = kwargs.pop("_lance_grid_thw", None)
+        if grid_thw is not None and hasattr(self.vit_model, "set_pending_grid_thw"):
+            self.vit_model.set_pending_grid_thw(grid_thw)
+        return super().forward_cache_update_vit(*args, **kwargs)
+
+    def prepare_video_latent(self, curr_kvlens, curr_rope, video_shapes, new_token_ids):
+        """3-D analogue of :meth:`prepare_vae_latent`.
+
+        ``video_shapes`` is a list of ``(T, H, W)`` per request (RGB pixel
+        space).  We package one packed-init-noise tensor over ``T_lat × H_lat ×
+        W_lat`` latent tokens per video, plus 1-D indices into the 3-D position
+        embedding table maintained by :class:`LancePositionEmbedding3D`
+        (`bagel.latent_pos_embed`).  Latent geometry:
+
+        - spatial: ``H_lat = H // latent_downsample`` (``=16`` for Lance)
+        - temporal: ``T_lat = (T - 1) // downsample_temporal + 1`` (``=4`` for Wan2.2)
+        - channels: ``latent_channel = 48``
+
+        Position ids are flattened ``t * max_per_side² + h * max_per_side + w``
+        so they index directly into the ``(max_num_frames * max_per_side²,
+        hidden_size)`` table.
+        """
+        downsample_t = int(getattr(self.config.vae_config, "downsample_temporal", 4))
+        max_per_side = int(self.max_latent_size)
+
+        packed_text_ids, packed_text_indexes = [], []
+        packed_vae_position_ids, packed_vae_token_indexes, packed_init_noises = [], [], []
+        packed_seqlens, packed_indexes = [], []
+        packed_key_value_indexes = []
+        # 3-D position ids per token (mRoPE).  ``per_axis_pos[i]`` is the list
+        # of (t, h, w) for the i-th query token; we stack into ``(3, S)`` at the
+        # end.  Text framing tokens broadcast scalar to all axes; video latent
+        # tokens carry per-token ``(P + t, P + h, P + w)``.
+        per_axis_pos: list[tuple[int, int, int]] = []
+
+        query_curr = curr = 0
+        for (T, H, W), curr_kvlen, curr_position_id in zip(video_shapes, curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+
+            packed_text_ids.append(new_token_ids["start_of_image"])
+            packed_text_indexes.append(query_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            query_curr += 1
+            per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id))
+
+            h = H // self.latent_downsample
+            w = W // self.latent_downsample
+            t = (T - 1) // downsample_t + 1
+            num_video_tokens = t * h * w
+
+            # 1-D position ids into the 3-D table (frame-major, then row, then col).
+            tt, hh, ww = torch.meshgrid(torch.arange(t), torch.arange(h), torch.arange(w), indexing="ij")
+            tt_flat = tt.flatten().tolist()
+            hh_flat = hh.flatten().tolist()
+            ww_flat = ww.flatten().tolist()
+            vae_position_ids = (tt * (max_per_side * max_per_side) + hh * max_per_side + ww).flatten()
+            packed_vae_position_ids.append(vae_position_ids)
+
+            packed_init_noises.append(torch.randn(num_video_tokens, self.latent_channel * self.latent_patch_size**2))
+
+            packed_vae_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
+            # Per-token 3-D mRoPE positions, matching upstream Lance's
+            # ``get_rope_index``: the temporal axis is amplified by
+            # ``tokens_per_second * second_per_grid_t`` (see Qwen2.5-VL docs)
+            # so neighbouring latent frames are well-separated in rope space.
+            # ``second_per_grid_t = 1.0`` matches upstream's default;
+            # ``tokens_per_second = 2`` is Lance's vit config.
+            t_scale = LANCE_TOKENS_PER_SECOND * LANCE_SECONDS_PER_GRID
+            for ti, hi, wi in zip(tt_flat, hh_flat, ww_flat):
+                per_axis_pos.append(
+                    (
+                        curr_position_id + int(ti * t_scale),
+                        curr_position_id + hi,
+                        curr_position_id + wi,
+                    )
+                )
+
+            packed_text_ids.append(new_token_ids["end_of_image"])
+            packed_text_indexes.append(query_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            query_curr += 1
+            # end_of_image sits "past" the latent block; bump along all axes by
+            # max of the latent block size so subsequent text resumes cleanly.
+            end_p = curr_position_id + max(t - 1, h - 1, w - 1) + 1
+            per_axis_pos.append((end_p, end_p, end_p))
+
+            packed_seqlens.append(num_video_tokens + 2)
+
+        # Stack into (3, total_query_tokens) for mRoPE.
+        pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long)
+        pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long)
+        pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long)
+        packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0)
+
+        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_vae_position_ids": torch.cat(packed_vae_position_ids, dim=0),
+            "packed_vae_token_indexes": torch.tensor(packed_vae_token_indexes, dtype=torch.long),
+            "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int),
+            "packed_position_ids": packed_position_ids_3d,
+            "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_video_latent_cfg(self, curr_kvlens, curr_rope, video_shapes):
+        """3-D analogue of :meth:`prepare_vae_latent_cfg` (CFG side).
+
+        Mirrors :meth:`prepare_video_latent`'s mRoPE 3-D position layout
+        (text frame ⇒ scalar; video latent block ⇒ per-token ``(P+t, P+h, P+w)``).
+        """
+        downsample_t = int(getattr(self.config.vae_config, "downsample_temporal", 4))
+
+        packed_indexes, packed_key_value_indexes = [], []
+        per_axis_pos: list[tuple[int, int, int]] = []
+        query_curr = curr = 0
+        for (T, H, W), curr_kvlen, curr_position_id in zip(video_shapes, curr_kvlens, curr_rope):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            curr += curr_kvlen
+
+            packed_indexes.append(curr)
+            curr += 1
+            query_curr += 1
+            per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id))
+
+            h = H // self.latent_downsample
+            w = W // self.latent_downsample
+            t = (T - 1) // downsample_t + 1
+            num_video_tokens = t * h * w
+            packed_indexes.extend(range(curr, curr + num_video_tokens))
+            curr += num_video_tokens
+            query_curr += num_video_tokens
+            tt, hh, ww = torch.meshgrid(torch.arange(t), torch.arange(h), torch.arange(w), indexing="ij")
+            for ti, hi, wi in zip(tt.flatten().tolist(), hh.flatten().tolist(), ww.flatten().tolist()):
+                per_axis_pos.append((curr_position_id + ti, curr_position_id + hi, curr_position_id + wi))
+
+            packed_indexes.append(curr)
+            curr += 1
+            query_curr += 1
+            end_p = curr_position_id + max(t - 1, h - 1, w - 1) + 1
+            per_axis_pos.append((end_p, end_p, end_p))
+
+        pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long)
+        pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long)
+        pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long)
+        cfg_packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0)
+
+        return {
+            "cfg_packed_position_ids": cfg_packed_position_ids_3d,
+            "cfg_packed_query_indexes": torch.tensor(packed_indexes, dtype=torch.long),
+            "cfg_key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
+            "cfg_packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
+        }
+
+    # ------------------------------------------------------------------ #
+    # Image-edit / video-edit VAE prefill
+    # ------------------------------------------------------------------ #
+    def _lance_native_prepare_vae_images(
+        self,
+        curr_kvlens,
+        curr_rope,
+        images,
+        transforms,
+        new_token_ids,
+        timestep=0,
+        is_video: bool = False,
+    ):
+        """Lance-native ``prepare_vae_images`` with 3-D mRoPE positions.
+
+        Mirrors :meth:`Bagel.prepare_vae_images` but emits per-token mRoPE
+        ``(curr + t, curr + h, curr + w)`` positions for VAE latent tokens
+        (BAGEL emits scalar ``curr_position_id`` for all of them, which
+        triggers a CUDA index-out-of-bounds gather on Lance's Qwen2.5-VL
+        backbone).  For still images ``t = 0``; for video clips ``t``
+        ranges over ``T_lat = (T - 1) // downsample_temporal + 1``.
+        """
+        downsample_t = int(getattr(self.config.vae_config, "downsample_temporal", 4))
+        max_per_side = int(self.max_latent_size)
+
+        packed_text_ids: list[int] = []
+        packed_text_indexes: list[int] = []
+        packed_vae_position_ids: list[torch.Tensor] = []
+        packed_vae_token_indexes: list[int] = []
+        packed_indexes: list[int] = []
+        packed_key_value_indexes: list[int] = []
+        packed_seqlens: list[int] = []
+        patchified_vae_latent_shapes: list[tuple] = []
+        vae_image_tensors: list[torch.Tensor] = []
+        per_axis_pos: list[tuple[int, int, int]] = []
+
+        query_curr = curr = 0
+        newlens: list[int] = []
+        new_rope: list[int] = []
+        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
+
+            packed_text_ids.append(new_token_ids["start_of_image"])
+            packed_text_indexes.append(query_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            query_curr += 1
+            per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id))
+
+            image_tensor = transforms(image)
+            vae_image_tensors.append(image_tensor)
+            # ``image_tensor`` is ``(C, H, W)`` for image input or ``(C, T, H, W)``
+            # for video input.  We derive the *latent* grid from the RGB spatial
+            # dims and our Wan2.2 VAE downsample factor.
+            if image_tensor.dim() == 3:
+                H, W = image_tensor.shape[1:]
+                T = 1
+            elif image_tensor.dim() == 4:
+                _, T, H, W = image_tensor.shape
+            else:
+                raise ValueError(f"vae transforms must return 3-D or 4-D tensor; got {image_tensor.shape}")
+            h_lat = H // self.latent_downsample
+            w_lat = W // self.latent_downsample
+            t_lat = ((T - 1) // downsample_t + 1) if is_video else 1
+            patchified_vae_latent_shapes.append((t_lat, h_lat, w_lat) if is_video else (h_lat, w_lat))
+
+            # 1-D index into the (3-D when video) ``latent_pos_embed`` table.
+            tt, hh, ww = torch.meshgrid(torch.arange(t_lat), torch.arange(h_lat), torch.arange(w_lat), indexing="ij")
+            vae_pos_ids = (tt * (max_per_side * max_per_side) + hh * max_per_side + ww).flatten()
+            packed_vae_position_ids.append(vae_pos_ids)
+
+            num_image_tokens = t_lat * h_lat * w_lat
+            packed_vae_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
+            for ti, hi, wi in zip(tt.flatten().tolist(), hh.flatten().tolist(), ww.flatten().tolist()):
+                per_axis_pos.append((curr_position_id + ti, curr_position_id + hi, curr_position_id + wi))
+
+            packed_text_ids.append(new_token_ids["end_of_image"])
+            packed_text_indexes.append(query_curr)
+            packed_indexes.append(curr)
+            curr += 1
+            query_curr += 1
+            end_p = curr_position_id + max(t_lat - 1, h_lat - 1, w_lat - 1) + 1
+            per_axis_pos.append((end_p, end_p, end_p))
+
+            packed_seqlens.append(num_image_tokens + 2)
+            newlens.append(curr_kvlen + num_image_tokens + 2)
+            new_rope.append(end_p + 1)
+
+        image_sizes = [item.shape for item in vae_image_tensors]
+        max_image_size = [max(item) for item in list(zip(*image_sizes))]
+        padded_images = torch.zeros(size=(len(vae_image_tensors), *max_image_size))
+        for i, image_tensor in enumerate(vae_image_tensors):
+            if image_tensor.dim() == 3:
+                padded_images[i, :, : image_tensor.shape[1], : image_tensor.shape[2]] = image_tensor
+            else:
+                padded_images[i, :, : image_tensor.shape[1], : image_tensor.shape[2], : image_tensor.shape[3]] = (
+                    image_tensor
+                )
+
+        pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long)
+        pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long)
+        pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long)
+        packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0)
+
+        generation_input = {
+            "padded_images": padded_images,
+            "patchified_vae_latent_shapes": patchified_vae_latent_shapes,
+            "packed_vae_position_ids": torch.cat(packed_vae_position_ids, dim=0),
+            "packed_timesteps": torch.tensor([timestep]),
+            "packed_vae_token_indexes": torch.tensor(packed_vae_token_indexes, dtype=torch.long),
+            "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long),
+            "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long),
+            "packed_position_ids": packed_position_ids_3d,
+            "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_vae_images(self, curr_kvlens, curr_rope, images, transforms, new_token_ids, timestep=0):
+        """No-op VAE prefill for the x2t / x2t_video path.
+
+        Lance's understanding paths route image / video context through the
+        Qwen2.5-VL ViT only.  BAGEL's pipeline runs ``prepare_vae_images``
+        unconditionally for any image input, so this default override
+        short-circuits it.  :meth:`_forward_image_edit` /
+        :meth:`_forward_video_edit` call :meth:`_lance_native_prepare_vae_images`
+        directly when they actually need Lance-style VAE prefill.
+        """
+        generation_input = {
+            "padded_images": torch.empty(0, 3, 0, 0),
+            "patchified_vae_latent_shapes": [],
+            "packed_vae_position_ids": torch.empty(0, dtype=torch.long),
+            "packed_timesteps": torch.tensor([timestep]),
+            "packed_vae_token_indexes": torch.empty(0, dtype=torch.long),
+            "packed_text_ids": torch.empty(0, dtype=torch.long),
+            "packed_text_indexes": torch.empty(0, dtype=torch.long),
+            "packed_position_ids": torch.empty(0, dtype=torch.long),
+            "packed_seqlens": torch.empty(0, dtype=torch.int),
+            "packed_indexes": torch.empty(0, dtype=torch.long),
+            "packed_key_value_indexes": torch.empty(0, dtype=torch.long),
+            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
+        }
+        return generation_input, list(curr_kvlens), list(curr_rope)
+
+    def forward_cache_update_vae(
+        self,
+        vae_model,
+        past_key_values,
+        padded_images=None,
+        patchified_vae_latent_shapes=None,
+        packed_vae_position_ids=None,
+        packed_timesteps=None,
+        packed_vae_token_indexes=None,
+        packed_text_ids=None,
+        packed_text_indexes=None,
+        packed_position_ids=None,
+        packed_seqlens=None,
+        packed_indexes=None,
+        key_values_lens=None,
+        packed_key_value_indexes=None,
+    ):
+        """Lance-native VAE prefill that *actually* scatters the encoded
+        latents into the LLM query sequence.
+
+        :meth:`Bagel.forward_cache_update_vae` in vllm-omni computes
+        ``packed_latent = vae2llm(...) + time_embed + pos_embed`` and then
+        passes only ``packed_text_ids`` to the LLM — the VAE embeddings
+        never enter the query sequence (the LLM builds it from
+        ``embed_tokens(packed_text_ids)`` which is just the 2 framing
+        tokens).  The mismatch between ``query_lens = [num_vae + 2]`` and
+        the resulting 2-token sequence is what crashes the gather inside
+        attention.
+
+        We scatter both pieces explicitly: text framing tokens at
+        ``packed_text_indexes`` and the VAE latent embeddings at
+        ``packed_vae_token_indexes``, producing a full-length
+        ``(sum(packed_seqlens), hidden)`` sequence the LLM can attend
+        over.  Empty prep data (legacy x2t / x2t_video no-op path) is
+        short-circuited.
+        """
+        if (
+            packed_text_ids is None
+            or packed_text_ids.numel() == 0
+            or padded_images is None
+            or padded_images.numel() == 0
+        ):
+            return past_key_values
+
+        padded_latent = vae_model.encode(padded_images)
+        p = self.latent_patch_size
+        packed_latent_list = []
+        for latent, shape in zip(padded_latent, patchified_vae_latent_shapes):
+            if isinstance(shape, tuple) and len(shape) == 3:
+                t_lat, h_lat, w_lat = shape
+                latent = latent[:, :t_lat, : h_lat * p, : w_lat * p].reshape(
+                    self.latent_channel, t_lat, h_lat, p, w_lat, p
+                )
+                latent = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, p * p * self.latent_channel)
+            else:
+                h_lat, w_lat = shape
+                latent = latent[:, : h_lat * p, : w_lat * p].reshape(self.latent_channel, h_lat, p, w_lat, p)
+                latent = torch.einsum("chpwq->hwpqc", latent).reshape(-1, p * p * self.latent_channel)
+            packed_latent_list.append(latent)
+        packed_latent = torch.cat(packed_latent_list, dim=0)
+
+        packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
+        packed_timestep_embeds = self.time_embedder(packed_timesteps)
+        packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + packed_pos_embed
+
+        # ---- Scatter text + VAE embeds into the full query sequence ---- #
+        total_len = int(packed_seqlens.sum().item())
+        packed_text_embed = self.language_model.model.embed_tokens(packed_text_ids)
+        packed_query_sequence = packed_latent.new_zeros((total_len, packed_latent.shape[-1]))
+        packed_query_sequence[packed_text_indexes] = packed_text_embed.to(packed_query_sequence.dtype)
+        packed_query_sequence[packed_vae_token_indexes] = packed_latent.to(packed_query_sequence.dtype)
+
+        extra_inputs = {}
+        if self.use_moe:
+            extra_inputs = {
+                "mode": "gen",
+                "packed_vae_token_indexes": packed_vae_token_indexes,
+                "packed_text_indexes": packed_text_indexes,
+            }
+
+        output = self.language_model.forward(
+            packed_query_sequence=packed_query_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=True,
+            is_causal=False,
+            **extra_inputs,
+        )
+        return output.past_key_values

vllm_omni/diffusion/models/lance/prompts.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance chat / system prompts.
+
+Matches the upstream Lance training distribution (see
+``bytedance/Lance/data/system_prompt_render.py`` and
+``data/common.py::generate_system_prompt``):
+
+  * each task has a specific system prompt describing what to attend to;
+  * the prompt is wrapped in the Qwen chat template
+    ``<|im_start|>system\\n…<|im_end|>\\n<|im_start|>user\\n…<|im_end|>\\n``
+    ``<|im_start|>assistant\\n``;
+  * vision tokens are framed by ``<|vision_start|><|video_pad|><|vision_end|>``
+    (upstream uses ``<|video_pad|>`` even for image type by default).
+
+These are pure string formatting helpers — the runtime pipeline still does
+the actual VAE / ViT prefill separately.
+"""
+
+from __future__ import annotations
+
+# Task-specific system prompts (verbatim from upstream's
+# ``generate_system_prompt`` — first variant of each pool).
+SYSTEM_PROMPTS: dict[tuple[str, str], str] = {
+    # (task, vision_type) -> system prompt
+    ("t2i", "image"): (
+        "Describe the image by detailing the color, quantity, text, shape, "
+        "size, texture, spatial relationships of the objects and background:"
+    ),
+    ("t2v", "video"): (
+        "Describe the video by detailing the color, quantity, visible text, "
+        "shape, size, texture, spatial relationships and motion/camera "
+        "movements of the objects and background:"
+    ),
+    ("image_edit", "image"): (
+        "Describe the key features of the input image (color, shape, size, "
+        "texture, objects, background), then explain how the user's text "
+        "instruction should alter or modify the image. Generate a new image "
+        "that meets the user's requirements while maintaining consistency "
+        "with the original input where appropriate."
+    ),
+    ("video_edit", "video"): (
+        "Describe the key features of the input video (color, shape, size, "
+        "texture, objects, background), then explain how the user's text "
+        "instruction should alter or modify the video. Generate a new video "
+        "that meets the user's requirements while maintaining consistency "
+        "with the original input where appropriate."
+    ),
+    ("x2t_image", "image"): (
+        "Generate a detailed and accurate description of the image, including all the key moments and visual details."
+    ),
+    ("x2t_video", "video"): (
+        "Generate a detailed and accurate description of the video, including all the key moments and visual details."
+    ),
+}
+
+# Vision-content placeholders (Qwen2.5-VL convention).  Upstream uses
+# ``<|video_pad|>`` for image inputs too unless ``force_video_pad=False``
+# (counter-intuitively named; see ``render_qwenvl_prompt`` upstream).
+VISION_START = "<|vision_start|>"
+VISION_END = "<|vision_end|>"
+VIDEO_PAD = "<|video_pad|>"
+IMAGE_PAD = "<|image_pad|>"
+
+
+def render_lance_prompt(
+    task: str,
+    user_text: str,
+    *,
+    vision_token: str | None = None,
+) -> str:
+    """Render a Lance-compatible user-side prompt.
+
+    ``task`` is one of ``t2i``, ``t2v``, ``image_edit``, ``video_edit``,
+    ``x2t_image``, ``x2t_video``.  ``vision_token`` should be the visual
+    placeholder to embed inside the user message (e.g. ``VIDEO_PAD`` or a
+    ``VISION_START + VIDEO_PAD + VISION_END`` block); pass ``None`` for
+    text-only inputs.
+
+    Returns a single string ready to be tokenized; no further wrapping is
+    needed by the caller.
+    """
+    # ``t2v`` / ``x2t_video`` / ``video_edit`` use video system prompts;
+    # the rest are image.
+    vision_type = "video" if task in {"t2v", "x2t_video", "video_edit"} else "image"
+    sys_prompt = SYSTEM_PROMPTS.get((task, vision_type), "You are a helpful assistant.")
+
+    if vision_token is None:
+        user_msg = user_text
+    else:
+        # Upstream concatenates the vision block immediately before the text.
+        user_msg = f"{vision_token}{user_text}"
+
+    return (
+        f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n"
+    )
+
+
+__all__ = [
+    "IMAGE_PAD",
+    "SYSTEM_PROMPTS",
+    "VIDEO_PAD",
+    "VISION_END",
+    "VISION_START",
+    "render_lance_prompt",
+]

vllm_omni/diffusion/models/lance/wan_vae.py

@@ -0,0 +1,864 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+#
+# Ported from ByteDance Lance upstream (https://github.com/bytedance/Lance,
+# modeling/vae/wan/{vae2_2.py,model.py}). Upstream copyright:
+#
+#   Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
+#   Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+#   Licensed under the Apache License, Version 2.0.
+#
+# Lance ships its VAE as a raw ``Wan2.2_VAE.pth`` whose state dict is keyed for
+# the upstream ``WanVAE_`` nn.Module here (not for diffusers' ``AutoencoderKLWan``).
+# We port that module verbatim so we can load the .pth directly, and wrap it with
+# :class:`LanceWanVAE` which exposes BAGEL's 4-D ``encode(BCHW)/decode(BCHW)``
+# surface (and a 5-D path for the video checkpoint).
+"""Wan2.2 VAE used by Lance, ported from upstream so ``Wan2.2_VAE.pth`` loads
+natively without state-dict surgery."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+    """Causal 3D conv with feature-map caching across temporal chunks."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (
+            self.padding[2],
+            self.padding[2],
+            self.padding[1],
+            self.padding[1],
+            2 * self.padding[0],
+            0,
+        )
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
+
+    def forward(self, x):
+        return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+    def forward(self, x):
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+    def __init__(self, dim, mode):
+        assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        if mode == "upsample2d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+        elif mode == "upsample3d":
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+                nn.Conv2d(dim, dim, 3, padding=1),
+            )
+            self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+        elif mode == "downsample2d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == "downsample3d":
+            self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == "upsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = "Rep"
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
+                        cache_x = torch.cat(
+                            [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                            dim=2,
+                        )
+                    if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
+                        cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
+                    if feat_cache[idx] == "Rep":
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.resample(x)
+        x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+        if self.mode == "downsample3d":
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1),
+        )
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """Single-head causal self-attention over spatial tokens, per frame."""
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, "b c t h w -> (b t) c h w")
+        x = self.norm(x)
+        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+        x = self.proj(x)
+        x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
+        return x + identity
+
+
+def _patchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+    if x.dim() == 5:
+        return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size)
+    raise ValueError(f"Invalid input shape: {x.shape}")
+
+
+def _unpatchify(x, patch_size):
+    if patch_size == 1:
+        return x
+    if x.dim() == 4:
+        return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+    if x.dim() == 5:
+        return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size)
+    return x
+
+
+class AvgDown3D(nn.Module):
+    def __init__(self, in_channels, out_channels, factor_t, factor_s=1):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+        assert in_channels * self.factor % out_channels == 0
+        self.group_size = in_channels * self.factor // out_channels
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+        x = F.pad(x, (0, 0, 0, 0, pad_t, 0))
+        B, C, T, H, W = x.shape
+        x = x.view(
+            B,
+            C,
+            T // self.factor_t,
+            self.factor_t,
+            H // self.factor_s,
+            self.factor_s,
+            W // self.factor_s,
+            self.factor_s,
+        )
+        x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+        x = x.view(B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s)
+        x = x.view(B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s)
+        return x.mean(dim=2)
+
+
+class DupUp3D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.factor_t = factor_t
+        self.factor_s = factor_s
+        self.factor = self.factor_t * self.factor_s * self.factor_s
+        assert out_channels * self.factor % in_channels == 0
+        self.repeats = out_channels * self.factor // in_channels
+
+    def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+        x = x.repeat_interleave(self.repeats, dim=1)
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            self.factor_t,
+            self.factor_s,
+            self.factor_s,
+            x.size(2),
+            x.size(3),
+            x.size(4),
+        )
+        x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+        x = x.view(
+            x.size(0),
+            self.out_channels,
+            x.size(2) * self.factor_t,
+            x.size(4) * self.factor_s,
+            x.size(6) * self.factor_s,
+        )
+        if first_chunk:
+            x = x[:, :, self.factor_t - 1 :, :, :]
+        return x
+
+
+class Down_ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False):
+        super().__init__()
+        self.avg_shortcut = AvgDown3D(
+            in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1
+        )
+        downsamples = []
+        for _ in range(mult):
+            downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+        if down_flag:
+            mode = "downsample3d" if temperal_downsample else "downsample2d"
+            downsamples.append(Resample(out_dim, mode=mode))
+        self.downsamples = nn.Sequential(*downsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        x_copy = x.clone()
+        for module in self.downsamples:
+            x = module(x, feat_cache, feat_idx)
+        return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+    def __init__(self, in_dim, out_dim, dropout, mult, temporal_upsample=False, up_flag=False):
+        super().__init__()
+        if up_flag:
+            self.avg_shortcut = DupUp3D(
+                in_dim, out_dim, factor_t=2 if temporal_upsample else 1, factor_s=2 if up_flag else 1
+            )
+        else:
+            self.avg_shortcut = None
+        upsamples = []
+        for _ in range(mult):
+            upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+            in_dim = out_dim
+        if up_flag:
+            mode = "upsample3d" if temporal_upsample else "upsample2d"
+            upsamples.append(Resample(out_dim, mode=mode))
+        self.upsamples = nn.Sequential(*upsamples)
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        x_main = x.clone()
+        for module in self.upsamples:
+            x_main = module(x_main, feat_cache, feat_idx)
+        if self.avg_shortcut is not None:
+            return x_main + self.avg_shortcut(x, first_chunk)
+        return x_main
+
+
+class Encoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=(1, 2, 4, 4),
+        num_res_blocks=2,
+        attn_scales=(),
+        temperal_downsample=(True, True, False),
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = list(dim_mult)
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = list(attn_scales)
+        self.temperal_downsample = list(temperal_downsample)
+
+        dims = [dim * u for u in [1] + self.dim_mult]
+        self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_down_flag = self.temperal_downsample[i] if i < len(self.temperal_downsample) else False
+            downsamples.append(
+                Down_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks,
+                    temperal_downsample=t_down_flag,
+                    down_flag=i != len(self.dim_mult) - 1,
+                )
+            )
+        self.downsamples = nn.Sequential(*downsamples)
+
+        self.middle = nn.Sequential(
+            ResidualBlock(out_dim, out_dim, dropout),
+            AttentionBlock(out_dim),
+            ResidualBlock(out_dim, out_dim, dropout),
+        )
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, z_dim, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+    def __init__(
+        self,
+        dim=128,
+        z_dim=4,
+        dim_mult=(1, 2, 4, 4),
+        num_res_blocks=2,
+        attn_scales=(),
+        temporal_upsample=(False, True, True),
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = list(dim_mult)
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = list(attn_scales)
+        self.temporal_upsample = list(temporal_upsample)
+
+        dims = [dim * u for u in [self.dim_mult[-1]] + self.dim_mult[::-1]]
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+        self.middle = nn.Sequential(
+            ResidualBlock(dims[0], dims[0], dropout),
+            AttentionBlock(dims[0]),
+            ResidualBlock(dims[0], dims[0], dropout),
+        )
+
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            t_up_flag = self.temporal_upsample[i] if i < len(self.temporal_upsample) else False
+            upsamples.append(
+                Up_ResidualBlock(
+                    in_dim=in_dim,
+                    out_dim=out_dim,
+                    dropout=dropout,
+                    mult=num_res_blocks + 1,
+                    temporal_upsample=t_up_flag,
+                    up_flag=i != len(self.dim_mult) - 1,
+                )
+            )
+        self.upsamples = nn.Sequential(*upsamples)
+        self.head = nn.Sequential(
+            RMS_norm(out_dim, images=False),
+            nn.SiLU(),
+            CausalConv3d(out_dim, 12, 3, padding=1),
+        )
+
+    def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                cache_x = torch.cat(
+                    [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                    dim=2,
+                )
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        for layer in self.middle:
+            if isinstance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx, first_chunk)
+            else:
+                x = layer(x)
+
+        for layer in self.head:
+            if isinstance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    cache_x = torch.cat(
+                        [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
+                        dim=2,
+                    )
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+def _count_conv3d(model):
+    return sum(1 for m in model.modules() if isinstance(m, CausalConv3d))
+
+
+class WanVAE_(nn.Module):
+    """Upstream Wan2.2 VAE module — encoder3d/decoder3d sandwich with 2x patchify
+    on input. State-dict-compatible with ``Wan2.2_VAE.pth``."""
+
+    def __init__(
+        self,
+        dim=160,
+        dec_dim=256,
+        z_dim=16,
+        dim_mult=(1, 2, 4, 4),
+        num_res_blocks=2,
+        attn_scales=(),
+        temperal_downsample=(True, True, False),
+        dropout=0.0,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = list(dim_mult)
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = list(attn_scales)
+        self.temperal_downsample = list(temperal_downsample)
+        self.temporal_upsample = list(temperal_downsample)[::-1]
+
+        self.encoder = Encoder3d(
+            dim, z_dim * 2, self.dim_mult, num_res_blocks, self.attn_scales, self.temperal_downsample, dropout
+        )
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(
+            dec_dim, z_dim, self.dim_mult, num_res_blocks, self.attn_scales, self.temporal_upsample, dropout
+        )
+
+    def encode(self, x, scale):
+        self.clear_cache()
+        x = _patchify(x, patch_size=2)
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+        out = None
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
+            else:
+                out_ = self.encoder(
+                    x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
+                    feat_cache=self._enc_feat_map,
+                    feat_idx=self._enc_conv_idx,
+                )
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
+        else:
+            mu = (mu - scale[0]) * scale[1]
+        self.clear_cache()
+        return mu, log_var
+
+    def decode(self, z, scale):
+        self.clear_cache()
+        if isinstance(scale[0], torch.Tensor):
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
+        else:
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        out = None
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out = self.decoder(
+                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True
+                )
+            else:
+                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+        out = _unpatchify(out, patch_size=2)
+        self.clear_cache()
+        return out
+
+    def clear_cache(self):
+        self._conv_num = _count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        self._enc_conv_num = _count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+# ----------------------------------------------------------------------------
+# Wan2.2 VAE wrapper + BAGEL-surface adapter
+# ----------------------------------------------------------------------------
+
+
+# Per-channel mean/std for the 48-channel Wan2.2 latent space (upstream constants).
+_WAN22_LATENT_MEAN: tuple[float, ...] = (
+    -0.2289,
+    -0.0052,
+    -0.1323,
+    -0.2339,
+    -0.2799,
+    0.0174,
+    0.1838,
+    0.1557,
+    -0.1382,
+    0.0542,
+    0.2813,
+    0.0891,
+    0.1570,
+    -0.0098,
+    0.0375,
+    -0.1825,
+    -0.2246,
+    -0.1207,
+    -0.0698,
+    0.5109,
+    0.2665,
+    -0.2108,
+    -0.2158,
+    0.2502,
+    -0.2055,
+    -0.0322,
+    0.1109,
+    0.1567,
+    -0.0729,
+    0.0899,
+    -0.2799,
+    -0.1230,
+    -0.0313,
+    -0.1649,
+    0.0117,
+    0.0723,
+    -0.2839,
+    -0.2083,
+    -0.0520,
+    0.3748,
+    0.0152,
+    0.1957,
+    0.1433,
+    -0.2944,
+    0.3573,
+    -0.0548,
+    -0.1681,
+    -0.0667,
+)
+_WAN22_LATENT_STD: tuple[float, ...] = (
+    0.4765,
+    1.0364,
+    0.4514,
+    1.1677,
+    0.5313,
+    0.4990,
+    0.4818,
+    0.5013,
+    0.8158,
+    1.0344,
+    0.5894,
+    1.0901,
+    0.6885,
+    0.6165,
+    0.8454,
+    0.4978,
+    0.5759,
+    0.3523,
+    0.7135,
+    0.6804,
+    0.5833,
+    1.4146,
+    0.8986,
+    0.5659,
+    0.7069,
+    0.5338,
+    0.4889,
+    0.4917,
+    0.4069,
+    0.4999,
+    0.6866,
+    0.4093,
+    0.5709,
+    0.6065,
+    0.6415,
+    0.4944,
+    0.5726,
+    1.2042,
+    0.5458,
+    1.6887,
+    0.3971,
+    1.0600,
+    0.3943,
+    0.5537,
+    0.5444,
+    0.4089,
+    0.7468,
+    0.7744,
+)
+
+
+class LanceWanVAE(nn.Module):
+    """Wan2.2 VAE wrapped for BAGEL's pipeline.
+
+    Exposes BAGEL's image-VAE surface — ``encode(BCHW) -> BC_zHW`` and
+    ``decode(BC_zHW) -> BCHW`` — by treating each image as a 1-frame video clip.
+    A 5-D ``encode_video``/``decode_video`` path is also provided for the
+    ``Lance_3B_Video`` checkpoint.
+
+    Construction is lazy: the heavy ``WanVAE_`` and ``Wan2.2_VAE.pth`` are not
+    materialized until first use. Once built, the inner module is registered as
+    a submodule so ``self.parameters()``, ``self.to(device)`` and ``vae_dtype =
+    next(vae.parameters()).dtype`` (used by BAGEL's decode path) all behave.
+    """
+
+    z_channels: int = 48
+    downsample_spatial: int = 16
+    downsample_temporal: int = 4
+
+    def __init__(
+        self,
+        vae_path: str,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device | str | None = None,
+        *,
+        lazy: bool = False,
+    ):
+        super().__init__()
+        self._vae_path = vae_path
+        self._dtype = dtype
+        self._device = torch.device(device) if device is not None else None
+        self._built = False
+        # Registered lazily once the .pth is loaded:
+        # self.model: WanVAE_   (set in _ensure_built)
+        if not lazy:
+            # BAGEL's decode path inspects ``next(vae.parameters()).dtype`` before
+            # the first encode/decode call, so we have to register the submodule
+            # eagerly when a real device is provided.  ``lazy=True`` is reserved
+            # for config-only / no-GPU paths.
+            self._ensure_built()
+
+    def _ensure_built(self) -> None:
+        if self._built:
+            return
+        device = self._device or (
+            next(self.parameters(), torch.empty(0)).device
+            if any(True for _ in self.parameters())
+            else torch.device("cpu")
+        )
+        # Wan2.2 VAE config from upstream Wan2_2_VAE defaults:
+        #   z_dim=48 (channels), c_dim=160 (encoder width), dec_dim=256 (decoder width),
+        #   dim_mult=[1,2,4,4], temperal_downsample=[False,True,True].
+        model = WanVAE_(
+            dim=160,
+            dec_dim=256,
+            z_dim=self.z_channels,
+            dim_mult=(1, 2, 4, 4),
+            num_res_blocks=2,
+            attn_scales=(),
+            temperal_downsample=(False, True, True),
+            dropout=0.0,
+        )
+        logger.info("Loading Wan2.2 VAE state dict from %s", self._vae_path)
+        state = torch.load(self._vae_path, map_location="cpu", weights_only=True)
+        missing, unexpected = model.load_state_dict(state, strict=False)
+        if missing:
+            logger.warning("Wan2.2 VAE missing keys (%d): %s", len(missing), missing[:5])
+        if unexpected:
+            logger.warning("Wan2.2 VAE unexpected keys (%d): %s", len(unexpected), unexpected[:5])
+        model = model.to(device=device, dtype=self._dtype).eval()
+        model.requires_grad_(False)
+        # Register as submodule so .to() / .parameters() pick it up.
+        self.model = model
+        # Per-channel normalization tensors (matches upstream Wan2_2_VAE.scale = [mean, 1/std]).
+        mean = torch.tensor(_WAN22_LATENT_MEAN, dtype=self._dtype, device=device)
+        inv_std = torch.tensor([1.0 / s for s in _WAN22_LATENT_STD], dtype=self._dtype, device=device)
+        # Buffers, so .to(device) follows along.
+        self.register_buffer("_latent_mean", mean, persistent=False)
+        self.register_buffer("_latent_inv_std", inv_std, persistent=False)
+        self._built = True
+
+    @torch.inference_mode()
+    def encode_video(self, video: torch.Tensor, *, use_sample: bool = True) -> torch.Tensor:
+        """Encode a 5-D clip ``[B, 3, T, H, W]`` -> latent ``[B, 48, t, h, w]``."""
+        self._ensure_built()
+        video = video.to(self.model.encoder.conv1.weight.dtype)
+        mu, log_var = self.model.encode(video, [self._latent_mean, self._latent_inv_std])
+        if use_sample:
+            std = torch.exp(0.5 * log_var)
+            mu = mu + std * torch.randn_like(std)
+        return mu
+
+    @torch.inference_mode()
+    def decode_video(self, latent: torch.Tensor) -> torch.Tensor:
+        """Decode a 5-D latent ``[B, 48, t, h, w]`` -> video ``[B, 3, T, H, W]``."""
+        self._ensure_built()
+        latent = latent.to(self.model.decoder.conv1.weight.dtype)
+        out = self.model.decode(latent, [self._latent_mean, self._latent_inv_std])
+        return out.clamp_(-1.0, 1.0)
+
+    # ----- BAGEL image-VAE surface (4-D) -------------------------------- #
+
+    @torch.inference_mode()
+    def encode(self, padded_images: torch.Tensor) -> torch.Tensor:
+        """``[B, 3, H, W]`` -> ``[B, 48, H/16, W/16]`` (single-frame image path).
+
+        Each image is wrapped as a 1-frame clip, encoded, and the temporal axis
+        is squeezed back out so the result matches BAGEL's iteration pattern.
+        """
+        if padded_images.dim() == 4:
+            video = padded_images.unsqueeze(2)  # B,C,1,H,W
+        elif padded_images.dim() == 5:
+            video = padded_images
+        else:
+            raise ValueError(f"LanceWanVAE.encode expects 4-D BCHW or 5-D BCTHW, got {tuple(padded_images.shape)}")
+        latent = self.encode_video(video, use_sample=True)
+        # B,48,t,h,w -> B,48,h,w when t == 1 (image path)
+        if latent.shape[2] == 1:
+            return latent.squeeze(2)
+        return latent
+
+    @torch.inference_mode()
+    def decode(self, latent: torch.Tensor) -> torch.Tensor:
+        """``[B, 48, h, w]`` -> ``[B, 3, H, W]`` (single-frame image path)."""
+        if latent.dim() == 4:
+            latent = latent.unsqueeze(2)  # B,48,1,h,w
+        elif latent.dim() != 5:
+            raise ValueError(f"LanceWanVAE.decode expects 4-D BCHW or 5-D BCTHW, got {tuple(latent.shape)}")
+        video = self.decode_video(latent)  # B,3,T,H,W
+        if video.shape[2] == 1:
+            return video.squeeze(2)
+        return video
+
+
+def build_wan22_vae(vae_path: str, dtype: torch.dtype = torch.bfloat16, device=None) -> LanceWanVAE:
+    """Convenience factory: lazy-construct a :class:`LanceWanVAE` adapter."""
+    return LanceWanVAE(vae_path=vae_path, dtype=dtype, device=device)
+
+
+__all__ = ["LanceWanVAE", "WanVAE_", "build_wan22_vae"]

vllm_omni/diffusion/registry.py

@@ -141,6 +141,11 @@ _DIFFUSION_MODELS = {
         "pipeline_bagel",
         "BagelPipeline",
     ),
+    "LancePipeline": (
+        "lance",
+        "pipeline_lance",
+        "LancePipeline",
+    ),
     "InternVLAA1Pipeline": (
         "internvla_a1",
         "pipeline_internvla_a1",
@@ -462,6 +467,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
     "LongCatImagePipeline": "get_longcat_image_post_process_func",
     "BagelPipeline": "get_bagel_post_process_func",
+    "LancePipeline": "get_lance_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",
     "StableDiffusion3Pipeline": "get_sd3_image_post_process_func",

vllm_omni/model_executor/models/lance/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance (ByteDance) — vLLM-Omni model_executor package.
+
+The runnable, registered path is the **single-stage** topology: a
+self-contained diffusion stage (``LancePipeline``) that owns the Qwen2-MoT
+LLM, the Qwen2.5-VL ViT, the Wan2.2 VAE and the tokenizer — directly analogous
+to ``bagel_single_stage``.
+
+The two-stage (separate AR "thinker" + DiT) topology additionally needs a
+``LanceConfig`` / ``LanceProcessor`` registered in the ``vllm`` package
+(mirroring ``vllm.transformers_utils.configs.bagel`` /
+``vllm.model_executor.models.bagel``); that is an explicit follow-up and is
+intentionally not wired here to avoid shipping a half-working AR stage.
+"""
+
+from vllm_omni.model_executor.models.lance.lance import (
+    OmniLanceForConditionalGeneration,
+)
+
+__all__ = ["OmniLanceForConditionalGeneration"]

vllm_omni/model_executor/models/lance/lance.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance AR-stage model wrapper (two-stage topology).
+
+Lance is BAGEL-lineage, so the AR/thinker stage is the same Qwen2-MoT model
+BAGEL uses.  This subclass exists so the two-stage pipeline has a distinct
+arch name to register against; behaviour is inherited unchanged from
+:class:`OmniBagelForConditionalGeneration`.
+
+IMPORTANT: the two-stage Lance topology is a follow-up.  It additionally
+requires a ``LanceConfig`` + ``LanceProcessor`` registered in the ``vllm``
+package (the base ``BagelForConditionalGeneration`` resolves its HF config /
+processor by ``model_type == "bagel"``; Lance ships ``model_type ==
+"qwen2_5_vl"`` in ``llm_config.json`` and no BAGEL-style top-level config).
+Until that lands, use the **single-stage** topology (``model_type: lance`` ->
+``LANCE_SINGLE_STAGE_PIPELINE`` -> ``LancePipeline``), which is fully wired and
+needs no ``vllm``-submodule changes.
+"""
+
+from __future__ import annotations
+
+from vllm_omni.model_executor.models.bagel.bagel import (
+    OmniBagelForConditionalGeneration,
+)
+
+
+class OmniLanceForConditionalGeneration(OmniBagelForConditionalGeneration):
+    """Lance AR-stage model.  Identical to BAGEL's (same Qwen2-MoT weights);
+    distinct class only so the two-stage pipeline can register a ``lance``
+    arch separately.  See module docstring for the two-stage follow-up."""
+
+    # Same packed→sublayer LoRA / MoT mapping as BAGEL (weight names match).
+    packed_modules_mapping = OmniBagelForConditionalGeneration.packed_modules_mapping

vllm_omni/model_executor/models/lance/pipeline.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance pipeline topologies.
+
+Single-stage (runnable / registered):
+  Stage 0: DiT — self-contained diffusion stage that handles all modalities
+           (text2img, image_edit, x2t_image) internally via its own
+           Qwen2-MoT LLM, Qwen2.5-VL ViT, Wan2.2 VAE, and tokenizer.
+           Mirrors ``bagel_single_stage``.
+
+Two-stage (follow-up — see vllm_omni/model_executor/models/lance/lance.py):
+  Stage 0: Thinker (AR) -> Stage 1: DiT.  Needs a LanceConfig/LanceProcessor
+  in the ``vllm`` package; defined here for completeness but gated behind the
+  follow-up so it is not selected by default.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.lance"
+
+# --- Single-stage (default, fully wired) -------------------------------- #
+LANCE_SINGLE_STAGE_PIPELINE = PipelineConfig(
+    model_type="lance",
+    model_arch="LancePipeline",
+    hf_architectures=(),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="dit",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(),
+            final_output=True,
+            final_output_type="image",
+        ),
+    ),
+)
+
+# --- Two-stage (follow-up; not registered as default) ------------------- #
+LANCE_PIPELINE = PipelineConfig(
+    model_type="lance_two_stage",
+    model_arch="OmniLanceForConditionalGeneration",
+    hf_architectures=(),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="thinker",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            model_arch="OmniLanceForConditionalGeneration",
+            engine_output_type="text",
+            prompt_expand_func=f"{_PROC}.expand_cfg_prompts",
+            omni_kv_config={
+                "need_send_cache": True,
+                "kv_transfer_criteria": {"type": "prefill_finished"},
+            },
+            sampling_constraints={"detokenize": True},
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            execution_type=StageExecutionType.DIFFUSION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            cfg_kv_collect_func=f"{_PROC}.collect_cfg_kv_caches",
+            omni_kv_config={"need_recv_cache": True},
+        ),
+    ),
+)

vllm_omni/model_executor/models/registry.py

@@ -134,6 +134,12 @@ _OMNI_MODELS = {
         "bagel",
         "OmniBagelForConditionalGeneration",
     ),
+    ## lance (ByteDance) — BAGEL-lineage; two-stage AR arch (follow-up)
+    "OmniLanceForConditionalGeneration": (
+        "lance",
+        "lance",
+        "OmniLanceForConditionalGeneration",
+    ),
     "HunyuanImage3ForCausalMM": (
         "hunyuan_image3",
         "hunyuan_image3",

vllm_omni/model_executor/stage_input_processors/lance.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stage input processor for Lance.
+
+Lance's classifier-free-guidance prompt expansion and KV-cache collection are
+model-agnostic and identical to BAGEL's (3-branch CFG: gen / cfg_text /
+cfg_img).  Re-export BAGEL's implementations so the two-stage pipeline config
+can reference ``...stage_input_processors.lance.{expand_cfg_prompts,
+collect_cfg_kv_caches}`` without duplicating logic.
+"""
+
+from vllm_omni.model_executor.stage_input_processors.bagel import (  # noqa: F401
+    CFG_IMG_SUFFIX,
+    CFG_TEXT_SUFFIX,
+    GEN_THINK_SYSTEM_PROMPT,
+    VLM_THINK_SYSTEM_PROMPT,
+    ExpandedPrompt,
+    collect_cfg_kv_caches,
+    expand_cfg_prompts,
+    expand_cfg_prompts_think,
+)
+
+__all__ = [
+    "CFG_IMG_SUFFIX",
+    "CFG_TEXT_SUFFIX",
+    "GEN_THINK_SYSTEM_PROMPT",
+    "VLM_THINK_SYSTEM_PROMPT",
+    "ExpandedPrompt",
+    "collect_cfg_kv_caches",
+    "expand_cfg_prompts",
+    "expand_cfg_prompts_think",
+]