[Refactor][Diffusion] Move single-model layers into their model directories (#5908)

Signed-off-by: congw729 <115451386+congw729@users.noreply.github.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/layers/norm.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from importlib.util import find_spec
 
 import torch
@@ -163,75 +166,3 @@ class RMSNorm(CustomOp):
         out = self.weight.to(torch.float32) * out
         out = out.to(input_dtype)
         return (out, residual) if residual is not None else out
-
-
-class RMSNormVAE(CustomOp):
-    """Root Mean Square Layer Normalization for Channel-First or Last"""
-
-    def __init__(
-        self,
-        dim: int,
-        channel_first: bool = True,
-        images: bool = True,
-        bias: bool = False,
-        epsilon: float = 1e-6,
-    ) -> None:
-        super().__init__()
-        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
-        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
-
-        self.channel_first = channel_first
-        self.scale = dim**0.5
-        self.gamma = nn.Parameter(torch.ones(shape))
-        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
-        self.epsilon = epsilon
-
-        self.gamma_rmsnorm = None
-
-    def forward_cuda(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.forward_native(x)
-
-    def forward_hip(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.forward_native(x)
-
-    def forward_npu(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        import torch_npu
-
-        if self.gamma_rmsnorm is None:
-            self.gamma_rmsnorm = self.gamma.reshape(-1)
-
-        if self.channel_first:
-            x = x.transpose(1, -1)
-            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
-        else:
-            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
-
-        if self.bias is not None:
-            out = out + self.bias
-        return out
-
-    def forward_native(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        out = (
-            F.normalize(
-                x,
-                dim=(1 if self.channel_first else -1),
-                eps=self.epsilon,
-            )
-            * self.scale
-            * self.gamma
-        )
-        if self.bias is not None:
-            out = out + self.bias
-        return out