fix: Model hashing fixes (#62)

* make sure hashing is done only when the models are selected

* fix cache file not being written, if it did not exist

---------

Co-authored-by: Frost Ming <me@frostming.com>

nodes/api.py

@@ -15,14 +15,14 @@ from fnmatch import fnmatch
 from pathlib import Path
 from typing import Any, Union
 
-import folder_paths
 from aiohttp import web
-from server import PromptServer
-
 from comfy_pack.hash import async_batch_get_sha256
 from comfy_pack.model_helper import alookup_model_source
 from comfy_pack.package import build_bento
 
+import folder_paths
+from server import PromptServer
+
 ZPath = Union[Path, zipfile.Path]
 TEMP_FOLDER = Path(__file__).parent.parent / "temp"
 COMFY_PACK_DIR = Path(__file__).parent.parent / "src" / "comfy_pack"
@@ -132,16 +132,31 @@ async def _get_models(
 
     models = []
     model_filenames = [
-        os.path.abspath(line)
+        os.path.abspath(line.strip().strip('"').replace('\\\\', '/'))
         for line in stdout.decode().splitlines()
-        if not os.path.basename(line).startswith(".")
+        if not os.path.basename(line.strip().strip('"').replace('\\\\', '/')).startswith(".")
     ]
+
+    # Only compute hashes for referenced models
+    to_include = []
+    if model_filter:
+        for m1 in model_filter:
+            for m2 in model_filenames:
+                if m1 in m2:
+                    to_include.append(m2)
+    else:
+        to_include = model_filenames
+
     model_hashes = await async_batch_get_sha256(
-        model_filenames,
+        to_include,
         cache_only=not (ensure_sha or store_models),
     )
 
-    for filename in model_filenames:
+    for filename in to_include:
+        # Skip if file doesn't exist
+        if not os.path.exists(filename):
+            continue
+            
         relpath = os.path.relpath(filename, folder_paths.base_path)
 
         model_data = {

nodes/nodes.py

@@ -1,4 +1,6 @@
+import glob
 import hashlib
+import io
 import json
 import os
 import shutil
@@ -10,7 +12,6 @@ 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.PngImagePlugin import PngInfo
 
@@ -450,31 +451,299 @@ class AnyInput:
         return True
 
 
-class OutputAudio(SaveAudio):
+class OutputZip:
+    CATEGORY = "ComfyPack/output"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {}
+
+    RETURN_TYPES = ()
+    FUNCTION = "null_op"
+
+    def null_op(self):
+        return ()
+
+
+class OutputAudio:
     CPACK_NODE = True
     CATEGORY = "ComfyPack/output"
+    RETURN_TYPES = ()
+    FUNCTION = "save_audio"
+    OUTPUT_NODE = True
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    def save_audio(
+        self,
+        audio,
+        filename_prefix,
+        format="flac",
+        prompt=None,
+        extra_pnginfo=None,
+        quality="128k",
+    ):
+        import av
+        import torchaudio
+
+        filename_prefix += self.prefix_append
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            get_save_image_path(filename_prefix, self.output_dir)
+        )
+        results = []
+
+        # Prepare metadata dictionary
+        metadata = {}
+        if prompt is not None:
+            metadata["prompt"] = json.dumps(prompt)
+        if extra_pnginfo is not None:
+            for x in extra_pnginfo:
+                metadata[x] = json.dumps(extra_pnginfo[x])
+
+        # Opus supported sample rates
+        OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
+
+        for batch_number, waveform in enumerate(audio["waveform"].cpu()):
+            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+            file = f"{filename_with_batch_num}_{counter:05}.{format}"
+            output_path = os.path.join(full_output_folder, file)
+
+            # Use original sample rate initially
+            sample_rate = audio["sample_rate"]
+
+            # Handle Opus sample rate requirements
+            if format == "opus":
+                if sample_rate > 48000:
+                    sample_rate = 48000
+                elif sample_rate not in OPUS_RATES:
+                    # Find the next highest supported rate
+                    for rate in sorted(OPUS_RATES):
+                        if rate > sample_rate:
+                            sample_rate = rate
+                            break
+                    if sample_rate not in OPUS_RATES:  # Fallback if still not supported
+                        sample_rate = 48000
+
+                # Resample if necessary
+                if sample_rate != audio["sample_rate"]:
+                    waveform = torchaudio.functional.resample(
+                        waveform, audio["sample_rate"], sample_rate
+                    )
+
+            # Create in-memory WAV buffer
+            wav_buffer = io.BytesIO()
+            torchaudio.save(wav_buffer, waveform, sample_rate, format="WAV")
+            wav_buffer.seek(0)  # Rewind for reading
+
+            # Use PyAV to convert and add metadata
+            input_container = av.open(wav_buffer)
+
+            # Create output with specified format
+            output_buffer = io.BytesIO()
+            output_container = av.open(output_buffer, mode="w", format=format)
+
+            # Set metadata on the container
+            for key, value in metadata.items():
+                output_container.metadata[key] = value
+
+            # Set up the output stream with appropriate properties
+            input_container.streams.audio[0]
+            if format == "opus":
+                out_stream = output_container.add_stream("libopus", rate=sample_rate)
+                if quality == "64k":
+                    out_stream.bit_rate = 64000
+                elif quality == "96k":
+                    out_stream.bit_rate = 96000
+                elif quality == "128k":
+                    out_stream.bit_rate = 128000
+                elif quality == "192k":
+                    out_stream.bit_rate = 192000
+                elif quality == "320k":
+                    out_stream.bit_rate = 320000
+            elif format == "mp3":
+                out_stream = output_container.add_stream("libmp3lame", rate=sample_rate)
+                if quality == "V0":
+                    # TODO i would really love to support V3 and V5 but there doesn't seem to be a way to set the qscale level, the property below is a bool
+                    out_stream.codec_context.qscale = 1
+                elif quality == "128k":
+                    out_stream.bit_rate = 128000
+                elif quality == "320k":
+                    out_stream.bit_rate = 320000
+            else:  # format == "flac":
+                out_stream = output_container.add_stream("flac", rate=sample_rate)
+
+            # Copy frames from input to output
+            for frame in input_container.decode(audio=0):
+                frame.pts = None  # Let PyAV handle timestamps
+                output_container.mux(out_stream.encode(frame))
+
+            # Flush encoder
+            output_container.mux(out_stream.encode(None))
+
+            # Close containers
+            output_container.close()
+            input_container.close()
+
+            # Write the output to file
+            output_buffer.seek(0)
+            with open(output_path, "wb") as f:
+                f.write(output_buffer.getbuffer())
+
+            results.append(
+                {"filename": file, "subfolder": subfolder, "type": self.type}
+            )
+            counter += 1
+
+        return {"ui": {"audio": results}}
 
     @classmethod
     def INPUT_TYPES(s):
         return {
             "required": {
                 "audio": ("AUDIO",),
-                "filename_prefix": ("STRING", {"default": "cpack_output_"}),
+                "filename_prefix": ("STRING", {"default": "audio/comfypack"}),
             },
             "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
         }
 
 
+class OutputVideo:
+    CPACK_NODE = True
+    CATEGORY = "ComfyPack/output"
+
+    RETURN_TYPES = ()
+    FUNCTION = "save_video"
+    DESCRIPTION = "Saves the input video to your ComfyUI output directory."
+    OUTPUT_NODE = True
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video": ("VIDEO", {"tooltip": "The video to save."}),
+                "filename_prefix": (
+                    "STRING",
+                    {
+                        "default": "video/comfypack",
+                        "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes.",
+                    },
+                ),
+                "format": (
+                    ["auto", "mp4"],
+                    {"default": "auto", "tooltip": "The format to save the video as."},
+                ),
+                "codec": (
+                    ["auto", "h264"],
+                    {"default": "auto", "tooltip": "The codec to use for the video."},
+                ),
+            },
+            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    def save_video(
+        self,
+        video,
+        filename_prefix,
+        format,
+        codec,
+        prompt=None,
+        extra_pnginfo=None,
+    ):
+        filename_prefix += self.prefix_append
+        width, height = video.get_dimensions()
+        full_output_folder, filename, counter, subfolder, filename_prefix = (
+            get_save_image_path(filename_prefix, self.output_dir, width, height)
+        )
+        results = list()
+        saved_metadata = None
+        metadata = {}
+        if extra_pnginfo is not None:
+            metadata.update(extra_pnginfo)
+        if prompt is not None:
+            metadata["prompt"] = prompt
+        if len(metadata) > 0:
+            saved_metadata = metadata
+        file = f"{filename}_{counter:05}.mp4"
+        video.save_to(
+            os.path.join(full_output_folder, file),
+            format=format,
+            codec=codec,
+            metadata=saved_metadata,
+        )
+
+        results.append({"filename": file, "subfolder": subfolder, "type": self.type})
+        counter += 1
+
+        return {"ui": {"images": results, "animated": (True,)}}
+
+
+class OutputTextFile:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "text": ("STRING", {"forceInput": True}),
+                "filename_prefix": ("STRING", {"default": "cpack_output_"}),
+            },
+            "optional": {
+                "file_extension": ("STRING", {"default": ".txt"}),
+            },
+        }
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ()
+    FUNCTION = "save_text_file"
+    CATEGORY = "ComfyPack/output"
+    CPACK_NODE = True
+
+    def save_text_file(
+        self, text: str, filename_prefix: str, file_extension: str = ".txt"
+    ):
+        subfolder, filename_prefix = os.path.split(os.path.normpath(filename_prefix))
+        output_dir = folder_paths.get_output_directory()
+        full_output_folder = os.path.join(output_dir, subfolder)
+
+        full_output_filename = self.get_output_filename(
+            full_output_folder, filename_prefix, file_extension
+        )
+        with open(full_output_filename, "w", encoding="utf-8", newline="\n") as f:
+            f.write(text)
+        return (text, {"ui": {"string": text}})
+
+    @staticmethod
+    def get_output_filename(folder: str, prefix: str, extension: str) -> str:
+        matched_files = [
+            os.path.basename(f)[len(prefix) + 1 : -len(extension)]
+            for f in glob.glob(os.path.join(folder, f"{prefix}_*{extension}"))
+        ]
+        print("MATCHING", matched_files)
+        max_count = max(
+            (int(name) for name in matched_files if name.isdigit()), default=0
+        )
+        return os.path.join(folder, f"{prefix}_{max_count + 1:04d}{extension}")
+
+
 NODE_CLASS_MAPPINGS = {
     "CPackOutputFile": OutputFile,
     "CPackOutputImage": OutputImage,
     "CPackOutputAudio": OutputAudio,
+    "CPackOutputVideo": OutputVideo,
     "CPackOutputZip": OutputImageWithStringTxt,
+    "CPackOutputZipSwitch": OutputZip,
     "CPackInputImage": ImageInput,
     "CPackInputString": StringInput,
     "CPackInputInt": IntInput,
     "CPackInputFile": FileInput,
     "CPackInputAny": AnyInput,
+    "CPackOutputTextFile": OutputTextFile,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -485,6 +754,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CPackInputAny": "Any Input",
     "CPackOutputImage": "Image Output",
     "CPackOutputAudio": "Audio Output",
+    "CPackOutputVideo": "Video Output",
     "CPackOutputFile": "File Output",
     "CPackOutputZip": "Zip Output(img + txt file)",
+    "CPackOutputZipSwitch": "Enable Zip Output",
+    "CPackOutputTextFile": "Output Text to File",
 }

src/comfy_pack/cli.py

@@ -302,15 +302,18 @@ def unpack_cmd(
         no_venv=no_venv,
         preheat=preheat,
     )
-    rich.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
-    rich.print(f"{dir}")
+    rich.print("\n[green]✓ ComfyUI Workspace is restored at:[/green]")
+    rich.print(os.path.abspath(dir))
+    steps = [f"Change directory to the restored workspace: `cd {dir}`"]
+    if not no_venv:
+        steps.append(
+            "Source the virtual environment by running `source .venv/bin/activate`"
+        )
+    steps.append("Run the ComfyUI project by running `python main.py`")
 
-    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"
-        "3. Run the ComfyUI project by running `python main.py`"
-    )
+    rich.print(f"\n[green]⏭️ Next steps: [/green]\n1. {steps[0]}\n2. {steps[1]}")
+    if len(steps) > 2:
+        rich.print(f"3. {steps[2]}")
 
 
 def _print_schema(schema, verbose: int = 0):

src/comfy_pack/hash.py

@@ -113,7 +113,12 @@ async def async_batch_get_sha256(
     try:
         with SHA_CACHE_FILE.open("r") as f:
             cache = json.load(f)
-        cache.update(new_cache)
+    except (IOError, OSError):
+        cache = {}
+
+    cache.update(new_cache)
+
+    try:
         with SHA_CACHE_FILE.open("w") as f:
             json.dump(cache, f, indent=2)
     except (IOError, OSError):

src/comfy_pack/package.py

@@ -150,7 +150,7 @@ def install_dependencies(
         print("Using the current Python environment")
         venv_py = Path(sys.executable)
     else:
-        venv = workspace / ".venv"
+        venv = (workspace / ".venv").absolute()
         venv_py = (
             venv / "Scripts" / "python.exe"
             if os.name == "nt"
@@ -222,9 +222,7 @@ def download_file(url: str, dest_path: Path, progress_callback=None):
     # prepare auth token from huggingface if possible
     if (token := os.getenv("HF_TOKEN")) and ("huggingface" in url):
         bearer = f"Bearer {token}"
-        urllib_request = urllib.request.Request(
-            url, headers={"Authorization": bearer}
-        )
+        urllib_request = urllib.request.Request(url, headers={"Authorization": bearer})
         curl_auth = ["-H", f"Authorization: {bearer}"]
     else:
         curl_auth = []
@@ -449,9 +447,16 @@ def install(
             no_venv=no_venv,
             verbose=verbose,
         )
-        cm_cli = workspace / "custom_nodes" / "ComfyUI-Manager" / "cm-cli.py"
+        cm_cli = Path("custom_nodes", "ComfyUI-Manager", "cm-cli.py")
         subprocess.check_call(
-            [str(py), str(cm_cli), "restore-snapshot", str(pack_dir / "snapshot.json")],
+            [
+                str(py),
+                str(cm_cli),
+                "restore-snapshot",
+                "--pip-non-url",
+                "--pip-non-local-url",
+                str(pack_dir / "snapshot.json"),
+            ],
             cwd=workspace,
         )
 

src/comfy_pack/run.py

@@ -5,6 +5,7 @@ import json
 import logging
 import os
 import random
+import shlex
 import shutil
 import socket
 import subprocess
@@ -115,8 +116,6 @@ class ComfyUIServer:
         subprocess.run(command, check=True, stdout=stdout, env=env)
         logger.info("Successfully disabled Comfy CLI tracking")
 
-        logger.info("Preparing directories required by ComfyUI...")
-
         logger.info("Starting ComfyUI in the background...")
         command = [
             "python",
@@ -133,6 +132,8 @@ class ComfyUIServer:
 
         if self.host != "localhost":
             command.extend(["--listen", self.host])
+        if options := env.pop("COMFYUI_OPTIONS", None):
+            command.extend(shlex.split(options))
 
         def preexec_fn():
             os.setpgrp()

src/comfy_pack/service.py

@@ -80,8 +80,9 @@ class ComfyService:
         logger = logging.getLogger("comfy_pack")
         logger.setLevel(logging.INFO)
         if not EXISTING_COMFYUI_SERVER:
+            workspace_override = os.environ.get("COMFYUI_PATH", None)
             self.server = comfy_pack.run.ComfyUIServer(
-                str(_get_workspace()),
+                workspace_override if workspace_override else str(_get_workspace()),
                 str(INPUT_DIR),
                 verbose=int("BENTOML_DEBUG" in os.environ),
             )

src/comfy_pack/utils.py

@@ -1,19 +1,15 @@
 from __future__ import annotations
 
-import subprocess
-from pathlib import Path
 import re
+import subprocess
 import sys
+import zipfile
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Union
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
 
-CPACK_OUTPUT_NODES = {
-    "CPackOutputFile",
-    "CPackOutputImage",
-    "CPackOutputZip"
-}
 
 CPACK_PATH_INPUT_NODES = {
     "CPackInputFile",
@@ -211,7 +207,7 @@ def populate_workflow(
 
     for _, node in output_spec.items():
         node_id = node["id"]
-        if node["class_type"] in CPACK_OUTPUT_NODES:
+        if node["class_type"].startswith("CPackOutput"):
             workflow[node_id]["inputs"]["filename_prefix"] = (
                 output_path / f"{session_id}{node_id}_"
             ).as_posix()
@@ -222,7 +218,7 @@ def retrieve_workflow_outputs(
     workflow: dict,
     output_path: Path,
     session_id: str = "",
-) -> Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]:
+) -> Union[Path, list[Path], dict[str, Path | list[Path]]]:
     """
     Gets the output file(s) from the workflow.
 
@@ -231,7 +227,7 @@ def retrieve_workflow_outputs(
         output_path (Path): The path where output files are saved.
 
     Returns:
-        Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]:
+        Union[Path, list[Path], dict[str, Path | list[Path]]]:
             - A single Path if there's only one output file.
             - A list of Paths if there are multiple files for a single output.
             - A dictionary mapping output names to Paths or lists of Paths for multiple outputs.
@@ -240,26 +236,47 @@ def retrieve_workflow_outputs(
         ValueError: If the output node is not of the expected type.
     """
     _, outputs = _parse_workflow(workflow)
+    should_zip = any(
+        node["class_type"] == "CPackOutputZipSwitch" for node in workflow.values()
+    )
+    zip_paths: list[tuple[Path, str]] = []
     if len(outputs) != 1:
         value_map = {}
         for k, node in outputs.items():
             node_id = node["id"]
-            path_strs = list(output_path.glob(f"{node_id}_*"))
+            path_strs = list(output_path.glob(f"{session_id}{node_id}_*"))
+            zip_paths.extend(
+                (p, p.name.replace(f"{session_id}{node_id}", k)) for p in path_strs
+            )
             if len(path_strs) == 1:
                 value_map[k] = path_strs[0]
             else:
                 value_map[k] = path_strs
-        return value_map
+        if not should_zip:
+            return value_map
+    else:
+        name, node = next(iter(outputs.items()))
+        if not node["class_type"].startswith("CPackOutput"):
+            raise ValueError(f"Node {name} is not a comfy-pack output node")
+        node_id = node["id"]
 
-    name, node = next(iter(outputs.items()))
-    if node["class_type"] not in CPACK_OUTPUT_NODES:
-        raise ValueError(f"Output node {name} is not of type {CPACK_OUTPUT_NODES}")
-    node_id = node["id"]
-
-    outs = list(output_path.glob(f"{session_id}{node_id}_*"))
-    if len(outs) == 1:
-        return outs[0]
-    return outs
+        outs = list(output_path.glob(f"{session_id}{node_id}_*"))
+        zip_paths.extend(
+            (p, p.name.replace(f"{session_id}{node_id}", name)) for p in outs
+        )
+        if not should_zip:
+            if len(outs) == 1:
+                return outs[0]
+            return outs
+    if len(zip_paths) == 1:
+        return zip_paths[0][0]
+    # Make a zipball from the collected files
+    output_zip = output_path / f"{session_id}_output.zip"
+    print(f"Creating zip file: {output_zip}")
+    with zipfile.ZipFile(output_zip, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for path, name in zip_paths:
+            zipf.write(path, arcname=name)
+    return output_zip
 
 
 def get_self_git_commit() -> str | None:

uv.lock

@@ -209,7 +209,7 @@ wheels = [
 
 [[package]]
 name = "bentoml"
-version = "1.4.13"
+version = "1.4.15"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "a2wsgi" },
@@ -255,9 +255,9 @@ dependencies = [
     { name = "uvicorn" },
     { name = "watchfiles" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/fb/15/9fd0bcfb2b4b407142a2362f69a1a5980820a77204405b402830cbdc6424/bentoml-1.4.13.tar.gz", hash = "sha256:9ba3ff502726a1496ebd9733943896797d8d9500a88f88da7fae19d51e833fc5", size = 971767, upload-time = "2025-05-09T07:59:01.918Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/53/c9/47f188d06196526173e51f1bfa2e461d73bfa9ed5a6f5f18974c9332f184/bentoml-1.4.15.tar.gz", hash = "sha256:e16fd5cbed6cccdafaa8d9c42cacb2d9b39ac2b335eb7b6051930688e6fff4f2", size = 972188, upload-time = "2025-05-23T07:34:23.377Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/a8/01/d64702692bfe11988c4ae0be5b103af71bf9dc7c2c78d1207f47a6e1e71c/bentoml-1.4.13-py3-none-any.whl", hash = "sha256:06859ec789f4f17c3744e1ab1d73da80ee0dc5ce0dd8400040b4b5c725af6567", size = 1151642, upload-time = "2025-05-09T07:58:58.649Z" },
+    { url = "https://files.pythonhosted.org/packages/5a/2d/53b0f577abb1d72606b73ac4377b2ce5bafaebc7d7c0a5ee691647790919/bentoml-1.4.15-py3-none-any.whl", hash = "sha256:3e9c411236ddb8a36346e3ab63a0ca742c63dc68140a610501f7b8a9c4ed02a5", size = 1151960, upload-time = "2025-05-23T07:34:20.972Z" },
 ]
 
 [[package]]