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