chore: initial commit

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

pyproject.toml

@@ -0,0 +1,26 @@
+[project]
+name = "bentoml-comfyui"
+description = "BentoML extensions for ComfyUI"
+readme = "README.md"
+authors = [
+    { name = "Frost Ming", email = "frost@bentoml.com" }
+]
+requires-python = ">=3.8"
+dependencies = [
+    "bentoml>=1.3.5",
+    "comfyui-idl>=0.0.1",
+]
+dynamic = ["version"]
+
+[project.urls]
+Homepage = "https://github.com/bentoml/bentoml-comfyui"
+
+[project.entry-points."bentoml.commands"]
+comfyui = "bentoml_comfyui.cli:comfyui_command"
+
+[build-system]
+requires = ["pdm-backend"]
+build-backend = "pdm.backend"
+
+[tool.pdm.version]
+source = "scm"

src/bentoml_comfyui/_core.py

@@ -0,0 +1,67 @@
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import bentoml
+
+
+# Function to ignore the 'input' and 'output' directories during copy
+def _ignore_dirs(src, names):
+    ignore_list = ["input", "output", ".venv", ".git", "__pycache__"]
+    return [item for item in names if item in ignore_list]
+
+
+def pack_model(name: str, workspace: str) -> str:
+    """Pack the ComfyUI source to a BentoML model
+
+    Args:
+        name (str): The name of the BentoML model
+        workspace (str): The path to the ComfyUI workspace
+
+    Returns:
+        str: Model tag
+    """
+    with bentoml.models.create(name=name) as model:
+        # Copy the entire directory tree from source to destination, ignoring 'input' and 'output'
+        shutil.copytree(workspace, model.path, ignore=_ignore_dirs, dirs_exist_ok=True)
+
+        # Create empty input, output, and output/exp_data directories because they are required by ComfyUI
+        os.makedirs(os.path.join(model.path, "input"), exist_ok=True)
+        os.makedirs(os.path.join(model.path, "output"), exist_ok=True)
+        os.makedirs(os.path.join(model.path, "output", "exp_data"), exist_ok=True)
+
+    return str(model.tag)
+
+
+def _ensure_virtualenv(python: str | None) -> None:
+    from bentoml.exceptions import BentoMLConfigException
+
+    if python:
+        pyvenv_cfg = Path(python).parent.parent / "pyvenv.cfg"
+    else:
+        pyvenv_cfg = Path(sys.prefix, "pyvenv.cfg")
+
+    if not pyvenv_cfg.exists():
+        raise BentoMLConfigException("ComfyUI must be installed in a virtualenv.")
+
+
+def get_requirements(python: str | None) -> str:
+    _ensure_virtualenv(python)
+    freeze_cmd = [
+        sys.executable,
+        "-m",
+        "uv",
+        "pip",
+        "freeze",
+        "--exclude-editable",
+        "-p",
+        python or sys.executable,
+    ]
+    output = subprocess.run(
+        freeze_cmd, capture_output=True, text=True, check=True
+    ).stdout
+    # Exclude bentoml from the requirements
+    lines = [line for line in output.splitlines() if not line.startswith("bentoml==")]
+    return "\n".join(lines)

src/bentoml_comfyui/_service.tpl

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import bentoml
+import comfyui_idl
+import comfyui_idl.run
+
+REQUEST_TIMEOUT = 360
+WORKFLOW_FILE = os.path.join(os.path.dirname(__file__), "workflow.json")
+
+with open(WORKFLOW_FILE, "r") as f:
+    workflow = json.load(f)
+
+InputModel = comfyui_idl.generate_input_model(workflow)
+
+
+@bentoml.service(name={name!r})
+class ComfyUIService:
+    pipeline = bentoml.models.BentoModel({model_tag!r})
+
+    def __init__(self):
+        comfy_output_dir = os.path.join(os.getcwd(), "comfy_output")
+        comfy_temp_dir = os.path.join(os.getcwd(), "comfy_temp")
+
+        self.comfy_proc = comfyui_idl.run.WorkflowRunner(
+            self.pipeline.path,
+            comfy_output_dir,
+            comfy_temp_dir,
+        )
+        self.comfy_proc.start()
+
+    @bentoml.api(input_spec=InputModel)
+    def generate(
+        self,
+        *,
+        ctx: bentoml.Context,
+        **kwargs: t.Any,
+    ) -> Path:
+        return self.comfy_proc.run_workflow(
+            workflow, temp_dir=ctx.temp_dir, timeout=REQUEST_TIMEOUT, **kwargs
+        )
+
+    @bentoml.on_shutdown
+    def on_shutdown(self):
+        self.comfy_proc.stop()

src/bentoml_comfyui/cli.py

@@ -0,0 +1,109 @@
+import shutil
+import tempfile
+from pathlib import Path
+
+import click
+import rich
+from bentoml_cli.utils import BentoMLCommandGroup
+
+
+def _check_comfyui_workspace(workspace: str) -> None:
+    from bentoml.exceptions import InvalidArgument
+
+    comfy_fingerprints = ["comfy", "comfy_execution", "comfy_extras"]
+
+    for fingerprint in comfy_fingerprints:
+        if not Path(workspace, fingerprint).exists():
+            raise InvalidArgument(
+                f"{workspace!r} does not look like a ComfyUI workspace. Please give a correct path."
+            )
+
+
+@click.group(name="comfyui", cls=BentoMLCommandGroup)
+def comfyui_command():
+    """ComfyUI Subcommands Groups."""
+
+
+@comfyui_command.command()
+@click.option(
+    "--name",
+    type=str,
+    help="The name of the model, defaults to `comfyui`",
+    default="comfyui",
+)
+@click.option(
+    "--version", type=str, help="The version of the model, or generated if not provided"
+)
+@click.argument("workspace", type=click.Path(exists=True), default=".")
+def pack(name: str, version: str | None, workspace: str):
+    """Pack the ComfyUI workspace to a BentoML model"""
+    from ._core import pack_model
+
+    _check_comfyui_workspace(workspace)
+
+    if version:
+        name = f"{name}:{version}"
+    tag = pack_model(name, workspace)
+    rich.print(
+        f"✅ [green]Successfully packed ComfyUI workspace {workspace!r} to BentoML model {tag}[/]"
+    )
+
+
+@comfyui_command.command()
+@click.option(
+    "--name",
+    type=str,
+    help="The name of the bento, defaults to `comfyui-service`",
+    default="comfyui-service",
+)
+@click.option(
+    "--version", type=str, help="The version of the bento, or generated if not provided"
+)
+@click.option(
+    "--model",
+    type=str,
+    help="The model tag to use. Defaults to `comfyui`",
+    default="comfyui",
+)
+@click.option(
+    "-p",
+    "--python",
+    type=str,
+    help="The Python interpreter path where ComfyUI is running. Defaults to the current Python interpreter",
+)
+@click.argument("workflow", required=True, type=click.Path(dir_okay=False, exists=True))
+def build(
+    name: str, version: str | None, model: str, python: str | None, workflow: str
+):
+    """Build a BentoML service from a ComfyUI workspace"""
+    from importlib.resources import read_text
+
+    import bentoml
+
+    from ._core import get_requirements
+
+    service_template = read_text(__package__, "_service.tpl")
+
+    with tempfile.TemporaryDirectory(
+        prefix="bentoml-comfyui-", suffix="-bento"
+    ) as temp_dir:
+        parent = Path(temp_dir)
+        rich.print("📂 [blue]Creating requirements.txt[/]")
+        with open(parent.joinpath("requirements.txt"), "w") as f:
+            f.write(get_requirements(python))
+        rich.print("📂 [blue]Creating service.py[/]")
+        with open(parent.joinpath("service.py"), "w") as f:
+            f.write(service_template.format(name=name, model_tag=model))
+        rich.print("📂 [blue]Creating workflow.json[/]")
+        shutil.copy2(workflow, parent.joinpath("workflow.json"))
+        bento = bentoml.build(
+            "service:ComfyUIService",
+            name=name,
+            version=version,
+            build_ctx=temp_dir,
+            python={"requirements_txt": "requirements.txt"},
+            include=["service.py", "workflow.json", "requirements.txt"],
+        )
+    rich.print(
+        f"✅ [green]Successfully built Bento {bento.tag} from ComfyUI workflow {workflow!r}[/]"
+    )