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)