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>
This commit is contained in:
Robin Huang 2025-02-13 14:26:35 -08:00 committed by GitHub
parent 6ff6e05408
commit eebace1652
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 21 additions and 13 deletions

View File

@ -1012,11 +1012,16 @@ def save_snapshot(
user_directory: str = typer.Option( user_directory: str = typer.Option(
None, None,
help="user directory" 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) cmd_ctx.set_user_directory(user_directory)
path = core.save_snapshot_with_postfix('snapshot', output, not full_snapshot)
path = asyncio.run(core.save_snapshot_with_postfix('snapshot', output))
print(f"Current snapshot is saved as `{path}`") print(f"Current snapshot is saved as `{path}`")

View File

@ -389,12 +389,13 @@ def apply_snapshot(path):
git_custom_node_infos = info['git_custom_nodes'] git_custom_node_infos = info['git_custom_nodes']
file_custom_node_infos = info['file_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) checkout_custom_node_hash(git_custom_node_infos)
invalidate_custom_node_file(file_custom_node_infos) invalidate_custom_node_file(file_custom_node_infos)
print("APPLY SNAPSHOT: True") print("APPLY SNAPSHOT: True")
if 'pips' in info: if 'pips' in info and info['pips']:
return info['pips'] return info['pips']
else: else:
return None return None

View File

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

View File

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