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")