Update app.py

app.py

@@ -281,13 +281,16 @@ image_mask_list = sorted([os.path.join(image_mask_dir, f) for f in os.listdir(im
 
 
 @spaces.GPU
-def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt):
+def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_option, ref_mask_option, base_text_prompt, ref_text_prompt):
     if base_mask_option == "Draw Mask":
         tar_image = base_image["background"]
         tar_mask = base_image["layers"][0]
-    else:
+    elif base_mask_option == "Upload with Mask":
         tar_image = base_image["background"]
         tar_mask = base_mask["background"]
+    else:  # Label to Mask
+        tar_image = base_image["background"]
+        tar_mask = get_mask(tar_image, base_text_prompt)
 
     if ref_mask_option == "Draw Mask":
         ref_image = reference_image["background"]
@@ -295,9 +298,9 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     elif ref_mask_option == "Upload with Mask":
         ref_image = reference_image["background"]
         ref_mask = ref_mask["background"]
-    else:
+    else:  # Label to Mask
         ref_image = reference_image["background"]
-        ref_mask = get_mask(ref_image, text_prompt)
+        ref_mask = get_mask(ref_image, ref_text_prompt)
 
     tar_image = tar_image.convert("RGB")
     tar_mask = tar_mask.convert("L")
@@ -393,10 +396,16 @@ def run_local(base_image, base_mask, reference_image, ref_mask, seed, base_mask_
     edited_image = crop_back(edited_image, old_tar_image, np.array([H1, W1, H2, W2]), np.array(tar_box_yyxx_crop))
     edited_image = Image.fromarray(edited_image)
 
-    if ref_mask_option != "Label to Mask":
-        return [show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
-    else:
-        return [return_ref_mask, show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
+    # Determine which masks to show as "generated" in output
+    masks_to_return = []
+    if base_mask_option == "Label to Mask":
+        masks_to_return.append(received_tar_mask)  # Show generated background mask
+    if ref_mask_option == "Label to Mask":
+        masks_to_return.append(received_ref_mask)   # Show generated reference mask
+    
+    # Build return list: generated_masks + diptych + final_image + received_masks
+    return_list = masks_to_return + [show_diptych_ref_tar, edited_image, received_tar_mask, received_ref_mask]
+    return return_list
 
 
 def update_ui(option):
@@ -420,8 +429,11 @@ with gr.Blocks() as demo:
                 base_mask = gr.ImageEditor(label="Background Mask", sources="upload", type="pil",
                                            layers=False, brush=False, eraser=False)
             with gr.Row():
-                base_mask_option = gr.Radio(["Draw Mask", "Upload with Mask"], label="Background Mask Input Option",
+                base_mask_option = gr.Radio(["Draw Mask", "Upload with Mask", "Label to Mask"], label="Background Mask Input Option",
                                             value="Upload with Mask")
+            with gr.Row():
+                base_text_prompt = gr.Textbox(label="Background Label",
+                                             placeholder="Enter the category to mask in background, e.g., sofa, table, person, etc.")
 
             with gr.Row():
                 ref_image = gr.ImageEditor(label="Reference Image", sources="upload", type="pil",
@@ -434,8 +446,8 @@ with gr.Blocks() as demo:
                 ref_mask_option = gr.Radio(["Draw Mask", "Upload with Mask", "Label to Mask"],
                                            label="Reference Mask Input Option", value="Upload with Mask")
             with gr.Row():
-                text_prompt = gr.Textbox(label="Label",
-                                         placeholder="Enter the category of the reference object, e.g., car, dress, toy, etc.")
+                ref_text_prompt = gr.Textbox(label="Reference Label",
+                                            placeholder="Enter the category of the reference object, e.g., car, dress, toy, etc.")
 
         with gr.Column(scale=1):
             baseline_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", height=695, columns=1)
@@ -446,6 +458,7 @@ with gr.Blocks() as demo:
                 gr.Markdown(" Draw Mask means manually drawing a mask on the original image.")
                 gr.Markdown(" Upload with Mask means uploading a mask file.")
                 gr.Markdown(" Label to Mask means simply inputting a label to automatically extract the mask and obtain the result.")
+                gr.Markdown(" Both background and reference images now support all three masking options including automatic mask generation from labels.")
 
     run_local_button = gr.Button(value="Run")
 
@@ -468,7 +481,7 @@ with gr.Blocks() as demo:
 
     run_local_button.click(
         fn=run_local,
-        inputs=[base_image, base_mask, ref_image, ref_mask, seed, base_mask_option, ref_mask_option, text_prompt],
+        inputs=[base_image, base_mask, ref_image, ref_mask, seed, base_mask_option, ref_mask_option, base_text_prompt, ref_text_prompt],
         outputs=[baseline_gallery]
     )
 demo.launch()