From 09422da704de8f90f6d9fe9e853c2f45780c5fbc Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Sat, 27 Apr 2024 13:09:01 +0900
Subject: [PATCH] feat: support .yaml snapshot

---
 cm-cli.py            | 15 +++++++++++++--
 git_helper.py        | 18 ++++++++++++++++--
 glob/manager_core.py | 33 ++++++++++++++++++++++++---------
 3 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/cm-cli.py b/cm-cli.py
index dbb799a1..5a2c8e65 100644
--- a/cm-cli.py
+++ b/cm-cli.py
@@ -108,7 +108,9 @@ def restore_dependencies():
 def restore_snapshot(snapshot_name):
     global processed_install
 
-    if not os.path.exists(snapshot_name):
+    if os.path.exists(snapshot_name):
+        snapshot_path = os.path.abspath(snapshot_name)
+    else:
         snapshot_path = os.path.join(core.comfyui_manager_path, 'snapshots', snapshot_name)
         if not os.path.exists(snapshot_path):
             print(f"ERROR: `{snapshot_path}` is not exists.")
@@ -486,6 +488,15 @@ def cancel():
         os.remove(restore_snapshot_path)
 
 
+def save_snapshot():
+    output_path = None
+    for i in range(len(sys.argv)):
+        if sys.argv[i] == '--output':
+            if len(sys.argv) >= i:
+                output_path = sys.argv[i+1]
+
+    return core.save_snapshot_with_postfix('snapshot', output_path)
+
 def for_each_nodes(act, allow_all=True):
     global nodes
 
@@ -566,7 +577,7 @@ elif op == 'cli-only-mode':
         print(f"\ninvalid value for cli-only-mode: {sys.argv[2]}\n")
 
 elif op == 'save-snapshot':
-    path = core.save_snapshot_with_postfix('snapshot')
+    path = save_snapshot()
     print(f"Current snapshot is saved as `{path}`")
 
 elif op == 'restore-snapshot':
diff --git a/git_helper.py b/git_helper.py
index 37d6dcc7..f7b2cc32 100644
--- a/git_helper.py
+++ b/git_helper.py
@@ -4,6 +4,7 @@ import git
 import configparser
 import re
 import json
+import yaml
 from torchvision.datasets.utils import download_url
 from tqdm.auto import tqdm
 from git.remote import RemoteProgress
@@ -278,8 +279,21 @@ def apply_snapshot(target):
     try:
         path = os.path.join(os.path.dirname(__file__), 'snapshots', f"{target}")
         if os.path.exists(path):
-            with open(path, 'r', encoding="UTF-8") as json_file:
-                info = json.load(json_file)
+            if not target.endswith('.json') and not target.endswith('.yaml'):
+                print(f"Snapshot file not found: `{path}`")
+                print("APPLY SNAPSHOT: False")
+                return
+
+            with open(path, 'r', encoding="UTF-8") as snapshot_file:
+                if target.endswith('.json'):
+                    info = json.load(snapshot_file)
+                elif target.endswith('.yaml'):
+                    info = yaml.load(snapshot_file, Loader=yaml.SafeLoader)
+                    info = info['custom_nodes']
+                else:
+                    # impossible case
+                    print("APPLY SNAPSHOT: False")
+                    return
 
                 comfyui_hash = info['comfyui']
                 git_custom_node_infos = info['git_custom_nodes']
diff --git a/glob/manager_core.py b/glob/manager_core.py
index 6207cbcd..066ec20d 100644
--- a/glob/manager_core.py
+++ b/glob/manager_core.py
@@ -14,6 +14,7 @@ import aiohttp
 import threading
 import json
 import time
+import yaml
 
 glob_path = os.path.join(os.path.dirname(__file__))  # ComfyUI-Manager/glob
 sys.path.append(glob_path)
@@ -21,7 +22,7 @@ sys.path.append(glob_path)
 import cm_global
 from manager_util import *
 
-version = [2, 23, 1]
+version = [2, 24]
 version_str = f"V{version[0]}.{version[1]}" + (f'.{version[2]}' if len(version) > 2 else '')
 
 comfyui_manager_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
@@ -982,15 +983,29 @@ def get_current_snapshot():
     }
 
 
-def save_snapshot_with_postfix(postfix):
-    now = datetime.now()
+def save_snapshot_with_postfix(postfix, path=None):
+    if path is None:
+        now = datetime.now()
 
-    date_time_format = now.strftime("%Y-%m-%d_%H-%M-%S")
-    file_name = f"{date_time_format}_{postfix}"
+        date_time_format = now.strftime("%Y-%m-%d_%H-%M-%S")
+        file_name = f"{date_time_format}_{postfix}"
 
-    path = os.path.join(comfyui_manager_path, 'snapshots', f"{file_name}.json")
-    with open(path, "w") as json_file:
-        json.dump(get_current_snapshot(), json_file, indent=4)
+        path = os.path.join(comfyui_manager_path, 'snapshots', f"{file_name}.json")
+    else:
+        file_name = path.replace('\\', '/').split('/')[-1]
+        file_name = file_name.split('.')[-2]
 
-    return file_name+'.json'
+    snapshot = get_current_snapshot()
+    if path.endswith('.json'):
+        with open(path, "w") as json_file:
+            json.dump(snapshot, json_file, indent=4)
+
+        return file_name + '.json'
+
+    elif path.endswith('.yaml'):
+        with open(path, "w") as yaml_file:
+            snapshot = {'custom_nodes': snapshot}
+            yaml.dump(snapshot, yaml_file, allow_unicode=True)
+
+        return path