| 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): |
| 824 | super().__init__() |
| 825 | self.encode_enc_1 = ResBlock(2*in_ch, in_ch) |
| 826 | self.encode_enc_2 = make_layer(RRDB, num_block, num_feat=in_ch, num_grow_ch=num_grow_ch) |
| 827 | self.encode_enc_3 = ResBlock(in_ch, out_ch) |
| 828 | |
| 829 | def forward(self, enc_feat, dec_feat, w=1): |
| 830 | enc_feat = self.encode_enc_1(torch.cat([enc_feat, dec_feat], dim=1)) |
| 831 | enc_feat = self.encode_enc_2(enc_feat) |
| 832 | enc_feat = self.encode_enc_3(enc_feat) |
| 833 | residual = w * enc_feat |
| 834 | out = dec_feat + residual |
| 835 | return out |
| 836 | |
| 837 | class SimpleDecoder(nn.Module): |
| 838 | def __init__(self, in_channels, out_channels, *args, **kwargs): |