Update ui_manager.py

ui_manager.py

@@ -635,15 +635,27 @@ class UIManager:
                 logger.warning(f"Unexpected editor output type: {type(editor_output)}")
                 return None
 
-            # Convert to grayscale if needed
+            # Convert to grayscale mask
             if len(mask_array.shape) == 3:
                 if mask_array.shape[2] == 4:
-                    # RGBA - use alpha channel
-                    mask_gray = mask_array[:, :, 3]
-                else:
-                    # RGB - convert to grayscale
+                    # RGBA format - extract white brush strokes from RGB channels
+                    # White brush strokes have high RGB values AND high alpha
+                    rgb_part = mask_array[:, :, :3]
+                    alpha_part = mask_array[:, :, 3]
+
+                    # Convert RGB to grayscale to detect white areas
+                    gray = cv2.cvtColor(rgb_part, cv2.COLOR_RGB2GRAY)
+
+                    # Combine: white areas (high gray value) with opacity (high alpha)
+                    # This captures white brush strokes
+                    mask_gray = np.minimum(gray, alpha_part)
+                elif mask_array.shape[2] == 3:
+                    # RGB - convert to grayscale (white areas become white in mask)
                     mask_gray = cv2.cvtColor(mask_array, cv2.COLOR_RGB2GRAY)
+                else:
+                    mask_gray = mask_array[:, :, 0]
             else:
+                # Already grayscale
                 mask_gray = mask_array
 
             return Image.fromarray(mask_gray.astype(np.uint8), mode='L')
@@ -840,7 +852,9 @@ class UIManager:
                         height=300,
                         brush=gr.Brush(colors=["#FFFFFF"], default_size=20),
                         eraser=gr.Eraser(default_size=20),
-                        layers=False
+                        layers=True,
+                        sources=["upload"],
+                        image_mode="RGBA"
                     )
 
                     # Template selection