vis: b c n text: b c n
(self, x, text)
| 618 | self.adapter = nn.Linear(in_dim, in_dim*out_dim*kernel_size*kernel_size+out_dim) |
| 619 | |
| 620 | def forward(self, x, text): |
| 621 | ''' |
| 622 | vis: b c n |
| 623 | text: b c n |
| 624 | ''' |
| 625 | |
| 626 | B, C, N = x.size() |
| 627 | x = rearrange(x, 'b c n -> (b n) c') |
| 628 | x = x.reshape(1, -1, 1, 1) |
| 629 | # b c h w -> b c*c h w |
| 630 | text = rearrange(text, 'b c n -> (b n) c') |
| 631 | adapter = self.adapter(text) # Eq. 8 |
| 632 | weight, bias = adapter[:, :-self.out_dim], text[:, -self.out_dim:] |
| 633 | weight = weight.reshape(-1, C, self.kernel_size, self.kernel_size) |
| 634 | # print(x.shape, weight.shape, bias.shape) |
| 635 | bias = bias.reshape(-1) |
| 636 | # 1 (b n c) 1 1 -> 1 (b n c) 1 1 |
| 637 | out = F.conv2d(x, |
| 638 | weight, |
| 639 | padding=0, |
| 640 | groups=B*N, |
| 641 | bias=bias |
| 642 | ) |
| 643 | out = rearrange(out.squeeze(), '(b n c) -> b n c', b=B, c=self.out_dim, n=N) |
| 644 | out = out.permute(0, 2, 1) |
| 645 | return out |
| 646 | |
| 647 | |
| 648 | class PWAM(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected