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