Dynamic worker cap based on GPU VRAM instead of hard limit, bump to 0.4.1

Calculates max safe workers from actual GPU memory: CUDA context (~400MB)
+ model weights + JIT trace + batch inference per worker. Falls back to
CPU core count / 2 on non-CUDA systems. Also prints VRAM size at startup.

Author: nadermx

backgroundremover/__init__.py

@@ -4,6 +4,6 @@
 A library to remove background from videos and images
 """
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 __author__ = 'Johnathan Nader'
 __credits__ = 'BackgroundRemoverAI.com'

backgroundremover/bg.py

@@ -27,7 +27,9 @@
 try:
     if torch.cuda.is_available():
         DEVICE = torch.device('cuda:0')
-        print(f"Device: CUDA ({torch.cuda.get_device_name(0)})")
+        _gpu_name = torch.cuda.get_device_name(0)
+        _gpu_mem = torch.cuda.get_device_properties(0).total_memory
+        print(f"Device: CUDA ({_gpu_name}, {_gpu_mem // (1024**2)}MB VRAM)")
     elif torch.backends.mps.is_available():
         DEVICE = torch.device('mps')
         print("Device: MPS (Apple Silicon GPU)")
@@ -38,6 +40,44 @@
     print(f"Device: CPU (Setting CUDA or MPS failed: {e})")
     DEVICE = torch.device('cpu')
 
+
+def max_workers(model_name="u2net", gpu_batchsize=2):
+    """Estimate max safe worker processes based on available GPU/system memory.
+
+    Each worker spawns a separate process that loads its own copy of the model
+    plus a CUDA context. This estimates how many can fit in VRAM.
+    """
+    if torch.cuda.is_available():
+        try:
+            total_mem = torch.cuda.get_device_properties(0).total_memory
+        except Exception:
+            return 1
+
+        # Per-worker VRAM estimate:
+        #   CUDA context per process:  ~400MB
+        #   Model weights (float32):   ~175MB (u2net/human_seg), ~5MB (u2netp)
+        #   JIT traced copy:           same as model weights
+        #   Batch inference tensors:   ~30MB per frame in batch
+        if model_name == "u2netp":
+            model_bytes = 5 * 1024 * 1024
+        else:
+            model_bytes = 175 * 1024 * 1024
+
+        per_worker = (
+            400 * 1024 * 1024       # CUDA context overhead
+            + model_bytes * 2       # model + JIT trace
+            + gpu_batchsize * 30 * 1024 * 1024  # inference tensors
+        )
+
+        # Reserve 512MB for OS/driver/display
+        usable = total_mem - 512 * 1024 * 1024
+        calculated = max(1, int(usable // per_worker))
+        return calculated
+
+    # CPU/MPS: limit by CPU cores (inference is compute-bound)
+    cpu_count = os.cpu_count() or 2
+    return max(1, cpu_count // 2)
+
 class Net(torch.nn.Module):
     def __init__(self, model_name):
         super(Net, self).__init__()

backgroundremover/cmd/cli.py

@@ -2,7 +2,7 @@
 import os
 from distutils.util import strtobool
 from .. import utilities
-from ..bg import remove
+from ..bg import remove, max_workers
 
 
 def main():
@@ -248,11 +248,11 @@ def main():
         print("Example: backgroundremover -i video.mp4 -tgwb -bi background.png -o output.gif")
         exit(1)
 
-    # Cap worker count to prevent hangs and resource exhaustion (see issue #181)
-    if args.workernodes > 4:
-        print(f"Warning: Requested {args.workernodes} workers, capping at 4. Higher values cause hangs and GPU memory exhaustion.")
-        print("Use -wn 1 through -wn 4 for best results.")
-        args.workernodes = 4
+    # Dynamically cap worker count based on available GPU memory (see issue #181)
+    safe_max = max_workers(model_name=args.model, gpu_batchsize=args.gpubatchsize)
+    if args.workernodes > safe_max:
+        print(f"Warning: Requested {args.workernodes} workers, capping at {safe_max} based on available memory.")
+        args.workernodes = safe_max
 
     # Parse background color if provided
     background_color = None

setup.py

@@ -11,7 +11,7 @@
 
 setup(
     name="backgroundremover",
-    version="0.4.0",
+    version="0.4.1",
     description="Background remover from image and video using AI",
     long_description=long_description,
     long_description_content_type="text/markdown",