PyTorch SDPA's is_causal=True with Q != K uses tril(diagonal=0) which is
top-left aligned. For a single decode step (Q=1) against a long ViT+text
cache (K=574+), only q_0 -> k_0 is unmasked — producing degenerate output
("TheTheThe..." token loops).
Bypass DiffusionAttention (whose _maybe_reshape_attn_mask only handles
2-D (B, K) masks) and call torch.nn.functional.scaled_dot_product_attention
directly with an explicit 4-D (1, 1, Q, K) bottom-right-aligned causal mask
when cache_len > 0. Also cast q/k to fp32 before RMSNorm for precision.
Verified all 7 tasks on Lance:
- x2t_image: "No" (valid VQA answer)
- x2t_video: "The purple sphere ... right direc..." (valid VQA answer)
- t2i / i2v / t2v / image_edit / video_edit all produce expected media.
Signed-off-by: lishunyang12 <lishunyang12@163.com>
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>
Resolves conflicts in:
- vllm_omni/diffusion/registry.py: keep both LancePipeline (HEAD)
and MingImagePipeline (main) entries.
- vllm_omni/diffusion/models/bagel/bagel_transformer.py: take main's
simpler cfg_batched payload; HEAD referenced undefined locals.
Also strips vllm_omni/deploy/lance.yaml -- Lance is single-stage
diffusion, so its engine knobs (pipeline=lance, enforce_eager,
trust_remote_code, max_num_seqs=1, ...) flow directly to Omni(...)
/ vllm-omni serve instead of through a deploy YAML.
create_default_diffusion builds the stage config from those kwargs.
Touched: examples/offline_inference/lance/{end2end,gradio_demo}.py,
tests/e2e/online_serving/test_lance.py, tests/helpers/stage_config.py.
Signed-off-by: lishunyang12 <lishunyang12@163.com>
- Add ``examples/offline_inference/lance/gradio_demo.py`` — a unified 7-task UI
(t2i / t2v / i2v / image_edit / video_edit / x2t_image / x2t_video) that
loads ``Lance_3B`` and ``Lance_3B_Video`` on two GPUs so image and video
tasks run concurrently. Surfaces official Lance examples per task; aspect
ratio is auto-snapped to one of upstream's 6 buckets from the input
image/video.
- Fix x2t_image / x2t_video output format: route x2t_image through a new
``_forward_x2t_image`` handler (was falling through to the img2img path
and emitting garbage), build the chat prefill in three segments
(prefix → ViT → suffix) so the ``<|im_end|>`` markers line up with
upstream's template, and start generation from a newline token instead of
``<|im_start|>`` so the assistant turn doesn't close itself with an
immediate EOS. Add a ``system_prompt`` override to ``render_lance_prompt``
so per-example QA instructions surface ("Look at the image carefully and
answer the question.") instead of the caption-style default.
- ``_forward_image_edit``: also stash the PIL image into ``custom_output``
so callers reading ``outputs[0].output`` *or* ``custom_output["image"]``
both work — the orchestrator IPC strips the bare ``output`` field across
process boundaries.
- ``pipeline_registry.py``, ``model_executor/models/lance/*``, registry
cleanup: remove the dead Lance-specific autoregressive shim
(``lance.py``) and its stage-input processor; the diffusion pipeline owns
the full flow now.
- Aspect-ratio buckets aligned with upstream's ``BucketResize`` (6 ratios:
21:9 / 16:9 / 4:3 / 1:1 / 3:4 / 9:16, sized for ``resolution_vae`` ∈
{480, 640, 768}). 3:2 / 2:3 dropped — upstream snaps those to 4:3 / 3:4.
- Gradio defaults match upstream's official scripts (cfg_text_scale=4.0,
validation_num_timesteps=30, validation_timestep_shift=3.5, seed=42); only
the demo *duration* is intentionally shorter than upstream's 121-frame
default to keep the UI snappy (user can extend to 10s via slider for
byte-identical reproduction).
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>
lishunyang12