Upload Inspyrenet_Rembg2.py

Inspyrenet_Rembg2.py

@@ -1,6 +1,9 @@
 from PIL import Image
 import os
 import urllib.request
+import gc
+import threading
+from typing import Dict, Tuple
 
 import torch
 import numpy as np
@@ -22,7 +25,6 @@ def _ensure_ckpt_base():
         if os.path.isfile(CKPT_PATH) and os.path.getsize(CKPT_PATH) > 0:
             return
     except Exception:
-        # If getsize fails for any reason, fall through to download attempt.
         pass
 
     os.makedirs(os.path.dirname(CKPT_PATH), exist_ok=True)
@@ -57,7 +59,6 @@ def _ensure_ckpt_base():
         os.replace(tmp_path, CKPT_PATH)
 
     finally:
-        # Clean up partial download if something went wrong
         if os.path.isfile(tmp_path):
             try:
                 os.remove(tmp_path)
@@ -68,7 +69,6 @@ def _ensure_ckpt_base():
 # Tensor to PIL
 def tensor2pil(image: torch.Tensor) -> Image.Image:
     arr = image.detach().cpu().numpy()
-    # Handle accidental singleton batch dim
     if arr.ndim == 4 and arr.shape[0] == 1:
         arr = arr[0]
     arr = np.clip(255.0 * arr, 0, 255).astype(np.uint8)
@@ -82,11 +82,11 @@ def pil2tensor(image: Image.Image) -> torch.Tensor:
 
 def _rgba_to_rgb_on_white(pil_img: Image.Image) -> Image.Image:
     """
-    5) If input is RGBA:
-         - alpha composite over WHITE background
-         - convert to RGB (drop alpha)
-       If input is RGB:
-         - carry on
+    If input is RGBA:
+      - alpha composite over WHITE background
+      - convert to RGB (drop alpha)
+    If input is RGB:
+      - carry on
     """
     if pil_img.mode == "RGBA":
         bg = Image.new("RGBA", pil_img.size, (255, 255, 255, 255))
@@ -99,9 +99,134 @@ def _rgba_to_rgb_on_white(pil_img: Image.Image) -> Image.Image:
     return pil_img
 
 
+# -----------------------------------------------------------------------------
+# Process-wide singleton Remover + OOM guard
+# -----------------------------------------------------------------------------
+
+# One cached Remover per (jit_flag,) for the entire Python process.
+_REMOVER_CACHE: Dict[Tuple[bool], Remover] = {}
+# Lock per remover to avoid concurrent .process() calls (prevents VRAM spikes).
+_REMOVER_RUN_LOCKS: Dict[Tuple[bool], threading.Lock] = {}
+# Protects cache/lock creation.
+_CACHE_LOCK = threading.Lock()
+
+
+def _is_oom_error(e: BaseException) -> bool:
+    # torch.cuda.OutOfMemoryError only exists on CUDA builds; guard it.
+    oom_cls = getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None)
+    if oom_cls is not None and isinstance(e, oom_cls):
+        return True
+
+    msg = str(e).lower()
+    # Covers common RuntimeError("CUDA out of memory") patterns too.
+    return ("out of memory" in msg) and ("cuda" in msg or "cublas" in msg or "hip" in msg)
+
+
+def _cuda_soft_cleanup() -> None:
+    """
+    Best-effort cleanup that should NOT evict "important" VRAM like model weights.
+
+    What it does:
+      - gc.collect(): drop dead Python objects sooner
+      - torch.cuda.empty_cache(): releases *unused* cached blocks back to the driver
+      - torch.cuda.ipc_collect(): helps in some multi-process cases
+
+    What it does NOT do:
+      - It does not unload models still referenced
+      - It does not free tensors that still have live references
+    """
+    try:
+        gc.collect()
+    except Exception:
+        pass
+
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.synchronize()
+        except Exception:
+            pass
+        try:
+            torch.cuda.empty_cache()
+        except Exception:
+            pass
+        try:
+            torch.cuda.ipc_collect()
+        except Exception:
+            pass
+
+
+def _get_remover(jit: bool = False) -> tuple[Remover, threading.Lock]:
+    """
+    Returns a cached Remover instance + a lock to serialize .process() calls.
+
+    - Only one Remover is constructed per jit setting for the entire process.
+    - If construction OOMs, we soft-clean and re-raise (and do NOT cache a broken instance).
+    """
+    key = (jit,)
+    with _CACHE_LOCK:
+        inst = _REMOVER_CACHE.get(key)
+        if inst is None:
+            _ensure_ckpt_base()
+            try:
+                inst = Remover(jit=jit) if jit else Remover()
+            except BaseException as e:
+                if _is_oom_error(e):
+                    _cuda_soft_cleanup()
+                raise
+            _REMOVER_CACHE[key] = inst
+
+        run_lock = _REMOVER_RUN_LOCKS.get(key)
+        if run_lock is None:
+            run_lock = threading.Lock()
+            _REMOVER_RUN_LOCKS[key] = run_lock
+
+    return inst, run_lock
+
+
+def _remover_process_safe(
+    remover: Remover,
+    run_lock: threading.Lock,
+    pil_img: Image.Image,
+    *,
+    out_type: str,
+    retries: int = 1,
+) -> Image.Image:
+    """
+    Runs remover.process() under:
+      - a lock (avoid concurrent VRAM spikes),
+      - torch.inference_mode() (less VRAM),
+      - OOM catch -> soft cleanup -> retry.
+
+    If it still OOMs after retries, it raises.
+    """
+    last_err: BaseException | None = None
+
+    for attempt in range(retries + 1):
+        try:
+            with run_lock:
+                with torch.inference_mode():
+                    return remover.process(pil_img, type=out_type)
+        except BaseException as e:
+            last_err = e
+            if _is_oom_error(e):
+                _cuda_soft_cleanup()
+                if attempt < retries:
+                    continue
+            raise
+
+    # Shouldn't hit, but keeps type-checkers happy.
+    raise last_err if last_err is not None else RuntimeError("Unknown failure in _remover_process_safe()")
+
+
+# -----------------------------------------------------------------------------
+# Nodes
+# -----------------------------------------------------------------------------
+
 class InspyrenetRembg2:
     """
-    Original node kept (unchanged behavior/output), except it now ensures ckpt exists.
+    Original node behavior/output kept, but:
+      - Remover is now a process-wide singleton (per jit flag)
+      - OOM is caught -> soft CUDA cleanup -> retry once -> then raise
     """
     def __init__(self):
         pass
@@ -120,18 +245,29 @@ class InspyrenetRembg2:
     CATEGORY = "image"
 
     def remove_background(self, image, torchscript_jit):
-        _ensure_ckpt_base()
-
-        if (torchscript_jit == "default"):
-            remover = Remover()
-        else:
-            remover = Remover(jit=True)
+        jit = (torchscript_jit != "default")
+        remover, run_lock = _get_remover(jit=jit)
 
         img_list = []
-        for img in tqdm(image, "Inspyrenet Rembg"):
-            mid = remover.process(tensor2pil(img), type='rgba')
-            out = pil2tensor(mid)
-            img_list.append(out)
+        try:
+            for img in tqdm(image, "Inspyrenet Rembg"):
+                pil_in = tensor2pil(img)
+                mid = _remover_process_safe(remover, run_lock, pil_in, out_type="rgba", retries=1)
+                out = pil2tensor(mid)
+                img_list.append(out)
+
+                # Help Python drop refs earlier (mostly relevant if exceptions occur).
+                del pil_in, mid, out
+
+        except BaseException as e:
+            if _is_oom_error(e):
+                # Ensure cache cleanup already happened; do another best-effort pass.
+                _cuda_soft_cleanup()
+                raise RuntimeError(
+                    "InspyrenetRembg2: CUDA out of memory during background removal. "
+                    "Freed PyTorch CUDA cache (unused blocks) and retried once; still failed."
+                ) from e
+            raise
 
         img_stack = torch.cat(img_list, dim=0)
         mask = img_stack[:, :, :, 3]
@@ -140,12 +276,9 @@ class InspyrenetRembg2:
 
 class InspyrenetRembg3:
     """
-    New node per requested changes:
-      - ensures ckpt_base.pth exists (downloads if missing)
-      - torchscript_jit hardcoded to "default" (no input, no JIT)
-      - NO MASK output (IMAGE only)
-      - if input is RGBA: composite over white, convert to RGB, then run remover
-      - output remains RGBA (type='rgba')
+    New node per your existing changes, plus:
+      - singleton Remover reuse
+      - OOM catch -> soft cleanup -> retry once -> then raise
     """
     def __init__(self):
         pass
@@ -163,24 +296,28 @@ class InspyrenetRembg3:
     CATEGORY = "image"
 
     def remove_background(self, image):
-        _ensure_ckpt_base()
-
-        # 3) hardcode torchscript_jit == "default"
-        remover = Remover()
+        remover, run_lock = _get_remover(jit=False)
 
         img_list = []
-        for img in tqdm(image, "Inspyrenet Rembg3"):
-            pil_in = tensor2pil(img)
-
-            # 5) normalize input to RGB for the model:
-            #    - if RGBA -> alpha composite on white -> RGB
-            #    - if RGB  -> keep
-            pil_rgb = _rgba_to_rgb_on_white(pil_in)
-
-            # do functionality as usual, output RGBA
-            mid = remover.process(pil_rgb, type="rgba")
-            out = pil2tensor(mid)
-            img_list.append(out)
+        try:
+            for img in tqdm(image, "Inspyrenet Rembg3"):
+                pil_in = tensor2pil(img)
+                pil_rgb = _rgba_to_rgb_on_white(pil_in)
+
+                mid = _remover_process_safe(remover, run_lock, pil_rgb, out_type="rgba", retries=1)
+                out = pil2tensor(mid)
+                img_list.append(out)
+
+                del pil_in, pil_rgb, mid, out
+
+        except BaseException as e:
+            if _is_oom_error(e):
+                _cuda_soft_cleanup()
+                raise RuntimeError(
+                    "InspyrenetRembg3: CUDA out of memory during background removal. "
+                    "Freed PyTorch CUDA cache (unused blocks) and retried once; still failed."
+                ) from e
+            raise
 
         img_stack = torch.cat(img_list, dim=0)
         return (img_stack,)