import os
import random
import gradio as gr
import requests
from PIL import Image

from utils import read_css_from_file
from inference import generate_image_from_text, generate_image_from_text_with_persistent_storage

# Read CSS from file
css = read_css_from_file("style.css")

DESCRIPTION = '''
<div id="content_align">
  <span style="color:darkred;font-size:32px;font-weight:bold">  
    WordCraft : Visuals from Verbs
  </span>
</div>
<div id="content_align">
  <span style="color:blue;font-size:18px;font-weight:bold;">  
    <br>A small, lighting fast efficient AI image generator
  </span>
</div>
<div id="content_align" style="margin-top: 10px;font-weight:bold;">
    <br>This 💻 demo uses the EfficientCLIP-GAN model trained on CUB and CC12M dataset. 
    <br>Keep your prompt coherent to domain of the selected model.
    <br>If you like the demo, don't forget to click on the like 💖 button.
</div>
'''
available_models = [
    ("EfficientCLIP-GAN trained on CUB dataset (Restricted to birds)", "CUB"),
    ("EfficientCLIP-GAN trained on CC12M dataset (More flexible)", "CC12M")
]

# Creating Gradio interface
with gr.Blocks(css=css) as app:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            text_prompt = gr.Textbox(label="Input Prompt", value="this tiny bird has a very small bill, a belly covered with white delicate feathers and has a set of black rounded eyes.", lines=3)
    model_selector = gr.Dropdown(choices=available_models, value="CUB", label="Select Model", info="Select a model with which you want to generate images")
    generate_button = gr.Button("Generate Images", variant='primary')

    with gr.Row():
        with gr.Column():
            image_output1 = gr.Image(label="Generated Image 1")
            image_output2 = gr.Image(label="Generated Image 2")

        with gr.Column():
            image_output3 = gr.Image(label="Generated Image 3")
            image_output4 = gr.Image(label="Generated Image 4")

    generate_button.click(generate_image_from_text_with_persistent_storage, inputs=[text_prompt, model_selector], outputs=[image_output1, image_output2, image_output3, image_output4])

# Launch the app
app.launch()