64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
|
|
def __init__(self, *, dtype: torch.dtype):
|
|
torch.nn.Module.__init__(self)
|
|
self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
|
|
|
|
|
|
def test_wan_vae_execution_context_handles_fp32():
|
|
model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
|
|
with model._execution_context():
|
|
output = model.dummy_weight + 1
|
|
assert output.dtype == torch.float32
|
|
|
|
|
|
def test_wan_vae_execution_context_handles_bf16():
|
|
model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
|
|
with model._execution_context():
|
|
output = model.dummy_weight + 1
|
|
assert output.dtype == torch.bfloat16
|
|
|
|
|
|
def test_wan_vae_execution_context_uses_platform_autocast(mocker):
|
|
sentinel = object()
|
|
platform = mocker.Mock()
|
|
platform.create_autocast_context.return_value = sentinel
|
|
mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
|
|
|
|
model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
|
|
|
|
assert model._execution_context() is sentinel
|
|
platform.create_autocast_context.assert_called_once_with(
|
|
device_type=model.dummy_weight.device.type,
|
|
dtype=torch.bfloat16,
|
|
enabled=True,
|
|
)
|
|
|
|
|
|
def test_prepare_pipeline_installs_data_movement_before_compile(monkeypatch):
|
|
model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
|
|
installed = []
|
|
monkeypatch.setattr(
|
|
"vllm_omni.diffusion.offloader.module_collector.ModuleDiscovery.discover",
|
|
lambda pipeline: SimpleNamespace(vaes=[model]),
|
|
)
|
|
monkeypatch.setattr(
|
|
wan_vae_module,
|
|
"install_wan_vae_data_movement",
|
|
lambda vae: installed.append(vae),
|
|
)
|
|
|
|
wan_vae_module.prepare_pipeline_wan_vae_data_movement(torch.nn.Identity())
|
|
|
|
assert installed == [model]
|