The final layer of video DiT.
| 857 | |
| 858 | |
| 859 | class FinalLayer(nn.Module): |
| 860 | """ |
| 861 | The final layer of video DiT. |
| 862 | """ |
| 863 | |
| 864 | def __init__( |
| 865 | self, |
| 866 | hidden_size: int, |
| 867 | spatial_patch_size: int, |
| 868 | temporal_patch_size: int, |
| 869 | out_channels: int, |
| 870 | use_adaln_lora: bool = False, |
| 871 | adaln_lora_dim: int = 256, |
| 872 | ): |
| 873 | super().__init__() |
| 874 | self.layer_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 875 | self.linear = nn.Linear( |
| 876 | hidden_size, spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels, bias=False |
| 877 | ) |
| 878 | self.hidden_size = hidden_size |
| 879 | self.n_adaln_chunks = 2 |
| 880 | self.use_adaln_lora = use_adaln_lora |
| 881 | self.adaln_lora_dim = adaln_lora_dim |
| 882 | if use_adaln_lora: |
| 883 | self.adaln_modulation = nn.Sequential( |
| 884 | nn.SiLU(), |
| 885 | nn.Linear(hidden_size, adaln_lora_dim, bias=False), |
| 886 | nn.Linear(adaln_lora_dim, self.n_adaln_chunks * hidden_size, bias=False), |
| 887 | ) |
| 888 | else: |
| 889 | self.adaln_modulation = nn.Sequential( |
| 890 | nn.SiLU(), nn.Linear(hidden_size, self.n_adaln_chunks * hidden_size, bias=False) |
| 891 | ) |
| 892 | |
| 893 | self.init_weights() |
| 894 | |
| 895 | def init_weights(self) -> None: |
| 896 | std = 1.0 / math.sqrt(self.hidden_size) |
| 897 | torch.nn.init.trunc_normal_(self.linear.weight, std=std, a=-3 * std, b=3 * std) |
| 898 | if self.use_adaln_lora: |
| 899 | torch.nn.init.trunc_normal_(self.adaln_modulation[1].weight, std=std, a=-3 * std, b=3 * std) |
| 900 | torch.nn.init.zeros_(self.adaln_modulation[2].weight) |
| 901 | else: |
| 902 | torch.nn.init.zeros_(self.adaln_modulation[1].weight) |
| 903 | |
| 904 | self.layer_norm.reset_parameters() |
| 905 | |
| 906 | def forward( |
| 907 | self, |
| 908 | x_B_T_H_W_D: torch.Tensor, |
| 909 | emb_B_T_D: torch.Tensor, |
| 910 | adaln_lora_B_T_3D: Optional[torch.Tensor] = None, |
| 911 | ): |
| 912 | if self.use_adaln_lora: |
| 913 | assert adaln_lora_B_T_3D is not None |
| 914 | shift_B_T_D, scale_B_T_D = ( |
| 915 | self.adaln_modulation(emb_B_T_D) + adaln_lora_B_T_3D[:, :, : 2 * self.hidden_size] |
| 916 | ).chunk(2, dim=-1) |