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