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

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=[]
    for i,ea in enumerate(loaded_model):
        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:
            out = Image.open(io.BytesIO(r.content))
            out_box.append(out)
        yield out_box

with gr.Blocks() as app:
    inp=gr.Textbox()
    outp=gr.Gallery()
    btn=gr.Button()
    btn.click(run_dif,inp,outp)
app.launch()