| 397 | |
| 398 | |
| 399 | class FreqCondInjection(nn.Module): |
| 400 | def __init__( |
| 401 | self, |
| 402 | fea_dim, |
| 403 | cond_dim, |
| 404 | qkv_dim, |
| 405 | dim_out, |
| 406 | groups=32, |
| 407 | nheads=8, |
| 408 | drop_path_prob=0.2, |
| 409 | ) -> None: |
| 410 | super().__init__() |
| 411 | assert fea_dim % nheads == 0, "@dim must be divisible by @nheads" |
| 412 | |
| 413 | self.prenorm_x = nn.GroupNorm(groups, fea_dim) |
| 414 | # self.prenorm_cond = nn.GroupNorm(groups // 4, cond_dim) |
| 415 | |
| 416 | self.q = nn.Sequential( |
| 417 | nn.Conv2d(fea_dim, fea_dim, 3, 1, 1, bias=False, groups=fea_dim), |
| 418 | nn.Conv2d(fea_dim, qkv_dim, 1, bias=True), |
| 419 | ) |
| 420 | self.kv = nn.Sequential( |
| 421 | nn.Conv2d(cond_dim, cond_dim, 3, 1, 1, bias=False, groups=cond_dim), |
| 422 | nn.Conv2d(cond_dim, qkv_dim * 2, 1, bias=True), |
| 423 | ) |
| 424 | self.nheads = nheads |
| 425 | self.scale = 1 / math.sqrt(qkv_dim // nheads) |
| 426 | |
| 427 | self.attn_out = nn.Conv2d(qkv_dim, dim_out, 1, bias=True) |
| 428 | self.attn_res = ( |
| 429 | nn.Conv2d(fea_dim, dim_out, 1, bias=True) |
| 430 | if fea_dim != dim_out |
| 431 | else nn.Identity() |
| 432 | ) |
| 433 | |
| 434 | self.ffn = nn.Sequential( |
| 435 | nn.Conv2d(dim_out, dim_out * 2, 3, 1, 1, bias=False), |
| 436 | nn.SiLU(), |
| 437 | nn.Conv2d(dim_out * 2, dim_out, 3, 1, 1, bias=False), |
| 438 | nn.Conv2d(dim_out, dim_out, 1, bias=True), |
| 439 | ) |
| 440 | self.ffn_drop_path = DropPath(drop_prob=drop_path_prob) |
| 441 | |
| 442 | def forward(self, x, cond): |
| 443 | x = self.prenorm_x(x) |
| 444 | # cond = self.prenorm_cond(cond) |
| 445 | |
| 446 | q = self.q(x) |
| 447 | k, v = self.kv(cond).chunk(2, dim=1) |
| 448 | |
| 449 | q, k, v = map(lambda in_qkv: F.normalize(in_qkv, dim=1), (q, k, v)) |
| 450 | |
| 451 | # convert to freq space |
| 452 | q = torch.fft.rfft2(q, dim=(-2, -1), norm="ortho") # b, c, h, w/2+1 |
| 453 | k = torch.fft.rfft2(k, dim=(-2, -1), norm="ortho") # b, c, h, w/2+1 |
| 454 | v = torch.fft.rfft2(v, dim=(-2, -1), norm="ortho") # b, c, h, w/2+1 |
| 455 | |
| 456 | # amp and phas attention |
nothing calls this directly
no outgoing calls
no test coverage detected