From 0e488f27fdcd345902eff4c2030fbf730a991209 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Tue, 31 Dec 2024 17:49:06 +0800 Subject: [PATCH] feat: build bento from an existing bento or a cpack file Signed-off-by: Frost Ming --- nodes/api.py | 54 ++------------ pyproject.toml | 1 - requirements.txt | 1 - src/comfy_pack/cli.py | 56 ++++++++++++-- src/comfy_pack/package.py | 77 ++++++++++++++++++-- {nodes => src/comfy_pack}/service.py | 4 +- {nodes => src/comfy_pack}/setup_workspace.py | 0 uv.lock | 4 +- 8 files changed, 130 insertions(+), 67 deletions(-) rename {nodes => src/comfy_pack}/service.py (98%) rename {nodes => src/comfy_pack}/setup_workspace.py (100%) diff --git a/nodes/api.py b/nodes/api.py index e9dcd1f..5676910 100644 --- a/nodes/api.py +++ b/nodes/api.py @@ -20,7 +20,7 @@ from comfy_pack.hash import async_batch_get_sha256 from comfy_pack.model_helper import alookup_model_source -from comfy_pack.utils import get_self_git_commit +from comfy_pack.package import build_bento ZPath = Union[Path, zipfile.Path] TEMP_FOLDER = Path(__file__).parent.parent / "temp" @@ -248,10 +248,7 @@ async def pack_workspace(request): with zipfile.ZipFile(TEMP_FOLDER / zip_filename, "w") as zf: path = zipfile.Path(zf) - await _prepare_bento_project( - path, - data, - ) + await _prepare_pack(path, data) return web.json_response({"download_url": f"/bentoml/download/{zip_filename}"}) @@ -447,11 +444,11 @@ async def download_workspace(request): return web.FileResponse(TEMP_FOLDER / zip_filename) -async def _prepare_bento_project( +async def _prepare_pack( working_dir: ZPath, data: dict, store_models: bool = False, -): +) -> None: model_filter = set(data.get("models", [])) models = await _get_models( store_models=store_models, @@ -462,20 +459,6 @@ async def _prepare_bento_project( await _write_snapshot(working_dir, data, models) await _write_workflow(working_dir, data) await _write_inputs(working_dir, data) - with working_dir.joinpath("service.py").open("w") as f: - f.write(Path(__file__).with_name("service.py").read_text()) - - # Copy comfy_pack directory - if isinstance(working_dir, Path): - shutil.copytree(COMFY_PACK_DIR, working_dir / COMFY_PACK_DIR.name) - else: # zipfile.Path - for src in COMFY_PACK_DIR.rglob("*"): - if src.is_file(): - rel_path = src.relative_to(COMFY_PACK_DIR.parent) - with working_dir.joinpath(rel_path).open("wb") as f: - f.write(src.read_bytes()) - models = [m for m in models if not m["disabled"]] # filter out disabled models - return models @PromptServer.instance.routes.post("/bentoml/model/query") @@ -490,7 +473,7 @@ async def get_models(request): @PromptServer.instance.routes.post("/bentoml/build") -async def build_bento(request): +async def build_bento_api(request): """Request body: { workflow_api: dict, workflow: dict, @@ -509,33 +492,12 @@ async def build_bento(request): with tempfile.TemporaryDirectory(suffix="-bento", prefix="comfy-pack-") as temp_dir: temp_dir_path = Path(temp_dir) - models = await _prepare_bento_project(temp_dir_path, data, store_models=True) + await _prepare_pack(temp_dir_path, data, store_models=True) # create a bento try: - bento = bentoml.build( - "service:ComfyService", - name=data["bento_name"], - build_ctx=temp_dir, - labels={"comfy-pack-version": get_self_git_commit() or "unknown"}, - models=[m["model_tag"] for m in models if "model_tag" in m], - docker={ - "python_version": f"{sys.version_info.major}.{sys.version_info.minor}", - "system_packages": [ - "git", - "libglib2.0-0", - "libsm6", - "libxrender1", - "libxext6", - "ffmpeg", - "libstdc++-12-dev", - *data.get("system_packages", []), - ], - "setup_script": Path(__file__) - .with_name("setup_workspace.py") - .as_posix(), - }, - python={"requirements_txt": "requirements.txt", "lock_packages": True}, + bento = build_bento( + data["bento_name"], temp_dir_path, data.get("system_packages") ) except bentoml.exceptions.BentoMLException as e: return web.json_response( diff --git a/pyproject.toml b/pyproject.toml index 8369db7..ca352ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ dependencies = [ "bentoml>=1.3.13", "click>=8.1.7", "comfy-cli>=1.2.8", - "pydantic>=2.9", ] dynamic = ["version"] diff --git a/requirements.txt b/requirements.txt index 1c5b78b..89b64fb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,3 @@ bentoml fastapi comfy-cli duckduckgo-search -uv diff --git a/src/comfy_pack/cli.py b/src/comfy_pack/cli.py index 78a6c97..870bacc 100644 --- a/src/comfy_pack/cli.py +++ b/src/comfy_pack/cli.py @@ -1,12 +1,15 @@ -import click import functools import json -from pathlib import Path +import os import shutil import subprocess import sys import tempfile -from .const import WORKSPACE_DIR, COMFYUI_REPO, COMFY_PACK_REPO, COMFYUI_MANAGER_REPO +from pathlib import Path + +import click + +from .const import COMFY_PACK_REPO, COMFYUI_MANAGER_REPO, COMFYUI_REPO, WORKSPACE_DIR from .hash import get_sha256 from .utils import get_self_git_commit @@ -50,9 +53,10 @@ def main(): help="Increase verbosity level", ) def init(dir: str, verbose: int): - from rich.console import Console import os + from rich.console import Console + console = Console() # Check if directory path is valid @@ -258,9 +262,10 @@ def init(dir: str, verbose: int): help="Increase verbosity level (use multiple times for more verbosity)", ) def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int): - from .package import install from rich.console import Console + from .package import install + console = Console() install(cpack, dir, verbose=verbose, all_models=include_disabled_models) @@ -276,8 +281,8 @@ def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int def _print_schema(schema, verbose: int = 0): - from rich.table import Table from rich.console import Console + from rich.table import Table table = Table(title="") @@ -335,10 +340,11 @@ def _get_cache_workspace(cpack: str): ) @click.pass_context def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int): - from .utils import generate_input_model from pydantic import ValidationError from rich.console import Console + from .utils import generate_input_model + inputs = dict( zip([k.lstrip("-").replace("-", "_") for k in ctx.args[::2]], ctx.args[1::2]) ) @@ -411,3 +417,39 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int): console.print(f"{i}: {value}") else: console.print(results) + + +@main.command(name="build-bento") +@click.option("--name", help="Name of the bento service") +@click.option("--version", help="Version of the bento service") +@click.argument("source") +def bento_cmd(source: str, name: str | None, version: str | None): + """Build a bento from the source, which can be either a .cpack.zip file or a bento tag.""" + import bentoml + from bentoml.bentos import BentoBuildConfig + + from .package import build_bento + + with tempfile.TemporaryDirectory() as temp_dir: + if source.endswith(".cpack.zip"): + name = name or os.path.basename(source).replace(".cpack.zip", "") + shutil.unpack_archive(source, temp_dir) + system_packages = None + include_default_system_packages = True + else: + existing_bento = bentoml.get(source) + name = name or existing_bento.tag.name + shutil.copytree(existing_bento.path, temp_dir, dirs_exist_ok=True) + build_config = BentoBuildConfig.from_bento_dir( + existing_bento.path_of("src") + ) + system_packages = build_config.docker.system_packages + include_default_system_packages = False + + build_bento( + name, + Path(temp_dir), + version=version, + system_packages=system_packages, + include_default_system_packages=include_default_system_packages, + ) diff --git a/src/comfy_pack/package.py b/src/comfy_pack/package.py index 0d9bc07..7dfdc35 100644 --- a/src/comfy_pack/package.py +++ b/src/comfy_pack/package.py @@ -3,15 +3,23 @@ import json import os import shutil -import urllib.parse import subprocess import sys import tempfile import threading +import urllib.parse import urllib.request -from .hash import get_sha256 from pathlib import Path -from .const import MODEL_DIR, COMFYUI_REPO +from typing import TYPE_CHECKING + +from .const import COMFYUI_REPO, MODEL_DIR +from .hash import get_sha256 +from .utils import get_self_git_commit + +if TYPE_CHECKING: + import bentoml + +COMFY_PACK_DIR = Path(__file__).parent def _clone_commit(url: str, commit: str, dir: Path, verbose: int = 0): @@ -238,7 +246,7 @@ def create_model_symlink(global_path: Path, sha: str, target_path: Path, filenam os.symlink(source, target) -def retrive_models( +def retrieve_models( snapshot: dict, workspace: Path, download: bool = True, @@ -387,7 +395,7 @@ def install( elif f.is_dir(): shutil.copytree(f, workspace / "input" / f.name, dirs_exist_ok=True) - retrive_models( + retrieve_models( snapshot, workspace, verbose=verbose, @@ -405,4 +413,61 @@ def install( ) as _: pass - retrive_models(snapshot, workspace, verbose=verbose, all_models=all_models) + retrieve_models(snapshot, workspace, verbose=verbose, all_models=all_models) + + +required_files = ["snapshot.json", "requirements.txt"] + + +def build_bento( + bento_name: str, + source_dir: Path, + *, + version: str | None = None, + system_packages: list[str] | None = None, + include_default_system_packages: bool = True, +) -> bentoml.Bento: + import bentoml + + for f in required_files: + if not (source_dir / f).exists(): + raise FileNotFoundError(f"Not a valid comfy-pack package: missing `{f}`") + + if include_default_system_packages: + system_packages = [ + "git", + "libglib2.0-0", + "libsm6", + "libxrender1", + "libxext6", + "ffmpeg", + "libstdc++-12-dev", + *(system_packages or []), + ] + else: + system_packages = system_packages or [] + + shutil.copy2(Path(__file__).with_name("service.py"), source_dir / "service.py") + # Copy comfy-pack package + shutil.copytree( + COMFY_PACK_DIR, source_dir / COMFY_PACK_DIR.name, dirs_exist_ok=True + ) + snapshot = json.loads((source_dir / "snapshot.json").read_text()) + return bentoml.build( + "service:ComfyService", + name=bento_name, + version=version, + build_ctx=str(source_dir), + labels={"comfy-pack-version": get_self_git_commit() or "unknown"}, + models=[ + m["model_tag"] + for m in snapshot["models"] + if "model_tag" in m and not m.get("disabled", False) + ], + docker={ + "python_version": f"{sys.version_info.major}.{sys.version_info.minor}", + "system_packages": system_packages, + "setup_script": Path(__file__).with_name("setup_workspace.py").as_posix(), + }, + python={"requirements_txt": "requirements.txt", "lock_packages": True}, + ) diff --git a/nodes/service.py b/src/comfy_pack/service.py similarity index 98% rename from nodes/service.py rename to src/comfy_pack/service.py index 96896aa..ae507db 100644 --- a/nodes/service.py +++ b/src/comfy_pack/service.py @@ -177,10 +177,8 @@ def prepare_models(): ) -if not EXISTING_COMFYUI_SERVER: +if False and not EXISTING_COMFYUI_SERVER: for model in snapshot["models"]: - if True: - continue if model.get("disabled"): continue source = model["source"] diff --git a/nodes/setup_workspace.py b/src/comfy_pack/setup_workspace.py similarity index 100% rename from nodes/setup_workspace.py rename to src/comfy_pack/setup_workspace.py diff --git a/uv.lock b/uv.lock index 91014b5..f979429 100644 --- a/uv.lock +++ b/uv.lock @@ -523,13 +523,12 @@ wheels = [ [[package]] name = "comfy-pack" -version = "0.1.4.dev10+g1d19c05" +version = "0.2.2.dev7+g5238115.d20241231" source = { editable = "." } dependencies = [ { name = "bentoml" }, { name = "click" }, { name = "comfy-cli" }, - { name = "pydantic" }, ] [package.metadata] @@ -537,7 +536,6 @@ requires-dist = [ { name = "bentoml", specifier = ">=1.3.13" }, { name = "click", specifier = ">=8.1.7" }, { name = "comfy-cli", specifier = ">=1.2.8" }, - { name = "pydantic", specifier = ">=2.9" }, ] [[package]]