chore: add packaging progress bar

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

.gitignore

@@ -160,3 +160,4 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+/temp/

__init__.py

@@ -1,3 +1,6 @@
+from .nodes import api  # noqa
 from .nodes.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+WEB_DIRECTORY = "./web"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

nodes/api.py

@@ -0,0 +1,159 @@
+import asyncio
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+import uuid
+import zipfile
+from pathlib import Path
+
+import folder_paths
+from aiohttp import web
+from server import PromptServer
+
+TEMP_FOLDER = Path(__file__).parent.parent / "temp"
+
+
+async def _write_requirements(zf: zipfile.ZipFile) -> None:
+    print("Package => Writing requirements.txt")
+    with zf.open("requirements.txt", "w") as f:
+        proc = await asyncio.subprocess.create_subprocess_exec(
+            sys.executable,
+            "-m",
+            "pip",
+            "freeze",
+            "--exclude-editable",
+            stdout=subprocess.PIPE,
+        )
+        stdout, _ = await proc.communicate()
+        f.write(stdout)
+
+
+async def _write_snapshot(zf: zipfile.ZipFile) -> None:
+    proc = await asyncio.subprocess.create_subprocess_exec(
+        "git", "rev-parse", "HEAD", stdout=subprocess.PIPE, cwd=folder_paths.base_path
+    )
+    stdout, _ = await proc.communicate()
+    with zf.open("snapshot.json", "w") as f:
+        data = {
+            "comfyui": stdout.decode().strip(),
+            "models": await _get_models(),
+            "custom_nodes": await _get_custom_nodes(),
+        }
+        f.write(json.dumps(data, indent=2).encode())
+
+
+async def _get_models() -> list:
+    print("Package => Writing models")
+    proc = await asyncio.subprocess.create_subprocess_exec(
+        "git",
+        "ls-files",
+        "--others",
+        folder_paths.models_dir,
+        stdout=subprocess.PIPE,
+    )
+    stdout, _ = await proc.communicate()
+    models = []
+    for line in stdout.decode().splitlines():
+        if os.path.basename(line).startswith("."):
+            continue
+        filename = os.path.abspath(line)
+        relpath = os.path.relpath(filename, folder_paths.base_path)
+        with open(filename, "rb") as model:
+            models.append(
+                {
+                    "filename": relpath,
+                    "sha256": hashlib.sha256(model.read()).hexdigest(),
+                }
+            )
+    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"),
+        }
+
+    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(zf: zipfile.ZipFile, data: dict) -> None:
+    print("Package => Writing workflow")
+    with zf.open("workflow_api.json", "w") as f:
+        f.write(json.dumps(data["workflow_api"], indent=2).encode())
+    with zf.open("workflow.json", "w") as f:
+        f.write(json.dumps(data["workflow"], indent=2).encode())
+
+
+async def _write_inputs(zf: zipfile.ZipFile) -> None:
+    print("Package => Writing inputs")
+    input_dir = folder_paths.get_input_directory()
+
+    for root_path, _, files in os.walk(input_dir):
+        for file in files:
+            file_path = os.path.join(root_path, file)
+            relpath = os.path.relpath(file_path, input_dir)
+            with zf.open(f"inputs/{relpath}", "w") as f:
+                with open(file_path, "rb") as input_file:
+                    shutil.copyfileobj(input_file, f)
+
+
+@PromptServer.instance.routes.post("/bentoml/pack")
+async def pack_workspace(request):
+    data = await request.json()
+    TEMP_FOLDER.mkdir(exist_ok=True)
+    older_than_1h = time.time() - 60 * 60
+    for file in TEMP_FOLDER.iterdir():
+        if file.is_file() and file.stat().st_ctime < older_than_1h:
+            file.unlink()
+
+    zip_filename = f"{uuid.uuid4()}.zip"
+
+    with zipfile.ZipFile(TEMP_FOLDER / zip_filename, "w") as zf:
+        await _write_requirements(zf)
+        await _write_snapshot(zf)
+        await _write_workflow(zf, data)
+        await _write_inputs(zf)
+
+    return web.json_response({"download_url": f"/bentoml/download/{zip_filename}"})
+
+
+@PromptServer.instance.routes.get("/bentoml/download/{zip_filename}")
+async def download_workspace(request):
+    zip_filename = request.match_info["zip_filename"]
+    return web.FileResponse(TEMP_FOLDER / zip_filename)

nodes/nodes.py

@@ -22,7 +22,7 @@ class AnyType(str):
 anytype = AnyType("*")  # when a != b is called, it will always return False
 
 
-class OutputPath:
+class OutputFile:
     COLOR = (142, 36, 170)
     OUTPUT_NODE = True
 
@@ -180,7 +180,6 @@ class ImageInput:
     CATEGORY = "ComfyUI-IDL/input"
     BENTOML_NODE = True
     RETURN_TYPES = ("IMAGE", "MASK")
-    RETUEN_NAMES = ("image", "mask")
     FUNCTION = "load_image"
 
     def load_image(self, image):
@@ -243,111 +242,7 @@ class ImageInput:
         return True
 
 
-class StringInput:
-    COLOR = (142, 36, 170)
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "value": ("STRING", {"default": ""}),
-            }
-        }
-
-    RETUEN_NAMES = ("string",)
-    RETURN_TYPES = (anytype,)
-    BENTOML_NODE = True
-    FUNCTION = "string_input"
-    CATEGORY = "ComfyUI-IDL/input"
-
-    def string_input(self, value):
-        return (value,)
-
-    @classmethod
-    def VALIDATE_INPUTS(s, value):
-        set_bentoml_output([(value,)])
-        return True
-
-
-class IntegerInput:
-    COLOR = (142, 36, 170)
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "value": ("INT", {"default": 0}),
-            }
-        }
-
-    RETUEN_NAMES = ("int",)
-    RETURN_TYPES = (anytype,)
-    FUNCTION = "identity"
-    BENTOML_NODE = True
-    CATEGORY = "ComfyUI-IDL/input"
-
-    def identity(self, value):
-        return (value,)
-
-    @classmethod
-    def VALIDATE_INPUTS(s, value):
-        set_bentoml_output([(value,)])
-        return True
-
-
-class FloatInput:
-    COLOR = (142, 36, 170)
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "value": ("FLOAT", {"default": 0.0}),
-            }
-        }
-
-    RETUEN_NAMES = ("float",)
-    RETURN_TYPES = (anytype,)
-    FUNCTION = "identity"
-    BENTOML_NODE = True
-    CATEGORY = "ComfyUI-IDL/input"
-
-    def identity(self, value):
-        return (value,)
-
-    @classmethod
-    def VALIDATE_INPUTS(s, value):
-        set_bentoml_output([(value,)])
-        return True
-
-
-class BooleanInput:
-    COLOR = (142, 36, 170)
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "value": ("BOOLEAN", {"default": False}),
-            }
-        }
-
-    RETUEN_NAMES = ("bool",)
-    RETURN_TYPES = (anytype,)
-    FUNCTION = "identity"
-    CATEGORY = "ComfyUI-IDL/input"
-    BENTOML_NODE = True
-
-    def identity(self, value):
-        return (value,)
-
-    @classmethod
-    def VALIDATE_INPUTS(s, value):
-        set_bentoml_output([(value,)])
-        return True
-
-
-class PathInput:
+class FileInput:
     COLOR = (142, 36, 170)
 
     @classmethod
@@ -359,7 +254,7 @@ class PathInput:
         }
 
     RETURN_TYPES = (anytype,)
-    RETUEN_NAMES = ("path",)
+    RETURN_NAMES = ("path",)
     FUNCTION = "identity"
     BENTOML_NODE = True
     CATEGORY = "ComfyUI-IDL/input"
@@ -373,24 +268,44 @@ class PathInput:
         return True
 
 
+class ValueInput:
+    COLOR = (142, 36, 170)
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input": ("*", {"default": ""}),
+            }
+        }
+
+    RETURN_TYPES = (anytype,)
+    RENAME = ("value",)
+    FUNCTION = "identity"
+    BENTOML_NODE = True
+    CATEGORY = "ComfyUI-IDL/input"
+
+    def identity(self, input):
+        return (input,)
+
+    @classmethod
+    def VALIDATE_INPUTS(s, input):
+        set_bentoml_output([(input,)])
+        return True
+
+
 NODE_CLASS_MAPPINGS = {
-    "BentoOutputPath": OutputPath,
+    "BentoOutputFile": OutputFile,
     "BentoOutputImage": OutputImage,
     "BentoInputImage": ImageInput,
-    "BentoInputString": StringInput,
-    "BentoInputInteger": IntegerInput,
-    "BentoInputFloat": FloatInput,
-    "BentoInputBoolean": BooleanInput,
-    "BentoInputPath": PathInput,
+    "BentoInputFile": FileInput,
+    "BentoInputValue": ValueInput,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "BentoInputImage": "Image Input",
-    "BentoInputPath": "Path Input",
-    "BentoInputInteger": "Integer Input",
-    "BentoInputFloat": "Float Input",
-    "BentoInputBoolean": "Boolean Input",
-    "BentoInputString": "String Input",
+    "BentoInputFile": "File Input",
     "BentoOutputImage": "Image Output",
-    "BentoOutputPath": "Path Output",
+    "BentoOutputFile": "File Output",
+    "BentoInputValue": "Plain Value Input",
 }

src/comfyui_idl/__init__.py

@@ -1,12 +1,11 @@
+from comfyui_idl.run import WorkflowRunner
 from comfyui_idl.utils import (
-    parse_workflow,
     generate_input_model,
+    parse_workflow,
     populate_workflow,
     retrieve_workflow_outputs,
 )
 
-from comfyui_idl.run import WorkflowRunner
-
 __all__ = [
     "WorkflowRunner",
     "parse_workflow",

src/comfyui_idl/utils.py

@@ -2,20 +2,11 @@ from __future__ import annotations
 
 import re
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Literal, Union
 
 if TYPE_CHECKING:
     from pydantic import BaseModel
 
-CLASS_TYPES = {
-    "BentoInputString": str,
-    "BentoInputBoolean": bool,
-    "BentoInputInteger": int,
-    "BentoInputFloat": float,
-    "BentoInputPath": Path,
-    "BentoInputImage": Path,
-}
-
 BENTO_OUTPUT_NODES = {
     "BentoOutputPath",
     "BentoOutputImage",
@@ -68,6 +59,13 @@ def _get_node_identifier(node, dep_map=None) -> str:
     return _normalize_to_identifier(title)
 
 
+API_WORKFLOW_MESSAGE = """
+It seems you are trying to parse an ordinary Workflow json file.
+Please save the Workflow with the Save (API Format).
+If you don't have this button, you must enable the "Dev mode Options" by clicking the Settings button on the top right (gear icon). Check the setting option "Enable Dev Mode options". After that, the Button Save (API Format) should appear.
+"""
+
+
 def _parse_workflow(workflow: dict) -> tuple[dict[str, Any], dict[str, Any]]:
     """
     Parse the workflow template and return the input and output definition
@@ -76,6 +74,9 @@ def _parse_workflow(workflow: dict) -> tuple[dict[str, Any], dict[str, Any]]:
     outputs = {}
     dep_map = {}
 
+    if "last_node_id" in workflow:
+        raise ValueError(API_WORKFLOW_MESSAGE)
+
     for id, node in workflow.items():
         for input_name, v in node["inputs"].items():
             if isinstance(v, list) and len(v) == 2:  # is a link
@@ -118,21 +119,40 @@ def generate_input_model(workflow: dict) -> type[BaseModel]:
         ValueError: If an unsupported class type is encountered in the workflow.
     """
     from pydantic import Field, create_model
+    from pydantic_core import PydanticUndefined
 
     inputs, _ = _parse_workflow(workflow)
 
     input_fields = {}
     for name, node in inputs.items():
         class_type = node["class_type"]
-        if class_type in CLASS_TYPES:
-            ann = CLASS_TYPES[class_type]
-            if class_type in BENTO_PATH_INPUT_NODES:
-                field = (ann, Field())
+        if class_type in BENTO_PATH_INPUT_NODES:
+            field = (Path, Field())
+        elif class_type == "BentoInputValue":
+            options = node.get("_meta", {}).get("options")
+            value = _get_node_value(node)
+            if not options:
+                field = (type(value), Field(default=value))
             else:
-                field = (ann, Field(default=_get_node_value(node)))
-            input_fields[name] = field
+                if values := options.get("values"):  # combo type
+                    field = (Literal[tuple(values)], Field(default=value))
+                elif any(
+                    f in options for f in ("min", "max", "round", "precision", "step")
+                ):  # must be number types
+                    type_ = float if options.get("round", 1) < 1 else int
+                    field = (
+                        type_,
+                        Field(
+                            default=value,
+                            ge=options.get("min", PydanticUndefined),
+                            le=options.get("max", PydanticUndefined),
+                        ),
+                    )
+                else:
+                    field = (type(value), Field(default=value))
         else:
             raise ValueError(f"Unsupported class type: {class_type}")
+        input_fields[name] = field
     return create_model("ParsedWorkflowTemplate", **input_fields)
 
 

web/dynamic.js

@@ -0,0 +1,75 @@
+import { app } from "../../scripts/app.js";
+
+function mimicNode(node, target, slot) {
+  const getWidget = (widget) => {
+    if (typeof widget.origType === "undefined") return widget;
+    const newWidget = Object.assign({}, widget);
+    for (const key in widget) {
+      if (key.startsWith("orig") || key === "computeSize") {
+        delete newWidget[key];
+      }
+    }
+    newWidget.type = widget.origType;
+    newWidget.name = inputName;
+    return newWidget;
+  }
+  const input = target.inputs[slot];
+  const inputName = node.inputs.length > 0 ? node.inputs[0].name : node.widgets[0].name;
+  if (typeof input.widget === "undefined") {  // This is an input
+    node.inputs = [{ ...input, name: inputName }];
+    node.widgets = [];
+  } else {  // This is a widget
+    node.widgets = [getWidget(target.widgets.find(w => w.name === input.name))];
+    node.inputs = [];
+    node.size = node.computeSize();
+    requestAnimationFrame(() => {
+      if (node.onResize) {
+        node.onResize(node.size);
+      }
+    });
+  }
+}
+
+app.registerExtension({
+	name: "Comfy.DynamicInput",
+  extensionNodes: ["BentoInputValue", "BentoInputFile"],
+
+	async beforeRegisterNodeDef(nodeType, nodeData) {
+    if (!this.extensionNodes.includes(nodeData.name)) return;
+
+    nodeType.prototype.onConnectOutput = function () {
+      if (this.outputs[0].links?.length > 0) return false;
+      const target = arguments[3];
+      if (this.type === "BentoInputFile" && !["COMBO", "STRING"].includes(target.inputs[arguments[4]].type))
+        return false;
+      mimicNode(this, target, arguments[4]);
+      this.title = target.inputs[arguments[4]].name;
+    }
+  },
+
+  async setup(app) {
+    app.graph.nodes.forEach((node) => {
+      if (!this.extensionNodes.includes(node.type)) return;
+      if (node.outputs.length > 0 && node.outputs[0].links.length > 0) {
+        const link = node.graph.links[node.outputs[0].links[0]];
+        mimicNode(node, app.graph.getNodeById(link.target_id), link.target_slot);
+      }
+    })
+  },
+
+  async init(app) {
+    const originalToPrompt = app.graphToPrompt;
+    const self = this;
+    app.graphToPrompt = async function(graph = app.graph, clean = true) {
+      const { workflow, output } = await originalToPrompt(graph, clean);
+      Object.entries(output).forEach(([id, nodeData]) => {
+        if (!self.extensionNodes.includes(nodeData.class_type)) return;
+        const node = graph.getNodeById(parseInt(id));
+        if (node.widgets.length === 0) return;
+        const widget = node.widgets[0];
+        nodeData["_meta"] = Object.assign({}, nodeData["_meta"] || {}, { options: widget.options });
+      });
+      return { workflow, output };
+    };
+  }
+});

web/main.js

@@ -0,0 +1,119 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+function createDownloadModal() {
+  const modal = document.createElement("div");
+  modal.style.cssText = `
+    position: fixed;
+    top: 50%;
+    left: 50%;
+    transform: translate(-50%, -50%);
+    background: #1a1a1a;
+    padding: 20px;
+    border-radius: 8px;
+    z-index: 1000;
+    min-width: 300px;
+  `;
+
+  const title = document.createElement("h3");
+  title.textContent = "Packaging...";
+  title.style.marginBottom = "15px";
+  title.style.color = "#fff";
+
+  const progress = document.createElement("div");
+  progress.style.cssText = `
+    width: 100%;
+    height: 20px;
+    background: #333;
+    border-radius: 10px;
+    overflow: hidden;
+  `;
+
+  const progressBar = document.createElement("div");
+  progressBar.style.cssText = `
+    width: 0%;
+    height: 100%;
+    background: #00a67d;
+    transition: width 0.3s ease;
+  `;
+
+  progress.appendChild(progressBar);
+  modal.appendChild(title);
+  modal.appendChild(progress);
+
+  const overlay = document.createElement("div");
+  overlay.style.cssText = `
+    position: fixed;
+    top: 0;
+    left: 0;
+    right: 0;
+    bottom: 0;
+    background: rgba(0, 0, 0, 0.5);
+    z-index: 999;
+  `;
+
+  document.body.appendChild(overlay);
+  document.body.appendChild(modal);
+
+  return {
+    modal,
+    overlay,
+    progressBar,
+    updateProgress: (percent) => {
+      progressBar.style.width = `${percent}%`;
+    },
+    close: () => {
+      modal.remove();
+      overlay.remove();
+    }
+  };
+}
+
+app.registerExtension({
+  name: "Comfy.BentoExtension",
+
+  async setup() {
+    const menu = document.querySelector(".comfy-menu");
+    const separator = document.createElement("hr");
+
+    separator.style.margin = "20px 0";
+    separator.style.width = "100%";
+    menu.append(separator);
+    const packButton = document.createElement("button");
+    packButton.textContent = "Package";
+    packButton.onclick = async () => {
+      packButton.disabled = true;
+      const downloadModal = createDownloadModal();
+      
+      try {
+        downloadModal.updateProgress(20);
+        const { workflow, output: workflow_api } = await app.graphToPrompt();
+        
+        downloadModal.updateProgress(40);
+        const body = JSON.stringify({ workflow, workflow_api });
+        
+        downloadModal.updateProgress(60);
+        const resp = await api.fetchApi("/bentoml/pack", { method: "POST", body, headers: { "Content-Type": "application/json" } });
+        
+        downloadModal.updateProgress(80);
+        const downloadUrl = (await resp.json())["download_url"];
+
+        downloadModal.updateProgress(100);
+        const link = document.createElement("a");
+        link.href = downloadUrl;
+        link.download = "workspace.zip";
+        link.click();
+
+        setTimeout(() => {
+          downloadModal.close();
+        }, 1000);
+      } catch (error) {
+        console.error("Package failed:", error);
+        downloadModal.close();
+      } finally {
+        packButton.disabled = false;
+      }
+    }
+    menu.append(packButton);
+  }
+});