(self)
| 188 | |
| 189 | class CogVAEDecoder(torch.nn.Module): |
| 190 | def __init__(self): |
| 191 | super().__init__() |
| 192 | self.scaling_factor = 0.7 |
| 193 | self.conv_in = CachedConv3d(16, 512, kernel_size=3, stride=1, padding=(0, 1, 1)) |
| 194 | |
| 195 | self.blocks = torch.nn.ModuleList([ |
| 196 | Resnet3DBlock(512, 512, 16, 32), |
| 197 | Resnet3DBlock(512, 512, 16, 32), |
| 198 | Resnet3DBlock(512, 512, 16, 32), |
| 199 | Resnet3DBlock(512, 512, 16, 32), |
| 200 | Resnet3DBlock(512, 512, 16, 32), |
| 201 | Resnet3DBlock(512, 512, 16, 32), |
| 202 | Upsample3D(512, 512, compress_time=True), |
| 203 | Resnet3DBlock(512, 256, 16, 32), |
| 204 | Resnet3DBlock(256, 256, 16, 32), |
| 205 | Resnet3DBlock(256, 256, 16, 32), |
| 206 | Resnet3DBlock(256, 256, 16, 32), |
| 207 | Upsample3D(256, 256, compress_time=True), |
| 208 | Resnet3DBlock(256, 256, 16, 32), |
| 209 | Resnet3DBlock(256, 256, 16, 32), |
| 210 | Resnet3DBlock(256, 256, 16, 32), |
| 211 | Resnet3DBlock(256, 256, 16, 32), |
| 212 | Upsample3D(256, 256, compress_time=False), |
| 213 | Resnet3DBlock(256, 128, 16, 32), |
| 214 | Resnet3DBlock(128, 128, 16, 32), |
| 215 | Resnet3DBlock(128, 128, 16, 32), |
| 216 | Resnet3DBlock(128, 128, 16, 32), |
| 217 | ]) |
| 218 | |
| 219 | self.norm_out = CogVideoXSpatialNorm3D(128, 16, 32) |
| 220 | self.conv_act = torch.nn.SiLU() |
| 221 | self.conv_out = CachedConv3d(128, 3, kernel_size=3, stride=1, padding=(0, 1, 1)) |
| 222 | |
| 223 | |
| 224 | def forward(self, sample): |
nothing calls this directly
no test coverage detected