import gradio as gr
import os
import base64

stable_diffusion = gr.Blocks.load(name="spaces/runwayml/stable-diffusion-v1-5")
clip_interrogator_2 = gr.Blocks.load(name="spaces/fffiloni/CLIP-Interrogator-2")

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

def get_new_prompt(img):
    interrogate = clip_interrogator_2(img, 'best', 4, fn_index=1)
    print(interrogate)
    return "done"

def infer(prompt):
    img = get_images(prompt)
    
    with open(img, "rb") as img_file:
        my_string = base64.b64encode(img_file.read())
    #print(my_string)
    print("img path: " + img)
    #img_to_send = "data:image/png;base64," + my_string
    #print(img_to_send)
    result = get_new_prompt(img)
    print(result)
    return "prompt written"
prompt_input = gr.Textbox(lines=4)
gallery_output = gr.Image(type='filepath')

gr.Interface(fn=infer, inputs=[prompt_input], outputs=['text']).queue(max_size=10).launch(enable_queue=True)