import torch  # for model
import numpy as np
import torch.nn as nn
import torch.optim as optim
from PIL import Image  #for importing images
import torchvision.models as models  #to load vgg 19 model
import torchvision.transforms as transforms  #to transform the images
from torchvision.utils import save_image #to save the generated images
from tqdm import tqdm # progress bar
import matplotlib.pyplot as plt
import gradio as gr
import spaces

from styleTransfer import style_transfer
from dataTransform import tensor_to_image

device = 'cuda' 
print(device)

@spaces.GPU(duration = 242)
def gradio_style_transfer(steps, content_image, style_image):
    generated_tensor = style_transfer(content_image, style_image, total_steps= steps)
    generated_image = tensor_to_image(generated_tensor)
    return generated_image


interface = gr.Interface(
    fn=gradio_style_transfer,
    inputs=[
        gr.Slider(minimum=100, maximum=3000, step=100, label="Number of Steps", value=500),
        gr.Image(type='filepath', label='Content Image'),
        gr.Image(type='filepath', label='Style Image')
    ],
    outputs=gr.Image(type='pil', label='Generated Image'),
    title="Neural Style Transfer",
    description="Upload a content image and a style image, and the app will generate a stylized image.",
)

interface.launch(debug = True)