| 480 | """ |
| 481 | |
| 482 | def __init__( |
| 483 | self, |
| 484 | in_dim: int, |
| 485 | out_dim: Optional[int] = None, |
| 486 | dropout: float = 0.0, |
| 487 | norm_num_groups: int = 32, |
| 488 | ): |
| 489 | super().__init__() |
| 490 | out_dim = out_dim or in_dim |
| 491 | self.in_dim = in_dim |
| 492 | self.out_dim = out_dim |
| 493 | |
| 494 | # conv layers |
| 495 | self.conv1 = nn.Sequential( |
| 496 | nn.GroupNorm(norm_num_groups, in_dim), |
| 497 | nn.SiLU(), |
| 498 | nn.Conv3d(in_dim, out_dim, (3, 1, 1), padding=(1, 0, 0)), |
| 499 | ) |
| 500 | self.conv2 = nn.Sequential( |
| 501 | nn.GroupNorm(norm_num_groups, out_dim), |
| 502 | nn.SiLU(), |
| 503 | nn.Dropout(dropout), |
| 504 | nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)), |
| 505 | ) |
| 506 | self.conv3 = nn.Sequential( |
| 507 | nn.GroupNorm(norm_num_groups, out_dim), |
| 508 | nn.SiLU(), |
| 509 | nn.Dropout(dropout), |
| 510 | nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)), |
| 511 | ) |
| 512 | self.conv4 = nn.Sequential( |
| 513 | nn.GroupNorm(norm_num_groups, out_dim), |
| 514 | nn.SiLU(), |
| 515 | nn.Dropout(dropout), |
| 516 | nn.Conv3d(out_dim, in_dim, (3, 1, 1), padding=(1, 0, 0)), |
| 517 | ) |
| 518 | |
| 519 | # zero out the last layer params,so the conv block is identity |
| 520 | nn.init.zeros_(self.conv4[-1].weight) |
| 521 | nn.init.zeros_(self.conv4[-1].bias) |
| 522 | |
| 523 | def forward(self, hidden_states: torch.Tensor, num_frames: int = 1) -> torch.Tensor: |
| 524 | hidden_states = ( |