Method forward

models/sr3.py:526–567

(self, x, cond)

        self.ffn_drop_path = DropPath(drop_prob=drop_path_prob)

    def forward(self, x, cond):
        x = self.prenorm_x(x)
        # cond = self.prenorm_cond(cond)

        q = self.q(x)
        k, v = self.kv(cond).chunk(2, dim=1)

        q, k, v = map(lambda in_qkv: F.normalize(in_qkv, dim=1), (q, k, v))

        q = q.softmax(dim=-2)
        k = k.softmax(dim=-1)

        # convert to freq space
        q = torch.fft.rfft2(q, dim=(-2, -1), norm="ortho")  # b, c, h, w/2+1
        k = torch.fft.rfft2(k, dim=(-2, -1), norm="ortho")  # b, c, h, w/2+1
        v = torch.fft.rfft2(v, dim=(-2, -1), norm="ortho")  # b, c, h, w/2+1

        b, c, xf, yf = q.shape

        q, k, v = map(
            lambda in_x: rearrange(
                in_x, "b (h c) xf yf -> b h c (xf yf)", h=self.nheads
            ),
            (q, k, v),
        )
        q = q * self.scale
        # c x c attn map
        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)
        # h w fused feature map
        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
        out = rearrange(
            out, "n h c (xf yf) -> n (h c) xf yf", xf=xf, yf=yf, h=self.nheads
        )

        # convert to rgb space
        out = torch.fft.irfft2(out, dim=(-2, -1), norm="ortho")

        attn_out = self.attn_out(out) + self.attn_res(x)

        # ffn
        ffn_out = self.ffn_drop_path(self.ffn(attn_out)) + attn_out
        return ffn_out


class WrappedCondInj(nn.Module):