Update gradio_demo.py

gradio_demo.py

@@ -743,7 +743,8 @@ class MaskMover:
     def __init__(self):
         self.extracted_fg = None
         self.original_fg = None  # Store original foreground
-        
+        self.current_bg = None
+
     def set_extracted_fg(self, fg_image):
         """Store the extracted foreground with alpha channel"""
         if isinstance(fg_image, np.ndarray):
@@ -753,11 +754,19 @@ class MaskMover:
             self.extracted_fg = np.array(fg_image)
             self.original_fg = np.array(fg_image)
         return self.extracted_fg
+
+    def update_background(self, new_bg):
+        """Update the current background without affecting the mask"""
+        if new_bg is not None:
+            self.current_bg = new_bg.copy()  # Overwrite the current background
+            print("Background updated successfully.")  # Debugging
+
     
     def create_composite(self, background, x_pos, y_pos, scale=1.0):
         """Create composite with foreground at specified position"""
-        if self.original_fg is None or background is None:
-            return background
+        
+        if self.original_fg is None or self.current_bg is None:
+            return self.current_bg  # Return current background if no foreground
         
         # Convert inputs to PIL Images
         if isinstance(background, np.ndarray):
@@ -814,6 +823,9 @@ def extract_foreground(image):
             
             return result, gr.update(visible=True), gr.update(visible=True)
 
+def on_background_upload(new_bg):
+    mask_mover.update_background(new_bg) 
+
 
 def compress_image(image):
     # Convert Gradio image (numpy array) to PIL Image
@@ -1213,6 +1225,7 @@ with block:
                     with gr.Group():
                         gr.Markdown("Position on Background")
                         input_bg = gr.Image(type="numpy", label="Background Image", height=480)
+                        reset_button = gr.Button(value="Reset Background")
                         
                         with gr.Row():
                             x_slider = gr.Slider(
@@ -1371,7 +1384,13 @@ with block:
                     self.original_bg = new_bg.copy()
                     self.current_mask = None  # Reset current mask when background changes
                     print("Background updated successfully.")  # Debugging
-        
+
+            def reset_background(self):
+                """Reset the background to its original state"""
+                if self.original_bg is not None:
+                    self.current_bg = self.original_bg.copy()  # Reset to original background
+                    print("Background reset to original state.")  # Debugging
+            
             def update_position(self, x_pos, y_pos, scale):
                 """Update composite when position changes."""
                 if self.original_bg is None:
@@ -1437,12 +1456,12 @@ with block:
         # )
 
                 # Update the input_bg.change function to reset the mask
-        input_bg.change(
-            fn=lambda new_bg: bg_manager.update_background(new_bg) or new_bg,
-            inputs=[input_bg],
-            outputs=[input_bg],
-            show_progress=False
-        )
+        # input_bg.change(
+        #     fn=lambda new_bg: bg_manager.update_background(new_bg) or new_bg,
+        #     inputs=[input_bg],
+        #     outputs=[input_bg],
+        #     show_progress=False
+        # )
         
         # Update the slider change functions to apply the current mask
         x_slider.change(
@@ -1565,6 +1584,13 @@ with block:
         inputs=ips_bg,
         outputs=[result_gallery], show_progress=True
     )
+
+    reset_button.click(
+        fn=lambda: bg_manager.reset_background(),  # Call the reset function
+        inputs=[],  # No inputs needed for reset
+        outputs=[input_bg],  # Update the displayed background
+        show_progress=True
+    )
     
     example_quick_prompts.click(lambda x, y: ', '.join(y.split(', ')[:2] + [x[0]]), inputs=[example_quick_prompts, prompt], outputs=prompt, show_progress=False, queue=False)
     example_quick_subjects.click(lambda x: x[0], inputs=example_quick_subjects, outputs=prompt, show_progress=False, queue=False)