| 795 | return h |
| 796 | |
| 797 | class ResBlock(nn.Module): |
| 798 | def __init__(self, in_channels, out_channels=None): |
| 799 | super(ResBlock, self).__init__() |
| 800 | self.in_channels = in_channels |
| 801 | self.out_channels = in_channels if out_channels is None else out_channels |
| 802 | self.norm1 = Normalize(in_channels) |
| 803 | self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| 804 | self.norm2 = Normalize(out_channels) |
| 805 | self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| 806 | if self.in_channels != self.out_channels: |
| 807 | self.conv_out = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) |
| 808 | |
| 809 | def forward(self, x_in): |
| 810 | x = x_in |
| 811 | x = self.norm1(x) |
| 812 | x = nonlinearity(x) |
| 813 | x = self.conv1(x) |
| 814 | x = self.norm2(x) |
| 815 | x = nonlinearity(x) |
| 816 | x = self.conv2(x) |
| 817 | if self.in_channels != self.out_channels: |
| 818 | x_in = self.conv_out(x_in) |
| 819 | |
| 820 | return x + x_in |
| 821 | |
| 822 | class Fuse_sft_block_RRDB(nn.Module): |
| 823 | def __init__(self, in_ch, out_ch, num_block=1, num_grow_ch=32): |