| 223 | |
| 224 | |
| 225 | class DoubleStreamBlock(nn.Module): |
| 226 | def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, backend = 'pytorch'): |
| 227 | super().__init__() |
| 228 | |
| 229 | mlp_hidden_dim = int(hidden_size * mlp_ratio) |
| 230 | self.num_heads = num_heads |
| 231 | self.hidden_size = hidden_size |
| 232 | self.img_mod = Modulation(hidden_size, double=True) |
| 233 | self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 234 | self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 235 | |
| 236 | self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 237 | self.img_mlp = nn.Sequential( |
| 238 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 239 | nn.GELU(approximate="tanh"), |
| 240 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 241 | ) |
| 242 | |
| 243 | self.backend = backend |
| 244 | |
| 245 | self.txt_mod = Modulation(hidden_size, double=True) |
| 246 | self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 247 | self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias) |
| 248 | |
| 249 | self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) |
| 250 | self.txt_mlp = nn.Sequential( |
| 251 | nn.Linear(hidden_size, mlp_hidden_dim, bias=True), |
| 252 | nn.GELU(approximate="tanh"), |
| 253 | nn.Linear(mlp_hidden_dim, hidden_size, bias=True), |
| 254 | ) |
| 255 | |
| 256 | |
| 257 | |
| 258 | |
| 259 | def forward(self, x: Tensor, vec: Tensor, pe: Tensor, mask: Tensor = None, txt_length = None): |
| 260 | img_mod1, img_mod2 = self.img_mod(vec) |
| 261 | txt_mod1, txt_mod2 = self.txt_mod(vec) |
| 262 | |
| 263 | txt, img = x[:, :txt_length], x[:, txt_length:] |
| 264 | |
| 265 | # prepare image for attention |
| 266 | img_modulated = self.img_norm1(img) |
| 267 | img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift |
| 268 | img_qkv = self.img_attn.qkv(img_modulated) |
| 269 | img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) |
| 270 | img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) |
| 271 | # prepare txt for attention |
| 272 | txt_modulated = self.txt_norm1(txt) |
| 273 | txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift |
| 274 | txt_qkv = self.txt_attn.qkv(txt_modulated) |
| 275 | txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) |
| 276 | txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) |
| 277 | |
| 278 | # run actual attention |
| 279 | q = torch.cat((txt_q, img_q), dim=2) |
| 280 | k = torch.cat((txt_k, img_k), dim=2) |
| 281 | v = torch.cat((txt_v, img_v), dim=2) |
| 282 | if mask is not None: |