[Lance] Fix x2t cache-offset SDPA mask for Q != K

PyTorch SDPA's is_causal=True with Q != K uses tril(diagonal=0) which is
top-left aligned. For a single decode step (Q=1) against a long ViT+text
cache (K=574+), only q_0 -> k_0 is unmasked — producing degenerate output
("TheTheThe..." token loops).

Bypass DiffusionAttention (whose _maybe_reshape_attn_mask only handles
2-D (B, K) masks) and call torch.nn.functional.scaled_dot_product_attention
directly with an explicit 4-D (1, 1, Q, K) bottom-right-aligned causal mask
when cache_len > 0. Also cast q/k to fp32 before RMSNorm for precision.

Verified all 7 tasks on Lance:
- x2t_image:  "No" (valid VQA answer)
- x2t_video:  "The purple sphere ... right direc..." (valid VQA answer)
- t2i / i2v / t2v / image_edit / video_edit all produce expected media.

Signed-off-by: lishunyang12 <lishunyang12@163.com>

.gitignore

@@ -270,3 +270,6 @@ vllm_omni/_version.py
 .worktrees/
 # CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
 tests/.ci_generated/
+
+# Local Lance precision-matching dump scripts (not part of the repo)
+lance_compare/

examples/offline_inference/lance/README.md

@@ -0,0 +1,52 @@
+# 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 \
+    --height 1024 --width 1024 \
+    --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, 1024×1024 (override with `--height` / `--width`).  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,333 @@
+# 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", "image2video"],
+        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("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
+    p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
+    p.add_argument(
+        "--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
+    )
+    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(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help=(
+            "Override the task system prompt (used by x2t_image / x2t_video for "
+            'per-example QA instructions like "Look at the image carefully and '
+            'answer the question."; without it x2t falls back to a '
+            "caption-style default and the model describes instead of answering)."
+        ),
+    )
+    # Lance is single-stage diffusion — no deploy YAML needed.  Required
+    # engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
+    # ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
+    # and ``create_default_diffusion`` materializes the stage config.
+
+    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_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
+    """Image-to-Video (no first-frame pin): image + text → long video.
+
+    Passes the input image via ``multi_modal_data.first_frame``.  The pipeline
+    treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
+    multi-frame video at the requested ``num_frames`` × ``video_h`` ×
+    ``video_w`` shape.
+    """
+    import os as _os
+
+    if not image_path or not _os.path.exists(image_path):
+        raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
+    from PIL import Image as _Image
+
+    img = _Image.open(image_path).convert("RGB")
+    extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
+    return [
+        {
+            "prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
+            "multi_modal_data": {"first_frame": img},
+            "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}")
+
+    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_path},
+            "modalities": ["video"],
+            "extra_args": extra,
+        }
+        for p in prompts
+    ]
+
+
+def _format_video2text_prompts(prompts, video_path, system_prompt=None):
+    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, system_prompt=system_prompt),
+            "multi_modal_data": {"video": video},
+            "modalities": ["text"],
+        }
+        for p in prompts
+    ]
+
+
+def _format_img2text_prompts(prompts, image_path, system_prompt=None):
+    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, system_prompt=system_prompt),
+            "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
+
+    # Video modalities require the Lance_3B_Video checkpoint (3D
+    # latent_pos_embed shaped (126976, 2048)).  Without it, vllm-omni silently
+    # loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
+    # immediately indexes out of bounds.  Rewrite the model path so
+    # LancePipeline._select_video_variant picks the video subfolder.
+    resolved_model = args.model
+    if args.modality in ("text2video", "video2video", "video2text", "image2video"):
+        if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
+            candidate = os.path.join(resolved_model, "Lance_3B_Video")
+            if os.path.isdir(candidate):
+                resolved_model = candidate
+
+    omni_kwargs = vars(args).copy()
+    omni_kwargs["model"] = resolved_model
+    # Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
+    omni_kwargs.setdefault("pipeline", "lance")
+    omni_kwargs.setdefault("max_num_batched_tokens", 32768)
+    omni_kwargs.setdefault("max_num_seqs", 1)
+    omni_kwargs.setdefault("enforce_eager", True)
+    omni_kwargs.setdefault("trust_remote_code", True)
+    omni_kwargs.setdefault("enable_prefix_caching", False)
+    omni_kwargs.setdefault("async_chunk", False)
+    omni = Omni(**omni_kwargs)
+
+    if args.modality == "img2text":
+        formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
+    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, system_prompt=args.system_prompt)
+    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
+        )
+    elif args.modality == "image2video":
+        formatted = _format_image2video_prompts(
+            prompts, args.image_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
+    if args.height is not None:
+        diffusion_params.height = args.height  # type: ignore
+    if args.width is not None:
+        diffusion_params.width = args.width  # 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", "image2video"):
+        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"

recipes/ByteDance/Lance.md

@@ -0,0 +1,177 @@
+# Lance
+
+> Unified autoregressive + diffusion multimodal (text / image / video)
+
+## Summary
+
+- Vendor: ByteDance
+- Model: [`bytedance-research/Lance`](https://huggingface.co/bytedance-research/Lance) (`Lance_3B`, `Lance_3B_Video`)
+- Task: text2img, text2video, img2img (image edit), video2video (video edit),
+  img2text (image understanding), video2text (video understanding)
+- Mode: Offline inference, Online serving (OpenAI-compatible API)
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to run Lance — a 3B unified AR + diffusion multimodal model
+on a Qwen2.5-VL backbone — via vLLM-Omni. Lance is **BAGEL-lineage**: the
+released checkpoint uses the same `*_moe_gen` Mixture-of-Transformers weight
+layout as BAGEL, so vLLM-Omni reuses BAGEL's transformer core and specializes
+only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
+
+All six single-stage modalities are supported on the `Lance_3B` checkpoint;
+`text2video` and `video2video` additionally need `Lance_3B_Video` for the 3-D
+`latent_pos_embed` table. The two-stage AR + DiT topology is deferred and
+needs `LanceConfig` / `LanceProcessor` registered in the `vllm` package (a
+separate upstream PR).
+
+## References
+
+- Offline example:
+  [`examples/offline_inference/lance/end2end.py`](../../examples/offline_inference/lance/end2end.py)
+- Online serving:
+  [`examples/online_serving/lance/`](../../examples/online_serving/lance/)
+- Deploy config:
+  [`vllm_omni/deploy/lance.yaml`](../../vllm_omni/deploy/lance.yaml)
+- E2E tests:
+  [`tests/e2e/online_serving/test_lance.py`](../../tests/e2e/online_serving/test_lance.py)
+- HuggingFace model page:
+  [bytedance-research/Lance](https://huggingface.co/bytedance-research/Lance)
+- PR: [vllm-project/vllm-omni#3710](https://github.com/vllm-project/vllm-omni/pull/3710)
+  (closes [#3697](https://github.com/vllm-project/vllm-omni/issues/3697))
+
+## Hardware Support
+
+## GPU
+
+### 1x NVIDIA B300 / A100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12
+- Driver / runtime: CUDA ≥ 12.4
+- vLLM-Omni version: 0.18.x.dev
+
+#### Command — text-to-image (default)
+
+```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 \
+    --height 1024 --width 1024 \
+    --seed 42 --output ./out
+```
+
+Defaults match upstream `inference_lance.sh`: 30 denoising steps,
+`timestep-shift 3.5`, text CFG 4.0, seed 42, 1024×1024.
+
+#### Command — image edit (img2img)
+
+```bash
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance --modality img2img \
+    --image-path /path/to/input.png \
+    --prompts "Convert this into a vibrant cartoon-style illustration" \
+    --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
+    --output ./out
+```
+
+The Lance-native VAE prefill scatters Wan2.2 latents into the LLM query
+sequence; no separate image encoder is needed.
+
+#### Command — text-to-video
+
+```bash
+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
+```
+
+Use the `Lance_3B_Video` subfolder for any video path so the 3-D
+`latent_pos_embed` table is loaded; image / understanding paths can point at
+the top-level repo and resolve the right sub-checkpoint automatically.
+
+#### Command — image / video understanding
+
+```bash
+# Image → text (caption / VQA)
+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." \
+    --do-sample --text-temperature 0.8
+
+# Video → text
+python examples/offline_inference/lance/end2end.py \
+    --model bytedance-research/Lance --modality video2text \
+    --video-path /path/to/clip.mp4 \
+    --prompts "What is happening in this video?"
+```
+
+Sampling is enabled by default at `--text-temperature 0.8` for the
+understanding paths because Lance's greedy decoder emits an immediate EOS for
+many prompts.
+
+#### Verification
+
+```bash
+pytest -s -v tests/e2e/online_serving/test_lance.py
+```
+
+#### Notes
+
+- BF16 footprint: ~7 GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE; comfortably fits
+  on a single 16 GB+ GPU for `Lance_3B`.
+- `rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired
+  through `BagelRotaryEmbedding`. `text2img` uses scalar position ids
+  (BAGEL-equivalent); `img2text` / `video` / edit paths thread per-axis 3-D
+  position ids.
+- `video_edit` quality is more abstract than `text2video` at the same
+  resolution (known position-id offset between VAE-prefill and gen-latent
+  blocks). Functionally correct end-to-end.
+
+## Online Serving
+
+Lance supports all single-stage modalities via the OpenAI-compatible
+`/v1/chat/completions` API.
+
+### Launch
+
+```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
+```
+
+For `text2video` / `video2video`, set
+`MODEL=bytedance-research/Lance/Lance_3B_Video`.
+
+### Send requests
+
+```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
+
+# Image understanding
+python examples/online_serving/lance/openai_chat_client.py \
+    --prompt "Describe this image" \
+    --modality img2text --image-url photo.jpg
+```
+
+The client is shared with BAGEL — same OpenAI message format, same
+`modalities` / `num_inference_steps` / `seed` / `height` / `width` knobs.

recipes/README.md

@@ -29,6 +29,7 @@ recipes/
 |--------|------|----------|
 | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
 | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
+| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
 | [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
 | [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
 | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |

tests/e2e/online_serving/test_lance.py

@@ -0,0 +1,144 @@
+# 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 \\
+        --pipeline lance --enforce-eager --trust-remote-code --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 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"
+
+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}"
+
+# Lance is single-stage diffusion — no deploy YAML.  Pass the engine
+# knobs that used to live in ``vllm_omni/deploy/lance.yaml`` via CLI
+# flags so ``create_default_diffusion`` builds the stage config.
+_LANCE_SERVE_ARGS = [
+    "--pipeline",
+    "lance",
+    "--max-num-batched-tokens",
+    "32768",
+    "--max-num-seqs",
+    "1",
+    "--enforce-eager",
+    "--trust-remote-code",
+    "--no-enable-prefix-caching",
+    "--no-async-chunk",
+]
+
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        server_args=_LANCE_SERVE_ARGS,
+        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)

vllm_omni/config/pipeline_registry.py

@@ -65,6 +65,11 @@ _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 DiT.
+    "lance": (
+        "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,20 @@
+# Lance (ByteDance) — single-stage diffusion pipeline selector.
+#
+# Lance's HF config.json is descriptive metadata (no ``model_type``),
+# so the loader can't auto-detect from the model dir alone.  This YAML
+# only tells the loader "use the Lance pipeline" — one stage, no multi-
+# stage orchestration.
+
+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

@@ -847,6 +847,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.
 
@@ -900,6 +915,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))
@@ -910,6 +934,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

@@ -106,7 +106,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):
@@ -120,10 +129,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":
+        if rope_type in ("default", "mrope"):
-            # transformers>=5.0 removed the 'default' entry from
+            # mrope shares the default sinusoidal frequency basis; the
-            # ROPE_INIT_FUNCTIONS; use the plain sinusoidal formula.
+            # multimodal split happens at forward time, not at init.
             rope_theta = (
                 getattr(config, "rope_theta", None)
                 or rope_parameters.get("rope_theta")
@@ -147,11 +163,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)
@@ -541,8 +582,12 @@ class PackedAttentionMoT(nn.Module):
         q = q.view(-1, self.num_heads, self.head_dim)
         k = k.view(-1, self.num_kv_heads, self.head_dim)
         v = v.view(-1, self.num_kv_heads, self.head_dim)
-        q = self.q_norm(q)
+        # Pre-merge code cast to float32 before q_norm/k_norm — bf16
-        k = self.k_norm(k)
+        # RMSNorm accumulation drift compounds across segmented prefill
+        # (x2t builds the cache via 3+ sequential forward_cache_update_*
+        # calls), producing degenerate token repetition in generate_text.
+        q = self.q_norm(q.to(torch.float32))
+        k = self.k_norm(k.to(torch.float32))
 
         cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
         q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
@@ -557,16 +602,52 @@ class PackedAttentionMoT(nn.Module):
             cache_v = past_key_values.value_cache[self.layer_idx]
             full_k = torch.cat([cache_k, k], dim=0)
             full_v = torch.cat([cache_v, v], dim=0)
+            cache_len = cache_k.shape[0]
         else:
             full_k = k
             full_v = v
+            cache_len = 0
 
-        attn = self.attn_causal if is_causal else self.attn_noncausal
+        if is_causal and cache_len > 0:
-        attn_out = attn(
+            # PyTorch SDPA's ``is_causal=True`` with ``Q != K`` uses
-            q.unsqueeze(0),
+            # ``tril(diagonal=0)`` — top-left aligned, so for Q=1 against a
-            full_k.unsqueeze(0),
+            # long cache only ``q_0 -> k_0`` is unmasked (degenerate
-            full_v.unsqueeze(0),
+            # generate_text output).  Build the correct bottom-right
-        )
+            # aligned mask manually: ``mask[i, j] = -inf if j > cache_len
+            # + i`` so each new query attends to all of ``cache + self
+            # up to i``.  Bypass ``DiffusionAttention`` entirely — its
+            # ``_maybe_reshape_attn_mask`` helper only handles 2-D
+            # ``(B, K)`` shape, not ``(Q, K)``, so the reshape path
+            # silently drops the per-query mask for Q>1.
+            Q_len = q.shape[0]
+            K_len = full_k.shape[0]
+            arange_q = torch.arange(Q_len, device=q.device).unsqueeze(1)
+            arange_k = torch.arange(K_len, device=q.device).unsqueeze(0)
+            mask = arange_k > (cache_len + arange_q)  # (Q, K) bool
+            attn_bias = torch.zeros(Q_len, K_len, dtype=q.dtype, device=q.device)
+            attn_bias.masked_fill_(mask, float("-inf"))
+            # Permute to (B=1, H, S, D) for SDPA.
+            q_4d = q.unsqueeze(0).permute(0, 2, 1, 3)
+            k_4d = full_k.unsqueeze(0).permute(0, 2, 1, 3)
+            v_4d = full_v.unsqueeze(0).permute(0, 2, 1, 3)
+            attn_out_4d = torch.nn.functional.scaled_dot_product_attention(
+                q_4d,
+                k_4d,
+                v_4d,
+                attn_mask=attn_bias.unsqueeze(0).unsqueeze(0),  # (1, 1, Q, K)
+                dropout_p=0.0,
+                is_causal=False,
+                scale=1.0 / (self.head_dim**0.5),
+                enable_gqa=(self.num_heads != self.num_kv_heads),
+            )
+            attn_out = attn_out_4d.permute(0, 2, 1, 3)
+        else:
+            attn = self.attn_causal if is_causal else self.attn_noncausal
+            attn_out = attn(
+                q.unsqueeze(0),
+                full_k.unsqueeze(0),
+                full_v.unsqueeze(0),
+            )
 
         attn_out = attn_out.squeeze(0).reshape(-1, self.q_size)
         attn_out, _ = self.o_proj(attn_out)
@@ -1614,10 +1695,26 @@ class Bagel(nn.Module):
         return_trajectory_latents: bool = False,
         scheduler: object | None = None,
         scheduler_kwargs: dict | None = None,
+        # Lance i2v: tokens to freeze at their initial (encoded-image)
+        # value throughout the denoise loop.  Matches upstream's
+        # ``mse_loss_indexes``-exclusion behaviour from PR #33: cond
+        # positions are never updated and get ``timestep=0``.
+        frame_condition_token_indexes: torch.LongTensor | None = None,
     ):
         x_t = packed_init_noises
+        # Snapshot the pinned subtensor BEFORE the denoise loop touches
+        # x_t; the cond positions in packed_init_noises hold the
+        # VAE-encoded conditioning latent that must be preserved verbatim.
+        pinned_x_t = None
+        if frame_condition_token_indexes is not None:
+            frame_condition_token_indexes = frame_condition_token_indexes.to(x_t.device).long()
+            pinned_x_t = x_t[frame_condition_token_indexes].clone()
 
-        timesteps = torch.linspace(1, 0, num_timesteps, device=x_t.device)
+        # Use num_timesteps + 1 sample points so we get `num_timesteps` denoise
+        # steps after dropping the terminal t=0 (which has no dt).  Upstream
+        # Lance / BAGEL both use this convention; without the +1 we silently
+        # run one fewer denoise iteration than the user asked for.
+        timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device)
         timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps)
         dts = timesteps[:-1] - timesteps[1:]
         timesteps = timesteps[:-1]
@@ -1669,6 +1766,12 @@ class Bagel(nn.Module):
                 trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
+            if frame_condition_token_indexes is not None:
+                # Cond positions stay at t=0 (clean signal).  Matches upstream
+                # PR #33 lance.py line 1605:
+                #     timestep[current_vae_mse_indexes_local_in_vae] = t
+                # (cond positions remain at the ``torch.zeros`` init value).
+                timestep[frame_condition_token_indexes] = 0.0
                 in_cfg_window = t > cfg_interval[0] and t <= cfg_interval[1]
                 cfg_text_scale_ = cfg_text_scale if in_cfg_window else 1.0
                 cfg_img_scale_ = cfg_img_scale if in_cfg_window else 1.0
@@ -1732,6 +1835,12 @@ class Bagel(nn.Module):
                 trajectory_latents.append(x_t.clone())
             for i, t in enumerate(timesteps):
                 timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
+            if frame_condition_token_indexes is not None:
+                # Cond positions stay at t=0 (clean signal).  Matches upstream
+                # PR #33 lance.py line 1605:
+                #     timestep[current_vae_mse_indexes_local_in_vae] = t
+                # (cond positions remain at the ``torch.zeros`` init value).
+                timestep[frame_condition_token_indexes] = 0.0
                 v_t = self.forward_single_branch(
                     x_t=x_t,
                     timestep=timestep,
@@ -1783,7 +1892,11 @@ 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),
+                # Concat along the *sequence* (last) dim — works for both
+                # scalar 1-D position ids ``(S,)`` and Lance's 3-D mRoPE
+                # ``(3, S)``.  Default ``dim=0`` collapses ``(3, S) + (3, S)``
+                # to ``(6, S)`` and breaks ``BagelRotaryEmbedding``.
+                "batched_position_ids": torch.cat(branches_pid, dim=-1),
                 "batched_text_indexes": torch.cat(
                     [packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)]
                 ),
@@ -1798,6 +1911,12 @@ class Bagel(nn.Module):
 
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
+            if frame_condition_token_indexes is not None:
+                # Cond positions stay at t=0 (clean signal).  Matches upstream
+                # PR #33 lance.py line 1605:
+                #     timestep[current_vae_mse_indexes_local_in_vae] = t
+                # (cond positions remain at the ``torch.zeros`` init value).
+                timestep[frame_condition_token_indexes] = 0.0
             if t > cfg_interval[0] and t <= cfg_interval[1]:
                 cfg_text_scale_ = cfg_text_scale
                 cfg_img_scale_ = cfg_img_scale
@@ -1828,6 +1947,12 @@ class Bagel(nn.Module):
                     trajectory_log_probs.append(out.log_prob)
             else:
                 x_t = x_t - v_t.to(x_t.device) * dts[i]  # velocity pointing from data to noise
+                if pinned_x_t is not None:
+                    # i2v: restore cond positions to their encoded-image
+                    # latent.  Matches upstream PR #33 lance.py line 1712:
+                    #     x_t[mse_indexes] = x_t[mse_indexes] - v_t[mse_indexes] * dts[i]
+                    # (cond positions are excluded from the update).
+                    x_t[frame_condition_token_indexes] = pinned_x_t
             if return_trajectory_latents:
                 trajectory_latents.append(x_t.clone())
                 trajectory_timesteps.append(timesteps[i])
@@ -1859,6 +1984,7 @@ class Bagel(nn.Module):
         return_trajectory_latents: bool = False,
         scheduler: object | None = None,
         scheduler_kwargs: dict | None = None,
+        frame_condition_token_indexes: torch.LongTensor | None = None,
     ):
         """CFG parallel denoising loop: each rank computes one CFG branch.
 
@@ -1918,6 +2044,12 @@ class Bagel(nn.Module):
 
         for i, t in enumerate(timesteps):
             timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
+            if frame_condition_token_indexes is not None:
+                # Cond positions stay at t=0 (clean signal).  Matches upstream
+                # PR #33 lance.py line 1605:
+                #     timestep[current_vae_mse_indexes_local_in_vae] = t
+                # (cond positions remain at the ``torch.zeros`` init value).
+                timestep[frame_condition_token_indexes] = 0.0
             use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
 
             if use_cfg_this_step:
@@ -2070,7 +2202,7 @@ class Bagel(nn.Module):
             packed_sequence = packed_text_embedding.new_zeros((int(local_seqlens.sum()), self.hidden_size))
             packed_sequence[local_text_indexes] = packed_text_embedding
 
-            assert timestep.unique().shape[0] == 1
+            # i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
             packed_pos_embed = self.latent_pos_embed(local_vae_pos_ids)
             local_timestep = timestep[: local_x_t.shape[0]]
             packed_timestep_embeds = self.time_embedder(local_timestep)
@@ -2107,7 +2239,7 @@ class Bagel(nn.Module):
         packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
         packed_sequence[packed_text_indexes] = packed_text_embedding
 
-        assert timestep.unique().shape[0] == 1
+        # i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
         packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
         packed_timestep_embeds = self.time_embedder(timestep)
         x_t_emb = self.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed
@@ -2159,7 +2291,7 @@ class Bagel(nn.Module):
         packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
         packed_sequence[packed_text_indexes] = packed_text_embedding
 
-        assert timestep.unique().shape[0] == 1
+        # i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
         packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
         packed_timestep_embeds = self.time_embedder(timestep)
         x_t = self.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -355,7 +355,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
 
         gen_params = BagelGenParams(
             num_timesteps=int(req.sampling_params.num_inference_steps or 50),
-            timestep_shift=3.0,
+            timestep_shift=float(extra_args.get("timestep_shift", 3.0)),
             cfg_text_scale=cfg_text_scale,
             cfg_img_scale=cfg_img_scale,
             cfg_interval=cfg_interval,
@@ -820,6 +820,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         custom = {}
         if think_text is not None:
             custom["think_text"] = think_text
+        # Mirror the PIL image into ``custom_output`` so callers reading via
+        # the orchestrator IPC boundary (which strips the bare ``output``
+        # field) can still recover the result.  ``video_frames`` already
+        # uses this pattern.
+        custom["image"] = img
 
         return DiffusionOutput(
             output=img,

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/prompts.py

@@ -0,0 +1,121 @@
+# 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:"
+    ),
+    # i2v reuses the t2v system prompt (matches upstream Lance PR #33's
+    # ``data/common.py`` where ``system_prompt_type in ("t2v", "i2v")``
+    # selects the same prompt pool).
+    ("i2v", "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,
+    system_prompt: 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.  ``system_prompt`` overrides the default task system
+    prompt — required for x2t QA examples whose upstream JSON carries a
+    per-example instruction (e.g. ``"Look at the image carefully and answer
+    the question."``); without it x2t falls back to the caption-style
+    default and the model describes instead of answering.
+
+    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", "i2v", "x2t_video", "video_edit"} else "image"
+    sys_prompt = (
+        system_prompt if system_prompt else 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,870 @@
+# 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
+            squeeze_time = True  # image path: caller iterates 3-D latents
+        elif padded_images.dim() == 5:
+            video = padded_images
+            # video path: caller expects 4-D latents (C,T,H,W), even when T==1
+            # (e.g. video2video with a 1-frame reference clip).  Squeezing here
+            # would drop the temporal axis and break the downstream indexing
+            # ``latent[:, :t_lat, ...]`` in ``forward_cache_update_vae``.
+            squeeze_time = False
+        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 only)
+        if squeeze_time and 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",
+    ),
     "MingImagePipeline": (
         "ming_flash_omni",
         "pipeline_ming_imagegen",
@@ -471,6 +476,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",
     "MingImagePipeline": "get_ming_image_post_process_func",
     "InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
     "LongCatImageEditPipeline": "get_longcat_image_post_process_func",

vllm_omni/model_executor/models/lance/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance (ByteDance) — vLLM-Omni model_executor package.
+
+Single-stage topology: a self-contained diffusion stage (``LancePipeline``
+defined in ``vllm_omni.diffusion.models.lance``) 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 ``PipelineConfig`` lives in :mod:`.pipeline`.
+"""

vllm_omni/model_executor/models/lance/pipeline.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Lance pipeline topology.
+
+Single-stage DiT — self-contained diffusion stage that handles all modalities
+(text2img, image_edit, x2t_image, video) internally via its own Qwen2-MoT
+LLM, Qwen2.5-VL ViT, Wan2.2 VAE, and tokenizer. Mirrors ``bagel_single_stage``.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+LANCE_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",
+        ),
+    ),
+)