diff --git a/LANCE_PR.md b/LANCE_PR.md new file mode 100644 index 000000000..d42ccce20 --- /dev/null +++ b/LANCE_PR.md @@ -0,0 +1,76 @@ +# [New Model] Lance (ByteDance) — #3697 + +Adds Lance (`bytedance-research/Lance`), a 3B unified AR+diffusion +multimodal model on a Qwen2.5-VL backbone. + +## Why this is BAGEL, not WAN + +The issue thread suggested remapping against `WanTransformer3DModel`. The +released `Lance_3B` checkpoint is instead **BAGEL Mixture-of-Transformers**: +every layer carries `mlp_moe_gen` / `q_proj_moe_gen` / +`input_layernorm_moe_gen` twins plus `q_norm`/`k_norm`, with top-level +`vae2llm` / `llm2vae` / `time_embedder` / `latent_pos_embed` connectors — +identical naming to vLLM-Omni's BAGEL. The only WAN element is the VAE. +Lance therefore reuses BAGEL's transformer core verbatim and specializes +only the LLM (Qwen2.5-VL-MoT), ViT (Qwen2.5-VL vision), VAE (Wan2.2) and +checkpoint layout. + +## Scope and status + +All six single-stage Lance modalities, verified end-to-end on 1× B300: + +| Modality | Status | +|----------|--------| +| `t2i` (text→image) | working — 1024×1024, 0 missing keys | +| `t2v` (text→video) | working — photorealistic frames with motion; 768×768 trained bucket | +| `image_edit` (img2img) | working — Lance-native VAE prefill scatters Wan2.2 latents into the LLM query sequence | +| `video_edit` (video2video) | working — multi-frame VAE prefill + 3-D mRoPE; needs the `Lance_3B_Video` checkpoint | +| `x2t_image` (image understanding) | working — Qwen2.5-VL ViT prefill + 3-D mRoPE | +| `x2t_video` (video understanding) | working — Qwen2-VL video processor + 3-D `video_grid_thw` | + +`rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired +through `BagelRotaryEmbedding`: t2i uses scalar position ids +(BAGEL-equivalent); x2t / video / edit thread per-axis 3-D position ids. +Geometry constants (`latent_patch_size_spatial = 1`, `max_latent_size = 64`, +untied `lm_head`, Wan2.2 VAE `z=48 /16 spatial /4 temporal`) were locked in +by inspecting the released checkpoint directly. + +The two-stage AR+DiT topology is the only deferred path (needs +`LanceConfig` / `LanceProcessor` in the `vllm` package — a separate +human-authored PR per vLLM's contribution policy). + +## Test plan + +Offline (`examples/offline_inference/lance/end2end.py`), e.g. t2i: + +```bash +python examples/offline_inference/lance/end2end.py \ + --model bytedance-research/Lance \ + --prompts "a corgi astronaut on the moon, cinematic" \ + --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 --output ./out +``` + +`--modality {text2video,img2img,video2video,img2text,video2text}` selects +the other paths; video paths use `--model bytedance-research/Lance/Lance_3B_Video`. + +Online (OpenAI-compatible, `examples/online_serving/lance/`): + +```bash +bash examples/online_serving/lance/run_server.sh # vllm serve … --omni +python examples/online_serving/lance/openai_chat_client.py \ + --prompt "A cute corgi astronaut on the moon, cinematic" \ + --modality text2img --output corgi.png +``` + +E2E coverage: `tests/e2e/online_serving/test_lance.py`. Tested on 1× +NVIDIA B300, `torch 2.11.0+cu130`, stock attention backend. + +## Follow-ups + +1. **Two-stage AR + DiT** — register `LanceConfig` / `LanceProcessor` in + the `vllm` package; separate human-authored PR per vLLM policy. +2. **`video_edit` quality polish** — output is more abstract than `t2v` + at the same resolution (position-id offset between the VAE-prefill and + gen-latent blocks). Functionally correct end-to-end. + +Closes #3697. diff --git a/examples/offline_inference/lance/README.md b/examples/offline_inference/lance/README.md new file mode 100644 index 000000000..0697c50c7 --- /dev/null +++ b/examples/offline_inference/lance/README.md @@ -0,0 +1,50 @@ +# Lance (ByteDance) — offline inference + +[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified +autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is +**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B` +checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni +implements it by reusing the BAGEL transformer core and specializing only the +ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout. + +This example covers all six Lance modalities from the upstream HF model card: +`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and +`x2t_video` (video understanding). + +## Hardware + +Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4. + +## Run + +```bash +# Text-to-image +python examples/offline_inference/lance/end2end.py \ + --model bytedance-research/Lance \ + --prompts "a corgi astronaut on the moon, cinematic" \ + --steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \ + --output ./out + +# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality`` +# choices for all six task variants) +python examples/offline_inference/lance/end2end.py \ + --model bytedance-research/Lance/Lance_3B_Video --modality text2video \ + --num-frames 25 --video-height 480 --video-width 768 \ + --prompts "a cat playing piano, cinematic" \ + --steps 30 --fps 8 --output ./out +``` + +`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so +the 3-D `latent_pos_embed` table is loaded; the other paths can point at +the top-level `bytedance-research/Lance` repo and resolve the right +sub-checkpoint automatically. + +The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`, +`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required. + +## Defaults + +Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5, +text CFG 4.0, seed 42, 768×768. For the understanding paths (`img2text` / +`video2text`), sampling is enabled by default at `--text-temperature 0.8` +because Lance's greedy decoder emits an immediate EOS for many prompts. diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py new file mode 100644 index 000000000..2f0a94dea --- /dev/null +++ b/examples/offline_inference/lance/end2end.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance-3B offline inference (single-stage). + +Lance is BAGEL-lineage; this example covers the single-stage paths wired in +this PR: + + text2img — image generation from text + img2text — image understanding (caption / VQA) + +Examples: + + python examples/offline_inference/lance/end2end.py \ + --model bytedance-research/Lance \ + --prompts "a corgi astronaut on the moon, cinematic" \ + --output ./out + + python examples/offline_inference/lance/end2end.py \ + --model bytedance-research/Lance --modality img2text \ + --image-path /path/to/photo.jpg \ + --prompts "Describe this image in detail." +""" + +import argparse +import os + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.") + p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") + p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.") + p.add_argument( + "--modality", + default="text2img", + choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video"], + help="Lance single-stage modality.", + ) + p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.") + p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.") + p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).") + p.add_argument("--video-height", type=int, default=480, help="Video frame height.") + p.add_argument("--video-width", type=int, default=768, help="Video frame width.") + p.add_argument("--fps", type=int, default=8, help="Output video FPS when saving MP4.") + p.add_argument("--output", type=str, default=".", help="Output directory.") + p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).") + p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).") + p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).") + p.add_argument("--negative-prompt", type=str, default=None) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.") + p.add_argument( + "--do-sample", + action="store_true", + default=True, + help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.", + ) + p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).") + p.add_argument( + "--text-temperature", + type=float, + default=0.8, + help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.", + ) + p.add_argument( + "--deploy-config", + type=str, + default="vllm_omni/deploy/lance.yaml", + help="Deploy YAML (single-stage Lance).", + ) + + from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults + + nullify_stage_engine_defaults(p) + return p.parse_args() + + +from vllm_omni.diffusion.models.lance.prompts import ( # noqa: E402 + VIDEO_PAD, + VISION_END, + VISION_START, + render_lance_prompt, +) + +# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|> +# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.) +_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}" + + +def _format_text2img_prompts(prompts): + return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts] + + +def _format_text2video_prompts(prompts, num_frames, video_h, video_w): + extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w} + return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts] + + +def _format_img2img_prompts(prompts, image_path): + import os as _os + + if not image_path or not _os.path.exists(image_path): + raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}") + from PIL import Image as _Image + + img = _Image.open(image_path).convert("RGB") + return [ + { + "prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK), + "multi_modal_data": {"img2img": img}, + "modalities": ["image"], + } + for p in prompts + ] + + +def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w): + import os as _os + + if not video_path or not _os.path.exists(video_path): + raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}") + import imageio.v3 as iio + + video = iio.imread(video_path) + extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w} + return [ + { + "prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK), + "multi_modal_data": {"video": video}, + "modalities": ["video"], + "extra_args": extra, + } + for p in prompts + ] + + +def _format_video2text_prompts(prompts, video_path): + import os as _os + + if not video_path or not _os.path.exists(video_path): + raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}") + import imageio.v3 as iio + + video = iio.imread(video_path) + return [ + { + "prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK), + "multi_modal_data": {"video": video}, + "modalities": ["text"], + } + for p in prompts + ] + + +def _format_img2text_prompts(prompts, image_path): + from PIL import Image + + if not image_path or not os.path.exists(image_path): + raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}") + img = Image.open(image_path).convert("RGB") + return [ + { + "prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK), + "multi_modal_data": {"image": img}, + "modalities": ["text"], + } + for p in prompts + ] + + +def main(): + args = parse_args() + os.makedirs(args.output, exist_ok=True) + + if args.txt_prompts: + with open(args.txt_prompts, encoding="utf-8") as f: + prompts = [ln.strip() for ln in f if ln.strip()] + elif args.modality == "img2text": + prompts = args.prompts or ["Describe this image in detail."] + else: + prompts = args.prompts or ["A cute cat, studio lighting, 4k"] + + from vllm_omni.entrypoints.omni import Omni + + omni_kwargs = vars(args).copy() + omni_kwargs["deploy_config"] = args.deploy_config + omni_kwargs["model"] = args.model + omni = Omni(**omni_kwargs) + + if args.modality == "img2text": + formatted = _format_img2text_prompts(prompts, args.image_path) + elif args.modality == "text2video": + formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width) + elif args.modality == "video2text": + formatted = _format_video2text_prompts(prompts, args.video_path) + elif args.modality == "img2img": + formatted = _format_img2img_prompts(prompts, args.image_path) + elif args.modality == "video2video": + formatted = _format_video2video_prompts( + prompts, args.video_path, args.num_frames, args.video_height, args.video_width + ) + else: + formatted = _format_text2img_prompts(prompts) + + params_list = omni.default_sampling_params_list + diffusion_params = params_list[0] # single-stage: one param set + diffusion_params.num_inference_steps = args.steps # type: ignore + if args.seed is not None: + diffusion_params.seed = args.seed # type: ignore + extra = getattr(diffusion_params, "extra_args", {}) or {} + extra["cfg_text_scale"] = args.cfg_text_scale + extra["timestep_shift"] = args.timestep_shift + if args.modality in ("img2text", "video2text"): + extra["max_think_tokens"] = args.max_text_tokens + extra["do_sample"] = args.do_sample + extra["text_temperature"] = args.text_temperature + if args.negative_prompt is not None: + extra["negative_prompt"] = args.negative_prompt + diffusion_params.extra_args = extra # type: ignore + + outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list)) + + if args.modality in ("img2text", "video2text"): + for i, req_output in enumerate(outputs): + text = getattr(req_output, "output", None) or getattr(req_output, "text", None) + if text is None and hasattr(req_output, "custom_output"): + text = (req_output.custom_output or {}).get("text_output") + print(f"[Output {i}] {text}") + return + + if args.modality in ("text2video", "video2video"): + for i, req_output in enumerate(outputs): + custom = getattr(req_output, "custom_output", None) or {} + frames = custom.get("video_frames") + if not frames: + print(f"[Output {i}] no video frames returned") + continue + mp4_path = os.path.join(args.output, f"lance_{i}.mp4") + try: + import imageio.v3 as iio + import numpy as np + + arr = np.stack([np.asarray(f) for f in frames], axis=0) + iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8) + print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}") + except Exception as exc: + print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead") + frame_dir = os.path.join(args.output, f"lance_{i}_frames") + os.makedirs(frame_dir, exist_ok=True) + for j, frame in enumerate(frames): + frame.save(os.path.join(frame_dir, f"{j:04d}.png")) + print(f" → {frame_dir} ({len(frames)} frames)") + return + + idx = 0 + for req_output in outputs: + images = getattr(req_output, "images", None) + if not images: + continue + for j, img in enumerate(images): + path = os.path.join(args.output, f"lance_{idx}_{j}.png") + img.save(path) + print(f"[Output] Saved image to {path}") + idx += 1 + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/lance/README.md b/examples/online_serving/lance/README.md new file mode 100644 index 000000000..5f2febce9 --- /dev/null +++ b/examples/online_serving/lance/README.md @@ -0,0 +1,40 @@ +# Lance — online serving + +OpenAI-compatible chat completions API for Lance, served by `vllm-omni`. + +## Start the server + +```bash +bash examples/online_serving/lance/run_server.sh +# or, with overrides +MODEL=bytedance-research/Lance \ +DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \ +PORT=8091 \ + bash examples/online_serving/lance/run_server.sh +``` + +The default deploy config is the single-stage `lance.yaml`; for `video_edit` +or `text2video` workloads, point ``DEPLOY_CONFIG`` at a video-checkpoint +config (or pass ``--model bytedance-research/Lance/Lance_3B_Video`` to +`run_server.sh`). + +## Send a request + +```bash +# text-to-image +python examples/online_serving/lance/openai_chat_client.py \ + --prompt "A cute corgi astronaut on the moon, cinematic" \ + --modality text2img \ + --output corgi.png + +# image edit +python examples/online_serving/lance/openai_chat_client.py \ + --prompt "Convert this into a vibrant cartoon-style illustration" \ + --modality img2img \ + --image-url path/to/photo.png \ + --output edited.png +``` + +The client is shared with BAGEL — same OpenAI message format, same +``modalities`` and ``num_inference_steps`` / ``seed`` / ``height`` / +``width`` knobs. diff --git a/examples/online_serving/lance/openai_chat_client.py b/examples/online_serving/lance/openai_chat_client.py new file mode 100755 index 000000000..1f470069b --- /dev/null +++ b/examples/online_serving/lance/openai_chat_client.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +""" +Lance OpenAI-compatible chat client for image generation and multimodal tasks. + +Usage: + python examples/online_serving/lance/openai_chat_client.py --prompt "A cute cat" --output output.png + python examples/online_serving/lance/openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png +""" + +import argparse +import base64 +from pathlib import Path + +import requests + + +def generate_image( + prompt: str, + server_url: str = "http://localhost:8091", + image_url: str | None = None, + height: int | None = None, + width: int | None = None, + steps: int | None = None, + seed: int | None = None, + negative_prompt: str | None = None, + modality: str = "text2img", # "text2img" (default), "img2img", "img2text", "text2text" +) -> bytes | str | None: + """Generate an image or text using the chat completions API. + + Args: + prompt: Text description or prompt + server_url: Server URL + image_url: URL or path to input image (for img2img/img2text) + height: Image height in pixels + width: Image width in pixels + steps: Number of inference steps + seed: Random seed + negative_prompt: Negative prompt + modality: Task modality hint + + Returns: + Image bytes (for image outputs) or Text string (for text outputs) or None if failed + """ + + # Construct Message Content + content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}] + + if image_url: + # Check if local file + if Path(image_url).exists(): + with open(image_url, "rb") as f: + b64_data = base64.b64encode(f.read()).decode("utf-8") + final_image_url = f"data:image/jpeg;base64,{b64_data}" + else: + final_image_url = image_url + + content.append({"type": "image_url", "image_url": {"url": final_image_url}}) + + messages = [{"role": "user", "content": content}] + + # Build request payload with all parameters at top level + # Note: vLLM ignores "extra_body", so we put parameters directly in the payload + payload = {"messages": messages} + + # Set output modalities at top level + if modality == "text2img" or modality == "img2img": + payload["modalities"] = ["image"] + elif modality == "img2text" or modality == "text2text": + payload["modalities"] = ["text"] + + # Add generation parameters directly to payload + if height is not None: + payload["height"] = height + if width is not None: + payload["width"] = width + if steps is not None: + payload["num_inference_steps"] = steps + if seed is not None: + payload["seed"] = seed + if negative_prompt: + payload["negative_prompt"] = negative_prompt + + # Send request + try: + print(f"Sending request to {server_url} with modality {modality}...") + response = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=300, + ) + response.raise_for_status() + data = response.json() + + # Extract content - check ALL choices since server may return multiple + # (e.g., text in choices[0], image in choices[1]) + choices = data.get("choices", []) + + # First pass: look for image output in any choice + for choice in choices: + choice_content = choice.get("message", {}).get("content") + + # Handle Image Output + if isinstance(choice_content, list) and len(choice_content) > 0: + first_item = choice_content[0] + if isinstance(first_item, dict) and "image_url" in first_item: + img_url_str = first_item["image_url"].get("url", "") + if img_url_str.startswith("data:image"): + _, b64_data = img_url_str.split(",", 1) + return base64.b64decode(b64_data) + + # Second pass: look for text output if no image found + for choice in choices: + choice_content = choice.get("message", {}).get("content") + if isinstance(choice_content, str) and choice_content: + return choice_content + + print(f"Unexpected response format: {choices}") + return None + + except Exception as e: + print(f"Error: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="Bagel multimodal chat client") + parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt") + parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)") + parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL") + + # Modality Control + parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path") + parser.add_argument( + "--modality", + "-m", + default="text2img", + choices=["text2img", "img2img", "img2text", "text2text"], + help="Task modality", + ) + + # Generation Params + parser.add_argument("--height", type=int, default=512, help="Image height") + parser.add_argument("--width", type=int, default=512, help="Image width") + parser.add_argument("--steps", type=int, default=25, help="Inference steps") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + parser.add_argument("--negative", help="Negative prompt") + + args = parser.parse_args() + + print(f"Mode: {args.modality}") + if args.image_url: + print(f"Input Image: {args.image_url}") + + result = generate_image( + prompt=args.prompt, + server_url=args.server, + image_url=args.image_url, + height=args.height, + width=args.width, + steps=args.steps, + seed=args.seed, + negative_prompt=args.negative, + modality=args.modality, + ) + + if result: + if isinstance(result, bytes): + # It's an image + output_path = Path(args.output) + output_path.write_bytes(result) + print(f"Image saved to: {output_path}") + print(f"Size: {len(result) / 1024:.1f} KB") + elif isinstance(result, str): + # It's text + print("Response:") + print(result) + else: + print("Failed to generate response") + exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/lance/run_server.sh b/examples/online_serving/lance/run_server.sh new file mode 100755 index 000000000..bba316677 --- /dev/null +++ b/examples/online_serving/lance/run_server.sh @@ -0,0 +1,15 @@ +#!/bin/bash +# Lance online serving startup script + +MODEL="${MODEL:-bytedance-research/Lance}" +PORT="${PORT:-8091}" +DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/lance.yaml}" + +echo "Starting Lance server..." +echo "Model: $MODEL" +echo "Deploy config: $DEPLOY_CONFIG" +echo "Port: $PORT" + +vllm serve "$MODEL" --omni \ + --deploy-config "$DEPLOY_CONFIG" \ + --port "$PORT" diff --git a/tests/e2e/online_serving/test_lance.py b/tests/e2e/online_serving/test_lance.py new file mode 100644 index 000000000..197c35746 --- /dev/null +++ b/tests/e2e/online_serving/test_lance.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end online serving test for Lance (single-stage). + +Verifies that the Lance pipeline can serve text-to-image and image-edit +requests via the OpenAI-compatible chat completions API exposed by +``vllm-omni serve``. + +Equivalent to running: + + vllm-omni serve "bytedance-research/Lance" --omni \\ + --deploy-config vllm_omni/deploy/lance.yaml --port 8091 + + # text2img + python3 examples/online_serving/lance/openai_chat_client.py \\ + --prompt "A cute corgi astronaut" --modality text2img +""" + +import base64 +import os +from io import BytesIO + +import pytest +from vllm.assets.image import ImageAsset + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.diffusion.models.lance.prompts import ( + VIDEO_PAD, + VISION_END, + VISION_START, + render_lance_prompt, +) + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "bytedance-research/Lance" +STAGE_CONFIGS_PATH = get_deploy_config_path("ci/lance.yaml") + +TEXT2IMG_PROMPT = "A cute corgi astronaut on the moon, cinematic" +IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration" + +# Lance vision block: <|vision_start|><|video_pad|><|vision_end|> +_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}" + +test_params = [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIGS_PATH, + stage_init_timeout=300, + ), +] + + +def _build_text2img_messages(prompt: str) -> list[dict]: + """Build OpenAI-format messages for text2img generation.""" + rendered = render_lance_prompt("t2i", prompt) + return [ + { + "role": "user", + "content": [{"type": "text", "text": rendered}], + } + ] + + +def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]: + """Build OpenAI-format messages for img2img generation.""" + rendered = render_lance_prompt("image_edit", prompt, vision_token=_VISION_BLOCK) + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": rendered}, + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}, + }, + ], + } + ] + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_lance_text2img_online(omni_server, openai_client) -> None: + """Lance text2img via the OpenAI-compatible chat completions API.""" + request_config = { + "model": omni_server.model, + "messages": _build_text2img_messages(TEXT2IMG_PROMPT), + "modalities": ["image"], + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "guidance_scale": 0.0, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_lance_img2img_online(omni_server, openai_client) -> None: + """Lance image_edit via the OpenAI-compatible chat completions API.""" + input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") + buffer = BytesIO() + input_image.save(buffer, format="JPEG") + image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8") + + request_config = { + "model": omni_server.model, + "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64), + "modalities": ["image"], + "extra_body": { + "num_inference_steps": 2, + "guidance_scale": 0.0, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index c3f9563cb..265c42821 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -495,6 +495,15 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = { }, ], }, + "lance": { + "base_config": "lance.yaml", + "stages": [ + { + "stage_id": 0, + "max_num_seqs": 1, + }, + ], + }, "bagel_mooncake": { "base_config": "bagel.yaml", "stages": [ diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 555f35e17..257babba8 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -65,6 +65,17 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.bagel.pipeline", "BAGEL_SINGLE_STAGE_PIPELINE", ), + # Lance (ByteDance) — BAGEL-lineage unified AR+diffusion. + # Single-stage is the runnable/registered topology; lance_two_stage is a + # follow-up (needs LanceConfig/LanceProcessor in the vllm package). + "lance": ( + "vllm_omni.model_executor.models.lance.pipeline", + "LANCE_SINGLE_STAGE_PIPELINE", + ), + "lance_two_stage": ( + "vllm_omni.model_executor.models.lance.pipeline", + "LANCE_PIPELINE", + ), "glm_image": ( "vllm_omni.model_executor.models.glm_image.pipeline", "GLM_IMAGE_PIPELINE", diff --git a/vllm_omni/deploy/lance.yaml b/vllm_omni/deploy/lance.yaml new file mode 100644 index 000000000..bf8c6906f --- /dev/null +++ b/vllm_omni/deploy/lance.yaml @@ -0,0 +1,29 @@ +# Lance-3B (ByteDance) deploy — single-stage, CUDA. Target: 1x NVIDIA A100 80GB. +# +# Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion). The single-stage +# DiT stage is self-contained: it owns the Qwen2-MoT LLM, the Qwen2.5-VL ViT, +# the Wan2.2 VAE and the tokenizer (text2img / image_edit / x2t_image). +# +# Lance-3B BF16 footprint: ~7GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE — comfortably +# within one A100 80GB. The video checkpoint (Lance_3B_Video) is a follow-up. +# +# Select via: +# vllm serve bytedance-research/Lance --omni \ +# --deploy-config vllm_omni/deploy/lance.yaml +# or programmatically: +# Omni(model="bytedance-research/Lance", +# deploy_config_path="vllm_omni/deploy/lance.yaml") + +pipeline: lance +async_chunk: false + +stages: + - stage_id: 0 + max_num_batched_tokens: 32768 + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + devices: "0" + default_sampling_params: + seed: 42 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c45ee63e7..3ba8b24f2 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -790,6 +790,21 @@ class OmniDiffusionConfig: self.supports_multimodal_inputs = metadata.supports_multimodal_inputs self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs + @staticmethod + def _looks_like_lance_subfolder(model: str | None) -> bool: + """Return True when ``--model`` points at a Lance per-component subfolder. + + Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and + ``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may + reasonably hand the AR-style sub-checkpoint path directly. The + ``LancePipeline`` constructor knows to walk up to the repo root from + either subfolder name. + """ + if not model: + return False + base = os.path.basename(str(model).rstrip("/")) + return base in {"Lance_3B", "Lance_3B_Video"} + def enrich_config(self) -> None: """Load model metadata from HuggingFace and populate config fields. @@ -843,6 +858,15 @@ class OmniDiffusionConfig: else: cfg = get_hf_file_to_dict("config.json", self.model) if cfg is None: + # Lance ships its top-level config.json one directory above + # the per-checkpoint subfolders (``Lance_3B/`` or + # ``Lance_3B_Video/``). Try to recover that case before + # raising. + if self._looks_like_lance_subfolder(self.model): + self.model_class_name = "LancePipeline" + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() + return raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") self.set_tf_model_config(TransformerConfig.from_dict(cfg)) @@ -853,6 +877,19 @@ class OmniDiffusionConfig: self.model_class_name = "BagelPipeline" self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() + elif ( + model_type == "lance" + or "LancePipeline" in architectures + or cfg.get("model_name") == "Lance" + or self._looks_like_lance_subfolder(self.model) + ): + # Lance ships a non-HF top-level config.json (model_name only) + # plus per-component subfolders; resolve to the Lance pipeline. + # Also accept --model pointing directly at the ``Lance_3B`` or + # ``Lance_3B_Video`` subfolder by walking up to the repo root. + self.model_class_name = "LancePipeline" + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() elif model_type == "neo_chat": self.model_class_name = "SenseNovaU1Pipeline" self.tf_model_config = TransformerConfig() diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index d33b26f60..22e7506e7 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -108,7 +108,16 @@ class BagelRotaryEmbedding(nn.Module): linear, dynamic-NTK, YaRN, …), we delegate the ``inv_freq`` / ``attention_scaling`` computation to HF's ``ROPE_INIT_FUNCTIONS`` so that the frequency basis and scaling factor are identical to the - original checkpoint. This module has no learnable parameters. + original checkpoint. + + For Qwen2.5-VL-style multimodal RoPE (``rope_scaling.rope_type == "mrope"``) + the ``inv_freq`` basis is the standard default-rope one; the difference + is that position ids are 3-D ``(t, h, w)`` per token and the + ``mrope_section`` describes how the head dimension is split across axes. + This module accepts either 2-D scalar position ids ``(B, S)`` or 3-D + multimodal position ids ``(B, 3, S)`` and dispatches accordingly so the + same module works for both BAGEL (1-D rope) and Lance (Qwen2.5-VL mrope). + This module has no learnable parameters. """ def __init__(self, config): @@ -122,10 +131,17 @@ class BagelRotaryEmbedding(nn.Module): rope_type = ( rope_scaling.get("rope_type") or rope_scaling.get("type") or rope_parameters.get("rope_type") or "default" ) + # Cache mrope_section for the forward-time section-split. + self._mrope_section: list[int] | None = None + if rope_type == "mrope": + section = rope_scaling.get("mrope_section") or rope_parameters.get("mrope_section") + if section is None: + raise ValueError("rope_scaling.rope_type == 'mrope' requires 'mrope_section'.") + self._mrope_section = list(section) - if rope_type == "default": - # transformers>=5.0 removed the 'default' entry from - # ROPE_INIT_FUNCTIONS; use the plain sinusoidal formula. + if rope_type in ("default", "mrope"): + # mrope shares the default sinusoidal frequency basis; the + # multimodal split happens at forward time, not at init. rope_theta = ( getattr(config, "rope_theta", None) or rope_parameters.get("rope_theta") @@ -149,11 +165,36 @@ class BagelRotaryEmbedding(nn.Module): Args: x: Input tensor (only used for dtype inference). - position_ids: Position indices, shape (batch_size, seq_len). + position_ids: Either 2-D scalar ``(batch_size, seq_len)`` for plain + 1-D RoPE, or 3-D multimodal ``(batch_size, 3, seq_len)`` for + Qwen2.5-VL-style mRoPE. The latter is auto-detected from + ``position_ids.ndim``. Returns: cos, sin: Rotary embeddings, each of shape (batch_size, seq_len, dim). """ + if position_ids.ndim == 3 and self._mrope_section is not None: + # multimodal path: position_ids is (B, 3, S) with rows = (t, h, w). + # Compute per-axis frequencies, then assemble the per-section + # rotary basis matching Qwen2-VL's ``apply_multimodal_rotary_pos_emb``. + B = position_ids.shape[0] + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(B, 3, -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3) + # ``freqs`` is (B, 3, S, head_dim/2); double along the last axis + # to get the full ``head_dim`` rotary basis. + emb = torch.cat((freqs, freqs), dim=-1) # (B, 3, S, head_dim) + cos_per_axis = emb.cos() * self.attention_scaling + sin_per_axis = emb.sin() * self.attention_scaling + # ``mrope_section`` (e.g. [16, 24, 24] for Qwen2.5-VL) sums to + # head_dim/2; doubled it sums to head_dim and cycles axis = i % 3. + sec_full = self._mrope_section * 2 + cos_split = cos_per_axis.split(sec_full, dim=-1) + sin_split = sin_per_axis.split(sec_full, dim=-1) + cos = torch.cat([c[:, i % 3] for i, c in enumerate(cos_split)], dim=-1) + sin = torch.cat([s[:, i % 3] for i, s in enumerate(sin_split)], dim=-1) + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) @@ -1982,7 +2023,9 @@ class Bagel(nn.Module): "num_branches": num_branches, "seq_len": seq_len, "batched_query_lens": packed_seqlens.repeat(num_branches), - "batched_position_ids": torch.cat(branches_pid), + # Concatenate along the sequence (last) dim so this works for both + # scalar position ids ``(S,)`` and 3-D mRoPE ``(3, S)`` layouts. + "batched_position_ids": torch.cat(branches_pid, dim=-1), "batched_kv_lens": torch.cat(branches_kvl), "batched_query_indexes": torch.cat( [qi + merged_offsets[b_idx] for b_idx, qi in enumerate(branches_qi)] diff --git a/vllm_omni/diffusion/models/lance/__init__.py b/vllm_omni/diffusion/models/lance/__init__.py new file mode 100644 index 000000000..ca5e1ce66 --- /dev/null +++ b/vllm_omni/diffusion/models/lance/__init__.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance (ByteDance) diffusion model components. + +Lance is a unified autoregressive + diffusion multimodal model on a +Qwen2.5-VL-3B backbone. Architecturally it is the BAGEL family (ByteDance +Mixture-of-Transformers): the released ``Lance_3B`` checkpoint uses the exact +same ``*_moe_gen`` MoT weight layout as BAGEL, plus ``vae2llm`` / ``llm2vae`` / +``time_embedder`` / ``latent_pos_embed`` connectors. The deltas vs BAGEL are: + + * backbone is Qwen2.5-VL (mRoPE) instead of Qwen2, + * understanding ViT is Qwen2.5-VL vision (not SigLIP), loaded from the base + ``Qwen/Qwen2.5-VL-3B-Instruct`` rather than from the Lance checkpoint, + * VAE is Wan2.2 (reused from the vLLM-Omni WAN path) instead of the BAGEL AE, + * video path adds 3D latent position embeddings (follow-up; this module + implements the image path first). + +Because Lance is BAGEL-lineage, the transformer core is reused verbatim from +``vllm_omni.diffusion.models.bagel.bagel_transformer`` and only the pipeline +wiring (ViT / VAE / checkpoint layout) is specialized here. +""" + +from vllm_omni.diffusion.models.lance.pipeline_lance import ( + LancePipeline, + get_lance_post_process_func, +) + +__all__ = [ + "LancePipeline", + "get_lance_post_process_func", +] diff --git a/vllm_omni/diffusion/models/lance/lance_transformer.py b/vllm_omni/diffusion/models/lance/lance_transformer.py new file mode 100644 index 000000000..2000508f6 --- /dev/null +++ b/vllm_omni/diffusion/models/lance/lance_transformer.py @@ -0,0 +1,944 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance transformer pieces. + +The Lance LLM is BAGEL's Qwen2-MoT transformer verbatim — the released +``Lance_3B`` checkpoint uses the identical ``*_moe_gen`` / ``q_norm`` / +``vae2llm`` / ``llm2vae`` / ``time_embedder`` / ``latent_pos_embed`` layout, so +``Bagel`` / ``Qwen2MoTForCausalLM`` / ``Qwen2MoTConfig`` / ``NaiveCache`` are +re-exported unchanged. Only the understanding ViT differs (Qwen2.5-VL vision +instead of SigLIP) and the video path adds a 3D latent position embedding. + +This module also provides :class:`LanceBagel`, a thin :class:`Bagel` subclass +that overrides the two ViT entry points to consume Qwen2.5-VL's packed +``pixel_values`` + ``image_grid_thw`` layout directly (BAGEL itself assumes +SigLIP-style ``(C, H, W)`` tensors that get patchified inside the model). + +NOTE ON mRoPE: Lance's backbone is Qwen2.5-VL and its understanding / +video paths use multimodal RoPE (``mrope_section=[16,24,24]``). BAGEL's +``BagelRotaryEmbedding`` is plain 1-D RoPE on scalar position ids. For the +text2img generation path Lance assigns scalar positions (the gen latent block +shares a single rope position, same as BAGEL), so the reused rotary is +correct there. Full mRoPE for the x2t / video understanding path is a +follow-up — see ``LancePositionEmbedding3D``. +""" + +from __future__ import annotations + +import numpy as np +import torch +from torch import nn + +# Lance LLM == BAGEL Qwen2-MoT: reuse verbatim (checkpoint weight names match). +from vllm_omni.diffusion.models.bagel.bagel_transformer import ( # noqa: F401 + Bagel, + BaseNavitOutputWithPast, + MLPconnector, + NaiveCache, + PackedAttentionMoT, + PositionEmbedding, + Qwen2MoTConfig, + Qwen2MoTDecoderLayer, + Qwen2MoTForCausalLM, + Qwen2MoTModel, + TimestepEmbedder, + get_1d_sincos_pos_embed_from_grid, + patchify, +) + +__all__ = [ + "Bagel", + "BaseNavitOutputWithPast", + "LanceBagel", + "LanceIdentityConnector", + "LancePositionEmbedding3D", + "LanceQwen2_5_VLNaViTWrapper", + "LanceZeroVitPosEmbed", + "MLPconnector", + "NaiveCache", + "PackedAttentionMoT", + "PositionEmbedding", + "Qwen2MoTConfig", + "Qwen2MoTDecoderLayer", + "Qwen2MoTForCausalLM", + "Qwen2MoTModel", + "TimestepEmbedder", + "patchify", +] + +# mRoPE temporal scaling constants for Lance (Qwen2.5-VL backbone). +# Upstream Lance computes vision rope positions as +# ``t_index = frame_idx * tokens_per_second * second_per_grid_t`` so adjacent +# latent frames sit at well-separated rope coordinates. Lance bundles a +# ``Qwen2.5-VL-ViT/config.json`` with ``tokens_per_second: 2`` and uses +# ``second_per_grid_ts = 1.0`` in its inference script. +LANCE_TOKENS_PER_SECOND = 2 +LANCE_SECONDS_PER_GRID = 1.0 + + +def get_3d_sincos_pos_embed_from_grid(embed_dim: int, grid: np.ndarray) -> np.ndarray: + """3-D sin-cos positional embedding (t, h, w), matching the upstream Lance + ``modeling/lance/modeling_utils.py`` dimension split exactly.""" + assert embed_dim % 2 == 0, "Embedding dimension must be even for 3D embeddings" + d = embed_dim // 3 + d = d if d % 2 == 0 else d - 1 + dim_t, dim_h = d, d + dim_w = embed_dim - 2 * d + assert dim_w % 2 == 0 + emb_t = get_1d_sincos_pos_embed_from_grid(dim_t, grid[0]) + emb_h = get_1d_sincos_pos_embed_from_grid(dim_h, grid[1]) + emb_w = get_1d_sincos_pos_embed_from_grid(dim_w, grid[2]) + return np.concatenate([emb_t, emb_h, emb_w], axis=1) + + +def get_3d_sincos_pos_embed(embed_dim: int, t: int, h: int, w: int) -> np.ndarray: + grid_t = np.arange(t, dtype=np.float32) + grid_h = np.arange(h, dtype=np.float32) + grid_w = np.arange(w, dtype=np.float32) + tt, hh, ww = np.meshgrid(grid_t, grid_h, grid_w, indexing="ij") + grid = np.stack([tt, hh, ww], axis=0) + return get_3d_sincos_pos_embed_from_grid(embed_dim, grid) + + +class LancePositionEmbedding3D(nn.Module): + """Frozen 3-D sin-cos latent position embedding for the video path. + + BAGEL only ships a 2-D ``PositionEmbedding`` (image latents). Lance's + ``Lance_3B_Video`` checkpoint adds a temporal axis; this mirrors upstream + ``modeling/lance/modeling_utils.py::PositionEmbedding3D``. The image path + uses ``t=1`` and is numerically equivalent to the 2-D embedding. + """ + + def __init__(self, max_num_frames: int, max_num_patch_per_side: int, hidden_size: int): + super().__init__() + self.max_num_frames = max_num_frames + self.max_num_patch_per_side = max_num_patch_per_side + self.hidden_size = hidden_size + n = max_num_frames * max_num_patch_per_side * max_num_patch_per_side + self.pos_embed = nn.Parameter(torch.zeros(n, hidden_size), requires_grad=False) + self._init_weights() + + def _init_weights(self) -> None: + pe = get_3d_sincos_pos_embed( + self.hidden_size, + self.max_num_frames, + self.max_num_patch_per_side, + self.max_num_patch_per_side, + ) + self.pos_embed.data.copy_(torch.from_numpy(pe).float()) + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: + return self.pos_embed[position_ids] + + +class LanceIdentityConnector(nn.Module): + """No-op connector for Lance. + + BAGEL's ``connector`` projects the ViT hidden size to the LLM hidden size. + Qwen2.5-VL's vision tower (which Lance uses) already projects to the LLM + hidden size internally via ``merger`` (``out_hidden_size = hidden_size``), + and the released Lance safetensors carry no ``connector.*`` weights. We + therefore plug in an ``Identity`` connector so ``forward_cache_update_vit`` + keeps its existing call site without a separate code path. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: # noqa: D401 + return x + + +class LanceZeroVitPosEmbed(nn.Module): + """No-op positional embedding for Lance's ViT tokens. + + BAGEL adds an extra 2-D sin-cos ``vit_pos_embed`` on top of the ViT output. + Qwen2.5-VL's vision tower already carries its own (rotary) positional + encoding, and the released Lance safetensors carry no + ``vit_pos_embed.*`` weights. This module returns a broadcast-friendly + zero so the addition in ``forward_cache_update_vit`` is a no-op without + requiring a code-path branch. + """ + + def forward(self, position_ids: torch.Tensor) -> torch.Tensor: # noqa: D401 + return torch.zeros((), device=position_ids.device) + + +class LanceQwen2_5_VLNaViTWrapper(nn.Module): + """Packed (NaViT-style) wrapper around the Qwen2.5-VL vision tower. + + Bridges BAGEL's ``vit(packed_pixel_values, ...) -> [num_tokens, vit_hidden]`` + surface to the HF ``Qwen2_5_VisionTransformerPretrainedModel`` which + consumes ``(hidden_states, grid_thw)``. The packed call additionally needs + a per-image ``image_grid_thw`` so non-square images (and the + spatial-merge token count) line up — :class:`LanceBagel` stashes the grid + on the wrapper before invoking the ViT. + """ + + def __init__(self, vision_model: nn.Module, spatial_merge_size: int = 2): + super().__init__() + # Accept either the full Qwen2_5_VLForConditionalGeneration.visual or + # the bare vision transformer. + self.vision_model = getattr(vision_model, "visual", vision_model) + self.spatial_merge_size = spatial_merge_size + # Set by ``LanceBagel.forward_cache_update_vit`` before each call so + # the wrapper can pass true per-image (T, H, W) to the HF ViT. + self._pending_grid_thw: torch.Tensor | None = None + + @property + def config(self): # parity with SiglipNaViTWrapper.vision_model.config access + return self.vision_model.config + + def set_pending_grid_thw(self, grid_thw: torch.Tensor) -> None: + self._pending_grid_thw = grid_thw + + def forward( + self, + packed_pixel_values: torch.Tensor, + packed_flattened_position_ids: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + ) -> torch.Tensor: + if self._pending_grid_thw is not None: + grid_thw_t = self._pending_grid_thw.to(packed_pixel_values.device) + self._pending_grid_thw = None + else: + # Fallback for square images when grid_thw was not pre-stashed. + cu = cu_seqlens.tolist() + grid_thw = [] + for i in range(len(cu) - 1): + n = cu[i + 1] - cu[i] + side = int(round(float(n) ** 0.5)) + grid_thw.append([1, side, side]) + grid_thw_t = torch.tensor(grid_thw, dtype=torch.long, device=packed_pixel_values.device) + + hidden = packed_pixel_values + if hasattr(self.vision_model, "dtype"): + hidden = hidden.to(self.vision_model.dtype) + out = self.vision_model(hidden_states=hidden, grid_thw=grid_thw_t) + # Qwen2.5-VL returns ``BaseModelOutputWithPooling(last_hidden_state=..., + # pooler_output=merged)``; the merged sequence (post spatial-merger, + # shape ``[sum_tokens_after_merge, out_hidden_size]``) is what the LLM + # consumes. + if hasattr(out, "pooler_output") and out.pooler_output is not None: + return out.pooler_output + if isinstance(out, tuple): + return out[0] + return out + + +class LanceBagel(Bagel): + """Bagel subclass with Lance-specific ViT handling. + + The released Lance checkpoint pairs BAGEL's Qwen2-MoT trunk with the + Qwen2.5-VL vision tower (whose ``merger`` already projects to LLM + ``hidden_size`` and which carries its own rotary positional encoding). + Two BAGEL assumptions therefore break for Lance: + + 1. ``prepare_vit_images`` calls ``transforms(image) -> (C, H, W)`` and + then ``patchify(...)``; Lance's image processor returns + ``(num_patches_flat, patch_features) + image_grid_thw`` already. + 2. ``forward_cache_update_vit`` adds a ``connector`` projection plus a + 2-D ``vit_pos_embed`` that has no checkpoint weights and would + also double-count the ViT's own positional encoding. + + This subclass overrides exactly those two methods. Everything else — + LLM trunk, VAE flow, generation loop, ``forward_cache_update_text`` / + ``forward_cache_update_vae`` — is reused unchanged. + """ + + @staticmethod + def _qwen_vl_processor_call(processor, image): + """Invoke the Qwen2-VL image processor and return ``(pixel_values, grid_thw)``.""" + proc_out = processor(images=image, return_tensors="pt") + pixel_values = proc_out["pixel_values"] + grid_thw = proc_out["image_grid_thw"] + return pixel_values, grid_thw + + # ------------------------------------------------------------------ # + # 3-D mRoPE plumbing + # ------------------------------------------------------------------ # + @staticmethod + def _mrope_broadcast(position_ids: torch.Tensor) -> torch.Tensor: + """Convert scalar ``(S,)`` position ids to mrope ``(3, S)`` by + broadcasting the scalar value to all three (t, h, w) axes. + + Lance/Qwen2.5-VL's mRoPE expects 3-D positions per token; for text + tokens (and image-edit context tokens) all three axes share the + same scalar position id. Vision-generation paths override this with + per-token ``(t, h, w)`` indices. + """ + if position_ids is None: + return position_ids + if position_ids.ndim == 1: + return position_ids.unsqueeze(0).expand(3, -1).contiguous() + return position_ids + + def prepare_prompts(self, *args, **kwargs): + gen_input, newlens, new_rope = super().prepare_prompts(*args, **kwargs) + if "packed_text_position_ids" in gen_input: + gen_input["packed_text_position_ids"] = self._mrope_broadcast(gen_input["packed_text_position_ids"]) + return gen_input, newlens, new_rope + + def prepare_vae_latent(self, *args, **kwargs): + gen_input = super().prepare_vae_latent(*args, **kwargs) + if "packed_position_ids" in gen_input: + gen_input["packed_position_ids"] = self._mrope_broadcast(gen_input["packed_position_ids"]) + return gen_input + + def prepare_vae_latent_cfg(self, *args, **kwargs): + gen_input = super().prepare_vae_latent_cfg(*args, **kwargs) + if "cfg_packed_position_ids" in gen_input: + gen_input["cfg_packed_position_ids"] = self._mrope_broadcast(gen_input["cfg_packed_position_ids"]) + return gen_input + + def prepare_start_tokens(self, *args, **kwargs): + gen_input = super().prepare_start_tokens(*args, **kwargs) + if isinstance(gen_input, dict): + for k in ( + "packed_query_position_ids", + "packed_text_position_ids", + "packed_position_ids", + ): + if k in gen_input and torch.is_tensor(gen_input[k]): + gen_input[k] = self._mrope_broadcast(gen_input[k]) + return gen_input + + def prepare_vit_videos(self, curr_kvlens, curr_rope, videos, new_token_ids): + """Multi-frame ViT prefill for the ``x2t_video`` understanding path. + + ``videos`` is a list of per-request video tensors / numpy arrays of + shape ``(T, H, W, 3)``. We call the Qwen2-VL **video** processor + (stored on ``self._lance_video_processor``) which produces + ``pixel_values_videos`` already in packed + ``(num_patches_flat_3d, patch_features)`` layout plus a 3-D + ``video_grid_thw`` (``[T_lat, H_patches, W_patches]``). We forward the + same pixel-values tensor to the Qwen2.5-VL ViT (wrapped by + :class:`LanceQwen2_5_VLNaViTWrapper`), set the wrapper's pending + grid, and place the post-merger tokens in the LLM sequence with + proper 3-D mRoPE positions ``(P + t, P + h, P + w)`` over the + ``T_lat × H_lat × W_lat`` post-merge grid. + """ + processor = getattr(self, "_lance_video_processor", None) + if processor is None: + raise RuntimeError( + "LanceBagel.prepare_vit_videos requires the pipeline to set " + "``bagel._lance_video_processor`` (a Qwen2-VL-compatible video processor)." + ) + + packed_vit_token_indexes: list[int] = [] + vit_token_seqlens: list[int] = [] + packed_vit_tokens: list[torch.Tensor] = [] + packed_vit_position_ids: list[torch.Tensor] = [] + packed_text_ids: list[int] = [] + packed_text_indexes: list[int] = [] + packed_seqlens: list[int] = [] + packed_indexes: list[int] = [] + packed_key_value_indexes: list[int] = [] + grid_thw_list: list[torch.Tensor] = [] + per_axis_pos: list[tuple[int, int, int]] = [] + + merge_size = int(getattr(self.vit_model, "spatial_merge_size", 2)) + merge_factor = merge_size * merge_size + + _curr = curr = 0 + newlens: list[int] = [] + new_rope: list[int] = [] + for video, curr_kvlen, curr_position_id in zip(videos, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_text_ids.append(new_token_ids["start_of_image"]) + packed_text_indexes.append(_curr) + packed_indexes.append(curr) + curr += 1 + _curr += 1 + per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id)) + + proc_out = processor(videos=video, return_tensors="pt") + # Qwen2-VL video processor: pixel_values_videos is the same packed + # (num_patches_flat, patch_features) layout as the image processor + # but with a temporal axis ``T_lat`` in video_grid_thw. + pixel_values = proc_out["pixel_values_videos"] + grid_thw = proc_out["video_grid_thw"] + T_lat, H, W = (int(v) for v in grid_thw[0].tolist()) + num_patches_pre_merge = T_lat * H * W + assert num_patches_pre_merge == pixel_values.shape[0], ( + f"pixel_values rows ({pixel_values.shape[0]}) != T_lat*H*W ({num_patches_pre_merge}); " + f"video_grid_thw={grid_thw.tolist()}" + ) + num_vit_tokens = num_patches_pre_merge // merge_factor + + packed_vit_tokens.append(pixel_values) + vit_token_seqlens.append(num_patches_pre_merge) + packed_vit_position_ids.append(torch.arange(num_patches_pre_merge, dtype=torch.long)) + grid_thw_list.append(grid_thw[0].to(torch.long)) + + packed_vit_token_indexes.extend(range(_curr, _curr + num_vit_tokens)) + packed_indexes.extend(range(curr, curr + num_vit_tokens)) + curr += num_vit_tokens + _curr += num_vit_tokens + h_merged = H // merge_size + w_merged = W // merge_size + t_scale = LANCE_TOKENS_PER_SECOND * LANCE_SECONDS_PER_GRID + for ti in range(T_lat): + for hi in range(h_merged): + for wi in range(w_merged): + per_axis_pos.append( + ( + curr_position_id + int(ti * t_scale), + curr_position_id + hi, + curr_position_id + wi, + ) + ) + + packed_text_ids.append(new_token_ids["end_of_image"]) + packed_text_indexes.append(_curr) + packed_indexes.append(curr) + curr += 1 + _curr += 1 + end_p = curr_position_id + max(int((T_lat - 1) * t_scale), h_merged - 1, w_merged - 1) + 1 + per_axis_pos.append((end_p, end_p, end_p)) + + packed_seqlens.append(num_vit_tokens + 2) + newlens.append(curr_kvlen + num_vit_tokens + 2) + new_rope.append(end_p + 1) + + pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long) + pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long) + pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long) + packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0) + + generation_input = { + "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long), + "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long), + "vit_token_seqlens": torch.tensor(vit_token_seqlens, dtype=torch.int), + "packed_vit_tokens": torch.cat(packed_vit_tokens, dim=0), + "packed_vit_position_ids": torch.cat(packed_vit_position_ids, dim=0), + "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long), + "packed_position_ids": packed_position_ids_3d, + "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int), + "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + "_lance_grid_thw": torch.stack(grid_thw_list, dim=0), + } + return generation_input, newlens, new_rope + + def prepare_vit_images(self, curr_kvlens, curr_rope, images, transforms, new_token_ids): + # ``transforms`` is set up by ``BagelPipeline.forward`` to be + # ``processor(images=img, return_tensors='pt').pixel_values[0]``, which + # discards the ``image_grid_thw``. For Lance we need both, so we ignore + # the supplied lambda and re-call the processor here directly. Lance's + # pipeline always sets ``self.image_processor`` to ``Qwen2VLImageProcessor``. + if not hasattr(self, "_lance_image_processor"): + raise RuntimeError( + "LanceBagel.prepare_vit_images requires the pipeline to set " + "``bagel._lance_image_processor`` (a Qwen2-VL-compatible processor)." + ) + processor = self._lance_image_processor + + packed_vit_token_indexes: list[int] = [] + vit_token_seqlens: list[int] = [] + packed_vit_tokens: list[torch.Tensor] = [] + packed_vit_position_ids: list[torch.Tensor] = [] + packed_text_ids: list[int] = [] + packed_text_indexes: list[int] = [] + packed_seqlens: list[int] = [] + packed_indexes: list[int] = [] + packed_key_value_indexes: list[int] = [] + grid_thw_list: list[torch.Tensor] = [] + + merge_size = int(getattr(self.vit_model, "spatial_merge_size", 2)) + merge_factor = merge_size * merge_size + + _curr = curr = 0 + newlens: list[int] = [] + new_rope: list[int] = [] + # 3-D position ids per query token (mRoPE) — text framing tokens use + # scalar broadcast, image patches use per-token ``(P + t, P + h_merged, + # P + w_merged)`` so the Qwen2.5-VL backbone sees the spatial layout + # along the (h, w) axes of mRoPE. + per_axis_pos: list[tuple[int, int, int]] = [] + for image, curr_kvlen, curr_position_id in zip(images, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_text_ids.append(new_token_ids["start_of_image"]) + packed_text_indexes.append(_curr) + packed_indexes.append(curr) + curr += 1 + _curr += 1 + per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id)) + + pixel_values, grid_thw = self._qwen_vl_processor_call(processor, image) + # pixel_values: (num_patches_flat, patch_features). grid_thw: (1, 3) with (T, H, W) in patch units. + T, H, W = (int(v) for v in grid_thw[0].tolist()) + num_patches_pre_merge = T * H * W + assert num_patches_pre_merge == pixel_values.shape[0], ( + f"pixel_values rows ({pixel_values.shape[0]}) != T*H*W ({num_patches_pre_merge}); " + f"image_grid_thw={grid_thw.tolist()}" + ) + num_img_tokens = num_patches_pre_merge // merge_factor + + packed_vit_tokens.append(pixel_values) + vit_token_seqlens.append(num_patches_pre_merge) + packed_vit_position_ids.append(torch.arange(num_patches_pre_merge, dtype=torch.long)) + grid_thw_list.append(grid_thw[0].to(torch.long)) + + packed_vit_token_indexes.extend(range(_curr, _curr + num_img_tokens)) + packed_indexes.extend(range(curr, curr + num_img_tokens)) + curr += num_img_tokens + _curr += num_img_tokens + # 3-D mRoPE positions for image patches: scan in row-major over the + # post-merge grid (H/merge × W/merge). Temporal axis stays at the + # block's scalar position (single still image). + h_merged = H // merge_size + w_merged = W // merge_size + for hi in range(h_merged): + for wi in range(w_merged): + per_axis_pos.append((curr_position_id, curr_position_id + hi, curr_position_id + wi)) + + packed_text_ids.append(new_token_ids["end_of_image"]) + packed_text_indexes.append(_curr) + packed_indexes.append(curr) + curr += 1 + _curr += 1 + end_p = curr_position_id + max(h_merged - 1, w_merged - 1) + 1 + per_axis_pos.append((end_p, end_p, end_p)) + + packed_seqlens.append(num_img_tokens + 2) + newlens.append(curr_kvlen + num_img_tokens + 2) + new_rope.append(end_p + 1) + + pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long) + pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long) + pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long) + packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0) + + generation_input = { + "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long), + "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long), + "vit_token_seqlens": torch.tensor(vit_token_seqlens, dtype=torch.int), + "packed_vit_tokens": torch.cat(packed_vit_tokens, dim=0), + "packed_vit_position_ids": torch.cat(packed_vit_position_ids, dim=0), + "packed_vit_token_indexes": torch.tensor(packed_vit_token_indexes, dtype=torch.long), + "packed_position_ids": packed_position_ids_3d, + "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int), + "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + # Lance-extra: per-image (T, H, W) so the ViT wrapper can pass the + # true grid to the HF tower (non-square images would otherwise be + # reconstructed incorrectly as ``sqrt(n)`` squares). + "_lance_grid_thw": torch.stack(grid_thw_list, dim=0), + } + return generation_input, newlens, new_rope + + def forward_cache_update_vit(self, *args, **kwargs): + # Pluck the Lance-specific extra and stash on the wrapper for the + # upcoming ViT call. Everything else delegates to the parent so we + # don't have to copy 60 lines of LLM-call boilerplate. + grid_thw = kwargs.pop("_lance_grid_thw", None) + if grid_thw is not None and hasattr(self.vit_model, "set_pending_grid_thw"): + self.vit_model.set_pending_grid_thw(grid_thw) + return super().forward_cache_update_vit(*args, **kwargs) + + def prepare_video_latent(self, curr_kvlens, curr_rope, video_shapes, new_token_ids): + """3-D analogue of :meth:`prepare_vae_latent`. + + ``video_shapes`` is a list of ``(T, H, W)`` per request (RGB pixel + space). We package one packed-init-noise tensor over ``T_lat × H_lat × + W_lat`` latent tokens per video, plus 1-D indices into the 3-D position + embedding table maintained by :class:`LancePositionEmbedding3D` + (`bagel.latent_pos_embed`). Latent geometry: + + - spatial: ``H_lat = H // latent_downsample`` (``=16`` for Lance) + - temporal: ``T_lat = (T - 1) // downsample_temporal + 1`` (``=4`` for Wan2.2) + - channels: ``latent_channel = 48`` + + Position ids are flattened ``t * max_per_side² + h * max_per_side + w`` + so they index directly into the ``(max_num_frames * max_per_side², + hidden_size)`` table. + """ + downsample_t = int(getattr(self.config.vae_config, "downsample_temporal", 4)) + max_per_side = int(self.max_latent_size) + + packed_text_ids, packed_text_indexes = [], [] + packed_vae_position_ids, packed_vae_token_indexes, packed_init_noises = [], [], [] + packed_seqlens, packed_indexes = [], [] + packed_key_value_indexes = [] + # 3-D position ids per token (mRoPE). ``per_axis_pos[i]`` is the list + # of (t, h, w) for the i-th query token; we stack into ``(3, S)`` at the + # end. Text framing tokens broadcast scalar to all axes; video latent + # tokens carry per-token ``(P + t, P + h, P + w)``. + per_axis_pos: list[tuple[int, int, int]] = [] + + query_curr = curr = 0 + for (T, H, W), curr_kvlen, curr_position_id in zip(video_shapes, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_text_ids.append(new_token_ids["start_of_image"]) + packed_text_indexes.append(query_curr) + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id)) + + h = H // self.latent_downsample + w = W // self.latent_downsample + t = (T - 1) // downsample_t + 1 + num_video_tokens = t * h * w + + # 1-D position ids into the 3-D table (frame-major, then row, then col). + tt, hh, ww = torch.meshgrid(torch.arange(t), torch.arange(h), torch.arange(w), indexing="ij") + tt_flat = tt.flatten().tolist() + hh_flat = hh.flatten().tolist() + ww_flat = ww.flatten().tolist() + vae_position_ids = (tt * (max_per_side * max_per_side) + hh * max_per_side + ww).flatten() + packed_vae_position_ids.append(vae_position_ids) + + packed_init_noises.append(torch.randn(num_video_tokens, self.latent_channel * self.latent_patch_size**2)) + + packed_vae_token_indexes.extend(range(query_curr, query_curr + num_video_tokens)) + packed_indexes.extend(range(curr, curr + num_video_tokens)) + curr += num_video_tokens + query_curr += num_video_tokens + # Per-token 3-D mRoPE positions, matching upstream Lance's + # ``get_rope_index``: the temporal axis is amplified by + # ``tokens_per_second * second_per_grid_t`` (see Qwen2.5-VL docs) + # so neighbouring latent frames are well-separated in rope space. + # ``second_per_grid_t = 1.0`` matches upstream's default; + # ``tokens_per_second = 2`` is Lance's vit config. + t_scale = LANCE_TOKENS_PER_SECOND * LANCE_SECONDS_PER_GRID + for ti, hi, wi in zip(tt_flat, hh_flat, ww_flat): + per_axis_pos.append( + ( + curr_position_id + int(ti * t_scale), + curr_position_id + hi, + curr_position_id + wi, + ) + ) + + packed_text_ids.append(new_token_ids["end_of_image"]) + packed_text_indexes.append(query_curr) + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + # end_of_image sits "past" the latent block; bump along all axes by + # max of the latent block size so subsequent text resumes cleanly. + end_p = curr_position_id + max(t - 1, h - 1, w - 1) + 1 + per_axis_pos.append((end_p, end_p, end_p)) + + packed_seqlens.append(num_video_tokens + 2) + + # Stack into (3, total_query_tokens) for mRoPE. + pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long) + pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long) + pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long) + packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0) + + generation_input = { + "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long), + "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long), + "packed_init_noises": torch.cat(packed_init_noises, dim=0), + "packed_vae_position_ids": torch.cat(packed_vae_position_ids, dim=0), + "packed_vae_token_indexes": torch.tensor(packed_vae_token_indexes, dtype=torch.long), + "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int), + "packed_position_ids": packed_position_ids_3d, + "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + } + return generation_input + + def prepare_video_latent_cfg(self, curr_kvlens, curr_rope, video_shapes): + """3-D analogue of :meth:`prepare_vae_latent_cfg` (CFG side). + + Mirrors :meth:`prepare_video_latent`'s mRoPE 3-D position layout + (text frame ⇒ scalar; video latent block ⇒ per-token ``(P+t, P+h, P+w)``). + """ + downsample_t = int(getattr(self.config.vae_config, "downsample_temporal", 4)) + + packed_indexes, packed_key_value_indexes = [], [] + per_axis_pos: list[tuple[int, int, int]] = [] + query_curr = curr = 0 + for (T, H, W), curr_kvlen, curr_position_id in zip(video_shapes, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id)) + + h = H // self.latent_downsample + w = W // self.latent_downsample + t = (T - 1) // downsample_t + 1 + num_video_tokens = t * h * w + packed_indexes.extend(range(curr, curr + num_video_tokens)) + curr += num_video_tokens + query_curr += num_video_tokens + tt, hh, ww = torch.meshgrid(torch.arange(t), torch.arange(h), torch.arange(w), indexing="ij") + for ti, hi, wi in zip(tt.flatten().tolist(), hh.flatten().tolist(), ww.flatten().tolist()): + per_axis_pos.append((curr_position_id + ti, curr_position_id + hi, curr_position_id + wi)) + + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + end_p = curr_position_id + max(t - 1, h - 1, w - 1) + 1 + per_axis_pos.append((end_p, end_p, end_p)) + + pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long) + pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long) + pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long) + cfg_packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0) + + return { + "cfg_packed_position_ids": cfg_packed_position_ids_3d, + "cfg_packed_query_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "cfg_key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + "cfg_packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + } + + # ------------------------------------------------------------------ # + # Image-edit / video-edit VAE prefill + # ------------------------------------------------------------------ # + def _lance_native_prepare_vae_images( + self, + curr_kvlens, + curr_rope, + images, + transforms, + new_token_ids, + timestep=0, + is_video: bool = False, + ): + """Lance-native ``prepare_vae_images`` with 3-D mRoPE positions. + + Mirrors :meth:`Bagel.prepare_vae_images` but emits per-token mRoPE + ``(curr + t, curr + h, curr + w)`` positions for VAE latent tokens + (BAGEL emits scalar ``curr_position_id`` for all of them, which + triggers a CUDA index-out-of-bounds gather on Lance's Qwen2.5-VL + backbone). For still images ``t = 0``; for video clips ``t`` + ranges over ``T_lat = (T - 1) // downsample_temporal + 1``. + """ + downsample_t = int(getattr(self.config.vae_config, "downsample_temporal", 4)) + max_per_side = int(self.max_latent_size) + + packed_text_ids: list[int] = [] + packed_text_indexes: list[int] = [] + packed_vae_position_ids: list[torch.Tensor] = [] + packed_vae_token_indexes: list[int] = [] + packed_indexes: list[int] = [] + packed_key_value_indexes: list[int] = [] + packed_seqlens: list[int] = [] + patchified_vae_latent_shapes: list[tuple] = [] + vae_image_tensors: list[torch.Tensor] = [] + per_axis_pos: list[tuple[int, int, int]] = [] + + query_curr = curr = 0 + newlens: list[int] = [] + new_rope: list[int] = [] + for image, curr_kvlen, curr_position_id in zip(images, curr_kvlens, curr_rope): + packed_key_value_indexes.extend(range(curr, curr + curr_kvlen)) + curr += curr_kvlen + + packed_text_ids.append(new_token_ids["start_of_image"]) + packed_text_indexes.append(query_curr) + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + per_axis_pos.append((curr_position_id, curr_position_id, curr_position_id)) + + image_tensor = transforms(image) + vae_image_tensors.append(image_tensor) + # ``image_tensor`` is ``(C, H, W)`` for image input or ``(C, T, H, W)`` + # for video input. We derive the *latent* grid from the RGB spatial + # dims and our Wan2.2 VAE downsample factor. + if image_tensor.dim() == 3: + H, W = image_tensor.shape[1:] + T = 1 + elif image_tensor.dim() == 4: + _, T, H, W = image_tensor.shape + else: + raise ValueError(f"vae transforms must return 3-D or 4-D tensor; got {image_tensor.shape}") + h_lat = H // self.latent_downsample + w_lat = W // self.latent_downsample + t_lat = ((T - 1) // downsample_t + 1) if is_video else 1 + patchified_vae_latent_shapes.append((t_lat, h_lat, w_lat) if is_video else (h_lat, w_lat)) + + # 1-D index into the (3-D when video) ``latent_pos_embed`` table. + tt, hh, ww = torch.meshgrid(torch.arange(t_lat), torch.arange(h_lat), torch.arange(w_lat), indexing="ij") + vae_pos_ids = (tt * (max_per_side * max_per_side) + hh * max_per_side + ww).flatten() + packed_vae_position_ids.append(vae_pos_ids) + + num_image_tokens = t_lat * h_lat * w_lat + packed_vae_token_indexes.extend(range(query_curr, query_curr + num_image_tokens)) + packed_indexes.extend(range(curr, curr + num_image_tokens)) + curr += num_image_tokens + query_curr += num_image_tokens + for ti, hi, wi in zip(tt.flatten().tolist(), hh.flatten().tolist(), ww.flatten().tolist()): + per_axis_pos.append((curr_position_id + ti, curr_position_id + hi, curr_position_id + wi)) + + packed_text_ids.append(new_token_ids["end_of_image"]) + packed_text_indexes.append(query_curr) + packed_indexes.append(curr) + curr += 1 + query_curr += 1 + end_p = curr_position_id + max(t_lat - 1, h_lat - 1, w_lat - 1) + 1 + per_axis_pos.append((end_p, end_p, end_p)) + + packed_seqlens.append(num_image_tokens + 2) + newlens.append(curr_kvlen + num_image_tokens + 2) + new_rope.append(end_p + 1) + + image_sizes = [item.shape for item in vae_image_tensors] + max_image_size = [max(item) for item in list(zip(*image_sizes))] + padded_images = torch.zeros(size=(len(vae_image_tensors), *max_image_size)) + for i, image_tensor in enumerate(vae_image_tensors): + if image_tensor.dim() == 3: + padded_images[i, :, : image_tensor.shape[1], : image_tensor.shape[2]] = image_tensor + else: + padded_images[i, :, : image_tensor.shape[1], : image_tensor.shape[2], : image_tensor.shape[3]] = ( + image_tensor + ) + + pos_t = torch.tensor([p[0] for p in per_axis_pos], dtype=torch.long) + pos_h = torch.tensor([p[1] for p in per_axis_pos], dtype=torch.long) + pos_w = torch.tensor([p[2] for p in per_axis_pos], dtype=torch.long) + packed_position_ids_3d = torch.stack([pos_t, pos_h, pos_w], dim=0) + + generation_input = { + "padded_images": padded_images, + "patchified_vae_latent_shapes": patchified_vae_latent_shapes, + "packed_vae_position_ids": torch.cat(packed_vae_position_ids, dim=0), + "packed_timesteps": torch.tensor([timestep]), + "packed_vae_token_indexes": torch.tensor(packed_vae_token_indexes, dtype=torch.long), + "packed_text_ids": torch.tensor(packed_text_ids, dtype=torch.long), + "packed_text_indexes": torch.tensor(packed_text_indexes, dtype=torch.long), + "packed_position_ids": packed_position_ids_3d, + "packed_seqlens": torch.tensor(packed_seqlens, dtype=torch.int), + "packed_indexes": torch.tensor(packed_indexes, dtype=torch.long), + "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), + "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + } + return generation_input, newlens, new_rope + + def prepare_vae_images(self, curr_kvlens, curr_rope, images, transforms, new_token_ids, timestep=0): + """No-op VAE prefill for the x2t / x2t_video path. + + Lance's understanding paths route image / video context through the + Qwen2.5-VL ViT only. BAGEL's pipeline runs ``prepare_vae_images`` + unconditionally for any image input, so this default override + short-circuits it. :meth:`_forward_image_edit` / + :meth:`_forward_video_edit` call :meth:`_lance_native_prepare_vae_images` + directly when they actually need Lance-style VAE prefill. + """ + generation_input = { + "padded_images": torch.empty(0, 3, 0, 0), + "patchified_vae_latent_shapes": [], + "packed_vae_position_ids": torch.empty(0, dtype=torch.long), + "packed_timesteps": torch.tensor([timestep]), + "packed_vae_token_indexes": torch.empty(0, dtype=torch.long), + "packed_text_ids": torch.empty(0, dtype=torch.long), + "packed_text_indexes": torch.empty(0, dtype=torch.long), + "packed_position_ids": torch.empty(0, dtype=torch.long), + "packed_seqlens": torch.empty(0, dtype=torch.int), + "packed_indexes": torch.empty(0, dtype=torch.long), + "packed_key_value_indexes": torch.empty(0, dtype=torch.long), + "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), + } + return generation_input, list(curr_kvlens), list(curr_rope) + + def forward_cache_update_vae( + self, + vae_model, + past_key_values, + padded_images=None, + patchified_vae_latent_shapes=None, + packed_vae_position_ids=None, + packed_timesteps=None, + packed_vae_token_indexes=None, + packed_text_ids=None, + packed_text_indexes=None, + packed_position_ids=None, + packed_seqlens=None, + packed_indexes=None, + key_values_lens=None, + packed_key_value_indexes=None, + ): + """Lance-native VAE prefill that *actually* scatters the encoded + latents into the LLM query sequence. + + :meth:`Bagel.forward_cache_update_vae` in vllm-omni computes + ``packed_latent = vae2llm(...) + time_embed + pos_embed`` and then + passes only ``packed_text_ids`` to the LLM — the VAE embeddings + never enter the query sequence (the LLM builds it from + ``embed_tokens(packed_text_ids)`` which is just the 2 framing + tokens). The mismatch between ``query_lens = [num_vae + 2]`` and + the resulting 2-token sequence is what crashes the gather inside + attention. + + We scatter both pieces explicitly: text framing tokens at + ``packed_text_indexes`` and the VAE latent embeddings at + ``packed_vae_token_indexes``, producing a full-length + ``(sum(packed_seqlens), hidden)`` sequence the LLM can attend + over. Empty prep data (legacy x2t / x2t_video no-op path) is + short-circuited. + """ + if ( + packed_text_ids is None + or packed_text_ids.numel() == 0 + or padded_images is None + or padded_images.numel() == 0 + ): + return past_key_values + + padded_latent = vae_model.encode(padded_images) + p = self.latent_patch_size + packed_latent_list = [] + for latent, shape in zip(padded_latent, patchified_vae_latent_shapes): + if isinstance(shape, tuple) and len(shape) == 3: + t_lat, h_lat, w_lat = shape + latent = latent[:, :t_lat, : h_lat * p, : w_lat * p].reshape( + self.latent_channel, t_lat, h_lat, p, w_lat, p + ) + latent = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, p * p * self.latent_channel) + else: + h_lat, w_lat = shape + latent = latent[:, : h_lat * p, : w_lat * p].reshape(self.latent_channel, h_lat, p, w_lat, p) + latent = torch.einsum("chpwq->hwpqc", latent).reshape(-1, p * p * self.latent_channel) + packed_latent_list.append(latent) + packed_latent = torch.cat(packed_latent_list, dim=0) + + packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids) + packed_timestep_embeds = self.time_embedder(packed_timesteps) + packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + packed_pos_embed + + # ---- Scatter text + VAE embeds into the full query sequence ---- # + total_len = int(packed_seqlens.sum().item()) + packed_text_embed = self.language_model.model.embed_tokens(packed_text_ids) + packed_query_sequence = packed_latent.new_zeros((total_len, packed_latent.shape[-1])) + packed_query_sequence[packed_text_indexes] = packed_text_embed.to(packed_query_sequence.dtype) + packed_query_sequence[packed_vae_token_indexes] = packed_latent.to(packed_query_sequence.dtype) + + extra_inputs = {} + if self.use_moe: + extra_inputs = { + "mode": "gen", + "packed_vae_token_indexes": packed_vae_token_indexes, + "packed_text_indexes": packed_text_indexes, + } + + output = self.language_model.forward( + packed_query_sequence=packed_query_sequence, + query_lens=packed_seqlens, + packed_query_position_ids=packed_position_ids, + packed_query_indexes=packed_indexes, + past_key_values=past_key_values, + key_values_lens=key_values_lens, + packed_key_value_indexes=packed_key_value_indexes, + update_past_key_values=True, + is_causal=False, + **extra_inputs, + ) + return output.past_key_values diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py new file mode 100644 index 000000000..83998e2f0 --- /dev/null +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -0,0 +1,1274 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""LancePipeline — Lance (ByteDance) packaged for the vLLM-Omni diffusion engine. + +Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion), so the transformer +core and the entire generation/forward machinery are inherited unchanged from +:class:`BagelPipeline`. Only model *construction* differs, and only in three +well-localized places: + + 1. Checkpoint layout — the HF repo ``bytedance-research/Lance`` bundles + ``Lance_3B/`` (image) and ``Lance_3B_Video/`` (video) LLM checkpoints, + ``Qwen2.5-VL-ViT/`` (understanding ViT) and ``Wan2.2_VAE.pth`` (VAE) in a + single repo. There is no BAGEL-style top-level ``config.json`` carrying + ``vae_config`` / ``vit_config`` / ``latent_patch_size`` — those are Lance + constants taken from upstream ``config/config_factory.py`` and + ``inference_lance.sh`` and hardcoded in :data:`LANCE_DEFAULTS`. + 2. Understanding ViT — Qwen2.5-VL vision tower (bundled + ``Qwen2.5-VL-ViT/vit.safetensors``) instead of SigLIP. + 3. VAE — Wan2.2 (``Wan2.2_VAE.pth``) instead of the BAGEL autoencoder. + +Scope: this lands the **image path** (``t2i`` / ``image_edit`` / ``x2t_image``) +which is the direct BAGEL analogue. The ``Lance_3B_Video`` path needs the 3-D +latent position embedding (:class:`LancePositionEmbedding3D`) and temporal VAE +handling and is an explicit follow-up — see the PR description. + +Bring-up status (verified on a B300 against ``bytedance-research/Lance``): + * t2i: end-to-end working (1024x1024 image in ~6 s, 0 missing keys). + * x2t (image understanding): plumbing wired (Qwen2.5-VL ViT + Qwen2-VL + image processor + no-op connector/vit_pos_embed; VAE prefill skipped + for Lance via :class:`LanceBagel`). The pipeline runs to completion + without crashes but currently emits an immediate EOS because the + Qwen2.5-VL backbone needs mRoPE position ids on the image+text + sequence and we presently force scalar positions + (``rope_scaling = None``); enabling mRoPE end-to-end is the next + follow-up. + * image_edit (img2img): blocked on the same VAE-prefill issue + (Wan2.2 latents do not map onto BAGEL's ``latent_pos_embed`` grid); + needs a Lance-specific ``prepare_vae_images``. + * video (``Lance_3B_Video``): needs the 3-D latent position embedding + wired into ``Bagel`` plus a multi-frame VAE decode path in the + pipeline. :class:`LanceWanVAE.decode_video` and + :class:`LancePositionEmbedding3D` are already in place. + * Two-stage (AR thinker + DiT): needs ``LanceConfig`` / ``LanceProcessor`` + registered in the ``vllm`` package; tracked in a separate PR. +""" + +from __future__ import annotations + +import json +import os +from dataclasses import dataclass + +import torch +from PIL import Image +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.models.bagel.pipeline_bagel import ( + BagelPipeline, + add_special_tokens, +) +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) + +from .lance_transformer import ( + LanceBagel, + LanceIdentityConnector, + LancePositionEmbedding3D, + LanceQwen2_5_VLNaViTWrapper, + LanceZeroVitPosEmbed, + Qwen2MoTConfig, + Qwen2MoTForCausalLM, +) +from .wan_vae import LanceWanVAE + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class LanceDefaults: + """Lance constants that upstream keeps in ``config/config_factory.py`` / + ``inference_lance.sh`` rather than in any shipped JSON. + + Verified against ``bytedance-research/Lance``'s released ``Lance_3B/model.safetensors``: + ``vae2llm.weight = (2048, 48)`` ⇒ ``patch_latent_dim = latent_patch_size**2 * z_channels = 1 * 48``, + i.e. Lance does *not* unfold the Wan latent into a 2×2 patch the way BAGEL does (Wan2.2 already + patchifies internally), and ``latent_pos_embed.pos_embed = (4096, 2048)`` ⇒ ``max_latent_size = 64``. + """ + + # latent / patch geometry (pt, ph, pw); image path only uses the spatial 1. + latent_patch_size_spatial: int = 1 + latent_patch_size_temporal: int = 1 + max_latent_size: int = 64 + # Lance_3B_Video ships ``latent_pos_embed.pos_embed`` of shape + # ``(31 * 64 * 64, 2048) = (126976, 2048)`` ⇒ max_num_frames in latent + # space is 31, equivalent to ``(num_frames - 1) // downsample_temporal + 1`` + # for up to 121 RGB frames. + max_num_video_latent_frames: int = 31 + vit_max_num_patch_per_side: int = 70 + connector_act: str = "gelu_pytorch_tanh" + timestep_shift: float = 3.5 + num_timesteps: int = 30 + cfg_text_scale: float = 4.0 + # Wan2.2 VAE + vae_z_channels: int = 48 + vae_downsample_spatial: int = 16 + vae_downsample_temporal: int = 4 + + +LANCE_DEFAULTS = LanceDefaults() + +# Subdirectories inside the HF repo (see config.json::checkpoint_directories). +_IMAGE_CKPT_DIR = "Lance_3B" +_VIDEO_CKPT_DIR = "Lance_3B_Video" +_VIT_DIR = "Qwen2.5-VL-ViT" +_VAE_FILE = "Wan2.2_VAE.pth" + + +def get_lance_post_process_func(od_config: OmniDiffusionConfig): + """Lance returns PIL.Image.Image directly, same as BAGEL.""" + + def post_process_func(x): + return x + + return post_process_func + + +def get_lance_pre_process_func(od_config: OmniDiffusionConfig): + def pre_process_func(x): + return x + + return pre_process_func + + +@dataclass +class _LanceVaeCfg: + z_channels: int = LANCE_DEFAULTS.vae_z_channels + # BAGEL core computes ``latent_downsample = vae.downsample * latent_patch_size``. + # For Lance image latents that is ``16 * 2 = 32``. + downsample: int = LANCE_DEFAULTS.vae_downsample_spatial + + +@dataclass +class _LanceVitCfg: + patch_size: int = 14 + hidden_size: int = 1280 # Qwen2.5-VL vision hidden (out_hidden_size=2048) + + +class LancePipeline(BagelPipeline): + """Lance pipeline. Inherits BAGEL's forward/generation; overrides only + construction (checkpoint layout, Qwen2.5-VL ViT, Wan2.2 VAE).""" + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + # Intentionally do NOT call BagelPipeline.__init__ — its assumptions + # about config.json / vit_config.json / SigLIP / the BAGEL AE do not + # hold for Lance. We replicate the BAGEL construction sequence with + # Lance-specific component builders, then reuse every inherited method. + import torch.nn as nn + from vllm.transformers_utils.configs.bagel import BagelConfig + + from vllm_omni.diffusion.distributed.utils import get_local_device + from vllm_omni.diffusion.model_loader.diffusers_loader import ( + DiffusersPipelineLoader, + ) + + nn.Module.__init__(self) + self.od_config = od_config + self.device = get_local_device() + self.scheduler = None + self.scheduler_kwargs = {} + + model = od_config.model + if os.path.exists(model): + repo_root = model + else: + repo_root = download_weights_from_hf_specific(model, od_config.revision, ["*"]) + + # If --model points directly at ``Lance_3B`` / ``Lance_3B_Video`` (or + # ``Qwen2.5-VL-ViT``), walk up so ``repo_root`` is the bundled top-level + # dir that owns the sibling components (VAE, ViT). + base = os.path.basename(repo_root.rstrip("/")) + if base in {_IMAGE_CKPT_DIR, _VIDEO_CKPT_DIR, _VIT_DIR}: + parent = os.path.dirname(repo_root.rstrip("/")) + if os.path.isdir(parent) and os.path.isfile(os.path.join(parent, _VAE_FILE)): + repo_root = parent + + is_video = self._select_video_variant(od_config) or base == _VIDEO_CKPT_DIR + ckpt_dir = _VIDEO_CKPT_DIR if is_video else _IMAGE_CKPT_DIR + ckpt_path = os.path.join(repo_root, ckpt_dir) + if not os.path.isdir(ckpt_path): + # Some users point --model directly at a single checkpoint dir. + ckpt_path = repo_root + self._repo_root = repo_root + self._ckpt_path = ckpt_path + self._is_video = is_video + if is_video: + logger.info( + "Lance video checkpoint selected (%s) — wired with " + "LancePositionEmbedding3D and Wan2.2 multi-frame VAE decode " + "(text-to-video). Image-edit / x2t_video remain follow-ups.", + ckpt_dir, + ) + + # ---- LLM (Qwen2-MoT; identical weight layout to BAGEL) ---- + llm_cfg_path = os.path.join(ckpt_path, "llm_config.json") + llm_config = Qwen2MoTConfig.from_json_file(llm_cfg_path) + llm_config.qk_norm = True + # The released Lance checkpoint ships a separate ``language_model.lm_head.weight`` + # tensor even though ``llm_config.json`` says ``tie_word_embeddings=True``; keep + # the head untied so the checkpoint loads with zero missing/unexpected keys. + llm_config.tie_word_embeddings = False + llm_config.layer_module = od_config.override_transformer_cls_name or "Qwen2MoTDecoderLayer" + # Lance is Qwen2.5-VL-MoT and ships ``rope_scaling = {"type": "mrope", + # "mrope_section": [16, 24, 24]}``. Keep the mRoPE configuration: + # :class:`BagelRotaryEmbedding` now auto-dispatches on ``rope_type`` and + # consumes 3-D ``(t, h, w)`` position ids, while :class:`LanceBagel` + # broadcasts scalar positions to ``(3, S)`` for text-only blocks and + # emits true per-token ``(t, h, w)`` for video latent blocks. Required + # for the Qwen2.5-VL backbone to produce coherent x2t / t2v output. + if not isinstance(getattr(llm_config, "rope_scaling", None), dict): + llm_config.rope_scaling = {"rope_type": "mrope", "mrope_section": [16, 24, 24]} + else: + llm_config.rope_scaling.setdefault("rope_type", llm_config.rope_scaling.get("type", "mrope")) + llm_config.rope_scaling.setdefault("mrope_section", [16, 24, 24]) + + self.tokenizer = self._load_tokenizer(ckpt_path) + self.tokenizer, self.new_token_ids, _ = add_special_tokens(self.tokenizer) + # Image / video preprocessor for the understanding paths + # (x2t_image / x2t_video / image_edit / video_edit). Lance's HF repo + # does not bundle a ``preprocessor_config.json``; we use transformers' + # ``Qwen2VLImageProcessor`` / ``Qwen2VLVideoProcessor`` with their + # default CLIP-style normalization, ``patch_size=14``, ``merge_size=2``, + # which matches the bundled ``Qwen2.5-VL-ViT/config.json`` exactly. + self.image_processor = self._build_image_processor() + self.video_processor = self._build_video_processor() + tok_len = len(self.tokenizer) + required_max_id = max(int(v) for v in self.new_token_ids.values()) + llm_config.vocab_size = max( + int(getattr(llm_config, "vocab_size", tok_len)), + int(tok_len), + int(required_max_id + 1), + ) + + parallel_config = od_config.parallel_config if od_config else None + quant_config = od_config.quantization_config + + self.language_model = Qwen2MoTForCausalLM( + llm_config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix="bagel.language_model", + ) + self.transformer = self.language_model.model + + # ---- Understanding ViT: Qwen2.5-VL vision (bundled) ---- + self.vit_model = self._build_qwen2_5_vl_vit(repo_root) + + # ---- VAE: Wan2.2 (bundled .pth) ---- + self.vae = self._build_wan22_vae(repo_root) + + vae_cfg = _LanceVaeCfg() + vit_cfg = _LanceVitCfg( + patch_size=14, + hidden_size=getattr(getattr(self.vit_model, "config", None), "hidden_size", 1280), + ) + + # Lance uses Qwen2.5-VL's vision tower whose ``merger`` already projects + # to the LLM hidden size and which carries its own positional encoding; + # there is no separate BAGEL-style ``connector`` / ``vit_pos_embed`` in + # the released checkpoint (Lance_3B safetensors only contain + # ``vit_model.*``, never ``connector.*`` / ``vit_pos_embed.*``). We keep + # BAGEL's ``visual_und=True`` so ``vit_model`` is registered as a child + # of ``self.bagel`` and BAGEL's ``forward_cache_update_vit`` works + # unchanged, then immediately replace ``connector`` with an identity and + # ``vit_pos_embed`` with a zero op so the strict load check does not + # demand those phantom weights and the addition in + # ``forward_cache_update_vit`` is numerically a no-op. + und_enabled = self.vit_model is not None + self.bagel = LanceBagel( + language_model=self.language_model, + vit_model=self.vit_model, + parallel_config=parallel_config, + quant_config=quant_config, + prefix="bagel", + config=BagelConfig( + llm_config=llm_config, + vae_config=vae_cfg, + vit_config=vit_cfg, + vit_max_num_patch_per_side=LANCE_DEFAULTS.vit_max_num_patch_per_side, + connector_act=LANCE_DEFAULTS.connector_act, + interpolate_pos=False, + latent_patch_size=LANCE_DEFAULTS.latent_patch_size_spatial, + max_latent_size=LANCE_DEFAULTS.max_latent_size, + timestep_shift=LANCE_DEFAULTS.timestep_shift, + visual_gen=True, + visual_und=und_enabled, + ), + ) + if und_enabled: + self.bagel.connector = LanceIdentityConnector() + self.bagel.vit_pos_embed = LanceZeroVitPosEmbed() + # Hand the Qwen2-VL processors to LanceBagel.prepare_vit_images / + # prepare_vit_videos so they can re-call them and recover + # ``image_grid_thw`` / ``video_grid_thw`` (BAGEL's lambda drops the + # grid). + self.bagel._lance_image_processor = self.image_processor + self.bagel._lance_video_processor = self.video_processor + + # Lance_3B_Video ships a 3-D latent positional table + # (``latent_pos_embed.pos_embed`` shape ``(max_num_frames * max_latent_size**2, hidden_size)``) + # in place of BAGEL's 2-D ``PositionEmbedding(max_latent_size, hidden)``; + # swap in :class:`LancePositionEmbedding3D` so the video checkpoint loads + # cleanly. ``max_num_frames`` is derived from the table size. + if is_video: + self.bagel.latent_pos_embed = LancePositionEmbedding3D( + max_num_frames=LANCE_DEFAULTS.max_num_video_latent_frames, + max_num_patch_per_side=LANCE_DEFAULTS.max_latent_size, + hidden_size=llm_config.hidden_size, + ) + + # Weight sources. Verified against bytedance-research/Lance: + # * Lance_3B/model.safetensors carries the LLM + connectors with + # keys ``language_model.* / vae2llm.* / llm2vae.* / time_embedder.* / + # latent_pos_embed.*`` — load under the ``bagel.`` namespace. + # * Lance_3B_Video/model.safetensors additionally includes + # ``vit_model.*`` (390 tensors). + # * Qwen2.5-VL-ViT/vit.safetensors (image checkpoint only) carries + # the bare ``blocks.* / merger.* / patch_embed.*`` and must land + # under ``bagel.vit_model.vision_model.*`` to match the + # :class:`LanceQwen2_5_VLNaViTWrapper` hierarchy. + # Always use the resolved repo_root for weight sources so subfolder paths + # resolve correctly even when the user passed a per-checkpoint subdir + # (we already walked ``repo_root`` up to the bundled root above). + weights_model = repo_root if os.path.isdir(repo_root) else od_config.model + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=weights_model, + subfolder=ckpt_dir if ckpt_path != repo_root else None, + revision=od_config.revision, + prefix="bagel.", + fall_back_to_pt=False, + ), + ] + if und_enabled and not is_video and os.path.isdir(os.path.join(repo_root, _VIT_DIR)): + self.weights_sources.append( + DiffusersPipelineLoader.ComponentSource( + model_or_path=weights_model, + subfolder=_VIT_DIR, + revision=od_config.revision, + prefix="bagel.vit_model.vision_model.", + fall_back_to_pt=False, + ) + ) + + if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp): + self.to(self.device) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + # ------------------------------------------------------------------ # + # Lance-specific component builders + # ------------------------------------------------------------------ # + @staticmethod + def _select_video_variant(od_config: OmniDiffusionConfig) -> bool: + """Pick Lance_3B vs Lance_3B_Video. Defaults to image; a user can + force video via the model path ending in the video dir or an + od_config extra flag.""" + model = str(od_config.model) + if model.rstrip("/").endswith(_VIDEO_CKPT_DIR): + return True + extra = getattr(od_config, "extra", None) or {} + return bool(extra.get("lance_video", False)) + + @staticmethod + def _load_tokenizer(ckpt_path: str): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(ckpt_path, local_files_only=True, trust_remote_code=True) + + @staticmethod + def _build_image_processor(): + """Construct a Qwen2.5-VL-compatible image preprocessor. + + Defaults mirror ``Qwen/Qwen2.5-VL-3B-Instruct`` (CLIP normalization, + 14-pixel patches, 2x spatial merge). + """ + from transformers import Qwen2VLImageProcessor + + return Qwen2VLImageProcessor() + + @staticmethod + def _build_video_processor(): + """Construct a Qwen2.5-VL-compatible video preprocessor. + + Used by the ``x2t_video`` understanding path. Returns + ``pixel_values_videos`` + 3-D ``video_grid_thw = [T_lat, H, W]``. + """ + from transformers import Qwen2VLVideoProcessor + + return Qwen2VLVideoProcessor() + + def _build_qwen2_5_vl_vit(self, repo_root: str): + """Build the bundled Qwen2.5-VL vision tower and wrap it NaViT-style. + + Lance's bundled ``Qwen2.5-VL-ViT/config.json`` sets + ``_attn_implementation = "flash_attention_2"``; we force ``sdpa`` + instead so the tower constructs on hardware (e.g. Blackwell) where + ``flash-attn`` is not present, falling back to PyTorch's native + scaled-dot-product attention. Numerically equivalent for inference. + """ + vit_dir = os.path.join(repo_root, _VIT_DIR) + cfg_path = os.path.join(vit_dir, "config.json") + with open(cfg_path, encoding="utf-8") as f: + vit_cfg_dict = json.load(f) + vit_cfg_dict = dict(vit_cfg_dict) + vit_cfg_dict["_attn_implementation"] = "sdpa" + try: + from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( + Qwen2_5_VLVisionConfig, + ) + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VisionTransformerPretrainedModel, + ) + + vit_conf = Qwen2_5_VLVisionConfig(**vit_cfg_dict) + vision = Qwen2_5_VisionTransformerPretrainedModel(vit_conf) + except Exception as e: + logger.warning( + "Could not instantiate Qwen2.5-VL vision tower (%s). " + "Understanding (x2t) path will be unavailable; generation (t2i) " + "does not use the ViT.", + e, + ) + return None + return LanceQwen2_5_VLNaViTWrapper(vision, spatial_merge_size=int(vit_cfg_dict.get("spatial_merge_size", 2))) + + # ------------------------------------------------------------------ # + # Lance text-to-video forward path + # ------------------------------------------------------------------ # + def forward(self, req): # type: ignore[override] + """Dispatch on prompt modality. + + - ``modalities == ["video"]`` (text-to-video) → :meth:`_forward_t2v` + (3-D latents + ``LanceWanVAE.decode_video``). + - ``modalities == ["text"]`` + ``multi_modal_data.video`` (x2t_video) → + :meth:`_forward_x2t_video` (multi-frame Qwen2.5-VL ViT prefill). + - ``modalities == ["image"]`` + ``multi_modal_data.img2img`` (image_edit) + → :meth:`_forward_image_edit` (Lance-native VAE+ViT prefill + image gen). + - ``modalities == ["video"]`` + ``multi_modal_data.video`` (video_edit) + → :meth:`_forward_video_edit` (Lance-native multi-frame VAE+ViT prefill + + video gen). + - Everything else falls through to :meth:`BagelPipeline.forward` + (t2i, x2t_image). + """ + first_prompt = req.prompts[0] if req.prompts else None + modalities: list[str] = [] + mm_data: dict = {} + 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("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 "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) + return super().forward(req) + + @torch.inference_mode() + def _forward_t2v(self, req): + """Minimal text-to-video forward: text prefill + 3-D latent denoising + + Wan2.2 multi-frame decode. Mirrors :meth:`BagelPipeline.forward`'s t2i + branch but with 3-D latents (no image input, no x2t).""" + from copy import deepcopy + + from vllm_omni.diffusion.data import DiffusionOutput + + from .lance_transformer import NaiveCache + + first_prompt = req.prompts[0] + if isinstance(first_prompt, dict): + prompt = first_prompt.get("prompt") or "" + extra_args = first_prompt.get("extra_args") or {} + else: + prompt = str(first_prompt) + extra_args = {} + # Sampling-side extras override prompt-side. + sp_extra = getattr(req.sampling_params, "extra_args", {}) or {} + extra_args = {**extra_args, **sp_extra} + + # Video shape. T = number of RGB frames (1..121), H/W in pixels. + T = int(extra_args.get("num_frames", 25)) + H = int(req.sampling_params.height or extra_args.get("video_height", 480)) + W = int(req.sampling_params.width or extra_args.get("video_width", 768)) + max_lat = self.bagel.max_latent_size + max_hw = max_lat * self.bagel.latent_downsample + if H > max_hw or W > max_hw: + raise ValueError(f"Requested video resolution {H}x{W} exceeds Lance limit {max_hw}x{max_hw}") + downsample_t = int(getattr(self.bagel.config.vae_config, "downsample_temporal", 4)) + max_T_lat = LANCE_DEFAULTS.max_num_video_latent_frames + if (T - 1) // downsample_t + 1 > max_T_lat: + raise ValueError( + f"Requested num_frames={T} exceeds Lance video temporal limit " + f"{(max_T_lat - 1) * downsample_t + 1} (max_num_video_latent_frames={max_T_lat})" + ) + video_shape = (T, H, W) + logger.info("Lance t2v: video_shape=%s", video_shape) + + cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.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) + + 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) + + # ---- Build positive-prompt KV cache ---- + gen_context = { + "kv_lens": [0], + "ropes": [0], + "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), + } + cfg_text_context = deepcopy(gen_context) + + 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 + + # ---- Build CFG text-unconditional KV cache (empty prompt) ---- + if cfg_text_scale > 1.0: + neg_prompt = str(extra_args.get("negative_prompt") or "") + neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + prompts=[neg_prompt], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in neg_input.items(): + if torch.is_tensor(v): + neg_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_text_context["past_key_values"], **neg_input + ) + cfg_text_context["kv_lens"] = neg_newlens + cfg_text_context["ropes"] = neg_rope + + # ---- 3-D latent init + CFG side metadata ---- + gen_input_lat = self.bagel.prepare_video_latent( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + video_shapes=[video_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_text_context["ropes"], + video_shapes=[video_shape], + ) + cfg_img_lat = self.bagel.prepare_video_latent_cfg( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + video_shapes=[video_shape], + ) + for k, v in cfg_text_lat.items(): + if torch.is_tensor(v): + cfg_text_lat[k] = v.to(self.device) + for k, v in cfg_img_lat.items(): + if torch.is_tensor(v): + cfg_img_lat[k] = v.to(self.device) + + # ---- Denoising loop (Bagel.generate_image is rank-agnostic over packed tokens) ---- + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + latents, *_ = self.bagel.generate_image( + past_key_values=gen_context["past_key_values"], + cfg_text_past_key_values=cfg_text_context["past_key_values"], + cfg_img_past_key_values=gen_context["past_key_values"], # no img CFG branch + num_timesteps=num_timesteps, + timestep_shift=timestep_shift, + cfg_text_scale=cfg_text_scale, + cfg_img_scale=1.0, + **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"], + ) + + frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) + # Convert numpy frames to PIL.Image list for downstream serialization. + frames = [Image.fromarray(f) for f in frames_np] + logger.info("Lance t2v: decoded %d frames at %dx%d", len(frames), frames[0].width, frames[0].height) + return DiffusionOutput( + output=frames[0], + custom_output={"video_frames": frames, "video_shape": video_shape}, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + @torch.inference_mode() + def _forward_image_edit(self, req): + """Image edit (img2img): reference image + text prompt → modified image. + + Mirrors :meth:`_forward_t2v` for the text/VAE prefill structure but + also runs ``_lance_native_prepare_vae_images`` for the reference + image (so the LLM sees both the Wan2.2-encoded latents and the + Qwen2.5-VL ViT context) and then a Lance image-gen loop emitting + 2-D latents for the new image. + """ + + import numpy as np + + from vllm_omni.diffusion.data import DiffusionOutput + + from .lance_transformer import NaiveCache + + first_prompt = req.prompts[0] + assert isinstance(first_prompt, dict), "image_edit requires dict-style prompt" + prompt = first_prompt.get("prompt") or "" + mm_data = first_prompt.get("multi_modal_data") or {} + image_input = mm_data.get("img2img") or mm_data.get("image") + if image_input is None: + raise ValueError("image_edit requires multi_modal_data.img2img (or .image).") + if not isinstance(image_input, list): + image_input = [image_input] + image_input = [Image.open(im) if isinstance(im, str) else im for im in 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} + cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.0)) + cfg_img_scale = float(extra_args.get("cfg_img_scale", 1.5)) + 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) + + # Resize the reference image to a multiple of latent_downsample and + # within the max latent grid. + stride = self.bagel.latent_downsample + max_hw = int(self.bagel.max_latent_size * stride) + + def _resize_to_stride(img): + if img.mode != "RGB": + img = img.convert("RGB") + w, h = img.size + scale = min(max_hw / max(w, h), 1.0) + scale = max(scale, min(256, max_hw) / min(w, h)) + new_w = max(stride, int(round(w * scale / stride) * stride)) + new_h = max(stride, int(round(h * scale / stride) * stride)) + new_w = min(new_w, max_hw) + new_h = min(new_h, max_hw) + if new_w != w or new_h != h: + img = img.resize((new_w, new_h), Image.BICUBIC) + return img + + image_input = [_resize_to_stride(im) for im in image_input] + resized_w, resized_h = image_input[0].size + image_shape = (resized_h, resized_w) + logger.info("Lance image_edit: ref image %dx%d", resized_w, resized_h) + + def vae_transforms(img): + arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0 + return arr.permute(2, 0, 1) # (C, H, W) + + def vit_transforms(img): + return self.image_processor(images=img, return_tensors="pt").pixel_values[0] + + 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) + + # Match upstream Lance's edit position-id convention: input reference + # VAE block and output gen latent share rope range ``[0, max_axis]``, + # while text / ViT prefill are shifted by ``_EDIT_SHIFT`` so they + # never collide. Without this, the model sees the reference at + # different rope coordinates than the gen target and effectively + # regenerates instead of editing. + _EDIT_SHIFT = 10000 + gen_context = { + "kv_lens": [0], + "ropes": [0], + "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), + } + + # VAE prefill (Lance-native, 3-D mRoPE positions) — rope starts at 0 + # to align with the gen latent below. + vae_input, newlens_vae, _ = self.bagel._lance_native_prepare_vae_images( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + images=image_input, + transforms=vae_transforms, + new_token_ids=self.new_token_ids, + is_video=False, + ) + for k, v in vae_input.items(): + if torch.is_tensor(v): + vae_input[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_vae( + self.vae, gen_context["past_key_values"], **vae_input + ) + gen_context["kv_lens"] = newlens_vae + # Shift rope out of the gen-latent range for the following text / ViT + # prefills; we snap back to 0 when constructing the gen latent. + gen_context["ropes"] = [_EDIT_SHIFT] + + # ViT prefill (Lance Qwen2.5-VL with packed flat patches). + vit_input, newlens_vit, new_rope_vit = self.bagel.prepare_vit_images( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + images=image_input, + transforms=vit_transforms, + new_token_ids=self.new_token_ids, + ) + for k, v in vit_input.items(): + if torch.is_tensor(v): + vit_input[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_vit( + gen_context["past_key_values"], **vit_input + ) + gen_context["kv_lens"] = newlens_vit + gen_context["ropes"] = new_rope_vit + + # Text prefill. + text_input, newlens_text, new_rope_text = 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 text_input.items(): + if torch.is_tensor(v): + text_input[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"], **text_input + ) + gen_context["kv_lens"] = newlens_text + gen_context["ropes"] = new_rope_text + + # CFG text branch (unconditional text). + cfg_text_context = { + "kv_lens": [0], + "ropes": [0], + "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), + } + if cfg_text_scale > 1.0: + neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + prompts=[str(extra_args.get("negative_prompt") or "")], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in neg_input.items(): + if torch.is_tensor(v): + neg_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_text_context["past_key_values"], **neg_input + ) + cfg_text_context["kv_lens"] = neg_newlens + cfg_text_context["ropes"] = neg_rope + + # ---- 2-D latent init (image output) ---- + # Gen latent uses rope=[0] so its tokens share position coordinates + # with the reference image's VAE block above (edit alignment). + gen_input_lat = self.bagel.prepare_vae_latent( + curr_kvlens=gen_context["kv_lens"], + curr_rope=[0], + image_sizes=[image_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_vae_latent_cfg( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=[0], + image_sizes=[image_shape], + ) + cfg_img_lat = self.bagel.prepare_vae_latent_cfg( + curr_kvlens=gen_context["kv_lens"], + curr_rope=[0], + image_sizes=[image_shape], + ) + for k, v in cfg_text_lat.items(): + if torch.is_tensor(v): + cfg_text_lat[k] = v.to(self.device) + for k, v in cfg_img_lat.items(): + if torch.is_tensor(v): + cfg_img_lat[k] = v.to(self.device) + + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + latents, *_ = self.bagel.generate_image( + past_key_values=gen_context["past_key_values"], + cfg_text_past_key_values=cfg_text_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, + **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"], + ) + + img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) + return DiffusionOutput( + output=img, + custom_output={"image_shape": image_shape}, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + @torch.inference_mode() + def _forward_video_edit(self, req): + """Video edit: reference video + text prompt → modified video. + + Lance-native VAE prefill (multi-frame Wan2.2 encoding) + Qwen2.5-VL + ViT prefill + 3-D latent denoising + ``LanceWanVAE.decode_video``. + Mirrors :meth:`_forward_image_edit` but with the temporal axis + threaded through both prefill and gen latent. + """ + + from vllm_omni.diffusion.data import DiffusionOutput + + from .lance_transformer import NaiveCache + + first_prompt = req.prompts[0] + assert isinstance(first_prompt, dict), "video_edit requires dict-style prompt" + prompt = first_prompt.get("prompt") or "" + mm_data = first_prompt.get("multi_modal_data") or {} + video_input = mm_data.get("video") + if video_input is None: + raise ValueError("video_edit requires multi_modal_data.video.") + if isinstance(video_input, str): + import imageio.v3 as iio + + video_input = iio.imread(video_input) + # ``video_input`` now: ndarray (T, H, W, 3). Permute to (C, T, H, W) for VAE. + import numpy as _np + + if isinstance(video_input, _np.ndarray): + video_chw = torch.from_numpy(video_input).float() / 127.5 - 1.0 # (T, H, W, 3) + video_chw = video_chw.permute(3, 0, 1, 2).contiguous() # (C, T, H, W) + elif torch.is_tensor(video_input): + video_chw = video_input.float() + else: + raise ValueError(f"Unsupported video_input type {type(video_input)}") + T = int(video_chw.shape[1]) + H = int(video_chw.shape[2]) + W = int(video_chw.shape[3]) + + extra_args = first_prompt.get("extra_args") or {} + sp_extra = getattr(req.sampling_params, "extra_args", {}) or {} + extra_args = {**extra_args, **sp_extra} + cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.0)) + cfg_img_scale = float(extra_args.get("cfg_img_scale", 1.5)) + 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) + video_shape = (T, H, W) + logger.info("Lance video_edit: video %dx%dx%d (T,H,W)", T, H, W) + + def vae_transforms(_): + return video_chw + + 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) + + # Position-id layout for video editing — matches upstream Lance's + # ``shift_position_ids(pos_shift=1000)`` + ``modality1=modality2`` + # convention: + # * input video VAE block + output gen video latent share the same + # rope range [0, max(T_lat, H_lat, W_lat)] so the model can align + # "edit *this* token" between reference and output; + # * text + ViT prefill use a far-shifted rope start (``_EDIT_SHIFT``) + # so their positions never collide with the VAE/gen range. + _EDIT_SHIFT = 10000 + gen_context = { + "kv_lens": [0], + "ropes": [0], + "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), + } + + # Multi-frame VAE prefill — rope starts at 0 (aligned with gen latent below). + vae_input, newlens_vae, _ = self.bagel._lance_native_prepare_vae_images( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + images=[video_chw], + transforms=vae_transforms, + new_token_ids=self.new_token_ids, + is_video=True, + ) + for k, v in vae_input.items(): + if torch.is_tensor(v): + vae_input[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_vae( + self.vae, gen_context["past_key_values"], **vae_input + ) + gen_context["kv_lens"] = newlens_vae + # Reset rope to the post-VAE shifted range for the text / ViT prefill; + # we'll snap back to 0 when constructing the gen latent so it shares + # positions with the VAE-prefilled reference. + gen_context["ropes"] = [_EDIT_SHIFT] + + # Multi-frame ViT prefill (reuse x2t_video prep). + vit_input, newlens_vit, new_rope_vit = self.bagel.prepare_vit_videos( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + videos=[video_input], + new_token_ids=self.new_token_ids, + ) + for k, v in vit_input.items(): + if torch.is_tensor(v): + vit_input[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_vit( + gen_context["past_key_values"], **vit_input + ) + gen_context["kv_lens"] = newlens_vit + gen_context["ropes"] = new_rope_vit + + # Text prefill. + text_input, newlens_text, new_rope_text = 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 text_input.items(): + if torch.is_tensor(v): + text_input[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"], **text_input + ) + gen_context["kv_lens"] = newlens_text + gen_context["ropes"] = new_rope_text + + # CFG text branch. + cfg_text_context = { + "kv_lens": [0], + "ropes": [0], + "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), + } + if cfg_text_scale > 1.0: + neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( + curr_kvlens=cfg_text_context["kv_lens"], + curr_rope=cfg_text_context["ropes"], + prompts=[str(extra_args.get("negative_prompt") or "")], + tokenizer=self.tokenizer, + new_token_ids=self.new_token_ids, + ) + for k, v in neg_input.items(): + if torch.is_tensor(v): + neg_input[k] = v.to(self.device) + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text( + cfg_text_context["past_key_values"], **neg_input + ) + cfg_text_context["kv_lens"] = neg_newlens + cfg_text_context["ropes"] = neg_rope + + # 3-D latent init (video output). ``curr_rope=[0]`` so the gen latent + # tokens share positions with the input VAE block above — that's what + # tells the model "edit token (t,h,w) in the input is token (t,h,w) in + # the output". The kv-cache still contains the text/ViT blocks at the + # shifted rope range, so attending over them works. + gen_input_lat = self.bagel.prepare_video_latent( + curr_kvlens=gen_context["kv_lens"], + curr_rope=[0], + video_shapes=[video_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=[0], + video_shapes=[video_shape], + ) + cfg_img_lat = self.bagel.prepare_video_latent_cfg( + curr_kvlens=gen_context["kv_lens"], + curr_rope=[0], + video_shapes=[video_shape], + ) + for k, v in cfg_text_lat.items(): + if torch.is_tensor(v): + cfg_text_lat[k] = v.to(self.device) + for k, v in cfg_img_lat.items(): + if torch.is_tensor(v): + cfg_img_lat[k] = v.to(self.device) + + with torch.autocast( + device_type=self.device.type, + enabled=self.device.type != "cpu", + dtype=self.od_config.dtype, + ): + latents, *_ = self.bagel.generate_image( + past_key_values=gen_context["past_key_values"], + cfg_text_past_key_values=cfg_text_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, + **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"], + ) + + frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) + frames = [Image.fromarray(f) for f in frames_np] + return DiffusionOutput( + output=frames[0], + custom_output={"video_frames": frames, "video_shape": video_shape}, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + @torch.inference_mode() + 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). + """ + + from vllm_omni.diffusion.data import DiffusionOutput + + from .lance_transformer import NaiveCache + + 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.") + # Accept a path → load via imageio; else assume tensor / numpy array (T, H, W, 3). + if isinstance(video_input, str): + import imageio.v3 as iio + + video_input = iio.imread(video_input) + # Wrap single video as a list for the prep function. + videos = [video_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", 200)) + do_sample = bool(extra_args.get("do_sample", False)) + text_temperature = float(extra_args.get("text_temperature", 0.3)) + + 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 ---- + 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 ---- + gen_input_vit, newlens_vit, new_rope_vit = self.bagel.prepare_vit_videos( + curr_kvlens=gen_context["kv_lens"], + curr_rope=gen_context["ropes"], + videos=videos, + 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( + 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_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) + 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, + ): + 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] + logger.info("Lance x2t_video: 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 _decode_video_from_latent( + bagel, + vae: LanceWanVAE, + latent: torch.Tensor, + video_shape: tuple[int, int, int], + ): + """Pack a flat Lance latent into ``(B, 48, t, h, w)`` and decode to a + ``(T_lat, H, W, 3)`` numpy frame list via :meth:`LanceWanVAE.decode_video`. + """ + + T, H, W = video_shape + downsample_t = int(getattr(bagel.config.vae_config, "downsample_temporal", 4)) + t_lat = (T - 1) // downsample_t + 1 + h_lat = H // bagel.latent_downsample + w_lat = W // bagel.latent_downsample + p = bagel.latent_patch_size + c = bagel.latent_channel + # ``latent`` is flat ``(t_lat*h_lat*w_lat, c * p * p)`` (Lance image + # path uses ``p=1`` so this is ``(N, c)``). Reshape into the 5-D layout + # the Wan2.2 VAE decoder consumes. + latent = latent.reshape(t_lat, h_lat, w_lat, p, p, c) + latent = torch.einsum("thwpqc->cthpwq", latent) + latent = latent.reshape(1, c, t_lat, h_lat * p, w_lat * p) + + vae_dtype = next(vae.parameters()).dtype + latent = latent.to(vae_dtype) + video = vae.decode_video(latent) # (1, 3, T_out, H_out, W_out) + # Scale [-1,1] -> [0,255] uint8 frames; the upstream VAE already clamps. + video = (video[0] * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 3, 0) * 255 + return video.to(torch.uint8).cpu().numpy() # (T_out, H_out, W_out, 3) + + def _build_wan22_vae(self, repo_root: str) -> LanceWanVAE: + """Wrap the bundled ``Wan2.2_VAE.pth`` with a BAGEL-compatible + ``.encode(images)`` / ``.decode(latent)`` surface. + + Uses the upstream Wan2.2 VAE module (ported in :mod:`wan_vae`) since + the released ``.pth`` is keyed for that module — *not* for the diffusers + ``AutoencoderKLWan`` state dict that ``OmniAutoencoderKLWan`` expects. + Single images are treated as 1-frame clips; video latents pass through + the same module via the 5-D ``encode_video``/``decode_video`` path. + Wan2.2 VAE constants: z=48ch, /16 spatial, /4 temporal. + """ + vae_path = os.path.join(repo_root, _VAE_FILE) + return LanceWanVAE(vae_path=vae_path, dtype=torch.bfloat16, device=self.device) diff --git a/vllm_omni/diffusion/models/lance/prompts.py b/vllm_omni/diffusion/models/lance/prompts.py new file mode 100644 index 000000000..a3e8611ca --- /dev/null +++ b/vllm_omni/diffusion/models/lance/prompts.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance chat / system prompts. + +Matches the upstream Lance training distribution (see +``bytedance/Lance/data/system_prompt_render.py`` and +``data/common.py::generate_system_prompt``): + + * each task has a specific system prompt describing what to attend to; + * the prompt is wrapped in the Qwen chat template + ``<|im_start|>system\\n…<|im_end|>\\n<|im_start|>user\\n…<|im_end|>\\n`` + ``<|im_start|>assistant\\n``; + * vision tokens are framed by ``<|vision_start|><|video_pad|><|vision_end|>`` + (upstream uses ``<|video_pad|>`` even for image type by default). + +These are pure string formatting helpers — the runtime pipeline still does +the actual VAE / ViT prefill separately. +""" + +from __future__ import annotations + +# Task-specific system prompts (verbatim from upstream's +# ``generate_system_prompt`` — first variant of each pool). +SYSTEM_PROMPTS: dict[tuple[str, str], str] = { + # (task, vision_type) -> system prompt + ("t2i", "image"): ( + "Describe the image by detailing the color, quantity, text, shape, " + "size, texture, spatial relationships of the objects and background:" + ), + ("t2v", "video"): ( + "Describe the video by detailing the color, quantity, visible text, " + "shape, size, texture, spatial relationships and motion/camera " + "movements of the objects and background:" + ), + ("image_edit", "image"): ( + "Describe the key features of the input image (color, shape, size, " + "texture, objects, background), then explain how the user's text " + "instruction should alter or modify the image. Generate a new image " + "that meets the user's requirements while maintaining consistency " + "with the original input where appropriate." + ), + ("video_edit", "video"): ( + "Describe the key features of the input video (color, shape, size, " + "texture, objects, background), then explain how the user's text " + "instruction should alter or modify the video. Generate a new video " + "that meets the user's requirements while maintaining consistency " + "with the original input where appropriate." + ), + ("x2t_image", "image"): ( + "Generate a detailed and accurate description of the image, including all the key moments and visual details." + ), + ("x2t_video", "video"): ( + "Generate a detailed and accurate description of the video, including all the key moments and visual details." + ), +} + +# Vision-content placeholders (Qwen2.5-VL convention). Upstream uses +# ``<|video_pad|>`` for image inputs too unless ``force_video_pad=False`` +# (counter-intuitively named; see ``render_qwenvl_prompt`` upstream). +VISION_START = "<|vision_start|>" +VISION_END = "<|vision_end|>" +VIDEO_PAD = "<|video_pad|>" +IMAGE_PAD = "<|image_pad|>" + + +def render_lance_prompt( + task: str, + user_text: str, + *, + vision_token: str | None = None, +) -> str: + """Render a Lance-compatible user-side prompt. + + ``task`` is one of ``t2i``, ``t2v``, ``image_edit``, ``video_edit``, + ``x2t_image``, ``x2t_video``. ``vision_token`` should be the visual + placeholder to embed inside the user message (e.g. ``VIDEO_PAD`` or a + ``VISION_START + VIDEO_PAD + VISION_END`` block); pass ``None`` for + text-only inputs. + + Returns a single string ready to be tokenized; no further wrapping is + needed by the caller. + """ + # ``t2v`` / ``x2t_video`` / ``video_edit`` use video system prompts; + # the rest are image. + vision_type = "video" if task in {"t2v", "x2t_video", "video_edit"} else "image" + sys_prompt = SYSTEM_PROMPTS.get((task, vision_type), "You are a helpful assistant.") + + if vision_token is None: + user_msg = user_text + else: + # Upstream concatenates the vision block immediately before the text. + user_msg = f"{vision_token}{user_text}" + + return ( + f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n" + ) + + +__all__ = [ + "IMAGE_PAD", + "SYSTEM_PROMPTS", + "VIDEO_PAD", + "VISION_END", + "VISION_START", + "render_lance_prompt", +] diff --git a/vllm_omni/diffusion/models/lance/wan_vae.py b/vllm_omni/diffusion/models/lance/wan_vae.py new file mode 100644 index 000000000..aad244d30 --- /dev/null +++ b/vllm_omni/diffusion/models/lance/wan_vae.py @@ -0,0 +1,864 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Ported from ByteDance Lance upstream (https://github.com/bytedance/Lance, +# modeling/vae/wan/{vae2_2.py,model.py}). Upstream copyright: +# +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates. +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +# Licensed under the Apache License, Version 2.0. +# +# Lance ships its VAE as a raw ``Wan2.2_VAE.pth`` whose state dict is keyed for +# the upstream ``WanVAE_`` nn.Module here (not for diffusers' ``AutoencoderKLWan``). +# We port that module verbatim so we can load the .pth directly, and wrap it with +# :class:`LanceWanVAE` which exposes BAGEL's 4-D ``encode(BCHW)/decode(BCHW)`` +# surface (and a 5-D path for the video checkpoint). +"""Wan2.2 VAE used by Lance, ported from upstream so ``Wan2.2_VAE.pth`` loads +natively without state-dict surgery.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from vllm.logger import init_logger + +logger = init_logger(__name__) + +CACHE_T = 2 + + +class CausalConv3d(nn.Conv3d): + """Causal 3D conv with feature-map caching across temporal chunks.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = ( + self.padding[2], + self.padding[2], + self.padding[1], + self.padding[1], + 2 * self.padding[0], + 0, + ) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + return super().forward(x) + + +class RMS_norm(nn.Module): + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0 + + def forward(self, x): + return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + def forward(self, x): + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + def __init__(self, dim, mode): + assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d") + super().__init__() + self.dim = dim + self.mode = mode + + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == "downsample3d": + self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0)) + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = "Rep" + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep": + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep": + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if feat_cache[idx] == "Rep": + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.resample(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +class ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), + nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), + nn.SiLU(), + nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1), + ) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """Single-head causal self-attention over spatial tokens, per frame.""" + + def __init__(self, dim): + super().__init__() + self.dim = dim + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, "b c t h w -> (b t) c h w") + x = self.norm(x) + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1) + x = F.scaled_dot_product_attention(q, k, v) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + x = self.proj(x) + x = rearrange(x, "(b t) c h w-> b c t h w", t=t) + return x + identity + + +def _patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + if x.dim() == 5: + return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size) + raise ValueError(f"Invalid input shape: {x.shape}") + + +def _unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + if x.dim() == 5: + return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size) + return x + + +class AvgDown3D(nn.Module): + def __init__(self, in_channels, out_channels, factor_t, factor_s=1): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + x = F.pad(x, (0, 0, 0, 0, pad_t, 0)) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view(B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s) + x = x.view(B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s) + return x.mean(dim=2) + + +class DupUp3D(nn.Module): + def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False): + super().__init__() + self.avg_shortcut = AvgDown3D( + in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1 + ) + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample(out_dim, mode=mode)) + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__(self, in_dim, out_dim, dropout, mult, temporal_upsample=False, up_flag=False): + super().__init__() + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, out_dim, factor_t=2 if temporal_upsample else 1, factor_s=2 if up_flag else 1 + ) + else: + self.avg_shortcut = None + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + if up_flag: + mode = "upsample3d" if temporal_upsample else "upsample2d" + upsamples.append(Resample(out_dim, mode=mode)) + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + return x_main + self.avg_shortcut(x, first_chunk) + return x_main + + +class Encoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(True, True, False), + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = list(dim_mult) + self.num_res_blocks = num_res_blocks + self.attn_scales = list(attn_scales) + self.temperal_downsample = list(temperal_downsample) + + dims = [dim * u for u in [1] + self.dim_mult] + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = self.temperal_downsample[i] if i < len(self.temperal_downsample) else False + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(self.dim_mult) - 1, + ) + ) + self.downsamples = nn.Sequential(*downsamples) + + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + def __init__( + self, + dim=128, + z_dim=4, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temporal_upsample=(False, True, True), + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = list(dim_mult) + self.num_res_blocks = num_res_blocks + self.attn_scales = list(attn_scales) + self.temporal_upsample = list(temporal_upsample) + + dims = [dim * u for u in [self.dim_mult[-1]] + self.dim_mult[::-1]] + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + self.middle = nn.Sequential( + ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout), + ) + + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = self.temporal_upsample[i] if i < len(self.temporal_upsample) else False + upsamples.append( + Up_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temporal_upsample=t_up_flag, + up_flag=i != len(self.dim_mult) - 1, + ) + ) + self.upsamples = nn.Sequential(*upsamples) + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1), + ) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def _count_conv3d(model): + return sum(1 for m in model.modules() if isinstance(m, CausalConv3d)) + + +class WanVAE_(nn.Module): + """Upstream Wan2.2 VAE module — encoder3d/decoder3d sandwich with 2x patchify + on input. State-dict-compatible with ``Wan2.2_VAE.pth``.""" + + def __init__( + self, + dim=160, + dec_dim=256, + z_dim=16, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(True, True, False), + dropout=0.0, + ): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = list(dim_mult) + self.num_res_blocks = num_res_blocks + self.attn_scales = list(attn_scales) + self.temperal_downsample = list(temperal_downsample) + self.temporal_upsample = list(temperal_downsample)[::-1] + + self.encoder = Encoder3d( + dim, z_dim * 2, self.dim_mult, num_res_blocks, self.attn_scales, self.temperal_downsample, dropout + ) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d( + dec_dim, z_dim, self.dim_mult, num_res_blocks, self.attn_scales, self.temporal_upsample, dropout + ) + + def encode(self, x, scale): + self.clear_cache() + x = _patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + out = None + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx, + ) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + mu = (mu - scale[0]) * scale[1] + self.clear_cache() + return mu, log_var + + def decode(self, z, scale): + self.clear_cache() + if isinstance(scale[0], torch.Tensor): + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1) + else: + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + out = None + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder( + x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True + ) + else: + out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) + out = _unpatchify(out, patch_size=2) + self.clear_cache() + return out + + def clear_cache(self): + self._conv_num = _count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + self._enc_conv_num = _count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +# ---------------------------------------------------------------------------- +# Wan2.2 VAE wrapper + BAGEL-surface adapter +# ---------------------------------------------------------------------------- + + +# Per-channel mean/std for the 48-channel Wan2.2 latent space (upstream constants). +_WAN22_LATENT_MEAN: tuple[float, ...] = ( + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, +) +_WAN22_LATENT_STD: tuple[float, ...] = ( + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, +) + + +class LanceWanVAE(nn.Module): + """Wan2.2 VAE wrapped for BAGEL's pipeline. + + Exposes BAGEL's image-VAE surface — ``encode(BCHW) -> BC_zHW`` and + ``decode(BC_zHW) -> BCHW`` — by treating each image as a 1-frame video clip. + A 5-D ``encode_video``/``decode_video`` path is also provided for the + ``Lance_3B_Video`` checkpoint. + + Construction is lazy: the heavy ``WanVAE_`` and ``Wan2.2_VAE.pth`` are not + materialized until first use. Once built, the inner module is registered as + a submodule so ``self.parameters()``, ``self.to(device)`` and ``vae_dtype = + next(vae.parameters()).dtype`` (used by BAGEL's decode path) all behave. + """ + + z_channels: int = 48 + downsample_spatial: int = 16 + downsample_temporal: int = 4 + + def __init__( + self, + vae_path: str, + dtype: torch.dtype = torch.bfloat16, + device: torch.device | str | None = None, + *, + lazy: bool = False, + ): + super().__init__() + self._vae_path = vae_path + self._dtype = dtype + self._device = torch.device(device) if device is not None else None + self._built = False + # Registered lazily once the .pth is loaded: + # self.model: WanVAE_ (set in _ensure_built) + if not lazy: + # BAGEL's decode path inspects ``next(vae.parameters()).dtype`` before + # the first encode/decode call, so we have to register the submodule + # eagerly when a real device is provided. ``lazy=True`` is reserved + # for config-only / no-GPU paths. + self._ensure_built() + + def _ensure_built(self) -> None: + if self._built: + return + device = self._device or ( + next(self.parameters(), torch.empty(0)).device + if any(True for _ in self.parameters()) + else torch.device("cpu") + ) + # Wan2.2 VAE config from upstream Wan2_2_VAE defaults: + # z_dim=48 (channels), c_dim=160 (encoder width), dec_dim=256 (decoder width), + # dim_mult=[1,2,4,4], temperal_downsample=[False,True,True]. + model = WanVAE_( + dim=160, + dec_dim=256, + z_dim=self.z_channels, + dim_mult=(1, 2, 4, 4), + num_res_blocks=2, + attn_scales=(), + temperal_downsample=(False, True, True), + dropout=0.0, + ) + logger.info("Loading Wan2.2 VAE state dict from %s", self._vae_path) + state = torch.load(self._vae_path, map_location="cpu", weights_only=True) + missing, unexpected = model.load_state_dict(state, strict=False) + if missing: + logger.warning("Wan2.2 VAE missing keys (%d): %s", len(missing), missing[:5]) + if unexpected: + logger.warning("Wan2.2 VAE unexpected keys (%d): %s", len(unexpected), unexpected[:5]) + model = model.to(device=device, dtype=self._dtype).eval() + model.requires_grad_(False) + # Register as submodule so .to() / .parameters() pick it up. + self.model = model + # Per-channel normalization tensors (matches upstream Wan2_2_VAE.scale = [mean, 1/std]). + mean = torch.tensor(_WAN22_LATENT_MEAN, dtype=self._dtype, device=device) + inv_std = torch.tensor([1.0 / s for s in _WAN22_LATENT_STD], dtype=self._dtype, device=device) + # Buffers, so .to(device) follows along. + self.register_buffer("_latent_mean", mean, persistent=False) + self.register_buffer("_latent_inv_std", inv_std, persistent=False) + self._built = True + + @torch.inference_mode() + def encode_video(self, video: torch.Tensor, *, use_sample: bool = True) -> torch.Tensor: + """Encode a 5-D clip ``[B, 3, T, H, W]`` -> latent ``[B, 48, t, h, w]``.""" + self._ensure_built() + video = video.to(self.model.encoder.conv1.weight.dtype) + mu, log_var = self.model.encode(video, [self._latent_mean, self._latent_inv_std]) + if use_sample: + std = torch.exp(0.5 * log_var) + mu = mu + std * torch.randn_like(std) + return mu + + @torch.inference_mode() + def decode_video(self, latent: torch.Tensor) -> torch.Tensor: + """Decode a 5-D latent ``[B, 48, t, h, w]`` -> video ``[B, 3, T, H, W]``.""" + self._ensure_built() + latent = latent.to(self.model.decoder.conv1.weight.dtype) + out = self.model.decode(latent, [self._latent_mean, self._latent_inv_std]) + return out.clamp_(-1.0, 1.0) + + # ----- BAGEL image-VAE surface (4-D) -------------------------------- # + + @torch.inference_mode() + def encode(self, padded_images: torch.Tensor) -> torch.Tensor: + """``[B, 3, H, W]`` -> ``[B, 48, H/16, W/16]`` (single-frame image path). + + Each image is wrapped as a 1-frame clip, encoded, and the temporal axis + is squeezed back out so the result matches BAGEL's iteration pattern. + """ + if padded_images.dim() == 4: + video = padded_images.unsqueeze(2) # B,C,1,H,W + elif padded_images.dim() == 5: + video = padded_images + else: + raise ValueError(f"LanceWanVAE.encode expects 4-D BCHW or 5-D BCTHW, got {tuple(padded_images.shape)}") + latent = self.encode_video(video, use_sample=True) + # B,48,t,h,w -> B,48,h,w when t == 1 (image path) + if latent.shape[2] == 1: + return latent.squeeze(2) + return latent + + @torch.inference_mode() + def decode(self, latent: torch.Tensor) -> torch.Tensor: + """``[B, 48, h, w]`` -> ``[B, 3, H, W]`` (single-frame image path).""" + if latent.dim() == 4: + latent = latent.unsqueeze(2) # B,48,1,h,w + elif latent.dim() != 5: + raise ValueError(f"LanceWanVAE.decode expects 4-D BCHW or 5-D BCTHW, got {tuple(latent.shape)}") + video = self.decode_video(latent) # B,3,T,H,W + if video.shape[2] == 1: + return video.squeeze(2) + return video + + +def build_wan22_vae(vae_path: str, dtype: torch.dtype = torch.bfloat16, device=None) -> LanceWanVAE: + """Convenience factory: lazy-construct a :class:`LanceWanVAE` adapter.""" + return LanceWanVAE(vae_path=vae_path, dtype=dtype, device=device) + + +__all__ = ["LanceWanVAE", "WanVAE_", "build_wan22_vae"] diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index b6a5a3700..d341c2928 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -141,6 +141,11 @@ _DIFFUSION_MODELS = { "pipeline_bagel", "BagelPipeline", ), + "LancePipeline": ( + "lance", + "pipeline_lance", + "LancePipeline", + ), "InternVLAA1Pipeline": ( "internvla_a1", "pipeline_internvla_a1", @@ -462,6 +467,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func", "LongCatImagePipeline": "get_longcat_image_post_process_func", "BagelPipeline": "get_bagel_post_process_func", + "LancePipeline": "get_lance_post_process_func", "InternVLAA1Pipeline": "get_internvla_a1_post_process_func", "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", diff --git a/vllm_omni/model_executor/models/lance/__init__.py b/vllm_omni/model_executor/models/lance/__init__.py new file mode 100644 index 000000000..ee47a467a --- /dev/null +++ b/vllm_omni/model_executor/models/lance/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance (ByteDance) — vLLM-Omni model_executor package. + +The runnable, registered path is the **single-stage** topology: a +self-contained diffusion stage (``LancePipeline``) that owns the Qwen2-MoT +LLM, the Qwen2.5-VL ViT, the Wan2.2 VAE and the tokenizer — directly analogous +to ``bagel_single_stage``. + +The two-stage (separate AR "thinker" + DiT) topology additionally needs a +``LanceConfig`` / ``LanceProcessor`` registered in the ``vllm`` package +(mirroring ``vllm.transformers_utils.configs.bagel`` / +``vllm.model_executor.models.bagel``); that is an explicit follow-up and is +intentionally not wired here to avoid shipping a half-working AR stage. +""" + +from vllm_omni.model_executor.models.lance.lance import ( + OmniLanceForConditionalGeneration, +) + +__all__ = ["OmniLanceForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/lance/lance.py b/vllm_omni/model_executor/models/lance/lance.py new file mode 100644 index 000000000..7ae9dd6f7 --- /dev/null +++ b/vllm_omni/model_executor/models/lance/lance.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance AR-stage model wrapper (two-stage topology). + +Lance is BAGEL-lineage, so the AR/thinker stage is the same Qwen2-MoT model +BAGEL uses. This subclass exists so the two-stage pipeline has a distinct +arch name to register against; behaviour is inherited unchanged from +:class:`OmniBagelForConditionalGeneration`. + +IMPORTANT: the two-stage Lance topology is a follow-up. It additionally +requires a ``LanceConfig`` + ``LanceProcessor`` registered in the ``vllm`` +package (the base ``BagelForConditionalGeneration`` resolves its HF config / +processor by ``model_type == "bagel"``; Lance ships ``model_type == +"qwen2_5_vl"`` in ``llm_config.json`` and no BAGEL-style top-level config). +Until that lands, use the **single-stage** topology (``model_type: lance`` -> +``LANCE_SINGLE_STAGE_PIPELINE`` -> ``LancePipeline``), which is fully wired and +needs no ``vllm``-submodule changes. +""" + +from __future__ import annotations + +from vllm_omni.model_executor.models.bagel.bagel import ( + OmniBagelForConditionalGeneration, +) + + +class OmniLanceForConditionalGeneration(OmniBagelForConditionalGeneration): + """Lance AR-stage model. Identical to BAGEL's (same Qwen2-MoT weights); + distinct class only so the two-stage pipeline can register a ``lance`` + arch separately. See module docstring for the two-stage follow-up.""" + + # Same packed→sublayer LoRA / MoT mapping as BAGEL (weight names match). + packed_modules_mapping = OmniBagelForConditionalGeneration.packed_modules_mapping diff --git a/vllm_omni/model_executor/models/lance/pipeline.py b/vllm_omni/model_executor/models/lance/pipeline.py new file mode 100644 index 000000000..fa7d14fa7 --- /dev/null +++ b/vllm_omni/model_executor/models/lance/pipeline.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Lance pipeline topologies. + +Single-stage (runnable / registered): + Stage 0: DiT — self-contained diffusion stage that handles all modalities + (text2img, image_edit, x2t_image) internally via its own + Qwen2-MoT LLM, Qwen2.5-VL ViT, Wan2.2 VAE, and tokenizer. + Mirrors ``bagel_single_stage``. + +Two-stage (follow-up — see vllm_omni/model_executor/models/lance/lance.py): + Stage 0: Thinker (AR) -> Stage 1: DiT. Needs a LanceConfig/LanceProcessor + in the ``vllm`` package; defined here for completeness but gated behind the + follow-up so it is not selected by default. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.lance" + +# --- Single-stage (default, fully wired) -------------------------------- # +LANCE_SINGLE_STAGE_PIPELINE = PipelineConfig( + model_type="lance", + model_arch="LancePipeline", + hf_architectures=(), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="image", + ), + ), +) + +# --- Two-stage (follow-up; not registered as default) ------------------- # +LANCE_PIPELINE = PipelineConfig( + model_type="lance_two_stage", + model_arch="OmniLanceForConditionalGeneration", + hf_architectures=(), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + model_arch="OmniLanceForConditionalGeneration", + engine_output_type="text", + prompt_expand_func=f"{_PROC}.expand_cfg_prompts", + omni_kv_config={ + "need_send_cache": True, + "kv_transfer_criteria": {"type": "prefill_finished"}, + }, + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + final_output=True, + final_output_type="image", + cfg_kv_collect_func=f"{_PROC}.collect_cfg_kv_caches", + omni_kv_config={"need_recv_cache": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index cc900d471..568d6621b 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -134,6 +134,12 @@ _OMNI_MODELS = { "bagel", "OmniBagelForConditionalGeneration", ), + ## lance (ByteDance) — BAGEL-lineage; two-stage AR arch (follow-up) + "OmniLanceForConditionalGeneration": ( + "lance", + "lance", + "OmniLanceForConditionalGeneration", + ), "HunyuanImage3ForCausalMM": ( "hunyuan_image3", "hunyuan_image3", diff --git a/vllm_omni/model_executor/stage_input_processors/lance.py b/vllm_omni/model_executor/stage_input_processors/lance.py new file mode 100644 index 000000000..5ab029a50 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/lance.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage input processor for Lance. + +Lance's classifier-free-guidance prompt expansion and KV-cache collection are +model-agnostic and identical to BAGEL's (3-branch CFG: gen / cfg_text / +cfg_img). Re-export BAGEL's implementations so the two-stage pipeline config +can reference ``...stage_input_processors.lance.{expand_cfg_prompts, +collect_cfg_kv_caches}`` without duplicating logic. +""" + +from vllm_omni.model_executor.stage_input_processors.bagel import ( # noqa: F401 + CFG_IMG_SUFFIX, + CFG_TEXT_SUFFIX, + GEN_THINK_SYSTEM_PROMPT, + VLM_THINK_SYSTEM_PROMPT, + ExpandedPrompt, + collect_cfg_kv_caches, + expand_cfg_prompts, + expand_cfg_prompts_think, +) + +__all__ = [ + "CFG_IMG_SUFFIX", + "CFG_TEXT_SUFFIX", + "GEN_THINK_SYSTEM_PROMPT", + "VLM_THINK_SYSTEM_PROMPT", + "ExpandedPrompt", + "collect_cfg_kv_caches", + "expand_cfg_prompts", + "expand_cfg_prompts_think", +]