Shaman
113 lines
4.2 KiB
Python
113 lines
4.2 KiB
Python
"""Model registry loading: read each workflow, validate links/tokens, and
|
|
decide which model an endpoint may run.
|
|
|
|
Startup validation of a workflow is local here (broken links + token check);
|
|
the collected graphs are then dry-injected and sent to the server elsewhere
|
|
(see main.lifespan).
|
|
"""
|
|
import json
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
from .config import Settings
|
|
from .errors import AdapterError, ConfigError, InvalidRequest
|
|
from .injector import extract_tokens
|
|
|
|
|
|
class TemplateError(AdapterError):
|
|
"""Malformed workflow template; blocks startup."""
|
|
|
|
def __init__(self, message: str):
|
|
super().__init__(500, message, "server_error", "template_error")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Model:
|
|
id: str
|
|
graph: dict
|
|
declared: frozenset[str]
|
|
owned_by: str = "comfy-adapter"
|
|
workflow: Path = field(default=Path("."))
|
|
|
|
|
|
# Token sets belong to the endpoint, not the model. A model is valid for an
|
|
# endpoint iff required(endpoint) <= declared(graph) <= supplied(endpoint).
|
|
ENDPOINT_VOCAB: dict[str, dict[str, frozenset[str]]] = {
|
|
"generations": {
|
|
"supplied": frozenset({"prompt", "negative", "seed", "width", "height"}),
|
|
"required": frozenset({"prompt"}),
|
|
},
|
|
"edits": {
|
|
"supplied": frozenset({"prompt", "negative", "seed", "image"}),
|
|
"required": frozenset({"prompt", "image"}),
|
|
},
|
|
}
|
|
|
|
|
|
def find_broken_links(graph: dict) -> list[str]:
|
|
"""References to node ids that do not exist in the graph."""
|
|
errors = []
|
|
for node_id, node in graph.items():
|
|
for field, value in node.get("inputs", {}).items():
|
|
if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str):
|
|
if value[0] not in graph:
|
|
errors.append(f"node {node_id}: input '{field}' links to missing node '{value[0]}'")
|
|
return errors
|
|
|
|
|
|
def is_compatible(declared: frozenset[str], endpoint: str) -> bool:
|
|
vocab = ENDPOINT_VOCAB[endpoint]
|
|
return vocab["required"] <= declared and declared <= vocab["supplied"]
|
|
|
|
|
|
def select_model(models: dict[str, Model], model: str | None, endpoint: str) -> Model:
|
|
"""Pick the model for a request by id, or the first compatible one.
|
|
|
|
An explicit id that is unknown or incompatible with the endpoint is a 400,
|
|
as is a missing id when no registered model is compatible.
|
|
"""
|
|
if model is not None:
|
|
found = models.get(model)
|
|
if found is None:
|
|
raise InvalidRequest(f"unknown model '{model}'")
|
|
if not is_compatible(found.declared, endpoint):
|
|
raise InvalidRequest(f"model '{model}' is not compatible with this endpoint")
|
|
return found
|
|
for candidate in models.values(): # declaration order = priority
|
|
if is_compatible(candidate.declared, endpoint):
|
|
return candidate
|
|
raise InvalidRequest(f"no model in the registry is compatible with this endpoint")
|
|
|
|
|
|
def _read_graph(path: Path, model_id: str) -> dict:
|
|
try:
|
|
with open(path, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except FileNotFoundError as e:
|
|
raise TemplateError(f"model '{model_id}': workflow file not found: {path}") from e
|
|
except json.JSONDecodeError as e:
|
|
raise TemplateError(f"model '{model_id}': workflow is not valid JSON ({path}): {e}") from e
|
|
|
|
|
|
def load_model(model_id: str, workflow_path: Path, owned_by: str) -> Model:
|
|
graph = _read_graph(workflow_path, model_id)
|
|
errors = find_broken_links(graph)
|
|
declared, unknown = extract_tokens(graph)
|
|
errors += unknown
|
|
if errors:
|
|
raise TemplateError(
|
|
f"workflow for model '{model_id}' ({workflow_path}) is invalid:\n " + "\n ".join(errors)
|
|
)
|
|
return Model(id=model_id, graph=graph, declared=frozenset(declared), owned_by=owned_by, workflow=workflow_path)
|
|
|
|
|
|
def load_registry(settings: Settings) -> dict[str, Model]:
|
|
"""Load every model in the registry, preserving order for default selection."""
|
|
base_dir = settings.models_file.parent
|
|
registry = settings.load_registry()
|
|
models: dict[str, Model] = {}
|
|
for model_id, entry in registry.items():
|
|
path = entry.workflow if entry.workflow.is_absolute() else base_dir / entry.workflow
|
|
models[model_id] = load_model(model_id, path, entry.owned_by)
|
|
return models
|