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/sd-turbo", torch_dtype=torch.float32, use_safetensors=True, 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,negative_prompt=["bad quality"]*2,num_inference_steps=5,width=512,height=512,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 SD-Turbo CPU", description="Running on CPU, very slow! by JoPmt")
iface.queue(max_size=1,api_open=False)
iface.launch(max_threads=1)