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