Fixing work with onnxruntime for rembg

install_init.py

@@ -172,13 +172,16 @@ def install(pkg_spec: str):
 def uninstall(pkg: str):
     subprocess.check_call([sys.executable, '-m', 'pip', 'uninstall', '-y', pkg])
 
-def install_onnxruntime():
-    if has_nvidia_cuda():
+def install_onnxruntime_and_rembg():
+    gpu = has_nvidia_cuda()
+    print(f"Found NVIDIA GPU: {gpu}")
+    if gpu:
         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).")
+            print(f"LP >>> Close ComfyUI and run the script at .\\ComfyUI\\custom_nodes\\ComfyUI-LevelPixel\\scripts\\remove_onnxruntime.bat")
         if not is_installed('onnxruntime-gpu'):
             install("onnxruntime-gpu>=1.22")
     else:
@@ -187,6 +190,11 @@ def install_onnxruntime():
         if not is_installed('onnxruntime'):
             install("onnxruntime>=1.22")
 
+    rembg_pkg = 'rembg[gpu]'      if gpu else 'rembg[cpu]'
+
+    if not is_installed('rembg', '2.0.66'):
+        install(f'{rembg_pkg}>={'2.0.66'}')
+
 def init(check_imports=None):
     log("Init")
 
@@ -200,5 +208,5 @@ def init(check_imports=None):
                 return False
 
     install_js()
-    install_onnxruntime()
+    install_onnxruntime_and_rembg()
     return True

nodes/image/image_utils_LP.py

@@ -4,9 +4,7 @@ import torch
 from PIL import Image, ImageOps
 import matplotlib.pyplot as plt
 import comfy.utils
-import subprocess
-import sys
-from importlib import metadata
+
     
 color_mapping = {
     "white": (255, 255, 255),
@@ -146,32 +144,6 @@ 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):
         pass
@@ -190,32 +162,7 @@ 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,)

scripts/remove_onnxruntime.bat

@@ -9,6 +9,10 @@ if /I "%user_input%"=="Y" (
 
     cd ..\..\..\..
     .\python_embeded\python.exe -m pip uninstall -y onnxruntime
+    .\python_embeded\python.exe -m pip uninstall -y onnxruntime-gpu
+
+    rmdir /s /q "python_embeded\Lib\site-packages\onnxruntime"
+    for /d %%D in ("python_embeded\Lib\site-packages\onnxruntime-*") do rmdir /s /q "%%~fD"
 ) else (
     echo Please close the ComfyUI console first and run this script again.
 )