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

def open_ims(choice):
    print(choice)
    imnames = os.listdir('images/'+'Photo_portrait_of_a_'+choice.replace(' ',' ')+'/Seed_46267/')
    images = [(Image.open(name)) for name in imnames]
    return images

def random_image():
    cat = random.choice(categories)
    print(cat)
    imnames = os.listdir('images/'+'Photo_portrait_of_a_'+choice.replace(' ',' ')+'/Seed_46267/')
    print(imnames)
    ims = [(Image.open(name)) for name in imnames]
    print(ims)
    return ims

categories=os.listdir('images/')
categories = [c.replace('Photo_portrait_of_a_','').replace('_',' ') for c in categories]
print(categories)

with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion Explorer")
    gr.Markdown("## TODO")

    with gr.Row():
        with gr.Column():
              choice = gr.Dropdown(categories, label = "Choose a first category", interactive=True)
              images = gr.Gallery(label="Images"),
        with gr.Column():
              gr.Textbox(label="Text prompt"),
              gr.Image(optional=True, label="Image prompt", type='filepath'),

    demo.load(random_image, None, images)          
    choice.change(open_ims, [prompt], [images])

demo.launch()