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

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

# Function to execute
def style_transfer(content_img, style_img):
    content_img = tf.image.resize(tf.convert_to_tensor(content_img,tf.float32)[tf.newaxis,...] / 255.,(512,512),preserve_aspect_ratio=True)
    style_img = tf.convert_to_tensor(style_img,tf.float32)[tf.newaxis,...] / 255.

    generated_img = model(content_img, style_img)[0]
    return Image.fromarray(np.uint8(generated_img[0]*255))


# Metadata
title = "Neural Style Transfer"
description = "Transform your photos into stunning works of art with our Neural Style Transfer app. Simply upload your content image and style image, and watch as the AI-powered model applies the artistic style of your choice to your photos. Unleash your creativity and turn ordinary images into extraordinary masterpieces!"

# Define Inputs with labels
inputs = [
    gr.Image(label="Content/Base Image"),
    gr.Image(label="Style Image")
]

# Example images
examples = [
    ["content_example_1.jpg", "style_example_1.jpg"],
    ["content_example_2.jpg", "style_example_2.jpg"]
]


interface = gr.Interface(fn=style_transfer,
                            inputs=inputs,
                            outputs=["image"],
                            title=title,
                            description=description,
                            examples=examples
                        )

interface.launch()