a473ae78c3
Signed-off-by: Asthenia <asthenia0412@gmail.com> Signed-off-by: Gao Han <hgaoaf@connect.ust.hk> Co-authored-by: Asthenia <asthenia0412@gmail.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Isolate the upstream numerical reference from Wan's global norm patch."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
|
|
import pytest
|
|
from diffusers.models.autoencoders import autoencoder_kl_wan
|
|
from torch import nn
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def original_wan_rms_norm() -> type[nn.Module]:
|
|
# Wan's package import replaces the public class, including aliases held
|
|
# by already collected tests. Load the installed upstream source under a
|
|
# private name instead of reloading (and changing) the live module/classes.
|
|
spec = importlib.util.spec_from_file_location(
|
|
"diffusers.models.autoencoders._wan_fastpath_reference", autoencoder_kl_wan.__file__
|
|
)
|
|
assert spec is not None and spec.loader is not None
|
|
reference = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(reference)
|
|
return reference.WanRMS_norm
|
|
|
|
|
|
@pytest.fixture
|
|
def unpatched_wan_rms_norm(monkeypatch, request, original_wan_rms_norm):
|
|
# Import the patching package before temporarily restoring the reference,
|
|
# so importing RMSNormVAE in a test cannot change it halfway through.
|
|
from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
|
|
|
|
assert original_wan_rms_norm is not RMSNormVAE
|
|
monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", original_wan_rms_norm)
|
|
if hasattr(request.module, "WanRMS_norm"):
|
|
monkeypatch.setattr(request.module, "WanRMS_norm", original_wan_rms_norm)
|