app.py

import os
import subprocess
import sys
import tempfile

import gradio as gr
from cached_path import cached_path
from PIL import Image

match sys.platform:
    case "darwin":
        latest_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip"
    case "linux":
        latest_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip"
    case "win32":
        latest_url = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip"
    case _:
        raise NotImplementedError(f"Unsupported platform: {sys.platform}")

path = cached_path(latest_url, extract_archive=True)
exe_name = path / "realesrgan-ncnn-vulkan"
if sys.platform == "win32":
    exe_name += ".exe"
else:
    os.chmod(exe_name, 0o755)

models = [
    p.replace(".param", "") for p in os.listdir(path / "models") if p.endswith(".param")
]


def process(filename: str, model: str, scale: int, tta: bool):
    assert scale in [2, 3, 4], f"Invalid scale: {scale}"
    assert model in models, f"Invalid model: {model}"
    with tempfile.NamedTemporaryFile(suffix=".webp") as temp_file:
        args = [
            exe_name,
            "-i",
            filename,
            "-o",
            temp_file.name,
            "-n",
            model,
            "-s",
            str(int(scale)),
        ]
        if tta:
            args.append("-x")

        process = subprocess.Popen(
            args,
            stderr=subprocess.PIPE,
            text=True,
            cwd=path,
        )
        _, err = process.communicate()
        assert process.returncode == 0, err
        return Image.open(temp_file)


demo = gr.Interface(
    fn=process,
    title="Real-ESRGAN Demo",
    description="Thanks to [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) for the amazing model!",
    inputs=[
        gr.Image(type="filepath", format="webp"),
        gr.Dropdown(models, label="Model", value="realesrgan-x4plus-anime"),
        gr.Slider(2, 4, label="Scaling factor", step=1),
        gr.Checkbox(False, label="Use TTA (8x slower, may improve quality)"),
    ],
    outputs=gr.Image(type="filepath", format="webp"),
    analytics_enabled=False,
    flagging_mode="never",
)

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