ShoeGenv2 / app.py
MaxMilan1
kgjk
ad9ba71
raw
history blame
844 Bytes
import spaces
import gradio as gr
import torch
from diffusers import DiffusionPipeline
model_id = "stabilityai/stable-diffusion-2-1"
pipe = DiffusionPipeline.from_pretrained(model_id)
pipe.to("cuda")
# Function to generate an image from text using diffusion
@spaces.GPU
def generate_image(prompt):
images = pipe(prompt).images
return images
_TITLE = "Shoe Generator"
with gr.Blocks(_TITLE) as ShoeGen:
with gr.Row():
with gr.Column():
prompt = gr.Textbox(label="Enter a prompt")
button_gen = gr.Button("Generate Image")
with gr.Column():
# show images
gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
button_gen.click(generate_image, inputs=[prompt], outputs=gallery)
ShoeGen.launch()