diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_pipelines_txt_seq_lens.py b/tests/diffusion/models/qwen_image/test_qwen_image_pipelines_txt_seq_lens.py new file mode 100644 index 000000000..25a8b0183 --- /dev/null +++ b/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 diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py deleted file mode 100644 index 2218db8d6..000000000 --- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py +++ /dev/null @@ -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 diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py index 336ae2ed4..52f1b6aa3 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py +++ b/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) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 3f0d16002..34e5733fd 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/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,