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

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/__init__.py

@@ -1,7 +1,8 @@
-import pathlib
 import sys
+import pathlib
+
 
 SRC_DIR = pathlib.Path(__file__).parent.parent / "src"
 
 if str(SRC_DIR) not in sys.path:
-    sys.path.insert(0, str(SRC_DIR))
+    sys.path.append(str(SRC_DIR))

nodes/api.py

@@ -11,7 +11,7 @@ import tempfile
 import time
 import uuid
 import zipfile
-from fnmatch import fnmatch
+from importlib.metadata import Distribution, distributions
 from pathlib import Path
 from typing import Any, Union
 
@@ -26,63 +26,53 @@ from comfy_pack.package import build_bento
 ZPath = Union[Path, zipfile.Path]
 TEMP_FOLDER = Path(__file__).parent.parent / "temp"
 COMFY_PACK_DIR = Path(__file__).parent.parent / "src" / "comfy_pack"
-EXCLUDE_PACKAGES = ["bentoml", "onnxruntime", "conda", "nvidia-*"]
+EXCLUDE_PACKAGES = ["bentoml", "onnxruntime", "conda"]  # TODO: standardize this
 
 
-def normalize_name(name: str) -> str:
+def _get_requirement_string(dist: Distribution) -> str:
-    import re
+    direct_url_text = dist.read_text("direct_url.json")
-
+    pinned_str = f'{dist.metadata["Name"]}=={dist.version}'
-    return re.sub(r"[-_.]+", "-", name).lower()
+    if not direct_url_text:
+        return pinned_str
+    direct_url = json.loads(direct_url_text)
+    if url := direct_url.get("url"):
+        if url.startswith("file://"):
+            # we are not able to share local files
+            return pinned_str
+        if vcs_info := direct_url.get("vcs_info"):
+            url = f"{vcs_info['vcs']}+{url}@{vcs_info['commit_id']}"
+        if subdirectory := direct_url.get("subdirectory"):
+            url += f"#subdirectory={subdirectory}"
+        return f"{dist.metadata['Name']} @ {url}"
+    else:
+        return pinned_str
 
 
-def get_snapshot_path() -> Path | None:
+async def _write_requirements(path: ZPath, extras: list[str] | None = None) -> None:
-    manager_file_path = Path(
+    print("Package => Writing requirements.txt")
-        folder_paths.get_user_directory(), "default", "ComfyUI-Manager"
+    with path.joinpath("requirements.txt").open("w") as f:
+        for dist in distributions():
+            f.write(_get_requirement_string(dist) + "\n")
+        if extras:
+            f.write("\n".join(extras) + "\n")
+
+
+async def _write_snapshot(path: ZPath, data: dict, models: list | None = None) -> None:
+    proc = await asyncio.subprocess.create_subprocess_exec(
+        "git", "rev-parse", "HEAD", stdout=subprocess.PIPE, cwd=folder_paths.base_path
     )
-    return manager_file_path / "snapshots"
+    stdout, _ = await proc.communicate()
-
+    if models is None:
-
+        print("Package => Writing models")
-async def _save_snapshot() -> dict[str, Any]:
+        models = await _get_models()
-    save_snapshot_route = next(
-        (
-            route
-            for route in PromptServer.instance.routes
-            if route.path == "/snapshot/save"
-        ),
-        None,
-    )
-    if not save_snapshot_route:
-        raise RuntimeError("ComfyUI-Manager must be installed to save snapshot")
-    await save_snapshot_route.handler(None)
-    snapshot_path = get_snapshot_path()
-    if not snapshot_path.exists():
-        raise RuntimeError("Snapshot save failed")
-
-    most_recent = max(
-        snapshot_path.glob("*.json"), key=lambda x: x.stat().st_mtime, default=None
-    )
-    if not most_recent:
-        raise RuntimeError("Snapshot save failed")
-    with most_recent.open("r") as f:
-        return json.load(f)
-
-
-async def _write_snapshot(path: ZPath, data: dict, models: list) -> None:
-    snapshot = await _save_snapshot()
-    for package in list(snapshot["pips"]):
-        if any(
-            fnmatch(normalize_name(package.split("==")[0]), pat)
-            for pat in EXCLUDE_PACKAGES
-        ):
-            del snapshot["pips"][package]
     with path.joinpath("snapshot.json").open("w") as f:
-        snapshot.update(
+        data = {
-            {
+            "python": f"{sys.version_info.major}.{sys.version_info.minor}",
-                "python": f"{sys.version_info.major}.{sys.version_info.minor}",
+            "comfyui": stdout.decode().strip(),
-                "models": models,
+            "models": models,
-            }
+            "custom_nodes": await _get_custom_nodes(),
-        )
+        }
-        f.write(json.dumps(snapshot, indent=2))
+        f.write(json.dumps(data, indent=2))
 
 
 def _is_port_in_use(port: int | str, host="localhost"):
@@ -168,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:
@@ -184,6 +174,47 @@ async def _get_models(
     return models
 
 
+async def _get_custom_nodes() -> list:
+    print("Package => Writing custom nodes")
+    custom_nodes = os.path.join(folder_paths.base_path, "custom_nodes")
+    coros = []
+
+    async def get_node_info(subdir: Path) -> dict:
+        proc = await asyncio.subprocess.create_subprocess_exec(
+            "git",
+            "config",
+            "--get",
+            "remote.origin.url",
+            cwd=subdir,
+            stdout=subprocess.PIPE,
+        )
+        stdout, _ = await proc.communicate()
+        url = stdout.decode().strip()
+
+        proc = await asyncio.subprocess.create_subprocess_exec(
+            "git",
+            "rev-parse",
+            "HEAD",
+            cwd=subdir,
+            stdout=subprocess.PIPE,
+        )
+        stdout, _ = await proc.communicate()
+        commit_hash = stdout.decode().strip()
+        return {
+            "url": url,
+            "commit_hash": commit_hash,
+            "disabled": subdir.name.endswith(".disabled"),
+            "path": str(subdir.relative_to(custom_nodes)),
+        }
+
+    for subdir in Path(custom_nodes).iterdir():
+        if not subdir.is_dir() or not subdir.joinpath(".git").exists():
+            continue
+        coros.append(get_node_info(subdir))
+
+    return await asyncio.gather(*coros)
+
+
 async def _write_workflow(path: ZPath, data: dict) -> None:
     print("Package => Writing workflow")
     with path.joinpath("workflow_api.json").open("w") as f:
@@ -246,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
@@ -256,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)
@@ -407,7 +436,7 @@ async def serve(request):
         )
 
 
-@PromptServer.instance.routes.get("/bentoml/serve/heartbeat")
+@PromptServer.instance.routes.post("/bentoml/serve/heartbeat")
 async def heartbeat(_):
     running = DevServer.feed_watch_dog()
 
@@ -433,15 +462,14 @@ 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_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)
@@ -507,14 +535,14 @@ 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:
             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,19 +1,14 @@
-import glob
 import hashlib
 import json
 import os
-import shutil
 import sys
-import zipfile
+import shutil
-from io import BytesIO
 
 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
-from comfy_extras.nodes_video import SaveVideo
-from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
 from PIL.PngImagePlugin import PngInfo
 
 from .monkeypatch import set_bentoml_output
@@ -168,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)
 
@@ -452,99 +353,14 @@ 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(SaveAudio):
-    CPACK_NODE = True
-    CATEGORY = "ComfyPack/output"
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "audio": ("AUDIO",),
-                "filename_prefix": ("STRING", {"default": "cpack_output_"}),
-            },
-            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-        }
-
-
-class OutputVideo(SaveVideo):
-    CPACK_NODE = True
-    CATEGORY = "ComfyPack/output"
-
-
-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 = {
@@ -554,10 +370,5 @@ 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",
 }

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,7 +5,6 @@ import shutil
 import subprocess
 import sys
 import tempfile
-import zipfile
 from pathlib import Path
 
 import click
@@ -31,7 +30,6 @@ def _ensure_uv() -> None:
 
 
 @click.group()
-@click.version_option()
 def main():
     """comfy-pack CLI"""
     pass
@@ -57,25 +55,27 @@ def main():
 def init(dir: str, verbose: int):
     import os
 
-    import rich
+    from rich.console import Console
+
+    console = Console()
 
     # Check if directory path is valid
     try:
         install_dir = Path(dir).absolute()
         if install_dir.exists() and not install_dir.is_dir():
-            rich.print(f"[red]Error: {dir} exists but is not a directory[/red]")
+            console.print(f"[red]Error: {dir} exists but is not a directory[/red]")
             return 1
 
         # Check if directory is empty or contains ComfyUI
         if install_dir.exists():
             contents = list(install_dir.iterdir())
             if contents and not (install_dir / ".git").exists():
-                rich.print(
+                console.print(
                     f"[red]Error: Directory {dir} is not empty and doesn't appear to be a ComfyUI installation[/red]"
                 )
                 return 1
     except Exception as e:
-        rich.print(f"[red]Error: Invalid directory path - {str(e)}[/red]")
+        console.print(f"[red]Error: Invalid directory path - {str(e)}[/red]")
         return 1
 
     # Check git installation
@@ -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"
@@ -241,7 +241,7 @@ def init(dir: str, verbose: int):
     name="unpack",
     help="Restore the ComfyUI workspace to specified directory",
 )
-@click.argument("cpack", type=click.Path(exists=True))
+@click.argument("cpack", type=click.Path(exists=True, dir_okay=False))
 @click.option(
     "--dir",
     "-d",
@@ -255,69 +255,33 @@ def init(dir: str, verbose: int):
     type=click.BOOL,
     is_flag=True,
 )
-@click.option(
-    "--no-models",
-    default=False,
-    type=click.BOOL,
-    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(
+def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int):
-    "--preheat",
+    from rich.console import Console
-    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,
-):
-    import rich
 
     from .package import install
 
-    install(
+    console = Console()
-        cpack,
-        dir,
-        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(f"\n[green]⏭️ Next steps: [/green]\n1. {steps[0]}\n2. {steps[1]}")
+    install(cpack, dir, verbose=verbose, all_models=include_disabled_models)
-    if len(steps) > 2:
+    console.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
-        rich.print(f"3. {steps[2]}")
+    console.print(f"{dir}")
+
+    console.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):
-    import rich
+    from rich.console import Console
     from rich.table import Table
 
     table = Table(title="")
@@ -348,7 +312,7 @@ def _print_schema(schema, verbose: int = 0):
             range_str,
         )
 
-    rich.print(table)
+    Console().print(table)
 
 
 @functools.lru_cache
@@ -376,8 +340,8 @@ def _get_cache_workspace(cpack: str):
 )
 @click.pass_context
 def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
-    import rich
     from pydantic import ValidationError
+    from rich.console import Console
 
     from .utils import generate_input_model
 
@@ -385,31 +349,35 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
         zip([k.lstrip("-").replace("-", "_") for k in ctx.args[::2]], ctx.args[1::2])
     )
 
-    with zipfile.ZipFile(cpack) as z:
+    console = Console()
-        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)
 
     # If help is requested, show command help and input schema
     if help:
-        rich.print(
+        console.print(
             'Usage: comfy-pack run [OPTIONS] CPACK --input1 "value1" --input2 "value2" ...'
         )
-        rich.print("Run a ComfyUI package with the given inputs:")
+        console.print("Run a ComfyUI package with the given inputs:")
         _print_schema(input_model.model_json_schema(), verbose)
         return 0
 
     try:
         validated_data = input_model(**inputs)
-        rich.print("[green]✓ Input is valid![/green]")
+        console.print("[green]✓ Input is valid![/green]")
         for field, value in validated_data.model_dump().items():
-            rich.print(f"{field}: {value}")
+            console.print(f"{field}: {value}")
     except ValidationError as e:
-        rich.print("[red]✗ Validation failed![/red]")
+        console.print("[red]✗ Validation failed![/red]")
         for error in e.errors():
-            rich.print(f"- {error['loc'][0]}: {error['msg']}")
+            console.print(f"- {error['loc'][0]}: {error['msg']}")
 
-        rich.print("\n[yellow]Expected inputs:[/yellow]")
+        console.print("\n[yellow]Expected inputs:[/yellow]")
         _print_schema(input_model.model_json_schema(), verbose)
         return 1
 
@@ -417,39 +385,38 @@ 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]")
+    console.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
-    rich.print(f"{workspace}")
+    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,
             workflow,
             Path(output_dir).absolute(),
             verbose=verbose,
-            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")
@@ -476,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
 
@@ -493,59 +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
-
-    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)
-
-    install(bento_obj.path_of("src"), workspace, verbose=verbose, prepare_models=False)
-
-    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/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,
+                                        "download_url": download_url,
-                                        "url": path,
+                                        "url": url,
                                         "repo": repo,
                                         "commit": commit,
                                         "path": path,

src/comfy_pack/package.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import contextlib
 import json
 import os
 import shutil
@@ -26,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,
     )
 
 
@@ -67,18 +61,6 @@ def install_comfyui(snapshot, workspace: Path, verbose: int = 0):
                 return
         shutil.rmtree(workspace)
     _clone_commit(COMFYUI_REPO, comfyui_commit, workspace, verbose=verbose)
-    manager_node = next(
-        url for url in snapshot.get("git_custom_nodes", {}) if "ComfyUI-Manager" in url
-    )
-    if manager_node:
-        manager_commit = snapshot["git_custom_nodes"][manager_node]["hash"].strip()
-        print("Installing ComfyUI-Manager")
-        _clone_commit(
-            manager_node,
-            manager_commit,
-            workspace / "custom_nodes" / "ComfyUI-Manager",
-            verbose=verbose,
-        )
     with open(workspace / ".DONE", "w") as f:
         f.write(comfyui_commit)
 
@@ -101,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)
 
@@ -136,39 +117,39 @@ def install_custom_modules(snapshot, workspace: Path, verbose: int = 0):
 
 
 def install_dependencies(
-    python_version: str,
+    snapshot: dict,
-    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")
+    python_version = snapshot["python"]
     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")
+    subprocess.check_call(
-        venv_py = Path(sys.executable)
+        ["uv", "python", "install", python_version],
-    else:
+        cwd=workspace,
-        venv = workspace / ".venv"
+        stdout=stdout,
-        venv_py = (
+        stderr=stderr,
-            venv / "Scripts" / "python.exe"
+    )
-            if os.name == "nt"
+    venv = workspace / ".venv"
-            else venv / "bin" / "python"
+    if (venv / "DONE").exists():
-        )
+        return
-        if (venv / "DONE").exists():
+    venv_py = (
-            return venv_py
+        venv / "Scripts" / "python.exe" if os.name == "nt" else venv / "bin" / "python"
-        subprocess.check_call(
+    )
-            [
+    subprocess.check_call(
-                "uv",
+        [
-                "venv",
+            "uv",
-                "--python",
+            "venv",
-                python_version,
+            "--python",
-                venv,
+            python_version,
-            ],
+            venv,
-            stdout=stdout,
+        ],
-            stderr=stderr,
+        stdout=stdout,
-        )
+        stderr=stderr,
+    )
     subprocess.check_call(
         [
             "uv",
@@ -182,28 +163,28 @@ 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,
+        "--no-deps",
     ]
-    for req_file in req_files:
+    if STRICT_MODE:
-        install_cmd.extend(["-r", str(req_file)])
+        pass
-    if not STRICT_MODE:
+    else:
         install_cmd.extend(["--index-strategy", "unsafe-best-match"])
-    if no_deps:
-        install_cmd.append("--no-deps")
     subprocess.check_call(
         install_cmd,
         stdout=stdout,
         stderr=stderr,
     )
-    if not no_venv:
+    with open(venv / "DONE", "w") as f:
-        with open(venv / "DONE", "w") as f:
+        f.write("DONE")
-            f.write("DONE")
     return venv_py
 
 
@@ -218,23 +199,13 @@ def get_search_url(sha: str) -> str:
 
 def download_file(url: str, dest_path: Path, progress_callback=None):
     """Download file with progress tracking"""
-
+    if subprocess.call(["curl", "--version"], stdout=subprocess.DEVNULL) == 0:
-    # prepare auth token from huggingface if possible
+        subprocess.check_call(
-    if (token := os.getenv("HF_TOKEN")) and ("huggingface" in url):
+            ["curl", "-L", url, "-o", str(dest_path)],
-        bearer = f"Bearer {token}"
+        )
-        urllib_request = urllib.request.Request(url, headers={"Authorization": bearer})
+        return True
-        curl_auth = ["-H", f"Authorization: {bearer}"]
-    else:
-        curl_auth = []
-        urllib_request = url
-
     try:
-        if shutil.which("curl"):
+        with urllib.request.urlopen(url) as response:
-            subprocess.check_call(
-                ["curl", "-L", url, *curl_auth, "--fail", "-o", str(dest_path)],
-            )
-            return True
-        with urllib.request.urlopen(urllib_request) as response:
             total_size = int(response.headers.get("content-length", 0))
             block_size = 8192
             downloaded = 0
@@ -411,82 +382,49 @@ def retrieve_models(
 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)
-    cpack = Path(cpack)
     print(f"Installing package {cpack} to {workspace} (verbose={verbose})")
-    with contextlib.ExitStack() as stack:
+    with tempfile.TemporaryDirectory() as temp_dir:
-        if cpack.is_file():
+        pack_dir = Path(temp_dir) / ".cpack"
-            temp_dir = stack.enter_context(tempfile.TemporaryDirectory())
+        shutil.unpack_archive(cpack, pack_dir)
-            pack_dir = Path(temp_dir) / ".cpack"
-            shutil.unpack_archive(cpack, pack_dir)
-        else:
-            pack_dir = cpack
         snapshot = json.loads((pack_dir / "snapshot.json").read_text())
-        if "pips" not in snapshot:
+        req_txt_file = pack_dir / "requirements.txt"
-            raise RuntimeError(
-                "This cpack is generated by an old version of comfy-pack, "
-                "please use comfy-pack<0.4.0 to unpack it."
-            )
 
         install_comfyui(snapshot, workspace, verbose=verbose)
-        py = install_dependencies(
+        install_dependencies(snapshot, str(req_txt_file), workspace, verbose=verbose)
-            snapshot["python"],
-            [
-                str(workspace / "requirements.txt"),
-                str(
-                    workspace / "custom_nodes" / "ComfyUI-Manager" / "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,
-        )
 
         for f in (pack_dir / "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)
-        if prepare_models:
-            retrieve_models(
-                snapshot,
-                workspace,
-                verbose=verbose,
-                download=False,
-            )
 
+        retrieve_models(
+            snapshot,
+            workspace,
+            verbose=verbose,
+            download=False,
+        )
+
+        install_custom_modules(snapshot, workspace, verbose=verbose)
         if preheat:
             from .run import ComfyUIServer
 
             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:
+
-            retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models)
+        retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models)
 
 
-required_files = ["snapshot.json"]
+required_files = ["snapshot.json", "requirements.txt"]
 
 
 def build_bento(
@@ -518,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()
+    # Copy comfy-pack package
-    setup_script = source_dir / "setup_workspace.sh"
+    shutil.copytree(
-    with Path(__file__).with_name("setup_workspace.sh").open() as f:
+        COMFY_PACK_DIR, source_dir / COMFY_PACK_DIR.name, dirs_exist_ok=True
-        setup_script.write_text(f.read().replace("<SNAPSHOT>", snapshot_text))
+    )
-    # Make setup script executable in a cross-platform way
+    snapshot = json.loads((source_dir / "snapshot.json").read_text())
-    if os.name in ("posix", "mac"):
-        setup_script.chmod(setup_script.stat().st_mode | 0o755)
-    snapshot = json.loads(snapshot_text)
     return bentoml.build(
         "service:ComfyService",
         name=bento_name,
@@ -540,10 +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.sh").as_posix(),
+            "setup_script": Path(__file__).with_name("setup_workspace.py").as_posix(),
-        },
-        python={
-            "lock_packages": True,
-            "packages": ["fastapi", "comfy-pack>=0.4.0a0"],
         },
+        python={"requirements_txt": "requirements.txt", "lock_packages": True},
     )

src/comfy_pack/run.py

@@ -81,7 +81,7 @@ class ComfyUIServer:
                 self.port = 8188
         else:
             self.port = port
-        self.venv = os.path.abspath(venv) if venv else None
+        self.venv = venv
 
     def start(self) -> None:
         """
@@ -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

@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+# ruff: noqa: E402
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+virtualenv = os.environ.get("VIRTUAL_ENV")
+if virtualenv and "--reload" not in sys.argv:
+    print("Re-executing in virtualenv:", virtualenv)
+    venv_python = os.path.join(virtualenv, "bin/python3")
+    os.execl(venv_python, venv_python, *sys.argv, "--reload")
+
+# The script path is ./env/docker/setup_script
+SRC_DIR = Path(__file__).parent.parent.parent / "src"
+INPUT_DIR = SRC_DIR / "input"
+sys.path.append(str(SRC_DIR))
+
+
+def _get_workspace() -> tuple[Path, dict]:
+    import hashlib
+    import json
+
+    from bentoml._internal.configuration.containers import BentoMLContainer
+
+    snapshot = SRC_DIR / "snapshot.json"
+    checksum = hashlib.md5(snapshot.read_bytes()).hexdigest()
+    wp = (
+        Path(BentoMLContainer.bentoml_home.get()) / "run" / "comfy_workspace" / checksum
+    )
+    wp.parent.mkdir(parents=True, exist_ok=True)
+    return wp, json.loads(snapshot.read_text())
+
+
+def prepare_comfy_workspace():
+    import shutil
+
+    from comfy_pack.package import install_comfyui, install_custom_modules
+
+    verbose = int("BENTOML_DEBUG" in os.environ)
+    comfy_workspace, snapshot = _get_workspace()
+
+    if not comfy_workspace.joinpath(".DONE").exists():
+        if comfy_workspace.exists():
+            print("Removing existing 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(
+            ["chown", "-R", "bentoml:bentoml", str(comfy_workspace)], check=True
+        )
+
+
+if __name__ == "__main__":
+    prepare_comfy_workspace()

src/comfy_pack/setup_workspace.sh

@@ -1,23 +0,0 @@
-#!/bin/bash
-
-set -eo pipefail
-
-CPACK=/tmp/bento-cpack
-mkdir -p "$CPACK"
-
-cat <<EOF | head -c -1 > "$CPACK"/snapshot.json
-<SNAPSHOT>
-EOF
-
-checksum=$(md5sum "$CPACK"/snapshot.json | awk '{print $1}')
-workspace="${BENTOML_HOME:-$HOME/bentoml}/run/comfy_workspace/${checksum}"
-
-if [ -n "$VIRTUAL_ENV" ]; then
-    # shellcheck disable=SC1091
-    source "$VIRTUAL_ENV"/bin/activate
-fi
-
-set -x
-comfy-pack unpack "$CPACK" -d "$workspace" --no-models --no-venv -v
-chown -R bentoml:bentoml "$workspace"
-set +x

src/comfy_pack/utils.py

@@ -1,15 +1,18 @@
 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",
+}
 
 CPACK_PATH_INPUT_NODES = {
     "CPackInputFile",
@@ -207,7 +210,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 +221,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 +230,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,48 +239,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()
-    )
-    print("Should create zip:", should_zip, outputs, workflow)
-    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}_*"))
+            path_strs = list(output_path.glob(f"{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
-        if not should_zip:
+        return value_map
-            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"]
 
-        outs = list(output_path.glob(f"{session_id}{node_id}_*"))
+    name, node = next(iter(outputs.items()))
-        zip_paths.extend(
+    if node["class_type"] not in CPACK_OUTPUT_NODES:
-            (p, p.name.replace(f"{session_id}{node_id}", name)) for p in outs
+        raise ValueError(f"Output node {name} is not of type {CPACK_OUTPUT_NODES}")
-        )
+    node_id = node["id"]
-        if not should_zip:
+
-            if len(outs) == 1:
+    outs = list(output_path.glob(f"{session_id}{node_id}_*"))
-                return outs[0]
+    if len(outs) == 1:
-            return outs
+        return outs[0]
-    if len(zip_paths) == 1:
+    return outs
-        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:

web/main.js

@@ -247,7 +247,7 @@ class FileTreeList {
     try {
       const { workflow, output: workflow_api } = await app.graphToPrompt();
       const files = await this.getInputFiles(workflow, workflow_api);
-
+      
       // 预先将默认选中的文件添加到选中列表
       files.forEach(file => {
         const path = file.path || file;
@@ -255,14 +255,14 @@ class FileTreeList {
           this.state.toggle(path, true);
         }
       });
-
+      
       this.renderTree(this.buildTree(files));
-
+    
     // 更新所有父目录的状态
     this.container.querySelectorAll("[data-action='check-dir']").forEach(checkbox => {
       this.updateFolderState(checkbox);
     });
-
+    
     // 更新总数
     this.updateCount();
     } catch(e) {
@@ -282,12 +282,12 @@ class FileTreeList {
 
   buildTree(files) {
     const root = { name: 'root', children: {}, files: [] };
-
+    
     for (const file of files) {
       const filePath = file.path || file;
       const parts = filePath.split(/[\/\\]/);
       let current = root;
-
+      
       for (let i = 0; i < parts.length; i++) {
         const part = parts[i];
         if (i === parts.length - 1) {
@@ -313,24 +313,24 @@ class FileTreeList {
         }
       }
     }
-
+    
     // 对目录按名称排序
     for (const dir in root.children) {
       root.children[dir].files.sort((a, b) => a.name.localeCompare(b.name));
     }
-
+    
     return root;
   }
 
   renderTree(node, level = 0) {
     const dirs = Object.values(node.children);
     const hasChildren = dirs.length > 0 || node.files.length > 0;
-
+    
     this.container.innerHTML = this.renderNode(node, true);
-
+    
     this.setupEventListeners();
     this.updateCount();
-
+  
   }
 
   renderNode(node, isRoot = false) {
@@ -354,35 +354,35 @@ class FileTreeList {
 
     if (hasChildren) {
       html += `<div class="cpack-tree-children" style="display: ${isRoot ? '' : 'none'}">`;
-
+      
       // 渲染子目录
       for (const dir of dirs) {
         html += this.renderNode(dir);
       }
-
+      
       // 渲染文件
       for (const file of node.files) {
         html += `
           <div class="cpack-tree-item">
             <label style="padding-left: 17px">
-              <input type="checkbox" name="files" value="${file.path}"
+              <input type="checkbox" name="files" value="${file.path}" 
                 ${this.state.isSelected(file.path) || file.checked ? 'checked' : ''} />
               <span>${file.name}</span>
-              ${file.badges ? file.badges.map(badge =>
+              ${file.badges ? file.badges.map(badge => 
-                `<span style="background: ${badge.color || '#00a67d33'};
+                `<span style="background: ${badge.color || '#00a67d33'}; 
-                  color: ${badge.textColor || '#00a67d'};
+                  color: ${badge.textColor || '#00a67d'}; 
-                  padding: 2px 6px;
+                  padding: 2px 6px; 
-                  border-radius: 4px;
+                  border-radius: 4px; 
-                  font-size: 0.8em;
+                  font-size: 0.8em; 
                   cursor: help;
-                  white-space: nowrap;"
+                  white-space: nowrap;" 
                   title="${badge.tooltip || ''}">${badge.text}</span>`
               ).join('') : ''}
             </label>
           </div>
         `;
       }
-
+      
       html += '</div>';
     }
 
@@ -422,7 +422,7 @@ class FileTreeList {
   }
 
   setupEventListeners() {
-
+    
 
     // 目录选择功能
     this.container.querySelectorAll("[data-action='check-dir']").forEach(checkbox => {
@@ -433,7 +433,7 @@ class FileTreeList {
           const childFiles = Array.from(children.querySelectorAll("input[name='files']"))
             .map(input => input.value);
           this.state.toggleMultiple(childFiles, e.target.checked);
-
+          
           // 更新UI
           children.querySelectorAll("input[type='checkbox']").forEach(child => {
             child.checked = e.target.checked;
@@ -442,7 +442,7 @@ class FileTreeList {
           }
           });
         }
-
+        
       });
     });
 
@@ -464,7 +464,7 @@ class FileTreeList {
     this.container.querySelectorAll("input[name='files']").forEach(checkbox => {
       checkbox.addEventListener('change', (e) => {
         this.state.toggle(e.target.value, e.target.checked);
-
+      
       // 更新父文件夹状态
       const parentFolder = checkbox.closest('.cpack-tree-children')
         ?.previousElementSibling
@@ -477,7 +477,7 @@ class FileTreeList {
   }
 
   updateCount() {
-
+    
     const countSpan = document.querySelector(`[data-files-count='${this.countId}']`);
     if (countSpan) {
       countSpan.textContent = this.state.getSelectedCount();
@@ -616,7 +616,7 @@ class ModelList {
     // 只返回用户选中的模型
     return Array.from(this.container.querySelectorAll("input[name='models']:checked"))
       .map(input => input.value);
-
+  
   }
 }
 
@@ -679,7 +679,7 @@ async function createPackModal() {
     const packageOptionsContainer = form.querySelector("#package-options-container");
     const packageOptions = new PackageOptions(form, "pack-models-list", "pack-files-list", true);
     packageOptionsContainer.innerHTML = packageOptions.getHtml();
-
+    
     packageOptions.init().then(() => {
       confirmButton.disabled = false;
     });
@@ -775,7 +775,7 @@ async function unpackAction() {
   content.innerHTML = `
     <div style="margin-bottom: 20px;">
       <p style="margin-bottom: 15px;">To unpack a workflow package, follow these steps:</p>
-
+      
       <div style="margin-bottom: 15px;">
         <p style="margin-bottom: 10px;">1. Install comfy-pack CLI:</p>
         <div class="cpack-copyable" style="margin: 10px 0;">
@@ -804,13 +804,13 @@ async function unpackAction() {
       const closeButton = document.createElement("button");
       closeButton.textContent = "Close";
       closeButton.className = "cpack-btn";
-
+      
       buttonContainer.appendChild(closeButton);
-
+      
       modal.appendChild(title);
       modal.appendChild(content);
       modal.appendChild(buttonContainer);
-
+      
       const { close } = createModal(modal);
       closeButton.onclick = close;
 }
@@ -1001,13 +1001,13 @@ class PackageOptions {
   async init() {
     const modelsList = this.container.querySelector(`#${this.modelsListId}`);
     const filesList = this.container.querySelector(`#${this.filesListId}`);
-
+    
     this.modelListComponent = new ModelList(modelsList, this.modelsListId);
     this.fileListComponent = new FileTreeList(filesList, this.filesListId);
 
     const addButton = this.container.querySelector("#add-button");
     const systemPackagesArray = this.container.querySelector("#system-packages-array");
-
+    
     addButton.addEventListener("click", (e) => {
       e.preventDefault();
       const row = document.createElement("div");
@@ -1063,7 +1063,7 @@ const buildForm = `
       <p style="font-size: 0.85em; color: #888; margin-top: 5px;">
         Get your API Token at <a href="https://cloud.bentoml.com/signup?from=comfy-pack" target="_blank" style="color: #00a67d;">cloud.bentoml.com</a>
       </p>
-
+  
     </div>
 
   </div>
@@ -1112,7 +1112,7 @@ function createBuildModal() {
     const packageOptionsContainer = form.querySelector("#package-options-container");
     const packageOptions = new PackageOptions(form, "build-models-list", "build-files-list", false);
     packageOptionsContainer.innerHTML = packageOptions.getHtml();
-
+    
     packageOptions.init().then(() => {
       confirmButton.disabled = false;
     });
@@ -1234,7 +1234,7 @@ async function createServeStatusModal(url) {
   // Start status checking
   const checkInterval = setInterval(async () => {
     try {
-      const resp = await api.fetchApi("/bentoml/serve/heartbeat");
+      const resp = await api.fetchApi("/bentoml/serve/heartbeat", { method: "POST" });
       const status = await resp.json();
       if (status.error) {
         title.textContent = "Server Error";