(self,
dim=128,
z_dim=4,
dim_mult=[1, 2, 4, 4],
num_res_blocks=2,
attn_scales=[],
temperal_downsample=[True, True, False],
dropout=0.0)
| 483 | class WanVAE_(nn.Module): |
| 484 | |
| 485 | def __init__(self, |
| 486 | dim=128, |
| 487 | z_dim=4, |
| 488 | dim_mult=[1, 2, 4, 4], |
| 489 | num_res_blocks=2, |
| 490 | attn_scales=[], |
| 491 | temperal_downsample=[True, True, False], |
| 492 | dropout=0.0): |
| 493 | super().__init__() |
| 494 | self.dim = dim |
| 495 | self.z_dim = z_dim |
| 496 | self.dim_mult = dim_mult |
| 497 | self.num_res_blocks = num_res_blocks |
| 498 | self.attn_scales = attn_scales |
| 499 | self.temperal_downsample = temperal_downsample |
| 500 | self.temperal_upsample = temperal_downsample[::-1] |
| 501 | |
| 502 | # modules |
| 503 | self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, |
| 504 | attn_scales, self.temperal_downsample, dropout) |
| 505 | self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) |
| 506 | self.conv2 = CausalConv3d(z_dim, z_dim, 1) |
| 507 | self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, |
| 508 | attn_scales, self.temperal_upsample, dropout) |
| 509 | |
| 510 | def forward(self, x): |
| 511 | mu, log_var = self.encode(x) |
nothing calls this directly
no test coverage detected