|
|
import spaces |
|
|
import gradio as gr |
|
|
import torch |
|
|
from diffusers import DiffusionPipeline |
|
|
import rembg |
|
|
|
|
|
model_id = "stabilityai/stable-diffusion-2-1" |
|
|
pipe = DiffusionPipeline.from_pretrained(model_id) |
|
|
|
|
|
pipe.to("cuda") |
|
|
|
|
|
|
|
|
@spaces.GPU |
|
|
def generate_image(prompt): |
|
|
|
|
|
image = pipe(prompt).images |
|
|
image = rembg.remove(image) |
|
|
return image |
|
|
|
|
|
_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(): |
|
|
image = gr.Image(label="Generated Image", show_download_button=True) |
|
|
|
|
|
button_gen.click(generate_image, inputs=[prompt], outputs=image) |
|
|
|
|
|
ShoeGen.launch() |
|
|
|
