import gradio as gr
from huggingface_hub import from_pretrained_keras
import tensorflow as tf
import numpy as np
from PIL import Image
import io

# Load pre-trained deblurring model
model = from_pretrained_keras("google/maxim-s3-deblurring-reds")

def deblur_image(input_image):
    # Preprocess the input image
    image = np.array(input_image)
    image = tf.convert_to_tensor(image)
    image = tf.image.resize(image, (256, 256))
    
    # Predict with the model
    predictions = model.predict(tf.expand_dims(image, 0))
    
    # Convert back to image
    output_image = predictions[0].numpy().astype(np.uint8)
    output_image = Image.fromarray(output_image)
    
    # Save the result in memory as a byte array
    byte_arr = io.BytesIO()
    output_image.save(byte_arr, format='PNG')
    byte_arr.seek(0)
    
    return byte_arr

# Set up the Gradio interface
iface = gr.Interface(fn=deblur_image,
                     inputs=gr.inputs.Image(type="pil"),
                     outputs=gr.outputs.Image(type="file"),
                     live=True)

# Launch the Gradio interface
iface.launch()