Update inpainting_module.py

inpainting_module.py

@@ -471,7 +471,8 @@ class InpaintingModule:
     def prepare_control_image(
         self,
         image: Image.Image,
-        mode: str = "canny"
+        mode: str = "canny",
+        mask: Optional[Image.Image] = None
     ) -> Image.Image:
         """
         Generate ControlNet conditioning image.
@@ -482,6 +483,9 @@ class InpaintingModule:
             Input image
         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.
 
         Returns
         -------
@@ -497,7 +501,24 @@ class InpaintingModule:
         img_array = np.array(image)
 
         if mode == "canny":
-            return self._generate_canny_edges(img_array)
+            canny_image = self._generate_canny_edges(img_array)
+
+            # Mask-aware processing: weaken edges in masked region
+            if mask is not None:
+                canny_array = np.array(canny_image)
+                mask_array = np.array(mask.convert('L'))
+
+                # In masked region, reduce Canny edge strength by 80%
+                # This prevents ControlNet from locking original colors
+                mask_region = mask_array > 128  # White = masked area
+                canny_array[mask_region] = (canny_array[mask_region] * 0.2).astype(np.uint8)
+                # ↑ Keep 20% edge strength for basic shape preservation
+
+                canny_image = Image.fromarray(canny_array)
+                logger.info("Applied mask-aware ControlNet: weakened edges in masked region by 80%")
+
+            return canny_image
+
         elif mode == "depth":
             return self._generate_depth_map(image)
         else:
@@ -890,7 +911,8 @@ class InpaintingModule:
 
             control_image = self.prepare_control_image(
                 image,
-                self._current_conditioning_type
+                self._current_conditioning_type,
+                mask=processed_mask  # Pass mask for mask-aware ControlNet processing
             )
 
             # Conditional prompt enhancement based on template