vllm_omni.diffusion.models.wan2_2.norm ¶

Normalization layer used by the Wan VAE patch.

RMSNormVAE ¶

Bases: CustomOp

Root Mean Square Layer Normalization for Channel-First or Last

bias = nn.Parameter(torch.zeros(shape)) if bias else None

channel_first = channel_first

epsilon = epsilon

gamma = nn.Parameter(torch.ones(shape))

gamma_rmsnorm = None

scale = dim ** 0.5

forward_cuda(x: Tensor) -> Tensor

forward_hip(x: Tensor) -> Tensor

forward_native(x: Tensor) -> Tensor

forward_npu(x: Tensor) -> Tensor