| 1664 | |
| 1665 | |
| 1666 | class SkipDownBlock2D(nn.Module): |
| 1667 | def __init__( |
| 1668 | self, |
| 1669 | in_channels: int, |
| 1670 | out_channels: int, |
| 1671 | temb_channels: int, |
| 1672 | dropout: float = 0.0, |
| 1673 | num_layers: int = 1, |
| 1674 | resnet_eps: float = 1e-6, |
| 1675 | resnet_time_scale_shift: str = "default", |
| 1676 | resnet_act_fn: str = "swish", |
| 1677 | resnet_pre_norm: bool = True, |
| 1678 | output_scale_factor: float = np.sqrt(2.0), |
| 1679 | add_downsample: bool = True, |
| 1680 | downsample_padding: int = 1, |
| 1681 | ): |
| 1682 | super().__init__() |
| 1683 | self.resnets = nn.ModuleList([]) |
| 1684 | |
| 1685 | for i in range(num_layers): |
| 1686 | in_channels = in_channels if i == 0 else out_channels |
| 1687 | self.resnets.append( |
| 1688 | ResnetBlock2D( |
| 1689 | in_channels=in_channels, |
| 1690 | out_channels=out_channels, |
| 1691 | temb_channels=temb_channels, |
| 1692 | eps=resnet_eps, |
| 1693 | groups=min(in_channels // 4, 32), |
| 1694 | groups_out=min(out_channels // 4, 32), |
| 1695 | dropout=dropout, |
| 1696 | time_embedding_norm=resnet_time_scale_shift, |
| 1697 | non_linearity=resnet_act_fn, |
| 1698 | output_scale_factor=output_scale_factor, |
| 1699 | pre_norm=resnet_pre_norm, |
| 1700 | ) |
| 1701 | ) |
| 1702 | |
| 1703 | if add_downsample: |
| 1704 | self.resnet_down = ResnetBlock2D( |
| 1705 | in_channels=out_channels, |
| 1706 | out_channels=out_channels, |
| 1707 | temb_channels=temb_channels, |
| 1708 | eps=resnet_eps, |
| 1709 | groups=min(out_channels // 4, 32), |
| 1710 | dropout=dropout, |
| 1711 | time_embedding_norm=resnet_time_scale_shift, |
| 1712 | non_linearity=resnet_act_fn, |
| 1713 | output_scale_factor=output_scale_factor, |
| 1714 | pre_norm=resnet_pre_norm, |
| 1715 | use_in_shortcut=True, |
| 1716 | down=True, |
| 1717 | kernel="fir", |
| 1718 | ) |
| 1719 | self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)]) |
| 1720 | self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1)) |
| 1721 | else: |
| 1722 | self.resnet_down = None |
| 1723 | self.downsamplers = None |
no outgoing calls
no test coverage detected
searching dependent graphs…