24importmath25fromtypingimportOptional,Tuple,Union,List2627importtorch28fromtorchimportnn2930fromlabml\_helpers.moduleimportModule

33classSwish(Module):

40defforward(self,x):41returnx\*torch.sigmoid(x)

44classTimeEmbedding(nn.Module):

49def\_\_init\_\_(self,n\_channels:int):

53super().\_\_init\_\_()54self.n\_channels=n\_channels

56self.lin1=nn.Linear(self.n\_channels//4,self.n\_channels)

58self.act=Swish()

60self.lin2=nn.Linear(self.n\_channels,self.n\_channels)

62defforward(self,t:torch.Tensor):

72half\_dim=self.n\_channels//873emb=math.log(10\_000)/(half\_dim-1)74emb=torch.exp(torch.arange(half\_dim,device=t.device)\*-emb)75emb=t[:,None]\*emb[None,:]76emb=torch.cat((emb.sin(),emb.cos()),dim=1)

79emb=self.act(self.lin1(emb))80emb=self.lin2(emb)

83returnemb

86classResidualBlock(Module):

94def\_\_init\_\_(self,in\_channels:int,out\_channels:int,time\_channels:int,95n\_groups:int=32,dropout:float=0.1):

103super().\_\_init\_\_()

105self.norm1=nn.GroupNorm(n\_groups,in\_channels)106self.act1=Swish()107self.conv1=nn.Conv2d(in\_channels,out\_channels,kernel\_size=(3,3),padding=(1,1))

110self.norm2=nn.GroupNorm(n\_groups,out\_channels)111self.act2=Swish()112self.conv2=nn.Conv2d(out\_channels,out\_channels,kernel\_size=(3,3),padding=(1,1))

116ifin\_channels!=out\_channels:117self.shortcut=nn.Conv2d(in\_channels,out\_channels,kernel\_size=(1,1))118else:119self.shortcut=nn.Identity()

122self.time\_emb=nn.Linear(time\_channels,out\_channels)123self.time\_act=Swish()124125self.dropout=nn.Dropout(dropout)

127defforward(self,x:torch.Tensor,t:torch.Tensor):

133h=self.conv1(self.act1(self.norm1(x)))

135h+=self.time\_emb(self.time\_act(t))[:,:,None,None]

137h=self.conv2(self.dropout(self.act2(self.norm2(h))))

140returnh+self.shortcut(x)

143classAttentionBlock(Module):

150def\_\_init\_\_(self,n\_channels:int,n\_heads:int=1,d\_k:int=None,n\_groups:int=32):

157super().\_\_init\_\_()

160ifd\_kisNone:161d\_k=n\_channels

163self.norm=nn.GroupNorm(n\_groups,n\_channels)

165self.projection=nn.Linear(n\_channels,n\_heads\*d\_k\*3)

167self.output=nn.Linear(n\_heads\*d\_k,n\_channels)

169self.scale=d\_k\*\*-0.5

171self.n\_heads=n\_heads172self.d\_k=d\_k

174defforward(self,x:torch.Tensor,t:Optional[torch.Tensor]=None):

181\_=t

183batch\_size,n\_channels,height,width=x.shape

185x=x.view(batch\_size,n\_channels,-1).permute(0,2,1)

187qkv=self.projection(x).view(batch\_size,-1,self.n\_heads,3\*self.d\_k)

189q,k,v=torch.chunk(qkv,3,dim=-1)

191attn=torch.einsum('bihd,bjhd-\>bijh',q,k)\*self.scale

193attn=attn.softmax(dim=2)

195res=torch.einsum('bijh,bjhd-\>bihd',attn,v)

197res=res.view(batch\_size,-1,self.n\_heads\*self.d\_k)

199res=self.output(res)

202res+=x

205res=res.permute(0,2,1).view(batch\_size,n\_channels,height,width)

208returnres

211classDownBlock(Module):

218def\_\_init\_\_(self,in\_channels:int,out\_channels:int,time\_channels:int,has\_attn:bool):219super().\_\_init\_\_()220self.res=ResidualBlock(in\_channels,out\_channels,time\_channels)221ifhas\_attn:222self.attn=AttentionBlock(out\_channels)223else:224self.attn=nn.Identity()

226defforward(self,x:torch.Tensor,t:torch.Tensor):227x=self.res(x,t)228x=self.attn(x)229returnx

232classUpBlock(Module):

239def\_\_init\_\_(self,in\_channels:int,out\_channels:int,time\_channels:int,has\_attn:bool):240super().\_\_init\_\_()

243self.res=ResidualBlock(in\_channels+out\_channels,out\_channels,time\_channels)244ifhas\_attn:245self.attn=AttentionBlock(out\_channels)246else:247self.attn=nn.Identity()

249defforward(self,x:torch.Tensor,t:torch.Tensor):250x=self.res(x,t)251x=self.attn(x)252returnx

255classMiddleBlock(Module):

263def\_\_init\_\_(self,n\_channels:int,time\_channels:int):264super().\_\_init\_\_()265self.res1=ResidualBlock(n\_channels,n\_channels,time\_channels)266self.attn=AttentionBlock(n\_channels)267self.res2=ResidualBlock(n\_channels,n\_channels,time\_channels)

269defforward(self,x:torch.Tensor,t:torch.Tensor):270x=self.res1(x,t)271x=self.attn(x)272x=self.res2(x,t)273returnx

276classUpsample(nn.Module):

281def\_\_init\_\_(self,n\_channels):282super().\_\_init\_\_()283self.conv=nn.ConvTranspose2d(n\_channels,n\_channels,(4,4),(2,2),(1,1))

285defforward(self,x:torch.Tensor,t:torch.Tensor):

288\_=t289returnself.conv(x)

292classDownsample(nn.Module):

297def\_\_init\_\_(self,n\_channels):298super().\_\_init\_\_()299self.conv=nn.Conv2d(n\_channels,n\_channels,(3,3),(2,2),(1,1))

301defforward(self,x:torch.Tensor,t:torch.Tensor):

304\_=t305returnself.conv(x)

308classUNet(Module):

313def\_\_init\_\_(self,image\_channels:int=3,n\_channels:int=64,314ch\_mults:Union[Tuple[int,...],List[int]]=(1,2,2,4),315is\_attn:Union[Tuple[bool,...],List[int]]=(False,False,True,True),316n\_blocks:int=2):

324super().\_\_init\_\_()

327n\_resolutions=len(ch\_mults)

330self.image\_proj=nn.Conv2d(image\_channels,n\_channels,kernel\_size=(3,3),padding=(1,1))

333self.time\_emb=TimeEmbedding(n\_channels\*4)

336down=[]

338out\_channels=in\_channels=n\_channels

340foriinrange(n\_resolutions):

342out\_channels=in\_channels\*ch\_mults[i]

344for\_inrange(n\_blocks):345down.append(DownBlock(in\_channels,out\_channels,n\_channels\*4,is\_attn[i]))346in\_channels=out\_channels

348ifi\<n\_resolutions-1:349down.append(Downsample(in\_channels))

352self.down=nn.ModuleList(down)

355self.middle=MiddleBlock(out\_channels,n\_channels\*4,)

358up=[]

360in\_channels=out\_channels

362foriinreversed(range(n\_resolutions)):

364out\_channels=in\_channels365for\_inrange(n\_blocks):366up.append(UpBlock(in\_channels,out\_channels,n\_channels\*4,is\_attn[i]))

368out\_channels=in\_channels//ch\_mults[i]369up.append(UpBlock(in\_channels,out\_channels,n\_channels\*4,is\_attn[i]))370in\_channels=out\_channels

372ifi\>0:373up.append(Upsample(in\_channels))

376self.up=nn.ModuleList(up)

379self.norm=nn.GroupNorm(8,n\_channels)380self.act=Swish()381self.final=nn.Conv2d(in\_channels,image\_channels,kernel\_size=(3,3),padding=(1,1))

383defforward(self,x:torch.Tensor,t:torch.Tensor):

390t=self.time\_emb(t)

393x=self.image\_proj(x)

396h=[x]

398forminself.down:399x=m(x,t)400h.append(x)

403x=self.middle(x,t)

406forminself.up:407ifisinstance(m,Upsample):408x=m(x,t)409else:

411s=h.pop()412x=torch.cat((x,s),dim=1)

414x=m(x,t)

417returnself.final(self.act(self.norm(x)))