From ff4f1073ba6a1b871fb2027058b2c65ecfe35741 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 2 Aug 2026 20:31:20 +0200 Subject: [PATCH] docs(architectures): document the registry and add a scaffolder for new architectures The integration guide described the world before the registry, and described it incompletely. Its checklist never mentioned step_callback.py, dependencies.py's safe_globals list, or the AnyVariant widening -- which is exactly why forgetting them was easy. Those three are now either derived from the registry or checked in CI, so the guide can stop being a memory test. What the guide gains is a section on the one file that now holds most of it, architectures/defs/.py, with a table of what moved and, more importantly, of when each mistake used to surface. That is the point of the change: none of those omissions failed at boot. They failed on the first generation that happened to use the new architecture. The starter-models section was stale in a way that would have sent someone to a file that no longer exists. It now describes the package, says where shared components belong, and warns that the order of STARTER_MODELS is what users see and is not sorted -- the constraint that shaped that split. The file tree gained the files a new architecture actually touches. It was missing fields.py, primitives.py, conditioning_data.py and invocation_api, all of which the last real architecture addition (ERNIE-Image, #9115) edited. That list was derived from that commit rather than from memory, then filtered against what PRs 0-6 absorbed. scripts/new_architecture.py generates the two files that can be generated -- the definition module and the starter-models module -- inserts the import line that makes the first one load, and prints the residual edits with the reason each one cannot be generated. It defaults to a dry run. The templates are tested against the fifteen architectures that already exist, not against themselves. A generator that drifts from what the registry expects is worse than none, because it produces files that look right and fail at boot: so the import line it would emit must be the line already present in architectures/__init__.py, and the module path it computes must match registry.defs_module_path(). One test asserts the template scaffolds every facet that declares itself REQUIRED, so adding a required facet without updating the scaffolder fails rather than producing definition modules that cannot boot. Starter models turned out to be optional -- SD2 ships none -- so that test skips where there is no module, with a second test pinning SD2 as the only such case so the skip cannot quietly spread. The docs build was not run: docs/node_modules is absent and installing it needs network. The MDX was checked by hand instead -- the components used are the two already imported by the file, and every aside type used is already used elsewhere in these docs. openapi.json is unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- .../contributing/new-model-integration.mdx | 103 +++++++- scripts/new_architecture.py | 228 ++++++++++++++++++ tests/test_new_architecture_scaffold.py | 128 ++++++++++ 3 files changed, 454 insertions(+), 5 deletions(-) create mode 100644 scripts/new_architecture.py create mode 100644 tests/test_new_architecture_scaffold.py diff --git a/docs/src/content/docs/contributing/new-model-integration.mdx b/docs/src/content/docs/contributing/new-model-integration.mdx index c20c933bc83..bc0744e4ab9 100644 --- a/docs/src/content/docs/contributing/new-model-integration.mdx +++ b/docs/src/content/docs/contributing/new-model-integration.mdx @@ -12,6 +12,59 @@ This guide describes all the steps required to integrate a new model type into I The code examples use a hypothetical `NewModel` architecture. The implementations of FLUX.1, FLUX.2 Klein, SD3, SDXL, and Z-Image in the InvokeAI codebase serve as excellent real-world references. ::: +:::tip[Start here] +Run the scaffolder first — it creates the files below that can be generated and prints the ones that +cannot: + +```sh +python scripts/new_architecture.py --enum-name NewModel --enum-value new-model --write +``` +::: + +--- + +## The architecture registry + +Most per-architecture facts live in **one file**, `invokeai/backend/architectures/defs/[newmodel].py`, +rather than scattered across core modules: + +```python title="invokeai/backend/architectures/defs/newmodel.py" +register( + BaseModelType.NewModel, + LatentSpaceFacet(FLUX_16), # required: preview projection + spatial scale + ConditioningFacet(NewModelConditioningInfo), # required: what the text encoder produces + VariantFacet({ModelType.Main: NewModelVariantType}), # optional + LoaderFlagsFacet(supports_fp8_storage=False), # optional +) +``` + +Add one line to `invokeai/backend/architectures/__init__.py` to import it, and the registration runs. + +This matters because of *when* mistakes surface. Before the registry, the facts below were spread +over core files and none of them failed at boot — a missing entry raised `Unsupported base model` on +the first preview, or broke conditioning deserialization mid-graph, on the first generation that +happened to use the new architecture. Now `validate()` refuses to start the app while a required +facet is missing, and names the file to edit. + +| Fact | Was | Now | +|------|-----|-----| +| Preview factors, spatial scale | an `elif` chain in `step_callback.py` | `LatentSpaceFacet` | +| `safe_globals` for conditioning | a hand-kept list in `dependencies.py` | derived from `ConditioningFacet` | +| Variant enum | four hand-widened unions | `VariantFacet`, with all four CI-checked | +| fp8 storage exceptions | `config.base ==` inside the generic loader | `LoaderFlagsFacet` | + +Reusing an existing `LatentSpace` is the normal case, not a shortcut: eleven of the fifteen shipped +architectures introduce no new preview data at all. + +:::tip[Checklist: Architecture Registry]{icon="approve-check"} +- [ ] Create `invokeai/backend/architectures/defs/[newmodel].py` +- [ ] Declare `LatentSpaceFacet` — reuse an existing `LatentSpace` unless the VAE genuinely differs +- [ ] Declare `ConditioningFacet` +- [ ] Declare `VariantFacet` / `LoaderFlagsFacet` / `UNetDownscaleFacet` only if they apply +- [ ] Add the import line to `invokeai/backend/architectures/__init__.py` +- [ ] Run `pytest tests/backend/architectures` — it names anything still missing +::: + --- ## 1. Backend: Model Manager @@ -990,9 +1043,11 @@ const recallNewmodelEncoderModel = async (metadata: CoreMetadata) => { ## 10. Starter Models -To allow users to easily download your model from the Model Manager UI, add it to the starter models list. +To allow users to easily download your model from the Model Manager UI, add it to the starter models +catalog. Entries live one module per architecture; only the display order and the bundles are +central. -```python title="invokeai/backend/model_manager/starter_models.py" +```python title="invokeai/backend/model_manager/starter_models/newmodel.py" # Main Model newmodel_main = StarterModel( name="NewModel Main", @@ -1030,7 +1085,21 @@ newmodel_fp8 = StarterModel( dependencies=[newmodel_vae, newmodel_encoder], # Dependencies! ) -# Add to STARTER_MODELS list: +Components shared with other architectures — text encoders, and VAEs more than one architecture +installs — belong in `starter_models/common.py`, not in your module. Import them from there. +`dependencies` holds object references rather than ids, so a name that does not exist is an +`ImportError` at startup instead of a broken install later. + +Then wire it into the catalog: + +```python title="invokeai/backend/model_manager/starter_models/__init__.py" +from invokeai.backend.model_manager.starter_models.newmodel import ( + newmodel_encoder, + newmodel_fp8, + newmodel_main, + newmodel_vae, +) + STARTER_MODELS: list[StarterModel] = [ # ... existing models newmodel_main, @@ -1040,13 +1109,20 @@ STARTER_MODELS: list[StarterModel] = [ ] ``` +:::caution[The order of `STARTER_MODELS` is what users see] +The frontend does not sort this list. Its order is curated — it follows neither base nor model type +— so put your entries where they should appear, not simply at the end. +::: + :::tip[Checklist: Starter Models]{icon="approve-check"} +- [ ] Create `starter_models/[newmodel].py` - [ ] Define main model StarterModel - [ ] Define VAE StarterModel if separate -- [ ] Define text encoder StarterModel if separate +- [ ] Define text encoder StarterModel if separate — in `common.py` if other architectures use it too - [ ] Define quantized variants (FP8, GGUF, etc.) - [ ] Set dependencies correctly -- [ ] Add to `STARTER_MODELS` list +- [ ] Import the module in `starter_models/__init__.py` and place the entries in `STARTER_MODELS` +- [ ] Add a `STARTER_BUNDLES` entry if the architecture ships a launchpad bundle ::: --- @@ -1185,11 +1261,16 @@ For a **minimal txt2img integration**, the following files are required: - invokeai - app/invocations - metadata.py + - fields.py + - primitives.py - `[newmodel]_model_loader.py` - `[newmodel]_text_encoder.py` - `[newmodel]_denoise.py` - `[newmodel]_vae_decode.py` - backend + - architectures + - `defs/[newmodel].py` + - `__init__.py` - model_manager - taxonomy.py - configs @@ -1197,9 +1278,15 @@ For a **minimal txt2img integration**, the following files are required: - factory.py - load/model_loaders - `[newmodel].py` + - starter_models + - `[newmodel].py` + - `__init__.py` + - stable_diffusion/diffusion + - conditioning_data.py - `[newmodel]` - sampling_utils.py - denoise.py + - `invocation_api/__init__.py` - frontend/web/src/features - nodes/util/graph - generation/buildNewModelGraph.ts @@ -1208,6 +1295,12 @@ For a **minimal txt2img integration**, the following files are required: - controlLayers/store/paramsSlice.ts +:::note[What you no longer touch] +`app/util/step_callback.py` and `app/api/dependencies.py` used to need an edit per architecture and +are now derived from the registry. `starter_models.py` is a package rather than one 2400-line file. +The frontend column is unchanged — `webv2` has no OpenAPI tooling, so nothing there is generated. +::: + For **img2img / inpaint / outpaint**, additionally: diff --git a/scripts/new_architecture.py b/scripts/new_architecture.py new file mode 100644 index 00000000000..efaa3989929 --- /dev/null +++ b/scripts/new_architecture.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python +"""Scaffold the per-architecture files a new `BaseModelType` needs, and list what is left by hand. + +Adding an architecture used to mean editing ~16 core files, most of which failed at generation time +rather than at boot when forgotten. The architecture registry moved those facts into one file per +architecture; this generates that file, its starter-model module, and the two import lines that make +them load -- then prints the edits that genuinely cannot be generated. + +Usage: + + python scripts/new_architecture.py --enum-name NewModel --enum-value new-model + python scripts/new_architecture.py --enum-name NewModel --enum-value new-model --write + +Without ``--write`` it only reports what it would do. + +The rendering functions are the testable part; `tests/test_new_architecture_scaffold.py` checks them +against the fifteen architectures that already exist, so this cannot drift from what the registry +actually expects. +""" + +from __future__ import annotations + +import argparse +import re +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +ARCHITECTURES_INIT = REPO_ROOT / "invokeai" / "backend" / "architectures" / "__init__.py" +STARTER_MODELS_INIT = REPO_ROOT / "invokeai" / "backend" / "model_manager" / "starter_models" / "__init__.py" + + +def module_name(enum_value: str) -> str: + """The module stem for an architecture: its enum *value*, dashes to underscores. + + Derived rather than chosen, so `registry.defs_module_path()` and this script cannot disagree + about where a definition lives. + """ + return enum_value.replace("-", "_") + + +def render_defs_module(enum_name: str, enum_value: str) -> str: + """`invokeai/backend/architectures/defs/.py`. + + Both required facets are present but obviously unfinished: `validate()` refuses to start the app + while either is missing, so leaving them out would only move the error later. + """ + return f"""from invokeai.backend.architectures.facets.conditioning import ConditioningFacet +from invokeai.backend.architectures.facets.latent_space import LatentSpaceFacet +from invokeai.backend.architectures.registry import register +from invokeai.backend.model_manager.taxonomy import BaseModelType + +# TODO({enum_value}): replace the two placeholders below. +# +# LatentSpaceFacet - reuse an existing LatentSpace from facets/latent_space.py if this +# architecture shares a VAE with one that already ships (most do: eleven of +# the fifteen introduce no new preview data at all). Only add a new one if +# the channel count or the projection factors genuinely differ. +# ConditioningFacet - the *ConditioningInfo dataclass this architecture's text encoder puts in a +# ConditioningFieldData. Add it to conditioning_data.py and to the +# `conditionings` union there; the safe_globals allowlist is derived from +# this registration, so nothing else needs updating. +# +# Optional facets, only if they apply: +# VariantFacet({{ModelType.Main: {enum_name}VariantType}}) - if the architecture has variants +# LoaderFlagsFacet(supports_fp8_storage=False) - only to opt out of a general loader rule +# UNetDownscaleFacet(max_unet_downscale=8) - UNet architectures with T2I-Adapter support + +register( + BaseModelType.{enum_name}, + LatentSpaceFacet(...), + ConditioningFacet(...), +) +""" + + +def render_starter_models_module(enum_name: str, enum_value: str) -> str: + """`invokeai/backend/model_manager/starter_models/.py`.""" + return f'''"""Starter models for the {enum_value} architecture.""" + +from invokeai.backend.model_manager.starter_models.types import StarterModel +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType + +# Shared components -- encoders, and VAEs used by more than one architecture -- belong in +# `common.py`, not here. Import them from there; `dependencies` holds object references, so an +# unknown name is an ImportError rather than a broken install. + +{module_name(enum_value)}_main = StarterModel( + name="TODO", + base=BaseModelType.{enum_name}, + source="TODO/repo-id", + description="TODO. ~0GB", + type=ModelType.Main, +) +''' + + +def defs_import_line(enum_value: str) -> str: + """The line `architectures/__init__.py` must contain for this architecture to register.""" + return f" {module_name(enum_value)}, # noqa: F401" + + +def starter_import_line(enum_value: str) -> str: + """The import `starter_models/__init__.py` must contain to reach the new entries.""" + return f"from invokeai.backend.model_manager.starter_models.{module_name(enum_value)} import (" + + +RESIDUAL_EDITS = [ + ( + "invokeai/backend/model_manager/taxonomy.py", + "Add the `BaseModelType` member. If the architecture has variants, add its variant enum too, " + "and widen `AnyVariant` plus both halves of `variant_type_adapter` -- CI checks all three " + "against the registry, and checks that variant *values* stay globally unique.", + ), + ( + "invokeai/app/services/model_records/model_records_base.py", + "Widen `ModelRecordChanges.variant` with the new variant enum. Same CI check as above.", + ), + ( + "invokeai/backend/model_manager/configs/main.py (and lora.py / vae.py / controlnet.py)", + "Add the config classes and their probes. This is also where `MainModelDefaultSettings." + "from_base` gains a case if the architecture needs non-default steps/CFG.", + ), + ( + "invokeai/backend/model_manager/configs/factory.py", + "Add the config classes to the `AnyModelConfig` union. Deliberately explicit: a dynamically " + "built union loses type information in IDEs.", + ), + ( + "invokeai/backend/stable_diffusion/diffusion/conditioning_data.py", + "Add the `*ConditioningInfo` dataclass and list it in the `conditionings` union. The " + "`safe_globals` allowlist is derived from the registry, so there is nothing to update in " + "dependencies.py -- but CI asserts the union and the registry agree.", + ), + ( + "invokeai/app/invocations/metadata.py", + "Add the `GENERATION_MODES` literals. These strings are persisted in image metadata and " + "cannot be changed later.", + ), + ( + "invokeai/backend/model_manager/starter_models/__init__.py", + "Import the new module, add the entries to `STARTER_MODELS` where they should appear -- that " + "order is what the model manager shows and is not sorted -- and add a bundle if the " + "architecture ships one.", + ), + ( + "invokeai/backend/model_manager/configs/lora.py", + "If the architecture has LoRAs: add the probe, including the negative clauses that keep other " + "architectures' probes from matching it. Known wart, out of scope of the registry.", + ), + ( + "invokeai/invocation_api/__init__.py", + "Export the new `*ConditioningInfo` so custom nodes can build conditioning. CI checks this " + "against the registry.", + ), +] + + +def _insert_sorted(source: str, line: str, block_start: str) -> str: + """Insert `line` into the alphabetically sorted import block beginning at `block_start`.""" + lines = source.splitlines() + try: + start = next(i for i, text in enumerate(lines) if text.startswith(block_start)) + except StopIteration as exc: + raise SystemExit(f"could not find the import block starting with {block_start!r}") from exc + end = next(i for i in range(start, len(lines)) if lines[i].rstrip() == ")") + body = lines[start + 1 : end] + if any(entry.strip() == line.strip() for entry in body): + return source + body = sorted([*body, line], key=lambda text: text.strip().lstrip("#").strip()) + return "\n".join([*lines[: start + 1], *body, *lines[end:]]) + "\n" + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--enum-name", required=True, help="the BaseModelType member name, e.g. ZImage") + parser.add_argument("--enum-value", required=True, help="its value, e.g. z-image") + parser.add_argument("--write", action="store_true", help="write the files instead of only reporting") + args = parser.parse_args(argv) + + if not re.fullmatch(r"[A-Z][A-Za-z0-9]*", args.enum_name): + raise SystemExit(f"--enum-name must be a CamelCase identifier, got {args.enum_name!r}") + if not re.fullmatch(r"[a-z0-9]+(-[a-z0-9]+)*", args.enum_value): + raise SystemExit(f"--enum-value must be lowercase with dashes, got {args.enum_value!r}") + + stem = module_name(args.enum_value) + targets = { + REPO_ROOT / "invokeai" / "backend" / "architectures" / "defs" / f"{stem}.py": render_defs_module( + args.enum_name, args.enum_value + ), + REPO_ROOT / "invokeai" / "backend" / "model_manager" / "starter_models" / f"{stem}.py": ( + render_starter_models_module(args.enum_name, args.enum_value) + ), + } + + for path, content in targets.items(): + rel = path.relative_to(REPO_ROOT).as_posix() + if path.exists(): + print(f" exists, left alone {rel}") + continue + print(f" {'write' if args.write else 'would write'} {rel}") + if args.write: + path.write_text(content, encoding="utf-8") + + init_edit = ( + ARCHITECTURES_INIT, + defs_import_line(args.enum_value), + "from invokeai.backend.architectures.defs import", + ) + path, line, block = init_edit + rel = path.relative_to(REPO_ROOT).as_posix() + updated = _insert_sorted(path.read_text(encoding="utf-8"), line, block) + if updated == path.read_text(encoding="utf-8"): + print(f" already imported {rel}") + else: + print(f" {'patch' if args.write else 'would patch'} {rel} -> {line.strip()}") + if args.write: + path.write_text(updated, encoding="utf-8") + + print("\nStill to do by hand -- these cannot be generated:\n") + for target, why in RESIDUAL_EDITS: + print(f" {target}\n {why}\n") + print("Then run: pytest tests/backend/architectures -- it will name anything still missing.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_new_architecture_scaffold.py b/tests/test_new_architecture_scaffold.py new file mode 100644 index 00000000000..2db2218b136 --- /dev/null +++ b/tests/test_new_architecture_scaffold.py @@ -0,0 +1,128 @@ +"""The scaffolder is checked against the fifteen architectures that already exist. + +A generator that drifts from what the registry expects is worse than none: it produces files that +look right and fail at boot. So rather than testing the templates against themselves, these assert +that what the scaffolder would emit for an existing architecture matches what that architecture +actually has. +""" + +import ast +import importlib.util +from pathlib import Path + +import pytest + +import invokeai +from invokeai.backend.architectures import generative_bases +from invokeai.backend.architectures.registry import defs_module_path +from invokeai.backend.model_manager.taxonomy import BaseModelType + +REPO_ROOT = Path(invokeai.__file__).parent.parent + +_spec = importlib.util.spec_from_file_location("new_architecture", REPO_ROOT / "scripts" / "new_architecture.py") +assert _spec and _spec.loader +scaffold = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(scaffold) + +BASES = sorted(generative_bases(), key=lambda b: b.value) + + +@pytest.mark.parametrize("base", BASES, ids=lambda b: b.value) +def test_module_name_matches_the_registry(base: BaseModelType) -> None: + """The scaffolder and `registry.defs_module_path()` must agree on where a definition lives.""" + assert defs_module_path(base).endswith(f"/defs/{scaffold.module_name(base.value)}.py") + + +@pytest.mark.parametrize("base", BASES, ids=lambda b: b.value) +def test_the_generated_import_line_is_the_one_actually_used(base: BaseModelType) -> None: + """What the scaffolder would add to `architectures/__init__.py` is what is already there. + + Catches the drift that would otherwise go unnoticed: someone changes the import style in the + aggregate, and the scaffolder keeps emitting the old shape. + """ + aggregate = (REPO_ROOT / "invokeai" / "backend" / "architectures" / "__init__.py").read_text(encoding="utf-8") + expected = scaffold.defs_import_line(base.value).strip() + + assert expected in {line.strip() for line in aggregate.splitlines()} + + +STARTER_MODELS_DIR = REPO_ROOT / "invokeai" / "backend" / "model_manager" / "starter_models" + + +@pytest.mark.parametrize("base", BASES, ids=lambda b: b.value) +def test_the_starter_models_import_line_is_the_one_actually_used(base: BaseModelType) -> None: + """Starter models are optional -- SD2 ships none -- so this only binds architectures that have a + module. Where one exists, the scaffolder must emit the import the catalog actually uses. + """ + module = STARTER_MODELS_DIR / f"{scaffold.module_name(base.value)}.py" + if not module.exists(): + pytest.skip(f"{base.value} ships no starter models") + + catalog = (STARTER_MODELS_DIR / "__init__.py").read_text(encoding="utf-8") + + assert scaffold.starter_import_line(base.value) in catalog + + +def test_sd2_is_the_only_architecture_without_starter_models() -> None: + """Pinned so the skip above stays a fact rather than a way for modules to go missing unnoticed.""" + without = { + base.value for base in BASES if not (STARTER_MODELS_DIR / f"{scaffold.module_name(base.value)}.py").exists() + } + + assert without == {"sd-2"} + + +def test_the_rendered_defs_module_is_valid_python() -> None: + """It is a template with placeholders, but it must still parse -- a contributor edits it, and a + syntax error in generated code is a bad first impression of the pattern. + """ + source = scaffold.render_defs_module("NewModel", "new-model") + + tree = ast.parse(source) + calls = [n for n in ast.walk(tree) if isinstance(n, ast.Call) and getattr(n.func, "id", None) == "register"] + + assert len(calls) == 1 + declared = {getattr(arg.func, "id", None) for arg in calls[0].args if isinstance(arg, ast.Call)} + assert declared == {"LatentSpaceFacet", "ConditioningFacet"}, ( + "the template must scaffold exactly the required facets: leaving one out moves the failure " + "from `validate()` at boot to somewhere later" + ) + + +def test_the_rendered_starter_module_is_valid_python() -> None: + ast.parse(scaffold.render_starter_models_module("NewModel", "new-model")) + + +def test_the_template_scaffolds_every_required_facet() -> None: + """Whatever declares itself REQUIRED must appear in the template. + + Adding a required facet without updating the scaffolder would produce definition modules that + cannot boot the app. + """ + from invokeai.backend.architectures.facet import Facet + + required = {facet.__name__ for facet in Facet.FACET_TYPES if facet.REQUIRED} + source = scaffold.render_defs_module("NewModel", "new-model") + missing = sorted(name for name in required if f"{name}(" not in source) + + assert missing == [], f"scripts/new_architecture.py does not scaffold: {missing}" + + +@pytest.mark.parametrize( + ("enum_name", "enum_value"), + [("newmodel", "new-model"), ("New_Model", "new-model"), ("NewModel", "New-Model"), ("NewModel", "new_model")], +) +def test_bad_identifiers_are_rejected(enum_name: str, enum_value: str) -> None: + with pytest.raises(SystemExit): + scaffold.main(["--enum-name", enum_name, "--enum-value", enum_value]) + + +def test_a_dry_run_writes_nothing(tmp_path: Path, capsys: pytest.CaptureFixture[str]) -> None: + before = (REPO_ROOT / "invokeai" / "backend" / "architectures" / "__init__.py").read_text(encoding="utf-8") + + scaffold.main(["--enum-name", "ScaffoldProbe", "--enum-value", "scaffold-probe"]) + + after = (REPO_ROOT / "invokeai" / "backend" / "architectures" / "__init__.py").read_text(encoding="utf-8") + assert after == before + assert not (REPO_ROOT / "invokeai" / "backend" / "architectures" / "defs" / "scaffold_probe.py").exists() + assert "would write" in capsys.readouterr().out