[Perf] Support request-level batching for Wan2.2 pipelines (#5676)
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
1 parent
d1e230c95b
commit
2a315e1a2d
18 files changed
+1871
-374
No files matched your search
@@ -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)
|
||||
@@ -93,6 +93,29 @@ def _make_pipeline() -> Wan22Pipeline:
|
||||
return pipeline
|
||||
|
||||
|
||||
def _make_sampling(**overrides):
|
||||
values = {
|
||||
"height": None,
|
||||
"width": None,
|
||||
"num_frames": 1,
|
||||
"num_inference_steps": 2,
|
||||
"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": 32,
|
||||
"latents": None,
|
||||
"output_type": "latent",
|
||||
"extra_args": {},
|
||||
}
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sampling_params_kwargs", "expected_low", "expected_high"),
|
||||
[
|
||||
@@ -129,9 +152,10 @@ def test_forward_delegates_denoising_to_diffuse(
|
||||
)
|
||||
batch = DiffusionRequestBatch(requests=[mock_req])
|
||||
|
||||
output = pipeline.forward(batch)
|
||||
outputs = pipeline.forward(batch)
|
||||
|
||||
assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8)))
|
||||
assert len(outputs) == 1
|
||||
assert torch.equal(outputs[0].output, torch.ones((1, 4, 1, 8, 8)))
|
||||
assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8))
|
||||
assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps)
|
||||
assert captured["guidance_low"] == expected_low
|
||||
@@ -142,6 +166,134 @@ def test_forward_delegates_denoising_to_diffuse(
|
||||
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
|
||||
|
||||
|
||||
def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
|
||||
pipeline = _make_pipeline()
|
||||
encode_call = {}
|
||||
prepare_call = {}
|
||||
|
||||
def _fake_encode_prompt(**kwargs):
|
||||
encode_call.update(kwargs)
|
||||
batch_size = len(kwargs["prompt"])
|
||||
n = batch_size * kwargs["num_videos_per_prompt"]
|
||||
return torch.arange(n, dtype=torch.float32).view(n, 1, 1), torch.zeros(n, 1, 1)
|
||||
|
||||
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, 1, 2, 2)
|
||||
latents_b = torch.ones(2, 4, 1, 2, 2)
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={"prompt": "first", "negative_prompt": "bad first"},
|
||||
sampling_params=_make_sampling(
|
||||
generator=gen_a,
|
||||
latents=latents_a,
|
||||
num_outputs_per_prompt=2,
|
||||
),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={"prompt": "second", "negative_prompt": "bad second"},
|
||||
sampling_params=_make_sampling(
|
||||
generator=gen_b,
|
||||
latents=latents_b,
|
||||
num_outputs_per_prompt=2,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
outputs = pipeline.forward(batch)
|
||||
|
||||
assert encode_call["prompt"] == ["first", "second"]
|
||||
assert encode_call["negative_prompt"] == ["bad first", "bad second"]
|
||||
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_forward_batches_precomputed_prompt_embeddings() -> None:
|
||||
pipeline = _make_pipeline()
|
||||
diffuse_call = {}
|
||||
pipeline.encode_prompt = lambda **kwargs: pytest.fail("text encoder must not run") # type: ignore[method-assign]
|
||||
pipeline.prepare_latents = lambda **kwargs: torch.zeros(kwargs["batch_size"], 4, 1, 2, 2) # type: ignore[method-assign]
|
||||
|
||||
def _fake_diffuse(**kwargs):
|
||||
diffuse_call.update(kwargs)
|
||||
return kwargs["latents"]
|
||||
|
||||
pipeline.diffuse = _fake_diffuse # type: ignore[method-assign]
|
||||
embeds_a = torch.zeros(3, 4)
|
||||
embeds_b = torch.ones(3, 4)
|
||||
negative_a = torch.full((3, 4), 2.0)
|
||||
negative_b = torch.full((3, 4), 3.0)
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={"prompt_embeds": embeds_a, "negative_prompt_embeds": negative_a},
|
||||
sampling_params=_make_sampling(guidance_scale=4.0),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={"prompt_embeds": embeds_b, "negative_prompt_embeds": negative_b},
|
||||
sampling_params=_make_sampling(guidance_scale=4.0),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
outputs = pipeline.forward(batch)
|
||||
|
||||
torch.testing.assert_close(diffuse_call["prompt_embeds"], torch.stack([embeds_a, embeds_b]))
|
||||
torch.testing.assert_close(diffuse_call["negative_prompt_embeds"], torch.stack([negative_a, negative_b]))
|
||||
assert len(outputs) == 2
|
||||
|
||||
|
||||
def test_prepare_latents_with_request_generators_matches_single_generation() -> None:
|
||||
pipeline = _make_pipeline()
|
||||
kwargs = {
|
||||
"num_channels_latents": 4,
|
||||
"height": 16,
|
||||
"width": 16,
|
||||
"num_frames": 5,
|
||||
"dtype": torch.float32,
|
||||
"device": torch.device("cpu"),
|
||||
}
|
||||
|
||||
batched = Wan22Pipeline.prepare_latents(
|
||||
pipeline,
|
||||
batch_size=2,
|
||||
generator=[torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)],
|
||||
**kwargs,
|
||||
)
|
||||
singles = torch.cat(
|
||||
[
|
||||
Wan22Pipeline.prepare_latents(
|
||||
pipeline,
|
||||
batch_size=1,
|
||||
generator=torch.Generator().manual_seed(seed),
|
||||
**kwargs,
|
||||
)
|
||||
for seed in (1, 2)
|
||||
]
|
||||
)
|
||||
|
||||
torch.testing.assert_close(batched, singles)
|
||||
assert not torch.equal(batched[0], batched[1])
|
||||
|
||||
|
||||
def test_diffuse_runs_prediction_and_scheduler_for_each_timestep() -> None:
|
||||
pipeline = _make_pipeline()
|
||||
latents = torch.zeros((1, 1, 1, 2, 2), dtype=torch.float32)
|
||||
|
||||
@@ -10,6 +10,7 @@ import torch
|
||||
import vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 as wan22_module
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
create_transformer_from_config,
|
||||
get_wan22_post_process_func,
|
||||
load_transformer_config,
|
||||
retrieve_latents,
|
||||
)
|
||||
@@ -26,6 +27,17 @@ class _LatentDist:
|
||||
return torch.tensor([2.0])
|
||||
|
||||
|
||||
def test_wan22_postprocess_honors_request_output_type() -> None:
|
||||
video = torch.zeros(1, 4, 1, 2, 2)
|
||||
|
||||
output = get_wan22_post_process_func(SimpleNamespace())(
|
||||
video,
|
||||
sampling_params=SimpleNamespace(output_type="latent"),
|
||||
)
|
||||
|
||||
assert output is video
|
||||
|
||||
|
||||
def test_retrieve_latents_supports_sample_mode_argmax_and_direct_latents() -> None:
|
||||
generator = torch.Generator(device="cpu")
|
||||
|
||||
|
||||
@@ -1,14 +1,307 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
|
||||
Wan22S2VPipeline,
|
||||
_make_clip_generators,
|
||||
)
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_s2v_sampling(**overrides):
|
||||
values = {
|
||||
"height": 16,
|
||||
"width": 16,
|
||||
"num_frames": 8,
|
||||
"num_inference_steps": 1,
|
||||
"guidance_scale_provided": True,
|
||||
"guidance_scale": 1.0,
|
||||
"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 _make_s2v_validation_pipeline() -> Wan22S2VPipeline:
|
||||
pipeline = object.__new__(Wan22S2VPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.transformer = SimpleNamespace(dtype=torch.float32)
|
||||
pipeline.vae_scale_factor_spatial = 8
|
||||
pipeline.resolution_divisor = 16
|
||||
pipeline.motion_frames = 7
|
||||
pipeline.drop_first_motion = True
|
||||
pipeline._DEFAULT_INFER_FRAMES = 8
|
||||
pipeline._guidance_scale = None
|
||||
pipeline.check_inputs = lambda *args, **kwargs: None
|
||||
pipeline.encode_prompt = lambda **kwargs: (torch.zeros(2, 2, 3), None) # type: ignore[method-assign]
|
||||
return pipeline
|
||||
|
||||
|
||||
def test_s2v_predict_noise_keeps_batch_dimension() -> None:
|
||||
pipeline = object.__new__(Wan22S2VPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
transformer = MagicMock()
|
||||
transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.bfloat16)])
|
||||
transformer.return_value = (torch.zeros(1, 16, 2, 2, 2),)
|
||||
pipeline.transformer = transformer
|
||||
|
||||
result = pipeline.predict_noise(hidden_states=torch.zeros(1, 16, 2, 2, 2))
|
||||
|
||||
assert result.shape == (1, 16, 2, 2, 2)
|
||||
|
||||
|
||||
def test_s2v_clip_generators_preserve_main_seed_per_clip_behavior() -> None:
|
||||
explicit_generator = torch.Generator(device="cpu").manual_seed(999)
|
||||
|
||||
generators = _make_clip_generators(
|
||||
seeds=[1234, 5678, None, None],
|
||||
request_generators=[
|
||||
torch.Generator(device="cpu").manual_seed(1234),
|
||||
torch.Generator(device="cpu").manual_seed(5678),
|
||||
explicit_generator,
|
||||
None,
|
||||
],
|
||||
clip_index=2,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert [generator.initial_seed() for generator in generators if generator is not None] == [1236, 5680, 999, 2]
|
||||
assert generators[2] is explicit_generator
|
||||
|
||||
|
||||
def test_s2v_forward_rejects_different_num_repeat() -> None:
|
||||
pipeline = _make_s2v_validation_pipeline()
|
||||
pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 2, 16) # type: ignore[method-assign]
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
audio = np.zeros(16000, dtype=np.float32)
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={
|
||||
"prompt": "first",
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
"additional_information": {"num_repeat": 1},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={
|
||||
"prompt": "second",
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
"additional_information": {"num_repeat": 2},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="num_repeat"):
|
||||
pipeline.forward(batch)
|
||||
|
||||
|
||||
def test_s2v_forward_rejects_different_audio_time_lengths() -> None:
|
||||
pipeline = _make_s2v_validation_pipeline()
|
||||
pipeline.encode_audio = lambda audio, **kwargs: ( # type: ignore[method-assign]
|
||||
torch.zeros(1, 1, 2, 16 if audio[0] == 0 else 24),
|
||||
1,
|
||||
16 if audio[0] == 0 else 24,
|
||||
)
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={
|
||||
"prompt": "first",
|
||||
"multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={
|
||||
"prompt": "second",
|
||||
"multi_modal_data": {"image": image, "audio": np.ones(16000, dtype=np.float32)},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="audio embedding time lengths"):
|
||||
pipeline.forward(batch)
|
||||
|
||||
|
||||
def test_s2v_forward_rejects_different_init_first_frame() -> None:
|
||||
pipeline = _make_s2v_validation_pipeline()
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
audio = np.zeros(16000, dtype=np.float32)
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={
|
||||
"prompt": "first",
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
"additional_information": {"init_first_frame": True},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={
|
||||
"prompt": "second",
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
"additional_information": {"init_first_frame": False},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="init_first_frame"):
|
||||
pipeline.forward(batch)
|
||||
|
||||
|
||||
def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
|
||||
pipeline = _make_s2v_validation_pipeline()
|
||||
pipeline.encode_audio = lambda *args, **kwargs: (torch.zeros(1, 1, 2, 16), 1, 16) # type: ignore[method-assign]
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={
|
||||
"prompt": "first",
|
||||
"multi_modal_data": {"image": image, "audio": np.zeros(16000, dtype=np.float32)},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={
|
||||
"prompt": "second",
|
||||
"multi_modal_data": {"image": image, "audio": np.zeros(8000, dtype=np.float32)},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="raw audio shapes and sample rates"):
|
||||
pipeline.forward(batch)
|
||||
|
||||
|
||||
def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
|
||||
pipeline = object.__new__(Wan22S2VPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.transformer = MagicMock()
|
||||
pipeline.transformer.dtype = torch.float32
|
||||
pipeline.transformer.parameters.return_value = iter([torch.zeros(1, dtype=torch.float32)])
|
||||
pipeline.transformer.casual_audio_encoder = None
|
||||
pipeline.transformer.encode_audio.side_effect = lambda audio, _motion: {"audio_emb": audio}
|
||||
pipeline.vae = MagicMock()
|
||||
pipeline.vae.dtype = torch.float32
|
||||
pipeline.vae.decode.return_value = (torch.zeros(4, 3, 8, 16, 16),)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=False,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
)
|
||||
pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
|
||||
pipeline.vae_scale_factor_spatial = 8
|
||||
pipeline.resolution_divisor = 16
|
||||
pipeline.motion_frames = 7
|
||||
pipeline.drop_first_motion = True
|
||||
pipeline._DEFAULT_INFER_FRAMES = 8
|
||||
pipeline._guidance_scale = None
|
||||
pipeline._num_timesteps = None
|
||||
pipeline.check_inputs = lambda *args, **kwargs: None
|
||||
pipeline.encode_prompt = MagicMock(return_value=(torch.zeros(4, 2, 3), torch.ones(4, 2, 3)))
|
||||
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 1, 2, 8), 1, 8))
|
||||
pipeline.encode_ref_image = MagicMock(return_value=torch.zeros(1, 16, 1, 2, 2))
|
||||
pipeline.prepare_motion_latents = MagicMock(
|
||||
side_effect=lambda pixels, **_: torch.zeros(pixels.shape[0], 16, 2, 2, 2)
|
||||
)
|
||||
pipeline._denormalize_latents = lambda latents: latents
|
||||
pipeline.diffuse = MagicMock(side_effect=lambda **kwargs: kwargs["latents"])
|
||||
|
||||
image = PIL.Image.new("RGB", (16, 16))
|
||||
audio_a = np.zeros(16000, dtype=np.float32)
|
||||
audio_b = np.ones(16000, dtype=np.float32)
|
||||
latents_a = torch.zeros(2, 16, 2, 2, 2)
|
||||
latents_b = torch.ones(2, 16, 2, 2, 2)
|
||||
generators_a = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
|
||||
generators_b = [torch.Generator().manual_seed(3), torch.Generator().manual_seed(4)]
|
||||
batch = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="a",
|
||||
prompt={
|
||||
"prompt": "first",
|
||||
"negative_prompt": "negative first",
|
||||
"multi_modal_data": {"image": image, "audio": audio_a},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(
|
||||
num_outputs_per_prompt=2,
|
||||
generator=generators_a,
|
||||
latents=latents_a,
|
||||
),
|
||||
),
|
||||
SimpleNamespace(
|
||||
request_id="b",
|
||||
prompt={
|
||||
"prompt": "second",
|
||||
"negative_prompt": "negative second",
|
||||
"multi_modal_data": {"image": image, "audio": audio_b},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(
|
||||
num_outputs_per_prompt=2,
|
||||
generator=generators_b,
|
||||
latents=latents_b,
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as platform:
|
||||
platform.is_available.return_value = False
|
||||
outputs = pipeline.forward(batch)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0].output[0].shape[0] == 2
|
||||
assert outputs[1].output[0].shape[0] == 2
|
||||
assert outputs[0].output[1].shape == (16000,)
|
||||
assert outputs[1].output[1].shape == (16000,)
|
||||
np.testing.assert_array_equal(outputs[0].output[1], audio_a)
|
||||
np.testing.assert_array_equal(outputs[1].output[1], audio_b)
|
||||
assert np.shares_memory(outputs[0].output[1], audio_a)
|
||||
assert np.shares_memory(outputs[1].output[1], audio_b)
|
||||
torch.testing.assert_close(pipeline.diffuse.call_args.kwargs["latents"], torch.cat([latents_a, latents_b]))
|
||||
assert pipeline.diffuse.call_args.kwargs["clip_generator"] == generators_a + generators_b
|
||||
assert pipeline.encode_prompt.call_args.kwargs["prompt"] == ["first", "second"]
|
||||
assert pipeline.encode_prompt.call_args.kwargs["negative_prompt"] == ["negative first", "negative second"]
|
||||
|
||||
|
||||
def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks():
|
||||
model = object.__new__(WanS2VTransformer3DModel)
|
||||
nn.Module.__init__(model)
|
||||
@@ -260,31 +553,37 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
|
||||
pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean
|
||||
|
||||
# -- Build request --
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
|
||||
sampling_params = MagicMock()
|
||||
sampling_params.height = 704
|
||||
sampling_params.width = 1024
|
||||
sampling_params.num_frames = 80
|
||||
sampling_params.num_inference_steps = 5
|
||||
sampling_params.guidance_scale = 4.5
|
||||
sampling_params.guidance_scale_provided = True
|
||||
sampling_params.generator = None
|
||||
sampling_params.seed = 42
|
||||
sampling_params = _make_s2v_sampling(
|
||||
height=704,
|
||||
width=1024,
|
||||
num_frames=80,
|
||||
num_inference_steps=5,
|
||||
guidance_scale=4.5,
|
||||
generator=torch.Generator(device="cpu").manual_seed(42),
|
||||
max_sequence_length=512,
|
||||
)
|
||||
|
||||
ref_image = PIL.Image.new("RGB", (1024, 704))
|
||||
audio_data = np.zeros(16000, dtype=np.float32)
|
||||
|
||||
req = MagicMock(spec=OmniDiffusionRequest)
|
||||
req.prompts = [
|
||||
{
|
||||
"prompt": "test prompt",
|
||||
"negative_prompt": "bad quality",
|
||||
"multi_modal_data": {"image": ref_image, "audio": audio_data},
|
||||
"additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False},
|
||||
}
|
||||
]
|
||||
req.sampling_params = sampling_params
|
||||
req = DiffusionRequestBatch(
|
||||
requests=[
|
||||
SimpleNamespace(
|
||||
request_id="test",
|
||||
prompt={
|
||||
"prompt": "test prompt",
|
||||
"negative_prompt": "bad quality",
|
||||
"multi_modal_data": {"image": ref_image, "audio": audio_data},
|
||||
"additional_information": {
|
||||
"audio_path": audio_data,
|
||||
"pose_video": None,
|
||||
"init_first_frame": False,
|
||||
},
|
||||
},
|
||||
sampling_params=sampling_params,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
# -- Mock methods that use complex internal state --
|
||||
pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80))
|
||||
@@ -295,7 +594,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
|
||||
pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False)
|
||||
|
||||
# Mock predict_noise_maybe_with_cfg from CFGParallelMixin
|
||||
pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128))
|
||||
pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(1, 16, 20, 88, 128))
|
||||
pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline)
|
||||
|
||||
# Mock platform methods
|
||||
@@ -306,7 +605,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
|
||||
with patch(
|
||||
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000)
|
||||
):
|
||||
result = Wan22S2VPipeline.forward(pipeline, req=req)
|
||||
result = Wan22S2VPipeline.forward(pipeline, req=req)[0]
|
||||
|
||||
# -- Assertions --
|
||||
# Text encoder was called
|
||||
|
||||
@@ -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)
|
||||
Reference in new issue
Block a user