diff --git a/nodes/api.py b/nodes/api.py index 5d1cf31..3155857 100644 --- a/nodes/api.py +++ b/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}"})