600 lines
21 KiB
Python
600 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import importlib
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from torch import nn
|
|
|
|
from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange
|
|
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler
|
|
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _cpu_platform(monkeypatch):
|
|
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
|
|
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
|
|
|
|
|
|
class _StubTransformer(nn.Module):
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return torch.float32
|
|
|
|
|
|
class _StubTextEncoder(nn.Module):
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
return torch.float32
|
|
|
|
|
|
class _StubVaeConfig:
|
|
latents_mean = [0.0, 0.0, 0.0, 0.0]
|
|
latents_std = [1.0, 1.0, 1.0, 1.0]
|
|
z_dim = 4
|
|
|
|
|
|
class _StubVae(nn.Module):
|
|
dtype = torch.float32
|
|
config = _StubVaeConfig()
|
|
|
|
def decode(self, latents, return_dict=False):
|
|
del return_dict
|
|
batch, _, frames, height, width = latents.shape
|
|
return (torch.zeros(batch, 3, frames, height, width),)
|
|
|
|
|
|
class _StubScheduler:
|
|
def __init__(self, timesteps: list[int]) -> None:
|
|
self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
|
|
self.config = SimpleNamespace(num_train_timesteps=1000)
|
|
self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = []
|
|
|
|
def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
|
|
self.set_timesteps_calls.append((num_steps, device, shift))
|
|
|
|
|
|
@contextmanager
|
|
def _noop_progress_bar(*args, **kwargs):
|
|
del args, kwargs
|
|
|
|
class _Bar:
|
|
def update(self) -> None:
|
|
return None
|
|
|
|
yield _Bar()
|
|
|
|
|
|
def _stub_encode_prompt(
|
|
prompt,
|
|
negative_prompt=None,
|
|
do_classifier_free_guidance=True,
|
|
num_videos_per_prompt=1,
|
|
max_sequence_length=512,
|
|
device=None,
|
|
dtype=None,
|
|
):
|
|
del negative_prompt, do_classifier_free_guidance, device, dtype
|
|
batch_size = 1 if isinstance(prompt, str) else len(prompt)
|
|
n = batch_size * num_videos_per_prompt
|
|
hidden_size = 8
|
|
prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size)
|
|
return prompt_embeds, None
|
|
|
|
|
|
def _make_pipeline() -> Wan22Pipeline:
|
|
pipeline = object.__new__(Wan22Pipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline.device = torch.device("cpu")
|
|
pipeline.transformer = _StubTransformer()
|
|
pipeline.transformer_2 = None
|
|
pipeline.text_encoder = _StubTextEncoder()
|
|
pipeline.vae = _StubVae()
|
|
pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4)
|
|
pipeline.scheduler = _StubScheduler([9, 5])
|
|
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
|
|
pipeline._sample_solver = "unipc"
|
|
pipeline._flow_shift = 5.0
|
|
pipeline.vae_scale_factor_temporal = 4
|
|
pipeline.vae_scale_factor_spatial = 8
|
|
pipeline.boundary_ratio = 0.875
|
|
pipeline.expand_timesteps = False
|
|
pipeline.is_dmd = 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 = _stub_encode_prompt # type: ignore[method-assign]
|
|
pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32)
|
|
pipeline.progress_bar = _noop_progress_bar
|
|
return pipeline
|
|
|
|
|
|
def _make_sampling(**overrides):
|
|
values: dict[str, object] = {
|
|
"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"),
|
|
[
|
|
({}, 4.0, 4.0),
|
|
({"guidance_scale": 0.0}, 0.0, 0.0),
|
|
({"guidance_scale": 1.0}, 1.0, 1.0),
|
|
({"guidance_scale": 3.0, "guidance_scale_2": 5.0}, 3.0, 5.0),
|
|
],
|
|
)
|
|
def test_forward_delegates_denoising_to_diffuse(
|
|
sampling_params_kwargs: dict[str, float],
|
|
expected_low: float,
|
|
expected_high: float,
|
|
) -> None:
|
|
pipeline = _make_pipeline()
|
|
captured: dict[str, object] = {}
|
|
|
|
def _fake_diffuse(**kwargs):
|
|
captured.update(kwargs)
|
|
return kwargs["latents"] + 1
|
|
|
|
pipeline.diffuse = _fake_diffuse # type: ignore[method-assign]
|
|
|
|
mock_req = OmniDiffusionRequest(
|
|
prompt="prompt",
|
|
request_id="test-req",
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
num_frames=1,
|
|
num_inference_steps=2,
|
|
max_sequence_length=32,
|
|
output_type="latent",
|
|
**sampling_params_kwargs,
|
|
),
|
|
)
|
|
batch = DiffusionRequestBatch(requests=[mock_req])
|
|
|
|
outputs = pipeline.forward(batch)
|
|
|
|
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
|
|
assert captured["guidance_high"] == expected_high
|
|
assert captured["boundary_timestep"] == pytest.approx(875.0)
|
|
assert captured["latent_condition"] is None
|
|
assert captured["first_frame_mask"] is None
|
|
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)]
|
|
|
|
|
|
@pytest.mark.parametrize("solver", ["unipc", "euler"])
|
|
def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None:
|
|
pipeline = _make_pipeline()
|
|
pipeline.diffuse = lambda **kwargs: kwargs["latents"]
|
|
for shift in (3.0, 12.0, 5.0):
|
|
sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift})
|
|
request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling)
|
|
pipeline.forward(DiffusionRequestBatch(requests=[request]))
|
|
reference = build_wan_scheduler(solver, shift)
|
|
if solver == "unipc":
|
|
reference.set_timesteps(5, device="cpu", shift=shift)
|
|
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
|
|
else:
|
|
reference.set_timesteps(5, device="cpu")
|
|
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
|
|
|
|
|
|
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_emits_request_local_typed_media_after_vae_decode() -> None:
|
|
pipeline = _make_pipeline()
|
|
|
|
def _fake_diffuse(
|
|
*,
|
|
latents,
|
|
timesteps,
|
|
prompt_embeds,
|
|
negative_prompt_embeds,
|
|
guidance_low,
|
|
guidance_high,
|
|
boundary_timestep,
|
|
dtype,
|
|
attention_kwargs,
|
|
latent_condition,
|
|
first_frame_mask,
|
|
generator,
|
|
):
|
|
del (
|
|
timesteps,
|
|
prompt_embeds,
|
|
negative_prompt_embeds,
|
|
guidance_low,
|
|
guidance_high,
|
|
boundary_timestep,
|
|
dtype,
|
|
attention_kwargs,
|
|
latent_condition,
|
|
first_frame_mask,
|
|
generator,
|
|
)
|
|
return torch.zeros_like(latents)
|
|
|
|
pipeline.diffuse = _fake_diffuse # type: ignore[method-assign]
|
|
batch = DiffusionRequestBatch(
|
|
requests=[
|
|
OmniDiffusionRequest(
|
|
prompt="prompt",
|
|
request_id="request-0",
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
num_frames=1,
|
|
num_inference_steps=2,
|
|
max_sequence_length=32,
|
|
output_type="np",
|
|
),
|
|
)
|
|
]
|
|
)
|
|
|
|
outputs = pipeline.forward(batch)
|
|
|
|
assert len(outputs) == 1
|
|
assert outputs[0].output is None
|
|
assert outputs[0].media is not None
|
|
assert outputs[0].media.prepared_for_transport is False
|
|
assert outputs[0].media.video.tensor.shape == (1, 3, 1, 8, 8)
|
|
assert outputs[0].media.video.spec.layout is VideoTensorLayout.BCTHW
|
|
assert outputs[0].media.video.spec.encoding is VideoTensorEncoding.NORMALIZED_FLOAT
|
|
assert outputs[0].media.video.spec.value_range is VideoValueRange.NEGATIVE_ONE_TO_ONE
|
|
|
|
|
|
def test_forward_keeps_legacy_output_on_non_owner_vae_rank() -> None:
|
|
# Distributed VAE decode uses broadcast_result=False, so non-owner ranks get
|
|
# an empty placeholder instead of the full video. Wrapping that as typed media
|
|
# would fail split_diffusion_output_by_request's batch check on every non-owner
|
|
# rank, so the pipeline must keep the placeholder on the legacy output field.
|
|
pipeline = _make_pipeline()
|
|
pipeline.vae.decode = lambda latents, return_dict=False: (torch.empty(0),) # type: ignore[assignment]
|
|
pipeline.diffuse = lambda **kwargs: torch.zeros_like(kwargs["latents"]) # type: ignore[method-assign]
|
|
|
|
batch = DiffusionRequestBatch(
|
|
requests=[
|
|
OmniDiffusionRequest(
|
|
prompt="prompt",
|
|
request_id="request-0",
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
num_frames=1,
|
|
num_inference_steps=2,
|
|
max_sequence_length=32,
|
|
output_type="np",
|
|
),
|
|
)
|
|
]
|
|
)
|
|
|
|
outputs = pipeline.forward(batch)
|
|
|
|
assert len(outputs) == 1
|
|
assert outputs[0].media is None
|
|
assert outputs[0].output is not None
|
|
assert outputs[0].output.numel() == 0
|
|
|
|
|
|
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)
|
|
timesteps = torch.tensor([7, 3], dtype=torch.int64)
|
|
prompt_embeds = torch.randn(1, 8)
|
|
|
|
predict_calls: list[dict[str, object]] = []
|
|
scheduler_calls: list[tuple[float, int, float, bool]] = []
|
|
|
|
def _fake_predict_noise_maybe_with_cfg(**kwargs):
|
|
predict_calls.append(kwargs)
|
|
timestep = kwargs["positive_kwargs"]["timestep"]
|
|
assert isinstance(timestep, torch.Tensor)
|
|
return torch.full_like(latents, float(timestep[0].item()))
|
|
|
|
def _fake_scheduler_step_maybe_with_cfg(noise_pred, t, current_latents, do_true_cfg):
|
|
scheduler_calls.append(
|
|
(float(noise_pred[0, 0, 0, 0, 0]), int(t.item()), float(current_latents.sum()), do_true_cfg)
|
|
)
|
|
return current_latents + noise_pred
|
|
|
|
pipeline.predict_noise_maybe_with_cfg = _fake_predict_noise_maybe_with_cfg # type: ignore[method-assign]
|
|
pipeline.scheduler_step_maybe_with_cfg = _fake_scheduler_step_maybe_with_cfg # type: ignore[method-assign]
|
|
|
|
result = pipeline.diffuse(
|
|
latents=latents,
|
|
timesteps=timesteps,
|
|
prompt_embeds=prompt_embeds,
|
|
negative_prompt_embeds=None,
|
|
guidance_low=1.0,
|
|
guidance_high=2.0,
|
|
boundary_timestep=5.0,
|
|
dtype=torch.float32,
|
|
attention_kwargs={},
|
|
)
|
|
|
|
assert len(predict_calls) == 2
|
|
assert predict_calls[0]["true_cfg_scale"] == 1.0
|
|
assert predict_calls[1]["true_cfg_scale"] == 2.0
|
|
assert scheduler_calls == [
|
|
(7.0, 7, 0.0, False),
|
|
(3.0, 3, 28.0, False),
|
|
]
|
|
assert torch.equal(result, torch.full_like(latents, 10.0))
|
|
|
|
|
|
class _StubDMDScheduler:
|
|
def __init__(self) -> None:
|
|
self.predict_clean_calls: list[tuple[float, float, float]] = []
|
|
self.add_noise_calls: list[tuple[float, float, float]] = []
|
|
|
|
def predict_clean(self, model_output, sample, timestep):
|
|
self.predict_clean_calls.append((float(model_output.mean()), float(sample.mean()), float(timestep)))
|
|
return sample - model_output
|
|
|
|
def add_noise(self, clean_sample, noise, timestep):
|
|
self.add_noise_calls.append((float(clean_sample.mean()), float(noise.mean()), float(timestep)))
|
|
return clean_sample + 10.0
|
|
|
|
|
|
def test_diffuse_dmd_predicts_clean_and_renoises_between_steps(monkeypatch) -> None:
|
|
pipeline = _make_pipeline()
|
|
pipeline.is_dmd = True
|
|
pipeline.scheduler = _StubDMDScheduler()
|
|
latents = torch.zeros((1, 1, 1, 1, 1), dtype=torch.float32)
|
|
timesteps = torch.tensor([1000.0, 757.0, 522.0])
|
|
|
|
pipeline.predict_noise_maybe_with_cfg = lambda **kwargs: torch.ones_like(latents) # type: ignore[method-assign]
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.randn_tensor",
|
|
lambda *args, **kwargs: torch.full(args[0], 2.0, dtype=kwargs["dtype"]),
|
|
)
|
|
|
|
result = pipeline.diffuse(
|
|
latents=latents,
|
|
timesteps=timesteps,
|
|
prompt_embeds=torch.zeros(1, 8),
|
|
negative_prompt_embeds=None,
|
|
guidance_low=1.0,
|
|
guidance_high=1.0,
|
|
boundary_timestep=None,
|
|
dtype=torch.float32,
|
|
attention_kwargs={},
|
|
generator=torch.Generator(device="cpu").manual_seed(1),
|
|
)
|
|
|
|
assert pipeline.scheduler.predict_clean_calls == [
|
|
(1.0, 0.0, 1000.0),
|
|
(1.0, 9.0, 757.0),
|
|
(1.0, 18.0, 522.0),
|
|
]
|
|
assert pipeline.scheduler.add_noise_calls == [
|
|
(-1.0, 2.0, 757.0),
|
|
(8.0, 2.0, 522.0),
|
|
]
|
|
torch.testing.assert_close(result, torch.tensor([[[[[17.0]]]]]))
|
|
|
|
|
|
def _make_gate_loading_pipeline():
|
|
pipeline = Wan22Pipeline.__new__(Wan22Pipeline)
|
|
nn.Module.__init__(pipeline)
|
|
gate = WanSelfAttention.__new__(WanSelfAttention)
|
|
nn.Module.__init__(gate)
|
|
gate.to_gate_compress = nn.Linear(1, 1)
|
|
pipeline.gate_holder = gate
|
|
return pipeline, gate
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("module_name", "class_name"),
|
|
[
|
|
("pipeline_wan2_2", "Wan22Pipeline"),
|
|
("pipeline_wan2_2_i2v", "Wan22I2VPipeline"),
|
|
("pipeline_wan2_2_s2v", "Wan22S2VPipeline"),
|
|
("pipeline_wan2_2_vace", "Wan22VACEPipeline"),
|
|
],
|
|
)
|
|
def test_wan_pipeline_loaders_share_optional_gate_cleanup(monkeypatch, module_name, class_name) -> None:
|
|
module = importlib.import_module(f"vllm_omni.diffusion.models.wan2_2.{module_name}")
|
|
pipeline_cls = getattr(module, class_name)
|
|
pipeline = pipeline_cls.__new__(pipeline_cls)
|
|
expected = {"loaded"}
|
|
|
|
def fake_loader(model, weights):
|
|
assert model is pipeline
|
|
assert list(weights) == [("weight", torch.ones(1))]
|
|
return expected
|
|
|
|
monkeypatch.setattr(module, "load_wan_weights_with_optional_gate", fake_loader)
|
|
|
|
assert pipeline_cls.load_weights(pipeline, iter((("weight", torch.ones(1)),))) is expected
|
|
|
|
|
|
def test_load_weights_removes_unloaded_vsa_gate(monkeypatch) -> None:
|
|
pipeline, gate = _make_gate_loading_pipeline()
|
|
|
|
class _Loader:
|
|
def __init__(self, model):
|
|
del model
|
|
|
|
def load_weights(self, weights):
|
|
return {name for name, _ in weights}
|
|
|
|
monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
|
|
pipeline.load_weights(iter((("other.weight", torch.ones(1)),)))
|
|
|
|
assert pipeline.has_gate_compress_weights is False
|
|
assert gate.to_gate_compress is None
|
|
|
|
|
|
def test_load_weights_keeps_trained_vsa_gate(monkeypatch) -> None:
|
|
pipeline, gate = _make_gate_loading_pipeline()
|
|
original_gate = gate.to_gate_compress
|
|
|
|
class _Loader:
|
|
def __init__(self, model):
|
|
del model
|
|
|
|
def load_weights(self, weights):
|
|
return {name for name, _ in weights}
|
|
|
|
monkeypatch.setattr("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2.AutoWeightsLoader", _Loader)
|
|
pipeline.load_weights(iter((("gate_holder.to_gate_compress.weight", torch.ones(1)),)))
|
|
|
|
assert pipeline.has_gate_compress_weights is True
|
|
assert gate.to_gate_compress is original_gate
|