2026-04-24 02:25:10 +08:00
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
|
|
|
|
|
|
import sys
|
|
|
|
|
|
|
|
|
|
from vllm_omni.diffusion.layers.norm import RMSNormVAE
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def patch_wan_rms_norm():
|
|
|
|
|
"""Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
|
|
|
|
|
|
2026-05-07 11:54:16 +08:00
|
|
|
# Probe `__dict__` directly instead of `hasattr`: the latter triggers
|
|
|
|
|
# custom `__getattr__` hooks (e.g. transformers' image_processing alias
|
|
|
|
|
# modules), which both emit spurious deprecation warnings and can lazily
|
|
|
|
|
# import submodules, mutating sys.modules mid-iteration.
|
|
|
|
|
for module in list(sys.modules.values()):
|
|
|
|
|
module_dict = getattr(module, "__dict__", None)
|
|
|
|
|
if module_dict is not None and "WanRMS_norm" in module_dict:
|
|
|
|
|
module_dict["WanRMS_norm"] = RMSNormVAE
|
