import gradio as gr
import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from PIL import Image

# Load the pre-trained style transfer model
model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')

def style_transfer(input_image, style_image, image_size=(512, 512)):  # Updated to a larger size
    # Preprocess the images
    input_image = preprocess_image(input_image, image_size)
    style_image = preprocess_image(style_image, image_size)
    
    # Perform style transfer
    stylized_image = model(tf.constant(input_image), tf.constant(style_image))[0]
    
    # Postprocess the image
    stylized_image = postprocess_image(stylized_image)
    
    return stylized_image

def preprocess_image(image, image_size):
    image = Image.fromarray(image.astype('uint8'), 'RGB')
    image = image.resize(image_size)  # Updated to variable image size
    image = np.array(image).astype('float32')
    image = image / 255.0
    image = np.expand_dims(image, axis=0)
    return image

def postprocess_image(image):
    image = image.numpy()
    image = np.squeeze(image)
    image = np.clip(image * 255, 0, 255).astype('uint8')
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=style_transfer,
    inputs=["image", "image"],
    outputs="image",
    live=False,
)

# Launch the Gradio app
iface.launch()