19fromtypingimportOptional2021importtorch22importtorch.nn.functionalasF23fromtorchimportnn

26classSpatialTransformer(nn.Module):

31def\_\_init\_\_(self,channels:int,n\_heads:int,n\_layers:int,d\_cond:int):

38super().\_\_init\_\_()

40self.norm=torch.nn.GroupNorm(num\_groups=32,num\_channels=channels,eps=1e-6,affine=True)

42self.proj\_in=nn.Conv2d(channels,channels,kernel\_size=1,stride=1,padding=0)

45self.transformer\_blocks=nn.ModuleList(46[BasicTransformerBlock(channels,n\_heads,channels//n\_heads,d\_cond=d\_cond)for\_inrange(n\_layers)]47)

50self.proj\_out=nn.Conv2d(channels,channels,kernel\_size=1,stride=1,padding=0)

52defforward(self,x:torch.Tensor,cond:torch.Tensor):

58b,c,h,w=x.shape

60x\_in=x

62x=self.norm(x)

64x=self.proj\_in(x)

67x=x.permute(0,2,3,1).view(b,h\*w,c)

69forblockinself.transformer\_blocks:70x=block(x,cond)

73x=x.view(b,h,w,c).permute(0,3,1,2)

75x=self.proj\_out(x)

77returnx+x\_in

80classBasicTransformerBlock(nn.Module):

85def\_\_init\_\_(self,d\_model:int,n\_heads:int,d\_head:int,d\_cond:int):

92super().\_\_init\_\_()

94self.attn1=CrossAttention(d\_model,d\_model,n\_heads,d\_head)95self.norm1=nn.LayerNorm(d\_model)

97self.attn2=CrossAttention(d\_model,d\_cond,n\_heads,d\_head)98self.norm2=nn.LayerNorm(d\_model)

100self.ff=FeedForward(d\_model)101self.norm3=nn.LayerNorm(d\_model)

103defforward(self,x:torch.Tensor,cond:torch.Tensor):

109x=self.attn1(self.norm1(x))+x

111x=self.attn2(self.norm2(x),cond=cond)+x

113x=self.ff(self.norm3(x))+x

115returnx

118classCrossAttention(nn.Module):

125use\_flash\_attention:bool=False

127def\_\_init\_\_(self,d\_model:int,d\_cond:int,n\_heads:int,d\_head:int,is\_inplace:bool=True):

136super().\_\_init\_\_()137138self.is\_inplace=is\_inplace139self.n\_heads=n\_heads140self.d\_head=d\_head

143self.scale=d\_head\*\*-0.5

146d\_attn=d\_head\*n\_heads147self.to\_q=nn.Linear(d\_model,d\_attn,bias=False)148self.to\_k=nn.Linear(d\_cond,d\_attn,bias=False)149self.to\_v=nn.Linear(d\_cond,d\_attn,bias=False)

152self.to\_out=nn.Sequential(nn.Linear(d\_attn,d\_model))

157try:

161fromflash\_attn.flash\_attentionimportFlashAttention162self.flash=FlashAttention()

164self.flash.softmax\_scale=self.scale

166exceptImportError:167self.flash=None

169defforward(self,x:torch.Tensor,cond:Optional[torch.Tensor]=None):

176has\_cond=condisnotNone177ifnothas\_cond:178cond=x

181q=self.to\_q(x)182k=self.to\_k(cond)183v=self.to\_v(cond)

186ifCrossAttention.use\_flash\_attentionandself.flashisnotNoneandnothas\_condandself.d\_head\<=128:187returnself.flash\_attention(q,k,v)

189else:190returnself.normal\_attention(q,k,v)

192defflash\_attention(self,q:torch.Tensor,k:torch.Tensor,v:torch.Tensor):

202batch\_size,seq\_len,\_=q.shape

206qkv=torch.stack((q,k,v),dim=2)

208qkv=qkv.view(batch\_size,seq\_len,3,self.n\_heads,self.d\_head)

212ifself.d\_head\<=32:213pad=32-self.d\_head214elifself.d\_head\<=64:215pad=64-self.d\_head216elifself.d\_head\<=128:217pad=128-self.d\_head218else:219raiseValueError(f'Head size ${self.d\_head} too large for Flash Attention')

222ifpad:223qkv=torch.cat((qkv,qkv.new\_zeros(batch\_size,seq\_len,3,self.n\_heads,pad)),dim=-1)

228out,\_=self.flash(qkv)

230out=out[:,:,:,:self.d\_head]

232out=out.reshape(batch\_size,seq\_len,self.n\_heads\*self.d\_head)

235returnself.to\_out(out)

237defnormal\_attention(self,q:torch.Tensor,k:torch.Tensor,v:torch.Tensor):

247q=q.view(\*q.shape[:2],self.n\_heads,-1)248k=k.view(\*k.shape[:2],self.n\_heads,-1)249v=v.view(\*v.shape[:2],self.n\_heads,-1)

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

256ifself.is\_inplace:257half=attn.shape[0]//2258attn[half:]=attn[half:].softmax(dim=-1)259attn[:half]=attn[:half].softmax(dim=-1)260else:261attn=attn.softmax(dim=-1)

265out=torch.einsum('bhij,bjhd-\>bihd',attn,v)

267out=out.reshape(\*out.shape[:2],-1)

269returnself.to\_out(out)

272classFeedForward(nn.Module):

277def\_\_init\_\_(self,d\_model:int,d\_mult:int=4):

282super().\_\_init\_\_()283self.net=nn.Sequential(284GeGLU(d\_model,d\_model\*d\_mult),285nn.Dropout(0.),286nn.Linear(d\_model\*d\_mult,d\_model)287)

289defforward(self,x:torch.Tensor):290returnself.net(x)

293classGeGLU(nn.Module):

300def\_\_init\_\_(self,d\_in:int,d\_out:int):301super().\_\_init\_\_()

303self.proj=nn.Linear(d\_in,d\_out\*2)

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

307x,gate=self.proj(x).chunk(2,dim=-1)

309returnx\*F.gelu(gate)