| 126 | |
| 127 | |
| 128 | class Resnet3DBlock(torch.nn.Module): |
| 129 | def __init__(self, in_channels, out_channels, spatial_norm_dim, groups, eps=1e-6, use_conv_shortcut=False): |
| 130 | super().__init__() |
| 131 | self.nonlinearity = torch.nn.SiLU() |
| 132 | if spatial_norm_dim is None: |
| 133 | self.norm1 = torch.nn.GroupNorm(num_channels=in_channels, num_groups=groups, eps=eps) |
| 134 | self.norm2 = torch.nn.GroupNorm(num_channels=out_channels, num_groups=groups, eps=eps) |
| 135 | else: |
| 136 | self.norm1 = CogVideoXSpatialNorm3D(in_channels, spatial_norm_dim, groups) |
| 137 | self.norm2 = CogVideoXSpatialNorm3D(out_channels, spatial_norm_dim, groups) |
| 138 | |
| 139 | self.conv1 = CachedConv3d(in_channels, out_channels, kernel_size=3, padding=(0, 1, 1)) |
| 140 | |
| 141 | self.conv2 = CachedConv3d(out_channels, out_channels, kernel_size=3, padding=(0, 1, 1)) |
| 142 | |
| 143 | if in_channels != out_channels: |
| 144 | if use_conv_shortcut: |
| 145 | self.conv_shortcut = CachedConv3d(in_channels, out_channels, kernel_size=3, padding=(0, 1, 1)) |
| 146 | else: |
| 147 | self.conv_shortcut = torch.nn.Conv3d(in_channels, out_channels, kernel_size=1) |
| 148 | else: |
| 149 | self.conv_shortcut = lambda x: x |
| 150 | |
| 151 | |
| 152 | def forward(self, hidden_states, zq): |
| 153 | residual = hidden_states |
| 154 | |
| 155 | hidden_states = self.norm1(hidden_states, zq) if isinstance(self.norm1, CogVideoXSpatialNorm3D) else self.norm1(hidden_states) |
| 156 | hidden_states = self.nonlinearity(hidden_states) |
| 157 | hidden_states = self.conv1(hidden_states) |
| 158 | |
| 159 | hidden_states = self.norm2(hidden_states, zq) if isinstance(self.norm2, CogVideoXSpatialNorm3D) else self.norm2(hidden_states) |
| 160 | hidden_states = self.nonlinearity(hidden_states) |
| 161 | hidden_states = self.conv2(hidden_states) |
| 162 | |
| 163 | hidden_states = hidden_states + self.conv_shortcut(residual) |
| 164 | |
| 165 | return hidden_states |
| 166 | |
| 167 | |
| 168 | |