app.py

"""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"
# The distilled student ships beside it as its own TorchScript file: the same torch runtime the
# default engine already uses, so the student lands on the Space GPU. The ONNX build was dropped
# because the Space installs the CPU-only `onnxruntime` wheel, which pinned it to CPU.
DISTILLED_PT_PATH = Path(__file__).resolve().parent / "model_distilled.pt"
# If the Space folder does not carry model_distilled.pt, it is pulled from the distilled repo on first use.
DISTILLED_REPO_ID = "tailot/universal-upscale-distilled"

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

# The distilled student renders ONE factor natively. Any other factor is produced by resampling
# its output: area-downsample below the native factor, bicubic above it. The Space UI locks the
# scale controls to this factor when the distilled model is selected; the only remaining shrink is
# the global MAX_SIDE safety cap.
DISTILLED_NATIVE = 9

DEFAULT_SCALE = 4.0

# Plain strings (not (label, value) tuples) so the dropdown works on any Gradio 5 minor version;
# upscale() maps the label back to a model key.
MODEL_LIIF = "🔍 Universal Upscale - arbitrary scale 1x-9.2x (default)"
MODEL_DISTILLED = "⚡ Universal Upscale Distilled - fixed 9x, runs on GPU"
MODEL_CHOICES = [MODEL_LIIF, MODEL_DISTILLED]


def normalize_model_choice(value) -> str:
    """Map the dropdown selection to a model key.



    Gradio pads short example rows with None, and a Dropdown can hand back either the label or

    the value, so the mapping is deliberately loose: anything mentioning "distilled" (or the old

    ONNX labels) is the fixed-factor student, everything else stays on the published default.

    """
    text = str(value or "").strip().lower()
    return "distilled" if "distill" in text or text in ("onnx", "student", "torchscript") else "liif"

def scale_controls_for_model(model_choice):
    """Lock the scale controls to the distilled student's native factor."""
    if normalize_model_choice(model_choice) == "distilled":
        native = float(DISTILLED_NATIVE)
        return (
            gr.update(
                value=native,
                interactive=False,
                minimum=native,
                maximum=native,
                label=f"Scale Factor (fixed {DISTILLED_NATIVE:g}x)",
            ),
            gr.update(
                value=0,
                interactive=False,
                label="Target Long Side (px) - disabled for the distilled model",
            ),
        )
    return (
        gr.update(
            value=DEFAULT_SCALE,
            interactive=True,
            minimum=1.0,
            maximum=float(MAX_SCALE),
            label="Scale Factor",
        ),
        gr.update(
            value=0,
            interactive=True,
            label="Or Target Long Side (px) [0 = use scale factor]",
        ),
    )



def load_model(device: str):
    model = torch.jit.load(str(MODEL_PATH), map_location=device)
    model.eval()
    return model


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, ""


_DISTILLED = None  # (model, native factor, device), built once per container


def distilled_model():
    """(model, native factor, device) for the distilled student, built once per container.



    The student is loaded with the same torch runtime as the default engine, so it runs on the

    Space GPU whenever there is one. ONNX Runtime was dropped on purpose: the Space installs the

    CPU-only `onnxruntime` wheel, so the ONNX build of these weights could never reach the GPU.

    """
    global _DISTILLED
    if _DISTILLED is not None:
        return _DISTILLED
    path = DISTILLED_PT_PATH
    if not path.exists():
        from huggingface_hub import hf_hub_download

        path = Path(hf_hub_download(repo_id=DISTILLED_REPO_ID, filename="model_scripted.pt"))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = torch.jit.load(str(path), map_location=device)
    model.eval()
    with torch.no_grad():
        # Read the fixed factor from the model itself instead of trusting a constant.
        probe = model(torch.zeros(1, 3, 8, 9, device=device))
        native = int(round(max(probe.shape[-2] / 8.0, probe.shape[-1] / 9.0)))
        # Warm the kernels with one full window (tile 64 + 2*6 padding, the defaults of
        # upscale_distilled below) so the first real request does not pay for cuDNN kernel
        # selection, which costs seconds on a cold GPU container.
        model(torch.zeros(1, 3, 76, 76, device=device))
    _DISTILLED = (model, native, device)
    return _DISTILLED


def upscale_distilled(rgb, scale, native, out_w, out_h, model=None, device="cpu",

                      tile: int = 64, pad: int = 6):
    """Tiled inference in LR space for the fixed-factor student.



    Same recipe as scripts/upscale_onnx.py: run the padded tile, resample the WHOLE padded tile

    to its own target size (so the crop offset stays inside it), then paste the crop on the canvas.



    Each tile owns the canvas span up to the NEXT tile origin, so the spans butt together exactly.

    Sizing the paste from that grid instead of from round(tile * scale) is what closes the seams:

    at a fractional scale round(top * scale) can land a pixel past where the previous tile's

    round((bottom - top) * scale) stopped, and the pixel in between kept the canvas' initial

    value - a black hairline every few tiles whenever scale * tile is not an integer.

    """
    height, width = rgb.shape[:2]
    canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8)
    resample = Image.Resampling.LANCZOS if native > scale else Image.Resampling.BICUBIC
    tile = max(8, min(int(tile), height))
    # Every window keeps the SAME size: at the borders it slides inward instead of shrinking.
    # TorchScript re-specialises its fused convolutions for the first window of each new shape
    # (~0.3 s each), so a variable window would make the first request pay for up to nine shapes.
    win_h, win_w = min(tile + 2 * pad, height), min(tile + 2 * pad, width)
    tops, lefts = list(range(0, height, tile)), list(range(0, width, tile))
    # Canvas cuts: tile k paints [cut[k], cut[k+1]), so consecutive tiles share no pixel and
    # leave none out, whatever the scale.
    row_cuts = [int(round(t * scale)) for t in tops] + [out_h]
    col_cuts = [int(round(l * scale)) for l in lefts] + [out_w]
    written = np.zeros((out_h, out_w), dtype=bool)
    for k, top in enumerate(tops):
        for j, left in enumerate(lefts):
            y0, x0 = max(0, min(top - pad, height - win_h)), max(0, min(left - pad, width - win_w))
            y1, x1 = y0 + win_h, x0 + win_w
            patch = rgb[y0:y1, x0:x1].astype(np.float32) / 255.0
            tensor = torch.from_numpy(patch).permute(2, 0, 1).unsqueeze(0).to(device)
            with torch.no_grad():
                sr = model(tensor)[0].clamp(0.0, 1.0).permute(1, 2, 0).cpu().numpy()
            sr = (sr * 255.0).round().astype(np.uint8)
            full_h, full_w = int(round((y1 - y0) * scale)), int(round((x1 - x0) * scale))
            sr = np.asarray(Image.fromarray(sr).resize((full_w, full_h), resample))
            oy = min(int(round((top - y0) * scale)), max(0, full_h - 1))
            ox = min(int(round((left - x0) * scale)), max(0, full_w - 1))
            th = min(row_cuts[k + 1] - row_cuts[k], full_h - oy, out_h - row_cuts[k])
            tw = min(col_cuts[j + 1] - col_cuts[j], full_w - ox, out_w - col_cuts[j])
            if th <= 0 or tw <= 0:
                continue
            canvas[row_cuts[k]:row_cuts[k] + th, col_cuts[j]:col_cuts[j] + tw] = sr[oy:oy + th, ox:ox + tw]
            written[row_cuts[k]:row_cuts[k] + th, col_cuts[j]:col_cuts[j] + tw] = True
    if not written.all():
        # A window clipped at the far border can run out of rows before its span ends. Those
        # last pixels take the plain resample of the whole image instead of the canvas' black.
        canvas[~written] = np.asarray(Image.fromarray(rgb).resize((out_w, out_h), resample))[~written]
    return canvas


@gpu_decorator
def upscale(image, model_choice=None, scale: float = DEFAULT_SCALE, 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")

    choice = normalize_model_choice(model_choice)

    if choice == "distilled":
        try:
            model, native, device = distilled_model()
        except Exception as exc:
            return None, f"Distilled model unavailable: {exc}"
        native = native or DISTILLED_NATIVE
        factor = float(native)
        clamped = ""
    else:
        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)))
    capped = max(out_w, out_h) > MAX_SIDE
    if capped:
        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))
    # The size actually produced, which is what the status line should report.
    applied = max(out_w / float(img.width), out_h / float(img.height))

    if choice == "distilled":
        arr = upscale_distilled(np.asarray(img, dtype=np.uint8), applied, native, out_w, out_h,
                                model=model, device=device)
        model_note = f"Universal Upscale Distilled (TorchScript, native {native}x"
        if abs(applied - native) > 1e-3:
            model_note += ", output resampled to the Space size cap)"
            clamped += f" (distilled {native:g}x capped to {MAX_SIDE} px)"
        else:
            model_note += ")"
        device = f"torch/{device}"
    else:
        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)
        model_note = "Universal Upscale (LIIF, arbitrary scale)"

    res_img = Image.fromarray(arr)
    status = (f"Input: {img.width}x{img.height} | Output: {res_img.width}x{res_img.height} "
              f"(Scale: {applied:.2f}x) | Model: {model_note} | 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.Dropdown(choices=MODEL_CHOICES, value=MODEL_LIIF, label="Model"),
        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=(
        "Pick the engine in the **Model** menu: **Universal Upscale** renders any factor from 1x up to the tested "
        "maximum 9.2x (or an exact target resolution), **Universal Upscale Distilled** is the 1.3 M "
        "parameter student that renders 9x natively in one pass; selecting it locks the "
        "Scale Factor and Target Long Side controls to that native factor. Tiled inference keeps high-resolution "
        "upscaling inside memory (outputs are capped at "
        "4096 px on the long side); the status line reports which model ran and what it did."
    ),
    examples=example_files if example_files else None,
)


# The distilled student has one native factor: lock the scale controls to it when selected, and
# restore the arbitrary-scale controls for the default model.
_input_components = getattr(demo, "input_components", None) or getattr(demo, "inputs", None)
if _input_components is not None and len(_input_components) >= 4:
    with demo:
        _input_components[1].change(
            fn=scale_controls_for_model,
            inputs=_input_components[1],
            outputs=[_input_components[2], _input_components[3]],
        )

if __name__ == "__main__":
    demo.launch()