fix: manage node detection

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

README.md

@@ -73,13 +73,6 @@ To unpack:
 comfy-pack unpack workflow.cpack.zip
 ```
 
-Huggingface gated models can be accessed by setting your `HF_TOKEN` as an environment variable before unpacking:
-
-```bash
-export HF_TOKEN=hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXx
-comfy-pack unpack workflow.cpack.zip
-```
-
 For example cpack files, check our [examples folder](examples/).
 
 ### Deploy a workflow as an API

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"
@@ -67,8 +67,11 @@ async def _save_snapshot() -> dict[str, Any]:
         return json.load(f)
 
 
-async def _write_snapshot(path: ZPath, data: dict, models: list) -> None:
+async def _write_snapshot(path: ZPath, data: dict, models: list | None = None) -> None:
     snapshot = await _save_snapshot()
+    if models is None:
+        print("Package => Writing models")
+        models = await _get_models()
     for package in list(snapshot["pips"]):
         if any(
             fnmatch(normalize_name(package.split("==")[0]), pat)
@@ -132,31 +135,16 @@ async def _get_models(
 
     models = []
     model_filenames = [
-        os.path.abspath(line.strip().strip('"').replace('\\\\', '/'))
+        os.path.abspath(line)
         for line in stdout.decode().splitlines()
-        if not os.path.basename(line.strip().strip('"').replace('\\\\', '/')).startswith(".")
+        if not os.path.basename(line).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(
-        to_include,
+        model_filenames,
         cache_only=not (ensure_sha or store_models),
     )
 
-    for filename in to_include:
-        # Skip if file doesn't exist
-        if not os.path.exists(filename):
-            continue
-            
+    for filename in model_filenames:
         relpath = os.path.relpath(filename, folder_paths.base_path)
 
         model_data = {
@@ -174,6 +162,7 @@ async def _get_models(
             model_data["sha256"],
             cache_only=not ensure_source,
         )
+        model_data["source"] = {}
         # should_store = store_models and (
         #     model_data["source"].get("source") != "huggingface"
         #     or model_data["source"].get("repo", "").startswith("datasets/")
@@ -448,13 +437,11 @@ async def _prepare_pack(
     working_dir: ZPath,
     data: dict,
     store_models: bool = False,
-    ensure_source: bool = True,
 ) -> None:
     model_filter = set(data.get("models", []))
     models = await _get_models(
         store_models=store_models,
         model_filter=model_filter,
-        ensure_source=ensure_source,
     )
 
     await _write_snapshot(working_dir, data, models)
@@ -522,7 +509,7 @@ async def build_bento_api(request):
 
     with tempfile.TemporaryDirectory(suffix="-bento", prefix="comfy-pack-") as temp_dir:
         temp_dir_path = Path(temp_dir)
-        await _prepare_pack(temp_dir_path, data, store_models=True, ensure_source=False)
+        await _prepare_pack(temp_dir_path, data, store_models=True)
 
         # create a bento
         try:

nodes/nodes.py

@@ -1,12 +1,9 @@
-import glob
 import hashlib
-import io
 import json
 import os
-import shutil
 import sys
+import shutil
 import zipfile
-from io import BytesIO
 
 import folder_paths
 import node_helpers
@@ -14,6 +11,7 @@ import numpy as np
 import torch
 from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
 from PIL.PngImagePlugin import PngInfo
+from io import BytesIO
 
 from .monkeypatch import set_bentoml_output
 
@@ -196,18 +194,9 @@ class OutputImageWithStringTxt:
     OUTPUT_NODE = True
 
     CATEGORY = "ComfyPack/output"
-    DESCRIPTION = (
-        "Saves the input images (and optional text) to your ComfyUI output directory."
-    )
+    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,
-    ):
+    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(
@@ -220,7 +209,7 @@ class OutputImageWithStringTxt:
         zip_path = os.path.join(full_output_folder, zip_filename)
 
         # create ZIP file
-        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        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()
@@ -236,12 +225,7 @@ class OutputImageWithStringTxt:
 
                 # write img file to RAM buffer
                 img_buffer = BytesIO()
-                img.save(
-                    img_buffer,
-                    format="PNG",
-                    pnginfo=metadata,
-                    compress_level=self.compress_level,
-                )
+                img.save(img_buffer, format="PNG", pnginfo=metadata, compress_level=self.compress_level)
                 img_buffer.seek(0)
 
                 # write img into ZIP file
@@ -451,299 +435,15 @@ class AnyInput:
         return True
 
 
-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": "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 = {
@@ -753,10 +453,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CPackInputFile": "File Input",
     "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

@@ -86,7 +86,7 @@ def init(dir: str, verbose: int):
             capture_output=True,
         )
     except (subprocess.SubprocessError, FileNotFoundError):
-        rich.print("[red]Error: git is not installed or not in PATH[/red]")
+        console.print("[red]Error: git is not installed or not in PATH[/red]")
         return 1
 
     # Check if we have write permissions
@@ -97,36 +97,36 @@ def init(dir: str, verbose: int):
         test_file.touch()
         test_file.unlink()
     except (OSError, PermissionError) as e:
-        rich.print(f"[red]Error: No write permission in {dir} - {str(e)}[/red]")
+        console.print(f"[red]Error: No write permission in {dir} - {str(e)}[/red]")
         return 1
 
     # Check if Python version is compatible
     if sys.version_info < (3, 8):
-        rich.print("[red]Error: Python 3.8 or higher is required[/red]")
+        console.print("[red]Error: Python 3.8 or higher is required[/red]")
         return 1
 
     # Check if uv is installed
     try:
         _ensure_uv()
     except RuntimeError as e:
-        rich.print(f"[red]Error: {str(e)}[/red]")
+        console.print(f"[red]Error: {str(e)}[/red]")
         return 1
 
     # Check if enough disk space is available (rough estimate: 2GB)
     try:
         free_space = shutil.disk_usage(install_dir).free
         if free_space < 2 * 1024 * 1024 * 1024:  # 2GB in bytes
-            rich.print(
+            console.print(
                 "[yellow]Warning: Less than 2GB free disk space available[/yellow]"
             )
     except Exception as e:
-        rich.print(
+        console.print(
             f"[yellow]Warning: Could not check free disk space - {str(e)}[/yellow]"
         )
 
     # Clone ComfyUI if not exists
     if not (install_dir / ".git").exists():
-        rich.print("[green]Cloning ComfyUI...[/green]")
+        console.print("[green]Cloning ComfyUI...[/green]")
         subprocess.run(
             [
                 "git",
@@ -138,7 +138,7 @@ def init(dir: str, verbose: int):
         )
 
     # Update ComfyUI
-    rich.print("[green]Updating ComfyUI...[/green]")
+    console.print("[green]Updating ComfyUI...[/green]")
     subprocess.run(
         ["git", "pull"],
         cwd=install_dir,
@@ -147,7 +147,7 @@ def init(dir: str, verbose: int):
 
     # Create and activate venv
     venv_dir = install_dir / ".venv"
-    rich.print("[green]Creating virtual environment with uv...[/green]")
+    console.print("[green]Creating virtual environment with uv...[/green]")
     if venv_dir.exists():
         shutil.rmtree(venv_dir)
     subprocess.run(
@@ -163,7 +163,7 @@ def init(dir: str, verbose: int):
         python = str(venv_dir / "bin" / "python")
 
     # Install requirements with uv
-    rich.print("[green]Installing ComfyUI requirements with uv...[/green]")
+    console.print("[green]Installing ComfyUI requirements with uv...[/green]")
     subprocess.run(
         ["uv", "pip", "install", "pip", "--upgrade"],
         env={
@@ -182,7 +182,7 @@ def init(dir: str, verbose: int):
     )
 
     # Install comfy-pack as custom node
-    rich.print("[green]Installing comfy-pack custom nodes...[/green]")
+    console.print("[green]Installing comfy-pack custom nodes...[/green]")
     custom_nodes_dir = install_dir / "custom_nodes"
     custom_nodes_dir.mkdir(exist_ok=True)
 
@@ -224,12 +224,12 @@ def init(dir: str, verbose: int):
         )
 
     version = get_self_git_commit() or "unknown"
-    rich.print(
+    console.print(
         f"\n[green]✓ Installation completed! (comfy-pack version: {version})[/green]"
     )
-    rich.print(f"ComfyUI directory: {install_dir}")
+    console.print(f"ComfyUI directory: {install_dir}")
 
-    rich.print(
+    console.print(
         "\n[green]Next steps:[/green]\n"
         f"1. cd {dir}\n"
         "2. source .venv/bin/activate  # On Windows: .venv\\Scripts\\activate\n"
@@ -262,32 +262,14 @@ def init(dir: str, verbose: int):
     is_flag=True,
     help="Do not install models",
 )
-@click.option(
-    "--no-venv",
-    is_flag=True,
-    help="Do not create a virtual environment for ComfyUI",
-    default=False,
-)
 @click.option(
     "--verbose",
     "-v",
     count=True,
     help="Increase verbosity level (use multiple times for more verbosity)",
 )
-@click.option(
-    "--preheat",
-    is_flag=True,
-    help="Preheat the workspace after unpacking",
-    default=False,
-)
 def unpack_cmd(
-    cpack: str,
-    dir: str,
-    include_disabled_models: bool,
-    no_models: bool,
-    no_venv: bool,
-    verbose: int,
-    preheat: bool,
+    cpack: str, dir: str, include_disabled_models: bool, no_models: bool, verbose: int
 ):
     import rich
 
@@ -299,21 +281,16 @@ def unpack_cmd(
         verbose=verbose,
         all_models=include_disabled_models,
         prepare_models=not no_models,
-        no_venv=no_venv,
-        preheat=preheat,
     )
-    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]✓ ComfyUI Workspace is restored![/green]")
+    rich.print(f"{dir}")
 
-    rich.print(f"\n[green]⏭️ Next steps: [/green]\n1. {steps[0]}\n2. {steps[1]}")
-    if len(steps) > 2:
-        rich.print(f"3. {steps[2]}")
+    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`"
+    )
 
 
 def _print_schema(schema, verbose: int = 0):
@@ -417,19 +394,19 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
 
     workspace = _get_cache_workspace(cpack)
     if not (workspace / "DONE").exists():
-        rich.print("\n[green]✓ Restoring ComfyUI Workspace...[/green]")
+        console.print("\n[green]✓ Restoring ComfyUI Workspace...[/green]")
         if workspace.exists():
             shutil.rmtree(workspace)
         install(cpack, workspace, verbose=verbose)
         with open(workspace / "DONE", "w") as f:
             f.write("DONE")
-    rich.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
-    rich.print(f"{workspace}")
+    console.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
+    console.print(f"{workspace}")
 
     from .run import ComfyUIServer, run_workflow
 
     with ComfyUIServer(str(workspace.absolute()), verbose=verbose) as server:
-        rich.print("\n[green]✓ ComfyUI is launched in the background![/green]")
+        console.print("\n[green]✓ ComfyUI is launched in the background![/green]")
         results = run_workflow(
             server.host,
             server.port,
@@ -439,17 +416,17 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
             workspace=server.workspace,
             **validated_data.model_dump(),
         )
-        rich.print("\n[green]✓ Workflow is executed successfully![/green]")
+        console.print("\n[green]✓ Workflow is executed successfully![/green]")
         if results:
-            rich.print("\n[green]✓ Retrieved outputs:[/green]")
+            console.print("\n[green]✓ Retrieved outputs:[/green]")
         if isinstance(results, dict):
             for field, value in results.items():
-                rich.print(f"{field}: {value}")
+                console.print(f"{field}: {value}")
         elif isinstance(results, list):
             for i, value in enumerate(results):
-                rich.print(f"{i}: {value}")
+                console.print(f"{i}: {value}")
         else:
-            rich.print(results)
+            console.print(results)
 
 
 @main.command(name="build-bento")

src/comfy_pack/hash.py

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

src/comfy_pack/model_helper.py

@@ -1,9 +1,8 @@
+from .const import MODEL_SOURCE_CACHE_FILE
 import asyncio
 import json
 import re
 
-from .const import MODEL_SOURCE_CACHE_FILE
-
 # MODEL_NAME = r"[a-zA-Z0-9-._]+"
 # COMMIT = r"[a-f0-9]+"
 
@@ -15,8 +14,9 @@ PATH_PATTERN = re.compile(
 
 
 async def _lookup_huggingface_model(model_sha: str) -> dict:
-    import aiohttp
     from duckduckgo_search import DDGS
+    import aiohttp
+
     query = f"site:huggingface.co blob {model_sha}"
 
     try:
@@ -25,7 +25,7 @@ async def _lookup_huggingface_model(model_sha: str) -> dict:
 
             async with aiohttp.ClientSession(trust_env=True) as session:
                 for result in search_results:
-                    url = result['href']
+                    url = result['link']
                     if "blob" not in url:
                         continue
 
@@ -38,9 +38,11 @@ async def _lookup_huggingface_model(model_sha: str) -> dict:
                                 repo, commit = commit_match.groups()
                                 if path_match := PATH_PATTERN.search(text):
                                     path = path_match.group(1)
+                                    download_url = f"https://huggingface.co/{repo}/resolve/{commit}/{path}?download=true"
+                                    url = f"https://huggingface.co/{repo}/blob/{commit}/{path}"
                                     info = {
-                                        "download_url": path,
-                                        "url": path,
+                                        "download_url": download_url,
+                                        "url": url,
                                         "repo": repo,
                                         "commit": commit,
                                         "path": path,

src/comfy_pack/package.py

@@ -137,38 +137,32 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
 
 def install_dependencies(
     python_version: str,
-    req_files: list[str],
+    req_file: str,
     workspace: Path,
     verbose: int = 0,
     no_deps: bool = False,
-    no_venv: bool = False,
-) -> Path:
+):
     print("Installing Python dependencies")
     stdout = None if verbose > 0 else subprocess.DEVNULL
     stderr = None if verbose > 1 else subprocess.DEVNULL
-    if no_venv:
-        print("Using the current Python environment")
-        venv_py = Path(sys.executable)
-    else:
-        venv = (workspace / ".venv").absolute()
-        venv_py = (
-            venv / "Scripts" / "python.exe"
-            if os.name == "nt"
-            else venv / "bin" / "python"
-        )
-        if (venv / "DONE").exists():
-            return venv_py
-        subprocess.check_call(
-            [
-                "uv",
-                "venv",
-                "--python",
-                python_version,
-                venv,
-            ],
-            stdout=stdout,
-            stderr=stderr,
-        )
+
+    venv = workspace / ".venv"
+    if (venv / "DONE").exists():
+        return
+    venv_py = (
+        venv / "Scripts" / "python.exe" if os.name == "nt" else venv / "bin" / "python"
+    )
+    subprocess.check_call(
+        [
+            "uv",
+            "venv",
+            "--python",
+            python_version,
+            venv,
+        ],
+        stdout=stdout,
+        stderr=stderr,
+    )
     subprocess.check_call(
         [
             "uv",
@@ -182,16 +176,16 @@ def install_dependencies(
         stderr=stderr,
     )
     if verbose > 0:
-        print(f"Installing dependencies from {req_files}")
+        print(f"Installing dependencies from {req_file}")
     install_cmd = [
         "uv",
         "pip",
         "install",
         "-p",
         str(venv_py),
+        "-r",
+        req_file,
     ]
-    for req_file in req_files:
-        install_cmd.extend(["-r", str(req_file)])
     if not STRICT_MODE:
         install_cmd.extend(["--index-strategy", "unsafe-best-match"])
     if no_deps:
@@ -201,9 +195,8 @@ def install_dependencies(
         stdout=stdout,
         stderr=stderr,
     )
-    if not no_venv:
-        with open(venv / "DONE", "w") as f:
-            f.write("DONE")
+    with open(venv / "DONE", "w") as f:
+        f.write("DONE")
     return venv_py
 
 
@@ -218,23 +211,13 @@ def get_search_url(sha: str) -> str:
 
 def download_file(url: str, dest_path: Path, progress_callback=None):
     """Download file with progress tracking"""
-
-    # 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})
-        curl_auth = ["-H", f"Authorization: {bearer}"]
-    else:
-        curl_auth = []
-        urllib_request = url
-
+    if subprocess.call(["curl", "--version"], stdout=subprocess.DEVNULL) == 0:
+        subprocess.check_call(
+            ["curl", "-L", url, "-o", str(dest_path)],
+        )
+        return True
     try:
-        if shutil.which("curl"):
-            subprocess.check_call(
-                ["curl", "-L", url, *curl_auth, "--fail", "-o", str(dest_path)],
-            )
-            return True
-        with urllib.request.urlopen(urllib_request) as response:
+        with urllib.request.urlopen(url) as response:
             total_size = int(response.headers.get("content-length", 0))
             block_size = 8192
             downloaded = 0
@@ -408,13 +391,18 @@ def retrieve_models(
                 continue
 
 
+def _ensure_startup_script_path(workspace: Path) -> Path:
+    script_path = workspace / "user" / "default" / "ComfyUI-Manager" / "startup-scripts"
+    script_path.mkdir(parents=True, exist_ok=True)
+    return script_path
+
+
 def install(
     cpack: str | Path,
     workspace: str | Path = "workspace",
-    preheat: bool = False,
+    preheat: bool = True,
     prepare_models: bool = True,
     all_models: bool = False,
-    no_venv: bool = False,
     verbose: int = 0,
 ):
     workspace = Path(workspace)
@@ -435,29 +423,15 @@ def install(
             )
 
         install_comfyui(snapshot, workspace, verbose=verbose)
-        py = install_dependencies(
+        install_dependencies(
             snapshot["python"],
-            [
-                str(workspace / "requirements.txt"),
-                str(
-                    workspace / "custom_nodes" / "ComfyUI-Manager" / "requirements.txt"
-                ),
-            ],
+            str(workspace / "requirements.txt"),
             workspace,
-            no_venv=no_venv,
             verbose=verbose,
         )
-        cm_cli = Path("custom_nodes", "ComfyUI-Manager", "cm-cli.py")
-        subprocess.check_call(
-            [
-                str(py),
-                str(cm_cli),
-                "restore-snapshot",
-                "--pip-non-url",
-                "--pip-non-local-url",
-                str(pack_dir / "snapshot.json"),
-            ],
-            cwd=workspace,
+        shutil.copy2(
+            (pack_dir / "snapshot.json"),
+            _ensure_startup_script_path(workspace) / "restore-snapshot.json",
         )
 
         for f in (pack_dir / "input").glob("*"):
@@ -479,7 +453,7 @@ def install(
             with ComfyUIServer(
                 str(workspace),
                 verbose=verbose,
-                venv=str(workspace / ".venv") if not no_venv else None,
+                venv=str(workspace / ".venv"),
             ) as _:
                 pass
         if prepare_models:

src/comfy_pack/run.py

@@ -5,7 +5,6 @@ import json
 import logging
 import os
 import random
-import shlex
 import shutil
 import socket
 import subprocess
@@ -116,6 +115,8 @@ 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",
@@ -132,8 +133,6 @@ 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,9 +80,8 @@ 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(
-                workspace_override if workspace_override else str(_get_workspace()),
+                str(_get_workspace()),
                 str(INPUT_DIR),
                 verbose=int("BENTOML_DEBUG" in os.environ),
             )

src/comfy_pack/setup_workspace.sh

@@ -18,6 +18,6 @@ if [ -n "$VIRTUAL_ENV" ]; then
 fi
 
 set -x
-comfy-pack unpack "$CPACK" -d "$workspace" --no-models --no-venv -v
+comfy-pack unpack "$CPACK" -d "$workspace" --no-models -v
 chown -R bentoml:bentoml "$workspace"
 set +x

src/comfy_pack/utils.py

@@ -1,15 +1,19 @@
 from __future__ import annotations
 
-import re
 import subprocess
-import sys
-import zipfile
 from pathlib import Path
+import re
+import sys
 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",
@@ -207,7 +211,7 @@ def populate_workflow(
 
     for _, node in output_spec.items():
         node_id = node["id"]
-        if node["class_type"].startswith("CPackOutput"):
+        if node["class_type"] in CPACK_OUTPUT_NODES:
             workflow[node_id]["inputs"]["filename_prefix"] = (
                 output_path / f"{session_id}{node_id}_"
             ).as_posix()
@@ -218,7 +222,7 @@ def retrieve_workflow_outputs(
     workflow: dict,
     output_path: Path,
     session_id: str = "",
-) -> Union[Path, list[Path], dict[str, Path | list[Path]]]:
+) -> Union[Path, list[Path], dict[str, Path], dict[str, list[Path]]]:
     """
     Gets the output file(s) from the workflow.
 
@@ -227,7 +231,7 @@ def retrieve_workflow_outputs(
         output_path (Path): The path where output files are saved.
 
     Returns:
-        Union[Path, list[Path], dict[str, Path | list[Path]]]:
+        Union[Path, list[Path], dict[str, Path], dict[str, 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.
@@ -236,47 +240,26 @@ 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"{session_id}{node_id}_*"))
-            zip_paths.extend(
-                (p, p.name.replace(f"{session_id}{node_id}", k)) for p in path_strs
-            )
+            path_strs = list(output_path.glob(f"{node_id}_*"))
             if len(path_strs) == 1:
                 value_map[k] = path_strs[0]
             else:
                 value_map[k] = path_strs
-        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"]
+        return value_map
 
-        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
+    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
 
 
 def get_self_git_commit() -> str | None:

uv.lock

@@ -209,7 +209,7 @@ wheels = [
 
 [[package]]
 name = "bentoml"
-version = "1.4.15"
+version = "1.4.13"
 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/53/c9/47f188d06196526173e51f1bfa2e461d73bfa9ed5a6f5f18974c9332f184/bentoml-1.4.15.tar.gz", hash = "sha256:e16fd5cbed6cccdafaa8d9c42cacb2d9b39ac2b335eb7b6051930688e6fff4f2", size = 972188, upload-time = "2025-05-23T07:34:23.377Z" }
+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" }
 wheels = [
-    { 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" },
+    { 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" },
 ]
 
 [[package]]