3dc35694b3
Signed-off-by: NumberWan <wantszkin2003@gmail.com> Signed-off-by: SamitHuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: SamitHuang <285365963@qq.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""CPU lock: adjacent condition images stay separate token blocks.
|
|
|
|
Two condition images with no text between them are one run of True in
|
|
``image_pad_mask``. Block ids come from ``img_shapes`` token counts, so they
|
|
must not collapse into one bidirectional attention span.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.qwen_image_21.qwen_image_21_transformer import (
|
|
QwenImage21Transformer2DModel,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def test_adjacent_condition_images_stay_separate_blocks() -> None:
|
|
# text(2) | cond0(2) | cond1(3) | text(1) | target(2)
|
|
# cond0 and cond1 are one contiguous True run.
|
|
image_pad_mask = torch.tensor([False, False, True, True, True, True, True, False, True, True])
|
|
img_shapes = [(1, 1, 2), (1, 1, 3), (1, 1, 2)]
|
|
|
|
image_ids, target_token_mask = QwenImage21Transformer2DModel.build_token_metadata(image_pad_mask, img_shapes)
|
|
|
|
assert image_ids.tolist() == [-1, -1, 0, 0, 1, 1, 1, -1, 2, 2]
|
|
assert target_token_mask.tolist() == [
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
False,
|
|
True,
|
|
True,
|
|
]
|
|
assert QwenImage21Transformer2DModel._build_full_attn_spans(image_ids) == [
|
|
(2, 4),
|
|
(4, 7),
|
|
(8, 10),
|
|
]
|