from PIL import Image
import gradio as gr
import random, os, gc
import torch
from accelerate import Accelerator
from transformers import pipeline
from diffusers.utils import load_image
from diffusers import EulerDiscreteScheduler, DiffusionPipeline

accelerator = Accelerator(cpu=True)
pipe = accelerator.prepare(DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.bfloat16, use_safetensors=True, variant="fp16", safety_checker=None))
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.unet.to(memory_format=torch.channels_last)
pipe = accelerator.prepare(pipe.to("cpu"))
generator = torch.Generator("cpu").manual_seed(random.randint(1, 867346))
apol=[]

def plex(prompt):
    gc.collect()
    apol=[]
    imags = pipe(prompt=[prompt]*2,num_inference_steps=2,width=512,height=512,guidance_scale=0.0,generator=generator)
    for i, igs in enumerate(imags["images"]):
        apol.append(igs)
    return apol

iface = gr.Interface(fn=plex,inputs=gr.Textbox(), outputs=gr.Gallery(columns=2), title="Stabilityai SDXL-Turbo CPU", description="Running on CPU, very slow! by JoPmt")
iface.queue(max_size=1)
iface.launch(max_threads=1)