Compare commits

..
1 Commits
Author SHA1 Message Date
bojiang bf1937bc4f fix: build-bento 2025-01-11 12:01:41 +08:00
15 changed files with 1508 additions and 1651 deletions
-2
View File
@@ -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
+3 -2
View File
@@ -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))
+89 -61
View File
@@ -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:
import re
return re.sub(r"[-_.]+", "-", name).lower()
def _get_requirement_string(dist: Distribution) -> str:
direct_url_text = dist.read_text("direct_url.json")
pinned_str = f'{dist.metadata["Name"]}=={dist.version}'
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:
manager_file_path = Path(
folder_paths.get_user_directory(), "default", "ComfyUI-Manager"
async def _write_requirements(path: ZPath, extras: list[str] | None = None) -> None:
print("Package => Writing requirements.txt")
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"
async def _save_snapshot() -> dict[str, Any]:
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]
stdout, _ = await proc.communicate()
if models is None:
print("Package => Writing models")
models = await _get_models()
with path.joinpath("snapshot.json").open("w") as f:
snapshot.update(
{
"python": f"{sys.version_info.major}.{sys.version_info.minor}",
"models": models,
}
)
f.write(json.dumps(snapshot, indent=2))
data = {
"python": f"{sys.version_info.major}.{sys.version_info.minor}",
"comfyui": stdout.decode().strip(),
"models": models,
"custom_nodes": await _get_custom_nodes(),
}
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(
+1 -85
View File
@@ -3,15 +3,13 @@ import json
import os
import sys
import shutil
import zipfile
import folder_paths
import node_helpers
import numpy as np
import torch
from PIL import Image, ImageOps, ImageSequence, PngImagePlugin
from PIL import Image, ImageOps, ImageSequence
from PIL.PngImagePlugin import PngInfo
from io import BytesIO
from .monkeypatch import set_bentoml_output
@@ -165,86 +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)
@@ -438,7 +356,6 @@ class AnyInput:
NODE_CLASS_MAPPINGS = {
"CPackOutputFile": OutputFile,
"CPackOutputImage": OutputImage,
"CPackOutputZip": OutputImageWithStringTxt,
"CPackInputImage": ImageInput,
"CPackInputString": StringInput,
"CPackInputInt": IntInput,
@@ -454,5 +371,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"CPackInputAny": "Any Input",
"CPackOutputImage": "Image Output",
"CPackOutputFile": "File Output",
"CPackOutputZip": "Zip Output(img + txt file)",
}
+1 -1
View File
@@ -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",
]
+1 -1
View File
@@ -1,4 +1,4 @@
bentoml>=1.3.21
bentoml
fastapi
comfy-cli
duckduckgo-search
+54 -147
View File
@@ -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,57 +255,24 @@ 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(
"--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,
):
import rich
def unpack_cmd(cpack: str, dir: str, include_disabled_models: bool, verbose: int):
from rich.console import Console
from .package import install
install(
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![/green]")
rich.print(f"{dir}")
console = Console()
rich.print(
install(cpack, dir, verbose=verbose, all_models=include_disabled_models)
console.print("\n[green]✓ ComfyUI Workspace is restored![/green]")
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"
@@ -314,7 +281,7 @@ def unpack_cmd(
def _print_schema(schema, verbose: int = 0):
import rich
from rich.console import Console
from rich.table import Table
table = Table(title="")
@@ -345,7 +312,7 @@ def _print_schema(schema, verbose: int = 0):
range_str,
)
rich.print(table)
Console().print(table)
@functools.lru_cache
@@ -373,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
@@ -382,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:
workflow = json.loads(z.read("workflow_api.json"))
console = Console()
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
@@ -414,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]")
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,
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")
@@ -473,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
@@ -490,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",
)
+62 -113
View File
@@ -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,
req_files: list[str],
snapshot: dict,
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")
venv_py = Path(sys.executable)
else:
venv = workspace / ".venv"
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,
)
subprocess.check_call(
["uv", "python", "install", python_version],
cwd=workspace,
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,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:
install_cmd.extend(["-r", str(req_file)])
if not STRICT_MODE:
if STRICT_MODE:
pass
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:
f.write("DONE")
with open(venv / "DONE", "w") as f:
f.write("DONE")
return venv_py
@@ -401,75 +382,49 @@ def retrieve_models(
def install(
cpack: str | Path,
workspace: str | Path = "workspace",
preheat: bool = False,
prepare_models: bool = True,
preheat: 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:
if cpack.is_file():
temp_dir = stack.enter_context(tempfile.TemporaryDirectory())
pack_dir = Path(temp_dir) / ".cpack"
shutil.unpack_archive(cpack, pack_dir)
else:
pack_dir = cpack
with tempfile.TemporaryDirectory() as temp_dir:
pack_dir = Path(temp_dir) / ".cpack"
shutil.unpack_archive(cpack, pack_dir)
snapshot = json.loads((pack_dir / "snapshot.json").read_text())
if "pips" not in snapshot:
raise RuntimeError(
"This cpack is generated by an old version of comfy-pack, "
"please use comfy-pack<0.4.0 to unpack it."
)
req_txt_file = pack_dir / "requirements.txt"
install_comfyui(snapshot, workspace, verbose=verbose)
py = install_dependencies(
snapshot["python"],
[
str(workspace / "requirements.txt"),
str(
workspace / "custom_nodes" / "ComfyUI-Manager" / "requirements.txt"
),
],
workspace,
no_venv=no_venv,
verbose=verbose,
)
cm_cli = workspace / "custom_nodes" / "ComfyUI-Manager" / "cm-cli.py"
subprocess.check_call(
[str(py), str(cm_cli), "restore-snapshot", str(pack_dir / "snapshot.json")],
cwd=workspace,
)
install_dependencies(snapshot, str(req_txt_file), workspace, verbose=verbose)
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(
@@ -501,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()
setup_script = source_dir / "setup_workspace.sh"
with Path(__file__).with_name("setup_workspace.sh").open() as f:
setup_script.write_text(f.read().replace("<SNAPSHOT>", snapshot_text))
# Make setup script executable in a cross-platform way
if os.name in ("posix", "mac"):
setup_script.chmod(setup_script.stat().st_mode | 0o755)
snapshot = json.loads(snapshot_text)
# 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,
@@ -523,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(),
},
python={
"lock_packages": True,
"packages": ["fastapi", "comfy-pack>=0.4.0a0"],
"setup_script": Path(__file__).with_name("setup_workspace.py").as_posix(),
},
python={"requirements_txt": "requirements.txt", "lock_packages": True},
)
+1 -4
View File
@@ -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(),
+9 -3
View File
@@ -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():
+65
View File
@@ -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()
-23
View File
@@ -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
-1
View File
@@ -12,7 +12,6 @@ if TYPE_CHECKING:
CPACK_OUTPUT_NODES = {
"CPackOutputFile",
"CPackOutputImage",
"CPackOutputZip"
}
CPACK_PATH_INPUT_NODES = {
Generated
+1184 -1170
View File
File diff suppressed because it is too large Load Diff
+38 -38
View File
@@ -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 =>
`<span style="background: ${badge.color || '#00a67d33'};
color: ${badge.textColor || '#00a67d'};
padding: 2px 6px;
border-radius: 4px;
font-size: 0.8em;
${file.badges ? file.badges.map(badge =>
`<span style="background: ${badge.color || '#00a67d33'};
color: ${badge.textColor || '#00a67d'};
padding: 2px 6px;
border-radius: 4px;
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";