image-diffusion / app.py
paavansundar's picture
Update app.py
f45a960
raw
history blame
No virus
528 Bytes
import torch
from diffusers import StableDiffusionPipeline
import gradio as gr
pipe = StableDiffusionPipeline.from_pretrained("paavansundar/image-diffusion", torch_dtype=torch.float32)
def generate_img(prompt):
pipe.enable_attention_slicing()
return pipe(prompt).images[0]
with gr.Blocks() as demo:
im = gr.Image()
prompt = gr.Textbox(label="Input", lines=2)
btn = gr.Button(value="Generate Image")
btn.click(generate_img, inputs=[prompt], outputs=[im])
if __name__ == "__main__":
demo.launch()