import gradio as gr
import random, os
from PIL import Image
import pandas as pd

def open_ims(adj, group):
    if group != '':
        if adj != '':
            if adj[0] in vowels:
                dirname='images/'+'Photo_portrait_of_an_'+ adj+'_'+ group.replace(' ','_')+'/Seed_'+ str(seed_choice)+'/'
            else:
                dirname='images/'+'Photo_portrait_of_a_'+ adj+ '_'+ group.replace(' ','_')+'/Seed_'+ str(seed_choice)+'/'
        else:
            if group[0] in vowels:
                dirname='images/'+'Photo_portrait_of_an_'+ group.replace(' ','_')+'/Seed_'+ str(seed_choice)+'/'
            else:
                dirname='images/'+'Photo_portrait_of_a_'+ group.replace(' ','_')+'/Seed_'+ str(seed_choice)+'/'
        imnames= os.listdir(dirname)
        images = [(Image.open(dirname+name)) for name in imnames]
        return images[:9]


vowels = ["a","e","i","o","u"]
prompts = pd.read_csv('shortprompts.csv')

seeds = [46267, 48040, 51237, 54325, 60884, 64830, 67031, 72935, 92118, 93109]


m_adjectives = prompts['Masc-adj'].tolist()[:11]
f_adjectives = prompts['Fem-adj'].tolist()[:11]
adjectives = sorted(m_adjectives+f_adjectives)
#adjectives = ['attractive','strong']
adjectives.insert(0, '')
professions = sorted([p.lower() for p in prompts['Occupation-Noun'].tolist()])

with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion Explorer")
    gr.Markdown("## Choose from the prompts below to explore how the Stable Diffusion model represents different professions")
    seed_choice = gr.State(0)
    seed_choice = 93109
    print("Seed choice is: " + str(seed_choice))
    with gr.Row():
        with gr.Column():
            adj1 = gr.Dropdown(adjectives, label = "Choose a first adjective (or leave this blank!)", interactive=True)
            choice1 = gr.Dropdown(professions, label = "Choose a first group", interactive=True)
            images1 = gr.Gallery(label="Images").style(grid=[3], height="auto")
        with gr.Column():
            adj2 = gr.Dropdown(adjectives, label = "Choose a second adjective (or leave this blank!)", interactive=True)
            choice2 = gr.Dropdown(professions, label = "Choose a second group", interactive=True)
            images2 = gr.Gallery(label="Images").style(grid=[3], height="auto")

    #demo.load(random_image, None, [images])
    choice1.change(open_ims, [adj1,choice1], [images1])
    choice2.change(open_ims, [adj2,choice2], [images2])
    adj1.change(open_ims, [adj1,choice1], [images1])
    adj2.change(open_ims, [adj2,choice2], [images2])
demo.launch()