Spaces:
Runtime error
Runtime error
File size: 1,293 Bytes
5c80069 092a751 ea10931 092a751 5c80069 092a751 5c80069 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import gradio as gr
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline
model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token='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', torch_dtype=torch.float32, low_cpu_mem_usage=True)
has_cuda = torch.cuda.is_available()
device = torch.device('cpu' if not has_cuda else 'cuda')
pipe = pipe.to(device)
def inference(diffusion_prompt):
samples = 4
generator = torch.Generator(device=device)
torch.cuda.empty_cache()
with autocast("cuda"):
images_list = pipe(
[diffusion_prompt] * samples,
height=256, width=384,
num_inference_steps=50,
)
images = []
for i, image in enumerate(images_list["sample"]):
images.append(image)
return images
with gr.Blocks() as demo:
gr.Markdown("# Text to Image Generator")
with gr.Row():
prompt = gr.Textbox(
lines=1,
placeholder="Enter your prompt..",
interactive=True,
label="Prompt"
)
submit = gr.Button("Run")
submit.click(fn=inference, inputs=[prompt], outputs=[images])
demo.launch() |