| 77 | |
| 78 | |
| 79 | class ResnetBlock(nn.Module): |
| 80 | |
| 81 | def __init__(self, in_c, out_c, down, ksize=3, sk=False, use_conv=True): |
| 82 | super().__init__() |
| 83 | ps = ksize // 2 |
| 84 | if in_c != out_c or sk == False: |
| 85 | self.in_conv = nn.Conv2d(in_c, out_c, ksize, 1, ps) |
| 86 | else: |
| 87 | self.in_conv = None |
| 88 | self.block1 = nn.Conv2d(out_c, out_c, 3, 1, 1) |
| 89 | self.act = nn.ReLU() |
| 90 | self.block2 = nn.Conv2d(out_c, out_c, ksize, 1, ps) |
| 91 | if sk == False: |
| 92 | self.skep = nn.Conv2d(in_c, out_c, ksize, 1, ps) |
| 93 | else: |
| 94 | self.skep = None |
| 95 | |
| 96 | self.down = down |
| 97 | if self.down == True: |
| 98 | self.down_opt = Downsample(in_c, use_conv=use_conv) |
| 99 | |
| 100 | def forward(self, x): |
| 101 | if self.down == True: |
| 102 | x = self.down_opt(x) |
| 103 | if self.in_conv is not None: # edit |
| 104 | x = self.in_conv(x) |
| 105 | |
| 106 | h = self.block1(x) |
| 107 | h = self.act(h) |
| 108 | h = self.block2(h) |
| 109 | if self.skep is not None: |
| 110 | return h + self.skep(x) |
| 111 | else: |
| 112 | return h + x |
| 113 | |
| 114 | |
| 115 | class Adapter_XL(nn.Module): |