# '''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="""
    ## Neural Style Transfer
    Upload a content image and a style image to create a stylized version of your content image.
    
    ### How to use:
    1. Upload your content image (the image you want to style)
    2. Upload your style image (the image whose style you want to apply)
    3. Click 'Submit' and wait for the result
    
    ### Tips:
    - For best results, use clear, high-quality images
    - The style image's characteristics will be transferred to your content image
    - Processing may take a few seconds depending on image size
    """,
    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(
        server_name="0.0.0.0",
        server_port=7860,
        share=False  # Set to False for Hugging Face Spaces
    )