From 4de5776a1d10933530ee29a545c1f4bb1feef538 Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Sat, 4 May 2024 16:58:08 +0900
Subject: [PATCH] feat: cm-cli - node-in-workflow command

---
 cm-cli.py            |  23 +++++++++
 glob/manager_core.py | 115 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 137 insertions(+), 1 deletion(-)

diff --git a/cm-cli.py b/cm-cli.py
index 007f562c..eb17f9a1 100644
--- a/cm-cli.py
+++ b/cm-cli.py
@@ -27,6 +27,7 @@ if not (len(sys.argv) == 2 and sys.argv[1] in ['save-snapshot', 'restore-depende
           f"    restore-snapshot <snapshot .json/.yaml>\n"
           f"    cli-only-mode [enable|disable]\n"
           f"    restore-dependencies\n"
+          f"    node-in-workflow <workflow .json/.png> <output name>\n"
           f"    clear\n")
     exit(-1)
 
@@ -497,6 +498,21 @@ def save_snapshot():
 
     return core.save_snapshot_with_postfix('snapshot', output_path)
 
+
+def node_in_workflow(input_path, output_path):
+    used_exts, unknown_nodes = asyncio.run(core.extract_nodes_from_workflow(input_path))
+
+    res = {
+            'custom_nodes': list(used_exts),
+            'unknown_nodes': list(unknown_nodes)
+    }
+
+    with open(output_path, "w", encoding='utf-8') as output_file:
+        json.dump(res, output_file, indent=4)
+
+    print(f"Workflow dependencies are being saved into {output_path}.")
+
+
 def for_each_nodes(act, allow_all=True):
     global nodes
 
@@ -576,6 +592,13 @@ elif op == 'cli-only-mode':
     else:
         print(f"\ninvalid value for cli-only-mode: {sys.argv[2]}\n")
 
+elif op == "nodes-in-workflow":
+    if len(sys.argv) < 4:
+        print(f"missing arguments: python cm-cli.py <workflow file> <output path>")
+        exit(-1)
+
+    node_in_workflow(sys.argv[2], sys.argv[3])
+
 elif op == 'save-snapshot':
     path = save_snapshot()
     print(f"Current snapshot is saved as `{path}`")
diff --git a/glob/manager_core.py b/glob/manager_core.py
index e1f00eba..82c8ac4f 100644
--- a/glob/manager_core.py
+++ b/glob/manager_core.py
@@ -22,7 +22,7 @@ sys.path.append(glob_path)
 import cm_global
 from manager_util import *
 
-version = [2, 24, 1]
+version = [2, 25]
 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__), '..'))
@@ -1009,3 +1009,116 @@ def save_snapshot_with_postfix(postfix, path=None):
 
         return path
 
+
+async def extract_nodes_from_workflow(filepath):
+    # prepare json data
+    workflow = None
+    if filepath.endswith('.json'):
+        with open(filepath, "r", encoding="UTF-8", errors="ignore") as json_file:
+            try:
+                workflow = json.load(json_file)
+            except:
+                print(f"Invalid workflow file: {filepath}")
+                exit(-1)
+
+    elif filepath.endswith('.png'):
+        from PIL import Image
+        with Image.open(filepath) as img:
+            if 'workflow' not in img.info:
+                print(f"The specified .png file doesn't have a workflow: {filepath}")
+                exit(-1)
+            else:
+                try:
+                    workflow = json.loads(img.info['workflow'])
+                except:
+                    print(f"This is not a valid .png file containing a ComfyUI workflow: {filepath}")
+                    exit(-1)
+
+    if workflow is None:
+        print(f"Invalid workflow file: {filepath}")
+        exit(-1)
+
+    # extract nodes
+    used_nodes = set()
+
+    def extract_nodes(sub_workflow):
+        for x in sub_workflow['nodes']:
+            node_name = x.get('type')
+
+            # skip virtual nodes
+            if node_name in ['Reroute', 'Note']:
+                continue
+
+            if node_name is not None and not node_name.startswith('workflow/'):
+                used_nodes.add(node_name)
+
+    if 'nodes' in workflow:
+        extract_nodes(workflow)
+
+        if 'extra' in workflow:
+            if 'groupNodes' in workflow['extra']:
+                for x in workflow['extra']['groupNodes'].values():
+                    extract_nodes(x)
+
+    # lookup dependent custom nodes
+    ext_map = await get_data_by_mode('local', 'extension-node-map.json')
+
+    rext_map = {}
+    preemption_map = {}
+    patterns = []
+    for k, v in ext_map.items():
+        if k == 'https://github.com/comfyanonymous/ComfyUI':
+            for x in v[0]:
+                if x not in preemption_map:
+                    preemption_map[x] = []
+
+                preemption_map[x] = k
+            continue
+
+        for x in v[0]:
+            if x not in rext_map:
+                rext_map[x] = []
+
+            rext_map[x].append(k)
+
+        if 'preemptions' in v[1]:
+            for x in v[1]['preemptions']:
+                if x not in preemption_map:
+                    preemption_map[x] = []
+
+                preemption_map[x] = k
+
+        if 'nodename_pattern' in v[1]:
+            patterns.append((v[1]['nodename_pattern'], k))
+
+    # identify used extensions
+    used_exts = set()
+    unknown_nodes = set()
+
+    for node_name in used_nodes:
+        ext = preemption_map.get(node_name)
+
+        if ext is None:
+            ext = rext_map.get(node_name)
+            if ext is not None:
+                ext = ext[0]
+
+        if ext is None:
+            for pat_ext in patterns:
+                if re.search(pat_ext[0], node_name):
+                    ext = pat_ext[1]
+                    break
+
+        if ext == 'https://github.com/comfyanonymous/ComfyUI':
+            pass
+        elif ext is not None:
+            if 'Fooocus' in ext:
+                print(f">> {node_name}")
+
+            used_exts.add(ext)
+        else:
+            unknown_nodes.add(node_name)
+
+
+    return used_exts, unknown_nodes
+