# '''NEURAL STYLE TRANSFER '''

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

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

# def tensor_to_image(tensor):
#     tensor *= 255
#     tensor = np.array(tensor, dtype=np.uint8)
#     if tensor.ndim > 3:
#         tensor = tensor[0]
#     return Image.fromarray(tensor)

# def transform_my_model(content_image, style_image):
#     content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
#     style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
#     stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
#     return tensor_to_image(stylized_image)

# demo = gr.Interface(
#     fn=transform_my_model,
#     inputs=[gr.Image(label="Content Image"), gr.Image(label="Style Image")],
#     outputs=gr.Image(label="Result"),
#     title="Style Transfer",
#     examples=[
#         ["Content_Images/contnt12.jpg", "VG516.jpg"],
#         ["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
#         ["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
#     ],
#     article="References-\n\nExploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin."
# )

# demo.launch(share=True)



'''NEURAL STYLE TRANSFER '''

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

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

np.set_printoptions(suppress=True)

# Load model with error handling
try:
    logger.info("Loading TensorFlow Hub model...")
    model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
    logger.info("Model loaded successfully!")
except Exception as e:
    logger.error(f"Error loading model: {str(e)}")
    raise

def tensor_to_image(tensor):
    try:
        tensor *= 255
        tensor = np.array(tensor, dtype=np.uint8)
        if tensor.ndim > 3:
            tensor = tensor[0]
        return Image.fromarray(tensor)
    except Exception as e:
        logger.error(f"Error in tensor_to_image: {str(e)}")
        raise

def transform_my_model(content_image, style_image):
    try:
        if content_image is None or style_image is None:
            raise ValueError("Both content and style images are required")
            
        logger.info("Processing images...")
        content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
        style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
        stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
        logger.info("Style transfer completed successfully!")
        return tensor_to_image(stylized_image)
    except Exception as e:
        logger.error(f"Error in transform_my_model: {str(e)}")
        raise

# Create the Gradio interface with Hugging Face Spaces configuration
demo = gr.Interface(
    fn=transform_my_model,
    inputs=[
        gr.Image(type="numpy", label="Content Image"),
        gr.Image(type="numpy", label="Style Image")
    ],
    outputs=gr.Image(type="pil", label="Result"),
    title="Neural Style Transfer",
    description="""
    Upload a content image and a style image to create a stylized version of your content image.
    """,
    examples=[
        ["Content_Images/contnt12.jpg", "VG516.jpg"],
        ["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
        ["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
    ],
    cache_examples=True,
    allow_flagging=False,
    analytics_enabled=False
)

# For Hugging Face Spaces deployment
if __name__ == "__main__":
    demo.launch(
        share=True  # Set to False for Hugging Face Spaces
    )