249 lines
8.8 KiB
Python
249 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_PLUGIN_DIR = Path(__file__).resolve().parent
|
|
_COMFYUI_ROOT = _PLUGIN_DIR.parent.parent
|
|
|
|
def _workflow_ui_link_widget_order() -> list[str]:
|
|
order = ["form_label"]
|
|
for i in range(8):
|
|
order.append(f"type_{i}")
|
|
order.append(f"name_{i}")
|
|
for i in range(8):
|
|
order.append(f"input_text_{i}")
|
|
for i in range(8):
|
|
order.append(f"input_number_{i}")
|
|
for i in range(8):
|
|
order.append(f"input_image_{i}")
|
|
for i in range(8):
|
|
order.append(f"input_video_{i}")
|
|
for i in range(8):
|
|
order.append(f"input_audio_{i}")
|
|
for i in range(8):
|
|
order.append(f"input_boolean_{i}")
|
|
return order
|
|
|
|
_WIDGET_ORDER: dict[str, list[str]] = {
|
|
"WorkflowUILink": _workflow_ui_link_widget_order(),
|
|
"WorkflowUI Link": _workflow_ui_link_widget_order(),
|
|
"CLIPTextEncode": ["text"],
|
|
"Text Multiline": ["text"],
|
|
"PrimitiveStringMultiline": ["value"],
|
|
"EmptyLatentImage": ["width", "height", "batch_size"],
|
|
"EmptySD3LatentImage": ["width", "height", "batch_size"],
|
|
"SDXLEmptyLatentSizePicker+": ["resolution", "batch_size", "width_override", "height_override"],
|
|
"KSampler": ["seed", "steps", "cfg", "sampler_name", "scheduler", "denoise"],
|
|
"BasicScheduler": ["scheduler", "steps", "denoise"],
|
|
"UNETLoader": ["unet_name", "weight_dtype"],
|
|
"DiffusionModelLoader": ["unet_name", "weight_dtype"],
|
|
"LoadDiffusionModel": ["unet_name", "weight_dtype"],
|
|
"StableCascadeCheckpointLoader": ["key_opt_b", "key_opt_c", "cache_mode"],
|
|
"StableCascade_CheckpointLoader": ["key_opt_b", "key_opt_c", "cache_mode"],
|
|
"SD3CheckpointLoader": ["ckpt_name", "shift"],
|
|
"SD3LoadCheckpoint": ["ckpt_name", "shift"],
|
|
"SaveImage": ["filename_prefix"],
|
|
"Save Image": ["filename_prefix"],
|
|
"Save Image (api)": ["filename_prefix"],
|
|
"SaveImageNode": ["filename_prefix"],
|
|
}
|
|
|
|
|
|
def get_workflows_directory() -> Path | None:
|
|
env_path = os.environ.get("WORKFLOWUI_WORKFLOWS_DIR", "").strip()
|
|
if env_path:
|
|
p = Path(env_path).resolve()
|
|
if p.is_dir():
|
|
return p
|
|
return None
|
|
default = _COMFYUI_ROOT / "user" / "default" / "workflows"
|
|
if default.is_dir():
|
|
return default
|
|
return None
|
|
|
|
|
|
def list_workflows() -> list[dict[str, str]]:
|
|
wf_dir = get_workflows_directory()
|
|
if not wf_dir:
|
|
return []
|
|
result: list[dict[str, str]] = []
|
|
try:
|
|
for path in sorted(wf_dir.glob("*.json")):
|
|
if path.is_file():
|
|
stem = path.stem
|
|
label = stem.replace("_", " ").strip() or stem
|
|
result.append({"id": stem, "label": label})
|
|
except OSError as e:
|
|
logger.warning("WorkflowUIPlugin: list_workflows failed: %s", e)
|
|
return result
|
|
|
|
|
|
def _is_api_format(data: dict[str, Any]) -> bool:
|
|
if not data or not isinstance(data, dict):
|
|
return False
|
|
for v in data.values():
|
|
if isinstance(v, dict) and ("class_type" in v or "inputs" in v):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _extract_graph(workflow: dict[str, Any]) -> dict[str, Any]:
|
|
if not isinstance(workflow, dict):
|
|
return workflow
|
|
if workflow.get("nodes") is not None:
|
|
return workflow
|
|
if _is_api_format(workflow):
|
|
return workflow
|
|
for key in ("graph", "workflow", "prompt"):
|
|
inner = workflow.get(key)
|
|
if isinstance(inner, dict):
|
|
if inner.get("nodes") is not None or _is_api_format(inner):
|
|
return inner
|
|
deeper = inner.get("graph") or inner.get("workflow")
|
|
if isinstance(deeper, dict):
|
|
return deeper
|
|
return workflow
|
|
|
|
|
|
def _build_link_map(links: list[Any]) -> dict[int | str, tuple[str, int]]:
|
|
result: dict[int | str, tuple[str, int]] = {}
|
|
if not isinstance(links, list):
|
|
return result
|
|
for link in links:
|
|
if not isinstance(link, (list, tuple)) or len(link) < 6:
|
|
continue
|
|
link_id, origin_id, origin_slot = link[0], link[1], link[2]
|
|
try:
|
|
origin_slot_int = int(origin_slot)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
result[link_id] = (str(origin_id), origin_slot_int)
|
|
return result
|
|
|
|
|
|
def _inputs_from_ui_node(
|
|
node: dict[str, Any],
|
|
class_type: str,
|
|
link_map: dict[int | str, tuple[str, int]],
|
|
) -> dict[str, Any]:
|
|
inputs: dict[str, Any] = {}
|
|
node_inputs = node.get("inputs")
|
|
widgets_values = (
|
|
node.get("widgets_values")
|
|
or node.get("widgetsValues")
|
|
or node.get("widget_values")
|
|
)
|
|
widget_order = _WIDGET_ORDER.get(class_type)
|
|
|
|
if isinstance(widgets_values, dict):
|
|
inputs = dict(widgets_values)
|
|
elif isinstance(widgets_values, list) and widget_order:
|
|
for idx, field in enumerate(widget_order):
|
|
if idx < len(widgets_values):
|
|
inputs[field] = widgets_values[idx]
|
|
|
|
if isinstance(node_inputs, list):
|
|
widget_only_names: list[str] = []
|
|
for inp in node_inputs:
|
|
if not isinstance(inp, dict):
|
|
continue
|
|
name = inp.get("name")
|
|
if not name or not isinstance(name, str):
|
|
continue
|
|
link = inp.get("link")
|
|
if link is not None:
|
|
conn = link_map.get(link)
|
|
if conn is not None:
|
|
inputs[name] = [conn[0], conn[1]]
|
|
else:
|
|
widget_only_names.append(name)
|
|
for idx, name in enumerate(widget_only_names):
|
|
if name not in inputs and isinstance(widgets_values, list) and idx < len(widgets_values):
|
|
inputs[name] = widgets_values[idx]
|
|
elif isinstance(node_inputs, dict) and not inputs:
|
|
inputs = dict(node_inputs)
|
|
return inputs
|
|
|
|
|
|
def _ensure_saveimage_defaults(api_graph: dict[str, Any]) -> None:
|
|
save_types = ("SaveImage", "Save Image", "Save Image (api)", "SaveImageNode")
|
|
for node in api_graph.values():
|
|
if not isinstance(node, dict) or node.get("class_type") not in save_types:
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if isinstance(inputs, dict) and not inputs.get("filename_prefix"):
|
|
inputs["filename_prefix"] = "ComfyUI"
|
|
|
|
_NODE_MODE_DISABLED = 4
|
|
|
|
|
|
def ui_to_api_format(workflow: dict[str, Any]) -> dict[str, Any]:
|
|
workflow = _extract_graph(workflow)
|
|
nodes = workflow.get("nodes")
|
|
if not isinstance(nodes, list):
|
|
if _is_api_format(workflow):
|
|
return workflow
|
|
return {}
|
|
links = workflow.get("links")
|
|
link_map = _build_link_map(links) if isinstance(links, list) else {}
|
|
out: dict[str, Any] = {}
|
|
for node in nodes:
|
|
if not isinstance(node, dict):
|
|
continue
|
|
if node.get("mode") == _NODE_MODE_DISABLED:
|
|
continue
|
|
nid = node.get("id")
|
|
if nid is None:
|
|
continue
|
|
node_id = str(nid)
|
|
class_type = node.get("type") or node.get("class_type")
|
|
if not class_type:
|
|
continue
|
|
inputs = _inputs_from_ui_node(node, class_type, link_map)
|
|
meta = node.get("_meta")
|
|
if not meta and isinstance(node.get("properties"), dict):
|
|
s_r_name = (node.get("properties") or {}).get("Node name for S&R")
|
|
if isinstance(s_r_name, str) and s_r_name.strip():
|
|
meta = {"title": s_r_name.strip()}
|
|
normalized: dict[str, Any] = {"class_type": class_type, "inputs": inputs}
|
|
if meta:
|
|
normalized["_meta"] = meta
|
|
out[node_id] = normalized
|
|
_ensure_saveimage_defaults(out)
|
|
return out
|
|
|
|
|
|
def load_workflow(workflow_id: str) -> tuple[str, dict[str, Any]]:
|
|
wf_dir = get_workflows_directory()
|
|
if not wf_dir:
|
|
raise FileNotFoundError("Workflow directory not configured")
|
|
safe_id = (workflow_id or "").strip()
|
|
if not safe_id or "/" in safe_id or "\\" in safe_id or ".." in safe_id:
|
|
raise ValueError("Invalid workflow id")
|
|
path = (wf_dir / safe_id).with_suffix(".json")
|
|
if not path.is_file():
|
|
raise FileNotFoundError(f"Workflow not found: {safe_id}")
|
|
try:
|
|
raw = path.read_text(encoding="utf-8", errors="replace")
|
|
data = json.loads(raw)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid workflow JSON: {e}") from e
|
|
if not isinstance(data, dict):
|
|
raise ValueError("Workflow must be a JSON object")
|
|
name = data.get("name")
|
|
if not isinstance(name, str) and isinstance(data.get("workflow"), dict):
|
|
name = data["workflow"].get("name")
|
|
if not isinstance(name, str) or not name.strip():
|
|
name = safe_id.replace("_", " ").strip() or safe_id
|
|
|
|
graph = _extract_graph(data)
|
|
if not graph or (isinstance(graph, dict) and not graph.get("nodes") and not _is_api_format(graph)):
|
|
raise ValueError("Workflow has no executable nodes")
|
|
return name, graph
|