app.py

import sys
import importlib
import os

# =================================================================
# CRITICAL PYTHON 3.13 / GRADIO SDK CONFLICT MONKEY-PATCHES
# =================================================================

# 1. Safely inject HfFolder directly into the real huggingface_hub without blocking other imports
try:
    real_hf_hub = importlib.import_module('huggingface_hub')
    if not hasattr(real_hf_hub, 'HfFolder'):
        class MockHfFolder:
            @classmethod
            def get_token(cls): return os.environ.get("HF_TOKEN")
            @classmethod
            def save_token(cls, token): pass
            @classmethod
            def delete_token(cls): pass

        real_hf_hub.HfFolder = MockHfFolder
        sys.modules['huggingface_hub'].HfFolder = MockHfFolder
except Exception as patch_err_1:
    print(f"Pre-import HfFolder patch skipped or failed: {patch_err_1}")

# 2. Patch the Gradio client schema serialization parser to avoid the 'bool' type loop error
try:
    import gradio_client.utils as client_utils
    orig_get_type = getattr(client_utils, 'get_type', None)
    
    if orig_get_type:
        def safe_get_type(schema):
            if isinstance(schema, bool):
                return "boolean"
            return orig_get_type(schema)
        client_utils.get_type = safe_get_type
        print("Successfully patched gradio_client schema serialization engine.")
except Exception as patch_err_2:
    print(f"Schema serialization engine patch deferred: {patch_err_2}")
# =================================================================

import cv2
import gradio as gr
import numpy as np
import onnxruntime as ort
from huggingface_hub import hf_hub_download

# Configuration with separate tile size constraints matching the ONNX compilation targets
MODELS = {
    "RealESRGAN_x2plus (Faster 2x)": {
        "repo_id": "tidus2102/Real-ESRGAN",
        "filename": "Real-ESRGAN_x2plus.onnx",
        "scale": 2,
        "tile_size": 64
    },
    "RealESRGAN_x4plus (High Quality 4x)": {
        "repo_id": "KingPro100/real-esrgan-onxx",
        "filename": "Real-開ESRGAN-x4plus.onnx" if False else "Real-ESRGAN-x4plus.onnx",
        "scale": 4,
        "tile_size": 128
    }
}

current_model_name = None
ort_session = None

def load_model(model_choice):
    global current_model_name, ort_session
    if current_model_name == model_choice and ort_session is not None:
        return ort_session

    cfg = MODELS[model_choice]
    print(f"Loading weights for {model_choice}...")
    token = os.environ.get("HF_TOKEN")
    
    model_path = hf_hub_download(
        repo_id=cfg["repo_id"], 
        filename=cfg["filename"],
        token=token
    )

    session_options = ort.SessionOptions()
    session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    
    ort_session = ort.InferenceSession(model_path, session_options, providers=['CPUExecutionProvider'])
    current_model_name = model_choice
    return ort_session

def upscale_image(input_img, model_choice):
    """
    Tiled inference pipeline. Splits arbitrary sized images into model-specific blocks
    (64x64 or 128x128), runs inference on each tile, and puts them back together safely.
    """
    if input_img is None:
        return None
        
    try:
        session = load_model(model_choice)
        cfg = MODELS[model_choice]
        
        scale = cfg["scale"]
        tile = cfg["tile_size"]
        
        h, w, c = input_img.shape
        
        # Calculate how much padding we need to make the image divisible by the tile size
        pad_h = (tile - (h % tile)) % tile
        pad_w = (tile - (w % tile)) % tile
        
        # Pad the image with edge mirroring to avoid border artifacts
        padded_img = cv2.copyMakeBorder(input_img, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)
        ph, pw, _ = padded_img.shape
        
        # Allocate output canvas
        output_h = ph * scale
        output_w = pw * scale
        output_padded = np.zeros((output_h, output_w, c), dtype=np.uint8)
        
        input_name = session.get_inputs()[0].name
        output_name = session.get_outputs()[0].name
        
        # Iterate over the rows and columns of the specific tile size blocks
        for y in range(0, ph, tile):
            for x in range(0, pw, tile):
                # Crop tile
                tile_img = padded_img[y:y+tile, x:x+tile]
                
                # Preprocess tile dynamically to match target shape [1, 3, tile_size, tile_size]
                tile_float = tile_img.astype(np.float32) / 255.0
                tile_chw = np.transpose(tile_float, (2, 0, 1))
                tile_batch = np.expand_dims(tile_chw, axis=0)
                
                # Execute ONNX forward tensor calculations on CPU
                ort_outs = session.run([output_name], {input_name: tile_batch})
                output_tensor = ort_outs[0]
                
                # Postprocess tile back to target scale shape
                output_tensor = np.squeeze(output_tensor, axis=0)
                output_tensor = np.clip(output_tensor, 0.0, 1.0)
                output_hwc = np.transpose(output_tensor, (1, 2, 0))
                output_tile = (output_hwc * 255.0).astype(np.uint8)
                
                # Insert the processed super-resolution tile into output matrix
                output_padded[y*scale:(y+tile)*scale, x*scale:(x+tile)*scale] = output_tile
                
        # Unpad the final image back to original proportional dimensions multiplied by scale factor
        final_h = h * scale
        final_w = w * scale
        final_output = output_padded[0:final_h, 0:final_w]
        
        return final_output
        
    except Exception as e:
        print(f"Error executing ONNX runtime tensor transformation graph: {str(e)}")
        blank_err_img = np.zeros((300, 500, 3), dtype=np.uint8)
        cv2.putText(blank_err_img, f"Execution Error: {str(e)[:40]}", (20, 150), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA)
        return blank_err_img

# Define the user interface layout
with gr.Blocks(title="AI Lightweight Image Upscaler (ONNX)") as demo:
    gr.Markdown("# 🖼️ AI Image Resizer & Upscaler (ONNX Engine)")
    gr.Markdown("Running locally on Hugging Face Free CPU hardware using dynamic tiling maps.")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Source Image", type="numpy")
            model_dropdown = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="RealESRGAN_x2plus (Faster 2x)",
                label="Select AI Upscaling Engine"
            )
            submit_btn = gr.Button("Upscale Image", variant="primary")
            
        with gr.Column():
            output_image = gr.Image(label="Enhanced Super-Resolution Result", type="numpy")

    submit_btn.click(
        fn=upscale_image,
        inputs=[input_image, model_dropdown],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)