import os
import time
import spaces
import gradio as gr
from gradio_imageslider import ImageSlider
from PIL import Image
import numpy as np
from aura_sr import AuraSR
import torch
import devicetorch
DEVICE = devicetorch.get(torch)

# Force CPU usage
torch.set_default_tensor_type(torch.FloatTensor)

# Override torch.load to always use CPU
original_load = torch.load
torch.load = lambda *args, **kwargs: original_load(*args, **kwargs, map_location=torch.device('cpu'))

# Initialize the AuraSR model
aura_sr = AuraSR.from_pretrained("fal/AuraSR-v2", device=DEVICE)

# Restore original torch.load
torch.load = original_load

# Create output folder if not exists
output_folder = '../outputs'
os.makedirs(output_folder, exist_ok=True)

def generate_output_filename():
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    return f"upscaled_{timestamp}.png"

def process_image(input_image):
    if input_image is None:
        raise gr.Error("Please provide an image to upscale.")

    # Convert to PIL Image for resizing
    pil_image = Image.fromarray(input_image)

    # Upscale the image using AuraSR
    upscaled_image = process_image_on_gpu(pil_image)

    # Convert result to numpy array if it's not already
    result_array = np.array(upscaled_image)

    # Save result as PNG
    output_filename = generate_output_filename()
    output_path = os.path.join(output_folder, output_filename)
    upscaled_image.save(output_path, format="PNG")

    return [input_image, result_array], output_path

@spaces.GPU
def process_image_on_gpu(pil_image):
    return aura_sr.upscale_4x(pil_image)
    
title = """<h1 align="center">AuraSR-v2 - An open reproduction of the GigaGAN Upscaler from fal.ai</h1>
<p><center>
<a href="https://huggingface.co/fal/AuraSR-v2" target="_blank">[AuraSR-v2]</a>
<a href="https://blog.fal.ai/introducing-aurasr-an-open-reproduction-of-the-gigagan-upscaler-2/" target="_blank">[Blog Post]</a>
<a href="https://huggingface.co/fal-ai/AuraSR" target="_blank">[v1 Model Page]</a>
</center></p>
"""

css = """
    .img { max-height: 80vh !Important }
    #slider-container { overflow: hidden; display: flex; justify-content: center }
    #row-height { height: 65px !important }
"""

with gr.Blocks(css=css) as demo:
    
    gr.HTML(title)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Input Image", type="numpy", elem_classes="img")
            process_btn = gr.Button("Upscale Image")
        with gr.Column(scale=1):
            output_slider = ImageSlider(label="Before / After", type="numpy", elem_classes="img", elem_id="slider-container")
            download_btn = gr.File(label="Download Image", elem_id="row-height")

    process_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_slider, download_btn]
    )

    # Add examples
    gr.Examples(
        examples=[
            "image1.png",
            "image3.png"
        ],
        inputs=input_image,
        outputs=output_slider,
        fn=process_image,
#        cache_examples=True
    )

demo.launch(debug=True)