| """Gradio Space demo for the Universal Upscale model."""
|
| from pathlib import Path
|
|
|
| import gradio_client.utils as gutils
|
|
|
|
|
| _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
|
|
|
|
|
| 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()
|
|
|
