import torch
from PIL import Image
from RealESRGAN import RealESRGAN
import gradio as gr
from gradio_imageslider import ImageSlider
import spaces

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model2 = RealESRGAN(device, scale=2)
model2.load_weights('weights/RealESRGAN_x2.pth', download=True)
model4 = RealESRGAN(device, scale=4)
model4.load_weights('weights/RealESRGAN_x4.pth', download=True)
model8 = RealESRGAN(device, scale=8)
model8.load_weights('weights/RealESRGAN_x8.pth', download=True)


@spaces.GPU
def inference(image, size):
    global model2
    global model4
    global model8
    if image is None:
        raise gr.Error("Image not uploaded")
        
    # Store original image for comparison
    original_image = image.copy()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    if size == '2x':
        try:
            result = model2.predict(image.convert('RGB'))
        except torch.cuda.OutOfMemoryError as e:
            print(e)
            model2 = RealESRGAN(device, scale=2)
            model2.load_weights('weights/RealESRGAN_x2.pth', download=False)
            result = model2.predict(image.convert('RGB'))
    elif size == '4x':
        try:
            result = model4.predict(image.convert('RGB'))
        except torch.cuda.OutOfMemoryError as e:
            print(e)
            model4 = RealESRGAN(device, scale=4)
            model4.load_weights('weights/RealESRGAN_x4.pth', download=False)
            result = model2.predict(image.convert('RGB'))
    else:
        try:
            width, height = image.size
            if width >= 5000 or height >= 5000:
                raise gr.Error("The image is too large.")
            result = model8.predict(image.convert('RGB'))
        except torch.cuda.OutOfMemoryError as e:
            print(e)
            model8 = RealESRGAN(device, scale=8)
            model8.load_weights('weights/RealESRGAN_x8.pth', download=False)
            result = model2.predict(image.convert('RGB'))
            
    print(f"Image size ({device}): {size} ... OK")
    # Return tuple of original and processed images for the slider
    return (original_image, result)


title = """<h1 align="center">ProFaker</h1>"""

with gr.Blocks() as demo:
    gr.HTML(title)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            size_select = gr.Radio(
                ["2x", "4x", "8x"],
                type="value",
                value="2x",
                label="Resolution model"
            )
            process_btn = gr.Button("Upscale Image")
            
        with gr.Column():
            result_slider = ImageSlider(
                interactive=False,
                label="Before and After Comparison"
            )
    
    process_btn.click(
        fn=inference,
        inputs=[input_image, size_select],
        outputs=result_slider
    )

demo.queue(api_open=True).launch(show_error=True, show_api=True)