From ed491ab9baa019cb4b005434761d72371183b45f Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Wed, 28 Jan 2026 03:05:23 +0000
Subject: [PATCH] Fix AI Paint tool processing state and use full U2net model

- Fix brush_select isProcessing flag not resetting after first use
  (reset in on_leave() when switching tools)
- Add onnxruntime dependency for U2net background removal
- Use full U2net model (176MB) instead of lightweight for better quality

https://claude.ai/code/session_01CLedz6CanT9t46KBvng3vz
---
 backend/app/routers/tools.py          | 19 +++++++++++--------
 backend/requirements.txt              |  3 +++
 frontend/src/js/tools/brush_select.js |  1 +
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/backend/app/routers/tools.py b/backend/app/routers/tools.py
index 00d6f92..d5f664f 100644
--- a/backend/app/routers/tools.py
+++ b/backend/app/routers/tools.py
@@ -204,27 +204,30 @@ _u2net_model = None
 
 
 async def _download_u2net_model(models_dir):
-    """Auto-download U2Net model (lightweight version ~4MB)"""
+    """Auto-download full U2Net model (~176MB) for best quality background removal"""
     import urllib.request
     from pathlib import Path
 
     models_dir = Path(models_dir)
     models_dir.mkdir(parents=True, exist_ok=True)
 
-    # Download lightweight u2netp model (only 4MB)
-    url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2netp.onnx"
-    dest_path = models_dir / "u2netp.onnx"
+    # Download full U2Net model (176MB) for best quality
+    url = "https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx"
+    dest_path = models_dir / "u2net.onnx"
 
-    print(f"Downloading U2Net model from {url}...")
+    print(f"Downloading full U2Net model from {url} (~176MB)...")
+    print("This may take a few minutes...")
 
     def download_progress(count, block_size, total_size):
         if total_size > 0:
             percent = min(100, count * block_size * 100 // total_size)
-            if count % 100 == 0:
-                print(f"  Download progress: {percent}%")
+            downloaded_mb = (count * block_size) / (1024 * 1024)
+            total_mb = total_size / (1024 * 1024)
+            if count % 500 == 0:
+                print(f"  Download progress: {percent}% ({downloaded_mb:.1f}/{total_mb:.1f} MB)")
 
     urllib.request.urlretrieve(url, str(dest_path), download_progress)
-    print(f"U2Net model downloaded to {dest_path}")
+    print(f"Full U2Net model downloaded to {dest_path}")
 
     return dest_path
 
diff --git a/backend/requirements.txt b/backend/requirements.txt
index 6209ebc..413fbfb 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -18,6 +18,9 @@ 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)
+onnxruntime==1.16.3
+
 # NOTE: rembg (background removal) disabled due to dependency conflicts
 # rembg>=2.0.70 requires:
 #   - scikit-image>=0.26.0 which requires numpy>=2.0
diff --git a/frontend/src/js/tools/brush_select.js b/frontend/src/js/tools/brush_select.js
index b1522b9..00b4ad5 100644
--- a/frontend/src/js/tools/brush_select.js
+++ b/frontend/src/js/tools/brush_select.js
@@ -764,6 +764,7 @@ class Brush_select_class extends Base_tools_class {
 
     on_leave() {
         this.isDrawing = false;
+        this.isProcessing = false;  // Reset processing state when leaving tool
         this.brushPath = [];
         return [];
     }