18importmath19fromtypingimportList2021importnumpyasnp22importtorch23importtorch.nnasnn24importtorch.nn.functionalasF2526fromlabml\_nn.diffusion.stable\_diffusion.model.unet\_attentionimportSpatialTransformer

29classUNetModel(nn.Module):

34def\_\_init\_\_(35self,\*,36in\_channels:int,37out\_channels:int,38channels:int,39n\_res\_blocks:int,40attention\_levels:List[int],41channel\_multipliers:List[int],42n\_heads:int,43tf\_layers:int=1,44d\_cond:int=768):

54super().\_\_init\_\_()55self.channels=channels

58levels=len(channel\_multipliers)

60d\_time\_emb=channels\*461self.time\_embed=nn.Sequential(62nn.Linear(channels,d\_time\_emb),63nn.SiLU(),64nn.Linear(d\_time\_emb,d\_time\_emb),65)

68self.input\_blocks=nn.ModuleList()

75self.input\_blocks.append(TimestepEmbedSequential(76nn.Conv2d(in\_channels,channels,3,padding=1)))

78input\_block\_channels=[channels]

80channels\_list=[channels\*mforminchannel\_multipliers]

82foriinrange(levels):

84for\_inrange(n\_res\_blocks):

87layers=[ResBlock(channels,d\_time\_emb,out\_channels=channels\_list[i])]88channels=channels\_list[i]

90ifiinattention\_levels:91layers.append(SpatialTransformer(channels,n\_heads,tf\_layers,d\_cond))

94self.input\_blocks.append(TimestepEmbedSequential(\*layers))95input\_block\_channels.append(channels)

97ifi!=levels-1:98self.input\_blocks.append(TimestepEmbedSequential(DownSample(channels)))99input\_block\_channels.append(channels)

102self.middle\_block=TimestepEmbedSequential(103ResBlock(channels,d\_time\_emb),104SpatialTransformer(channels,n\_heads,tf\_layers,d\_cond),105ResBlock(channels,d\_time\_emb),106)

109self.output\_blocks=nn.ModuleList([])

111foriinreversed(range(levels)):

113forjinrange(n\_res\_blocks+1):

117layers=[ResBlock(channels+input\_block\_channels.pop(),d\_time\_emb,out\_channels=channels\_list[i])]118channels=channels\_list[i]

120ifiinattention\_levels:121layers.append(SpatialTransformer(channels,n\_heads,tf\_layers,d\_cond))

125ifi!=0andj==n\_res\_blocks:126layers.append(UpSample(channels))

128self.output\_blocks.append(TimestepEmbedSequential(\*layers))

131self.out=nn.Sequential(132normalization(channels),133nn.SiLU(),134nn.Conv2d(channels,out\_channels,3,padding=1),135)

137deftime\_step\_embedding(self,time\_steps:torch.Tensor,max\_period:int=10000):

145half=self.channels//2

147frequencies=torch.exp(148-math.log(max\_period)\*torch.arange(start=0,end=half,dtype=torch.float32)/half149).to(device=time\_steps.device)

151args=time\_steps[:,None].float()\*frequencies[None]

153returntorch.cat([torch.cos(args),torch.sin(args)],dim=-1)

155defforward(self,x:torch.Tensor,time\_steps:torch.Tensor,cond:torch.Tensor):

162x\_input\_block=[]

165t\_emb=self.time\_step\_embedding(time\_steps)166t\_emb=self.time\_embed(t\_emb)

169formoduleinself.input\_blocks:170x=module(x,t\_emb,cond)171x\_input\_block.append(x)

173x=self.middle\_block(x,t\_emb,cond)

175formoduleinself.output\_blocks:176x=torch.cat([x,x\_input\_block.pop()],dim=1)177x=module(x,t\_emb,cond)

180returnself.out(x)

183classTimestepEmbedSequential(nn.Sequential):

191defforward(self,x,t\_emb,cond=None):192forlayerinself:193ifisinstance(layer,ResBlock):194x=layer(x,t\_emb)195elifisinstance(layer,SpatialTransformer):196x=layer(x,cond)197else:198x=layer(x)199returnx

202classUpSample(nn.Module):

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

211super().\_\_init\_\_()

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

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

220x=F.interpolate(x,scale\_factor=2,mode="nearest")

222returnself.conv(x)

225classDownSample(nn.Module):

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

234super().\_\_init\_\_()

236self.op=nn.Conv2d(channels,channels,3,stride=2,padding=1)

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

243returnself.op(x)

246classResBlock(nn.Module):

251def\_\_init\_\_(self,channels:int,d\_t\_emb:int,\*,out\_channels=None):

257super().\_\_init\_\_()

259ifout\_channelsisNone:260out\_channels=channels

263self.in\_layers=nn.Sequential(264normalization(channels),265nn.SiLU(),266nn.Conv2d(channels,out\_channels,3,padding=1),267)

270self.emb\_layers=nn.Sequential(271nn.SiLU(),272nn.Linear(d\_t\_emb,out\_channels),273)

275self.out\_layers=nn.Sequential(276normalization(out\_channels),277nn.SiLU(),278nn.Dropout(0.),279nn.Conv2d(out\_channels,out\_channels,3,padding=1)280)

283ifout\_channels==channels:284self.skip\_connection=nn.Identity()285else:286self.skip\_connection=nn.Conv2d(channels,out\_channels,1)

288defforward(self,x:torch.Tensor,t\_emb:torch.Tensor):

294h=self.in\_layers(x)

296t\_emb=self.emb\_layers(t\_emb).type(h.dtype)

298h=h+t\_emb[:,:,None,None]

300h=self.out\_layers(h)

302returnself.skip\_connection(x)+h

305classGroupNorm32(nn.GroupNorm):

310defforward(self,x):311returnsuper().forward(x.float()).type(x.dtype)

314defnormalization(channels):

320returnGroupNorm32(32,channels)

323def\_test\_time\_embeddings():

327importmatplotlib.pyplotasplt328329plt.figure(figsize=(15,5))330m=UNetModel(in\_channels=1,out\_channels=1,channels=320,n\_res\_blocks=1,attention\_levels=[],331channel\_multipliers=[],332n\_heads=1,tf\_layers=1,d\_cond=1)333te=m.time\_step\_embedding(torch.arange(0,1000))334plt.plot(np.arange(1000),te[:,[50,100,190,260]].numpy())335plt.legend(["dim %d"%pforpin[50,100,190,260]])336plt.title("Time embeddings")337plt.show()

341if\_\_name\_\_=='\_\_main\_\_':342\_test\_time\_embeddings()