feat: unpack-bento command (#46)

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

pyproject.toml

@@ -8,7 +8,7 @@ authors = [{ name = "Frost Ming", email = "frost@bentoml.com" }]
 readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [
-    "bentoml>=1.3.13",
+    "bentoml>=1.3.21",
     "click>=8.1.7",
     "comfy-cli>=1.2.8",
 ]

requirements.txt

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

src/comfy_pack/cli.py

@@ -6,6 +6,7 @@ import subprocess
 import sys
 import tempfile
 from pathlib import Path
+from typing import cast
 
 import click
 
@@ -461,3 +462,82 @@ def bento_cmd(source: str, name: str | None, version: str | None):
             system_packages=system_packages,
             include_default_system_packages=include_default_system_packages,
         )
+
+
+def setup_cloud_client(
+    ctx: click.Context, param: click.Parameter, value: str | None
+) -> str | None:
+    from bentoml._internal.configuration.containers import BentoMLContainer
+
+    if value:
+        BentoMLContainer.cloud_context.set(value)
+        os.environ["BENTOML_CLOUD_CONTEXT"] = value
+    return value
+
+
+@main.command()
+@click.argument("bento")
+@click.option(
+    "-w",
+    "--workspace",
+    type=click.Path(file_okay=False),
+    default="workspace",
+    help="Workspace directory, defaults to './workspace'.",
+)
+@click.option("-v", "--verbose", count=True, help="Increase verbosity level")
+@click.option(
+    "--context",
+    help="BentoCloud context name.",
+    expose_value=False,
+    callback=setup_cloud_client,
+)
+def unpack_bento(bento: str, workspace: str, verbose: int):
+    """Restore the ComfyUI workspace from a given bento."""
+    import bentoml
+
+    from .package import install_comfyui, install_custom_modules, install_dependencies
+
+    bento_obj = bentoml.get(bento)
+    comfy_workspace = Path(workspace)
+    comfy_workspace.parent.mkdir(parents=True, exist_ok=True)
+    if not comfy_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, comfy_workspace, verbose=verbose)
+        reqs_txt = bento_obj.path_of("env/python/requirements.txt")
+        install_dependencies(snapshot, reqs_txt, comfy_workspace, verbose=verbose)
+
+        for f in Path(bento_obj.path_of("src/input")).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, dirs_exist_ok=True
+                )
+        install_custom_modules(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_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")
+        comfy_workspace.joinpath(".DONE").touch()
+
+    if os.name == "nt":
+        exe = "Scripts/python.exe"
+    else:
+        exe = "bin/python"
+    click.echo(
+        f"Workspace is ready at {comfy_workspace}\n"
+        f"You can start ComfyUI by running `cd {comfy_workspace} && .venv/{exe} main.py`",
+        color="green",
+    )

src/comfy_pack/package.py

@@ -25,32 +25,33 @@ COMFY_PACK_DIR = Path(__file__).parent
 def _clone_commit(url: str, commit: str, dir: Path, verbose: int = 0):
     stdout = None if verbose > 0 else subprocess.DEVNULL
     stderr = None if verbose > 1 else subprocess.DEVNULL
+    env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"}
     subprocess.check_call(
         ["git", "clone", "--recurse-submodules", "--filter=blob:none", url, dir],
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
+        env=env,
     )
     subprocess.check_call(
         ["git", "fetch", "-q", url, commit],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
+        env=env,
     )
     subprocess.check_call(
         ["git", "checkout", "FETCH_HEAD"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
+        env=env,
     )
     subprocess.check_call(
         ["git", "submodule", "update", "--init", "--recursive"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
+        env=env,
     )
 
 
@@ -131,12 +132,6 @@ def install_dependencies(
     stdout = None if verbose > 0 else subprocess.DEVNULL
     stderr = None if verbose > 1 else subprocess.DEVNULL
 
-    subprocess.check_call(
-        ["uv", "python", "install", python_version],
-        cwd=workspace,
-        stdout=stdout,
-        stderr=stderr,
-    )
     venv = workspace / ".venv"
     if (venv / "DONE").exists():
         return

src/comfy_pack/service.py

@@ -40,7 +40,7 @@ def _get_workspace() -> Path:
     from bentoml._internal.configuration.containers import BentoMLContainer
 
     snapshot = BASE_DIR / "snapshot.json"
-    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
+    checksum = hashlib.md5(snapshot.read_bytes().strip()).hexdigest()
     wp = (
         Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
     )
@@ -122,7 +122,7 @@ class ComfyService:
             output_dir=ctx.temp_dir,
             timeout=REQUEST_TIMEOUT,
             verbose=verbose,
-            workspace=self.server.workspace,
+            workspace=self.server.workspace if hasattr(self, "server") else ".",
             **kwargs,
         )
         if isinstance(ret, list):

src/comfy_pack/setup_workspace.py

@@ -22,12 +22,12 @@ def _get_workspace() -> tuple[Path, dict]:
     import hashlib
     import json
 
-    from bentoml._internal.configuration.containers import BentoMLContainer
+    bentoml_home = os.getenv(
+        "BENTOML_HOME", os.path.join(os.path.expanduser("~"), "bentoml")
+    )
 
     checksum = hashlib.md5(SNAPSHOT.strip().encode("utf8")).hexdigest()
-    wp = (
-        Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
-    )
+    wp = Path(bentoml_home) / "run" / "comfy_workspace" / checksum
     wp.parent.mkdir(parents=True, exist_ok=True)
     return wp, json.loads(SNAPSHOT)