fix: build-bento

nodes/api.py

@@ -31,7 +31,7 @@ EXCLUDE_PACKAGES = ["bentoml", "onnxruntime", "conda"]  # TODO: standardize this
 
 def _get_requirement_string(dist: Distribution) -> str:
     direct_url_text = dist.read_text("direct_url.json")
-    pinned_str = f"{dist.metadata['Name']}=={dist.version}"
+    pinned_str = f'{dist.metadata["Name"]}=={dist.version}'
     if not direct_url_text:
         return pinned_str
     direct_url = json.loads(direct_url_text)
@@ -48,21 +48,10 @@ def _get_requirement_string(dist: Distribution) -> str:
         return pinned_str
 
 
-def normalize_name(name: str) -> str:
-    import re
-
-    return re.sub(r"[-_.]+", "-", name).lower()
-
-
 async def _write_requirements(path: ZPath, extras: list[str] | None = None) -> None:
     print("Package => Writing requirements.txt")
-    seen_packages: set[str] = set(EXCLUDE_PACKAGES)
     with path.joinpath("requirements.txt").open("w") as f:
         for dist in distributions():
-            pkg_name = normalize_name(dist.metadata["Name"])
-            if pkg_name in seen_packages:
-                continue
-            seen_packages.add(pkg_name)
             f.write(_get_requirement_string(dist) + "\n")
         if extras:
             f.write("\n".join(extras) + "\n")
@@ -169,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:
@@ -288,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
@@ -298,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)
@@ -482,7 +469,7 @@ async def _prepare_pack(
         model_filter=model_filter,
     )
 
-    await _write_requirements(working_dir, ["comfy-cli", "fastapi", "comfy-pack"])
+    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)
@@ -555,7 +542,7 @@ async def build_bento_api(request):
             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(

src/comfy_pack/cli.py

@@ -404,7 +404,6 @@ def run(ctx, cpack: str, output_dir: str, help: bool, verbose: int):
             workflow,
             Path(output_dir).absolute(),
             verbose=verbose,
-            workspace=server.workspace,
             **validated_data.model_dump(),
         )
         console.print("\n[green]✓ Workflow is executed successfully![/green]")
@@ -444,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
 

src/comfy_pack/package.py

@@ -29,28 +29,24 @@ def _clone_commit(url: str, commit: str, dir: Path, verbose: int = 0):
         ["git", "clone", "--recurse-submodules", "--filter=blob:none", url, dir],
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
     )
     subprocess.check_call(
         ["git", "fetch", "-q", url, commit],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
     )
     subprocess.check_call(
         ["git", "checkout", "FETCH_HEAD"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
     )
     subprocess.check_call(
         ["git", "submodule", "update", "--init", "--recursive"],
         cwd=dir,
         stdout=stdout,
         stderr=stderr,
-        env={"GIT_TERMINAL_PROMPT": "0"},
     )
 
 
@@ -460,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.py"
-    with Path(__file__).with_name("setup_workspace.py").open() as f:
-        setup_script.write_text(f.read().format(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,
@@ -482,7 +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.py").as_posix(),
+            "setup_script": Path(__file__).with_name("setup_workspace.py").as_posix(),
         },
         python={"requirements_txt": "requirements.txt", "lock_packages": True},
     )

src/comfy_pack/run.py

@@ -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

@@ -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,
             **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

@@ -13,9 +13,10 @@ if virtualenv and "--reload" not in sys.argv:
     venv_python = os.path.join(virtualenv, "bin/python3")
     os.execl(venv_python, venv_python, *sys.argv, "--reload")
 
-SNAPSHOT = """\
-{snapshot}
-"""
+# 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]:
@@ -24,12 +25,13 @@ def _get_workspace() -> tuple[Path, dict]:
 
     from bentoml._internal.configuration.containers import BentoMLContainer
 
-    checksum = hashlib.md5(SNAPSHOT.strip().encode("utf8")).hexdigest()
+    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)
+    return wp, json.loads(snapshot.read_text())
 
 
 def prepare_comfy_workspace():
@@ -37,7 +39,7 @@ def prepare_comfy_workspace():
 
     from comfy_pack.package import install_comfyui, install_custom_modules
 
-    verbose = 2
+    verbose = int("BENTOML_DEBUG" in os.environ)
     comfy_workspace, snapshot = _get_workspace()
 
     if not comfy_workspace.joinpath(".DONE").exists():
@@ -46,6 +48,12 @@ def prepare_comfy_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(

uv.lock

@@ -456,7 +456,7 @@ name = "click"
 version = "8.1.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "colorama", marker = "platform_system == 'Windows'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 }
 wheels = [
@@ -523,6 +523,7 @@ wheels = [
 
 [[package]]
 name = "comfy-pack"
+version = "0.2.2.dev7+g5238115.d20241231"
 source = { editable = "." }
 dependencies = [
     { name = "bentoml" },