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. | |
| from typing import Callable | |
| import torch | |
| import torch.nn as nn | |
| class ModulateDiT(nn.Module): | |
| """Modulation layer for DiT.""" | |
| def __init__( | |
| self, | |
| hidden_size: int, | |
| factor: int, | |
| act_layer: Callable, | |
| dtype=None, | |
| device=None, | |
| ): | |
| factory_kwargs = {"dtype": dtype, "device": device} | |
| super().__init__() | |
| self.act = act_layer() | |
| self.linear = nn.Linear(hidden_size, factor * hidden_size, bias=True, **factory_kwargs) | |
| # Zero-initialize the modulation | |
| nn.init.zeros_(self.linear.weight) | |
| nn.init.zeros_(self.linear.bias) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.linear(self.act(x)) | |
| def modulate(x, shift=None, scale=None): | |
| """modulate by shift and scale | |
| Args: | |
| x (torch.Tensor): input tensor. | |
| shift (torch.Tensor, optional): shift tensor. Defaults to None. | |
| scale (torch.Tensor, optional): scale tensor. Defaults to None. | |
| Returns: | |
| torch.Tensor: the output tensor after modulate. | |
| """ | |
| if scale is None and shift is None: | |
| return x | |
| elif shift is None: | |
| return x * (1 + scale.unsqueeze(1)) | |
| elif scale is None: | |
| return x + shift.unsqueeze(1) | |
| else: | |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) | |
| def apply_gate(x, gate=None, tanh=False): | |
| """AI is creating summary for apply_gate | |
| Args: | |
| x (torch.Tensor): input tensor. | |
| gate (torch.Tensor, optional): gate tensor. Defaults to None. | |
| tanh (bool, optional): whether to use tanh function. Defaults to False. | |
| Returns: | |
| torch.Tensor: the output tensor after apply gate. | |
| """ | |
| if gate is None: | |
| return x | |
| if tanh: | |
| return x * gate.unsqueeze(1).tanh() | |
| else: | |
| return x * gate.unsqueeze(1) | |
| def ckpt_wrapper(module): | |
| def ckpt_forward(*inputs): | |
| outputs = module(*inputs) | |
| return outputs | |
| return ckpt_forward | |