feat(nodes): new node: Save Audio

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

.gitignore

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

nodes/nodes.py

@@ -1,14 +1,17 @@
 import hashlib
 import json
 import os
-import sys
 import shutil
+import sys
+import zipfile
+from io import BytesIO
 
 import folder_paths
 import node_helpers
 import numpy as np
 import torch
-from PIL import Image, ImageOps, ImageSequence
+from comfy_extras.nodes_audio import SaveAudio
+from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
 from PIL.PngImagePlugin import PngInfo
 
 from .monkeypatch import set_bentoml_output
@@ -163,6 +166,100 @@ 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)
 
@@ -356,6 +453,8 @@ class AnyInput:
 NODE_CLASS_MAPPINGS = {
     "CPackOutputFile": OutputFile,
     "CPackOutputImage": OutputImage,
+    "CPackOutputAudio": SaveAudio,
+    "CPackOutputZip": OutputImageWithStringTxt,
     "CPackInputImage": ImageInput,
     "CPackInputString": StringInput,
     "CPackInputInt": IntInput,
@@ -370,5 +469,7 @@ 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)",
 }

src/comfy_pack/cli.py

@@ -5,7 +5,9 @@ import shutil
 import subprocess
 import sys
 import tempfile
+import zipfile
 from pathlib import Path
+from typing import cast
 
 import click
 
@@ -30,6 +32,7 @@ def _ensure_uv() -> None:
 
 
 @click.group()
+@click.version_option()
 def main():
     """comfy-pack CLI"""
     pass
@@ -351,10 +354,8 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
 
     console = Console()
 
-    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())
+    with zipfile.ZipFile(cpack) as z:
+        workflow = json.loads(z.read("workflow_api.json"))
 
     input_model = generate_input_model(workflow)
 
@@ -461,3 +462,93 @@ 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,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,
     )
 
 
@@ -87,6 +88,7 @@ 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)
 
@@ -131,12 +133,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

@@ -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/utils.py

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