From 6ef2cb94e13fe13dd1d48ad11b298d938a1b9b41 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sat, 17 Oct 2026 03:33:50 +0200 Subject: [PATCH 01/22] feat(e2e): add per-family end-to-end regression runner Runs each installed model family through a launched scratch server and compares PSNR, time and driver-level VRAM against accepted references. Graphs come from the frontend builders via an offline graph export; the generate submission plan is shared with the workbench. --- .../docs/development/Guides/family-e2e.mdx | 119 ++ invokeai/frontend/webv2/ARCHITECTURE.md | 2 +- invokeai/frontend/webv2/package.json | 3 +- .../frontend/webv2/scripts/export-graphs.mjs | 80 ++ .../webv2/scripts/export-graphs.test.mjs | 74 ++ .../webv2/src/app/graphExport.test.ts | 220 ++++ .../frontend/webv2/src/app/graphExport.ts | 247 ++++ .../src/architecture/dependencyPolicy.test.ts | 1 + .../src/architecture/featureInterfaces.ts | 1 + .../generation/core/baseGenerationPolicies.ts | 12 +- .../generation/core/dynamicPrompts.test.ts | 69 ++ .../generation/core/dynamicPrompts.ts | 78 ++ .../features/generation/core/graphExport.ts | 492 ++++++++ .../generation/core/previewGraph.test.ts | 19 +- .../features/generation/core/previewGraph.ts | 39 +- .../src/features/generation/graphExport.ts | 10 + .../webv2/src/features/generation/settings.ts | 1 + .../webv2/src/features/queue/contracts.ts | 5 + .../src/features/queue/core/promptBatch.ts | 34 + .../webv2/src/features/queue/core/types.ts | 19 +- .../src/features/queue/data/submissionApi.ts | 22 +- .../webv2/src/workbench/workbenchState.ts | 52 +- scripts/family_e2e/__init__.py | 0 scripts/family_e2e/__main__.py | 245 ++++ scripts/family_e2e/client.py | 303 +++++ scripts/family_e2e/detach.py | 78 ++ scripts/family_e2e/gpu_memory.py | 580 +++++++++ scripts/family_e2e/graphs.py | 180 +++ scripts/family_e2e/manifest.example.yaml | 80 ++ scripts/family_e2e/manifest.py | 260 +++++ scripts/family_e2e/metrics.py | 136 +++ scripts/family_e2e/records.py | 284 +++++ scripts/family_e2e/references.py | 124 ++ scripts/family_e2e/report.py | 1036 +++++++++++++++++ scripts/family_e2e/residency.py | 85 ++ scripts/family_e2e/runner.py | 886 ++++++++++++++ scripts/family_e2e/server.py | 509 ++++++++ tests/family_e2e/__init__.py | 0 tests/family_e2e/fake_invokeai_server.py | 191 +++ tests/family_e2e/test_cli.py | 131 +++ tests/family_e2e/test_client.py | 72 ++ tests/family_e2e/test_flow.py | 229 ++++ tests/family_e2e/test_gpu_memory.py | 240 ++++ tests/family_e2e/test_graphs.py | 120 ++ tests/family_e2e/test_manifest.py | 140 +++ tests/family_e2e/test_metrics.py | 116 ++ tests/family_e2e/test_records.py | 287 +++++ tests/family_e2e/test_references.py | 117 ++ tests/family_e2e/test_report.py | 316 +++++ tests/family_e2e/test_residency.py | 87 ++ tests/family_e2e/test_runner.py | 195 ++++ tests/family_e2e/test_server.py | 310 +++++ 52 files changed, 8838 insertions(+), 98 deletions(-) create mode 100644 docs/src/content/docs/development/Guides/family-e2e.mdx create mode 100644 invokeai/frontend/webv2/scripts/export-graphs.mjs create mode 100644 invokeai/frontend/webv2/scripts/export-graphs.test.mjs create mode 100644 invokeai/frontend/webv2/src/app/graphExport.test.ts create mode 100644 invokeai/frontend/webv2/src/app/graphExport.ts create mode 100644 invokeai/frontend/webv2/src/features/generation/core/graphExport.ts create mode 100644 invokeai/frontend/webv2/src/features/generation/graphExport.ts create mode 100644 scripts/family_e2e/__init__.py create mode 100644 scripts/family_e2e/__main__.py create mode 100644 scripts/family_e2e/client.py create mode 100644 scripts/family_e2e/detach.py create mode 100644 scripts/family_e2e/gpu_memory.py create mode 100644 scripts/family_e2e/graphs.py create mode 100644 scripts/family_e2e/manifest.example.yaml create mode 100644 scripts/family_e2e/manifest.py create mode 100644 scripts/family_e2e/metrics.py create mode 100644 scripts/family_e2e/records.py create mode 100644 scripts/family_e2e/references.py create mode 100644 scripts/family_e2e/report.py create mode 100644 scripts/family_e2e/residency.py create mode 100644 scripts/family_e2e/runner.py create mode 100644 scripts/family_e2e/server.py create mode 100644 tests/family_e2e/__init__.py create mode 100644 tests/family_e2e/fake_invokeai_server.py create mode 100644 tests/family_e2e/test_cli.py create mode 100644 tests/family_e2e/test_client.py create mode 100644 tests/family_e2e/test_flow.py create mode 100644 tests/family_e2e/test_gpu_memory.py create mode 100644 tests/family_e2e/test_graphs.py create mode 100644 tests/family_e2e/test_manifest.py create mode 100644 tests/family_e2e/test_metrics.py create mode 100644 tests/family_e2e/test_records.py create mode 100644 tests/family_e2e/test_references.py create mode 100644 tests/family_e2e/test_report.py create mode 100644 tests/family_e2e/test_residency.py create mode 100644 tests/family_e2e/test_runner.py create mode 100644 tests/family_e2e/test_server.py diff --git a/docs/src/content/docs/development/Guides/family-e2e.mdx b/docs/src/content/docs/development/Guides/family-e2e.mdx new file mode 100644 index 00000000000..27580f6f40d --- /dev/null +++ b/docs/src/content/docs/development/Guides/family-e2e.mdx @@ -0,0 +1,119 @@ +--- +title: Family E2E Runs +lastUpdated: 2026-10-17 +--- + +`scripts/family_e2e` runs model families end to end through the HTTP API of a server it launches itself. Per case and size it records the image, the run time, the time per denoise step and the VRAM the server process used, measured at driver level. The report compares them with an explicitly accepted reference and with the previous run of the same profile. + +It is a developer tool. Memory sampling uses WDDM, so it runs on Windows only for now, with NVIDIA and AMD (ROCm) GPUs. + +## How a run works + +1. **Scratch root.** The template root's `invokeai.yaml` and a snapshot of its SQLite database go to `/scratch-roots/-default/`. The snapshot is taken with SQLite's backup API, so a WAL left by an unclean stop is included and nothing is written next to the template's database. + - The tool refuses while the database file is open in any process, or while anything listens on the template's configured port. Either means its server is still running. + - It also refuses while a model install is unfinished (see [Shared paths](#shared-paths)). + - Stores whose files the database owns start empty in the scratch root: outputs, uploaded fonts, style preset images, workflow thumbnails and profiles. The scratch server reconciles them with its own database; at the template it would, for example, delete uploaded fonts its database does not list. + - Every other path setting points back at the template root as an absolute path, so models, custom nodes and memory settings are the template's own. + - The tool overrides a few settings: a free port on `127.0.0.1`, a plain-format log file in the run directory, the image index off, and the queue cleared at startup. SSL and `base_url` are dropped. + - `INVOKEAI_*` variables of the calling environment are not passed on. Put settings the server needs into the manifest's `server.env`. +2. **Server.** The tool checks that `server.python` imports `invokeai` from `server.code_dir`. It then starts ` /scripts/invokeai-web.py --root ` and reads the actual port from the server log. Multiuser mode is refused. + - On Windows the server runs in a job object of the runner, so it dies with the runner however the runner ends. +3. **Graphs.** The real webv2 graph builders produce each case's batch through `pnpm graphs:export`, run in the tool checkout. Run `pnpm install` in `invokeai/frontend/webv2` there first. + - The runner disables the invocation cache on every node and saves the batch as `graphs/.json`. + - It hashes the graph and its batch data. Uploaded image names are replaced by the input file's hash, so the graph hash changes only when the graph or an input changes. + - Init and control images enter the canvas compiler as finished layers. One whose size differs from the case size is resized (LANCZOS) into `runs//inputs/` before upload. Reference images are uploaded as they are. +4. **Steps.** Per family the model cache is emptied first. Each case size then gets its warm-ups, then its measured runs. + - Per run, the tool waits up to 10 minutes for the queue to drain, samples the server's GPU memory, enqueues the batch and follows its socket events. + - When the run finishes, it saves the output image and the memory trace, reads the run's slice of the server log and appends a record to `records.jsonl`. + - If the server dies during a step, that step gets an error record and a new server starts on a new scratch root. A case whose steps crashed the server twice is given up; its remaining steps are not run again, not even by `resume`. +5. **Report.** The server is stopped, the scratch root removed, and `report.html` and `summary.md` are written. The exit code is not 0 when a case fails or errors. + +## Data directory + +``` +/inputs/ input images the manifest refers to +/references/// image.png, reference.json (written by `accept` only) +/runs// run.json, records.jsonl, heartbeat.json, runner.log, server-default.log, + graphs/, images/, traces/.csv.gz, report.html, summary.md +/scratch-roots/-default/ +``` + +`run_id` is `YYYYMMDD-HHMMSS-`, and `case_key` is `@x`. A trace holds the raw samples as `series,t,bytes`: series `dedicated` or `shared`, and `t` in seconds from the enqueue. + +One runner per GPU: a lock file per adapter lives in `%LOCALAPPDATA%\invokeai-family-e2e\locks\`, shared by every data directory on the machine. + +## Manifest + +A YAML file per machine and GPU, kept with the data rather than in the repository. Start from `scripts/family_e2e/manifest.example.yaml`. Relative paths are relative to the manifest file. + +| Field | Meaning | +| --- | --- | +| `profile` | Names runs and references. | +| `data_dir` | The data directory above. | +| `gpu` | Substring of the adapter name that matches exactly one adapter, e.g. `RTX 4090` or `9060`. `Radeon` alone also matches an AMD integrated GPU. | +| `server.template_root` | The root whose settings and database are copied. The scratch server shares its other paths; see [Shared paths](#shared-paths). | +| `server.python` | The interpreter that runs the server. | +| `server.code_dir` | The checkout the server runs. Defaults to the tool's checkout. | +| `server.env` | Extra server environment, e.g. `HIP_VISIBLE_DEVICES`. | +| `defaults` | `warmup` (at least 1), `repeats`, `seed`, `prompt`, `negative_prompt`, `psnr_min_db`, `time_tolerance` (relative), `peak_tolerance_mb`, `run_timeout_s`. | +| `families[].cases[]` | `id`, `model` (key), `mode` (`txt2img`, `img2img` or `controlnet`), `sizes` (a list of `[width, height]`), `steps`, optional `seed` and `prompt`, `components` (component slot to model key), `loras`, `reference_images`, `init_image` (`path`, `strength`), `control` (`model`, `image`, `weight`, `begin`, `end`), `settings` (generate settings), and per-case thresholds. | + +Input paths are relative to `/inputs` and must exist. `init_image` belongs to `img2img` cases only and `control` to `controlnet` cases only. The exporter rejects an unknown `settings` key, component slot or model key; that case then gets error records and the run goes on. Several `sizes` per case give a size sweep for working-memory calibration. + +## Commands + +Run the tool from the repository root with the checkout's environment, for example `uv run --no-sync python -m scripts.family_e2e …`. + +| Command | What it does | +| --- | --- | +| `start --manifest M [--family F]… [--case ID]…` | Checks the GPU lock, the template's server, unfinished model installs and foreign GPU load. It then creates the run and starts its runner in the background, and prints the run id and paths. | +| `status [RUN]` | Shows progress, the current step, the heartbeat age and whether the runner is alive. | +| `stop [RUN]` | Kills the runner and its server, removes the run's scratch roots and marks the run stopped. | +| `resume [RUN]` | Starts the runner of a stopped or crashed run again. | +| `report [RUN]` | Writes `report.html` and `summary.md` again. | +| `accept [RUN] [--case KEY]…` | Makes the run's measured results the references of its profile. Nothing else writes references. | +| `run --run-dir D` | The runner in the foreground: what `start` and `resume` launch. | + +`RUN` is a run directory, or a run id together with `--manifest M`. Without it, the profile's latest run is used. + +On Windows the background runner breaks away from the caller's job object, so a terminal or an agent harness that kills its job's processes does not end a long run. Where breakaway is refused, the runner is created through WMI. It then starts with the user's default environment instead of the caller's. + +A resumed run skips every step that already has a final record and keeps the measures of the earlier sessions. Its new server starts with a cold model cache, so each case size with measures left first repeats its last warm-up, once. The manifest, the graphs, and the server's commit and dirty flag must not have changed since the run started; further edits to an already dirty checkout cannot be detected. + +## Validity + +A run's record can be complete and still unfit for some comparisons. Each invalid flag comes with its reason in `invalid_reasons`. + +- **Residency.** The model cache logs `Loaded model '…' … VRAM: … (P%)` whenever a lock moves weights or leaves a model short of fully resident. A warm run usually logs nothing, because its models stay resident. The tool therefore keeps the latest line per model since the last cache flush. Time and memory are invalid when: + - a line in the run's log slice is below 100 %, + - a `VRAM budget … is short` warning appears, or + - no full load of the graph's models was seen since the flush ("residency unknown"). +- **fp8.** With fp8 weights in play, an incompletely resident run does not reproduce its image, so PSNR is invalid too. fp8 is in play when a graph model sets `fp8_storage` in its default settings, or when the server's `fp8_compute` is on. Model configs do not record which checkpoints are scaled fp8, so then every graph counts. +- **Shared memory.** If the server's shared GPU memory grows by more than 256 MiB during a run, it paged to system memory: time and memory are invalid. +- **Foreign load.** Another process above 1 GiB of dedicated GPU memory at either end of the run, or grown by more than 256 MiB, invalidates time and memory. So does a queue that is still busy after the 10-minute wait. The run is repeated, up to three attempts. `start` refuses outright when another process holds more than 1 GiB. + +## Shared paths + +The scratch server uses the template root's models directory, model caches, legacy model configs, custom fonts and custom nodes. Its startup touches them as starting the template's own server would: + +- `models/README.txt` is written. +- Install claims (`*.invokeai_install_active`) of processes that no longer run are removed. +- Missing model files are reported in the log. +- `configs/` receives the legacy model configs of the server's code when they differ. +- `fonts/README.txt` is written when missing. +- Models a node fetches on first use are cached in `models/.download_cache`. + +The server would also finish an incomplete install, but the model would then be registered in the scratch database only. So `start` and `resume` refuse while the models directory holds a `tmpinstall_*` entry: finish or cancel that install in the dev server first. + +Database migrations of a newer `code_dir` run against the scratch database only. A migration that also moves or rewrites files acts on the shared models directory and caches; keep that in mind before pointing a run at a checkout ahead of the template's own server. + +## Limits + +- Memory is sampled every 3 ms with `D3DKMTQueryVideoMemoryInfo`: shorter peaks are missed. It counts resident memory (`CurrentUsage`), not commit, including the CUDA or HIP context and every resident model. That is why records keep a baseline and a peak per node. +- With PyTorch's native allocator, reserved memory stays allocated between nodes, so a node's peak can stem from an earlier node. Compare peaks only between runs with the same allocator configuration; the run records the runtime config. +- On Windows with ROCm, HIP keeps freed blocks at driver level. Per-node peaks after the first large node are upper bounds there (observed on an RX 9060 XT). +- Time per step comes from the preview frames, which the server throttles to one per 0.1 s; a step's number is read from its progress percentage. + - For txt2img and controlnet the case's `steps` fixes the step count. + - img2img starts inside the schedule at a family-specific step, so the count is inferred. When steps are faster than 0.2 s that inference can be off by a factor, so the time per step is left empty. +- Ideogram 4 ignores `steps`; its step count follows `settings.ideogram4Steps` or `settings.ideogram4SamplerPreset`. `img2img` and `controlnet` are not available for ERNIE-Image and Ideogram 4. Dynamic prompts, canvas rasterization, video and upscale graphs are not covered. +- Linux memory sampling and A/B runs of two checkouts are not available yet. diff --git a/invokeai/frontend/webv2/ARCHITECTURE.md b/invokeai/frontend/webv2/ARCHITECTURE.md index d72f3f3a710..6caa7acb868 100644 --- a/invokeai/frontend/webv2/ARCHITECTURE.md +++ b/invokeai/frontend/webv2/ARCHITECTURE.md @@ -104,7 +104,7 @@ Feature interfaces stay ordinary module interfaces unless two adapters add lever | Nodes | Custom-node catalog/templates and execution lifecycle; Workflow editor and node manager | Platform; descriptors, catalog commands, execution snapshots and manager contribution | | Queue | Backend queue state, realtime interpretation, submission/reconciliation and Queue UI; App runtime, Workbench aggregate/shell, Generation/Gallery/Workflow | Platform and public envelopes/capabilities from Identity, Models, Generation, Workflow, Upscale; backend port, queue reads/commands and widget contribution | | Gallery | Image collection, selection, board policy and Gallery UI; App adapter, Generation, Workflow, Upscale, Models, Workbench image actions | Platform and Queue's public progress reads; gallery reads/commands and widget contribution | -| Generation | Generation settings, validation, graph compilation and Generate UI; Workbench invocation, Queue, Gallery, Workflow, Upscale | Platform and public Models/Gallery/Queue interfaces; settings/validation, compiled submission envelope and widget contribution | +| Generation | Generation settings, validation, graph compilation and Generate UI; Workbench invocation, Queue, Gallery, Workflow, Upscale | Platform and public Models/Gallery/Queue interfaces; settings/validation, compiled submission envelope, offline graph export and widget contribution | | Workflow | Workflow document, validation, compilation, preview/editor/library and explicit publication; App adapter, Workbench invocation, Nodes | Platform and public Nodes/Models/Gallery/Queue/Generation interfaces; document reads/commands, compiled submission envelope and widget contribution | | Upscale | Upscale settings, validation, compilation and UI; Workbench invocation and Queue | Platform and public Models/Gallery/Queue interfaces; value/validation reads, submission command and widget contribution | | Intermediates | Per-project intermediates read models, cleanup previews/operations and the Settings section that manages them; Workbench settings, project menus, account administration | Platform plus Identity's public session/capability reads; query keys, read models, the focus intent, the browser hold lease and the `settingsContribution` metadata whose loader owns the section body | diff --git a/invokeai/frontend/webv2/package.json b/invokeai/frontend/webv2/package.json index 8cc729ef6b6..069b1cc3701 100644 --- a/invokeai/frontend/webv2/package.json +++ b/invokeai/frontend/webv2/package.json @@ -17,6 +17,7 @@ "fix": "pnpm run lint:oxc:fix && pnpm run format:write", "format:check": "oxfmt --check .", "format:write": "oxfmt --write .", + "graphs:export": "node scripts/export-graphs.mjs", "i18n:report": "node scripts/i18n-report.mjs", "i18n:report:verbose": "node scripts/i18n-report.mjs --verbose", "lint": "pnpm run format:check && pnpm run lint:oxc && pnpm run lint:tsc && pnpm run architecture:check", @@ -32,7 +33,7 @@ "test:all": "pnpm run test && pnpm run test:browser && pnpm run test:fixtures", "test:browser": "vitest run --config vitest.browser.config.mts", "test:browser:update-goldens": "CANVAS_GOLDEN_UPDATE=1 vitest run --config vitest.browser.config.mts src/workbench/canvas-engine/render/golden", - "test:fixtures": "node --test scripts/mock-backend-fixtures.test.mjs scripts/preview-server.test.mjs scripts/run-project-file-journey.test.mjs scripts/locale-assets-plugin.test.mjs", + "test:fixtures": "node --test scripts/mock-backend-fixtures.test.mjs scripts/preview-server.test.mjs scripts/run-project-file-journey.test.mjs scripts/locale-assets-plugin.test.mjs scripts/export-graphs.test.mjs", "test:performance:architecture": "pnpm run test:performance:build && pnpm run test:performance:browser", "test:performance:browser": "node scripts/measure-architecture-performance.mjs", "test:performance:browser:update-baseline": "node scripts/measure-architecture-performance.mjs --update-baseline", diff --git a/invokeai/frontend/webv2/scripts/export-graphs.mjs b/invokeai/frontend/webv2/scripts/export-graphs.mjs new file mode 100644 index 00000000000..cc44c137198 --- /dev/null +++ b/invokeai/frontend/webv2/scripts/export-graphs.mjs @@ -0,0 +1,80 @@ +/** + * `pnpm graphs:export ` compiles offline export cases into the enqueue batches the UI + * would send (src/app/graphExport.ts). Vite's module runner loads the TypeScript sources with the app's aliases and + * opens no port. Exit 0 once the response is written, with per-case errors inside it; exit 2 for an unreadable or + * malformed request, with no response written. + */ + +import { readFile, writeFile } from 'node:fs/promises'; +import { dirname, resolve } from 'node:path'; +import process from 'node:process'; +import { fileURLToPath } from 'node:url'; +import { loadConfigFromFile, runnerImport } from 'vite'; + +const PACKAGE_ROOT = resolve(dirname(fileURLToPath(import.meta.url)), '..'); +const EXIT_MALFORMED_REQUEST = 2; + +const rejectRequest = (message) => { + process.stderr.write(`graphs:export: ${message}\n`); + process.exitCode = EXIT_MALFORMED_REQUEST; +}; + +const readRequest = async (path) => { + try { + return JSON.parse(await readFile(path, 'utf8')); + } catch (error) { + rejectRequest(`cannot read ${path}: ${error instanceof Error ? error.message : String(error)}`); + return undefined; + } +}; + +/** Only the aliases and defines: plugins transform the browser bundle, not these pure modules. */ +const loadExporter = async () => { + const { config } = await loadConfigFromFile( + { command: 'serve', mode: 'development' }, + resolve(PACKAGE_ROOT, 'vite.config.mts'), + PACKAGE_ROOT, + 'error' + ); + const { module } = await runnerImport(resolve(PACKAGE_ROOT, 'src/app/graphExport.ts'), { + define: config.define, + logLevel: 'error', + resolve: { alias: config.resolve.alias }, + root: PACKAGE_ROOT, + }); + + return module; +}; + +const main = async () => { + const [requestPath, responsePath, ...extra] = process.argv.slice(2); + + if (!requestPath || !responsePath || extra.length > 0) { + rejectRequest('usage: pnpm graphs:export '); + return; + } + + const request = await readRequest(requestPath); + + if (request === undefined) { + return; + } + + const { exportGraphs, GraphExportRequestError } = await loadExporter(); + let response; + + try { + response = exportGraphs(request); + } catch (error) { + if (error instanceof GraphExportRequestError) { + rejectRequest(`malformed request: ${error.message}`); + return; + } + + throw error; + } + + await writeFile(responsePath, `${JSON.stringify(response, null, 2)}\n`); +}; + +await main(); diff --git a/invokeai/frontend/webv2/scripts/export-graphs.test.mjs b/invokeai/frontend/webv2/scripts/export-graphs.test.mjs new file mode 100644 index 00000000000..607daf3856b --- /dev/null +++ b/invokeai/frontend/webv2/scripts/export-graphs.test.mjs @@ -0,0 +1,74 @@ +import assert from 'node:assert/strict'; +import { execFile } from 'node:child_process'; +import { access, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { dirname, resolve } from 'node:path'; +import test from 'node:test'; +import { fileURLToPath } from 'node:url'; +import { promisify } from 'node:util'; + +const PACKAGE_ROOT = resolve(dirname(fileURLToPath(import.meta.url)), '..'); +const CLI = resolve(PACKAGE_ROOT, 'scripts/export-graphs.mjs'); +const CAPABILITIES = resolve(PACKAGE_ROOT, 'src/features/generation/core/__fixtures__/architectureCapabilities.json'); + +const runCli = async (context, requestText) => { + const directory = await mkdtemp(resolve(tmpdir(), 'invokeai-graphs-export-')); + context.after(() => rm(directory, { force: true, recursive: true })); + const requestPath = resolve(directory, 'request.json'); + const responsePath = resolve(directory, 'response.json'); + await writeFile(requestPath, requestText); + + const { code, stderr } = await promisify(execFile)(process.execPath, [CLI, requestPath, responsePath], { + cwd: PACKAGE_ROOT, + }).then( + ({ stderr }) => ({ code: 0, stderr }), + (error) => ({ code: error.code, stderr: error.stderr }) + ); + const response = await access(responsePath).then( + async () => JSON.parse(await readFile(responsePath, 'utf8')), + () => null + ); + + return { code, response, stderr }; +}; + +test('writes per-case results through the real module graph and exits 0', async (context) => { + const capabilities = JSON.parse(await readFile(CAPABILITIES, 'utf8')); + const sdxl = { base: 'sdxl', format: 'diffusers', hash: 'blake3:sdxl', key: 'sdxl', name: 'SDXL', type: 'main' }; + const testCase = { height: 1024, mode: 'txt2img', positive_prompt: 'a fox', seed: 7, steps: 4, width: 1024 }; + const request = { + capabilities, + cases: [ + { ...testCase, id: 'ok', model_key: 'sdxl' }, + { ...testCase, id: 'missing', model_key: 'missing' }, + ], + models: [sdxl], + version: 1, + }; + + const { code, response } = await runCli(context, JSON.stringify(request)); + + assert.equal(code, 0); + assert.equal(response.version, 1); + assert.deepEqual( + response.cases.map((result) => [result.id, 'batch' in result]), + [ + ['ok', true], + ['missing', false], + ] + ); + assert.match(response.cases[1].error, /model_key "missing"/); +}); + +for (const [label, requestText] of [ + ['unparsable JSON', '{ "version": 1,'], + ['a malformed request', JSON.stringify({ cases: [], models: [], version: 1 })], +]) { + test(`rejects ${label} with exit 2 and writes no response`, async (context) => { + const { code, response, stderr } = await runCli(context, requestText); + + assert.equal(code, 2); + assert.equal(response, null); + assert.match(stderr, /^graphs:export: /); + }); +} diff --git a/invokeai/frontend/webv2/src/app/graphExport.test.ts b/invokeai/frontend/webv2/src/app/graphExport.test.ts new file mode 100644 index 00000000000..cd29c1b5864 --- /dev/null +++ b/invokeai/frontend/webv2/src/app/graphExport.test.ts @@ -0,0 +1,220 @@ +import type { BackendGraphContract, BackendInvocationContract } from '@features/generation/contracts'; + +import { architectureCapabilitiesFixture } from '@features/generation/core/architectureCapabilities.testing'; +import { generateGraphCases, satisfiedSettingsFor } from '@features/generation/core/graphCoverage.testing'; +import { withArchitectureCapabilities } from '@features/generation/graphExport'; +import { describe, expect, it } from 'vitest'; + +import { exportGraphs, GraphExportRequestError, type GraphExportCaseResult } from './graphExport'; + +const model = (key: string, base: string, type: string, extra: Record = {}) => ({ + base, + hash: `blake3:${key}`, + key, + name: `${key} name`, + type, + ...extra, +}); + +const MODELS = [ + model('sdxl', 'sdxl', 'main', { format: 'diffusers', variant: 'normal' }), + model('sdxl-lora', 'sdxl', 'lora', { format: 'lycoris' }), + model('sdxl-canny', 'sdxl', 'controlnet', { format: 'diffusers' }), + model('sdxl-ip-adapter', 'sdxl', 'ip_adapter', { format: 'invokeai' }), + model('z-image-gguf', 'z-image', 'main', { format: 'gguf_quantized' }), + model('qwen3-06b', 'any', 'qwen3_encoder', { format: 'checkpoint', variant: 'qwen3_06b' }), + model('qwen3-4b', 'any', 'qwen3_encoder', { format: 'checkpoint', variant: 'qwen3_4b' }), + model('qwen3-8b', 'any', 'qwen3_encoder', { format: 'checkpoint', variant: 'qwen3_8b' }), + model('sdxl-vae', 'sdxl', 'vae', { format: 'diffusers' }), + model('flux-vae', 'flux', 'vae', { format: 'checkpoint' }), +]; + +const sdxlCase = (id: string, overrides: Record = {}) => ({ + height: 832, + id, + mode: 'txt2img', + model_key: 'sdxl', + negative_prompt: 'blurry', + positive_prompt: 'a lighthouse at dusk', + seed: 4242, + steps: 12, + width: 1216, + ...overrides, +}); + +const request = (cases: unknown[], models: unknown[] = MODELS) => ({ + capabilities: architectureCapabilitiesFixture, + cases, + models, + version: 1, +}); + +const exportOne = (testCase: Record, models?: unknown[]) => { + const [result] = exportGraphs(request([testCase], models)).cases; + + if (!result || !('batch' in result)) { + throw new Error(`expected a batch, got ${JSON.stringify(result)}`); + } + + return result; +}; + +const nodesOfType = (graph: BackendGraphContract, type: string): BackendInvocationContract[] => + Object.values(graph.nodes).filter((node) => node.type === type); + +const errorOf = (result: GraphExportCaseResult | undefined): string => + result && 'error' in result ? result.error : ''; + +describe('graph export', () => { + it('sends the case through Generate: seed and prompts as batch data, size and steps in the graph', () => { + const { batch, id } = exportOne( + sdxlCase('sdxl-t2i', { loras: [{ key: 'sdxl-lora', weight: 0.65 }], reference_images: ['reference.png'] }) + ); + const { graph } = batch; + const noise = nodesOfType(graph, 'noise')[0]; + const denoise = graph.nodes.denoise_latents; + + expect(id).toBe('sdxl-t2i'); + expect(batch.data).toEqual([ + [ + { field_name: 'value', items: [4242], node_path: 'seed' }, + { field_name: 'value', items: ['a lighthouse at dusk'], node_path: 'positive_prompt' }, + { field_name: 'value', items: ['blurry'], node_path: 'negative_prompt' }, + ], + ]); + expect(batch.data[0]?.map((datum) => graph.nodes[datum.node_path]?.type)).toEqual(['integer', 'string', 'string']); + expect(batch).toMatchObject({ destination: 'gallery', origin: 'graphs-export:sdxl-t2i' }); + expect(graph.id).toBe('sdxl-t2i'); + expect([noise?.width, noise?.height, denoise?.steps]).toEqual([1216, 832, 12]); + expect( + nodesOfType(graph, 'lora_selector').map((node) => [(node.lora as { key: string }).key, node.weight]) + ).toEqual([['sdxl-lora', 0.65]]); + expect(nodesOfType(graph, 'ip_adapter')).toMatchObject([ + { image: { image_name: 'reference.png' }, ip_adapter_model: { key: 'sdxl-ip-adapter' } }, + ]); + }); + + it('sends img2img through Canvas with the uploaded image and strength', () => { + const { graph } = exportOne( + sdxlCase('sdxl-i2i', { init_image: { image_name: 'init.png', strength: 0.3 }, mode: 'img2img' }) + ).batch; + + expect(nodesOfType(graph, 'i2l').map((node) => node.image)).toEqual([{ image_name: 'init.png' }]); + expect(graph.nodes.denoise_latents?.denoising_start).toBeCloseTo(0.7); + expect(nodesOfType(graph, 'core_metadata')[0]).toMatchObject({ generation_mode: 'sdxl_img2img', strength: 0.3 }); + }); + + it('adds each control entry as a Canvas control layer feeding the denoiser', () => { + const { graph } = exportOne( + sdxlCase('sdxl-control', { + control: [{ begin: 0.1, end: 0.6, image_name: 'canny.png', model_key: 'sdxl-canny', weight: 0.9 }], + mode: 'controlnet', + }) + ).batch; + const [controlNet] = nodesOfType(graph, 'controlnet'); + + expect(controlNet).toMatchObject({ + begin_step_percent: 0.1, + control_model: { key: 'sdxl-canny' }, + control_weight: 0.9, + end_step_percent: 0.6, + image: { image_name: 'canny.png' }, + }); + expect(graph.edges).toContainEqual({ + destination: { field: 'control', node_id: 'denoise_latents' }, + source: { field: 'collection', node_id: 'control_net_collector' }, + }); + }); + + it('fills required components from the installed models and keeps explicit picks', () => { + const zImage = { ...sdxlCase('z-image'), model_key: 'z-image-gguf' }; + + // The 0.6B encoder is listed first but Z-Image cannot use it. + expect(exportOne(zImage).components).toEqual({ qwen3EncoderModel: 'qwen3-4b', vae: 'flux-vae' }); + expect(exportOne({ ...zImage, components: { qwen3EncoderModel: 'qwen3-8b' } }).components).toEqual({ + qwen3EncoderModel: 'qwen3-8b', + vae: 'flux-vae', + }); + }); + + it.each(generateGraphCases)('compiles $label from exactly the components it needs', ({ base, shape }) => { + const { + filled, + model: main, + settings, + } = withArchitectureCapabilities(architectureCapabilitiesFixture, () => satisfiedSettingsFor(base, shape)); + const components = Object.fromEntries(filled.map((key) => [key, (settings[key] as { key: string }).key])); + const testCase = sdxlCase(base, { height: settings.height, model_key: main.key, width: settings.width }); + + expect(exportOne(testCase, [main, ...filled.map((key) => settings[key])]).components).toEqual(components); + }); + + it('applies settings overrides and rejects values Generate would not keep', () => { + const { graph } = exportOne(sdxlCase('cfg', { settings: { cfgScale: 3.5, scheduler: 'dpmpp_2m' } })).batch; + const [badPid, badScheduler] = exportGraphs( + request([ + sdxlCase('pid', { settings: { pidMode: 'sideways' } }), + sdxlCase('scheduler', { settings: { scheduler: 'not-a-scheduler' } }), + ]) + ).cases; + + expect(graph.nodes.denoise_latents).toMatchObject({ cfg_scale: 3.5, scheduler: 'dpmpp_2m' }); + expect(errorOf(badPid)).toMatch(/^settings\.pidMode: Generate does not keep "sideways"/); + expect(errorOf(badScheduler)).toMatch(/^settings\.scheduler: .* Options: .*euler_a/); + }); + + it('reports unknown names per case, with the valid choices, and still exports the other cases', () => { + const response = exportGraphs( + request([ + sdxlCase('unknown-setting', { settings: { cfgscale: 3 } }), + sdxlCase('unknown-slot', { components: { t5EncoderModel: 'qwen3-4b' }, model_key: 'z-image-gguf' }), + sdxlCase('unknown-component', { components: { vae: 'missing-vae' }, model_key: 'z-image-gguf' }), + sdxlCase('unknown-model', { model_key: 'missing' }), + sdxlCase('unknown-lora', { loras: [{ key: 'missing-lora', weight: 1 }] }), + sdxlCase('unknown-control', { + control: [{ begin: 0, end: 1, image_name: 'canny.png', model_key: 'sdxl-lora', weight: 1 }], + mode: 'controlnet', + }), + sdxlCase('ok'), + ]) + ); + const errors = Object.fromEntries(response.cases.map((result) => [result.id, errorOf(result)])); + + expect(errors['unknown-setting']).toMatch(/^Unknown settings key "cfgscale"\. Valid keys: .*\bcfgScale\b/); + expect(errors['unknown-slot']).toMatch( + /^Unknown component slot "t5EncoderModel" .* Valid slots: .*qwen3EncoderModel, vae/ + ); + expect(errors['unknown-component']).toBe( + 'components.vae: model key "missing-vae" is not installed. Choices: flux-vae.' + ); + expect(errors['unknown-model']).toBe( + 'model_key "missing" is not a selectable main model. Choices: sdxl, z-image-gguf.' + ); + expect(errors['unknown-lora']).toMatch(/^loras\[0\]: "missing-lora" .* Choices: sdxl-lora\.$/); + expect(errors['unknown-control']).toMatch(/^control\[0\]: "sdxl-lora" cannot drive .* Choices: sdxl-canny\.$/); + expect(response.cases.find((result) => result.id === 'ok')).toHaveProperty('batch'); + }); + + it('exports identical requests identically, with stable node ids', () => { + const cases = [ + sdxlCase('t2i', { loras: [{ key: 'sdxl-lora', weight: 1 }], reference_images: ['reference.png'] }), + sdxlCase('i2i', { init_image: { image_name: 'init.png', strength: 0.5 }, mode: 'img2img' }), + { ...sdxlCase('z-image'), model_key: 'z-image-gguf' }, + ]; + const first = exportGraphs(request(cases)); + const nodeIds = first.cases.flatMap((result) => ('batch' in result ? Object.keys(result.batch.graph.nodes) : [])); + + expect(exportGraphs(request(cases))).toEqual(first); + expect(nodeIds.filter((id) => /_[a-z0-9]{7,9}_[a-z0-9]{4,6}$/.test(id))).toEqual([]); + }); + + it.each([ + ['a non-object request', []], + ['another version', { ...request([]), version: 2 }], + ['models that are not model configs', request([], [{ key: 'only-a-key' }])], + ['cases without ids', request([{ mode: 'txt2img' }])], + ['duplicate case ids', request([sdxlCase('same'), sdxlCase('same')])], + ])('rejects %s as a malformed request', (_label, malformed) => { + expect(() => exportGraphs(malformed)).toThrow(GraphExportRequestError); + }); +}); diff --git a/invokeai/frontend/webv2/src/app/graphExport.ts b/invokeai/frontend/webv2/src/app/graphExport.ts new file mode 100644 index 00000000000..ea53a234c9e --- /dev/null +++ b/invokeai/frontend/webv2/src/app/graphExport.ts @@ -0,0 +1,247 @@ +/** + * Composition root of `pnpm graphs:export` (scripts/export-graphs.mjs): installed model configs, the capability table + * and test cases become the enqueue batches the UI would send, offline, so end-to-end runs exercise the graphs users + * submit. Each batch is what Generate or Canvas puts into `enqueue_batch` minus `runs`, `project_id` and + * `idempotency_key`, with the case seed held fixed. Generation and Queue own every rule; this module maps the wire + * format and supplies the Workbench defaults the widgets would. + */ + +import type { GeneratePromptBatchPlan } from '@features/queue/contracts'; + +import { + compileGraphExportCase, + withArchitectureCapabilities, + type ArchitectureCapabilitiesRow, + type CompiledGraphExportCase, + type GraphExportCase, + type GraphExportEnvironment, + type GraphExportMode, +} from '@features/generation/graphExport'; +import { planGenerateSubmissionBatch } from '@features/queue/contracts'; +import { resolveRandDeviceMetadata } from '@features/queue/devices'; +import { SEED_MAX } from '@platform/core/seed'; +import { DEFAULT_CANVAS_COMPOSITING } from '@workbench/widgets/canvas/invoke/canvasCompositing'; +import { DEFAULT_CANVAS_SCALING } from '@workbench/widgets/canvas/invoke/canvasScaling'; +import { DEFAULT_CANVAS_DENOISING_STRENGTH } from '@workbench/widgets/canvas/invoke/canvasStrength'; + +const GRAPH_EXPORT_VERSION = 1; + +/** Durable outputs: Canvas staging would mark them intermediate. */ +const DESTINATION = 'gallery'; +/** Not a webv2 origin, so no open editor adopts the items as its own submissions. */ +const ORIGIN_PREFIX = 'graphs-export:'; +/** `DEFAULT_PROJECT_SETTINGS` (workbench/settings/store.ts); importing it would load the whole settings runtime. */ +const PROJECT_SETTINGS = { useCpuNoise: true }; +/** ROCm builds of torch report their GPUs as CUDA devices as well. */ +const GPU_DEVICE_OPTIONS = [{ device: 'cuda:0', name: 'GPU' }]; +const MODES: readonly GraphExportMode[] = ['txt2img', 'img2img', 'controlnet']; +const CASE_FIELDS = [ + 'components', + 'control', + 'height', + 'id', + 'init_image', + 'loras', + 'mode', + 'model_key', + 'negative_prompt', + 'positive_prompt', + 'reference_images', + 'seed', + 'settings', + 'steps', + 'width', +]; + +/** The request itself is unusable; no response is written. */ +export class GraphExportRequestError extends Error { + override name = 'GraphExportRequestError'; +} + +export interface GraphExportBatch { + data: GeneratePromptBatchPlan['data']; + destination: typeof DESTINATION; + graph: CompiledGraphExportCase['graph']; + origin: string; +} + +export type GraphExportCaseResult = + | { id: string; batch: GraphExportBatch; components: Record } + | { id: string; error: string }; + +export interface GraphExportResponse { + version: typeof GRAPH_EXPORT_VERSION; + cases: GraphExportCaseResult[]; +} + +type JsonRecord = Record; + +const isRecord = (value: unknown): value is JsonRecord => + typeof value === 'object' && value !== null && !Array.isArray(value); + +const fail = (message: string): never => { + throw new Error(message); +}; + +const readString = (record: JsonRecord, field: string, label = field): string => + typeof record[field] === 'string' ? record[field] : fail(`${label} must be a string.`); + +const readNumber = (record: JsonRecord, field: string, label = field): number => + typeof record[field] === 'number' && Number.isFinite(record[field]) + ? record[field] + : fail(`${label} must be a number.`); + +const readInteger = (record: JsonRecord, field: string, min: number, max = Number.MAX_SAFE_INTEGER): number => { + const value = record[field]; + + return typeof value === 'number' && Number.isInteger(value) && value >= min && value <= max + ? value + : fail(`${field} must be an integer from ${min} to ${max}.`); +}; + +/** Absent and null both read as empty. */ +const readList = (record: JsonRecord, field: string, read: (item: unknown, label: string) => T): T[] => { + const value = record[field] ?? []; + + return Array.isArray(value) + ? value.map((item, index) => read(item, `${field}[${index}]`)) + : fail(`${field} must be a list.`); +}; + +const readObject = (value: unknown, label: string): JsonRecord => + isRecord(value) ? value : fail(`${label} must be an object.`); + +const readCase = (record: JsonRecord): GraphExportCase => { + const unknownField = Object.keys(record).find((field) => !CASE_FIELDS.includes(field)); + + if (unknownField !== undefined) { + fail(`Unknown case field "${unknownField}". Valid fields: ${CASE_FIELDS.join(', ')}.`); + } + + const mode = record.mode; + const components = readObject(record.components ?? {}, 'components'); + const initImage = + record.init_image === undefined || record.init_image === null ? null : readObject(record.init_image, 'init_image'); + + return { + components: Object.fromEntries( + Object.keys(components).map((slot) => [slot, readString(components, slot, `components.${slot}`)]) + ), + control: readList(record, 'control', (item, label) => { + const layer = readObject(item, label); + + return { + begin: readNumber(layer, 'begin', `${label}.begin`), + end: readNumber(layer, 'end', `${label}.end`), + imageName: readString(layer, 'image_name', `${label}.image_name`), + modelKey: readString(layer, 'model_key', `${label}.model_key`), + weight: readNumber(layer, 'weight', `${label}.weight`), + }; + }), + height: readInteger(record, 'height', 1), + id: readString(record, 'id'), + initImage: initImage && { + imageName: readString(initImage, 'image_name', 'init_image.image_name'), + strength: readNumber(initImage, 'strength', 'init_image.strength'), + }, + loras: readList(record, 'loras', (item, label) => { + const lora = readObject(item, label); + + return { key: readString(lora, 'key', `${label}.key`), weight: readNumber(lora, 'weight', `${label}.weight`) }; + }), + mode: MODES.includes(mode as GraphExportMode) + ? (mode as GraphExportMode) + : fail(`mode must be one of ${MODES.join(', ')}.`), + modelKey: readString(record, 'model_key'), + negativePrompt: + record.negative_prompt === null || record.negative_prompt === undefined + ? '' + : readString(record, 'negative_prompt'), + positivePrompt: readString(record, 'positive_prompt'), + referenceImages: readList(record, 'reference_images', (item, label) => + typeof item === 'string' ? item : fail(`${label} must be an image name.`) + ), + seed: readInteger(record, 'seed', 0, SEED_MAX), + settings: readObject(record.settings ?? {}, 'settings'), + steps: readInteger(record, 'steps', 1), + width: readInteger(record, 'width', 1), + }; +}; + +const exportCase = (record: JsonRecord, env: GraphExportEnvironment): GraphExportCaseResult => { + const id = record.id as string; + + try { + const { components, graph, submission } = compileGraphExportCase(readCase(record), env); + const batch: GraphExportBatch = { + data: planGenerateSubmissionBatch(submission).data, + destination: DESTINATION, + graph, + origin: `${ORIGIN_PREFIX}${id}`, + }; + + return { batch, components, id }; + } catch (error) { + return { error: error instanceof Error ? error.message : String(error), id }; + } +}; + +const isModelConfig = (value: unknown): value is GraphExportEnvironment['models'][number] => + isRecord(value) && ['key', 'name', 'base', 'type'].every((field) => typeof value[field] === 'string'); + +const isCapabilitiesRow = (value: unknown): value is ArchitectureCapabilitiesRow => + isRecord(value) && typeof value.base === 'string' && isRecord(value.features); + +const readRequest = (request: unknown) => { + if (!isRecord(request)) { + throw new GraphExportRequestError('The request must be a JSON object.'); + } + + if (request.version !== GRAPH_EXPORT_VERSION) { + throw new GraphExportRequestError(`Unsupported request version ${JSON.stringify(request.version)}; expected 1.`); + } + + const { capabilities, cases, models } = request; + + if (!Array.isArray(models) || !models.every(isModelConfig)) { + throw new GraphExportRequestError('models must be the `models` list of GET /api/v2/models/.'); + } + + if (!Array.isArray(capabilities) || !capabilities.every(isCapabilitiesRow)) { + throw new GraphExportRequestError('capabilities must be the body of GET /api/v2/models/capabilities.'); + } + + if (!Array.isArray(cases) || !cases.every((item) => isRecord(item) && typeof item.id === 'string' && item.id)) { + throw new GraphExportRequestError('cases must be a list of objects with a non-empty string id.'); + } + + const ids = cases.map((item) => item.id as string); + const duplicate = ids.find((id, index) => ids.indexOf(id) !== index); + + if (duplicate !== undefined) { + throw new GraphExportRequestError(`Case id "${duplicate}" is not unique.`); + } + + return { capabilities, cases: cases as JsonRecord[], models }; +}; + +/** Throws GraphExportRequestError for a malformed request; a case that cannot compile carries its error instead. */ +export const exportGraphs = (request: unknown): GraphExportResponse => { + const { capabilities, cases, models } = readRequest(request); + const env: GraphExportEnvironment = { + canvas: { + compositing: DEFAULT_CANVAS_COMPOSITING, + scaling: DEFAULT_CANVAS_SCALING, + strength: DEFAULT_CANVAS_DENOISING_STRENGTH, + }, + destination: DESTINATION, + models, + projectSettings: PROJECT_SETTINGS, + randDevice: resolveRandDeviceMetadata(PROJECT_SETTINGS.useCpuNoise, GPU_DEVICE_OPTIONS), + }; + + return withArchitectureCapabilities(capabilities, () => ({ + cases: cases.map((record) => exportCase(record, env)), + version: GRAPH_EXPORT_VERSION, + })); +}; diff --git a/invokeai/frontend/webv2/src/architecture/dependencyPolicy.test.ts b/invokeai/frontend/webv2/src/architecture/dependencyPolicy.test.ts index 812832d1ce7..60e477c547b 100644 --- a/invokeai/frontend/webv2/src/architecture/dependencyPolicy.test.ts +++ b/invokeai/frontend/webv2/src/architecture/dependencyPolicy.test.ts @@ -235,6 +235,7 @@ describe('feature public-interface registry', () => { ]); expect(checkDependency('workbench/invocationSubmit.ts', '@features/generation/prompts')).toEqual([]); expect(checkDependency('app/GenerateWidgetSyncRuntime.tsx', '@features/generation/runtime')).toEqual([]); + expect(checkDependency('app/graphExport.ts', '@features/generation/graphExport')).toEqual([]); expect(checkDependency('workbench/widgets/canvas/invoke.ts', '@features/generation/canvasGraph')).toEqual([]); expect(checkDependency('workbench/WorkbenchContext.tsx', '@features/intermediates/holdLease')).toEqual([]); expect( diff --git a/invokeai/frontend/webv2/src/architecture/featureInterfaces.ts b/invokeai/frontend/webv2/src/architecture/featureInterfaces.ts index 9405bcfd449..c5fb3a9b00d 100644 --- a/invokeai/frontend/webv2/src/architecture/featureInterfaces.ts +++ b/invokeai/frontend/webv2/src/architecture/featureInterfaces.ts @@ -22,6 +22,7 @@ export const FEATURE_PUBLIC_INTERFACES: Readonly ({ +export const getComponentPolicyContext = ( + model: GenerateModelConfig, + settings: GenerateSettings +): ComponentPolicyContext => ({ model, settings, selectedComponents: { @@ -1185,6 +1188,9 @@ const COMPONENT_SETTING_LABELS: Record = { componentSourceModel: 'Component source', }; +/** Every settings key that holds a component model. */ +export const GENERATE_COMPONENT_VALUE_KEYS = Object.keys(COMPONENT_SETTING_LABELS) as GenerateComponentValueKey[]; + const addClearedLabel = (labels: string[], label: string) => { if (!labels.includes(label)) { labels.push(label); @@ -1438,7 +1444,7 @@ const getSettingsWithCompatibleModelSelections = ( const slotsByKey = new Map(policy.slots.map((slotPolicy) => [slotPolicy.key, slotPolicy])); const uiPolicy = getGenerationUiPolicy(model); - for (const key of Object.keys(COMPONENT_SETTING_LABELS) as GenerateComponentValueKey[]) { + for (const key of GENERATE_COMPONENT_VALUE_KEYS) { if (!nextSettings[key]) { continue; } @@ -1542,7 +1548,7 @@ export const getGenerationModelAvailabilityReasons = ( reasons.push(`Selected model "${model.name}" is no longer installed.`); } - for (const key of Object.keys(COMPONENT_SETTING_LABELS) as GenerateComponentValueKey[]) { + for (const key of GENERATE_COMPONENT_VALUE_KEYS) { const value = settings[key]; if (value && !hasModelKey(models, value.key, value.type)) { diff --git a/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.test.ts b/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.test.ts index 38d2ee38047..f1a513880ed 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.test.ts @@ -6,9 +6,11 @@ import { hasDynamicPromptSyntax, matchesKnownWildcard, normalizeWildcardValues, + planGenerateSubmission, scanWildcardReferences, sanitizeDynamicPromptsConfig, sanitizeMaxPrompts, + type GenerateSubmissionSource, } from './dynamicPrompts'; describe('hasDynamicPromptSyntax', () => { @@ -239,3 +241,70 @@ describe('getWildcardValuesError', () => { expect(getWildcardValuesError(['x'.repeat(2_000)])).toBeNull(); }); }); + +const source = (overrides: Partial = {}): GenerateSubmissionSource => ({ + batchCount: 2, + negativePrompt: 'blurry', + negativePromptEnabled: true, + positivePrompt: 'a cat', + promptTemplate: null, + seed: 10, + seedMode: 'increment', + ...overrides, +}); + +describe('planGenerateSubmission', () => { + it('merges the template and suppresses its negative side with the negative prompt', () => { + const promptTemplate = { + id: 't', + name: 'Photo', + negativePrompt: '{prompt}, lowres', + positivePrompt: 'photo of {prompt}', + }; + + expect(planGenerateSubmission(source({ promptTemplate })).submission).toMatchObject({ + negativePrompt: 'blurry, lowres', + positivePrompt: 'photo of a cat', + }); + + const suppressed = planGenerateSubmission(source({ negativePromptEnabled: false, promptTemplate })); + + expect(suppressed.submission.negativePrompt).toBe(''); + expect(suppressed.effectivePrompts.negativePrompt).toBe('blurry, lowres'); + }); + + it('sends an expansion and its seed behaviour only for a prompt with dynamic syntax', () => { + const expansion = ['a red cat', 'a green cat']; + const expanded = planGenerateSubmission( + source({ dynamicPromptsSeedBehaviour: 'per-image', positivePrompt: 'a {red|green} cat' }), + expansion + ); + const plain = planGenerateSubmission(source({ dynamicPromptsSeedBehaviour: 'per-image' }), expansion); + + expect(expanded.submission).toEqual({ + batchCount: 2, + negativePrompt: 'blurry', + positivePrompt: 'a {red|green} cat', + positivePrompts: expansion, + seed: 10, + seedBehaviour: 'per-image', + seedStep: 1, + }); + // Two prompts per image across two runs consume four seeds. + expect(expanded.seedPlan.nextSeed).toBe(14); + expect(plain.submission).toEqual({ + batchCount: 2, + negativePrompt: 'blurry', + positivePrompt: 'a cat', + seed: 10, + seedStep: 1, + }); + expect(plain.seedPlan.nextSeed).toBe(12); + }); + + it('holds a fixed seed', () => { + const { seedPlan, submission } = planGenerateSubmission(source({ seedMode: 'fixed' })); + + expect([submission.seed, submission.seedStep, seedPlan.nextSeed]).toEqual([10, 0, null]); + }); +}); diff --git a/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.ts b/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.ts index 3512f3c45e7..52967295237 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/dynamicPrompts.ts @@ -1,3 +1,9 @@ +import { planSeedSubmission, type SeedMode, type SeedStep, type SeedSubmissionPlan } from '@platform/core/seed'; + +import type { EffectivePrompts, PromptTemplateSnapshot } from './promptTemplates'; + +import { getEffectivePrompts } from './promptTemplates'; + const SEED_MAX = 4_294_967_295; export const DYNAMIC_PROMPTS_MIN_PROMPTS = 1; @@ -302,3 +308,75 @@ export const sanitizeSampleSeed = (value: unknown): number => typeof value === 'number' && Number.isFinite(value) ? Math.abs(Math.round(value)) % SEED_MAX : 0; export const createDynamicPromptsSampleSeed = (): number => Math.floor(Math.random() * SEED_MAX); + +/** The prompted, seeded values Generate, Canvas, Upscale and Video submit from. */ +export interface GenerateSubmissionSource { + batchCount: number; + negativePrompt: string; + negativePromptEnabled: boolean; + positivePrompt: string; + promptTemplate?: PromptTemplateSnapshot | null; + seed: number; + seedMode: SeedMode; + dynamicPromptsSeedBehaviour?: DynamicPromptsSeedBehaviour; +} + +/** The prompt and seed fields of a queued generate submission; the host adds the graph and its node ids. */ +export interface GenerateSubmissionPrompts { + batchCount: number; + negativePrompt: string; + positivePrompt: string; + /** Present when the positive prompt expanded; the expansion is sent instead of the prompt. */ + positivePrompts?: string[]; + seed: number; + seedBehaviour?: DynamicPromptsSeedBehaviour; + seedStep: SeedStep; +} + +export interface GenerateSubmissionPlan { + /** Template-merged prompts, before negative-prompt suppression. */ + effectivePrompts: EffectivePrompts; + seedPlan: SeedSubmissionPlan; + submission: GenerateSubmissionPrompts; +} + +/** + * The prompts and seeds a submission sends. `expandedPositivePrompts` are what the host expanded the merged positive + * prompt into; they count only when that prompt has dynamic syntax, so templates that introduce syntax count too. + */ +export const planGenerateSubmission = ( + source: GenerateSubmissionSource, + expandedPositivePrompts?: string[] +): GenerateSubmissionPlan => { + const effectivePrompts = getEffectivePrompts(source); + // Record expansions even when only one prompt results. + const positivePrompts = + expandedPositivePrompts && + expandedPositivePrompts.length > 0 && + hasDynamicPromptSyntax(effectivePrompts.positivePrompt) + ? expandedPositivePrompts + : undefined; + const seedBehaviour = positivePrompts ? source.dynamicPromptsSeedBehaviour : undefined; + const seedPlan = planSeedSubmission({ + batchCount: source.batchCount, + promptCount: positivePrompts?.length ?? 1, + seedBehaviour: seedBehaviour ?? 'per-iteration', + seedMode: source.seedMode, + startSeed: source.seed, + }); + + return { + effectivePrompts, + seedPlan, + submission: { + batchCount: source.batchCount, + // Disabling the negative prompt also suppresses the template's negative side. + negativePrompt: source.negativePromptEnabled ? effectivePrompts.negativePrompt : '', + positivePrompt: effectivePrompts.positivePrompt, + ...(positivePrompts ? { positivePrompts } : {}), + seed: source.seed, + ...(seedBehaviour ? { seedBehaviour } : {}), + seedStep: seedPlan.step, + }, + }; +}; diff --git a/invokeai/frontend/webv2/src/features/generation/core/graphExport.ts b/invokeai/frontend/webv2/src/features/generation/core/graphExport.ts new file mode 100644 index 00000000000..72ae821403a --- /dev/null +++ b/invokeai/frontend/webv2/src/features/generation/core/graphExport.ts @@ -0,0 +1,492 @@ +/** + * Offline export cases compile to the graph Generate or Canvas would submit: the model's defaults, the case's values + * and component picks, required components filled from the installed models the way a user fills them, and the same + * reconciliation and compilers the widget runs. Callers own the wire format and the Workbench canvas defaults. + */ + +import type { GenerationModelCatalogItem as ModelConfig } from '@features/generation/core/contracts'; + +import { areJsonValuesStructurallyEqual } from '@platform/core/json'; + +import type { ArchitectureCapabilitiesRow } from './architectureCapabilities'; +import type { ComponentPolicyContext, ComponentSlotPolicy } from './baseGenerationPolicies'; +import type { ControlLayerGraphInput } from './canvas/addControlLayers'; +import type { CanvasCompositingSettings, CanvasScalingSettings } from './canvas/types'; +import type { BackendGraphContract, GenerationProjectSettings, ResultDestination } from './contracts'; +import type { ComponentModelConfig, GenerateModelConfig, GenerateWidgetValues } from './types'; + +import { resetArchitectureCapabilities, setArchitectureCapabilities } from './architectureCapabilities'; +import { + coerceSchedulerForGraph, + GENERATE_COMPONENT_VALUE_KEYS, + getComponentPolicyContext, + getComponentSectionPolicy, + getDefaultGenerateSettings, + getDefaultReferenceImageConfig, + getGenerateModelSelectionResult, + getSchedulerOptions, + isArchitectureDescribed, + isGenerateModelSelectable, + isLoraSupported, +} from './baseGenerationPolicies'; +import { compileCanvasGraph } from './canvas/compileCanvasGraph'; +import { + CONTROL_ADAPTER_KINDS, + isControlKindSupportedForBase, + isControlModelUsableForKind, +} from './canvas/controlValidation'; +import { hasDynamicPromptSyntax, planGenerateSubmission, type GenerateSubmissionPrompts } from './dynamicPrompts'; +import { compileGenerateGraph } from './graph'; +import { stabilizeBackendGraphNodeIds } from './previewGraph'; +import { resolveGenerateWidgetValues } from './resolveGenerateWidgetValues'; +import { + deriveAspectRatioId, + GENERATE_UI_STATE_KEYS, + getModelDefaultVae, + hasModelDefaultVae, + isLoraCompatibleWithModel, + isLoraModelConfig, + isMainModelConfig, + isModelIdentifierConfig, + isVaeModelConfig, +} from './settings'; + +/** `txt2img` compiles like Generate; `img2img` and `controlnet` compile like a full-frame Canvas invoke. */ +export type GraphExportMode = 'txt2img' | 'img2img' | 'controlnet'; + +export interface GraphExportCase { + /** Becomes the graph id. */ + id: string; + mode: GraphExportMode; + modelKey: string; + /** Picks keyed by component slot; required slots left out are filled from the installed models. */ + components: Readonly>; + width: number; + height: number; + steps: number; + seed: number; + positivePrompt: string; + negativePrompt: string; + loras: readonly { key: string; weight: number }[]; + /** Uploaded image names. */ + referenceImages: readonly string[]; + /** An uploaded image already sized width × height, as the Canvas composite always is. */ + initImage: { imageName: string; strength: number } | null; + /** Each image is an uploaded width × height control composite. */ + control: readonly { modelKey: string; imageName: string; weight: number; begin: number; end: number }[]; + /** Generate settings overrides. Keys the case sets through its own fields are rejected. */ + settings: Readonly>; +} + +export interface GraphExportEnvironment { + models: readonly ModelConfig[]; + destination: ResultDestination; + projectSettings: GenerationProjectSettings; + randDevice: string; + /** What the Canvas widget supplies from its persisted values. */ + canvas: { compositing: CanvasCompositingSettings; scaling: CanvasScalingSettings; strength: number }; +} + +export interface CompiledGraphExportCase { + /** Node ids are stabilized: identical cases compile to identical graphs. */ + graph: BackendGraphContract; + /** The key of every component model the case compiled with, by settings key. */ + components: Record; + /** What Generate queues for the graph, with the prompt and seed node ids in `graph`. */ + submission: GenerateSubmissionPrompts & { + negativePromptNodeId: string; + positivePromptNodeId: string; + seedNodeId: string; + }; +} + +/** Keys a case sets through its own fields, picks as components, or that need backend catalogs or expansion. */ +const NON_OVERRIDABLE_SETTINGS: ReadonlySet = new Set([ + 'aspectRatioId', + 'aspectRatioValue', + 'dynamicPromptsCombinatorial', + 'dynamicPromptsMaxPrompts', + 'dynamicPromptsResample', + 'dynamicPromptsSampleSeed', + 'dynamicPromptsSeedBehaviour', + 'height', + 'loras', + 'model', + 'modelKey', + 'negativePrompt', + 'positivePrompt', + 'promptTemplate', + 'referenceImages', + 'seed', + 'seedMode', + 'steps', + 'width', + ...GENERATE_COMPONENT_VALUE_KEYS, + ...Object.keys(GENERATE_UI_STATE_KEYS), +]); + +const listKeys = (keys: readonly string[]): string => (keys.length > 0 ? keys.join(', ') : 'none'); + +const listModelKeys = (models: readonly ModelConfig[]): string => listKeys(models.map((model) => model.key)); + +/** Hold `rows` as the capability table, which every Generation policy reads, while `run` compiles. */ +export const withArchitectureCapabilities = (rows: readonly ArchitectureCapabilitiesRow[], run: () => T): T => { + setArchitectureCapabilities(rows); + + try { + return run(); + } finally { + resetArchitectureCapabilities(); + } +}; + +const getSelectableModel = (key: string, models: readonly ModelConfig[]): GenerateModelConfig => { + const selectable = models.filter(isGenerateModelSelectable).filter(isArchitectureDescribed); + const model = selectable.find((candidate) => candidate.key === key); + + if (!model) { + throw new Error(`model_key "${key}" is not a selectable main model. Choices: ${listModelKeys(selectable)}.`); + } + + return model; +}; + +/** Selecting the model, then "Reset to model defaults", which also applies the model's default VAE. */ +const getModelDefaultValues = (model: GenerateModelConfig, models: readonly ModelConfig[]): GenerateWidgetValues => { + const { settings } = getGenerateModelSelectionResult({ + currentValues: getDefaultGenerateSettings(model), + model, + models, + }); + const vae = hasModelDefaultVae(model) ? getModelDefaultVae(model, models.filter(isVaeModelConfig)) : settings.vae; + + return { ...settings, model, vae }; +}; + +const withCaseValues = (values: GenerateWidgetValues, input: GraphExportCase): GenerateWidgetValues => ({ + ...values, + aspectRatioId: deriveAspectRatioId(input.width, input.height), + aspectRatioValue: input.width / input.height, + height: input.height, + negativePrompt: input.negativePrompt, + positivePrompt: input.positivePrompt, + seed: input.seed, + seedMode: 'fixed', + steps: input.steps, + width: input.width, +}); + +const withSettingsOverrides = ( + values: GenerateWidgetValues, + overrides: GraphExportCase['settings'] +): GenerateWidgetValues => { + const validKeys = Object.keys(values) + .filter((key) => !NON_OVERRIDABLE_SETTINGS.has(key)) + .sort(); + const unknownKey = Object.keys(overrides).find((key) => !validKeys.includes(key)); + + if (unknownKey !== undefined) { + throw new Error(`Unknown settings key "${unknownKey}". Valid keys: ${listKeys(validKeys)}.`); + } + + return { ...values, ...overrides }; +}; + +const fitsSlot = (slot: ComponentSlotPolicy, candidate: ModelConfig, ctx: ComponentPolicyContext): boolean => + (slot.modelTypes as readonly string[]).includes(candidate.type) && (!slot.filter || slot.filter(candidate, ctx)); + +/** The value the component picker commits for `model`, by the slot's value kind. */ +const toSlotValue = (slot: ComponentSlotPolicy, model: ModelConfig): ComponentModelConfig | null => { + switch (slot.valueKind) { + case 'main': + return isMainModelConfig(model) ? model : null; + case 'vae': + return isVaeModelConfig(model) ? model : null; + case 'component': + return isModelIdentifierConfig(model) ? model : null; + } +}; + +const withComponentPicks = ( + values: GenerateWidgetValues, + picks: GraphExportCase['components'], + models: readonly ModelConfig[] +): GenerateWidgetValues => { + const { slots } = getComponentSectionPolicy(values.model, values); + let next = values; + + for (const [slotKey, modelKey] of Object.entries(picks)) { + const slot = slots.find((candidate) => candidate.key === slotKey); + + if (!slot) { + const validSlots = listKeys(slots.map((candidate) => candidate.key)); + + throw new Error(`Unknown component slot "${slotKey}" for ${values.model.name}. Valid slots: ${validSlots}.`); + } + + const ctx = getComponentPolicyContext(next.model, next); + const fitting = models.filter((candidate) => fitsSlot(slot, candidate, ctx)); + const pick = fitting.find((candidate) => candidate.key === modelKey); + const value = pick ? toSlotValue(slot, pick) : null; + + if (!value) { + const problem = models.some((candidate) => candidate.key === modelKey) + ? `"${modelKey}" cannot be the ${slot.label} of ${values.model.name}` + : `model key "${modelKey}" is not installed`; + + throw new Error(`components.${slotKey}: ${problem}. Choices: ${listModelKeys(fitting)}.`); + } + + next = { ...next, [slot.key]: value }; + } + + return next; +}; + +/** Pick the first installed model for the first required slot that is unfilled; null when none can be filled. */ +const withFirstRequiredComponent = ( + values: GenerateWidgetValues, + models: readonly ModelConfig[] +): GenerateWidgetValues | null => { + const ctx = getComponentPolicyContext(values.model, values); + + for (const slot of getComponentSectionPolicy(values.model, values).slots) { + const current = ctx.selectedComponents[slot.key]; + + if (!slot.required?.(ctx) || (current && (!slot.filter || slot.filter(current as ModelConfig, ctx)))) { + continue; + } + + for (const candidate of models) { + const value = fitsSlot(slot, candidate, ctx) ? toSlotValue(slot, candidate) : null; + + if (value) { + return { ...values, [slot.key]: value }; + } + } + } + + return null; +}; + +const withLoras = ( + values: GenerateWidgetValues, + loras: GraphExportCase['loras'], + models: readonly ModelConfig[] +): GenerateWidgetValues => { + if (loras.length === 0) { + return values; + } + + const { model } = values; + + if (!isLoraSupported(model)) { + throw new Error(`${model.name} does not load LoRAs.`); + } + + const compatible = models.filter(isLoraModelConfig).filter((lora) => isLoraCompatibleWithModel(lora, model)); + + return { + ...values, + loras: loras.map(({ key, weight }, index) => { + const lora = compatible.find((candidate) => candidate.key === key); + + if (!lora) { + const choices = listModelKeys(compatible); + + throw new Error(`loras[${index}]: "${key}" is not an installed LoRA for ${model.name}. Choices: ${choices}.`); + } + + return { isEnabled: true, model: lora, weight }; + }), + }; +}; + +const withReferenceImages = ( + values: GenerateWidgetValues, + input: GraphExportCase, + models: readonly ModelConfig[] +): GenerateWidgetValues => ({ + ...values, + referenceImages: input.referenceImages.map((imageName, index) => ({ + // Cases name images only; the output size stands in for dimensions that reach metadata, not inference. + config: getDefaultReferenceImageConfig(values.model, models, { + original: { image: { height: input.height, image_name: imageName, width: input.width } }, + }), + id: `reference_image_${index + 1}`, + isEnabled: true, + })), +}); + +/** The widget reconciles after every commit; reconciliation can pick a FLUX.2 component source. */ +const reconcile = (values: GenerateWidgetValues, models: readonly ModelConfig[]): GenerateWidgetValues => { + const resolved = resolveGenerateWidgetValues({ models, storedValues: values }); + + if (!resolved) { + throw new Error('Generate cannot resolve settings without the capability table.'); + } + + return resolved.values; +}; + +/** Report a requested value the widget would not keep instead of compiling a graph nobody asked for. */ +const assertRequestedValuesKept = (values: GenerateWidgetValues, input: GraphExportCase): void => { + for (const [key, requested] of Object.entries(input.settings)) { + const kept = values[key as keyof GenerateWidgetValues]; + + if (!areJsonValuesStructurallyEqual(kept, requested)) { + throw new Error( + `settings.${key}: Generate does not keep ${JSON.stringify(requested)}; it holds ${JSON.stringify(kept)}.` + ); + } + } + + if ('scheduler' in input.settings && coerceSchedulerForGraph(values.model, values.scheduler) !== values.scheduler) { + const options = getSchedulerOptions(values.model).map((option) => option.value); + + throw new Error( + `settings.scheduler: ${values.model.name} does not use "${values.scheduler}". Options: ${listKeys(options)}.` + ); + } + + for (const [key, requested] of Object.entries(input.components)) { + const kept = values[key as keyof GenerateWidgetValues] as ComponentModelConfig | null; + + if (kept?.key !== requested) { + throw new Error( + `components.${key}: Generate replaces "${requested}" with ${kept ? `"${kept.key}"` : 'nothing'}.` + ); + } + } +}; + +const getControlLayers = ( + input: GraphExportCase, + model: GenerateModelConfig, + models: readonly ModelConfig[] +): ControlLayerGraphInput[] => { + // A layer's kind is the one its model can drive on this base, the kind a user switches the layer to. + const getKind = (candidate: ModelConfig) => + candidate.base === model.base + ? CONTROL_ADAPTER_KINDS.find( + (kind) => isControlKindSupportedForBase(model.base, kind) && isControlModelUsableForKind(candidate, kind) + ) + : undefined; + + return input.control.map((layer, index) => { + const adapter = models.find((candidate) => candidate.key === layer.modelKey); + const kind = adapter ? getKind(adapter) : undefined; + + if (!adapter || !kind) { + const choices = listModelKeys(models.filter((candidate) => getKind(candidate) !== undefined)); + const problem = `"${layer.modelKey}" cannot drive a control layer of ${model.name}`; + + throw new Error(`control[${index}]: ${problem}. Choices: ${choices}.`); + } + + return { + beginEndStepPct: [layer.begin, layer.end], + // A new ControlNet layer starts "balanced", which the compiler sends for null; other kinds carry no mode. + controlMode: null, + id: `control_${index + 1}`, + imageName: layer.imageName, + kind, + model: { + base: adapter.base, + key: adapter.key, + name: adapter.name, + type: adapter.type, + ...(typeof adapter.hash === 'string' ? { hash: adapter.hash } : {}), + }, + modelCondInChannels: typeof adapter.cond_in_channels === 'number' ? adapter.cond_in_channels : null, + weight: layer.weight, + }; + }); +}; + +const assertModeInputs = (input: GraphExportCase): void => { + if (input.mode === 'txt2img' && (input.initImage || input.control.length > 0)) { + throw new Error('txt2img cases take no init_image or control; use img2img or controlnet.'); + } + + if (input.mode === 'img2img' && !input.initImage) { + throw new Error('img2img cases need init_image.'); + } + + if (input.mode === 'controlnet' && input.control.length === 0) { + throw new Error('controlnet cases need at least one control entry.'); + } +}; + +/** Throws a message naming the offending field and its valid choices when the case cannot compile. */ +export const compileGraphExportCase = ( + input: GraphExportCase, + env: GraphExportEnvironment +): CompiledGraphExportCase => { + assertModeInputs(input); + + const { models } = env; + let values = withCaseValues(getModelDefaultValues(getSelectableModel(input.modelKey, models), models), input); + values = withSettingsOverrides(values, input.settings); + values = withComponentPicks(values, input.components, models); + values = withLoras(values, input.loras, models); + values = reconcile(withReferenceImages(values, input, models), models); + + // Filling one slot can satisfy or require another; the bound stops a reconciliation that keeps undoing a pick. + for (let pass = 0; pass < GENERATE_COMPONENT_VALUE_KEYS.length; pass += 1) { + const filled = withFirstRequiredComponent(values, models); + + if (!filled) { + break; + } + + values = reconcile(filled, models); + } + + assertRequestedValuesKept(values, input); + + // No expansion: the UI expands dynamic prompts through the backend before it queues. + const { effectivePrompts, submission } = planGenerateSubmission(values); + + if (hasDynamicPromptSyntax(effectivePrompts.positivePrompt)) { + throw new Error( + 'positive_prompt has dynamic prompt syntax, which only the backend expands; export each expansion.' + ); + } + + const compiled = + input.mode === 'txt2img' + ? compileGenerateGraph(values, values.model, env.destination, env.projectSettings, env.randDevice) + : compileCanvasGraph({ + bbox: { height: input.height, width: input.width, x: 0, y: 0 }, + compositeImageName: input.initImage?.imageName ?? null, + compositing: env.canvas.compositing, + controlLayers: getControlLayers(input, values.model, models), + destination: env.destination, + mode: input.initImage ? 'img2img' : 'txt2img', + model: values.model, + projectSettings: env.projectSettings, + randDevice: env.randDevice, + scaling: env.canvas.scaling, + settings: values, + strength: input.initImage?.strength ?? env.canvas.strength, + }); + + const stable = stabilizeBackendGraphNodeIds(compiled.backendGraph, { + negativePromptNodeId: compiled.negativePromptNodeId, + positivePromptNodeId: compiled.positivePromptNodeId, + seedNodeId: compiled.seedNodeId, + }); + + return { + components: Object.fromEntries( + GENERATE_COMPONENT_VALUE_KEYS.flatMap((key) => { + const component = values[key]; + + return component ? [[key, component.key]] : []; + }) + ), + graph: { ...stable.graph, id: input.id }, + submission: { ...submission, ...stable.nodeIds }, + }; +}; diff --git a/invokeai/frontend/webv2/src/features/generation/core/previewGraph.test.ts b/invokeai/frontend/webv2/src/features/generation/core/previewGraph.test.ts index a9f8d2c127f..6e22c106b51 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/previewGraph.test.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/previewGraph.test.ts @@ -10,7 +10,7 @@ import { setArchitectureCapabilities, } from './architectureCapabilities'; import { getDefaultGenerateSettings } from './baseGenerationPolicies'; -import { compileGeneratePreviewGraph, stabilizeBackendGraphIds } from './previewGraph'; +import { compileGeneratePreviewGraph, stabilizeBackendGraphIds, stabilizeBackendGraphNodeIds } from './previewGraph'; const sdxlModel: MainModelConfig = { base: 'sdxl', key: 'sdxl-model', name: 'SDXL', type: 'main' }; @@ -171,4 +171,21 @@ describe('stabilizeBackendGraphIds', () => { source: { field: 'value', node_id: 'seed' }, }); }); + + it('follows fixed node ids that lose their name to an earlier created id', () => { + const graph = { + edges: [], + id: 'graph', + nodes: { + seed_lqx1abc_a1b2c3: { id: 'seed_lqx1abc_a1b2c3', type: 'noise' }, + seed: { id: 'seed', type: 'integer' }, + positive_prompt: { id: 'positive_prompt', type: 'string' }, + }, + }; + + const result = stabilizeBackendGraphNodeIds(graph, { positive: 'positive_prompt', seed: 'seed' }); + + expect(result.nodeIds).toEqual({ positive: 'positive_prompt', seed: 'seed_2' }); + expect(result.graph.nodes[result.nodeIds.seed]?.type).toBe('integer'); + }); }); diff --git a/invokeai/frontend/webv2/src/features/generation/core/previewGraph.ts b/invokeai/frontend/webv2/src/features/generation/core/previewGraph.ts index 105177ac794..6d2929959ed 100644 --- a/invokeai/frontend/webv2/src/features/generation/core/previewGraph.ts +++ b/invokeai/frontend/webv2/src/features/generation/core/previewGraph.ts @@ -25,7 +25,11 @@ export type GeneratePreviewResult = const CREATED_ID_PATTERN = /^(.+)_[a-z0-9]{7,9}_[a-z0-9]{4,6}$/; -export const stabilizeBackendGraphIds = (graph: BackendGraphContract): BackendGraphContract => { +/** Stabilize ids and follow `nodeIds` through the renames: a fixed id loses its name to an earlier created one. */ +export const stabilizeBackendGraphNodeIds = >( + graph: BackendGraphContract, + nodeIds: T +): { graph: BackendGraphContract; nodeIds: T } => { const rename = new Map(); const used = new Set(); @@ -44,26 +48,27 @@ export const stabilizeBackendGraphIds = (graph: BackendGraphContract): BackendGr rename.set(id, candidate); } + const renamed = (id: string): string => rename.get(id) ?? id; + return { - edges: graph.edges.map((edge) => ({ - destination: { - field: edge.destination.field, - node_id: rename.get(edge.destination.node_id) ?? edge.destination.node_id, - }, - source: { field: edge.source.field, node_id: rename.get(edge.source.node_id) ?? edge.source.node_id }, - ...(edge.type ? { type: edge.type } : {}), - })), - id: 'generate-preview', - nodes: Object.fromEntries( - Object.entries(graph.nodes).map(([id, node]) => { - const nextId = rename.get(id) ?? id; - - return [nextId, { ...node, id: nextId }]; - }) - ), + graph: { + edges: graph.edges.map((edge) => ({ + destination: { field: edge.destination.field, node_id: renamed(edge.destination.node_id) }, + source: { field: edge.source.field, node_id: renamed(edge.source.node_id) }, + ...(edge.type ? { type: edge.type } : {}), + })), + id: 'generate-preview', + nodes: Object.fromEntries( + Object.entries(graph.nodes).map(([id, node]) => [renamed(id), { ...node, id: renamed(id) }]) + ), + }, + nodeIds: Object.fromEntries(Object.entries(nodeIds).map(([key, id]) => [key, renamed(id)])) as T, }; }; +export const stabilizeBackendGraphIds = (graph: BackendGraphContract): BackendGraphContract => + stabilizeBackendGraphNodeIds(graph, {}).graph; + const toPreviewContract = (backendGraph: BackendGraphContract, label: string): GraphContract => ({ backendGraph, edges: backendGraph.edges.map((edge, index) => ({ diff --git a/invokeai/frontend/webv2/src/features/generation/graphExport.ts b/invokeai/frontend/webv2/src/features/generation/graphExport.ts new file mode 100644 index 00000000000..2d1915f7a74 --- /dev/null +++ b/invokeai/frontend/webv2/src/features/generation/graphExport.ts @@ -0,0 +1,10 @@ +/** Offline graph export for tooling such as `pnpm graphs:export`; the browser bundle never loads it. */ +export { + compileGraphExportCase, + withArchitectureCapabilities, + type CompiledGraphExportCase, + type GraphExportCase, + type GraphExportEnvironment, + type GraphExportMode, +} from './core/graphExport'; +export type { ArchitectureCapabilitiesRow } from './core/architectureCapabilities'; diff --git a/invokeai/frontend/webv2/src/features/generation/settings.ts b/invokeai/frontend/webv2/src/features/generation/settings.ts index 70160f0d43b..25c6c33a3fc 100644 --- a/invokeai/frontend/webv2/src/features/generation/settings.ts +++ b/invokeai/frontend/webv2/src/features/generation/settings.ts @@ -84,6 +84,7 @@ export { DYNAMIC_PROMPTS_MIN_PROMPTS, hasDynamicPromptSyntax, isDynamicPromptsSeedBehaviour, + planGenerateSubmission, sanitizeDynamicPromptsConfig, sanitizeMaxPrompts, sanitizeSampleSeed, diff --git a/invokeai/frontend/webv2/src/features/queue/contracts.ts b/invokeai/frontend/webv2/src/features/queue/contracts.ts index e858b6d8b08..828667aed54 100644 --- a/invokeai/frontend/webv2/src/features/queue/contracts.ts +++ b/invokeai/frontend/webv2/src/features/queue/contracts.ts @@ -57,6 +57,11 @@ export { shouldSubmitPendingQueueItem, } from './core/submissionRules'; export { extractGenerationMeta, getResultImageName, type QueueGenerationMeta } from './core/generationMeta'; +export { + planGenerateSubmissionBatch, + type GeneratePromptBatchPlan, + type GeneratePromptSubmission, +} from './core/promptBatch'; export { buildProjectQueueItemOriginPrefix, buildQueueItemOrigin, diff --git a/invokeai/frontend/webv2/src/features/queue/core/promptBatch.ts b/invokeai/frontend/webv2/src/features/queue/core/promptBatch.ts index 1a42d04dfbb..36339aac9bd 100644 --- a/invokeai/frontend/webv2/src/features/queue/core/promptBatch.ts +++ b/invokeai/frontend/webv2/src/features/queue/core/promptBatch.ts @@ -339,3 +339,37 @@ export const buildGeneratePromptBatchPlan = ({ runs: isHeld ? iterations : 1, }; }; + +/** The prompt and seed fields of a generate submission, from compilation to `enqueue_batch`. */ +export interface GeneratePromptSubmission { + batchCount: number; + negativePrompt: string; + negativePromptNodeId: string; + positivePrompt: string; + positivePromptNodeId: string; + /** + * Prompts to submit as a batch dimension, already expanded by the caller. + * Absent or empty submits `positivePrompt` verbatim. + */ + positivePrompts?: string[]; + seed: number; + seedBehaviour?: QueuePromptSeedBehaviour; + seedNodeId: string; + seedStep: QueueSeedStep; + /** Legacy recovery maps the old random toggle to seedStep and preserves its expansion rules. */ + legacySeedPlan?: true; +} + +/** The batch data and run count `enqueue_batch` receives for a generate submission. */ +export const planGenerateSubmissionBatch = (submission: GeneratePromptSubmission): GeneratePromptBatchPlan => + (submission.legacySeedPlan ? buildLegacyGeneratePromptBatchPlan : buildGeneratePromptBatchPlan)({ + batchCount: sanitizeBatchCount(submission.batchCount), + negativePrompt: submission.negativePrompt, + negativePromptNodeId: submission.negativePromptNodeId, + positivePromptNodeId: submission.positivePromptNodeId, + prompts: submission.positivePrompts?.length ? submission.positivePrompts : [submission.positivePrompt], + seed: submission.seed, + seedBehaviour: submission.seedBehaviour ?? 'per-iteration', + seedNodeId: submission.seedNodeId, + seedStep: submission.seedStep, + }); diff --git a/invokeai/frontend/webv2/src/features/queue/core/types.ts b/invokeai/frontend/webv2/src/features/queue/core/types.ts index 685e0071531..f043ceabfc7 100644 --- a/invokeai/frontend/webv2/src/features/queue/core/types.ts +++ b/invokeai/frontend/webv2/src/features/queue/core/types.ts @@ -1,4 +1,5 @@ import type { + GeneratePromptSubmission, QueuePromptSeedBehaviour, QueueSeedStep, QueueWorkflowBatchDatum, @@ -56,23 +57,7 @@ export interface QueueEnqueueWorkflowRequest extends QueueEnqueueRequestBase { batchData?: QueueWorkflowBatchDatum[][]; } -export interface QueueEnqueueGenerateRequest extends QueueEnqueueRequestBase { - negativePrompt: string; - negativePromptNodeId: string; - positivePrompt: string; - positivePromptNodeId: string; - /** - * Prompts to submit as a batch dimension, already expanded by the caller. - * Absent or empty submits `positivePrompt` verbatim. - */ - positivePrompts?: string[]; - seed: number; - seedBehaviour?: QueuePromptSeedBehaviour; - seedNodeId: string; - seedStep: QueueSeedStep; - /** Legacy recovery maps the old random toggle to seedStep and preserves its expansion rules. */ - legacySeedPlan?: true; -} +export interface QueueEnqueueGenerateRequest extends QueueEnqueueRequestBase, GeneratePromptSubmission {} export interface QueueEnqueueResult { batchId?: string; diff --git a/invokeai/frontend/webv2/src/features/queue/data/submissionApi.ts b/invokeai/frontend/webv2/src/features/queue/data/submissionApi.ts index 22a081e8a4e..2ecb618e3b4 100644 --- a/invokeai/frontend/webv2/src/features/queue/data/submissionApi.ts +++ b/invokeai/frontend/webv2/src/features/queue/data/submissionApi.ts @@ -8,12 +8,7 @@ import type { QueueResultVideoOptions, } from '@features/queue/core/types'; -import { - buildGeneratePromptBatchPlan, - buildLegacyGeneratePromptBatchPlan, - buildQueueWorkflowBatchPlan, - sanitizeBatchCount, -} from '@features/queue/core/promptBatch'; +import { buildQueueWorkflowBatchPlan, planGenerateSubmissionBatch } from '@features/queue/core/promptBatch'; import { mapWithConcurrency } from '@platform/core/concurrency'; import { addOutputImageNames } from '@platform/core/outputImages'; import { assertAccountScopeCurrent, captureAccountScope } from '@platform/state/accountLifecycle'; @@ -78,20 +73,7 @@ const mapEnqueueResult = (value: unknown, isReceipt = false): QueueEnqueueResult }; export const enqueueGenerate = async (request: QueueEnqueueGenerateRequest): Promise => { - const planInput = { - batchCount: sanitizeBatchCount(request.batchCount), - negativePrompt: request.negativePrompt, - negativePromptNodeId: request.negativePromptNodeId, - positivePromptNodeId: request.positivePromptNodeId, - prompts: request.positivePrompts?.length ? request.positivePrompts : [request.positivePrompt], - seed: request.seed, - seedBehaviour: request.seedBehaviour ?? 'per-iteration', - seedNodeId: request.seedNodeId, - }; - const plan = (request.legacySeedPlan ? buildLegacyGeneratePromptBatchPlan : buildGeneratePromptBatchPlan)({ - ...planInput, - seedStep: request.seedStep, - }); + const plan = planGenerateSubmissionBatch(request); const result = await apiFetchJson('/api/v1/queue/default/enqueue_batch', { body: JSON.stringify({ batch: { diff --git a/invokeai/frontend/webv2/src/workbench/workbenchState.ts b/invokeai/frontend/webv2/src/workbench/workbenchState.ts index 4a233f8afb0..b105ae91709 100644 --- a/invokeai/frontend/webv2/src/workbench/workbenchState.ts +++ b/invokeai/frontend/webv2/src/workbench/workbenchState.ts @@ -84,7 +84,6 @@ import { type GeneratedVideoContract, generatedVideoToGalleryItem, } from '@features/gallery/contracts'; -import { planSeedSubmission } from '@platform/core/seed'; import { describeError } from '@platform/logging/normalize'; import { WIDGET_REGIONS } from '@workbench/layoutContracts'; import { prependProjectEvent, PROJECT_EVENT_LIMIT } from '@workbench/projectEvents'; @@ -138,14 +137,13 @@ import { applyProjectPromptDraft, cloneGenerateWidgetValues, createDynamicPromptsSampleSeed, - getEffectivePrompts, getGenerationModelAvailabilityReasons, getPromptDraftFromValues, getPromptHistoryItemFromGenerateSettings, - hasDynamicPromptSyntax, migrateProjectPromptDraft, normalizeGenerateSettings, normalizeGenerateWidgetValues, + planGenerateSubmission, removePromptHistoryItem, sanitizeBatchCount, syncGenerateWidgetValuesWithModels, @@ -3466,36 +3464,21 @@ const enqueueCompiledSnapshot = ( : route.sourceId === 'video' ? videoSettings : null; - // Resolve template-merged prompts here, where Generate and Canvas submissions converge; Upscale merges to - // identity. - const effectivePrompts = sourceGenerateSettings ? getEffectivePrompts(sourceGenerateSettings) : null; - // Record expansions for Generate routes even when only one prompt results. Inspect merged prompts because - // templates may introduce dynamic syntax. - const expandedPositivePrompts = - route.sourceId !== 'upscale' && - route.sourceId !== 'video' && - compiled.positivePrompts && - compiled.positivePrompts.length > 0 && - effectivePrompts && - hasDynamicPromptSyntax(effectivePrompts.positivePrompt) - ? compiled.positivePrompts - : undefined; - const expandedSeedBehaviour = expandedPositivePrompts - ? (canvasGenerateSettings ?? generateSettings)?.dynamicPromptsSeedBehaviour - : undefined; + // Resolve template-merged prompts and seeds here, where Generate and Canvas submissions converge; Upscale and Video + // merge to identity and never expand. + const submissionPlan = sourceGenerateSettings + ? planGenerateSubmission( + sourceGenerateSettings, + route.sourceId === 'upscale' || route.sourceId === 'video' ? undefined : compiled.positivePrompts + ) + : null; + const effectivePrompts = submissionPlan?.effectivePrompts ?? null; + const expandedPositivePrompts = submissionPlan?.submission.positivePrompts; // The draw the expansion used. Generate expands before dispatch, so the project's sample seed may have rotated // since; recall and rotation follow the expansion, not the project. const expansionSampleSeed = expandedPositivePrompts ? (compiled.positivePromptsSampleSeed ?? null) : null; // Compilation already reserved the starting seed; failures and cancellations do not roll the sequence back. - const seedPlan = sourceGenerateSettings - ? planSeedSubmission({ - batchCount: sourceGenerateSettings.batchCount, - promptCount: expandedPositivePrompts?.length ?? 1, - seedBehaviour: expandedSeedBehaviour ?? 'per-iteration', - seedMode: sourceGenerateSettings.seedMode, - startSeed: sourceGenerateSettings.seed, - }) - : null; + const seedPlan = submissionPlan?.seedPlan ?? null; const backendSubmission: QueueCompiledSubmission = !backendGraph ? { error: `${route.sourceId} queue item is missing a compiled backend graph.`, kind: 'invalid' } : route.sourceId === 'workflow' @@ -3512,21 +3495,14 @@ const enqueueCompiledSnapshot = ( // active workflow changes. ...(compiled.projectWorkflowId ? { projectWorkflowId: compiled.projectWorkflowId } : {}), } - : sourceGenerateSettings && effectivePrompts + : submissionPlan ? { - batchCount: sourceGenerateSettings.batchCount, + ...submissionPlan.submission, graph: backendGraph, kind: 'generate', - // Disabling the negative prompt also suppresses the template's negative side. - negativePrompt: sourceGenerateSettings.negativePromptEnabled ? effectivePrompts.negativePrompt : '', negativePromptNodeId: generate?.negativePromptNodeId ?? 'negative_prompt', - positivePrompt: effectivePrompts.positivePrompt, positivePromptNodeId: generate?.positivePromptNodeId ?? 'positive_prompt', - ...(expandedPositivePrompts ? { positivePrompts: expandedPositivePrompts } : {}), - seed: sourceGenerateSettings.seed, - ...(expandedSeedBehaviour ? { seedBehaviour: expandedSeedBehaviour } : {}), seedNodeId: generate?.seedNodeId ?? 'seed', - seedStep: seedPlan?.step ?? 0, } : { error: `${route.sourceId} queue item is missing source submission metadata.`, kind: 'invalid' }; // A workflow compiles every node's board in (Auto included), so its board-less results stay Uncategorized. diff --git a/scripts/family_e2e/__init__.py b/scripts/family_e2e/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/scripts/family_e2e/__main__.py b/scripts/family_e2e/__main__.py new file mode 100644 index 00000000000..ddbf7a7040b --- /dev/null +++ b/scripts/family_e2e/__main__.py @@ -0,0 +1,245 @@ +"""Command line of the family E2E tool: `python -m scripts.family_e2e `. + +See docs/src/content/docs/development/Guides/family-e2e.mdx. +""" + +from __future__ import annotations + +import argparse +import logging +import os +import re +import sys +from datetime import datetime, timezone +from pathlib import Path + +from scripts.family_e2e.detach import RUNNER_LOG, DetachError, spawn_runner +from scripts.family_e2e.gpu_memory import UnsupportedPlatformError +from scripts.family_e2e.graphs import ExportError +from scripts.family_e2e.manifest import Manifest, ManifestError, load_manifest +from scripts.family_e2e.records import RUN_INFO_FILE, load_records, load_run_info, write_run_info +from scripts.family_e2e.references import AcceptError, accept +from scripts.family_e2e.report import write_report +from scripts.family_e2e.runner import ( + RunnerError, + gpu_lock, + is_alive, + kill_identified, + load_run_manifest, + new_run, + now_iso, + open_gpu, + preflight, + read_heartbeat, + remaining_steps, + remove_scratch_roots, + run, +) +from scripts.family_e2e.server import ServerError + +logger = logging.getLogger("family_e2e") +RUN_ID = re.compile(r"^\d{8}-\d{6}-(?P.+)$") +EXPECTED_ERRORS = ( + ManifestError, + RunnerError, + ServerError, + ExportError, + AcceptError, + DetachError, + UnsupportedPlatformError, + # load_records names the malformed line of a damaged records.jsonl. + ValueError, +) + + +class UsageError(ValueError): + pass + + +def _data_dir(run_dir: Path) -> Path: + # /runs/ + return run_dir.parent.parent + + +def _resolve_run(arg: str | None, manifest_path: Path | None) -> Path: + """A run directory, a run id of the manifest's data dir, or (no argument) the manifest profile's latest run.""" + if arg is not None and Path(arg).is_dir(): + return Path(arg).resolve() + if manifest_path is None: + raise UsageError("pass a run directory, or --manifest to look up run ids and the latest run") + manifest = load_manifest(manifest_path) + runs_dir = manifest.data_dir / "runs" + if arg is not None: + if not (runs_dir / arg / RUN_INFO_FILE).is_file(): + raise UsageError(f"no run {arg} in {runs_dir}") + return runs_dir / arg + runs = sorted( + d.name + for d in runs_dir.glob("*") + if (match := RUN_ID.match(d.name)) and match["profile"] == manifest.profile and (d / RUN_INFO_FILE).is_file() + ) + if not runs: + raise UsageError(f"no runs of profile {manifest.profile} in {runs_dir}") + return runs_dir / runs[-1] + + +def _check_startable(manifest: Manifest) -> None: + gpu = open_gpu(manifest.gpu) + with gpu_lock(gpu.adapter): + preflight(manifest, gpu) + + +def _print_started(run_dir: Path, pid: int) -> None: + print(f"run {run_dir.name}") + print(f"dir {run_dir}") + print(f"runner pid {pid}, log {run_dir / RUNNER_LOG}") + print(f"report {run_dir / 'report.html'} (when finished)") + print(f"status python -m scripts.family_e2e status {run_dir}") + + +def cmd_start(args: argparse.Namespace) -> int: + _check_startable(load_manifest(args.manifest)) + run_dir = new_run(args.manifest, args.family, args.case) + _print_started(run_dir, spawn_runner(run_dir)) + return 0 + + +def cmd_run(args: argparse.Namespace) -> int: + try: + return run(args.run_dir.resolve()) + except KeyboardInterrupt: + return 130 + except Exception: + logger.exception("the runner failed") + return 2 + + +def cmd_resume(args: argparse.Namespace) -> int: + run_dir = _resolve_run(args.run, args.manifest) + if is_alive(read_heartbeat(run_dir).get("runner")): + raise UsageError(f"the runner of {run_dir.name} is still running") + info = load_run_info(run_dir) + _check_startable(load_run_manifest(info)) + write_run_info(run_dir, info.model_copy(update={"status": "running", "finished_at": None})) + _print_started(run_dir, spawn_runner(run_dir)) + return 0 + + +def cmd_status(args: argparse.Namespace) -> int: + run_dir = _resolve_run(args.run, args.manifest) + info = load_run_info(run_dir) + records = load_records(run_dir) + left = remaining_steps(info.plan, records) + heartbeat = read_heartbeat(run_dir) + alive = is_alive(heartbeat.get("runner")) + print(f"run {info.run_id} ({info.status})") + print(f"dir {run_dir}") + print(f"progress {len(info.plan) - len(left)}/{len(info.plan)} steps final, {len(records)} records") + print(f"invalid {sum(1 for r in records if r.invalid_reasons)} records with invalid reasons") + print(f"errors {sum(1 for r in records if r.status != 'completed')} records not completed") + if heartbeat: + updated = datetime.fromisoformat(heartbeat["updated_at"]) + age = (datetime.now(timezone.utc) - updated).total_seconds() + state = "alive" if alive else "gone" + print(f"runner {state}, step: {heartbeat.get('step')}, heartbeat {age:.0f}s ago") + print(f"log {run_dir / RUNNER_LOG}") + if (run_dir / "report.html").is_file(): + print(f"report {run_dir / 'report.html'}") + if info.status == "running" and not alive: + print(f"the runner is gone; continue with: python -m scripts.family_e2e resume {run_dir}") + return 0 + + +def cmd_stop(args: argparse.Namespace) -> int: + run_dir = _resolve_run(args.run, args.manifest) + heartbeat = read_heartbeat(run_dir) + # The runner first: it would otherwise notice its server dying and record a crash. + for role in ("runner", "server_launcher", "server"): + if kill_identified(heartbeat.get(role)): + print(f"killed {role} (pid {int(heartbeat[role][0])}) and its descendants") + info = load_run_info(run_dir) + # Nothing of the run is alive any more, so its scratch roots are unused. + for leftover in remove_scratch_roots(_data_dir(run_dir), info): + print(f"could not remove {leftover}") + if info.status == "running": + write_run_info( + run_dir, + info.model_copy(update={"status": "stopped", "finished_at": now_iso()}), + ) + print(f"{info.run_id} stopped; continue with: python -m scripts.family_e2e resume {run_dir}") + return 0 + + +def cmd_accept(args: argparse.Namespace) -> int: + run_dir = _resolve_run(args.run, args.manifest) + for reference in accept(_data_dir(run_dir), run_dir, args.case or None): + print(f"accepted {reference.case_key} from {reference.source_run_id}") + return 0 + + +def cmd_report(args: argparse.Namespace) -> int: + run_dir = _resolve_run(args.run, args.manifest) + result = write_report(_data_dir(run_dir), run_dir) + print(f"report {result.report_html}") + print(f"summary {result.summary_md}") + print(f"{result.failures} failures, {result.outliers} outliers, {result.invalid} invalid") + return 1 if result.failures else 0 + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m scripts.family_e2e", description="End-to-end runs per model family." + ) + commands = parser.add_subparsers(dest="command", required=True) + + start = commands.add_parser("start", help="create a run and start its runner in the background") + start.add_argument("--manifest", type=Path, required=True) + start.add_argument("--family", action="append", default=[], help="only this family (repeatable)") + start.add_argument("--case", action="append", default=[], help="only this case id (repeatable)") + start.set_defaults(handler=cmd_start) + + runner = commands.add_parser("run", help="run the remaining steps of a run in the foreground (what start launches)") + runner.add_argument("--run-dir", type=Path, required=True) + runner.add_argument("--log-file", type=Path, help="append stdout and stderr to this file") + runner.set_defaults(handler=cmd_run) + + for name, handler, help_text in ( + ("resume", cmd_resume, "start the runner of a stopped or crashed run again"), + ("status", cmd_status, "show progress and whether the runner is alive"), + ("stop", cmd_stop, "kill the runner and its server"), + ("accept", cmd_accept, "make a run's measured results the references of its profile"), + ("report", cmd_report, "write report.html and summary.md of a run again"), + ): + command = commands.add_parser(name, help=help_text) + command.add_argument("run", nargs="?", help="run directory or run id; default: the profile's latest run") + command.add_argument("--manifest", type=Path, help="resolves run ids and the latest run") + if name == "accept": + command.add_argument("--case", action="append", default=[], help="only this case key (repeatable)") + command.set_defaults(handler=handler) + return parser + + +def _redirect_stdio(path: Path) -> None: + # A background runner has no usable stdio; tracebacks of a crash must land in its log too. + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_APPEND) + os.dup2(fd, 1) + os.dup2(fd, 2) + os.close(fd) + sys.stdout = open(1, "w", encoding="utf-8", buffering=1, closefd=False) + sys.stderr = open(2, "w", encoding="utf-8", buffering=1, closefd=False) + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + if getattr(args, "log_file", None): + _redirect_stdio(args.log_file) + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s", stream=sys.stderr) + try: + return args.handler(args) + except (UsageError, *EXPECTED_ERRORS) as e: + print(f"error: {e}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/family_e2e/client.py b/scripts/family_e2e/client.py new file mode 100644 index 00000000000..ed5949c1f1f --- /dev/null +++ b/scripts/family_e2e/client.py @@ -0,0 +1,303 @@ +"""HTTP and socket.io client for the server under test, and the run timing built from its events. + +This process must never import torch: a CUDA context here would land on the GPU being measured. Socket events are +stamped on arrival with `time.perf_counter()`, the clock of the memory sampler; the server's own event timestamps +have whole-second resolution, and `time.monotonic()` ticks only every 15.6 ms on Windows. +""" + +from __future__ import annotations + +import mimetypes +import threading +import time +from collections import defaultdict, deque +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import requests +import socketio + +from scripts.family_e2e.records import ModelLoad, NodeTiming, Timing + +QUEUE_ID = "default" +SOCKETIO_PATH = "/ws/socket.io" +REQUEST_TIMEOUT_S = 120.0 +# While waiting for a queue item, its status is also polled, in case a socket event was lost. +POLL_S = 2.0 +EVENT_GRACE_S = 1.0 +TERMINAL_STATUSES = ("completed", "failed", "canceled") +EVENT_NAMES = ( + "invocation_started", + "invocation_progress", + "invocation_complete", + "invocation_error", + "queue_item_status_changed", + "model_load_started", + "model_load_complete", +) +# Upper bound for the step count inferred from progress percentages. +MAX_INFERRED_STEPS = 1000 +# The server drops an image-bearing progress frame sent sooner than this after the previous one +# (invokeai/app/services/progress_previews/progress_previews_default.py). +PREVIEW_INTERVAL_S = 0.1 + + +class ApiError(RuntimeError): + pass + + +class InvokeClient: + def __init__(self, base_url: str) -> None: + self.base_url = base_url.rstrip("/") + self._http = requests.Session() + + def _request(self, method: str, path: str, **kwargs: Any) -> requests.Response: + response = self._http.request(method, self.base_url + path, timeout=REQUEST_TIMEOUT_S, **kwargs) + if response.status_code >= 400: + raise ApiError(f"{method} {path} -> {response.status_code}: {response.text[:2000]}") + return response + + def get(self, path: str, **params: Any) -> Any: + return self._request("GET", path, params=params or None).json() + + def post(self, path: str, body: Any = None) -> Any: + return self._request("POST", path, json=body).json() + + def image_bytes(self, image_name: str) -> bytes: + return self._request("GET", f"/api/v1/images/i/{image_name}/full").content + + def upload_image(self, path: Path) -> str: + content_type = mimetypes.guess_type(path.name)[0] or "image/png" + with path.open("rb") as f: + dto = self._request( + "POST", + "/api/v1/images/upload", + params={"image_category": "user", "is_intermediate": "true"}, + files={"file": (path.name, f, content_type)}, + ).json() + return dto["image_name"] + + def enqueue(self, batch: dict[str, Any]) -> int: + result = self.post(f"/api/v1/queue/{QUEUE_ID}/enqueue_batch", {"batch": batch, "prepend": False}) + if len(result["item_ids"]) != 1: + raise ApiError(f"expected one queue item from the batch, got {len(result['item_ids'])}") + return result["item_ids"][0] + + def queue_item(self, item_id: int) -> dict[str, Any]: + return self.get(f"/api/v1/queue/{QUEUE_ID}/i/{item_id}") + + def queue_busy(self) -> bool: + queue = self.get(f"/api/v1/queue/{QUEUE_ID}/status")["queue"] + return queue["pending"] + queue["in_progress"] > 0 + + def cancel(self, item_id: int) -> None: + self._request("PUT", f"/api/v1/queue/{QUEUE_ID}/i/{item_id}/cancel") + + +@dataclass(frozen=True) +class Event: + name: str + t: float + """`time.perf_counter()` at arrival.""" + data: dict[str, Any] + + +class EventRecorder: + """Records the queue and model-load events of a server from connect to close.""" + + def __init__(self, base_url: str) -> None: + self._base_url = base_url + # Polling only: websocket-client is not installed, and polling delivers within a few ms on localhost. + self._sio = socketio.Client(reconnection=True) + self._events: list[Event] = [] + self._changed = threading.Condition() + self._sio.on("connect", self._subscribe) + for name in EVENT_NAMES: + self._sio.on(name, self._handler(name)) + + def _subscribe(self) -> None: + # Also runs after a reconnect, which starts a new socket without room memberships. + self._sio.emit("subscribe_queue", {"queue_id": QUEUE_ID}) + + def _handler(self, name: str) -> Callable[[dict[str, Any]], None]: + def handle(data: dict[str, Any]) -> None: + t = time.perf_counter() + if name == "invocation_progress": + # Preview images are large and only their presence matters (it marks denoise steps). + data = {k: v for k, v in data.items() if k != "image"} | {"has_image": data.get("image") is not None} + with self._changed: + self._events.append(Event(name, t, data)) + self._changed.notify_all() + + return handle + + def connect(self) -> None: + self._sio.connect(self._base_url, socketio_path=SOCKETIO_PATH, transports=["polling"], wait_timeout=30) + # The connect handler's emit is not acknowledged; this call returns once the server joined the rooms. + self._sio.call("subscribe_queue", {"queue_id": QUEUE_ID}, timeout=30) + + def close(self) -> None: + self._sio.disconnect() + + def mark(self) -> int: + with self._changed: + return len(self._events) + + def since(self, mark: int) -> list[Event]: + with self._changed: + return self._events[mark:] + + def wait(self, timeout: float) -> None: + with self._changed: + self._changed.wait(timeout) + + +def wait_for_item( + client: InvokeClient, + recorder: EventRecorder, + mark: int, + item_id: int, + timeout: float, + tick: Callable[[], None], +) -> tuple[dict[str, Any], float]: + """Waits for the item to finish. Returns the item and the `perf_counter` time it finished. + + The finish time is the last arrival among the item's events once its terminal status event is in. A poll that + sees the item finished first waits EVENT_GRACE_S for that event, because the poll can overtake the events still + in transit; without it, the poll time stands in. `tick` runs before every check, at least every POLL_S seconds, + and may raise to abort the wait. Raises TimeoutError. + """ + deadline = time.perf_counter() + timeout + next_poll = time.perf_counter() + POLL_S + polled: tuple[dict[str, Any], float] | None = None + while True: + tick() + own = _item_events(recorder.since(mark), item_id) + latest = max((e.t for e in own), default=0.0) + if any(e.name == "queue_item_status_changed" and e.data.get("status") in TERMINAL_STATUSES for e in own): + return client.queue_item(item_id), latest + now = time.perf_counter() + if polled is not None and now - polled[1] >= EVENT_GRACE_S: + return polled[0], max(polled[1], latest) + if polled is None and now >= next_poll: + item = client.queue_item(item_id) + if item["status"] in TERMINAL_STATUSES: + polled = (item, now) + next_poll = now + POLL_S + if now >= deadline: + raise TimeoutError(f"queue item {item_id} did not finish within {timeout:.0f}s") + recorder.wait(min(POLL_S, EVENT_GRACE_S, max(0.0, deadline - now))) + + +def wait_for_idle_queue(client: InvokeClient, timeout: float, tick: Callable[[], None]) -> bool: + """Waits until the queue has nothing pending or in progress; False if it is still busy after `timeout`.""" + deadline = time.perf_counter() + timeout + while True: + tick() + if not client.queue_busy(): + return True + if time.perf_counter() >= deadline: + return False + time.sleep(POLL_S) + + +def _item_events(events: list[Event], item_id: int) -> list[Event]: + return [e for e in events if e.data.get("item_id") == item_id] + + +def per_step_seconds(progress: list[tuple[float, float]], expected_steps: int | None = None) -> float | None: + """Seconds per denoise step from (arrival time, percentage) pairs of one denoise node. + + Image-bearing progress frames are throttled server-side and the throttled steps emit nothing, so each frame's + step comes from its percentage (k / N); the slope over the first and last frame then needs no frame in between. + N is `expected_steps` (a run over the full schedule) when every percentage fits it. Otherwise it is the smallest + count every percentage fits, which may be a divisor of the real one when only every other step left a frame; a + step slower than two throttle intervals rules that out, and a faster one gives None. + """ + if len(progress) < 2: + return None + + def fits(n: int) -> bool: + return all(abs(p * n - round(p * n)) < 1e-6 for _, p in progress) + + trusted = expected_steps is not None and fits(expected_steps) + steps = expected_steps if trusted else next((n for n in range(1, MAX_INFERRED_STEPS + 1) if fits(n)), None) + if steps is None: + return None + (t_first, p_first), (t_last, p_last) = progress[0], progress[-1] + k_first, k_last = round(p_first * steps), round(p_last * steps) + if k_last <= k_first: + return None + seconds = (t_last - t_first) / (k_last - k_first) + if not trusted and seconds < 2 * PREVIEW_INTERVAL_S: + return None + return seconds + + +def build_timing( + events: list[Event], + item_id: int, + t_enqueue: float, + t_done: float, + item: dict[str, Any], + expected_steps: int | None = None, +) -> Timing: + """Run timing from the recorded events; every `t_*` is relative to the enqueue instant. + + `expected_steps` is the denoise step count when the run covers the whole schedule (not img2img). + """ + # The socket client runs every handler on its own thread, so the list order can differ from the arrival order. + events = sorted(events, key=lambda e: e.t) + own = _item_events(events, item_id) + starts: dict[str, tuple[str, str, float]] = {} + ends: dict[str, float] = {} + progress: dict[str, list[tuple[float, float]]] = defaultdict(list) + for event in own: + invocation = event.data.get("invocation") or {} + prepared_id = invocation.get("id") + if prepared_id is None: + continue + if event.name == "invocation_started": + starts[prepared_id] = (event.data["invocation_source_id"], invocation.get("type", ""), event.t) + elif event.name in ("invocation_complete", "invocation_error"): + ends[prepared_id] = event.t + elif event.name == "invocation_progress" and event.data.get("has_image"): + percentage = event.data.get("percentage") + if percentage is not None: + progress[prepared_id].append((event.t, percentage)) + + nodes = [ + NodeTiming( + node_id=source_id, + type=node_type, + t_start=t - t_enqueue, + t_end=ends[prepared_id] - t_enqueue if prepared_id in ends else None, + ) + for prepared_id, (source_id, node_type, t) in sorted(starts.items(), key=lambda kv: kv[1][2]) + ] + + # Model load events carry no queue item; within the run's window they belong to it. + pending: dict[tuple[str, str | None], deque[int]] = defaultdict(deque) + loads: list[dict[str, Any]] = [] + for event in events: + if event.name not in ("model_load_started", "model_load_complete") or not t_enqueue <= event.t <= t_done: + continue + key = ((event.data.get("config") or {}).get("key", ""), event.data.get("submodel_type")) + if event.name == "model_load_started": + pending[key].append(len(loads)) + loads.append({"model": key[0], "submodel": key[1], "t_start": event.t - t_enqueue, "t_end": None}) + elif pending[key]: + loads[pending[key].popleft()]["t_end"] = event.t - t_enqueue + + # The denoise node is the one with the most preview frames. + denoise = max(progress.values(), key=len, default=[]) + return Timing( + total_s=t_done - t_enqueue, + nodes=nodes, + model_loads=[ModelLoad(**load) for load in loads], + per_step_s=per_step_seconds(denoise, expected_steps), + queue_started_at=item.get("started_at"), + queue_completed_at=item.get("completed_at"), + ) diff --git a/scripts/family_e2e/detach.py b/scripts/family_e2e/detach.py new file mode 100644 index 00000000000..8b827a3fcdc --- /dev/null +++ b/scripts/family_e2e/detach.py @@ -0,0 +1,78 @@ +"""Starting the runner as a background process that outlives the caller. + +On Windows the runner breaks away from the caller's job object, so a caller that kills its job's processes (an agent +harness does after two hours) cannot take a long run with it. Where the job forbids breakaway, the runner is created +through WMI instead, which puts it in no job of the caller; such a process starts with the user's default +environment, not the caller's. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +from scripts.family_e2e.manifest import TOOL_CHECKOUT + +RUNNER_LOG = "runner.log" +WMI_TIMEOUT_S = 120 + +# Hidden window: WMI shows a console program's window otherwise. +_WMI_CREATE = ( + "$startup = New-CimInstance -ClassName Win32_ProcessStartup -ClientOnly -Property @{ShowWindow=[uint16]0}; " + "$r = Invoke-CimMethod -ClassName Win32_Process -MethodName Create -Arguments @{" + "CommandLine=$env:FAMILY_E2E_COMMAND; CurrentDirectory=$env:FAMILY_E2E_CWD; ProcessStartupInformation=$startup}; " + "if ($r.ReturnValue -ne 0) { [Console]::Error.WriteLine('Win32_Process.Create returned ' + $r.ReturnValue); exit 1 }; " + "$r.ProcessId" +) + + +class DetachError(RuntimeError): + pass + + +def runner_command(run_dir: Path) -> list[str]: + return [ + sys.executable, + "-X", + "utf8", + "-m", + "scripts.family_e2e", + "run", + "--run-dir", + str(run_dir), + "--log-file", + str(run_dir / RUNNER_LOG), + ] + + +def spawn_runner(run_dir: Path) -> int: + """Starts the runner for `run_dir` in the background and returns the started process's pid.""" + args = runner_command(run_dir) + stdio = {"stdin": subprocess.DEVNULL, "stdout": subprocess.DEVNULL, "stderr": subprocess.DEVNULL} + if sys.platform != "win32": + return subprocess.Popen(args, cwd=TOOL_CHECKOUT, start_new_session=True, **stdio).pid + # CREATE_NO_WINDOW, not DETACHED_PROCESS: a venv's python.exe is a launcher, and under DETACHED_PROCESS the + # interpreter it starts gets a new visible console, which the server and every tool call would then share. + flags = subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP | subprocess.CREATE_BREAKAWAY_FROM_JOB + try: + return subprocess.Popen(args, cwd=TOOL_CHECKOUT, creationflags=flags, close_fds=True, **stdio).pid + except OSError: + return _spawn_via_wmi(args) + + +def _spawn_via_wmi(args: list[str]) -> int: + # Passed through the environment, so neither PowerShell nor WMI has to parse quoting of their own. + env = os.environ | {"FAMILY_E2E_COMMAND": subprocess.list2cmdline(args), "FAMILY_E2E_CWD": str(TOOL_CHECKOUT)} + result = subprocess.run( + ["powershell.exe", "-NoProfile", "-NonInteractive", "-Command", _WMI_CREATE], + env=env, + capture_output=True, + text=True, + timeout=WMI_TIMEOUT_S, + check=False, + ) + if result.returncode != 0: + raise DetachError(f"starting the runner through WMI failed:\n{result.stderr.strip()}") + return int(result.stdout.strip().splitlines()[-1]) diff --git a/scripts/family_e2e/gpu_memory.py b/scripts/family_e2e/gpu_memory.py new file mode 100644 index 00000000000..f4993026140 --- /dev/null +++ b/scripts/family_e2e/gpu_memory.py @@ -0,0 +1,580 @@ +"""Driver-level video memory of another process, sampled without a GPU context of our own. + +The allocator's own figure (``torch.cuda.max_memory_allocated``) misses its pools and the cuDNN/MIOpen workspaces, so +the real peak runs 25-75 % above it. What the driver charges the process is what decides whether a card of a given size +copes, and Windows reports it per process and adapter through the WDDM kernel thunks: ``D3DKMTQueryVideoMemoryInfo`` +on another process's handle answers in about a microsecond, so the dedicated (local segment) usage can be sampled every +few milliseconds. The bytes Windows has paged out to system memory are only visible in the ``GPU Process Memory`` +performance counters (``Shared Usage``), read through PDH at a slower rate. + +The client never imports torch: a CUDA/HIP context of its own would sit on the same adapter and distort what the +server can allocate. Windows (WDDM) only for now; it covers NVIDIA and AMD/ROCm alike. +""" + +import bisect +import ctypes +import functools +import logging +import operator +import re +import sys +import threading +import time +from dataclasses import dataclass +from types import SimpleNamespace + +_STATUS_SUCCESS = 0x00000000 +# The target has no allocation context on the adapter: it never opened it, or has closed it again. +_STATUS_INVALID_PARAMETER = 0xC000000D +_STATUS_ACCESS_DENIED = 0xC0000022 +# D3DKMTQueryVideoMemoryInfo answers STATUS_ACCESS_DENIED for PROCESS_QUERY_LIMITED_INFORMATION; it needs the full +# query right (measured on Windows 10 22H2 with an NVIDIA and an AMD driver). +_PROCESS_QUERY_INFORMATION = 0x0400 +_STILL_ACTIVE = 259 +_ERROR_ACCESS_DENIED = 5 +_ERROR_INVALID_PARAMETER = 87 + +_KMTQAITYPE_ADAPTERREGISTRYINFO = 8 +_KMTQAITYPE_ADAPTERTYPE = 15 +# D3DKMT_KMD_DRIVER_VERSION: the kernel-mode driver's version, the one Device Manager shows. +_KMTQAITYPE_KMD_DRIVER_VERSION = 49 +# D3DKMT_ADAPTERTYPE bits. +_ADAPTER_RENDER_SUPPORTED = 1 << 0 +_ADAPTER_SOFTWARE_DEVICE = 1 << 2 +_ADAPTER_COMPUTE_ONLY = 1 << 11 +_MEMORY_SEGMENT_GROUP_LOCAL = 0 +_MAX_PATH = 260 + +_PDH_FMT_LARGE = 0x00000400 +_PDH_MORE_DATA = 0x800007D2 +_PDH_NO_DATA = 0x800007D5 +# PDH_CSTATUS_VALID_DATA and PDH_CSTATUS_NEW_DATA; any other status leaves `largeValue` undefined. +_PDH_VALID_DATA = (0x00000000, 0x00000001) +# English paths; PdhAddEnglishCounterW resolves them on a localized Windows too. +_PDH_DEDICATED_USAGE = r"\GPU Process Memory(*)\Dedicated Usage" +_PDH_SHARED_USAGE = r"\GPU Process Memory(*)\Shared Usage" +_PDH_INSTANCE = re.compile(r"pid_(\d+)_luid_0x([0-9a-f]{8})_0x([0-9a-f]{8})_phys_(\d+)", re.IGNORECASE) + +_logger = logging.getLogger(__name__) + +# Win32 types spelled as fixed-width ctypes, so this module imports on every platform. +_UINT = ctypes.c_uint32 +_LONG = ctypes.c_int32 +_NTSTATUS = ctypes.c_uint32 + + +class UnsupportedPlatformError(RuntimeError): + """GPU memory sampling is not implemented for this operating system.""" + + +@dataclass(frozen=True) +class Adapter: + luid: int + """``(HighPart << 32) | LowPart``, as D3DKMT and the PDH instance names spell it.""" + description: str + driver_version: str | None = None + """As Device Manager shows it, e.g. ``32.0.16.1047``; ``None`` when the driver does not report one.""" + + +@dataclass(frozen=True) +class ProcessUsage: + dedicated_bytes: int + shared_bytes: int + + +@dataclass(frozen=True) +class MemoryTrace: + """Samples of one process's usage on one adapter, stamped with the client's ``time.perf_counter()``. + + The trace is a sample-and-hold step function: a value stands from its timestamp until the next sample. + """ + + method: str + interval_ms: float + """Mean spacing of the dedicated samples; the configured interval while they span no measurable time.""" + dedicated: tuple[tuple[float, int], ...] + """Local-segment ``CurrentUsage`` in bytes.""" + shared: tuple[tuple[float, int], ...] + """PDH ``Shared Usage`` in bytes: what Windows keeps of this process's allocations in system memory.""" + + def peak(self, t0: float | None = None, t1: float | None = None) -> int | None: + """The highest dedicated usage in ``[t0, t1]``: the value in effect at ``t0`` and every sample up to ``t1``. + + Counting the value in effect at ``t0`` keeps a window shorter than the sampling interval from coming back + empty. ``None`` bounds are open; ``None`` is returned when no sample falls on or before ``t1``. + """ + if t0 is not None and t1 is not None and t1 < t0: + raise ValueError(f"window ends before it starts: [{t0}, {t1}]") + samples = self.dedicated + lo = 0 + if t0 is not None: + lo = bisect.bisect_left(samples, t0, key=_sample_time) + if lo > 0 and (lo == len(samples) or samples[lo][0] > t0): + lo -= 1 + hi = len(samples) if t1 is None else bisect.bisect_right(samples, t1, key=_sample_time) + return max((value for _, value in samples[lo:hi]), default=None) + + def at(self, t: float) -> int | None: + """The dedicated usage in effect at ``t``: the last sample at or before it.""" + index = bisect.bisect_right(self.dedicated, t, key=_sample_time) + return self.dedicated[index - 1][1] if index else None + + +_sample_time = operator.itemgetter(0) + + +class _Luid(ctypes.Structure): + _fields_ = [("LowPart", _UINT), ("HighPart", _LONG)] + + +class _AdapterInfo(ctypes.Structure): + """``D3DKMT_ADAPTERINFO``.""" + + _fields_ = [("hAdapter", _UINT), ("AdapterLuid", _Luid), ("NumOfSources", _UINT), ("bPrecise", _LONG)] + + +class _EnumAdapters2(ctypes.Structure): + """``D3DKMT_ENUMADAPTERS2``.""" + + _fields_ = [("NumAdapters", _UINT), ("pAdapters", ctypes.POINTER(_AdapterInfo))] + + +class _QueryAdapterInfo(ctypes.Structure): + """``D3DKMT_QUERYADAPTERINFO``.""" + + _fields_ = [("hAdapter", _UINT), ("Type", ctypes.c_int), ("pData", ctypes.c_void_p), ("DataSize", _UINT)] + + +class _AdapterRegistryInfo(ctypes.Structure): + """``D3DKMT_ADAPTERREGISTRYINFO``.""" + + _fields_ = [ + ("AdapterString", ctypes.c_wchar * _MAX_PATH), + ("BiosString", ctypes.c_wchar * _MAX_PATH), + ("DacType", ctypes.c_wchar * _MAX_PATH), + ("ChipType", ctypes.c_wchar * _MAX_PATH), + ] + + +class _OpenAdapterFromLuid(ctypes.Structure): + """``D3DKMT_OPENADAPTERFROMLUID``.""" + + _fields_ = [("AdapterLuid", _Luid), ("hAdapter", _UINT)] + + +class _CloseAdapter(ctypes.Structure): + """``D3DKMT_CLOSEADAPTER``.""" + + _fields_ = [("hAdapter", _UINT)] + + +class _QueryVideoMemoryInfo(ctypes.Structure): + """``D3DKMT_QUERYVIDEOMEMORYINFO``. A null ``hProcess`` means the calling process.""" + + _fields_ = [ + ("hProcess", ctypes.c_void_p), + ("hAdapter", _UINT), + ("MemorySegmentGroup", ctypes.c_int), + ("Budget", ctypes.c_uint64), + ("CurrentUsage", ctypes.c_uint64), + ("CurrentReservation", ctypes.c_uint64), + ("AvailableForReservation", ctypes.c_uint64), + ("PhysicalAdapterIndex", _UINT), + ] + + +class _PdhCounterValue(ctypes.Structure): + """``PDH_FMT_COUNTERVALUE`` read as ``PDH_FMT_LARGE``.""" + + _fields_ = [("CStatus", ctypes.c_uint32), ("largeValue", ctypes.c_int64)] + + +class _PdhCounterValueItem(ctypes.Structure): + """``PDH_FMT_COUNTERVALUE_ITEM_W``.""" + + _fields_ = [("szName", ctypes.c_wchar_p), ("FmtValue", _PdhCounterValue)] + + +def _require_windows() -> None: + if sys.platform != "win32": + raise UnsupportedPlatformError( + f"GPU memory sampling is implemented for Windows (WDDM) only, not {sys.platform}" + ) + + +@functools.cache +def _win32() -> SimpleNamespace: + """gdi32, kernel32 and pdh with the prototypes used here; loaded on first use, since only Windows has them.""" + gdi32 = ctypes.WinDLL("gdi32") + for name in ( + "D3DKMTEnumAdapters2", + "D3DKMTQueryAdapterInfo", + "D3DKMTOpenAdapterFromLuid", + "D3DKMTQueryVideoMemoryInfo", + "D3DKMTCloseAdapter", + ): + fn = getattr(gdi32, name) + fn.argtypes = [ctypes.c_void_p] + fn.restype = _NTSTATUS + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + kernel32.OpenProcess.argtypes = [ctypes.c_uint32, ctypes.c_int32, ctypes.c_uint32] + kernel32.OpenProcess.restype = ctypes.c_void_p + kernel32.GetExitCodeProcess.argtypes = [ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint32)] + kernel32.GetExitCodeProcess.restype = ctypes.c_int32 + kernel32.CloseHandle.argtypes = [ctypes.c_void_p] + kernel32.CloseHandle.restype = ctypes.c_int32 + + pdh = ctypes.WinDLL("pdh") + prototypes = { + "PdhOpenQueryW": [ctypes.c_wchar_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p)], + "PdhAddEnglishCounterW": [ctypes.c_void_p, ctypes.c_wchar_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p)], + "PdhCollectQueryData": [ctypes.c_void_p], + "PdhGetFormattedCounterArrayW": [ + ctypes.c_void_p, + ctypes.c_uint32, + ctypes.POINTER(ctypes.c_uint32), + ctypes.POINTER(ctypes.c_uint32), + ctypes.c_void_p, + ], + "PdhCloseQuery": [ctypes.c_void_p], + } + for name, argtypes in prototypes.items(): + fn = getattr(pdh, name) + fn.argtypes = argtypes + fn.restype = ctypes.c_uint32 + return SimpleNamespace(gdi32=gdi32, kernel32=kernel32, pdh=pdh) + + +def _check(status: int, call: str) -> None: + if status != _STATUS_SUCCESS: + raise OSError(f"{call} failed with status {status:#010x}") + + +def _luid_value(luid: _Luid) -> int: + return ((luid.HighPart & 0xFFFFFFFF) << 32) | luid.LowPart + + +def _luid_struct(value: int) -> _Luid: + return _Luid(value & 0xFFFFFFFF, value >> 32) + + +def _query_adapter_info(handle: int, info_type: int, data: ctypes.Structure | ctypes.c_uint32 | ctypes.c_uint64) -> int: + query = _QueryAdapterInfo(handle, info_type, ctypes.cast(ctypes.byref(data), ctypes.c_void_p), ctypes.sizeof(data)) + return _win32().gdi32.D3DKMTQueryAdapterInfo(ctypes.byref(query)) + + +def _describe(info: _AdapterInfo) -> Adapter | None: + """The adapter, or ``None`` for one that cannot hold a compute process's memory or cannot be named.""" + adapter_type = _UINT() + if _query_adapter_info(info.hAdapter, _KMTQAITYPE_ADAPTERTYPE, adapter_type) != _STATUS_SUCCESS: + return None + # The Microsoft Basic Render Driver is a software device; display-only adapters (indirect displays) run nothing. + if adapter_type.value & _ADAPTER_SOFTWARE_DEVICE: + return None + if not adapter_type.value & (_ADAPTER_RENDER_SUPPORTED | _ADAPTER_COMPUTE_ONLY): + return None + registry = _AdapterRegistryInfo() + if _query_adapter_info(info.hAdapter, _KMTQAITYPE_ADAPTERREGISTRYINFO, registry) != _STATUS_SUCCESS: + return None + version = ctypes.c_uint64() + status = _query_adapter_info(info.hAdapter, _KMTQAITYPE_KMD_DRIVER_VERSION, version) + return Adapter( + luid=_luid_value(info.AdapterLuid), + description=registry.AdapterString, + driver_version=_format_driver_version(version.value) if status == _STATUS_SUCCESS else None, + ) + + +def _format_driver_version(packed: int) -> str | None: + """A driver version packed as four 16-bit fields, most significant first, in dotted form; 0 is unknown.""" + if not packed: + return None + return ".".join(str(packed >> shift & 0xFFFF) for shift in (48, 32, 16, 0)) + + +def list_adapters() -> list[Adapter]: + """The hardware GPU adapters, in the order the kernel enumerates them.""" + _require_windows() + gdi32 = _win32().gdi32 + # A null array asks for the size: the call answers the most adapters it could return. + enum = _EnumAdapters2(0, None) + _check(gdi32.D3DKMTEnumAdapters2(ctypes.byref(enum)), "D3DKMTEnumAdapters2") + infos = (_AdapterInfo * enum.NumAdapters)() + enum.pAdapters = infos + _check(gdi32.D3DKMTEnumAdapters2(ctypes.byref(enum)), "D3DKMTEnumAdapters2") + opened = infos[: enum.NumAdapters] + try: + return [adapter for adapter in map(_describe, opened) if adapter is not None] + finally: + for info in opened: + gdi32.D3DKMTCloseAdapter(ctypes.byref(_CloseAdapter(info.hAdapter))) + + +def _select_adapter(adapters: list[Adapter], description_substring: str) -> Adapter: + needle = description_substring.casefold() + matches = [adapter for adapter in adapters if needle in adapter.description.casefold()] + if len(matches) == 1: + return matches[0] + available = "; ".join(f"{adapter.description} (LUID {adapter.luid:#x})" for adapter in adapters) or "none" + found = "No adapter matches" if not matches else f"{len(matches)} adapters match" + raise LookupError(f"{found} {description_substring!r}. Adapters: {available}") + + +def find_adapter(description_substring: str) -> Adapter: + """The one adapter whose description contains the substring, ignoring case. ``LookupError`` for none or several.""" + return _select_adapter(list_adapters(), description_substring) + + +def _parse_instance(name: str) -> tuple[int, int] | None: + """``(pid, luid)`` of a ``GPU Process Memory`` instance on physical adapter 0, else ``None``.""" + match = _PDH_INSTANCE.fullmatch(name) + if match is None or match[4] != "0": + return None + return int(match[1]), (int(match[2], 16) << 32) | int(match[3], 16) + + +class _GpuProcessMemoryCounters: + """One PDH query over every process's ``GPU Process Memory`` instances. + + Only the first query in a process is slow to open (about 0.35 s while PDH loads the provider); a collection costs + well under a millisecond. + """ + + def __init__(self) -> None: + pdh = _win32().pdh + self._query = ctypes.c_void_p() + _check(pdh.PdhOpenQueryW(None, None, ctypes.byref(self._query)), "PdhOpenQueryW") + try: + self._dedicated = self._add(_PDH_DEDICATED_USAGE) + self._shared = self._add(_PDH_SHARED_USAGE) + except BaseException: + self.close() + raise + + def _add(self, path: str) -> ctypes.c_void_p: + counter = ctypes.c_void_p() + status = _win32().pdh.PdhAddEnglishCounterW(self._query, path, None, ctypes.byref(counter)) + _check(status, f"PdhAddEnglishCounterW({path})") + return counter + + def collect(self) -> dict[tuple[int, int], ProcessUsage]: + """Usage keyed by ``(pid, luid)``, for every process that has opened an adapter.""" + status = _win32().pdh.PdhCollectQueryData(self._query) + if status == _PDH_NO_DATA: + return {} + _check(status, "PdhCollectQueryData") + dedicated = self._values(self._dedicated) + shared = self._values(self._shared) + return { + key: ProcessUsage(dedicated.get(key, 0), shared.get(key, 0)) for key in dedicated.keys() | shared.keys() + } + + def _values(self, counter: ctypes.c_void_p) -> dict[tuple[int, int], int]: + pdh = _win32().pdh + size, count = ctypes.c_uint32(0), ctypes.c_uint32(0) + status = pdh.PdhGetFormattedCounterArrayW( + counter, _PDH_FMT_LARGE, ctypes.byref(size), ctypes.byref(count), None + ) + if status == _STATUS_SUCCESS: + return {} + if status != _PDH_MORE_DATA: + _check(status, "PdhGetFormattedCounterArrayW") + buffer = (ctypes.c_byte * size.value)() + status = pdh.PdhGetFormattedCounterArrayW( + counter, _PDH_FMT_LARGE, ctypes.byref(size), ctypes.byref(count), buffer + ) + _check(status, "PdhGetFormattedCounterArrayW") + items = ctypes.cast(buffer, ctypes.POINTER(_PdhCounterValueItem)) + values: dict[tuple[int, int], int] = {} + for item in items[: count.value]: + key = _parse_instance(item.szName or "") + if key is not None and item.FmtValue.CStatus in _PDH_VALID_DATA: + values[key] = int(item.FmtValue.largeValue) + return values + + def close(self) -> None: + if self._query: + _win32().pdh.PdhCloseQuery(self._query) + self._query = ctypes.c_void_p() + + +def process_usage_on_adapter(adapter: Adapter) -> dict[int, ProcessUsage]: + """Every process that has opened the adapter, by pid, with its dedicated and shared usage.""" + _require_windows() + counters = _GpuProcessMemoryCounters() + try: + usage = counters.collect() + finally: + counters.close() + return {pid: value for (pid, luid), value in usage.items() if luid == adapter.luid} + + +class ProcessMemorySampler: + """Samples one process's usage on one adapter on a background thread, from ``start()`` until ``stop()``. + + Dedicated usage is read every ``interval_s``, shared usage every ``shared_interval_s``. A process that has not + opened the adapter yet reads as zero. Sampling ends early when the process exits. + + ``pid`` must be the process that allocates. On Windows the ``python.exe`` in a venv's ``Scripts`` folder and the + console-script ``.exe`` files there are launchers: they run the interpreter as a child process and read as zero. + """ + + def __init__(self, pid: int, adapter: Adapter, interval_s: float = 0.003, shared_interval_s: float = 0.1) -> None: + _require_windows() + if pid <= 0: + raise ValueError(f"invalid pid {pid}") + if interval_s <= 0 or shared_interval_s <= 0: + raise ValueError(f"sampling intervals must be positive, got {interval_s} and {shared_interval_s}") + self._pid = pid + self._adapter = adapter + self._interval_s = interval_s + self._shared_interval_s = shared_interval_s + self._dedicated: list[tuple[float, int]] = [] + self._shared: list[tuple[float, int]] = [] + self._process: int | None = None + self._adapter_handle: int | None = None + self._counters: _GpuProcessMemoryCounters | None = None + self._stopping = threading.Event() + self._thread: threading.Thread | None = None + self._trace: MemoryTrace | None = None + self._warned: set[str] = set() + + def start(self) -> None: + """Open the process and the adapter, take one sample and start the sampling thread.""" + if self._thread is not None or self._trace is not None: + raise RuntimeError("a sampler starts once") + try: + self._open() + self._sample_dedicated(raise_on_failure=True) + self._sample_shared(raise_on_failure=True) + except BaseException: + self._close() + raise + self._thread = threading.Thread(target=self._run, name=f"gpu-memory-sampler-{self._pid}", daemon=True) + self._thread.start() + + def stop(self) -> MemoryTrace: + """Stop sampling, release the handles and return the trace; later calls return the same trace.""" + if self._trace is None: + self._stopping.set() + if self._thread is not None: + self._thread.join() + self._close() + self._trace = MemoryTrace( + method="d3dkmt", + interval_ms=self._mean_interval_ms(), + dedicated=tuple(self._dedicated), + shared=tuple(self._shared), + ) + return self._trace + + def _open(self) -> None: + win32 = _win32() + process = win32.kernel32.OpenProcess(_PROCESS_QUERY_INFORMATION, 0, self._pid) + if not process: + error = ctypes.get_last_error() + if error == _ERROR_INVALID_PARAMETER: + raise ProcessLookupError(f"no process with pid {self._pid}") + if error == _ERROR_ACCESS_DENIED: + raise PermissionError(f"access denied opening process {self._pid} for memory queries") + raise ctypes.WinError(error) + self._process = process + if self._exited(): + raise ProcessLookupError(f"process {self._pid} has exited") + + opened = _OpenAdapterFromLuid(_luid_struct(self._adapter.luid), 0) + status = win32.gdi32.D3DKMTOpenAdapterFromLuid(ctypes.byref(opened)) + _check(status, f"D3DKMTOpenAdapterFromLuid for {self._adapter.description} (LUID {self._adapter.luid:#x})") + self._adapter_handle = opened.hAdapter + # The target's STATUS_INVALID_PARAMETER only means "no allocation on this adapter" once the same call has + # succeeded for this process, which proves the adapter handle and the struct layout. + status, _ = self._query_usage(None) + _check(status, "D3DKMTQueryVideoMemoryInfo for the sampling process") + self._counters = _GpuProcessMemoryCounters() + + def _close(self) -> None: + win32 = _win32() + if self._counters is not None: + self._counters.close() + self._counters = None + if self._adapter_handle is not None: + win32.gdi32.D3DKMTCloseAdapter(ctypes.byref(_CloseAdapter(self._adapter_handle))) + self._adapter_handle = None + if self._process is not None: + win32.kernel32.CloseHandle(self._process) + self._process = None + + def _query_usage(self, process: int | None) -> tuple[int, int]: + assert self._adapter_handle is not None + info = _QueryVideoMemoryInfo(process, self._adapter_handle, _MEMORY_SEGMENT_GROUP_LOCAL) + status = _win32().gdi32.D3DKMTQueryVideoMemoryInfo(ctypes.byref(info)) + return status, int(info.CurrentUsage) + + def _exited(self) -> bool: + code = ctypes.c_uint32() + if not _win32().kernel32.GetExitCodeProcess(self._process, ctypes.byref(code)): + raise ctypes.WinError(ctypes.get_last_error()) + return code.value != _STILL_ACTIVE + + def _sample_dedicated(self, raise_on_failure: bool = False) -> bool: + """Append one dedicated sample; ``False`` once the process has exited.""" + t = time.perf_counter() + status, usage = self._query_usage(self._process) + if status == _STATUS_SUCCESS: + self._dedicated.append((t, usage)) + return True + if self._exited(): + if raise_on_failure: + raise ProcessLookupError(f"process {self._pid} has exited") + return False + if status == _STATUS_INVALID_PARAMETER: + self._dedicated.append((t, 0)) + return True + if raise_on_failure: + if status == _STATUS_ACCESS_DENIED: + raise PermissionError(f"access denied querying video memory of process {self._pid}") + _check(status, f"D3DKMTQueryVideoMemoryInfo for process {self._pid}") + self._warn_once("dedicated", f"D3DKMTQueryVideoMemoryInfo failed with status {status:#010x}; sample skipped") + return True + + def _sample_shared(self, raise_on_failure: bool = False) -> None: + assert self._counters is not None + t = time.perf_counter() + try: + usage = self._counters.collect().get((self._pid, self._adapter.luid)) + except OSError as error: + if raise_on_failure: + raise + self._warn_once("shared", f"{error}; shared sample skipped") + return + self._shared.append((t, 0 if usage is None else usage.shared_bytes)) + + def _warn_once(self, series: str, message: str) -> None: + if series not in self._warned: + self._warned.add(series) + _logger.warning(f"GPU memory sampler for process {self._pid}: {message} (further failures not logged)") + + def _run(self) -> None: + clock = time.perf_counter + next_dedicated = clock() + self._interval_s + next_shared = clock() + self._shared_interval_s + while not self._stopping.is_set(): + delay = next_dedicated - clock() + if delay > 0: + time.sleep(delay) + if not self._sample_dedicated(): + return + now = clock() + if now >= next_shared: + self._sample_shared() + next_shared = now + self._shared_interval_s + # A fixed rate: a late wake-up shortens the next sleep instead of delaying every later sample, and after a + # stall (GIL, suspend) sampling resumes from now rather than bursting to catch up. + next_dedicated = max(next_dedicated + self._interval_s, now) + + def _mean_interval_ms(self) -> float: + span = self._dedicated[-1][0] - self._dedicated[0][0] if len(self._dedicated) > 1 else 0.0 + if span <= 0: + return self._interval_s * 1000 + return span / (len(self._dedicated) - 1) * 1000 diff --git a/scripts/family_e2e/graphs.py b/scripts/family_e2e/graphs.py new file mode 100644 index 00000000000..8b7f46813ea --- /dev/null +++ b/scripts/family_e2e/graphs.py @@ -0,0 +1,180 @@ +"""Graphs for the cases: the exporter request, the webv2 exporter call, and the runner's batch policy. + +The webv2 exporter builds each case's batch with the real frontend graph builders (contract 1 in +`.scratch/agents/family-e2e/contract.md`). The runner then disables the invocation cache on every node, so repeated +runs really execute, and hashes the result so a frontend graph change shows up against the reference. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +import shutil +import subprocess +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from scripts.family_e2e.manifest import TOOL_CHECKOUT, ResolvedCase +from scripts.family_e2e.server import spawn_in_job, terminate + +EXPORT_VERSION = 1 +WEBV2_DIR = TOOL_CHECKOUT / "invokeai" / "frontend" / "webv2" +EXPORT_TIMEOUT_S = 600 + + +class ExportError(RuntimeError): + pass + + +@dataclass(frozen=True) +class PreparedGraph: + batch: dict[str, Any] + """The `batch` body for enqueue_batch.""" + graph_hash: str + model_keys: frozenset[str] + """Keys of every model identifier in the graph (main model, components, LoRAs, control models).""" + components: dict[str, str] + """The exporter's resolved component choice, for provenance.""" + + +def build_export_request( + cases: list[ResolvedCase], + models: list[dict[str, Any]], + capabilities: Any, + image_names: dict[str, dict[tuple[str, bool], str]], +) -> dict[str, Any]: + """`image_names` maps each case key's (manifest input path, used at the case size) to the uploaded image name.""" + request_cases = [] + for case in cases: + spec = case.spec + names = image_names[case.case_key] + request_cases.append( + { + "id": case.case_key, + "mode": spec.mode, + "model_key": spec.model, + "components": spec.components, + "width": case.width, + "height": case.height, + "steps": spec.steps, + "seed": case.seed, + "positive_prompt": case.prompt, + "negative_prompt": case.negative_prompt, + "loras": [{"key": lora.key, "weight": lora.weight} for lora in spec.loras], + "reference_images": [names[(path, False)] for path in spec.reference_images], + "init_image": ( + None + if spec.init_image is None + else {"image_name": names[(spec.init_image.path, True)], "strength": spec.init_image.strength} + ), + "control": [ + { + "model_key": c.model, + "image_name": names[(c.image, True)], + "weight": c.weight, + "begin": c.begin, + "end": c.end, + } + for c in spec.control + ], + "settings": spec.settings, + } + ) + return {"version": EXPORT_VERSION, "models": models, "capabilities": capabilities, "cases": request_cases} + + +def run_pnpm_exporter(request: dict[str, Any], work_dir: Path) -> dict[str, Any]: + """Runs `pnpm graphs:export` in the tool checkout's webv2 package. Per-case errors stay in the response.""" + pnpm = shutil.which("pnpm") + if pnpm is None: + raise ExportError("pnpm is not on PATH; the graph exporter runs with pnpm from invokeai/frontend/webv2") + work_dir.mkdir(parents=True, exist_ok=True) + # Absolute: pnpm runs the exporter from the webv2 package directory. + request_path = (work_dir / "export-request.json").resolve() + response_path = (work_dir / "export-response.json").resolve() + request_path.write_text(json.dumps(request), encoding="utf-8") + response_path.unlink(missing_ok=True) + # pnpm is a shim (pnpm.cmd on Windows) around node: a timeout must kill the whole tree, not just the shim. + process, job = spawn_in_job( + [pnpm, "graphs:export", str(request_path), str(response_path)], + cwd=WEBV2_DIR, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + try: + stdout, stderr = process.communicate(timeout=EXPORT_TIMEOUT_S) + except subprocess.TimeoutExpired: + raise ExportError(f"pnpm graphs:export did not finish within {EXPORT_TIMEOUT_S}s") from None + finally: + terminate(process, job) + if process.returncode != 0 or not response_path.is_file(): + output = (stderr or stdout).decode("utf-8", errors="replace") + raise ExportError(f"pnpm graphs:export exited with {process.returncode}:\n{output[-4000:]}") + response = json.loads(response_path.read_text(encoding="utf-8")) + if response.get("version") != EXPORT_VERSION: + raise ExportError(f"exporter answered version {response.get('version')!r}, expected {EXPORT_VERSION}") + return response + + +def _model_keys(value: Any) -> Iterator[str]: + # A ModelIdentifierField: {"key", "hash", "name", "base", "type", ...}. + if isinstance(value, dict): + if isinstance(value.get("key"), str) and "base" in value and "type" in value: + yield value["key"] + for item in value.values(): + yield from _model_keys(item) + elif isinstance(value, list): + for item in value: + yield from _model_keys(item) + + +def _replace_strings(value: Any, replacements: dict[str, str]) -> Any: + if isinstance(value, str): + return replacements.get(value, value) + if isinstance(value, dict): + return {k: _replace_strings(v, replacements) for k, v in value.items()} + if isinstance(value, list): + return [_replace_strings(v, replacements) for v in value] + return value + + +def prepare_batch(batch: dict[str, Any], components: dict[str, str], input_ids: dict[str, str]) -> PreparedGraph: + """Applies the runner policy to an exported batch and hashes it. + + `input_ids` maps uploaded image names to a stable identity of the input file. Upload names differ on every + server start, so the hash covers the input's identity instead. + """ + batch = copy.deepcopy(batch) + for node in batch["graph"]["nodes"].values(): + node["use_cache"] = False + batch["runs"] = 1 + hashed = _replace_strings({"graph": batch["graph"], "data": batch.get("data")}, input_ids) + canonical = json.dumps(hashed, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + return PreparedGraph( + batch=batch, + graph_hash=hashlib.sha256(canonical.encode("utf-8")).hexdigest(), + model_keys=frozenset(_model_keys(batch["graph"])), + components=components, + ) + + +def output_image_name(graph: dict[str, Any], session: dict[str, Any]) -> str: + """The image of the graph's one non-intermediate node that produced an image.""" + mapping = session.get("source_prepared_mapping") or {} + results = session.get("results") or {} + names = [] + for node_id, node in graph["nodes"].items(): + # The server's default: a node without the field is not intermediate. + if node.get("is_intermediate", False): + continue + for prepared_id in mapping.get(node_id, []): + image = (results.get(prepared_id) or {}).get("image") + if isinstance(image, dict) and isinstance(image.get("image_name"), str): + names.append(image["image_name"]) + if len(names) != 1: + raise ValueError(f"expected one output image from the graph's non-intermediate nodes, found {len(names)}") + return names[0] diff --git a/scripts/family_e2e/manifest.example.yaml b/scripts/family_e2e/manifest.example.yaml new file mode 100644 index 00000000000..26351a9b6b4 --- /dev/null +++ b/scripts/family_e2e/manifest.example.yaml @@ -0,0 +1,80 @@ +# Family E2E manifest, example for the first wave. Copy it next to your data directory, replace every <...> (paths +# and the model keys of your install; GET /api/v2/models/ lists the keys), then: +# python -m scripts.family_e2e start --manifest +# Relative paths are relative to this file. + +profile: rtx4090 # names runs and references; one profile per machine and GPU +data_dir: # inputs/, references/, runs/, scratch-roots/ +gpu: RTX 4090 # substring of the adapter name; must match one adapter ("Radeon" also matches an AMD iGPU, use "9060") + +server: + # Its invokeai.yaml and a database snapshot are copied; its server must be off. The scratch server shares its + # models directory: it writes models/README.txt and removes install claims of processes that no longer run there. + template_root: