File size: 2,299 Bytes
40adac6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Gradio Space demo for the Universal Upscale model."""
from pathlib import Path
import gradio as gr
import numpy as np
import torch
from PIL import Image
MODEL_PATH = Path(__file__).resolve().parent / "model_scripted.pt"
MODEL = torch.jit.load(str(MODEL_PATH), map_location="cpu")
MODEL.eval()
# Free Spaces run on CPU: keep the output bounded so the demo stays responsive.
MAX_SIDE = 2048
def upscale(image, scale, target_long_side):
if image is None:
return None
img = image if isinstance(image, Image.Image) else Image.open(image)
img = img.convert("RGB")
if target_long_side:
factor = float(target_long_side) / float(max(img.size))
else:
factor = float(scale)
out_w, out_h = max(1, int(round(img.width * factor))), max(1, int(round(img.height * factor)))
if max(out_w, out_h) > MAX_SIDE:
shrink = MAX_SIDE / float(max(out_w, out_h))
out_w, out_h = max(1, int(out_w * shrink)), max(1, int(out_h * shrink))
lr = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0).permute(2, 0, 1).contiguous().unsqueeze(0)
out = MODEL(lr, out_h, out_w)
arr = (out.squeeze(0).permute(1, 2, 0).numpy() * 255.0).round().astype(np.uint8)
return Image.fromarray(arr)
with gr.Blocks(title="Universal Upscale") as demo:
gr.Markdown(
"### Universal Upscale\n"
"Arbitrary-scale super-resolution: one checkpoint, any scale (1x-10x) or any target size. "
"Tiled inference is built into the model, so large outputs do not blow up memory."
)
with gr.Row():
inp = gr.Image(type="pil", label="Input")
out = gr.Image(type="pil", label="Upscaled output")
with gr.Row():
scale = gr.Slider(1, 10, value=4, step=0.5, label="Scale")
long_side = gr.Number(value=None, label="Or target long side (px)", precision=0)
run = gr.Button("Upscale", variant="primary")
run.click(upscale, inputs=[inp, scale, long_side], outputs=out)
example_files = [
name
for name in ("examples/input.png", "examples/demo_input.png")
if (Path(__file__).resolve().parent / name).exists()
]
if example_files:
gr.Examples(examples=example_files, inputs=inp)
demo.launch()
|