Txt2Img_SD_Trb / app.py
JoPmt's picture
Update app.py
5a73097
raw
history blame
1.18 kB
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)
iface.launch(max_threads=1)