# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """LancePipeline — Lance (ByteDance) packaged for the vLLM-Omni diffusion engine. Lance is BAGEL-lineage (Qwen2-MoT unified AR+diffusion), so the transformer core and the entire generation/forward machinery are inherited unchanged from :class:`BagelPipeline`. Only model *construction* differs, and only in three well-localized places: 1. Checkpoint layout — the HF repo ``bytedance-research/Lance`` bundles ``Lance_3B/`` (image) and ``Lance_3B_Video/`` (video) LLM checkpoints, ``Qwen2.5-VL-ViT/`` (understanding ViT) and ``Wan2.2_VAE.pth`` (VAE) in a single repo. There is no BAGEL-style top-level ``config.json`` carrying ``vae_config`` / ``vit_config`` / ``latent_patch_size`` — those are Lance constants taken from upstream ``config/config_factory.py`` and ``inference_lance.sh`` and hardcoded in :data:`LANCE_DEFAULTS`. 2. Understanding ViT — Qwen2.5-VL vision tower (bundled ``Qwen2.5-VL-ViT/vit.safetensors``) instead of SigLIP. 3. VAE — Wan2.2 (``Wan2.2_VAE.pth``) instead of the BAGEL autoencoder. Scope: this lands the **image path** (``t2i`` / ``image_edit`` / ``x2t_image``) which is the direct BAGEL analogue. The ``Lance_3B_Video`` path needs the 3-D latent position embedding (:class:`LancePositionEmbedding3D`) and temporal VAE handling and is an explicit follow-up — see the PR description. Bring-up status (verified on a B300 against ``bytedance-research/Lance``): * t2i: end-to-end working (1024x1024 image in ~6 s, 0 missing keys). * x2t (image understanding): plumbing wired (Qwen2.5-VL ViT + Qwen2-VL image processor + no-op connector/vit_pos_embed; VAE prefill skipped for Lance via :class:`LanceBagel`). The pipeline runs to completion without crashes but currently emits an immediate EOS because the Qwen2.5-VL backbone needs mRoPE position ids on the image+text sequence and we presently force scalar positions (``rope_scaling = None``); enabling mRoPE end-to-end is the next follow-up. * image_edit (img2img): blocked on the same VAE-prefill issue (Wan2.2 latents do not map onto BAGEL's ``latent_pos_embed`` grid); needs a Lance-specific ``prepare_vae_images``. * video (``Lance_3B_Video``): needs the 3-D latent position embedding wired into ``Bagel`` plus a multi-frame VAE decode path in the pipeline. :class:`LanceWanVAE.decode_video` and :class:`LancePositionEmbedding3D` are already in place. * Two-stage (AR thinker + DiT): needs ``LanceConfig`` / ``LanceProcessor`` registered in the ``vllm`` package; tracked in a separate PR. """ from __future__ import annotations import json import os from dataclasses import dataclass import torch from PIL import Image from vllm.logger import init_logger from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.bagel.pipeline_bagel import ( BagelPipeline, add_special_tokens, ) from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) from .lance_transformer import ( LanceBagel, LanceIdentityConnector, LancePositionEmbedding3D, LanceQwen2_5_VLNaViTWrapper, LanceZeroVitPosEmbed, Qwen2MoTConfig, Qwen2MoTForCausalLM, ) from .wan_vae import LanceWanVAE logger = init_logger(__name__) def _extract_user_instruction(rendered: str) -> str: """Pull just the user instruction out of a chat-template-rendered prompt. ``render_lance_prompt`` wraps the instruction in: <|im_start|>system\\n{sys}<|im_end|>\\n <|im_start|>user\\n{vision_block}{user_text}<|im_end|>\\n <|im_start|>assistant\\n The ``image_edit`` pipeline rebuilds the chat template segment-by-segment (so the rope positions and segment boundaries match upstream Lance), and therefore needs only ``user_text`` — not the wrapped string. This helper extracts it; if the input doesn't look templated it's returned as-is. """ if "<|im_start|>user" not in rendered: return rendered.strip() # User block sits between the user header and the next ``<|im_end|>``. after_user = rendered.split("<|im_start|>user", 1)[1] after_user = after_user.lstrip("\n") user_block = after_user.split("<|im_end|>", 1)[0] # Strip any vision_start..vision_end blocks. out = [] i = 0 while i < len(user_block): if user_block.startswith("<|vision_start|>", i): j = user_block.find("<|vision_end|>", i) if j == -1: break i = j + len("<|vision_end|>") continue out.append(user_block[i]) i += 1 return "".join(out).strip() @dataclass(frozen=True) class LanceDefaults: """Lance constants that upstream keeps in ``config/config_factory.py`` / ``inference_lance.sh`` rather than in any shipped JSON. Verified against ``bytedance-research/Lance``'s released ``Lance_3B/model.safetensors``: ``vae2llm.weight = (2048, 48)`` ⇒ ``patch_latent_dim = latent_patch_size**2 * z_channels = 1 * 48``, i.e. Lance does *not* unfold the Wan latent into a 2×2 patch the way BAGEL does (Wan2.2 already patchifies internally), and ``latent_pos_embed.pos_embed = (4096, 2048)`` ⇒ ``max_latent_size = 64``. """ # latent / patch geometry (pt, ph, pw); image path only uses the spatial 1. latent_patch_size_spatial: int = 1 latent_patch_size_temporal: int = 1 max_latent_size: int = 64 # Lance_3B_Video ships ``latent_pos_embed.pos_embed`` of shape # ``(31 * 64 * 64, 2048) = (126976, 2048)`` ⇒ max_num_frames in latent # space is 31, equivalent to ``(num_frames - 1) // downsample_temporal + 1`` # for up to 121 RGB frames. max_num_video_latent_frames: int = 31 vit_max_num_patch_per_side: int = 70 connector_act: str = "gelu_pytorch_tanh" timestep_shift: float = 3.5 num_timesteps: int = 30 cfg_text_scale: float = 4.0 # Wan2.2 VAE vae_z_channels: int = 48 vae_downsample_spatial: int = 16 vae_downsample_temporal: int = 4 LANCE_DEFAULTS = LanceDefaults() # Subdirectories inside the HF repo (see config.json::checkpoint_directories). _IMAGE_CKPT_DIR = "Lance_3B" _VIDEO_CKPT_DIR = "Lance_3B_Video" _VIT_DIR = "Qwen2.5-VL-ViT" _VAE_FILE = "Wan2.2_VAE.pth" def get_lance_post_process_func(od_config: OmniDiffusionConfig): """Lance returns PIL.Image.Image directly, same as BAGEL.""" def post_process_func(x): return x return post_process_func def get_lance_pre_process_func(od_config: OmniDiffusionConfig): def pre_process_func(x): return x return pre_process_func @dataclass class _LanceVaeCfg: z_channels: int = LANCE_DEFAULTS.vae_z_channels # BAGEL core computes ``latent_downsample = vae.downsample * latent_patch_size``. # For Lance image latents that is ``16 * 2 = 32``. downsample: int = LANCE_DEFAULTS.vae_downsample_spatial @dataclass class _LanceVitCfg: patch_size: int = 14 hidden_size: int = 1280 # Qwen2.5-VL vision hidden (out_hidden_size=2048) class LancePipeline(BagelPipeline): """Lance pipeline. Inherits BAGEL's forward/generation; overrides only construction (checkpoint layout, Qwen2.5-VL ViT, Wan2.2 VAE).""" def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): # Intentionally do NOT call BagelPipeline.__init__ — its assumptions # about config.json / vit_config.json / SigLIP / the BAGEL AE do not # hold for Lance. We replicate the BAGEL construction sequence with # Lance-specific component builders, then reuse every inherited method. import torch.nn as nn from vllm.transformers_utils.configs.bagel import BagelConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import ( DiffusersPipelineLoader, ) nn.Module.__init__(self) self.od_config = od_config self.device = get_local_device() self.scheduler = None self.scheduler_kwargs = {} model = od_config.model if os.path.exists(model): repo_root = model else: repo_root = download_weights_from_hf_specific(model, od_config.revision, ["*"]) # If --model points directly at ``Lance_3B`` / ``Lance_3B_Video`` (or # ``Qwen2.5-VL-ViT``), walk up so ``repo_root`` is the bundled top-level # dir that owns the sibling components (VAE, ViT). base = os.path.basename(repo_root.rstrip("/")) if base in {_IMAGE_CKPT_DIR, _VIDEO_CKPT_DIR, _VIT_DIR}: parent = os.path.dirname(repo_root.rstrip("/")) if os.path.isdir(parent) and os.path.isfile(os.path.join(parent, _VAE_FILE)): repo_root = parent is_video = self._select_video_variant(od_config) or base == _VIDEO_CKPT_DIR ckpt_dir = _VIDEO_CKPT_DIR if is_video else _IMAGE_CKPT_DIR ckpt_path = os.path.join(repo_root, ckpt_dir) if not os.path.isdir(ckpt_path): # Some users point --model directly at a single checkpoint dir. ckpt_path = repo_root self._repo_root = repo_root self._ckpt_path = ckpt_path self._is_video = is_video if is_video: logger.info( "Lance video checkpoint selected (%s) — wired with " "LancePositionEmbedding3D and Wan2.2 multi-frame VAE decode " "(text-to-video). Image-edit / x2t_video remain follow-ups.", ckpt_dir, ) # ---- LLM (Qwen2-MoT; identical weight layout to BAGEL) ---- llm_cfg_path = os.path.join(ckpt_path, "llm_config.json") llm_config = Qwen2MoTConfig.from_json_file(llm_cfg_path) llm_config.qk_norm = True # The released Lance checkpoint ships a separate ``language_model.lm_head.weight`` # tensor even though ``llm_config.json`` says ``tie_word_embeddings=True``; keep # the head untied so the checkpoint loads with zero missing/unexpected keys. llm_config.tie_word_embeddings = False llm_config.layer_module = od_config.override_transformer_cls_name or "Qwen2MoTDecoderLayer" # Lance is Qwen2.5-VL-MoT and ships ``rope_scaling = {"type": "mrope", # "mrope_section": [16, 24, 24]}``. Keep the mRoPE configuration: # :class:`BagelRotaryEmbedding` now auto-dispatches on ``rope_type`` and # consumes 3-D ``(t, h, w)`` position ids, while :class:`LanceBagel` # broadcasts scalar positions to ``(3, S)`` for text-only blocks and # emits true per-token ``(t, h, w)`` for video latent blocks. Required # for the Qwen2.5-VL backbone to produce coherent x2t / t2v output. if not isinstance(getattr(llm_config, "rope_scaling", None), dict): llm_config.rope_scaling = {"rope_type": "mrope", "mrope_section": [16, 24, 24]} else: llm_config.rope_scaling.setdefault("rope_type", llm_config.rope_scaling.get("type", "mrope")) llm_config.rope_scaling.setdefault("mrope_section", [16, 24, 24]) self.tokenizer = self._load_tokenizer(ckpt_path) self.tokenizer, self.new_token_ids, _ = add_special_tokens(self.tokenizer) # Image / video preprocessor for the understanding paths # (x2t_image / x2t_video / image_edit / video_edit). Lance's HF repo # does not bundle a ``preprocessor_config.json``; we use transformers' # ``Qwen2VLImageProcessor`` / ``Qwen2VLVideoProcessor`` with their # default CLIP-style normalization, ``patch_size=14``, ``merge_size=2``, # which matches the bundled ``Qwen2.5-VL-ViT/config.json`` exactly. self.image_processor = self._build_image_processor() self.video_processor = self._build_video_processor() tok_len = len(self.tokenizer) required_max_id = max(int(v) for v in self.new_token_ids.values()) llm_config.vocab_size = max( int(getattr(llm_config, "vocab_size", tok_len)), int(tok_len), int(required_max_id + 1), ) parallel_config = od_config.parallel_config if od_config else None quant_config = od_config.quantization_config self.language_model = Qwen2MoTForCausalLM( llm_config, parallel_config=parallel_config, quant_config=quant_config, prefix="bagel.language_model", ) self.transformer = self.language_model.model # ---- Understanding ViT: Qwen2.5-VL vision (bundled) ---- self.vit_model = self._build_qwen2_5_vl_vit(repo_root) # ---- VAE: Wan2.2 (bundled .pth) ---- self.vae = self._build_wan22_vae(repo_root) vae_cfg = _LanceVaeCfg() vit_cfg = _LanceVitCfg( patch_size=14, hidden_size=getattr(getattr(self.vit_model, "config", None), "hidden_size", 1280), ) # Lance uses Qwen2.5-VL's vision tower whose ``merger`` already projects # to the LLM hidden size and which carries its own positional encoding; # there is no separate BAGEL-style ``connector`` / ``vit_pos_embed`` in # the released checkpoint (Lance_3B safetensors only contain # ``vit_model.*``, never ``connector.*`` / ``vit_pos_embed.*``). We keep # BAGEL's ``visual_und=True`` so ``vit_model`` is registered as a child # of ``self.bagel`` and BAGEL's ``forward_cache_update_vit`` works # unchanged, then immediately replace ``connector`` with an identity and # ``vit_pos_embed`` with a zero op so the strict load check does not # demand those phantom weights and the addition in # ``forward_cache_update_vit`` is numerically a no-op. und_enabled = self.vit_model is not None self.bagel = LanceBagel( language_model=self.language_model, vit_model=self.vit_model, parallel_config=parallel_config, quant_config=quant_config, prefix="bagel", config=BagelConfig( llm_config=llm_config, vae_config=vae_cfg, vit_config=vit_cfg, vit_max_num_patch_per_side=LANCE_DEFAULTS.vit_max_num_patch_per_side, connector_act=LANCE_DEFAULTS.connector_act, interpolate_pos=False, latent_patch_size=LANCE_DEFAULTS.latent_patch_size_spatial, max_latent_size=LANCE_DEFAULTS.max_latent_size, timestep_shift=LANCE_DEFAULTS.timestep_shift, visual_gen=True, visual_und=und_enabled, ), ) if und_enabled: self.bagel.connector = LanceIdentityConnector() self.bagel.vit_pos_embed = LanceZeroVitPosEmbed() # Hand the Qwen2-VL processors to LanceBagel.prepare_vit_images / # prepare_vit_videos so they can re-call them and recover # ``image_grid_thw`` / ``video_grid_thw`` (BAGEL's lambda drops the # grid). self.bagel._lance_image_processor = self.image_processor self.bagel._lance_video_processor = self.video_processor # Lance_3B_Video ships a 3-D latent positional table # (``latent_pos_embed.pos_embed`` shape ``(max_num_frames * max_latent_size**2, hidden_size)``) # in place of BAGEL's 2-D ``PositionEmbedding(max_latent_size, hidden)``; # swap in :class:`LancePositionEmbedding3D` so the video checkpoint loads # cleanly. ``max_num_frames`` is derived from the table size. if is_video: self.bagel.latent_pos_embed = LancePositionEmbedding3D( max_num_frames=LANCE_DEFAULTS.max_num_video_latent_frames, max_num_patch_per_side=LANCE_DEFAULTS.max_latent_size, hidden_size=llm_config.hidden_size, ) # Weight sources. Verified against bytedance-research/Lance: # * Lance_3B/model.safetensors carries the LLM + connectors with # keys ``language_model.* / vae2llm.* / llm2vae.* / time_embedder.* / # latent_pos_embed.*`` — load under the ``bagel.`` namespace. # * Lance_3B_Video/model.safetensors additionally includes # ``vit_model.*`` (390 tensors). # * Qwen2.5-VL-ViT/vit.safetensors (image checkpoint only) carries # the bare ``blocks.* / merger.* / patch_embed.*`` and must land # under ``bagel.vit_model.vision_model.*`` to match the # :class:`LanceQwen2_5_VLNaViTWrapper` hierarchy. # Always use the resolved repo_root for weight sources so subfolder paths # resolve correctly even when the user passed a per-checkpoint subdir # (we already walked ``repo_root`` up to the bundled root above). weights_model = repo_root if os.path.isdir(repo_root) else od_config.model self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=weights_model, subfolder=ckpt_dir if ckpt_path != repo_root else None, revision=od_config.revision, prefix="bagel.", fall_back_to_pt=False, ), ] # Always overlay ``Qwen2.5-VL-ViT/vit.safetensors`` when the bundle # contains it. Upstream Lance's ``inference_lance.py`` loads ViT # weights from ``vit.safetensors`` unconditionally (line 455-458), # using the released understanding-ViT checkpoint regardless of t2i # vs t2v / image_edit vs video_edit. Lance_3B_Video bundles its own # ``vit_model.*`` inside ``model.safetensors``, but loading those # produced byte-different ViT outputs from upstream (rel_l2 ≈ 9 # at the merger output) — the bundled video-checkpoint ViT diverges # from the standalone vit.safetensors. Mirroring upstream's load # order (vit.safetensors LAST so it wins) restores byte-identical # ViT outputs for both image_edit and video_edit. if und_enabled and os.path.isdir(os.path.join(repo_root, _VIT_DIR)): self.weights_sources.append( DiffusersPipelineLoader.ComponentSource( model_or_path=weights_model, subfolder=_VIT_DIR, revision=od_config.revision, prefix="bagel.vit_model.vision_model.", fall_back_to_pt=False, ) ) if quant_config is None and not (od_config.enable_layerwise_offload or od_config.parallel_config.use_hsdp): self.to(self.device) self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) # ------------------------------------------------------------------ # # Lance-specific component builders # ------------------------------------------------------------------ # @staticmethod def _select_video_variant(od_config: OmniDiffusionConfig) -> bool: """Pick Lance_3B vs Lance_3B_Video. Defaults to image; a user can force video via the model path ending in the video dir or an od_config extra flag.""" model = str(od_config.model) if model.rstrip("/").endswith(_VIDEO_CKPT_DIR): return True extra = getattr(od_config, "extra", None) or {} return bool(extra.get("lance_video", False)) @staticmethod def _load_tokenizer(ckpt_path: str): from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(ckpt_path, local_files_only=True, trust_remote_code=True) @staticmethod def _build_image_processor(): """Construct a Qwen2.5-VL-compatible image preprocessor. Defaults mirror ``Qwen/Qwen2.5-VL-3B-Instruct`` (CLIP normalization, 14-pixel patches, 2x spatial merge). """ from transformers import Qwen2VLImageProcessor return Qwen2VLImageProcessor() @staticmethod def _build_video_processor(): """Construct a Qwen2.5-VL-compatible video preprocessor. Used by the ``x2t_video`` understanding path. Returns ``pixel_values_videos`` + 3-D ``video_grid_thw = [T_lat, H, W]``. """ from transformers import Qwen2VLVideoProcessor return Qwen2VLVideoProcessor() def _build_qwen2_5_vl_vit(self, repo_root: str): """Build the bundled Qwen2.5-VL vision tower and wrap it NaViT-style. Lance's bundled ``Qwen2.5-VL-ViT/config.json`` sets ``_attn_implementation = "flash_attention_2"``; we force ``sdpa`` instead so the tower constructs on hardware (e.g. Blackwell) where ``flash-attn`` is not present, falling back to PyTorch's native scaled-dot-product attention. Numerically equivalent for inference. """ vit_dir = os.path.join(repo_root, _VIT_DIR) cfg_path = os.path.join(vit_dir, "config.json") with open(cfg_path, encoding="utf-8") as f: vit_cfg_dict = json.load(f) vit_cfg_dict = dict(vit_cfg_dict) vit_cfg_dict["_attn_implementation"] = "sdpa" try: from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( Qwen2_5_VLVisionConfig, ) from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( Qwen2_5_VisionTransformerPretrainedModel, ) vit_conf = Qwen2_5_VLVisionConfig(**vit_cfg_dict) vision = Qwen2_5_VisionTransformerPretrainedModel(vit_conf) except Exception as e: logger.warning( "Could not instantiate Qwen2.5-VL vision tower (%s). " "Understanding (x2t) path will be unavailable; generation (t2i) " "does not use the ViT.", e, ) return None return LanceQwen2_5_VLNaViTWrapper(vision, spatial_merge_size=int(vit_cfg_dict.get("spatial_merge_size", 2))) # ------------------------------------------------------------------ # # Lance text-to-video forward path # ------------------------------------------------------------------ # def forward(self, req): # type: ignore[override] """Dispatch on prompt modality. - ``modalities == ["video"]`` (text-to-video) → :meth:`_forward_t2v` (3-D latents + ``LanceWanVAE.decode_video``). - ``modalities == ["text"]`` + ``multi_modal_data.video`` (x2t_video) → :meth:`_forward_x2t_video` (multi-frame Qwen2.5-VL ViT prefill). - ``modalities == ["image"]`` + ``multi_modal_data.img2img`` (image_edit) → :meth:`_forward_image_edit` (Lance-native VAE+ViT prefill + image gen). - ``modalities == ["video"]`` + ``multi_modal_data.video`` (video_edit) → :meth:`_forward_video_edit` (Lance-native multi-frame VAE+ViT prefill + video gen). - Everything else falls through to :meth:`BagelPipeline.forward` (t2i, x2t_image). """ first_prompt = req.prompts[0] if req.prompts else None modalities: list[str] = [] mm_data: dict = {} if isinstance(first_prompt, dict): modalities = first_prompt.get("modalities") or [] mm_data = first_prompt.get("multi_modal_data") or {} if "video" in modalities and mm_data.get("video") is not None: return self._forward_video_edit(req) if "video" in modalities: return self._forward_t2v(req) if "text" in modalities and mm_data.get("video") is not None: return self._forward_x2t_video(req) if "image" in modalities and (mm_data.get("img2img") is not None or mm_data.get("image") is not None): return self._forward_image_edit(req) # t2i falls through to BAGEL parent. Inject Lance defaults if the # caller hasn't set them: upstream uses ``cfg_text_scale=4.0`` and # ``cfg_vit_scale=1.0`` (= cfg_img_scale=1.0 in BAGEL terms). sp = req.sampling_params if not hasattr(sp, "extra_args") or sp.extra_args is None: sp.extra_args = {} sp.extra_args.setdefault("cfg_img_scale", 1.0) return super().forward(req) @torch.inference_mode() def _forward_t2v(self, req): """Minimal text-to-video forward: text prefill + 3-D latent denoising + Wan2.2 multi-frame decode. Mirrors :meth:`BagelPipeline.forward`'s t2i branch but with 3-D latents (no image input, no x2t).""" from copy import deepcopy from vllm_omni.diffusion.data import DiffusionOutput from .lance_transformer import NaiveCache first_prompt = req.prompts[0] if isinstance(first_prompt, dict): prompt = first_prompt.get("prompt") or "" extra_args = first_prompt.get("extra_args") or {} else: prompt = str(first_prompt) extra_args = {} # Sampling-side extras override prompt-side. sp_extra = getattr(req.sampling_params, "extra_args", {}) or {} extra_args = {**extra_args, **sp_extra} # Video shape. T = number of RGB frames (1..121), H/W in pixels. T = int(extra_args.get("num_frames", 25)) H = int(req.sampling_params.height or extra_args.get("video_height", 480)) W = int(req.sampling_params.width or extra_args.get("video_width", 768)) max_lat = self.bagel.max_latent_size max_hw = max_lat * self.bagel.latent_downsample if H > max_hw or W > max_hw: raise ValueError(f"Requested video resolution {H}x{W} exceeds Lance limit {max_hw}x{max_hw}") downsample_t = int(getattr(self.bagel.config.vae_config, "downsample_temporal", 4)) max_T_lat = LANCE_DEFAULTS.max_num_video_latent_frames if (T - 1) // downsample_t + 1 > max_T_lat: raise ValueError( f"Requested num_frames={T} exceeds Lance video temporal limit " f"{(max_T_lat - 1) * downsample_t + 1} (max_num_video_latent_frames={max_T_lat})" ) video_shape = (T, H, W) logger.info("Lance t2v: video_shape=%s", video_shape) cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.0)) # Upstream Lance defaults cfg_interval=[0.4, 1.0] (turn off CFG below # t=0.4); without this vllm-omni keeps CFG on for late iters and the # last few denoise steps diverge from upstream. cfg_renorm_type / # cfg_renorm_min default to 'global' / 0 on both sides. cfg_interval = tuple(extra_args.get("cfg_interval", (0.4, 1.0))) cfg_renorm_type = str(extra_args.get("cfg_renorm_type", "global")) cfg_renorm_min = float(extra_args.get("cfg_renorm_min", 0.0)) timestep_shift = float(extra_args.get("timestep_shift", LANCE_DEFAULTS.timestep_shift)) num_timesteps = int(req.sampling_params.num_inference_steps or LANCE_DEFAULTS.num_timesteps) if req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.seed) # ---- Build positive-prompt KV cache ---- gen_context = { "kv_lens": [0], "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } cfg_text_context = deepcopy(gen_context) gen_input_text, newlens, new_rope = self.bagel.prepare_prompts( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], prompts=[prompt], tokenizer=self.tokenizer, new_token_ids=self.new_token_ids, ) for k, v in gen_input_text.items(): if torch.is_tensor(v): gen_input_text[k] = v.to(self.device) with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): gen_context["past_key_values"] = self.bagel.forward_cache_update_text( gen_context["past_key_values"], **gen_input_text ) gen_context["kv_lens"] = newlens gen_context["ropes"] = new_rope # ---- Build CFG text-unconditional KV cache (empty prompt) ---- if cfg_text_scale > 1.0: neg_prompt = str(extra_args.get("negative_prompt") or "") neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( curr_kvlens=cfg_text_context["kv_lens"], curr_rope=cfg_text_context["ropes"], prompts=[neg_prompt], tokenizer=self.tokenizer, new_token_ids=self.new_token_ids, ) for k, v in neg_input.items(): if torch.is_tensor(v): neg_input[k] = v.to(self.device) with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): cfg_text_context["past_key_values"] = self.bagel.forward_cache_update_text( cfg_text_context["past_key_values"], **neg_input ) cfg_text_context["kv_lens"] = neg_newlens cfg_text_context["ropes"] = neg_rope # ---- 3-D latent init + CFG side metadata ---- gen_input_lat = self.bagel.prepare_video_latent( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], video_shapes=[video_shape], new_token_ids=self.new_token_ids, ) for k, v in gen_input_lat.items(): if torch.is_tensor(v): gen_input_lat[k] = v.to(self.device) cfg_text_lat = self.bagel.prepare_video_latent_cfg( curr_kvlens=cfg_text_context["kv_lens"], curr_rope=cfg_text_context["ropes"], video_shapes=[video_shape], ) cfg_img_lat = self.bagel.prepare_video_latent_cfg( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], video_shapes=[video_shape], ) for k, v in cfg_text_lat.items(): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) for k, v in cfg_img_lat.items(): if torch.is_tensor(v): cfg_img_lat[k] = v.to(self.device) # ---- Denoising loop (Bagel.generate_image is rank-agnostic over packed tokens) ---- with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"], cfg_img_past_key_values=gen_context["past_key_values"], # no img CFG branch num_timesteps=num_timesteps, timestep_shift=timestep_shift, cfg_text_scale=cfg_text_scale, cfg_img_scale=1.0, cfg_interval=cfg_interval, cfg_renorm_type=cfg_renorm_type, cfg_renorm_min=cfg_renorm_min, **gen_input_lat, cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"], cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"], cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"], cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"], cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"], cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"], cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"], ) frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) # Convert numpy frames to PIL.Image list for downstream serialization. frames = [Image.fromarray(f) for f in frames_np] logger.info("Lance t2v: decoded %d frames at %dx%d", len(frames), frames[0].width, frames[0].height) return DiffusionOutput( output=frames[0], custom_output={"video_frames": frames, "video_shape": video_shape}, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @torch.inference_mode() def _forward_image_edit(self, req): """Image edit (img2img): reference image + text prompt → modified image. Matches upstream Lance's image_edit prefill structure (see ``Lance.validation_gen_KVcache`` and ``validation_dataset.py``). The sequence is laid out as 5 prefill segments + 1 noise QUERY: seg1 (causal, modality=system_prompt): ``<|im_start|>system\\n{sys}<|im_end|>\\n<|im_start|>user\\n`` seg2 (full, modality=ref_vit): ViT(ref) — ref image through Qwen2.5-VL ViT seg3 (full_noise, modality=ref_source): VAE(ref) — ref image through Wan2.2 VAE, time_embed(t=0) seg4 (causal, modality=text): the user instruction seg5 (causal, modality=system_prompt): ``<|im_end|>\\n<|im_start|>assistant\\n`` seg6 (noise QUERY): the gen latent to denoise cfg_text_context skips ONLY seg4 (the user text instruction) — every other segment, including the system header, ViT and VAE prefill, the separator and the noise query, is shared between gen_context and cfg_text_context. The rope counter still advances by seg4's length for cfg_text_context so that seg5's rope position matches between the two branches. cfg_img is OFF by default (``cfg_img_scale=1.0`` matches upstream's ``cfg_vit_scale=1.0``). """ import numpy as np from vllm_omni.diffusion.data import DiffusionOutput from .lance_transformer import NaiveCache from .prompts import SYSTEM_PROMPTS first_prompt = req.prompts[0] assert isinstance(first_prompt, dict), "image_edit requires dict-style prompt" # The caller may pass either a raw user instruction (e.g. "Remove the # hat from the painting.") OR an already-rendered template string # (with system + chat template). Extract just the user instruction # so we can rebuild the segments cleanly. user_text = first_prompt.get("user_text") rendered = first_prompt.get("prompt") or "" if not user_text: user_text = _extract_user_instruction(rendered) mm_data = first_prompt.get("multi_modal_data") or {} image_input = mm_data.get("img2img") or mm_data.get("image") if image_input is None: raise ValueError("image_edit requires multi_modal_data.img2img (or .image).") if not isinstance(image_input, list): image_input = [image_input] image_input = [Image.open(im) if isinstance(im, str) else im for im in image_input] extra_args = first_prompt.get("extra_args") or {} sp_extra = getattr(req.sampling_params, "extra_args", {}) or {} extra_args = {**extra_args, **sp_extra} cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.0)) cfg_img_scale = float(extra_args.get("cfg_img_scale", 1.0)) cfg_interval = tuple(extra_args.get("cfg_interval", (0.4, 1.0))) cfg_renorm_type = str(extra_args.get("cfg_renorm_type", "global")) cfg_renorm_min = float(extra_args.get("cfg_renorm_min", 0.0)) timestep_shift = float(extra_args.get("timestep_shift", LANCE_DEFAULTS.timestep_shift)) num_timesteps = int(req.sampling_params.num_inference_steps or LANCE_DEFAULTS.num_timesteps) # Resize the reference image to a multiple of latent_downsample and # within the max latent grid. stride = self.bagel.latent_downsample max_hw = int(self.bagel.max_latent_size * stride) def _resize_to_stride(img): if img.mode != "RGB": img = img.convert("RGB") w, h = img.size scale = min(max_hw / max(w, h), 1.0) scale = max(scale, min(256, max_hw) / min(w, h)) new_w = max(stride, int(round(w * scale / stride) * stride)) new_h = max(stride, int(round(h * scale / stride) * stride)) new_w = min(new_w, max_hw) new_h = min(new_h, max_hw) if new_w != w or new_h != h: img = img.resize((new_w, new_h), Image.BICUBIC) return img image_input = [_resize_to_stride(im) for im in image_input] resized_w, resized_h = image_input[0].size image_shape = (resized_h, resized_w) logger.info("Lance image_edit: ref image %dx%d, user_text=%r", resized_w, resized_h, user_text) def vae_transforms(img): arr = torch.from_numpy(np.array(img)).float() / 127.5 - 1.0 return arr.permute(2, 0, 1) # (C, H, W) def vit_transforms(img): return self.image_processor(images=img, return_tensors="pt").pixel_values[0] if req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.seed) gen_context = { "kv_lens": [0], "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } cfg_text_context = { "kv_lens": [0], "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } autocast_kwargs = dict( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ) # ----- prefill helpers (no bos/eos wrapping — that's only for free- # form prompts; the segmented template provides its own markers) ----- def _raw_text_prefill(ctx, text_str): """Prefill `text_str` as-is — no automatic bos/eos. ``Bagel.prepare_prompts`` wraps with ``[bos] + ids + [eos]`` which is fine for a single user prompt but breaks segment-by-segment templating where each segment already carries its own framing tokens. This helper tokenizes raw and feeds the LLM directly. """ text_ids = self.tokenizer.encode(text_str, add_special_tokens=False) if not text_ids: return curr_kvlen = ctx["kv_lens"][0] curr_rope = ctx["ropes"][0] seq_len = len(text_ids) inp = { "text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device), "packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device), "packed_text_position_ids": torch.arange( curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device ), "packed_text_indexes": torch.arange( curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device ), "packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device), "key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device), } with torch.autocast(**autocast_kwargs): ctx["past_key_values"] = self.bagel.forward_cache_update_text(ctx["past_key_values"], **inp) ctx["kv_lens"] = [curr_kvlen + seq_len] ctx["ropes"] = [curr_rope + seq_len] def _vit_prefill(ctx): inp, new_kvlens, new_rope = self.bagel.prepare_vit_images( curr_kvlens=ctx["kv_lens"], curr_rope=ctx["ropes"], images=image_input, transforms=vit_transforms, new_token_ids=self.new_token_ids, ) for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) with torch.autocast(**autocast_kwargs): ctx["past_key_values"] = self.bagel.forward_cache_update_vit(ctx["past_key_values"], **inp) ctx["kv_lens"] = new_kvlens ctx["ropes"] = new_rope # Encode the reference image ONCE so both gen and cfg branches share # the same posterior sample. ``LanceWanVAE.encode`` uses # ``mu + std * randn_like(std)`` (Wan2.2 VAE's reparameterize path) # which yields a different latent every call; calling it separately # per branch drifts gen-branch K, V cache from cfg-branch by ~6% # rel_l2 even though both should see the same ref. _ref_padded_latent_cache = {"v": None} def _vae_prefill(ctx): inp, new_kvlens, new_rope = self.bagel._lance_native_prepare_vae_images( curr_kvlens=ctx["kv_lens"], curr_rope=ctx["ropes"], images=image_input, transforms=vae_transforms, new_token_ids=self.new_token_ids, is_video=False, ) for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) # Compute the ref VAE latent on the first prefill call; reuse # for subsequent branches. if _ref_padded_latent_cache["v"] is None: _ref_padded_latent_cache["v"] = self.vae.encode(inp["padded_images"].to(self.device)) inp["precomputed_latent"] = _ref_padded_latent_cache["v"] with torch.autocast(**autocast_kwargs): ctx["past_key_values"] = self.bagel.forward_cache_update_vae(self.vae, ctx["past_key_values"], **inp) ctx["kv_lens"] = new_kvlens ctx["ropes"] = new_rope # ----- prefill segments in upstream order ----- sys_prompt = SYSTEM_PROMPTS[("image_edit", "image")] seg1_str = f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n" seg5_str = "<|im_end|>\n<|im_start|>assistant\n" # seg1: system + user header (shared) _raw_text_prefill(gen_context, seg1_str) if cfg_text_scale > 1.0: _raw_text_prefill(cfg_text_context, seg1_str) # seg2: ViT(ref) (shared) _vit_prefill(gen_context) if cfg_text_scale > 1.0: _vit_prefill(cfg_text_context) # seg3: VAE(ref) with timestep=0 (shared) # # IMPORTANT: upstream Lance places the gen latent (the noise QUERY) # at the SAME mRoPE positions as the ref VAE — they share the same # spatial grid (T_lat × H_lat × W_lat), and the model treats them # as two timesteps of the same latent location. Concretely, upstream # ``get_rope_index`` emits ``[base_t, base_h+hi, base_w+wi]`` for # both the VAE cond block AND the noise block, anchored at the SAME # base. This is the inductive bias that lets the model "edit" the # ref image — without this overlap the model has no way to map noise # tokens to ref tokens. # # We therefore snapshot ``ropes`` BEFORE the VAE prefill and reuse # that as the gen latent's ``curr_rope`` below. rope_before_vae = gen_context["ropes"][0] cfg_rope_before_vae = cfg_text_context["ropes"][0] if cfg_text_scale > 1.0 else rope_before_vae _vae_prefill(gen_context) if cfg_text_scale > 1.0: _vae_prefill(cfg_text_context) # seg4: user instruction (gen only). Upstream Lance's cfg_text # branch is a SHORTER sequence with the user_text segment REMOVED # (see ``uncond_split_pro_kvcache`` + ``get_rope_index`` on the # filtered text_ids), so cfg_text's rope counter MUST NOT advance # past the user_text region. seg5 (separator) therefore lands at # different absolute rope positions in gen vs cfg branches — that's # the intended behavior. _raw_text_prefill(gen_context, user_text) # seg5: separator + assistant header (gen + cfg). Each branch # places it at its OWN current rope. _raw_text_prefill(gen_context, seg5_str) if cfg_text_scale > 1.0: _raw_text_prefill(cfg_text_context, seg5_str) # -- (4) Gen latent at the SAME rope position as the ref VAE block. # See the comment on `rope_before_vae` above. Note the KV cache # still contains all prefilled segments (system, ViT, ref VAE, user # text, separator); only the gen latent's QUERY rope coordinates are # rewound to overlap with the ref VAE. gen_input_lat = self.bagel.prepare_vae_latent( curr_kvlens=gen_context["kv_lens"], curr_rope=[rope_before_vae], image_sizes=[image_shape], new_token_ids=self.new_token_ids, ) for k, v in gen_input_lat.items(): if torch.is_tensor(v): gen_input_lat[k] = v.to(self.device) cfg_text_lat = self.bagel.prepare_vae_latent_cfg( curr_kvlens=cfg_text_context["kv_lens"], curr_rope=[cfg_rope_before_vae], image_sizes=[image_shape], ) for k, v in cfg_text_lat.items(): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) # cfg_img branch: keep off (cfg_img_scale=1.0 matches upstream # ``cfg_vit_scale=1.0``). generate_image needs the cfg_img_* args # but skips the branch when scale<=1.0; just point them at the # cond's lat metadata so shapes/types are fine. cfg_img_lat = cfg_text_lat with torch.autocast(**autocast_kwargs): latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"] if cfg_text_scale > 1.0 else gen_context["past_key_values"], cfg_img_past_key_values=gen_context["past_key_values"], num_timesteps=num_timesteps, timestep_shift=timestep_shift, cfg_text_scale=cfg_text_scale, cfg_img_scale=cfg_img_scale, cfg_interval=cfg_interval, cfg_renorm_type=cfg_renorm_type, cfg_renorm_min=cfg_renorm_min, **gen_input_lat, cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"], cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"], cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"], cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"], cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"], cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"], cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"], ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) return DiffusionOutput( output=img, custom_output={"image_shape": image_shape}, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @torch.inference_mode() def _forward_video_edit(self, req): """Video edit: reference video + text prompt → modified video. Mirrors :meth:`_forward_image_edit`'s segmented prefill layout with a temporal axis threaded through both prefill and the gen latent. Verified against upstream Lance's ``Lance.validation_gen_KVcache`` for ``--task video_edit --resolution video_480p``: seg1 (causal, modality=system_prompt): ``<|im_start|>system\\n{sys}<|im_end|>\\n<|im_start|>user\\n`` seg2 (full, modality=ref_vit): ViT(video) — ref video through Qwen2.5-VL ViT, t-axis +1000 seg3 (full_noise, modality=ref_source): VAE(video) — ref video through Wan2.2 VAE, time_embed(t=0). Rope shares range with the gen noise latent (modality 1==2 trick) so the model can align ref↔output per ``(t,h,w)``. seg4 (causal, modality=text): the user instruction seg5 (causal, modality=system_prompt): ``<|im_end|>\\n<|im_start|>assistant\\n`` seg6 (noise QUERY): 3-D gen latent at the SAME rope as seg3. cfg_text_context skips ONLY seg4 (user text), matching upstream's ``uncond_split_pro_kvcache``. cfg_img defaults to off (``cfg_img_scale=1.0`` → upstream's ``cfg_vit_scale=1.0``). Replaces the previous shifted-rope layout (``_EDIT_SHIFT=10000``) which kept ViT/text at a far-away rope range — that caused 15.5% rel_l2 divergence at layer 0 output (cond) vs upstream because the noise query attended to a cached K/V at rope (0..47) while upstream attends at rope (111..159). """ import numpy as _np from vllm_omni.diffusion.data import DiffusionOutput from .lance_transformer import NaiveCache from .prompts import SYSTEM_PROMPTS first_prompt = req.prompts[0] assert isinstance(first_prompt, dict), "video_edit requires dict-style prompt" user_text = first_prompt.get("user_text") rendered = first_prompt.get("prompt") or "" if not user_text: user_text = _extract_user_instruction(rendered) mm_data = first_prompt.get("multi_modal_data") or {} video_input = mm_data.get("video") if video_input is None: raise ValueError("video_edit requires multi_modal_data.video.") # Resolve to raw (T, H, W, 3) uint8 ndarray + origin_fps so the # upstream-style bucket resize + frame sampler can replicate the # exact preprocessing path. If the caller passed an already-decoded # ndarray with no FPS, default to 24fps (matches the official # video_edit examples). extra_args = first_prompt.get("extra_args") or {} sp_extra = getattr(req.sampling_params, "extra_args", {}) or {} extra_args = {**extra_args, **sp_extra} origin_fps_default = float(extra_args.get("origin_fps", 24.0)) # Path-only fast path: use decord (matches upstream Lance's # ``data/datasets_custom/validation_dataset.py`` which decodes via # ``decord.VideoReader``). decord vs cv2 produce different YUV→RGB # conversions (~0.1 RGB units of max diff), and that propagates into # both the VAE and ViT pixel inputs. Falling back to cv2 only when # decord is unavailable keeps single-machine parity with upstream # while not requiring decord in production. if isinstance(video_input, str): try: import decord as _decord from decord import VideoReader as _VideoReader # type: ignore vr = _VideoReader(video_input, ctx=_decord.cpu(0)) origin_fps = float(vr.get_avg_fps()) or origin_fps_default idx_all = list(range(len(vr))) video_raw = vr.get_batch(idx_all).asnumpy() # (T, H, W, 3) uint8 RGB except ImportError: import cv2 as _cv2 cap = _cv2.VideoCapture(video_input) origin_fps = float(cap.get(_cv2.CAP_PROP_FPS) or origin_fps_default) frames_bgr = [] while True: ok, f = cap.read() if not ok: break frames_bgr.append(_cv2.cvtColor(f, _cv2.COLOR_BGR2RGB)) cap.release() video_raw = _np.stack(frames_bgr, axis=0) elif isinstance(video_input, _np.ndarray): video_raw = video_input origin_fps = origin_fps_default elif torch.is_tensor(video_input): arr = video_input.detach().cpu().numpy() if arr.max() <= 1.5: arr = (arr * 255.0).clip(0, 255).astype(_np.uint8) video_raw = arr origin_fps = origin_fps_default else: raise ValueError(f"Unsupported video_input type {type(video_input)}") cfg_text_scale = float(extra_args.get("cfg_text_scale", 4.0)) cfg_img_scale = float(extra_args.get("cfg_img_scale", 1.0)) # Upstream Lance default is ``cfg_interval=[0, 1]`` (CFG active on every # denoise step). vllm-omni's previous (0.4, 1.0) skipped CFG for the # last ~10 iterations after timestep_shift=3.5, which contributed # most of the late-trajectory drift from upstream. cfg_interval = tuple(extra_args.get("cfg_interval", (0.0, 1.0))) cfg_renorm_type = str(extra_args.get("cfg_renorm_type", "global")) cfg_renorm_min = float(extra_args.get("cfg_renorm_min", 0.0)) timestep_shift = float(extra_args.get("timestep_shift", LANCE_DEFAULTS.timestep_shift)) num_timesteps = int(req.sampling_params.num_inference_steps or LANCE_DEFAULTS.num_timesteps) # Run upstream's BucketResize + frame sampler so VAE / ViT see # byte-identical pixels to the reference implementation. Without this, # vllm-omni's ViT grid was (T=23, H=40, W=54) for the car video versus # upstream's (T=23, H=36, W=50) — different K/V at ViT positions # accounted for the residual ~9% rel_l2 at layer 0 output. from .lance_transformer import LanceBagel as _LanceBagel vae_video, vit_pixels, vit_grid_thw, sampled_T = _LanceBagel._lance_video_preprocess(video_raw, origin_fps) video_chw = vae_video # (C, T_sampled, H_vae, W_vae) — model-space [-1,1] T_sampled = int(video_chw.shape[1]) H_vae = int(video_chw.shape[2]) W_vae = int(video_chw.shape[3]) video_shape = (T_sampled, H_vae, W_vae) logger.info( "Lance video_edit: input %dx%dx%d → sampled %dx%dx%d (T,H,W) @ %.2ffps, ViT grid_thw=%s, user_text=%r", video_raw.shape[0], video_raw.shape[1], video_raw.shape[2], T_sampled, H_vae, W_vae, origin_fps, tuple(vit_grid_thw[0].tolist()), user_text, ) if req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.seed) gen_context = { "kv_lens": [0], "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } cfg_text_context = { "kv_lens": [0], "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } autocast_kwargs = dict( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ) def _raw_text_prefill(ctx, text_str): """Prefill ``text_str`` as raw token ids — no bos/eos wrapping.""" text_ids = self.tokenizer.encode(text_str, add_special_tokens=False) if not text_ids: return curr_kvlen = ctx["kv_lens"][0] curr_rope = ctx["ropes"][0] seq_len = len(text_ids) inp = { "text_token_lens": torch.tensor([seq_len], dtype=torch.int, device=self.device), "packed_text_ids": torch.tensor(text_ids, dtype=torch.long, device=self.device), "packed_text_position_ids": torch.arange( curr_rope, curr_rope + seq_len, dtype=torch.long, device=self.device ), "packed_text_indexes": torch.arange( curr_kvlen, curr_kvlen + seq_len, dtype=torch.long, device=self.device ), "packed_key_value_indexes": torch.arange(0, curr_kvlen, dtype=torch.long, device=self.device), "key_values_lens": torch.tensor([curr_kvlen], dtype=torch.int, device=self.device), } with torch.autocast(**autocast_kwargs): ctx["past_key_values"] = self.bagel.forward_cache_update_text(ctx["past_key_values"], **inp) ctx["kv_lens"] = [curr_kvlen + seq_len] ctx["ropes"] = [curr_rope + seq_len] def _vit_prefill(ctx): # Feed the bucket-resized + patchified ViT pixels directly so the # Qwen2VLImageProcessor path (smart-resize, wrong grid) is bypassed. inp, new_kvlens, new_rope = self.bagel.prepare_vit_videos( curr_kvlens=ctx["kv_lens"], curr_rope=ctx["ropes"], videos=[video_raw], # passed through for length/iter only new_token_ids=self.new_token_ids, precomputed_vit=[(vit_pixels, vit_grid_thw)], ) for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) with torch.autocast(**autocast_kwargs): ctx["past_key_values"] = self.bagel.forward_cache_update_vit(ctx["past_key_values"], **inp) ctx["kv_lens"] = new_kvlens ctx["ropes"] = new_rope # Encode the reference video ONCE so both gen and cfg branches share # the same posterior sample (see image_edit fix — Wan2.2 VAE's # ``mu + std * randn_like(std)`` would otherwise drift between # branches). _ref_padded_latent_cache = {"v": None} def _vae_transforms(_): return video_chw def _vae_prefill(ctx): inp, new_kvlens, new_rope = self.bagel._lance_native_prepare_vae_images( curr_kvlens=ctx["kv_lens"], curr_rope=ctx["ropes"], images=[video_chw], transforms=_vae_transforms, new_token_ids=self.new_token_ids, is_video=True, ) for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) if _ref_padded_latent_cache["v"] is None: _ref_padded_latent_cache["v"] = self.vae.encode(inp["padded_images"].to(self.device)) inp["precomputed_latent"] = _ref_padded_latent_cache["v"] with torch.autocast(**autocast_kwargs): ctx["past_key_values"] = self.bagel.forward_cache_update_vae(self.vae, ctx["past_key_values"], **inp) ctx["kv_lens"] = new_kvlens ctx["ropes"] = new_rope # ----- prefill segments in upstream order ----- sys_prompt = SYSTEM_PROMPTS[("video_edit", "video")] seg1_str = f"<|im_start|>system\n{sys_prompt}<|im_end|>\n<|im_start|>user\n" seg5_str = "<|im_end|>\n<|im_start|>assistant\n" # seg1: system + user header (shared) _raw_text_prefill(gen_context, seg1_str) if cfg_text_scale > 1.0: _raw_text_prefill(cfg_text_context, seg1_str) # seg2: ViT(ref) (shared) _vit_prefill(gen_context) if cfg_text_scale > 1.0: _vit_prefill(cfg_text_context) # seg3: VAE(ref) — snapshot rope BEFORE the prefill so the gen noise # latent can be placed at the SAME rope range (matches upstream's # modality==1≡2 trick). rope_before_vae = gen_context["ropes"][0] cfg_rope_before_vae = cfg_text_context["ropes"][0] if cfg_text_scale > 1.0 else rope_before_vae _vae_prefill(gen_context) if cfg_text_scale > 1.0: _vae_prefill(cfg_text_context) # seg4: user instruction (gen only — cfg branch skips per upstream). _raw_text_prefill(gen_context, user_text) # seg5: separator + assistant header (both gen and cfg, each at its # own current rope; cfg's lands earlier than gen's because cfg # skipped seg4). _raw_text_prefill(gen_context, seg5_str) if cfg_text_scale > 1.0: _raw_text_prefill(cfg_text_context, seg5_str) # seg6 (noise QUERY): gen latent at the SAME rope as the ref VAE # block. KV cache still contains sys+ViT+VAE_ref+user_text+sep at # their post-prefill rope positions; only the noise query's Q rope # is rewound to overlap with the VAE_ref block. gen_input_lat = self.bagel.prepare_video_latent( curr_kvlens=gen_context["kv_lens"], curr_rope=[rope_before_vae], video_shapes=[video_shape], new_token_ids=self.new_token_ids, ) for k, v in gen_input_lat.items(): if torch.is_tensor(v): gen_input_lat[k] = v.to(self.device) cfg_text_lat = self.bagel.prepare_video_latent_cfg( curr_kvlens=cfg_text_context["kv_lens"], curr_rope=[cfg_rope_before_vae], video_shapes=[video_shape], ) for k, v in cfg_text_lat.items(): if torch.is_tensor(v): cfg_text_lat[k] = v.to(self.device) # cfg_img branch is off (cfg_img_scale=1.0); generate_image needs the # arg surface but skips the branch when scale<=1.0. cfg_img_lat = cfg_text_lat with torch.autocast(**autocast_kwargs): latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"] if cfg_text_scale > 1.0 else gen_context["past_key_values"], cfg_img_past_key_values=gen_context["past_key_values"], num_timesteps=num_timesteps, timestep_shift=timestep_shift, cfg_text_scale=cfg_text_scale, cfg_img_scale=cfg_img_scale, cfg_interval=cfg_interval, cfg_renorm_type=cfg_renorm_type, cfg_renorm_min=cfg_renorm_min, **gen_input_lat, cfg_text_packed_position_ids=cfg_text_lat["cfg_packed_position_ids"], cfg_text_packed_query_indexes=cfg_text_lat["cfg_packed_query_indexes"], cfg_text_key_values_lens=cfg_text_lat["cfg_key_values_lens"], cfg_text_packed_key_value_indexes=cfg_text_lat["cfg_packed_key_value_indexes"], cfg_img_packed_position_ids=cfg_img_lat["cfg_packed_position_ids"], cfg_img_packed_query_indexes=cfg_img_lat["cfg_packed_query_indexes"], cfg_img_key_values_lens=cfg_img_lat["cfg_key_values_lens"], cfg_img_packed_key_value_indexes=cfg_img_lat["cfg_packed_key_value_indexes"], ) frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) frames = [Image.fromarray(f) for f in frames_np] return DiffusionOutput( output=frames[0], custom_output={"video_frames": frames, "video_shape": video_shape}, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @torch.inference_mode() def _forward_x2t_video(self, req): """Video understanding (x2t_video): video → text caption / VQA answer. Mirrors :meth:`_forward_t2v` for the prefill / generation bookkeeping but routes the video through the Qwen2.5-VL ViT (no VAE prefill — Lance's training has Wan2.2 latents only for image_edit / video_edit, not for understanding). """ from vllm_omni.diffusion.data import DiffusionOutput from .lance_transformer import NaiveCache first_prompt = req.prompts[0] assert isinstance(first_prompt, dict), "x2t_video requires dict-style prompt" prompt = first_prompt.get("prompt") or "" video_input = (first_prompt.get("multi_modal_data") or {}).get("video") if video_input is None: raise ValueError("x2t_video requires multi_modal_data.video to be a video tensor/array/path.") # Accept a path → load via imageio; else assume tensor / numpy array (T, H, W, 3). if isinstance(video_input, str): import imageio.v3 as iio video_input = iio.imread(video_input) # Wrap single video as a list for the prep function. videos = [video_input] extra_args = first_prompt.get("extra_args") or {} sp_extra = getattr(req.sampling_params, "extra_args", {}) or {} extra_args = {**extra_args, **sp_extra} max_text_tokens = int(extra_args.get("max_think_tokens", 200)) do_sample = bool(extra_args.get("do_sample", False)) text_temperature = float(extra_args.get("text_temperature", 0.3)) if req.sampling_params.seed is not None: torch.manual_seed(req.sampling_params.seed) if self.device.type == "cuda": torch.cuda.manual_seed(req.sampling_params.seed) # ---- Text-prompt prefill ---- gen_context = { "kv_lens": [0], "ropes": [0], "past_key_values": NaiveCache(self.bagel.config.llm_config.num_hidden_layers), } if prompt: gen_input_text, newlens, new_rope = self.bagel.prepare_prompts( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], prompts=[prompt], tokenizer=self.tokenizer, new_token_ids=self.new_token_ids, ) for k, v in gen_input_text.items(): if torch.is_tensor(v): gen_input_text[k] = v.to(self.device) with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): gen_context["past_key_values"] = self.bagel.forward_cache_update_text( gen_context["past_key_values"], **gen_input_text ) gen_context["kv_lens"] = newlens gen_context["ropes"] = new_rope # ---- Multi-frame ViT prefill ---- gen_input_vit, newlens_vit, new_rope_vit = self.bagel.prepare_vit_videos( curr_kvlens=gen_context["kv_lens"], curr_rope=gen_context["ropes"], videos=videos, new_token_ids=self.new_token_ids, ) for k, v in gen_input_vit.items(): if torch.is_tensor(v): gen_input_vit[k] = v.to(self.device) with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): gen_context["past_key_values"] = self.bagel.forward_cache_update_vit( gen_context["past_key_values"], **gen_input_vit ) gen_context["kv_lens"] = newlens_vit gen_context["ropes"] = new_rope_vit # ---- Text generation ---- start_input = self.bagel.prepare_start_tokens(gen_context["kv_lens"], gen_context["ropes"], self.new_token_ids) for k, v in start_input.items(): if torch.is_tensor(v): start_input[k] = v.to(self.device) with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): token_ids = self.bagel.generate_text( past_key_values=gen_context["past_key_values"], max_length=max_text_tokens, do_sample=do_sample, temperature=text_temperature, end_token_id=self.new_token_ids["eos_token_id"], **start_input, ) decoded = self.tokenizer.decode(token_ids[:, 0].tolist()) text_output = decoded.split("<|im_end|>")[0] if "<|im_start|>" in text_output: text_output = text_output.split("<|im_start|>")[-1] logger.info("Lance x2t_video: generated %d tokens", token_ids.shape[0]) return DiffusionOutput( output=text_output, custom_output={"text_output": text_output}, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @staticmethod def _decode_video_from_latent( bagel, vae: LanceWanVAE, latent: torch.Tensor, video_shape: tuple[int, int, int], ): """Pack a flat Lance latent into ``(B, 48, t, h, w)`` and decode to a ``(T_lat, H, W, 3)`` numpy frame list via :meth:`LanceWanVAE.decode_video`. """ T, H, W = video_shape downsample_t = int(getattr(bagel.config.vae_config, "downsample_temporal", 4)) t_lat = (T - 1) // downsample_t + 1 h_lat = H // bagel.latent_downsample w_lat = W // bagel.latent_downsample p = bagel.latent_patch_size c = bagel.latent_channel # ``latent`` is flat ``(t_lat*h_lat*w_lat, c * p * p)`` (Lance image # path uses ``p=1`` so this is ``(N, c)``). Reshape into the 5-D layout # the Wan2.2 VAE decoder consumes. latent = latent.reshape(t_lat, h_lat, w_lat, p, p, c) latent = torch.einsum("thwpqc->cthpwq", latent) latent = latent.reshape(1, c, t_lat, h_lat * p, w_lat * p) vae_dtype = next(vae.parameters()).dtype latent = latent.to(vae_dtype) video = vae.decode_video(latent) # (1, 3, T_out, H_out, W_out) # Scale [-1,1] -> [0,255] uint8 frames; the upstream VAE already clamps. video = (video[0] * 0.5 + 0.5).clamp(0, 1).permute(1, 2, 3, 0) * 255 return video.to(torch.uint8).cpu().numpy() # (T_out, H_out, W_out, 3) def _build_wan22_vae(self, repo_root: str) -> LanceWanVAE: """Wrap the bundled ``Wan2.2_VAE.pth`` with a BAGEL-compatible ``.encode(images)`` / ``.decode(latent)`` surface. Uses the upstream Wan2.2 VAE module (ported in :mod:`wan_vae`) since the released ``.pth`` is keyed for that module — *not* for the diffusers ``AutoencoderKLWan`` state dict that ``OmniAutoencoderKLWan`` expects. Single images are treated as 1-frame clips; video latents pass through the same module via the 5-D ``encode_video``/``decode_video`` path. Wan2.2 VAE constants: z=48ch, /16 spatial, /4 temporal. """ vae_path = os.path.join(repo_root, _VAE_FILE) return LanceWanVAE(vae_path=vae_path, dtype=torch.bfloat16, device=self.device)