After merging main into ``lance-model-support`` (commit c15f51a2),
Bagel's internal API was aggressively trimmed: ``forward_cache_update_text``,
``forward_cache_update_vit``, ``generate_image``, ``generate_text``, and
``Qwen2MoTForCausalLM.forward`` no longer accept ``packed_text_indexes`` /
``packed_key_value_indexes`` / ``key_values_lens`` / ``packed_query_indexes``
— they're now derived from ``past_key_values``. Lance's pipeline was
still passing those kwargs and crashed on every task.
Changes:
- ``vllm_omni/deploy/lance.yaml`` (restored): Lance's HF ``config.json``
is descriptive metadata with no ``model_type`` field, so the loader
can't auto-detect the pipeline. This YAML is the minimal hook that
tells the loader "use the Lance pipeline" — one stage, no multi-stage
orchestration.
- ``examples/offline_inference/lance/gradio_demo.py``: pass
``deploy_config`` again so ``Omni`` picks up the YAML.
- ``vllm_omni/diffusion/models/bagel/bagel_transformer.py``: in the
batched CFG setup, concat ``branches_pid`` along ``dim=-1`` (the
sequence axis) instead of the default ``dim=0`` so Lance's 3-D mRoPE
position ids ``(3, S)`` aren't collapsed into ``(6, S)`` —
``BagelRotaryEmbedding`` enforces exactly 3 axes for mRoPE.
- ``vllm_omni/diffusion/models/bagel/pipeline_bagel.py``: mirror the
generated PIL image into ``custom_output["image"]`` so callers reading
through the orchestrator IPC boundary (which strips the bare
``output`` field) can still recover it. ``video_frames`` already used
this pattern.
- ``vllm_omni/diffusion/models/lance/lance_transformer.py``: drop
``packed_query_indexes`` / ``key_values_lens`` /
``packed_key_value_indexes`` from the ``language_model.forward`` call
inside ``forward_cache_update_vae`` — derived now.
- ``vllm_omni/diffusion/models/lance/pipeline_lance.py``:
* drop the same 3 fields from the inline ``inp`` dicts that build
``forward_cache_update_text`` arguments in ``_forward_image_edit``,
``_forward_i2v``, ``_forward_video_edit``, and
``_x2t_raw_text_prefill``.
* drop the 6 ``cfg_text_packed_query_indexes`` /
``cfg_text_key_values_lens`` / ``cfg_text_packed_key_value_indexes``
(+ ``cfg_img_*`` triple) kwargs from the four ``generate_image``
call sites.
* filter ``key_values_lens`` / ``packed_indexes`` /
``packed_key_value_indexes`` out of the ``gen_input_lat`` /
``gen_input_vit`` dicts before unpacking into ``generate_image``
and ``forward_cache_update_vit`` (the ``prepare_*_latent`` /
``prepare_vit_*`` helpers still return them for the
backwards-compatible parent-class path).
* drop ``key_values_lens`` / ``packed_key_value_indexes`` from
``_x2t_prepare_assistant_start``'s returned dict so
``Bagel.generate_text`` doesn't reject them.
Verification (local online inference, single GPU pair, CFG=4, 10 steps):
| Task | Status |
|---------------|-------------------------------------------------|
| t2i | OK (~2 s) |
| image_edit | OK (~2.5 s) |
| i2v | OK (~8 s @ 25 frames) |
| t2v | OK (~7 s @ 25 frames) |
| video_edit | OK (~60 s) |
| x2t_image | RUNS but generates degenerate token repetition |
| x2t_video | RUNS but generates degenerate token repetition |
x2t degeneration is a deeper semantic mismatch with main's new
``past_key_values``-derived cache layout: Lance's x2t prefill stitches
the chat template across multiple ``forward_cache_update_*`` calls
(system → ViT → question → ``assistant\\n``), but the new API expects a
single contiguous KV layout. Fixing requires reworking the segmented
prefill to match main's new cache contract — left as a follow-up since
the 5 generative tasks now work.
Signed-off-by: lishunyang12 <lishunyang12@163.com>