Zhou TaichangandYuanheng Zhao
7f2b1e9386
Signed-off-by: tzhouam <tzhouam@connect.ust.hk> Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk> Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com> Co-authored-by: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com>
19 lines
703 B
Python
19 lines
703 B
Python
# 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."""
|
|
|
|
# NOTE: iterate over a snapshot of sys.modules. `hasattr` can trigger lazy
|
|
# submodule imports (e.g. transformers' `_LazyModule.__getattr__`), which
|
|
# mutate sys.modules during iteration and raise
|
|
# `RuntimeError: dictionary changed size during iteration`.
|
|
for module_name, module in list(sys.modules.items()):
|
|
if hasattr(module, "WanRMS_norm"):
|
|
setattr(module, "WanRMS_norm", RMSNormVAE)
|