Add stack Update, ComfyUI model manager, and slim ComfyUI stack.
Server UI gains Update on stack cards, ComfyUI Models tab with workflow scan and downloads, and centralized comfyui_config. Model catalog and download scripts move from stacks/comfyui to server-ui so ComfyUI stays a minimal Docker wrapper for easier image updates.
This commit is contained in:
tomasz-syn-grzegorza
@@ -0,0 +1,278 @@
|
||||
"""Scan ComfyUI workflow JSON files for required model filenames."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from comfyui_config import COMFYUI_CONTAINER, container_workflow_path
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# class_type -> list of (input_field, category)
|
||||
LOADER_NODES: dict[str, list[tuple[str, str]]] = {
|
||||
"CheckpointLoaderSimple": [("ckpt_name", "checkpoint")],
|
||||
"CheckpointLoader": [("ckpt_name", "checkpoint")],
|
||||
"unCLIPCheckpointLoader": [("ckpt_name", "checkpoint")],
|
||||
"ImageOnlyCheckpointLoader": [("ckpt_name", "checkpoint")],
|
||||
"VAELoader": [("vae_name", "vae")],
|
||||
"LoraLoader": [("lora_name", "lora")],
|
||||
"LoraLoaderModelOnly": [("lora_name", "lora")],
|
||||
"UNETLoader": [("unet_name", "unet")],
|
||||
"CLIPLoader": [("clip_name", "clip")],
|
||||
"DualCLIPLoader": [("clip_name1", "clip"), ("clip_name2", "clip")],
|
||||
"ControlNetLoader": [("control_net_name", "controlnet")],
|
||||
"UpscaleModelLoader": [("model_name", "upscale")],
|
||||
"StyleModelLoader": [("style_model_name", "style_models")],
|
||||
"LatentUpscaleModelLoader": [("model_name", "latent_upscale")],
|
||||
"LTXVAudioVAELoader": [("ckpt_name", "vae")],
|
||||
"LTXAVTextEncoderLoader": [("gemma_model", "text_encoder"), ("ltx_model", "checkpoint")],
|
||||
"DiffusersLoader": [("model_path", "diffusion_models")],
|
||||
}
|
||||
|
||||
# Nodes where every .safetensors in widgets_values should be collected
|
||||
MULTI_WIDGET_LOADERS = frozenset({"LTXAVTextEncoderLoader"})
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelRef:
|
||||
filename: str
|
||||
category: str
|
||||
workflows: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def _is_model_filename(value: Any) -> bool:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return False
|
||||
lower = value.lower()
|
||||
return lower.endswith(
|
||||
(".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".sft")
|
||||
)
|
||||
|
||||
|
||||
def _infer_category_from_node(class_type: str, filename: str) -> str:
|
||||
ct = class_type.lower()
|
||||
fn = filename.lower()
|
||||
if "lora" in ct or "lora" in fn:
|
||||
return "lora"
|
||||
if "upscale" in ct or "upscaler" in fn or "upscale" in fn:
|
||||
return "latent_upscale"
|
||||
if "vae" in ct:
|
||||
return "vae"
|
||||
if "clip" in ct or "text_encoder" in ct or "gemma" in fn:
|
||||
return "text_encoder"
|
||||
if "unet" in ct:
|
||||
return "unet"
|
||||
if "controlnet" in ct:
|
||||
return "controlnet"
|
||||
return "checkpoint"
|
||||
|
||||
|
||||
def _read_workflow_json(path: Path, workflows_dir: Path) -> dict[str, Any] | None:
|
||||
try:
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
except PermissionError:
|
||||
try:
|
||||
rel = path.relative_to(workflows_dir.resolve())
|
||||
except ValueError:
|
||||
log.warning("Workflow outside workflows dir, cannot use docker fallback: %s", path)
|
||||
return None
|
||||
container_path = container_workflow_path(rel.as_posix())
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "exec", COMFYUI_CONTAINER, "cat", container_path],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
check=False,
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||||
log.warning("Docker fallback failed for workflow %s: %s", path, exc)
|
||||
return None
|
||||
if result.returncode != 0:
|
||||
log.warning(
|
||||
"Docker fallback failed for workflow %s: %s",
|
||||
path,
|
||||
result.stderr.strip(),
|
||||
)
|
||||
return None
|
||||
log.debug("Read workflow via docker exec: %s", path.name)
|
||||
try:
|
||||
data = json.loads(result.stdout)
|
||||
except json.JSONDecodeError as exc:
|
||||
log.warning("Invalid JSON from docker fallback for %s: %s", path, exc)
|
||||
return None
|
||||
return data if isinstance(data, dict) else None
|
||||
except (OSError, json.JSONDecodeError) as exc:
|
||||
log.warning("Failed to read workflow %s: %s", path, exc)
|
||||
return None
|
||||
|
||||
|
||||
def _iter_workflow_nodes(data: dict[str, Any]) -> Iterator[dict[str, Any]]:
|
||||
nodes = data.get("nodes")
|
||||
if isinstance(nodes, list):
|
||||
for node in nodes:
|
||||
if isinstance(node, dict):
|
||||
yield node
|
||||
|
||||
definitions = data.get("definitions")
|
||||
if isinstance(definitions, dict):
|
||||
subgraphs = definitions.get("subgraphs")
|
||||
if isinstance(subgraphs, list):
|
||||
for subgraph in subgraphs:
|
||||
if not isinstance(subgraph, dict):
|
||||
continue
|
||||
inner = subgraph.get("nodes")
|
||||
if isinstance(inner, list):
|
||||
for node in inner:
|
||||
if isinstance(node, dict):
|
||||
yield node
|
||||
|
||||
|
||||
def _node_class_type(node: dict[str, Any]) -> str | None:
|
||||
ct = node.get("class_type") or node.get("type")
|
||||
if isinstance(ct, str) and not ct.isupper() and ct not in ("INT", "FLOAT", "STRING", "COMBO", "BOOLEAN", "*"):
|
||||
return ct
|
||||
return None
|
||||
|
||||
|
||||
def _extract_from_widgets(
|
||||
class_type: str,
|
||||
widgets: list[Any],
|
||||
) -> list[tuple[str, str]]:
|
||||
refs: list[tuple[str, str]] = []
|
||||
if class_type in MULTI_WIDGET_LOADERS:
|
||||
field_specs = LOADER_NODES.get(class_type, [])
|
||||
for idx, value in enumerate(widgets):
|
||||
if not _is_model_filename(value):
|
||||
continue
|
||||
if idx < len(field_specs):
|
||||
category = field_specs[idx][1]
|
||||
else:
|
||||
category = _infer_category_from_node(class_type, str(value))
|
||||
refs.append((str(value), category))
|
||||
return refs
|
||||
|
||||
if class_type in LOADER_NODES and widgets:
|
||||
first = widgets[0]
|
||||
if _is_model_filename(first):
|
||||
refs.append((str(first), LOADER_NODES[class_type][0][1]))
|
||||
return refs
|
||||
|
||||
|
||||
def _extract_generic_widgets(class_type: str, widgets: list[Any]) -> list[tuple[str, str]]:
|
||||
refs: list[tuple[str, str]] = []
|
||||
for value in widgets:
|
||||
if _is_model_filename(value):
|
||||
refs.append((str(value), _infer_category_from_node(class_type, str(value))))
|
||||
return refs
|
||||
|
||||
|
||||
def _extract_from_node(node: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
refs: list[tuple[str, str]] = []
|
||||
class_type = _node_class_type(node)
|
||||
if not class_type:
|
||||
return refs
|
||||
|
||||
widgets = node.get("widgets_values")
|
||||
if isinstance(widgets, list):
|
||||
if class_type in LOADER_NODES:
|
||||
refs.extend(_extract_from_widgets(class_type, widgets))
|
||||
else:
|
||||
refs.extend(_extract_generic_widgets(class_type, widgets))
|
||||
|
||||
inputs = node.get("inputs")
|
||||
if isinstance(inputs, dict) and class_type in LOADER_NODES:
|
||||
for field_name, category in LOADER_NODES[class_type]:
|
||||
value = inputs.get(field_name)
|
||||
if _is_model_filename(value):
|
||||
refs.append((str(value), category))
|
||||
|
||||
return refs
|
||||
|
||||
|
||||
def _extract_from_api_format(data: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
refs: list[tuple[str, str]] = []
|
||||
for _node_id, node in data.items():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
class_type = node.get("class_type")
|
||||
if not isinstance(class_type, str):
|
||||
continue
|
||||
if class_type not in LOADER_NODES:
|
||||
continue
|
||||
inputs = node.get("inputs")
|
||||
if not isinstance(inputs, dict):
|
||||
continue
|
||||
for field_name, category in LOADER_NODES[class_type]:
|
||||
value = inputs.get(field_name)
|
||||
if _is_model_filename(value):
|
||||
refs.append((str(value), category))
|
||||
return refs
|
||||
|
||||
|
||||
def _extract_from_ui_format(data: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
refs: list[tuple[str, str]] = []
|
||||
seen: set[tuple[str, str]] = set()
|
||||
for node in _iter_workflow_nodes(data):
|
||||
for ref in _extract_from_node(node):
|
||||
if ref not in seen:
|
||||
seen.add(ref)
|
||||
refs.append(ref)
|
||||
return refs
|
||||
|
||||
|
||||
def parse_workflow_data(data: dict[str, Any]) -> list[tuple[str, str]]:
|
||||
if "nodes" in data:
|
||||
return _extract_from_ui_format(data)
|
||||
return _extract_from_api_format(data)
|
||||
|
||||
|
||||
def parse_workflow_file(path: Path, workflows_dir: Path | None = None) -> list[tuple[str, str]]:
|
||||
"""Return list of (filename, category) referenced in a workflow file."""
|
||||
wf_dir = workflows_dir or path.parent
|
||||
data = _read_workflow_json(path, wf_dir)
|
||||
if not data:
|
||||
return []
|
||||
return parse_workflow_data(data)
|
||||
|
||||
|
||||
def find_model_file(models_root: Path, filename: str) -> Path | None:
|
||||
"""Find a model file by basename anywhere under models_root."""
|
||||
if not models_root.is_dir():
|
||||
return None
|
||||
matches: list[Path] = []
|
||||
for path in models_root.rglob(filename):
|
||||
if path.is_file() and path.name == filename:
|
||||
matches.append(path)
|
||||
if not matches:
|
||||
return None
|
||||
return min(matches, key=lambda p: len(p.parts))
|
||||
|
||||
|
||||
def scan_workflows(workflows_dir: Path) -> tuple[list[ModelRef], int]:
|
||||
"""Scan all JSON workflows and aggregate model references."""
|
||||
if not workflows_dir.is_dir():
|
||||
return [], 0
|
||||
|
||||
aggregated: dict[tuple[str, str], ModelRef] = {}
|
||||
file_count = 0
|
||||
resolved_dir = workflows_dir.resolve()
|
||||
for path in sorted(resolved_dir.rglob("*.json")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
file_count += 1
|
||||
wf_name = path.stem
|
||||
for filename, category in parse_workflow_file(path, resolved_dir):
|
||||
key = (filename, category)
|
||||
if key not in aggregated:
|
||||
aggregated[key] = ModelRef(filename=filename, category=category)
|
||||
if wf_name not in aggregated[key].workflows:
|
||||
aggregated[key].workflows.append(wf_name)
|
||||
|
||||
return sorted(aggregated.values(), key=lambda r: (r.category, r.filename)), file_count
|
||||
Reference in New Issue
Block a user