From ca3d8c16a8b9893e913a18fa9b3c2cc36159b811 Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" Date: Fri, 12 Apr 2024 22:22:28 +0900 Subject: [PATCH] feat: custom mapping for pip installation https://github.com/ltdrdata/ComfyUI-Manager/discussions/554 --- __init__.py | 16 +++++++++++++--- pip_overrides.json.template | 21 +++++++++++++++++++++ prestartup_script.py | 24 ++++++++++++++++++++---- 3 files changed, 54 insertions(+), 7 deletions(-) create mode 100644 pip_overrides.json.template diff --git a/__init__.py b/__init__.py index 1ae7235..adf4357 100644 --- a/__init__.py +++ b/__init__.py @@ -30,7 +30,7 @@ except: print(f"[WARN] ComfyUI-Manager: Your ComfyUI version is outdated. Please update to the latest version.") -version = [2, 16, 1] +version = [2, 17] version_str = f"V{version[0]}.{version[1]}" + (f'.{version[2]}' if len(version) > 2 else '') print(f"### Loading: ComfyUI-Manager ({version_str})") @@ -42,6 +42,15 @@ cache_lock = threading.Lock() pip_map = None +def remap_pip_package(pkg): + if pkg in cm_global.pip_overrides: + res = cm_global.pip_overrides[pkg] + print(f"[ComfyUI-Manager] '{pkg}' is remapped to '{res}'") + return res + else: + return pkg + + def get_installed_packages(): global pip_map @@ -1377,7 +1386,7 @@ def execute_install_script(url, repo_path, lazy_mode=False): print("Install: pip packages") with open(requirements_path, "r") as requirements_file: for line in requirements_file: - package_name = line.strip() + package_name = remap_pip_package(line.strip()) if package_name: install_cmd = [sys.executable, "-m", "pip", "install", package_name] if package_name.strip() != "": @@ -1641,7 +1650,8 @@ async def install_custom_node(request): if 'pip' in json_data: for pname in json_data['pip']: - install_cmd = [sys.executable, "-m", "pip", "install", pname] + pkg = remap_pip_package(pname) + install_cmd = [sys.executable, "-m", "pip", "install", pkg] try_install_script(json_data['files'][0], ".", install_cmd) clear_pip_cache() diff --git a/pip_overrides.json.template b/pip_overrides.json.template new file mode 100644 index 0000000..a6b1f26 --- /dev/null +++ b/pip_overrides.json.template @@ -0,0 +1,21 @@ +{ + "imageio-ffmpeg": "imageio", + "imageio[ffmpeg]": "imageio", + "imageio_ffmpeg": "imageio", + "diffusers~=0.21.4": "diffusers", + "huggingface_hub": "huggingface-hub", + "numpy<1.24>=1.18": "numpy", + "numpy>=1.18.5, <1.25.0": "numpy", + "opencv-contrib-python": "opencv-contrib-python-headless", + "opencv-python": "opencv-contrib-python-headless", + "opencv-python-headless": "opencv-contrib-python-headless", + "opencv-python-headless[ffmpeg]<=4.7.0.72": "opencv-contrib-python-headless", + "opencv-python>=4.7.0.72": "opencv-contrib-python-headless", + "pandas<=1.5.1": "pandas", + "scikit-image==0.20.0": "scikit-image", + "scipy>=1.11.4": "scipy", + "segment_anything": "segment-anything", + "timm==0.6.5": "timm", + "timm>=0.4.12": "timm", + "transformers==4.26.1": "transformers" +} \ No newline at end of file diff --git a/prestartup_script.py b/prestartup_script.py index 0134aec..f77465a 100644 --- a/prestartup_script.py +++ b/prestartup_script.py @@ -7,6 +7,7 @@ import threading import re import locale import platform +import json glob_path = os.path.join(os.path.dirname(__file__), "glob") @@ -69,6 +70,23 @@ custom_nodes_path = os.path.abspath(os.path.join(comfyui_manager_path, "..")) startup_script_path = os.path.join(comfyui_manager_path, "startup-scripts") restore_snapshot_path = os.path.join(startup_script_path, "restore-snapshot.json") git_script_path = os.path.join(comfyui_manager_path, "git_helper.py") +pip_overrides_path = os.path.join(comfyui_manager_path, "pip_overrides.json") + + +cm_global.pip_overrides = {} +if os.path.exists(pip_overrides_path): + with open(pip_overrides_path, 'r', encoding="UTF-8", errors="ignore") as json_file: + cm_global.pip_overrides = json.load(json_file) + + +def remap_pip_package(pkg): + if pkg in cm_global.pip_overrides: + res = cm_global.pip_overrides[pkg] + print(f"[ComfyUI-Manager] '{pkg}' is remapped to '{res}'") + return res + else: + return pkg + std_log_lock = threading.Lock() @@ -391,8 +409,6 @@ def is_installed(name): if os.path.exists(restore_snapshot_path): try: - import json - cloned_repos = [] def msg_capture(stream, prefix): @@ -436,7 +452,7 @@ if os.path.exists(restore_snapshot_path): if os.path.exists(requirements_path): with open(requirements_path, 'r', encoding="UTF-8", errors="ignore") as file: for line in file: - package_name = line.strip() + package_name = remap_pip_package(line.strip()) if package_name and not is_installed(package_name): install_cmd = [sys.executable, "-m", "pip", "install", package_name] this_exit_code += process_wrap(install_cmd, repo_path) @@ -476,7 +492,7 @@ def execute_lazy_install_script(repo_path, executable): print(f"Install: pip packages for '{repo_path}'") with open(requirements_path, "r") as requirements_file: for line in requirements_file: - package_name = line.strip() + package_name = remap_pip_package(line.strip()) if package_name and not is_installed(package_name): install_cmd = [executable, "-m", "pip", "install", package_name] process_wrap(install_cmd, repo_path)