Changed logic for using onnxruntime

install_init.py

@@ -3,6 +3,9 @@ import json
 import shutil
 import inspect
 from server import PromptServer
+import subprocess
+import sys
+from importlib import metadata
 
 config = None
 
@@ -143,6 +146,47 @@ def install_js():
     log("Copying JS files")
     shutil.copytree(src_dir, dst_dir, dirs_exist_ok=True)
 
+def has_nvidia_cuda():
+    try:
+        output = subprocess.check_output(
+            ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
+            stderr=subprocess.DEVNULL,
+            encoding='utf-8'
+        ).strip()
+        return bool(output)
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return False
+
+def is_installed(pkg_name: str, min_version: str = '') -> bool:
+    try:
+        ver = metadata.version(pkg_name)
+        if min_version:
+            return tuple(map(int, ver.split('.'))) >= tuple(map(int, min_version.split('.')))
+        return True
+    except metadata.PackageNotFoundError:
+        return False
+
+def install(pkg_spec: str):
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', pkg_spec])
+
+def uninstall(pkg: str):
+    subprocess.check_call([sys.executable, '-m', 'pip', 'uninstall', '-y', pkg])
+
+def install_onnxruntime():
+    if has_nvidia_cuda():
+        if is_installed('onnxruntime'):
+            #uninstall("onnxruntime")
+            print(f"LP >>> Your python has the 'onnxruntime' library installed, although your computer supports 'onnxruntime-gpu'.")
+            print(f"LP >>> Solution: If other node packages do not use the 'onnxruntime' library, then remove the 'onnxruntime' library for your python.")
+            print(f"LP >>> Otherwise the Image Remove Background node will not work effectively (the node will use the CPU instead of the GPU).")
+        if not is_installed('onnxruntime-gpu'):
+            install("onnxruntime-gpu>=1.22")
+    else:
+        if is_installed('onnxruntime-gpu'):
+            uninstall("onnxruntime-gpu")
+        if not is_installed('onnxruntime'):
+            install("onnxruntime>=1.22")
+
 def init(check_imports=None):
     log("Init")
 
@@ -156,4 +200,5 @@ def init(check_imports=None):
                 return False
 
     install_js()
+    install_onnxruntime()
     return True

nodes/image/image_utils_LP.py

@@ -3,9 +3,10 @@ import io
 import torch
 from PIL import Image, ImageOps
 import matplotlib.pyplot as plt
-from rembg import new_session, remove
-import comfy.sd
 import comfy.utils
+import subprocess
+import sys
+from importlib import metadata
     
 color_mapping = {
     "white": (255, 255, 255),
@@ -145,6 +146,31 @@ def tensor2pil(image):
 def list_model():
     return model_list
 
+def has_nvidia_cuda():
+    try:
+        output = subprocess.check_output(
+            ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
+            stderr=subprocess.DEVNULL,
+            encoding='utf-8'
+        ).strip()
+        return bool(output)
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return False
+    
+def is_installed(pkg_name: str, min_version: str = '') -> bool:
+    try:
+        ver = metadata.version(pkg_name)
+        if min_version:
+            return tuple(map(int, ver.split('.'))) >= tuple(map(int, min_version.split('.')))
+        return True
+    except metadata.PackageNotFoundError:
+        return False
+
+def install(pkg_spec: str):
+    subprocess.check_call([sys.executable, '-m', 'pip', 'install', pkg_spec])
+
+def uninstall(pkg: str):
+    subprocess.check_call([sys.executable, '-m', 'pip', 'uninstall', '-y', pkg])
 
 class ImageRemoveBackground:
     def __init__(self):
@@ -164,6 +190,32 @@ class ImageRemoveBackground:
     CATEGORY = "LevelPixel/Image"
 
     def remove_background(self, image, model_name):
+        gpu = has_nvidia_cuda()
+        print(f"Found NVIDIA GPU: {gpu}")
+
+        rembg_pkg = 'rembg[gpu]'      if gpu else 'rembg[cpu]'
+        rembg_ver = '2.0.66'
+
+        if gpu:
+            if is_installed('onnxruntime'):
+                #uninstall("onnxruntime")
+                print(f"LP >>> !!! WARNING about Image Remove Background:")
+                print(f"LP >>> Your python has the 'onnxruntime' library installed, although your computer supports 'onnxruntime-gpu'.")
+                print(f"LP >>> Solution: If other node packages do not use the 'onnxruntime' library, then remove the 'onnxruntime' library for your python.")
+                print(f"LP >>> Otherwise the Image Remove Background node will not work effectively (the node will use the CPU instead of the GPU).")
+            if not is_installed('onnxruntime-gpu'):
+                install("onnxruntime-gpu>=1.22")
+        else:
+            if is_installed('onnxruntime-gpu'):
+                uninstall("onnxruntime-gpu")
+            if not is_installed('onnxruntime'):
+                install("onnxruntime>=1.22")
+
+        if not is_installed('rembg', rembg_ver):
+            install(f'{rembg_pkg}>={rembg_ver}')
+
+        from rembg import new_session, remove
+
         session = new_session(model_name)
         image = pil2tensor(remove(tensor2pil(image), session = session))
         return (image,)

requirements.txt

@@ -1,14 +1,7 @@
 torch>=2.0.1
-torchvision>=0.15.2
-pillow>=9.4.0
 numpy
 matplotlib
-scikit-build-core>=0.10.7
 rembg>=2.0.59
-onnxruntime-gpu>=1.20.0
-onnxruntime>=1.20.0
-googletrans
 langdetect
-httpcore
 deep_translator
 wordninja

scripts/remove_onnxruntime.bat

@@ -0,0 +1,17 @@
+@echo off
+setlocal
+
+echo Please close the ComfyUI console before continuing.
+set /p user_input=Did you close the ComfyUI console? (Y/n): 
+
+if /I "%user_input%"=="Y" (
+    echo Uninstalling onnxruntime...
+
+    cd ..\..\..\..
+    .\python_embeded\python.exe -m pip uninstall -y onnxruntime
+) else (
+    echo Please close the ComfyUI console first and run this script again.
+)
+
+endlocal
+pause