Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: build bento from an existing bento or a cpack file #18

Merged
merged 1 commit into from
Jan 2, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 8 additions & 46 deletions nodes/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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}"})

Expand Down Expand Up @@ -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,
Expand All @@ -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")
Expand All @@ -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,
Expand All @@ -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(
Expand Down
1 change: 0 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ dependencies = [
"bentoml>=1.3.13",
"click>=8.1.7",
"comfy-cli>=1.2.8",
"pydantic>=2.9",
]
dynamic = ["version"]

Expand Down
1 change: 0 additions & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,3 @@ bentoml
fastapi
comfy-cli
duckduckgo-search
uv
56 changes: 49 additions & 7 deletions src/comfy_pack/cli.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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="")

Expand Down Expand Up @@ -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])
)
Expand Down Expand Up @@ -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,
)
77 changes: 71 additions & 6 deletions src/comfy_pack/package.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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},
)
4 changes: 1 addition & 3 deletions nodes/service.py → src/comfy_pack/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
File renamed without changes.
4 changes: 1 addition & 3 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.