Update app.py

app.py

@@ -6,8 +6,11 @@ from PIL import Image
 import requests
 import logging
 
-# Import the withoutBG library
-from withoutbg import WithoutBG
+# Import the withoutBG library (correct way from Qiita article)
+from withoutbg.core import WithoutBGOpenSource
+from huggingface_hub import hf_hub_download
+from pathlib import Path
+import shutil
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -16,10 +19,35 @@ logger = logging.getLogger(__name__)
 app = Flask(__name__)
 CORS(app)
 
+# Model directory
+MODEL_DIR = Path("/app/models")
+MODEL_DIR.mkdir(parents=True, exist_ok=True)
+
+def _ensure_model_file(filename: str) -> Path:
+    """Download model file from HuggingFace if not exists"""
+    target = MODEL_DIR / filename
+    if target.exists():
+        return target
+    logger.info(f"📥 Downloading model file: {filename}")
+    downloaded = Path(hf_hub_download(repo_id="withoutbg/focus", filename=filename))
+    shutil.copy2(downloaded, target)
+    logger.info(f"✅ Model file downloaded: {filename}")
+    return target
+
+def _create_model() -> WithoutBGOpenSource:
+    """Create WithoutBG model instance"""
+    logger.info("🚀 Creating WithoutBG model...")
+    return WithoutBGOpenSource(
+        depth_model_path=_ensure_model_file("depth_anything_v2_vits_slim.onnx"),
+        isnet_model_path=_ensure_model_file("isnet.onnx"),
+        matting_model_path=_ensure_model_file("focus_matting_1.0.0.onnx"),
+        refiner_model_path=_ensure_model_file("focus_refiner_1.0.0.onnx"),
+    )
+
 # Initialize the model once at startup
 try:
     logger.info("🚀 Loading withoutBG model...")
-    model = WithoutBG.opensource()
+    model = _create_model()
     logger.info("✅ Model loaded successfully!")
 except Exception as e:
     logger.error(f"❌ Failed to load model: {e}")
@@ -72,12 +100,12 @@ def remove_background():
         
         # Open image
         img = Image.open(image_data)
-        logger.info(f"🖼️ Image loaded: {img.size}")
+        logger.info(f"🖼️ Image loaded: {img.size}, mode: {img.mode}")
         
-        # Remove background using withoutBG
-        logger.info("🔄 Removing background...")
+        # Remove background using withoutBG (Qiita article method)
+        logger.info("🔄 Removing background with WithoutBGOpenSource...")
         result = model.remove_background(img)
-        logger.info(f"✅ Background removed successfully! Mode: {result.mode}, Size: {result.size}")
+        logger.info(f"✅ Background removed! Result mode: {result.mode}, Size: {result.size}")
         
         # Convert to RGBA first if not already
         if result.mode != 'RGBA':
@@ -110,6 +138,8 @@ def remove_background():
         
     except Exception as e:
         logger.error(f"❌ Error: {e}")
+        import traceback
+        traceback.print_exc()
         return jsonify({
             'success': False,
             'error': str(e)