Update main.py

main.py

@@ -4,6 +4,7 @@ from dotenv import load_dotenv
 load_dotenv()
 import prodia
 import requests
+from rembg import remove
 import random
 from datetime import datetime
 import os
@@ -46,6 +47,9 @@ def process_input_img2img(init, prompt, negative_prompt, steps, cfg_scale, numbe
                 f.write(img_data)
     return images
 
+def segment(image):
+     return remove(image)
+
 """
 def process_input_control(init, prompt, negative_prompt, steps, cfg_scale, number, seed, model, control_model, sampler):
     images = []
@@ -205,6 +209,30 @@ with gr.Blocks(theme=theme) as demo:
             outputs=[result_image],
         )
 
+    with gr.Tab("Remove Background"):
+        with gr.Row():
+        
+            with gr.Column():
+                i_img = gr.Image(label="Input Image")
+
+        with gr.Row():
+
+            with gr.Column():
+                run_btn = gr.Button("Remove", variant="primary")
+    
+        with gr.Row():
+
+            with gr.Column():
+                o_img = gr.Image(label="Result Image")
+
+        run_btn.click(
+                segment,
+                inputs=[
+                    i_img,
+                ],
+                outputs=[o_img],
+        )
+    
     with gr.Tab(label="Gallery"):
 
         gr.load("nateraw/stable_diffusion_gallery", src="spaces")