Spaces:
Running on Zero
Running on Zero
| # Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE | |
| # | |
| # Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any | |
| # output and results therefrom are provided "AS IS" without any express or implied warranties of | |
| # any kind including any warranties of title, merchantability, noninfringement, course of dealing, | |
| # usage of trade, or fitness for a particular purpose. You are solely responsible for determining the | |
| # appropriateness of using, reproducing, modifying, performing, displaying or distributing any of | |
| # the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a | |
| # third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise | |
| # of rights and permissions under this agreement. | |
| # See the License for the specific language governing permissions and limitations under the License. | |
| import torch | |
| import torch.nn as nn | |
| class RMSNorm(nn.Module): | |
| def __init__( | |
| self, | |
| dim: int, | |
| elementwise_affine=True, | |
| eps: float = 1e-6, | |
| device=None, | |
| dtype=None, | |
| ): | |
| """ | |
| Initialize the RMSNorm normalization layer. | |
| Args: | |
| dim (int): The dimension of the input tensor. | |
| eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. | |
| Attributes: | |
| eps (float): A small value added to the denominator for numerical stability. | |
| weight (nn.Parameter): Learnable scaling parameter. | |
| """ | |
| factory_kwargs = {"device": device, "dtype": dtype} | |
| super().__init__() | |
| self.eps = eps | |
| if elementwise_affine: | |
| self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) | |
| def _norm(self, x): | |
| """ | |
| Apply the RMSNorm normalization to the input tensor. | |
| Args: | |
| x (torch.Tensor): The input tensor. | |
| Returns: | |
| torch.Tensor: The normalized tensor. | |
| """ | |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| def reset_parameters(self): | |
| if hasattr(self, "weight"): | |
| self.weight.fill_(1) | |
| def forward(self, x): | |
| """ | |
| Forward pass through the RMSNorm layer. | |
| Args: | |
| x (torch.Tensor): The input tensor. | |
| Returns: | |
| torch.Tensor: The output tensor after applying RMSNorm. | |
| """ | |
| output = self._norm(x.float()).type_as(x) | |
| if hasattr(self, "weight"): | |
| output = output * self.weight | |
| return output | |
| def get_norm_layer(norm_layer): | |
| """ | |
| Get the normalization layer. | |
| Args: | |
| norm_layer (str): The type of normalization layer. | |
| Returns: | |
| norm_layer (nn.Module): The normalization layer. | |
| """ | |
| if norm_layer == "layer": | |
| return nn.LayerNorm | |
| elif norm_layer == "rms": | |
| return RMSNorm | |
| else: | |
| raise NotImplementedError(f"Norm layer {norm_layer} is not implemented") | |