[New Model] Add Lance (ByteDance) (#3710)
Signed-off-by: lishunyang <lishunyang03@gmail.com> Signed-off-by: lishunyang12 <lishunyang12@163.com> Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com> Co-authored-by: lishunyang <lishunyang03@gmail.com> Co-authored-by: Haco <923390377@qq.com>
This commit is contained in:
24 files changed
+7082
-114
No files matched your search
@@ -270,3 +270,6 @@ vllm_omni/_version.py
|
||||
.worktrees/
|
||||
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
|
||||
tests/.ci_generated/
|
||||
|
||||
# Local Lance precision-matching dump scripts (not part of the repo)
|
||||
lance_compare/
|
||||
@@ -0,0 +1,52 @@
|
||||
# Lance (ByteDance) — offline inference
|
||||
|
||||
[Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified
|
||||
autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is
|
||||
**BAGEL-lineage** (ByteDance Mixture-of-Transformers): the released `Lance_3B`
|
||||
checkpoint uses the same `*_moe_gen` MoT weight layout as BAGEL, so vLLM-Omni
|
||||
implements it by reusing the BAGEL transformer core and specializing only the
|
||||
ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
|
||||
|
||||
This example covers all six Lance modalities from the upstream HF model card:
|
||||
`t2i`, `t2v`, `image_edit`, `video_edit`, `x2t_image` (image understanding) and
|
||||
`x2t_video` (video understanding).
|
||||
|
||||
## Hardware
|
||||
|
||||
Single NVIDIA GPU with 16 GB+ VRAM in BF16 (we test on B300 / A100). CUDA ≥ 12.4.
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
# Text-to-image
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance \
|
||||
--prompts "a corgi astronaut on the moon, cinematic" \
|
||||
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
|
||||
--height 1024 --width 1024 \
|
||||
--output ./out
|
||||
|
||||
# Text-to-video (uses the Lance_3B_Video subfolder; see ``--modality``
|
||||
# choices for all six task variants)
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance/Lance_3B_Video --modality text2video \
|
||||
--num-frames 25 --video-height 480 --video-width 768 \
|
||||
--prompts "a cat playing piano, cinematic" \
|
||||
--steps 30 --fps 8 --output ./out
|
||||
```
|
||||
|
||||
`video_edit` requires `--model bytedance-research/Lance/Lance_3B_Video` so
|
||||
the 3-D `latent_pos_embed` table is loaded; the other paths can point at
|
||||
the top-level `bytedance-research/Lance` repo and resolve the right
|
||||
sub-checkpoint automatically.
|
||||
|
||||
The HF repo bundles everything (`Lance_3B/`, `Lance_3B_Video/`,
|
||||
`Qwen2.5-VL-ViT/`, `Wan2.2_VAE.pth`); no separate downloads are required.
|
||||
|
||||
## Defaults
|
||||
|
||||
Matches upstream `inference_lance.sh`: 30 denoising steps, timestep-shift 3.5,
|
||||
text CFG 4.0, seed 42, 1024×1024 (override with `--height` / `--width`). For
|
||||
the understanding paths (`img2text` / `video2text`), sampling is enabled by
|
||||
default at `--text-temperature 0.8` because Lance's greedy decoder emits an
|
||||
immediate EOS for many prompts.
|
||||
@@ -0,0 +1,339 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lance-3B offline inference (single-stage).
|
||||
|
||||
Lance is BAGEL-lineage; this example covers the single-stage paths wired in
|
||||
this PR:
|
||||
|
||||
text2img — image generation from text
|
||||
img2text — image understanding (caption / VQA)
|
||||
|
||||
Examples:
|
||||
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance \
|
||||
--prompts "a corgi astronaut on the moon, cinematic" \
|
||||
--output ./out
|
||||
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance --modality img2text \
|
||||
--image-path /path/to/photo.jpg \
|
||||
--prompts "Describe this image in detail."
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model", default="bytedance-research/Lance", help="HF repo or local path.")
|
||||
p.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
|
||||
p.add_argument("--txt-prompts", type=str, default=None, help="File with one prompt per line.")
|
||||
p.add_argument(
|
||||
"--modality",
|
||||
default="text2img",
|
||||
choices=["text2img", "img2text", "text2video", "video2text", "img2img", "video2video", "image2video"],
|
||||
help="Lance single-stage modality.",
|
||||
)
|
||||
p.add_argument("--image-path", type=str, default=None, help="Input image for img2text.")
|
||||
p.add_argument("--video-path", type=str, default=None, help="Input video file for video2text.")
|
||||
p.add_argument("--num-frames", type=int, default=25, help="Number of RGB frames for text2video (max 121).")
|
||||
p.add_argument("--video-height", type=int, default=480, help="Video frame height.")
|
||||
p.add_argument("--video-width", type=int, default=768, help="Video frame width.")
|
||||
p.add_argument("--height", type=int, default=None, help="Image height (t2i). Default = max_hw (1024).")
|
||||
p.add_argument("--width", type=int, default=None, help="Image width (t2i). Default = max_hw (1024).")
|
||||
p.add_argument(
|
||||
"--fps", type=int, default=12, help="Output video FPS when saving MP4 (matches upstream Lance's save_fps=12)."
|
||||
)
|
||||
p.add_argument("--output", type=str, default=".", help="Output directory.")
|
||||
p.add_argument("--steps", type=int, default=30, help="Denoising steps (Lance default 30).")
|
||||
p.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (Lance default 4.0).")
|
||||
p.add_argument("--timestep-shift", type=float, default=3.5, help="Flow-match timestep shift (Lance default 3.5).")
|
||||
p.add_argument("--negative-prompt", type=str, default=None)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--max-text-tokens", type=int, default=512, help="Max tokens for img2text generation.")
|
||||
p.add_argument(
|
||||
"--do-sample",
|
||||
action="store_true",
|
||||
default=True,
|
||||
help="Sample (vs greedy) for x2t generation; greedy frequently emits an immediate EOS for Lance.",
|
||||
)
|
||||
p.add_argument("--no-sample", dest="do_sample", action="store_false", help="Disable sampling (greedy decoding).")
|
||||
p.add_argument(
|
||||
"--text-temperature",
|
||||
type=float,
|
||||
default=0.8,
|
||||
help="Sampling temperature for x2t generation. Lance frequently emits an immediate EOS below ~0.7; 0.8 is a good default.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--system-prompt",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Override the task system prompt (used by x2t_image / x2t_video for "
|
||||
'per-example QA instructions like "Look at the image carefully and '
|
||||
'answer the question."; without it x2t falls back to a '
|
||||
"caption-style default and the model describes instead of answering)."
|
||||
),
|
||||
)
|
||||
# Lance is single-stage diffusion — no deploy YAML needed. Required
|
||||
# engine knobs (``pipeline``, ``enforce_eager``, ``trust_remote_code``,
|
||||
# ``max_num_seqs=1`` …) are passed as flat kwargs to ``Omni`` below
|
||||
# and ``create_default_diffusion`` materializes the stage config.
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(p)
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
from vllm_omni.diffusion.models.lance.prompts import ( # noqa: E402
|
||||
VIDEO_PAD,
|
||||
VISION_END,
|
||||
VISION_START,
|
||||
render_lance_prompt,
|
||||
)
|
||||
|
||||
# Visual block matching upstream Lance: <|vision_start|><|video_pad|><|vision_end|>
|
||||
# (Lance's renderer uses ``<|video_pad|>`` even for image inputs by default.)
|
||||
_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
|
||||
|
||||
|
||||
def _format_text2img_prompts(prompts):
|
||||
return [{"prompt": render_lance_prompt("t2i", p), "modalities": ["image"]} for p in prompts]
|
||||
|
||||
|
||||
def _format_text2video_prompts(prompts, num_frames, video_h, video_w):
|
||||
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
|
||||
return [{"prompt": render_lance_prompt("t2v", p), "modalities": ["video"], "extra_args": extra} for p in prompts]
|
||||
|
||||
|
||||
def _format_image2video_prompts(prompts, image_path, num_frames, video_h, video_w):
|
||||
"""Image-to-Video (no first-frame pin): image + text → long video.
|
||||
|
||||
Passes the input image via ``multi_modal_data.first_frame``. The pipeline
|
||||
treats it as a 1-frame reference (VAE+ViT prefill) and generates a fresh
|
||||
multi-frame video at the requested ``num_frames`` × ``video_h`` ×
|
||||
``video_w`` shape.
|
||||
"""
|
||||
import os as _os
|
||||
|
||||
if not image_path or not _os.path.exists(image_path):
|
||||
raise ValueError(f"image2video requires --image-path pointing to an existing file, got: {image_path}")
|
||||
from PIL import Image as _Image
|
||||
|
||||
img = _Image.open(image_path).convert("RGB")
|
||||
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("i2v", p, vision_token=_VISION_BLOCK),
|
||||
"multi_modal_data": {"first_frame": img},
|
||||
"modalities": ["video"],
|
||||
"extra_args": extra,
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_img2img_prompts(prompts, image_path):
|
||||
import os as _os
|
||||
|
||||
if not image_path or not _os.path.exists(image_path):
|
||||
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {image_path}")
|
||||
from PIL import Image as _Image
|
||||
|
||||
img = _Image.open(image_path).convert("RGB")
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("image_edit", p, vision_token=_VISION_BLOCK),
|
||||
"multi_modal_data": {"img2img": img},
|
||||
"modalities": ["image"],
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_video2video_prompts(prompts, video_path, num_frames, video_h, video_w):
|
||||
import os as _os
|
||||
|
||||
if not video_path or not _os.path.exists(video_path):
|
||||
raise ValueError(f"video2video requires --video-path pointing to an existing file, got: {video_path}")
|
||||
|
||||
extra = {"num_frames": num_frames, "video_height": video_h, "video_width": video_w}
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("video_edit", p, vision_token=_VISION_BLOCK),
|
||||
"multi_modal_data": {"video": video_path},
|
||||
"modalities": ["video"],
|
||||
"extra_args": extra,
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_video2text_prompts(prompts, video_path, system_prompt=None):
|
||||
import os as _os
|
||||
|
||||
if not video_path or not _os.path.exists(video_path):
|
||||
raise ValueError(f"video2text requires --video-path pointing to an existing file, got: {video_path}")
|
||||
import imageio.v3 as iio
|
||||
|
||||
video = iio.imread(video_path)
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("x2t_video", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
|
||||
"multi_modal_data": {"video": video},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def _format_img2text_prompts(prompts, image_path, system_prompt=None):
|
||||
from PIL import Image
|
||||
|
||||
if not image_path or not os.path.exists(image_path):
|
||||
raise ValueError(f"img2text requires --image-path pointing to an existing file, got: {image_path}")
|
||||
img = Image.open(image_path).convert("RGB")
|
||||
return [
|
||||
{
|
||||
"prompt": render_lance_prompt("x2t_image", p, vision_token=_VISION_BLOCK, system_prompt=system_prompt),
|
||||
"multi_modal_data": {"image": img},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
|
||||
# Lance aligns more naturally with upstream when using FLASH_ATTN
|
||||
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
|
||||
# numerical drift through the 36-layer Qwen2 stack on B300). Default
|
||||
# to flash_attn when the user hasn't explicitly set a backend.
|
||||
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
|
||||
if args.txt_prompts:
|
||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||
prompts = [ln.strip() for ln in f if ln.strip()]
|
||||
elif args.modality == "img2text":
|
||||
prompts = args.prompts or ["Describe this image in detail."]
|
||||
else:
|
||||
prompts = args.prompts or ["A cute cat, studio lighting, 4k"]
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
# Video modalities require the Lance_3B_Video checkpoint (3D
|
||||
# latent_pos_embed shaped (126976, 2048)). Without it, vllm-omni silently
|
||||
# loads Lance_3B (4096, 2048) — the image-only table — and any t_lat >= 1
|
||||
# immediately indexes out of bounds. Rewrite the model path so
|
||||
# LancePipeline._select_video_variant picks the video subfolder.
|
||||
resolved_model = args.model
|
||||
if args.modality in ("text2video", "video2video", "video2text", "image2video"):
|
||||
if not str(resolved_model).rstrip("/").endswith("Lance_3B_Video"):
|
||||
candidate = os.path.join(resolved_model, "Lance_3B_Video")
|
||||
if os.path.isdir(candidate):
|
||||
resolved_model = candidate
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs["model"] = resolved_model
|
||||
# Lance single-stage defaults (formerly in vllm_omni/deploy/lance.yaml):
|
||||
omni_kwargs.setdefault("pipeline", "lance")
|
||||
omni_kwargs.setdefault("max_num_batched_tokens", 32768)
|
||||
omni_kwargs.setdefault("max_num_seqs", 1)
|
||||
omni_kwargs.setdefault("enforce_eager", True)
|
||||
omni_kwargs.setdefault("trust_remote_code", True)
|
||||
omni_kwargs.setdefault("enable_prefix_caching", False)
|
||||
omni_kwargs.setdefault("async_chunk", False)
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
if args.modality == "img2text":
|
||||
formatted = _format_img2text_prompts(prompts, args.image_path, system_prompt=args.system_prompt)
|
||||
elif args.modality == "text2video":
|
||||
formatted = _format_text2video_prompts(prompts, args.num_frames, args.video_height, args.video_width)
|
||||
elif args.modality == "video2text":
|
||||
formatted = _format_video2text_prompts(prompts, args.video_path, system_prompt=args.system_prompt)
|
||||
elif args.modality == "img2img":
|
||||
formatted = _format_img2img_prompts(prompts, args.image_path)
|
||||
elif args.modality == "video2video":
|
||||
formatted = _format_video2video_prompts(
|
||||
prompts, args.video_path, args.num_frames, args.video_height, args.video_width
|
||||
)
|
||||
elif args.modality == "image2video":
|
||||
formatted = _format_image2video_prompts(
|
||||
prompts, args.image_path, args.num_frames, args.video_height, args.video_width
|
||||
)
|
||||
else:
|
||||
formatted = _format_text2img_prompts(prompts)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
diffusion_params = params_list[0] # single-stage: one param set
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
if args.seed is not None:
|
||||
diffusion_params.seed = args.seed # type: ignore
|
||||
if args.height is not None:
|
||||
diffusion_params.height = args.height # type: ignore
|
||||
if args.width is not None:
|
||||
diffusion_params.width = args.width # type: ignore
|
||||
extra = getattr(diffusion_params, "extra_args", {}) or {}
|
||||
extra["cfg_text_scale"] = args.cfg_text_scale
|
||||
extra["timestep_shift"] = args.timestep_shift
|
||||
if args.modality in ("img2text", "video2text"):
|
||||
extra["max_think_tokens"] = args.max_text_tokens
|
||||
extra["do_sample"] = args.do_sample
|
||||
extra["text_temperature"] = args.text_temperature
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
|
||||
outputs = list(omni.generate(prompts=formatted, sampling_params_list=params_list))
|
||||
|
||||
if args.modality in ("img2text", "video2text"):
|
||||
for i, req_output in enumerate(outputs):
|
||||
text = getattr(req_output, "output", None) or getattr(req_output, "text", None)
|
||||
if text is None and hasattr(req_output, "custom_output"):
|
||||
text = (req_output.custom_output or {}).get("text_output")
|
||||
print(f"[Output {i}] {text}")
|
||||
return
|
||||
|
||||
if args.modality in ("text2video", "video2video", "image2video"):
|
||||
for i, req_output in enumerate(outputs):
|
||||
custom = getattr(req_output, "custom_output", None) or {}
|
||||
frames = custom.get("video_frames")
|
||||
if not frames:
|
||||
print(f"[Output {i}] no video frames returned")
|
||||
continue
|
||||
mp4_path = os.path.join(args.output, f"lance_{i}.mp4")
|
||||
try:
|
||||
import imageio.v3 as iio
|
||||
import numpy as np
|
||||
|
||||
arr = np.stack([np.asarray(f) for f in frames], axis=0)
|
||||
iio.imwrite(mp4_path, arr, fps=args.fps, codec="libx264", quality=8)
|
||||
print(f"[Output {i}] Saved MP4 ({len(frames)} frames, {arr.shape[1]}x{arr.shape[2]}) to {mp4_path}")
|
||||
except Exception as exc:
|
||||
print(f"[Output {i}] MP4 encode failed ({exc!r}); writing per-frame ONGs instead")
|
||||
frame_dir = os.path.join(args.output, f"lance_{i}_frames")
|
||||
os.makedirs(frame_dir, exist_ok=True)
|
||||
for j, frame in enumerate(frames):
|
||||
frame.save(os.path.join(frame_dir, f"{j:04d}.png"))
|
||||
print(f" → {frame_dir} ({len(frames)} frames)")
|
||||
return
|
||||
|
||||
idx = 0
|
||||
for req_output in outputs:
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images:
|
||||
continue
|
||||
for j, img in enumerate(images):
|
||||
path = os.path.join(args.output, f"lance_{idx}_{j}.png")
|
||||
img.save(path)
|
||||
print(f"[Output] Saved image to {path}")
|
||||
idx += 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large.
Load diff
@@ -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.
|
||||
+184
@@ -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()
|
||||
Executable
+15
@@ -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"
|
||||
@@ -0,0 +1,177 @@
|
||||
# Lance
|
||||
|
||||
> Unified autoregressive + diffusion multimodal (text / image / video)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: ByteDance
|
||||
- Model: [`bytedance-research/Lance`](https://huggingface.co/bytedance-research/Lance) (`Lance_3B`, `Lance_3B_Video`)
|
||||
- Task: text2img, text2video, img2img (image edit), video2video (video edit),
|
||||
img2text (image understanding), video2text (video understanding)
|
||||
- Mode: Offline inference, Online serving (OpenAI-compatible API)
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to run Lance — a 3B unified AR + diffusion multimodal model
|
||||
on a Qwen2.5-VL backbone — via vLLM-Omni. Lance is **BAGEL-lineage**: the
|
||||
released checkpoint uses the same `*_moe_gen` Mixture-of-Transformers weight
|
||||
layout as BAGEL, so vLLM-Omni reuses BAGEL's transformer core and specializes
|
||||
only the ViT (Qwen2.5-VL vision), the VAE (Wan2.2) and the checkpoint layout.
|
||||
|
||||
All six single-stage modalities are supported on the `Lance_3B` checkpoint;
|
||||
`text2video` and `video2video` additionally need `Lance_3B_Video` for the 3-D
|
||||
`latent_pos_embed` table. The two-stage AR + DiT topology is deferred and
|
||||
needs `LanceConfig` / `LanceProcessor` registered in the `vllm` package (a
|
||||
separate upstream PR).
|
||||
|
||||
## References
|
||||
|
||||
- Offline example:
|
||||
[`examples/offline_inference/lance/end2end.py`](../../examples/offline_inference/lance/end2end.py)
|
||||
- Online serving:
|
||||
[`examples/online_serving/lance/`](../../examples/online_serving/lance/)
|
||||
- Deploy config:
|
||||
[`vllm_omni/deploy/lance.yaml`](../../vllm_omni/deploy/lance.yaml)
|
||||
- E2E tests:
|
||||
[`tests/e2e/online_serving/test_lance.py`](../../tests/e2e/online_serving/test_lance.py)
|
||||
- HuggingFace model page:
|
||||
[bytedance-research/Lance](https://huggingface.co/bytedance-research/Lance)
|
||||
- PR: [vllm-project/vllm-omni#3710](https://github.com/vllm-project/vllm-omni/pull/3710)
|
||||
(closes [#3697](https://github.com/vllm-project/vllm-omni/issues/3697))
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA B300 / A100 80GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12
|
||||
- Driver / runtime: CUDA ≥ 12.4
|
||||
- vLLM-Omni version: 0.18.x.dev
|
||||
|
||||
#### Command — text-to-image (default)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance \
|
||||
--prompts "a corgi astronaut on the moon, cinematic" \
|
||||
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
|
||||
--height 1024 --width 1024 \
|
||||
--seed 42 --output ./out
|
||||
```
|
||||
|
||||
Defaults match upstream `inference_lance.sh`: 30 denoising steps,
|
||||
`timestep-shift 3.5`, text CFG 4.0, seed 42, 1024×1024.
|
||||
|
||||
#### Command — image edit (img2img)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance --modality img2img \
|
||||
--image-path /path/to/input.png \
|
||||
--prompts "Convert this into a vibrant cartoon-style illustration" \
|
||||
--steps 30 --cfg-text-scale 4.0 --timestep-shift 3.5 \
|
||||
--output ./out
|
||||
```
|
||||
|
||||
The Lance-native VAE prefill scatters Wan2.2 latents into the LLM query
|
||||
sequence; no separate image encoder is needed.
|
||||
|
||||
#### Command — text-to-video
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance/Lance_3B_Video --modality text2video \
|
||||
--num-frames 25 --video-height 480 --video-width 768 \
|
||||
--prompts "a cat playing piano, cinematic" \
|
||||
--steps 30 --fps 8 --output ./out
|
||||
```
|
||||
|
||||
Use the `Lance_3B_Video` subfolder for any video path so the 3-D
|
||||
`latent_pos_embed` table is loaded; image / understanding paths can point at
|
||||
the top-level repo and resolve the right sub-checkpoint automatically.
|
||||
|
||||
#### Command — image / video understanding
|
||||
|
||||
```bash
|
||||
# Image → text (caption / VQA)
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance --modality img2text \
|
||||
--image-path /path/to/photo.jpg \
|
||||
--prompts "Describe this image in detail." \
|
||||
--do-sample --text-temperature 0.8
|
||||
|
||||
# Video → text
|
||||
python examples/offline_inference/lance/end2end.py \
|
||||
--model bytedance-research/Lance --modality video2text \
|
||||
--video-path /path/to/clip.mp4 \
|
||||
--prompts "What is happening in this video?"
|
||||
```
|
||||
|
||||
Sampling is enabled by default at `--text-temperature 0.8` for the
|
||||
understanding paths because Lance's greedy decoder emits an immediate EOS for
|
||||
many prompts.
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/e2e/online_serving/test_lance.py
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- BF16 footprint: ~7 GB LLM + Qwen2.5-VL ViT + Wan2.2 VAE; comfortably fits
|
||||
on a single 16 GB+ GPU for `Lance_3B`.
|
||||
- `rope_scaling = {"type": "mrope", "mrope_section": [16, 24, 24]}` is wired
|
||||
through `BagelRotaryEmbedding`. `text2img` uses scalar position ids
|
||||
(BAGEL-equivalent); `img2text` / `video` / edit paths thread per-axis 3-D
|
||||
position ids.
|
||||
- `video_edit` quality is more abstract than `text2video` at the same
|
||||
resolution (known position-id offset between VAE-prefill and gen-latent
|
||||
blocks). Functionally correct end-to-end.
|
||||
|
||||
## Online Serving
|
||||
|
||||
Lance supports all single-stage modalities via the OpenAI-compatible
|
||||
`/v1/chat/completions` API.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
bash examples/online_serving/lance/run_server.sh
|
||||
# or, with overrides
|
||||
MODEL=bytedance-research/Lance \
|
||||
DEPLOY_CONFIG=vllm_omni/deploy/lance.yaml \
|
||||
PORT=8091 \
|
||||
bash examples/online_serving/lance/run_server.sh
|
||||
```
|
||||
|
||||
For `text2video` / `video2video`, set
|
||||
`MODEL=bytedance-research/Lance/Lance_3B_Video`.
|
||||
|
||||
### Send requests
|
||||
|
||||
```bash
|
||||
# Text-to-image
|
||||
python examples/online_serving/lance/openai_chat_client.py \
|
||||
--prompt "A cute corgi astronaut on the moon, cinematic" \
|
||||
--modality text2img --output corgi.png
|
||||
|
||||
# Image edit
|
||||
python examples/online_serving/lance/openai_chat_client.py \
|
||||
--prompt "Convert this into a vibrant cartoon-style illustration" \
|
||||
--modality img2img --image-url path/to/photo.png \
|
||||
--output edited.png
|
||||
|
||||
# Image understanding
|
||||
python examples/online_serving/lance/openai_chat_client.py \
|
||||
--prompt "Describe this image" \
|
||||
--modality img2text --image-url photo.jpg
|
||||
```
|
||||
|
||||
The client is shared with BAGEL — same OpenAI message format, same
|
||||
`modalities` / `num_inference_steps` / `seed` / `height` / `width` knobs.
|
||||
@@ -29,6 +29,7 @@ recipes/
|
||||
|--------|------|----------|
|
||||
| [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB |
|
||||
| [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB |
|
||||
| [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB |
|
||||
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
|
||||
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
|
||||
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
|
||||
|
||||
@@ -19,8 +19,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
NUM_TOKENS = 8
|
||||
HIDDEN_DIM = 16
|
||||
NUM_TIMESTEPS = 5
|
||||
# generate_image uses timesteps[:-1], so actual steps = NUM_TIMESTEPS - 1
|
||||
EXPECTED_STEPS = NUM_TIMESTEPS - 1
|
||||
# generate_image samples (num_timesteps + 1) points then drops the last, matching
|
||||
# upstream lance: see Bagel.generate_image in bagel_transformer.py.
|
||||
EXPECTED_STEPS = NUM_TIMESTEPS
|
||||
|
||||
|
||||
def _make_mock_bagel(mocker: MockerFixture):
|
||||
@@ -178,7 +179,7 @@ class TestTrajectoryRecording:
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
# Recompute the expected timestep schedule (mirrors generate_image logic)
|
||||
ts = torch.linspace(1, 0, args["num_timesteps"])
|
||||
ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
|
||||
shift = args.get("timestep_shift", 1.0)
|
||||
ts = shift * ts / (1 + (shift - 1) * ts)
|
||||
expected_timesteps = ts[:-1] # last element is dropped
|
||||
@@ -196,7 +197,7 @@ class TestTrajectoryRecording:
|
||||
bagel, args = bagel_and_args
|
||||
|
||||
# Compute the number of denoising steps the same way generate_image does
|
||||
ts = torch.linspace(1, 0, args["num_timesteps"])
|
||||
ts = torch.linspace(1, 0, args["num_timesteps"] + 1)
|
||||
shift = args.get("timestep_shift", 1.0)
|
||||
ts = shift * ts / (1 + (shift - 1) * ts)
|
||||
num_steps = len(ts) - 1 # timesteps = ts[:-1]
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end online serving test for Lance (single-stage).
|
||||
|
||||
Verifies that the Lance pipeline can serve text-to-image and image-edit
|
||||
requests via the OpenAI-compatible chat completions API exposed by
|
||||
``vllm-omni serve``.
|
||||
|
||||
Equivalent to running:
|
||||
|
||||
vllm-omni serve "bytedance-research/Lance" --omni \\
|
||||
--pipeline lance --enforce-eager --trust-remote-code --port 8091
|
||||
|
||||
# text2img
|
||||
python3 examples/online_serving/lance/openai_chat_client.py \\
|
||||
--prompt "A cute corgi astronaut" --modality text2img
|
||||
"""
|
||||
|
||||
import base64
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
import pytest
|
||||
from vllm.assets.image import ImageAsset
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from vllm_omni.diffusion.models.lance.prompts import (
|
||||
VIDEO_PAD,
|
||||
VISION_END,
|
||||
VISION_START,
|
||||
render_lance_prompt,
|
||||
)
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
MODEL = "bytedance-research/Lance"
|
||||
|
||||
TEXT2IMG_PROMPT = "A cute corgi astronaut on the moon, cinematic"
|
||||
IMG2IMG_PROMPT = "Convert this into a vibrant cartoon-style illustration"
|
||||
|
||||
# Lance vision block: <|vision_start|><|video_pad|><|vision_end|>
|
||||
_VISION_BLOCK = f"{VISION_START}{VIDEO_PAD}{VISION_END}"
|
||||
|
||||
# Lance is single-stage diffusion — no deploy YAML. Pass the engine
|
||||
# knobs that used to live in ``vllm_omni/deploy/lance.yaml`` via CLI
|
||||
# flags so ``create_default_diffusion`` builds the stage config.
|
||||
_LANCE_SERVE_ARGS = [
|
||||
"--pipeline",
|
||||
"lance",
|
||||
"--max-num-batched-tokens",
|
||||
"32768",
|
||||
"--max-num-seqs",
|
||||
"1",
|
||||
"--enforce-eager",
|
||||
"--trust-remote-code",
|
||||
"--no-enable-prefix-caching",
|
||||
"--no-async-chunk",
|
||||
]
|
||||
|
||||
test_params = [
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
server_args=_LANCE_SERVE_ARGS,
|
||||
stage_init_timeout=300,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _build_text2img_messages(prompt: str) -> list[dict]:
|
||||
"""Build OpenAI-format messages for text2img generation."""
|
||||
rendered = render_lance_prompt("t2i", prompt)
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": rendered}],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
|
||||
"""Build OpenAI-format messages for img2img generation."""
|
||||
rendered = render_lance_prompt("image_edit", prompt, vision_token=_VISION_BLOCK)
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": rendered},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_lance_text2img_online(omni_server, openai_client) -> None:
|
||||
"""Lance text2img via the OpenAI-compatible chat completions API."""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_text2img_messages(TEXT2IMG_PROMPT),
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 0.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_lance_img2img_online(omni_server, openai_client) -> None:
|
||||
"""Lance image_edit via the OpenAI-compatible chat completions API."""
|
||||
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
|
||||
buffer = BytesIO()
|
||||
input_image.save(buffer, format="JPEG")
|
||||
image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64),
|
||||
"modalities": ["image"],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 2,
|
||||
"guidance_scale": 0.0,
|
||||
"seed": 42,
|
||||
},
|
||||
}
|
||||
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
@@ -65,6 +65,11 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
|
||||
"vllm_omni.model_executor.models.bagel.pipeline",
|
||||
"BAGEL_SINGLE_STAGE_PIPELINE",
|
||||
),
|
||||
# Lance (ByteDance) — BAGEL-lineage unified AR+diffusion, single-stage DiT.
|
||||
"lance": (
|
||||
"vllm_omni.model_executor.models.lance.pipeline",
|
||||
"LANCE_PIPELINE",
|
||||
),
|
||||
"glm_image": (
|
||||
"vllm_omni.model_executor.models.glm_image.pipeline",
|
||||
"GLM_IMAGE_PIPELINE",
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
# Lance (ByteDance) — single-stage diffusion pipeline selector.
|
||||
#
|
||||
# Lance's HF config.json is descriptive metadata (no ``model_type``),
|
||||
# so the loader can't auto-detect from the model dir alone. This YAML
|
||||
# only tells the loader "use the Lance pipeline" — one stage, no multi-
|
||||
# stage orchestration.
|
||||
|
||||
pipeline: lance
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_batched_tokens: 32768
|
||||
max_num_seqs: 1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
devices: "0"
|
||||
default_sampling_params:
|
||||
seed: 42
|
||||
@@ -847,6 +847,21 @@ class OmniDiffusionConfig:
|
||||
self.supports_multimodal_inputs = metadata.supports_multimodal_inputs
|
||||
self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs
|
||||
|
||||
@staticmethod
|
||||
def _looks_like_lance_subfolder(model: str | None) -> bool:
|
||||
"""Return True when ``--model`` points at a Lance per-component subfolder.
|
||||
|
||||
Lance's HF repo bundles ``Lance_3B/``, ``Lance_3B_Video/`` and
|
||||
``Qwen2.5-VL-ViT/`` under a single top-level ``config.json``; users may
|
||||
reasonably hand the AR-style sub-checkpoint path directly. The
|
||||
``LancePipeline`` constructor knows to walk up to the repo root from
|
||||
either subfolder name.
|
||||
"""
|
||||
if not model:
|
||||
return False
|
||||
base = os.path.basename(str(model).rstrip("/"))
|
||||
return base in {"Lance_3B", "Lance_3B_Video"}
|
||||
|
||||
def enrich_config(self) -> None:
|
||||
"""Load model metadata from HuggingFace and populate config fields.
|
||||
|
||||
@@ -900,6 +915,15 @@ class OmniDiffusionConfig:
|
||||
else:
|
||||
cfg = get_hf_file_to_dict("config.json", self.model)
|
||||
if cfg is None:
|
||||
# Lance ships its top-level config.json one directory above
|
||||
# the per-checkpoint subfolders (``Lance_3B/`` or
|
||||
# ``Lance_3B_Video/``). Try to recover that case before
|
||||
# raising.
|
||||
if self._looks_like_lance_subfolder(self.model):
|
||||
self.model_class_name = "LancePipeline"
|
||||
self.set_tf_model_config(TransformerConfig())
|
||||
self.update_multimodal_support()
|
||||
return
|
||||
raise ValueError(f"Could not find config.json or model_index.json for model {self.model}")
|
||||
|
||||
self.set_tf_model_config(TransformerConfig.from_dict(cfg))
|
||||
@@ -910,6 +934,19 @@ class OmniDiffusionConfig:
|
||||
self.model_class_name = "BagelPipeline"
|
||||
self.set_tf_model_config(TransformerConfig())
|
||||
self.update_multimodal_support()
|
||||
elif (
|
||||
model_type == "lance"
|
||||
or "LancePipeline" in architectures
|
||||
or cfg.get("model_name") == "Lance"
|
||||
or self._looks_like_lance_subfolder(self.model)
|
||||
):
|
||||
# Lance ships a non-HF top-level config.json (model_name only)
|
||||
# plus per-component subfolders; resolve to the Lance pipeline.
|
||||
# Also accept --model pointing directly at the ``Lance_3B`` or
|
||||
# ``Lance_3B_Video`` subfolder by walking up to the repo root.
|
||||
self.model_class_name = "LancePipeline"
|
||||
self.set_tf_model_config(TransformerConfig())
|
||||
self.update_multimodal_support()
|
||||
elif model_type == "neo_chat":
|
||||
self.model_class_name = "SenseNovaU1Pipeline"
|
||||
self.tf_model_config = TransformerConfig()
|
||||
|
||||
@@ -106,7 +106,16 @@ class BagelRotaryEmbedding(nn.Module):
|
||||
linear, dynamic-NTK, YaRN, …), we delegate the ``inv_freq`` /
|
||||
``attention_scaling`` computation to HF's ``ROPE_INIT_FUNCTIONS`` so
|
||||
that the frequency basis and scaling factor are identical to the
|
||||
original checkpoint. This module has no learnable parameters.
|
||||
original checkpoint.
|
||||
|
||||
For Qwen2.5-VL-style multimodal RoPE (``rope_scaling.rope_type == "mrope"``)
|
||||
the ``inv_freq`` basis is the standard default-rope one; the difference
|
||||
is that position ids are 3-D ``(t, h, w)`` per token and the
|
||||
``mrope_section`` describes how the head dimension is split across axes.
|
||||
This module accepts either 2-D scalar position ids ``(B, S)`` or 3-D
|
||||
multimodal position ids ``(B, 3, S)`` and dispatches accordingly so the
|
||||
same module works for both BAGEL (1-D rope) and Lance (Qwen2.5-VL mrope).
|
||||
This module has no learnable parameters.
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
@@ -120,10 +129,17 @@ class BagelRotaryEmbedding(nn.Module):
|
||||
rope_type = (
|
||||
rope_scaling.get("rope_type") or rope_scaling.get("type") or rope_parameters.get("rope_type") or "default"
|
||||
)
|
||||
# Cache mrope_section for the forward-time section-split.
|
||||
self._mrope_section: list[int] | None = None
|
||||
if rope_type == "mrope":
|
||||
section = rope_scaling.get("mrope_section") or rope_parameters.get("mrope_section")
|
||||
if section is None:
|
||||
raise ValueError("rope_scaling.rope_type == 'mrope' requires 'mrope_section'.")
|
||||
self._mrope_section = list(section)
|
||||
|
||||
if rope_type == "default":
|
||||
# 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")
|
||||
@@ -147,11 +163,36 @@ class BagelRotaryEmbedding(nn.Module):
|
||||
|
||||
Args:
|
||||
x: Input tensor (only used for dtype inference).
|
||||
position_ids: Position indices, shape (batch_size, seq_len).
|
||||
position_ids: Either 2-D scalar ``(batch_size, seq_len)`` for plain
|
||||
1-D RoPE, or 3-D multimodal ``(batch_size, 3, seq_len)`` for
|
||||
Qwen2.5-VL-style mRoPE. The latter is auto-detected from
|
||||
``position_ids.ndim``.
|
||||
|
||||
Returns:
|
||||
cos, sin: Rotary embeddings, each of shape (batch_size, seq_len, dim).
|
||||
"""
|
||||
if position_ids.ndim == 3 and self._mrope_section is not None:
|
||||
# multimodal path: position_ids is (B, 3, S) with rows = (t, h, w).
|
||||
# Compute per-axis frequencies, then assemble the per-section
|
||||
# rotary basis matching Qwen2-VL's ``apply_multimodal_rotary_pos_emb``.
|
||||
B = position_ids.shape[0]
|
||||
inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(B, 3, -1, 1)
|
||||
position_ids_expanded = position_ids[:, :, None, :].float()
|
||||
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(2, 3)
|
||||
# ``freqs`` is (B, 3, S, head_dim/2); double along the last axis
|
||||
# to get the full ``head_dim`` rotary basis.
|
||||
emb = torch.cat((freqs, freqs), dim=-1) # (B, 3, S, head_dim)
|
||||
cos_per_axis = emb.cos() * self.attention_scaling
|
||||
sin_per_axis = emb.sin() * self.attention_scaling
|
||||
# ``mrope_section`` (e.g. [16, 24, 24] for Qwen2.5-VL) sums to
|
||||
# head_dim/2; doubled it sums to head_dim and cycles axis = i % 3.
|
||||
sec_full = self._mrope_section * 2
|
||||
cos_split = cos_per_axis.split(sec_full, dim=-1)
|
||||
sin_split = sin_per_axis.split(sec_full, dim=-1)
|
||||
cos = torch.cat([c[:, i % 3] for i, c in enumerate(cos_split)], dim=-1)
|
||||
sin = torch.cat([s[:, i % 3] for i, s in enumerate(sin_split)], dim=-1)
|
||||
return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
|
||||
|
||||
inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
|
||||
position_ids_expanded = position_ids[:, None, :].float()
|
||||
freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
|
||||
@@ -541,8 +582,12 @@ class PackedAttentionMoT(nn.Module):
|
||||
q = q.view(-1, self.num_heads, self.head_dim)
|
||||
k = k.view(-1, self.num_kv_heads, self.head_dim)
|
||||
v = v.view(-1, self.num_kv_heads, self.head_dim)
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
# Pre-merge code cast to float32 before q_norm/k_norm — bf16
|
||||
# RMSNorm accumulation drift compounds across segmented prefill
|
||||
# (x2t builds the cache via 3+ sequential forward_cache_update_*
|
||||
# calls), producing degenerate token repetition in generate_text.
|
||||
q = self.q_norm(q.to(torch.float32))
|
||||
k = self.k_norm(k.to(torch.float32))
|
||||
|
||||
cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
|
||||
q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
|
||||
@@ -557,16 +602,52 @@ class PackedAttentionMoT(nn.Module):
|
||||
cache_v = past_key_values.value_cache[self.layer_idx]
|
||||
full_k = torch.cat([cache_k, k], dim=0)
|
||||
full_v = torch.cat([cache_v, v], dim=0)
|
||||
cache_len = cache_k.shape[0]
|
||||
else:
|
||||
full_k = k
|
||||
full_v = v
|
||||
cache_len = 0
|
||||
|
||||
attn = self.attn_causal if is_causal else self.attn_noncausal
|
||||
attn_out = attn(
|
||||
q.unsqueeze(0),
|
||||
full_k.unsqueeze(0),
|
||||
full_v.unsqueeze(0),
|
||||
)
|
||||
if is_causal and cache_len > 0:
|
||||
# PyTorch SDPA's ``is_causal=True`` with ``Q != K`` uses
|
||||
# ``tril(diagonal=0)`` — top-left aligned, so for Q=1 against a
|
||||
# long cache only ``q_0 -> k_0`` is unmasked (degenerate
|
||||
# generate_text output). Build the correct bottom-right
|
||||
# aligned mask manually: ``mask[i, j] = -inf if j > cache_len
|
||||
# + i`` so each new query attends to all of ``cache + self
|
||||
# up to i``. Bypass ``DiffusionAttention`` entirely — its
|
||||
# ``_maybe_reshape_attn_mask`` helper only handles 2-D
|
||||
# ``(B, K)`` shape, not ``(Q, K)``, so the reshape path
|
||||
# silently drops the per-query mask for Q>1.
|
||||
Q_len = q.shape[0]
|
||||
K_len = full_k.shape[0]
|
||||
arange_q = torch.arange(Q_len, device=q.device).unsqueeze(1)
|
||||
arange_k = torch.arange(K_len, device=q.device).unsqueeze(0)
|
||||
mask = arange_k > (cache_len + arange_q) # (Q, K) bool
|
||||
attn_bias = torch.zeros(Q_len, K_len, dtype=q.dtype, device=q.device)
|
||||
attn_bias.masked_fill_(mask, float("-inf"))
|
||||
# Permute to (B=1, H, S, D) for SDPA.
|
||||
q_4d = q.unsqueeze(0).permute(0, 2, 1, 3)
|
||||
k_4d = full_k.unsqueeze(0).permute(0, 2, 1, 3)
|
||||
v_4d = full_v.unsqueeze(0).permute(0, 2, 1, 3)
|
||||
attn_out_4d = torch.nn.functional.scaled_dot_product_attention(
|
||||
q_4d,
|
||||
k_4d,
|
||||
v_4d,
|
||||
attn_mask=attn_bias.unsqueeze(0).unsqueeze(0), # (1, 1, Q, K)
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=1.0 / (self.head_dim**0.5),
|
||||
enable_gqa=(self.num_heads != self.num_kv_heads),
|
||||
)
|
||||
attn_out = attn_out_4d.permute(0, 2, 1, 3)
|
||||
else:
|
||||
attn = self.attn_causal if is_causal else self.attn_noncausal
|
||||
attn_out = attn(
|
||||
q.unsqueeze(0),
|
||||
full_k.unsqueeze(0),
|
||||
full_v.unsqueeze(0),
|
||||
)
|
||||
|
||||
attn_out = attn_out.squeeze(0).reshape(-1, self.q_size)
|
||||
attn_out, _ = self.o_proj(attn_out)
|
||||
@@ -1503,23 +1584,6 @@ class Bagel(nn.Module):
|
||||
|
||||
return generation_input
|
||||
|
||||
@staticmethod
|
||||
def _merge_naive_caches(caches: list) -> NaiveCache:
|
||||
"""Merge multiple NaiveCache objects by concatenating KV tensors per layer."""
|
||||
if not caches:
|
||||
# Handle empty list case gracefully if desired,
|
||||
# though original code also crashed on this.
|
||||
return NaiveCache(0)
|
||||
|
||||
num_layers = len(caches[0].key_cache)
|
||||
merged = NaiveCache(num_layers)
|
||||
for layer_idx in range(num_layers):
|
||||
key_parts = [c.key_cache[layer_idx] for c in caches if c.key_cache[layer_idx] is not None]
|
||||
val_parts = [c.value_cache[layer_idx] for c in caches if c.value_cache[layer_idx] is not None]
|
||||
merged.key_cache[layer_idx] = torch.cat(key_parts, dim=0) if key_parts else None
|
||||
merged.value_cache[layer_idx] = torch.cat(val_parts, dim=0) if val_parts else None
|
||||
return merged
|
||||
|
||||
def prepare_start_tokens(self, curr_kvlens, curr_rope, new_token_ids):
|
||||
"""Prepare start tokens for autoregressive text generation.
|
||||
|
||||
@@ -1614,10 +1678,26 @@ class Bagel(nn.Module):
|
||||
return_trajectory_latents: bool = False,
|
||||
scheduler: object | None = None,
|
||||
scheduler_kwargs: dict | None = None,
|
||||
# Lance i2v: tokens to freeze at their initial (encoded-image)
|
||||
# value throughout the denoise loop. Matches upstream's
|
||||
# ``mse_loss_indexes``-exclusion behaviour from PR #33: cond
|
||||
# positions are never updated and get ``timestep=0``.
|
||||
frame_condition_token_indexes: torch.LongTensor | None = None,
|
||||
):
|
||||
x_t = packed_init_noises
|
||||
# Snapshot the pinned subtensor BEFORE the denoise loop touches
|
||||
# x_t; the cond positions in packed_init_noises hold the
|
||||
# VAE-encoded conditioning latent that must be preserved verbatim.
|
||||
pinned_x_t = None
|
||||
if frame_condition_token_indexes is not None:
|
||||
frame_condition_token_indexes = frame_condition_token_indexes.to(x_t.device).long()
|
||||
pinned_x_t = x_t[frame_condition_token_indexes].clone()
|
||||
|
||||
timesteps = torch.linspace(1, 0, num_timesteps, device=x_t.device)
|
||||
# Use num_timesteps + 1 sample points so we get `num_timesteps` denoise
|
||||
# steps after dropping the terminal t=0 (which has no dt). Upstream
|
||||
# Lance / BAGEL both use this convention; without the +1 we silently
|
||||
# run one fewer denoise iteration than the user asked for.
|
||||
timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device)
|
||||
timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps)
|
||||
dts = timesteps[:-1] - timesteps[1:]
|
||||
timesteps = timesteps[:-1]
|
||||
@@ -1669,6 +1749,12 @@ class Bagel(nn.Module):
|
||||
trajectory_latents.append(x_t.clone())
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
if frame_condition_token_indexes is not None:
|
||||
# Cond positions stay at t=0 (clean signal). Matches upstream
|
||||
# PR #33 lance.py line 1605:
|
||||
# timestep[current_vae_mse_indexes_local_in_vae] = t
|
||||
# (cond positions remain at the ``torch.zeros`` init value).
|
||||
timestep[frame_condition_token_indexes] = 0.0
|
||||
in_cfg_window = t > cfg_interval[0] and t <= cfg_interval[1]
|
||||
cfg_text_scale_ = cfg_text_scale if in_cfg_window else 1.0
|
||||
cfg_img_scale_ = cfg_img_scale if in_cfg_window else 1.0
|
||||
@@ -1732,6 +1818,12 @@ class Bagel(nn.Module):
|
||||
trajectory_latents.append(x_t.clone())
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
if frame_condition_token_indexes is not None:
|
||||
# Cond positions stay at t=0 (clean signal). Matches upstream
|
||||
# PR #33 lance.py line 1605:
|
||||
# timestep[current_vae_mse_indexes_local_in_vae] = t
|
||||
# (cond positions remain at the ``torch.zeros`` init value).
|
||||
timestep[frame_condition_token_indexes] = 0.0
|
||||
v_t = self.forward_single_branch(
|
||||
x_t=x_t,
|
||||
timestep=timestep,
|
||||
@@ -1758,46 +1850,31 @@ class Bagel(nn.Module):
|
||||
unpacked_latent = x_t.split((packed_seqlens - 2).tolist())
|
||||
return unpacked_latent, trajectory_latents, trajectory_timesteps, trajectory_log_probs
|
||||
|
||||
# ── Batched CFG mode (cfg_parallel_size=1, no SP) ──
|
||||
cfg_batched = None
|
||||
# ── Sequential CFG mode (cfg_parallel_size=1, no SP) ──
|
||||
# Each CFG branch runs its own LLM forward; we just need the
|
||||
# per-branch packed_position_ids and past_key_values for
|
||||
# ``Bagel.forward`` to dispatch through.
|
||||
cfg_branches: dict | None = None
|
||||
|
||||
if use_cfg_text:
|
||||
seq_len = int(packed_seqlens.sum())
|
||||
|
||||
# Branch 0: main (gen_context), always present
|
||||
branches_pid = [packed_position_ids]
|
||||
branches_cache = [past_key_values]
|
||||
|
||||
# Branch 1: cfg_text (unconditional text), always present when use_cfg_text
|
||||
branches_pid.append(cfg_text_packed_position_ids)
|
||||
branches_cache.append(cfg_text_past_key_values)
|
||||
|
||||
# Branch 2: cfg_img (text-only, no image), optional
|
||||
branches_pid = [packed_position_ids, cfg_text_packed_position_ids]
|
||||
branches_cache = [past_key_values, cfg_text_past_key_values]
|
||||
if use_cfg_img:
|
||||
branches_pid.append(cfg_img_packed_position_ids)
|
||||
branches_cache.append(cfg_img_past_key_values)
|
||||
|
||||
num_branches = len(branches_cache)
|
||||
|
||||
cfg_batched = {
|
||||
"num_branches": num_branches,
|
||||
"seq_len": seq_len,
|
||||
"batched_query_lens": packed_seqlens.repeat(num_branches),
|
||||
"batched_position_ids": torch.cat(branches_pid),
|
||||
"batched_text_indexes": torch.cat(
|
||||
[packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)]
|
||||
),
|
||||
"batched_vae_indexes": torch.cat(
|
||||
[packed_vae_token_indexes + b_idx * seq_len for b_idx in range(num_branches)]
|
||||
),
|
||||
"merged_cache": self._merge_naive_caches(branches_cache),
|
||||
}
|
||||
cfg_branches = {"pids": branches_pid, "caches": branches_cache}
|
||||
|
||||
if return_trajectory_latents and len(timesteps) > 0:
|
||||
trajectory_latents.append(x_t.clone())
|
||||
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
if frame_condition_token_indexes is not None:
|
||||
# Cond positions stay at t=0 (clean signal). Matches upstream
|
||||
# PR #33 lance.py line 1605:
|
||||
# timestep[current_vae_mse_indexes_local_in_vae] = t
|
||||
# (cond positions remain at the ``torch.zeros`` init value).
|
||||
timestep[frame_condition_token_indexes] = 0.0
|
||||
if t > cfg_interval[0] and t <= cfg_interval[1]:
|
||||
cfg_text_scale_ = cfg_text_scale
|
||||
cfg_img_scale_ = cfg_img_scale
|
||||
@@ -1818,7 +1895,7 @@ class Bagel(nn.Module):
|
||||
cfg_renorm_type=cfg_renorm_type,
|
||||
cfg_text_scale=cfg_text_scale_,
|
||||
cfg_img_scale=cfg_img_scale_,
|
||||
cfg_batched=cfg_batched,
|
||||
cfg_branches=cfg_branches,
|
||||
)
|
||||
|
||||
if scheduler is not None:
|
||||
@@ -1828,6 +1905,12 @@ class Bagel(nn.Module):
|
||||
trajectory_log_probs.append(out.log_prob)
|
||||
else:
|
||||
x_t = x_t - v_t.to(x_t.device) * dts[i] # velocity pointing from data to noise
|
||||
if pinned_x_t is not None:
|
||||
# i2v: restore cond positions to their encoded-image
|
||||
# latent. Matches upstream PR #33 lance.py line 1712:
|
||||
# x_t[mse_indexes] = x_t[mse_indexes] - v_t[mse_indexes] * dts[i]
|
||||
# (cond positions are excluded from the update).
|
||||
x_t[frame_condition_token_indexes] = pinned_x_t
|
||||
if return_trajectory_latents:
|
||||
trajectory_latents.append(x_t.clone())
|
||||
trajectory_timesteps.append(timesteps[i])
|
||||
@@ -1859,6 +1942,7 @@ class Bagel(nn.Module):
|
||||
return_trajectory_latents: bool = False,
|
||||
scheduler: object | None = None,
|
||||
scheduler_kwargs: dict | None = None,
|
||||
frame_condition_token_indexes: torch.LongTensor | None = None,
|
||||
):
|
||||
"""CFG parallel denoising loop: each rank computes one CFG branch.
|
||||
|
||||
@@ -1918,6 +2002,12 @@ class Bagel(nn.Module):
|
||||
|
||||
for i, t in enumerate(timesteps):
|
||||
timestep = torch.tensor([t] * x_t.shape[0], device=x_t.device)
|
||||
if frame_condition_token_indexes is not None:
|
||||
# Cond positions stay at t=0 (clean signal). Matches upstream
|
||||
# PR #33 lance.py line 1605:
|
||||
# timestep[current_vae_mse_indexes_local_in_vae] = t
|
||||
# (cond positions remain at the ``torch.zeros`` init value).
|
||||
timestep[frame_condition_token_indexes] = 0.0
|
||||
use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0
|
||||
|
||||
if use_cfg_this_step:
|
||||
@@ -2070,7 +2160,7 @@ class Bagel(nn.Module):
|
||||
packed_sequence = packed_text_embedding.new_zeros((int(local_seqlens.sum()), self.hidden_size))
|
||||
packed_sequence[local_text_indexes] = packed_text_embedding
|
||||
|
||||
assert timestep.unique().shape[0] == 1
|
||||
# i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
|
||||
packed_pos_embed = self.latent_pos_embed(local_vae_pos_ids)
|
||||
local_timestep = timestep[: local_x_t.shape[0]]
|
||||
packed_timestep_embeds = self.time_embedder(local_timestep)
|
||||
@@ -2107,7 +2197,7 @@ class Bagel(nn.Module):
|
||||
packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
|
||||
packed_sequence[packed_text_indexes] = packed_text_embedding
|
||||
|
||||
assert timestep.unique().shape[0] == 1
|
||||
# i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
|
||||
packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
|
||||
packed_timestep_embeds = self.time_embedder(timestep)
|
||||
x_t_emb = self.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed
|
||||
@@ -2149,7 +2239,7 @@ class Bagel(nn.Module):
|
||||
cfg_renorm_type: str = "global",
|
||||
cfg_text_scale: float = 1.0,
|
||||
cfg_img_scale: float = 1.0,
|
||||
cfg_batched: dict | None = None,
|
||||
cfg_branches: dict | None = None,
|
||||
):
|
||||
# Build query sequence (identical for all CFG branches)
|
||||
packed_text_embedding = self.language_model.forward(
|
||||
@@ -2159,7 +2249,7 @@ class Bagel(nn.Module):
|
||||
packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
|
||||
packed_sequence[packed_text_indexes] = packed_text_embedding
|
||||
|
||||
assert timestep.unique().shape[0] == 1
|
||||
# i2v relaxes this: per-token timestep (cond=0, noncond=t) is valid.
|
||||
packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
|
||||
packed_timestep_embeds = self.time_embedder(timestep)
|
||||
x_t = self.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed
|
||||
@@ -2170,55 +2260,40 @@ class Bagel(nn.Module):
|
||||
extra_inputs = {}
|
||||
if self.use_moe:
|
||||
extra_inputs["mode"] = "gen"
|
||||
extra_inputs["packed_vae_token_indexes"] = packed_vae_token_indexes
|
||||
extra_inputs["packed_text_indexes"] = packed_text_indexes
|
||||
|
||||
use_cfg = cfg_text_scale > 1.0
|
||||
cfg_text_v_t = None
|
||||
cfg_img_v_t = None
|
||||
|
||||
if use_cfg and cfg_batched is not None:
|
||||
# ── Batched CFG: single LLM forward for all branches ──
|
||||
seq_len = cfg_batched["seq_len"]
|
||||
num_branches = cfg_batched["num_branches"]
|
||||
if use_cfg and cfg_branches is not None:
|
||||
# Sequential per-branch CFG forwards (matches upstream lance.py).
|
||||
# The previous batched path concatenated cond + cfg into one LLM
|
||||
# forward, but the block-diagonal attention mask was lost when
|
||||
# PR #3728 dropped flash_attn_varlen, so branches leaked into
|
||||
# each other. Running each branch through its own forward is
|
||||
# numerically identical to upstream.
|
||||
def _run_branch(branch_pkv, branch_pids):
|
||||
out = self.language_model.forward(
|
||||
packed_query_sequence=packed_sequence,
|
||||
query_lens=packed_seqlens,
|
||||
packed_query_position_ids=branch_pids,
|
||||
past_key_values=branch_pkv,
|
||||
update_past_key_values=False,
|
||||
is_causal=False,
|
||||
**extra_inputs,
|
||||
)
|
||||
return self.llm2vae(out.packed_query_sequence)[packed_vae_token_indexes]
|
||||
|
||||
batched_sequence = packed_sequence.repeat(num_branches, 1)
|
||||
|
||||
if self.use_moe:
|
||||
extra_inputs["packed_text_indexes"] = cfg_batched["batched_text_indexes"]
|
||||
extra_inputs["packed_vae_token_indexes"] = cfg_batched["batched_vae_indexes"]
|
||||
|
||||
output = self.language_model.forward(
|
||||
packed_query_sequence=batched_sequence,
|
||||
query_lens=cfg_batched["batched_query_lens"],
|
||||
packed_query_position_ids=cfg_batched["batched_position_ids"],
|
||||
past_key_values=cfg_batched["merged_cache"],
|
||||
update_past_key_values=False,
|
||||
is_causal=False,
|
||||
**extra_inputs,
|
||||
)
|
||||
|
||||
# Extract per-branch velocities from batched output
|
||||
all_hidden = output.packed_query_sequence
|
||||
assert all_hidden.shape[0] == seq_len * num_branches, (
|
||||
f"Expected packed sequence length {seq_len * num_branches}, but got {all_hidden.shape[0]}"
|
||||
)
|
||||
|
||||
v_t = self.llm2vae(all_hidden[:seq_len])[packed_vae_token_indexes]
|
||||
|
||||
branch_idx = 1
|
||||
cfg_text_v_t = self.llm2vae(all_hidden[branch_idx * seq_len : (branch_idx + 1) * seq_len])[
|
||||
packed_vae_token_indexes
|
||||
]
|
||||
branch_idx += 1
|
||||
if cfg_img_scale > 1.0:
|
||||
cfg_img_v_t = self.llm2vae(all_hidden[branch_idx * seq_len : (branch_idx + 1) * seq_len])[
|
||||
packed_vae_token_indexes
|
||||
]
|
||||
branches_pids = cfg_branches["pids"]
|
||||
branches_caches = cfg_branches["caches"]
|
||||
v_t = _run_branch(branches_caches[0], branches_pids[0])
|
||||
cfg_text_v_t = _run_branch(branches_caches[1], branches_pids[1])
|
||||
if cfg_img_scale > 1.0 and len(branches_caches) > 2:
|
||||
cfg_img_v_t = _run_branch(branches_caches[2], branches_pids[2])
|
||||
else:
|
||||
# ── Single forward (no CFG or outside cfg_interval) ──
|
||||
if self.use_moe:
|
||||
extra_inputs["packed_vae_token_indexes"] = packed_vae_token_indexes
|
||||
extra_inputs["packed_text_indexes"] = packed_text_indexes
|
||||
|
||||
# Single forward (no CFG or outside cfg_interval).
|
||||
output = self.language_model.forward(
|
||||
packed_query_sequence=packed_sequence,
|
||||
query_lens=packed_seqlens,
|
||||
@@ -2228,8 +2303,7 @@ class Bagel(nn.Module):
|
||||
is_causal=False,
|
||||
**extra_inputs,
|
||||
)
|
||||
v_t = self.llm2vae(output.packed_query_sequence)
|
||||
v_t = v_t[packed_vae_token_indexes]
|
||||
v_t = self.llm2vae(output.packed_query_sequence)[packed_vae_token_indexes]
|
||||
|
||||
# ── CFG combination ──
|
||||
if use_cfg:
|
||||
|
||||
@@ -319,6 +319,27 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
|
||||
image = (image * 0.5 + 0.5).clamp(0, 1)[0].permute(1, 2, 0) * 255
|
||||
return Image.fromarray(image.to(torch.uint8).cpu().numpy())
|
||||
|
||||
def _regen_init_noise_on_device(self, gen_input: dict, seed: int | None) -> None:
|
||||
"""Resample ``gen_input["packed_init_noises"]`` on-device with a fresh
|
||||
per-call ``torch.Generator``.
|
||||
|
||||
``Bagel.prepare_input`` (and the Lance video equivalent) call
|
||||
``torch.randn`` with no device or generator, falling back to CPU+fp32
|
||||
via the global RNG. Upstream Lance samples directly on CUDA+bf16 via
|
||||
``torch.Generator(device=cuda).manual_seed(seed)`` (lance.py:1536),
|
||||
so for the same seed the two sides land on different noise streams.
|
||||
Mutates ``gen_input`` in place; no-op if seed is unset or device is CPU.
|
||||
"""
|
||||
if seed is None or self.device.type != "cuda":
|
||||
return
|
||||
ref = gen_input["packed_init_noises"]
|
||||
gen_input["packed_init_noises"] = torch.randn(
|
||||
ref.shape,
|
||||
generator=torch.Generator(device=self.device).manual_seed(int(seed)),
|
||||
device=self.device,
|
||||
dtype=self.od_config.dtype,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
if len(req.prompts) > 1:
|
||||
@@ -355,7 +376,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
|
||||
|
||||
gen_params = BagelGenParams(
|
||||
num_timesteps=int(req.sampling_params.num_inference_steps or 50),
|
||||
timestep_shift=3.0,
|
||||
timestep_shift=float(extra_args.get("timestep_shift", 3.0)),
|
||||
cfg_text_scale=cfg_text_scale,
|
||||
cfg_img_scale=cfg_img_scale,
|
||||
cfg_interval=cfg_interval,
|
||||
@@ -756,6 +777,8 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
|
||||
if torch.is_tensor(v):
|
||||
generation_input[k] = v.to(self.device)
|
||||
|
||||
self._regen_init_noise_on_device(generation_input, req.sampling_params.seed)
|
||||
|
||||
# text cfg
|
||||
generation_input_cfg_text = self.bagel.prepare_vae_latent_cfg(
|
||||
curr_kvlens=cfg_text_context["kv_lens"],
|
||||
@@ -820,6 +843,11 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
|
||||
custom = {}
|
||||
if think_text is not None:
|
||||
custom["think_text"] = think_text
|
||||
# Mirror the PIL image into ``custom_output`` so callers reading via
|
||||
# the orchestrator IPC boundary (which strips the bare ``output``
|
||||
# field) can still recover the result. ``video_frames`` already
|
||||
# uses this pattern.
|
||||
custom["image"] = img
|
||||
|
||||
return DiffusionOutput(
|
||||
output=img,
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,121 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lance chat / system prompts.
|
||||
|
||||
Matches the upstream Lance training distribution (see
|
||||
``bytedance/Lance/data/system_prompt_render.py`` and
|
||||
``data/common.py::generate_system_prompt``):
|
||||
|
||||
* each task has a specific system prompt describing what to attend to;
|
||||
* the prompt is wrapped in the Qwen chat template
|
||||
``<|im_start|>system\\n…<|im_end|>\\n<|im_start|>user\\n…<|im_end|>\\n``
|
||||
``<|im_start|>assistant\\n``;
|
||||
* vision tokens are framed by ``<|vision_start|><|video_pad|><|vision_end|>``
|
||||
(upstream uses ``<|video_pad|>`` even for image type by default).
|
||||
|
||||
These are pure string formatting helpers — the runtime pipeline still does
|
||||
the actual VAE / ViT prefill separately.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
# Task-specific system prompts (verbatim from upstream's
|
||||
# ``generate_system_prompt`` — first variant of each pool).
|
||||
SYSTEM_PROMPTS: dict[tuple[str, str], str] = {
|
||||
# (task, vision_type) -> system prompt
|
||||
("t2i", "image"): (
|
||||
"Describe the image by detailing the color, quantity, text, shape, "
|
||||
"size, texture, spatial relationships of the objects and background:"
|
||||
),
|
||||
("t2v", "video"): (
|
||||
"Describe the video by detailing the color, quantity, visible text, "
|
||||
"shape, size, texture, spatial relationships and motion/camera "
|
||||
"movements of the objects and background:"
|
||||
),
|
||||
# i2v reuses the t2v system prompt (matches upstream Lance PR #33's
|
||||
# ``data/common.py`` where ``system_prompt_type in ("t2v", "i2v")``
|
||||
# selects the same prompt pool).
|
||||
("i2v", "video"): (
|
||||
"Describe the video by detailing the color, quantity, visible text, "
|
||||
"shape, size, texture, spatial relationships and motion/camera "
|
||||
"movements of the objects and background:"
|
||||
),
|
||||
("image_edit", "image"): (
|
||||
"Describe the key features of the input image (color, shape, size, "
|
||||
"texture, objects, background), then explain how the user's text "
|
||||
"instruction should alter or modify the image. Generate a new image "
|
||||
"that meets the user's requirements while maintaining consistency "
|
||||
"with the original input where appropriate."
|
||||
),
|
||||
("video_edit", "video"): (
|
||||
"Describe the key features of the input video (color, shape, size, "
|
||||
"texture, objects, background), then explain how the user's text "
|
||||
"instruction should alter or modify the video. Generate a new video "
|
||||
"that meets the user's requirements while maintaining consistency "
|
||||
"with the original input where appropriate."
|
||||
),
|
||||
("x2t_image", "image"): (
|
||||
"Generate a detailed and accurate description of the image, including all the key moments and visual details."
|
||||
),
|
||||
("x2t_video", "video"): (
|
||||
"Generate a detailed and accurate description of the video, including all the key moments and visual details."
|
||||
),
|
||||
}
|
||||
|
||||
# Vision-content placeholders (Qwen2.5-VL convention). Upstream uses
|
||||
# ``<|video_pad|>`` for image inputs too unless ``force_video_pad=False``
|
||||
# (counter-intuitively named; see ``render_qwenvl_prompt`` upstream).
|
||||
VISION_START = "<|vision_start|>"
|
||||
VISION_END = "<|vision_end|>"
|
||||
VIDEO_PAD = "<|video_pad|>"
|
||||
IMAGE_PAD = "<|image_pad|>"
|
||||
|
||||
|
||||
def render_lance_prompt(
|
||||
task: str,
|
||||
user_text: str,
|
||||
*,
|
||||
vision_token: str | None = None,
|
||||
system_prompt: str | None = None,
|
||||
) -> str:
|
||||
"""Render a Lance-compatible user-side prompt.
|
||||
|
||||
``task`` is one of ``t2i``, ``t2v``, ``image_edit``, ``video_edit``,
|
||||
``x2t_image``, ``x2t_video``. ``vision_token`` should be the visual
|
||||
placeholder to embed inside the user message (e.g. ``VIDEO_PAD`` or a
|
||||
``VISION_START + VIDEO_PAD + VISION_END`` block); pass ``None`` for
|
||||
text-only inputs. ``system_prompt`` overrides the default task system
|
||||
prompt — required for x2t QA examples whose upstream JSON carries a
|
||||
per-example instruction (e.g. ``"Look at the image carefully and answer
|
||||
the question."``); without it x2t falls back to the caption-style
|
||||
default and the model describes instead of answering.
|
||||
|
||||
Returns a single string ready to be tokenized; no further wrapping is
|
||||
needed by the caller.
|
||||
"""
|
||||
# ``t2v`` / ``x2t_video`` / ``video_edit`` use video system prompts;
|
||||
# the rest are image.
|
||||
vision_type = "video" if task in {"t2v", "i2v", "x2t_video", "video_edit"} else "image"
|
||||
sys_prompt = (
|
||||
system_prompt if system_prompt else SYSTEM_PROMPTS.get((task, vision_type), "You are a helpful assistant.")
|
||||
)
|
||||
|
||||
if vision_token is None:
|
||||
user_msg = user_text
|
||||
else:
|
||||
# Upstream concatenates the vision block immediately before the text.
|
||||
user_msg = f"{vision_token}{user_text}"
|
||||
|
||||
return (
|
||||
f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"IMAGE_PAD",
|
||||
"SYSTEM_PROMPTS",
|
||||
"VIDEO_PAD",
|
||||
"VISION_END",
|
||||
"VISION_START",
|
||||
"render_lance_prompt",
|
||||
]
|
||||
@@ -0,0 +1,870 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Ported from ByteDance Lance upstream (https://github.com/bytedance/Lance,
|
||||
# modeling/vae/wan/{vae2_2.py,model.py}). Upstream copyright:
|
||||
#
|
||||
# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates.
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
# Licensed under the Apache License, Version 2.0.
|
||||
#
|
||||
# Lance ships its VAE as a raw ``Wan2.2_VAE.pth`` whose state dict is keyed for
|
||||
# the upstream ``WanVAE_`` nn.Module here (not for diffusers' ``AutoencoderKLWan``).
|
||||
# We port that module verbatim so we can load the .pth directly, and wrap it with
|
||||
# :class:`LanceWanVAE` which exposes BAGEL's 4-D ``encode(BCHW)/decode(BCHW)``
|
||||
# surface (and a 5-D path for the video checkpoint).
|
||||
"""Wan2.2 VAE used by Lance, ported from upstream so ``Wan2.2_VAE.pth`` loads
|
||||
natively without state-dict surgery."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
CACHE_T = 2
|
||||
|
||||
|
||||
class CausalConv3d(nn.Conv3d):
|
||||
"""Causal 3D conv with feature-map caching across temporal chunks."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._padding = (
|
||||
self.padding[2],
|
||||
self.padding[2],
|
||||
self.padding[1],
|
||||
self.padding[1],
|
||||
2 * self.padding[0],
|
||||
0,
|
||||
)
|
||||
self.padding = (0, 0, 0)
|
||||
|
||||
def forward(self, x, cache_x=None):
|
||||
padding = list(self._padding)
|
||||
if cache_x is not None and self._padding[4] > 0:
|
||||
cache_x = cache_x.to(x.device)
|
||||
x = torch.cat([cache_x, x], dim=2)
|
||||
padding[4] -= cache_x.shape[2]
|
||||
x = F.pad(x, padding)
|
||||
return super().forward(x)
|
||||
|
||||
|
||||
class RMS_norm(nn.Module):
|
||||
def __init__(self, dim, channel_first=True, images=True, bias=False):
|
||||
super().__init__()
|
||||
broadcastable_dims = (1, 1, 1) if not images else (1, 1)
|
||||
shape = (dim, *broadcastable_dims) if channel_first else (dim,)
|
||||
self.channel_first = channel_first
|
||||
self.scale = dim**0.5
|
||||
self.gamma = nn.Parameter(torch.ones(shape))
|
||||
self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.0
|
||||
|
||||
def forward(self, x):
|
||||
return F.normalize(x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma + self.bias
|
||||
|
||||
|
||||
class Upsample(nn.Upsample):
|
||||
def forward(self, x):
|
||||
return super().forward(x.float()).type_as(x)
|
||||
|
||||
|
||||
class Resample(nn.Module):
|
||||
def __init__(self, dim, mode):
|
||||
assert mode in ("none", "upsample2d", "upsample3d", "downsample2d", "downsample3d")
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.mode = mode
|
||||
|
||||
if mode == "upsample2d":
|
||||
self.resample = nn.Sequential(
|
||||
Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
|
||||
nn.Conv2d(dim, dim, 3, padding=1),
|
||||
)
|
||||
elif mode == "upsample3d":
|
||||
self.resample = nn.Sequential(
|
||||
Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
|
||||
nn.Conv2d(dim, dim, 3, padding=1),
|
||||
)
|
||||
self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
|
||||
elif mode == "downsample2d":
|
||||
self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
|
||||
elif mode == "downsample3d":
|
||||
self.resample = nn.Sequential(nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)))
|
||||
self.time_conv = CausalConv3d(dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
|
||||
else:
|
||||
self.resample = nn.Identity()
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=[0]):
|
||||
b, c, t, h, w = x.size()
|
||||
if self.mode == "upsample3d":
|
||||
if feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
if feat_cache[idx] is None:
|
||||
feat_cache[idx] = "Rep"
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] != "Rep":
|
||||
cache_x = torch.cat(
|
||||
[feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx] == "Rep":
|
||||
cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2)
|
||||
if feat_cache[idx] == "Rep":
|
||||
x = self.time_conv(x)
|
||||
else:
|
||||
x = self.time_conv(x, feat_cache[idx])
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
x = x.reshape(b, 2, c, t, h, w)
|
||||
x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3)
|
||||
x = x.reshape(b, c, t * 2, h, w)
|
||||
t = x.shape[2]
|
||||
x = rearrange(x, "b c t h w -> (b t) c h w")
|
||||
x = self.resample(x)
|
||||
x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
|
||||
|
||||
if self.mode == "downsample3d":
|
||||
if feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
if feat_cache[idx] is None:
|
||||
feat_cache[idx] = x.clone()
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
cache_x = x[:, :, -1:, :, :].clone()
|
||||
x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
return x
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, dropout=0.0):
|
||||
super().__init__()
|
||||
self.in_dim = in_dim
|
||||
self.out_dim = out_dim
|
||||
self.residual = nn.Sequential(
|
||||
RMS_norm(in_dim, images=False),
|
||||
nn.SiLU(),
|
||||
CausalConv3d(in_dim, out_dim, 3, padding=1),
|
||||
RMS_norm(out_dim, images=False),
|
||||
nn.SiLU(),
|
||||
nn.Dropout(dropout),
|
||||
CausalConv3d(out_dim, out_dim, 3, padding=1),
|
||||
)
|
||||
self.shortcut = CausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else nn.Identity()
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=[0]):
|
||||
h = self.shortcut(x)
|
||||
for layer in self.residual:
|
||||
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
||||
cache_x = torch.cat(
|
||||
[feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
x = layer(x, feat_cache[idx])
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = layer(x)
|
||||
return x + h
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
"""Single-head causal self-attention over spatial tokens, per frame."""
|
||||
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.norm = RMS_norm(dim)
|
||||
self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
|
||||
self.proj = nn.Conv2d(dim, dim, 1)
|
||||
nn.init.zeros_(self.proj.weight)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
b, c, t, h, w = x.size()
|
||||
x = rearrange(x, "b c t h w -> (b t) c h w")
|
||||
x = self.norm(x)
|
||||
q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(0, 1, 3, 2).contiguous().chunk(3, dim=-1)
|
||||
x = F.scaled_dot_product_attention(q, k, v)
|
||||
x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
|
||||
x = self.proj(x)
|
||||
x = rearrange(x, "(b t) c h w-> b c t h w", t=t)
|
||||
return x + identity
|
||||
|
||||
|
||||
def _patchify(x, patch_size):
|
||||
if patch_size == 1:
|
||||
return x
|
||||
if x.dim() == 4:
|
||||
return rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
|
||||
if x.dim() == 5:
|
||||
return rearrange(x, "b c f (h q) (w r) -> b (c r q) f h w", q=patch_size, r=patch_size)
|
||||
raise ValueError(f"Invalid input shape: {x.shape}")
|
||||
|
||||
|
||||
def _unpatchify(x, patch_size):
|
||||
if patch_size == 1:
|
||||
return x
|
||||
if x.dim() == 4:
|
||||
return rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
|
||||
if x.dim() == 5:
|
||||
return rearrange(x, "b (c r q) f h w -> b c f (h q) (w r)", q=patch_size, r=patch_size)
|
||||
return x
|
||||
|
||||
|
||||
class AvgDown3D(nn.Module):
|
||||
def __init__(self, in_channels, out_channels, factor_t, factor_s=1):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.factor_t = factor_t
|
||||
self.factor_s = factor_s
|
||||
self.factor = self.factor_t * self.factor_s * self.factor_s
|
||||
assert in_channels * self.factor % out_channels == 0
|
||||
self.group_size = in_channels * self.factor // out_channels
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
|
||||
x = F.pad(x, (0, 0, 0, 0, pad_t, 0))
|
||||
B, C, T, H, W = x.shape
|
||||
x = x.view(
|
||||
B,
|
||||
C,
|
||||
T // self.factor_t,
|
||||
self.factor_t,
|
||||
H // self.factor_s,
|
||||
self.factor_s,
|
||||
W // self.factor_s,
|
||||
self.factor_s,
|
||||
)
|
||||
x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
|
||||
x = x.view(B, C * self.factor, T // self.factor_t, H // self.factor_s, W // self.factor_s)
|
||||
x = x.view(B, self.out_channels, self.group_size, T // self.factor_t, H // self.factor_s, W // self.factor_s)
|
||||
return x.mean(dim=2)
|
||||
|
||||
|
||||
class DupUp3D(nn.Module):
|
||||
def __init__(self, in_channels: int, out_channels: int, factor_t, factor_s=1):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.out_channels = out_channels
|
||||
self.factor_t = factor_t
|
||||
self.factor_s = factor_s
|
||||
self.factor = self.factor_t * self.factor_s * self.factor_s
|
||||
assert out_channels * self.factor % in_channels == 0
|
||||
self.repeats = out_channels * self.factor // in_channels
|
||||
|
||||
def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
|
||||
x = x.repeat_interleave(self.repeats, dim=1)
|
||||
x = x.view(
|
||||
x.size(0),
|
||||
self.out_channels,
|
||||
self.factor_t,
|
||||
self.factor_s,
|
||||
self.factor_s,
|
||||
x.size(2),
|
||||
x.size(3),
|
||||
x.size(4),
|
||||
)
|
||||
x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
|
||||
x = x.view(
|
||||
x.size(0),
|
||||
self.out_channels,
|
||||
x.size(2) * self.factor_t,
|
||||
x.size(4) * self.factor_s,
|
||||
x.size(6) * self.factor_s,
|
||||
)
|
||||
if first_chunk:
|
||||
x = x[:, :, self.factor_t - 1 :, :, :]
|
||||
return x
|
||||
|
||||
|
||||
class Down_ResidualBlock(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False):
|
||||
super().__init__()
|
||||
self.avg_shortcut = AvgDown3D(
|
||||
in_dim, out_dim, factor_t=2 if temperal_downsample else 1, factor_s=2 if down_flag else 1
|
||||
)
|
||||
downsamples = []
|
||||
for _ in range(mult):
|
||||
downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
||||
in_dim = out_dim
|
||||
if down_flag:
|
||||
mode = "downsample3d" if temperal_downsample else "downsample2d"
|
||||
downsamples.append(Resample(out_dim, mode=mode))
|
||||
self.downsamples = nn.Sequential(*downsamples)
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=[0]):
|
||||
x_copy = x.clone()
|
||||
for module in self.downsamples:
|
||||
x = module(x, feat_cache, feat_idx)
|
||||
return x + self.avg_shortcut(x_copy)
|
||||
|
||||
|
||||
class Up_ResidualBlock(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, dropout, mult, temporal_upsample=False, up_flag=False):
|
||||
super().__init__()
|
||||
if up_flag:
|
||||
self.avg_shortcut = DupUp3D(
|
||||
in_dim, out_dim, factor_t=2 if temporal_upsample else 1, factor_s=2 if up_flag else 1
|
||||
)
|
||||
else:
|
||||
self.avg_shortcut = None
|
||||
upsamples = []
|
||||
for _ in range(mult):
|
||||
upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
|
||||
in_dim = out_dim
|
||||
if up_flag:
|
||||
mode = "upsample3d" if temporal_upsample else "upsample2d"
|
||||
upsamples.append(Resample(out_dim, mode=mode))
|
||||
self.upsamples = nn.Sequential(*upsamples)
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
|
||||
x_main = x.clone()
|
||||
for module in self.upsamples:
|
||||
x_main = module(x_main, feat_cache, feat_idx)
|
||||
if self.avg_shortcut is not None:
|
||||
return x_main + self.avg_shortcut(x, first_chunk)
|
||||
return x_main
|
||||
|
||||
|
||||
class Encoder3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim=128,
|
||||
z_dim=4,
|
||||
dim_mult=(1, 2, 4, 4),
|
||||
num_res_blocks=2,
|
||||
attn_scales=(),
|
||||
temperal_downsample=(True, True, False),
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.z_dim = z_dim
|
||||
self.dim_mult = list(dim_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.attn_scales = list(attn_scales)
|
||||
self.temperal_downsample = list(temperal_downsample)
|
||||
|
||||
dims = [dim * u for u in [1] + self.dim_mult]
|
||||
self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
|
||||
|
||||
downsamples = []
|
||||
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
||||
t_down_flag = self.temperal_downsample[i] if i < len(self.temperal_downsample) else False
|
||||
downsamples.append(
|
||||
Down_ResidualBlock(
|
||||
in_dim=in_dim,
|
||||
out_dim=out_dim,
|
||||
dropout=dropout,
|
||||
mult=num_res_blocks,
|
||||
temperal_downsample=t_down_flag,
|
||||
down_flag=i != len(self.dim_mult) - 1,
|
||||
)
|
||||
)
|
||||
self.downsamples = nn.Sequential(*downsamples)
|
||||
|
||||
self.middle = nn.Sequential(
|
||||
ResidualBlock(out_dim, out_dim, dropout),
|
||||
AttentionBlock(out_dim),
|
||||
ResidualBlock(out_dim, out_dim, dropout),
|
||||
)
|
||||
self.head = nn.Sequential(
|
||||
RMS_norm(out_dim, images=False),
|
||||
nn.SiLU(),
|
||||
CausalConv3d(out_dim, z_dim, 3, padding=1),
|
||||
)
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=[0]):
|
||||
if feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
||||
cache_x = torch.cat(
|
||||
[feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
x = self.conv1(x, feat_cache[idx])
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = self.conv1(x)
|
||||
|
||||
for layer in self.downsamples:
|
||||
if feat_cache is not None:
|
||||
x = layer(x, feat_cache, feat_idx)
|
||||
else:
|
||||
x = layer(x)
|
||||
|
||||
for layer in self.middle:
|
||||
if isinstance(layer, ResidualBlock) and feat_cache is not None:
|
||||
x = layer(x, feat_cache, feat_idx)
|
||||
else:
|
||||
x = layer(x)
|
||||
|
||||
for layer in self.head:
|
||||
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
||||
cache_x = torch.cat(
|
||||
[feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
x = layer(x, feat_cache[idx])
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class Decoder3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim=128,
|
||||
z_dim=4,
|
||||
dim_mult=(1, 2, 4, 4),
|
||||
num_res_blocks=2,
|
||||
attn_scales=(),
|
||||
temporal_upsample=(False, True, True),
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.z_dim = z_dim
|
||||
self.dim_mult = list(dim_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.attn_scales = list(attn_scales)
|
||||
self.temporal_upsample = list(temporal_upsample)
|
||||
|
||||
dims = [dim * u for u in [self.dim_mult[-1]] + self.dim_mult[::-1]]
|
||||
self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
|
||||
self.middle = nn.Sequential(
|
||||
ResidualBlock(dims[0], dims[0], dropout),
|
||||
AttentionBlock(dims[0]),
|
||||
ResidualBlock(dims[0], dims[0], dropout),
|
||||
)
|
||||
|
||||
upsamples = []
|
||||
for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
|
||||
t_up_flag = self.temporal_upsample[i] if i < len(self.temporal_upsample) else False
|
||||
upsamples.append(
|
||||
Up_ResidualBlock(
|
||||
in_dim=in_dim,
|
||||
out_dim=out_dim,
|
||||
dropout=dropout,
|
||||
mult=num_res_blocks + 1,
|
||||
temporal_upsample=t_up_flag,
|
||||
up_flag=i != len(self.dim_mult) - 1,
|
||||
)
|
||||
)
|
||||
self.upsamples = nn.Sequential(*upsamples)
|
||||
self.head = nn.Sequential(
|
||||
RMS_norm(out_dim, images=False),
|
||||
nn.SiLU(),
|
||||
CausalConv3d(out_dim, 12, 3, padding=1),
|
||||
)
|
||||
|
||||
def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
|
||||
if feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
||||
cache_x = torch.cat(
|
||||
[feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
x = self.conv1(x, feat_cache[idx])
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = self.conv1(x)
|
||||
|
||||
for layer in self.middle:
|
||||
if isinstance(layer, ResidualBlock) and feat_cache is not None:
|
||||
x = layer(x, feat_cache, feat_idx)
|
||||
else:
|
||||
x = layer(x)
|
||||
|
||||
for layer in self.upsamples:
|
||||
if feat_cache is not None:
|
||||
x = layer(x, feat_cache, feat_idx, first_chunk)
|
||||
else:
|
||||
x = layer(x)
|
||||
|
||||
for layer in self.head:
|
||||
if isinstance(layer, CausalConv3d) and feat_cache is not None:
|
||||
idx = feat_idx[0]
|
||||
cache_x = x[:, :, -CACHE_T:, :, :].clone()
|
||||
if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
|
||||
cache_x = torch.cat(
|
||||
[feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x],
|
||||
dim=2,
|
||||
)
|
||||
x = layer(x, feat_cache[idx])
|
||||
feat_cache[idx] = cache_x
|
||||
feat_idx[0] += 1
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
def _count_conv3d(model):
|
||||
return sum(1 for m in model.modules() if isinstance(m, CausalConv3d))
|
||||
|
||||
|
||||
class WanVAE_(nn.Module):
|
||||
"""Upstream Wan2.2 VAE module — encoder3d/decoder3d sandwich with 2x patchify
|
||||
on input. State-dict-compatible with ``Wan2.2_VAE.pth``."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim=160,
|
||||
dec_dim=256,
|
||||
z_dim=16,
|
||||
dim_mult=(1, 2, 4, 4),
|
||||
num_res_blocks=2,
|
||||
attn_scales=(),
|
||||
temperal_downsample=(True, True, False),
|
||||
dropout=0.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.z_dim = z_dim
|
||||
self.dim_mult = list(dim_mult)
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.attn_scales = list(attn_scales)
|
||||
self.temperal_downsample = list(temperal_downsample)
|
||||
self.temporal_upsample = list(temperal_downsample)[::-1]
|
||||
|
||||
self.encoder = Encoder3d(
|
||||
dim, z_dim * 2, self.dim_mult, num_res_blocks, self.attn_scales, self.temperal_downsample, dropout
|
||||
)
|
||||
self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
|
||||
self.conv2 = CausalConv3d(z_dim, z_dim, 1)
|
||||
self.decoder = Decoder3d(
|
||||
dec_dim, z_dim, self.dim_mult, num_res_blocks, self.attn_scales, self.temporal_upsample, dropout
|
||||
)
|
||||
|
||||
def encode(self, x, scale):
|
||||
self.clear_cache()
|
||||
x = _patchify(x, patch_size=2)
|
||||
t = x.shape[2]
|
||||
iter_ = 1 + (t - 1) // 4
|
||||
out = None
|
||||
for i in range(iter_):
|
||||
self._enc_conv_idx = [0]
|
||||
if i == 0:
|
||||
out = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
|
||||
else:
|
||||
out_ = self.encoder(
|
||||
x[:, :, 1 + 4 * (i - 1) : 1 + 4 * i, :, :],
|
||||
feat_cache=self._enc_feat_map,
|
||||
feat_idx=self._enc_conv_idx,
|
||||
)
|
||||
out = torch.cat([out, out_], 2)
|
||||
mu, log_var = self.conv1(out).chunk(2, dim=1)
|
||||
if isinstance(scale[0], torch.Tensor):
|
||||
mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1)
|
||||
else:
|
||||
mu = (mu - scale[0]) * scale[1]
|
||||
self.clear_cache()
|
||||
return mu, log_var
|
||||
|
||||
def decode(self, z, scale):
|
||||
self.clear_cache()
|
||||
if isinstance(scale[0], torch.Tensor):
|
||||
z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(1, self.z_dim, 1, 1, 1)
|
||||
else:
|
||||
z = z / scale[1] + scale[0]
|
||||
iter_ = z.shape[2]
|
||||
x = self.conv2(z)
|
||||
out = None
|
||||
for i in range(iter_):
|
||||
self._conv_idx = [0]
|
||||
if i == 0:
|
||||
out = self.decoder(
|
||||
x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True
|
||||
)
|
||||
else:
|
||||
out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
|
||||
out = torch.cat([out, out_], 2)
|
||||
out = _unpatchify(out, patch_size=2)
|
||||
self.clear_cache()
|
||||
return out
|
||||
|
||||
def clear_cache(self):
|
||||
self._conv_num = _count_conv3d(self.decoder)
|
||||
self._conv_idx = [0]
|
||||
self._feat_map = [None] * self._conv_num
|
||||
self._enc_conv_num = _count_conv3d(self.encoder)
|
||||
self._enc_conv_idx = [0]
|
||||
self._enc_feat_map = [None] * self._enc_conv_num
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------------
|
||||
# Wan2.2 VAE wrapper + BAGEL-surface adapter
|
||||
# ----------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Per-channel mean/std for the 48-channel Wan2.2 latent space (upstream constants).
|
||||
_WAN22_LATENT_MEAN: tuple[float, ...] = (
|
||||
-0.2289,
|
||||
-0.0052,
|
||||
-0.1323,
|
||||
-0.2339,
|
||||
-0.2799,
|
||||
0.0174,
|
||||
0.1838,
|
||||
0.1557,
|
||||
-0.1382,
|
||||
0.0542,
|
||||
0.2813,
|
||||
0.0891,
|
||||
0.1570,
|
||||
-0.0098,
|
||||
0.0375,
|
||||
-0.1825,
|
||||
-0.2246,
|
||||
-0.1207,
|
||||
-0.0698,
|
||||
0.5109,
|
||||
0.2665,
|
||||
-0.2108,
|
||||
-0.2158,
|
||||
0.2502,
|
||||
-0.2055,
|
||||
-0.0322,
|
||||
0.1109,
|
||||
0.1567,
|
||||
-0.0729,
|
||||
0.0899,
|
||||
-0.2799,
|
||||
-0.1230,
|
||||
-0.0313,
|
||||
-0.1649,
|
||||
0.0117,
|
||||
0.0723,
|
||||
-0.2839,
|
||||
-0.2083,
|
||||
-0.0520,
|
||||
0.3748,
|
||||
0.0152,
|
||||
0.1957,
|
||||
0.1433,
|
||||
-0.2944,
|
||||
0.3573,
|
||||
-0.0548,
|
||||
-0.1681,
|
||||
-0.0667,
|
||||
)
|
||||
_WAN22_LATENT_STD: tuple[float, ...] = (
|
||||
0.4765,
|
||||
1.0364,
|
||||
0.4514,
|
||||
1.1677,
|
||||
0.5313,
|
||||
0.4990,
|
||||
0.4818,
|
||||
0.5013,
|
||||
0.8158,
|
||||
1.0344,
|
||||
0.5894,
|
||||
1.0901,
|
||||
0.6885,
|
||||
0.6165,
|
||||
0.8454,
|
||||
0.4978,
|
||||
0.5759,
|
||||
0.3523,
|
||||
0.7135,
|
||||
0.6804,
|
||||
0.5833,
|
||||
1.4146,
|
||||
0.8986,
|
||||
0.5659,
|
||||
0.7069,
|
||||
0.5338,
|
||||
0.4889,
|
||||
0.4917,
|
||||
0.4069,
|
||||
0.4999,
|
||||
0.6866,
|
||||
0.4093,
|
||||
0.5709,
|
||||
0.6065,
|
||||
0.6415,
|
||||
0.4944,
|
||||
0.5726,
|
||||
1.2042,
|
||||
0.5458,
|
||||
1.6887,
|
||||
0.3971,
|
||||
1.0600,
|
||||
0.3943,
|
||||
0.5537,
|
||||
0.5444,
|
||||
0.4089,
|
||||
0.7468,
|
||||
0.7744,
|
||||
)
|
||||
|
||||
|
||||
class LanceWanVAE(nn.Module):
|
||||
"""Wan2.2 VAE wrapped for BAGEL's pipeline.
|
||||
|
||||
Exposes BAGEL's image-VAE surface — ``encode(BCHW) -> BC_zHW`` and
|
||||
``decode(BC_zHW) -> BCHW`` — by treating each image as a 1-frame video clip.
|
||||
A 5-D ``encode_video``/``decode_video`` path is also provided for the
|
||||
``Lance_3B_Video`` checkpoint.
|
||||
|
||||
Construction is lazy: the heavy ``WanVAE_`` and ``Wan2.2_VAE.pth`` are not
|
||||
materialized until first use. Once built, the inner module is registered as
|
||||
a submodule so ``self.parameters()``, ``self.to(device)`` and ``vae_dtype =
|
||||
next(vae.parameters()).dtype`` (used by BAGEL's decode path) all behave.
|
||||
"""
|
||||
|
||||
z_channels: int = 48
|
||||
downsample_spatial: int = 16
|
||||
downsample_temporal: int = 4
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae_path: str,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
device: torch.device | str | None = None,
|
||||
*,
|
||||
lazy: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self._vae_path = vae_path
|
||||
self._dtype = dtype
|
||||
self._device = torch.device(device) if device is not None else None
|
||||
self._built = False
|
||||
# Registered lazily once the .pth is loaded:
|
||||
# self.model: WanVAE_ (set in _ensure_built)
|
||||
if not lazy:
|
||||
# BAGEL's decode path inspects ``next(vae.parameters()).dtype`` before
|
||||
# the first encode/decode call, so we have to register the submodule
|
||||
# eagerly when a real device is provided. ``lazy=True`` is reserved
|
||||
# for config-only / no-GPU paths.
|
||||
self._ensure_built()
|
||||
|
||||
def _ensure_built(self) -> None:
|
||||
if self._built:
|
||||
return
|
||||
device = self._device or (
|
||||
next(self.parameters(), torch.empty(0)).device
|
||||
if any(True for _ in self.parameters())
|
||||
else torch.device("cpu")
|
||||
)
|
||||
# Wan2.2 VAE config from upstream Wan2_2_VAE defaults:
|
||||
# z_dim=48 (channels), c_dim=160 (encoder width), dec_dim=256 (decoder width),
|
||||
# dim_mult=[1,2,4,4], temperal_downsample=[False,True,True].
|
||||
model = WanVAE_(
|
||||
dim=160,
|
||||
dec_dim=256,
|
||||
z_dim=self.z_channels,
|
||||
dim_mult=(1, 2, 4, 4),
|
||||
num_res_blocks=2,
|
||||
attn_scales=(),
|
||||
temperal_downsample=(False, True, True),
|
||||
dropout=0.0,
|
||||
)
|
||||
logger.info("Loading Wan2.2 VAE state dict from %s", self._vae_path)
|
||||
state = torch.load(self._vae_path, map_location="cpu", weights_only=True)
|
||||
missing, unexpected = model.load_state_dict(state, strict=False)
|
||||
if missing:
|
||||
logger.warning("Wan2.2 VAE missing keys (%d): %s", len(missing), missing[:5])
|
||||
if unexpected:
|
||||
logger.warning("Wan2.2 VAE unexpected keys (%d): %s", len(unexpected), unexpected[:5])
|
||||
model = model.to(device=device, dtype=self._dtype).eval()
|
||||
model.requires_grad_(False)
|
||||
# Register as submodule so .to() / .parameters() pick it up.
|
||||
self.model = model
|
||||
# Per-channel normalization tensors (matches upstream Wan2_2_VAE.scale = [mean, 1/std]).
|
||||
mean = torch.tensor(_WAN22_LATENT_MEAN, dtype=self._dtype, device=device)
|
||||
inv_std = torch.tensor([1.0 / s for s in _WAN22_LATENT_STD], dtype=self._dtype, device=device)
|
||||
# Buffers, so .to(device) follows along.
|
||||
self.register_buffer("_latent_mean", mean, persistent=False)
|
||||
self.register_buffer("_latent_inv_std", inv_std, persistent=False)
|
||||
self._built = True
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode_video(self, video: torch.Tensor, *, use_sample: bool = True) -> torch.Tensor:
|
||||
"""Encode a 5-D clip ``[B, 3, T, H, W]`` -> latent ``[B, 48, t, h, w]``."""
|
||||
self._ensure_built()
|
||||
video = video.to(self.model.encoder.conv1.weight.dtype)
|
||||
mu, log_var = self.model.encode(video, [self._latent_mean, self._latent_inv_std])
|
||||
if use_sample:
|
||||
std = torch.exp(0.5 * log_var)
|
||||
mu = mu + std * torch.randn_like(std)
|
||||
return mu
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode_video(self, latent: torch.Tensor) -> torch.Tensor:
|
||||
"""Decode a 5-D latent ``[B, 48, t, h, w]`` -> video ``[B, 3, T, H, W]``."""
|
||||
self._ensure_built()
|
||||
latent = latent.to(self.model.decoder.conv1.weight.dtype)
|
||||
out = self.model.decode(latent, [self._latent_mean, self._latent_inv_std])
|
||||
return out.clamp_(-1.0, 1.0)
|
||||
|
||||
# ----- BAGEL image-VAE surface (4-D) -------------------------------- #
|
||||
|
||||
@torch.inference_mode()
|
||||
def encode(self, padded_images: torch.Tensor) -> torch.Tensor:
|
||||
"""``[B, 3, H, W]`` -> ``[B, 48, H/16, W/16]`` (single-frame image path).
|
||||
|
||||
Each image is wrapped as a 1-frame clip, encoded, and the temporal axis
|
||||
is squeezed back out so the result matches BAGEL's iteration pattern.
|
||||
"""
|
||||
if padded_images.dim() == 4:
|
||||
video = padded_images.unsqueeze(2) # B,C,1,H,W
|
||||
squeeze_time = True # image path: caller iterates 3-D latents
|
||||
elif padded_images.dim() == 5:
|
||||
video = padded_images
|
||||
# video path: caller expects 4-D latents (C,T,H,W), even when T==1
|
||||
# (e.g. video2video with a 1-frame reference clip). Squeezing here
|
||||
# would drop the temporal axis and break the downstream indexing
|
||||
# ``latent[:, :t_lat, ...]`` in ``forward_cache_update_vae``.
|
||||
squeeze_time = False
|
||||
else:
|
||||
raise ValueError(f"LanceWanVAE.encode expects 4-D BCHW or 5-D BCTHW, got {tuple(padded_images.shape)}")
|
||||
latent = self.encode_video(video, use_sample=True)
|
||||
# B,48,t,h,w -> B,48,h,w when t == 1 (image path only)
|
||||
if squeeze_time and latent.shape[2] == 1:
|
||||
return latent.squeeze(2)
|
||||
return latent
|
||||
|
||||
@torch.inference_mode()
|
||||
def decode(self, latent: torch.Tensor) -> torch.Tensor:
|
||||
"""``[B, 48, h, w]`` -> ``[B, 3, H, W]`` (single-frame image path)."""
|
||||
if latent.dim() == 4:
|
||||
latent = latent.unsqueeze(2) # B,48,1,h,w
|
||||
elif latent.dim() != 5:
|
||||
raise ValueError(f"LanceWanVAE.decode expects 4-D BCHW or 5-D BCTHW, got {tuple(latent.shape)}")
|
||||
video = self.decode_video(latent) # B,3,T,H,W
|
||||
if video.shape[2] == 1:
|
||||
return video.squeeze(2)
|
||||
return video
|
||||
|
||||
|
||||
def build_wan22_vae(vae_path: str, dtype: torch.dtype = torch.bfloat16, device=None) -> LanceWanVAE:
|
||||
"""Convenience factory: lazy-construct a :class:`LanceWanVAE` adapter."""
|
||||
return LanceWanVAE(vae_path=vae_path, dtype=dtype, device=device)
|
||||
|
||||
|
||||
__all__ = ["LanceWanVAE", "WanVAE_", "build_wan22_vae"]
|
||||
@@ -141,6 +141,11 @@ _DIFFUSION_MODELS = {
|
||||
"pipeline_bagel",
|
||||
"BagelPipeline",
|
||||
),
|
||||
"LancePipeline": (
|
||||
"lance",
|
||||
"pipeline_lance",
|
||||
"LancePipeline",
|
||||
),
|
||||
"MingImagePipeline": (
|
||||
"ming_flash_omni",
|
||||
"pipeline_ming_imagegen",
|
||||
@@ -471,6 +476,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
|
||||
"WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func",
|
||||
"LongCatImagePipeline": "get_longcat_image_post_process_func",
|
||||
"BagelPipeline": "get_bagel_post_process_func",
|
||||
"LancePipeline": "get_lance_post_process_func",
|
||||
"MingImagePipeline": "get_ming_image_post_process_func",
|
||||
"InternVLAA1Pipeline": "get_internvla_a1_post_process_func",
|
||||
"LongCatImageEditPipeline": "get_longcat_image_post_process_func",
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lance (ByteDance) — vLLM-Omni model_executor package.
|
||||
|
||||
Single-stage topology: a self-contained diffusion stage (``LancePipeline``
|
||||
defined in ``vllm_omni.diffusion.models.lance``) that owns the Qwen2-MoT LLM,
|
||||
the Qwen2.5-VL ViT, the Wan2.2 VAE and the tokenizer — directly analogous to
|
||||
``bagel_single_stage``. The ``PipelineConfig`` lives in :mod:`.pipeline`.
|
||||
"""
|
||||
@@ -0,0 +1,30 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Lance pipeline topology.
|
||||
|
||||
Single-stage DiT — self-contained diffusion stage that handles all modalities
|
||||
(text2img, image_edit, x2t_image, video) internally via its own Qwen2-MoT
|
||||
LLM, Qwen2.5-VL ViT, Wan2.2 VAE, and tokenizer. Mirrors ``bagel_single_stage``.
|
||||
"""
|
||||
|
||||
from vllm_omni.config.stage_config import (
|
||||
PipelineConfig,
|
||||
StageExecutionType,
|
||||
StagePipelineConfig,
|
||||
)
|
||||
|
||||
LANCE_PIPELINE = PipelineConfig(
|
||||
model_type="lance",
|
||||
model_arch="LancePipeline",
|
||||
hf_architectures=(),
|
||||
stages=(
|
||||
StagePipelineConfig(
|
||||
stage_id=0,
|
||||
model_stage="dit",
|
||||
execution_type=StageExecutionType.DIFFUSION,
|
||||
input_sources=(),
|
||||
final_output=True,
|
||||
final_output_type="image",
|
||||
),
|
||||
),
|
||||
)
|
||||
Reference in new issue
Block a user