Update inpainting_module.py

inpainting_module.py

@@ -473,7 +473,8 @@ class InpaintingModule:
         self,
         image: Image.Image,
         mode: str = "canny",
-        mask: Optional[Image.Image] = None
+        mask: Optional[Image.Image] = None,
+        preserve_structure: bool = False
     ) -> Image.Image:
         """
         Generate ControlNet conditioning image.
@@ -485,15 +486,17 @@ class InpaintingModule:
         mode : str
             Conditioning mode: "canny" or "depth"
         mask : PIL.Image, optional
-            If provided, weakens ControlNet edges in masked region to give prompt more control.
-            This is crucial for color transformation tasks where strong edges can lock colors.
+            If provided, can suppress edges in masked region (when preserve_structure=False).
+        preserve_structure : bool
+            If True, keep edges in masked region (for color change tasks).
+            If False, suppress edges in masked region (for replacement/removal tasks).
 
         Returns
         -------
         PIL.Image
             Generated control image (edges or depth map)
         """
-        logger.info(f"Preparing control image with mode: {mode}")
+        logger.info(f"Preparing control image with mode: {mode}, preserve_structure: {preserve_structure}")
 
         # Convert to RGB if needed
         if image.mode != 'RGB':
@@ -504,19 +507,20 @@ class InpaintingModule:
         if mode == "canny":
             canny_image = self._generate_canny_edges(img_array)
 
-            # Mask-aware processing: weaken edges in masked region
-            if mask is not None:
+            # Mask-aware processing: suppress edges in masked region ONLY if not preserving structure
+            if mask is not None and not preserve_structure:
                 canny_array = np.array(canny_image)
                 mask_array = np.array(mask.convert('L'))
 
                 # In masked region, completely suppress Canny edges
-                # This prevents ControlNet from locking original colors
+                # This allows complete replacement/removal of the object
                 mask_region = mask_array > 128  # White = masked area
                 canny_array[mask_region] = 0
-                # ↑ Complete suppression gives prompt maximum control over color transformation
 
                 canny_image = Image.fromarray(canny_array)
-                logger.info("Applied mask-aware ControlNet: suppressed edges in masked region completely")
+                logger.info("Suppressed edges in masked region for replacement/removal")
+            elif preserve_structure:
+                logger.info("Preserving edges in masked region for color change")
 
             return canny_image
 
@@ -866,6 +870,7 @@ class InpaintingModule:
             )
             feather_radius = kwargs.get('feather_radius', self.config.feather_radius)
             strength = kwargs.get('strength', self.config.strength)
+            preserve_structure = kwargs.get('preserve_structure_in_mask', False)
 
             if progress_callback:
                 progress_callback("Preparing images...", 5)
@@ -914,7 +919,8 @@ class InpaintingModule:
             control_image = self.prepare_control_image(
                 image,
                 self._current_conditioning_type,
-                mask=processed_mask  # Pass mask for mask-aware ControlNet processing
+                mask=processed_mask,
+                preserve_structure=preserve_structure  # True for color change, False for replacement/removal
             )
 
             # Conditional prompt enhancement based on template
@@ -1059,6 +1065,7 @@ class InpaintingModule:
                     "conditioning_type": self._current_conditioning_type,
                     "conditioning_scale": conditioning_scale,
                     "strength": strength,
+                    "preserve_structure": preserve_structure,
                     "num_inference_steps": num_steps,
                     "guidance_scale": guidance,
                     "feather_radius": feather_radius,