kjgfhjh

app.py

@@ -2,6 +2,7 @@ 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)
@@ -12,8 +13,9 @@ pipe.to("cuda")
 @spaces.GPU
 def generate_image(prompt):
     
-    images = pipe(prompt).images
-    return images
+    image = pipe(prompt).images
+    image = rembg.remove(image)
+    return image
 
 _TITLE = "Shoe Generator"
 with gr.Blocks(_TITLE) as ShoeGen:
@@ -22,10 +24,8 @@ with gr.Blocks(_TITLE) as ShoeGen:
             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")
-    
+            image = gr.Image(label="Generated Image", show_download_button=True)    
 
-    button_gen.click(generate_image, inputs=[prompt], outputs=gallery)
+    button_gen.click(generate_image, inputs=[prompt], outputs=image)
 
 ShoeGen.launch()