(
self,
in_channels: int,
out_channels: int | None = None,
temb_channels: int = 512,
eps: float = 1e-6,
)
| 556 | """ |
| 557 | |
| 558 | def __init__( |
| 559 | self, |
| 560 | in_channels: int, |
| 561 | out_channels: int | None = None, |
| 562 | temb_channels: int = 512, |
| 563 | eps: float = 1e-6, |
| 564 | ): |
| 565 | super().__init__() |
| 566 | self.in_channels = in_channels |
| 567 | out_channels = in_channels if out_channels is None else out_channels |
| 568 | self.out_channels = out_channels |
| 569 | |
| 570 | kernel_size = (3, 1, 1) |
| 571 | padding = [k // 2 for k in kernel_size] |
| 572 | |
| 573 | self.norm1 = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=eps, affine=True) |
| 574 | self.conv1 = nn.Conv3d( |
| 575 | in_channels, |
| 576 | out_channels, |
| 577 | kernel_size=kernel_size, |
| 578 | stride=1, |
| 579 | padding=padding, |
| 580 | ) |
| 581 | |
| 582 | if temb_channels is not None: |
| 583 | self.time_emb_proj = nn.Linear(temb_channels, out_channels) |
| 584 | else: |
| 585 | self.time_emb_proj = None |
| 586 | |
| 587 | self.norm2 = torch.nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=eps, affine=True) |
| 588 | |
| 589 | self.dropout = torch.nn.Dropout(0.0) |
| 590 | self.conv2 = nn.Conv3d( |
| 591 | out_channels, |
| 592 | out_channels, |
| 593 | kernel_size=kernel_size, |
| 594 | stride=1, |
| 595 | padding=padding, |
| 596 | ) |
| 597 | |
| 598 | self.nonlinearity = get_activation("silu") |
| 599 | |
| 600 | self.use_in_shortcut = self.in_channels != out_channels |
| 601 | |
| 602 | self.conv_shortcut = None |
| 603 | if self.use_in_shortcut: |
| 604 | self.conv_shortcut = nn.Conv3d( |
| 605 | in_channels, |
| 606 | out_channels, |
| 607 | kernel_size=1, |
| 608 | stride=1, |
| 609 | padding=0, |
| 610 | ) |
| 611 | |
| 612 | def forward(self, input_tensor: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: |
| 613 | hidden_states = input_tensor |
nothing calls this directly
no test coverage detected