app.py

import os
import math
import uuid
import numpy as np
import onnxruntime as ort
from PIL import Image
import gradio as gr
import tempfile
import requests


MODEL_DIR     = "model"
MODEL_X2_PATH = os.path.join(MODEL_DIR, "Real-ESRGAN_x2plus.onnx")
MODEL_X4_PATH = os.path.join(MODEL_DIR, "Real-ESRGAN-x4plus.onnx")


FILE_ID_X2 = "15xmXXZNH2wMyeQv4ie5hagT7eWK9MgP6"
FILE_ID_X4 = "1wDBHad9RCJgJDGsPdapLYl3cr8j-PMJ6"

def download_from_drive(file_id: str, dest_path: str):
  
    URL = "https://drive.google.com/uc?export=download"
    session = requests.Session()

    
    response = session.get(URL, params={"id": file_id}, stream=True)
    token = None
    for key, value in response.cookies.items():
        if key.startswith('download_warning'):
            token = value
            break

    
    if token:
        params = {"id": file_id, "confirm": token}
        response = session.get(URL, params=params, stream=True)

    
    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
    with open(dest_path, "wb") as f:
        for chunk in response.iter_content(chunk_size=32768):
            if chunk:
                f.write(chunk)
    print(f"Model telah diunduh dan disimpan di {dest_path}")
    return dest_path


if not os.path.isfile(MODEL_X2_PATH):
    download_from_drive(FILE_ID_X2, MODEL_X2_PATH)

# Unduh model ×4
if not os.path.isfile(MODEL_X4_PATH):
    download_from_drive(FILE_ID_X4, MODEL_X4_PATH)


sess_opts = ort.SessionOptions()
sess_opts.intra_op_num_threads = 2
sess_opts.inter_op_num_threads = 2

session_x2 = ort.InferenceSession(MODEL_X2_PATH, sess_options=sess_opts, providers=["CPUExecutionProvider"])
session_x4 = ort.InferenceSession(MODEL_X4_PATH, sess_options=sess_opts, providers=["CPUExecutionProvider"])


input_meta_x2 = session_x2.get_inputs()[0]
_, _, H_in_x2, W_in_x2 = tuple(input_meta_x2.shape)
H_in_x2, W_in_x2 = int(H_in_x2), int(W_in_x2)

input_meta_x4 = session_x4.get_inputs()[0]
_, _, H_in_x4, W_in_x4 = tuple(input_meta_x4.shape)
H_in_x4, W_in_x4 = int(H_in_x4), int(W_in_x4)


dummy_x2 = np.zeros((1, 3, H_in_x2, W_in_x2), dtype=np.float32)
dummy_out_x2 = session_x2.run(None, {input_meta_x2.name: dummy_x2})[0]
_, _, H_out_x2, W_out_x2 = dummy_out_x2.shape
SCALE_X2 = H_out_x2 // H_in_x2
if SCALE_X2 != 2:
    raise RuntimeError(f"Model ×2 menghasilkan scale = {SCALE_X2}, bukan 2")

dummy_x4 = np.zeros((1, 3, H_in_x4, W_in_x4), dtype=np.float32)
dummy_out_x4 = session_x4.run(None, {input_meta_x4.name: dummy_x4})[0]
_, _, H_out_x4, W_out_x4 = dummy_out_x4.shape
SCALE_X4 = H_out_x4 // H_in_x4
if SCALE_X4 != 4:
    raise RuntimeError(f"Model ×4 menghasilkan scale = {SCALE_X4}, bukan 4")


def run_tile_x2(tile_np: np.ndarray) -> np.ndarray:
    patch_nchw = np.transpose(tile_np, (2, 0, 1))[None, ...]
    out_nchw = session_x2.run(None, {input_meta_x2.name: patch_nchw})[0]
    out_nchw = np.squeeze(out_nchw, axis=0)
    out_hwc  = np.transpose(out_nchw, (1, 2, 0))
    return out_hwc

def run_tile_x4(tile_np: np.ndarray) -> np.ndarray:
    patch_nchw = np.transpose(tile_np, (2, 0, 1))[None, ...]
    out_nchw = session_x4.run(None, {input_meta_x4.name: patch_nchw})[0]
    out_nchw = np.squeeze(out_nchw, axis=0)
    out_hwc  = np.transpose(out_nchw, (1, 2, 0))
    return out_hwc


def tile_upscale(input_img: Image.Image, scale: int, max_dim=1024):
    if scale == 2:
        H_in, W_in, run_tile, SCALE = H_in_x2, W_in_x2, run_tile_x2, SCALE_X2
    else:
        H_in, W_in, run_tile, SCALE = H_in_x4, W_in_x4, run_tile_x4, SCALE_X4

    
    w, h = input_img.size
    if w > max_dim or h > max_dim:
        scale_factor = max_dim / float(max(w, h))
        new_w = int(w * scale_factor)
        new_h = int(h * scale_factor)
        input_img = input_img.resize((new_w, new_h), Image.LANCZOS)

    img_rgb = input_img.convert("RGB")
    arr = np.array(img_rgb).astype(np.float32) / 255.0
    h_orig, w_orig, _ = arr.shape

    tiles_h = math.ceil(h_orig / H_in)
    tiles_w = math.ceil(w_orig / W_in)
    pad_h = tiles_h * H_in - h_orig
    pad_w = tiles_w * W_in - w_orig

    arr_padded = np.pad(arr, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
    out_h = tiles_h * H_in * SCALE
    out_w = tiles_w * W_in * SCALE
    out_arr = np.zeros((out_h, out_w, 3), dtype=np.float32)

    for i in range(tiles_h):
        for j in range(tiles_w):
            y0, x0 = i * H_in, j * W_in
            tile = arr_padded[y0:y0+H_in, x0:x0+W_in, :]
            up_tile = run_tile(tile)
            oy0, ox0 = i * H_in * SCALE, j * W_in * SCALE
            out_arr[oy0:oy0 + H_in * SCALE, ox0:ox0 + W_in * SCALE, :] = up_tile

    final_arr = out_arr[0:h_orig * SCALE, 0:w_orig * SCALE, :]
    final_arr = np.clip(final_arr, 0.0, 1.0)
    final_uint8 = (final_arr * 255.0).round().astype(np.uint8)
    final_pil = Image.fromarray(final_uint8)

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    final_pil.save(tmp.name, format="PNG")
    tmp.close()
    return final_pil, tmp.name


def upscale_x2(input_img: Image.Image):
    return tile_upscale(input_img, scale=2)

def standard_upscale(input_img: Image.Image):
    return tile_upscale(input_img, scale=4)

def premium_upscale(input_img: Image.Image):
    final_4x, _ = tile_upscale(input_img, scale=4)
    w_orig, h_orig = input_img.size
    target_size = (w_orig * 8, h_orig * 8)
    final_8x = final_4x.resize(target_size, resample=Image.LANCZOS)
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    final_8x.save(tmp.name, format="PNG")
    tmp.close()
    return final_8x, tmp.name


css = """
#x2-btn {
    background-color: lightgreen !important;
    color: black !important;
}
#premium-btn {
    background-color: gold !important;
    color: black !important;
}
"""


with gr.Blocks(css=css, title="SpectraGAN Triple-Mode Upscaler") as demo:
    gr.Markdown(
        """
        # SpectraGAN Upscaler  
        **Upscale (×2)**, **Standard Upscale (×4)** atau **Premium Upscale 🚀 (×8)**.  
        
        """
    )

    with gr.Row():
        inp_image = gr.Image(type="pil", label="Upload Source Image")

    with gr.Row():
        btn_x2   = gr.Button("Upscale (×2)", elem_id="x2-btn")
        btn_std  = gr.Button("Standard Upscale (×4)", variant="primary", elem_id="std-btn")
        btn_prem = gr.Button("Premium Upscale 🚀 (×8)", elem_id="premium-btn")

    out_preview  = gr.Image(type="pil", label="Upscaled Preview")
    out_download = gr.DownloadButton("⬇️ Download PNG", visible=True)

    btn_x2.click(fn=upscale_x2, inputs=inp_image, outputs=[out_preview, out_download])
    btn_std.click(fn=standard_upscale, inputs=inp_image, outputs=[out_preview, out_download])
    btn_prem.click(fn=premium_upscale, inputs=inp_image, outputs=[out_preview, out_download])

demo.launch(server_name="0.0.0.0", server_port=7860)