From 2116e88638953990bbdea41d032789481cdbcac2 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Mon, 20 Apr 2026 11:14:12 +0800 Subject: [PATCH] [cherry-pick][release/v0.18.0.post1] cherry-pick #2847 #2780 #2840 #2876 #2877 (#2878) Signed-off-by: david6666666 <530634352@qq.com> Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- .../qwen_image/test_qwen_image_edit_plus.py | 38 +++ .../test_qwen_image_max_sequence_length.py | 260 ++++++++++++++++++ .../wan2_2/test_wan22_max_sequence_length.py | 190 +++++++++++++ .../test_async_omni_engine_stage_init.py | 1 + .../openai_api/test_image_server.py | 67 ++++- .../openai_api/test_video_api_utils.py | 9 +- .../openai_api/test_video_server.py | 5 +- .../test_async_omni_diffusion_config.py | 14 + tests/test_diffusion_config_propagation.py | 10 + vllm_omni/diffusion/data.py | 11 +- vllm_omni/diffusion/diffusion_engine.py | 5 +- vllm_omni/diffusion/model_metadata.py | 31 +++ .../models/qwen_image/pipeline_qwen_image.py | 44 ++- .../qwen_image/pipeline_qwen_image_edit.py | 49 +++- .../pipeline_qwen_image_edit_plus.py | 60 +++- .../qwen_image/pipeline_qwen_image_layered.py | 39 ++- .../models/wan2_2/pipeline_wan2_2.py | 69 ++++- .../models/wan2_2/pipeline_wan2_2_i2v.py | 69 ++++- .../models/wan2_2/pipeline_wan2_2_ti2v.py | 69 ++++- .../models/wan2_2/wan2_2_transformer.py | 7 +- .../postprocess/rife_interpolator.py | 9 +- vllm_omni/diffusion/utils/prompt_utils.py | 38 +++ vllm_omni/engine/async_omni_engine.py | 11 + vllm_omni/engine/stage_init_utils.py | 3 +- vllm_omni/entrypoints/openai/api_server.py | 47 +++- 25 files changed, 1069 insertions(+), 86 deletions(-) create mode 100644 tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py create mode 100644 tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py create mode 100644 tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py create mode 100644 vllm_omni/diffusion/model_metadata.py create mode 100644 vllm_omni/diffusion/utils/prompt_utils.py diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py b/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py new file mode 100644 index 000000000..873b52bf7 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +from PIL import Image + +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import ( + get_qwen_image_edit_plus_pre_process_func, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path): + vae_dir = tmp_path / "vae" + vae_dir.mkdir() + # Keep the mock config intentionally minimal: this test only needs the + # fields touched during pre-process initialization. + (vae_dir / "config.json").write_text(json.dumps({"z_dim": 16})) + + pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path))) + image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)) + request = SimpleNamespace( + prompts=[ + { + "prompt": "combine", + "multi_modal_data": {"image": [image, image, image, image, image]}, + } + ], + sampling_params=SimpleNamespace(height=None, width=None), + ) + + with pytest.raises(ValueError, match=r"At most 4 images are supported by this model"): + pre_process(request) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py new file mode 100644 index 000000000..f5676a005 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py @@ -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 diff --git a/tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py b/tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py new file mode 100644 index 000000000..17dec2ce0 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py @@ -0,0 +1,190 @@ +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, +) + +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, padded_length: int | None = None): + padded_length = padded_length or total_sequence_length + attention_mask = torch.zeros((1, padded_length), dtype=torch.long) + attention_mask[:, :total_sequence_length] = 1 + self.input_ids = attention_mask.clone() + self.attention_mask = attention_mask + + +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 _FakeTokenBatch(total_sequence_length, kwargs.get("max_length")) + + +class _RecordingTextEncoder: + dtype = torch.float32 + + def __init__(self): + self.input_lengths: list[int] = [] + + def __call__(self, input_ids, attention_mask): + self.input_lengths.append(input_ids.shape[1]) + return SimpleNamespace(last_hidden_state=torch.zeros((1, input_ids.shape[1], 4), dtype=torch.float32)) + + +PIPELINE_CASES = [ + pytest.param(Wan22Pipeline, id="wan22-t2v"), + pytest.param(Wan22I2VPipeline, id="wan22-i2v"), + pytest.param(Wan22TI2VPipeline, id="wan22-ti2v"), +] + + +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=rf"got {WAN22_MAX_SEQUENCE_LENGTH + 1} tokens, but `max_sequence_length` is {WAN22_MAX_SEQUENCE_LENGTH}", + ): + 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) + + +@pytest.mark.parametrize("pipeline_class", PIPELINE_CASES) +def test_encode_prompt_uses_actual_prompt_length_for_text_encoder_padding(pipeline_class: type): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RecordingTextEncoder() + # prompt validation, prompt encode, negative validation, negative encode + pipeline.tokenizer = _FakeTokenizer([17, 17, 9, 9]) + pipeline.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH + + prompt_embeds, negative_prompt_embeds = pipeline.encode_prompt( + prompt="prompt", + negative_prompt="neg", + do_classifier_free_guidance=True, + ) + + assert pipeline.text_encoder.input_lengths == [17, 17] + assert prompt_embeds.shape[1] == 17 + assert negative_prompt_embeds is not None + assert negative_prompt_embeds.shape[1] == 17 + + +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", + ), + ], +) +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 diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index f95aa53af..e5ba2f31d 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -70,6 +70,7 @@ def test_initialize_stages_restores_device_visibility_after_diffusion_init(monke else: os.environ[env_var] = old_env + def test_initialize_stages_uses_inline_diffusion_client_for_single_stage(monkeypatch): """Single-stage diffusion init should request the inline client path.""" import vllm_omni.engine.async_omni_engine as engine_mod diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 8695311fb..e49571a3f 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -171,7 +171,7 @@ def test_client(mock_async_diffusion): app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")] app.state.diffusion_model_name = "Qwen/Qwen-Image" # For models endpoint app.state.args = Namespace( - default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1024 * 1792, ) @@ -194,7 +194,7 @@ def async_omni_test_client(): SimpleNamespace(stage_type="diffusion"), ] app.state.args = Namespace( - default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1048576, # 1024*1024 to support resolution tests ) return TestClient(app) @@ -216,7 +216,7 @@ def async_omni_rgba_test_client(): SimpleNamespace(stage_type="diffusion"), ] app.state.args = Namespace( - default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1048576, ) return TestClient(app) @@ -238,7 +238,7 @@ def async_omni_stage_configs_only_client(): # Intentionally do not populate app.state.stage_configs. Refactored # AsyncOmni exposes stage_configs on the engine instance. app.state.args = Namespace( - default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1024 * 1792, ) return TestClient(app) @@ -776,6 +776,63 @@ def test_image_edit_rejects_multiple_images_when_model_does_not_support_them(asy assert engine.captured_prompt is None +def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511(async_omni_test_client): + engine = async_omni_test_client.app.state.engine_client + engine.get_diffusion_od_config = lambda: SimpleNamespace( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=4, + ) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[ + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ], + data={"prompt": "hello world."}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model." + assert engine.captured_prompt is None + + +def test_image_edit_rejects_too_many_images_for_qwen_image_edit_2511_before_loading( + async_omni_test_client, monkeypatch: pytest.MonkeyPatch +): + import vllm_omni.entrypoints.openai.api_server as api_server_module + + engine = async_omni_test_client.app.state.engine_client + engine.get_diffusion_od_config = lambda: SimpleNamespace( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=4, + ) + + def _fail_load(*args, **kwargs): + raise AssertionError("_load_input_images should not run for over-limit requests") + + monkeypatch.setattr(api_server_module, "_load_input_images", _fail_load) + + response = async_omni_test_client.post( + "/v1/images/edits", + files=[ + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ("image", make_test_image_bytes((16, 16))), + ], + data={"prompt": "hello world."}, + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "Received 5 input images. At most 4 images are supported by this model." + assert engine.captured_prompt is None + + def test_image_edit_parameter_pass(async_omni_test_client): img_bytes_1 = make_test_image_bytes((16, 16)) @@ -954,6 +1011,7 @@ def test_image_edit_parameter_default(async_omni_test_client): assert captured_sampling_params.num_outputs_per_prompt == 1 assert captured_sampling_params.num_inference_steps == 4 assert captured_sampling_params.guidance_scale == 7.5 + assert captured_sampling_params.generator_device == "cpu" # Test that a size exceeding max_generated_image_size returns 400 response = async_omni_test_client.post( @@ -987,6 +1045,7 @@ def test_image_edit_parameter_default_single_stage(test_client): assert captured_sampling_params.num_outputs_per_prompt == 1 assert captured_sampling_params.num_inference_steps == 4 assert captured_sampling_params.guidance_scale == 7.5 + assert captured_sampling_params.generator_device == "cpu" # Size exceeding max_generated_image_size (1024*1792) returns 400 response = test_client.post( diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 3868b6573..7a0494226 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -87,7 +87,9 @@ def test_encode_video_bytes_without_audio_uses_diffusers_export(monkeypatch): _install_fake_export_to_video(monkeypatch, export_calls) monkeypatch.setattr( "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes", - lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("no-audio path should not use mux_video_audio_bytes")), + lambda *args, **kwargs: (_ for _ in ()).throw( + AssertionError("no-audio path should not use mux_video_audio_bytes") + ), ) video = np.linspace(0.0, 1.0, num=4 * 2 * 2 * 3, dtype=np.float32).reshape(4, 2, 2, 3) @@ -123,7 +125,7 @@ def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch): assert torch.isfinite(output_video).all() -def test_frame_interpolator_prefers_input_tensor_device(monkeypatch): +def test_frame_interpolator_uses_platform_device_when_tensor_is_cpu(monkeypatch): chosen_devices = [] model = rife_interpolator.Model().eval() @@ -134,10 +136,11 @@ def test_frame_interpolator_prefers_input_tensor_device(monkeypatch): interpolator = rife_interpolator.FrameInterpolator() monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded) monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet) + monkeypatch.setattr(rife_interpolator, "_select_torch_device", lambda: torch.device("cuda")) video = torch.zeros(1, 3, 2, 32, 32) output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) - assert chosen_devices == [video.device] + assert chosen_devices == [torch.device("cuda")] assert multiplier == 2 assert output_video.shape == (1, 3, 3, 32, 32) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 9ebcb6897..258ea76f4 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -14,8 +14,7 @@ import time from types import SimpleNamespace import pytest -from fastapi import FastAPI -from fastapi import HTTPException +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture @@ -981,5 +980,3 @@ async def test_run_generation_maps_omni_request_error_to_http_exception(): "error_type": "OutOfMemoryError", "detail": {"retryable": False}, } - - diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index ca5624f2d..1e38c4438 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -69,6 +69,20 @@ def test_default_stage_config_propagates_ulysses_mode(): assert parallel_config.ulysses_mode == "advanced_uaa" +def test_default_stage_config_includes_default_sampling_params(): + """Ensure default sampling params survive the default diffusion-stage builder.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}', + } + )[0] + + assert stage_cfg["default_sampling_params"] == { + "generator_device": "cpu", + "guidance_scale": 7.5, + } + + def test_serve_cli_accepts_ulysses_mode(): """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config.""" parser = FlexibleArgumentParser() diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py index 58eb6097c..18c6135c0 100644 --- a/tests/test_diffusion_config_propagation.py +++ b/tests/test_diffusion_config_propagation.py @@ -14,6 +14,7 @@ from vllm_omni.diffusion.data import ( DiffusionParallelConfig, OmniDiffusionConfig, ) +from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig: @@ -106,3 +107,12 @@ class TestCreateDefaultDiffusion: ea = stages[0]["engine_args"] assert ea["enforce_eager"] is True assert ea["lora_path"] == "/tmp/lora" + + +def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): + od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline") + + od_config.update_multimodal_support() + + assert od_config.supports_multimodal_inputs is True + assert od_config.max_multimodal_image_inputs == QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 390231bf9..5a1121ca9 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) +from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available from vllm_omni.quantization import build_quant_config @@ -468,8 +469,10 @@ class OmniDiffusionConfig: # Scheduler flow_shift for Wan2.2 (12.0 for 480p, 5.0 for 720p) flow_shift: float | None = None - # support multi images input + # Support multi-image inputs and expose any model-specific request limit + # through a generic config field so serving code stays model-agnostic. supports_multimodal_inputs: bool = False + max_multimodal_image_inputs: int | None = None log_level: str = "info" @@ -615,7 +618,11 @@ class OmniDiffusionConfig: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") def update_multimodal_support(self) -> None: - self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"} + # Resolve serving-visible multimodal behavior from shared metadata + # instead of importing concrete pipeline modules into the config layer. + metadata = get_diffusion_model_metadata(self.model_class_name) + self.supports_multimodal_inputs = metadata.supports_multimodal_inputs + self.max_multimodal_image_inputs = metadata.max_multimodal_image_inputs def enrich_config(self) -> None: """Load model metadata from HuggingFace and populate config fields. diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 17dfe6d3a..f78ec1f05 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -308,7 +308,7 @@ class DiffusionEngine: def make_engine( config: OmniDiffusionConfig, scheduler: SchedulerInterface | None = None, - ) -> "DiffusionEngine": + ) -> DiffusionEngine: """Factory method to create a DiffusionEngine instance. Args: @@ -357,9 +357,10 @@ class DiffusionEngine: finished_req_ids = self.scheduler.update_from_output(sched_output, output) if output.error: + status_code = 400 if "`max_sequence_length`" in output.error else 500 raise OmniRequestError( output.error, - status_code=500, + status_code=status_code, error_type="DiffusionExecutionError", ) if target_sched_req_id in finished_req_ids: diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py new file mode 100644 index 000000000..ec133e738 --- /dev/null +++ b/vllm_omni/diffusion/model_metadata.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class DiffusionModelMetadata: + # Keep serving-facing capability metadata in a lightweight shared module so + # config/model plumbing can read it without importing concrete pipelines. + supports_multimodal_inputs: bool = False + max_multimodal_image_inputs: int | None = None + + +QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES = 4 + + +_DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { + "QwenImageEditPlusPipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, + ), +} + + +def get_diffusion_model_metadata(model_class_name: str | None) -> DiffusionModelMetadata: + # Unknown models fall back to "no special multimodal capabilities" so new + # pipelines do not accidentally inherit limits meant for other models. + if model_class_name is None: + return DiffusionModelMetadata() + return _DIFFUSION_MODEL_METADATA.get(model_class_name, DiffusionModelMetadata()) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 8a8697ebb..105445f31 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -34,6 +34,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) 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 ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -363,8 +366,10 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP "that was used to generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -378,6 +383,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP self, prompt: str | list[str] = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): dtype = dtype or self.text_encoder.dtype @@ -388,12 +395,27 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP txt = [template.format(e) for e in prompt] txt_tokens = self.tokenizer( txt, - max_length=self.tokenizer_max_length + drop_idx, padding=True, - truncation=True, + truncation=False, return_tensors="pt", ).to(self.device) - # print(f"attention mask: {txt_tokens.attention_mask}") + # Validate only the user prompt contribution. The Qwen template also + # adds a fixed suffix after the user text, so subtracting only + # prompt_template_encode_start_idx would overcount near-limit prompts. + template_tokens = self.tokenizer( + [template.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) encoder_hidden_states = self.text_encoder( input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, @@ -422,6 +444,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -439,7 +462,11 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) prompt_embeds = prompt_embeds[:, :max_sequence_length] prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] @@ -632,6 +659,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) else: negative_prompt_embeds = None @@ -703,7 +731,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1, generator=sampling.generator, true_cfg_scale=sampling.true_cfg_scale or 4.0, - max_sequence_length=sampling.max_sequence_length or 512, + max_sequence_length=sampling.max_sequence_length or self.tokenizer_max_length, attention_kwargs=kwargs.get("attention_kwargs"), ) @@ -934,7 +962,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, + max_sequence_length: int = 1024, ) -> DiffusionOutput: extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) prompt = extracted_prompt or prompt 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 59c635266..f73cefade 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 @@ -37,6 +37,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) 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 ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -323,8 +326,10 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi "that was used to generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -384,6 +389,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi prompt: str | list[str] = None, image: PIL.Image.Image | torch.Tensor | None = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): """Get prompt embeddings with image support for editing.""" dtype = dtype or self.text_encoder.dtype @@ -393,6 +400,33 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx txt = [template.format(e) for e in prompt] + txt_tokens = self.tokenizer( + txt, + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # The edit template contains fixed multimodal scaffolding around the + # instruction. Validate against the empty-template baseline so image + # placeholder text does not consume the user's text budget. + template_tokens = self.tokenizer( + [template.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # Qwen-Image-Edit expands image placeholders into many vision tokens + # inside the processor. `max_sequence_length` is meant to constrain the + # prompt text length, so validate on the text template before image + # token expansion. + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) # Use processor to handle both text and image inputs model_inputs = self.processor( @@ -434,6 +468,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -453,7 +488,12 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + image, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) @@ -624,7 +664,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, + max_sequence_length: int = 1024, ) -> DiffusionOutput: """Forward pass for image editing.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") @@ -739,6 +779,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) num_channels_latents = self.transformer.in_channels // 4 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 59b997543..beca8c8c7 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -25,6 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( QwenImageCFGParallelMixin, @@ -40,6 +41,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) 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 ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -53,6 +57,12 @@ logger = logging.getLogger(__name__) CONDITION_IMAGE_SIZE = 384 * 384 VAE_IMAGE_SIZE = 1024 * 1024 +# Keep this in sync with the practical conditioning-token budget for +# Qwen-Image-Edit-2511. Empirically, 4 images stays within the supported range +# while 5 images overflows the prompt/conditioning path and fails downstream. +# Re-export the shared metadata value locally so this pipeline keeps a nearby, +# descriptive constant for validation and tests without becoming the source of truth. +MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES = QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES def get_qwen_image_edit_plus_pre_process_func( @@ -90,6 +100,11 @@ def get_qwen_image_edit_plus_pre_process_func( if not isinstance(raw_image, list): raw_image = [raw_image] + if len(raw_image) > MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES: + raise ValueError( + f"Received {len(raw_image)} input images. " + f"At most {MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES} images are supported by this model." + ) image = [ PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im) for im in raw_image @@ -283,8 +298,10 @@ class QwenImageEditPlusPipeline( "that was used to generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -299,6 +316,8 @@ class QwenImageEditPlusPipeline( prompt: str | list[str], image: list[torch.Tensor] | torch.Tensor | None = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): """Get prompt embeddings with support for multiple images.""" dtype = dtype or self.text_encoder.dtype @@ -319,6 +338,32 @@ class QwenImageEditPlusPipeline( template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx txt = [template.format(base_img_prompt + e) for e in prompt] + txt_tokens = self.tokenizer( + txt, + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # Multi-image edit prepends "Picture N" placeholders before the user + # instruction. Subtract the placeholder-aware baseline so attached + # images do not reduce the remaining prompt budget. + template_tokens = self.tokenizer( + [template.format(base_img_prompt)], + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # The processor expands image placeholders into many vision tokens. + # `max_sequence_length` should guard the prompt text length before that + # multimodal expansion happens. + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) # Use processor to handle both text and image inputs model_inputs = self.processor( @@ -360,6 +405,7 @@ class QwenImageEditPlusPipeline( prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -379,7 +425,12 @@ class QwenImageEditPlusPipeline( batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + image, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) @@ -557,7 +608,7 @@ class QwenImageEditPlusPipeline( output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, + max_sequence_length: int = 1024, ) -> DiffusionOutput: """Forward pass for image editing with support for multiple images.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") @@ -692,6 +743,7 @@ class QwenImageEditPlusPipeline( prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) num_channels_latents = self.transformer.in_channels // 4 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index df59b4639..f670d8044 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -36,6 +36,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) 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 ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -340,8 +343,10 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as "generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -356,6 +361,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as prompt: str | list[str] | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): device = device or self.device dtype = dtype or self.text_encoder.dtype @@ -368,8 +375,26 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as txt_tokens = self.tokenizer( txt, padding=True, + truncation=False, return_tensors="pt", ).to(device) + # The layered template also appends fixed non-user tokens after the + # editable text, so use the empty-template tokenized baseline instead of + # counting everything after prompt_template_encode_start_idx. + template_tokens = self.tokenizer( + [template.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ).to(device) + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) encoder_hidden_states = self.text_encoder( input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, @@ -399,6 +424,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -419,7 +445,11 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) prompt_embeds = prompt_embeds[:, :max_sequence_length] prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] @@ -603,7 +633,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as negative_prompt_embeds_mask: torch.Tensor | None = None, output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int = 512, + max_sequence_length: int = 1024, resolution: int = 640, cfg_normalize: bool = False, use_en_prompt: bool = False, @@ -736,6 +766,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as device=self.device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) # 4. Prepare latent variables diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 4b8d01dfa..016c5177a 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -28,11 +28,15 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3 from vllm_omni.diffusion.postprocess import interpolate_video_tensor 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 ( + validate_prompt_sequence_lengths, +) from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) DEBUG_PERF = False +WAN22_MAX_SEQUENCE_LENGTH = 2048 def retrieve_latents( @@ -248,6 +252,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe pass self.boundary_ratio = od_config.boundary_ratio + self.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH # Determine which transformers to load based on boundary_ratio # boundary_ratio=1.0: only load transformer_2 (low-noise stage only) @@ -423,6 +428,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt_embeds=negative_prompt_embeds, guidance_scale_2=guidance_high if boundary_ratio is not None else None, boundary_ratio=boundary_ratio, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) if num_frames % self.vae_scale_factor_temporal != 1: @@ -456,7 +462,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) @@ -743,7 +749,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - max_sequence_length: int = 512, + max_sequence_length: int = WAN22_MAX_SEQUENCE_LENGTH, device: torch.device | None = None, dtype: torch.dtype | None = None, ): @@ -753,11 +759,51 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) + text_inputs_untruncated = self.tokenizer( + prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + error_context="for Wan2.2 text encoding", + ) + prompt_encode_length = max(int(text_inputs_untruncated.attention_mask.sum(dim=1).max().item()), 1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_clean = [self._prompt_clean(p) for p in negative_prompt] + neg_text_inputs_untruncated = self.tokenizer( + negative_prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + neg_text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name="negative_prompt", + error_context="for Wan2.2 text encoding", + ) + prompt_encode_length = max( + prompt_encode_length, + int(neg_text_inputs_untruncated.attention_mask.sum(dim=1).max().item()), + ) text_inputs = self.tokenizer( prompt_clean, padding="max_length", - max_length=max_sequence_length, + max_length=prompt_encode_length, truncation=True, add_special_tokens=True, return_attention_mask=True, @@ -770,21 +816,18 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + [torch.cat([u, u.new_zeros(prompt_encode_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - negative_prompt_embeds = None if do_classifier_free_guidance: - negative_prompt = negative_prompt or "" - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], + negative_prompt_clean, padding="max_length", - max_length=max_sequence_length, + max_length=prompt_encode_length, truncation=True, add_special_tokens=True, return_attention_mask=True, @@ -797,7 +840,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] negative_prompt_embeds = torch.stack( [ - torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + torch.cat([u, u.new_zeros(prompt_encode_length - u.size(0), u.size(1))]) for u in negative_prompt_embeds ], dim=0, @@ -854,6 +897,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt_embeds=None, guidance_scale_2=None, boundary_ratio=None, + max_sequence_length=None, ): if height % 16 != 0 or width % 16 != 0: raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") @@ -880,5 +924,10 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe ): raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) + if boundary_ratio is None and guidance_scale_2 is not None: raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.") diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 112dfdfdd..6166df78e 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -26,6 +26,7 @@ from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + WAN22_MAX_SEQUENCE_LENGTH, create_transformer_from_config, load_transformer_config, retrieve_latents, @@ -33,6 +34,9 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( from vllm_omni.diffusion.postprocess import interpolate_video_tensor 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 ( + validate_prompt_sequence_lengths, +) from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -211,6 +215,7 @@ class Wan22I2VPipeline( # Text encoder self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) @@ -394,6 +399,7 @@ class Wan22I2VPipeline( image_embeds=image_embeds, guidance_scale_2=guidance_high if boundary_ratio is not None else None, boundary_ratio=boundary_ratio, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) # Adjust num_frames to be compatible with VAE temporal scaling @@ -422,7 +428,7 @@ class Wan22I2VPipeline( negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) @@ -655,7 +661,7 @@ class Wan22I2VPipeline( negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - max_sequence_length: int = 512, + max_sequence_length: int = WAN22_MAX_SEQUENCE_LENGTH, device: torch.device | None = None, dtype: torch.dtype | None = None, ): @@ -666,11 +672,51 @@ class Wan22I2VPipeline( prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) + text_inputs_untruncated = self.tokenizer( + prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + error_context="for Wan2.2 text encoding", + ) + prompt_encode_length = max(int(text_inputs_untruncated.attention_mask.sum(dim=1).max().item()), 1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_clean = [self._prompt_clean(p) for p in negative_prompt] + neg_text_inputs_untruncated = self.tokenizer( + negative_prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + neg_text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name="negative_prompt", + error_context="for Wan2.2 text encoding", + ) + prompt_encode_length = max( + prompt_encode_length, + int(neg_text_inputs_untruncated.attention_mask.sum(dim=1).max().item()), + ) text_inputs = self.tokenizer( prompt_clean, padding="max_length", - max_length=max_sequence_length, + max_length=prompt_encode_length, truncation=True, add_special_tokens=True, return_attention_mask=True, @@ -683,21 +729,18 @@ class Wan22I2VPipeline( prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + [torch.cat([u, u.new_zeros(prompt_encode_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - negative_prompt_embeds = None if do_classifier_free_guidance: - negative_prompt = negative_prompt or "" - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], + negative_prompt_clean, padding="max_length", - max_length=max_sequence_length, + max_length=prompt_encode_length, truncation=True, add_special_tokens=True, return_attention_mask=True, @@ -710,7 +753,7 @@ class Wan22I2VPipeline( negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] negative_prompt_embeds = torch.stack( [ - torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + torch.cat([u, u.new_zeros(prompt_encode_length - u.size(0), u.size(1))]) for u in negative_prompt_embeds ], dim=0, @@ -837,6 +880,7 @@ class Wan22I2VPipeline( image_embeds=None, guidance_scale_2=None, boundary_ratio=None, + max_sequence_length=None, ): if image is None and image_embeds is None: raise ValueError("Provide either `image` or `image_embeds`. Cannot leave both undefined.") @@ -858,6 +902,11 @@ class Wan22I2VPipeline( if prompt is None and prompt_embeds is None: raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) + if boundary_ratio is None and guidance_scale_2 is not None: raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.") diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py index 4953a809d..0f4caf9fd 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py @@ -38,12 +38,16 @@ from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + WAN22_MAX_SEQUENCE_LENGTH, create_transformer_from_config, load_transformer_config, retrieve_latents, ) from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -183,6 +187,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress # Text encoder self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) @@ -303,6 +308,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress width=width, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) # Adjust num_frames to be compatible with VAE temporal scaling @@ -326,7 +332,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_scale > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) @@ -502,7 +508,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, - max_sequence_length: int = 512, + max_sequence_length: int = WAN22_MAX_SEQUENCE_LENGTH, device: torch.device | None = None, dtype: torch.dtype | None = None, ): @@ -513,11 +519,51 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) + text_inputs_untruncated = self.tokenizer( + prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + error_context="for Wan2.2 text encoding", + ) + prompt_encode_length = max(int(text_inputs_untruncated.attention_mask.sum(dim=1).max().item()), 1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_clean = [self._prompt_clean(p) for p in negative_prompt] + neg_text_inputs_untruncated = self.tokenizer( + negative_prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + neg_text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name="negative_prompt", + error_context="for Wan2.2 text encoding", + ) + prompt_encode_length = max( + prompt_encode_length, + int(neg_text_inputs_untruncated.attention_mask.sum(dim=1).max().item()), + ) text_inputs = self.tokenizer( prompt_clean, padding="max_length", - max_length=max_sequence_length, + max_length=prompt_encode_length, truncation=True, add_special_tokens=True, return_attention_mask=True, @@ -530,21 +576,18 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress prompt_embeds = prompt_embeds.to(dtype=dtype, device=device) prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)] prompt_embeds = torch.stack( - [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 + [torch.cat([u, u.new_zeros(prompt_encode_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0 ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - negative_prompt_embeds = None if do_classifier_free_guidance: - negative_prompt = negative_prompt or "" - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], + negative_prompt_clean, padding="max_length", - max_length=max_sequence_length, + max_length=prompt_encode_length, truncation=True, add_special_tokens=True, return_attention_mask=True, @@ -557,7 +600,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)] negative_prompt_embeds = torch.stack( [ - torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + torch.cat([u, u.new_zeros(prompt_encode_length - u.size(0), u.size(1))]) for u in negative_prompt_embeds ], dim=0, @@ -667,6 +710,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress width, prompt_embeds=None, negative_prompt_embeds=None, + max_sequence_length=None, ): if height % 16 != 0 or width % 16 != 0: raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") @@ -682,6 +726,11 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress if prompt is None and prompt_embeds is None: raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights using AutoWeightsLoader for vLLM integration.""" loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 726d3c3cd..006f2d671 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -31,6 +31,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelOutput, ) from vllm_omni.diffusion.forward_context import get_forward_context +from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm logger = init_logger(__name__) @@ -633,7 +634,7 @@ class WanTransformerBlock(nn.Module): # 1. Self-attention self.norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) - + self.attn1 = WanSelfAttention( dim=dim, num_heads=num_heads, @@ -694,9 +695,7 @@ class WanTransformerBlock(nn.Module): hidden_states = hidden_states + attn_output # 3. Feed-forward - norm_hidden_states = self.norm3(hidden_states, c_scale_msa, c_shift_msa).type_as( - hidden_states - ) + norm_hidden_states = self.norm3(hidden_states, c_scale_msa, c_shift_msa).type_as(hidden_states) ff_output = self.ffn(norm_hidden_states) hidden_states = (hidden_states + ff_output * c_gate_msa).type_as(hidden_states) diff --git a/vllm_omni/diffusion/postprocess/rife_interpolator.py b/vllm_omni/diffusion/postprocess/rife_interpolator.py index b2b4a9319..89297d0a4 100644 --- a/vllm_omni/diffusion/postprocess/rife_interpolator.py +++ b/vllm_omni/diffusion/postprocess/rife_interpolator.py @@ -412,9 +412,12 @@ class FrameInterpolator: return restore_layout(video), 1 video, restore_range = _normalize_video_tensor_range(video) - # Prefer the decoded video's current device so CPU-offloaded requests do - # not move the tensor back to GPU just for interpolation. - model = self._ensure_model_loaded(preferred_device=video.device) + # A CPU tensor may be transport/offload state rather than an execution + # choice, so only trust it when it is already on an accelerator. + preferred_device = video.device + if preferred_device.type == "cpu": + preferred_device = _select_torch_device() + model = self._ensure_model_loaded(preferred_device=preferred_device) video = video.to(model.device()) intermediates_per_pair = 2**exp // 2 diff --git a/vllm_omni/diffusion/utils/prompt_utils.py b/vllm_omni/diffusion/utils/prompt_utils.py new file mode 100644 index 000000000..fc1769f4d --- /dev/null +++ b/vllm_omni/diffusion/utils/prompt_utils.py @@ -0,0 +1,38 @@ +import torch + + +def validate_prompt_sequence_lengths( + attention_mask: torch.Tensor, + *, + max_sequence_length: int, + supported_max_sequence_length: int, + prompt_name: str = "prompt", + length_offset: int = 0, + baseline_attention_mask: torch.Tensor | None = None, + error_context: str, +) -> None: + sequence_lengths = attention_mask.sum(dim=1) + if baseline_attention_mask is not None: + # Some callers need to validate only the user-controlled portion of a + # templated prompt. In those cases we subtract the fully-tokenized + # template baseline instead of only removing a fixed prefix length, + # because the template may also contribute a suffix or image markers. + baseline_lengths = baseline_attention_mask.sum(dim=1) + if baseline_lengths.shape[0] == 1 and sequence_lengths.shape[0] > 1: + baseline_lengths = baseline_lengths.expand(sequence_lengths.shape[0]) + sequence_lengths = sequence_lengths - baseline_lengths + if length_offset: + sequence_lengths = sequence_lengths - length_offset + sequence_lengths = torch.clamp(sequence_lengths, min=0) + too_long = torch.nonzero(sequence_lengths > max_sequence_length, as_tuple=False) + if too_long.numel() == 0: + return + + batch_idx = int(too_long[0].item()) + actual_length = int(sequence_lengths[batch_idx].item()) + prompt_ref = f"`{prompt_name}` at batch index {batch_idx}" if attention_mask.shape[0] > 1 else f"`{prompt_name}`" + raise ValueError( + f"{prompt_ref} is too long {error_context}: got {actual_length} tokens, but " + f"`max_sequence_length` is {max_sequence_length}. Shorten the prompt or increase " + f"`max_sequence_length` up to {supported_max_sequence_length}." + ) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 34d9d2404..2828b05d0 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -800,6 +800,16 @@ class AsyncOmniEngine: # We temporally create a default config for diffusion stage. # In the future, we should merge the default config with the user-provided config. normalized_kwargs = dict(kwargs) + default_sampling_params = normalized_kwargs.get("default_sampling_params") + if isinstance(default_sampling_params, str): + try: + default_sampling_params = json.loads(default_sampling_params) + except json.JSONDecodeError: + logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") + default_sampling_params = None + if not isinstance(default_sampling_params, dict): + default_sampling_params = None + stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): @@ -888,6 +898,7 @@ class AsyncOmniEngine: else {} ), }, + "default_sampling_params": stage_default_sampling_params, "final_output": True, "final_output_type": "image", } diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index ffe90cb3b..e127a3050 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -450,9 +450,8 @@ def initialize_diffusion_stage( and ultimately to ``AsyncOmniDiffusion``. use_inline: If True, uses the inline diffusion client instead of subprocess. """ - from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client - from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client od_config = OmniDiffusionConfig.from_kwargs( model=model, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index cec9869d2..df1960dfb 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -81,6 +81,7 @@ from vllm.tool_parsers import ToolParserManager from vllm.utils import random_uuid from vllm.utils.system_utils import decorate_logs +from vllm_omni.diffusion.data import OmniRequestError from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError from vllm_omni.entrypoints.openai.image_api_utils import ( @@ -1432,12 +1433,24 @@ async def edit_images( input_images_list.extend(urls) if not input_images_list: raise HTTPException(status_code=422, detail="Field 'image' or 'url' is required") - pil_images = await _load_input_images(input_images_list) - if len(pil_images) > 1 and not _supports_multimodal_image_inputs(raw_request, engine_client): + # Reject oversized multi-image edit requests before fetching or decoding + # any inputs. This keeps over-limit URL requests from burning network, + # CPU, and memory on work that will be rejected anyway. + max_input_images = _get_max_edit_input_images(raw_request, engine_client) + if max_input_images is not None and len(input_images_list) > max_input_images: + detail = ( + "Received multiple input images. Only a single image is supported by this model." + if max_input_images == 1 + else ( + f"Received {len(input_images_list)} input images. " + f"At most {max_input_images} images are supported by this model." + ) + ) raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, - detail="Received multiple input images. Only a single image is supported by this model.", + detail=detail, ) + pil_images = await _load_input_images(input_images_list) prompt["multi_modal_data"] = {} prompt["multi_modal_data"]["image"] = pil_images @@ -1601,18 +1614,25 @@ def _get_engine_and_model(raw_request: Request): return engine_client, model_name, normalized_stage_configs -def _supports_multimodal_image_inputs(raw_request: Request, engine_client: Any) -> bool: +def _get_diffusion_od_config(raw_request: Request, engine_client: Any) -> Any: diffusion_engine = getattr(raw_request.app.state, "diffusion_engine", None) or engine_client get_diffusion_od_config = getattr(diffusion_engine, "get_diffusion_od_config", None) - od_config = ( + return ( get_diffusion_od_config() if callable(get_diffusion_od_config) else getattr(diffusion_engine, "od_config", None) ) + +def _get_max_edit_input_images(raw_request: Request, engine_client: Any) -> int | None: + od_config = _get_diffusion_od_config(raw_request, engine_client) if od_config is None: # Preserve the existing compatibility behavior when the diffusion # config is not exposed on the serving surface. - return True - return bool(getattr(od_config, "supports_multimodal_inputs", False)) + return None + + if not bool(getattr(od_config, "supports_multimodal_inputs", False)): + return 1 + + return getattr(od_config, "max_multimodal_image_inputs", None) def _get_lora_from_json_str(lora_body): @@ -1682,11 +1702,14 @@ async def _generate_with_async_omni( pass sampling_params_list.append(default_stage_params) - async for output in engine_client.generate( - sampling_params_list=sampling_params_list, - **kwargs, - ): - result = output + try: + async for output in engine_client.generate( + sampling_params_list=sampling_params_list, + **kwargs, + ): + result = output + except OmniRequestError as e: + raise HTTPException(status_code=e.status_code, detail=str(e)) from e if result is None: raise HTTPException(