| """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 |
|
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| REPO = "Phips/NERVE" |
| MODELS: dict[str, tuple[str, int]] = { |
| "NERVE 2x — OTF GAN (blurry / noisy / JPEG input)": ( |
| "models/2x_NERVE_OTF_gan.safetensors", |
| 2, |
| ), |
| "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 |
| 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).to(DEVICE) |
| with torch.autocast("cuda", dtype=torch.float16, enabled=DEVICE.type == "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).to(DEVICE) |
| 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**, and download the full-quality PNG. |
| |
| **What is NERVE?** A small (~1.8M params) pure-convolution network built for |
| *ease of use*, not benchmark chasing: one file, one variant, and a checkpoint |
| that converts cleanly to dynamic ONNX, TensorRT and ncnn — no fused/unfused |
| pairs, no headaches. |
| |
| - **2x / 4x release** → for **clean** images. |
| - **2x / 4x OTF GAN** → for **degraded** images (blurry, noisy, JPEG-compressed). |
| The **2x OTF GAN** is the default; other models stay selectable. |
| |
| Models, ONNX/ncnn files, training configs and the architecture live at |
| **[huggingface.co/Phips/NERVE](https://huggingface.co/Phips/NERVE)**. |
| |
| **Links:** [NERVE](https://huggingface.co/Phips/NERVE) · [HEART](https://huggingface.co/Phips/HEART) · [NERVE demo](https://huggingface.co/spaces/Phips/nerve) · [HEART demo](https://huggingface.co/spaces/Phips/heart) · [convert a checkpoint](https://huggingface.co/spaces/Phips/nerve-convert) · [local app](https://github.com/Phhofm/nerve-app) · [train NERVE](https://github.com/Phhofm/nerve-train) · [train HEART](https://github.com/Phhofm/heart-train) |
| [try it online (ZeroGPU)](https://huggingface.co/spaces/Phips/nerve) · |
| [convert a checkpoint](https://huggingface.co/spaces/Phips/nerve-convert) · |
| [local app](https://github.com/Phhofm/nerve-app) · |
| [train your own (Colab)](https://github.com/Phhofm/nerve-train) |
| """ |
| ) |
| 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)[0], |
| 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]) |
| gr.Examples( |
| examples=[ |
| ["examples/noisy.png"], |
| ["examples/realphoto.png"], |
| ["examples/ani.png"], |
| ["examples/text.png"], |
| ], |
| inputs=[inp], |
| label="Try one of these", |
| ) |
|
|
| if __name__ == "__main__": |
| demo.queue().launch() |
|
|
