Spaces:
Runtime error
Runtime error
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() |