(self, x: Tensor, vec: Tensor,
pe: Tensor, mask: Tensor = None,
txt_length=None,
edit_length=None)
| 428 | ) |
| 429 | |
| 430 | def forward(self, x: Tensor, vec: Tensor, |
| 431 | pe: Tensor, mask: Tensor = None, |
| 432 | txt_length=None, |
| 433 | edit_length=None): |
| 434 | if edit_length is not None: |
| 435 | txt, edit, img = x[:, :txt_length], x[:, txt_length:txt_length + edit_length], x[:, txt_length + edit_length:] |
| 436 | else: |
| 437 | txt, img = x[:, :txt_length], x[:, txt_length:] |
| 438 | img_mod1, img_mod2 = self.img_mod(vec) |
| 439 | txt_mod1, txt_mod2 = self.txt_mod(vec) |
| 440 | # prepare image for attention |
| 441 | img_modulated = self.img_norm1(img) |
| 442 | img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift |
| 443 | img_qkv = self.img_attn.qkv(img_modulated) |
| 444 | 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) |
| 445 | img_q, img_k = self.img_attn.norm(img_q, img_k, img_v) |
| 446 | # prepare txt for attention |
| 447 | txt_modulated = self.txt_norm1(txt) |
| 448 | txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift |
| 449 | txt_qkv = self.txt_attn.qkv(txt_modulated) |
| 450 | 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) |
| 451 | txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v) |
| 452 | |
| 453 | if edit_length is not None: |
| 454 | edit_mod1, edit_mod2 = self.edit_mod(vec) |
| 455 | # prepare edit for attention |
| 456 | edit_modulated = self.edit_norm1(edit) |
| 457 | edit_modulated = (1 + edit_mod1.scale) * edit_modulated + edit_mod1.shift |
| 458 | edit_qkv = self.edit_attn.qkv(edit_modulated) |
| 459 | edit_q, edit_k, edit_v = rearrange(edit_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads) |
| 460 | edit_q, edit_k = self.edit_attn.norm(edit_q, edit_k, edit_v) |
| 461 | else: |
| 462 | edit_q, edit_k, edit_v = None, None, None |
| 463 | |
| 464 | |
| 465 | # run actual attention |
| 466 | q = torch.cat((txt_q,) + ((edit_q,) if edit_q is not None else ()) + (img_q,), dim=2) |
| 467 | k = torch.cat((txt_k,) + ((edit_k,) if edit_k is not None else ()) + (img_k,), dim=2) |
| 468 | v = torch.cat((txt_v,) + ((edit_v,) if edit_v is not None else ()) + (img_v,), dim=2) |
| 469 | if mask is not None: |
| 470 | mask = repeat(mask, 'B L S-> B H L S', H=self.num_heads) |
| 471 | attn = attention(q, k, v, pe=pe, mask=mask, backend=self.backend) |
| 472 | if edit_length is not None: |
| 473 | txt_attn, edit_attn, img_attn = attn[:, : txt_length], attn[:, txt_length:txt_length + edit_length ], attn[:, txt_length + edit_length:] |
| 474 | else: |
| 475 | txt_attn, img_attn = attn[:, : txt_length], attn[:, txt_length:] |
| 476 | |
| 477 | # calculate the img bloks |
| 478 | img = img + img_mod1.gate * self.img_attn.proj(img_attn) |
| 479 | img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift) |
| 480 | |
| 481 | # calculate the txt bloks |
| 482 | txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn) |
| 483 | txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift) |
| 484 | |
| 485 | # calculate the img bloks |
| 486 | if edit_length is not None: |
| 487 | edit = edit + edit_mod1.gate * self.edit_attn.proj(edit_attn) |
nothing calls this directly
no test coverage detected