[Bugfix] enforce max_sequence_length for Qwen-Image and Wan2.2 series before encoding (#2847)
Signed-off-by: david6666666 <530634352@qq.com>
This commit is contained in:
WeiQing Chen authored and GitHub committed
1 parent
1e3bb36521
commit
3079e943de
11 files changed
+744
-27
No files matched your search
@@ -0,0 +1,260 @@
|
||||
import inspect
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import (
|
||||
QwenImagePipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import (
|
||||
QwenImageEditPipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import (
|
||||
QwenImageEditPlusPipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import (
|
||||
QwenImageLayeredPipeline,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _RejectingTextEncoder:
|
||||
dtype = torch.float32
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length")
|
||||
|
||||
|
||||
class _FakeModelInputs:
|
||||
def __init__(self, total_sequence_length: int):
|
||||
attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long)
|
||||
self.input_ids = attention_mask.clone()
|
||||
self.attention_mask = attention_mask
|
||||
self.pixel_values = None
|
||||
self.image_grid_thw = None
|
||||
|
||||
def to(self, device):
|
||||
return self
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def __init__(self, total_sequence_length: int | list[int]):
|
||||
if isinstance(total_sequence_length, list):
|
||||
self.total_sequence_lengths = list(total_sequence_length)
|
||||
else:
|
||||
self.total_sequence_lengths = [total_sequence_length]
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
if len(self.total_sequence_lengths) > 1:
|
||||
total_sequence_length = self.total_sequence_lengths.pop(0)
|
||||
else:
|
||||
total_sequence_length = self.total_sequence_lengths[0]
|
||||
return _FakeModelInputs(total_sequence_length)
|
||||
|
||||
|
||||
class _FakeProcessor(_FakeTokenizer):
|
||||
pass
|
||||
|
||||
|
||||
class _FakeScheduler:
|
||||
def __init__(self):
|
||||
self.begin_index = None
|
||||
|
||||
def set_begin_index(self, begin_index: int):
|
||||
self.begin_index = begin_index
|
||||
|
||||
|
||||
PIPELINE_CASES = [
|
||||
pytest.param(QwenImagePipeline, 34, "tokenizer", id="qwen-image"),
|
||||
pytest.param(QwenImageLayeredPipeline, 34, "tokenizer", id="qwen-image-layered"),
|
||||
pytest.param(QwenImageEditPipeline, 64, "processor", id="qwen-image-edit"),
|
||||
pytest.param(QwenImageEditPlusPipeline, 64, "processor", id="qwen-image-edit-plus"),
|
||||
]
|
||||
|
||||
|
||||
def _make_pipeline(
|
||||
pipeline_class: type,
|
||||
*,
|
||||
total_sequence_length: int,
|
||||
drop_idx: int,
|
||||
input_kind: str,
|
||||
):
|
||||
pipeline = object.__new__(pipeline_class)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.text_encoder = _RejectingTextEncoder()
|
||||
pipeline.tokenizer_max_length = 1024
|
||||
pipeline.prompt_template_encode = "{}"
|
||||
pipeline.prompt_template_encode_start_idx = drop_idx
|
||||
pipeline.tokenizer = _FakeTokenizer([total_sequence_length, 0])
|
||||
if input_kind == "processor":
|
||||
pipeline.processor = _FakeProcessor(total_sequence_length)
|
||||
return pipeline
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
|
||||
def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length(
|
||||
pipeline_class: type,
|
||||
drop_idx: int,
|
||||
input_kind: str,
|
||||
):
|
||||
pipeline = _make_pipeline(
|
||||
pipeline_class,
|
||||
total_sequence_length=1025,
|
||||
drop_idx=drop_idx,
|
||||
input_kind=input_kind,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"got 1025 tokens, but `max_sequence_length` is 1024"):
|
||||
pipeline.encode_prompt(prompt="prompt")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES)
|
||||
def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length(
|
||||
pipeline_class: type,
|
||||
drop_idx: int,
|
||||
input_kind: str,
|
||||
):
|
||||
pipeline = _make_pipeline(
|
||||
pipeline_class,
|
||||
total_sequence_length=17,
|
||||
drop_idx=drop_idx,
|
||||
input_kind=input_kind,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"):
|
||||
pipeline.encode_prompt(prompt="prompt", max_sequence_length=16)
|
||||
|
||||
|
||||
def test_prepare_encode_defaults_to_tokenizer_max_length():
|
||||
pipeline = object.__new__(QwenImagePipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.tokenizer_max_length = 1024
|
||||
pipeline.vae_scale_factor = 8
|
||||
pipeline.default_sample_size = 128
|
||||
pipeline.scheduler = _FakeScheduler()
|
||||
pipeline._extract_prompts = lambda prompts: (["prompt"], None)
|
||||
|
||||
captured = {}
|
||||
|
||||
def _fake_prepare_generation_context(**kwargs):
|
||||
captured["max_sequence_length"] = kwargs["max_sequence_length"]
|
||||
embeds = torch.ones((1, 1, 1))
|
||||
mask = torch.ones((1, 1), dtype=torch.long)
|
||||
return {
|
||||
"prompt_embeds": embeds,
|
||||
"prompt_embeds_mask": mask,
|
||||
"negative_prompt_embeds": None,
|
||||
"negative_prompt_embeds_mask": None,
|
||||
"latents": embeds,
|
||||
"timesteps": torch.tensor([1]),
|
||||
"do_true_cfg": False,
|
||||
"guidance": None,
|
||||
"img_shapes": [[(1, 1, 1)]],
|
||||
"txt_seq_lens": [1],
|
||||
"negative_txt_seq_lens": None,
|
||||
}
|
||||
|
||||
pipeline._prepare_generation_context = _fake_prepare_generation_context
|
||||
state = SimpleNamespace(
|
||||
prompts=["prompt"],
|
||||
sampling=SimpleNamespace(
|
||||
height=None,
|
||||
width=None,
|
||||
num_inference_steps=None,
|
||||
sigmas=None,
|
||||
guidance_scale_provided=False,
|
||||
num_outputs_per_prompt=0,
|
||||
generator=None,
|
||||
true_cfg_scale=None,
|
||||
max_sequence_length=None,
|
||||
),
|
||||
)
|
||||
|
||||
pipeline.prepare_encode(state)
|
||||
|
||||
assert captured["max_sequence_length"] == 1024
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pipeline_class", "drop_idx"),
|
||||
[
|
||||
pytest.param(QwenImageEditPipeline, 64, id="qwen-image-edit"),
|
||||
pytest.param(QwenImageEditPlusPipeline, 64, id="qwen-image-edit-plus"),
|
||||
],
|
||||
)
|
||||
def test_edit_pipelines_validate_text_prompt_length_before_image_token_expansion(
|
||||
pipeline_class: type,
|
||||
drop_idx: int,
|
||||
):
|
||||
pipeline = object.__new__(pipeline_class)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.text_encoder = _RejectingTextEncoder()
|
||||
pipeline.tokenizer_max_length = 1024
|
||||
pipeline.prompt_template_encode = "{}"
|
||||
pipeline.prompt_template_encode_start_idx = drop_idx
|
||||
pipeline.tokenizer = _FakeTokenizer([8, 0])
|
||||
pipeline.processor = _FakeProcessor(drop_idx + 1500)
|
||||
|
||||
with pytest.raises(AssertionError, match="text encoder should not run"):
|
||||
pipeline.encode_prompt(prompt="short prompt")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_class",
|
||||
[
|
||||
pytest.param(QwenImagePipeline, id="qwen-image"),
|
||||
pytest.param(QwenImageLayeredPipeline, id="qwen-image-layered"),
|
||||
],
|
||||
)
|
||||
def test_qwen_generation_validator_excludes_template_suffix_from_budget(pipeline_class: type):
|
||||
pipeline = object.__new__(pipeline_class)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.text_encoder = _RejectingTextEncoder()
|
||||
pipeline.tokenizer_max_length = 1024
|
||||
pipeline.prompt_template_encode = "{}"
|
||||
pipeline.prompt_template_encode_start_idx = 34
|
||||
pipeline.tokenizer = _FakeTokenizer([1029, 5])
|
||||
|
||||
with pytest.raises(AssertionError, match="text encoder should not run"):
|
||||
pipeline.encode_prompt(prompt="boundary prompt")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_class",
|
||||
[
|
||||
pytest.param(QwenImageEditPipeline, id="qwen-image-edit"),
|
||||
pytest.param(QwenImageEditPlusPipeline, id="qwen-image-edit-plus"),
|
||||
],
|
||||
)
|
||||
def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_class: type):
|
||||
pipeline = object.__new__(pipeline_class)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.text_encoder = _RejectingTextEncoder()
|
||||
pipeline.tokenizer_max_length = 1024
|
||||
pipeline.prompt_template_encode = "{}"
|
||||
pipeline.prompt_template_encode_start_idx = 64
|
||||
pipeline.tokenizer = _FakeTokenizer([30, 20])
|
||||
pipeline.processor = _FakeProcessor(1500)
|
||||
|
||||
with pytest.raises(AssertionError, match="text encoder should not run"):
|
||||
pipeline.encode_prompt(prompt="short prompt")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_class",
|
||||
[
|
||||
QwenImagePipeline,
|
||||
QwenImageLayeredPipeline,
|
||||
QwenImageEditPipeline,
|
||||
QwenImageEditPlusPipeline,
|
||||
],
|
||||
)
|
||||
def test_forward_max_sequence_length_default_is_1024(pipeline_class: type):
|
||||
assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024
|
||||
@@ -0,0 +1,150 @@
|
||||
from types import SimpleNamespace
|
||||
|
||||
import PIL.Image
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
WAN22_MAX_SEQUENCE_LENGTH,
|
||||
Wan22Pipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
|
||||
Wan22I2VPipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
|
||||
Wan22TI2VPipeline,
|
||||
)
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
|
||||
Wan22VACEPipeline,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _RejectingTextEncoder:
|
||||
dtype = torch.float32
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length")
|
||||
|
||||
|
||||
class _FakeTokenBatch:
|
||||
def __init__(self, total_sequence_length: int):
|
||||
attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long)
|
||||
self.input_ids = attention_mask.clone()
|
||||
self.attention_mask = attention_mask
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
def __init__(self, total_sequence_length: int):
|
||||
self.total_sequence_length = total_sequence_length
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return _FakeTokenBatch(self.total_sequence_length)
|
||||
|
||||
|
||||
PIPELINE_CASES = [
|
||||
pytest.param(Wan22Pipeline, id="wan22-t2v"),
|
||||
pytest.param(Wan22I2VPipeline, id="wan22-i2v"),
|
||||
pytest.param(Wan22TI2VPipeline, id="wan22-ti2v"),
|
||||
pytest.param(Wan22VACEPipeline, id="wan22-vace"),
|
||||
]
|
||||
|
||||
|
||||
def _make_pipeline(pipeline_class: type, *, total_sequence_length: int):
|
||||
pipeline = object.__new__(pipeline_class)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.text_encoder = _RejectingTextEncoder()
|
||||
pipeline.tokenizer = _FakeTokenizer(total_sequence_length)
|
||||
pipeline.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH
|
||||
return pipeline
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_class", PIPELINE_CASES)
|
||||
def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length(pipeline_class: type):
|
||||
pipeline = _make_pipeline(pipeline_class, total_sequence_length=WAN22_MAX_SEQUENCE_LENGTH + 1)
|
||||
|
||||
with pytest.raises(ValueError, match=r"got 513 tokens, but `max_sequence_length` is 512"):
|
||||
pipeline.encode_prompt(prompt="prompt")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pipeline_class", PIPELINE_CASES)
|
||||
def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length(pipeline_class: type):
|
||||
pipeline = _make_pipeline(pipeline_class, total_sequence_length=17)
|
||||
|
||||
with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"):
|
||||
pipeline.encode_prompt(prompt="prompt", max_sequence_length=16)
|
||||
|
||||
|
||||
def _sampling_params(**overrides):
|
||||
defaults = dict(
|
||||
height=None,
|
||||
width=None,
|
||||
num_frames=None,
|
||||
num_inference_steps=None,
|
||||
generator=None,
|
||||
guidance_scale_provided=False,
|
||||
guidance_scale_2=None,
|
||||
boundary_ratio=None,
|
||||
num_outputs_per_prompt=0,
|
||||
max_sequence_length=None,
|
||||
seed=None,
|
||||
extra_args={},
|
||||
prompt_embeds=None,
|
||||
negative_prompt_embeds=None,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return SimpleNamespace(**defaults)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("pipeline_class", "prompt_value", "forward_kwargs"),
|
||||
[
|
||||
pytest.param(Wan22Pipeline, "prompt", {}, id="wan22-t2v"),
|
||||
pytest.param(
|
||||
Wan22I2VPipeline,
|
||||
{"prompt": "prompt", "multi_modal_data": {"image": PIL.Image.new("RGB", (64, 64))}},
|
||||
{"image": PIL.Image.new("RGB", (64, 64))},
|
||||
id="wan22-i2v",
|
||||
),
|
||||
pytest.param(
|
||||
Wan22TI2VPipeline,
|
||||
{"prompt": "prompt", "multi_modal_data": {"image": PIL.Image.new("RGB", (64, 64))}},
|
||||
{"image": PIL.Image.new("RGB", (64, 64))},
|
||||
id="wan22-ti2v",
|
||||
),
|
||||
pytest.param(Wan22VACEPipeline, "prompt", {}, id="wan22-vace"),
|
||||
],
|
||||
)
|
||||
def test_forward_defaults_to_wan22_tokenizer_max_length(
|
||||
pipeline_class: type,
|
||||
prompt_value,
|
||||
forward_kwargs,
|
||||
):
|
||||
pipeline = object.__new__(pipeline_class)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH
|
||||
pipeline.boundary_ratio = None
|
||||
pipeline.vae_scale_factor_temporal = 4
|
||||
pipeline.vae_scale_factor_spatial = 8
|
||||
pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2))
|
||||
|
||||
captured = {}
|
||||
|
||||
def _fake_check_inputs(*args, **kwargs):
|
||||
captured["max_sequence_length"] = kwargs["max_sequence_length"]
|
||||
raise RuntimeError("stop after capture")
|
||||
|
||||
pipeline.check_inputs = _fake_check_inputs
|
||||
|
||||
req = SimpleNamespace(
|
||||
prompts=[prompt_value],
|
||||
sampling_params=_sampling_params(),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="stop after capture"):
|
||||
pipeline.forward(req, **forward_kwargs)
|
||||
|
||||
assert captured["max_sequence_length"] == WAN22_MAX_SEQUENCE_LENGTH
|
||||
Reference in new issue
Block a user