# 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 math import torch import torch.nn as nn from hyvideo.commons import to_2tuple class PatchEmbed(nn.Module): """2D Image to Patch Embedding Image to Patch Embedding using Conv2d A convolution based approach to patchifying a 2D image w/ embedding projection. Based on the impl in https://github.com/google-research/vision_transformer Hacked together by / Copyright 2020 Ross Wightman Remove the _assert function in forward function to be compatible with multi-resolution images. """ def __init__( self, patch_size=16, in_chans=3, embed_dim=768, is_reshape_temporal_channels=False, concat_condition=True, norm_layer=None, flatten=True, bias=True, dtype=None, device=None, ): factory_kwargs = {"dtype": dtype, "device": device} super().__init__() patch_size = to_2tuple(patch_size) self.patch_size = patch_size self.flatten = flatten # Only support concat mode (multitask mask training) orig_in_chans = in_chans if concat_condition: if is_reshape_temporal_channels: in_chans = in_chans + in_chans // 2 + 1 else: in_chans = in_chans * 2 + 1 self.proj = nn.Conv3d( in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias, **factory_kwargs, ) nn.init.xavier_uniform_( self.proj.weight[:, :orig_in_chans].view( self.proj.weight[:, :orig_in_chans].size(0), -1 ) ) # Special initialization for concat mode nn.init.zeros_( self.proj.weight[:, orig_in_chans:].view( self.proj.weight[:, orig_in_chans:].size(0), -1 ) ) if bias: nn.init.zeros_(self.proj.bias) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() def forward(self, x): x = self.proj(x) if self.flatten: x = x.flatten(2).transpose(1, 2) # BCHW -> BNC x = self.norm(x) return x class TextProjection(nn.Module): """ Projects text embeddings. Also handles dropout for classifier-free guidance. Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py """ def __init__(self, in_channels, hidden_size, act_layer, dtype=None, device=None): factory_kwargs = {"dtype": dtype, "device": device} super().__init__() self.linear_1 = nn.Linear( in_features=in_channels, out_features=hidden_size, bias=True, **factory_kwargs, ) self.act_1 = act_layer() self.linear_2 = nn.Linear( in_features=hidden_size, out_features=hidden_size, bias=True, **factory_kwargs, ) def forward(self, caption): hidden_states = self.linear_1(caption) hidden_states = self.act_1(hidden_states) hidden_states = self.linear_2(hidden_states) return hidden_states class VisionProjection(torch.nn.Module): def __init__(self, input_dim, output_dim): super().__init__() self.proj = torch.nn.Sequential( torch.nn.LayerNorm(input_dim), torch.nn.Linear(input_dim, input_dim), torch.nn.GELU(), torch.nn.Linear(input_dim, output_dim), torch.nn.LayerNorm(output_dim), ) def forward(self, vision_embeds): return self.proj(vision_embeds) class ClipVisionProjection(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.up = nn.Linear(in_channels, out_channels * 3) self.down = nn.Linear(out_channels * 3, out_channels) torch.nn.init.zeros_(self.down.weight) torch.nn.init.zeros_(self.down.bias) def forward(self, x): projected_x = self.down(nn.functional.silu(self.up(x))) return projected_x def timestep_embedding(t, dim, max_period=10000): """ Create sinusoidal timestep embeddings. Args: t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. dim (int): the dimension of the output. max_period (int): controls the minimum frequency of the embeddings. Returns: embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py """ half = dim // 2 freqs = torch.exp( -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half ).to(device=t.device) args = t[:, None].float() * freqs[None] embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) if dim % 2: embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) return embedding class TimestepEmbedder(nn.Module): """ Embeds scalar timesteps into vector representations. """ def __init__( self, hidden_size, act_layer, frequency_embedding_size=256, max_period=10000, out_size=None, dtype=None, device=None, ): factory_kwargs = {"dtype": dtype, "device": device} super().__init__() self.frequency_embedding_size = frequency_embedding_size self.max_period = max_period if out_size is None: out_size = hidden_size self.mlp = nn.Sequential( nn.Linear( frequency_embedding_size, hidden_size, bias=True, **factory_kwargs ), act_layer(), nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), ) nn.init.normal_(self.mlp[0].weight, std=0.02) nn.init.normal_(self.mlp[2].weight, std=0.02) def forward(self, t): t_freq = timestep_embedding( t, self.frequency_embedding_size, self.max_period ).type(self.mlp[0].weight.dtype) t_emb = self.mlp(t_freq) return t_emb