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

def open_ims(choice):
    dirname='images/'+choice.replace(' ','_')+'/Seed_'+ str(seed_choice)+'/'
    imnames= os.listdir(dirname)
    images = [(Image.open(dirname+name)) for name in imnames]
    return images[:9]

def random_image():
    cat = random.choice(categories)
    dirname = 'images/'+cat.replace(' ','_')+'/Seed_'+str(seed_choice)+'/'
    imnames = os.listdir(dirname)
    images = [(Image.open(dirname+name)) for name in imnames]
    return images[:9]

categories=os.listdir('images/')
categories = [c.replace('_',' ') for c in categories]
seeds = [46267, 48040, 51237, 54325, 60884, 64830, 67031, 72935, 92118, 93109]

with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion Explorer")
    gr.Markdown("## Choose from the prompts below to explore the way in which the Stable Diffusion model encodes societal biases")
    seed_choice = gr.State(0)
    seed_choice = random.choice(seeds)
    print("Seed choice is: " + str(seed_choice))
    with gr.Row():
        with gr.Column():
            choice1 = gr.Dropdown(categories, label = "Choose a first category", interactive=True)
            images1 = gr.Gallery(label="Images").style(grid=[3], height="auto")
        with gr.Column():
            choice2 = gr.Dropdown(categories, label = "Choose a second category", interactive=True)
            images2 = gr.Gallery(label="Images").style(grid=[3], height="auto")

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