File size: 3,944 Bytes
40adac6 7378a93 40adac6 7378a93 40adac6 7378a93 40adac6 7378a93 40adac6 7378a93 40adac6 7378a93 40adac6 7378a93 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """Gradio Space demo for the Universal Upscale model."""
from pathlib import Path
import gradio_client.utils as gutils
# Monkey-patch gradio_client schema parser to handle boolean additionalProperties in Gradio 5.
_orig_json_schema = gutils._json_schema_to_python_type
def _safe_json_schema_to_python_type(schema, defs=None):
if isinstance(schema, bool):
return "Any" if schema else "Never"
return _orig_json_schema(schema, defs)
gutils._json_schema_to_python_type = _safe_json_schema_to_python_type
import gradio as gr
import numpy as np
import torch
from PIL import Image
try:
import spaces
gpu_decorator = spaces.GPU
except ImportError:
def gpu_decorator(fn):
return fn
MODEL_PATH = Path(__file__).resolve().parent / "model_scripted.pt"
def load_model(device: str):
model = torch.jit.load(str(MODEL_PATH), map_location=device)
model.eval()
return model
MAX_SIDE = 4096
# Largest scale factor verified end-to-end for the published checkpoint. The slider is capped
# here; "Target Long Side" still accepts any size up to MAX_SIDE.
MAX_SCALE = 9.2
def resolve_factor(input_size, scale, target_long_side):
"""Scale factor to apply: an explicit target size wins, scale requests stay under MAX_SCALE."""
target = float(target_long_side) if target_long_side is not None else 0.0
if target > 0:
return target / float(max(input_size)), ""
factor = float(scale if scale is not None else 4.0)
if factor > MAX_SCALE:
return MAX_SCALE, f" (clamped to the tested maximum {MAX_SCALE:g}x)"
return factor, ""
@gpu_decorator
def upscale(image, scale: float = 4.0, target_long_side: float = 0.0):
if image is None:
return None, "No image provided."
img = image if isinstance(image, Image.Image) else Image.open(image)
img = img.convert("RGB")
factor, clamped = resolve_factor(img.size, scale, target_long_side)
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))
device = "cuda" if torch.cuda.is_available() else "cpu"
model = load_model(device)
lr = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0).permute(2, 0, 1).contiguous().unsqueeze(0).to(device)
with torch.no_grad():
out = model(lr, out_h, out_w)
arr = (out.squeeze(0).permute(1, 2, 0).cpu().numpy() * 255.0).round().astype(np.uint8)
res_img = Image.fromarray(arr)
status = f"Input: {img.width}x{img.height} | Output: {res_img.width}x{res_img.height} (Scale: {factor:.2f}x) | Device: {device}{clamped}"
return res_img, status
example_files = [
[str(Path(__file__).resolve().parent / name)]
for name in ("examples/input.png", "examples/demo_input.png")
if (Path(__file__).resolve().parent / name).exists()
]
demo = gr.Interface(
fn=upscale,
inputs=[
gr.Image(type="pil", label="Input Image"),
gr.Slider(minimum=1.0, maximum=9.2, value=4.0, step=0.2, label="Scale Factor"),
gr.Number(value=0, label="Or Target Long Side (px) [0 = use scale factor]"),
],
outputs=[
gr.Image(type="pil", label="Upscaled Image"),
gr.Textbox(label="Status / Image Info"),
],
title="Universal Upscale",
description=(
"Arbitrary-scale super-resolution: upscale images to any scale factor from 1x up to the tested maximum "
"9.2x, or to an exact target resolution. Built-in tiled inference enables high-resolution "
"upscaling without memory issues (outputs are capped at 4096 px on the long side)."
),
examples=example_files if example_files else None,
)
if __name__ == "__main__":
demo.launch()
|