412 lines
15 KiB
Python
412 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from dataclasses import dataclass, field
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
from torch import nn
|
|
|
|
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
|
|
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler
|
|
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]
|
|
|
|
|
|
def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeypatch) -> None:
|
|
@dataclass
|
|
class RuntimeConfig:
|
|
flow_shift: float | None = None
|
|
diffusion_offload_config: dict[str, object] = field(
|
|
default_factory=lambda: {
|
|
"mode": "layer",
|
|
"components": ["text_encoder"],
|
|
}
|
|
)
|
|
enable_layerwise_offload: bool = True
|
|
post_init_calls: int = 0
|
|
_resolved_diffusion_offload: object = None
|
|
_diffusion_offload_flags_materialized: bool = False
|
|
|
|
def __post_init__(self) -> None:
|
|
self.post_init_calls += 1
|
|
|
|
config = RuntimeConfig()
|
|
resolved_offload = object()
|
|
config._resolved_diffusion_offload = resolved_offload
|
|
config._diffusion_offload_flags_materialized = True
|
|
captured: dict[str, object] = {}
|
|
|
|
def capture_base_init(self, *, od_config, prefix="") -> None:
|
|
captured.update(od_config=od_config, prefix=prefix)
|
|
|
|
monkeypatch.setattr(Wan22VACEPipeline.__mro__[1], "__init__", capture_base_init)
|
|
|
|
Wan22VACEPipeline(od_config=config, prefix="stage")
|
|
|
|
forwarded = captured["od_config"]
|
|
assert isinstance(forwarded, RuntimeConfig)
|
|
assert forwarded is not config
|
|
assert config.flow_shift is None
|
|
assert forwarded.flow_shift == 3.0
|
|
assert forwarded.diffusion_offload_config is config.diffusion_offload_config
|
|
assert forwarded._resolved_diffusion_offload is resolved_offload
|
|
assert forwarded._diffusion_offload_flags_materialized is True
|
|
assert forwarded.post_init_calls == 1
|
|
assert captured["prefix"] == "stage"
|
|
|
|
|
|
def _make_vace_pipeline() -> Wan22VACEPipeline:
|
|
pipeline = object.__new__(Wan22VACEPipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline.device = torch.device("cpu")
|
|
pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
|
|
pipeline.transformer_config = pipeline.transformer.config
|
|
pipeline.vae = StubVAE(z_dim=4)
|
|
pipeline.vae_scale_factor_temporal = 4
|
|
pipeline.vae_scale_factor_spatial = 8
|
|
pipeline.progress_bar = noop_progress_bar
|
|
return pipeline
|
|
|
|
|
|
def _make_vace_sampling(**overrides):
|
|
values: dict[str, object] = {
|
|
"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")
|
|
frame = Image.new("RGB", (64, 64), "black")
|
|
mask = Image.new("L", (64, 64), 255)
|
|
request = SimpleNamespace(
|
|
prompt={
|
|
"prompt": "p",
|
|
"multi_modal_data": {
|
|
"image": ref,
|
|
"video": [frame],
|
|
"mask": mask,
|
|
},
|
|
},
|
|
sampling_params=SimpleNamespace(height=None, width=None),
|
|
)
|
|
|
|
result = preprocess(request)
|
|
additional_info = result.prompt["additional_information"]
|
|
|
|
assert result.sampling_params.height == 432
|
|
assert result.sampling_params.width == 880
|
|
assert additional_info["reference_images"] == [ref]
|
|
assert additional_info["source_video"] == [frame]
|
|
assert additional_info["mask"] == [mask]
|
|
|
|
|
|
def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None:
|
|
captured: dict[str, object] = {}
|
|
|
|
class FakeVACETransformer:
|
|
def __init__(self, **kwargs) -> None:
|
|
captured.update(kwargs)
|
|
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.WanVACETransformer3DModel",
|
|
FakeVACETransformer,
|
|
)
|
|
|
|
transformer = create_vace_transformer_from_config(
|
|
{
|
|
"patch_size": [1, 2, 2],
|
|
"in_channels": 96,
|
|
"out_channels": 16,
|
|
"vace_layers": [0, 1, 2],
|
|
"vace_in_channels": 132,
|
|
"unknown": "ignored",
|
|
}
|
|
)
|
|
|
|
assert isinstance(transformer, FakeVACETransformer)
|
|
assert captured == {
|
|
"patch_size": (1, 2, 2),
|
|
"in_channels": 96,
|
|
"out_channels": 16,
|
|
"vace_layers": [0, 1, 2],
|
|
"vace_in_channels": 132,
|
|
}
|
|
|
|
|
|
def test_vace_transformer_layerwise_offloads_conditioning_blocks() -> None:
|
|
assert WanVACETransformer3DModel._layerwise_offload_blocks_attrs == ["vace_blocks", "blocks"]
|
|
|
|
|
|
def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None:
|
|
pipeline = _make_vace_pipeline()
|
|
mask = torch.ones(1, 3, 5, 16, 16)
|
|
reference_images = [[torch.zeros(3, 16, 16), torch.zeros(3, 16, 16)]]
|
|
|
|
encoded = pipeline.prepare_masks(mask, reference_images)
|
|
|
|
assert encoded.shape == (1, 64, 4, 2, 2)
|
|
torch.testing.assert_close(encoded[:, :, :2], torch.zeros(1, 64, 2, 2, 2))
|
|
torch.testing.assert_close(encoded[:, :, 2:], torch.ones(1, 64, 2, 2, 2))
|
|
|
|
|
|
def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
|
|
pipeline = _make_vace_pipeline()
|
|
latents = torch.zeros(1, 4, 1, 2, 2)
|
|
vace_context = torch.ones(1, 12, 1, 2, 2)
|
|
calls = []
|
|
|
|
def fake_predict_noise_maybe_with_cfg(**kwargs):
|
|
calls.append(kwargs)
|
|
return torch.ones_like(latents)
|
|
|
|
pipeline.predict_noise_maybe_with_cfg = fake_predict_noise_maybe_with_cfg # type: ignore[method-assign]
|
|
pipeline.scheduler_step_maybe_with_cfg = lambda noise, t, current, cfg: current + noise # type: ignore[method-assign]
|
|
|
|
result = pipeline.diffuse(
|
|
latents=latents,
|
|
timesteps=torch.tensor([5]),
|
|
prompt_embeds=torch.zeros(1, 2, 3),
|
|
negative_prompt_embeds=torch.zeros(1, 2, 3),
|
|
guidance_scale=4.0,
|
|
dtype=torch.float32,
|
|
attention_kwargs={},
|
|
vace_context=vace_context,
|
|
vace_context_scale=0.75,
|
|
)
|
|
|
|
assert calls[0]["do_true_cfg"] is True
|
|
assert calls[0]["true_cfg_scale"] == 4.0
|
|
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))
|
|
|
|
|
|
@pytest.mark.parametrize("solver", ["unipc", "euler"])
|
|
@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
|
|
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> 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 = build_wan_scheduler(solver, shift)
|
|
pipeline.od_config = SimpleNamespace(flow_shift=shift)
|
|
pipeline._sample_solver = solver
|
|
pipeline._flow_shift = shift
|
|
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,
|
|
num_inference_steps=50,
|
|
extra_args={"sample_solver": solver, "flow_shift": shift},
|
|
),
|
|
),
|
|
SimpleNamespace(
|
|
request_id="b",
|
|
prompt={"prompt": "second"},
|
|
sampling_params=_make_vace_sampling(
|
|
generator=gen_b,
|
|
latents=latents_b,
|
|
num_outputs_per_prompt=2,
|
|
num_inference_steps=50,
|
|
extra_args={"sample_solver": solver, "flow_shift": shift},
|
|
),
|
|
),
|
|
]
|
|
)
|
|
|
|
outputs = pipeline.forward(batch)
|
|
if solver == "unipc":
|
|
sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
|
|
sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
|
|
torch.testing.assert_close(
|
|
pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
|
|
)
|
|
assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
|
|
else:
|
|
reference = build_wan_scheduler("euler", shift)
|
|
reference.set_timesteps(50, device="cpu")
|
|
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
|
|
|
|
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)
|