Add functionality to download the DIS ONNX model if not found locally in app.py. Updated remove_background function to call new ensure_dis_onnx_model method for improved error handling. Added requests library to requirements.txt for HTTP requests.

app.py

@@ -14,6 +14,7 @@ import json
 import os
 import json
 import argparse
+import requests
 
 from model import CRM
 from inference import generate3d
@@ -21,6 +22,7 @@ from dis_bg_remover import remove_background as dis_remove_background
 
 # Configurable ONNX model path (can be set via environment variable)
 DIS_ONNX_MODEL_PATH = os.environ.get("DIS_ONNX_MODEL_PATH", "isnet_dis.onnx")
+DIS_ONNX_MODEL_URL = "https://huggingface.co/stoned0651/isnet_dis.onnx/resolve/main/isnet_dis.onnx"
 
 pipeline = None
 
@@ -41,19 +43,32 @@ def check_input_image(input_image):
         raise gr.Error("No image uploaded!")
 
 
+def ensure_dis_onnx_model():
+    if not os.path.exists(DIS_ONNX_MODEL_PATH):
+        try:
+            print(f"Model file not found at {DIS_ONNX_MODEL_PATH}. Downloading from {DIS_ONNX_MODEL_URL}...")
+            response = requests.get(DIS_ONNX_MODEL_URL, stream=True)
+            response.raise_for_status()
+            with open(DIS_ONNX_MODEL_PATH, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+            print(f"Downloaded model to {DIS_ONNX_MODEL_PATH}")
+        except Exception as e:
+            raise gr.Error(
+                f"Failed to download DIS background remover model file: {e}\n"
+                f"Please manually download it from {DIS_ONNX_MODEL_URL} and place it in the project directory or set the DIS_ONNX_MODEL_PATH environment variable."
+            )
+
+
 def remove_background(
     image: PIL.Image.Image,
     rembg_session: Any = None,
     force: bool = False,
     **rembg_kwargs,
 ) -> PIL.Image.Image:
-    # Check if the ONNX model exists
-    if not os.path.exists(DIS_ONNX_MODEL_PATH):
-        raise gr.Error(
-            f"DIS background remover model file not found at {DIS_ONNX_MODEL_PATH}. "
-            "Please download it from https://huggingface.co/stoned0651/isnet_dis.onnx/resolve/main/isnet_dis.onnx "
-            "and place it in the project directory or set the DIS_ONNX_MODEL_PATH environment variable."
-        )
+    # Ensure the ONNX model exists (download if needed)
+    ensure_dis_onnx_model()
     extracted_img, mask = dis_remove_background(DIS_ONNX_MODEL_PATH, image)
     return extracted_img
 

requirements.txt

@@ -19,5 +19,5 @@ xatlas
 ninja
 pymeshlab
 onnxruntime
-
+requests