(self)
| 277 | |
| 278 | class CogVAEEncoder(torch.nn.Module): |
| 279 | def __init__(self): |
| 280 | super().__init__() |
| 281 | self.scaling_factor = 0.7 |
| 282 | self.conv_in = CachedConv3d(3, 128, kernel_size=3, stride=1, padding=(0, 1, 1)) |
| 283 | |
| 284 | self.blocks = torch.nn.ModuleList([ |
| 285 | Resnet3DBlock(128, 128, None, 32), |
| 286 | Resnet3DBlock(128, 128, None, 32), |
| 287 | Resnet3DBlock(128, 128, None, 32), |
| 288 | Downsample3D(128, 128, compress_time=True), |
| 289 | Resnet3DBlock(128, 256, None, 32), |
| 290 | Resnet3DBlock(256, 256, None, 32), |
| 291 | Resnet3DBlock(256, 256, None, 32), |
| 292 | Downsample3D(256, 256, compress_time=True), |
| 293 | Resnet3DBlock(256, 256, None, 32), |
| 294 | Resnet3DBlock(256, 256, None, 32), |
| 295 | Resnet3DBlock(256, 256, None, 32), |
| 296 | Downsample3D(256, 256, compress_time=False), |
| 297 | Resnet3DBlock(256, 512, None, 32), |
| 298 | Resnet3DBlock(512, 512, None, 32), |
| 299 | Resnet3DBlock(512, 512, None, 32), |
| 300 | Resnet3DBlock(512, 512, None, 32), |
| 301 | Resnet3DBlock(512, 512, None, 32), |
| 302 | ]) |
| 303 | |
| 304 | self.norm_out = torch.nn.GroupNorm(32, 512, eps=1e-06, affine=True) |
| 305 | self.conv_act = torch.nn.SiLU() |
| 306 | self.conv_out = CachedConv3d(512, 32, kernel_size=3, stride=1, padding=(0, 1, 1)) |
| 307 | |
| 308 | |
| 309 | def forward(self, sample): |
nothing calls this directly
no test coverage detected