Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import gradio as gr
 from load_image import load_img
 import spaces
@@ -10,17 +11,25 @@ import numpy as np
 
 torch.set_float32_matmul_precision(["high", "highest"][0])
 
-# load model
+# load 2 models
+
 birefnet = AutoModelForImageSegmentation.from_pretrained(
     "ZhengPeng7/BiRefNet", trust_remote_code=True
 )
 
-# Keep model in a dict for easy switching
+
+RMBG2 = AutoModelForImageSegmentation.from_pretrained(
+    "briaai/RMBG-2.0", trust_remote_code=True
+)
+
+# Keep them in a dict to switch easily
 models_dict = {
-    "BiRefNet": birefnet
+    "BiRefNet": birefnet,
+    "RMBG-2.0": RMBG2
 }
 
 # Transform
+
 transform_image = transforms.Compose(
     [
         transforms.Resize((1024, 1024)),
@@ -35,6 +44,7 @@ def process(image: Image.Image, model_choice: str):
     Runs inference to remove the background (adds alpha) 
     with the chosen segmentation model.
     """
+    # Select the model
     current_model = models_dict[model_choice]
 
     # Prepare image
@@ -43,6 +53,8 @@ def process(image: Image.Image, model_choice: str):
 
     # Inference
     with torch.no_grad():
+        # Each model returns a list of preds in its forward, 
+        # so we take the last element, apply sigmoid, and move to CPU
         preds = current_model(input_images)[-1].sigmoid().cpu()
 
     # Convert single-channel pred to a PIL mask
@@ -56,7 +68,6 @@ def process(image: Image.Image, model_choice: str):
     image.putalpha(mask)
     return image
 
-
 def fn(source: str, model_choice: str):
     """
     Used by Tab 1 & Tab 2 to produce a processed image with alpha.
@@ -64,6 +75,7 @@ def fn(source: str, model_choice: str):
       a URL string (textbox).
     - 'model_choice' is the user's selection from the radio.
     """
+    # Load from local path or URL
     im = load_img(source, output_type="pil")
     im = im.convert("RGB")
 
@@ -71,7 +83,6 @@ def fn(source: str, model_choice: str):
     processed_image = process(im, model_choice)
     return processed_image
 
-
 def process_file(file_path: str, model_choice: str):
     """
     For Tab 3 (file output).
@@ -82,16 +93,29 @@ def process_file(file_path: str, model_choice: str):
     im = load_img(file_path, output_type="pil")
     im = im.convert("RGB")
 
+    # Run the chosen model
     transparent = process(im, model_choice)
     transparent.save(name_path)
     return name_path
 
 
-# Gradio UI
-model_selector = gr.Radio(
-choices=["BiRefNet"],
-value="BiRefNet",
-label="Select Model")
+# GRadio UI
+
+model_selector_1 = gr.Radio(
+    choices=["BiRefNet","RMBG-2.0"]
+    value="BiRefNet",
+    label="Select Model"
+)
+model_selector_2 = gr.Radio(
+    choices=["BiRefNet","RMBG-2.0"],
+    value="BiRefNet",
+    label="Select Model"
+)
+model_selector_3 = gr.Radio(
+    choices=["BiRefNet", "RMBG-2.0"],
+    value="BiRefNet",
+    label="Select Model"
+)
 
 # Outputs for tabs 1 & 2: single processed image
 processed_img_upload = gr.Image(label="Processed Image (Upload)", type="pil")
@@ -110,7 +134,7 @@ output_file = gr.File(label="Output PNG File")
 # Tab 1: local image -> processed image
 tab1 = gr.Interface(
     fn=fn,
-    inputs=[image_upload, model_selector],
+    inputs=[image_upload, model_selector_1],
     outputs=processed_img_upload,
     api_name="image",
     description="Upload an image and choose your background removal model."
@@ -119,7 +143,7 @@ tab1 = gr.Interface(
 # Tab 2: URL input -> processed image
 tab2 = gr.Interface(
     fn=fn,
-    inputs=[url_input, model_selector],
+    inputs=[url_input, model_selector_2],
     outputs=processed_img_url,
     api_name="text",
     description="Paste an image URL and choose your background removal model."
@@ -128,36 +152,20 @@ tab2 = gr.Interface(
 # Tab 3: file output -> returns path to .png
 tab3 = gr.Interface(
     fn=process_file,
-    inputs=[image_file_upload, model_selector],
+    inputs=[image_file_upload, model_selector_3],
     outputs=output_file,
     api_name="png",
     description="Upload an image, choose a model, and get a transparent PNG."
 )
 
 # Combine all tabs
-with gr.Blocks() as demo:
-    tabs = gr.Tabs()
-    with tabs:
-        with gr.TabItem("Image Upload"):
-            inp1 = gr.Image(type="filepath", label="Upload an image")
-            mdl1 = gr.Radio(["BiRefNet"], value="BiRefNet", label="Select Model")
-            out1 = gr.Image(type="pil", label="Processed Image (Upload)")
-            btn1 = gr.Button("Run")
-            btn1.click(fn=fn, inputs=[inp1,mdl1], outputs=out1)
-
-        with gr.TabItem("URL Input"):
-            inp2 = gr.Textbox(label="Paste an image URL")
-            mdl2 = gr.Radio(["BiRefNet"], value="BiRefNet", label="Select Model")
-            out2 = gr.Image(type="pil", label="Processed Image (URL)")
-            btn2 = gr.Button("Run")
-            btn2.click(fn=fn, inputs=[inp2,mdl2], outputs=out2)
-
-        with gr.TabItem("File Output"):
-            inp3 = gr.Image(type="filepath", label="Upload an image")
-            mdl3 = gr.Radio(["BiRefNet"], value="BiRefNet", label="Select Model")
-            out3 = gr.File(label="Output PNG File")
-            btn3 = gr.Button("Run")
-            btn3.click(fn=process_file, inputs=[inp3,mdl3], outputs=out3)
-
-demo.launch(show_error=True, share=True)
+demo = gr.TabbedInterface(
+    [tab1, tab2, tab3],
+    ["Image Upload", "URL Input", "File Output"],
+    title="Background Removal Tool"
+)
+
+if __name__ == "__main__":
+    demo.launch(show_error=True, share=True)
+