Add support for custom node only snapshots (#4) (#1542)

* Add support for custom node only snapshots (#4)

* Fix ruff lint.

---------

Co-authored-by: pythongosssss <125205205+pythongosssss@users.noreply.github.com>

cm-cli.py

@@ -1012,17 +1012,22 @@ def save_snapshot(
         user_directory: str = typer.Option(
             None,
             help="user directory"
-        )
+        ),
+        full_snapshot: Annotated[
+            bool,
+            typer.Option(
+                show_default=False, help="If the snapshot should include custom node, ComfyUI version and pip versions (default), or only custom node details"
+            ),
+        ] = True,
 ):
     cmd_ctx.set_user_directory(user_directory)
-
-    path = asyncio.run(core.save_snapshot_with_postfix('snapshot', output))
+    path = core.save_snapshot_with_postfix('snapshot', output, not full_snapshot)
     print(f"Current snapshot is saved as `{path}`")
 
 
 @app.command("restore-snapshot", help="Restore snapshot from snapshot file")
 def restore_snapshot(
-        snapshot_name: str,
+        snapshot_name: str, 
         pip_non_url: Optional[bool] = typer.Option(
             default=None,
             show_default=False,

git_helper.py

@@ -389,12 +389,13 @@ def apply_snapshot(path):
                 git_custom_node_infos = info['git_custom_nodes']
                 file_custom_node_infos = info['file_custom_nodes']
 
-                checkout_comfyui_hash(comfyui_hash)
+                if comfyui_hash:
+                    checkout_comfyui_hash(comfyui_hash)
                 checkout_custom_node_hash(git_custom_node_infos)
                 invalidate_custom_node_file(file_custom_node_infos)
 
                 print("APPLY SNAPSHOT: True")
-                if 'pips' in info:
+                if 'pips' in info and info['pips']:
                     return info['pips']
                 else:
                     return None

glob/manager_core.py

@@ -2478,7 +2478,7 @@ def get_installed_pip_packages():
     return res
 
 
-async def get_current_snapshot():
+async def get_current_snapshot(custom_nodes_only = False):
     await unified_manager.reload('cache')
     await unified_manager.get_custom_nodes('default', 'cache')
 
@@ -2489,8 +2489,10 @@ async def get_current_snapshot():
         print("ComfyUI update fail: The installed ComfyUI does not have a Git repository.")
         return {}
 
-    repo = git.Repo(repo_path)
-    comfyui_commit_hash = repo.head.commit.hexsha
+    comfyui_commit_hash = None
+    if not custom_nodes_only:
+        repo = git.Repo(repo_path)
+        comfyui_commit_hash = repo.head.commit.hexsha
 
     git_custom_nodes = {}
     cnr_custom_nodes = {}
@@ -2556,7 +2558,7 @@ async def get_current_snapshot():
 
                 file_custom_nodes.append(item)
 
-    pip_packages = get_installed_pip_packages()
+    pip_packages = None if custom_nodes_only else get_installed_pip_packages()
 
     return {
         'comfyui': comfyui_commit_hash,
@@ -2567,7 +2569,7 @@ async def get_current_snapshot():
     }
 
 
-async def save_snapshot_with_postfix(postfix, path=None):
+async def save_snapshot_with_postfix(postfix, path=None, custom_nodes_only = False):
     if path is None:
         now = datetime.now()
 
@@ -2579,7 +2581,7 @@ async def save_snapshot_with_postfix(postfix, path=None):
         file_name = path.replace('\\', '/').split('/')[-1]
         file_name = file_name.split('.')[-2]
 
-    snapshot = await get_current_snapshot()
+    snapshot = await get_current_snapshot(custom_nodes_only)
     if path.endswith('.json'):
         with open(path, "w") as json_file:
             json.dump(snapshot, json_file, indent=4)

glob/manager_server.py

@@ -447,7 +447,7 @@ async def task_worker():
         except Exception:
             traceback.print_exc()
 
-        return f"An error occurred while updating 'comfyui'."
+        return "An error occurred while updating 'comfyui'."
 
     async def do_fix(item) -> str:
         ui_id, node_name, node_ver = item