fix: no venv when unpacking bento

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

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,7 +11,7 @@ import tempfile
 import time
 import uuid
 import zipfile
-from importlib.metadata import Distribution, distributions
+from fnmatch import fnmatch
 from pathlib import Path
 from typing import Any, Union
 
@@ -26,26 +26,7 @@ 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", "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
+EXCLUDE_PACKAGES = ["bentoml", "onnxruntime", "conda", "nvidia-*"]
 
 
 def normalize_name(name: str) -> str:
@@ -54,36 +35,54 @@ def normalize_name(name: str) -> str:
     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:
-        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")
-
-
-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"):
@@ -185,47 +184,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"),
-            "path": str(subdir.relative_to(custom_nodes)),
-        }
-
-    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:
@@ -449,7 +407,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()
 
@@ -475,14 +433,15 @@ 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", "comfy-pack"])
     await _write_snapshot(working_dir, data, models)
     await _write_workflow(working_dir, data)
     await _write_inputs(working_dir, data)
@@ -548,7 +507,7 @@ async def build_bento_api(request):
 
     with tempfile.TemporaryDirectory(suffix="-bento", prefix="comfy-pack-") as temp_dir:
         temp_dir_path = Path(temp_dir)
-        await _prepare_pack(temp_dir_path, data, store_models=True)
+        await _prepare_pack(temp_dir_path, data, store_models=True, ensure_source=False)
 
         # create a bento
         try:

src/comfy_pack/cli.py

@@ -7,7 +7,6 @@ import sys
 import tempfile
 import zipfile
 from pathlib import Path
-from typing import cast
 
 import click
 
@@ -58,27 +57,25 @@ def main():
 def init(dir: str, verbose: int):
     import os
 
-    from rich.console import Console
-
-    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
@@ -244,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",
@@ -258,24 +255,49 @@ 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):
-    from rich.console import Console
+def unpack_cmd(
+    cpack: str,
+    dir: str,
+    include_disabled_models: bool,
+    no_models: bool,
+    no_venv: bool,
+    verbose: int,
+):
+    import rich
 
     from .package import install
 
-    console = Console()
+    install(
+        cpack,
+        dir,
+        verbose=verbose,
+        all_models=include_disabled_models,
+        prepare_models=not no_models,
+        no_venv=no_venv,
+    )
+    rich.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
+    rich.print(f"{dir}")
 
-    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(
+    rich.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"
@@ -284,7 +306,7 @@ def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int
 
 
 def _print_schema(schema, verbose: int = 0):
-    from rich.console import Console
+    import rich
     from rich.table import Table
 
     table = Table(title="")
@@ -315,7 +337,7 @@ def _print_schema(schema, verbose: int = 0):
             range_str,
         )
 
-    Console().print(table)
+    rich.print(table)
 
 
 @functools.lru_cache
@@ -343,8 +365,8 @@ def _get_cache_workspace(cpack: str):
 )
 @click.pass_context
 def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
+    import rich
     from pydantic import ValidationError
-    from rich.console import Console
 
     from .utils import generate_input_model
 
@@ -352,8 +374,6 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
         zip([k.lstrip("-").replace("-", "_") for k in ctx.args[::2]], ctx.args[1::2])
     )
 
-    console = Console()
-
     with zipfile.ZipFile(cpack) as z:
         workflow = json.loads(z.read("workflow_api.json"))
 
@@ -361,24 +381,24 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
 
     # 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
 
@@ -495,7 +515,7 @@ def unpack_bento(bento: str, workspace: Path, verbose: int):
     """Restore the ComfyUI workspace from a given bento."""
     import bentoml
 
-    from .package import install_comfyui, install_custom_modules, install_dependencies
+    from .package import install
 
     try:
         bento_obj = bentoml.get(bento)
@@ -506,42 +526,8 @@ def unpack_bento(bento: str, workspace: Path, verbose: int):
         )
         bentoml.pull(bento)
         bento_obj = bentoml.get(bento)
-    workspace.parent.mkdir(parents=True, exist_ok=True)
-    if not workspace.joinpath(".DONE").exists():
-        for model in bento_obj.info.models:
-            model.to_model().resolve()
-        snapshot = json.loads(Path(bento_obj.path_of("src/snapshot.json")).read_text())
-        install_comfyui(snapshot, workspace, verbose=verbose)
-        workspace.joinpath(".DONE").unlink()  # created by install_comfyui
-        reqs_txt = bento_obj.path_of("env/python/requirements.txt")
-        if sys.platform != "linux":
-            src_reqs_txt = bento_obj.path_of("src/requirements.txt")
-            if os.path.exists(src_reqs_txt):
-                click.echo("Using requirements.txt from src directory")
-                reqs_txt = src_reqs_txt
-        install_dependencies(snapshot, reqs_txt, workspace, verbose=verbose)
 
-        for f in Path(bento_obj.path_of("src/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)
-        install_custom_modules(snapshot, workspace, verbose=verbose)
-        for model in snapshot["models"]:
-            if model.get("disabled", False):
-                continue
-            model_path = 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")
-                click.echo(f"Copying {model_file} to {model_path}")
-                model_path.symlink_to(model_file)
-            else:
-                click.echo("WARN: Unrecognized model source, the model may be missing")
-        workspace.joinpath(".DONE").write_text(snapshot["comfyui"])
+    install(bento_obj.path_of("src"), workspace, verbose=verbose, prepare_models=False)
 
     if os.name == "nt":
         exe = "Scripts/python.exe"

src/comfy_pack/package.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import json
 import os
 import shutil
@@ -66,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)
 
@@ -123,33 +136,39 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
 
 
 def install_dependencies(
-    snapshot: dict,
+    python_version: str,
     req_file: str,
     workspace: Path,
     verbose: int = 0,
+    no_deps: bool = False,
+    no_venv: bool = False,
 ):
     print("Installing Python dependencies")
-    python_version = snapshot["python"]
     stdout = None if verbose > 0 else subprocess.DEVNULL
     stderr = None if verbose > 1 else subprocess.DEVNULL
-
-    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"
+        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,
+        )
     subprocess.check_call(
         [
             "uv",
@@ -172,12 +191,11 @@ def install_dependencies(
         str(venv_py),
         "-r",
         req_file,
-        "--no-deps",
     ]
-    if STRICT_MODE:
-        pass
-    else:
+    if not STRICT_MODE:
         install_cmd.extend(["--index-strategy", "unsafe-best-match"])
+    if no_deps:
+        install_cmd.append("--no-deps")
     subprocess.check_call(
         install_cmd,
         stdout=stdout,
@@ -379,52 +397,78 @@ def retrieve_models(
                 continue
 
 
+def _ensure_startup_script_path(workspace: Path) -> Path:
+    script_path = workspace / "user" / "default" / "ComfyUI-Manager" / "startup-scripts"
+    script_path.mkdir(parents=True, exist_ok=True)
+    return script_path
+
+
 def install(
     cpack: str | Path,
     workspace: str | Path = "workspace",
     preheat: bool = True,
+    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)
+        install_dependencies(
+            snapshot["python"],
+            str(workspace / "requirements.txt"),
+            workspace,
+            no_venv=no_venv,
+            verbose=verbose,
+        )
+        shutil.copy2(
+            (pack_dir / "snapshot.json"),
+            _ensure_startup_script_path(workspace) / "restore-snapshot.json",
+        )
 
         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,
+            )
 
-        retrieve_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
-
-        retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models)
+        if prepare_models:
+            retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models)
 
 
-required_files = ["snapshot.json", "requirements.txt"]
+required_files = ["snapshot.json"]
 
 
 def build_bento(
@@ -457,9 +501,9 @@ def build_bento(
 
     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))
+    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)
@@ -478,7 +522,10 @@ def build_bento(
         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(),
+            "setup_script": source_dir.joinpath("setup_workspace.sh").as_posix(),
+        },
+        python={
+            "lock_packages": True,
+            "packages": ["fastapi", "comfy-pack>=0.4.0a0"],
         },
-        python={"requirements_txt": "requirements.txt", "lock_packages": True},
     )

src/comfy_pack/run.py

@@ -81,7 +81,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:
         """

src/comfy_pack/setup_workspace.py

@@ -1,57 +0,0 @@
-#!/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
-
-    bentoml_home = os.getenv(
-        "BENTOML_HOME", os.path.join(os.path.expanduser("~"), "bentoml")
-    )
-
-    checksum = hashlib.md5(SNAPSHOT.strip().encode("utf8")).hexdigest()
-    wp = Path(bentoml_home) / "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()

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

web/main.js

@@ -247,7 +247,7 @@ class FileTreeList {
     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;
@@ -255,14 +255,14 @@ class FileTreeList {
           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) {
@@ -282,12 +282,12 @@ class FileTreeList {
 
   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) {
@@ -313,24 +313,24 @@ class FileTreeList {
         }
       }
     }
-    
+
     // 对目录按名称排序
     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) {
@@ -354,35 +354,35 @@ class FileTreeList {
 
     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}" 
+              <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; 
+              ${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;" 
+                  white-space: nowrap;"
                   title="${badge.tooltip || ''}">${badge.text}</span>`
               ).join('') : ''}
             </label>
           </div>
         `;
       }
-      
+
       html += '</div>';
     }
 
@@ -422,7 +422,7 @@ class FileTreeList {
   }
 
   setupEventListeners() {
-    
+
 
     // 目录选择功能
     this.container.querySelectorAll("[data-action='check-dir']").forEach(checkbox => {
@@ -433,7 +433,7 @@ class FileTreeList {
           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;
@@ -442,7 +442,7 @@ class FileTreeList {
           }
           });
         }
-        
+
       });
     });
 
@@ -464,7 +464,7 @@ class FileTreeList {
     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
@@ -477,7 +477,7 @@ class FileTreeList {
   }
 
   updateCount() {
-    
+
     const countSpan = document.querySelector(`[data-files-count='${this.countId}']`);
     if (countSpan) {
       countSpan.textContent = this.state.getSelectedCount();
@@ -616,7 +616,7 @@ class ModelList {
     // 只返回用户选中的模型
     return Array.from(this.container.querySelectorAll("input[name='models']:checked"))
       .map(input => input.value);
-  
+
   }
 }
 
@@ -679,7 +679,7 @@ async function createPackModal() {
     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;
     });
@@ -775,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;">
@@ -804,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;
 }
@@ -1001,13 +1001,13 @@ class PackageOptions {
   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");
@@ -1063,7 +1063,7 @@ 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>
@@ -1112,7 +1112,7 @@ function createBuildModal() {
     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;
     });
@@ -1234,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";