anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Perf] Support request-level batching for Wan2.2 pipelines (#5676)

Browse Source 
Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>
This commit is contained in:
hurukawa authored and GitHub committed 2026-08-15 23:58:36 +08:00
1 parent d1e230c95b
commit 2a315e1a2d
18 files changed
+1871 -374
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py
+285 -1
View File
@@ -8,15 +8,28 @@ import torch
from PIL import Image
from torch import nn
from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
Wan22I2VPipeline,
get_wan22_i2v_post_process_func,
get_wan22_i2v_pre_process_func,
)
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
def test_wan22_i2v_postprocess_honors_request_output_type() -> None:
video = torch.zeros(1, 4, 1, 2, 2)
output = get_wan22_i2v_post_process_func(SimpleNamespace())(
video,
sampling_params=SimpleNamespace(output_type="latent"),
)
assert output is video
def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
pipeline = object.__new__(Wan22I2VPipeline)
nn.Module.__init__(pipeline)
@@ -31,6 +44,29 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
return pipeline
def _make_i2v_sampling(**overrides):
values = {
"height": 16,
"width": 16,
"num_frames": 5,
"num_inference_steps": 1,
"guidance_scale_provided": True,
"guidance_scale": 1.0,
"guidance_scale_2": None,
"guidance_scale_2_provided": False,
"boundary_ratio": None,
"generator": None,
"seed": None,
"num_outputs_per_prompt": 1,
"max_sequence_length": 8,
"latents": None,
"output_type": "latent",
"extra_args": {},
}
values.update(overrides)
return SimpleNamespace(**values)
def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
request = SimpleNamespace(
@@ -53,6 +89,88 @@ def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None:
preprocess(missing_image)
def _make_i2v_preprocess_request(last_image):
return SimpleNamespace(
prompt={
"prompt": "p",
"multi_modal_data": {
"image": Image.new("RGB", (320, 160), "red"),
"last_image": last_image,
},
},
sampling_params=SimpleNamespace(height=16, width=16),
)
def test_i2v_preprocess_treats_empty_last_image_list_as_absent() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
result = preprocess(_make_i2v_preprocess_request([]))
assert result.prompt["multi_modal_data"]["last_image"] is None
assert result.batch_compatibility_key == ("wan22_i2v_last_image", False)
def test_i2v_preprocess_unwraps_single_last_image_list() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
last_image = Image.new("RGB", (16, 16), "blue")
result = preprocess(_make_i2v_preprocess_request([last_image]))
assert result.prompt["multi_modal_data"]["last_image"] is last_image
assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
@pytest.mark.parametrize(
"last_image",
[
Image.new("RGB", (16, 16), "blue"),
torch.zeros(3, 16, 16),
],
)
def test_i2v_preprocess_preserves_supported_last_image(last_image) -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
result = preprocess(_make_i2v_preprocess_request(last_image))
assert result.prompt["multi_modal_data"]["last_image"] is last_image
assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
def test_i2v_preprocess_loads_last_image_path(tmp_path) -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
path = tmp_path / "last.png"
Image.new("RGB", (16, 16), "blue").save(path)
result = preprocess(_make_i2v_preprocess_request(str(path)))
last_image = result.prompt["multi_modal_data"]["last_image"]
assert isinstance(last_image, Image.Image)
assert last_image.mode == "RGB"
assert result.batch_compatibility_key == ("wan22_i2v_last_image", True)
def test_i2v_preprocess_rejects_multiple_last_images() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
with pytest.raises(ValueError, match="at most one last_image"):
preprocess(
_make_i2v_preprocess_request(
[
Image.new("RGB", (16, 16), "blue"),
Image.new("RGB", (16, 16), "green"),
]
)
)
def test_i2v_preprocess_rejects_unsupported_last_image_type() -> None:
preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace())
with pytest.raises(TypeError, match="Unsupported last_image format"):
preprocess(_make_i2v_preprocess_request({"not": "an image"}))
def test_i2v_diffuse_selects_stage_guidance_and_expands_timesteps() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
latents = torch.zeros(1, 4, 2, 4, 4)
@@ -126,3 +244,169 @@ def test_i2v_prepare_latents_builds_expand_condition_and_first_frame_mask() -> N
assert first_frame_mask.shape == (1, 1, 2, 2, 2)
assert first_frame_mask[:, :, 0].sum() == 0
assert first_frame_mask[:, :, 1].sum() == 4
def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
generators = [
torch.Generator(device="cpu").manual_seed(1),
torch.Generator(device="cpu").manual_seed(2),
]
latents, condition, first_frame_mask = pipeline.prepare_latents(
image=torch.zeros(2, 3, 16, 16),
batch_size=2,
num_channels_latents=4,
height=16,
width=16,
num_frames=5,
dtype=torch.float32,
device=torch.device("cpu"),
generator=generators,
)
assert latents.shape == (2, 4, 2, 2, 2)
assert condition.shape == (2, 4, 1, 2, 2)
assert first_frame_mask.shape == (2, 1, 2, 2, 2)
def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_i2v_pipeline(expand_timesteps=True)
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 5.0
pipeline.boundary_ratio = 0.875
pipeline.has_image_encoder = False
pipeline._guidance_scale = None
pipeline._guidance_scale_2 = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
prepare_call = {}
def _fake_encode_prompt(**kwargs):
batch_size = len(kwargs["prompt"])
n = batch_size * kwargs["num_videos_per_prompt"]
return torch.zeros(n, 2, 3), None
def _fake_prepare_latents(**kwargs):
prepare_call.update(kwargs)
batch_size = kwargs["batch_size"]
return (
kwargs["latents"],
torch.zeros(batch_size, 4, 1, 2, 2),
torch.ones(batch_size, 1, 2, 2, 2),
)
pipeline.encode_prompt = _fake_encode_prompt # type: ignore[method-assign]
pipeline.prepare_latents = _fake_prepare_latents # type: ignore[method-assign]
pipeline.diffuse = lambda **kwargs: kwargs["latents"] # type: ignore[method-assign]
gen_a = torch.Generator(device="cpu").manual_seed(1)
gen_b = torch.Generator(device="cpu").manual_seed(2)
latents_a = torch.zeros(2, 4, 2, 2, 2)
latents_b = torch.ones(2, 4, 2, 2, 2)
image_a = torch.zeros(1, 3, 16, 16)
image_b = torch.ones(1, 3, 16, 16)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first", "multi_modal_data": {"image": image_a}},
sampling_params=_make_i2v_sampling(
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": image_b}},
sampling_params=_make_i2v_sampling(
generator=gen_b,
latents=latents_b,
num_outputs_per_prompt=2,
),
),
]
)
outputs = pipeline.forward(batch)
assert prepare_call["batch_size"] == 4
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
torch.testing.assert_close(prepare_call["latents"], torch.cat([latents_a, latents_b]))
torch.testing.assert_close(
prepare_call["image"],
torch.cat([image_a, image_a, image_b, image_b]),
)
assert len(outputs) == 2
torch.testing.assert_close(outputs[0].output, latents_a)
torch.testing.assert_close(outputs[1].output, latents_b)
def test_i2v_forward_rejects_mismatched_tensor_condition_shapes(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_i2v_pipeline(expand_timesteps=True)
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 5.0
pipeline.boundary_ratio = 0.875
pipeline.has_image_encoder = False
pipeline._guidance_scale = None
pipeline._guidance_scale_2 = None
pipeline._num_timesteps = None
pipeline._current_timestep = None
pipeline.check_inputs = lambda **kwargs: None
pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 16)}},
sampling_params=_make_i2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": torch.zeros(1, 3, 16, 8)}},
sampling_params=_make_i2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="image condition"):
pipeline.forward(batch)
def test_i2v_forward_rejects_mixed_last_image_presence() -> None:
pipeline = _make_i2v_pipeline(expand_timesteps=True)
image = torch.zeros(1, 3, 16, 16)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"multi_modal_data": {"image": image, "last_image": image},
},
sampling_params=_make_i2v_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second", "multi_modal_data": {"image": image}},
sampling_params=_make_i2v_sampling(),
),
]
)
with pytest.raises(ValueError, match="mix of provided and missing last_image"):
pipeline.forward(batch)