FlameF0X/Nano-RestoreNet

RestoreNet

Restormer-style image restoration model with a task embedding. Tasks: denoise, upscale, enhance. Input and output are RGB images in [0, 1].

task     method       PSNR↑   SSIM↑  LPIPS↓
denoise  baseline     20.69  0.4472  0.4006
denoise  model        28.46  0.8144  0.1353
upscale  baseline     28.18  0.8609  0.1985
upscale  model        27.77  0.8193  0.2270
enhance  baseline     26.05  0.7697  0.2820
enhance  model        26.46  0.7821  0.3105

# pip install torch torchvision huggingface_hub safetensors pillow
import math
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.functional as TF
from PIL import Image
from huggingface_hub import snapshot_download
from safetensors.torch import load_file

REPO_ID = "FlameF0X/Nano-RestoreNet"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
USE_AMP = DEVICE == "cuda"
TASKS = ["denoise", "upscale", "enhance"]

# ----------------------------------------------------------------------------
# Model definition (copied from the training script)
# ----------------------------------------------------------------------------
class LayerNorm2d(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        return self.norm(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

class MDTA(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.temp = nn.Parameter(torch.ones(heads, 1, 1))
        self.qkv = nn.Conv2d(dim, dim * 3, 1)
        self.qkv_dw = nn.Conv2d(dim * 3, dim * 3, 3, padding=1, groups=dim * 3)
        self.proj = nn.Conv2d(dim, dim, 1)

    def forward(self, x):
        b, c, h, w = x.shape
        q, k, v = self.qkv_dw(self.qkv(x)).chunk(3, dim=1)
        q = F.normalize(q.reshape(b, self.heads, c // self.heads, h * w), dim=-1)
        k = F.normalize(k.reshape(b, self.heads, c // self.heads, h * w), dim=-1)
        v = v.reshape(b, self.heads, c // self.heads, h * w)
        attn = ((q @ k.transpose(-2, -1)) * self.temp).softmax(dim=-1)
        return self.proj((attn @ v).reshape(b, c, h, w))

class GDFN(nn.Module):
    def __init__(self, dim, mult=2.66):
        super().__init__()
        hidden = int(dim * mult)
        self.proj_in = nn.Conv2d(dim, hidden * 2, 1)
        self.dw = nn.Conv2d(hidden * 2, hidden * 2, 3, padding=1, groups=hidden * 2)
        self.proj_out = nn.Conv2d(hidden, dim, 1)

    def forward(self, x):
        x1, x2 = self.dw(self.proj_in(x)).chunk(2, dim=1)
        return self.proj_out(F.gelu(x1) * x2)

class TBlock(nn.Module):
    def __init__(self, dim, emb_dim):
        super().__init__()
        self.norm1 = LayerNorm2d(dim)
        self.attn = MDTA(dim, heads=max(1, dim // 16))
        self.norm2 = LayerNorm2d(dim)
        self.ffn = GDFN(dim)
        self.film = nn.Linear(emb_dim, 4 * dim)
        nn.init.zeros_(self.film.weight)
        nn.init.zeros_(self.film.bias)

    def forward(self, x, t):
        s1, b1, s2, b2 = self.film(t)[:, :, None, None].chunk(4, dim=1)
        x = x + self.attn(self.norm1(x) * (1 + s1) + b1)
        x = x + self.ffn(self.norm2(x) * (1 + s2) + b2)
        return x

class Stage(nn.Module):
    def __init__(self, dim, n, emb_dim):
        super().__init__()
        self.blocks = nn.ModuleList([TBlock(dim, emb_dim) for _ in range(n)])

    def forward(self, x, t):
        for b in self.blocks:
            x = b(x, t)
        return x

class RestoreNet(nn.Module):
    def __init__(self, C=64, blocks=(4, 6, 8), emb_dim=64, num_tasks=3):
        super().__init__()
        self.task_emb = nn.Embedding(num_tasks, emb_dim)
        self.stem = nn.Conv2d(3, C, 3, padding=1)
        self.enc1 = Stage(C, blocks[0], emb_dim)
        self.down1 = nn.Sequential(nn.PixelUnshuffle(2), nn.Conv2d(4 * C, 2 * C, 1))
        self.enc2 = Stage(2 * C, blocks[1], emb_dim)
        self.down2 = nn.Sequential(nn.PixelUnshuffle(2), nn.Conv2d(8 * C, 4 * C, 1))
        self.bottleneck = Stage(4 * C, blocks[2], emb_dim)
        self.up2 = nn.Sequential(nn.Conv2d(4 * C, 8 * C, 1), nn.PixelShuffle(2))
        self.fuse2 = nn.Conv2d(4 * C, 2 * C, 1)
        self.dec2 = Stage(2 * C, blocks[1], emb_dim)
        self.up1 = nn.Sequential(nn.Conv2d(2 * C, 4 * C, 1), nn.PixelShuffle(2))
        self.fuse1 = nn.Conv2d(2 * C, C, 1)
        self.dec1 = Stage(C, blocks[0], emb_dim)
        self.out = nn.Conv2d(C, 3, 3, padding=1)

    def forward(self, x, task):
        t = self.task_emb(task)
        s = self.stem(x)
        e1 = self.enc1(s, t)
        e2 = self.enc2(self.down1(e1), t)
        b = self.bottleneck(self.down2(e2), t)
        d2 = self.dec2(self.fuse2(torch.cat([self.up2(b), e2], dim=1)), t)
        d1 = self.dec1(self.fuse1(torch.cat([self.up1(d2), e1], dim=1)), t)
        return x + self.out(d1)

# ----------------------------------------------------------------------------
# Loading from the Hub
# ----------------------------------------------------------------------------
def load_restore_net(repo_id: str = REPO_ID) -> RestoreNet:
    model_dir = snapshot_download(repo_id)  # downloads/caches config.json, model.safetensors

    with open(f"{model_dir}/config.json") as f:
        c = json.load(f)

    model = RestoreNet(
        C=c["base_channels"],
        blocks=tuple(c["blocks"]),
        emb_dim=c["emb_dim"],
        num_tasks=len(c["tasks"]),
    )
    state = load_file(f"{model_dir}/model.safetensors", device="cpu")
    model.load_state_dict(state)
    model.to(DEVICE).eval()
    print(f"Loaded {repo_id} on {DEVICE} (trained for {c.get('training_steps')} steps)")
    return model

# ----------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------
@torch.no_grad()
def restore(model, img: Image.Image, task: str, scale: int = 2) -> Image.Image:
    assert task in TASKS, f"task must be one of {TASKS}"
    img = img.convert("RGB")
    if task == "upscale":
        img = img.resize((img.width * scale, img.height * scale), Image.BICUBIC)

    x = TF.pil_to_tensor(img).float().div(255).unsqueeze(0).to(DEVICE)
    H, W = x.shape[-2:]
    ph, pw = (4 - H % 4) % 4, (4 - W % 4) % 4
    x = F.pad(x, (0, pw, 0, ph), mode="reflect")

    t = torch.tensor([TASKS.index(task)], device=DEVICE)
    with torch.autocast(device_type=DEVICE, dtype=torch.float16, enabled=USE_AMP):
        y = model(x, t)
    y = y[..., :H, :W].float().clamp(0, 1)[0]
    return TF.to_pil_image(y.cpu())

# ----------------------------------------------------------------------------
# Example usage
# ----------------------------------------------------------------------------
if __name__ == "__main__":
    model = load_restore_net()
    img = Image.open("input.jpg")
    out = restore(model, img, task="denoise")   # or "upscale", "enhance"
    out.save("output.png")