postponed processing for cnr switch & migration

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

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:
+            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)
+            moves.append((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)
 
         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)
+            moves.append((v[1], new_path))
-            fixes[x] = v[0], new_path
 
-        self.active_nodes.update(fixes)
+        self.reserve_migration(moves)
 
-        print(f"DONE")
+        print(f"DONE (Migration reserved)")
 
 
 unified_manager = UnifiedManager()

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('<', '&lt;').replace('>', '&gt;')
+
+
+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

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 <node_id>@<cur_ver>
+    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 <node_id>@<prev_ver> ==> <node_id>@<cur_ver>
+    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]