import os
import sys
import sys
import torch
import gradio as gr
from stable_diffusion import StableDiffusion
from utils import invert_loss, get_style_embeddings, show_images

torch.manual_seed(1)
torch_device = "cuda" if torch.cuda.is_available() else "cpu" 


def gradio_interface(prompt:str, style_name:str):
    stable_diffuser = StableDiffusion(torch_device)
    #outputs_1 = []
    #outputs_2 = []
    seed_values = [1,2,3,4,5]
    custom_loss_scale = 150.0
    num_styles = len(style_files)
    style_file = style_files[style_name]
    style_token_embedding = get_style_embeddings(style_file) if style_file is not None else None
    this_generated_img_1 = stable_diffuser.generate_image_with_custom_style(prompt,
                                                          style_token_embedding = style_token_embedding,
                                                          random_seed = 42,
                                                          custom_loss_fn = None)
    #outputs_1.append(this_generated_img_1)
    this_generated_img_2 = stable_diffuser.generate_image_with_custom_style(prompt,
                                                          style_token_embedding = style_token_embedding,
                                                          random_seed = 42,
                                                          custom_loss_fn = invert_loss,
                                                          custom_loss_scale = custom_loss_scale)
    #outputs_2.append(this_generated_img_2)
    return this_generated_img_1, this_generated_img_2
    

style_files = {'No_style': None,
               'watercolor': 'learned_embeds_watercolor.bin',
               'strip': 'learned_embeds_strip_style.bin',
               'oil_paint': 'learned_embeds_oil_paint.bin',
               'kaleido': 'learned_embeds_kaleido.bin',
                'doodle': 'learned_embeds_doodle.bin'}

torch_device = "cuda" if torch.cuda.is_available() else "cpu"

prompt_examples = [["A person wearing a hat and smoking a cigar", 'watercolor']]
#["An oil painting of a bear playing guitar"]]
#["An oil painting of a lion eating pizza and drinking coke"]]

# Define Interface
description = 'A Stable Diffusion based Generative AI tool to generate images. Also generates grayscale images.'
title = 'Image Generation using Stable Diffusion with styles.'
style_names = ['No_style','watercolor','strip','oil_paint','kaleido','doodle']
demo = gr.Interface(gradio_interface,
                    inputs = [gr.Textbox('A deer crossing the street', label="Text Prompt"),
                             gr.Dropdown(style_names,value='watercolor',label="Display Style")],
                    outputs = [gr.Image(shape=(512, 512),label='Generated Image').style(width=512, height=512),
                               gr.Image(shape=(512, 512),label='Generated Images with custom loss').style(width=512, height=512)
                              ],
                    examples=prompt_examples,
                    title = title,
                    description = description
                   )
demo.launch(debug=True)