import torch
import gradio as gr
from diffusers import UniDiffuserPipeline
from diffusers.utils import load_image
from accelerate import Accelerator
accelerator = Accelerator(cpu=True)

pipe = accelerator.prepare(UniDiffuserPipeline.from_pretrained("thu-ml/unidiffuser-v1", torch_dtype=torch.bfloat16))
pipe = pipe.to("cpu")
apol=[]
def plex(image_url,stips):
    init_image = load_image(image_url).resize((512, 512))
    sample = pipe(image=init_image, num_inference_steps=stips, guidance_scale=8.0)
    i2t_text = sample.text[0]
    sample = pipe(prompt=i2t_text, num_inference_steps=stips, guidance_scale=8.0)
    for i, imge in enumerate(sample["images"]):
        apol.append(imge)
    return apol

iface = gr.Interface(fn=plex, inputs=[gr.Image(label="img",type="filepath"), gr.Slider(label="num inference steps", minimum=1, step=1, maximum=5, value=5)], outputs=gr.Gallery(label="out", columns=2))
iface.queue(max_size=1,api_open=False)
iface.launch(max_threads=1)