98 lines
3.4 KiB
Python
98 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
from torch import nn
|
|
|
|
from tests.diffusion.models.wan2_2.conftest import StubTransformer, StubVAE, noop_progress_bar
|
|
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import (
|
|
Wan22TI2VPipeline,
|
|
get_wan22_ti2v_pre_process_func,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def _make_ti2v_pipeline() -> Wan22TI2VPipeline:
|
|
pipeline = object.__new__(Wan22TI2VPipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline.device = torch.device("cpu")
|
|
pipeline.transformer = StubTransformer(in_channels=4, out_channels=4)
|
|
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 test_ti2v_preprocess_uses_720p_area_for_image_condition() -> None:
|
|
preprocess = get_wan22_ti2v_pre_process_func(SimpleNamespace())
|
|
request = SimpleNamespace(
|
|
prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "blue")}}],
|
|
sampling_params=SimpleNamespace(height=None, width=None),
|
|
)
|
|
|
|
result = preprocess(request)
|
|
|
|
assert result.sampling_params.height == 672
|
|
assert result.sampling_params.width == 1344
|
|
assert result.prompts[0]["multi_modal_data"]["image"].size == (1344, 672)
|
|
|
|
|
|
def test_ti2v_diffuse_without_image_condition_expands_patch_timesteps() -> None:
|
|
pipeline = _make_ti2v_pipeline()
|
|
latents = torch.zeros(1, 4, 2, 4, 4)
|
|
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([7]),
|
|
prompt_embeds=torch.zeros(1, 2, 3),
|
|
negative_prompt_embeds=torch.zeros(1, 2, 3),
|
|
guidance_scale=3.0,
|
|
dtype=torch.float32,
|
|
attention_kwargs={"a": "b"},
|
|
num_latent_frames=2,
|
|
latent_height=4,
|
|
latent_width=4,
|
|
)
|
|
|
|
positive = calls[0]["positive_kwargs"]
|
|
assert calls[0]["do_true_cfg"] is True
|
|
assert positive["timestep"].shape == (1, 8)
|
|
torch.testing.assert_close(positive["timestep"], torch.full((1, 8), 7, dtype=positive["timestep"].dtype))
|
|
torch.testing.assert_close(positive["hidden_states"], latents)
|
|
torch.testing.assert_close(result, torch.ones_like(latents))
|
|
|
|
|
|
def test_ti2v_prepare_i2v_latents_encodes_condition_and_masks_first_frame() -> None:
|
|
pipeline = _make_ti2v_pipeline()
|
|
latents, latent_condition, first_frame_mask = pipeline.prepare_i2v_latents(
|
|
image=torch.zeros(1, 3, 16, 16),
|
|
batch_size=1,
|
|
num_channels_latents=4,
|
|
height=16,
|
|
width=16,
|
|
num_frames=5,
|
|
dtype=torch.float32,
|
|
device=torch.device("cpu"),
|
|
generator=None,
|
|
latents=torch.zeros(1, 4, 2, 2, 2),
|
|
)
|
|
|
|
torch.testing.assert_close(latents, torch.zeros(1, 4, 2, 2, 2))
|
|
assert latent_condition.shape == (1, 4, 1, 2, 2)
|
|
assert first_frame_mask[:, :, 0].sum() == 0
|
|
assert first_frame_mask[:, :, 1].sum() == 4
|