| import gradio as gr | |
| import torch | |
| import diffusers | |
| pipe=diffusers.StableDiffusionPipeline.from_pretrained('stabilityai/stable-diffusion-2-1',torch_dtype=torch.float16) | |
| pipe.to('cuda') | |
| def gen(prompt): | |
| img=pipe(prompt)[0][0] | |
| return img | |
| with gr.Blocks() as demo: | |
| prompt = gr.Textbox(label="Prompt") | |
| output = gr.Image(label="Output Box") | |
| greet_btn = gr.Button("Generate") | |
| greet_btn.click(fn=gen, inputs=prompt, outputs=output, api_name="gen") | |
| demo.launch(debug=True) |