import numpy as np
import tensorflow as tf
import onnxruntime as ort
import gradio as gr
import time
import logging
from PIL import Image

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

model_path = "model.onnx"

try:
    start_time = time.perf_counter()
    model = ort.InferenceSession(model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
    logger.info(f"Model successfully loaded from {model_path} in {time.perf_counter()-start_time:.2f} sec")

except Exception as e:
    logger.error(f"Failed to load model from {model_path}: {e}")

def preprocess_image(image):
    "Convert to rgb, normalize and add batch dimension"
    img = Image.open(image).convert("RGB")
    img_arr = np.array(img, dtype=np.float32) / 255.0
    img_arr = np.expand_dims(img_arr, axis=0)

    return img_arr

def super_resolution_image(image):
    try:
        model_inputs = model.get_inputs()[0].name
        model_outputs = model.get_outputs()[0].name
        # Inference
        sr_img = model.run([model_outputs], {model_inputs: image})[0]
        # Postprocess
        sr_img = (np.clip(sr_img, 0, 1)*255).astype(np.uint8)
        return Image.fromarray(sr_img[0])
    
    except Exception as e:
        raise RuntimeError(f"Model inference failed, {str(e)}")
    
def gradio_inference(image):
    img_arr = preprocess_image(image)
    sr_img = super_resolution_image(img_arr)
    return sr_img

demo = gr.Interface(
    fn=gradio_inference,
    inputs=gr.Image(type="filepath"),  
    outputs="image",
    title="Image Upscaling 4x with EDSR",
    description=("Upscale your images 4× using an ONNX EDSR model.\n\n"
        "**⚠️ CPU-only demo. Images larger than 512×512 may take significantly longer.**"
    ),
    examples=[
        "examples/comic.png",
        "examples/bird.png"
    ]
)

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