18fromtypingimportList1920importtorch21importtorch.nn.functionalasF22fromtorchimportnn

25classAutoencoder(nn.Module):

32def\_\_init\_\_(self,encoder:'Encoder',decoder:'Decoder',emb\_channels:int,z\_channels:int):

39super().\_\_init\_\_()40self.encoder=encoder41self.decoder=decoder

44self.quant\_conv=nn.Conv2d(2\*z\_channels,2\*emb\_channels,1)

47self.post\_quant\_conv=nn.Conv2d(emb\_channels,z\_channels,1)

49defencode(self,img:torch.Tensor)-\>'GaussianDistribution':

56z=self.encoder(img)

58moments=self.quant\_conv(z)

60returnGaussianDistribution(moments)

62defdecode(self,z:torch.Tensor):

69z=self.post\_quant\_conv(z)

71returnself.decoder(z)

74classEncoder(nn.Module):

79def\_\_init\_\_(self,\*,channels:int,channel\_multipliers:List[int],n\_resnet\_blocks:int,80in\_channels:int,z\_channels:int):

89super().\_\_init\_\_()

93n\_resolutions=len(channel\_multipliers)

96self.conv\_in=nn.Conv2d(in\_channels,channels,3,stride=1,padding=1)

99channels\_list=[m\*channelsformin[1]+channel\_multipliers]

102self.down=nn.ModuleList()

104foriinrange(n\_resolutions):

106resnet\_blocks=nn.ModuleList()

108for\_inrange(n\_resnet\_blocks):109resnet\_blocks.append(ResnetBlock(channels,channels\_list[i+1]))110channels=channels\_list[i+1]

112down=nn.Module()113down.block=resnet\_blocks

115ifi!=n\_resolutions-1:116down.downsample=DownSample(channels)117else:118down.downsample=nn.Identity()

120self.down.append(down)

123self.mid=nn.Module()124self.mid.block\_1=ResnetBlock(channels,channels)125self.mid.attn\_1=AttnBlock(channels)126self.mid.block\_2=ResnetBlock(channels,channels)

129self.norm\_out=normalization(channels)130self.conv\_out=nn.Conv2d(channels,2\*z\_channels,3,stride=1,padding=1)

132defforward(self,img:torch.Tensor):

138x=self.conv\_in(img)

141fordowninself.down:

143forblockindown.block:144x=block(x)

146x=down.downsample(x)

149x=self.mid.block\_1(x)150x=self.mid.attn\_1(x)151x=self.mid.block\_2(x)

154x=self.norm\_out(x)155x=swish(x)156x=self.conv\_out(x)

159returnx

162classDecoder(nn.Module):

167def\_\_init\_\_(self,\*,channels:int,channel\_multipliers:List[int],n\_resnet\_blocks:int,168out\_channels:int,z\_channels:int):

177super().\_\_init\_\_()

181num\_resolutions=len(channel\_multipliers)

184channels\_list=[m\*channelsforminchannel\_multipliers]

187channels=channels\_list[-1]

190self.conv\_in=nn.Conv2d(z\_channels,channels,3,stride=1,padding=1)

193self.mid=nn.Module()194self.mid.block\_1=ResnetBlock(channels,channels)195self.mid.attn\_1=AttnBlock(channels)196self.mid.block\_2=ResnetBlock(channels,channels)

199self.up=nn.ModuleList()

201foriinreversed(range(num\_resolutions)):

203resnet\_blocks=nn.ModuleList()

205for\_inrange(n\_resnet\_blocks+1):206resnet\_blocks.append(ResnetBlock(channels,channels\_list[i]))207channels=channels\_list[i]

209up=nn.Module()210up.block=resnet\_blocks

212ifi!=0:213up.upsample=UpSample(channels)214else:215up.upsample=nn.Identity()

217self.up.insert(0,up)

220self.norm\_out=normalization(channels)221self.conv\_out=nn.Conv2d(channels,out\_channels,3,stride=1,padding=1)

223defforward(self,z:torch.Tensor):

229h=self.conv\_in(z)

232h=self.mid.block\_1(h)233h=self.mid.attn\_1(h)234h=self.mid.block\_2(h)

237forupinreversed(self.up):

239forblockinup.block:240h=block(h)

242h=up.upsample(h)

245h=self.norm\_out(h)246h=swish(h)247img=self.conv\_out(h)

250returnimg

253classGaussianDistribution:

258def\_\_init\_\_(self,parameters:torch.Tensor):

264self.mean,log\_var=torch.chunk(parameters,2,dim=1)

266self.log\_var=torch.clamp(log\_var,-30.0,20.0)

268self.std=torch.exp(0.5\*self.log\_var)

270defsample(self):

272returnself.mean+self.std\*torch.randn\_like(self.std)

275classAttnBlock(nn.Module):

280def\_\_init\_\_(self,channels:int):

284super().\_\_init\_\_()

286self.norm=normalization(channels)

288self.q=nn.Conv2d(channels,channels,1)289self.k=nn.Conv2d(channels,channels,1)290self.v=nn.Conv2d(channels,channels,1)

292self.proj\_out=nn.Conv2d(channels,channels,1)

294self.scale=channels\*\*-0.5

296defforward(self,x:torch.Tensor):

301x\_norm=self.norm(x)

303q=self.q(x\_norm)304k=self.k(x\_norm)305v=self.v(x\_norm)

310b,c,h,w=q.shape311q=q.view(b,c,h\*w)312k=k.view(b,c,h\*w)313v=v.view(b,c,h\*w)

316attn=torch.einsum('bci,bcj-\>bij',q,k)\*self.scale317attn=F.softmax(attn,dim=2)

320out=torch.einsum('bij,bcj-\>bci',attn,v)

323out=out.view(b,c,h,w)

325out=self.proj\_out(out)

328returnx+out

331classUpSample(nn.Module):

335def\_\_init\_\_(self,channels:int):

339super().\_\_init\_\_()

341self.conv=nn.Conv2d(channels,channels,3,padding=1)

343defforward(self,x:torch.Tensor):

348x=F.interpolate(x,scale\_factor=2.0,mode="nearest")

350returnself.conv(x)

353classDownSample(nn.Module):

357def\_\_init\_\_(self,channels:int):

361super().\_\_init\_\_()

363self.conv=nn.Conv2d(channels,channels,3,stride=2,padding=0)

365defforward(self,x:torch.Tensor):

370x=F.pad(x,(0,1,0,1),mode="constant",value=0)

372returnself.conv(x)

375classResnetBlock(nn.Module):

379def\_\_init\_\_(self,in\_channels:int,out\_channels:int):

384super().\_\_init\_\_()

386self.norm1=normalization(in\_channels)387self.conv1=nn.Conv2d(in\_channels,out\_channels,3,stride=1,padding=1)

389self.norm2=normalization(out\_channels)390self.conv2=nn.Conv2d(out\_channels,out\_channels,3,stride=1,padding=1)

392ifin\_channels!=out\_channels:393self.nin\_shortcut=nn.Conv2d(in\_channels,out\_channels,1,stride=1,padding=0)394else:395self.nin\_shortcut=nn.Identity()

397defforward(self,x:torch.Tensor):

402h=x

405h=self.norm1(h)406h=swish(h)407h=self.conv1(h)

410h=self.norm2(h)411h=swish(h)412h=self.conv2(h)

415returnself.nin\_shortcut(x)+h

418defswish(x:torch.Tensor):

424returnx\*torch.sigmoid(x)

427defnormalization(channels:int):

433returnnn.GroupNorm(num\_groups=32,num\_channels=channels,eps=1e-6)