(self, f_channels, zq_channels, groups)
| 102 | |
| 103 | class CogVideoXSpatialNorm3D(torch.nn.Module): |
| 104 | def __init__(self, f_channels, zq_channels, groups): |
| 105 | super().__init__() |
| 106 | self.norm_layer = torch.nn.GroupNorm(num_channels=f_channels, num_groups=groups, eps=1e-6, affine=True) |
| 107 | self.conv_y = torch.nn.Conv3d(zq_channels, f_channels, kernel_size=1, stride=1) |
| 108 | self.conv_b = torch.nn.Conv3d(zq_channels, f_channels, kernel_size=1, stride=1) |
| 109 | |
| 110 | |
| 111 | def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor: |