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

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)