[Lance] Add Gradio demo and align inference with upstream

- Add ``examples/offline_inference/lance/gradio_demo.py`` — a unified 7-task UI
  (t2i / t2v / i2v / image_edit / video_edit / x2t_image / x2t_video) that
  loads ``Lance_3B`` and ``Lance_3B_Video`` on two GPUs so image and video
  tasks run concurrently.  Surfaces official Lance examples per task; aspect
  ratio is auto-snapped to one of upstream's 6 buckets from the input
  image/video.
- Fix x2t_image / x2t_video output format: route x2t_image through a new
  ``_forward_x2t_image`` handler (was falling through to the img2img path
  and emitting garbage), build the chat prefill in three segments
  (prefix → ViT → suffix) so the ``<|im_end|>`` markers line up with
  upstream's template, and start generation from a newline token instead of
  ``<|im_start|>`` so the assistant turn doesn't close itself with an
  immediate EOS.  Add a ``system_prompt`` override to ``render_lance_prompt``
  so per-example QA instructions surface ("Look at the image carefully and
  answer the question.") instead of the caption-style default.
- ``_forward_image_edit``: also stash the PIL image into ``custom_output``
  so callers reading ``outputs[0].output`` *or* ``custom_output["image"]``
  both work — the orchestrator IPC strips the bare ``output`` field across
  process boundaries.
- ``pipeline_registry.py``, ``model_executor/models/lance/*``, registry
  cleanup: remove the dead Lance-specific autoregressive shim
  (``lance.py``) and its stage-input processor; the diffusion pipeline owns
  the full flow now.
- Aspect-ratio buckets aligned with upstream's ``BucketResize`` (6 ratios:
  21:9 / 16:9 / 4:3 / 1:1 / 3:4 / 9:16, sized for ``resolution_vae`` ∈
  {480, 640, 768}).  3:2 / 2:3 dropped — upstream snaps those to 4:3 / 3:4.
- Gradio defaults match upstream's official scripts (cfg_text_scale=4.0,
  validation_num_timesteps=30, validation_timestep_shift=3.5, seed=42); only
  the demo *duration* is intentionally shorter than upstream's 121-frame
  default to keep the UI snappy (user can extend to 10s via slider for
  byte-identical reproduction).

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

Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>

examples/offline_inference/lance/end2end.py

@@ -33,7 +33,7 @@ def parse_args():
     p.add_argument(
         "--modality",
         default="text2img",
-        choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video"],
+        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.")
@@ -43,7 +43,9 @@ def parse_args():
     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=8, help="Output video FPS when saving MP4.")
+    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).")
@@ -64,6 +66,17 @@ def parse_args():
         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)."
+        ),
+    )
     p.add_argument(
         "--deploy-config",
         type=str,
@@ -98,6 +111,33 @@ def _format_text2video_prompts(prompts, num_frames, video_h, 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
 
@@ -121,14 +161,12 @@ def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_
 
     if not video_path or not _os.path.exists(video_path):
         raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
-    import imageio.v3 as iio
 
-    video = iio.imread(video_path)
     extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
     return [
         {
             "prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
-            "multi_modal_data": {"video": video},
+            "multi_modal_data": {"video": video_path},
             "modalities": ["video"],
             "extra_args": extra,
         }
@@ -136,7 +174,7 @@ def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_
     ]
 
 
-def _format_video2text_prompts(prompts, video_path):
+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):
@@ -146,7 +184,7 @@ def _format_video2text_prompts(prompts, video_path):
     video = iio.imread(video_path)
     return [
         {
-            "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK),
+            "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
             "multi_modal_data": {"video": video},
             "modalities": ["text"],
         }
@@ -154,7 +192,7 @@ def _format_video2text_prompts(prompts, video_path):
     ]
 
 
-def _format_img2text_prompts(prompts, image_path):
+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):
@@ -162,7 +200,7 @@ def _format_img2text_prompts(prompts, image_path):
     img = Image.open(image_path).convert("RGB")
     return [
         {
-            "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK),
+            "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
             "multi_modal_data": {"image": img},
             "modalities": ["text"],
         }
@@ -184,23 +222,39 @@ def main():
 
     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["deploy_config"] = args.deploy_config
-    omni_kwargs["model"] = args.model
+    omni_kwargs["model"] = resolved_model
     omni = Omni(**omni_kwargs)
 
     if args.modality == "img2text":
-        formatted = _format_img2text_prompts(prompts, args.image_path)
+        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)
+        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)
 
@@ -234,7 +288,7 @@ def main():
             print(f"[Output {i}] {text}")
         return
 
-    if args.modality in ("text2video", "video2video"):
+    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")

vllm_omni/config/pipeline_registry.py

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

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1819,8 +1819,20 @@ 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()
 
         # Use num_timesteps + 1 sample points so we get `num_timesteps` denoise
         # steps after dropping the terminal t=0 (which has no dt).  Upstream
@@ -1887,6 +1899,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
@@ -1959,6 +1977,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,
@@ -2051,6 +2075,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
@@ -2084,6 +2114,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])
@@ -2124,6 +2160,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.
 
@@ -2192,6 +2229,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:
@@ -2353,7 +2396,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)
@@ -2411,7 +2454,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
@@ -2469,7 +2512,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/lance/pipeline_lance.py

@@ -506,12 +506,16 @@ class LancePipeline(BagelPipeline):
         if isinstance(first_prompt, dict):
             modalities = first_prompt.get("modalities") or []
             mm_data = first_prompt.get("multi_modal_data") or {}
+        if "video" in modalities and mm_data.get("first_frame") is not None:
+            return self._forward_i2v(req)
         if "video" in modalities and mm_data.get("video") is not None:
             return self._forward_video_edit(req)
         if "video" in modalities:
             return self._forward_t2v(req)
         if "text" in modalities and mm_data.get("video") is not None:
             return self._forward_x2t_video(req)
+        if "text" in modalities and mm_data.get("image") is not None:
+            return self._forward_x2t_image(req)
         if "image" in modalities and (mm_data.get("img2img") is not None or mm_data.get("image") is not None):
             return self._forward_image_edit(req)
         # t2i falls through to BAGEL parent.  Inject Lance defaults if the
@@ -1000,13 +1004,386 @@ class LancePipeline(BagelPipeline):
             )
 
         img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
+        # Put the PIL image into ``custom_output`` too — the orchestrator
+        # serializes ``custom_output`` across the IPC boundary but strips
+        # the bare ``output`` field, so consumers that only see
+        # ``outputs[0].output`` (e.g. gradio_demo) get ``None`` otherwise.
         return DiffusionOutput(
             output=img,
-            custom_output={"image_shape": image_shape},
+            custom_output={"image": img, "image_shape": image_shape},
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
     @torch.inference_mode()
+    @torch.inference_mode()
+    def _forward_i2v(self, req):
+        """Image-to-Video (upstream-faithful first-frame conditioning).
+
+        Mirrors upstream Lance PR #33's ``ff2v_sample`` + ``validation_gen_KVcache``
+        math byte-for-byte at the algorithmic level:
+
+          1. Treat the input image as 4 replicated pixel frames
+             ``(3, 4, H, W)``.  This matches upstream
+             ``vae_tensor[:, :4, :, :] = vae_image_tensor[:, 0:1, :, :].repeat(1, 4, 1, 1)``
+             — required because Wan2.2 VAE has temporal stride 4, so 4
+             pixel frames collapse to 1 latent frame.
+
+          2. VAE-encode → first latent frame at the image content.
+
+          3. Inject the encoded latent into ``packed_init_noises`` at
+             the first ``h_lat × w_lat`` token slice (frame 0).
+
+          4. Pass ``frame_condition_token_indexes = arange(0, h_lat*w_lat)``
+             to ``Bagel.generate_image`` so that:
+               * ``timestep[cond] = 0`` every iter (clean signal),
+               * ``x_t[cond]`` is restored verbatim after each
+                 velocity update (never denoised).
+
+        Plus a ViT+VAE-ref prefill of the same image (so the LLM also
+        attends to it through KV cache, matching upstream's
+        gen-pre-condition pass).
+        """
+
+        import numpy as _np
+        from PIL import Image as _PILImage
+
+        from vllm_omni.diffusion.data import DiffusionOutput
+
+        from .lance_transformer import NaiveCache
+        from .prompts import SYSTEM_PROMPTS
+
+        first_prompt = req.prompts[0]
+        assert isinstance(first_prompt, dict), "i2v requires dict-style prompt"
+        user_text = first_prompt.get("user_text")
+        rendered = first_prompt.get("prompt") or ""
+        if not user_text:
+            user_text = _extract_user_instruction(rendered)
+        mm_data = first_prompt.get("multi_modal_data") or {}
+        image_input = mm_data.get("first_frame")
+        if image_input is None:
+            raise ValueError("i2v requires multi_modal_data.first_frame.")
+
+        # Resolve image to (H, W, 3) uint8 RGB, then wrap as 1-frame video.
+        if isinstance(image_input, _PILImage.Image):
+            image_raw = _np.array(image_input.convert("RGB"))
+        elif isinstance(image_input, str):
+            image_raw = _np.array(_PILImage.open(image_input).convert("RGB"))
+        elif isinstance(image_input, _np.ndarray):
+            image_raw = image_input if image_input.ndim == 3 else image_input.squeeze()
+        else:
+            raise ValueError(f"Unsupported first_frame type {type(image_input)}")
+        video_raw = image_raw[None]  # (1, H, W, 3)
+
+        extra_args = first_prompt.get("extra_args") or {}
+        sp_extra = getattr(req.sampling_params, "extra_args", {}) or {}
+        extra_args = {**extra_args, **sp_extra}
+
+        # Output video shape — independent of input image dims.
+        num_frames_out = int(extra_args.get("num_frames", 61))
+        out_H = int(extra_args.get("video_height", 480))
+        out_W = int(extra_args.get("video_width", 848))
+        origin_fps = float(extra_args.get("origin_fps", 24.0))
+
+        # Upstream Lance i2v default is cfg_text_scale=4.0.  Note: vllm-omni's
+        # bf16 attention trajectory (vllm-native kernels vs upstream's
+        # HF-native ones) compounds per-layer drift that may DAMPEN prompt-
+        # directed motion for subtle prompts (e.g. micro facial expression).
+        # For such cases, boost cfg_text_scale to 10-15 via
+        # ``extra_args["cfg_text_scale"]`` to amplify prompt influence.
+        # WARNING: too-high cfg (e.g. 15) can cause anatomical distortion
+        # for natural-motion prompts (e.g. animals walking).  Tune per case.
+        cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.0))
+        cfg_img_scale = float(extra_args.get("cfg_img_scale", 1.0))
+        # Upstream Lance i2v default (config/config_factory.py:cfg_interval):
+        # [0.4, 1.0] — CFG active in high-noise phase only.  Always-on CFG
+        # ``(0, 1)`` over-constrains the late-step velocity toward the
+        # conditioned first frame, suppressing subtle motion (e.g. micro
+        # facial expression changes in ex5).
+        cfg_interval = tuple(extra_args.get("cfg_interval", (0.4, 1.0)))
+        cfg_renorm_type = str(extra_args.get("cfg_renorm_type", "global"))
+        cfg_renorm_min = float(extra_args.get("cfg_renorm_min", 0.0))
+        timestep_shift = float(extra_args.get("timestep_shift", LANCE_DEFAULTS.timestep_shift))
+        num_timesteps = int(req.sampling_params.num_inference_steps or LANCE_DEFAULTS.num_timesteps)
+
+        # Bucket-resize + VAE/ViT preprocess the 1-frame "video" (image).
+        from .lance_transformer import LanceBagel as _LanceBagel
+
+        vae_video, vit_pixels, vit_grid_thw, _ = _LanceBagel._lance_video_preprocess(video_raw, origin_fps)
+        T_ref = int(vae_video.shape[1])
+        H_ref = int(vae_video.shape[2])
+        W_ref = int(vae_video.shape[3])
+
+        # Target output shape (independent of ref) — pass through the same
+        # VAE-aligned bucket so the gen latent dims are valid.  We pick H/W
+        # from extra_args and snap to the VAE downsample grid (×8 spatial,
+        # ×4 temporal for Wan2.2 VAE).  num_frames_out controls T.
+        # Wan2.2 latent dims: t_lat = (T - 1) // 4 + 1, h_lat = H // 8, w_lat = W // 8.
+        out_shape = (num_frames_out, out_H, out_W)
+
+        logger.info(
+            "Lance i2v: ref image %dx%d → 1-frame VAE %dx%dx%d, target output %dx%dx%d (T,H,W), user_text=%r",
+            image_raw.shape[0],
+            image_raw.shape[1],
+            T_ref,
+            H_ref,
+            W_ref,
+            num_frames_out,
+            out_H,
+            out_W,
+            user_text,
+        )
+
+        if req.sampling_params.seed is not None:
+            torch.manual_seed(req.sampling_params.seed)
+            if self.device.type == "cuda":
+                torch.cuda.manual_seed(req.sampling_params.seed)
+
+        gen_context = {
+            "kv_lens": [0],
+            "ropes": [0],
+            "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers),
+        }
+        cfg_text_context = {
+            "kv_lens": [0],
+            "ropes": [0],
+            "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers),
+        }
+
+        autocast_kwargs = dict(
+            device_type=self.device.type,
+            enabled=self.device.type != "cpu",
+            dtype=self.od_config.dtype,
+        )
+
+        def _raw_text_prefill(ctx, text_str):
+            text_ids = self.tokenizer.encode(text_str, add_special_tokens=False)
+            if not text_ids:
+                return
+            curr_kvlen = ctx["kv_lens"][0]
+            curr_rope = ctx["ropes"][0]
+            seq_len = len(text_ids)
+            inp = {
+                "text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device),
+                "packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device),
+                "packed_text_position_ids": torch.arange(
+                    curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device
+                ),
+                "packed_text_indexes": torch.arange(
+                    curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device
+                ),
+                "packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device),
+                "key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device),
+            }
+            with torch.autocast(**autocast_kwargs):
+                ctx["past_key_values"] = self.bagel.forward_cache_update_text(ctx["past_key_values"], **inp)
+            ctx["kv_lens"] = [curr_kvlen + seq_len]
+            ctx["ropes"] = [curr_rope + seq_len]
+
+        def _vit_prefill(ctx):
+            inp, new_kvlens, new_rope = self.bagel.prepare_vit_videos(
+                curr_kvlens=ctx["kv_lens"],
+                curr_rope=ctx["ropes"],
+                videos=[video_raw],
+                new_token_ids=self.new_token_ids,
+                precomputed_vit=[(vit_pixels, vit_grid_thw)],
+            )
+            for k, v in inp.items():
+                if torch.is_tensor(v):
+                    inp[k] = v.to(self.device)
+            with torch.autocast(**autocast_kwargs):
+                ctx["past_key_values"] = self.bagel.forward_cache_update_vit(ctx["past_key_values"], **inp)
+            ctx["kv_lens"] = new_kvlens
+            ctx["ropes"] = new_rope
+
+        _ref_padded_latent_cache = {"v": None}
+
+        def _vae_transforms(_):
+            return vae_video
+
+        def _vae_prefill(ctx):
+            inp, new_kvlens, new_rope = self.bagel._lance_native_prepare_vae_images(
+                curr_kvlens=ctx["kv_lens"],
+                curr_rope=ctx["ropes"],
+                images=[vae_video],
+                transforms=_vae_transforms,
+                new_token_ids=self.new_token_ids,
+                is_video=True,
+            )
+            for k, v in inp.items():
+                if torch.is_tensor(v):
+                    inp[k] = v.to(self.device)
+            if _ref_padded_latent_cache["v"] is None:
+                _ref_padded_latent_cache["v"] = self.vae.encode(inp["padded_images"].to(self.device))
+            inp["precomputed_latent"] = _ref_padded_latent_cache["v"]
+            with torch.autocast(**autocast_kwargs):
+                ctx["past_key_values"] = self.bagel.forward_cache_update_vae(self.vae, ctx["past_key_values"], **inp)
+            ctx["kv_lens"] = new_kvlens
+            ctx["ropes"] = new_rope
+
+        # i2v system prompt (matches t2v wording per upstream).
+        sys_prompt = SYSTEM_PROMPTS[("i2v", "video")]
+        seg1_str = f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n"
+        seg5_str = "<|im_end|>\n<|im_start|>assistant\n"
+
+        # Upstream Lance PR #33 i2v does NOT prefill the input image via
+        # ViT/VAE as a "reference".  The image conditioning happens entirely
+        # through the first-frame latent pin (cond positions in x_t set to
+        # VAE-encoded image, never updated through denoise loop).  Prefill
+        # is text-only: system prompt → user instruction → assistant header.
+        # Adding ViT/VAE prefill (which video_edit does) over-constrains the
+        # generation and prevents the prompt-driven motion from developing.
+        _raw_text_prefill(gen_context, seg1_str)
+        if cfg_text_scale > 1.0:
+            _raw_text_prefill(cfg_text_context, seg1_str)
+
+        # Snapshot rope BEFORE adding user text; gen latent will use this
+        # as its starting rope (matches upstream's modality==1≡2 trick).
+        rope_before_vae = gen_context["ropes"][0]
+        cfg_rope_before_vae = cfg_text_context["ropes"][0] if cfg_text_scale > 1.0 else rope_before_vae
+
+        _raw_text_prefill(gen_context, user_text)
+        _raw_text_prefill(gen_context, seg5_str)
+        if cfg_text_scale > 1.0:
+            _raw_text_prefill(cfg_text_context, seg5_str)
+
+        # GEN LATENT — at OUTPUT shape (not ref shape).  Place at rope
+        # starting from rope_before_vae (the modality==1≡2 trick).  The KV
+        # cache still holds the ref VAE at its original (smaller) rope
+        # range; the noise query's own rope just starts at the same
+        # ``rope_before_vae`` and extends through num_vid_tokens.
+        gen_input_lat = self.bagel.prepare_video_latent(
+            curr_kvlens=gen_context["kv_lens"],
+            curr_rope=[rope_before_vae],
+            video_shapes=[out_shape],
+            new_token_ids=self.new_token_ids,
+        )
+        for k, v in gen_input_lat.items():
+            if torch.is_tensor(v):
+                gen_input_lat[k] = v.to(self.device)
+        cfg_text_lat = self.bagel.prepare_video_latent_cfg(
+            curr_kvlens=cfg_text_context["kv_lens"],
+            curr_rope=[cfg_rope_before_vae],
+            video_shapes=[out_shape],
+        )
+        for k, v in cfg_text_lat.items():
+            if torch.is_tensor(v):
+                cfg_text_lat[k] = v.to(self.device)
+        cfg_img_lat = cfg_text_lat
+
+        # Deterministic init-noise regeneration.  ``prepare_video_latent``
+        # builds packed_init_noises via ``torch.randn(...)`` against the
+        # GLOBAL CUDA RNG state, which is contaminated by everything the
+        # prefill phase (text + ViT + VAE) consumed before this point.
+        # vllm-omni and upstream interleave RNG-consuming ops differently
+        # in prefill, so even with ``torch.manual_seed(seed)`` at the top
+        # the two sides land at different RNG positions when randn is called.
+        # Upstream uses ``torch.Generator(device).manual_seed(seed)`` AT the
+        # randn call (modeling/lance/lance.py:1533) so its noise depends
+        # only on the seed. Match that here.
+        if req.sampling_params.seed is not None:
+            noise_gen = torch.Generator(device=self.device).manual_seed(int(req.sampling_params.seed))
+            init_noises_tensor = gen_input_lat["packed_init_noises"]
+            gen_input_lat["packed_init_noises"] = torch.randn(
+                init_noises_tensor.shape,
+                generator=noise_gen,
+                device=self.device,
+                dtype=init_noises_tensor.dtype,
+            )
+
+        # ── First-frame conditioning: VAE-encode the FULL pixel-space
+        # target tensor (matches upstream Lance PR #33's ff2v_sample):
+        #
+        #   vae_tensor = torch.randn([3, num_frames, H, W])
+        #   vae_tensor[:, :4, :, :] = image.repeat(1, 4, 1, 1)
+        #
+        # Then VAE-encode the whole thing and take latent frame 0.  This is
+        # important — Wan2.2 VAE has temporal convolutions, so the first
+        # latent frame's value depends on neighbouring pixel frames (not
+        # just frames 0-3).  Encoding only the 4 image-replicated frames
+        # (which is what we did before) gives a *different* latent than
+        # upstream's full-sequence encode, and that delta produces an
+        # over-anchoring effect at i2v inference time (ex5 portrait case
+        # was visibly frozen because of this).
+        h_lat = int(out_H // self.bagel.latent_downsample)
+        w_lat = int(out_W // self.bagel.latent_downsample)
+        first_slice_n = h_lat * w_lat
+
+        from torchvision.transforms.functional import resize as _tv_resize
+
+        img_chw = torch.from_numpy(image_raw).permute(2, 0, 1).float() / 127.5 - 1.0  # (3, H, W) in [-1, 1]
+        img_chw = _tv_resize(img_chw, [out_H, out_W], antialias=True)
+
+        # Build the full pixel-space target tensor: random noise everywhere,
+        # except the first 4 pixel frames are the input image (replicated).
+        # Wan2.2 VAE collapses every 4 pixel frames → 1 latent frame, so the
+        # first latent slice will represent the conditioning image.
+        full_pixel_tensor = torch.randn((3, num_frames_out, out_H, out_W), dtype=torch.float32, device=self.device)
+        full_pixel_tensor[:, :4, :, :] = img_chw.to(self.device).unsqueeze(1).repeat(1, 4, 1, 1)
+        full_pixel_5d = full_pixel_tensor.unsqueeze(0)  # (1, 3, num_frames_out, H, W)
+
+        with torch.autocast(**autocast_kwargs):
+            image_latent_5d = self.vae.encode(full_pixel_5d)  # (1, 48, T_lat, h_lat, w_lat)
+
+        # Extract first latent frame in (h_lat, w_lat, 48) layout, then
+        # flatten to (h_lat*w_lat, 48) to match packed_init_noises[:first_slice_n].
+        if image_latent_5d.dim() == 5:
+            image_latent = image_latent_5d[0, :, 0, :, :]  # (48, h_lat, w_lat)
+        elif image_latent_5d.dim() == 4:
+            # encode may squeeze T=1 if the caller passed 4-D input — we
+            # passed 5-D so this branch shouldn't fire, but keep defensive.
+            image_latent = image_latent_5d[0]
+        else:
+            raise RuntimeError(f"Unexpected VAE-encoded shape {tuple(image_latent_5d.shape)}")
+        image_latent_packed = image_latent.permute(1, 2, 0).reshape(first_slice_n, -1)
+
+        init_noises = gen_input_lat["packed_init_noises"]
+        image_latent_packed = image_latent_packed.to(device=init_noises.device, dtype=init_noises.dtype)
+        init_noises[:first_slice_n] = image_latent_packed
+        gen_input_lat["packed_init_noises"] = init_noises
+
+        frame_condition_token_indexes = torch.arange(first_slice_n, dtype=torch.long, device=self.device)
+
+        logger.info(
+            "Lance i2v cond: latent grid h_lat=%d w_lat=%d -> pinning first %d tokens (frame 0) at VAE-encoded image",
+            h_lat,
+            w_lat,
+            first_slice_n,
+        )
+
+        with torch.autocast(**autocast_kwargs):
+            latents, *_ = self.bagel.generate_image(
+                past_key_values=gen_context["past_key_values"],
+                cfg_text_past_key_values=cfg_text_context["past_key_values"]
+                if cfg_text_scale > 1.0
+                else gen_context["past_key_values"],
+                cfg_img_past_key_values=gen_context["past_key_values"],
+                num_timesteps=num_timesteps,
+                timestep_shift=timestep_shift,
+                cfg_text_scale=cfg_text_scale,
+                cfg_img_scale=cfg_img_scale,
+                cfg_interval=cfg_interval,
+                cfg_renorm_type=cfg_renorm_type,
+                cfg_renorm_min=cfg_renorm_min,
+                **gen_input_lat,
+                cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"],
+                cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"],
+                cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"],
+                cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"],
+                cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"],
+                cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"],
+                cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"],
+                cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"],
+                frame_condition_token_indexes=frame_condition_token_indexes,
+            )
+
+        frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], out_shape)
+        frames = [Image.fromarray(f) for f in frames_np]
+        return DiffusionOutput(
+            output=frames[0],
+            custom_output={"video_frames": frames, "video_shape": out_shape},
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
     def _forward_video_edit(self, req):
         """Video edit: reference video + text prompt → modified video.
 
@@ -1335,10 +1712,10 @@ class LancePipeline(BagelPipeline):
     def _forward_x2t_video(self, req):
         """Video understanding (x2t_video): video → text caption / VQA answer.
 
-        Mirrors :meth:`_forward_t2v` for the prefill / generation
-        bookkeeping but routes the video through the Qwen2.5-VL ViT (no VAE
-        prefill — Lance's training has Wan2.2 latents only for image_edit /
-        video_edit, not for understanding).
+        Builds the prefill in three segments matching upstream's chat
+        template — see :meth:`_forward_x2t_image` for the rationale on
+        why naïve concatenation (text prompt → ViT) lands the EOS / vit
+        markers in the wrong place and degrades the answer format.
         """
 
         from vllm_omni.diffusion.data import DiffusionOutput
@@ -1347,7 +1724,6 @@ class LancePipeline(BagelPipeline):
 
         first_prompt = req.prompts[0]
         assert isinstance(first_prompt, dict), "x2t_video requires dict-style prompt"
-        prompt = first_prompt.get("prompt") or ""
         video_input = (first_prompt.get("multi_modal_data") or {}).get("video")
         if video_input is None:
             raise ValueError("x2t_video requires multi_modal_data.video to be a video tensor/array/path.")
@@ -1366,40 +1742,32 @@ class LancePipeline(BagelPipeline):
         do_sample = bool(extra_args.get("do_sample", False))
         text_temperature = float(extra_args.get("text_temperature", 0.3))
 
+        prefix_text, suffix_text = self._split_x2t_prompt(
+            first_prompt.get("prompt") or "",
+            user_text_override=extra_args.get("user_text"),
+            system_prompt_override=extra_args.get("system_prompt"),
+            default_system_prompt="Watch the video carefully and answer the question.",
+        )
+
         if req.sampling_params.seed is not None:
             torch.manual_seed(req.sampling_params.seed)
             if self.device.type == "cuda":
                 torch.cuda.manual_seed(req.sampling_params.seed)
 
-        # ---- Text-prompt prefill ----
+        autocast_kwargs = dict(
+            device_type=self.device.type,
+            enabled=self.device.type != "cpu",
+            dtype=self.od_config.dtype,
+        )
+
         gen_context = {
             "kv_lens": [0],
             "ropes": [0],
             "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers),
         }
-        if prompt:
-            gen_input_text, newlens, new_rope = self.bagel.prepare_prompts(
-                curr_kvlens=gen_context["kv_lens"],
-                curr_rope=gen_context["ropes"],
-                prompts=[prompt],
-                tokenizer=self.tokenizer,
-                new_token_ids=self.new_token_ids,
-            )
-            for k, v in gen_input_text.items():
-                if torch.is_tensor(v):
-                    gen_input_text[k] = v.to(self.device)
-            with torch.autocast(
-                device_type=self.device.type,
-                enabled=self.device.type != "cpu",
-                dtype=self.od_config.dtype,
-            ):
-                gen_context["past_key_values"] = self.bagel.forward_cache_update_text(
-                    gen_context["past_key_values"], **gen_input_text
-                )
-            gen_context["kv_lens"] = newlens
-            gen_context["ropes"] = new_rope
 
-        # ---- Multi-frame ViT prefill ----
+        self._x2t_raw_text_prefill(gen_context, prefix_text, autocast_kwargs)
+
         gen_input_vit, newlens_vit, new_rope_vit = self.bagel.prepare_vit_videos(
             curr_kvlens=gen_context["kv_lens"],
             curr_rope=gen_context["ropes"],
@@ -1409,27 +1777,20 @@ class LancePipeline(BagelPipeline):
         for k, v in gen_input_vit.items():
             if torch.is_tensor(v):
                 gen_input_vit[k] = v.to(self.device)
-        with torch.autocast(
-            device_type=self.device.type,
-            enabled=self.device.type != "cpu",
-            dtype=self.od_config.dtype,
-        ):
+        with torch.autocast(**autocast_kwargs):
             gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
                 gen_context["past_key_values"], **gen_input_vit
             )
         gen_context["kv_lens"] = newlens_vit
         gen_context["ropes"] = new_rope_vit
 
-        # ---- Text generation ----
-        start_input = self.bagel.prepare_start_tokens(gen_context["kv_lens"], gen_context["ropes"], self.new_token_ids)
+        self._x2t_raw_text_prefill(gen_context, suffix_text, autocast_kwargs)
+
+        start_input = self._x2t_prepare_assistant_start(gen_context)
         for k, v in start_input.items():
             if torch.is_tensor(v):
                 start_input[k] = v.to(self.device)
-        with torch.autocast(
-            device_type=self.device.type,
-            enabled=self.device.type != "cpu",
-            dtype=self.od_config.dtype,
-        ):
+        with torch.autocast(**autocast_kwargs):
             token_ids = self.bagel.generate_text(
                 past_key_values=gen_context["past_key_values"],
                 max_length=max_text_tokens,
@@ -1442,6 +1803,9 @@ class LancePipeline(BagelPipeline):
         text_output = decoded.split("<|im_end|>")[0]
         if "<|im_start|>" in text_output:
             text_output = text_output.split("<|im_start|>")[-1]
+        # Strip the leading newline that was fed as the start token (see
+        # ``_x2t_prepare_assistant_start``); see x2t_image for rationale.
+        text_output = text_output.lstrip("\n")
         logger.info("Lance x2t_video: generated %d tokens", token_ids.shape[0])
         return DiffusionOutput(
             output=text_output,
@@ -1449,6 +1813,231 @@ class LancePipeline(BagelPipeline):
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
+    @torch.inference_mode()
+    def _forward_x2t_image(self, req):
+        """Image understanding (x2t_image): image + question → text answer.
+
+        Builds the prefill in three text/vit segments matching upstream's
+        chat-template layout::
+
+            <|im_start|>system\\n{system}<|im_end|>\\n<|im_start|>user\\n
+            <|vision_start|>{N×ViT}<|vision_end|>
+            {question}<|im_end|>\\n<|im_start|>assistant\\n
+
+        ``prepare_prompts`` wraps the *whole* string with ``[bos]…[eos]``
+        and ``prepare_vit_images`` adds ``<|vision_start|>``/``<|vision_end|>``
+        of its own — concatenating them naïvely (text first, then vit)
+        produced ``…assistant\\n[eos]<|vision_start|>{ViT}<|vision_end|>``,
+        so the model saw an EOS right after the assistant tag and a
+        bare ``<|video_pad|>`` token inside the user message.  Segmenting
+        the text into prefix/suffix and sliding the vit prefill between
+        them puts the markers exactly where Lance was trained on.
+        """
+
+        from PIL import Image as _PIL_Image
+
+        from vllm_omni.diffusion.data import DiffusionOutput
+
+        from .lance_transformer import NaiveCache
+
+        first_prompt = req.prompts[0]
+        assert isinstance(first_prompt, dict), "x2t_image requires dict-style prompt"
+        image_input = (first_prompt.get("multi_modal_data") or {}).get("image")
+        if image_input is None:
+            raise ValueError("x2t_image requires multi_modal_data.image to be a PIL image / array / path.")
+        if isinstance(image_input, str):
+            image_input = _PIL_Image.open(image_input).convert("RGB")
+        images = [image_input]
+
+        extra_args = first_prompt.get("extra_args") or {}
+        sp_extra = getattr(req.sampling_params, "extra_args", {}) or {}
+        extra_args = {**extra_args, **sp_extra}
+        max_text_tokens = int(extra_args.get("max_think_tokens", 256))
+        do_sample = bool(extra_args.get("do_sample", False))
+        text_temperature = float(extra_args.get("text_temperature", 0.3))
+
+        prefix_text, suffix_text = self._split_x2t_prompt(
+            first_prompt.get("prompt") or "",
+            user_text_override=extra_args.get("user_text"),
+            system_prompt_override=extra_args.get("system_prompt"),
+            default_system_prompt="Look at the image carefully and answer the question.",
+        )
+
+        if req.sampling_params.seed is not None:
+            torch.manual_seed(req.sampling_params.seed)
+            if self.device.type == "cuda":
+                torch.cuda.manual_seed(req.sampling_params.seed)
+
+        def vit_transforms(img):
+            return self.image_processor(images=img, return_tensors="pt").pixel_values[0]
+
+        autocast_kwargs = dict(
+            device_type=self.device.type,
+            enabled=self.device.type != "cpu",
+            dtype=self.od_config.dtype,
+        )
+
+        gen_context = {
+            "kv_lens": [0],
+            "ropes": [0],
+            "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers),
+        }
+
+        self._x2t_raw_text_prefill(gen_context, prefix_text, autocast_kwargs)
+
+        gen_input_vit, newlens_vit, new_rope_vit = self.bagel.prepare_vit_images(
+            curr_kvlens=gen_context["kv_lens"],
+            curr_rope=gen_context["ropes"],
+            images=images,
+            transforms=vit_transforms,
+            new_token_ids=self.new_token_ids,
+        )
+        for k, v in gen_input_vit.items():
+            if torch.is_tensor(v):
+                gen_input_vit[k] = v.to(self.device)
+        with torch.autocast(**autocast_kwargs):
+            gen_context["past_key_values"] = self.bagel.forward_cache_update_vit(
+                gen_context["past_key_values"], **gen_input_vit
+            )
+        gen_context["kv_lens"] = newlens_vit
+        gen_context["ropes"] = new_rope_vit
+
+        self._x2t_raw_text_prefill(gen_context, suffix_text, autocast_kwargs)
+
+        start_input = self._x2t_prepare_assistant_start(gen_context)
+        for k, v in start_input.items():
+            if torch.is_tensor(v):
+                start_input[k] = v.to(self.device)
+        with torch.autocast(**autocast_kwargs):
+            token_ids = self.bagel.generate_text(
+                past_key_values=gen_context["past_key_values"],
+                max_length=max_text_tokens,
+                do_sample=do_sample,
+                temperature=text_temperature,
+                end_token_id=self.new_token_ids["eos_token_id"],
+                **start_input,
+            )
+        decoded = self.tokenizer.decode(token_ids[:, 0].tolist())
+        text_output = decoded.split("<|im_end|>")[0]
+        if "<|im_start|>" in text_output:
+            text_output = text_output.split("<|im_start|>")[-1]
+        # Strip the leading newline that was fed as the start token (see
+        # ``_x2t_prepare_assistant_start``); it always sits at index 0 of
+        # the generated sequence and is just plumbing, not model output.
+        text_output = text_output.lstrip("\n")
+        logger.info("Lance x2t_image: generated %d tokens", token_ids.shape[0])
+        return DiffusionOutput(
+            output=text_output,
+            custom_output={"text_output": text_output},
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    @staticmethod
+    def _split_x2t_prompt(
+        rendered_prompt: str,
+        *,
+        user_text_override: str | None,
+        system_prompt_override: str | None,
+        default_system_prompt: str,
+    ) -> tuple[str, str]:
+        """Split a rendered Lance x2t chat template at the vision block.
+
+        Accepts either a pre-rendered template containing
+        ``<|vision_start|>…<|vision_end|>`` (the historical
+        ``render_lance_prompt`` output) or recovers prefix/suffix from
+        ``user_text_override`` + ``system_prompt_override`` when the
+        template was not pre-rendered.  Returns ``(prefix, suffix)`` such
+        that the segmented prefill is ``prefix + ViT + suffix``.
+
+        ``suffix`` deliberately ends at ``assistant`` (NOT ``assistant\\n``):
+        :meth:`Bagel.prepare_start_tokens` injects ``<|im_start|>`` (Qwen
+        bos == 151644) as the first generation step, and ``\\n<|im_start|>``
+        right after ``assistant\\n`` tells the model "begin a new
+        message" — it answers with an immediate ``<|im_end|>`` and
+        produces an empty caption.  Holding the trailing ``\\n`` out of
+        the prefill (it gets re-applied via ``_x2t_prepare_assistant_start``)
+        lets the next forward step start *inside* the assistant turn
+        instead of opening a fresh one.
+        """
+        sys_prompt = system_prompt_override or default_system_prompt
+
+        if user_text_override is not None:
+            user_text = user_text_override
+        elif "<|vision_start|>" in rendered_prompt and "<|vision_end|>" in rendered_prompt:
+            _, _, after_vis_end = rendered_prompt.partition("<|vision_end|>")
+            user_text, _, _ = after_vis_end.partition("<|im_end|>")
+        else:
+            # No vision markers and no override — best-effort fallback:
+            # treat the entire prompt as the user question and rebuild
+            # the chat scaffolding.
+            user_text = rendered_prompt
+
+        prefix = f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n"
+        suffix = f"{user_text}<|im_end|>\n<|im_start|>assistant"
+        return prefix, suffix
+
+    def _x2t_prepare_assistant_start(self, gen_context):
+        """Build a ``generate_text`` start input where the first fed token
+        is the newline that closes ``<|im_start|>assistant\\n``.
+
+        :meth:`Bagel.prepare_start_tokens` would instead feed
+        ``<|im_start|>``, opening a fresh message right after the
+        assistant tag and triggering an immediate EOS on Lance_3B (the
+        image-only checkpoint is stricter about chat-turn framing than
+        the video model, which happens to recover).  Feeding ``\\n``
+        keeps the model *inside* the assistant turn so the next argmax
+        is the first answer token.
+        """
+        newline_ids = self.tokenizer.encode("\n", add_special_tokens=False)
+        assert len(newline_ids) == 1, f"expected single token for '\\n', got {newline_ids!r}"
+        nl_id = newline_ids[0]
+
+        curr_kvlens = gen_context["kv_lens"]
+        curr_ropes = gen_context["ropes"]
+        packed_key_value_indexes: list[int] = []
+        packed_query_position_ids: list[int] = []
+        packed_start_tokens: list[int] = []
+        curr = 0
+        for curr_kvlen, curr_position_id in zip(curr_kvlens, curr_ropes):
+            packed_key_value_indexes.extend(range(curr, curr + curr_kvlen))
+            packed_start_tokens.append(nl_id)
+            packed_query_position_ids.append(curr_position_id)
+            curr += curr_kvlen
+        return {
+            "packed_start_tokens": torch.tensor(packed_start_tokens, dtype=torch.long),
+            "packed_query_position_ids": torch.tensor(packed_query_position_ids, dtype=torch.long),
+            "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int),
+            "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long),
+        }
+
+    def _x2t_raw_text_prefill(self, gen_context, text_str, autocast_kwargs):
+        """Tokenize ``text_str`` *raw* (no bos/eos wrapping) and append it
+        to ``gen_context``'s KV cache.  ``prepare_prompts`` is unusable
+        here because it would inject ``[bos] + ids + [eos]`` around each
+        segment, breaking the chat template framing."""
+        if not text_str:
+            return
+        text_ids = self.tokenizer.encode(text_str, add_special_tokens=False)
+        if not text_ids:
+            return
+        curr_kvlen = gen_context["kv_lens"][0]
+        curr_rope = gen_context["ropes"][0]
+        seq_len = len(text_ids)
+        inp = {
+            "text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device),
+            "packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device),
+            "packed_text_position_ids": torch.arange(
+                curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device
+            ),
+            "packed_text_indexes": torch.arange(curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device),
+            "packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device),
+            "key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device),
+        }
+        with torch.autocast(**autocast_kwargs):
+            gen_context["past_key_values"] = self.bagel.forward_cache_update_text(gen_context["past_key_values"], **inp)
+        gen_context["kv_lens"] = [curr_kvlen + seq_len]
+        gen_context["ropes"] = [curr_rope + seq_len]
+
     @staticmethod
     def _decode_video_from_latent(
         bagel,

vllm_omni/diffusion/models/lance/prompts.py

@@ -32,6 +32,14 @@ SYSTEM_PROMPTS: dict[tuple[str, str], str] = {
         "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 "
@@ -68,6 +76,7 @@ def render_lance_prompt(
     user_text: str,
     *,
     vision_token: str | None = None,
+    system_prompt: str | None = None,
 ) -> str:
     """Render a Lance-compatible user-side prompt.
 
@@ -75,15 +84,21 @@ def render_lance_prompt(
     ``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.
+    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", "x2t_video", "video_edit"} else "image"
-    sys_prompt = SYSTEM_PROMPTS.get((task, vision_type), "You are a helpful assistant.")
+    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

vllm_omni/diffusion/models/lance/wan_vae.py

@@ -833,13 +833,19 @@ class LanceWanVAE(nn.Module):
         """
         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)
-        if latent.shape[2] == 1:
+        # 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
 

vllm_omni/model_executor/models/lance/__init__.py

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

vllm_omni/model_executor/models/lance/lance.py

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

vllm_omni/model_executor/models/lance/pipeline.py

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

vllm_omni/model_executor/models/registry.py

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

vllm_omni/model_executor/stage_input_processors/lance.py

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