initial NERVE upscale Space

README.md

@@ -1,13 +1,20 @@
 ---
-title: Nerve
-emoji: ⚡
-colorFrom: red
-colorTo: pink
+title: NERVE
+colorFrom: indigo
+colorTo: green
 sdk: gradio
-sdk_version: 6.28.0
-python_version: '3.13'
+sdk_version: 5.25.2
 app_file: app.py
 pinned: false
+suggested_hardware: zero-a10g
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# NERVE — image upscaling & restoration
+
+A tiny demo for the released [NERVE](https://huggingface.co/Phips/NERVE) models.
+
+- **2x / 4x release** — for clean images.
+- **4x OTF GAN** — for degraded images (blurry, noisy, JPEG-compressed).
+
+Input: an image. Output: an uncompressed PNG you can download.
+Runs on ZeroGPU; large images are processed in tiles.

app.py

@@ -0,0 +1,104 @@
+"""NERVE — upscale & restore images with the released NERVE models.
+
+Simple Space: pick a model, drop an image, get an uncompressed PNG back.
+"""
+
+from __future__ import annotations
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from safetensors.torch import load_file
+
+from nerve_arch import nerve
+
+REPO = "Phips/NERVE"
+MODELS: dict[str, tuple[str, int]] = {
+    "NERVE 2x — release (clean input)": ("models/2x_NERVE_release.safetensors", 2),
+    "NERVE 4x — release (clean input)": ("models/4x_NERVE_release.safetensors", 4),
+    "NERVE 4x — OTF GAN (blurry / noisy / JPEG input)": (
+        "models/4x_NERVE_OTF_gan.safetensors",
+        4,
+    ),
+}
+
+_cache: dict[str, torch.nn.Module] = {}
+
+
+def get_model(key: str) -> torch.nn.Module:
+    if key not in _cache:
+        path, scale = MODELS[key]
+        model = nerve(scale=scale)
+        model.load_state_dict(
+            load_file(hf_hub_download(REPO, path), device="cpu"), strict=True
+        )
+        _cache[key] = model.eval()
+    return _cache[key]
+
+
+@torch.inference_mode()
+def _run(model: torch.nn.Module, img: np.ndarray, scale: int) -> np.ndarray:
+    """img: float32 HWC RGB in [0,1]; returns float32 HWC RGB in [0,1]."""
+    h, w = img.shape[:2]
+    tile = 512  # LR tile size; keeps VRAM low on ZeroGPU
+    overlap = 32
+    out = np.zeros((h * scale, w * scale, 3), dtype=np.float32)
+    for y in range(0, h, tile - overlap):
+        for x in range(0, w, tile - overlap):
+            y1, x1 = min(y + tile, h), min(x + tile, w)
+            y0, x0 = max(0, y1 - tile), max(0, x1 - tile)
+            crop = img[y0:y1, x0:x1]
+            t = torch.from_numpy(crop.transpose(2, 0, 1)).unsqueeze(0).cuda()
+            y_hat = model(t).clamp(0, 1).squeeze(0).permute(1, 2, 0).cpu().numpy()
+            out[y0 * scale : y1 * scale, x0 * scale : x1 * scale] = y_hat
+    return out
+
+
+@spaces.GPU(duration=90)
+def upscale(image: Image.Image, model_key: str) -> Image.Image:
+    if image is None:
+        raise gr.Error("Please upload an image first.")
+    model = get_model(model_key).cuda()
+    scale = MODELS[model_key][1]
+
+    rgb = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0
+    result = _run(model, rgb, scale)
+    return Image.fromarray((result * 255.0 + 0.5).astype(np.uint8), mode="RGB")
+
+
+with gr.Blocks(title="NERVE", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # NERVE
+        Lightweight super-resolution / restoration. Drop an image, pick a model,
+        press **Upscale**, download the full-quality PNG.
+
+        - **release** models are for **clean** images.
+        - the **OTF GAN** model is for **degraded** images (blurry, noisy, JPEG-y).
+        """
+    )
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(type="pil", label="Input image", height=320)
+            model = gr.Dropdown(
+                choices=list(MODELS),
+                value=list(MODELS)[1],
+                label="Model",
+            )
+            btn = gr.Button("Upscale", variant="primary")
+        with gr.Column():
+            out = gr.Image(type="pil", label="Output (uncompressed PNG)", height=320)
+            dl = gr.File(label="Download PNG")
+
+    def _wrap(image, model_key):
+        result = upscale(image, model_key)
+        result.save("/tmp/nerve_output.png")
+        return result, "/tmp/nerve_output.png"
+
+    btn.click(_wrap, inputs=[inp, model], outputs=[out, dl])
+
+if __name__ == "__main__":
+    demo.queue().launch()

nerve_arch.py

@@ -0,0 +1,97 @@
+# NERVE — Norm-free Efficient Restoration for Various Edge devices.
+#
+# A lightweight super-resolution network designed for ease of use rather than
+# benchmark chasing: pure convolution (Conv-ReLU-Conv residual blocks, no
+# normalization, no attention, no gating), a single readable file, and standard
+# ops only (Conv / Add / ReLU / PixelShuffle / bicubic Resize) so the checkpoint
+# exports to dynamic-shape ONNX and converts to NCNN/CoreML without custom ops
+# or fused/unfused checkpoint pairs.
+#
+# A 5x5 stem opens the receptive field for real-world degradations; a global
+# residual and a bicubic input residual let the net learn only what bicubic
+# upscaling misses. Default dim=64 / n_blocks=24 (~1.8M params at 4x).
+
+from __future__ import annotations
+
+import math
+
+import torch
+from spandrel.util import store_hyperparameters
+from torch import Tensor, nn
+from torch.nn import functional as F  # noqa: N812
+
+def _icnr_init(conv: nn.Conv2d, upscale: int) -> None:
+    """ICNR initialization for a Conv2d that feeds a PixelShuffle(upscale).
+
+    Every sub-pixel phase starts from the same base kernel, so the upsampling
+    head cannot begin with a checkerboard bias. Init-only: no graph or runtime
+    cost, and no effect once a checkpoint is loaded.
+    """
+    out_c = conv.out_channels
+    sub = upscale * upscale
+    assert out_c % sub == 0
+    base_c = out_c // sub
+    base = nn.init.kaiming_uniform_(conv.weight[:base_c].clone(), a=math.sqrt(5))
+    with torch.no_grad():
+        conv.weight.data.zero_()
+        for i in range(sub):
+            conv.weight.data[i::sub] = base
+    if conv.bias is not None:
+        nn.init.zeros_(conv.bias)
+
+
+class NerveBlock(nn.Module):
+    """Conv-ReLU-Conv residual block. No normalization, no shortcuts."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(dim, dim, 3, padding=1, bias=False)
+        self.conv2 = nn.Conv2d(dim, dim, 3, padding=1, bias=False)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.conv2(F.relu(self.conv1(x)))
+
+
+@store_hyperparameters()
+class NERVE(nn.Module):
+    """5x5 stem -> n_blocks residual blocks -> PixelShuffle head -> + bicubic(x)."""
+
+    hyperparameters = {}  # noqa: RUF012
+
+    def __init__(
+        self,
+        *,
+        in_ch: int = 3,
+        out_ch: int = 3,
+        dim: int = 64,
+        n_blocks: int = 24,
+        upscale: int = 4,
+    ) -> None:
+        super().__init__()
+        self.upscale = upscale
+
+        self.stem = nn.Conv2d(in_ch, dim, 5, padding=2, bias=False)
+        self.blocks = nn.Sequential(*[NerveBlock(dim) for _ in range(n_blocks)])
+        head = nn.Conv2d(dim, out_ch * upscale**2, 3, padding=1, bias=False)
+        self.upsampler = nn.Sequential(head, nn.PixelShuffle(upscale))
+        _icnr_init(head, upscale)
+
+    def forward(self, x: Tensor) -> Tensor:
+        feat = self.stem(x)
+        out = feat + self.blocks(feat)
+        out = self.upsampler(out)
+        base = F.interpolate(
+            x, scale_factor=self.upscale, mode="bicubic", align_corners=False
+        )
+        return out + base
+
+
+def nerve(
+    scale: int = 4,
+    *,
+    in_ch: int = 3,
+    out_ch: int = 3,
+    dim: int = 64,
+    n_blocks: int = 24,
+) -> NERVE:
+    return NERVE(in_ch=in_ch, out_ch=out_ch, dim=dim, n_blocks=n_blocks, upscale=scale)

requirements.txt

@@ -0,0 +1,7 @@
+gradio==5.25.2
+spaces
+torch
+safetensors
+huggingface_hub
+Pillow
+numpy