fix: Model hashing fixes (#62)

* make sure hashing is done only when the models are selected

* fix cache file not being written, if it did not exist

---------

Co-authored-by: Frost Ming <me@frostming.com>

.github/workflows/release.yml

@@ -22,7 +22,7 @@ jobs:
       id-token: write
     environment:
       name: pypi
-      url: https://pypi.org/project/comfyui-idl/${{ github.ref_name }}
+      url: https://pypi.org/project/comfy-pack/${{ github.ref_name }}
     runs-on: ubuntu-latest
     needs: build-package
 

.gitignore

@@ -161,3 +161,5 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 /temp/
+src/.DS_Store
+.DS_Store

README.md

@@ -1,73 +1,214 @@
-# ComfyPack
+# Comfy-Pack: Making ComfyUI Workflows Shareable
 
-A comprehensive toolkit for standardizing, packaging and deploying ComfyUI workflows as reproducible environments and production-ready REST services.
+![banner2](https://github.com/user-attachments/assets/14a7e469-6683-4818-9d54-5e5a8d0aa454)
 
-## Features
-- **Package Everything**: Create reproducible `.cpack.zip` files containing your workflow, custom nodes, model versions, and all dependencies
-- **Standardize Parameters**: Define and validate workflow inputs through UI nodes for images, text, numbers and more
-- **CLI Support**: Restore environment and run inference from command line
-- **REST API Generation**: Auto-convert any workflow into REST service with OpenAPI docs
 
-## Quick Start
+comfy-pack is a comprehensive toolkit for reliably packing and unpacking environments for ComfyUI workflows. 
+
+
+- 📦 **Pack workflow environments as artifacts:** Saves the workflow environment in a `.cpack.zip` artifact with Python package versions, ComfyUI and custom node revisions, and model hashes.
+- ✨ **Unpack artifacts to recreate workflow environments:** Unpacks the `.cpack.zip` artifact to recreate the same environment with the exact Python package versions, ComfyUI and custom node revisions, and model weights.
+- 🚀 **Deploy workflows as APIs:** Deploys the workflow as a RESTful API with customizable input and output parameters.
+
+## Motivations
+ComfyUI Manager is great for find missing custom nodes. But when sharing ComfyUI workflows to others(your audience or team members), you've still likely heard these responses:
+
+- "Custom Node not found"
+- "Cannot find the correct model file"
+- "Missing Python dependencies"
+
+These are fundamental challenges in workflow sharing – every component should match exactly: custom nodes, model files, and Python dependencies. Modern pacakge managers like npm and poetry introduced "lock" feature, which means record the exact version for every requirement. ComfyUI Manager isn't designed for that.
+
+We learned it from our community and developed comfy-pack to address these problems. With a single click, it captures and locks your entire workflow environment into a `.cpack.zip` file, including Python packages, custom nodes, model hashes, and required assets.
+
+Users can recreate the exact environment with one command:
+
+```bash
+comfy-pack unpack workflow.cpack.zip
+```
+
+This means you can focus on your creative work while comfy-pack handles the rest.
+
+## Usages
 
 ### Installation
+
+We recommend you use ComfyUI Manager to install comfy-pack. Simply search for `comfy-pack` and click **Install**. Restart the server and refresh your ComfyUI interface to apply changes.
+
+![install_node](https://github.com/user-attachments/assets/dbfb730d-edff-4a52-b6c4-695e3ec70368)
+
+Alternatively, clone the project repository through `git`.
+
+```bash
+cd ComfyUI/custom_nodes
+git clone https://github.com/bentoml/comfy-pack.git
+```
+
+To install the comfy-pack CLI, run:
+
 ```bash
 pip install comfy-pack
 ```
 
-### Create a Pack
-1. Install ComfyPack custom nodes in ComfyUI
-2. Design your workflow with parameter nodes
-3. Click "Package" button to create `.cpack.zip`
+### Pack a ComfyUI workflow and its environment
+
+You can package a workflow and the environment required to run the workflow into an artifact that can be unpacked elsewhere.
+
+1. Click the **Package** button to create a `.cpack.zip` artifact.
+2. (Optional) Select the models that you want to include (only model hash will be recorded, so you won't get a 100GB zip file).
+
+![pack](https://github.com/user-attachments/assets/e08bbed2-84dc-474e-a701-6c6db16edf76)
+
+### Unpack the ComfyUI environments
+
+Unpacking a `.cpack.zip` artifact will restore the ComfyUI environment for the workflow. During unpacking, comfy-pack will perform the following steps.
+
+1. Prepare a Python virtual environment with the exact packages used to run the workflow.
+2. Clone ComfyUI and custom nodes from the exact revisions required by the workflow.
+3. Search for and download models from common registries like Hugging Face and Civitai. Unpacking workflows using the same model will not cause the model to be downloaded multiple times. Instead, model weights will be symbolically linked.
+
+To unpack:
 
-### Restore to a ComfyUI project
 ```bash
-# Restore environment from pack, will install everything needed except the models.
-comfy-pack restore workflow.cpack.zip --dir ./
+comfy-pack unpack workflow.cpack.zip
 ```
 
-### Run Inference
+Huggingface gated models can be accessed by setting your `HF_TOKEN` as an environment variable before unpacking:
+
+```bash
+export HF_TOKEN=hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXx
+comfy-pack unpack workflow.cpack.zip
+```
+
+For example cpack files, check our [examples folder](examples/).
+
+### Deploy a workflow as an API
+
+You can turn a ComfyUI workflow into an API endpoint callable using any clients through HTTP.
+
+<details>
+<summary> 1. Annotate input & output </summary>
+
+Use custom nodes provided by comfy-pack to annotate the fields to be used as input and output parameters. To add a comfy-pack node, right-click and select **Add Node** > **ComfyPack** > **output/input** > [Select a type]
+
+Input nodes:
+
+- ImageInput: Accepts `image` type input, similar to the official `LoadImage` node
+- StringInput: Accepts `string` type input (e.g., prompts)
+- IntInput: Accepts `int` type input (e.g., dimensions, seeds)
+- AnyInput: Accepts `combo` type and more input (e.g., custom nodes)
+
+![input](https://github.com/user-attachments/assets/44264007-0ac8-4e23-8dc0-e60aa0ebcea2)
+
+Output nodes:
+
+- ImageOutput: Outputs `image` type, similar to the official `SaveImage` node
+- FileOutput: Outputs file path as `string` type and saves the file under that path
+
+![output](https://github.com/user-attachments/assets/a4526661-8930-4575-bacc-33b6887f6271)
+
+More field types are under way.
+</details>
+
+<details>
+<summary> 2. Serve the workflow </summary>
+
+Start an HTTP server at `http://127.0.0.1:3000` (default) to serve the workflow under the `/generate` path.
+
+![serve](https://github.com/user-attachments/assets/8d4c92c5-d6d7-485e-bc71-e4fc0fe8bf35)
+
+You can call the `/generate` endpoint by specifying parameters configured through your comfy-pack nodes, such as prompt, width, height, and seed.
+
+> [!NOTE]
+> The name of a comfy-pack node is the parameter name used for API calls.
+
+Examples to call the endpoint:
+
+CURL
+
+```bash
+curl -X 'POST' \
+  'http://127.0.0.1:3000/generate' \
+  -H 'accept: application/octet-stream' \
+  -H 'Content-Type: application/json' \
+  -d '{
+  "prompt": "rocks in a bottle",
+  "width": 512, 
+  "height": 512,
+  "seed": 1
+}'
+```
+
+BentoML client
+
+Under the hood, comfy-pack leverages [BentoML](https://github.com/bentoml/BentoML), the unified model serving framework. You can invoke the endpoint using [the BentoML Python client](https://docs.bentoml.com/en/latest/build-with-bentoml/clients.html):
+
+```python
+import bentoml
+
+with bentoml.SyncHTTPClient("http://127.0.0.1:3000") as client:
+        result = client.generate(
+            prompt="rocks in a bottle",
+            width=512,
+            height=512,
+            seed=1
+        )
+```
+
+</details>
+
+<details>
+<summary> 3. (Optional) Pack the workflow and environment </summary>
+
+Pack the workflow and environment into an artifact that can be unpacked elsewhere to recreate the workflow.
+
 ```bash
 # Get the workflow input spec
-comfy-pack info workflow.cpack.zip
+comfy-pack run workflow.cpack.zip --help
 
 # Run
-comfy-pack run workflow.cpack.zip --src-image image.png --video video.mp4 
+comfy-pack run workflow.cpack.zip --src-image image.png --video video.mp4
 ```
+</details>
 
-### Online REST service
-under development
+<details> 
+<summary> 4. (Optional) Deploy to the cloud </summary>
 
-## Parameter Nodes
+Deploy to [BentoCloud](https://www.bentoml.com/) with access to a variety of GPUs and blazing fast scaling.
 
-ComfyPack provides custom nodes for standardizing inputs:
-- ImageInput
-- StringInput  
-- IntInput
-- AnyInput
-- ImageOutput
-- FileOutput
-- ...
+Follow [the instructions here](https://docs.bentoml.com/en/latest/scale-with-bentocloud/manage-api-tokens.html) to get your BentoCloud access token. If you don’t have a BentoCloud account, you can [sign up for free](https://bentoml.com/).
 
-These nodes help define clear interfaces for your workflow.
+![image](https://github.com/user-attachments/assets/1ffa31fc-1f50-4ea7-a47e-7dae3b874273)
 
-## Docker Support
-Under development
+</details>
+
+## Security Guidelines
+
+A cpack file only contains the metadata of the workflow environment, such as Python package versions, ComfyUI and custom node revisions, and model hashes. It does not contain any sensitive information like API keys, passwords, or user data. However, unpacking a cpack file will install custom nodes and Python dependencies. It is recommended to unpack cpack files from trusted sources.
+
+comfy-pack has a strict mode for unpacking. You can enable it by setting the `CPACK_STRICT_MODE` environment variable to `true`. It will sacrifice some flexibility and compatibility for security. For now, comfy-pack will:
+
+* Use more strict index strategy in Python package installation
+
+More security features are under way.
 
 
-## Examples
+## Roadmap
 
-Check our [examples folder](examples/) for:
-- Basic workflow packaging
-- Parameter configuration
-- API integration
-- Docker deployment
+This project is under active development. Currently we are working on:
 
-## License
-MIT License
+- Enhanced user experience
+- Docker support
+- Local `.cpack` file management with version control
+- Enhanced service capabilities
 
 ## Community
-- Issues & Feature Requests: GitHub Issues
-- Questions & Discussion: Discord Server
 
-Detailed documentation: under development
+comfy-pack is actively maintained by the BentoML team. Feel free to reach out 👉 [Join our Slack community!](https://l.bentoml.com/join-slack)
+
+## Contributing
+
+As an open-source project, we welcome contributions of all kinds, such as new features, bug fixes, and documentation. Here are some of the ways to contribute:
+
+- Repost a bug by creating a [GitHub issue](https://github.com/bentoml/comfy-pack/issues).
+- Submit a [pull request](https://github.com/bentoml/comfy-pack/pulls) or help review other developers’ pull requests.

nodes/__init__.py

@@ -0,0 +1,7 @@
+import pathlib
+import sys
+
+SRC_DIR = pathlib.Path(__file__).parent.parent / "src"
+
+if str(SRC_DIR) not in sys.path:
+    sys.path.insert(0, str(SRC_DIR))

nodes/api.py

@@ -1,102 +1,126 @@
 from __future__ import annotations
 
 import asyncio
-import hashlib
 import json
 import os
 import shutil
+import socket
 import subprocess
 import sys
 import tempfile
 import time
 import uuid
 import zipfile
+from fnmatch import fnmatch
 from pathlib import Path
-from typing import Union
+from typing import Any, Union
+
+from aiohttp import web
+from comfy_pack.hash import async_batch_get_sha256
+from comfy_pack.model_helper import alookup_model_source
+from comfy_pack.package import build_bento
 
 import folder_paths
-from aiohttp import web
 from server import PromptServer
 
 ZPath = Union[Path, zipfile.Path]
 TEMP_FOLDER = Path(__file__).parent.parent / "temp"
-COMFY_PACK_DIR = Path(__file__).parent.parent / "src/comfy_pack"
-EXCLUDE_PACKAGES = ["bentoml", "onnxruntime"]  # TODO: standardize this
+COMFY_PACK_DIR = Path(__file__).parent.parent / "src" / "comfy_pack"
+EXCLUDE_PACKAGES = ["bentoml", "onnxruntime", "conda", "nvidia-*"]
 
 
-async def _write_requirements(path: ZPath, extras: list[str] | None = None) -> None:
-    print("Package => Writing requirements.txt")
-    with path.joinpath("requirements.txt").open("w") as f:
-        proc = await asyncio.subprocess.create_subprocess_exec(
-            sys.executable,
-            "-m",
-            "pip",
-            "freeze",
-            "--exclude-editable",
-            *[f"--exclude={p}" for p in EXCLUDE_PACKAGES],
-            stdout=subprocess.PIPE,
-        )
-        stdout, _ = await proc.communicate()
-        f.write(stdout.decode().rstrip("\n") + "\n")
-        if extras:
-            f.write("\n".join(extras) + "\n")
+def normalize_name(name: str) -> str:
+    import re
+
+    return re.sub(r"[-_.]+", "-", name).lower()
 
 
-async def _write_snapshot(path: ZPath, data: dict, models: list | None = None) -> None:
-    proc = await asyncio.subprocess.create_subprocess_exec(
-        "git", "rev-parse", "HEAD", stdout=subprocess.PIPE, cwd=folder_paths.base_path
+def get_snapshot_path() -> Path | None:
+    manager_file_path = Path(
+        folder_paths.get_user_directory(), "default", "ComfyUI-Manager"
     )
-    stdout, _ = await proc.communicate()
-    if models is None:
-        models = await _get_models(data)
-    with path.joinpath("snapshot.json").open("w") as f:
-        data = {
-            "python": f"{sys.version_info.major}.{sys.version_info.minor}",
-            "comfyui": stdout.decode().strip(),
-            "models": models,
-            "custom_nodes": await _get_custom_nodes(),
-        }
-        f.write(json.dumps(data, indent=2))
+    return manager_file_path / "snapshots"
 
 
-def _get_file_info_key(file_path: Path) -> tuple:
-    # use mtime, size to determine if a file has changed
-    return (file_path.stat().st_mtime, file_path.stat().st_size)
+async def _save_snapshot() -> dict[str, Any]:
+    save_snapshot_route = next(
+        (
+            route
+            for route in PromptServer.instance.routes
+            if route.path == "/snapshot/save"
+        ),
+        None,
+    )
+    if not save_snapshot_route:
+        raise RuntimeError("ComfyUI-Manager must be installed to save snapshot")
+    await save_snapshot_route.handler(None)
+    snapshot_path = get_snapshot_path()
+    if not snapshot_path.exists():
+        raise RuntimeError("Snapshot save failed")
+
+    most_recent = max(
+        snapshot_path.glob("*.json"), key=lambda x: x.stat().st_mtime, default=None
+    )
+    if not most_recent:
+        raise RuntimeError("Snapshot save failed")
+    with most_recent.open("r") as f:
+        return json.load(f)
 
 
-def _get_hash(file_path: Path) -> str:
-    cpack_hash_cache_file = TEMP_FOLDER / "model_sha_cache.json"
-    try:
-        with cpack_hash_cache_file.open("r") as f:
-            cache = json.load(f)
-    except (FileNotFoundError, json.JSONDecodeError):
-        cache = {}
-
-    mtime, size = _get_file_info_key(file_path)
-    if file_path.absolute().as_posix() in cache:
-        if (
-            cache[file_path.absolute().as_posix()]["mtime"] == mtime
-            and cache[file_path.absolute().as_posix()]["size"] == size
+async def _write_snapshot(path: ZPath, data: dict, models: list) -> None:
+    snapshot = await _save_snapshot()
+    for package in list(snapshot["pips"]):
+        if any(
+            fnmatch(normalize_name(package.split("==")[0]), pat)
+            for pat in EXCLUDE_PACKAGES
         ):
-            return cache[file_path.absolute().as_posix()]["sha256"]
-
-    with cpack_hash_cache_file.open("w") as f:
-        hasher = hashlib.sha256()
-        with open(file_path, "rb") as file:
-            for chunk in iter(lambda: file.read(4096), b""):
-                hasher.update(chunk)
-        sha = hasher.hexdigest()
-        cache[file_path.absolute().as_posix()] = {
-            "sha256": sha,
-            "mtime": mtime,
-            "size": size,
-        }
-        json.dump(cache, f)
-        return sha
+            del snapshot["pips"][package]
+    with path.joinpath("snapshot.json").open("w") as f:
+        snapshot.update(
+            {
+                "python": f"{sys.version_info.major}.{sys.version_info.minor}",
+                "models": models,
+            }
+        )
+        f.write(json.dumps(snapshot, indent=2))
 
 
-async def _get_models(data: dict, store_models: bool = False) -> list:
-    print("Package => Writing models")
+def _is_port_in_use(port: int | str, host="localhost"):
+    if isinstance(port, str):
+        port = int(port)
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        try:
+            s.connect((host, port))
+            return True
+        except ConnectionRefusedError:
+            return False
+        except Exception:
+            return True
+
+
+def _is_file_refered(file_path: Path, workflow_api: dict) -> bool:
+    """ """
+    used_inputs = set()
+    for node in workflow_api.values():
+        for _, v in node["inputs"].items():
+            if isinstance(v, str):
+                used_inputs.add(v)
+    all_inputs = "\n".join(used_inputs)
+    file_path = file_path.absolute().relative_to(folder_paths.base_path)
+    if file_path.parts[0] == "input":
+        relpath = Path(*file_path.parts[1:])
+    else:  # models
+        relpath = Path(*file_path.parts[2:])
+    return str(relpath) in all_inputs
+
+
+async def _get_models(
+    store_models: bool = False,
+    workflow_api: dict | None = None,
+    model_filter: set[str] | None = None,
+    ensure_sha=True,
+    ensure_source=True,
+) -> list:
     proc = await asyncio.subprocess.create_subprocess_exec(
         "git",
         "ls-files",
@@ -106,84 +130,75 @@ async def _get_models(data: dict, store_models: bool = False) -> list:
     )
     stdout, _ = await proc.communicate()
 
-    used_inputs = set()
-    for node in data["workflow_api"].values():
-        for _, v in node["inputs"].items():
-            if isinstance(v, str):
-                used_inputs.add(v)
-
     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)
+    model_filenames = [
+        os.path.abspath(line.strip().strip('"').replace('\\\\', '/'))
+        for line in stdout.decode().splitlines()
+        if not os.path.basename(line.strip().strip('"').replace('\\\\', '/')).startswith(".")
+    ]
 
-        relpath_path = Path(relpath)
+    # Only compute hashes for referenced models
+    to_include = []
+    if model_filter:
+        for m1 in model_filter:
+            for m2 in model_filenames:
+                if m1 in m2:
+                    to_include.append(m2)
+    else:
+        to_include = model_filenames
+
+    model_hashes = await async_batch_get_sha256(
+        to_include,
+        cache_only=not (ensure_sha or store_models),
+    )
+
+    for filename in to_include:
+        # Skip if file doesn't exist
+        if not os.path.exists(filename):
+            continue
+            
+        relpath = os.path.relpath(filename, folder_paths.base_path)
 
         model_data = {
             "filename": relpath,
-            "sha256": _get_hash(Path(filename)),
-            "explicit": relpath_path.name in used_inputs,
             "size": os.path.getsize(filename),
+            "atime": os.path.getatime(filename),
+            "ctime": os.path.getctime(filename),
+            "disabled": relpath not in model_filter
+            if model_filter is not None
+            else False,
+            "sha256": model_hashes.get(filename),
         }
-        if store_models:
+
+        model_data["source"] = await alookup_model_source(
+            model_data["sha256"],
+            cache_only=not ensure_source,
+        )
+        # should_store = store_models and (
+        #     model_data["source"].get("source") != "huggingface"
+        #     or model_data["source"].get("repo", "").startswith("datasets/")
+        # )  # TODO: sort this out
+        should_store = store_models
+
+        if should_store:
             import bentoml
 
-            model_tag = f'cpack-model:{model_data["sha256"][:16]}'
+            model_tag = f"cpack-model:{model_data['sha256'][:16]}"
             try:
                 model = bentoml.models.get(model_tag)
             except bentoml.exceptions.NotFound:
                 with bentoml.models.create(
-                    model_tag,
-                    labels={"filename": filename},
+                    model_tag, labels={"filename": relpath}
                 ) as model:
                     shutil.copy(filename, model.path_of("model.bin"))
             model_data["model_tag"] = model_tag
         models.append(model_data)
+    if workflow_api:
+        for model in models:
+            model["refered"] = _is_file_refered(Path(model["filename"]), workflow_api)
     return models
 
 
-async def _get_custom_nodes() -> list:
-    print("Package => Writing custom nodes")
-    custom_nodes = os.path.join(folder_paths.base_path, "custom_nodes")
-    coros = []
-
-    async def get_node_info(subdir: Path) -> dict:
-        proc = await asyncio.subprocess.create_subprocess_exec(
-            "git",
-            "config",
-            "--get",
-            "remote.origin.url",
-            cwd=subdir,
-            stdout=subprocess.PIPE,
-        )
-        stdout, _ = await proc.communicate()
-        url = stdout.decode().strip()
-
-        proc = await asyncio.subprocess.create_subprocess_exec(
-            "git",
-            "rev-parse",
-            "HEAD",
-            cwd=subdir,
-            stdout=subprocess.PIPE,
-        )
-        stdout, _ = await proc.communicate()
-        commit_hash = stdout.decode().strip()
-        return {
-            "url": url,
-            "commit_hash": commit_hash,
-            "disabled": subdir.name.endswith(".disabled"),
-        }
-
-    for subdir in Path(custom_nodes).iterdir():
-        if not subdir.is_dir() or not subdir.joinpath(".git").exists():
-            continue
-        coros.append(get_node_info(subdir))
-
-    return await asyncio.gather(*coros)
-
-
 async def _write_workflow(path: ZPath, data: dict) -> None:
     print("Package => Writing workflow")
     with path.joinpath("workflow_api.json").open("w") as f:
@@ -199,15 +214,16 @@ async def _write_inputs(path: ZPath, data: dict) -> None:
 
     input_dir = folder_paths.get_input_directory()
 
-    used_inputs = set()
-    for node in data["workflow_api"].values():
-        for _, v in node["inputs"].items():
-            if isinstance(v, str):
-                used_inputs.add(v)
+    if "files" in data:
+        selected = "\n".join(set(data.get("files", [])))
+    else:
+        selected = None
 
     src_root = Path(input_dir).absolute()
     for src in src_root.glob("**/*"):
         rel = src.relative_to(src_root)
+        if selected is not None and str(rel) not in selected:
+            continue
         if src.is_dir():
             if isinstance(path, Path):
                 path.joinpath("input").joinpath(rel).mkdir(parents=True, exist_ok=True)
@@ -230,22 +246,264 @@ async def pack_workspace(request):
 
     with zipfile.ZipFile(TEMP_FOLDER / zip_filename, "w") as zf:
         path = zipfile.Path(zf)
-        await _write_requirements(path)
-        await _write_snapshot(path, data)
-        await _write_workflow(path, data)
-        await _write_inputs(path, data)
+        await _prepare_pack(path, data)
 
     return web.json_response({"download_url": f"/bentoml/download/{zip_filename}"})
 
 
+class DevServer:
+    TIMEOUT = 3600 * 24
+    proc: Union[None, subprocess.Popen] = None
+    watch_dog_task: asyncio.Task | None = None
+    last_feed = 0
+    run_dir: Path | None = None
+    port = 0
+
+    @classmethod
+    def start(cls, workflow_api: dict, port: int = 3000):
+        from comfy_pack import __file__ as comfy_pack_file
+
+        cls.stop()
+
+        cls.port = port
+        # prepare a temporary directory
+        cls.run_dir = Path(tempfile.mkdtemp(suffix="-bento", prefix="comfy-pack-"))
+        with cls.run_dir.joinpath("workflow_api.json").open("w") as f:
+            f.write(json.dumps(workflow_api, indent=2))
+        shutil.copy(
+            Path(comfy_pack_file).with_name("service.py"),
+            cls.run_dir / "service.py",
+        )
+        shutil.copytree(COMFY_PACK_DIR, cls.run_dir / COMFY_PACK_DIR.name)
+
+        # find a free port
+        self_port = 8188
+        for i, arg in enumerate(sys.argv):
+            if arg == "--port" or arg == "-p":
+                self_port = int(sys.argv[i + 1])
+                break
+
+        print(f"Starting dev server at port {port}, comfyui at port {self_port}")
+        cls.proc = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "bentoml",
+                "serve",
+                "service:ComfyService",
+                "--port",
+                str(port),
+            ],
+            cwd=str(cls.run_dir.absolute()),
+            env={
+                **os.environ,
+                "COMFYUI_SERVER": f"localhost:{self_port}",
+            },
+        )
+        cls.watch_dog_task = asyncio.create_task(cls.watch_dog())
+        cls.last_feed = time.time()
+
+    @classmethod
+    async def watch_dog(cls):
+        while True:
+            await asyncio.sleep(0.1)
+            if cls.last_feed + cls.TIMEOUT < time.time():
+                cls.stop()
+                break
+
+    @classmethod
+    def feed_watch_dog(cls):
+        if cls.proc:
+            if cls.proc.poll() is None:
+                cls.last_feed = time.time()
+                return True
+            else:
+                cls.stop()
+                return False
+        return False
+
+    @classmethod
+    def stop(cls):
+        if cls.proc:
+            cls.proc.terminate()
+            cls.proc.wait()
+            cls.proc = None
+            time.sleep(1)
+            print("Dev server stopped")
+        if cls.watch_dog_task:
+            cls.watch_dog_task.cancel()
+            cls.watch_dog_task = None
+            cls.last_feed = 0
+        if cls.run_dir:
+            shutil.rmtree(cls.run_dir)
+            cls.run_dir = None
+
+
+def _parse_workflow(workflow: dict) -> tuple[dict[str, Any], dict[str, Any]]:
+    inputs = {}
+    outputs = {}
+    dep_map = {}
+
+    for id, node in workflow.items():
+        for input_name, v in node["inputs"].items():
+            if isinstance(v, list) and len(v) == 2:  # is a link
+                dep_map[tuple(v)] = node, input_name
+
+    for id, node in workflow.items():
+        node["id"] = id
+        if node["class_type"].startswith("CPackInput"):
+            if not node.get("inputs"):
+                continue
+            inputs[id] = node
+        elif node["class_type"].startswith("CPackOutput"):
+            if not node.get("inputs"):
+                continue
+            outputs[id] = node
+
+    return inputs, outputs
+
+
+def _validate_workflow(data: dict):
+    workflow = data.get("workflow_api", {})
+    if not workflow:
+        return web.json_response(
+            {
+                "result": "error",
+                "error": "empty workflow",
+            },
+        )
+    input_spec, output_spec = _parse_workflow(workflow)
+    if not input_spec:
+        return web.json_response(
+            {
+                "result": "error",
+                "error": "At least one ComfyPack input node is required",
+            },
+        )
+    if not output_spec:
+        return web.json_response(
+            {
+                "result": "error",
+                "error": "At least one ComfyPack output node is required",
+            },
+        )
+
+
+@PromptServer.instance.routes.post("/bentoml/serve")
+async def serve(request):
+    data = await request.json()
+
+    if (error := _validate_workflow(data)) is not None:
+        return error
+
+    DevServer.stop()
+
+    if _is_port_in_use(data.get("port", 3000), host=data.get("host", "localhost")):
+        return web.json_response(
+            {
+                "result": "error",
+                "error": "Port is already in use",
+            },
+        )
+    try:
+        DevServer.start(workflow_api=data["workflow_api"], port=data.get("port", 3000))
+        return web.json_response(
+            {
+                "result": "success",
+                "url": f"http://{data.get('host', 'localhost')}:{data.get('port', 3000)}",
+            },
+        )
+    except Exception as e:
+        return web.json_response(
+            {
+                "result": "error",
+                "error": f"Build failed: {e.__class__.__name__}: {e}",
+            },
+        )
+
+
+@PromptServer.instance.routes.get("/bentoml/serve/heartbeat")
+async def heartbeat(_):
+    running = DevServer.feed_watch_dog()
+
+    if running:
+        return web.json_response({"ready": True})
+    else:
+        return web.json_response({"error": "Server is not running"})
+
+
+@PromptServer.instance.routes.post("/bentoml/serve/terminate")
+async def terminate(_):
+    DevServer.stop()
+    return web.json_response({"result": "success"})
+
+
 @PromptServer.instance.routes.get("/bentoml/download/{zip_filename}")
 async def download_workspace(request):
     zip_filename = request.match_info["zip_filename"]
     return web.FileResponse(TEMP_FOLDER / zip_filename)
 
 
+async def _prepare_pack(
+    working_dir: ZPath,
+    data: dict,
+    store_models: bool = False,
+    ensure_source: bool = True,
+) -> None:
+    model_filter = set(data.get("models", []))
+    models = await _get_models(
+        store_models=store_models,
+        model_filter=model_filter,
+        ensure_source=ensure_source,
+    )
+
+    await _write_snapshot(working_dir, data, models)
+    await _write_workflow(working_dir, data)
+    await _write_inputs(working_dir, data)
+
+
+@PromptServer.instance.routes.post("/bentoml/model/query")
+async def get_models(request):
+    data = await request.json()
+    models = await _get_models(
+        workflow_api=data.get("workflow_api"),
+        ensure_sha=False,
+        ensure_source=False,
+    )
+    return web.json_response({"models": models})
+
+
+async def _get_inputs(workflow_api):
+    input_dir = folder_paths.get_input_directory()
+    inputs = []
+    for src in Path(input_dir).rglob("*"):
+        if src.is_file():
+            rel = src.relative_to(input_dir)
+            badges = []
+            checked = False
+            if _is_file_refered(src, workflow_api):
+                badges.append({"text": "Referenced"})
+                checked = True
+            data = {
+                "path": str(rel),
+                "badges": badges,
+                "checked": checked,
+            }
+            inputs.append(data)
+    return inputs
+
+
+@PromptServer.instance.routes.post("/bentoml/file/query")
+async def get_inputs(request):
+    data = await request.json()
+    inputs = await _get_inputs(
+        workflow_api=data.get("workflow_api"),
+    )
+    return web.json_response({"files": inputs})
+
+
 @PromptServer.instance.routes.post("/bentoml/build")
-async def build_bento(request):
+async def build_bento_api(request):
     """Request body: {
         workflow_api: dict,
         workflow: dict,
@@ -259,42 +517,19 @@ async def build_bento(request):
 
     data = await request.json()
 
+    if (error := _validate_workflow(data)) is not None:
+        return error
+
     with tempfile.TemporaryDirectory(suffix="-bento", prefix="comfy-pack-") as temp_dir:
         temp_dir_path = Path(temp_dir)
-        # copy comfy_pack source code into the bento
-        shutil.copytree(COMFY_PACK_DIR, temp_dir_path / COMFY_PACK_DIR.name)
-        models = await _get_models(data, store_models=True)
-
-        await _write_requirements(temp_dir_path, ["comfy-cli", "fastapi"])
-        await _write_snapshot(temp_dir_path, data, models)
-        await _write_workflow(temp_dir_path, data)
-        await _write_inputs(temp_dir_path, data)
-        shutil.copy(
-            Path(__file__).with_name("service.py"),
-            temp_dir_path / "service.py",
-        )
+        await _prepare_pack(temp_dir_path, data, store_models=True, ensure_source=False)
 
         # create a bento
         try:
-            bento = bentoml.build(
-                "service:ComfyService",
-                name=data["bento_name"],
-                build_ctx=temp_dir,
-                models=[m["model_tag"] for m in models if "model_tag" in m],
-                docker={
-                    "python_version": f"{sys.version_info.major}.{sys.version_info.minor}",
-                    "system_packages": [
-                        "git",
-                        "libglib2.0-0",
-                        "libsm6",
-                        "libxrender1",
-                        "libxext6",
-                        "ffmpeg",
-                        "libstdc++-12-dev",
-                        *data.get("system_packages", []),
-                    ],
-                },
-                python={"requirements_txt": "requirements.txt", "lock_packages": True},
+            bento = build_bento(
+                data["bento_name"],
+                temp_dir_path,
+                system_packages=data.get("system_packages"),
             )
         except bentoml.exceptions.BentoMLException as e:
             return web.json_response(

nodes/nodes.py

@@ -1,14 +1,18 @@
+import glob
 import hashlib
+import io
 import json
 import os
-import sys
 import shutil
+import sys
+import zipfile
+from io import BytesIO
 
 import folder_paths
 import node_helpers
 import numpy as np
 import torch
-from PIL import Image, ImageOps, ImageSequence
+from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
 from PIL.PngImagePlugin import PngInfo
 
 from .monkeypatch import set_bentoml_output
@@ -163,6 +167,100 @@ class OutputImage:
         return {"ui": {"images": results}}
 
 
+class OutputImageWithStringTxt:
+    COLOR = (142, 36, 170)
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+        self.compress_level = 4
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", {"tooltip": "The images to save."}),
+                "filename_prefix": ("STRING", {"default": "cpack_output_"}),
+                "text": ("STRING", {"default": ""}),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "extra_pnginfo": "EXTRA_PNGINFO",
+            },
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "save_images"
+    CPACK_NODE = True
+    OUTPUT_NODE = True
+
+    CATEGORY = "ComfyPack/output"
+    DESCRIPTION = (
+        "Saves the input images (and optional text) to your ComfyUI output directory."
+    )
+
+    def save_images(
+        self,
+        images,
+        filename_prefix="cpack_output_",
+        text="",
+        prompt=None,
+        extra_pnginfo=None,
+    ):
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            get_save_image_path(
+                filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
+            )
+        )
+
+        base_counter = counter  # use for name zip
+        zip_filename = f"{filename}_batch_{base_counter:05}.zip"
+        zip_path = os.path.join(full_output_folder, zip_filename)
+
+        # create ZIP file
+        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+            for batch_number, image in enumerate(images):
+                # temp store img to RAM
+                i = 255.0 * image.cpu().numpy()
+                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+
+                # write meta data
+                metadata = PngImagePlugin.PngInfo()
+                if prompt is not None:
+                    metadata.add_text("prompt", json.dumps(prompt))
+                if extra_pnginfo is not None:
+                    for x in extra_pnginfo:
+                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
+
+                # write img file to RAM buffer
+                img_buffer = BytesIO()
+                img.save(
+                    img_buffer,
+                    format="PNG",
+                    pnginfo=metadata,
+                    compress_level=self.compress_level,
+                )
+                img_buffer.seek(0)
+
+                # write img into ZIP file
+                image_filename = f"image_{batch_number:05}.png"
+                zipf.writestr(image_filename, img_buffer.read())
+
+                # write txt into ZIP file
+                text_filename = f"text_{batch_number:05}.txt"
+                zipf.writestr(text_filename, text)
+
+        # return zip as output
+        out = [{"filename": zip_filename, "subfolder": subfolder, "type": "zip"}]
+        return {
+            "ui": {
+                "zip": out,
+            }
+        }
+
+
 class ImageInput:
     COLOR = (142, 36, 170)
 
@@ -353,14 +451,299 @@ class AnyInput:
         return True
 
 
+class OutputZip:
+    CATEGORY = "ComfyPack/output"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {}
+
+    RETURN_TYPES = ()
+    FUNCTION = "null_op"
+
+    def null_op(self):
+        return ()
+
+
+class OutputAudio:
+    CPACK_NODE = True
+    CATEGORY = "ComfyPack/output"
+    RETURN_TYPES = ()
+    FUNCTION = "save_audio"
+    OUTPUT_NODE = True
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    def save_audio(
+        self,
+        audio,
+        filename_prefix,
+        format="flac",
+        prompt=None,
+        extra_pnginfo=None,
+        quality="128k",
+    ):
+        import av
+        import torchaudio
+
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            get_save_image_path(filename_prefix, self.output_dir)
+        )
+        results = []
+
+        # Prepare metadata dictionary
+        metadata = {}
+        if prompt is not None:
+            metadata["prompt"] = json.dumps(prompt)
+        if extra_pnginfo is not None:
+            for x in extra_pnginfo:
+                metadata[x] = json.dumps(extra_pnginfo[x])
+
+        # Opus supported sample rates
+        OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
+
+        for batch_number, waveform in enumerate(audio["waveform"].cpu()):
+            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+            file = f"{filename_with_batch_num}_{counter:05}.{format}"
+            output_path = os.path.join(full_output_folder, file)
+
+            # Use original sample rate initially
+            sample_rate = audio["sample_rate"]
+
+            # Handle Opus sample rate requirements
+            if format == "opus":
+                if sample_rate > 48000:
+                    sample_rate = 48000
+                elif sample_rate not in OPUS_RATES:
+                    # Find the next highest supported rate
+                    for rate in sorted(OPUS_RATES):
+                        if rate > sample_rate:
+                            sample_rate = rate
+                            break
+                    if sample_rate not in OPUS_RATES:  # Fallback if still not supported
+                        sample_rate = 48000
+
+                # Resample if necessary
+                if sample_rate != audio["sample_rate"]:
+                    waveform = torchaudio.functional.resample(
+                        waveform, audio["sample_rate"], sample_rate
+                    )
+
+            # Create in-memory WAV buffer
+            wav_buffer = io.BytesIO()
+            torchaudio.save(wav_buffer, waveform, sample_rate, format="WAV")
+            wav_buffer.seek(0)  # Rewind for reading
+
+            # Use PyAV to convert and add metadata
+            input_container = av.open(wav_buffer)
+
+            # Create output with specified format
+            output_buffer = io.BytesIO()
+            output_container = av.open(output_buffer, mode="w", format=format)
+
+            # Set metadata on the container
+            for key, value in metadata.items():
+                output_container.metadata[key] = value
+
+            # Set up the output stream with appropriate properties
+            input_container.streams.audio[0]
+            if format == "opus":
+                out_stream = output_container.add_stream("libopus", rate=sample_rate)
+                if quality == "64k":
+                    out_stream.bit_rate = 64000
+                elif quality == "96k":
+                    out_stream.bit_rate = 96000
+                elif quality == "128k":
+                    out_stream.bit_rate = 128000
+                elif quality == "192k":
+                    out_stream.bit_rate = 192000
+                elif quality == "320k":
+                    out_stream.bit_rate = 320000
+            elif format == "mp3":
+                out_stream = output_container.add_stream("libmp3lame", rate=sample_rate)
+                if quality == "V0":
+                    # TODO i would really love to support V3 and V5 but there doesn't seem to be a way to set the qscale level, the property below is a bool
+                    out_stream.codec_context.qscale = 1
+                elif quality == "128k":
+                    out_stream.bit_rate = 128000
+                elif quality == "320k":
+                    out_stream.bit_rate = 320000
+            else:  # format == "flac":
+                out_stream = output_container.add_stream("flac", rate=sample_rate)
+
+            # Copy frames from input to output
+            for frame in input_container.decode(audio=0):
+                frame.pts = None  # Let PyAV handle timestamps
+                output_container.mux(out_stream.encode(frame))
+
+            # Flush encoder
+            output_container.mux(out_stream.encode(None))
+
+            # Close containers
+            output_container.close()
+            input_container.close()
+
+            # Write the output to file
+            output_buffer.seek(0)
+            with open(output_path, "wb") as f:
+                f.write(output_buffer.getbuffer())
+
+            results.append(
+                {"filename": file, "subfolder": subfolder, "type": self.type}
+            )
+            counter += 1
+
+        return {"ui": {"audio": results}}
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "audio": ("AUDIO",),
+                "filename_prefix": ("STRING", {"default": "audio/comfypack"}),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+
+class OutputVideo:
+    CPACK_NODE = True
+    CATEGORY = "ComfyPack/output"
+
+    RETURN_TYPES = ()
+    FUNCTION = "save_video"
+    DESCRIPTION = "Saves the input video to your ComfyUI output directory."
+    OUTPUT_NODE = True
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video": ("VIDEO", {"tooltip": "The video to save."}),
+                "filename_prefix": (
+                    "STRING",
+                    {
+                        "default": "video/comfypack",
+                        "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes.",
+                    },
+                ),
+                "format": (
+                    ["auto", "mp4"],
+                    {"default": "auto", "tooltip": "The format to save the video as."},
+                ),
+                "codec": (
+                    ["auto", "h264"],
+                    {"default": "auto", "tooltip": "The codec to use for the video."},
+                ),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    def save_video(
+        self,
+        video,
+        filename_prefix,
+        format,
+        codec,
+        prompt=None,
+        extra_pnginfo=None,
+    ):
+        filename_prefix += self.prefix_append
+        width, height = video.get_dimensions()
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            get_save_image_path(filename_prefix, self.output_dir, width, height)
+        )
+        results = list()
+        saved_metadata = None
+        metadata = {}
+        if extra_pnginfo is not None:
+            metadata.update(extra_pnginfo)
+        if prompt is not None:
+            metadata["prompt"] = prompt
+        if len(metadata) > 0:
+            saved_metadata = metadata
+        file = f"{filename}_{counter:05}.mp4"
+        video.save_to(
+            os.path.join(full_output_folder, file),
+            format=format,
+            codec=codec,
+            metadata=saved_metadata,
+        )
+
+        results.append({"filename": file, "subfolder": subfolder, "type": self.type})
+        counter += 1
+
+        return {"ui": {"images": results, "animated": (True,)}}
+
+
+class OutputTextFile:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING", {"forceInput": True}),
+                "filename_prefix": ("STRING", {"default": "cpack_output_"}),
+            },
+            "optional": {
+                "file_extension": ("STRING", {"default": ".txt"}),
+            },
+        }
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+    FUNCTION = "save_text_file"
+    CATEGORY = "ComfyPack/output"
+    CPACK_NODE = True
+
+    def save_text_file(
+        self, text: str, filename_prefix: str, file_extension: str = ".txt"
+    ):
+        subfolder, filename_prefix = os.path.split(os.path.normpath(filename_prefix))
+        output_dir = folder_paths.get_output_directory()
+        full_output_folder = os.path.join(output_dir, subfolder)
+
+        full_output_filename = self.get_output_filename(
+            full_output_folder, filename_prefix, file_extension
+        )
+        with open(full_output_filename, "w", encoding="utf-8", newline="\n") as f:
+            f.write(text)
+        return (text, {"ui": {"string": text}})
+
+    @staticmethod
+    def get_output_filename(folder: str, prefix: str, extension: str) -> str:
+        matched_files = [
+            os.path.basename(f)[len(prefix) + 1 : -len(extension)]
+            for f in glob.glob(os.path.join(folder, f"{prefix}_*{extension}"))
+        ]
+        print("MATCHING", matched_files)
+        max_count = max(
+            (int(name) for name in matched_files if name.isdigit()), default=0
+        )
+        return os.path.join(folder, f"{prefix}_{max_count + 1:04d}{extension}")
+
+
 NODE_CLASS_MAPPINGS = {
     "CPackOutputFile": OutputFile,
     "CPackOutputImage": OutputImage,
+    "CPackOutputAudio": OutputAudio,
+    "CPackOutputVideo": OutputVideo,
+    "CPackOutputZip": OutputImageWithStringTxt,
+    "CPackOutputZipSwitch": OutputZip,
     "CPackInputImage": ImageInput,
     "CPackInputString": StringInput,
     "CPackInputInt": IntInput,
     "CPackInputFile": FileInput,
     "CPackInputAny": AnyInput,
+    "CPackOutputTextFile": OutputTextFile,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -370,5 +753,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CPackInputFile": "File Input",
     "CPackInputAny": "Any Input",
     "CPackOutputImage": "Image Output",
+    "CPackOutputAudio": "Audio Output",
+    "CPackOutputVideo": "Video Output",
     "CPackOutputFile": "File Output",
+    "CPackOutputZip": "Zip Output(img + txt file)",
+    "CPackOutputZipSwitch": "Enable Zip Output",
+    "CPackOutputTextFile": "Output Text to File",
 }

nodes/service.py

@@ -1,121 +0,0 @@
-from __future__ import annotations
-
-import json
-import logging
-import os
-import shutil
-from functools import lru_cache
-from pathlib import Path
-from typing import Any, cast
-
-import bentoml
-import fastapi
-
-import comfy_pack
-import comfy_pack.run
-
-REQUEST_TIMEOUT = 360
-BASE_DIR = Path(__file__).parent
-WORKFLOW_FILE = BASE_DIR / "workflow_api.json"
-COPY_THRESHOLD = 10 * 1024 * 1024
-INPUT_DIR = BASE_DIR / "input"
-logger = logging.getLogger("bentoml.service")
-
-with open(WORKFLOW_FILE, "r") as f:
-    workflow = json.load(f)
-
-InputModel = comfy_pack.generate_input_model(workflow)
-app = fastapi.FastAPI()
-
-
-@lru_cache
-def _get_workspace() -> Path:
-    import hashlib
-
-    from bentoml._internal.configuration.containers import BentoMLContainer
-
-    snapshot = BASE_DIR / "snapshot.json"
-    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
-    wp = Path(BentoMLContainer.bentoml_home.get()) / "run/comfy_workspace" / checksum
-    wp.parent.mkdir(parents=True, exist_ok=True)
-    return wp
-
-
-@app.get("/workflow.json")
-def workflow_json():
-    return workflow
-
-
-@bentoml.mount_asgi_app(app, path="/comfy")
-@bentoml.service(traffic={"timeout": REQUEST_TIMEOUT * 2})
-class ComfyService:
-    def __init__(self):
-        self.comfy_proc = comfy_pack.run.WorkflowRunner(
-            str(_get_workspace()),
-            str(INPUT_DIR),
-        )
-        self.comfy_proc.start(verbose=int("BENTOML_DEBUG" in os.environ))
-
-    @bentoml.api(input_spec=InputModel)
-    def generate(
-        self,
-        *,
-        ctx: bentoml.Context,
-        **kwargs: Any,
-    ) -> Path:
-        verbose = int("BENTOML_DEBUG" in os.environ)
-        ret = self.comfy_proc.run_workflow(
-            workflow,
-            output_dir=ctx.temp_dir,
-            timeout=REQUEST_TIMEOUT,
-            verbose=verbose,
-            **kwargs,
-        )
-        if isinstance(ret, list):
-            ret = ret[-1]
-        return ret
-
-    @bentoml.on_shutdown
-    def on_shutdown(self):
-        self.comfy_proc.stop()
-
-    @bentoml.on_deployment
-    @staticmethod
-    def prepare_comfy_workspace():
-        from comfy_pack.package import install_comfyui, install_custom_modules
-
-        verbose = int("BENTOML_DEBUG" in os.environ)
-        comfy_workspace = _get_workspace()
-
-        with BASE_DIR.joinpath("snapshot.json").open("rb") as f:
-            snapshot = json.load(f)
-
-        if not comfy_workspace.joinpath(".DONE").exists():
-            if comfy_workspace.exists():
-                logger.info("Removing existing workspace")
-                shutil.rmtree(comfy_workspace, ignore_errors=True)
-            install_comfyui(snapshot, comfy_workspace, verbose=verbose)
-
-            for model in snapshot["models"]:
-                model_tag = model.get("model_tag")
-                if not model_tag:
-                    logger.warning(
-                        "Model %s is not in model store, the workflow may not work",
-                        model["filename"],
-                    )
-                    continue
-                model_path = comfy_workspace / cast(str, model["filename"])
-                model_path.parent.mkdir(parents=True, exist_ok=True)
-                bento_model = bentoml.models.get(model_tag)
-                model_file = bento_model.path_of("model.bin")
-                logger.info("Copying %s to %s", model_file, model_path)
-                model_path.symlink_to(model_file)
-
-            for f in INPUT_DIR.glob("*"):
-                if f.is_file():
-                    shutil.copy(f, comfy_workspace / "input" / f.name)
-                elif f.is_dir():
-                    shutil.copytree(f, comfy_workspace / "input" / f.name)
-
-            install_custom_modules(snapshot, comfy_workspace, verbose=verbose)
-            comfy_workspace.joinpath(".DONE").touch()

pyproject.toml

@@ -1,14 +1,16 @@
 [project]
 name = "comfy-pack"
-description = "ComfyUI Interface Definition Language"
+description = """\
+    A comprehensive toolkit for standardizing, packaging and deploying ComfyUI workflows \
+    as reproducible environments and production-ready REST services\
+"""
 authors = [{ name = "Frost Ming", email = "frost@bentoml.com" }]
 readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [
-    "bentoml>=1.3.13",
+    "bentoml>=1.3.21",
     "click>=8.1.7",
     "comfy-cli>=1.2.8",
-    "pydantic>=2.9",
 ]
 dynamic = ["version"]
 
@@ -19,7 +21,7 @@ classifiers = [
 ]
 
 [project.urls]
-Homepage = "https://github.com/bentoml/ComfyUI-IDL"
+Homepage = "https://github.com/bentoml/comfy-pack"
 
 [project.scripts]
 comfy-pack = "comfy_pack.cli:main"

requirements.txt

@@ -1,2 +1,4 @@
-bentoml
+bentoml>=1.3.21
 fastapi
+comfy-cli
+duckduckgo-search

src/comfy_pack/__init__.py

@@ -1,4 +1,4 @@
-from .run import WorkflowRunner
+from .run import ComfyUIServer, run_workflow
 from .utils import (
     generate_input_model,
     parse_workflow,
@@ -7,7 +7,8 @@ from .utils import (
 )
 
 __all__ = [
-    "WorkflowRunner",
+    "ComfyUIServer",
+    "run_workflow",
     "parse_workflow",
     "generate_input_model",
     "populate_workflow",

src/comfy_pack/cli.py

@@ -1,31 +1,247 @@
-import click
-import hashlib
 import functools
 import json
-from pathlib import Path
+import os
 import shutil
+import subprocess
+import sys
 import tempfile
+import zipfile
+from pathlib import Path
+
+import click
+
+from .const import COMFY_PACK_REPO, COMFYUI_MANAGER_REPO, COMFYUI_REPO, WORKSPACE_DIR
+from .hash import get_sha256
+from .utils import get_self_git_commit
+
+
+def _ensure_uv() -> None:
+    """Ensure uv is installed, raise error if not."""
+    try:
+        subprocess.run(
+            ["uv", "--version"],
+            check=True,
+            capture_output=True,
+        )
+    except (subprocess.SubprocessError, FileNotFoundError):
+        raise RuntimeError(
+            "uv is not installed. Please install it first:\n"
+            "curl -LsSf https://astral.sh/uv/install.sh | sh"
+        )
 
 
 @click.group()
+@click.version_option()
+def main():
+    """comfy-pack CLI"""
+    pass
+
+
+@main.command(
+    name="init",
+    help="Install latest ComfyUI and comfy-pack custom nodes and create a virtual environment",
+)
+@click.option(
+    "--dir",
+    "-d",
+    default="ComfyUI",
+    help="Target directory to install ComfyUI",
+    type=click.Path(file_okay=False),
+)
 @click.option(
     "--verbose",
     "-v",
     count=True,
-    help="Increase verbosity level (use multiple times for more verbosity)",
+    help="Increase verbosity level",
 )
-@click.pass_context
-def main(ctx, verbose):
-    """ComfyUI Pack CLI"""
-    ctx.ensure_object(dict)
-    ctx.obj["verbose"] = verbose
+def init(dir: str, verbose: int):
+    import os
+
+    import rich
+
+    # Check if directory path is valid
+    try:
+        install_dir = Path(dir).absolute()
+        if install_dir.exists() and not install_dir.is_dir():
+            rich.print(f"[red]Error: {dir} exists but is not a directory[/red]")
+            return 1
+
+        # Check if directory is empty or contains ComfyUI
+        if install_dir.exists():
+            contents = list(install_dir.iterdir())
+            if contents and not (install_dir / ".git").exists():
+                rich.print(
+                    f"[red]Error: Directory {dir} is not empty and doesn't appear to be a ComfyUI installation[/red]"
+                )
+                return 1
+    except Exception as e:
+        rich.print(f"[red]Error: Invalid directory path - {str(e)}[/red]")
+        return 1
+
+    # Check git installation
+    try:
+        subprocess.run(
+            ["git", "--version"],
+            check=True,
+            capture_output=True,
+        )
+    except (subprocess.SubprocessError, FileNotFoundError):
+        rich.print("[red]Error: git is not installed or not in PATH[/red]")
+        return 1
+
+    # Check if we have write permissions
+    try:
+        if not install_dir.exists():
+            install_dir.mkdir(parents=True)
+        test_file = install_dir / ".write_test"
+        test_file.touch()
+        test_file.unlink()
+    except (OSError, PermissionError) as e:
+        rich.print(f"[red]Error: No write permission in {dir} - {str(e)}[/red]")
+        return 1
+
+    # Check if Python version is compatible
+    if sys.version_info < (3, 8):
+        rich.print("[red]Error: Python 3.8 or higher is required[/red]")
+        return 1
+
+    # Check if uv is installed
+    try:
+        _ensure_uv()
+    except RuntimeError as e:
+        rich.print(f"[red]Error: {str(e)}[/red]")
+        return 1
+
+    # Check if enough disk space is available (rough estimate: 2GB)
+    try:
+        free_space = shutil.disk_usage(install_dir).free
+        if free_space < 2 * 1024 * 1024 * 1024:  # 2GB in bytes
+            rich.print(
+                "[yellow]Warning: Less than 2GB free disk space available[/yellow]"
+            )
+    except Exception as e:
+        rich.print(
+            f"[yellow]Warning: Could not check free disk space - {str(e)}[/yellow]"
+        )
+
+    # Clone ComfyUI if not exists
+    if not (install_dir / ".git").exists():
+        rich.print("[green]Cloning ComfyUI...[/green]")
+        subprocess.run(
+            [
+                "git",
+                "clone",
+                COMFYUI_REPO,
+                str(install_dir),
+            ],
+            check=True,
+        )
+
+    # Update ComfyUI
+    rich.print("[green]Updating ComfyUI...[/green]")
+    subprocess.run(
+        ["git", "pull"],
+        cwd=install_dir,
+        check=True,
+    )
+
+    # Create and activate venv
+    venv_dir = install_dir / ".venv"
+    rich.print("[green]Creating virtual environment with uv...[/green]")
+    if venv_dir.exists():
+        shutil.rmtree(venv_dir)
+    subprocess.run(
+        ["uv", "venv", str(venv_dir)],
+        check=True,
+    )
+
+    # Get python path for future use
+    if sys.platform == "win32":
+        python = str(venv_dir / "Scripts" / "python.exe")
+
+    else:
+        python = str(venv_dir / "bin" / "python")
+
+    # Install requirements with uv
+    rich.print("[green]Installing ComfyUI requirements with uv...[/green]")
+    subprocess.run(
+        ["uv", "pip", "install", "pip", "--upgrade"],
+        env={
+            "VIRTUAL_ENV": str(venv_dir),
+            "PATH": str(venv_dir / "bin") + os.pathsep + os.environ["PATH"],
+        },
+        check=True,
+    )
+    subprocess.run(
+        ["uv", "pip", "install", "-r", str(install_dir / "requirements.txt")],
+        env={
+            "VIRTUAL_ENV": str(venv_dir),
+            "PATH": str(venv_dir / "bin") + os.pathsep + os.environ["PATH"],
+        },
+        check=True,
+    )
+
+    # Install comfy-pack as custom node
+    rich.print("[green]Installing comfy-pack custom nodes...[/green]")
+    custom_nodes_dir = install_dir / "custom_nodes"
+    custom_nodes_dir.mkdir(exist_ok=True)
+
+    comfyui_manager_dir = custom_nodes_dir / "ComfyUI-Manager"
+    if not (comfyui_manager_dir / ".git").exists():
+        # Clone ComfyUI-Manager
+        subprocess.run(
+            ["git", "clone", COMFYUI_MANAGER_REPO, str(comfyui_manager_dir)],
+            check=True,
+        )
+
+    comfy_pack_dir = custom_nodes_dir / "comfy-pack"
+    if not (comfy_pack_dir / ".git").exists():
+        # Clone comfy-pack
+        subprocess.run(
+            ["git", "clone", COMFY_PACK_REPO, str(comfy_pack_dir)],
+            check=True,
+        )
+
+    # Update comfy-pack
+    subprocess.run(
+        ["git", "pull"],
+        cwd=comfy_pack_dir,
+        check=True,
+    )
+
+    # Install comfy-pack requirements
+    if (comfy_pack_dir / "requirements.txt").exists():
+        subprocess.run(
+            [
+                python,
+                "-m",
+                "pip",
+                "install",
+                "-r",
+                str(comfy_pack_dir / "requirements.txt"),
+            ],
+            check=True,
+        )
+
+    version = get_self_git_commit() or "unknown"
+    rich.print(
+        f"\n[green]✓ Installation completed! (comfy-pack version: {version})[/green]"
+    )
+    rich.print(f"ComfyUI directory: {install_dir}")
+
+    rich.print(
+        "\n[green]Next steps:[/green]\n"
+        f"1. cd {dir}\n"
+        "2. source .venv/bin/activate  # On Windows: .venv\\Scripts\\activate\n"
+        "3. python main.py"
+    )
 
 
 @main.command(
-    name="restore",
+    name="unpack",
     help="Restore the ComfyUI workspace to specified directory",
 )
-@click.argument("cpack", type=click.Path(exists=True, dir_okay=False))
+@click.argument("cpack", type=click.Path(exists=True))
 @click.option(
     "--dir",
     "-d",
@@ -33,28 +249,76 @@ def main(ctx, verbose):
     help="target directory to restore the ComfyUI project",
     type=click.Path(file_okay=False),
 )
-@click.pass_context
-def restore_cmd(ctx, cpack: str, dir: str):
+@click.option(
+    "--include-disabled-models",
+    default=False,
+    type=click.BOOL,
+    is_flag=True,
+)
+@click.option(
+    "--no-models",
+    default=False,
+    type=click.BOOL,
+    is_flag=True,
+    help="Do not install models",
+)
+@click.option(
+    "--no-venv",
+    is_flag=True,
+    help="Do not create a virtual environment for ComfyUI",
+    default=False,
+)
+@click.option(
+    "--verbose",
+    "-v",
+    count=True,
+    help="Increase verbosity level (use multiple times for more verbosity)",
+)
+@click.option(
+    "--preheat",
+    is_flag=True,
+    help="Preheat the workspace after unpacking",
+    default=False,
+)
+def unpack_cmd(
+    cpack: str,
+    dir: str,
+    include_disabled_models: bool,
+    no_models: bool,
+    no_venv: bool,
+    verbose: int,
+    preheat: bool,
+):
+    import rich
+
     from .package import install
-    from rich.console import Console
 
-    console = Console()
-
-    install(cpack, dir, ctx.obj["verbose"])
-    console.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
-    console.print(f"{dir}")
-
-    console.print(
-        "\n[green] Next steps: [/green]\n"
-        "1. Change directory to the restored workspace\n"
-        "2. Source the virtual environment by running `source .venv/bin/activate`\n"
-        "3. Run the ComfyUI project by running `python main.py`"
+    install(
+        cpack,
+        dir,
+        verbose=verbose,
+        all_models=include_disabled_models,
+        prepare_models=not no_models,
+        no_venv=no_venv,
+        preheat=preheat,
     )
+    rich.print("\n[green]✓ ComfyUI Workspace is restored at:[/green]")
+    rich.print(os.path.abspath(dir))
+    steps = [f"Change directory to the restored workspace: `cd {dir}`"]
+    if not no_venv:
+        steps.append(
+            "Source the virtual environment by running `source .venv/bin/activate`"
+        )
+    steps.append("Run the ComfyUI project by running `python main.py`")
+
+    rich.print(f"\n[green]⏭️ Next steps: [/green]\n1. {steps[0]}\n2. {steps[1]}")
+    if len(steps) > 2:
+        rich.print(f"3. {steps[2]}")
 
 
 def _print_schema(schema, verbose: int = 0):
+    import rich
     from rich.table import Table
-    from rich.console import Console
 
     table = Table(title="")
 
@@ -84,36 +348,13 @@ def _print_schema(schema, verbose: int = 0):
             range_str,
         )
 
-    Console().print(table)
-
-
-@main.command(name="info")
-@click.argument("cpack", type=click.Path(exists=True, dir_okay=False))
-@click.pass_context
-def info_cmd(ctx, cpack: str):
-    """
-    Display information about the ComfyUI package.
-
-    Example:
-        $ comfy_pack info workspace.cpack.zip
-    """
-    from .utils import generate_input_model
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        pack_dir = Path(temp_dir) / ".cpack"
-        shutil.unpack_archive(cpack, pack_dir)
-        workflow = json.loads((pack_dir / "workflow_api.json").read_text())
-
-    inputs = generate_input_model(workflow)
-    _print_schema(inputs.model_json_schema(), ctx.obj["verbose"])
+    rich.print(table)
 
 
 @functools.lru_cache
 def _get_cache_workspace(cpack: str):
-    m = hashlib.sha256()
-    with open(cpack, "rb") as f:
-        m.update(f.read())
-    return Path.home() / ".comfypack" / "workspace" / m.hexdigest()[0:8]
+    sha = get_sha256(cpack)
+    return WORKSPACE_DIR / sha[0:8]
 
 
 @main.command(
@@ -122,76 +363,189 @@ def _get_cache_workspace(cpack: str):
         "allow_extra_args": True,
     },
     help="Run a ComfyUI package with the given inputs",
+    add_help_option=False,
 )
 @click.argument("cpack", type=click.Path(exists=True, dir_okay=False))
 @click.option("--output-dir", "-o", type=click.Path(), default=".")
+@click.option("--help", "-h", is_flag=True, help="Show this message and input schema")
+@click.option(
+    "--verbose",
+    "-v",
+    count=True,
+    help="Increase verbosity level (use multiple times for more verbosity)",
+)
 @click.pass_context
-def run(ctx, cpack: str, output_dir: str):
-    from .utils import generate_input_model
+def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
+    import rich
     from pydantic import ValidationError
-    from rich.console import Console
+
+    from .utils import generate_input_model
 
     inputs = dict(
         zip([k.lstrip("-").replace("-", "_") for k in ctx.args[::2]], ctx.args[1::2])
     )
 
-    console = Console()
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        pack_dir = Path(temp_dir) / ".cpack"
-        shutil.unpack_archive(cpack, pack_dir)
-        workflow = json.loads((pack_dir / "workflow_api.json").read_text())
+    with zipfile.ZipFile(cpack) as z:
+        workflow = json.loads(z.read("workflow_api.json"))
 
     input_model = generate_input_model(workflow)
 
+    # If help is requested, show command help and input schema
+    if help:
+        rich.print(
+            'Usage: comfy-pack run [OPTIONS] CPACK --input1 "value1" --input2 "value2" ...'
+        )
+        rich.print("Run a ComfyUI package with the given inputs:")
+        _print_schema(input_model.model_json_schema(), verbose)
+        return 0
+
     try:
         validated_data = input_model(**inputs)
-        console.print("[green]✓ Input is valid![/green]")
+        rich.print("[green]✓ Input is valid![/green]")
         for field, value in validated_data.model_dump().items():
-            console.print(f"{field}: {value}")
+            rich.print(f"{field}: {value}")
     except ValidationError as e:
-        console.print("[red]✗ Validation failed![/red]")
+        rich.print("[red]✗ Validation failed![/red]")
         for error in e.errors():
-            console.print(f"- {error['loc'][0]}: {error['msg']}")
+            rich.print(f"- {error['loc'][0]}: {error['msg']}")
+
+        rich.print("\n[yellow]Expected inputs:[/yellow]")
+        _print_schema(input_model.model_json_schema(), verbose)
         return 1
 
     from .package import install
 
     workspace = _get_cache_workspace(cpack)
     if not (workspace / "DONE").exists():
-        console.print("\n[green]✓ Restoring ComfyUI Workspace...[/green]")
+        rich.print("\n[green]✓ Restoring ComfyUI Workspace...[/green]")
         if workspace.exists():
             shutil.rmtree(workspace)
-        install(cpack, workspace)
+        install(cpack, workspace, verbose=verbose)
         with open(workspace / "DONE", "w") as f:
             f.write("DONE")
-    console.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
-    console.print(f"{workspace}")
+    rich.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
+    rich.print(f"{workspace}")
 
-    from .run import WorkflowRunner
+    from .run import ComfyUIServer, run_workflow
 
-    runner = None
-    try:
-        runner = WorkflowRunner(str(workspace.absolute()))
-        runner.start()
-        console.print("\n[green]✓ ComfyUI is launched in the background![/green]")
-        results = runner.run_workflow(
+    with ComfyUIServer(str(workspace.absolute()), verbose=verbose) as server:
+        rich.print("\n[green]✓ ComfyUI is launched in the background![/green]")
+        results = run_workflow(
+            server.host,
+            server.port,
             workflow,
             Path(output_dir).absolute(),
-            verbose=ctx.obj["verbose"],
+            verbose=verbose,
+            workspace=server.workspace,
             **validated_data.model_dump(),
         )
-        console.print("\n[green]✓ Workflow is executed successfully![/green]")
+        rich.print("\n[green]✓ Workflow is executed successfully![/green]")
         if results:
-            console.print("\n[green]✓ Retrieved outputs:[/green]")
+            rich.print("\n[green]✓ Retrieved outputs:[/green]")
         if isinstance(results, dict):
             for field, value in results.items():
-                console.print(f"{field}: {value}")
+                rich.print(f"{field}: {value}")
         elif isinstance(results, list):
             for i, value in enumerate(results):
-                console.print(f"{i}: {value}")
+                rich.print(f"{i}: {value}")
         else:
-            console.print(results)
-    finally:
-        if runner:
-            runner.stop(verbose=ctx.obj["verbose"])
+            rich.print(results)
+
+
+@main.command(name="build-bento")
+@click.argument("source")
+@click.option("--name", help="Name of the bento service")
+@click.option("--version", help="Version of the bento service")
+def bento_cmd(source: str, name: str | None, version: str | None):
+    """Build a bento from the source, which can be either a .cpack.zip file or a bento tag."""
+    import bentoml
+    from bentoml.bentos import BentoBuildConfig
+
+    from .package import build_bento
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        if source.endswith(".cpack.zip"):
+            name = name or os.path.basename(source).replace(".cpack.zip", "")
+            shutil.unpack_archive(source, temp_dir)
+            system_packages = None
+            include_default_system_packages = True
+        else:
+            existing_bento = bentoml.get(source)
+            name = name or existing_bento.tag.name
+            shutil.copytree(existing_bento.path_of("src"), temp_dir, dirs_exist_ok=True)
+            build_config = BentoBuildConfig.from_bento_dir(
+                existing_bento.path_of("src")
+            )
+            requirements_txt = Path(temp_dir) / "requirements.txt"
+            if (
+                requirements_txt.exists()
+                and "comfy-pack" not in requirements_txt.read_text()
+            ):
+                with open(requirements_txt, "a") as f:
+                    f.write("\ncomfy-pack")
+            system_packages = build_config.docker.system_packages
+            include_default_system_packages = False
+
+        build_bento(
+            name,
+            Path(temp_dir),
+            version=version,
+            system_packages=system_packages,
+            include_default_system_packages=include_default_system_packages,
+        )
+
+
+def setup_cloud_client(
+    ctx: click.Context, param: click.Parameter, value: str | None
+) -> str | None:
+    from bentoml._internal.configuration.containers import BentoMLContainer
+
+    if value:
+        BentoMLContainer.cloud_context.set(value)
+        os.environ["BENTOML_CLOUD_CONTEXT"] = value
+    return value
+
+
+@main.command()
+@click.argument("bento")
+@click.option(
+    "-w",
+    "--workspace",
+    type=click.Path(file_okay=False, path_type=Path),
+    default="workspace",
+    help="Workspace directory, defaults to './workspace'.",
+)
+@click.option("-v", "--verbose", count=True, help="Increase verbosity level")
+@click.option(
+    "--context",
+    help="BentoCloud context name.",
+    expose_value=False,
+    callback=setup_cloud_client,
+)
+def unpack_bento(bento: str, workspace: Path, verbose: int):
+    """Restore the ComfyUI workspace from a given bento."""
+    import bentoml
+
+    from .package import install
+
+    try:
+        bento_obj = bentoml.get(bento)
+    except bentoml.exceptions.NotFound:
+        click.echo(
+            f"Bento {bento} not found in the local repository, trying to pull from BentoCloud",
+            err=True,
+        )
+        bentoml.pull(bento)
+        bento_obj = bentoml.get(bento)
+
+    install(bento_obj.path_of("src"), workspace, verbose=verbose, prepare_models=False)
+
+    if os.name == "nt":
+        exe = "Scripts/python.exe"
+    else:
+        exe = "bin/python"
+    click.echo(
+        f"Workspace is ready at {workspace}\n"
+        f"You can start ComfyUI by running `cd {workspace} && .venv/{exe} main.py`",
+        color="green",
+    )

src/comfy_pack/const.py

@@ -0,0 +1,22 @@
+import pathlib
+import os
+
+
+CPACK_HOME = (
+    pathlib.Path.home() / ".comfypack"
+    if not os.environ.get("CPACK_HOME", "")
+    else pathlib.Path(os.environ.get("CPACK_HOME", ""))
+)
+if not CPACK_HOME.exists():
+    CPACK_HOME.mkdir()
+
+MODEL_DIR = CPACK_HOME / "models"
+WORKSPACE_DIR = CPACK_HOME / "workspace"
+SHA_CACHE_FILE = CPACK_HOME / "sha_cache.json"
+MODEL_SOURCE_CACHE_FILE = CPACK_HOME / "model_source_cache.json"
+
+COMFYUI_REPO = "https://github.com/comfyanonymous/ComfyUI.git"
+COMFY_PACK_REPO = "https://github.com/bentoml/comfy-pack.git"
+COMFYUI_MANAGER_REPO = "https://github.com/ltdrdata/ComfyUI-Manager.git"
+
+STRICT_MODE = os.environ.get("CPACK_STRICT_MODE", "0") in ["1", "true", "True"]

src/comfy_pack/hash.py

@@ -0,0 +1,127 @@
+import asyncio
+import json
+import os
+import subprocess
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from functools import partial
+from typing import Dict, List
+
+from .const import SHA_CACHE_FILE
+
+CALC_CMD = """
+import hashlib
+import sys
+
+filepath = sys.argv[1]
+chunk_size = int(sys.argv[2])
+
+sha256 = hashlib.sha256()
+with open(filepath, "rb") as f:
+    for chunk in iter(lambda: f.read(chunk_size), b""):
+        sha256.update(chunk)
+print(sha256.hexdigest())
+"""
+
+
+def calculate_sha256_worker(filepath: str, chunk_size: int = 4 * 1024 * 1024) -> str:
+    """Calculate SHA-256 in a separate process"""
+    result = subprocess.run(
+        [sys.executable, "-c", CALC_CMD, filepath, str(chunk_size)],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+    assert result.returncode == 0, result.stderr
+    return result.stdout.strip()
+
+
+def get_sha256(filepath: str) -> str:
+    return batch_get_sha256([filepath])[filepath]
+
+
+def async_get_sha256(filepath: str) -> str:
+    return asyncio.run(async_batch_get_sha256([filepath]))[filepath]
+
+
+def batch_get_sha256(filepaths: List[str], cache_only: bool = False) -> Dict[str, str]:
+    return asyncio.run(async_batch_get_sha256(filepaths, cache_only=cache_only))
+
+
+async def async_batch_get_sha256(
+    filepaths: List[str],
+    cache_only: bool = False,
+) -> Dict[str, str]:
+    # Load cache
+    cache = {}
+    if SHA_CACHE_FILE.exists():
+        try:
+            with SHA_CACHE_FILE.open("r") as f:
+                cache = json.load(f)
+        except (json.JSONDecodeError, IOError):
+            pass
+
+    # Initialize process pool
+    max_workers = max(1, (os.cpu_count() or 1))
+
+    # Process files
+    results = {}
+    new_cache = {}
+    async with asyncio.Lock():
+        with ThreadPoolExecutor(max_workers=max_workers) as pool:
+            loop = asyncio.get_event_loop()
+
+            for filepath in filepaths:
+                if not os.path.exists(filepath):
+                    results[filepath] = None
+                    continue
+
+                # Get file info
+                stat = os.stat(filepath)
+                current_size = stat.st_size
+                current_time = stat.st_ctime
+
+                # Check cache
+                cache_entry = cache.get(filepath)
+                if cache_entry:
+                    if (
+                        cache_entry["size"] == current_size
+                        and cache_entry["birthtime"] == current_time
+                    ):
+                        results[filepath] = cache_entry["sha256"]
+                        continue
+
+                if cache_only:
+                    results[filepath] = ""
+                    continue
+
+                # Calculate new SHA
+                calc_func = partial(calculate_sha256_worker, filepath)
+                sha256 = await loop.run_in_executor(pool, calc_func)
+
+                # Update cache and results
+                new_cache[filepath] = {
+                    "sha256": sha256,
+                    "size": current_size,
+                    "birthtime": current_time,
+                    "last_verified": datetime.now().isoformat(),
+                }
+                results[filepath] = sha256
+
+    # Save cache
+    try:
+        with SHA_CACHE_FILE.open("r") as f:
+            cache = json.load(f)
+    except (IOError, OSError):
+        cache = {}
+
+    cache.update(new_cache)
+
+    try:
+        with SHA_CACHE_FILE.open("w") as f:
+            json.dump(cache, f, indent=2)
+    except (IOError, OSError):
+        pass
+
+    return results

src/comfy_pack/model_helper.py

@@ -0,0 +1,168 @@
+import asyncio
+import json
+import re
+
+from .const import MODEL_SOURCE_CACHE_FILE
+
+# MODEL_NAME = r"[a-zA-Z0-9-._]+"
+# COMMIT = r"[a-f0-9]+"
+
+COMMIT_PATTERN = re.compile(r'href="/([a-zA-Z0-9-._/]+)/commit/([a-f0-9]+)"')
+
+PATH_PATTERN = re.compile(
+    r'data-target="CopyButton" data-props="{"value":"([^&]+)"'
+)
+
+
+async def _lookup_huggingface_model(model_sha: str) -> dict:
+    import aiohttp
+    from duckduckgo_search import DDGS
+    query = f"site:huggingface.co blob {model_sha}"
+
+    try:
+        with DDGS() as ddgs:
+            search_results = ddgs.text(query, max_results=5)
+
+            async with aiohttp.ClientSession(trust_env=True) as session:
+                for result in search_results:
+                    url = result['href']
+                    if "blob" not in url:
+                        continue
+
+                    try:
+                        async with session.get(url) as resp:
+                            if resp.status != 200:
+                                continue
+                            text = await resp.text()
+                            if commit_match := COMMIT_PATTERN.search(text):
+                                repo, commit = commit_match.groups()
+                                if path_match := PATH_PATTERN.search(text):
+                                    path = path_match.group(1)
+                                    info = {
+                                        "download_url": path,
+                                        "url": path,
+                                        "repo": repo,
+                                        "commit": commit,
+                                        "path": path,
+                                        "source": "huggingface",
+                                    }
+                                    return info
+                    except aiohttp.ClientError:
+                        continue
+    except Exception:
+        pass
+
+    return {}
+
+
+async def _loopup_civitai_model(model_sha: str) -> dict:
+    import aiohttp
+
+    async with aiohttp.ClientSession(trust_env=True) as session:
+        async with session.post(
+            "https://meilisearch-v1-9.civitai.com/multi-search",
+            headers={
+                "accept": "*/*",
+                "accept-language": "en,zh;q=0.9,zh-CN;q=0.8",
+                "cache-control": "no-cache",
+                "content-type": "application/json",
+                "origin": "https://civitai.com",
+                "pragma": "no-cache",
+                "priority": "u=1, i",
+                "referer": "https://civitai.com/",
+                "sec-ch-ua": '"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
+                "sec-ch-ua-mobile": "?0",
+                "sec-ch-ua-platform": '"macOS"',
+                "sec-fetch-dest": "empty",
+                "sec-fetch-mode": "cors",
+                "sec-fetch-site": "same-site",
+                "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
+                "x-meilisearch-client": "Meilisearch instant-meilisearch (v0.13.5) ; Meilisearch JavaScript (v0.34.0)",
+            },
+            json={
+                "queries": [
+                    {
+                        "q": model_sha,
+                        "indexUid": "models_v9",
+                        "facets": [
+                            "category.name",
+                            "checkpointType",
+                            "fileFormats",
+                            "lastVersionAtUnix",
+                            "tags.name",
+                            "type",
+                            "user.username",
+                            "version.baseModel",
+                        ],
+                        "attributesToHighlight": [],
+                        "highlightPreTag": "__ais-highlight__",
+                        "highlightPostTag": "__/ais-highlight__",
+                        "limit": 51,
+                        "offset": 0,
+                        "filter": ["nsfwLevel=1"],
+                    }
+                ]
+            },
+        ) as resp:
+            if resp.status != 200:
+                return {}
+            data = await resp.json()
+
+            if len(data.get("results", [])) == 0:
+                return {}
+
+            if len(data["results"][0]["hits"]) == 0:
+                return {}
+
+            hit = data["results"][0]["hits"][0]
+            repo_id = hit["id"]
+            repo_name = hit["name"]
+            versions = hit["versions"]
+
+            for version in versions:
+                if model_sha.upper() in version["hashes"]:
+                    break
+            else:
+                return {}
+            version_id = version["id"]
+            version_name = version["name"]
+            return {
+                "download_url": f"https://civitai.com/api/download/models/{version_id}",
+                "url": f"https://civitai.com/models/{repo_id}?modelVersionId={version_id}",
+                "repo": repo_id,
+                "commit": version_id,
+                "source": "civitai",
+                "repo_name": repo_name,
+                "version_name": version_name,
+            }
+
+
+async def alookup_model_source(model_sha: str, cache_only=False) -> dict:
+    if not model_sha:
+        return {}
+    try:
+        model_source_cache = json.loads(MODEL_SOURCE_CACHE_FILE.read_text())
+    except Exception:
+        with open(MODEL_SOURCE_CACHE_FILE, "w") as f:
+            f.write("{}")
+        model_source_cache = {}
+    if model_source_cache.get(model_sha):
+        return model_source_cache[model_sha]
+    if cache_only:
+        return {}
+
+    info = await _lookup_huggingface_model(model_sha)
+    if not info:
+        info = await _loopup_civitai_model(model_sha)
+
+    # elemental read and write
+    model_source_cache = json.loads(MODEL_SOURCE_CACHE_FILE.read_text())
+    model_source_cache[model_sha] = info
+    with open(MODEL_SOURCE_CACHE_FILE, "w") as f:
+        json.dump(model_source_cache, f)
+
+    return info
+
+
+def lookup_model_source(model_sha: str, cache_only=False) -> dict:
+    return asyncio.run(alookup_model_source(model_sha, cache_only=cache_only))

src/comfy_pack/package.py

@@ -1,41 +1,58 @@
 from __future__ import annotations
 
+import contextlib
 import json
 import os
 import shutil
 import subprocess
 import sys
 import tempfile
+import threading
+import urllib.parse
+import urllib.request
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-COMFYUI_REPO = "https://github.com/comfyanonymous/ComfyUI.git"
+from .const import COMFYUI_REPO, MODEL_DIR, STRICT_MODE
+from .hash import get_sha256
+from .utils import get_self_git_commit
+
+if TYPE_CHECKING:
+    import bentoml
+
+COMFY_PACK_DIR = Path(__file__).parent
 
 
 def _clone_commit(url: str, commit: str, dir: Path, verbose: int = 0):
     stdout = None if verbose > 0 else subprocess.DEVNULL
     stderr = None if verbose > 1 else subprocess.DEVNULL
+    env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
     subprocess.check_call(
         ["git", "clone", "--recurse-submodules", "--filter=blob:none", url, dir],
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
     subprocess.check_call(
         ["git", "fetch", "-q", url, commit],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
     subprocess.check_call(
         ["git", "checkout", "FETCH_HEAD"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
     subprocess.check_call(
         ["git", "submodule", "update", "--init", "--recursive"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
 
 
@@ -50,6 +67,18 @@ def install_comfyui(snapshot, workspace: Path, verbose: int = 0):
                 return
         shutil.rmtree(workspace)
     _clone_commit(COMFYUI_REPO, comfyui_commit, workspace, verbose=verbose)
+    manager_node = next(
+        url for url in snapshot.get("git_custom_nodes", {}) if "ComfyUI-Manager" in url
+    )
+    if manager_node:
+        manager_commit = snapshot["git_custom_nodes"][manager_node]["hash"].strip()
+        print("Installing ComfyUI-Manager")
+        _clone_commit(
+            manager_node,
+            manager_commit,
+            workspace / "custom_nodes" / "ComfyUI-Manager",
+            verbose=verbose,
+        )
     with open(workspace / ".DONE", "w") as f:
         f.write(comfyui_commit)
 
@@ -58,6 +87,9 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
     print("Installing custom nodes")
     for module in snapshot["custom_nodes"]:
         url = module["url"]
+        if not url.strip():
+            print(f"Skipping invalid custom node: {module}")
+            continue
         directory = url.split("/")[-1].split(".")[0]
         module_dir = workspace / "custom_nodes" / directory
 
@@ -69,10 +101,12 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
                     continue
             shutil.rmtree(module_dir)
 
+        print(f"Installing custom node {url}")
         commit_hash = module["commit_hash"]
         _clone_commit(url, commit_hash, module_dir, verbose=verbose)
 
         if module_dir.joinpath("install.py").exists():
+            env = os.environ.copy()
             venv = workspace / ".venv"
             if venv.exists():
                 python = (
@@ -80,8 +114,17 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
                     if os.name == "nt"
                     else venv / "bin" / "python"
                 )
+                if "PATH" in env:
+                    env["PATH"] = f"{str(python.parent)}:{env['PATH']}"
+                else:
+                    env["PATH"] = str(python.parent)
+                env["VIRTUAL_ENV"] = str(venv)
             else:
                 python = Path(sys.executable)
+
+            if verbose > 0:
+                print(f"Installing {directory} custom node")
+                print(f"$ {python.absolute()} install.py")
             subprocess.check_call(
                 [str(python.absolute()), "install.py"],
                 cwd=module_dir,
@@ -93,40 +136,39 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
 
 
 def install_dependencies(
-    snapshot: dict,
-    req_file: str,
+    python_version: str,
+    req_files: list[str],
     workspace: Path,
     verbose: int = 0,
-):
-    if verbose > 0:
-        print("Installing Python dependencies")
-    python_version = snapshot["python"]
+    no_deps: bool = False,
+    no_venv: bool = False,
+) -> Path:
+    print("Installing Python dependencies")
     stdout = None if verbose > 0 else subprocess.DEVNULL
     stderr = None if verbose > 1 else subprocess.DEVNULL
-
-    subprocess.check_call(
-        ["uv", "python", "install", python_version],
-        cwd=workspace,
-        stdout=stdout,
-        stderr=stderr,
-    )
-    venv = workspace / ".venv"
-    if (venv / "DONE").exists():
-        return
-    venv_py = (
-        venv / "Scripts" / "python.exe" if os.name == "nt" else venv / "bin" / "python"
-    )
-    subprocess.check_call(
-        [
-            "uv",
-            "venv",
-            "--python",
-            python_version,
-            venv,
-        ],
-        stdout=stdout,
-        stderr=stderr,
-    )
+    if no_venv:
+        print("Using the current Python environment")
+        venv_py = Path(sys.executable)
+    else:
+        venv = (workspace / ".venv").absolute()
+        venv_py = (
+            venv / "Scripts" / "python.exe"
+            if os.name == "nt"
+            else venv / "bin" / "python"
+        )
+        if (venv / "DONE").exists():
+            return venv_py
+        subprocess.check_call(
+            [
+                "uv",
+                "venv",
+                "--python",
+                python_version,
+                venv,
+            ],
+            stdout=stdout,
+            stderr=stderr,
+        )
     subprocess.check_call(
         [
             "uv",
@@ -139,39 +181,369 @@ def install_dependencies(
         stdout=stdout,
         stderr=stderr,
     )
+    if verbose > 0:
+        print(f"Installing dependencies from {req_files}")
+    install_cmd = [
+        "uv",
+        "pip",
+        "install",
+        "-p",
+        str(venv_py),
+    ]
+    for req_file in req_files:
+        install_cmd.extend(["-r", str(req_file)])
+    if not STRICT_MODE:
+        install_cmd.extend(["--index-strategy", "unsafe-best-match"])
+    if no_deps:
+        install_cmd.append("--no-deps")
     subprocess.check_call(
-        [
-            "uv",
-            "pip",
-            "install",
-            "-p",
-            str(venv_py),
-            "-r",
-            req_file,
-            "--no-deps",
-        ],
+        install_cmd,
         stdout=stdout,
         stderr=stderr,
     )
-    with open(venv / "DONE", "w") as f:
-        f.write("DONE")
+    if not no_venv:
+        with open(venv / "DONE", "w") as f:
+            f.write("DONE")
     return venv_py
 
 
-def install(cpack: str | Path, workspace: str | Path = "workspace", verbose: int = 0):
+def get_search_url(sha: str) -> str:
+    """Generate custom search URLs for model on HuggingFace and CivitAI"""
+    base_url = "https://duckduckgo.com"
+    sha = sha.upper()
+    hf_query = f"{sha} OR {sha[:10]}"
+    hf_query = urllib.parse.quote(hf_query)
+    return f"{base_url}?q={hf_query}"
+
+
+def download_file(url: str, dest_path: Path, progress_callback=None):
+    """Download file with progress tracking"""
+
+    # prepare auth token from huggingface if possible
+    if (token := os.getenv("HF_TOKEN")) and ("huggingface" in url):
+        bearer = f"Bearer {token}"
+        urllib_request = urllib.request.Request(url, headers={"Authorization": bearer})
+        curl_auth = ["-H", f"Authorization: {bearer}"]
+    else:
+        curl_auth = []
+        urllib_request = url
+
+    try:
+        if shutil.which("curl"):
+            subprocess.check_call(
+                ["curl", "-L", url, *curl_auth, "--fail", "-o", str(dest_path)],
+            )
+            return True
+        with urllib.request.urlopen(urllib_request) as response:
+            total_size = int(response.headers.get("content-length", 0))
+            block_size = 8192
+            downloaded = 0
+
+            with open(dest_path, "wb") as f:
+                while True:
+                    buffer = response.read(block_size)
+                    if not buffer:
+                        break
+                    downloaded += len(buffer)
+                    f.write(buffer)
+                    if progress_callback:
+                        progress = (
+                            (downloaded / total_size) * 100 if total_size > 0 else 0
+                        )
+                        progress_callback(progress)
+        return True
+    except Exception as e:
+        print(f"Download failed: {e}")
+        if dest_path.exists():
+            dest_path.unlink()
+        return False
+
+
+def show_progress(filename: str):
+    """Progress callback function"""
+
+    def callback(progress):
+        print(f"\rDownloading {filename}: {progress:.1f}%", end="")
+
+    return callback
+
+
+def create_model_symlink(global_path: Path, sha: str, target_path: Path, filename: str):
+    """Create symlink from global storage to workspace"""
+    source = global_path / sha
+    target = target_path / filename
+
+    if target.exists():
+        if target.is_symlink():
+            target.unlink()
+        else:
+            raise RuntimeError(f"File {target} already exists and is not a symlink")
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    os.symlink(source, target)
+
+
+def retrieve_models(
+    snapshot: dict,
+    workspace: Path,
+    download: bool = True,
+    all_models: bool = False,
+    verbose: int = 0,
+):
+    """Retrieve models from user downloads"""
+    print("Retrieving models")
+    models = snapshot.get("models", [])
+    if not models:
+        return
+
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+
+    for model in models:
+        sha = model["sha256"]
+        filename = model["filename"]
+        disabled = model.get("disabled", False)
+        if (workspace / filename).exists():
+            if not (MODEL_DIR / sha).exists() and (workspace / filename).is_file():
+                shutil.move(workspace / filename, MODEL_DIR / sha)
+                create_model_symlink(MODEL_DIR, sha, workspace, filename)
+            continue
+
+        if (MODEL_DIR / sha).exists():
+            print(f"Model {filename} already exists in cache")
+            create_model_symlink(MODEL_DIR, sha, workspace, filename)
+            continue
+
+        if disabled and not all_models:
+            continue
+
+        if not download:
+            continue
+
+        print(f"\nModel {filename} is never downloaded before")
+        if source := model.get("source"):
+            url = source["download_url"]
+            target_path = MODEL_DIR / sha
+            download_thread = threading.Thread(
+                target=download_file,
+                args=(url, target_path, show_progress(filename)),
+            )
+            download_thread.start()
+            download_thread.join()
+
+            if not target_path.exists():
+                print("\nDownload failed!")
+                continue
+
+            print("\nDownload completed! Verifying SHA256...")
+            terget_sha = get_sha256(str(target_path))
+            if terget_sha != sha:
+                print("SHA256 verification failed! File may be corrupted or incorrect.")
+                target_path.unlink()
+            else:
+                continue
+
+        search_url = get_search_url(sha)
+        print(f"Search URL: {search_url}")
+        print(f"Path: {workspace / filename}")
+
+        while True:
+            path = input("Enter path to downloaded file (or 'skip' to skip): ")
+            if path.lower() == "skip":
+                break
+
+            try:
+                # Check if input is a URL
+                if path.startswith(("http://", "https://")):
+                    url = path
+                    target_path = MODEL_DIR / sha
+
+                    # Start download in a separate thread
+                    download_thread = threading.Thread(
+                        target=download_file,
+                        args=(url, target_path, show_progress(filename)),
+                    )
+                    download_thread.start()
+                    download_thread.join()
+
+                    if not target_path.exists():
+                        print("\nDownload failed!")
+                        continue
+
+                    print("\nDownload completed! Verifying SHA256...")
+                    terget_sha = get_sha256(str(target_path))
+                    if terget_sha != sha:
+                        print(
+                            "SHA256 verification failed! File may be corrupted or incorrect."
+                        )
+                        target_path.unlink()
+                        continue
+
+                    print("SHA256 verification successful!")
+                else:
+                    # Handle local file
+                    downloaded_path = Path(path)
+                    if not downloaded_path.exists():
+                        print("File does not exist!")
+                        continue
+
+                    # Verify SHA256 before copying
+                    print("Verifying SHA256...")
+                    target_sha = get_sha256(str(downloaded_path))
+                    if target_sha != sha:
+                        print(
+                            f"Downloaded file SHA256 does not match expected SHA256: {target_sha} != {sha}"
+                        )
+                        continue
+
+                    print("SHA256 verification successful!")
+                    # Copy to global storage
+                    shutil.copy2(downloaded_path, MODEL_DIR / sha)
+
+                # Create symlink
+                create_model_symlink(MODEL_DIR, sha, workspace, filename)
+                print(f"Model {filename} installed successfully")
+                break
+            except Exception as e:
+                print(f"Error processing file: {e}")
+                continue
+
+
+def install(
+    cpack: str | Path,
+    workspace: str | Path = "workspace",
+    preheat: bool = False,
+    prepare_models: bool = True,
+    all_models: bool = False,
+    no_venv: bool = False,
+    verbose: int = 0,
+):
     workspace = Path(workspace)
-    with tempfile.TemporaryDirectory() as temp_dir:
-        pack_dir = Path(temp_dir) / ".cpack"
-        shutil.unpack_archive(cpack, pack_dir)
+    cpack = Path(cpack)
+    print(f"Installing package {cpack} to {workspace} (verbose={verbose})")
+    with contextlib.ExitStack() as stack:
+        if cpack.is_file():
+            temp_dir = stack.enter_context(tempfile.TemporaryDirectory())
+            pack_dir = Path(temp_dir) / ".cpack"
+            shutil.unpack_archive(cpack, pack_dir)
+        else:
+            pack_dir = cpack
         snapshot = json.loads((pack_dir / "snapshot.json").read_text())
-        req_txt_file = pack_dir / "requirements.txt"
+        if "pips" not in snapshot:
+            raise RuntimeError(
+                "This cpack is generated by an old version of comfy-pack, "
+                "please use comfy-pack<0.4.0 to unpack it."
+            )
 
         install_comfyui(snapshot, workspace, verbose=verbose)
-        install_dependencies(snapshot, str(req_txt_file), workspace, verbose=verbose)
-        install_custom_modules(snapshot, workspace, verbose=verbose)
+        py = install_dependencies(
+            snapshot["python"],
+            [
+                str(workspace / "requirements.txt"),
+                str(
+                    workspace / "custom_nodes" / "ComfyUI-Manager" / "requirements.txt"
+                ),
+            ],
+            workspace,
+            no_venv=no_venv,
+            verbose=verbose,
+        )
+        cm_cli = Path("custom_nodes", "ComfyUI-Manager", "cm-cli.py")
+        subprocess.check_call(
+            [
+                str(py),
+                str(cm_cli),
+                "restore-snapshot",
+                "--pip-non-url",
+                "--pip-non-local-url",
+                str(pack_dir / "snapshot.json"),
+            ],
+            cwd=workspace,
+        )
 
         for f in (pack_dir / "input").glob("*"):
             if f.is_file():
                 shutil.copy(f, workspace / "input" / f.name)
             elif f.is_dir():
-                shutil.copytree(f, workspace / "input" / f.name)
+                shutil.copytree(f, workspace / "input" / f.name, dirs_exist_ok=True)
+        if prepare_models:
+            retrieve_models(
+                snapshot,
+                workspace,
+                verbose=verbose,
+                download=False,
+            )
+
+        if preheat:
+            from .run import ComfyUIServer
+
+            with ComfyUIServer(
+                str(workspace),
+                verbose=verbose,
+                venv=str(workspace / ".venv") if not no_venv else None,
+            ) as _:
+                pass
+        if prepare_models:
+            retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models)
+
+
+required_files = ["snapshot.json"]
+
+
+def build_bento(
+    bento_name: str,
+    source_dir: Path,
+    *,
+    version: str | None = None,
+    system_packages: list[str] | None = None,
+    include_default_system_packages: bool = True,
+) -> bentoml.Bento:
+    import bentoml
+
+    for f in required_files:
+        if not (source_dir / f).exists():
+            raise FileNotFoundError(f"Not a valid comfy-pack package: missing `{f}`")
+
+    if include_default_system_packages:
+        system_packages = [
+            "git",
+            "libglib2.0-0",
+            "libsm6",
+            "libxrender1",
+            "libxext6",
+            "ffmpeg",
+            "libstdc++-12-dev",
+            *(system_packages or []),
+        ]
+    else:
+        system_packages = system_packages or []
+
+    shutil.copy2(Path(__file__).with_name("service.py"), source_dir / "service.py")
+    snapshot_text = (source_dir / "snapshot.json").read_text()
+    setup_script = source_dir / "setup_workspace.sh"
+    with Path(__file__).with_name("setup_workspace.sh").open() as f:
+        setup_script.write_text(f.read().replace("<SNAPSHOT>", snapshot_text))
+    # Make setup script executable in a cross-platform way
+    if os.name in ("posix", "mac"):
+        setup_script.chmod(setup_script.stat().st_mode | 0o755)
+    snapshot = json.loads(snapshot_text)
+    return bentoml.build(
+        "service:ComfyService",
+        name=bento_name,
+        version=version,
+        build_ctx=str(source_dir),
+        labels={"comfy-pack-version": get_self_git_commit() or "unknown"},
+        models=[
+            m["model_tag"]
+            for m in snapshot["models"]
+            if "model_tag" in m and not m.get("disabled", False)
+        ],
+        docker={
+            "python_version": f"{sys.version_info.major}.{sys.version_info.minor}",
+            "system_packages": system_packages,
+            "setup_script": source_dir.joinpath("setup_workspace.sh").as_posix(),
+        },
+        python={
+            "lock_packages": True,
+            "packages": ["fastapi", "comfy-pack>=0.4.0a0"],
+        },
+    )

src/comfy_pack/run.py

@@ -1,12 +1,16 @@
 from __future__ import annotations
 
-import random
 import copy
 import json
 import logging
 import os
+import random
+import shlex
 import shutil
+import socket
 import subprocess
+import time
+import uuid
 from pathlib import Path
 from typing import Any, Union
 
@@ -24,28 +28,63 @@ def _probe_comfyui_server(port: int) -> None:
     req = request.Request(full_url)
     _ = request.urlopen(req)
 
+    full_url = f"http://127.0.0.1:{port}/api/object_info"
+    req = request.Request(full_url)
+    _ = request.urlopen(req)
 
-class WorkflowRunner:
+
+def _is_port_in_use(port: int | str, host="localhost"):
+    if isinstance(port, str):
+        port = int(port)
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        try:
+            s.connect((host, port))
+            return True
+        except ConnectionRefusedError:
+            return False
+        except Exception:
+            return True
+
+
+class ComfyUIServer:
     def __init__(
         self,
         workspace: str,
         input_dir: str | None = None,
+        host: str = "localhost",
         port: int | None = None,
+        venv: str | None = None,
+        verbose: int = 0,
     ) -> None:
         """
-        Initialize the WorkflowRunner.
-
         Args:
-            workspace (str): The workspace path for ComfyUI.
+            workspace (str, optional): The workspace path for ComfyUI. If not specified, runner will try to connect to an existing ComfyUI server.
+            input_dir (str, optional): The input directory for ComfyUI. Defaults to None.
+            port (int, optional): The port number for ComfyUI. Defaults to None. If 8188 is in use, a random port will be chosen.
         """
         self.workspace = workspace
-        self.temp_dir = Path(workspace) / "cli_run" / "temp"
-        self.output_dir = Path(workspace) / "cli_run" / "output"
         self.input_dir = input_dir
-        self.is_running = False
-        self.port = port if port else random.randint(58000, 58999)
+        self.verbose = verbose
+        self.host = host
+        self.server_proc: subprocess.Popen | None = None
 
-    def start(self, verbose: int = 0) -> None:
+        run_dir = (Path(workspace) / "cli_run").absolute()
+        self.temp_dir = run_dir / "temp"
+        self.output_dir = run_dir / "output"
+
+        self.temp_dir.mkdir(parents=True, exist_ok=True)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+
+        if port is None:
+            if _is_port_in_use(8188):
+                self.port = port if port else random.randint(58000, 58999)
+            else:
+                self.port = 8188
+        else:
+            self.port = port
+        self.venv = os.path.abspath(venv) if venv else None
+
+    def start(self) -> None:
         """
         Start the ComfyUI process.
 
@@ -58,30 +97,29 @@ class WorkflowRunner:
         Raises:
             RuntimeError: If ComfyUI is already running.
         """
-        if self.is_running:
-            raise RuntimeError("ComfyUI Runner is already started")
-
         logger.info(
             "Disable tracking from Comfy CLI, not for privacy concerns, but to workaround a bug"
         )
         self.temp_dir.mkdir(parents=True, exist_ok=True)
         self.output_dir.mkdir(parents=True, exist_ok=True)
 
-        stdout = None if verbose > 0 else subprocess.DEVNULL
-        command = ["comfy", "--skip-prompt", "tracking", "disable"]
-        subprocess.run(command, check=True, stdout=stdout)
-        logger.info("Successfully disabled Comfy CLI tracking")
+        env = os.environ.copy()
+        if self.venv:
+            env["VIRTUAL_ENV"] = self.venv
+            if os.name == "nt":
+                env["PATH"] = f"{self.venv}\\Scripts;{env.get('PATH', '')}"
+            else:
+                env["PATH"] = f"{self.venv}/bin:{env.get('PATH', '')}"
 
-        logger.info("Preparing directories required by ComfyUI...")
+        stdout = None if self.verbose > 0 else subprocess.DEVNULL
+        command = ["comfy", "--skip-prompt", "tracking", "disable"]
+        subprocess.run(command, check=True, stdout=stdout, env=env)
+        logger.info("Successfully disabled Comfy CLI tracking")
 
         logger.info("Starting ComfyUI in the background...")
         command = [
-            "comfy",
-            "--workspace",
-            self.workspace,
-            "launch",
-            "--background",
-            "--",
+            "python",
+            "main.py",
             "--output-directory",
             self.output_dir,
             "--temp-directory",
@@ -91,14 +129,38 @@ class WorkflowRunner:
         ]
         if self.input_dir:
             command.extend(["--input-directory", self.input_dir])
-        if subprocess.run(command, check=True, stdout=stdout):
-            self.is_running = True
+
+        if self.host != "localhost":
+            command.extend(["--listen", self.host])
+        if options := env.pop("COMFYUI_OPTIONS", None):
+            command.extend(shlex.split(options))
+
+        def preexec_fn():
+            os.setpgrp()
+
+        self.server_proc = subprocess.Popen(
+            command,
+            stdout=stdout,
+            stderr=None,
+            preexec_fn=preexec_fn,
+            env=env,
+            cwd=self.workspace,
+        )
+
+        if _wait_for_startup(self.host, self.port):
             _probe_comfyui_server(self.port)
             logger.info("Successfully started ComfyUI in the background")
         else:
             logger.error("Failed to start ComfyUI in the background")
 
-    def stop(self, verbose: int = 0) -> None:
+    def is_running(self) -> bool:
+        if self.server_proc is None:
+            return False
+        if self.server_proc.poll() is not None:
+            return False
+        return True
+
+    def stop(self) -> None:
         """
         Stop the ComfyUI process.
 
@@ -107,13 +169,13 @@ class WorkflowRunner:
         Raises:
             RuntimeError: If ComfyUI is not currently running.
         """
-        if not self.is_running:
-            raise RuntimeError("ComfyUI Runner is not started yet")
-
+        if self.server_proc is None:
+            raise RuntimeError("ComfyUI server is not started yet")
+        proc = self.server_proc
+        self.server_proc = None
         logger.info("Stopping ComfyUI...")
-        command = ["comfy", "stop"]
-        stdout = None if verbose > 0 else subprocess.DEVNULL
-        subprocess.run(command, check=True, stdout=stdout)
+        proc.terminate()
+        proc.wait()
         logger.info("Successfully stopped ComfyUI")
 
         logger.info("Cleaning up temporary directory...")
@@ -121,79 +183,106 @@ class WorkflowRunner:
         shutil.rmtree(self.output_dir, ignore_errors=True)
         logger.info("Successfully cleaned up temporary directory")
 
-        self.is_running = False
+    def __enter__(self):
+        self.start()
+        return self
 
-    def run_workflow(
-        self,
-        workflow: dict,
-        output_dir: Union[str, Path, None] = None,
-        timeout: int = 300,
-        verbose: int = 0,
-        **kwargs: Any,
-    ) -> Any:
-        """
-        Run a ComfyUI workflow.
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
 
-        This method executes a given workflow, populates it with input data,
-        and retrieves the output.
 
-        Args:
-            workflow (dict): The workflow to run.
-            output_dir (Union[str, Path, None], optional): Temporary directory for the workflow. Defaults to None.
-            timeout (int, optional): Timeout for the workflow execution in seconds. Defaults to 300.
-            **kwargs: Additional keyword arguments for workflow population.
+def _wait_for_startup(host: str, port: int, timeout: int = 1800) -> bool:
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        if _is_port_in_use(port, host):
+            return True
+        time.sleep(1)
+    return False
 
-        Returns:
-            Any: The output of the workflow.
 
-        Raises:
-            RuntimeError: If ComfyUI is not started.
-        """
-        if not self.is_running:
-            raise RuntimeError("ComfyUI Runner is not started yet")
+def run_workflow(
+    host: str,
+    port: int,
+    workflow: dict,
+    output_dir: Union[str, Path, None] = None,
+    timeout: int = 300,
+    verbose: int = 0,
+    workspace: str = ".",
+    **kwargs: Any,
+) -> Any:
+    """
+    Run a ComfyUI workflow.
 
-        workflow_copy = copy.deepcopy(workflow)
-        if output_dir is None:
-            output_dir = Path(".")
-        if isinstance(output_dir, str):
-            output_dir = Path(output_dir)
+    This method executes a given workflow, populates it with input data,
+    and retrieves the output.
 
-        run_id = os.urandom(8).hex()
-        populate_workflow(
-            workflow_copy,
-            output_dir,
-            session_id=run_id,
-            **kwargs,
-        )
+    Args:
+        workflow (dict): The workflow to run.
+        output_dir (Union[str, Path, None], optional): Temporary directory for the workflow. Defaults to None.
+        timeout (int, optional): Timeout for the workflow execution in seconds. Defaults to 300.
+        **kwargs: Additional keyword arguments for workflow population.
 
-        workflow_file_path = Path(self.workspace) / "workflow.json"
-        with open(workflow_file_path, "w") as file:
-            json.dump(workflow_copy, file)
+    Returns:
+        Any: The output of the workflow.
 
-        extra_args = []
-        if verbose > 0:
-            extra_args.append("--verbose")
+    Raises:
+        RuntimeError: If ComfyUI is not started.
+    """
+    run_id = uuid.uuid4().hex[:8]
 
-        # Execute the workflow
-        command = [
-            "comfy",
-            "run",
-            "--workflow",
-            workflow_file_path.as_posix(),
-            "--port",
-            str(self.port),
-            "--timeout",
-            str(timeout),
-            "--wait",
-            *extra_args,
-        ]
-        env = os.environ.copy()
-        env["NO_COLOR"] = "1"
-        subprocess.run(command, check=True, env=env)
+    workflow_copy = copy.deepcopy(workflow)
+    if output_dir is None:
+        output_dir = Path(".")
+    if isinstance(output_dir, str):
+        output_dir = Path(output_dir)
 
-        # retrieve the output
-        return retrieve_workflow_outputs(
-            workflow_copy,
-            output_dir,
-            session_id=run_id,
-        )
+    run_id = os.urandom(8).hex()
+    populate_workflow(
+        workflow_copy,
+        output_dir,
+        session_id=run_id,
+        **kwargs,
+    )
+
+    workflow_file_path = output_dir / f"workflow_{run_id}.json"
+    with open(workflow_file_path, "w") as file:
+        json.dump(workflow_copy, file)
+
+    extra_args = []
+    if verbose > 0:
+        extra_args.append("--verbose")
+
+    stdout = None if verbose > 0 else subprocess.DEVNULL
+    command = ["comfy", "--skip-prompt", "tracking", "disable"]
+    subprocess.run(command, check=True, stdout=stdout)
+
+    # Execute the workflow
+    command = [
+        "comfy",
+        "--skip-prompt",
+        "--workspace",
+        workspace,
+        "run",
+        "--workflow",
+        workflow_file_path.as_posix(),
+        "--port",
+        str(port),
+        "--host",
+        host,
+        "--timeout",
+        str(timeout),
+        "--wait",
+        *extra_args,
+    ]
+    env = os.environ.copy()
+    env["NO_COLOR"] = "1"
+    subprocess.run(command, check=True, env=env)
+
+    workflow_file_path.unlink()
+
+    # retrieve the output
+    return retrieve_workflow_outputs(
+        workflow_copy,
+        output_dir,
+        session_id=run_id,
+    )

src/comfy_pack/service.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import signal
+import threading
+import time
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, cast
+
+import bentoml
+import fastapi
+from bentoml.models import HuggingFaceModel
+
+import comfy_pack
+import comfy_pack.run
+
+REQUEST_TIMEOUT = 3600
+BASE_DIR = Path(__file__).parent
+COPY_THRESHOLD = 10 * 1024 * 1024
+INPUT_DIR = BASE_DIR / "input"
+logger = logging.getLogger("bentoml.service")
+
+
+EXISTING_COMFYUI_SERVER = os.environ.get("COMFYUI_SERVER")
+
+with BASE_DIR.joinpath("workflow_api.json").open() as f:
+    workflow = json.load(f)
+
+InputModel = comfy_pack.generate_input_model(workflow)
+app = fastapi.FastAPI()
+
+
+@lru_cache
+def _get_workspace() -> Path:
+    import hashlib
+
+    from bentoml._internal.configuration.containers import BentoMLContainer
+
+    snapshot = BASE_DIR / "snapshot.json"
+    checksum = hashlib.md5(snapshot.read_bytes().strip()).hexdigest()
+    wp = (
+        Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
+    )
+    wp.parent.mkdir(parents=True, exist_ok=True)
+    return wp
+
+
+@app.get("/workflow.json")
+def workflow_json():
+    return workflow
+
+
+def _watch_server(server: comfy_pack.run.ComfyUIServer):
+    while True:
+        time.sleep(1)
+        if not server.is_running():
+            if server.server_proc is not None:
+                logger.warning(
+                    "Server exited with code %s", server.server_proc.returncode
+                )
+                os.kill(os.getpid(), signal.SIGTERM)
+            break
+
+
+if not EXISTING_COMFYUI_SERVER:
+    # register models
+    with BASE_DIR.joinpath("snapshot.json").open("rb") as f:
+        snapshot = json.load(f)
+else:
+    snapshot = {}
+
+
+@bentoml.asgi_app(app, path="/comfy")
+@bentoml.service(traffic={"timeout": REQUEST_TIMEOUT * 2}, resources={"gpu": 1})
+class ComfyService:
+    def __init__(self):
+        logger = logging.getLogger("comfy_pack")
+        logger.setLevel(logging.INFO)
+        if not EXISTING_COMFYUI_SERVER:
+            workspace_override = os.environ.get("COMFYUI_PATH", None)
+            self.server = comfy_pack.run.ComfyUIServer(
+                workspace_override if workspace_override else str(_get_workspace()),
+                str(INPUT_DIR),
+                verbose=int("BENTOML_DEBUG" in os.environ),
+            )
+            self.server.start()
+            logger.info(
+                "ComfyUI Server started at %s:%s", self.server.host, self.server.port
+            )
+            self.host = self.server.host
+            self.port = self.server.port
+            self.watch_thread = threading.Thread(
+                target=_watch_server,
+                args=(self.server,),
+                daemon=True,
+            )
+            self.watch_thread.start()
+            logger.info("Watch thread started")
+        else:
+            logger.info("Attaching to ComfyUI server: %s", EXISTING_COMFYUI_SERVER)
+            if ":" in EXISTING_COMFYUI_SERVER:
+                self.host, port = EXISTING_COMFYUI_SERVER.split(":")
+                self.port = int(port)
+            else:
+                self.host = EXISTING_COMFYUI_SERVER
+                self.port = 80
+
+    @bentoml.api(input_spec=InputModel)
+    def generate(
+        self,
+        *,
+        ctx: bentoml.Context,
+        **kwargs: Any,
+    ) -> Path:
+        verbose = int("BENTOML_DEBUG" in os.environ)
+        ret = comfy_pack.run_workflow(
+            self.host,
+            self.port,
+            workflow,
+            output_dir=ctx.temp_dir,
+            timeout=REQUEST_TIMEOUT,
+            verbose=verbose,
+            workspace=self.server.workspace if hasattr(self, "server") else ".",
+            **kwargs,
+        )
+        if isinstance(ret, list):
+            ret = ret[-1]
+        return ret
+
+    @bentoml.on_deployment
+    @staticmethod
+    def prepare_models():
+        if EXISTING_COMFYUI_SERVER:
+            return
+        comfy_workspace = _get_workspace()
+        if not comfy_workspace.joinpath(".DONE").exists():
+            raise RuntimeError("ComfyUI workspace is not ready")
+        for model in snapshot["models"]:
+            if model.get("disabled", False):
+                continue
+            model_path = comfy_workspace / cast(str, model["filename"])
+            if model_path.exists():
+                continue
+            if model_tag := model.get("model_tag"):
+                model_path.parent.mkdir(parents=True, exist_ok=True)
+                bento_model = bentoml.models.get(model_tag)
+                model_file = bento_model.path_of("model.bin")
+                print(f"Copying {model_file} to {model_path}")
+                model_path.symlink_to(model_file)
+            elif (source := model["source"]).get("source") == "huggingface":
+                matched = next(
+                    (
+                        m
+                        for m in ComfyService.models
+                        if isinstance(m, HuggingFaceModel)
+                        and m.model_id.lower() == source["repo"].lower()
+                        and source["commit"].lower() == m.revision.lower()
+                    ),
+                    None,
+                )
+                if matched is not None:
+                    model_file = os.path.join(matched.resolve(), source["path"])
+                    model_path.parent.mkdir(parents=True, exist_ok=True)
+                    print(f"Copying {model_file} to {model_path}")
+                    model_path.symlink_to(model_file)
+            else:
+                print(
+                    f"WARN: Unrecognized model source: {source}, the model may be missing"
+                )
+
+
+if False and not EXISTING_COMFYUI_SERVER:
+    for model in snapshot["models"]:
+        if model.get("disabled"):
+            continue
+        source = model["source"]
+        if source.get("source") != "huggingface" or source["repo"].startswith(
+            "datasets/"
+        ):
+            continue
+        ComfyService.models.append(HuggingFaceModel(source["repo"], source["commit"]))

src/comfy_pack/setup_workspace.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+set -eo pipefail
+
+CPACK=/tmp/bento-cpack
+mkdir -p "$CPACK"
+
+cat <<EOF | head -c -1 > "$CPACK"/snapshot.json
+<SNAPSHOT>
+EOF
+
+checksum=$(md5sum "$CPACK"/snapshot.json | awk '{print $1}')
+workspace="${BENTOML_HOME:-$HOME/bentoml}/run/comfy_workspace/${checksum}"
+
+if [ -n "$VIRTUAL_ENV" ]; then
+    # shellcheck disable=SC1091
+    source "$VIRTUAL_ENV"/bin/activate
+fi
+
+set -x
+comfy-pack unpack "$CPACK" -d "$workspace" --no-models --no-venv -v
+chown -R bentoml:bentoml "$workspace"
+set +x

src/comfy_pack/utils.py

@@ -1,17 +1,15 @@
 from __future__ import annotations
 
 import re
+import subprocess
 import sys
+import zipfile
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Union
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
 
-CPACK_OUTPUT_NODES = {
-    "CPackOutputFile",
-    "CPackOutputImage",
-}
 
 CPACK_PATH_INPUT_NODES = {
     "CPackInputFile",
@@ -209,7 +207,7 @@ def populate_workflow(
 
     for _, node in output_spec.items():
         node_id = node["id"]
-        if node["class_type"] in CPACK_OUTPUT_NODES:
+        if node["class_type"].startswith("CPackOutput"):
             workflow[node_id]["inputs"]["filename_prefix"] = (
                 output_path / f"{session_id}{node_id}_"
             ).as_posix()
@@ -220,7 +218,7 @@ def retrieve_workflow_outputs(
     workflow: dict,
     output_path: Path,
     session_id: str = "",
-) -> Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]:
+) -> Union[Path, list[Path], dict[str, Path | list[Path]]]:
     """
     Gets the output file(s) from the workflow.
 
@@ -229,7 +227,7 @@ def retrieve_workflow_outputs(
         output_path (Path): The path where output files are saved.
 
     Returns:
-        Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]:
+        Union[Path, list[Path], dict[str, Path | list[Path]]]:
             - A single Path if there's only one output file.
             - A list of Paths if there are multiple files for a single output.
             - A dictionary mapping output names to Paths or lists of Paths for multiple outputs.
@@ -238,23 +236,85 @@ def retrieve_workflow_outputs(
         ValueError: If the output node is not of the expected type.
     """
     _, outputs = _parse_workflow(workflow)
+    should_zip = any(
+        node["class_type"] == "CPackOutputZipSwitch" for node in workflow.values()
+    )
+    zip_paths: list[tuple[Path, str]] = []
     if len(outputs) != 1:
         value_map = {}
         for k, node in outputs.items():
             node_id = node["id"]
-            path_strs = list(output_path.glob(f"{node_id}_*"))
+            path_strs = list(output_path.glob(f"{session_id}{node_id}_*"))
+            zip_paths.extend(
+                (p, p.name.replace(f"{session_id}{node_id}", k)) for p in path_strs
+            )
             if len(path_strs) == 1:
                 value_map[k] = path_strs[0]
             else:
                 value_map[k] = path_strs
-        return value_map
+        if not should_zip:
+            return value_map
+    else:
+        name, node = next(iter(outputs.items()))
+        if not node["class_type"].startswith("CPackOutput"):
+            raise ValueError(f"Node {name} is not a comfy-pack output node")
+        node_id = node["id"]
 
-    name, node = next(iter(outputs.items()))
-    if node["class_type"] not in CPACK_OUTPUT_NODES:
-        raise ValueError(f"Output node {name} is not of type {CPACK_OUTPUT_NODES}")
-    node_id = node["id"]
+        outs = list(output_path.glob(f"{session_id}{node_id}_*"))
+        zip_paths.extend(
+            (p, p.name.replace(f"{session_id}{node_id}", name)) for p in outs
+        )
+        if not should_zip:
+            if len(outs) == 1:
+                return outs[0]
+            return outs
+    if len(zip_paths) == 1:
+        return zip_paths[0][0]
+    # Make a zipball from the collected files
+    output_zip = output_path / f"{session_id}_output.zip"
+    print(f"Creating zip file: {output_zip}")
+    with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for path, name in zip_paths:
+            zipf.write(path, arcname=name)
+    return output_zip
 
-    outs = list(output_path.glob(f"{session_id}{node_id}_*"))
-    if len(outs) == 1:
-        return outs[0]
-    return outs
+
+def get_self_git_commit() -> str | None:
+    """Get current git commit of the repository.
+
+    Returns:
+        str | None: Git commit hash in format "{hash}[-dirty]" or None if not in a git repo
+    """
+    try:
+        repo_root = Path(__file__).parent.parent.parent
+
+        # Check if we're in a git repo
+        subprocess.run(
+            ["git", "rev-parse", "--git-dir"],
+            cwd=repo_root,
+            check=True,
+            capture_output=True,
+        )
+
+        # Get current commit hash
+        commit_hash = subprocess.run(
+            ["git", "rev-parse", "--short", "HEAD"],
+            cwd=repo_root,
+            check=True,
+            capture_output=True,
+            text=True,
+        ).stdout.strip()
+
+        # Check if working directory is clean
+        is_dirty = (
+            subprocess.run(
+                ["git", "diff", "--quiet"],
+                cwd=repo_root,
+                check=False,
+            ).returncode
+            != 0
+        )
+
+        return f"{commit_hash}-dirty" if is_dirty else commit_hash
+    except (subprocess.SubprocessError, FileNotFoundError):
+        return None