fix: build-bento

.gitignore

@@ -161,5 +161,3 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 /temp/
-src/.DS_Store
-.DS_Store

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,21 +48,10 @@ 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")
@@ -169,7 +158,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:
@@ -288,8 +277,6 @@ 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
@@ -298,7 +285,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(comfy_pack_file).with_name("service.py"),
+            Path(__file__).with_name("service.py"),
             cls.run_dir / "service.py",
         )
         shutil.copytree(COMFY_PACK_DIR, cls.run_dir / COMFY_PACK_DIR.name)
@@ -482,7 +469,7 @@ async def _prepare_pack(
         model_filter=model_filter,
     )
 
-    await _write_requirements(working_dir, ["comfy-cli", "fastapi", "comfy-pack"])
+    await _write_requirements(working_dir, ["comfy-cli", "fastapi"])
     await _write_snapshot(working_dir, data, models)
     await _write_workflow(working_dir, data)
     await _write_inputs(working_dir, data)
@@ -555,7 +542,7 @@ async def build_bento_api(request):
             bento = build_bento(
                 data["bento_name"],
                 temp_dir_path,
-                system_packages=data.get("system_packages"),
+                data.get("system_packages"),
             )
         except bentoml.exceptions.BentoMLException as e:
             return web.json_response(

nodes/nodes.py

@@ -1,17 +1,14 @@
 import hashlib
 import json
 import os
-import shutil
 import sys
-import zipfile
-from io import BytesIO
+import shutil
 
 import folder_paths
 import node_helpers
 import numpy as np
 import torch
-from comfy_extras.nodes_audio import SaveAudio
-from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
+from PIL import Image, ImageOps, ImageSequence
 from PIL.PngImagePlugin import PngInfo
 
 from .monkeypatch import set_bentoml_output
@@ -166,100 +163,6 @@ class OutputImage:
         return {"ui": {"images": results}}
 
 
-class OutputImageWithStringTxt:
-    COLOR = (142, 36, 170)
-
-    def __init__(self):
-        self.output_dir = folder_paths.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
-        self.compress_level = 4
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "images": ("IMAGE", {"tooltip": "The images to save."}),
-                "filename_prefix": ("STRING", {"default": "cpack_output_"}),
-                "text": ("STRING", {"default": ""}),
-            },
-            "hidden": {
-                "prompt": "PROMPT",
-                "extra_pnginfo": "EXTRA_PNGINFO",
-            },
-        }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_images"
-    CPACK_NODE = True
-    OUTPUT_NODE = True
-
-    CATEGORY = "ComfyPack/output"
-    DESCRIPTION = (
-        "Saves the input images (and optional text) to your ComfyUI output directory."
-    )
-
-    def save_images(
-        self,
-        images,
-        filename_prefix="cpack_output_",
-        text="",
-        prompt=None,
-        extra_pnginfo=None,
-    ):
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = (
-            get_save_image_path(
-                filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
-            )
-        )
-
-        base_counter = counter  # use for name zip
-        zip_filename = f"{filename}_batch_{base_counter:05}.zip"
-        zip_path = os.path.join(full_output_folder, zip_filename)
-
-        # create ZIP file
-        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
-            for batch_number, image in enumerate(images):
-                # temp store img to RAM
-                i = 255.0 * image.cpu().numpy()
-                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
-
-                # write meta data
-                metadata = PngImagePlugin.PngInfo()
-                if prompt is not None:
-                    metadata.add_text("prompt", json.dumps(prompt))
-                if extra_pnginfo is not None:
-                    for x in extra_pnginfo:
-                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
-
-                # write img file to RAM buffer
-                img_buffer = BytesIO()
-                img.save(
-                    img_buffer,
-                    format="PNG",
-                    pnginfo=metadata,
-                    compress_level=self.compress_level,
-                )
-                img_buffer.seek(0)
-
-                # write img into ZIP file
-                image_filename = f"image_{batch_number:05}.png"
-                zipf.writestr(image_filename, img_buffer.read())
-
-                # write txt into ZIP file
-                text_filename = f"text_{batch_number:05}.txt"
-                zipf.writestr(text_filename, text)
-
-        # return zip as output
-        out = [{"filename": zip_filename, "subfolder": subfolder, "type": "zip"}]
-        return {
-            "ui": {
-                "zip": out,
-            }
-        }
-
-
 class ImageInput:
     COLOR = (142, 36, 170)
 
@@ -453,8 +356,6 @@ class AnyInput:
 NODE_CLASS_MAPPINGS = {
     "CPackOutputFile": OutputFile,
     "CPackOutputImage": OutputImage,
-    "CPackOutputAudio": SaveAudio,
-    "CPackOutputZip": OutputImageWithStringTxt,
     "CPackInputImage": ImageInput,
     "CPackInputString": StringInput,
     "CPackInputInt": IntInput,
@@ -469,7 +370,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CPackInputFile": "File Input",
     "CPackInputAny": "Any Input",
     "CPackOutputImage": "Image Output",
-    "CPackOutputAudio": "Audio Output",
     "CPackOutputFile": "File Output",
-    "CPackOutputZip": "Zip Output(img + txt file)",
 }

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.21",
+    "bentoml>=1.3.13",
     "click>=8.1.7",
     "comfy-cli>=1.2.8",
 ]

requirements.txt

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

src/comfy_pack/cli.py

@@ -5,9 +5,7 @@ import shutil
 import subprocess
 import sys
 import tempfile
-import zipfile
 from pathlib import Path
-from typing import cast
 
 import click
 
@@ -32,7 +30,6 @@ def _ensure_uv() -> None:
 
 
 @click.group()
-@click.version_option()
 def main():
     """comfy-pack CLI"""
     pass
@@ -354,8 +351,10 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
 
     console = Console()
 
-    with zipfile.ZipFile(cpack) as z:
-        workflow = json.loads(z.read("workflow_api.json"))
+    with tempfile.TemporaryDirectory() as temp_dir:
+        pack_dir = Path(temp_dir) / ".cpack"
+        shutil.unpack_archive(cpack, pack_dir)
+        workflow = json.loads((pack_dir / "workflow_api.json").read_text())
 
     input_model = generate_input_model(workflow)
 
@@ -405,7 +404,6 @@ 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]")
@@ -445,13 +443,6 @@ 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
 
@@ -462,93 +453,3 @@ 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, path_type=Path),
-    default="workspace",
-    help="Workspace directory, defaults to './workspace'.",
-)
-@click.option("-v", "--verbose", count=True, help="Increase verbosity level")
-@click.option(
-    "--context",
-    help="BentoCloud context name.",
-    expose_value=False,
-    callback=setup_cloud_client,
-)
-def unpack_bento(bento: str, workspace: Path, verbose: int):
-    """Restore the ComfyUI workspace from a given bento."""
-    import bentoml
-
-    from .package import install_comfyui, install_custom_modules, install_dependencies
-
-    try:
-        bento_obj = bentoml.get(bento)
-    except bentoml.exceptions.NotFound:
-        click.echo(
-            f"Bento {bento} not found in the local repository, trying to pull from BentoCloud",
-            err=True,
-        )
-        bentoml.pull(bento)
-        bento_obj = bentoml.get(bento)
-    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"])
-
-    if os.name == "nt":
-        exe = "Scripts/python.exe"
-    else:
-        exe = "bin/python"
-    click.echo(
-        f"Workspace is ready at {workspace}\n"
-        f"You can start ComfyUI by running `cd {workspace} && .venv/{exe} main.py`",
-        color="green",
-    )

src/comfy_pack/package.py

@@ -25,33 +25,28 @@ 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,
     )
 
 
@@ -88,7 +83,6 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
                     continue
             shutil.rmtree(module_dir)
 
-        print(f"Installing custom node {url}")
         commit_hash = module["commit_hash"]
         _clone_commit(url, commit_hash, module_dir, verbose=verbose)
 
@@ -133,6 +127,12 @@ 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,14 +456,11 @@ def build_bento(
         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)
+    # 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())
     return bentoml.build(
         "service:ComfyService",
         name=bento_name,
@@ -478,7 +475,7 @@ 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": Path(__file__).with_name("setup_workspace.py").as_posix(),
         },
         python={"requirements_txt": "requirements.txt", "lock_packages": True},
     )

src/comfy_pack/run.py

@@ -206,7 +206,6 @@ def run_workflow(
     output_dir: Union[str, Path, None] = None,
     timeout: int = 300,
     verbose: int = 0,
-    workspace: str = ".",
     **kwargs: Any,
 ) -> Any:
     """
@@ -259,8 +258,6 @@ 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().strip()).hexdigest()
+    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
     wp = (
         Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
     )
@@ -73,7 +73,7 @@ else:
     snapshot = {}
 
 
-@bentoml.asgi_app(app, path="/comfy")
+@bentoml.mount_asgi_app(app, path="/comfy")
 @bentoml.service(traffic={"timeout": REQUEST_TIMEOUT * 2}, resources={"gpu": 1})
 class ComfyService:
     def __init__(self):
@@ -122,13 +122,19 @@ 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,23 +13,25 @@ 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")
 
-SNAPSHOT = """\
-{snapshot}
-"""
+# 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))
 
 
 def _get_workspace() -> tuple[Path, dict]:
     import hashlib
     import json
 
-    bentoml_home = os.getenv(
-        "BENTOML_HOME", os.path.join(os.path.expanduser("~"), "bentoml")
-    )
+    from bentoml._internal.configuration.containers import BentoMLContainer
 
-    checksum = hashlib.md5(SNAPSHOT.strip().encode("utf8")).hexdigest()
-    wp = Path(bentoml_home) / "run" / "comfy_workspace" / checksum
+    snapshot = SRC_DIR / "snapshot.json"
+    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
+    wp = (
+        Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
+    )
     wp.parent.mkdir(parents=True, exist_ok=True)
-    return wp, json.loads(SNAPSHOT)
+    return wp, json.loads(snapshot.read_text())
 
 
 def prepare_comfy_workspace():
@@ -37,7 +39,7 @@ def prepare_comfy_workspace():
 
     from comfy_pack.package import install_comfyui, install_custom_modules
 
-    verbose = 2
+    verbose = int("BENTOML_DEBUG" in os.environ)
     comfy_workspace, snapshot = _get_workspace()
 
     if not comfy_workspace.joinpath(".DONE").exists():
@@ -46,6 +48,12 @@ 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(

src/comfy_pack/utils.py

@@ -12,7 +12,6 @@ if TYPE_CHECKING:
 CPACK_OUTPUT_NODES = {
     "CPackOutputFile",
     "CPackOutputImage",
-    "CPackOutputZip"
 }
 
 CPACK_PATH_INPUT_NODES = {

uv.lock

@@ -456,7 +456,7 @@ name = "click"
 version = "8.1.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "colorama", marker = "platform_system == 'Windows'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 }
 wheels = [
@@ -523,6 +523,7 @@ wheels = [
 
 [[package]]
 name = "comfy-pack"
+version = "0.2.2.dev7+g5238115.d20241231"
 source = { editable = "." }
 dependencies = [
     { name = "bentoml" },