Add configurable ONNX model path in app.py for background removal. Implement error handling for missing model file and ensure valid image processing. Updated remove_background function to utilize environment variable for model path.

app.py

@@ -19,6 +19,9 @@ from model import CRM
 from inference import generate3d
 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")
+
 pipeline = None
 
 
@@ -44,8 +47,14 @@ def remove_background(
     force: bool = False,
     **rembg_kwargs,
 ) -> PIL.Image.Image:
-    model_path = "isnet_dis.onnx"  # Update with the correct path to the ONNX model
-    extracted_img, mask = dis_remove_background(model_path, 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."
+        )
+    extracted_img, mask = dis_remove_background(DIS_ONNX_MODEL_PATH, image)
     return extracted_img
 
 def do_resize_content(original_image: Image, scale_rate):
@@ -79,6 +88,8 @@ def preprocess_image(image, background_choice, foreground_ratio, backgroud_color
         image = Image.alpha_composite(background, image)
     else:
         image = remove_background(image, force=True)
+        if image is None:
+            raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.")
     image = do_resize_content(image, foreground_ratio)
     image = expand_to_square(image)
     image = add_background(image, backgroud_color)