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_vace_pipeline.py
+209 -1
View File
@@ -8,13 +8,14 @@ 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_vace import (
Wan22VACEPipeline,
create_vace_transformer_from_config,
get_wan22_vace_pre_process_func,
)
from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -32,6 +33,27 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
return pipeline
def _make_vace_sampling(**overrides):
values = {
"height": 16,
"width": 16,
"num_frames": 5,
"num_inference_steps": 1,
"guidance_scale_provided": True,
"guidance_scale": 1.0,
"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_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
preprocess = get_wan22_vace_pre_process_func(SimpleNamespace())
ref = Image.new("RGB", (320, 160), "green")
@@ -138,3 +160,189 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
assert calls[0]["positive_kwargs"]["vace_context"] is vace_context
assert calls[0]["negative_kwargs"]["vace_context_scale"] == 0.75
torch.testing.assert_close(result, torch.ones_like(latents))
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = None
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.boundary_ratio = None
pipeline._guidance_scale = 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)
return kwargs["latents"]
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)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={"prompt": "first"},
sampling_params=_make_vace_sampling(
generator=gen_a,
latents=latents_a,
num_outputs_per_prompt=2,
),
),
SimpleNamespace(
request_id="b",
prompt={"prompt": "second"},
sampling_params=_make_vace_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]))
assert len(outputs) == 2
torch.testing.assert_close(outputs[0].output, latents_a)
torch.testing.assert_close(outputs[1].output, latents_b)
def test_vace_forward_rejects_mismatched_condition_shapes(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = object()
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.boundary_ratio = None
pipeline._guidance_scale = 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]
def _fake_preprocess(video, **kwargs):
del kwargs
mask = torch.ones_like(video)
return video, mask, [[]]
pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign]
pipeline.prepare_video_latents = ( # type: ignore[method-assign]
lambda video, mask, refs, generator, device: torch.zeros(1, 8, video.shape[2], 2, 2)
)
pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, mask.shape[2], 2, 2) # type: ignore[method-assign]
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"additional_information": {"source_video": torch.zeros(1, 3, 5, 16, 16)},
},
sampling_params=_make_vace_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"additional_information": {"source_video": torch.zeros(1, 3, 3, 16, 16)},
},
sampling_params=_make_vace_sampling(),
),
]
)
with pytest.raises(ValueError, match="VACE condition"):
pipeline.forward(batch)
def test_vace_forward_rejects_mismatched_reference_image_counts(monkeypatch) -> None:
monkeypatch.setattr(
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
SimpleNamespace(is_available=lambda: False),
)
pipeline = _make_vace_pipeline()
pipeline.transformer.vace_patch_embedding = object()
pipeline.transformer_2 = None
pipeline.scheduler = StubScheduler([9])
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
pipeline._sample_solver = "unipc"
pipeline._flow_shift = 3.0
pipeline.boundary_ratio = None
pipeline._guidance_scale = 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]
def _fake_preprocess(video, mask, reference_images, **kwargs):
del kwargs
return video, mask, [reference_images]
pipeline.preprocess_conditions = _fake_preprocess # type: ignore[method-assign]
pipeline.prepare_video_latents = ( # type: ignore[method-assign]
lambda video, mask, refs, generator, device: torch.zeros(1, 8, 5, 2, 2)
)
pipeline.prepare_masks = lambda mask, refs: torch.zeros(1, 4, 5, 2, 2) # type: ignore[method-assign]
video = torch.zeros(1, 3, 5, 16, 16)
mask = torch.ones_like(video)
reference = torch.zeros(3, 16, 16)
batch = DiffusionRequestBatch(
requests=[
SimpleNamespace(
request_id="a",
prompt={
"prompt": "first",
"additional_information": {
"source_video": video,
"mask": mask,
"reference_images": [reference],
},
},
sampling_params=_make_vace_sampling(),
),
SimpleNamespace(
request_id="b",
prompt={
"prompt": "second",
"additional_information": {
"source_video": video,
"mask": mask,
"reference_images": [reference, reference],
},
},
sampling_params=_make_vace_sampling(),
),
]
)
with pytest.raises(ValueError, match="same number of reference images"):
pipeline.forward(batch)