diff --git a/glob/cnr_utils.py b/glob/cnr_utils.py index 5b2e9021..ea9a1a93 100644 --- a/glob/cnr_utils.py +++ b/glob/cnr_utils.py @@ -99,14 +99,3 @@ def all_versions_of_node(node_id): else: return None - -def extract_package_as_zip(file_path, extract_path): - try: - with zipfile.ZipFile(file_path, "r") as zip_ref: - zip_ref.extractall(extract_path) - extracted_files = zip_ref.namelist() - print(f"Extracted zip file to {extract_path}") - return extracted_files - except zipfile.BadZipFile: - print(f"File '{file_path}' is not a zip or is corrupted.") - return None diff --git a/glob/manager_core.py b/glob/manager_core.py index 07bfe393..c24249e6 100644 --- a/glob/manager_core.py +++ b/glob/manager_core.py @@ -21,7 +21,6 @@ from rich import print from packaging import version import uuid -import requests glob_path = os.path.join(os.path.dirname(__file__)) # ComfyUI-Manager/glob sys.path.append(glob_path) @@ -37,24 +36,6 @@ version_str = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' i DEFAULT_CHANNEL = "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main" -def download_url(url, dest_folder, filename): - # Ensure the destination folder exists - if not os.path.exists(dest_folder): - os.makedirs(dest_folder) - - # Full path to save the file - dest_path = os.path.join(dest_folder, filename) - - # Download the file - response = requests.get(url, stream=True) - if response.status_code == 200: - with open(dest_path, 'wb') as file: - for chunk in response.iter_content(chunk_size=1024): - if chunk: - file.write(chunk) - else: - raise Exception(f"Failed to download file from {url}") - custom_nodes_path = os.path.abspath(os.path.join(comfyui_manager_path, '..')) @@ -767,6 +748,24 @@ class UnifiedManager: return True + def reserve_cnr_switch(self, target, zip_url, from_path, to_path, no_deps): + script_path = os.path.join(startup_script_path, "install-scripts.txt") + with open(script_path, "a") as file: + obj = [target, "#LAZY-CNR-SWITCH-SCRIPT", zip_url, from_path, to_path, no_deps, custom_nodes_path, sys.executable] + file.write(f"{obj}\n") + + print(f"Installation reserved: {target}") + + return True + + def reserve_migration(self, moves): + script_path = os.path.join(startup_script_path, "install-scripts.txt") + with open(script_path, "a") as file: + obj = ["", "#LAZY-MIGRATION", moves] + file.write(f"{obj}\n") + + return True + def unified_fix(self, node_id, version_spec, instant_execution=False, no_deps=False): """ fix dependencies @@ -783,6 +782,44 @@ class UnifiedManager: return result def cnr_switch_version(self, node_id, version_spec=None, instant_execution=False, no_deps=False, return_postinstall=False): + if instant_execution: + return self.cnr_switch_version_instant(node_id, version_spec, instant_execution, no_deps, return_postinstall) + else: + return self.cnr_switch_version_lazy(node_id, version_spec, no_deps, return_postinstall) + + def cnr_switch_version_lazy(self, node_id, version_spec=None, no_deps=False, return_postinstall=False): + """ + switch between cnr version (lazy mode) + """ + + result = ManagedResult('switch-cnr') + + node_info = cnr_utils.install_node(node_id, version_spec) + if node_info is None or not node_info.download_url: + return result.fail(f'not available node: {node_id}@{version_spec}') + + version_spec = node_info.version + + if self.active_nodes[node_id][0] == version_spec: + return ManagedResult('skip').with_msg("Up to date") + + zip_url = node_info.download_url + from_path = self.active_nodes[node_id][1] + target = f"{node_id}@{version_spec.replace('.', '_')}" + to_path = os.path.join(custom_nodes_path, target) + + def postinstall(): + return self.reserve_cnr_switch(target, zip_url, from_path, to_path, no_deps) + + if return_postinstall: + return result.with_postinstall(postinstall) + else: + if not postinstall(): + return result.fail(f"Failed to execute install script: {node_id}@{version_spec}") + + return result + + def cnr_switch_version_instant(self, node_id, version_spec=None, instant_execution=True, no_deps=False, return_postinstall=False): """ switch between cnr version """ @@ -809,7 +846,9 @@ class UnifiedManager: os.remove(download_path) if extracted is None: - shutil.rmtree(install_path) + if len(os.listdir(install_path)) == 0: + shutil.rmtree(install_path) + return result.fail(f'Empty archive file: {node_id}@{version_spec}') # 3. calculate garbage files (.tracking - extracted) @@ -1284,48 +1323,18 @@ class UnifiedManager: await self.get_custom_nodes('default', 'cache') print(f"Migration: STAGE 1") + moves = [] + # migrate nightly inactive - fixes = {} for x, v in self.nightly_inactive_nodes.items(): if v.endswith('@nightly'): continue new_path = os.path.join(custom_nodes_path, '.disabled', f"{x}@nightly") - shutil.move(v, new_path) - fixes[x] = new_path - - self.nightly_inactive_nodes.update(fixes) - - # NOTE: Don't migration unknown node - keep original name as possible - # print(f"Migration: STAGE 2") - # # migrate unknown inactive - # fixes = {} - # for x, v in self.unknown_inactive_nodes.items(): - # if v[1].endswith('@unknown'): - # continue - # - # new_path = os.path.join(custom_nodes_path, '.disabled', f"{x}@unknown") - # shutil.move(v[1], new_path) - # fixes[x] = v[0], new_path - # - # self.unknown_inactive_nodes.update(fixes) - - # print(f"Migration: STAGE 3") - # migrate unknown active nodes - # fixes = {} - # for x, v in self.unknown_active_nodes.items(): - # if v[1].endswith('@unknown'): - # continue - # - # new_path = os.path.join(custom_nodes_path, f"{x}@unknown") - # shutil.move(v[1], new_path) - # fixes[x] = v[0], new_path - # - # self.unknown_active_nodes.update(fixes) + moves.append((v, new_path)) print(f"Migration: STAGE 2") # migrate active nodes - fixes = {} for x, v in self.active_nodes.items(): if v[0] not in ['nightly']: continue @@ -1334,12 +1343,11 @@ class UnifiedManager: continue new_path = os.path.join(custom_nodes_path, f"{x}@nightly") - shutil.move(v[1], new_path) - fixes[x] = v[0], new_path + moves.append((v[1], new_path)) - self.active_nodes.update(fixes) + self.reserve_migration(moves) - print(f"DONE") + print(f"DONE (Migration reserved)") unified_manager = UnifiedManager() diff --git a/glob/manager_util.py b/glob/manager_util.py index 7ff0b1f3..f86a30b6 100644 --- a/glob/manager_util.py +++ b/glob/manager_util.py @@ -138,3 +138,37 @@ async def get_data_with_cache(uri, silent=False, cache_mode=True): def sanitize_tag(x): return x.replace('<', '<').replace('>', '>') + + +def download_url(url, dest_folder, filename): + import requests + + # Ensure the destination folder exists + if not os.path.exists(dest_folder): + os.makedirs(dest_folder) + + # Full path to save the file + dest_path = os.path.join(dest_folder, filename) + + # Download the file + response = requests.get(url, stream=True) + if response.status_code == 200: + with open(dest_path, 'wb') as file: + for chunk in response.iter_content(chunk_size=1024): + if chunk: + file.write(chunk) + else: + raise Exception(f"Failed to download file from {url}") + + +def extract_package_as_zip(file_path, extract_path): + import zipfile + try: + with zipfile.ZipFile(file_path, "r") as zip_ref: + zip_ref.extractall(extract_path) + extracted_files = zip_ref.namelist() + print(f"Extracted zip file to {extract_path}") + return extracted_files + except zipfile.BadZipFile: + print(f"File '{file_path}' is not a zip or is corrupted.") + return None \ No newline at end of file diff --git a/prestartup_script.py b/prestartup_script.py index 7265957a..728eabc9 100644 --- a/prestartup_script.py +++ b/prestartup_script.py @@ -560,6 +560,65 @@ def execute_lazy_install_script(repo_path, executable): process_wrap(install_cmd, repo_path, env=new_env) +def execute_lazy_cnr_switch(target, zip_url, from_path, to_path, no_deps, custom_nodes_path): + import uuid + import shutil + + # 1. download + archive_name = f"CNR_temp_{str(uuid.uuid4())}.zip" # should be unpredictable name - security precaution + download_path = os.path.join(custom_nodes_path, archive_name) + download_url(zip_url, custom_nodes_path, archive_name) + + # 2. extract files into @ + extracted = extract_package_as_zip(download_path, from_path) + os.remove(download_path) + + if extracted is None: + if len(os.listdir(from_path)) == 0: + shutil.rmtree(from_path) + + print(f'Empty archive file: {target}') + return False + + + # 3. calculate garbage files (.tracking - extracted) + tracking_info_file = os.path.join(from_path, '.tracking') + prev_files = set() + with open(tracking_info_file, 'r') as f: + for line in f: + prev_files.add(line.strip()) + garbage = prev_files.difference(extracted) + garbage = [os.path.join(custom_nodes_path, x) for x in garbage] + + # 4-1. remove garbage files + for x in garbage: + if os.path.isfile(x): + os.remove(x) + + # 4-2. remove garbage dir if empty + for x in garbage: + if os.path.isdir(x): + if not os.listdir(x): + os.rmdir(x) + + # 5. rename dir name @ ==> @ + print(f"'{from_path}' is moved to '{to_path}'") + shutil.move(from_path, to_path) + + # 6. create .tracking file + tracking_info_file = os.path.join(to_path, '.tracking') + with open(tracking_info_file, "w", encoding='utf-8') as file: + file.write('\n'.join(list(extracted))) + + +def execute_migration(moves): + import shutil + for x in moves: + if os.path.exists(x[0]) and not os.path.exists(x[1]): + shutil.move(x[0], x[1]) + print(f"[ComfyUI-Manager] MIGRATION: '{x[0]}' -> '{x[1]}'") + + # Check if script_list_path exists if os.path.exists(script_list_path): print("\n#######################################################################") @@ -581,6 +640,13 @@ if os.path.exists(script_list_path): if script[1] == "#LAZY-INSTALL-SCRIPT": execute_lazy_install_script(script[0], script[2]) + elif script[1] == "#LAZY-CNR-SWITCH-SCRIPT": + execute_lazy_cnr_switch(script[0], script[2], script[3], script[4], script[5], script[6]) + execute_lazy_install_script(script[3], script[7]) + + elif script[1] == "#LAZY-MIGRATION": + execute_migration(script[2]) + elif os.path.exists(script[0]): if script[1] == "#FORCE": del script[1]