import gradio as gr
from gradio_client import Client, handle_file
from PIL import Image
from io import BytesIO
import os
import tempfile

def upscale_image(url):
    client = Client("doevent/Face-Real-ESRGAN")
    result = client.predict(
        image=handle_file(url),  # Directly pass the image
        size="4x",
        api_name="/predict"
    )
    # print("\nTask Completed!")
    # return result
    # Read the image from the file path returned by the model
    if os.path.exists(result):
        with open(result, 'rb') as img_file:
            img_data = img_file.read()

        # Convert result to PNG
        img = Image.open(BytesIO(img_data))
        
        # Save the converted image to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as temp_file:
            img.save(temp_file, format="JPEG", quality=95)
            temp_file_path = temp_file.name

        # Optionally, delete the temp file after processing (you can remove this line if not needed)
        os.remove(result)

        print("\nTask Completed!")
        return temp_file_path


app = gr.Interface(upscale_image,
             inputs = [gr.Textbox(label="Url")],
             outputs = [gr.Image(label="Upscaled Image", format='png')])

app.launch(debug=True)