import gradio as gr
from neural_style_transfer import NeuralStyleTransfer

def style_transfer(content_image, style_image):
    # Initialize the neural style transfer model
    nst = NeuralStyleTransfer(content_image, style_image)
    
    # Generate the stylized image
    output_image = nst.transfer_style()
    
    return output_image

# Gradio interface setup
iface = gr.Interface(
    fn=style_transfer, 
    inputs=[
        gr.Image(type="pil", label="Content Image"), 
        gr.Image(type="pil", label="Style Image")
    ],
    outputs=gr.Image(label="Stylized Image"),
    title="Neural Style Transfer",
    description="Upload a content image and a style image to apply neural style transfer."
)

if __name__ == "__main__":
    iface.launch()