From 6999dd366e7c04dbbc9aebffb1b55b2571a589df Mon Sep 17 00:00:00 2001 From: Level Pixel Dev Date: Wed, 28 May 2025 06:18:56 +0600 Subject: [PATCH] Changed logic for using onnxruntime --- install_init.py | 45 +++++++++++++++++++++++++++ nodes/image/image_utils_LP.py | 56 ++++++++++++++++++++++++++++++++-- requirements.txt | 7 ----- scripts/remove_onnxruntime.bat | 17 +++++++++++ 4 files changed, 116 insertions(+), 9 deletions(-) create mode 100644 scripts/remove_onnxruntime.bat diff --git a/install_init.py b/install_init.py index 45a3023..1a054b0 100644 --- a/install_init.py +++ b/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 diff --git a/nodes/image/image_utils_LP.py b/nodes/image/image_utils_LP.py index 0b5f21a..f380ea2 100644 --- a/nodes/image/image_utils_LP.py +++ b/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,) diff --git a/requirements.txt b/requirements.txt index 6a28c24..a60a96f 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file diff --git a/scripts/remove_onnxruntime.bat b/scripts/remove_onnxruntime.bat new file mode 100644 index 0000000..17de093 --- /dev/null +++ b/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 \ No newline at end of file