from PIL import Image
import gradio as gr
import numpy as np
import random, os, gc, base64, io
import cv2
import torch
from accelerate import Accelerator
from transformers import pipeline, DiffusionModel
from diffusers.utils import load_image
from diffusers import EulerDiscreteScheduler
from gradio_client import Client

accelerator = Accelerator(cpu=True)
pipe = accelerator.prepare(DiffusionModel.from_pretrained("stabilityai/sd-turbo", torch_dtype=torch.float32, use_safetensors=True, safety_checker=None))
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
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,negative_prompt="bad quality",scheduler=pipe.scheduler,num_inference_steps=5,width=512,height=512,generator=generator).images[0]
    apol.append(imags)
    return apol

iface = gr.Interface(fn=plex,inputs=[gr.Image(type="filepath"),gr.Textbox()], outputs=gr.Gallery(columns=2), title="Img2Img_SkyV22CntrlNet_CPU", description="Running on CPU, very slow!")
iface.queue(max_size=1)
iface.launch(max_threads=1)