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>

.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,9 +1,9 @@
-# Comfy-Pack: Package and Deploy ComfyUI Workflows
-> ⚠️ **WARNING**: This project is under active development. 
+# Comfy-Pack: Making ComfyUI Workflows Shareable
 
-<img width="952" alt="banner" src="https://github.com/user-attachments/assets/1ab2c7a7-55da-4fa4-a821-db7ba720fbc9" />
+![banner2](https://github.com/user-attachments/assets/14a7e469-6683-4818-9d54-5e5a8d0aa454)
 
-A comprehensive toolkit for reliably packing and unpacking environments for ComfyUI workflows. 
+
+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.
@@ -11,70 +11,157 @@ A comprehensive toolkit for reliably packing and unpacking environments for Comf
 - 🚀 **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:
 
-When sharing ComfyUI workflows to others, your audiences or team member, you've likely heard these responses:
 - "Custom Node not found"
-- "Cannot find the model file"
+- "Cannot find the correct model file"
 - "Missing Python dependencies"
 
-These are fundamental challenges in workflow sharing – every component must match exactly: custom nodes, model files, and 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 communit, and developed Comfy-Pack to address this directly. 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.
+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
 ```
 
-Focus on creating. Let Comfy-Pack handle the rest.
+This means you can focus on your creative work while comfy-pack handles the rest.
 
-## Quick Start
+## Usages
 
 ### Installation
 
-Search `comfy-pack` in ComfyUI Manager (Recommended)
+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)
 
-or install from Git:
+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
+```
+
+### 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 a ComfyUI workflow
-1. Click "Package" button to create `.cpack.zip`
-2. (Optional) select the models that you want to include (only model hash will be recorded, so you wont 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:
 
-### Unpack a ComfyUI workflow
 ```bash
-# Restore the full ComfyUI directory from cpack files.
 comfy-pack unpack workflow.cpack.zip
 ```
 
+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>
-  
+<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 and test locally </summary>
-  
+<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 & run anywhere </summary>
-  
+<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 run workflow.cpack.zip --help
@@ -85,46 +172,43 @@ comfy-pack run workflow.cpack.zip --src-image image.png --video video.mp4
 </details>
 
 <details> 
-<summary> 4. (Optional) deploy to cloud * </summary>
+<summary> 4. (Optional) Deploy to the cloud </summary>
+
+Deploy to [BentoCloud](https://www.bentoml.com/) with access to a variety of GPUs and blazing fast scaling.
+
+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/).
 
 ![image](https://github.com/user-attachments/assets/1ffa31fc-1f50-4ea7-a47e-7dae3b874273)
 
-
 </details>
 
+## Security Guidelines
 
-## Custom Node List
+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.
 
-ComfyPack provides custom nodes for standardizing inputs:
-- ImageInput: provides `image` type input, similar to official `LoadImage` node
-- StringInput: provides `string` type input, nice for prompts
-- IntInput: provides `int` type input, suitable for size or seeds
-- AnyInput: provides `combo` type and more input, suitable for custom nodes
-- ImageOutput: takes `image` type inputs, similar to official `SaveImage` node, take an image of a bunch of images
-- FileOutput: takes file path as `string` type, save and output the file under that path
-- ...
+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:
 
-These nodes help define clear interfaces for your workflow.
+* Use more strict index strategy in Python package installation
 
-## 🚀 Roadmap
-- Docker Support
-- local cpack manager and Version Control
+More security features are under way.
+
+
+## Roadmap
+
+This project is under active development. Currently we are working on:
+
+- Enhanced user experience
+- Docker support
+- Local `.cpack` file management with version control
 - Enhanced service capabilities
 
-
-## Examples
-
-Check our [examples folder](examples/) for:
-- Basic workflow packaging
-- Parameter configuration
-- API integration
-- Docker deployment
-
-## License
-MIT License
-
 ## 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

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

nodes/api.py

@@ -11,59 +11,78 @@ import tempfile
 import time
 import uuid
 import zipfile
+from fnmatch import fnmatch
 from pathlib import Path
 from typing import Any, Union
 
-import folder_paths
 from aiohttp import web
-from server import PromptServer
-
 from comfy_pack.hash import async_batch_get_sha256
 from comfy_pack.model_helper import alookup_model_source
-from comfy_pack.utils import get_self_git_commit
+from comfy_pack.package import build_bento
+
+import folder_paths
+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
+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",
-            "list",
-            "--format",
-            "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:
-        print("Package => Writing models")
-        models = await _get_models()
+    return manager_file_path / "snapshots"
+
+
+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)
+
+
+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
+        ):
+            del snapshot["pips"][package]
     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))
+        snapshot.update(
+            {
+                "python": f"{sys.version_info.major}.{sys.version_info.minor}",
+                "models": models,
+            }
+        )
+        f.write(json.dumps(snapshot, indent=2))
 
 
 def _is_port_in_use(port: int | str, host="localhost"):
@@ -88,7 +107,10 @@ def _is_file_refered(file_path: Path, workflow_api: dict) -> bool:
                 used_inputs.add(v)
     all_inputs = "\n".join(used_inputs)
     file_path = file_path.absolute().relative_to(folder_paths.base_path)
-    relpath = Path(*file_path.parts[2:])
+    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
 
 
@@ -110,16 +132,31 @@ async def _get_models(
 
     models = []
     model_filenames = [
-        os.path.abspath(line)
+        os.path.abspath(line.strip().strip('"').replace('\\\\', '/'))
         for line in stdout.decode().splitlines()
-        if not os.path.basename(line).startswith(".")
+        if not os.path.basename(line.strip().strip('"').replace('\\\\', '/')).startswith(".")
     ]
+
+    # 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(
-        model_filenames,
+        to_include,
         cache_only=not (ensure_sha or store_models),
     )
 
-    for filename in model_filenames:
+    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 = {
@@ -146,7 +183,7 @@ async def _get_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:
@@ -162,46 +199,6 @@ async def _get_models(
     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:
@@ -217,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)
@@ -248,10 +246,7 @@ async def pack_workspace(request):
 
     with zipfile.ZipFile(TEMP_FOLDER / zip_filename, "w") as zf:
         path = zipfile.Path(zf)
-        await _prepare_bento_project(
-            path,
-            data,
-        )
+        await _prepare_pack(path, data)
 
     return web.json_response({"download_url": f"/bentoml/download/{zip_filename}"})
 
@@ -266,6 +261,8 @@ class DevServer:
 
     @classmethod
     def start(cls, workflow_api: dict, port: int = 3000):
+        from comfy_pack import __file__ as comfy_pack_file
+
         cls.stop()
 
         cls.port = port
@@ -274,7 +271,7 @@ class DevServer:
         with cls.run_dir.joinpath("workflow_api.json").open("w") as f:
             f.write(json.dumps(workflow_api, indent=2))
         shutil.copy(
-            Path(__file__).with_name("service.py"),
+            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)
@@ -425,7 +422,7 @@ async def serve(request):
         )
 
 
-@PromptServer.instance.routes.post("/bentoml/serve/heartbeat")
+@PromptServer.instance.routes.get("/bentoml/serve/heartbeat")
 async def heartbeat(_):
     running = DevServer.feed_watch_dog()
 
@@ -447,35 +444,22 @@ async def download_workspace(request):
     return web.FileResponse(TEMP_FOLDER / zip_filename)
 
 
-async def _prepare_bento_project(
+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_requirements(working_dir, ["comfy-cli", "fastapi"])
     await _write_snapshot(working_dir, data, models)
     await _write_workflow(working_dir, data)
     await _write_inputs(working_dir, data)
-    with working_dir.joinpath("service.py").open("w") as f:
-        f.write(Path(__file__).with_name("service.py").read_text())
-
-    # Copy comfy_pack directory
-    if isinstance(working_dir, Path):
-        shutil.copytree(COMFY_PACK_DIR, working_dir / COMFY_PACK_DIR.name)
-    else:  # zipfile.Path
-        for src in COMFY_PACK_DIR.rglob("*"):
-            if src.is_file():
-                rel_path = src.relative_to(COMFY_PACK_DIR.parent)
-                with working_dir.joinpath(rel_path).open("wb") as f:
-                    f.write(src.read_bytes())
-    models = [m for m in models if not m["disabled"]]  # filter out disabled models
-    return models
 
 
 @PromptServer.instance.routes.post("/bentoml/model/query")
@@ -489,8 +473,37 @@ async def get_models(request):
     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,
@@ -509,30 +522,14 @@ async def build_bento(request):
 
     with tempfile.TemporaryDirectory(suffix="-bento", prefix="comfy-pack-") as temp_dir:
         temp_dir_path = Path(temp_dir)
-        models = await _prepare_bento_project(temp_dir_path, data, store_models=True)
+        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,
-                labels={"comfy-pack-version": get_self_git_commit() or "unknown"},
-                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",
 }

pyproject.toml

@@ -8,10 +8,9 @@ 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"]
 

requirements.txt

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

src/comfy_pack/cli.py

@@ -1,12 +1,16 @@
-import click
 import functools
 import json
-from pathlib import Path
+import os
 import shutil
 import subprocess
 import sys
 import tempfile
-from .const import WORKSPACE_DIR, COMFYUI_REPO, COMFY_PACK_REPO, COMFYUI_MANAGER_REPO
+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
 
@@ -27,6 +31,7 @@ def _ensure_uv() -> None:
 
 
 @click.group()
+@click.version_option()
 def main():
     """comfy-pack CLI"""
     pass
@@ -50,28 +55,27 @@ def main():
     help="Increase verbosity level",
 )
 def init(dir: str, verbose: int):
-    from rich.console import Console
     import os
 
-    console = Console()
+    import rich
 
     # Check if directory path is valid
     try:
         install_dir = Path(dir).absolute()
         if install_dir.exists() and not install_dir.is_dir():
-            console.print(f"[red]Error: {dir} exists but is not a directory[/red]")
+            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():
-                console.print(
+                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:
-        console.print(f"[red]Error: Invalid directory path - {str(e)}[/red]")
+        rich.print(f"[red]Error: Invalid directory path - {str(e)}[/red]")
         return 1
 
     # Check git installation
@@ -82,7 +86,7 @@ def init(dir: str, verbose: int):
             capture_output=True,
         )
     except (subprocess.SubprocessError, FileNotFoundError):
-        console.print("[red]Error: git is not installed or not in PATH[/red]")
+        rich.print("[red]Error: git is not installed or not in PATH[/red]")
         return 1
 
     # Check if we have write permissions
@@ -93,36 +97,36 @@ def init(dir: str, verbose: int):
         test_file.touch()
         test_file.unlink()
     except (OSError, PermissionError) as e:
-        console.print(f"[red]Error: No write permission in {dir} - {str(e)}[/red]")
+        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):
-        console.print("[red]Error: Python 3.8 or higher is required[/red]")
+        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:
-        console.print(f"[red]Error: {str(e)}[/red]")
+        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
-            console.print(
+            rich.print(
                 "[yellow]Warning: Less than 2GB free disk space available[/yellow]"
             )
     except Exception as e:
-        console.print(
+        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():
-        console.print("[green]Cloning ComfyUI...[/green]")
+        rich.print("[green]Cloning ComfyUI...[/green]")
         subprocess.run(
             [
                 "git",
@@ -134,7 +138,7 @@ def init(dir: str, verbose: int):
         )
 
     # Update ComfyUI
-    console.print("[green]Updating ComfyUI...[/green]")
+    rich.print("[green]Updating ComfyUI...[/green]")
     subprocess.run(
         ["git", "pull"],
         cwd=install_dir,
@@ -143,7 +147,7 @@ def init(dir: str, verbose: int):
 
     # Create and activate venv
     venv_dir = install_dir / ".venv"
-    console.print("[green]Creating virtual environment with uv...[/green]")
+    rich.print("[green]Creating virtual environment with uv...[/green]")
     if venv_dir.exists():
         shutil.rmtree(venv_dir)
     subprocess.run(
@@ -159,7 +163,7 @@ def init(dir: str, verbose: int):
         python = str(venv_dir / "bin" / "python")
 
     # Install requirements with uv
-    console.print("[green]Installing ComfyUI requirements with uv...[/green]")
+    rich.print("[green]Installing ComfyUI requirements with uv...[/green]")
     subprocess.run(
         ["uv", "pip", "install", "pip", "--upgrade"],
         env={
@@ -178,7 +182,7 @@ def init(dir: str, verbose: int):
     )
 
     # Install comfy-pack as custom node
-    console.print("[green]Installing comfy-pack custom nodes...[/green]")
+    rich.print("[green]Installing comfy-pack custom nodes...[/green]")
     custom_nodes_dir = install_dir / "custom_nodes"
     custom_nodes_dir.mkdir(exist_ok=True)
 
@@ -220,12 +224,12 @@ def init(dir: str, verbose: int):
         )
 
     version = get_self_git_commit() or "unknown"
-    console.print(
+    rich.print(
         f"\n[green]✓ Installation completed! (comfy-pack version: {version})[/green]"
     )
-    console.print(f"ComfyUI directory: {install_dir}")
+    rich.print(f"ComfyUI directory: {install_dir}")
 
-    console.print(
+    rich.print(
         "\n[green]Next steps:[/green]\n"
         f"1. cd {dir}\n"
         "2. source .venv/bin/activate  # On Windows: .venv\\Scripts\\activate\n"
@@ -237,7 +241,7 @@ def init(dir: str, verbose: int):
     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",
@@ -251,33 +255,70 @@ def init(dir: str, verbose: int):
     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)",
 )
-def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int):
+@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, verbose=verbose, all_models=include_disabled_models)
-    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="")
 
@@ -307,7 +348,7 @@ def _print_schema(schema, verbose: int = 0):
             range_str,
         )
 
-    Console().print(table)
+    rich.print(table)
 
 
 @functools.lru_cache
@@ -335,43 +376,40 @@ def _get_cache_workspace(cpack: str):
 )
 @click.pass_context
 def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
-    from .utils import generate_input_model
+    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:
-        console.print(
+        rich.print(
             'Usage: comfy-pack run [OPTIONS] CPACK --input1 "value1" --input2 "value2" ...'
         )
-        console.print("Run a ComfyUI package with the given inputs:")
+        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']}")
 
-        console.print("\n[yellow]Expected inputs:[/yellow]")
+        rich.print("\n[yellow]Expected inputs:[/yellow]")
         _print_schema(input_model.model_json_schema(), verbose)
         return 1
 
@@ -379,35 +417,135 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
 
     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, 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 ComfyUIServer, run_workflow
 
     with ComfyUIServer(str(workspace.absolute()), verbose=verbose) as server:
-        console.print("\n[green]✓ ComfyUI is launched in the background![/green]")
+        rich.print("\n[green]✓ ComfyUI is launched in the background![/green]")
         results = run_workflow(
             server.host,
             server.port,
             workflow,
             Path(output_dir).absolute(),
             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)
+            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

@@ -18,3 +18,5 @@ 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

@@ -113,7 +113,12 @@ async def async_batch_get_sha256(
     try:
         with SHA_CACHE_FILE.open("r") as f:
             cache = json.load(f)
-        cache.update(new_cache)
+    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):

src/comfy_pack/model_helper.py

@@ -1,8 +1,9 @@
-from .const import MODEL_SOURCE_CACHE_FILE
 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]+"
 
@@ -14,9 +15,8 @@ PATH_PATTERN = re.compile(
 
 
 async def _lookup_huggingface_model(model_sha: str) -> dict:
-    from duckduckgo_search import DDGS
     import aiohttp
-
+    from duckduckgo_search import DDGS
     query = f"site:huggingface.co blob {model_sha}"
 
     try:
@@ -25,7 +25,7 @@ async def _lookup_huggingface_model(model_sha: str) -> dict:
 
             async with aiohttp.ClientSession(trust_env=True) as session:
                 for result in search_results:
-                    url = result['link']
+                    url = result['href']
                     if "blob" not in url:
                         continue
 
@@ -38,11 +38,9 @@ async def _lookup_huggingface_model(model_sha: str) -> dict:
                                 repo, commit = commit_match.groups()
                                 if path_match := PATH_PATTERN.search(text):
                                     path = path_match.group(1)
-                                    download_url = f"https://huggingface.co/{repo}/resolve/{commit}/{path}?download=true"
-                                    url = f"https://huggingface.co/{repo}/blob/{commit}/{path}"
                                     info = {
-                                        "download_url": download_url,
-                                        "url": url,
+                                        "download_url": path,
+                                        "url": path,
                                         "repo": repo,
                                         "commit": commit,
                                         "path": path,

src/comfy_pack/package.py

@@ -1,44 +1,58 @@
 from __future__ import annotations
 
+import contextlib
 import json
 import os
 import shutil
-import urllib.parse
 import subprocess
 import sys
 import tempfile
 import threading
+import urllib.parse
 import urllib.request
-from .hash import get_sha256
 from pathlib import Path
-from .const import MODEL_DIR, COMFYUI_REPO
+from typing import TYPE_CHECKING
+
+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,
     )
 
 
@@ -53,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)
 
@@ -61,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
 
@@ -72,6 +101,7 @@ 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)
 
@@ -106,39 +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,
-):
+    no_deps: bool = False,
+    no_venv: bool = False,
+) -> Path:
     print("Installing Python dependencies")
-    python_version = snapshot["python"]
     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",
@@ -152,23 +182,28 @@ def install_dependencies(
         stderr=stderr,
     )
     if verbose > 0:
-        print(f"Installing dependencies from {req_file}")
+        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
 
 
@@ -183,13 +218,23 @@ def get_search_url(sha: str) -> str:
 
 def download_file(url: str, dest_path: Path, progress_callback=None):
     """Download file with progress tracking"""
-    if subprocess.call(["curl", "--version"], stdout=subprocess.DEVNULL) == 0:
-        subprocess.check_call(
-            ["curl", "-L", url, "-o", str(dest_path)],
-        )
-        return True
+
+    # 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:
-        with urllib.request.urlopen(url) as response:
+        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
@@ -238,7 +283,7 @@ def create_model_symlink(global_path: Path, sha: str, target_path: Path, filenam
     os.symlink(source, target)
 
 
-def retrive_models(
+def retrieve_models(
     snapshot: dict,
     workspace: Path,
     download: bool = True,
@@ -366,43 +411,139 @@ def retrive_models(
 def install(
     cpack: str | Path,
     workspace: str | Path = "workspace",
-    preheat: bool = True,
+    preheat: bool = False,
+    prepare_models: bool = True,
     all_models: bool = False,
+    no_venv: bool = False,
     verbose: int = 0,
 ):
     workspace = Path(workspace)
+    cpack = Path(cpack)
     print(f"Installing package {cpack} to {workspace} (verbose={verbose})")
-    with tempfile.TemporaryDirectory() as temp_dir:
-        pack_dir = Path(temp_dir) / ".cpack"
-        shutil.unpack_archive(cpack, pack_dir)
+    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)
+        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, dirs_exist_ok=True)
+        if prepare_models:
+            retrieve_models(
+                snapshot,
+                workspace,
+                verbose=verbose,
+                download=False,
+            )
 
-        retrive_models(
-            snapshot,
-            workspace,
-            verbose=verbose,
-            download=False,
-        )
-
-        install_custom_modules(snapshot, workspace, verbose=verbose)
         if preheat:
             from .run import ComfyUIServer
 
             with ComfyUIServer(
                 str(workspace),
                 verbose=verbose,
-                venv=str(workspace / ".venv"),
+                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)
 
-        retrive_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

@@ -5,6 +5,7 @@ import json
 import logging
 import os
 import random
+import shlex
 import shutil
 import socket
 import subprocess
@@ -13,8 +14,6 @@ import uuid
 from pathlib import Path
 from typing import Any, Union
 
-import psutil
-
 from .utils import populate_workflow, retrieve_workflow_outputs
 
 logger = logging.getLogger(__name__)
@@ -68,7 +67,6 @@ class ComfyUIServer:
         self.verbose = verbose
         self.host = host
         self.server_proc: subprocess.Popen | None = None
-        self._children_procs: list[psutil.Process] = []
 
         run_dir = (Path(workspace) / "cli_run").absolute()
         self.temp_dir = run_dir / "temp"
@@ -84,7 +82,7 @@ class ComfyUIServer:
                 self.port = 8188
         else:
             self.port = port
-        self.venv = venv
+        self.venv = os.path.abspath(venv) if venv else None
 
     def start(self) -> None:
         """
@@ -118,15 +116,10 @@ class ComfyUIServer:
         subprocess.run(command, check=True, stdout=stdout, env=env)
         logger.info("Successfully disabled Comfy CLI tracking")
 
-        logger.info("Preparing directories required by ComfyUI...")
-
         logger.info("Starting ComfyUI in the background...")
         command = [
-            "comfy",
-            "--workspace",
-            self.workspace,
-            "launch",
-            "--",
+            "python",
+            "main.py",
             "--output-directory",
             self.output_dir,
             "--temp-directory",
@@ -139,6 +132,8 @@ class ComfyUIServer:
 
         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()
@@ -149,6 +144,7 @@ class ComfyUIServer:
             stderr=None,
             preexec_fn=preexec_fn,
             env=env,
+            cwd=self.workspace,
         )
 
         if _wait_for_startup(self.host, self.port):
@@ -156,18 +152,12 @@ class ComfyUIServer:
             logger.info("Successfully started ComfyUI in the background")
         else:
             logger.error("Failed to start ComfyUI in the background")
-        self._children_procs = psutil.Process(self.server_proc.pid).children(
-            recursive=True
-        )
 
     def is_running(self) -> bool:
         if self.server_proc is None:
             return False
         if self.server_proc.poll() is not None:
             return False
-        for child in self._children_procs:
-            if not child.is_running():
-                return False
         return True
 
     def stop(self) -> None:
@@ -184,12 +174,8 @@ class ComfyUIServer:
         proc = self.server_proc
         self.server_proc = None
         logger.info("Stopping ComfyUI...")
-        for child in self._children_procs:
-            child.terminate()
-            child.wait()
         proc.terminate()
         proc.wait()
-        self._children_procs.clear()
         logger.info("Successfully stopped ComfyUI")
 
         logger.info("Cleaning up temporary directory...")
@@ -221,6 +207,7 @@ def run_workflow(
     output_dir: Union[str, Path, None] = None,
     timeout: int = 300,
     verbose: int = 0,
+    workspace: str = ".",
     **kwargs: Any,
 ) -> Any:
     """
@@ -273,6 +260,8 @@ def run_workflow(
     command = [
         "comfy",
         "--skip-prompt",
+        "--workspace",
+        workspace,
         "run",
         "--workflow",
         workflow_file_path.as_posix(),

src/comfy_pack/service.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import json
 import logging
 import os
-import shutil
 import signal
 import threading
 import time
@@ -41,7 +40,7 @@ def _get_workspace() -> Path:
     from bentoml._internal.configuration.containers import BentoMLContainer
 
     snapshot = BASE_DIR / "snapshot.json"
-    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
+    checksum = hashlib.md5(snapshot.read_bytes().strip()).hexdigest()
     wp = (
         Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
     )
@@ -74,15 +73,16 @@ else:
     snapshot = {}
 
 
-@bentoml.mount_asgi_app(app, path="/comfy")
+@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(
-                str(_get_workspace()),
+                workspace_override if workspace_override else str(_get_workspace()),
                 str(INPUT_DIR),
                 verbose=int("BENTOML_DEBUG" in os.environ),
             )
@@ -123,84 +123,57 @@ class ComfyService:
             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_shutdown
-    def on_shutdown(self):
-        logger.info("Shutting down")
-        if not EXISTING_COMFYUI_SERVER:
-            self.server.stop()
-            logger.info("server stopped")
-            self.watch_thread.join()
-            logger.info("Watch thread finished")
-
     @bentoml.on_deployment
     @staticmethod
-    def prepare_comfy_workspace():
+    def prepare_models():
         if EXISTING_COMFYUI_SERVER:
             return
-
-        from comfy_pack.package import install_comfyui, install_custom_modules
-
-        verbose = int("BENTOML_DEBUG" in os.environ)
         comfy_workspace = _get_workspace()
-
         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"]:
-                if model.get("disabled", False):
-                    continue
-                model_path = comfy_workspace / cast(str, model["filename"])
-                if model_tag := model.get("model_tag"):
+            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)
-                    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)
+                    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)
-                        logger.info("Copying %s to %s", model_file, model_path)
-                        model_path.symlink_to(model_file)
-                else:
-                    logger.warning(
-                        "Unrecognized model source: %s, the model may be missing",
-                        source,
-                    )
-
-            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()
+            else:
+                print(
+                    f"WARN: Unrecognized model source: {source}, the model may be missing"
+                )
 
 
-if not EXISTING_COMFYUI_SERVER:
+if False and not EXISTING_COMFYUI_SERVER:
     for model in snapshot["models"]:
-        if True:
-            continue
         if model.get("disabled"):
             continue
         source = model["source"]

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,18 +1,15 @@
 from __future__ import annotations
 
-import subprocess
-from pathlib import Path
 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",
@@ -210,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()
@@ -221,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.
 
@@ -230,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.
@@ -239,26 +236,47 @@ 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}_*"))
-    if len(outs) == 1:
-        return outs[0]
-    return outs
+        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
 
 
 def get_self_git_commit() -> str | None:

web/main.js

@@ -9,6 +9,48 @@ const spinner = `<svg xmlns="http://www.w3.org/2000/svg" width="32" height="32"
 </svg>`
 
 const style = `
+.cpack-tree-list {
+  max-height: 300px;
+  overflow-y: auto;
+  border: 1px solid #444;
+  border-radius: 4px;
+  padding: 5px;
+}
+
+.cpack-tree-item {
+  padding: 3px 0;
+}
+
+.cpack-tree-item label {
+  display: flex;
+  align-items: center;
+  gap: 5px;
+  min-height: 20px;
+}
+
+.cpack-tree-children {
+  margin-left: 12px;
+}
+
+.cpack-tree-toggle {
+  width: 12px;
+  height: 12px;
+  cursor: pointer;
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  color: #888;
+font-size: 0.8em;
+min-width: 12px;
+}
+
+.cpack-tree-toggle:hover {
+  color: #fff;
+}
+
+.cpack-tree-toggle.empty {
+  visibility: hidden;
+}
 .cpack-modal {
   position: fixed;
   top: 50%;
@@ -84,6 +126,10 @@ const style = `
   margin-bottom: 15px;
 }
 
+.cpack-form-item:last-child {
+  margin-bottom: -5px;
+}
+
 .cpack-form-item label {
   margin-bottom: 5px;
 }
@@ -131,10 +177,323 @@ const style = `
 }
 `
 
+class TreeState {
+  constructor() {
+    this.selectedFiles = new Set();
+    this.subscribers = new Set();
+  }
+
+  subscribe(callback) {
+    this.subscribers.add(callback);
+    return () => this.subscribers.delete(callback);
+  }
+
+  notify() {
+    this.subscribers.forEach(callback => callback(this.selectedFiles));
+  }
+
+  toggle(path, checked) {
+    if (checked) {
+      this.selectedFiles.add(path);
+    } else {
+      this.selectedFiles.delete(path);
+    }
+    this.notify();
+  }
+
+  toggleMultiple(paths, checked) {
+    paths.forEach(path => {
+      if (checked) {
+        this.selectedFiles.add(path);
+      } else {
+        this.selectedFiles.delete(path);
+      }
+    });
+    this.notify();
+  }
+
+  isSelected(path) {
+    return this.selectedFiles.has(path);
+  }
+
+  clear() {
+    this.selectedFiles.clear();
+    this.notify();
+  }
+
+  getSelectedCount() {
+    return this.selectedFiles.size;
+  }
+
+  getSelectedFiles() {
+    return Array.from(this.selectedFiles);
+  }
+}
+
+class FileTreeList {
+  constructor(container, countId) {
+    this.container = container;
+    this.countId = countId;
+    this.state = new TreeState();
+    this.init();
+  }
+
+  init() {
+    this.container.classList.add('cpack-tree-list');
+  this.state.subscribe(() => this.updateCount());
+  }
+
+  async load() {
+    try {
+      const { workflow, output: workflow_api } = await app.graphToPrompt();
+      const files = await this.getInputFiles(workflow, workflow_api);
+
+      // 预先将默认选中的文件添加到选中列表
+      files.forEach(file => {
+        const path = file.path || file;
+        if (file.checked) {
+          this.state.toggle(path, true);
+        }
+      });
+
+      this.renderTree(this.buildTree(files));
+
+    // 更新所有父目录的状态
+    this.container.querySelectorAll("[data-action='check-dir']").forEach(checkbox => {
+      this.updateFolderState(checkbox);
+    });
+
+    // 更新总数
+    this.updateCount();
+    } catch(e) {
+      this.container.innerHTML = `<div style="color: #ff8383">Failed to load files: ${e.message}</div>`;
+    }
+  }
+
+  async getInputFiles(workflow, workflow_api) {
+    const resp = await api.fetchApi("/bentoml/file/query", {
+      method: "POST",
+      body: JSON.stringify({ workflow, workflow_api }),
+      headers: { "Content-Type": "application/json" }
+    });
+    const data = await resp.json();
+    return Array.isArray(data) ? data : (data.files || []);
+  }
+
+  buildTree(files) {
+    const root = { name: 'root', children: {}, files: [] };
+
+    for (const file of files) {
+      const filePath = file.path || file;
+      const parts = filePath.split(/[\/\\]/);
+      let current = root;
+
+      for (let i = 0; i < parts.length; i++) {
+        const part = parts[i];
+        if (i === parts.length - 1) {
+          // 这是文件
+          current.files.push({
+            name: part,
+            path: file.path || file,
+            badges: file.badges || [],
+            checked: file.checked || false
+          });
+        // 对文件按名称排序
+        current.files.sort((a, b) => a.name.localeCompare(b.name));
+        } else {
+          // 这是目录
+          if (!current.children[part]) {
+            current.children[part] = {
+              name: part,
+              children: {},
+              files: []
+            };
+          }
+          current = current.children[part];
+        }
+      }
+    }
+
+    // 对目录按名称排序
+    for (const dir in root.children) {
+      root.children[dir].files.sort((a, b) => a.name.localeCompare(b.name));
+    }
+
+    return root;
+  }
+
+  renderTree(node, level = 0) {
+    const dirs = Object.values(node.children);
+    const hasChildren = dirs.length > 0 || node.files.length > 0;
+
+    this.container.innerHTML = this.renderNode(node, true);
+
+    this.setupEventListeners();
+    this.updateCount();
+
+  }
+
+  renderNode(node, isRoot = false) {
+    const dirs = Object.values(node.children);
+    const hasChildren = dirs.length > 0 || node.files.length > 0;
+    let html = '';
+
+    if (!isRoot) {
+      html += `
+        <div class="cpack-tree-item">
+          <label>
+            <span class="cpack-tree-toggle ${hasChildren ? '' : 'empty'}" data-action="toggle">
+              ${hasChildren ? '▶' : '▶'}
+            </span>
+            <input type="checkbox" data-action="check-dir" data-path="${node.name}" />
+            <span>${node.name}/</span>
+          </label>
+        </div>
+      `;
+    }
+
+    if (hasChildren) {
+      html += `<div class="cpack-tree-children" style="display: ${isRoot ? '' : 'none'}">`;
+
+      // 渲染子目录
+      for (const dir of dirs) {
+        html += this.renderNode(dir);
+      }
+
+      // 渲染文件
+      for (const file of node.files) {
+        html += `
+          <div class="cpack-tree-item">
+            <label style="padding-left: 17px">
+              <input type="checkbox" name="files" value="${file.path}"
+                ${this.state.isSelected(file.path) || file.checked ? 'checked' : ''} />
+              <span>${file.name}</span>
+              ${file.badges ? file.badges.map(badge =>
+                `<span style="background: ${badge.color || '#00a67d33'};
+                  color: ${badge.textColor || '#00a67d'};
+                  padding: 2px 6px;
+                  border-radius: 4px;
+                  font-size: 0.8em;
+                  cursor: help;
+                  white-space: nowrap;"
+                  title="${badge.tooltip || ''}">${badge.text}</span>`
+              ).join('') : ''}
+            </label>
+          </div>
+        `;
+      }
+
+      html += '</div>';
+    }
+
+    return html;
+  }
+
+  updateFolderState(folderCheckbox) {
+    const treeItem = folderCheckbox.closest('.cpack-tree-item');
+    const children = treeItem.nextElementSibling;
+    if (!children) return;
+
+    const childFiles = children.querySelectorAll("input[name='files']");
+    const checkedCount = Array.from(childFiles).filter(cb => cb.checked).length;
+
+    if (checkedCount === 0) {
+      folderCheckbox.checked = false;
+      folderCheckbox.indeterminate = false;
+    } else if (checkedCount === childFiles.length) {
+      folderCheckbox.checked = true;
+      folderCheckbox.indeterminate = false;
+    } else {
+      folderCheckbox.checked = false;
+      folderCheckbox.indeterminate = true;
+    }
+
+    // 递归更新父文件夹状态
+    const parentFolder = treeItem.parentElement.closest('.cpack-tree-item');
+    if (parentFolder) {
+      const parentCheckbox = parentFolder.querySelector("[data-action='check-dir']");
+      if (parentCheckbox) {
+        this.updateFolderState(parentCheckbox);
+      }
+    }
+
+    // 更新选中文件总数
+    this.updateCount();
+  }
+
+  setupEventListeners() {
+
+
+    // 目录选择功能
+    this.container.querySelectorAll("[data-action='check-dir']").forEach(checkbox => {
+      checkbox.addEventListener('change', (e) => {
+        const treeItem = e.target.closest('.cpack-tree-item');
+        const children = treeItem.nextElementSibling;
+        if (children) {
+          const childFiles = Array.from(children.querySelectorAll("input[name='files']"))
+            .map(input => input.value);
+          this.state.toggleMultiple(childFiles, e.target.checked);
+
+          // 更新UI
+          children.querySelectorAll("input[type='checkbox']").forEach(child => {
+            child.checked = e.target.checked;
+          if (child.hasAttribute('data-action')) {
+            child.indeterminate = false;
+          }
+          });
+        }
+
+      });
+    });
+
+    // 折叠功能
+    this.container.querySelectorAll("[data-action='toggle']").forEach(toggle => {
+      toggle.addEventListener('click', (e) => {
+        e.preventDefault();
+        e.stopPropagation();
+        const treeItem = e.target.closest('.cpack-tree-item');
+        const children = treeItem.nextElementSibling;
+        if (children) {
+          children.style.display = children.style.display === 'none' ? '' : 'none';
+          e.target.textContent = children.style.display === 'none' ? '▶' : '▼';
+        }
+      });
+    });
+
+    // 文件选择功能
+    this.container.querySelectorAll("input[name='files']").forEach(checkbox => {
+      checkbox.addEventListener('change', (e) => {
+        this.state.toggle(e.target.value, e.target.checked);
+
+      // 更新父文件夹状态
+      const parentFolder = checkbox.closest('.cpack-tree-children')
+        ?.previousElementSibling
+        ?.querySelector("[data-action='check-dir']");
+      if (parentFolder) {
+        this.updateFolderState(parentFolder);
+      }
+      });
+    });
+  }
+
+  updateCount() {
+
+    const countSpan = document.querySelector(`[data-files-count='${this.countId}']`);
+    if (countSpan) {
+      countSpan.textContent = this.state.getSelectedCount();
+    }
+  }
+
+  getSelectedFiles() {
+    return this.state.getSelectedFiles();
+  }
+}
+
 class ModelList {
   constructor(container, countId) {
     this.container = container;
     this.countId = countId;
+    this.selectedModels = new Set();
     this.init();
   }
 
@@ -254,8 +613,10 @@ class ModelList {
   }
 
   getSelectedModels() {
+    // 只返回用户选中的模型
     return Array.from(this.container.querySelectorAll("input[name='models']:checked"))
       .map(input => input.value);
+
   }
 }
 
@@ -292,14 +653,7 @@ async function createPackModal() {
         <label for="filename">Name</label>
         <input type="text" class="cpack-input" name="filename" value="${localStorage.getItem('cpack-bento-name') || 'comfy-pack-pkg'}" />
       </div>
-      <div class="cpack-form-item">
-        <details>
-          <summary style="cursor: pointer; margin-bottom: 10px;">Models (<span data-models-count="models-list">0</span> selected)</summary>
-          <div id="models-list">
-            ${spinner}
-          </div>
-        </details>
-      </div>
+      <div id="package-options-container"></div>
     `;
 
     const buttonContainer = document.createElement("div");
@@ -322,9 +676,11 @@ async function createPackModal() {
 
     const { close } = createModal(modal);
 
-    const modelsList = form.querySelector("#models-list");
-    const modelListComponent = new ModelList(modelsList, "models-list");
-    modelListComponent.load().then(() => {
+    const packageOptionsContainer = form.querySelector("#package-options-container");
+    const packageOptions = new PackageOptions(form, "pack-models-list", "pack-files-list", true);
+    packageOptionsContainer.innerHTML = packageOptions.getHtml();
+
+    packageOptions.init().then(() => {
       confirmButton.disabled = false;
     });
 
@@ -333,10 +689,13 @@ async function createPackModal() {
       if (filename) {
         // Save filename to localStorage
         localStorage.setItem('cpack-bento-name', filename);
+        const selectedData = packageOptions.getSelectedData();
         close();
         resolve({
           filename,
-          models: Array.from(form.querySelectorAll("input[name='models']:checked")).map(input => input.value)
+          models: selectedData.models,
+          files: selectedData.files,
+          systemPackages: selectedData.systemPackages
         });
       }
     };
@@ -416,7 +775,7 @@ async function unpackAction() {
   content.innerHTML = `
     <div style="margin-bottom: 20px;">
       <p style="margin-bottom: 15px;">To unpack a workflow package, follow these steps:</p>
-      
+
       <div style="margin-bottom: 15px;">
         <p style="margin-bottom: 10px;">1. Install comfy-pack CLI:</p>
         <div class="cpack-copyable" style="margin: 10px 0;">
@@ -445,13 +804,13 @@ async function unpackAction() {
       const closeButton = document.createElement("button");
       closeButton.textContent = "Close";
       closeButton.className = "cpack-btn";
-      
+
       buttonContainer.appendChild(closeButton);
-      
+
       modal.appendChild(title);
       modal.appendChild(content);
       modal.appendChild(buttonContainer);
-      
+
       const { close } = createModal(modal);
       closeButton.onclick = close;
 }
@@ -472,7 +831,9 @@ async function packageAction() {
     const body = JSON.stringify({
       workflow,
       workflow_api,
-      models: result.models
+      models: result.models,
+      files: result.files,
+      system_packages: result.systemPackages
     });
 
     downloadModal.updateProgress(60);
@@ -591,28 +952,102 @@ const serveForm = `
 </div>
 `
 
+class PackageOptions {
+  constructor(container, modelsListId, filesListId, defaultOpen = true) {
+    this.container = container;
+    this.modelsListId = modelsListId;
+    this.filesListId = filesListId;
+    this.modelListComponent = null;
+    this.fileListComponent = null;
+  this.defaultOpen = defaultOpen;
+  }
+
+  getHtml() {
+    return `
+      <div class="cpack-form-item">
+        <details ${this.defaultOpen ? 'open' : ''}>
+          <summary style="cursor: pointer; margin-bottom: 10px;">Package Options</summary>
+          <div style="padding: 10px; background: #1d1d1d; border-radius: 4px;">
+            <div class="cpack-form-item">
+              <details>
+                <summary style="cursor: pointer; margin-bottom: 10px;">Models (<span data-models-count="${this.modelsListId}">0</span> selected)</summary>
+                <div id="${this.modelsListId}">
+                  ${spinner}
+                </div>
+              </details>
+            </div>
+            <div class="cpack-form-item">
+              <details>
+                <summary style="cursor: pointer; margin-bottom: 10px;">Input Files (<span data-files-count="${this.filesListId}">0</span> selected)</summary>
+                <div id="${this.filesListId}">
+                  ${spinner}
+                </div>
+              </details>
+            </div>
+            <div class="cpack-form-item">
+              <details>
+                <summary style="cursor: pointer; margin-bottom: 10px;">System Packages</summary>
+                <div id="system-packages-array" style="padding: 10px;">
+                  <button class="cpack-btn" id="add-button" style="margin: 5px 0px">Add Package</button>
+                </div>
+              </details>
+            </div>
+          </div>
+        </details>
+      </div>
+    `;
+  }
+
+  async init() {
+    const modelsList = this.container.querySelector(`#${this.modelsListId}`);
+    const filesList = this.container.querySelector(`#${this.filesListId}`);
+
+    this.modelListComponent = new ModelList(modelsList, this.modelsListId);
+    this.fileListComponent = new FileTreeList(filesList, this.filesListId);
+
+    const addButton = this.container.querySelector("#add-button");
+    const systemPackagesArray = this.container.querySelector("#system-packages-array");
+
+    addButton.addEventListener("click", (e) => {
+      e.preventDefault();
+      const row = document.createElement("div");
+      row.className = "cpack-input-row";
+      row.innerHTML = `
+        <div style="flex: 1"><input type="text" class="cpack-input" name="systemPackages" placeholder="package name in Ubuntu" /></div>
+        <button class="cpack-btn" style="margin-left: 10px">Remove</button>
+      `
+      systemPackagesArray.appendChild(row);
+      row.querySelector("button").onclick = (e) => {
+        e.preventDefault();
+        row.remove();
+      }
+    });
+
+    await Promise.all([
+      this.modelListComponent.load(),
+      this.fileListComponent.load()
+    ]);
+  }
+
+  getSelectedData() {
+    return {
+      models: this.modelListComponent.getSelectedModels(),
+      files: this.fileListComponent.getSelectedFiles(),
+      systemPackages: Array.from(this.container.querySelectorAll("input[name='systemPackages']"))
+        .map(input => input.value)
+        .filter(Boolean)
+    };
+  }
+}
+
 const buildForm = `
   <p style="font-size: 0.85em; color: #888; margin-top: 5px;">
     This feature is powered by <a href="https://www.bentoml.com/?from=comfy-pack" target="_blank" style="color: #00a67d;">BentoCloud</a>, a platform for deploying <br>and managing ML services in customizable clusters
   </p>
 <div class="cpack-form-item required">
-  <label for="bentoName">Bento Name</label>
+  <label for="bentoName">Name</label>
   <input type="text" class="cpack-input" name="bentoName" placeholder="comfy-pack-app" />
-  <div class="error-message">Bento name is required</div>
-</div>
-<div class="cpack-form-item">
-  <details>
-    <summary style="cursor: pointer; margin-bottom: 10px;">Models (<span data-models-count="build-models-list">0</span> selected)</summary>
-    <div id="build-models-list">
-      ${spinner}
-    </div>
-  </details>
-</div>
-<div class="cpack-form-item">
-  <label for="systemPackages">System Packages</label>
-  <div id="system-packages-array">
-    <button class="cpack-btn" id="add-button" style="margin: 5px 0px">Add</button>
-  </div>
+  <div class="error-message">Name is required</div>
 </div>
 <div class="cpack-form-item required">
   <label>BentoCloud API</label>
@@ -628,8 +1063,11 @@ const buildForm = `
       <p style="font-size: 0.85em; color: #888; margin-top: 5px;">
         Get your API Token at <a href="https://cloud.bentoml.com/signup?from=comfy-pack" target="_blank" style="color: #00a67d;">cloud.bentoml.com</a>
       </p>
+
+    </div>
+
   </div>
-</div>
+  <div id="package-options-container"></div>
 `
 
 function createBuildModal() {
@@ -644,23 +1082,6 @@ function createBuildModal() {
   const form = document.createElement("form");
   form.innerHTML = buildForm;
 
-  const addButton = form.querySelector("#add-button");
-  const systemPackagesArray = form.querySelector("#system-packages-array");
-  addButton.addEventListener("click", (e) => {
-    e.preventDefault();
-    const row = document.createElement("div");
-    row.className = "cpack-input-row";
-    row.innerHTML = `
-      <div style="flex: 1"><input type="text" class="cpack-input" name="systemPackages" placeholder="package name in Ubuntu" /></div>
-      <button class="cpack-btn" style="margin-left: 10px">Remove</button>
-    `
-    systemPackagesArray.appendChild(row);
-    row.querySelector("button").onclick = (e) => {
-      e.preventDefault();
-      row.remove();
-    }
-  });
-
 
   const buttonContainer = document.createElement("div");
   buttonContainer.className = "cpack-btn-container";
@@ -688,9 +1109,11 @@ function createBuildModal() {
   return new Promise((resolve) => {
     form.querySelector("input[name='bentoName']").select();
 
-    const modelsList = form.querySelector("#build-models-list");
-    const modelListComponent = new ModelList(modelsList, "build-models-list");
-    modelListComponent.load().then(() => {
+    const packageOptionsContainer = form.querySelector("#package-options-container");
+    const packageOptions = new PackageOptions(form, "build-models-list", "build-files-list", false);
+    packageOptionsContainer.innerHTML = packageOptions.getHtml();
+
+    packageOptions.init().then(() => {
       confirmButton.disabled = false;
     });
 
@@ -729,13 +1152,15 @@ function createBuildModal() {
       localStorage.setItem('cpack-api-key', apiKey);
 
       const { workflow, output: workflow_api } = await app.graphToPrompt();
+      const selectedData = packageOptions.getSelectedData();
       const data = {
         bento_name: bentoName,
-        system_packages: Array.from(formData.getAll("systemPackages").filter(Boolean)),
+        system_packages: selectedData.systemPackages,
         push: true,
         api_key: apiKey,
         endpoint: endpoint,
-        models: Array.from(form.querySelectorAll("input[name='models']:checked")).map(input => input.value),
+        models: selectedData.models,
+        files: selectedData.files,
         workflow,
         workflow_api
       };
@@ -809,7 +1234,7 @@ async function createServeStatusModal(url) {
   // Start status checking
   const checkInterval = setInterval(async () => {
     try {
-      const resp = await api.fetchApi("/bentoml/serve/heartbeat", { method: "POST" });
+      const resp = await api.fetchApi("/bentoml/serve/heartbeat");
       const status = await resp.json();
       if (status.error) {
         title.textContent = "Server Error";