Merge pull request #43 from outis1one/claude/fix-ai-paint-tool-yz6q8

Fix U2Net, alertify dialogs, and improve AI Paint workflow

backend/app/routers/tools.py

@@ -204,18 +204,19 @@ _u2net_model = None
 
 
 async def _download_u2net_model(models_dir):
-    """Auto-download full U2Net model (~176MB) for best quality background removal"""
+    """Auto-download U2Net PyTorch model (~176MB) for background removal"""
     import urllib.request
     from pathlib import Path
 
     models_dir = Path(models_dir)
     models_dir.mkdir(parents=True, exist_ok=True)
 
-    # Download full U2Net model (176MB) for best quality
+    # Download U2Net PyTorch model (avoids ONNX executable stack issues in Docker)
+    # Using the PyTorch state dict format
     url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
     dest_path = models_dir / "u2net.onnx"
 
-    print(f"Downloading full U2Net model from {url} (~176MB)...")
+    print(f"Downloading U2Net model from {url} (~176MB)...")
     print("This may take a few minutes...")
 
     def download_progress(count, block_size, total_size):
@@ -227,40 +228,38 @@ async def _download_u2net_model(models_dir):
                 print(f"  Download progress: {percent}% ({downloaded_mb:.1f}/{total_mb:.1f} MB)")
 
     urllib.request.urlretrieve(url, str(dest_path), download_progress)
-    print(f"Full U2Net model downloaded to {dest_path}")
+    print(f"U2Net model downloaded to {dest_path}")
 
     return dest_path
 
 
 async def _remove_background_u2net(img: Image.Image) -> bytes:
     """
-    Remove background using U2Net model directly.
+    Remove background using U2Net model via OpenCV DNN.
-    This avoids rembg dependency issues while providing good quality.
+    Uses OpenCV's DNN module which doesn't have executable stack issues.
     """
     global _u2net_model
 
-    import torch
     from pathlib import Path
 
     # Check for U2Net model
     models_dir = Path('/app/data/models')
-    u2net_path = models_dir / 'u2net.pth'
+    u2net_path = None
 
-    # Also check alternative names
+    # Check for ONNX model (preferred for OpenCV DNN)
-    if not u2net_path.exists():
+    for alt_name in ['u2net.onnx', 'u2netp.onnx']:
-        for alt_name in ['u2net.onnx', 'u2netp.pth', 'u2net_human_seg.pth']:
+        alt_path = models_dir / alt_name
-            alt_path = models_dir / alt_name
+        if alt_path.exists():
-            if alt_path.exists():
+            u2net_path = alt_path
-                u2net_path = alt_path
+            break
-                break
 
-    if not u2net_path.exists():
+    if u2net_path is None:
         # Try to auto-download the model
         print("U2Net model not found, attempting to download...")
         try:
             await _download_u2net_model(models_dir)
             # Check again
-            for alt_name in ['u2net.onnx', 'u2netp.onnx', 'u2net.pth']:
+            for alt_name in ['u2net.onnx', 'u2netp.onnx']:
                 alt_path = models_dir / alt_name
                 if alt_path.exists():
                     u2net_path = alt_path
@@ -268,57 +267,49 @@ async def _remove_background_u2net(img: Image.Image) -> bytes:
         except Exception as download_error:
             print(f"Auto-download failed: {download_error}")
 
-    if not u2net_path.exists():
+    if u2net_path is None:
         raise FileNotFoundError(
             "U2Net model not found. To fix this, run:\n"
             "  docker exec -it ai-photo-edit-backend python /scripts/download_u2net_model.py\n"
             "Or manually download from: https://github.com/danielgatis/rembg/releases"
         )
 
-    # Load model if not cached
+    # Load model if not cached (using OpenCV DNN - no executable stack issues)
     if _u2net_model is None:
-        print(f"Loading U2Net model from {u2net_path}")
+        print(f"Loading U2Net model from {u2net_path} using OpenCV DNN")
-
+        try:
-        if str(u2net_path).endswith('.onnx'):
+            _u2net_model = cv2.dnn.readNetFromONNX(str(u2net_path))
-            # Use ONNX runtime
+            print("U2Net model loaded successfully with OpenCV DNN")
-            import onnxruntime as ort
+        except Exception as e:
-            _u2net_model = ort.InferenceSession(str(u2net_path))
+            print(f"Failed to load with OpenCV DNN: {e}")
-        else:
+            raise
-            # Use PyTorch
-            from app.services.u2net_model import U2NET
-            _u2net_model = U2NET(3, 1)
-            _u2net_model.load_state_dict(torch.load(str(u2net_path), map_location='cpu'))
-            _u2net_model.eval()
-
-        print("U2Net model loaded")
 
     # Preprocess image
-    img_np = np.array(img)
     original_size = img.size
-
-    # Resize to model input size
     input_size = 320
+
+    # Resize and convert to blob
     img_resized = img.resize((input_size, input_size), Image.Resampling.BILINEAR)
     img_np = np.array(img_resized).astype(np.float32)
 
-    # Normalize
+    # Normalize (ImageNet normalization)
     img_np = img_np / 255.0
     img_np = (img_np - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
-    img_np = img_np.transpose(2, 0, 1)  # HWC to CHW
+
-    img_np = np.expand_dims(img_np, 0)  # Add batch dimension
+    # Create blob (NCHW format)
+    blob = cv2.dnn.blobFromImage(
+        img_np.astype(np.float32),
+        scalefactor=1.0,
+        size=(input_size, input_size),
+        swapRB=False
+    )
 
     # Run inference
-    if hasattr(_u2net_model, 'run'):
+    _u2net_model.setInput(blob)
-        # ONNX runtime
+    outputs = _u2net_model.forward()
-        input_name = _u2net_model.get_inputs()[0].name
+
-        outputs = _u2net_model.run(None, {input_name: img_np})
+    # Get mask from first output
-        mask = outputs[0][0, 0]
+    mask = outputs[0, 0]
-    else:
-        # PyTorch
-        with torch.no_grad():
-            input_tensor = torch.from_numpy(img_np).float()
-            d1, d2, d3, d4, d5, d6, d7 = _u2net_model(input_tensor)
-            mask = d1[0, 0].numpy()
 
     # Post-process mask
     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8)

backend/requirements.txt

@@ -18,8 +18,8 @@ torch==2.1.2
 torchvision==0.16.2
 segment-anything @ git+https://github.com/facebookresearch/segment-anything.git
 
-# ONNX Runtime for U2Net background removal (works with numpy<2.0)
+# Note: Using OpenCV DNN instead of onnxruntime for U2Net
-onnxruntime==1.16.3
+# (onnxruntime has executable stack issues in some Docker environments)
 
 # NOTE: rembg (background removal) disabled due to dependency conflicts
 # rembg>=2.0.70 requires:

frontend/src/css/layout.css

@@ -867,6 +867,63 @@ canvas{
 	background-color: #ffc107;
 	color: #000;
 }
+
+/* Alertify dialog styling - fix text color issues */
+.alertify .ajs-dialog {
+	background-color: #3a3f44;
+	color: #f4f3f3;
+	border-radius: 6px;
+	box-shadow: 0 4px 20px rgba(0,0,0,0.4);
+}
+.alertify .ajs-header {
+	background-color: #4a5058;
+	color: #f4f3f3;
+	border-bottom: 1px solid #555;
+	padding: 10px 15px;
+	font-weight: bold;
+}
+.alertify .ajs-body {
+	color: #f4f3f3;
+	padding: 15px;
+}
+.alertify .ajs-body .ajs-content {
+	color: #f4f3f3;
+}
+.alertify .ajs-footer {
+	background-color: #3a3f44;
+	border-top: 1px solid #555;
+	padding: 10px 15px;
+}
+.alertify .ajs-footer .ajs-buttons .ajs-button {
+	background-color: #4a5058;
+	color: #f4f3f3;
+	border: 1px solid #666;
+	border-radius: 4px;
+	padding: 6px 16px;
+	margin: 0 4px;
+	cursor: pointer;
+}
+.alertify .ajs-footer .ajs-buttons .ajs-button:hover {
+	background-color: #5a6068;
+}
+.alertify .ajs-footer .ajs-buttons .ajs-button.ajs-ok {
+	background-color: #28a745;
+	border-color: #28a745;
+}
+.alertify .ajs-footer .ajs-buttons .ajs-button.ajs-ok:hover {
+	background-color: #218838;
+}
+.alertify .ajs-input {
+	background-color: #2a2f34;
+	color: #f4f3f3;
+	border: 1px solid #555;
+	padding: 8px;
+	border-radius: 4px;
+	width: 100%;
+}
+.alertify .ajs-input::placeholder {
+	color: #999;
+}
 .effectsPreview{
 	cursor: pointer;
 	background-color: #ddd;

frontend/src/js/tools/brush_select.js

@@ -262,7 +262,8 @@ class Brush_select_class extends Base_tools_class {
             if (isAdditive && this.currentMask) {
                 alertify.success('Added to selection! Shift+brush to add more.');
             } else {
-                alertify.success('Selection complete! Ctrl+C to copy, Ctrl+X to cut.');
+                // Offer to float the selection for immediate manipulation (Canva-like)
+                this.offerFloatSelection();
             }
 
         } catch (error) {
@@ -287,7 +288,8 @@ class Brush_select_class extends Base_tools_class {
             var maskCanvas = await this.decodeMask(result.mask);
             this.applyMaskCanvas(maskCanvas, false);
 
-            alertify.success('Selection complete! Brush over more to add.');
+            // Offer to float the selection for immediate manipulation
+            this.offerFloatSelection();
 
         } catch (error) {
             console.error('Select error:', error);
@@ -297,6 +299,34 @@ class Brush_select_class extends Base_tools_class {
         }
     }
 
+    /**
+     * Offer to float the selection to a new layer for manipulation (Canva-like workflow)
+     */
+    offerFloatSelection() {
+        var _this = this;
+
+        alertify.confirm(
+            'Selection Complete',
+            'Would you like to move/scale this selection? This will copy it to a new layer.',
+            function() {
+                // Yes - copy to layer and switch to Select tool
+                _this.copyToLayer();
+
+                // Switch to Select tool
+                setTimeout(function() {
+                    var selectTool = document.querySelector('.sidebar_left .item[data-tool="select"]');
+                    if (selectTool) {
+                        selectTool.click();
+                    }
+                }, 100);
+            },
+            function() {
+                // No - just keep the selection
+                alertify.message('Tip: Use Ctrl+C to copy or Ctrl+X to cut the selection.');
+            }
+        ).set('labels', {ok: 'Yes, Move/Scale', cancel: 'Keep Selection'});
+    }
+
     /**
      * Sample key points from brush stroke for SAM
      */