fix: remove on shutdown hook (#40)

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

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,150 @@ 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
 ```
 
+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 +165,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/api.py

@@ -11,6 +11,7 @@ import tempfile
 import time
 import uuid
 import zipfile
+from importlib.metadata import Distribution, distributions
 from pathlib import Path
 from typing import Any, Union
 
@@ -20,30 +21,49 @@ 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
 
 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"]  # TODO: standardize this
+
+
+def _get_requirement_string(dist: Distribution) -> str:
+    direct_url_text = dist.read_text("direct_url.json")
+    pinned_str = f"{dist.metadata['Name']}=={dist.version}"
+    if not direct_url_text:
+        return pinned_str
+    direct_url = json.loads(direct_url_text)
+    if url := direct_url.get("url"):
+        if url.startswith("file://"):
+            # we are not able to share local files
+            return pinned_str
+        if vcs_info := direct_url.get("vcs_info"):
+            url = f"{vcs_info['vcs']}+{url}@{vcs_info['commit_id']}"
+        if subdirectory := direct_url.get("subdirectory"):
+            url += f"#subdirectory={subdirectory}"
+        return f"{dist.metadata['Name']} @ {url}"
+    else:
+        return pinned_str
+
+
+def normalize_name(name: str) -> str:
+    import re
+
+    return re.sub(r"[-_.]+", "-", name).lower()
 
 
 async def _write_requirements(path: ZPath, extras: list[str] | None = None) -> None:
     print("Package => Writing requirements.txt")
+    seen_packages: set[str] = set(EXCLUDE_PACKAGES)
     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")
+        for dist in distributions():
+            pkg_name = normalize_name(dist.metadata["Name"])
+            if pkg_name in seen_packages:
+                continue
+            seen_packages.add(pkg_name)
+            f.write(_get_requirement_string(dist) + "\n")
         if extras:
             f.write("\n".join(extras) + "\n")
 
@@ -88,7 +108,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
 
 
@@ -146,7 +169,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:
@@ -192,6 +215,7 @@ async def _get_custom_nodes() -> list:
             "url": url,
             "commit_hash": commit_hash,
             "disabled": subdir.name.endswith(".disabled"),
+            "path": str(subdir.relative_to(custom_nodes)),
         }
 
     for subdir in Path(custom_nodes).iterdir():
@@ -217,15 +241,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 +273,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 +288,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 +298,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)
@@ -447,35 +471,21 @@ 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,
-):
+) -> None:
     model_filter = set(data.get("models", []))
     models = await _get_models(
         store_models=store_models,
         model_filter=model_filter,
     )
 
-    await _write_requirements(working_dir, ["comfy-cli", "fastapi"])
+    await _write_requirements(working_dir, ["comfy-cli", "fastapi", "comfy-pack"])
     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 +499,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 +548,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)
 
         # 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(

pyproject.toml

@@ -11,7 +11,6 @@ dependencies = [
     "bentoml>=1.3.13",
     "click>=8.1.7",
     "comfy-cli>=1.2.8",
-    "pydantic>=2.9",
 ]
 dynamic = ["version"]
 

requirements.txt

@@ -2,4 +2,3 @@ bentoml
 fastapi
 comfy-cli
 duckduckgo-search
-uv

src/comfy_pack/cli.py

@@ -1,12 +1,15 @@
-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
+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
 
@@ -50,9 +53,10 @@ def main():
     help="Increase verbosity level",
 )
 def init(dir: str, verbose: int):
-    from rich.console import Console
     import os
 
+    from rich.console import Console
+
     console = Console()
 
     # Check if directory path is valid
@@ -258,9 +262,10 @@ def init(dir: str, verbose: int):
     help="Increase verbosity level (use multiple times for more verbosity)",
 )
 def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int):
-    from .package import install
     from rich.console import Console
 
+    from .package import install
+
     console = Console()
 
     install(cpack, dir, verbose=verbose, all_models=include_disabled_models)
@@ -276,8 +281,8 @@ def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int
 
 
 def _print_schema(schema, verbose: int = 0):
-    from rich.table import Table
     from rich.console import Console
+    from rich.table import Table
 
     table = Table(title="")
 
@@ -335,10 +340,11 @@ 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
     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])
     )
@@ -398,6 +404,7 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
             workflow,
             Path(output_dir).absolute(),
             verbose=verbose,
+            workspace=server.workspace,
             **validated_data.model_dump(),
         )
         console.print("\n[green]✓ Workflow is executed successfully![/green]")
@@ -411,3 +418,46 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
                 console.print(f"{i}: {value}")
         else:
             console.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,
+        )

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/package.py

@@ -3,15 +3,23 @@ from __future__ import annotations
 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):
@@ -21,24 +29,28 @@ def _clone_commit(url: str, commit: str, dir: Path, verbose: int = 0):
         ["git", "clone", "--recurse-submodules", "--filter=blob:none", url, dir],
         stdout=stdout,
         stderr=stderr,
+        env={"GIT_TERMINAL_PROMPT": "0"},
     )
     subprocess.check_call(
         ["git", "fetch", "-q", url, commit],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env={"GIT_TERMINAL_PROMPT": "0"},
     )
     subprocess.check_call(
         ["git", "checkout", "FETCH_HEAD"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env={"GIT_TERMINAL_PROMPT": "0"},
     )
     subprocess.check_call(
         ["git", "submodule", "update", "--init", "--recursive"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env={"GIT_TERMINAL_PROMPT": "0"},
     )
 
 
@@ -61,6 +73,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
 
@@ -153,17 +168,22 @@ def install_dependencies(
     )
     if verbose > 0:
         print(f"Installing dependencies from {req_file}")
+    install_cmd = [
+        "uv",
+        "pip",
+        "install",
+        "-p",
+        str(venv_py),
+        "-r",
+        req_file,
+        "--no-deps",
+    ]
+    if STRICT_MODE:
+        pass
+    else:
+        install_cmd.extend(["--index-strategy", "unsafe-best-match"])
     subprocess.check_call(
-        [
-            "uv",
-            "pip",
-            "install",
-            "-p",
-            str(venv_py),
-            "-r",
-            req_file,
-            "--no-deps",
-        ],
+        install_cmd,
         stdout=stdout,
         stderr=stderr,
     )
@@ -238,7 +258,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,
@@ -387,7 +407,7 @@ def install(
             elif f.is_dir():
                 shutil.copytree(f, workspace / "input" / f.name, dirs_exist_ok=True)
 
-        retrive_models(
+        retrieve_models(
             snapshot,
             workspace,
             verbose=verbose,
@@ -405,4 +425,64 @@ def install(
             ) as _:
                 pass
 
-        retrive_models(snapshot, workspace, verbose=verbose, all_models=all_models)
+        retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models)
+
+
+required_files = ["snapshot.json", "requirements.txt"]
+
+
+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.py"
+    with Path(__file__).with_name("setup_workspace.py").open() as f:
+        setup_script.write_text(f.read().format(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.py").as_posix(),
+        },
+        python={"requirements_txt": "requirements.txt", "lock_packages": True},
+    )

src/comfy_pack/run.py

@@ -13,8 +13,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 +66,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"
@@ -122,11 +119,8 @@ class ComfyUIServer:
 
         logger.info("Starting ComfyUI in the background...")
         command = [
-            "comfy",
-            "--workspace",
-            self.workspace,
-            "launch",
-            "--",
+            "python",
+            "main.py",
             "--output-directory",
             self.output_dir,
             "--temp-directory",
@@ -149,6 +143,7 @@ class ComfyUIServer:
             stderr=None,
             preexec_fn=preexec_fn,
             env=env,
+            cwd=self.workspace,
         )
 
         if _wait_for_startup(self.host, self.port):
@@ -156,18 +151,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 +173,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 +206,7 @@ def run_workflow(
     output_dir: Union[str, Path, None] = None,
     timeout: int = 300,
     verbose: int = 0,
+    workspace: str = ".",
     **kwargs: Any,
 ) -> Any:
     """
@@ -273,6 +259,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
@@ -74,7 +73,7 @@ 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):
@@ -123,84 +122,57 @@ class ComfyService:
             output_dir=ctx.temp_dir,
             timeout=REQUEST_TIMEOUT,
             verbose=verbose,
+            workspace=self.server.workspace,
             **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.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+# ruff: noqa: E402
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+virtualenv = os.environ.get("VIRTUAL_ENV")
+if virtualenv and "--reload" not in sys.argv:
+    print("Re-executing in virtualenv:", virtualenv)
+    venv_python = os.path.join(virtualenv, "bin/python3")
+    os.execl(venv_python, venv_python, *sys.argv, "--reload")
+
+SNAPSHOT = """\
+{snapshot}
+"""
+
+
+def _get_workspace() -> tuple[Path, dict]:
+    import hashlib
+    import json
+
+    from bentoml._internal.configuration.containers import BentoMLContainer
+
+    checksum = hashlib.md5(SNAPSHOT.strip().encode("utf8")).hexdigest()
+    wp = (
+        Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
+    )
+    wp.parent.mkdir(parents=True, exist_ok=True)
+    return wp, json.loads(SNAPSHOT)
+
+
+def prepare_comfy_workspace():
+    import shutil
+
+    from comfy_pack.package import install_comfyui, install_custom_modules
+
+    verbose = 2
+    comfy_workspace, snapshot = _get_workspace()
+
+    if not comfy_workspace.joinpath(".DONE").exists():
+        if comfy_workspace.exists():
+            print("Removing existing workspace")
+            shutil.rmtree(comfy_workspace, ignore_errors=True)
+        install_comfyui(snapshot, comfy_workspace, verbose=verbose)
+
+        install_custom_modules(snapshot, comfy_workspace, verbose=verbose)
+        comfy_workspace.joinpath(".DONE").touch()
+        subprocess.run(
+            ["chown", "-R", "bentoml:bentoml", str(comfy_workspace)], check=True
+        )
+
+
+if __name__ == "__main__":
+    prepare_comfy_workspace()

uv.lock

@@ -456,7 +456,7 @@ name = "click"
 version = "8.1.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "platform_system == 'Windows'" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 }
 wheels = [
@@ -523,13 +523,11 @@ wheels = [
 
 [[package]]
 name = "comfy-pack"
-version = "0.1.4.dev10+g1d19c05"
 source = { editable = "." }
 dependencies = [
     { name = "bentoml" },
     { name = "click" },
     { name = "comfy-cli" },
-    { name = "pydantic" },
 ]
 
 [package.metadata]
@@ -537,7 +535,6 @@ requires-dist = [
     { name = "bentoml", specifier = ">=1.3.13" },
     { name = "click", specifier = ">=8.1.7" },
     { name = "comfy-cli", specifier = ">=1.2.8" },
-    { name = "pydantic", specifier = ">=2.9" },
 ]
 
 [[package]]

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
         });
       }
     };
@@ -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
       };