--- library_name: pytorch tags: - image-restoration - denoising - super-resolution - image-enhancement license: apache-2.0 datasets: - FlameF0X/nano-banana-pro-gen-zh-en --- # 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 ``` ```py # 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") ```