diff --git a/examples/offline_inference/lance/gradio_demo.py b/examples/offline_inference/lance/gradio_demo.py index a8297288a..d2628e58a 100644 --- a/examples/offline_inference/lance/gradio_demo.py +++ b/examples/offline_inference/lance/gradio_demo.py @@ -288,18 +288,11 @@ def _init_models( f" replicas={replicas_per_omni}...", flush=True, ) - # Lance is single-stage diffusion — no deploy YAML. The defaults - # below match the (now-removed) ``vllm_omni/deploy/lance.yaml`` so - # ``create_default_diffusion`` produces the same stage config. - lance_engine_defaults = dict( - pipeline="lance", - max_num_batched_tokens=32768, - max_num_seqs=1, - enforce_eager=True, - trust_remote_code=True, - enable_prefix_caching=False, - async_chunk=False, - ) + # Lance needs ``vllm_omni/deploy/lance.yaml`` as the pipeline-selector + # (Lance HF config has no ``model_type`` field for auto-detect). Pass + # via ``deploy_config`` — flat kwargs alone don't trigger pipeline + # resolution. + lance_engine_defaults = dict(deploy_config="vllm_omni/deploy/lance.yaml") img_kwargs: dict = {"model": img_ckpt, **lance_engine_defaults} if img_devices != "0": diff --git a/vllm_omni/deploy/lance.yaml b/vllm_omni/deploy/lance.yaml new file mode 100644 index 000000000..25e0afd4d --- /dev/null +++ b/vllm_omni/deploy/lance.yaml @@ -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 diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index e1bd93f94..84d74c187 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -1852,7 +1852,11 @@ class Bagel(nn.Module): "num_branches": num_branches, "seq_len": seq_len, "batched_query_lens": packed_seqlens.repeat(num_branches), - "batched_position_ids": torch.cat(branches_pid), + # Concat along the *sequence* (last) dim — works for both + # scalar 1-D position ids ``(S,)`` and Lance's 3-D mRoPE + # ``(3, S)``. Default ``dim=0`` collapses ``(3, S) + (3, S)`` + # to ``(6, S)`` and breaks ``BagelRotaryEmbedding``. + "batched_position_ids": torch.cat(branches_pid, dim=-1), "batched_text_indexes": torch.cat( [packed_text_indexes + b_idx * seq_len for b_idx in range(num_branches)] ), diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 137ba3b3c..cc554f2f5 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -820,6 +820,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, diff --git a/vllm_omni/diffusion/models/lance/lance_transformer.py b/vllm_omni/diffusion/models/lance/lance_transformer.py index bdc495055..012ba3f42 100644 --- a/vllm_omni/diffusion/models/lance/lance_transformer.py +++ b/vllm_omni/diffusion/models/lance/lance_transformer.py @@ -1535,14 +1535,14 @@ class LanceBagel(Bagel): "packed_text_indexes": packed_text_indexes, } + # ``Qwen2MoTForCausalLM.forward`` post-main-merge no longer accepts + # ``packed_query_indexes`` / ``key_values_lens`` / + # ``packed_key_value_indexes`` — derived internally now. output = self.language_model.forward( packed_query_sequence=packed_query_sequence, query_lens=packed_seqlens, packed_query_position_ids=packed_position_ids, - packed_query_indexes=packed_indexes, past_key_values=past_key_values, - key_values_lens=key_values_lens, - packed_key_value_indexes=packed_key_value_indexes, update_past_key_values=True, is_causal=False, **extra_inputs, diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index e1c051db3..eb7ca29ef 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -669,6 +669,12 @@ class LancePipeline(BagelPipeline): enabled=self.device.type != "cpu", dtype=self.od_config.dtype, ): + # ``prepare_video_latent`` / ``prepare_vae_latent`` still return + # ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes`` + # but post-main-merge ``generate_image`` doesn't accept them — drop + # before unpacking via ``**gen_input_lat``. + for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"): + gen_input_lat.pop(_drop, None) latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"], @@ -682,13 +688,11 @@ class LancePipeline(BagelPipeline): 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_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` + # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge + # (derived from ``cfg_text_past_key_values``). 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"], + # ``cfg_img_*`` index/lens kwargs removed — same as above. ) frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) @@ -836,17 +840,17 @@ class LancePipeline(BagelPipeline): curr_kvlen = ctx["kv_lens"][0] curr_rope = ctx["ropes"][0] seq_len = len(text_ids) + # ``forward_cache_update_text`` post-main-merge derives the + # remaining positions from ``past_key_values`` + the query + # length; the old ``packed_text_indexes`` / + # ``packed_key_value_indexes`` / ``key_values_lens`` kwargs + # were dropped from its signature. 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) @@ -864,6 +868,12 @@ class LancePipeline(BagelPipeline): for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) + # ``prepare_vit_images`` still returns ``packed_indexes`` / + # ``packed_key_value_indexes`` / ``key_values_lens`` for + # historical reasons; ``forward_cache_update_vit`` post-main + # no longer accepts them — pop before unpacking. + for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"): + inp.pop(_drop, None) 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 @@ -979,6 +989,12 @@ class LancePipeline(BagelPipeline): cfg_img_lat = cfg_text_lat with torch.autocast(**autocast_kwargs): + # ``prepare_video_latent`` / ``prepare_vae_latent`` still return + # ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes`` + # but post-main-merge ``generate_image`` doesn't accept them — drop + # before unpacking via ``**gen_input_lat``. + for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"): + gen_input_lat.pop(_drop, None) latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"] @@ -994,13 +1010,11 @@ class LancePipeline(BagelPipeline): 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_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` + # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge + # (derived from ``cfg_text_past_key_values``). 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"], + # ``cfg_img_*`` index/lens kwargs removed — same as above. ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) @@ -1162,17 +1176,17 @@ class LancePipeline(BagelPipeline): curr_kvlen = ctx["kv_lens"][0] curr_rope = ctx["ropes"][0] seq_len = len(text_ids) + # ``forward_cache_update_text`` post-main-merge derives the + # remaining positions from ``past_key_values`` + the query + # length; the old ``packed_text_indexes`` / + # ``packed_key_value_indexes`` / ``key_values_lens`` kwargs + # were dropped from its signature. 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) @@ -1190,6 +1204,12 @@ class LancePipeline(BagelPipeline): for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) + # ``prepare_vit_images`` still returns ``packed_indexes`` / + # ``packed_key_value_indexes`` / ``key_values_lens`` for + # historical reasons; ``forward_cache_update_vit`` post-main + # no longer accepts them — pop before unpacking. + for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"): + inp.pop(_drop, None) 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 @@ -1351,6 +1371,12 @@ class LancePipeline(BagelPipeline): ) with torch.autocast(**autocast_kwargs): + # ``prepare_video_latent`` / ``prepare_vae_latent`` still return + # ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes`` + # but post-main-merge ``generate_image`` doesn't accept them — drop + # before unpacking via ``**gen_input_lat``. + for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"): + gen_input_lat.pop(_drop, None) latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"] @@ -1366,13 +1392,11 @@ class LancePipeline(BagelPipeline): 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_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` + # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge + # (derived from ``cfg_text_past_key_values``). 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"], + # ``cfg_img_*`` index/lens kwargs removed — same as above. frame_condition_token_indexes=frame_condition_token_indexes, ) @@ -1551,17 +1575,17 @@ class LancePipeline(BagelPipeline): curr_kvlen = ctx["kv_lens"][0] curr_rope = ctx["ropes"][0] seq_len = len(text_ids) + # ``forward_cache_update_text`` post-main-merge derives the + # remaining positions from ``past_key_values`` + the query + # length; the old ``packed_text_indexes`` / + # ``packed_key_value_indexes`` / ``key_values_lens`` kwargs + # were dropped from its signature. 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) @@ -1581,6 +1605,12 @@ class LancePipeline(BagelPipeline): for k, v in inp.items(): if torch.is_tensor(v): inp[k] = v.to(self.device) + # ``prepare_vit_images`` still returns ``packed_indexes`` / + # ``packed_key_value_indexes`` / ``key_values_lens`` for + # historical reasons; ``forward_cache_update_vit`` post-main + # no longer accepts them — pop before unpacking. + for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"): + inp.pop(_drop, None) 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 @@ -1676,6 +1706,12 @@ class LancePipeline(BagelPipeline): cfg_img_lat = cfg_text_lat with torch.autocast(**autocast_kwargs): + # ``prepare_video_latent`` / ``prepare_vae_latent`` still return + # ``key_values_lens`` / ``packed_indexes`` / ``packed_key_value_indexes`` + # but post-main-merge ``generate_image`` doesn't accept them — drop + # before unpacking via ``**gen_input_lat``. + for _drop in ("key_values_lens", "packed_indexes", "packed_key_value_indexes"): + gen_input_lat.pop(_drop, None) latents, *_ = self.bagel.generate_image( past_key_values=gen_context["past_key_values"], cfg_text_past_key_values=cfg_text_context["past_key_values"] @@ -1691,13 +1727,11 @@ class LancePipeline(BagelPipeline): 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_text_packed_query_indexes`` / ``cfg_text_key_values_lens`` + # / ``cfg_text_packed_key_value_indexes`` removed post-main-merge + # (derived from ``cfg_text_past_key_values``). 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"], + # ``cfg_img_*`` index/lens kwargs removed — same as above. ) frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) @@ -1777,6 +1811,11 @@ class LancePipeline(BagelPipeline): for k, v in gen_input_vit.items(): if torch.is_tensor(v): gen_input_vit[k] = v.to(self.device) + # Post-main-merge ``forward_cache_update_vit`` drops the + # legacy ``packed_indexes`` / ``packed_key_value_indexes`` / + # ``key_values_lens`` kwargs. + for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"): + gen_input_vit.pop(_drop, None) with torch.autocast(**autocast_kwargs): gen_context["past_key_values"] = self.bagel.forward_cache_update_vit( gen_context["past_key_values"], **gen_input_vit @@ -1895,6 +1934,11 @@ class LancePipeline(BagelPipeline): for k, v in gen_input_vit.items(): if torch.is_tensor(v): gen_input_vit[k] = v.to(self.device) + # Post-main-merge ``forward_cache_update_vit`` drops the + # legacy ``packed_indexes`` / ``packed_key_value_indexes`` / + # ``key_values_lens`` kwargs. + for _drop in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"): + gen_input_vit.pop(_drop, None) with torch.autocast(**autocast_kwargs): gen_context["past_key_values"] = self.bagel.forward_cache_update_vit( gen_context["past_key_values"], **gen_input_vit @@ -2003,11 +2047,12 @@ class LancePipeline(BagelPipeline): packed_start_tokens.append(nl_id) packed_query_position_ids.append(curr_position_id) curr += curr_kvlen + # ``Bagel.generate_text`` post-main-merge no longer accepts + # ``key_values_lens`` / ``packed_key_value_indexes`` — they're + # derived from ``past_key_values``. return { "packed_start_tokens": torch.tensor(packed_start_tokens, dtype=torch.long), "packed_query_position_ids": torch.tensor(packed_query_position_ids, dtype=torch.long), - "key_values_lens": torch.tensor(curr_kvlens, dtype=torch.int), - "packed_key_value_indexes": torch.tensor(packed_key_value_indexes, dtype=torch.long), } def _x2t_raw_text_prefill(self, gen_context, text_str, autocast_kwargs): @@ -2029,9 +2074,9 @@ class LancePipeline(BagelPipeline): "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), + # ``packed_text_indexes`` / ``packed_key_value_indexes`` / + # ``key_values_lens`` removed from ``forward_cache_update_text`` + # post-main-merge (derived from ``past_key_values``). } with torch.autocast(**autocast_kwargs): gen_context["past_key_values"] = self.bagel.forward_cache_update_text(gen_context["past_key_values"], **inp)