feat: unpack-bento command (#46)

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

nodes/api.py

@@ -31,7 +31,7 @@ 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}'
+    pinned_str = f"{dist.metadata['Name']}=={dist.version}"
     if not direct_url_text:
         return pinned_str
     direct_url = json.loads(direct_url_text)
@@ -48,10 +48,21 @@ def _get_requirement_string(dist: Distribution) -> str:
         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:
         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")
@@ -158,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:
@@ -277,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
@@ -285,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)
@@ -469,7 +482,7 @@ async def _prepare_pack(
         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)
@@ -542,7 +555,7 @@ async def build_bento_api(request):
             bento = build_bento(
                 data["bento_name"],
                 temp_dir_path,
-                data.get("system_packages"),
+                system_packages=data.get("system_packages"),
             )
         except bentoml.exceptions.BentoMLException as e:
             return web.json_response(

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
 
@@ -404,6 +405,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]")
@@ -443,6 +445,13 @@ def bento_cmd(source: str, name: str | None, version: str | None):
             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
 
@@ -453,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,28 +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=env,
     )
     subprocess.check_call(
         ["git", "fetch", "-q", url, commit],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
     subprocess.check_call(
         ["git", "checkout", "FETCH_HEAD"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
     subprocess.check_call(
         ["git", "submodule", "update", "--init", "--recursive"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
+        env=env,
     )
 
 
@@ -127,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
@@ -456,11 +455,14 @@ def build_bento(
         system_packages = system_packages or []
 
     shutil.copy2(Path(__file__).with_name("service.py"), source_dir / "service.py")
-    # Copy comfy-pack package
-    shutil.copytree(
-        COMFY_PACK_DIR, source_dir / COMFY_PACK_DIR.name, dirs_exist_ok=True
-    )
-    snapshot = json.loads((source_dir / "snapshot.json").read_text())
+    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,
@@ -475,7 +477,7 @@ def build_bento(
         docker={
             "python_version": f"{sys.version_info.major}.{sys.version_info.minor}",
             "system_packages": system_packages,
-            "setup_script": Path(__file__).with_name("setup_workspace.py").as_posix(),
+            "setup_script": source_dir.joinpath("setup_workspace.py").as_posix(),
         },
         python={"requirements_txt": "requirements.txt", "lock_packages": True},
     )

src/comfy_pack/run.py

@@ -206,6 +206,7 @@ def run_workflow(
     output_dir: Union[str, Path, None] = None,
     timeout: int = 300,
     verbose: int = 0,
+    workspace: str = ".",
     **kwargs: Any,
 ) -> Any:
     """
@@ -258,6 +259,8 @@ def run_workflow(
     command = [
         "comfy",
         "--skip-prompt",
+        "--workspace",
+        workspace,
         "run",
         "--workflow",
         workflow_file_path.as_posix(),

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
     )
@@ -73,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):
@@ -122,19 +122,13 @@ class ComfyService:
             output_dir=ctx.temp_dir,
             timeout=REQUEST_TIMEOUT,
             verbose=verbose,
+            workspace=self.server.workspace if hasattr(self, "server") else ".",
             **kwargs,
         )
         if isinstance(ret, list):
             ret = ret[-1]
         return ret
 
-    @bentoml.on_shutdown
-    def on_shutdown(self):
-        logger.info("Shutting down")
-        if not EXISTING_COMFYUI_SERVER:
-            self.watch_thread.join()
-            logger.info("Watch thread finished")
-
     @bentoml.on_deployment
     @staticmethod
     def prepare_models():

src/comfy_pack/setup_workspace.py

@@ -13,25 +13,23 @@ if virtualenv and "--reload" not in sys.argv:
     venv_python = os.path.join(virtualenv, "bin/python3")
     os.execl(venv_python, venv_python, *sys.argv, "--reload")
 
-# The script path is ./env/docker/setup_script
-SRC_DIR = Path(__file__).parent.parent.parent / "src"
-INPUT_DIR = SRC_DIR / "input"
-sys.path.append(str(SRC_DIR))
+SNAPSHOT = """\
+{snapshot}
+"""
 
 
 def _get_workspace() -> tuple[Path, dict]:
     import hashlib
     import json
 
-    from bentoml._internal.configuration.containers import BentoMLContainer
-
-    snapshot = SRC_DIR / "snapshot.json"
-    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
-    wp = (
-        Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
+    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.read_text())
+    return wp, json.loads(SNAPSHOT)
 
 
 def prepare_comfy_workspace():
@@ -39,7 +37,7 @@ def prepare_comfy_workspace():
 
     from comfy_pack.package import install_comfyui, install_custom_modules
 
-    verbose = int("BENTOML_DEBUG" in os.environ)
+    verbose = 2
     comfy_workspace, snapshot = _get_workspace()
 
     if not comfy_workspace.joinpath(".DONE").exists():
@@ -48,12 +46,6 @@ def prepare_comfy_workspace():
             shutil.rmtree(comfy_workspace, ignore_errors=True)
         install_comfyui(snapshot, comfy_workspace, verbose=verbose)
 
-        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()
         subprocess.run(

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,7 +523,6 @@ wheels = [
 
 [[package]]
 name = "comfy-pack"
-version = "0.2.2.dev7+g5238115.d20241231"
 source = { editable = "." }
 dependencies = [
     { name = "bentoml" },