r""" A Resnet block. Parameters: in_channels (`int`): The number of channels in the input. out_channels (`int`, *optional*, default to be `None`): The number of output channels for the first conv2d layer. If None, same as `in_channels`. temb_channels (`in
| 544 | |
| 545 | |
| 546 | class TemporalResnetBlock(nn.Module): |
| 547 | r""" |
| 548 | A Resnet block. |
| 549 | |
| 550 | Parameters: |
| 551 | in_channels (`int`): The number of channels in the input. |
| 552 | out_channels (`int`, *optional*, default to be `None`): |
| 553 | The number of output channels for the first conv2d layer. If None, same as `in_channels`. |
| 554 | temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. |
| 555 | eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. |
| 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: |
no outgoing calls
no test coverage detected
searching dependent graphs…