(self, in_channels, out_channels, spatial_norm_dim, groups, eps=1e-6, use_conv_shortcut=False)
| 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): |
nothing calls this directly
no test coverage detected