diff --git a/install_init.py b/install_init.py index 1a054b0..7a86f86 100644 --- a/install_init.py +++ b/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 diff --git a/nodes/image/image_utils_LP.py b/nodes/image/image_utils_LP.py index f380ea2..b911b7b 100644 --- a/nodes/image/image_utils_LP.py +++ b/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,) diff --git a/scripts/remove_onnxruntime.bat b/scripts/remove_onnxruntime.bat index 17de093..d19dc7a 100644 --- a/scripts/remove_onnxruntime.bat +++ b/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. )