Fix OpenMP fork issues in background removal by using multiprocessing with spawn method. Maintains high-quality alpha matting while avoiding threading conflicts.

hy3dshape/hy3dshape/rembg.py

@@ -12,35 +12,109 @@
 # fine-tuning enabling code and other elements of the foregoing made publicly available
 # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
 
+import os
+import io
+import multiprocessing
+import traceback
+import numpy as np
+
+# Set OpenMP environment variables to prevent fork() errors
+os.environ["OMP_NUM_THREADS"] = "1"  # Limit OpenMP to single thread
+os.environ["PYTHONFAULTHANDLER"] = "1"  # Better error reporting
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"  # Avoid Intel MKL errors
+os.environ["OMP_THREAD_LIMIT"] = "1"  # Limit total number of OpenMP threads
+
 from PIL import Image
 from rembg import remove, new_session
 
 
-class BackgroundRemover():
-    def __init__(self):
-        self.session = new_session()
-
-    def __call__(self, image: Image.Image):
-        try:
-            # First try with alpha matting for better quality
+def _remove_bg_in_process(image_bytes, use_alpha_matting=True):
+    """Function to run in a separate process to avoid OpenMP issues"""
+    try:
+        # Create a new session in this process
+        session = new_session()
+        
+        # Convert bytes back to PIL Image
+        image = Image.open(io.BytesIO(image_bytes))
+        
+        if use_alpha_matting:
+            # High quality mode with alpha matting
             output = remove(
                 image,
-                session=self.session,
+                session=session,
                 bgcolor=[255, 255, 255, 0],
                 alpha_matting=True,
-                alpha_matting_foreground_threshold=220,   # Lower threshold for foreground
-                alpha_matting_background_threshold=5,     # Lower background threshold
-                alpha_matting_erode_structure_size=5,     # Smaller erode size
-                alpha_matting_base_size=512,               # Smaller scale for matting
+                alpha_matting_foreground_threshold=220,
+                alpha_matting_background_threshold=5,
+                alpha_matting_erode_structure_size=5,
+                alpha_matting_base_size=512,
             )
-            return output
-        except ZeroDivisionError:
-            # If alpha matting fails, fall back to simpler method
-            print("Alpha matting failed, falling back to standard background removal")
+        else:
+            # Fallback mode without alpha matting
             output = remove(
                 image,
-                session=self.session,
+                session=session,
                 bgcolor=[255, 255, 255, 0],
-                alpha_matting=False,  # Disable alpha matting
+                alpha_matting=False,
             )
-            return output
+        
+        # Convert back to bytes
+        output_bytes = io.BytesIO()
+        output.save(output_bytes, format='PNG')
+        return output_bytes.getvalue()
+    
+    except Exception as e:
+        print(f"Background removal process error: {e}")
+        traceback.print_exc()
+        return None
+
+
+class BackgroundRemover():
+    def __init__(self):
+        # Set multiprocessing start method to 'spawn' to avoid OpenMP fork issues
+        # Only set if not already set
+        if multiprocessing.get_start_method(allow_none=True) != 'spawn':
+            try:
+                multiprocessing.set_start_method('spawn', force=True)
+            except RuntimeError:
+                # If already set and not 'spawn', we'll have to work with what we have
+                print("Warning: Could not set multiprocessing start method to 'spawn'")
+        
+        # We don't initialize the session here anymore as it will be created in the subprocess
+        self.session = None
+
+    def __call__(self, image: Image.Image):
+        try:
+            # Convert image to bytes for passing to subprocess
+            img_bytes = io.BytesIO()
+            image.save(img_bytes, format='PNG')
+            img_bytes = img_bytes.getvalue()
+            
+            # Create a process context
+            ctx = multiprocessing.get_context('spawn')
+            
+            # Try with alpha matting first (high quality)
+            with ctx.Pool(processes=1) as pool:
+                result = pool.apply(_remove_bg_in_process, args=(img_bytes, True))
+            
+            if result is not None:
+                return Image.open(io.BytesIO(result))
+            
+            # If alpha matting fails, try without it
+            print("Alpha matting failed, falling back to standard background removal")
+            with ctx.Pool(processes=1) as pool:
+                result = pool.apply(_remove_bg_in_process, args=(img_bytes, False))
+            
+            if result is not None:
+                return Image.open(io.BytesIO(result))
+            
+            # If all else fails, return the original image
+            print("Background removal completely failed, returning original image")
+            return image
+            
+        except Exception as e:
+            print(f"Background removal error in main process: {e}")
+            traceback.print_exc()
+            # If all else fails, just return the original image
+            # This ensures the pipeline doesn't break completely
+            return image