[BugFix][Qwen-Image-Edit] restore txt_seq_lens RoPE width after #3588 regression (#5586)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>

tests/diffusion/models/qwen_image/test_qwen_image_pipelines_txt_seq_lens.py

@@ -0,0 +1,150 @@
+"""Behavioral regression tests for Qwen-Image Edit ``txt_seq_lens`` wiring.
+
+RoPE text length must follow padded embed width, not valid-token count from
+``mask.sum()``. This exercises the Edit pipeline call site: when masks are
+shorter than padded embeds, ``forward()`` must pass padded width into
+``diffuse()``.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
+    QwenImageEditPipeline,
+)
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+PADDED_WIDTH = 64
+VALID_TOKENS = 10
+EMBED_DIM = 8
+
+
+def _make_padded_prompt_pair():
+    prompt_embeds = torch.zeros(1, PADDED_WIDTH, EMBED_DIM)
+    prompt_embeds_mask = torch.zeros(1, PADDED_WIDTH, dtype=torch.bool)
+    prompt_embeds_mask[:, :VALID_TOKENS] = True
+    return prompt_embeds, prompt_embeds_mask
+
+
+def _make_edit_forward_batch(*, negative_prompt: str | None) -> DiffusionRequestBatch:
+    preprocessed = torch.zeros(3, 32, 32)
+    return DiffusionRequestBatch(
+        requests=[
+            SimpleNamespace(
+                request_id="edit-txt-seq-lens",
+                prompt={
+                    "prompt": "edit prompt",
+                    "negative_prompt": negative_prompt,
+                    "additional_information": {
+                        "prompt_image": preprocessed,
+                        "preprocessed_image": preprocessed,
+                        "calculated_height": 32,
+                        "calculated_width": 32,
+                    },
+                },
+                sampling_params=SimpleNamespace(
+                    height=32,
+                    width=32,
+                    num_inference_steps=2,
+                    sigmas=None,
+                    max_sequence_length=1024,
+                    generator=None,
+                    true_cfg_scale=4.0,
+                    guidance_scale_provided=False,
+                    guidance_scale=1.0,
+                    num_outputs_per_prompt=1,
+                    latents=None,
+                    output_type="latent",
+                ),
+            )
+        ]
+    )
+
+
+def _make_stubbed_edit_pipeline():
+    pipeline = object.__new__(QwenImageEditPipeline)
+    nn.Module.__init__(pipeline)
+    pipeline.device = torch.device("cpu")
+
+    pos_embeds, pos_mask = _make_padded_prompt_pair()
+    neg_embeds, neg_mask = _make_padded_prompt_pair()
+
+    pipeline.check_inputs = lambda *args, **kwargs: None
+    pipeline.check_cfg_parallel_validity = lambda *args, **kwargs: None
+    pipeline.vae_scale_factor = 8
+
+    class _GuidanceFreeTransformer:
+        guidance_embeds = False
+        in_channels = 16
+
+    pipeline.transformer = _GuidanceFreeTransformer()
+    pipeline._attention_kwargs = None
+
+    def _fake_encode_prompt(*, prompt_name="prompt", **kwargs):
+        if prompt_name == "negative_prompt":
+            return neg_embeds, neg_mask
+        return pos_embeds, pos_mask
+
+    pipeline.encode_prompt = _fake_encode_prompt
+
+    latents = torch.zeros(1, 4, 8, 8)
+    image_latents = torch.zeros(1, 4, 8, 8)
+    pipeline.prepare_latents = lambda *args, **kwargs: (latents, image_latents)
+    pipeline.prepare_timesteps = lambda *args, **kwargs: (torch.tensor([1.0, 0.5]), 2)
+
+    captured: dict[str, object] = {}
+
+    def _fake_diffuse(
+        _prompt_embeds,
+        _prompt_embeds_mask,
+        _negative_prompt_embeds,
+        _negative_prompt_embeds_mask,
+        latents_arg,
+        _img_shapes,
+        txt_seq_lens,
+        negative_txt_seq_lens,
+        _timesteps,
+        do_true_cfg,
+        _guidance,
+        _true_cfg_scale,
+        **kwargs,
+    ):
+        captured["txt_seq_lens"] = txt_seq_lens
+        captured["negative_txt_seq_lens"] = negative_txt_seq_lens
+        captured["do_true_cfg"] = do_true_cfg
+        return latents_arg
+
+    pipeline.diffuse = _fake_diffuse
+    return pipeline, captured, pos_mask, neg_mask
+
+
+def test_qwen_image_edit_forward_passes_padded_txt_seq_lens_to_diffuse():
+    pipeline, captured, pos_mask, neg_mask = _make_stubbed_edit_pipeline()
+    batch = _make_edit_forward_batch(negative_prompt="bad prompt")
+
+    pipeline.forward(batch)
+
+    assert pos_mask.sum(dim=1).tolist() == [VALID_TOKENS]
+    assert neg_mask.sum(dim=1).tolist() == [VALID_TOKENS]
+    assert captured["txt_seq_lens"] == [PADDED_WIDTH]
+    assert captured["negative_txt_seq_lens"] == [PADDED_WIDTH]
+    assert captured["do_true_cfg"] is True
+
+
+def test_qwen_image_edit_forward_skips_negative_txt_seq_lens_without_cfg():
+    pipeline, captured, pos_mask, _neg_mask = _make_stubbed_edit_pipeline()
+    batch = _make_edit_forward_batch(negative_prompt=None)
+
+    pipeline.forward(batch)
+
+    assert pos_mask.sum(dim=1).tolist() == [VALID_TOKENS]
+    assert captured["txt_seq_lens"] == [PADDED_WIDTH]
+    assert captured["negative_txt_seq_lens"] is None
+    assert captured["do_true_cfg"] is False

tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py

@@ -1,44 +0,0 @@
-"""Regression tests for #4443: padded prompt width vs RoPE table length.
-
-When ``max(txt_seq_lens) < prompt_embeds.shape[1]``, ``QwenEmbedRope`` builds a
-text frequency table that is too short for the padded encoder width.
-"""
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope
-from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-# Representative request-batch padding: valid tokens 10/20, shared width 32.
-PADDED_WIDTH = 32
-VALID_LENS = [10, 20]
-
-
-def test_pre_pr_txt_seq_lens_is_shorter_than_padded_width():
-    mask = torch.zeros(len(VALID_LENS), PADDED_WIDTH, dtype=torch.bool)
-    mask[0, : VALID_LENS[0]] = True
-    mask[1, : VALID_LENS[1]] = True
-
-    pre_pr_txt_seq_lens = mask.sum(dim=1).tolist()
-    assert pre_pr_txt_seq_lens == VALID_LENS
-    assert max(pre_pr_txt_seq_lens) < PADDED_WIDTH
-
-
-def test_post_pr_txt_seq_lens_matches_padded_width():
-    prompt_embeds = torch.zeros(len(VALID_LENS), PADDED_WIDTH, 8)
-    assert txt_seq_lens_from_embeds(prompt_embeds) == [PADDED_WIDTH, PADDED_WIDTH]
-
-
-def test_rope_table_length_matches_padded_width_after_fix():
-    rope = QwenEmbedRope(theta=10000, axes_dim=[16, 56, 56], scale_rope=True)
-    img_shapes = [[(1, 16, 16)]]
-
-    _, freqs_pre_pr = rope(img_shapes, VALID_LENS, device="cpu")
-    _, freqs_post_pr = rope(img_shapes, [PADDED_WIDTH, PADDED_WIDTH], device="cpu")
-
-    assert freqs_pre_pr.shape[0] == max(VALID_LENS)
-    assert freqs_pre_pr.shape[0] < PADDED_WIDTH
-    assert freqs_post_pr.shape[0] == PADDED_WIDTH

tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py

@@ -1,6 +1,7 @@
 import pytest
 import torch
 
+from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope
 from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -15,6 +16,17 @@ def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count():
     assert txt_seq_lens_from_embeds(prompt_embeds) == [64, 64]
 
 
+def test_txt_seq_lens_from_embeds_builds_rope_table_for_padded_width():
+    padded_width = 32
+    prompt_embeds = torch.zeros(2, padded_width, 8)
+    txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
+    rope = QwenEmbedRope(theta=10000, axes_dim=[16, 56, 56], scale_rope=True)
+
+    _, txt_freqs = rope([[(1, 16, 16)]], txt_seq_lens, device="cpu")
+
+    assert txt_freqs.shape[0] == padded_width
+
+
 def test_txt_seq_lens_from_embeds_supports_2d_embeds():
     prompt_embeds = torch.zeros(48, 16)
 

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_
 from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import (
     QwenImageTransformer2DModel,
 )
+from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.utils.prompt_utils import (
@@ -815,10 +816,8 @@ class QwenImageEditPipeline(
         if self.attention_kwargs is None:
             self._attention_kwargs = {}
 
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        negative_txt_seq_lens = (
-            negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None
-        )
+        txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds)
+        negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds)
 
         latents = self.diffuse(
             prompt_embeds,