Доработка управления моделями

app/config.py

@@ -1,30 +1,20 @@
-"""Environment settings and the parameter -> workflow node mapping."""
+"""Environment settings and the model registry loaded from a JSON file."""
 import json
 from pathlib import Path
 
+from pydantic import BaseModel, ValidationError
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
+from .errors import ConfigError
+
 REPO_ROOT = Path(__file__).resolve().parent.parent
 
-CREATE_MODEL = "z-image-turbo"
-EDIT_MODEL = "qwen-image-edit"
-AVAILABLE_MODELS = [CREATE_MODEL, EDIT_MODEL]
 
-# Node ids refer to workflow_create.json / workflow_edit.json.
-# Override the whole map via MAPPING_FILE (JSON with "create" and "edit" keys).
-DEFAULT_MAPPING: dict[str, dict[str, dict[str, str]]] = {
-    "create": {
-        "prompt": {"node": "83:27", "field": "text"},
-        "size": {"node": "83:13", "width_field": "width", "height_field": "height"},
-        "seed": {"node": "83:3", "field": "seed"},
-    },
-    "edit": {
-        "prompt": {"node": "115:111", "field": "prompt"},
-        "negative": {"node": "115:110", "field": "prompt"},
-        "seed": {"node": "115:3", "field": "seed"},
-        "image": {"node": "78", "field": "image"},
-    },
-}
+class ModelEntry(BaseModel):
+    """A registry record: an API model id maps to one workflow file."""
+
+    workflow: Path
+    owned_by: str = "comfy-adapter"
 
 
 class Settings(BaseSettings):
@@ -34,21 +24,31 @@ class Settings(BaseSettings):
     adapter_api_key: str = ""
     execution_timeout: float = 300.0
     poll_interval: float = 1.5
-    create_template_path: Path = REPO_ROOT / "workflow_create.json"
-    edit_template_path: Path = REPO_ROOT / "workflow_edit.json"
-    mapping_file: Path | None = None
-    # Validate graphs on the ComfyUI server (POST /prompt?front=false) at startup.
+    # JSON file mapping an API model id to its workflow. Required; no defaults in code.
+    models_file: Path = REPO_ROOT / "models.json"
+    # Validate collected graphs on the ComfyUI server (POST /prompt?front=false) at startup.
     validate_templates_on_server: bool = True
 
-    def mapping(self) -> dict[str, dict[str, dict[str, str]]]:
-        if self.mapping_file is None:
-            return DEFAULT_MAPPING
-        with open(self.mapping_file, encoding="utf-8") as f:
-            loaded = json.load(f)
-        merged = {k: dict(v) for k, v in DEFAULT_MAPPING.items()}
-        for section, section_map in loaded.items():
-            merged.setdefault(section, {}).update(section_map)
-        return merged
+    def load_registry(self) -> dict[str, ModelEntry]:
+        """Load the model registry, preserving declaration order (first = preferred).
+
+        Fail-fast: a missing or malformed file raises ConfigError and blocks startup.
+        """
+        try:
+            raw = json.loads(self.models_file.read_text(encoding="utf-8"))
+        except FileNotFoundError as e:
+            raise ConfigError(f"model registry file not found: {self.models_file}") from e
+        except json.JSONDecodeError as e:
+            raise ConfigError(f"model registry is not valid JSON ({self.models_file}): {e}") from e
+        if not isinstance(raw, dict) or not raw:
+            raise ConfigError(f"model registry must be a non-empty object of model entries: {self.models_file}")
+        registry: dict[str, ModelEntry] = {}
+        for model_id, spec in raw.items():
+            try:
+                registry[model_id] = ModelEntry.model_validate(spec)
+            except ValidationError as e:
+                raise ConfigError(f"model '{model_id}' has an invalid entry: {e}") from e
+        return registry
 
 
 def get_settings() -> Settings:

app/errors.py

@@ -31,6 +31,13 @@ class ComfyUIError(AdapterError):
         super().__init__(502, message, "server_error", code)
 
 
+class ConfigError(AdapterError):
+    """Invalid model registry configuration; blocks startup."""
+
+    def __init__(self, message: str):
+        super().__init__(500, message, "server_error", "config_error")
+
+
 def error_body(message: str, err_type: str, code: str | None = None) -> dict:
     return {"error": {"message": message, "type": err_type, "param": None, "code": code}}
 

app/injector.py

@@ -1,5 +1,12 @@
-"""Build a concrete workflow graph from a template + request parameters."""
+"""Build a concrete workflow graph from a template + request values.
+
+Values are injected through placeholder tokens written directly inside the
+workflow graph (e.g. "{prompt}", "{seed}"), so the graph itself declares which
+node input receives which request field. The set of recognised tokens is fixed
+and owned by the adapter, not by the model.
+"""
 import copy
+import re
 import secrets
 from typing import Any
 
@@ -7,6 +14,21 @@ from .errors import InvalidRequest
 
 SEED_MAX = 2**63 - 1
 
+# Fixed token vocabulary owned by the adapter. A model's workflow may only use
+# these tokens; anything else is rejected. Types drive typed vs string injection.
+TOKEN_TYPES: dict[str, str] = {
+    "prompt": "str",
+    "negative": "str",
+    "seed": "int",
+    "width": "int",
+    "height": "int",
+    "image": "str",
+}
+KNOWN_TOKENS = frozenset(TOKEN_TYPES)
+
+_TOKEN = re.compile(r"\{(\w+)\}")
+_WHOLE = re.compile(r"^\{(\w+)\}$")
+
 
 def parse_size(size: str) -> tuple[int, int]:
     """Parse OpenAI `WxH` size strings (e.g. '1024x1024')."""
@@ -28,40 +50,55 @@ def new_seed(seed: int | None) -> int:
     return seed
 
 
-def build_workflow(
-    template: dict[str, Any],
-    mapping: dict[str, dict[str, str]],
-    *,
-    prompt: str,
-    seed: int,
-    size: str | None = None,
-    image_name: str | None = None,
-    negative: str = "",
-) -> dict[str, Any]:
-    """Deep-copy the template and inject request values per the node mapping."""
+def extract_tokens(graph: dict[str, Any]) -> tuple[set[str], list[str]]:
+    """Token names referenced by the graph.
+
+    Returns (known_tokens, unknown_token_errors). A token is any `{name}` inside
+    a string node input. Unknown names are reported so startup can reject the
+    workflow before it ever reaches a request.
+    """
+    declared: set[str] = set()
+    unknown: set[str] = set()
+    for node_id, node in graph.items():
+        for field, value in node.get("inputs", {}).items():
+            if isinstance(value, str):
+                for name in _TOKEN.findall(value):
+                    if name in KNOWN_TOKENS:
+                        declared.add(name)
+                    else:
+                        unknown.add(f"node {node_id}: input '{field}' uses unknown token '{{{name}}}'")
+    return declared, sorted(unknown)
+
+
+def build_workflow(template: dict[str, Any], values: dict[str, Any]) -> dict[str, Any]:
+    """Deep-copy the template and substitute placeholder tokens with values.
+
+    Injection is type-preserving: a value that is exactly one token (e.g.
+    "{seed}") is replaced by the typed object (int for seed/width/height, str
+    otherwise); a token embedded in surrounding text is string-interpolated.
+    A referenced token that is unrecognized or missing from ``values`` rejects
+    the build, so a malformed graph never reaches ComfyUI.
+    """
     graph = copy.deepcopy(template)
-
-    def put(key: str, value: Any) -> None:
-        spec = mapping.get(key)
-        if spec is None:
-            return
-        node = graph.get(spec["node"])
-        if node is None:
-            raise InvalidRequest(f"mapping key '{key}' points to missing node '{spec['node']}'")
-        node["inputs"][spec["field"]] = value
-
-    put("prompt", prompt)
-    put("negative", negative)
-    put("seed", seed)
-    if image_name is not None:
-        put("image", image_name)
-    if size is not None:
-        spec = mapping.get("size")
-        if spec is not None:
-            width, height = parse_size(size)
-            node = graph.get(spec["node"])
-            if node is None:
-                raise InvalidRequest(f"mapping key 'size' points to missing node '{spec['node']}'")
-            node["inputs"][spec["width_field"]] = width
-            node["inputs"][spec["height_field"]] = height
+    for node in graph.values():
+        inputs = node.get("inputs", {})
+        for field, value in inputs.items():
+            if isinstance(value, str):
+                inputs[field] = _substitute(value, values)
     return graph
+
+
+def _substitute(value: str, values: dict[str, Any]) -> Any:
+    for name in _TOKEN.findall(value):
+        if name not in KNOWN_TOKENS:
+            raise InvalidRequest(f"workflow uses unrecognized token '{{{name}}}'")
+        if name not in values:
+            raise InvalidRequest(f"no value supplied for token '{{{name}}}'")
+    whole = _WHOLE.fullmatch(value)
+    if whole:
+        return _coerce(whole.group(1), values[whole.group(1)])
+    return _TOKEN.sub(lambda m: str(values[m.group(1)]), value)
+
+
+def _coerce(name: str, value: Any) -> Any:
+    return int(value) if TOKEN_TYPES[name] == "int" else str(value)

app/main.py

@@ -8,9 +8,9 @@ from fastapi import FastAPI, Request
 
 from .auth import require_key
 from .comfy_client import ComfyClient, extract_output_image
-from .config import AVAILABLE_MODELS, CREATE_MODEL, EDIT_MODEL, Settings, get_settings
+from .config import Settings, get_settings
 from .errors import InvalidRequest, install_error_handlers
-from .injector import build_workflow, new_seed
+from .injector import build_workflow, new_seed, parse_size
 from .schemas import (
     ImageEditJsonRequest,
     ImageGenerationRequest,
@@ -19,32 +19,42 @@ from .schemas import (
     ModelCard,
     ModelsResponse,
 )
-from .templates import load_all
+from .templates import load_registry, select_model
 
 logger = logging.getLogger("comfy-adapter")
 
 PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
 JPEG_MAGIC = b"\xff\xd8\xff"
 
+# Type-correct dummy values used to dry-inject every workflow for server validation.
+_DUMMY_VALUES = {
+    "prompt": "x",
+    "negative": "",
+    "seed": 1,
+    "width": 1024,
+    "height": 1024,
+    "image": "dummy.png",
+}
+
 
 def create_app(settings: Settings | None = None) -> FastAPI:
     settings = settings or get_settings()
 
     @asynccontextmanager
     async def lifespan(app: FastAPI):
-        create_graph, edit_graph = load_all(settings)
+        registry = load_registry(settings)
         client = ComfyClient(settings.comfy_base_url)
         if settings.validate_templates_on_server:
-            try:
-                await client.validate_graph(create_graph)
-                await client.validate_graph(edit_graph)
-                logger.info("workflow templates validated on ComfyUI server")
-            except Exception as e:  # noqa: BLE001 - startup check is advisory, logged
-                logger.error("server-side template validation failed: %s", e)
+            for model in registry.values():
+                try:
+                    await client.validate_graph(build_workflow(model.graph, _DUMMY_VALUES))
+                    logger.info("workflow validated on ComfyUI server: %s", model.id)
+                except Exception as e:  # noqa: BLE001 - startup check is advisory, logged
+                    logger.error("server-side validation failed for model '%s': %s", model.id, e)
         app.state.settings = settings
         app.state.comfy = client
-        app.state.templates = {"create": create_graph, "edit": edit_graph}
-        logger.info("comfy-adapter ready (server=%s)", settings.comfy_base_url)
+        app.state.registry = registry
+        logger.info("comfy-adapter ready (server=%s, models=%s)", settings.comfy_base_url, list(registry))
         try:
             yield
         finally:
@@ -75,7 +85,10 @@ def create_app(settings: Settings | None = None) -> FastAPI:
     @app.get("/v1/models")
     async def models(request: Request) -> ModelsResponse:
         require_key(request, _s(request))
-        return ModelsResponse(data=[ModelCard(id=m) for m in AVAILABLE_MODELS])
+        registry = request.app.state.registry
+        return ModelsResponse(
+            data=[ModelCard(id=model_id, owned_by=model.owned_by) for model_id, model in registry.items()]
+        )
 
     @app.post("/v1/images/generations")
     async def generations(request: Request) -> ImagesResponse:
@@ -86,14 +99,11 @@ def create_app(settings: Settings | None = None) -> FastAPI:
             raise InvalidRequest("prompt must not be empty")
         if body.n not in (None, 1):
             raise InvalidRequest("only n=1 is supported")
-        if body.model is not None and body.model != CREATE_MODEL:
-            raise InvalidRequest(f"unknown model '{body.model}', expected '{CREATE_MODEL}'")
+        model = select_model(request.app.state.registry, body.model, "generations")
+        width, height = parse_size(body.size)
         graph = build_workflow(
-            app.state.templates["create"],
-            settings.mapping()["create"],
-            prompt=body.prompt,
-            seed=new_seed(body.seed),
-            size=body.size,
+            model.graph,
+            {"prompt": body.prompt, "seed": new_seed(body.seed), "negative": "", "width": width, "height": height},
         )
         return await _run(graph)
 
@@ -112,18 +122,13 @@ def create_app(settings: Settings | None = None) -> FastAPI:
             raise InvalidRequest("expected multipart/form-data or application/json body")
         if not prompt.strip():
             raise InvalidRequest("prompt must not be empty")
-        if model is not None and model != EDIT_MODEL:
-            raise InvalidRequest(f"unknown model '{model}' for edits, expected '{EDIT_MODEL}'")
         _validate_image_bytes(image_bytes)
+        model = select_model(request.app.state.registry, model, "edits")
         comfy: ComfyClient = app.state.comfy
         name = await comfy.upload_image(image_bytes)
         graph = build_workflow(
-            app.state.templates["edit"],
-            settings.mapping()["edit"],
-            prompt=prompt,
-            seed=new_seed(seed),
-            image_name=name,
-            negative="",
+            model.graph,
+            {"prompt": prompt, "seed": new_seed(seed), "negative": "", "image": name},
         )
         return await _run(graph)
 

app/templates.py

@@ -1,9 +1,17 @@
-"""Template loading with startup validation of graph links and mappings."""
+"""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
+from .errors import AdapterError, ConfigError, InvalidRequest
+from .injector import extract_tokens
 
 
 class TemplateError(AdapterError):
@@ -13,6 +21,29 @@ class TemplateError(AdapterError):
         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 = []
@@ -24,39 +55,58 @@ def find_broken_links(graph: dict) -> list[str]:
     return errors
 
 
-def check_mapping(graph: dict, mapping: dict, name: str) -> list[str]:
-    """Every mapped node id must exist with a scalar input field."""
-    errors = []
-    for key, spec in mapping.items():
-        node = graph.get(spec["node"])
-        if node is None:
-            errors.append(f"mapping '{name}.{key}': node '{spec['node']}' not found in template")
-            continue
-        for field in _mapped_fields(spec):
-            if field not in node.get("inputs", {}) or isinstance(node["inputs"][field], list):
-                errors.append(f"mapping '{name}.{key}': node '{spec['node']}' has no scalar input '{field}'")
-    return errors
+def is_compatible(declared: frozenset[str], endpoint: str) -> bool:
+    vocab = ENDPOINT_VOCAB[endpoint]
+    return vocab["required"] <= declared and declared <= vocab["supplied"]
 
 
-def _mapped_fields(spec: dict) -> list[str]:
-    if "width_field" in spec:
-        return [spec["width_field"], spec["height_field"]]
-    return [spec["field"]]
+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 load_template(path: Path, graph_name: str, mapping: dict, mapping_name: str) -> dict:
-    with open(path, encoding="utf-8") as f:
-        graph = json.load(f)
-    errors = find_broken_links(graph) + check_mapping(graph, mapping[mapping_name], mapping_name)
+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 template '{graph_name}' ({path}) is invalid:\n  " + "\n  ".join(errors))
-    return graph
+        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_all(settings: Settings) -> tuple[dict, dict]:
-    mapping = settings.mapping()
-    create = load_template(
-        settings.create_template_path, "workflow_create.json", mapping, "create"
-    )
-    edit = load_template(settings.edit_template_path, "workflow_edit.json", mapping, "edit")
-    return create, edit
+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