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

style_transfer_model = hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")


def perform_style_transfer(content_image, style_image):
    
    content_image = tf.convert_to_tensor(content_image, np.float32)[tf.newaxis, ...] / 255.
    style_image = tf.convert_to_tensor(style_image, np.float32)[tf.newaxis, ...] / 255.
    
    output = style_transfer_model(content_image, style_image)
    stylized_image = output[0]
    
    return Image.fromarray(np.uint8(stylized_image[0] * 255))


content_image_input = gr.inputs.Image()
style_image_input = gr.inputs.Image(shape=(256, 256))

interface = gr.Interface(fn=perform_style_transfer,
                  inputs=[content_image_input, style_image_input],
                  outputs="image")
Interface.launch()