import gradio as gr
import os


img_to_text = gr.Interface.load("spaces/pharma/CLIP-Interrogator")

stable_diffusion = gr.Blocks.load(name="spaces/stabilityai/stable-diffusion")

def get_images(prompt):
    gallery_dir = stable_diffusion(prompt, fn_index=2)
    return [os.path.join(gallery_dir, image) for image in os.listdir(gallery_dir)]

def get_prompts(uploaded_image):
    return img_to_text(uploaded_image)


with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown(
            """
            ## Stable Diffusion Perception 🎆🌌
            Input an image and see how the model perceives it! 👀
            """
        )

        gr.Markdown(
            """
            Upload your image below 👇
            """
        )
        
        with gr.Column():
            input_img = gr.Image(
                                type="filepath",
                            )
            with gr.Row():
                translate_img = gr.Button("Upload Image!", elem_id = "img2text")
                see_prompts = gr.Button("Check how your image prompts your model!", elem_id="check_btn_1")              
                #translate_img_directly = gr.Button("Translate your image now.", elem_id="translate_img")
        
        with gr.Accordion(label="Stable Diffusion Settings", elem_id="sd_settings", visible=False):
            with gr.Row():
                guidance_scale = gr.Slider(2, 15, value = 7, label = 'Guidance Scale')
                nb_iterations = gr.Slider(10, 50, value = 25, step = 1, label = 'Steps')
                seed = gr.Slider(label = "Seed", minimum = 0, maximum = 2147483647, step = 1, randomize = True)
        
        gr.Markdown(
            """
            ## Check what the model thinks of the image 📜
            """
        )
        
        with gr.Column():
            img2text_output = gr.Textbox(
                                    label="Convert your image to text!", 
                                    lines=4,
                                    elem_id="translated"
                                )
            with gr.Row():
                clear_btn = gr.Button(value="Clear")
                diffuse_btn = gr.Button(value="Diffuse it!", elem_id="diffuse_btn")
                
                clear_btn.click(fn=lambda value: gr.update(value=""), inputs=clear_btn, outputs=img2text_output)
    
                
                    
        gr.Markdown("""
            ## 3. Diffuse it! 🧑‍🎨
            """
            ) 
        
        sd_output = gr.Gallery().style(grid=2, height="auto")
                
                
        see_prompts.click(img_to_text, 
                                inputs = input_img, 
                                outputs = [
                                    img2text_output
                                ])
        
        def translate_directly(img):
            images = get_images(get_prompts(img))
            return images

        """
        translate_img_directly.click(translate_directly, 
                                inputs = [
                                    input_img
                                ], 
                                outputs = [
                                    img2text_output,
                                    sd_output
                                ])
        
        """
        diffuse_btn.click(get_images, 
                              inputs = [
                                  img2text_output
                                  ], 
                              outputs = sd_output
                          )




demo.launch()