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

import torch
from diffusers import AutoPipelineForImage2Image
from diffusers.utils import make_image_grid, load_image

base_url=f'https://omnibus-top-20-img-img-basic.hf.space/file='
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)

pipeline = AutoPipelineForImage2Image.from_pretrained("runwayml/stable-diffusion-v1-5", safety_checker=None, variant="fp16", use_safetensors=True).to("cpu")

def load_model(model_drop):
    pipeline = AutoPipelineForImage2Image.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float32, use_safetensors=True)

def run_dif(prompt,im_path,model_drop,cnt,strength,guidance,infer):
    out_box=[]
    for i in range(int(cnt)):
        yield out_box,f"Working on {i} of {int(cnt)}"
        url = base_url+im_path
        print(url)
        init_image=load_image(url)
        #image = pipeline(prompt, image=init_image, strength=0.8,guidance_scale=8.0,negative_prompt=negative_prompt,num_inference_steps=50).images[0]
        image = pipeline(prompt, image=init_image, strength=float(strength),guidance_scale=float(guidance),num_inference_steps=int(infer)).images[0]
        out_box.append(image)
    yield out_box,"Complete"


css="""
.grid_class{
display:flex;
height:100%;
}
.img_class{
min-width:200px;
}

"""

with gr.Blocks(css=css) as app:
    with gr.Row():
        with gr.Column():
            inp=gr.Textbox(label="Prompt")
            strength=gr.Slider(label="Strength",minimum=0,maximum=1,step=0.1,value=0.2)
            guidance=gr.Slider(label="Guidance",minimum=0,maximum=10,step=0.1,value=8.0)
            infer=gr.Slider(label="Inference Steps",minimum=0,maximum=50,step=1,value=10)
            
            with gr.Row():
                btn=gr.Button()
                stop_btn=gr.Button("Stop")
        with gr.Column():
            inp_im=gr.Image(type='filepath')
            im_btn=gr.Button("Image Grid")
    with gr.Row():
        model_drop=gr.Dropdown(label="Models", choices=models, type='index', value=models[0])
        cnt = gr.Number(value=1)
    out_html=gr.HTML()
    outp=gr.Gallery(columns=10)
    #fingal=gr.Gallery(columns=10)
    #im_list=gr.Textbox()    
    #im_btn.click(load_im,inp_im,[outp,im_list])
    go_btn = btn.click(run_dif,[inp,inp_im,model_drop,cnt,strength,guidance,infer],[outp,out_html])
    stop_btn.click(None,None,None,cancels=[go_btn])
app.queue().launch()