Txt2Img_SD_Trb / app.py
JoPmt's picture
Update app.py
7592d36
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)