From 7c934b61888a487ef26e9904c776de8dab34709e Mon Sep 17 00:00:00 2001 From: lishunyang12 Date: Sun, 31 May 2026 11:40:40 +0000 Subject: [PATCH] [Lance] Add Gradio demo and align inference with upstream MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 Signed-off-by: WeiQing Chen --- examples/offline_inference/lance/end2end.py | 80 +- .../offline_inference/lance/gradio_demo.py | 1280 +++++++++++++++++ vllm_omni/config/pipeline_registry.py | 8 +- .../models/bagel/bagel_transformer.py | 49 +- .../diffusion/models/lance/pipeline_lance.py | 671 ++++++++- vllm_omni/diffusion/models/lance/prompts.py | 21 +- vllm_omni/diffusion/models/lance/wan_vae.py | 10 +- .../model_executor/models/lance/__init__.py | 20 +- .../model_executor/models/lance/lance.py | 33 - .../model_executor/models/lance/pipeline.py | 57 +- vllm_omni/model_executor/models/registry.py | 6 - .../stage_input_processors/lance.py | 32 - 12 files changed, 2059 insertions(+), 208 deletions(-) create mode 100644 examples/offline_inference/lance/gradio_demo.py delete mode 100644 vllm_omni/model_executor/models/lance/lance.py delete mode 100644 vllm_omni/model_executor/stage_input_processors/lance.py diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py index 344699fc8..ec7985304 100644 --- a/examples/offline_inference/lance/end2end.py +++ b/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") diff --git a/examples/offline_inference/lance/gradio_demo.py b/examples/offline_inference/lance/gradio_demo.py new file mode 100644 index 000000000..3d6b240d3 --- /dev/null +++ b/examples/offline_inference/lance/gradio_demo.py @@ -0,0 +1,1280 @@ +"""Unified Gradio demo for Lance × vllm-omni — all 7 tasks. + +Style mirrors upstream Lance's lance_gradio.py: + * header with Lance + vllm-omni logos + shield badges + * radio task selector (Video Generation / Video Edit / Video Understanding / + Image-to-Video / Image Generation / Image Edit / Image Understanding) + * aspect-ratio + resolution dropdowns (not free H/W sliders) + * two-column layout (inputs left, output right) + * advanced-params accordion at the bottom + +Backend: routes to one of two Omni instances based on task: + * Lance_3B (t2i, image_edit, x2t_image) + * Lance_3B_Video (t2v, i2v, video_edit, x2t_video) + +Run: + CUDA_VISIBLE_DEVICES=6,7 \\ + /home/zjy/code/vllm-src/.venv/bin/python \\ + examples/offline_inference/lance/gradio_demo.py \\ + --model /path/to/Lance/snapshots/ +""" + +from __future__ import annotations + +import argparse +import base64 +import os +import tempfile +from pathlib import Path + +import gradio as gr +import imageio.v3 as iio +import numpy as np +from PIL import Image + +# --------------------------------------------------------- constants ----- + +# Mirror upstream Lance's resolution/aspect tables. +# +# Per ``data/datasets_custom/validation_dataset.py``, upstream restricts +# aspect-ratio buckets to **6 fixed ratios** and computes the per-bucket +# canonical (W, H) via: +# +# max_area = resolution_vae ** 2 +# for each ar = w/h: +# w1 = round(sqrt(max_area * ar) / 16) * 16 +# h1 = round(w1 / ar / 16) * 16 +# (… pick the candidate whose actual ratio is closer to ``ar``) +# +# with ``resolution_vae = 480`` for ``video_360p``, ``640`` for +# ``video_480p``, ``768`` for ``image_768res``. The numbers below are +# computed exactly that way so this gradio demo lands on the same canvas +# upstream would for the same input image — no aspect-ratio drift. +ASPECT_RATIO_CHOICES = ["21:9", "16:9", "4:3", "1:1", "3:4", "9:16"] +DEFAULT_VIDEO_AR = "16:9" +DEFAULT_IMAGE_AR = "1:1" + +VIDEO_360P_AR_TO_SIZE = { + "21:9": (752, 320), + "16:9": (624, 352), + "4:3": (560, 416), + "1:1": (480, 480), + "3:4": (416, 560), + "9:16": (352, 624), +} +VIDEO_480P_AR_TO_SIZE = { + "21:9": (976, 416), + "16:9": (848, 480), + "4:3": (752, 560), + "1:1": (640, 640), + "3:4": (560, 752), + "9:16": (480, 848), +} +IMAGE_AR_TO_SIZE = { + "21:9": (1152, 496), + "16:9": (1024, 576), + "4:3": (896, 672), + "1:1": (768, 768), + "3:4": (672, 896), + "9:16": (576, 1024), +} + +VIDEO_RES_CHOICES = [("360p", "video_360p"), ("480p", "video_480p")] +VIDEO_RES_MAP = {"video_360p": VIDEO_360P_AR_TO_SIZE, "video_480p": VIDEO_480P_AR_TO_SIZE} +DEFAULT_VIDEO_RES = "video_480p" +DEFAULT_TIMESTEPS = 30 # upstream: --validation_num_timesteps 30 +DEFAULT_TIMESTEP_SHIFT = 3.5 # upstream: --validation_timestep_shift 3.5 +DEFAULT_CFG = 4.0 # upstream: --cfg_text_scale 4.0 (NOT 8.0) +DEFAULT_SEED = 42 # upstream: --validation_data_seed 42 +DEFAULT_VIDEO_DURATION_SEC = 5 # 5s × 12 fps + 1 = 61 frames — snappy for interactive demo +MAX_VIDEO_DURATION_SEC = 10 +# Per-task duration defaults — quality knobs (cfg, steps, ts_shift) are +# locked to upstream values, but the **demo duration** is intentionally +# shorter than upstream's per-script ``num_frames`` to keep gradio +# responsive (upstream uses 121 frames for t2v/i2v = 10s ≈ 2 min/req on +# B300; demo defaults give the same per-frame fidelity in ~60 s). +# Users can extend up to 10s with the slider when they want +# byte-identical reproduction of upstream's official output. +DURATION_DEFAULT_PER_TASK: dict[str, int] = { + "t2v": 5, # ~61 frames; upstream uses 121 + "i2v": 5, # ~61 frames; upstream uses 121 + "video_edit": 4, # ~49 frames; matches upstream's 50 +} + +# Task radio choices (label, internal_key) +TASK_VIDEO_GEN = "Video Generation" +TASK_I2V_LABEL = "Image to Video" +TASK_VIDEO_EDIT = "Video Edit" +TASK_VIDEO_UND = "Video Understanding" +TASK_IMAGE_GEN = "Image Generation" +TASK_IMAGE_EDIT = "Image Edit" +TASK_IMAGE_UND = "Image Understanding" + +LABEL_TO_TASK = { + TASK_VIDEO_GEN: "t2v", + TASK_I2V_LABEL: "i2v", + TASK_VIDEO_EDIT: "video_edit", + TASK_VIDEO_UND: "x2t_video", + TASK_IMAGE_GEN: "t2i", + TASK_IMAGE_EDIT: "image_edit", + TASK_IMAGE_UND: "x2t_image", +} + +LANCE_HOMEPAGE_URL = "https://lance-project.github.io/" +LANCE_PAPER_URL = "http://arxiv.org/abs/2605.18678" +LANCE_HF_URL = "https://huggingface.co/bytedance-research/Lance" +LANCE_GITHUB_URL = "https://github.com/bytedance/Lance" +VLLM_OMNI_GITHUB_URL = "https://github.com/vllm-project/vllm-omni" + + +# --------------------------------------------------- module handles ----- + +_OMNI_IMAGE = None +_OMNI_VIDEO = None +_PARAMS_IMAGE = None +_PARAMS_VIDEO = None + + +# ----------------------------------------------------------- args ----- + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", required=True, help="Parent snapshot dir containing Lance_3B and Lance_3B_Video subdirs.") + p.add_argument("--deploy-config", default="vllm_omni/deploy/lance.yaml", help="Lance deploy YAML.") + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=7860) + p.add_argument("--share", action="store_true") + p.add_argument( + "--examples-root", + default="/home/zjy/code/lance-upstream", + help=( + "Path to a Lance upstream checkout containing ``config/examples/*.json`` " + "and the referenced media assets. When present, each task tab gets a " + "clickable list of official examples that auto-fill the prompt + media. " + "Set to '' to disable." + ), + ) + p.add_argument( + "--ulysses-degree", + type=int, + default=1, + help=( + "Ulysses sequence-parallel degree for *each* Omni instance " + "(image and video). Needs ``2 * ulysses_degree`` GPUs exposed via " + "CUDA_VISIBLE_DEVICES. NOTE: Lance currently lacks an ``_sp_plan`` " + "registry and the i2v first-frame pin is not SP-aware, so SP >1 is " + "experimental — t2i/t2v/edit/x2t work but i2v crashes. Prefer " + "``--replicas-per-omni`` for throughput scaling instead." + ), + ) + p.add_argument( + "--replicas-per-omni", + type=int, + default=1, + help=( + "Spawn ``N`` parallel replicas of each Omni (image + video) across " + "``2*N`` GPUs. Default ``1`` = the P0 layout (2 GPUs total, image " + "on GPU 0, video on GPU 1). Set ``4`` to use all 8 GPUs for " + "throughput; each replica loads a fresh model copy so startup time " + "and memory scale linearly." + ), + ) + return p.parse_args() + + +# -------------------------------------------------- model loading ----- + + +def _resolve_ckpts(snapshot: str): + img = Path(snapshot) / "Lance_3B" + vid = Path(snapshot) / "Lance_3B_Video" + if not img.is_dir(): + raise FileNotFoundError(f"Lance_3B not found under {snapshot}") + if not vid.is_dir(): + raise FileNotFoundError(f"Lance_3B_Video not found under {snapshot}") + return str(img), str(vid) + + +def _init_models( + snapshot: str, + deploy_config: str, + ulysses_degree: int = 1, + replicas_per_omni: int = 1, +): + """Load the two Lance checkpoints onto distinct GPU sets so image + and video tasks can run truly in parallel, optionally with Ulysses + sequence parallel inside each Omni. + + Device layout (with ``CUDA_VISIBLE_DEVICES=$DEVICES``): + + - ``ulysses_degree=1`` (P0): image Omni → logical 0, video Omni → + logical 1. Concurrent throughput = 2x serial; needs 2 GPUs total. + - ``ulysses_degree=2`` (P1): image Omni → logical ``0,1``, video + Omni → logical ``2,3``. Each Omni internally splits the DiT + denoising sequence across 2 GPUs (Ulysses SP). Total = 4 GPUs; + single-request t2v/i2v ~1.5-1.8x faster *and* image+video still + concurrent. + - ``ulysses_degree=N``: image → ``0..N-1``, video → ``N..2N-1``. + Needs ``2*N`` GPUs. + + Falls back gracefully when there aren't enough GPUs: prints a + warning and downgrades to the largest feasible degree. + """ + global _OMNI_IMAGE, _OMNI_VIDEO, _PARAMS_IMAGE, _PARAMS_VIDEO + from vllm_omni.entrypoints.omni import Omni + + img_ckpt, vid_ckpt = _resolve_ckpts(snapshot) + + import os + + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + visible = [d for d in cvd.split(",") if d.strip()] if cvd else [] + n_visible = len(visible) + + # Resolve ulysses & replicas against available device count. + if ulysses_degree < 1: + ulysses_degree = 1 + if replicas_per_omni < 1: + replicas_per_omni = 1 + # Per-Omni GPU count = ulysses_degree × replicas_per_omni + # (ulysses splits a single replica across GPUs; replicas multiply by GPU). + per_omni_gpus = ulysses_degree * replicas_per_omni + needed = 2 * per_omni_gpus + if needed > max(n_visible, 1): + # Downgrade replicas first (cheaper than losing SP), then ulysses. + feasible_replicas = max(1, n_visible // (2 * ulysses_degree)) + if feasible_replicas < replicas_per_omni: + print( + f"[unified] WARN: --replicas-per-omni={replicas_per_omni} × " + f"ulysses_degree={ulysses_degree} needs {needed} GPUs but only " + f"{n_visible} are visible. Falling back to replicas={feasible_replicas}.", + flush=True, + ) + replicas_per_omni = feasible_replicas + per_omni_gpus = ulysses_degree * replicas_per_omni + if 2 * per_omni_gpus > n_visible: + feasible_sp = max(1, n_visible // 2) + if ulysses_degree > feasible_sp: + print( + f"[unified] WARN: --ulysses-degree={ulysses_degree} needs " + f"{2 * ulysses_degree} GPUs. Falling back to ulysses_degree={feasible_sp}.", + flush=True, + ) + ulysses_degree = feasible_sp + per_omni_gpus = ulysses_degree * replicas_per_omni + + can_split = n_visible >= 2 + use_sp = ulysses_degree > 1 + use_replicas = replicas_per_omni > 1 + + print( + f"[unified] CUDA_VISIBLE_DEVICES={cvd!r} n_visible={n_visible} " + f"can_split={can_split} ulysses_degree={ulysses_degree} " + f"replicas_per_omni={replicas_per_omni} per_omni_gpus={per_omni_gpus}", + flush=True, + ) + + # Build per-Omni device strings. Image takes the lower half, video + # the upper half. + if per_omni_gpus > 1: + img_devices = ",".join(str(i) for i in range(per_omni_gpus)) + vid_devices = ",".join(str(i) for i in range(per_omni_gpus, 2 * per_omni_gpus)) + elif can_split: + img_devices, vid_devices = "0", "1" + else: + img_devices, vid_devices = "0", "0" + + print( + f"[unified] Loading Lance_3B (image weights) on logical devices {img_devices!r}" + f" replicas={replicas_per_omni}...", + flush=True, + ) + img_kwargs: dict = {"model": img_ckpt, "deploy_config": deploy_config} + if img_devices != "0": + img_kwargs["stage_0_devices"] = img_devices + if use_sp: + # Per the post-#3483 API, parallel knobs flow through *flat* + # ``ulysses_degree`` / ``ring_degree`` kwargs on the Omni + # constructor; ``_apply_diffusion_parallel_runtime_overrides`` + # moves them into a nested ``parallel_config`` dict for the stage. + img_kwargs["ulysses_degree"] = ulysses_degree + if use_replicas: + img_kwargs["stage_0_num_replicas"] = replicas_per_omni + _OMNI_IMAGE = Omni(**img_kwargs) + _PARAMS_IMAGE = _OMNI_IMAGE.default_sampling_params_list + + print( + f"[unified] Loading Lance_3B_Video (video weights) on logical devices {vid_devices!r}" + f" replicas={replicas_per_omni}...", + flush=True, + ) + vid_kwargs: dict = {"model": vid_ckpt, "deploy_config": deploy_config} + vid_kwargs["stage_0_devices"] = vid_devices + if use_sp: + vid_kwargs["ulysses_degree"] = ulysses_degree + if use_replicas: + vid_kwargs["stage_0_num_replicas"] = replicas_per_omni + _OMNI_VIDEO = Omni(**vid_kwargs) + _PARAMS_VIDEO = _OMNI_VIDEO.default_sampling_params_list + print("[unified] Both models ready.", flush=True) + + +def _omni_for(task: str): + if task in {"t2v", "i2v", "video_edit", "x2t_video"}: + return _OMNI_VIDEO, _PARAMS_VIDEO + return _OMNI_IMAGE, _PARAMS_IMAGE + + +# ------------------------------------------ size resolution helpers ----- + +_AR_RATIOS: dict[str, float] = { + # Mirror upstream's 6 buckets exactly — *no* 3:2 / 2:3, those are + # snapped to 4:3 / 3:4 by upstream's bucket selector. Must match + # ``ASPECT_RATIO_CHOICES`` or gradio will reject the auto-picked + # value as "not in the list of choices". + "21:9": 21 / 9, + "16:9": 16 / 9, + "4:3": 4 / 3, + "1:1": 1.0, + "3:4": 3 / 4, + "9:16": 9 / 16, +} + + +def _closest_aspect_ratio(width: int, height: int) -> str: + """Snap an arbitrary ``W×H`` to the nearest preset in + ``ASPECT_RATIO_CHOICES`` (used to auto-select the dropdown when the + user uploads / clicks an i2v example image). Distance is computed + in log-space so 4:3 isn't unfairly favored over 1:1 just because of + linear-scale proximity.""" + import math + + if not width or not height: + return DEFAULT_VIDEO_AR + img_log = math.log(width / height) + best = min(_AR_RATIOS.items(), key=lambda kv: abs(math.log(kv[1]) - img_log)) + return best[0] + + +def get_size_for_task(task: str, ar: str, resolution: str): + """Return (width, height) for a task + aspect-ratio + resolution preset.""" + if task in {"t2v", "i2v", "video_edit"}: + m = VIDEO_RES_MAP.get(resolution, VIDEO_480P_AR_TO_SIZE) + return m.get(ar, m["16:9"]) + elif task in {"t2i", "image_edit"}: + return IMAGE_AR_TO_SIZE.get(ar, IMAGE_AR_TO_SIZE["1:1"]) + return (640, 352) + + +def _format_size_md(task: str, w: int, h: int) -> str: + unit = "frames" if task in {"t2v", "i2v", "video_edit"} else "px" + return f"{w} × {h} {unit}" + + +# ------------------------------------------ per-task prompt formats ----- + + +def _vision_block(kind: str = "video"): + from vllm_omni.diffusion.models.lance.prompts import ( + IMAGE_PAD, + VIDEO_PAD, + VISION_END, + VISION_START, + ) + + pad = VIDEO_PAD if kind == "video" else IMAGE_PAD + return f"{VISION_START}{pad}{VISION_END}" + + +def _format(task, prompt, image, video, num_frames, height, width, cfg, ts_shift): + from vllm_omni.diffusion.models.lance.prompts import render_lance_prompt + + if task == "t2i": + return [ + { + "prompt": render_lance_prompt("t2i", prompt), + "modalities": ["image"], + "extra_args": {"height": int(height), "width": int(width), "cfg_text_scale": float(cfg)}, + } + ] + if task == "image_edit": + return [ + { + "prompt": render_lance_prompt("image_edit", prompt, vision_token=_vision_block("image")), + "multi_modal_data": {"img2img": image}, + "modalities": ["image"], + "extra_args": {"cfg_text_scale": float(cfg)}, + } + ] + if task == "x2t_image": + q = (prompt or "").strip() or "Describe this image in detail." + # Upstream's official x2t_image examples carry a QA-style + # per-example system prompt ("Look at the image carefully and + # answer the question."). Without it, Lance's caption-default + # system prompt makes the model describe the image instead of + # answering — matching that here so gradio out-of-the-box agrees + # with the upstream CLI on VQA cases. + return [ + { + "prompt": render_lance_prompt( + "x2t_image", + q, + vision_token=_vision_block("image"), + system_prompt="Look at the image carefully and answer the question.", + ), + "multi_modal_data": {"image": image}, + "modalities": ["text"], + "extra_args": {"max_think_tokens": 500}, + } + ] + if task == "t2v": + return [ + { + "prompt": render_lance_prompt("t2v", prompt), + "modalities": ["video"], + "extra_args": { + "num_frames": int(num_frames), + "video_height": int(height), + "video_width": int(width), + "cfg_text_scale": float(cfg), + "timestep_shift": float(ts_shift), + }, + } + ] + if task == "i2v": + return [ + { + "prompt": render_lance_prompt("i2v", prompt, vision_token=_vision_block("video")), + "multi_modal_data": {"first_frame": image}, + "modalities": ["video"], + "extra_args": { + "num_frames": int(num_frames), + "video_height": int(height), + "video_width": int(width), + "cfg_text_scale": float(cfg), + "timestep_shift": float(ts_shift), + }, + } + ] + if task == "video_edit": + return [ + { + "prompt": render_lance_prompt("video_edit", prompt, vision_token=_vision_block("video")), + "multi_modal_data": {"video": video}, + "modalities": ["video"], + "extra_args": { + "num_frames": int(num_frames), + "video_height": int(height), + "video_width": int(width), + "cfg_text_scale": float(cfg), + "timestep_shift": float(ts_shift), + }, + } + ] + if task == "x2t_video": + q = (prompt or "").strip() or "Describe this video in detail." + return [ + { + "prompt": render_lance_prompt( + "x2t_video", + q, + vision_token=_vision_block("video"), + system_prompt="Watch the video carefully and answer the question.", + ), + "multi_modal_data": {"video": video}, + "modalities": ["text"], + "extra_args": {"max_think_tokens": 500}, + } + ] + raise ValueError(f"Unknown task: {task}") + + +def _extract_demo_sources(root: str, cache_dir: str = "/tmp/lance_demo_sources") -> list[tuple[str, str]]: + """Split upstream Lance's ``video-editing-demo-*.mp4`` and the + ``multi-turn-editing-demo-01.mp4`` filmstrip into runnable source + videos (the *input* half of each side-by-side demo) so they can be + surfaced as clickable ``gr.Examples`` for ``video_edit``. + + Upstream's 8 video-editing demos are **horizontally split** + ``(source | edit)`` — left half is the input video, right half is + the edited result. (A naïve top↔bottom RGB diff comparison is + misleading: the demos are landscape scenes where natural sky/ground + variance dominates the metric. Confirmed visually that the split + runs vertically through the middle.) The multi-turn-editing demo + is a horizontal filmstrip ``(panel0 | panel1 | … | panelN)`` where + panel 0 is the source and each subsequent panel adds one more edit + step. + + Returns ``[(source_video_path, suggested_prompt), ...]``. Generated + files are cached under ``cache_dir``; if the cache already has them + we skip re-encoding. + """ + from pathlib import Path as _Path + + import imageio.v3 as _iio + import numpy as _np + + if not root or not _Path(root).is_dir(): + return [] + root_p = _Path(root) + cache_p = _Path(cache_dir) + cache_p.mkdir(parents=True, exist_ok=True) + out: list[tuple[str, str]] = [] + + # Hand-curated prompts inferred by visually diffing each demo's left + # (source) vs right (edit) half — upstream's README ships the result + # videos but doesn't publish the prompts that produced them, so these + # are best-effort descriptions of the visible transformation. + demo_prompts = { + 1: "Add a roaring bonfire with bright orange flames behind the bear.", + 2: "Change the lush green tropical forest to a dry autumn forest with warm sunlight.", + 3: "Add several colorful balloons floating around the giraffe.", + 4: "Change her hair from red to natural chestnut brown and have her gently touch it with one hand.", + 5: "Replace the dog with a tabby cat, keeping the same red collar and car window.", + 6: "Remove the black square patch from her cheek.", + 7: "Change the boy to a girl with long black hair, blue eyes, and a black dress.", + 8: "Place her in a snowy forest and add a red striped scarf around her neck.", + } + + # ---- 8 horizontally-split video-editing demos ---- + for i in range(1, 9): + src_demo = root_p / "assets" / "video-editing" / "videos" / f"video-editing-demo-{i:02d}.mp4" + if not src_demo.exists(): + continue + out_path = cache_p / f"video-edit-demo-{i:02d}-source.mp4" + if not out_path.exists(): + arr = _iio.imread(str(src_demo), plugin="pyav") + # Horizontal split: take left half (width axis). + left = arr[:, :, : arr.shape[2] // 2, :] + left = _np.ascontiguousarray(left) + _iio.imwrite(str(out_path), left, fps=12, codec="libx264", quality=8) + out.append((str(out_path), demo_prompts.get(i, "Describe the edit you'd like to apply."))) + + # ---- multi-turn editing filmstrip → first panel as source ---- + mt_src = root_p / "assets" / "multi-turn-editing" / "videos" / "multi-turn-editing-demo-01.mp4" + if mt_src.exists(): + out_path = cache_p / "multi-turn-edit-source.mp4" + if not out_path.exists(): + arr = _iio.imread(str(mt_src), plugin="pyav") + T, H, W, _C = arr.shape + # The filmstrip is roughly ``N`` square-ish panels of width H + # each; pick the first panel as the source. W/H ≈ 11.7 on the + # shipped clip → panel_width ≈ H to keep aspect. + panel_w = H + first = arr[:, :, :panel_w, :] + first = _np.ascontiguousarray(first) + _iio.imwrite(str(out_path), first, fps=12, codec="libx264", quality=8) + out.append( + ( + str(out_path), + "Apply a multi-turn consistency edit — try sequential edits like " + "'change clothing color, then change background, then add accessories'.", + ) + ) + + return out + + +def _load_lance_showcase(root: str) -> dict[str, list[str]]: + """Collect upstream Lance's README showcase videos as visual + reference (full side-by-side ``before|after`` demos). Distinct from + :func:`_extract_demo_sources`, which crops just the source half so + they can be run. + + Surfaces: + * ``video_edit``: 8 ``video-editing-demo-{01..08}.mp4`` + the + ``multi-turn-editing-demo-01.mp4`` filmstrip. + """ + from pathlib import Path as _Path + + if not root or not _Path(root).is_dir(): + return {} + root_p = _Path(root) + out: dict[str, list[str]] = {} + + edit_videos: list[str] = [] + for i in range(1, 9): + p = root_p / "assets" / "video-editing" / "videos" / f"video-editing-demo-{i:02d}.mp4" + if p.exists(): + edit_videos.append(str(p)) + mt = root_p / "assets" / "multi-turn-editing" / "videos" / "multi-turn-editing-demo-01.mp4" + if mt.exists(): + edit_videos.append(str(mt)) + if edit_videos: + out["video_edit"] = edit_videos + + return out + + +def _load_lance_examples(root: str) -> dict[str, list[list]]: + """Parse upstream Lance's ``config/examples/*.json`` into clickable + presets keyed by task name. + + Each value is the list of rows that :class:`gr.Examples` expects, with + one row per official sample. The row column order matches the + ``inputs`` we'll bind below: ``[prompt, image_path | None, video_path | None]``. + + Path resolution: example JSONs reference media via repo-relative paths + (``config/examples/...``). We join against ``root`` so the gradio + process can serve the files directly. Examples whose media is + missing on disk are skipped silently so a partial upstream checkout + still works. + """ + import json as _json + from pathlib import Path as _Path + + if not root or not _Path(root).is_dir(): + return {} + root_p = _Path(root) + out: dict[str, list[list]] = {} + + def _resolve(rel: str) -> str | None: + if not isinstance(rel, str): + return None + p = root_p / rel + return str(p) if p.exists() else None + + def _safe_json(name: str): + f = root_p / "config" / "examples" / name + if not f.exists(): + return None + try: + return _json.loads(f.read_text()) + except Exception: # malformed json → skip + return None + + # ---- t2i / t2v: ``{filename: prompt}`` maps; only prompt is useful as + # an input preset (the filename is the *expected output*, not an + # input). ---- + for task, fname in (("t2i", "t2i_example.json"), ("t2v", "t2v_example.json")): + data = _safe_json(fname) + if isinstance(data, dict): + out[task] = [[prompt, None, None] for _, prompt in data.items() if isinstance(prompt, str)] + + # ---- i2v / video_edit: ``interleave_array = [prompt, media_path]`` ---- + # NOTE: image_edit handled separately below because upstream ships + # only one source image with 5 different prompts; we substitute + # diverse source images so each gradio example is visually distinct. + for task, fname, modality in ( + ("i2v", "i2v_example.json", "image"), + ("video_edit", "video_edit_example.json", "video"), + ): + data = _safe_json(fname) + if not isinstance(data, dict): + continue + rows: list[list] = [] + for _, sample in data.items(): + ia = sample.get("interleave_array") if isinstance(sample, dict) else None + if not (isinstance(ia, list) and len(ia) >= 2): + continue + prompt, media = ia[0], ia[1] + media_path = _resolve(media) + if media_path is None: + continue + row = [prompt, media_path, None] if modality == "image" else [prompt, None, media_path] + rows.append(row) + if rows: + out[task] = rows + + # ---- image_edit: curated visually-distinct pairs ---- + # Upstream ``image_edit_example.json`` ships 5 *different prompts* against + # the *same* source image (``edit_img.jpg``), which makes for a confusing + # gradio gallery (5 thumbnails that look identical). We substitute 5 + # visually-distinct sources from the i2v frame pool — same upstream + # repo, just paired differently — and write a fresh edit prompt for + # each so users see what kinds of transforms image_edit supports. + image_edit_pairs = [ + # (relative source path, prompt) + ("config/examples/image_edit_examples/edit_img.jpg", "Remove the hat from the painting."), + ( + "config/examples/text_image_to_video_examples/00001.png", + "Replace the icy glacier background with a sandy desert.", + ), + ( + "config/examples/text_image_to_video_examples/00005.png", + "Change her hair to bright red curly hair and add red lipstick.", + ), + ("config/examples/text_image_to_video_examples/00004.png", "Convert the scene into a 3D cartoon render style."), + ( + "config/examples/text_image_to_video_examples/00003.png", + "Add a vivid sunset sky with orange and purple clouds behind the penguin.", + ), + ] + ie_rows: list[list] = [] + for rel, prompt in image_edit_pairs: + path = _resolve(rel) + if path is not None: + ie_rows.append([prompt, path, None]) + if ie_rows: + out["image_edit"] = ie_rows + + # ---- Append the cropped video-editing showcase demos to ``video_edit`` ---- + # ``_extract_demo_sources`` slices the *source* half out of upstream's + # vertically-stacked side-by-side clips so they're runnable as + # ``video_edit`` examples (the bare side-by-side mp4s have the edit + # output baked in and would re-edit the wrong frames). + demo_rows = _extract_demo_sources(root) + if demo_rows: + existing = out.get("video_edit", []) + for src_path, prompt in demo_rows: + existing.append([prompt, None, src_path]) + out["video_edit"] = existing + + # ---- x2t_image / x2t_video: ``interleave_array = [media_path, [system, question, answer]]`` ---- + for task, fname, modality in ( + ("x2t_image", "x2t_image_example.json", "image"), + ("x2t_video", "x2t_video_example.json", "video"), + ): + data = _safe_json(fname) + if not isinstance(data, dict): + continue + rows: list[list] = [] + for _, sample in data.items(): + ia = sample.get("interleave_array") if isinstance(sample, dict) else None + if not (isinstance(ia, list) and len(ia) >= 2 and isinstance(ia[1], list) and len(ia[1]) >= 2): + continue + media_path = _resolve(ia[0]) + question = ia[1][1] + if media_path is None or not isinstance(question, str): + continue + row = [question, media_path, None] if modality == "image" else [question, None, media_path] + rows.append(row) + if rows: + out[task] = rows + + return out + + +def _to_pil(x): + if x is None: + return None + if isinstance(x, Image.Image): + return x.convert("RGB") + if isinstance(x, np.ndarray): + return Image.fromarray(x).convert("RGB") + if isinstance(x, str): + return Image.open(x).convert("RGB") + raise ValueError(f"Unsupported image type {type(x)}") + + +def _save_video(frames): + arr = np.stack([np.asarray(f) for f in frames], axis=0).astype(np.uint8) + out = Path(tempfile.mkdtemp(prefix="lance_grad_")) / "out.mp4" + iio.imwrite(str(out), arr, fps=12, codec="libx264", quality=8) + return str(out) + + +def _save_image(img): + out = Path(tempfile.mkdtemp(prefix="lance_grad_")) / "out.png" + if isinstance(img, np.ndarray): + Image.fromarray(img).save(out) + elif isinstance(img, Image.Image): + img.save(out) + else: + raise ValueError(type(img)) + return str(out) + + +# ---------------------------------------------------- generation ----- + + +def run_task( + task_label, + prompt, + image, + video, + aspect_ratio, + resolution, + video_duration_sec, + steps, + cfg, + ts_shift, + seed, + num_frames_override, +): + task = LABEL_TO_TASK[task_label] + omni, params = _omni_for(task) + + image = _to_pil(image) if image is not None else None + + # Resolve W/H from aspect + resolution + width, height = get_size_for_task(task, aspect_ratio, resolution) + + # Frames: for video tasks, use duration × fps + if task in {"t2v", "i2v", "video_edit"}: + fps = 12 + if num_frames_override and num_frames_override > 0: + num_frames = int(num_frames_override) + else: + num_frames = int(video_duration_sec * fps + 1) # +1 for first frame + else: + num_frames = 1 + + sp = params[0] + if seed is not None: + try: + sp.seed = int(seed) + except Exception: + sp.seed = DEFAULT_SEED + sp.num_inference_steps = int(steps) + + formatted = _format(task, prompt or "", image, video, num_frames, height, width, cfg, ts_shift) + print(f"[unified] task={task} {width}x{height} frames={num_frames} cfg={cfg} steps={steps} seed={seed}", flush=True) + + outputs = list(omni.generate(prompts=formatted, sampling_params_list=params)) + custom = getattr(outputs[0], "custom_output", None) or {} + # Lance/Bagel return text at outputs[0].output (string) or + # custom_output["text_output"]. + text_out = None + if task in {"x2t_image", "x2t_video"}: + text_out = custom.get("text_output") or getattr(outputs[0], "output", None) + if not isinstance(text_out, str): + text_out = None + + # NOTE: previously this returned a 4-tuple ending in a status string + # rendered into a ``gr.Markdown`` slot below the outputs. That slot + # left an empty grey bar in the UI before the first run and was + # dropped — Gradio's built-in error toast already surfaces failures, + # and successful runs are self-evident from the output panel. + if task in {"x2t_image", "x2t_video"}: + return ( + gr.update(visible=False), + gr.update(visible=False), + gr.update(value=str(text_out or "(empty)"), visible=True), + ) + if task in {"t2v", "i2v", "video_edit"}: + frames = custom.get("video_frames") + if frames is None: + raise gr.Error("No video frames in output.") + return (gr.update(visible=False), gr.update(value=_save_video(frames), visible=True), gr.update(visible=False)) + # image tasks — orchestrator strips ``outputs[0].output`` across the + # IPC boundary but preserves ``custom_output``, so prefer the latter. + img_out = custom.get("image") or getattr(outputs[0], "output", None) + if img_out is None: + raise gr.Error("No image in output.") + return (gr.update(value=_save_image(img_out), visible=True), gr.update(visible=False), gr.update(visible=False)) + + +# -------------------------------------------------- UI building ----- + +# Custom CSS — minimal, matches Lance branding (rounded panels + dark header). +APP_CSS = """ +.lance-hero { + display: flex; flex-direction: column; align-items: center; + gap: 12px; padding: 18px 12px; margin-bottom: 8px; +} +.lance-hero .logos { + display: flex; align-items: center; gap: 24px; justify-content: center; +} +.lance-hero .logos img { height: 64px; } +.lance-hero .logos .x { font-size: 28px; color: #888; font-weight: 300; } +.lance-hero .vllm-logo { height: 44px; } +.lance-hero .lance-title { + margin: 6px 0 0 0; font-size: 1.45rem; font-weight: 600; + text-align: center; color: #475569; +} +.lance-hero .lance-badges { display: flex; gap: 8px; flex-wrap: wrap; justify-content: center; } +.lance-panel { padding: 12px; border-radius: 12px; } +.task-selector .wrap { display: flex; flex-wrap: wrap; gap: 6px; } +.lance-run-button { font-size: 1.05rem !important; padding: 12px 18px !important; } +/* Reserve a single line of vertical space so the layout doesn't jump + when status text appears/disappears — but ``:empty`` collapses to 0 + so a placeholder ``gr.Markdown("")`` (e.g. the run-status slot before + the first generation) doesn't render as a phantom grey bar. */ +.lance-run-status { text-align: center; color: #475569; min-height: 1.4em; } +.lance-run-status:empty, +.lance-run-status .prose:empty, +.lance-run-status .prose p:empty { min-height: 0 !important; padding: 0 !important; margin: 0 !important; } +.lance-display-frame video, .lance-display-frame img { max-height: 560px; } +""" + + +def _embed_logo(path: str, mime: str) -> str: + if not os.path.exists(path): + return "" + with open(path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("ascii") + return f"data:{mime};base64,{b64}" + + +def build_header_html() -> str: + lance_uri = _embed_logo("/home/zjy/code/lance-upstream/assets/logo/lance-logo.webp", "image/webp") + vllm_uri = _embed_logo("/home/zjy/code/vllm-omni/docs/source/logos/vllm-omni-logo.png", "image/png") + logos_html = "" + if lance_uri: + logos_html += f'Lance' + if vllm_uri: + logos_html += '×' + logos_html += f'' + return f""" +
+
{logos_html}
+

Lance × vllm-omni · Unified Multimodal Inference

+
+ Homepage + Paper + HF + Lance GitHub + vllm-omni GitHub +
+
+ """ + + +def build_ui( + examples_by_task: dict[str, list[list]] | None = None, + showcase_by_task: dict[str, list[str]] | None = None, +): + examples_by_task = examples_by_task or {} + showcase_by_task = showcase_by_task or {} + with gr.Blocks(title="Lance × vllm-omni", css=APP_CSS) as demo: + gr.HTML(build_header_html()) + + # ── Task selector ── + task_radio = gr.Radio( + label="Task", + show_label=False, + choices=list(LABEL_TO_TASK.keys()), + value=TASK_I2V_LABEL, + elem_classes=["task-selector"], + ) + + with gr.Row(): + # ─── LEFT: inputs ─── + with gr.Column(scale=1): + with gr.Group(elem_classes=["lance-panel"]): + prompt_in = gr.Textbox( + label="Prompt", + lines=5, + placeholder="Describe the desired output...", + ) + # Aspect ratio is **not user-configurable** — upstream + # Lance always snaps to one of 6 hard-coded buckets based + # on the input image / video. We keep the component in + # the tree (so existing event wiring still works) but + # render it invisibly; ``_auto_aspect_from_image`` / + # ``_auto_aspect_from_video`` are the only writers. + aspect_in = gr.Dropdown( + label="Aspect Ratio (auto)", + choices=ASPECT_RATIO_CHOICES, + value=DEFAULT_VIDEO_AR, + visible=False, + interactive=False, + ) + with gr.Row(): + resolution_in = gr.Dropdown( + label="Resolution", + choices=VIDEO_RES_CHOICES, + value=DEFAULT_VIDEO_RES, + visible=True, + ) + duration_in = gr.Slider( + 1, + MAX_VIDEO_DURATION_SEC, + value=DEFAULT_VIDEO_DURATION_SEC, + step=1, + label="Duration (sec)", + ) + size_md = gr.Markdown(_format_size_md("i2v", 848, 480), elem_classes=["lance-run-status"]) + + image_in = gr.Image(label="Input Image", type="pil", visible=True, elem_classes=["lance-display-frame"]) + video_in = gr.Video(label="Input Video", visible=False, elem_classes=["lance-display-frame"]) + + with gr.Accordion("Advanced Parameters", open=False): + seed_in = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0) + steps_in = gr.Slider(5, 80, value=DEFAULT_TIMESTEPS, step=5, label="Denoise steps") + with gr.Row(): + cfg_in = gr.Number(label="CFG text scale", value=DEFAULT_CFG) + ts_shift_in = gr.Number(label="Timestep shift", value=DEFAULT_TIMESTEP_SHIFT) + num_frames_override = gr.Number( + label="Override num_frames (0 = derive from duration)", + value=0, + precision=0, + ) + + run_btn = gr.Button("🚀 Generate", variant="primary", elem_classes=["lance-run-button"]) + + # ─── RIGHT: output ─── + with gr.Column(scale=1): + with gr.Group(elem_classes=["lance-panel"]): + image_out = gr.Image(label="Output Image", visible=False, elem_classes=["lance-display-frame"]) + video_out = gr.Video( + label="Output Video", + autoplay=True, + loop=True, + visible=True, + elem_classes=["lance-display-frame"], + ) + text_out = gr.Textbox(label="Output Caption", lines=6, visible=False) + # Status text was previously rendered as ``gr.Markdown("", + # ...)`` which left an empty grey bar in the UI before the + # first run. Dropped — Gradio's built-in toast / error + # bubble already surfaces success / failure, so the extra + # placeholder added no signal but stole vertical real estate. + + # ─── Official Lance examples per task ─── + # Click an example row → prompt + image/video auto-fill. Each + # task has its own ``gr.Examples`` block; only the active task's + # block is shown. Order matches LABEL_TO_TASK (radio order). + _initial_task = LABEL_TO_TASK[TASK_I2V_LABEL] # default selection + _example_groups: dict[str, gr.Group] = {} + if examples_by_task: + # Per-task input bindings — Gradio's ``gr.Video`` doesn't reliably + # update from a programmatic example click when the Examples row + # also includes ``None`` placeholders for ``Image`` (the player + # stays blank). Binding each task's Examples block to *only* the + # components it actually fills sidesteps that issue. + task_input_components: dict[str, list] = { + "t2i": [prompt_in], + "t2v": [prompt_in], + "i2v": [prompt_in, image_in], + "image_edit": [prompt_in, image_in], + "x2t_image": [prompt_in, image_in], + "video_edit": [prompt_in, video_in], + "x2t_video": [prompt_in, video_in], + } + + def _project_row(task: str, row: list) -> list: + # ``examples_by_task`` rows are always ``[prompt, image, video]`` + # (3 columns). Project down to whatever this task binds. + prompt, img, vid = row + comps = task_input_components[task] + out = [] + for c in comps: + if c is prompt_in: + out.append(prompt) + elif c is image_in: + out.append(img) + elif c is video_in: + out.append(vid) + return out + + for task in ("t2i", "t2v", "i2v", "image_edit", "video_edit", "x2t_image", "x2t_video"): + rows = examples_by_task.get(task) + if not rows: + continue + projected_rows = [_project_row(task, r) for r in rows] + with gr.Group(visible=(task == _initial_task), elem_classes=["lance-panel"]) as grp: + gr.Markdown(f"### 📁 Official Lance examples — `{task}` (click to load)") + # Bind to only the components this task actually populates. + gr.Examples( + examples=projected_rows, + inputs=task_input_components[task], + label="", + examples_per_page=4, + cache_examples=False, + ) + # NOTE: Upstream's 8 ``video-editing-demo-*.mp4`` + the + # ``multi-turn-editing-demo-01.mp4`` filmstrip are merged + # into the runnable examples list above via + # ``_extract_demo_sources`` (left-half / first-panel crops + # become the source video, paired with a generic editable + # prompt). No separate showcase gallery is needed. + _example_groups[task] = grp + + gr.Markdown("""--- +💡 **Tips:** +* **Video Generation**: text → 5-second video at 480p +* **Image to Video**: upload first frame + describe motion; for subtle facial expressions try CFG 10-15 +* **Image Edit**: upload image + describe the edit; outputs a new image +* **Video Edit**: upload video + describe the edit; outputs a modified video +* **Understanding tasks**: upload media; outputs caption text""") + + # ─────────────── UI logic ─────────────── + + # Fixed task order for the example-group visibility tail of on_task's + # outputs — used both to build the return list and to register the + # corresponding output components below. Must stay in sync with the + # group-creation loop above. + _example_task_order = [ + t + for t in ("t2i", "t2v", "i2v", "image_edit", "video_edit", "x2t_image", "x2t_video") + if t in _example_groups + ] + + def on_task(label): + task = LABEL_TO_TASK[label] + is_video_task = task in {"t2v", "i2v", "video_edit"} + is_image_task = task in {"t2i", "image_edit"} + is_understanding = task in {"x2t_image", "x2t_video"} + # Understanding tasks don't need a prompt — system prompt handles + # captioning. Generation/edit tasks need a prompt. + need_text = not is_understanding + need_image = task in {"i2v", "image_edit", "x2t_image"} + need_video = task in {"video_edit", "x2t_video"} + default_ar = DEFAULT_VIDEO_AR if (is_video_task or is_understanding) else DEFAULT_IMAGE_AR + w, h = get_size_for_task(task, default_ar, DEFAULT_VIDEO_RES) + base = [ + gr.update(visible=need_text), # prompt_in + gr.update(visible=need_image, value=None), # image_in + gr.update(visible=need_video, value=None), # video_in + gr.update(value=default_ar, visible=False), # aspect (always hidden, auto-detected) + gr.update(visible=is_video_task, value=DEFAULT_VIDEO_RES), # resolution + gr.update( + visible=is_video_task, + value=DURATION_DEFAULT_PER_TASK.get(task, DEFAULT_VIDEO_DURATION_SEC), + ), # duration — matches upstream's per-script num_frames default + gr.update(value=_format_size_md(task, w, h)), # size_md + gr.update(visible=is_image_task and not is_video_task), # image_out (only show for image-gen tasks) + gr.update(visible=is_video_task), # video_out + gr.update(visible=is_understanding), # text_out + ] + return base + [gr.update(visible=(t == task)) for t in _example_task_order] + + task_radio.change( + on_task, + [task_radio], + [ + prompt_in, + image_in, + video_in, + aspect_in, + resolution_in, + duration_in, + size_md, + image_out, + video_out, + text_out, + ] + + [_example_groups[t] for t in _example_task_order], + ) + + def on_aspect_or_res(task_label, ar, resolution): + task = LABEL_TO_TASK[task_label] + w, h = get_size_for_task(task, ar, resolution) + return gr.update(value=_format_size_md(task, w, h)) + + aspect_in.change(on_aspect_or_res, [task_radio, aspect_in, resolution_in], [size_md]) + resolution_in.change(on_aspect_or_res, [task_radio, aspect_in, resolution_in], [size_md]) + + # ── Auto-pick aspect ratio from the input media ── + # For tasks where the model is conditioned on an input + # image/video (image_edit, i2v, video_edit), the output should + # preserve the source aspect ratio — otherwise the model has to + # squeeze / stretch the scene to fill a different frame, and the + # geometry drifts (e.g. the snow-leopard "stepping into the + # void" case: a 3:2 input rendered into a 16:9 output pushed the + # landing platform off-frame). When the user uploads or clicks + # an example, snap the aspect dropdown to whichever preset is + # closest to the source. + def _auto_aspect_from_image(task_label, img): + task = LABEL_TO_TASK[task_label] + if task not in {"image_edit", "i2v"} or img is None: + return gr.update(), gr.update() + try: + pil = _to_pil(img) + w, h = pil.size + except Exception: + return gr.update(), gr.update() + ar = _closest_aspect_ratio(w, h) + new_w, new_h = get_size_for_task(task, ar, DEFAULT_VIDEO_RES) + return ( + gr.update(value=ar), + gr.update(value=_format_size_md(task, new_w, new_h)), + ) + + def _auto_aspect_from_video(task_label, vid_path): + task = LABEL_TO_TASK[task_label] + if task != "video_edit" or not vid_path: + return gr.update(), gr.update() + try: + import imageio.v3 as _iio + + meta = _iio.immeta(vid_path, plugin="pyav") + # ``immeta`` ships ``size = (w, h)`` for video readers. + if "size" in meta: + w, h = meta["size"] + else: + # Fall back to a single-frame read. + arr = _iio.imread(vid_path, plugin="pyav", index=0) + h, w = arr.shape[:2] + except Exception: + return gr.update(), gr.update() + ar = _closest_aspect_ratio(w, h) + new_w, new_h = get_size_for_task(task, ar, DEFAULT_VIDEO_RES) + return ( + gr.update(value=ar), + gr.update(value=_format_size_md(task, new_w, new_h)), + ) + + image_in.change( + _auto_aspect_from_image, + [task_radio, image_in], + [aspect_in, size_md], + ) + video_in.change( + _auto_aspect_from_video, + [task_radio, video_in], + [aspect_in, size_md], + ) + + run_btn.click( + run_task, + inputs=[ + task_radio, + prompt_in, + image_in, + video_in, + aspect_in, + resolution_in, + duration_in, + steps_in, + cfg_in, + ts_shift_in, + seed_in, + num_frames_override, + ], + outputs=[image_out, video_out, text_out], + ) + + return demo + + +def main(): + args = parse_args() + print("=" * 60, flush=True) + print("[unified] Lance × vllm-omni demo", flush=True) + print(f" snapshot: {args.model}", flush=True) + print(f" gradio: {args.host}:{args.port} (share={args.share})", flush=True) + print("=" * 60, flush=True) + _init_models( + args.model, + args.deploy_config, + ulysses_degree=args.ulysses_degree, + replicas_per_omni=args.replicas_per_omni, + ) + examples_by_task = _load_lance_examples(args.examples_root) + if examples_by_task: + total = sum(len(rows) for rows in examples_by_task.values()) + per_task = ", ".join(f"{t}={len(r)}" for t, r in examples_by_task.items()) + print(f"[unified] Loaded {total} official examples ({per_task}) from {args.examples_root}", flush=True) + else: + print(f"[unified] No official examples loaded (examples_root={args.examples_root!r})", flush=True) + demo = build_ui(examples_by_task=examples_by_task) + # Concurrency budget = (image replicas + video replicas) so all the + # available GPU slots can run in parallel. P0 (1 replica each) = 2; + # ``--replicas-per-omni 4`` on 8 GPUs gives 8 concurrent slots. + queue_limit = 2 * max(1, args.replicas_per_omni) + print(f"[unified] Gradio queue concurrency limit = {queue_limit}", flush=True) + demo.queue(default_concurrency_limit=queue_limit) + # Whitelist the upstream checkout so gr.Examples can stage media that + # lives outside the gradio working dir / temp. Without this, clicks + # raise ``InvalidPathError`` and the example silently fails to load. + launch_kwargs: dict = dict(server_name=args.host, server_port=args.port, share=args.share) + allowed = [] + if args.examples_root and Path(args.examples_root).is_dir(): + allowed.append(str(Path(args.examples_root).resolve())) + # ``_extract_demo_sources`` caches cropped video-editing demo sources under + # ``/tmp/lance_demo_sources`` so they survive across launches; whitelist it + # too or the click-to-load will hit ``InvalidPathError``. + demo_cache = Path("/tmp/lance_demo_sources") + if demo_cache.is_dir(): + allowed.append(str(demo_cache.resolve())) + if allowed: + launch_kwargs["allowed_paths"] = allowed + demo.launch(**launch_kwargs) + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 257babba8..00df0dff8 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/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", ), diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 7ed65a1af..d333aa481 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index 24a3631fa..e1c051db3 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/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, diff --git a/vllm_omni/diffusion/models/lance/prompts.py b/vllm_omni/diffusion/models/lance/prompts.py index a3e8611ca..efdd007e8 100644 --- a/vllm_omni/diffusion/models/lance/prompts.py +++ b/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 diff --git a/vllm_omni/diffusion/models/lance/wan_vae.py b/vllm_omni/diffusion/models/lance/wan_vae.py index aad244d30..a8e167534 100644 --- a/vllm_omni/diffusion/models/lance/wan_vae.py +++ b/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 diff --git a/vllm_omni/model_executor/models/lance/__init__.py b/vllm_omni/model_executor/models/lance/__init__.py index ee47a467a..2ecd5df1f 100644 --- a/vllm_omni/model_executor/models/lance/__init__.py +++ b/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"] diff --git a/vllm_omni/model_executor/models/lance/lance.py b/vllm_omni/model_executor/models/lance/lance.py deleted file mode 100644 index 7ae9dd6f7..000000000 --- a/vllm_omni/model_executor/models/lance/lance.py +++ /dev/null @@ -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 diff --git a/vllm_omni/model_executor/models/lance/pipeline.py b/vllm_omni/model_executor/models/lance/pipeline.py index fa7d14fa7..777f8942d 100644 --- a/vllm_omni/model_executor/models/lance/pipeline.py +++ b/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}, - ), - ), -) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 0fa2d4c62..4f29c0ae3 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/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", diff --git a/vllm_omni/model_executor/stage_input_processors/lance.py b/vllm_omni/model_executor/stage_input_processors/lance.py deleted file mode 100644 index 5ab029a50..000000000 --- a/vllm_omni/model_executor/stage_input_processors/lance.py +++ /dev/null @@ -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", -]