Add DIS background remover integration and model download functionality

- Introduced a new function to ensure the DIS ONNX model is downloaded if not present.
- Updated the background removal process to utilize the DIS background remover instead of the previous method.
- Implemented temporary file handling for image processing to enhance performance and maintainability.
- Improved error handling for background removal failures, ensuring user feedback on model availability.

app.py

@@ -16,12 +16,18 @@ import json
 import os
 import json
 import argparse
+import requests
+import tempfile
 
 from model import CRM
 from inference import generate3d
+from dis_bg_remover import remove_background as dis_remove_background
+
+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
-rembg_session = rembg.new_session()
 
 
 def expand_to_square(image, bg_color=(0, 0, 0, 0)):
@@ -39,24 +45,37 @@ def check_input_image(input_image):
     if input_image is None:
         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:
-    do_remove = True
-    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
-        # explain why current do not rm bg
-        print("alhpa channl not enpty, skip remove background, using alpha channel as mask")
-        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
-        image = Image.alpha_composite(background, image)
-        do_remove = False
-    do_remove = do_remove or force
-    if do_remove:
-        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
-    return image
+    ensure_dis_onnx_model()
+    # Create a temporary file to save the image
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=True) as temp:
+        # Save the PIL image to the temporary file
+        image.save(temp.name)
+        # Call the background removal function with the file path
+        extracted_img, _ = dis_remove_background(DIS_ONNX_MODEL_PATH, temp.name)
+    # The function should return the PIL image directly
+    return extracted_img
 
 def do_resize_content(original_image: Image, scale_rate):
     # resize image content wile retain the original image size
@@ -88,11 +107,9 @@ def preprocess_image(image, background_choice, foreground_ratio, backgroud_color
         background = Image.new("RGBA", image.size, (0, 0, 0, 0))
         image = Image.alpha_composite(background, image)
     else:
-        image = remove_background(image, rembg_session, force=True)
-    image = do_resize_content(image, foreground_ratio)
-    image = expand_to_square(image)
-    image = add_background(image, backgroud_color)
-    return image.convert("RGB")
+        image = remove_background(image)
+        if image is None:
+            raise gr.Error("Background removal failed. Please check the input image and ensure the model file exists and is valid.")
 
 @spaces.GPU
 def gen_image(input_image, seed, scale, step):