feat: collect comfyui hash, models, custom nodes into snapshot.json (#2)

Signed-off-by: Frost Ming <me@frostming.com>

nodes/api.py

@@ -2,7 +2,6 @@ import asyncio
 import hashlib
 import json
 import os
-import re
 import shutil
 import subprocess
 import sys
@@ -19,7 +18,7 @@ TEMP_FOLDER = Path(__file__).parent.parent / "temp"
 
 
 async def _write_requirements(zf: zipfile.ZipFile) -> None:
-    print("Pack => Writing requirements.txt")
+    print("Package => Writing requirements.txt")
     with zf.open("requirements.txt", "w") as f:
         proc = await asyncio.subprocess.create_subprocess_exec(
             sys.executable,
@@ -33,8 +32,22 @@ async def _write_requirements(zf: zipfile.ZipFile) -> None:
         f.write(stdout)
 
 
-async def _write_models(zf: zipfile.ZipFile) -> None:
-    print("Pack => Writing models")
+async def _write_snapshot(zf: zipfile.ZipFile) -> None:
+    proc = await asyncio.subprocess.create_subprocess_exec(
+        "git", "rev-parse", "HEAD", stdout=subprocess.PIPE, cwd=folder_paths.base_path
+    )
+    stdout, _ = await proc.communicate()
+    with zf.open("snapshot.json", "w") as f:
+        data = {
+            "comfyui": stdout.decode().strip(),
+            "models": await _get_models(),
+            "custom_nodes": await _get_custom_nodes(),
+        }
+        f.write(json.dumps(data, indent=2).encode())
+
+
+async def _get_models() -> list:
+    print("Package => Writing models")
     proc = await asyncio.subprocess.create_subprocess_exec(
         "git",
         "ls-files",
@@ -43,36 +56,24 @@ async def _write_models(zf: zipfile.ZipFile) -> None:
         stdout=subprocess.PIPE,
     )
     stdout, _ = await proc.communicate()
-    with zf.open("models.json", "w") as f:
-        data = []
-        for line in stdout.decode().splitlines():
-            if os.path.basename(line).startswith("."):
-                continue
-            filename = os.path.abspath(line)
-            relpath = os.path.relpath(filename, folder_paths.base_path)
-            with open(filename, "rb") as model:
-                data.append(
-                    {
-                        "filename": relpath,
-                        "sha256": hashlib.sha256(model.read()).hexdigest(),
-                    }
-                )
-        f.write(json.dumps(data, indent=2).encode())
+    models = []
+    for line in stdout.decode().splitlines():
+        if os.path.basename(line).startswith("."):
+            continue
+        filename = os.path.abspath(line)
+        relpath = os.path.relpath(filename, folder_paths.base_path)
+        with open(filename, "rb") as model:
+            models.append(
+                {
+                    "filename": relpath,
+                    "sha256": hashlib.sha256(model.read()).hexdigest(),
+                }
+            )
+    return models
 
 
-async def _write_workflow(zf: zipfile.ZipFile, data: dict) -> None:
-    print("Pack => Writing workflow")
-    with zf.open("workflow_api.json", "w") as f:
-        f.write(json.dumps(data["workflow_api"], indent=2).encode())
-    with zf.open("workflow.json", "w") as f:
-        f.write(json.dumps(data["workflow"], indent=2).encode())
-
-
-GIT_URL_REGEX = re.compile(r"(?<=:|/)[^:/]+/[^:/]+(?=\.git|$)")
-
-
-async def _write_custom_nodes(zf: zipfile.ZipFile) -> None:
-    print("Pack => Writing custom nodes")
+async def _get_custom_nodes() -> list:
+    print("Package => Writing custom nodes")
     custom_nodes = os.path.join(folder_paths.base_path, "custom_nodes")
     coros = []
 
@@ -108,13 +109,19 @@ async def _write_custom_nodes(zf: zipfile.ZipFile) -> None:
             continue
         coros.append(get_node_info(subdir))
 
-    with zf.open("custom_nodes.json", "w") as f:
-        data = await asyncio.gather(*coros)
-        f.write(json.dumps(data, indent=2).encode())
+    return await asyncio.gather(*coros)
+
+
+async def _write_workflow(zf: zipfile.ZipFile, data: dict) -> None:
+    print("Package => Writing workflow")
+    with zf.open("workflow_api.json", "w") as f:
+        f.write(json.dumps(data["workflow_api"], indent=2).encode())
+    with zf.open("workflow.json", "w") as f:
+        f.write(json.dumps(data["workflow"], indent=2).encode())
 
 
 async def _write_inputs(zf: zipfile.ZipFile) -> None:
-    print("Pack => Writing inputs")
+    print("Package => Writing inputs")
     input_dir = folder_paths.get_input_directory()
 
     for root_path, _, files in os.walk(input_dir):
@@ -139,9 +146,8 @@ async def pack_workspace(request):
 
     with zipfile.ZipFile(TEMP_FOLDER / zip_filename, "w") as zf:
         await _write_requirements(zf)
-        await _write_models(zf)
+        await _write_snapshot(zf)
         await _write_workflow(zf, data)
-        await _write_custom_nodes(zf)
         await _write_inputs(zf)
 
     return web.json_response({"download_url": f"/bentoml/download/{zip_filename}"})