| 108 | return self.conv_out(self.norm_out(x).swish()) |
| 109 | |
| 110 | class Encoder: |
| 111 | def __init__(self): |
| 112 | sz = [(128, 128), (128, 256), (256, 512), (512, 512)] |
| 113 | self.conv_in = Conv2d(3,128,3, padding=1) |
| 114 | |
| 115 | arr = [] |
| 116 | for i,s in enumerate(sz): |
| 117 | arr.append({"block": |
| 118 | [ResnetBlock(s[0], s[1]), |
| 119 | ResnetBlock(s[1], s[1])]}) |
| 120 | if i != 3: arr[-1]['downsample'] = {"conv": Conv2d(s[1], s[1], 3, stride=2, padding=(0,1,0,1))} |
| 121 | self.down = arr |
| 122 | |
| 123 | self.mid = Mid(512) |
| 124 | self.norm_out = GroupNorm(32, 512) |
| 125 | self.conv_out = Conv2d(512, 8, 3, padding=1) |
| 126 | |
| 127 | def __call__(self, x): |
| 128 | x = self.conv_in(x) |
| 129 | |
| 130 | for i, l in enumerate(self.down): |
| 131 | for b in l['block']: x = b(x) |
| 132 | if 'downsample' in l: x = l['downsample']['conv'](x) |
| 133 | |
| 134 | x = self.mid(x) |
| 135 | return self.conv_out(self.norm_out(x).swish()) |
| 136 | |
| 137 | class AutoencoderKL: |
| 138 | def __init__(self): |