tomasz-syn-grzegorza
73e4fc005e
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.
279 lines
9.7 KiB
Python
279 lines
9.7 KiB
Python
"""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
|