Real World Super Resolution With Rectified Flow Degradation

"""
Unsupervised real-world image super-resolution via rectified flow degradation modeling.
Reference PyTorch implementation of the pipeline in
    Zhou, Qin, Chen, He, Zhu. Computational Visual Media 12(4), 1035-1051, 2026.
    https://doi.org/10.26599/CVM.2026.9450559

Pipeline (Algorithm 1 of the paper)
    1. DT-LR      repeated bilinear up and down sampling (N = 10) pushes differently degraded
                  LR images toward a common, "degradation transformed" state
    2. FGDM       Fourier prior guided degradation module. AENet (conv + residual state space
                  blocks) processes DT-LR, its FFT amplitude is kept, the FFT phase is taken
                  from a structure reference, and the IFFT gives a preliminary degraded LR.
                  Trained with L1 against the real LR image.
    3. RFDM       rectified flow degradation module. A time conditioned UNet learns the
                  velocity from X0 = FGDM output + lambda * noise to X1 = real LR (Eq. 5-7).
                  Synthesis integrates the ODE with K = 20 Euler steps (Eq. 8).
    4. SR model   any off-the-shelf SR network is trained on (synthetic LR, HR) pairs.

Honest notes
    * The residual state space block (RSSB) follows MambaIR. Its selective scan is written
      in plain PyTorch with a sequential loop, which is exact but slow. For full size training
      swap `selective_scan` for the CUDA kernel in the `mamba_ssm` package.
    * Items the paper does not specify, and which we therefore chose: the up/down scale
      factor inside DT-LR (2), the AENet global residual, the UNet widths, learning rates,
      and the SR network used in the smoke test (a tiny residual CNN standing in for SwinIR,
      Real-ESRGAN or StableSR).
    * LPIPS and FID need pretrained networks; `evaluate` reports PSNR and SSIM on the
      Y channel of YCbCr over a central crop, as the paper does for fidelity metrics.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# ---------------------------------------------------------------------------
# 1. Degradation transformed LR (DT-LR)
# ---------------------------------------------------------------------------

def dt_lr(x: torch.Tensor, n: int = 10, factor: int = 2, mode: str = "bilinear") -> torch.Tensor:
    """Apply n rounds of (upsample by `factor`, then downsample back). Output keeps x's size."""
    h, w = x.shape[-2:]
    for _ in range(n):
        x = F.interpolate(x, scale_factor=factor, mode=mode, align_corners=False)
        x = F.interpolate(x, size=(h, w), mode=mode, align_corners=False)
    return x


def bilinear_down(y: torch.Tensor, scale: int = 4) -> torch.Tensor:
    return F.interpolate(y, scale_factor=1 / scale, mode="bilinear", align_corners=False, antialias=True)

# ---------------------------------------------------------------------------
# 2. Fourier helpers
# ---------------------------------------------------------------------------

def amp_phase(x: torch.Tensor, eps: float = 1e-8):
    X = torch.fft.fft2(x, norm="ortho")
    amp = torch.sqrt(X.real ** 2 + X.imag ** 2 + eps)
    pha = torch.atan2(X.imag, X.real + 0.0)
    return amp, pha


def combine(amp: torch.Tensor, pha: torch.Tensor) -> torch.Tensor:
    return torch.fft.ifft2(torch.polar(amp, pha), norm="ortho").real

# ---------------------------------------------------------------------------
# 3. Residual state space block (MambaIR style) for AENet
# ---------------------------------------------------------------------------

def selective_scan(u, delta, A, B, C, D):
    """Mamba S6 recurrence, h_t = exp(delta_t A) h_{t-1} + delta_t B_t u_t,  y_t = C_t h_t + D u_t.
    u, delta (b, d, L)   A (d, n)   B, C (b, n, L)   D (d,)"""
    b, d, L = u.shape
    dA = torch.exp(delta.unsqueeze(-1) * A[None, :, None, :])      # (b, d, L, n)
    dBu = delta.unsqueeze(-1) * B.transpose(1, 2).unsqueeze(1) * u.unsqueeze(-1)
    h = u.new_zeros(b, d, A.shape[1])
    ys = []
    for t in range(L):
        h = dA[:, :, t] * h + dBu[:, :, t]
        ys.append((h * C[:, :, t].unsqueeze(1)).sum(-1))
    return torch.stack(ys, -1) + u * D[None, :, None]


class SS2D(nn.Module):
    """2D selective scan over four directions (row major, column major, and their reverses)."""

    def __init__(self, dim, d_state=16, expand=2):
        super().__init__()
        di = dim * expand
        self.di, self.n = di, d_state
        self.dt_rank = math.ceil(dim / 16)
        self.in_proj = nn.Linear(dim, 2 * di)
        self.dwconv = nn.Conv2d(di, di, 3, padding=1, groups=di)
        self.x_proj = nn.Parameter(torch.randn(4, self.dt_rank + 2 * d_state, di) * di ** -0.5)
        self.dt_proj_w = nn.Parameter(torch.randn(4, di, self.dt_rank) * self.dt_rank ** -0.5)
        dt = torch.exp(torch.rand(4, di) * (math.log(0.1) - math.log(0.001)) + math.log(0.001))
        self.dt_proj_b = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))    # inverse softplus
        self.A_log = nn.Parameter(torch.log(torch.arange(1, d_state + 1).float()).repeat(4 * di, 1))
        self.D = nn.Parameter(torch.ones(4 * di))
        self.out_norm = nn.LayerNorm(di)
        self.out_proj = nn.Linear(di, dim)

    def forward(self, x):                                                   # x (b, h, w, c)
        b, h, w, _ = x.shape
        xz = self.in_proj(x)
        xi, z = xz.chunk(2, -1)
        xi = F.silu(self.dwconv(xi.permute(0, 3, 1, 2)))                    # (b, di, h, w)
        L = h * w
        row = xi.flatten(2)
        col = xi.transpose(2, 3).flatten(2)
        seqs = torch.stack([row, col, row.flip(-1), col.flip(-1)], 1)       # (b, 4, di, L)
        A = -torch.exp(self.A_log).view(4, self.di, self.n)
        outs = []
        for k in range(4):
            u = seqs[:, k]
            proj = torch.einsum("bdl,cd->bcl", u, self.x_proj[k])
            dt_r, Bm, Cm = torch.split(proj, [self.dt_rank, self.n, self.n], 1)
            delta = F.softplus(torch.einsum("brl,dr->bdl", dt_r, self.dt_proj_w[k]) + self.dt_proj_b[k][None, :, None])
            outs.append(selective_scan(u, delta, A[k], Bm, Cm, self.D.view(4, self.di)[k]))
        y = outs[0] + outs[2].flip(-1)
        yc = outs[1] + outs[3].flip(-1)
        y = y + yc.view(b, self.di, w, h).transpose(2, 3).flatten(2)
        y = self.out_norm(y.transpose(1, 2).view(b, h, w, self.di))
        return self.out_proj(y * F.silu(z))


class ChannelAttention(nn.Module):
    def __init__(self, c, r=16):
        super().__init__()
        self.net = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Conv2d(c, max(c // r, 4), 1), nn.ReLU(inplace=True),
                                 nn.Conv2d(max(c // r, 4), c, 1), nn.Sigmoid())

    def forward(self, x):
        return x * self.net(x)


class RSSB(nn.Module):
    """Two residual branches with learnable scales (MambaIR).
    Branch 1, LayerNorm then SS2D for long range context.
    Branch 2, LayerNorm then conv then channel attention for local detail."""

    def __init__(self, dim):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
        self.ss2d = SS2D(dim)
        self.conv = nn.Sequential(nn.Conv2d(dim, dim // 3, 3, padding=1), nn.GELU(), nn.Conv2d(dim // 3, dim, 3, padding=1))
        self.ca = ChannelAttention(dim)
        self.s1 = nn.Parameter(torch.ones(dim))
        self.s2 = nn.Parameter(torch.ones(dim))

    def forward(self, x):                                                   # (b, c, h, w)
        t = x.permute(0, 2, 3, 1)
        t = t * self.s1 + self.ss2d(self.ln1(t))
        c = self.ln2(t).permute(0, 3, 1, 2)
        c = self.ca(self.conv(c)).permute(0, 2, 3, 1)
        t = t * self.s2 + c
        return t.permute(0, 3, 1, 2)


class AENet(nn.Module):
    """Amplitude enhancement network. Conv, n RSSBs, conv. Paper default n = 6, width 96."""

    def __init__(self, width=96, n_blocks=6, ch=3):
        super().__init__()
        self.head = nn.Conv2d(ch, width, 3, padding=1)
        self.body = nn.Sequential(*[RSSB(width) for _ in range(n_blocks)])
        self.tail = nn.Conv2d(width, ch, 3, padding=1)

    def forward(self, x):
        f = self.head(x)
        return x + self.tail(self.body(f) + f)                            # global residual (our choice)


class FGDM(nn.Module):
    """Fourier prior guided degradation module.
    amplitude  <- FFT(AENet(DT-LR))       learns what the degradation does to the spectrum
    phase      <- FFT(structure reference) real LR during training, bilinear LR of HR at synthesis
    output      = IFFT(amplitude, phase)"""

    def __init__(self, width=96, n_blocks=6):
        super().__init__()
        self.aenet = AENet(width, n_blocks)

    def forward(self, x_dt, phase_ref):
        amp, _ = amp_phase(self.aenet(x_dt))
        _, pha = amp_phase(phase_ref)
        return combine(amp, pha)

# ---------------------------------------------------------------------------
# 4. RFDM: rectified flow with a time conditioned UNet velocity field
# ---------------------------------------------------------------------------

def time_embedding(t, dim):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
    a = t[:, None] * 1000 * freqs[None]
    return torch.cat([a.sin(), a.cos()], -1)


class ResBlock(nn.Module):
    def __init__(self, cin, cout, tdim):
        super().__init__()
        self.n1, self.n2 = nn.GroupNorm(8, cin), nn.GroupNorm(8, cout)
        self.c1, self.c2 = nn.Conv2d(cin, cout, 3, padding=1), nn.Conv2d(cout, cout, 3, padding=1)
        self.t = nn.Linear(tdim, cout)
        self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity()

    def forward(self, x, temb):
        h = self.c1(F.silu(self.n1(x))) + self.t(temb)[:, :, None, None]
        h = self.c2(F.silu(self.n2(h)))
        return h + self.skip(x)


class VelocityUNet(nn.Module):
    def __init__(self, ch=3, base=64, mults=(1, 2, 4)):
        super().__init__()
        self.tdim = base * 4
        self.tmlp = nn.Sequential(nn.Linear(base, self.tdim), nn.SiLU(), nn.Linear(self.tdim, self.tdim))
        self.base = base
        self.inp = nn.Conv2d(ch, base, 3, padding=1)
        chs = [base * m for m in mults]
        self.down, self.pools = nn.ModuleList(), nn.ModuleList()
        c = base
        for i, co in enumerate(chs):
            self.down.append(ResBlock(c, co, self.tdim))
            self.pools.append(nn.Conv2d(co, co, 3, stride=2, padding=1) if i < len(chs) - 1 else nn.Identity())
            c = co
        self.mid = ResBlock(c, c, self.tdim)
        self.up, self.ups = nn.ModuleList(), nn.ModuleList()
        for i, co in reversed(list(enumerate(chs))):
            self.up.append(ResBlock(c + co, co, self.tdim))
            self.ups.append(nn.Upsample(scale_factor=2, mode="nearest") if i > 0 else nn.Identity())
            c = co
        self.out = nn.Sequential(nn.GroupNorm(8, c), nn.SiLU(), nn.Conv2d(c, ch, 3, padding=1))

    def forward(self, x, t):
        temb = self.tmlp(time_embedding(t, self.base))
        h = self.inp(x)
        skips = []
        for blk, pool in zip(self.down, self.pools):
            h = blk(h, temb)
            skips.append(h)
            h = pool(h)
        h = self.mid(h, temb)
        for blk, up in zip(self.up, self.ups):
            h = blk(torch.cat([h, skips.pop()], 1), temb)
            h = up(h)
        return self.out(h)


class RFDM(nn.Module):
    def __init__(self, base=64, lam=0.1, steps=20):
        super().__init__()
        self.v = VelocityUNet(base=base)
        self.lam, self.steps = lam, steps

    def loss(self, x_bar, x_real):
        """Eq. 5-7. X0 = FGDM output + lambda * n,  X1 = real LR,  L2 on the constant velocity X1 - X0."""
        x0 = x_bar + self.lam * torch.randn_like(x_bar)
        t = torch.rand(x0.shape[0], device=x0.device)
        xt = t[:, None, None, None] * x_real + (1 - t[:, None, None, None]) * x0
        return F.mse_loss(self.v(xt, t), x_real - x0)

    @torch.no_grad()
    def sample(self, x_bar, steps=None):
        """Eq. 8. Euler integration from X0 = x_bar + lambda * n to X1 in K steps."""
        K = steps or self.steps
        x = x_bar + self.lam * torch.randn_like(x_bar)
        for i in range(K):
            t = torch.full((x.shape[0],), i / K, device=x.device)
            x = x + self.v(x, t) / K
        return x

# ---------------------------------------------------------------------------
# 5. The full degradation pipeline
# ---------------------------------------------------------------------------

class RFDegradation(nn.Module):
    def __init__(self, width=96, n_blocks=6, unet_base=64, lam=0.1, steps=20, n_updown=10, scale=4):
        super().__init__()
        self.fgdm = FGDM(width, n_blocks)
        self.rfdm = RFDM(unet_base, lam, steps)
        self.n_updown, self.scale = n_updown, scale

    def fgdm_loss(self, x_real):
        """Stage 1. L1 between FGDM(DT-LR(real), phase(real)) and the real LR."""
        x_bar = self.fgdm(dt_lr(x_real, self.n_updown), x_real)
        return F.l1_loss(x_bar, x_real)

    def rfdm_loss(self, x_real):
        """Stage 2. FGDM is frozen, RFDM learns the flow from its output to the real LR."""
        with torch.no_grad():
            x_bar = self.fgdm(dt_lr(x_real, self.n_updown), x_real)
        return self.rfdm.loss(x_bar, x_real)

    @torch.no_grad()
    def synthesize(self, y_hr):
        """HR -> bilinear LR -> DT-LR -> FGDM (phase of the bilinear LR) -> RFDM ODE -> synthetic LR."""
        lr_bi = bilinear_down(y_hr, self.scale)
        x_bar = self.fgdm(dt_lr(lr_bi, self.n_updown), lr_bi)
        return self.rfdm.sample(x_bar).clamp(0, 1)

# ---------------------------------------------------------------------------
# 6. A stand in SR network, training loops and evaluation
# ---------------------------------------------------------------------------

class TinySR(nn.Module):
    """Small residual CNN with pixel shuffle. Replace with SwinIR, Real-ESRGAN or StableSR."""

    def __init__(self, scale=4, c=48, n=6):
        super().__init__()
        self.head = nn.Conv2d(3, c, 3, padding=1)
        self.body = nn.Sequential(*[nn.Sequential(nn.Conv2d(c, c, 3, padding=1), nn.ReLU(True), nn.Conv2d(c, c, 3, padding=1)) for _ in range(n)])
        self.up = nn.Sequential(nn.Conv2d(c, 3 * scale * scale, 3, padding=1), nn.PixelShuffle(scale))
        self.scale = scale

    def forward(self, x):
        f = self.head(x)
        for blk in self.body:
            f = f + blk(f)
        return self.up(f) + F.interpolate(x, scale_factor=self.scale, mode="bilinear", align_corners=False)


def augment(x):
    """Random rotation by 0, 90, 180 or 270 degrees and random horizontal flip (Sec. 4.1)."""
    k = torch.randint(4, (1,)).item()
    x = torch.rot90(x, k, (-2, -1))
    return x.flip(-1) if torch.rand(1).item() < 0.5 else x


def train_degradation(model, real_lr, epochs_fgdm=1000, epochs_rfdm=1000, bs=8, lr=2e-4, log=100):
    """Paper setting, 1000 epochs each with Adam(beta1 0.9, beta2 0.999). Learning rate not reported."""
    opt = torch.optim.Adam(model.fgdm.parameters(), lr=lr, betas=(0.9, 0.999))
    for ep in range(1, epochs_fgdm + 1):
        for i in range(0, len(real_lr), bs):
            loss = model.fgdm_loss(augment(real_lr[i:i + bs]))
            opt.zero_grad(); loss.backward(); opt.step()
        if log and ep % log == 0:
            print(f"FGDM epoch {ep}  L1 {loss.item():.4f}")
    model.fgdm.requires_grad_(False)
    opt = torch.optim.Adam(model.rfdm.parameters(), lr=lr, betas=(0.9, 0.999))
    for ep in range(1, epochs_rfdm + 1):
        for i in range(0, len(real_lr), bs):
            loss = model.rfdm_loss(augment(real_lr[i:i + bs]))
            opt.zero_grad(); loss.backward(); opt.step()
        if log and ep % log == 0:
            print(f"RFDM epoch {ep}  L2 {loss.item():.4f}")
    return model


def train_sr(sr, lr_imgs, hr_imgs, epochs=100, bs=8, lr=2e-4):
    opt = torch.optim.Adam(sr.parameters(), lr=lr)
    for _ in range(epochs):
        perm = torch.randperm(len(lr_imgs))
        for i in range(0, len(lr_imgs), bs):
            idx = perm[i:i + bs]
            loss = F.l1_loss(sr(lr_imgs[idx]), hr_imgs[idx])
            opt.zero_grad(); loss.backward(); opt.step()
    return sr


def rgb_to_y(x):
    return (16 + 65.481 * x[:, 0] + 128.553 * x[:, 1] + 24.966 * x[:, 2]) / 255.0


def center_crop(x, frac=0.5):
    h, w = x.shape[-2:]
    ch, cw = int(h * frac), int(w * frac)
    return x[..., (h - ch) // 2:(h - ch) // 2 + ch, (w - cw) // 2:(w - cw) // 2 + cw]


def psnr_y(a, b):
    return (10 * torch.log10(1.0 / F.mse_loss(rgb_to_y(a), rgb_to_y(b), reduction="none").mean((-2, -1)))).mean()


def ssim_y(a, b, win=11):
    a, b = rgb_to_y(a)[:, None], rgb_to_y(b)[:, None]
    g = torch.exp(-((torch.arange(win) - win // 2) ** 2) / (2 * 1.5 ** 2)); g = g / g.sum()
    k = (g[:, None] * g[None])[None, None]
    f = lambda t: F.conv2d(t, k, padding=win // 2)
    ma, mb = f(a), f(b)
    va, vb, cab = f(a * a) - ma ** 2, f(b * b) - mb ** 2, f(a * b) - ma * mb
    C1, C2 = 0.01 ** 2, 0.03 ** 2
    return (((2 * ma * mb + C1) * (2 * cab + C2)) / ((ma ** 2 + mb ** 2 + C1) * (va + vb + C2))).mean()


@torch.no_grad()
def evaluate(sr, lr_test, hr_test):
    out = sr(lr_test).clamp(0, 1)
    out, hr = center_crop(out), center_crop(hr_test)
    return psnr_y(out, hr).item(), ssim_y(out, hr).item()

# ---------------------------------------------------------------------------
# 7. Smoke test on dummy data
# ---------------------------------------------------------------------------

def smooth_textures(n, size, seed):
    g = torch.Generator().manual_seed(seed)
    x = torch.rand(n, 3, size // 8, size // 8, generator=g)
    x = F.interpolate(x, size=size, mode="bicubic", align_corners=False)
    stripes = torch.sin(torch.linspace(0, 12 * math.pi, size))[None, None, None] * torch.rand(n, 3, 1, 1, generator=g)
    return (0.7 * x + 0.3 * (stripes + 1) / 2).clamp(0, 1)


def hidden_camera(y, seed=0):
    """The 'unknown' real world degradation, used only to create real LR and test pairs.
    Gaussian blur, 4x decimation, a slight color cast and signal dependent noise."""
    g = torch.Generator().manual_seed(seed)
    k = torch.exp(-(torch.arange(7) - 3.0) ** 2 / (2 * 1.6 ** 2)); k = (k[:, None] * k[None]) / k.sum() ** 2
    yb = F.conv2d(F.pad(y, (3, 3, 3, 3), mode="reflect"), k.expand(3, 1, 7, 7), groups=3)
    x = yb[..., ::4, ::4] * torch.tensor([1.03, 1.0, 0.95])[None, :, None, None]
    return (x + 0.03 * torch.sqrt(x.clamp(min=1e-4)) * torch.randn(x.shape, generator=g)).clamp(0, 1)


def hf_energy(x):
    return (x[..., 1:, :] - x[..., :-1, :]).abs().mean().item() + (x[..., 1:] - x[..., :-1]).abs().mean().item()


if __name__ == "__main__":
    torch.manual_seed(0)
    S = 64                                                                   # HR size, LR is S/4

    # Unpaired training data: real LR images come from different scenes than the HR set
    hr_train = smooth_textures(24, S, seed=1)
    real_lr = hidden_camera(smooth_textures(24, S, seed=2), seed=3)
    hr_test = smooth_textures(8, S, seed=4)
    lr_test = hidden_camera(hr_test, seed=5)

    # (a) DT-LR convergence, the observation behind Fig. 3
    y = smooth_textures(8, S, seed=6)
    lr_real, lr_bi = hidden_camera(y, seed=7), bilinear_down(y)
    for n in (0, 1, 5, 10):
        p = psnr_y(dt_lr(lr_real, n), dt_lr(lr_bi, n)).item()
        print(f"DT-LR rounds {n:2d}  PSNR(real, bilinear) {p:6.2f} dB  high frequency energy {hf_energy(dt_lr(lr_real, n)):.4f}")

    # (b) Fourier amplitude and phase swap, the observation behind Fig. 5
    a_amp, a_pha = amp_phase(lr_real)
    b_amp, b_pha = amp_phase(lr_bi)
    swapped = combine(b_amp, a_pha)
    print(f"amplitude of bilinear + phase of real, PSNR to bilinear {psnr_y(swapped, lr_bi).item():.2f} dB, to real {psnr_y(swapped, lr_real).item():.2f} dB")

    # (c) Train the degradation modules (tiny sizes and few epochs for the smoke test)
    model = RFDegradation(width=24, n_blocks=2, unet_base=32, lam=0.1, steps=20)
    print("parameters FGDM", sum(p.numel() for p in model.fgdm.parameters()),
          " RFDM", sum(p.numel() for p in model.rfdm.parameters()))
    train_degradation(model, real_lr, epochs_fgdm=20, epochs_rfdm=40, bs=8, lr=1e-3, log=10)

    # (d) Synthesize LR for the HR set and compare statistics with real LR
    syn_lr = model.synthesize(hr_train)
    bic_lr = bilinear_down(hr_train)
    print(f"high frequency energy  real LR {hf_energy(real_lr):.4f}  synthetic LR {hf_energy(syn_lr):.4f}  bilinear LR {hf_energy(bic_lr):.4f}")
    print(f"mean color  real {real_lr.mean((0, 2, 3)).tolist()}")
    print(f"mean color  synthetic {syn_lr.mean((0, 2, 3)).tolist()}")

    # (e) Train the same SR network on synthetic pairs and on bilinear pairs, test on hidden camera LR
    torch.manual_seed(0)
    sr_syn = train_sr(TinySR(), syn_lr, hr_train, epochs=60)
    torch.manual_seed(0)
    sr_bic = train_sr(TinySR(), bic_lr, hr_train, epochs=60)
    print("SR trained on synthetic pairs, test PSNR/SSIM (Y, center crop):", [round(v, 4) for v in evaluate(sr_syn, lr_test, hr_test)])
    print("SR trained on bilinear pairs,  test PSNR/SSIM (Y, center crop):", [round(v, 4) for v in evaluate(sr_bic, lr_test, hr_test)])
    print("shapes", tuple(real_lr.shape), tuple(syn_lr.shape), tuple(hr_train.shape))