| 481 | |
| 482 | |
| 483 | class FastFreqCondInjection(nn.Module): |
| 484 | def __init__( |
| 485 | self, |
| 486 | fea_dim, |
| 487 | cond_dim, |
| 488 | qkv_dim, |
| 489 | dim_out, |
| 490 | groups=32, |
| 491 | nheads=8, |
| 492 | drop_path_prob=0.2, |
| 493 | ) -> None: |
| 494 | super().__init__() |
| 495 | assert fea_dim % nheads == 0, "@dim must be divisible by @nheads" |
| 496 | |
| 497 | self.prenorm_x = nn.GroupNorm(groups, fea_dim) |
| 498 | # self.prenorm_cond = nn.GroupNorm(groups // 4, cond_dim) |
| 499 | |
| 500 | self.q = nn.Sequential( |
| 501 | nn.Conv2d(fea_dim, fea_dim, 3, 1, 1, bias=False, groups=fea_dim), |
| 502 | nn.Conv2d(fea_dim, qkv_dim, 1, bias=True), |
| 503 | ) |
| 504 | self.kv = nn.Sequential( |
| 505 | nn.Conv2d(cond_dim, cond_dim, 3, 1, 1, bias=False, groups=cond_dim), |
| 506 | nn.Conv2d(cond_dim, qkv_dim * 2, 1, bias=True), |
| 507 | ) |
| 508 | self.nheads = nheads |
| 509 | self.scale = 1 / math.sqrt(qkv_dim // nheads) |
| 510 | |
| 511 | self.attn_out = nn.Conv2d(qkv_dim, dim_out, 1, bias=True) |
| 512 | self.attn_res = ( |
| 513 | nn.Conv2d(fea_dim, dim_out, 1, bias=True) |
| 514 | if fea_dim != dim_out |
| 515 | else nn.Identity() |
| 516 | ) |
| 517 | |
| 518 | self.ffn = nn.Sequential( |
| 519 | nn.Conv2d(dim_out, dim_out * 2, 3, 1, 1, bias=False), |
| 520 | nn.GELU(), |
| 521 | nn.Conv2d(dim_out * 2, dim_out, 3, 1, 1, bias=False), |
| 522 | nn.Conv2d(dim_out, dim_out, 1, bias=True), |
| 523 | ) |
| 524 | self.ffn_drop_path = DropPath(drop_prob=drop_path_prob) |
| 525 | |
| 526 | def forward(self, x, cond): |
| 527 | x = self.prenorm_x(x) |
| 528 | # cond = self.prenorm_cond(cond) |
| 529 | |
| 530 | q = self.q(x) |
| 531 | k, v = self.kv(cond).chunk(2, dim=1) |
| 532 | |
| 533 | q, k, v = map(lambda in_qkv: F.normalize(in_qkv, dim=1), (q, k, v)) |
| 534 | |
| 535 | q = q.softmax(dim=-2) |
| 536 | k = k.softmax(dim=-1) |
| 537 | |
| 538 | # convert to freq space |
| 539 | q = torch.fft.rfft2(q, dim=(-2, -1), norm="ortho") # b, c, h, w/2+1 |
| 540 | k = torch.fft.rfft2(k, dim=(-2, -1), norm="ortho") # b, c, h, w/2+1 |
nothing calls this directly
no outgoing calls
no test coverage detected