import gradio as gr
from models import models
from PIL import Image
import requests
import uuid
import io 
import base64


loaded_model=[]
for i,model in enumerate(models):
    try:
        loaded_model.append(gr.load(f'models/{model}'))
    except Exception as e:
        print(e)
        pass
print (loaded_model)

def run_dif(out_prompt):
    out_box=[]
    out_html=""
    for i,ea in enumerate(loaded_model):
        try:
            print (ea)
            model=loaded_model[int(i)]
            out_img=model(out_prompt)
            print(out_img)
            url=f'https://omnibus-top-20.hf.space/file={out_img}'
            print(url)
            uid = uuid.uuid4()
            #urllib.request.urlretrieve(image, 'tmp.png')
            #out=Image.open('tmp.png')
            r = requests.get(url, stream=True)
            
            if r.status_code == 200:
                img_buffer = io.BytesIO(r.content)
                print (f'bytes:: {io.BytesIO(r.content)}')
                str_equivalent_image = base64.b64encode(img_buffer.getvalue()).decode()
                img_tag = "<img src='data:image/png;base64," + str_equivalent_image + "'/>"                
                out_html+=f"<div><a href='https://huggingface.co/models/{ea}'>{ea}</a><br>"+img_tag+"</div>"
                out = Image.open(io.BytesIO(r.content))
                out_box.append(out)
            html_out = "<div class='grid_class'>"+out_html+"</div>"
            yield out_box,html_out
        except Exception as e:
            out_html+=e
            html_out = "<div class='grid_class'>"+out_html+"</div>"
            
            yield out_box,html_out
css="""
.grid_class{
display:flex;
}

"""
with gr.Blocks(css=css) as app:
    inp=gr.Textbox()
    out_html=gr.HTML()
    outp=gr.Gallery()
    btn=gr.Button()
    btn.click(run_dif,inp,[outp,out_html])
app.launch()