chore: import upstream snapshot with attribution
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,357 @@
|
||||
# InvokeAI Graph - Design Overview
|
||||
|
||||
High-level design for the graph module. Focuses on responsibilities, data flow, and how traversal works.
|
||||
|
||||
## 1) Purpose
|
||||
|
||||
Provide a typed, acyclic workflow model (**Graph**) plus a runtime scheduler (**GraphExecutionState**) that expands
|
||||
iterator patterns, tracks readiness via indegree (the number of incoming edges to a node in the directed graph), and
|
||||
executes nodes in class-grouped batches. In normal execution, runtime expansion happens in a separate execution graph
|
||||
instead of mutating the source graph.
|
||||
|
||||
## 2) Major Data Types
|
||||
|
||||
### EdgeConnection
|
||||
|
||||
- Fields: `node_id: str`, `field: str`.
|
||||
- Hashable; printed as `node.field` for readable diagnostics.
|
||||
|
||||
### Edge
|
||||
|
||||
- Fields: `source: EdgeConnection`, `destination: EdgeConnection`.
|
||||
- One directed connection from a specific output port to a specific input port.
|
||||
|
||||
### AnyInvocation / AnyInvocationOutput
|
||||
|
||||
- Pydantic wrappers that carry concrete invocation models and outputs.
|
||||
- No registry logic in this file; they are permissive containers for heterogeneous nodes.
|
||||
|
||||
### IterateInvocation / CollectInvocation
|
||||
|
||||
- Control nodes used by validation and execution:
|
||||
|
||||
- **IterateInvocation**: input `collection`, outputs include `item` (and index/total).
|
||||
- **CollectInvocation**: many `item` inputs aggregated to one `collection` output.
|
||||
|
||||
## 3) Graph (author-time model)
|
||||
|
||||
A container for declared nodes and edges. Does **not** perform iteration expansion.
|
||||
|
||||
### 3.1 Data
|
||||
|
||||
- `nodes: dict[str, AnyInvocation]` - key must equal `node.id`.
|
||||
- `edges: list[Edge]` - zero or more.
|
||||
- Utility: `_get_input_edges(node_id, field?)`, `_get_output_edges(node_id, field?)` These scan `self.edges` (no
|
||||
adjacency indices in the current code).
|
||||
|
||||
### 3.2 Validation (`validate_self`)
|
||||
|
||||
Runs a sequence of checks:
|
||||
|
||||
1. **Node ID uniqueness** No duplicate IDs; map key equals `node.id`.
|
||||
|
||||
1. **Endpoint existence** Source and destination node IDs must exist.
|
||||
|
||||
1. **Port existence** Input ports must exist on the node class; output ports on the node's output model.
|
||||
|
||||
1. **DAG constraint** Build a *flat* `DiGraph` (no runtime expansion) and assert acyclicity.
|
||||
|
||||
1. **Type compatibility** `get_output_field_type` vs `get_input_field_type` and `are_connection_types_compatible`.
|
||||
|
||||
Special case:
|
||||
|
||||
- `call_saved_workflow` currently accepts dynamic destination handles of the form
|
||||
`saved_workflow_input::{childNodeId}::{childFieldName}` as part of its temporary call-boundary contract.
|
||||
- Those handles are allowed through graph validation even though they are not static Python model fields on the
|
||||
invocation class.
|
||||
- Runtime later validates them against the selected child workflow's exposed callable interface before applying
|
||||
values to the child graph.
|
||||
- The editor preserves dynamic caller values only while the exposed field type remains compatible; type drift at the
|
||||
same child node/field path resets to the selected workflow's current initial value.
|
||||
- Saved-workflow picker search is server-backed so large workflow libraries do not require scrolling every page
|
||||
before selecting a workflow by name.
|
||||
|
||||
1. **Iterator / collector structure** Enforce special rules:
|
||||
|
||||
- Iterator's input must be `collection`; its outgoing edges use `item`.
|
||||
- Collector accepts many `item` inputs; outputs a single `collection`.
|
||||
- Edge fan-in to a non-collector input is rejected.
|
||||
|
||||
### 3.3 Edge admission (`_validate_edge`)
|
||||
|
||||
Checks a single prospective edge before insertion:
|
||||
|
||||
- Endpoints/ports exist.
|
||||
- Destination port is not already occupied unless it's a collector `item`.
|
||||
- Adding the edge to the flat DAG must keep it acyclic.
|
||||
- Iterator/collector constraints re-checked when the edge creates relevant patterns.
|
||||
|
||||
### 3.4 Topology utilities
|
||||
|
||||
- `nx_graph()` - DiGraph of declared nodes and edges.
|
||||
- `nx_graph_flat()` - "flattened" DAG (still author-time; no runtime copies). Used in validation and in `_prepare()`
|
||||
during execution planning.
|
||||
|
||||
### 3.5 Mutation helpers
|
||||
|
||||
- `add_node`, `update_node` (preserve edges, rewrite endpoints if id changes), `delete_node`.
|
||||
- `add_edge`, `delete_edge` (with validation).
|
||||
|
||||
## 4) GraphExecutionState (runtime)
|
||||
|
||||
Holds the state for a single run. Keeps the source graph intact and materializes a separate execution graph.
|
||||
`GraphExecutionState` is still the public runtime entry point, but most execution behavior is now delegated to a small
|
||||
set of internal helper classes.
|
||||
|
||||
The source graph is treated as stable during normal execution, but the runtime object still exposes guarded graph
|
||||
mutation helpers. Those helpers reject changes once the affected nodes have already been prepared or executed.
|
||||
|
||||
### 4.1 Data
|
||||
|
||||
- `graph: Graph` - source graph for the run; treated as stable during normal execution.
|
||||
- `execution_graph: Graph` - materialized runtime nodes/edges. This is mutable runtime state, not an immutable audit
|
||||
log. Lazy `If` pruning may remove unselected input edges during execution, so persisted failed/completed session
|
||||
snapshots can contain a structurally pruned execution graph. Retry paths rebuild from `graph`, not from a previously
|
||||
persisted `execution_graph`.
|
||||
- `executed: set[str]`, `executed_history: list[str]`.
|
||||
- `results: dict[str, AnyInvocationOutput]`, `errors: dict[str, str]`.
|
||||
- `prepared_source_mapping: dict[str, str]` - exec id -> source id.
|
||||
- `source_prepared_mapping: dict[str, set[str]]` - source id -> exec ids.
|
||||
- `indegree: dict[str, int]` - unmet inputs per exec node.
|
||||
- Workflow-call runtime state:
|
||||
- `workflow_call_stack` - active parent call frames.
|
||||
- `workflow_call_history` - completed or failed workflow-call relationships observed by this execution state.
|
||||
- `workflow_call_parent` - parent workflow-call relationship metadata when this execution state is a child session.
|
||||
- `waiting_workflow_call` - the call frame currently suspending this execution state, if any.
|
||||
- `waiting_workflow_call_execution` - the active parent/child workflow-call relationship record for the waiting call.
|
||||
- `waiting_workflow_call_child_session` - attached child execution state for the waiting workflow call, if any.
|
||||
- `max_workflow_call_depth` - runtime guardrail for nested or recursive workflow calls.
|
||||
- Prepared exec metadata caches:
|
||||
- source node id
|
||||
- iteration path
|
||||
- runtime state such as pending, ready, executed, or skipped
|
||||
- **Ready queues grouped by class** (private attrs): `_ready_queues: dict[class_name, deque[str]]`,
|
||||
`_active_class: Optional[str]`. Optional `ready_order: list[str]` to prioritize classes. Queues are rebuilt from
|
||||
persisted execution state when a session is deserialized.
|
||||
|
||||
### 4.2 Core methods
|
||||
|
||||
- `next()` Returns the next ready exec node. If none are ready, it asks the materializer to expand more source nodes and
|
||||
then retries. If the execution state is paused on a workflow call boundary, it returns `None` without scheduling more
|
||||
work. Before returning a node, the runtime helper deep-copies inbound values into the node fields.
|
||||
- `complete(node_id, output)` Records the result, marks the exec node executed, marks the source node executed once all
|
||||
of its prepared exec copies are done, then decrements downstream indegrees and enqueues newly ready nodes.
|
||||
|
||||
Workflow-call note:
|
||||
|
||||
- `GraphExecutionState` can represent a paused parent execution plus an attached child execution state, but it does not
|
||||
itself orchestrate child execution.
|
||||
- In the current implementation, `DefaultSessionRunner.run_node()` establishes the workflow call boundary and attaches
|
||||
the child execution state, while `WorkflowCallCoordinator` handles call-specific setup and
|
||||
`WorkflowCallQueueLifecycle` later resumes or fails the parent based on that child queue row's outcome.
|
||||
- Child `SessionQueueItem` rows created by the coordinator now carry explicit relationship metadata such as
|
||||
`workflow_call_id`, `parent_item_id`, `parent_session_id`, `root_item_id`, and `workflow_call_depth`, even though the
|
||||
higher-level scheduler semantics are still evolving.
|
||||
- The `session_queue` schema now has matching columns for those relationship fields, and parent queue items can enter a
|
||||
`waiting` status while suspended on a child workflow execution.
|
||||
- Queue lifecycle semantics are now partially defined for workflow-call chains:
|
||||
- child success resumes the waiting parent
|
||||
- multiple child queue rows may complete under one waiting parent when the called workflow contains direct batch
|
||||
nodes; the parent resumes only after all expected child rows complete
|
||||
- child failure fails the waiting parent and can cascade upward through ancestors
|
||||
- failing child rows cancel their remaining workflow-call siblings before the parent is failed
|
||||
- cancelation is chain-aware across parents and children, including nested descendants of batched siblings
|
||||
- "all except current" queue actions preserve the active current item plus its workflow-call chain, while still
|
||||
canceling or deleting unrelated waiting chains
|
||||
- startup recovery cancels interrupted `in_progress` or `waiting` workflow-call chains, including pending descendants
|
||||
- deleting a workflow-call queue row currently deletes the whole parent/child chain rather than leaving orphaned rows
|
||||
behind
|
||||
- retry is root-oriented and should not be exposed directly on child queue rows in the UI
|
||||
- child queue-row creation is cleaned up on boundary-setup failure and child fan-out is bounded by remaining queue
|
||||
capacity
|
||||
- child workflows that mix supported batch nodes with unrelated generator nodes are rejected for now
|
||||
- This is still an intermediate architecture step and should eventually be replaced by a more general parent/child
|
||||
execution mechanism rather than workflow-call-specific queue lifecycle handling.
|
||||
|
||||
### 4.3 Runtime helper classes
|
||||
|
||||
`GraphExecutionState` now delegates most runtime behavior to internal helpers:
|
||||
|
||||
- `_PreparedExecRegistry` Owns the relationship between source graph nodes and prepared execution graph nodes, plus
|
||||
cached metadata such as iteration path and runtime state.
|
||||
- `_ExecutionMaterializer` Expands source graph nodes into concrete execution graph nodes when the scheduler runs out of
|
||||
ready work. It owns iterator expansion, collector grouping, prepared-parent selection, and creation of execution-graph
|
||||
edges. When matching prepared parents for a downstream exec node, skipped prepared exec nodes are ignored and cannot
|
||||
be selected as live inputs.
|
||||
- `_ExecutionScheduler` Owns indegree transitions, ready queues, class batching, and downstream release on completion.
|
||||
- `_ExecutionRuntime` Owns iteration-path lookup, collect input ordering, and input hydration for prepared exec nodes.
|
||||
- `_IfBranchScheduler` Applies lazy `If` semantics by deferring branch-local work until the condition is known, then
|
||||
releasing the selected branch and skipping the unselected branch.
|
||||
|
||||
`GraphExecutionState.model_post_init()` rehydrates private runtime helpers and caches after normal construction or a
|
||||
JSON/model round trip. Rehydration reconstructs prepared exec metadata, cached iteration paths, resolved `If` branch
|
||||
state when the condition is already available, and ready queues from `execution_graph`, `indegree`, `executed`, and
|
||||
`results`. This keeps persisted sessions resumable without persisting private helper objects.
|
||||
|
||||
### 4.4 Preparation (`_prepare()`)
|
||||
|
||||
- Build a flat DAG from the **source** graph.
|
||||
|
||||
- Choose the **next source node** in topological order that:
|
||||
|
||||
1. has not been prepared,
|
||||
1. if it is an iterator, *its inputs are already executed*,
|
||||
1. it has *no unexecuted iterator ancestors*.
|
||||
|
||||
- If the node is a **CollectInvocation**: group prepared parent exec nodes by iteration path and create one collector
|
||||
exec node per group. A collector collapses the immediate iterator that feeds its `item` input, but preserves enclosing
|
||||
iterator paths. This lets a shape such as `outer_iter -> inner_collection -> inner_iter -> collect -> consumer`
|
||||
produce one collected result per outer iteration instead of mixing all inner items into one global collection.
|
||||
Incoming `collection` inputs are treated as ancestor groups and are copied into each matching descendant item group.
|
||||
|
||||
- Otherwise: compute all combinations of prepared iterator ancestors. For each combination, choose the prepared parent
|
||||
for each upstream by matching iterator ancestry, then create **one** exec node. If a node no longer has visible
|
||||
iterator ancestors because the source path crosses a collector, prepared parent iteration paths are still used to
|
||||
materialize one downstream exec node for each preserved collector path.
|
||||
|
||||
- For each new exec node:
|
||||
|
||||
- Deep-copy the source node; assign a fresh ID (and `index` for iterators).
|
||||
- Cache the preserved iteration path when the materializer has one, such as for grouped collectors.
|
||||
- Wire edges from chosen prepared parents.
|
||||
- Set `indegree = number of unmet inputs` (i.e., parents not yet executed).
|
||||
- Try to resolve any `If`-specific scheduling state.
|
||||
- If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue.
|
||||
|
||||
### 4.5 Readiness and batching
|
||||
|
||||
- `_enqueue_if_ready(nid)` enqueues by class name only when `indegree == 0`, the node has not already executed, and the
|
||||
node is not deferred by an unresolved `If`.
|
||||
- `_get_next_node()` drains the `_active_class` queue; when empty, selects the next nonempty class queue (by
|
||||
`ready_order` if set, else alphabetical), and continues. Within each class queue, ready exec nodes are ordered by
|
||||
iteration path so expanded iterator work runs in a stable outer-to-inner order. Optional fairness knobs can limit
|
||||
batch size per class; default is drain fully.
|
||||
|
||||
#### 4.5.1 Indegree (what it is and how it's used)
|
||||
|
||||
**Indegree** is the number of incoming edges to a node in the execution graph that are still unmet. In this engine:
|
||||
|
||||
- For every materialized exec node, `indegree[node]` equals the count of its prerequisite parents that have **not**
|
||||
finished yet.
|
||||
- A node is "ready" exactly when `indegree[node] == 0`; only then is it enqueued.
|
||||
- When a node completes, the scheduler decrements `indegree[child]` for each outgoing edge. Any child that reaches 0 is
|
||||
enqueued.
|
||||
|
||||
Example: edges `A->C`, `B->C`, `C->D`. Start: `A:0, B:0, C:2, D:1`. Run `A` -> `C:1`. Run `B` -> `C:0` -> enqueue `C`.
|
||||
Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done.
|
||||
|
||||
### 4.6 Input hydration (`_prepare_inputs()`)
|
||||
|
||||
- For **CollectInvocation**: gather the materialized incoming `item` values into `collection`, sorting inputs by
|
||||
iteration path so collected results are stable across expanded iterations. Incoming `collection` values are merged
|
||||
first, then incoming `item` values are appended. By the time hydration runs, the materializer has already selected the
|
||||
iteration group for this collector exec node, so hydration only sees inputs that belong to that group.
|
||||
- For **IfInvocation**: hydrate only `condition` and the selected branch input. As a defensive guard against
|
||||
inconsistent runtime or deserialized session state, the runtime raises if the selected input edge points at an exec
|
||||
node with no stored runtime output. In normal scheduling this path should be unreachable.
|
||||
- For all others: deep-copy each incoming edge's value into the destination field. This prevents cross-node mutation
|
||||
through shared references.
|
||||
|
||||
### 4.7 Lazy `If` semantics
|
||||
|
||||
`IfInvocation` now acts as a lazy branch boundary rather than a simple value multiplexer.
|
||||
|
||||
- The `condition` input must resolve first.
|
||||
- Nodes that are exclusive to the true or false branch can remain deferred even when their indegree is zero.
|
||||
- Once the prepared `If` node resolves its condition:
|
||||
- the selected branch is released
|
||||
- the unselected branch is marked skipped
|
||||
- unselected input edges on the prepared `If` exec node are pruned from the execution graph so they no longer
|
||||
participate in downstream indegree accounting
|
||||
- branch-exclusive ancestors of the unselected branch are never executed
|
||||
- Skipped branch-local exec nodes may still be treated as executed for scheduling purposes, but they do not create
|
||||
entries in `results`.
|
||||
- Shared ancestors still execute if they are required by the selected branch or by any other live path in the graph.
|
||||
|
||||
This behavior is implemented in the runtime scheduler, not in the invocation body itself.
|
||||
|
||||
## 5) Traversal Summary
|
||||
|
||||
1. Author builds a valid **Graph**.
|
||||
|
||||
1. Create **GraphExecutionState** with that graph.
|
||||
|
||||
1. Loop:
|
||||
|
||||
- `node = state.next()` -> may trigger `_prepare()` expansion.
|
||||
- Execute node externally -> `output`.
|
||||
- `state.complete(node.id, output)` -> updates indegrees, `If` state, and ready queues.
|
||||
|
||||
1. Finish when `next()` returns `None` and the execution state is not paused waiting on a workflow call boundary.
|
||||
|
||||
In normal execution, all runtime expansion occurs in `execution_graph` with traceability back to source nodes.
|
||||
|
||||
## 6) Invariants
|
||||
|
||||
- Source **Graph** remains a DAG and type-consistent.
|
||||
- `execution_graph` remains a DAG.
|
||||
- Nodes are enqueued only when `indegree == 0` and they are not deferred by an unresolved `If`.
|
||||
- `results` and `errors` are keyed by **exec node id**.
|
||||
- Collectors aggregate `item` inputs and may also merge incoming `collection` inputs during runtime hydration.
|
||||
Collectors nested under iterators preserve enclosing iteration paths, so downstream consumers materialize per enclosing
|
||||
iteration instead of receiving a mixed collection from unrelated outer iterations.
|
||||
- Branch-exclusive nodes behind an unselected `If` branch are skipped, not failed.
|
||||
|
||||
## 7) Extensibility
|
||||
|
||||
- **New node types**: implement as Pydantic models with typed fields and outputs. Register per your invocation system;
|
||||
this file accepts them as `AnyInvocation`.
|
||||
- **Scheduling policy**: adjust `ready_order` to batch by class; add a batch cap for fairness without changing
|
||||
complexity.
|
||||
- **Dynamic behaviors** (future): can be added in `GraphExecutionState` by creating exec nodes and edges at `complete()`
|
||||
time, as long as the DAG invariant holds.
|
||||
- **Workflow call boundaries**: `GraphExecutionState` can suspend a parent execution state on a workflow call, attach a
|
||||
child execution state, and later resume the parent without mutating the source graph.
|
||||
|
||||
Current limitation:
|
||||
|
||||
- Child workflow executions are now represented as first-class queue items. Parent resume/failure is intentionally
|
||||
handled by a dedicated workflow-call queue lifecycle component for this PR because no other feature currently needs a
|
||||
generalized dependent-queue scheduler.
|
||||
- Called workflows currently require exactly one valid `workflow_return` node to be callable at all.
|
||||
- A single `workflow_return_value.value` may connect directly to `workflow_return.values`; multiple named return members
|
||||
should be collected and then connected to `workflow_return.values`.
|
||||
- Direct batch-special child workflows are now supported by expanding them into multiple child queue rows.
|
||||
- Batch outputs may feed a named `workflow_return_value.value` directly. Parent resume aggregates named return maps as
|
||||
`values: dict[str, list[Any]]`, and all rows in one batch call must return the same key set.
|
||||
- Generator-backed batch child workflows are now supported when the batch node is fed directly by a supported integer,
|
||||
float, string, or image generator.
|
||||
- Connected batch child inputs produced by ordinary non-generator upstream nodes are still rejected before any child
|
||||
queue row is created.
|
||||
- Workflow library API responses now include compatibility metadata so the frontend can disable unsupported callees
|
||||
before execution rather than failing only at runtime.
|
||||
- Workflow library list compatibility uses structural generator-backed batch validation so list and picker rendering do
|
||||
not enumerate every image in board-backed generators; workflow detail and runtime execution still resolve real
|
||||
generator values.
|
||||
- Batch-specific compatibility failures, including multiple connected inputs to one batch field, are reported as
|
||||
`unsupported_batch_input` rather than generic unsupported-node failures.
|
||||
- The workflow library list also surfaces that metadata as an informational unsupported state; workflows remain
|
||||
viewable/editable even when they are not currently callable by `call_saved_workflow`.
|
||||
- Single-user workflow CRUD socket events emit only to the admin room because every single-user socket already joins
|
||||
that room, avoiding duplicate delivery through both `user:system` and `admin`.
|
||||
|
||||
## 8) Error Model (selected)
|
||||
|
||||
- `DuplicateNodeIdError`, `NodeAlreadyInGraphError`
|
||||
- `NodeNotFoundError`, `NodeFieldNotFoundError`
|
||||
- `InvalidEdgeError`, `CyclicalGraphError`
|
||||
- `NodeInputError` (raised when preparing inputs for execution)
|
||||
|
||||
Messages favor short, precise diagnostics (node id, field, and failing condition).
|
||||
|
||||
## 9) Rationale
|
||||
|
||||
- **Two-graph approach** isolates authoring from execution expansion and keeps validation simple.
|
||||
- **Indegree + queues** gives O(1) scheduling decisions with clear batching semantics.
|
||||
- **Iterator/collector separation** keeps fan-out/fan-in explicit and testable.
|
||||
- **Deep-copy hydration** avoids incidental aliasing bugs between nodes.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,808 @@
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Callable, Optional, Union
|
||||
|
||||
from PIL.Image import Image
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
from torch import Tensor
|
||||
|
||||
from invokeai.app.invocations.constants import IMAGE_MODES
|
||||
from invokeai.app.invocations.fields import MetadataField, WithBoard, WithMetadata
|
||||
from invokeai.app.services.board_records.board_records_common import BoardRecordOrderBy
|
||||
from invokeai.app.services.boards.boards_common import BoardDTO
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory, ResourceOrigin
|
||||
from invokeai.app.services.images.images_common import ImageDTO
|
||||
from invokeai.app.services.invocation_services import InvocationServices
|
||||
from invokeai.app.services.model_records.model_records_base import UnknownModelException
|
||||
from invokeai.app.services.session_processor.session_processor_common import ProgressImage
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.util.step_callback import diffusion_step_callback
|
||||
from invokeai.backend.model_manager.configs.base import Config_Base
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig
|
||||
from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
|
||||
from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
|
||||
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.invocations.baseinvocation import BaseInvocation
|
||||
from invokeai.app.invocations.model import ModelIdentifierField
|
||||
from invokeai.app.services.session_queue.session_queue_common import SessionQueueItem
|
||||
|
||||
"""
|
||||
The InvocationContext provides access to various services and data about the current invocation.
|
||||
|
||||
We do not provide the invocation services directly, as their methods are both dangerous and
|
||||
inconvenient to use.
|
||||
|
||||
For example:
|
||||
- The `images` service allows nodes to delete or unsafely modify existing images.
|
||||
- The `configuration` service allows nodes to change the app's config at runtime.
|
||||
- The `events` service allows nodes to emit arbitrary events.
|
||||
|
||||
Wrapping these services provides a simpler and safer interface for nodes to use.
|
||||
|
||||
When a node executes, a fresh `InvocationContext` is built for it, ensuring nodes cannot interfere
|
||||
with each other.
|
||||
|
||||
Many of the wrappers have the same signature as the methods they wrap. This allows us to write
|
||||
user-facing docstrings and not need to go and update the internal services to match.
|
||||
|
||||
Note: The docstrings are in weird places, but that's where they must be to get IDEs to see them.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class InvocationContextData:
|
||||
queue_item: "SessionQueueItem"
|
||||
"""The queue item that is being executed."""
|
||||
invocation: "BaseInvocation"
|
||||
"""The invocation that is being executed."""
|
||||
source_invocation_id: str
|
||||
"""The ID of the invocation from which the currently executing invocation was prepared."""
|
||||
|
||||
|
||||
class InvocationContextInterface:
|
||||
def __init__(self, services: InvocationServices, data: InvocationContextData) -> None:
|
||||
self._services = services
|
||||
self._data = data
|
||||
|
||||
|
||||
class BoardsInterface(InvocationContextInterface):
|
||||
def create(self, board_name: str) -> BoardDTO:
|
||||
"""Creates a board for the current user.
|
||||
|
||||
Args:
|
||||
board_name: The name of the board to create.
|
||||
|
||||
Returns:
|
||||
The created board DTO.
|
||||
"""
|
||||
user_id = self._data.queue_item.user_id
|
||||
return self._services.boards.create(board_name, user_id)
|
||||
|
||||
def get_dto(self, board_id: str) -> BoardDTO:
|
||||
"""Gets a board DTO.
|
||||
|
||||
Args:
|
||||
board_id: The ID of the board to get.
|
||||
|
||||
Returns:
|
||||
The board DTO.
|
||||
"""
|
||||
return self._services.boards.get_dto(board_id)
|
||||
|
||||
def get_all(self) -> list[BoardDTO]:
|
||||
"""Gets all boards accessible to the current user.
|
||||
|
||||
Returns:
|
||||
A list of all boards accessible to the current user.
|
||||
"""
|
||||
user_id = self._data.queue_item.user_id
|
||||
return self._services.boards.get_all(
|
||||
user_id, order_by=BoardRecordOrderBy.CreatedAt, direction=SQLiteDirection.Descending
|
||||
)
|
||||
|
||||
def add_image_to_board(self, board_id: str, image_name: str) -> None:
|
||||
"""Adds an image to a board.
|
||||
|
||||
Args:
|
||||
board_id: The ID of the board to add the image to.
|
||||
image_name: The name of the image to add to the board.
|
||||
"""
|
||||
return self._services.board_images.add_image_to_board(board_id, image_name)
|
||||
|
||||
def get_all_image_names_for_board(self, board_id: str) -> list[str]:
|
||||
"""Gets all image names for a board.
|
||||
|
||||
Args:
|
||||
board_id: The ID of the board to get the image names for.
|
||||
|
||||
Returns:
|
||||
A list of all image names for the board.
|
||||
"""
|
||||
return self._services.board_images.get_all_board_image_names_for_board(
|
||||
board_id,
|
||||
categories=None,
|
||||
is_intermediate=None,
|
||||
)
|
||||
|
||||
|
||||
class LoggerInterface(InvocationContextInterface):
|
||||
def debug(self, message: str) -> None:
|
||||
"""Logs a debug message.
|
||||
|
||||
Args:
|
||||
message: The message to log.
|
||||
"""
|
||||
self._services.logger.debug(message)
|
||||
|
||||
def info(self, message: str) -> None:
|
||||
"""Logs an info message.
|
||||
|
||||
Args:
|
||||
message: The message to log.
|
||||
"""
|
||||
self._services.logger.info(message)
|
||||
|
||||
def warning(self, message: str) -> None:
|
||||
"""Logs a warning message.
|
||||
|
||||
Args:
|
||||
message: The message to log.
|
||||
"""
|
||||
self._services.logger.warning(message)
|
||||
|
||||
def error(self, message: str) -> None:
|
||||
"""Logs an error message.
|
||||
|
||||
Args:
|
||||
message: The message to log.
|
||||
"""
|
||||
self._services.logger.error(message)
|
||||
|
||||
|
||||
class ImagesInterface(InvocationContextInterface):
|
||||
def __init__(self, services: InvocationServices, data: InvocationContextData, util: "UtilInterface") -> None:
|
||||
super().__init__(services, data)
|
||||
self._util = util
|
||||
|
||||
def save(
|
||||
self,
|
||||
image: Image,
|
||||
board_id: Optional[str] = None,
|
||||
image_category: ImageCategory = ImageCategory.GENERAL,
|
||||
metadata: Optional[MetadataField] = None,
|
||||
) -> ImageDTO:
|
||||
"""Saves an image, returning its DTO.
|
||||
|
||||
If the current queue item has a workflow or metadata, it is automatically saved with the image.
|
||||
|
||||
Args:
|
||||
image: The image to save, as a PIL image.
|
||||
board_id: The board ID to add the image to, if it should be added. It the invocation \
|
||||
inherits from `WithBoard`, that board will be used automatically. **Use this only if \
|
||||
you want to override or provide a board manually!**
|
||||
image_category: The category of the image. Only the GENERAL category is added \
|
||||
to the gallery.
|
||||
metadata: The metadata to save with the image, if it should have any. If the \
|
||||
invocation inherits from `WithMetadata`, that metadata will be used automatically. \
|
||||
**Use this only if you want to override or provide metadata manually!**
|
||||
|
||||
Returns:
|
||||
The saved image DTO.
|
||||
"""
|
||||
|
||||
self._util.signal_progress("Saving image")
|
||||
|
||||
# If `metadata` is provided directly, use that. Else, use the metadata provided by `WithMetadata`, falling back to None.
|
||||
metadata_ = None
|
||||
if metadata:
|
||||
metadata_ = metadata.model_dump_json()
|
||||
elif isinstance(self._data.invocation, WithMetadata) and self._data.invocation.metadata:
|
||||
metadata_ = self._data.invocation.metadata.model_dump_json()
|
||||
|
||||
# If `board_id` is provided directly, use that. Else, use the board provided by `WithBoard`, falling back to None.
|
||||
board_id_ = None
|
||||
if board_id:
|
||||
board_id_ = board_id
|
||||
elif isinstance(self._data.invocation, WithBoard) and self._data.invocation.board:
|
||||
board_id_ = self._data.invocation.board.board_id
|
||||
|
||||
workflow_ = None
|
||||
if self._data.queue_item.workflow:
|
||||
workflow_ = self._data.queue_item.workflow.model_dump_json()
|
||||
|
||||
graph_ = None
|
||||
if self._data.queue_item.session.graph:
|
||||
graph_ = self._data.queue_item.session.graph.model_dump_json()
|
||||
|
||||
return self._services.images.create(
|
||||
image=image,
|
||||
is_intermediate=self._data.invocation.is_intermediate,
|
||||
image_category=image_category,
|
||||
board_id=board_id_,
|
||||
metadata=metadata_,
|
||||
image_origin=ResourceOrigin.INTERNAL,
|
||||
workflow=workflow_,
|
||||
graph=graph_,
|
||||
session_id=self._data.queue_item.session_id,
|
||||
node_id=self._data.invocation.id,
|
||||
user_id=self._data.queue_item.user_id,
|
||||
)
|
||||
|
||||
def get_pil(self, image_name: str, mode: IMAGE_MODES | None = None) -> Image:
|
||||
"""Gets an image as a PIL Image object. This method returns a copy of the image.
|
||||
|
||||
Args:
|
||||
image_name: The name of the image to get.
|
||||
mode: The color mode to convert the image to. If None, the original mode is used.
|
||||
|
||||
Returns:
|
||||
The image as a PIL Image object.
|
||||
"""
|
||||
image = self._services.images.get_pil_image(image_name)
|
||||
if mode and mode != image.mode:
|
||||
try:
|
||||
# convert makes a copy!
|
||||
image = image.convert(mode)
|
||||
except ValueError:
|
||||
self._services.logger.warning(
|
||||
f"Could not convert image from {image.mode} to {mode}. Using original mode instead."
|
||||
)
|
||||
else:
|
||||
# copy the image to prevent the user from modifying the original
|
||||
image = image.copy()
|
||||
return image
|
||||
|
||||
def get_metadata(self, image_name: str) -> Optional[MetadataField]:
|
||||
"""Gets an image's metadata, if it has any.
|
||||
|
||||
Args:
|
||||
image_name: The name of the image to get the metadata for.
|
||||
|
||||
Returns:
|
||||
The image's metadata, if it has any.
|
||||
"""
|
||||
return self._services.images.get_metadata(image_name)
|
||||
|
||||
def get_dto(self, image_name: str) -> ImageDTO:
|
||||
"""Gets an image as an ImageDTO object.
|
||||
|
||||
Args:
|
||||
image_name: The name of the image to get.
|
||||
|
||||
Returns:
|
||||
The image as an ImageDTO object.
|
||||
"""
|
||||
return self._services.images.get_dto(image_name)
|
||||
|
||||
def get_path(self, image_name: str, thumbnail: bool = False) -> Path:
|
||||
"""Gets the internal path to an image or thumbnail.
|
||||
|
||||
Args:
|
||||
image_name: The name of the image to get the path of.
|
||||
thumbnail: Get the path of the thumbnail instead of the full image
|
||||
|
||||
Returns:
|
||||
The local path of the image or thumbnail.
|
||||
"""
|
||||
return Path(self._services.images.get_path(image_name, thumbnail))
|
||||
|
||||
|
||||
class TensorsInterface(InvocationContextInterface):
|
||||
def save(self, tensor: Tensor) -> str:
|
||||
"""Saves a tensor, returning its name.
|
||||
|
||||
Args:
|
||||
tensor: The tensor to save.
|
||||
|
||||
Returns:
|
||||
The name of the saved tensor.
|
||||
"""
|
||||
|
||||
name = self._services.tensors.save(obj=tensor)
|
||||
return name
|
||||
|
||||
def load(self, name: str) -> Tensor:
|
||||
"""Loads a tensor by name. This method returns a copy of the tensor.
|
||||
|
||||
Args:
|
||||
name: The name of the tensor to load.
|
||||
|
||||
Returns:
|
||||
The tensor.
|
||||
"""
|
||||
return self._services.tensors.load(name).clone()
|
||||
|
||||
|
||||
class ConditioningInterface(InvocationContextInterface):
|
||||
def save(self, conditioning_data: ConditioningFieldData) -> str:
|
||||
"""Saves a conditioning data object, returning its name.
|
||||
|
||||
Args:
|
||||
conditioning_data: The conditioning data to save.
|
||||
|
||||
Returns:
|
||||
The name of the saved conditioning data.
|
||||
"""
|
||||
|
||||
name = self._services.conditioning.save(obj=conditioning_data)
|
||||
return name
|
||||
|
||||
def load(self, name: str) -> ConditioningFieldData:
|
||||
"""Loads conditioning data by name. This method returns a copy of the conditioning data.
|
||||
|
||||
Args:
|
||||
name: The name of the conditioning data to load.
|
||||
|
||||
Returns:
|
||||
The conditioning data.
|
||||
"""
|
||||
|
||||
return deepcopy(self._services.conditioning.load(name))
|
||||
|
||||
|
||||
class ModelsInterface(InvocationContextInterface):
|
||||
"""Common API for loading, downloading and managing models."""
|
||||
|
||||
def __init__(self, services: InvocationServices, data: InvocationContextData, util: "UtilInterface") -> None:
|
||||
super().__init__(services, data)
|
||||
self._util = util
|
||||
|
||||
def exists(self, identifier: Union[str, "ModelIdentifierField"]) -> bool:
|
||||
"""Check if a model exists.
|
||||
|
||||
Args:
|
||||
identifier: The key or ModelField representing the model.
|
||||
|
||||
Returns:
|
||||
True if the model exists, False if not.
|
||||
"""
|
||||
if isinstance(identifier, str):
|
||||
return self._services.model_manager.store.exists(identifier)
|
||||
else:
|
||||
return self._services.model_manager.store.exists(identifier.key)
|
||||
|
||||
def load(
|
||||
self, identifier: Union[str, "ModelIdentifierField"], submodel_type: Optional[SubModelType] = None
|
||||
) -> LoadedModel:
|
||||
"""Load a model.
|
||||
|
||||
Args:
|
||||
identifier: The key or ModelField representing the model.
|
||||
submodel_type: The submodel of the model to get.
|
||||
|
||||
Returns:
|
||||
An object representing the loaded model.
|
||||
"""
|
||||
|
||||
# The model manager emits events as it loads the model. It needs the context data to build
|
||||
# the event payloads.
|
||||
|
||||
if isinstance(identifier, str):
|
||||
model = self._services.model_manager.store.get_model(identifier)
|
||||
else:
|
||||
submodel_type = submodel_type or identifier.submodel_type
|
||||
model = self._services.model_manager.store.get_model(identifier.key)
|
||||
|
||||
self._raise_if_external(model)
|
||||
|
||||
message = f"Loading model {model.name}"
|
||||
if submodel_type:
|
||||
message += f" ({submodel_type.value})"
|
||||
self._util.signal_progress(message)
|
||||
return self._services.model_manager.load.load_model(model, submodel_type)
|
||||
|
||||
def load_by_attrs(
|
||||
self, name: str, base: BaseModelType, type: ModelType, submodel_type: Optional[SubModelType] = None
|
||||
) -> LoadedModel:
|
||||
"""Load a model by its attributes.
|
||||
|
||||
Args:
|
||||
name: Name of the model.
|
||||
base: The models' base type, e.g. `BaseModelType.StableDiffusion1`, `BaseModelType.StableDiffusionXL`, etc.
|
||||
type: Type of the model, e.g. `ModelType.Main`, `ModelType.Vae`, etc.
|
||||
submodel_type: The type of submodel to load, e.g. `SubModelType.UNet`, `SubModelType.TextEncoder`, etc. Only main
|
||||
models have submodels.
|
||||
|
||||
Returns:
|
||||
An object representing the loaded model.
|
||||
"""
|
||||
|
||||
configs = self._services.model_manager.store.search_by_attr(model_name=name, base_model=base, model_type=type)
|
||||
if len(configs) == 0:
|
||||
raise UnknownModelException(f"No model found with name {name}, base {base}, and type {type}")
|
||||
|
||||
if len(configs) > 1:
|
||||
raise ValueError(f"More than one model found with name {name}, base {base}, and type {type}")
|
||||
|
||||
self._raise_if_external(configs[0])
|
||||
message = f"Loading model {name}"
|
||||
if submodel_type:
|
||||
message += f" ({submodel_type.value})"
|
||||
self._util.signal_progress(message)
|
||||
return self._services.model_manager.load.load_model(configs[0], submodel_type)
|
||||
|
||||
@staticmethod
|
||||
def _raise_if_external(model: AnyModelConfig) -> None:
|
||||
if model.base == BaseModelType.External or model.format == ModelFormat.ExternalApi:
|
||||
raise ValueError("External API models cannot be loaded from disk")
|
||||
|
||||
def get_config(self, identifier: Union[str, "ModelIdentifierField"]) -> AnyModelConfig:
|
||||
"""Get a model's config.
|
||||
|
||||
Args:
|
||||
identifier: The key or ModelField representing the model.
|
||||
|
||||
Returns:
|
||||
The model's config.
|
||||
"""
|
||||
if isinstance(identifier, str):
|
||||
return self._services.model_manager.store.get_model(identifier)
|
||||
else:
|
||||
return self._services.model_manager.store.get_model(identifier.key)
|
||||
|
||||
def search_by_path(self, path: Path) -> list[AnyModelConfig]:
|
||||
"""Search for models by path.
|
||||
|
||||
Args:
|
||||
path: The path to search for.
|
||||
|
||||
Returns:
|
||||
A list of models that match the path.
|
||||
"""
|
||||
return self._services.model_manager.store.search_by_path(path)
|
||||
|
||||
def search_by_attrs(
|
||||
self,
|
||||
name: Optional[str] = None,
|
||||
base: Optional[BaseModelType] = None,
|
||||
type: Optional[ModelType] = None,
|
||||
format: Optional[ModelFormat] = None,
|
||||
) -> list[AnyModelConfig]:
|
||||
"""Search for models by attributes.
|
||||
|
||||
Args:
|
||||
name: The name to search for (exact match).
|
||||
base: The base to search for, e.g. `BaseModelType.StableDiffusion1`, `BaseModelType.StableDiffusionXL`, etc.
|
||||
type: Type type of model to search for, e.g. `ModelType.Main`, `ModelType.Vae`, etc.
|
||||
format: The format of model to search for, e.g. `ModelFormat.Checkpoint`, `ModelFormat.Diffusers`, etc.
|
||||
|
||||
Returns:
|
||||
A list of models that match the attributes.
|
||||
"""
|
||||
|
||||
return self._services.model_manager.store.search_by_attr(
|
||||
model_name=name,
|
||||
base_model=base,
|
||||
model_type=type,
|
||||
model_format=format,
|
||||
)
|
||||
|
||||
def download_and_cache_model(
|
||||
self,
|
||||
source: str | AnyHttpUrl,
|
||||
) -> Path:
|
||||
"""
|
||||
Download the model file located at source to the models cache and return its Path.
|
||||
|
||||
This can be used to single-file install models and other resources of arbitrary types
|
||||
which should not get registered with the database. If the model is already
|
||||
installed, the cached path will be returned. Otherwise it will be downloaded.
|
||||
|
||||
Args:
|
||||
source: A URL that points to the model, or a huggingface repo_id.
|
||||
|
||||
Returns:
|
||||
Path to the downloaded model
|
||||
"""
|
||||
self._util.signal_progress(f"Downloading model {source}")
|
||||
return self._services.model_manager.install.download_and_cache_model(source=source)
|
||||
|
||||
def load_local_model(
|
||||
self,
|
||||
model_path: Path,
|
||||
loader: Optional[Callable[[Path], AnyModel]] = None,
|
||||
) -> LoadedModelWithoutConfig:
|
||||
"""
|
||||
Load the model file located at the indicated path
|
||||
|
||||
If a loader callable is provided, it will be invoked to load the model. Otherwise,
|
||||
`safetensors.torch.load_file()` or `torch.load()` will be called to load the model.
|
||||
|
||||
Be aware that the LoadedModelWithoutConfig object has no `config` attribute
|
||||
|
||||
Args:
|
||||
path: A model Path
|
||||
loader: A Callable that expects a Path and returns a dict[str|int, Any]
|
||||
|
||||
Returns:
|
||||
A LoadedModelWithoutConfig object.
|
||||
"""
|
||||
|
||||
self._util.signal_progress(f"Loading model {model_path.name}")
|
||||
return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
|
||||
|
||||
def load_remote_model(
|
||||
self,
|
||||
source: str | AnyHttpUrl,
|
||||
loader: Optional[Callable[[Path], AnyModel]] = None,
|
||||
) -> LoadedModelWithoutConfig:
|
||||
"""
|
||||
Download, cache, and load the model file located at the indicated URL or repo_id.
|
||||
|
||||
If the model is already downloaded, it will be loaded from the cache.
|
||||
|
||||
If the a loader callable is provided, it will be invoked to load the model. Otherwise,
|
||||
`safetensors.torch.load_file()` or `torch.load()` will be called to load the model.
|
||||
|
||||
Be aware that the LoadedModelWithoutConfig object has no `config` attribute
|
||||
|
||||
Args:
|
||||
source: A URL or huggingface repoid.
|
||||
loader: A Callable that expects a Path and returns a dict[str|int, Any]
|
||||
|
||||
Returns:
|
||||
A LoadedModelWithoutConfig object.
|
||||
"""
|
||||
model_path = self._services.model_manager.install.download_and_cache_model(source=str(source))
|
||||
|
||||
self._util.signal_progress(f"Loading model {source}")
|
||||
return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
|
||||
|
||||
def get_absolute_path(self, config_or_path: AnyModelConfig | Path | str) -> Path:
|
||||
"""Gets the absolute path for a given model config or path.
|
||||
|
||||
For example, if the model's path is `flux/main/FLUX Dev.safetensors`, and the models path is
|
||||
`/home/username/InvokeAI/models`, this method will return
|
||||
`/home/username/InvokeAI/models/flux/main/FLUX Dev.safetensors`.
|
||||
|
||||
Args:
|
||||
config_or_path: The model config or path.
|
||||
|
||||
Returns:
|
||||
The absolute path to the model.
|
||||
"""
|
||||
|
||||
model_path = Path(config_or_path.path) if isinstance(config_or_path, Config_Base) else Path(config_or_path)
|
||||
|
||||
if model_path.is_absolute():
|
||||
return model_path.resolve()
|
||||
|
||||
base_models_path = self._services.configuration.models_path
|
||||
joined_path = base_models_path / model_path
|
||||
resolved_path = joined_path.resolve()
|
||||
return resolved_path
|
||||
|
||||
|
||||
class ConfigInterface(InvocationContextInterface):
|
||||
def get(self) -> InvokeAIAppConfig:
|
||||
"""Gets the app's config.
|
||||
|
||||
Returns:
|
||||
The app's config.
|
||||
"""
|
||||
|
||||
return self._services.configuration
|
||||
|
||||
|
||||
class UtilInterface(InvocationContextInterface):
|
||||
def __init__(
|
||||
self, services: InvocationServices, data: InvocationContextData, is_canceled: Callable[[], bool]
|
||||
) -> None:
|
||||
super().__init__(services, data)
|
||||
self._is_canceled = is_canceled
|
||||
|
||||
def is_canceled(self) -> bool:
|
||||
"""Checks if the current session has been canceled.
|
||||
|
||||
Returns:
|
||||
True if the current session has been canceled, False if not.
|
||||
"""
|
||||
return self._is_canceled()
|
||||
|
||||
def sd_step_callback(self, intermediate_state: PipelineIntermediateState, base_model: BaseModelType) -> None:
|
||||
"""
|
||||
The step callback emits a progress event with the current step, the total number of
|
||||
steps, a preview image, and some other internal metadata.
|
||||
|
||||
This should be called after each denoising step.
|
||||
|
||||
Args:
|
||||
intermediate_state: The intermediate state of the diffusion pipeline.
|
||||
base_model: The base model for the current denoising step.
|
||||
"""
|
||||
|
||||
diffusion_step_callback(
|
||||
signal_progress=self.signal_progress,
|
||||
intermediate_state=intermediate_state,
|
||||
base_model=base_model,
|
||||
is_canceled=self.is_canceled,
|
||||
)
|
||||
|
||||
def flux_step_callback(self, intermediate_state: PipelineIntermediateState) -> None:
|
||||
"""
|
||||
The step callback emits a progress event with the current step, the total number of
|
||||
steps, a preview image, and some other internal metadata.
|
||||
|
||||
This should be called after each denoising step.
|
||||
|
||||
Args:
|
||||
intermediate_state: The intermediate state of the diffusion pipeline.
|
||||
"""
|
||||
|
||||
diffusion_step_callback(
|
||||
signal_progress=self.signal_progress,
|
||||
intermediate_state=intermediate_state,
|
||||
base_model=BaseModelType.Flux,
|
||||
is_canceled=self.is_canceled,
|
||||
)
|
||||
|
||||
def flux2_step_callback(self, intermediate_state: PipelineIntermediateState) -> None:
|
||||
"""
|
||||
The step callback for FLUX.2 Klein models (32-channel VAE).
|
||||
|
||||
Args:
|
||||
intermediate_state: The intermediate state of the diffusion pipeline.
|
||||
"""
|
||||
|
||||
diffusion_step_callback(
|
||||
signal_progress=self.signal_progress,
|
||||
intermediate_state=intermediate_state,
|
||||
base_model=BaseModelType.Flux2,
|
||||
is_canceled=self.is_canceled,
|
||||
)
|
||||
|
||||
def signal_progress(
|
||||
self,
|
||||
message: str,
|
||||
percentage: float | None = None,
|
||||
image: Image | None = None,
|
||||
image_size: tuple[int, int] | None = None,
|
||||
) -> None:
|
||||
"""Signals the progress of some long-running invocation. The progress is displayed in the UI.
|
||||
|
||||
If a percentage is provided, the UI will display a progress bar and automatically append the percentage to the
|
||||
message. You should not include the percentage in the message.
|
||||
|
||||
Example:
|
||||
```py
|
||||
total_steps = 10
|
||||
for i in range(total_steps):
|
||||
percentage = i / (total_steps - 1)
|
||||
context.util.signal_progress("Doing something cool", percentage)
|
||||
```
|
||||
|
||||
If an image is provided, the UI will display it. If your image should be displayed at a different size, provide
|
||||
a tuple of `(width, height)` for the `image_size` parameter. The image will be displayed at the specified size
|
||||
in the UI.
|
||||
|
||||
For example, SD denoising progress images are 1/8 the size of the original image, so you'd do this to ensure the
|
||||
image is displayed at the correct size:
|
||||
```py
|
||||
# Calculate the output size of the image (8x the progress image's size)
|
||||
width = progress_image.width * 8
|
||||
height = progress_image.height * 8
|
||||
# Signal the progress with the image and output size
|
||||
signal_progress("Denoising", percentage, progress_image, (width, height))
|
||||
```
|
||||
|
||||
If your progress image is very large, consider downscaling it to reduce the payload size and provide the original
|
||||
size to the `image_size` parameter. The PIL `thumbnail` method is useful for this, as it maintains the aspect
|
||||
ratio of the image:
|
||||
```py
|
||||
# `thumbnail` modifies the image in-place, so we need to first make a copy
|
||||
thumbnail_image = progress_image.copy()
|
||||
# Resize the image to a maximum of 256x256 pixels, maintaining the aspect ratio
|
||||
thumbnail_image.thumbnail((256, 256))
|
||||
# Signal the progress with the thumbnail, passing the original size
|
||||
signal_progress("Denoising", percentage, thumbnail, progress_image.size)
|
||||
```
|
||||
|
||||
Args:
|
||||
message: A message describing the current status. Do not include the percentage in this message.
|
||||
percentage: The current percentage completion for the process. Omit for indeterminate progress.
|
||||
image: An optional image to display.
|
||||
image_size: The optional size of the image to display. If omitted, the image will be displayed at its
|
||||
original size.
|
||||
"""
|
||||
|
||||
self._services.events.emit_invocation_progress(
|
||||
queue_item=self._data.queue_item,
|
||||
invocation=self._data.invocation,
|
||||
message=message,
|
||||
percentage=percentage,
|
||||
image=ProgressImage.build(image, image_size) if image else None,
|
||||
)
|
||||
|
||||
|
||||
class InvocationContext:
|
||||
"""Provides access to various services and data for the current invocation.
|
||||
|
||||
Attributes:
|
||||
images (ImagesInterface): Methods to save, get and update images and their metadata.
|
||||
tensors (TensorsInterface): Methods to save and get tensors, including image, noise, masks, and masked images.
|
||||
conditioning (ConditioningInterface): Methods to save and get conditioning data.
|
||||
models (ModelsInterface): Methods to check if a model exists, get a model, and get a model's info.
|
||||
logger (LoggerInterface): The app logger.
|
||||
config (ConfigInterface): The app config.
|
||||
util (UtilInterface): Utility methods, including a method to check if an invocation was canceled and step callbacks.
|
||||
boards (BoardsInterface): Methods to interact with boards.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
images: ImagesInterface,
|
||||
tensors: TensorsInterface,
|
||||
conditioning: ConditioningInterface,
|
||||
models: ModelsInterface,
|
||||
logger: LoggerInterface,
|
||||
config: ConfigInterface,
|
||||
util: UtilInterface,
|
||||
boards: BoardsInterface,
|
||||
data: InvocationContextData,
|
||||
services: InvocationServices,
|
||||
) -> None:
|
||||
self.images = images
|
||||
"""Methods to save, get and update images and their metadata."""
|
||||
self.tensors = tensors
|
||||
"""Methods to save and get tensors, including image, noise, masks, and masked images."""
|
||||
self.conditioning = conditioning
|
||||
"""Methods to save and get conditioning data."""
|
||||
self.models = models
|
||||
"""Methods to check if a model exists, get a model, and get a model's info."""
|
||||
self.logger = logger
|
||||
"""The app logger."""
|
||||
self.config = config
|
||||
"""The app config."""
|
||||
self.util = util
|
||||
"""Utility methods, including a method to check if an invocation was canceled and step callbacks."""
|
||||
self.boards = boards
|
||||
"""Methods to interact with boards."""
|
||||
self._data = data
|
||||
"""An internal API providing access to data about the current queue item and invocation. You probably shouldn't use this. It may change without warning."""
|
||||
self._services = services
|
||||
"""An internal API providing access to all application services. You probably shouldn't use this. It may change without warning."""
|
||||
|
||||
|
||||
def build_invocation_context(
|
||||
services: InvocationServices,
|
||||
data: InvocationContextData,
|
||||
is_canceled: Callable[[], bool],
|
||||
) -> InvocationContext:
|
||||
"""Builds the invocation context for a specific invocation execution.
|
||||
|
||||
Args:
|
||||
services: The invocation services to wrap.
|
||||
data: The invocation context data.
|
||||
|
||||
Returns:
|
||||
The invocation context.
|
||||
"""
|
||||
|
||||
logger = LoggerInterface(services=services, data=data)
|
||||
tensors = TensorsInterface(services=services, data=data)
|
||||
config = ConfigInterface(services=services, data=data)
|
||||
util = UtilInterface(services=services, data=data, is_canceled=is_canceled)
|
||||
conditioning = ConditioningInterface(services=services, data=data)
|
||||
models = ModelsInterface(services=services, data=data, util=util)
|
||||
images = ImagesInterface(services=services, data=data, util=util)
|
||||
boards = BoardsInterface(services=services, data=data)
|
||||
|
||||
ctx = InvocationContext(
|
||||
images=images,
|
||||
logger=logger,
|
||||
config=config,
|
||||
tensors=tensors,
|
||||
models=models,
|
||||
data=data,
|
||||
util=util,
|
||||
conditioning=conditioning,
|
||||
services=services,
|
||||
boards=boards,
|
||||
)
|
||||
|
||||
return ctx
|
||||
@@ -0,0 +1,41 @@
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
GenericBaseModel = TypeVar("GenericBaseModel", bound=BaseModel)
|
||||
|
||||
|
||||
class CursorPaginatedResults(BaseModel, Generic[GenericBaseModel]):
|
||||
"""
|
||||
Cursor-paginated results
|
||||
Generic must be a Pydantic model
|
||||
"""
|
||||
|
||||
limit: int = Field(..., description="Limit of items to get")
|
||||
has_more: bool = Field(..., description="Whether there are more items available")
|
||||
items: list[GenericBaseModel] = Field(..., description="Items")
|
||||
|
||||
|
||||
class OffsetPaginatedResults(BaseModel, Generic[GenericBaseModel]):
|
||||
"""
|
||||
Offset-paginated results
|
||||
Generic must be a Pydantic model
|
||||
"""
|
||||
|
||||
limit: int = Field(description="Limit of items to get")
|
||||
offset: int = Field(description="Offset from which to retrieve items")
|
||||
total: int = Field(description="Total number of items in result")
|
||||
items: list[GenericBaseModel] = Field(description="Items")
|
||||
|
||||
|
||||
class PaginatedResults(BaseModel, Generic[GenericBaseModel]):
|
||||
"""
|
||||
Paginated results
|
||||
Generic must be a Pydantic model
|
||||
"""
|
||||
|
||||
page: int = Field(description="Current Page")
|
||||
pages: int = Field(description="Total number of pages")
|
||||
per_page: int = Field(description="Number of items per page")
|
||||
total: int = Field(description="Total number of items in result")
|
||||
items: list[GenericBaseModel] = Field(description="Items")
|
||||
@@ -0,0 +1,10 @@
|
||||
from enum import Enum
|
||||
|
||||
from invokeai.app.util.metaenum import MetaEnum
|
||||
|
||||
sqlite_memory = ":memory:"
|
||||
|
||||
|
||||
class SQLiteDirection(str, Enum, metaclass=MetaEnum):
|
||||
Ascending = "ASC"
|
||||
Descending = "DESC"
|
||||
@@ -0,0 +1,93 @@
|
||||
import sqlite3
|
||||
import threading
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from logging import Logger
|
||||
from pathlib import Path
|
||||
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import sqlite_memory
|
||||
|
||||
|
||||
class SqliteDatabase:
|
||||
"""
|
||||
Manages a connection to an SQLite database.
|
||||
|
||||
:param db_path: Path to the database file. If None, an in-memory database is used.
|
||||
:param logger: Logger to use for logging.
|
||||
:param verbose: Whether to log SQL statements. Provides `logger.debug` as the SQLite trace callback.
|
||||
|
||||
This is a light wrapper around the `sqlite3` module, providing a few conveniences:
|
||||
- The database file is written to disk if it does not exist.
|
||||
- Foreign key constraints are enabled by default.
|
||||
- The connection is configured to use the `sqlite3.Row` row factory.
|
||||
|
||||
In addition to the constructor args, the instance provides the following attributes and methods:
|
||||
- `conn`: A `sqlite3.Connection` object. Note that the connection must never be closed if the database is in-memory.
|
||||
- `lock`: A shared re-entrant lock, used to approximate thread safety.
|
||||
- `clean()`: Runs the SQL `VACUUM;` command and reports on the freed space.
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: Path | None, logger: Logger, verbose: bool = False) -> None:
|
||||
"""Initializes the database. This is used internally by the class constructor."""
|
||||
self._logger = logger
|
||||
self._db_path = db_path
|
||||
self._verbose = verbose
|
||||
self._lock = threading.RLock()
|
||||
|
||||
if not self._db_path:
|
||||
logger.info("Initializing in-memory database")
|
||||
else:
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._logger.info(f"Initializing database at {self._db_path}")
|
||||
|
||||
self._conn = sqlite3.connect(database=self._db_path or sqlite_memory, check_same_thread=False)
|
||||
self._conn.row_factory = sqlite3.Row
|
||||
|
||||
if self._verbose:
|
||||
self._conn.set_trace_callback(self._logger.debug)
|
||||
|
||||
# Enable foreign key constraints
|
||||
self._conn.execute("PRAGMA foreign_keys = ON;")
|
||||
|
||||
# Enable Write-Ahead Logging (WAL) mode for better concurrency
|
||||
self._conn.execute("PRAGMA journal_mode = WAL;")
|
||||
|
||||
# Set a busy timeout to prevent database lockups during writes
|
||||
self._conn.execute("PRAGMA busy_timeout = 5000;") # 5 seconds
|
||||
|
||||
def clean(self) -> None:
|
||||
"""
|
||||
Cleans the database by running the VACUUM command, reporting on the freed space.
|
||||
"""
|
||||
# No need to clean in-memory database
|
||||
if not self._db_path:
|
||||
return
|
||||
try:
|
||||
with self._conn as conn:
|
||||
initial_db_size = Path(self._db_path).stat().st_size
|
||||
conn.execute("VACUUM;")
|
||||
conn.commit()
|
||||
final_db_size = Path(self._db_path).stat().st_size
|
||||
freed_space_in_mb = round((initial_db_size - final_db_size) / 1024 / 1024, 2)
|
||||
if freed_space_in_mb > 0:
|
||||
self._logger.info(f"Cleaned database (freed {freed_space_in_mb}MB)")
|
||||
except Exception as e:
|
||||
self._logger.error(f"Error cleaning database: {e}")
|
||||
raise
|
||||
|
||||
@contextmanager
|
||||
def transaction(self) -> Generator[sqlite3.Cursor, None, None]:
|
||||
"""
|
||||
Thread-safe context manager for DB work.
|
||||
Acquires the RLock, yields a Cursor, then commits or rolls back.
|
||||
"""
|
||||
with self._lock:
|
||||
cursor = self._conn.cursor()
|
||||
try:
|
||||
yield cursor
|
||||
self._conn.commit()
|
||||
except Exception:
|
||||
self._conn.rollback()
|
||||
raise
|
||||
finally:
|
||||
cursor.close()
|
||||
@@ -0,0 +1,32 @@
|
||||
from logging import Logger
|
||||
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
from invokeai.app.services.shared.sqlite_migrator.migration_loader import MigrationBuildContext, build_migrations
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator
|
||||
|
||||
|
||||
def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileStorageBase) -> SqliteDatabase:
|
||||
"""
|
||||
Initializes the SQLite database.
|
||||
|
||||
:param config: The app config
|
||||
:param logger: The logger
|
||||
:param image_files: The image files service (used by migration 2)
|
||||
|
||||
This function:
|
||||
- Instantiates a :class:`SqliteDatabase`
|
||||
- Instantiates a :class:`SqliteMigrator` and registers all migrations
|
||||
- Runs all migrations
|
||||
"""
|
||||
db_path = None if config.use_memory_db else config.db_path
|
||||
db = SqliteDatabase(db_path=db_path, logger=logger, verbose=config.log_sql)
|
||||
|
||||
migrator = SqliteMigrator(db=db)
|
||||
migration_context = MigrationBuildContext(app_config=config, logger=logger, image_files=image_files)
|
||||
for migration in build_migrations(migration_context):
|
||||
migrator.register_migration(migration)
|
||||
migrator.run_migrations()
|
||||
|
||||
return db
|
||||
@@ -0,0 +1,136 @@
|
||||
import importlib
|
||||
import inspect
|
||||
import pkgutil
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from logging import Logger
|
||||
from types import ModuleType
|
||||
from typing import Any
|
||||
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
DEFAULT_MIGRATIONS_PACKAGE = "invokeai.app.services.shared.sqlite_migrator.migrations"
|
||||
_MIGRATION_MODULE_RE = re.compile(r"^migration_(\d+)$")
|
||||
_DATED_MIGRATION_MODULE_RE = re.compile(r"^migration_(\d{4}_\d{2}_\d{2}_[a-z0-9][a-z0-9_]*)$")
|
||||
|
||||
|
||||
class MigrationLoaderError(RuntimeError):
|
||||
"""Raised when migration discovery or construction fails."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MigrationBuildContext:
|
||||
app_config: InvokeAIAppConfig
|
||||
logger: Logger
|
||||
image_files: ImageFileStorageBase
|
||||
|
||||
def get_dependency(self, name: str) -> Any:
|
||||
if name == "app_config":
|
||||
return self.app_config
|
||||
if name == "config":
|
||||
return self.app_config
|
||||
if name == "logger":
|
||||
return self.logger
|
||||
if name == "image_files":
|
||||
return self.image_files
|
||||
raise MigrationLoaderError(f"Migration builder requested unknown dependency '{name}'")
|
||||
|
||||
|
||||
MigrationBuilder = Callable[..., Migration]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DiscoveredMigrationBuilder:
|
||||
module_name: str
|
||||
sort_key: tuple[int, int | str]
|
||||
expected_migration_id: str
|
||||
builder: MigrationBuilder
|
||||
|
||||
|
||||
def discover_migration_builders(package_name: str = DEFAULT_MIGRATIONS_PACKAGE) -> list[DiscoveredMigrationBuilder]:
|
||||
try:
|
||||
package = importlib.import_module(package_name)
|
||||
except Exception as e:
|
||||
raise MigrationLoaderError(f"Unable to import migrations package '{package_name}': {e}") from e
|
||||
|
||||
package_path = getattr(package, "__path__", None)
|
||||
if package_path is None:
|
||||
raise MigrationLoaderError(f"Migrations package '{package_name}' is not a package")
|
||||
|
||||
discovered: list[DiscoveredMigrationBuilder] = []
|
||||
for module_info in pkgutil.iter_modules(package_path):
|
||||
module_short_name = module_info.name
|
||||
if not module_short_name.startswith("migration_"):
|
||||
continue
|
||||
numeric_match = _MIGRATION_MODULE_RE.match(module_short_name)
|
||||
dated_match = _DATED_MIGRATION_MODULE_RE.match(module_short_name)
|
||||
if numeric_match is None and dated_match is None:
|
||||
raise MigrationLoaderError(f"Malformed migration module name '{module_short_name}'")
|
||||
|
||||
if numeric_match is not None:
|
||||
sort_key: tuple[int, int | str] = (0, int(numeric_match.group(1)))
|
||||
expected_builder_name = f"build_migration_{numeric_match.group(1)}"
|
||||
expected_migration_id = f"migration_{numeric_match.group(1)}"
|
||||
else:
|
||||
sort_key = (1, module_short_name)
|
||||
expected_builder_name = "build_migration"
|
||||
expected_migration_id = module_short_name.removeprefix("migration_")
|
||||
|
||||
full_module_name = f"{package_name}.{module_short_name}"
|
||||
try:
|
||||
module = importlib.import_module(full_module_name)
|
||||
except Exception as e:
|
||||
raise MigrationLoaderError(f"Unable to import migration module '{full_module_name}': {e}") from e
|
||||
|
||||
builder = _get_builder(module, expected_builder_name)
|
||||
discovered.append(DiscoveredMigrationBuilder(full_module_name, sort_key, expected_migration_id, builder))
|
||||
|
||||
return sorted(discovered, key=lambda discovered_builder: discovered_builder.sort_key)
|
||||
|
||||
|
||||
def build_migrations(context: MigrationBuildContext, package_name: str = DEFAULT_MIGRATIONS_PACKAGE) -> list[Migration]:
|
||||
migrations: list[Migration] = []
|
||||
for discovered_builder in discover_migration_builders(package_name):
|
||||
kwargs = _get_builder_kwargs(discovered_builder.builder, context, discovered_builder.module_name)
|
||||
migration = discovered_builder.builder(**kwargs)
|
||||
if not isinstance(migration, Migration):
|
||||
raise MigrationLoaderError(
|
||||
f"Migration builder '{discovered_builder.builder.__name__}' in "
|
||||
f"'{discovered_builder.module_name}' must return Migration"
|
||||
)
|
||||
if migration.id != discovered_builder.expected_migration_id:
|
||||
raise MigrationLoaderError(
|
||||
f"Migration builder '{discovered_builder.builder.__name__}' in "
|
||||
f"'{discovered_builder.module_name}' must return migration id "
|
||||
f"'{discovered_builder.expected_migration_id}', got '{migration.id}'"
|
||||
)
|
||||
migrations.append(migration)
|
||||
return migrations
|
||||
|
||||
|
||||
def _get_builder(module: ModuleType, builder_name: str) -> MigrationBuilder:
|
||||
builder = getattr(module, builder_name, None)
|
||||
if builder is None:
|
||||
raise MigrationLoaderError(f"Migration module '{module.__name__}' must define '{builder_name}'")
|
||||
if not callable(builder):
|
||||
raise MigrationLoaderError(f"Migration builder '{module.__name__}.{builder_name}' is not callable")
|
||||
return builder
|
||||
|
||||
|
||||
def _get_builder_kwargs(builder: MigrationBuilder, context: MigrationBuildContext, module_name: str) -> dict[str, Any]:
|
||||
signature = inspect.signature(builder)
|
||||
kwargs: dict[str, Any] = {}
|
||||
for parameter in signature.parameters.values():
|
||||
if parameter.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
|
||||
raise MigrationLoaderError(
|
||||
f"Migration builder '{module_name}.{builder.__name__}' must not use *args or **kwargs"
|
||||
)
|
||||
if parameter.kind == inspect.Parameter.POSITIONAL_ONLY:
|
||||
raise MigrationLoaderError(
|
||||
f"Migration builder '{module_name}.{builder.__name__}' must not use positional-only parameters"
|
||||
)
|
||||
kwargs[parameter.name] = context.get_dependency(parameter.name)
|
||||
return kwargs
|
||||
@@ -0,0 +1,374 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration1Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Migration callback for database version 1."""
|
||||
|
||||
self._create_board_images(cursor)
|
||||
self._create_boards(cursor)
|
||||
self._create_images(cursor)
|
||||
self._create_model_config(cursor)
|
||||
self._create_session_queue(cursor)
|
||||
self._create_workflow_images(cursor)
|
||||
self._create_workflows(cursor)
|
||||
|
||||
def _create_board_images(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the `board_images` table, indices and triggers."""
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS board_images (
|
||||
board_id TEXT NOT NULL,
|
||||
image_name TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Soft delete, currently unused
|
||||
deleted_at DATETIME,
|
||||
-- enforce one-to-many relationship between boards and images using PK
|
||||
-- (we can extend this to many-to-many later)
|
||||
PRIMARY KEY (image_name),
|
||||
FOREIGN KEY (board_id) REFERENCES boards (board_id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (image_name) REFERENCES images (image_name) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
indices = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_board_images_board_id ON board_images (board_id);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_board_images_board_id_created_at ON board_images (board_id, created_at);",
|
||||
]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_board_images_updated_at
|
||||
AFTER UPDATE
|
||||
ON board_images FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE board_images SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE board_id = old.board_id AND image_name = old.image_name;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _create_boards(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the `boards` table, indices and triggers."""
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS boards (
|
||||
board_id TEXT NOT NULL PRIMARY KEY,
|
||||
board_name TEXT NOT NULL,
|
||||
cover_image_name TEXT,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Soft delete, currently unused
|
||||
deleted_at DATETIME,
|
||||
FOREIGN KEY (cover_image_name) REFERENCES images (image_name) ON DELETE SET NULL
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
indices = ["CREATE INDEX IF NOT EXISTS idx_boards_created_at ON boards (created_at);"]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_boards_updated_at
|
||||
AFTER UPDATE
|
||||
ON boards FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE boards SET updated_at = current_timestamp
|
||||
WHERE board_id = old.board_id;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _create_images(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the `images` table, indices and triggers. Adds the `starred` column."""
|
||||
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS images (
|
||||
image_name TEXT NOT NULL PRIMARY KEY,
|
||||
-- This is an enum in python, unrestricted string here for flexibility
|
||||
image_origin TEXT NOT NULL,
|
||||
-- This is an enum in python, unrestricted string here for flexibility
|
||||
image_category TEXT NOT NULL,
|
||||
width INTEGER NOT NULL,
|
||||
height INTEGER NOT NULL,
|
||||
session_id TEXT,
|
||||
node_id TEXT,
|
||||
metadata TEXT,
|
||||
is_intermediate BOOLEAN DEFAULT FALSE,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Soft delete, currently unused
|
||||
deleted_at DATETIME
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
indices = [
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS idx_images_image_name ON images(image_name);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_images_image_origin ON images(image_origin);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_images_image_category ON images(image_category);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_images_created_at ON images(created_at);",
|
||||
]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_images_updated_at
|
||||
AFTER UPDATE
|
||||
ON images FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE images SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE image_name = old.image_name;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
# Add the 'starred' column to `images` if it doesn't exist
|
||||
cursor.execute("PRAGMA table_info(images)")
|
||||
columns = [column[1] for column in cursor.fetchall()]
|
||||
|
||||
if "starred" not in columns:
|
||||
tables.append("ALTER TABLE images ADD COLUMN starred BOOLEAN DEFAULT FALSE;")
|
||||
indices.append("CREATE INDEX IF NOT EXISTS idx_images_starred ON images(starred);")
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _create_model_config(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the `model_config` table, `model_manager_metadata` table, indices and triggers."""
|
||||
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS model_config (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
-- The next 3 fields are enums in python, unrestricted string here
|
||||
base TEXT NOT NULL,
|
||||
type TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
path TEXT NOT NULL,
|
||||
original_hash TEXT, -- could be null
|
||||
-- Serialized JSON representation of the whole config object,
|
||||
-- which will contain additional fields from subclasses
|
||||
config TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- unique constraint on combo of name, base and type
|
||||
UNIQUE(name, base, type)
|
||||
);
|
||||
""",
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS model_manager_metadata (
|
||||
metadata_key TEXT NOT NULL PRIMARY KEY,
|
||||
metadata_value TEXT NOT NULL
|
||||
);
|
||||
""",
|
||||
]
|
||||
|
||||
# Add trigger for `updated_at`.
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS model_config_updated_at
|
||||
AFTER UPDATE
|
||||
ON model_config FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE model_config SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
# Add indexes for searchable fields
|
||||
indices = [
|
||||
"CREATE INDEX IF NOT EXISTS base_index ON model_config(base);",
|
||||
"CREATE INDEX IF NOT EXISTS type_index ON model_config(type);",
|
||||
"CREATE INDEX IF NOT EXISTS name_index ON model_config(name);",
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS path_index ON model_config(path);",
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _create_session_queue(self, cursor: sqlite3.Cursor) -> None:
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS session_queue (
|
||||
item_id INTEGER PRIMARY KEY AUTOINCREMENT, -- used for ordering, cursor pagination
|
||||
batch_id TEXT NOT NULL, -- identifier of the batch this queue item belongs to
|
||||
queue_id TEXT NOT NULL, -- identifier of the queue this queue item belongs to
|
||||
session_id TEXT NOT NULL UNIQUE, -- duplicated data from the session column, for ease of access
|
||||
field_values TEXT, -- NULL if no values are associated with this queue item
|
||||
session TEXT NOT NULL, -- the session to be executed
|
||||
status TEXT NOT NULL DEFAULT 'pending', -- the status of the queue item, one of 'pending', 'in_progress', 'completed', 'failed', 'canceled'
|
||||
priority INTEGER NOT NULL DEFAULT 0, -- the priority, higher is more important
|
||||
error TEXT, -- any errors associated with this queue item
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')), -- updated via trigger
|
||||
started_at DATETIME, -- updated via trigger
|
||||
completed_at DATETIME -- updated via trigger, completed items are cleaned up on application startup
|
||||
-- Ideally this is a FK, but graph_executions uses INSERT OR REPLACE, and REPLACE triggers the ON DELETE CASCADE...
|
||||
-- FOREIGN KEY (session_id) REFERENCES graph_executions (id) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
indices = [
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS idx_session_queue_item_id ON session_queue(item_id);",
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS idx_session_queue_session_id ON session_queue(session_id);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_batch_id ON session_queue(batch_id);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_created_priority ON session_queue(priority);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_created_status ON session_queue(status);",
|
||||
]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_session_queue_completed_at
|
||||
AFTER UPDATE OF status ON session_queue
|
||||
FOR EACH ROW
|
||||
WHEN
|
||||
NEW.status = 'completed'
|
||||
OR NEW.status = 'failed'
|
||||
OR NEW.status = 'canceled'
|
||||
BEGIN
|
||||
UPDATE session_queue
|
||||
SET completed_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE item_id = NEW.item_id;
|
||||
END;
|
||||
""",
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_session_queue_started_at
|
||||
AFTER UPDATE OF status ON session_queue
|
||||
FOR EACH ROW
|
||||
WHEN
|
||||
NEW.status = 'in_progress'
|
||||
BEGIN
|
||||
UPDATE session_queue
|
||||
SET started_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE item_id = NEW.item_id;
|
||||
END;
|
||||
""",
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_session_queue_updated_at
|
||||
AFTER UPDATE
|
||||
ON session_queue FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE session_queue
|
||||
SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE item_id = old.item_id;
|
||||
END;
|
||||
""",
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _create_workflow_images(self, cursor: sqlite3.Cursor) -> None:
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS workflow_images (
|
||||
workflow_id TEXT NOT NULL,
|
||||
image_name TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Soft delete, currently unused
|
||||
deleted_at DATETIME,
|
||||
-- enforce one-to-many relationship between workflows and images using PK
|
||||
-- (we can extend this to many-to-many later)
|
||||
PRIMARY KEY (image_name),
|
||||
FOREIGN KEY (workflow_id) REFERENCES workflows (workflow_id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (image_name) REFERENCES images (image_name) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
indices = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_images_workflow_id ON workflow_images (workflow_id);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_images_workflow_id_created_at ON workflow_images (workflow_id, created_at);",
|
||||
]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_workflow_images_updated_at
|
||||
AFTER UPDATE
|
||||
ON workflow_images FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE workflow_images SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE workflow_id = old.workflow_id AND image_name = old.image_name;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _create_workflows(self, cursor: sqlite3.Cursor) -> None:
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS workflows (
|
||||
workflow TEXT NOT NULL,
|
||||
workflow_id TEXT GENERATED ALWAYS AS (json_extract(workflow, '$.id')) VIRTUAL NOT NULL UNIQUE, -- gets implicit index
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')) -- updated via trigger
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_workflows_updated_at
|
||||
AFTER UPDATE
|
||||
ON workflows FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE workflows
|
||||
SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE workflow_id = old.workflow_id;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
for stmt in tables + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
|
||||
def build_migration_1() -> Migration:
|
||||
"""
|
||||
Builds the migration from database version 0 (init) to 1.
|
||||
|
||||
This migration represents the state of the database circa InvokeAI v3.4.0, which was the last
|
||||
version to not use migrations to manage the database.
|
||||
|
||||
As such, this migration does include some ALTER statements, and the SQL statements are written
|
||||
to be idempotent.
|
||||
|
||||
- Create `board_images` junction table
|
||||
- Create `boards` table
|
||||
- Create `images` table, add `starred` column
|
||||
- Create `model_config` table
|
||||
- Create `session_queue` table
|
||||
- Create `workflow_images` junction table
|
||||
- Create `workflows` table
|
||||
"""
|
||||
|
||||
migration_1 = Migration(
|
||||
id="migration_1",
|
||||
depends_on=None,
|
||||
from_version=0,
|
||||
to_version=1,
|
||||
callback=Migration1Callback(),
|
||||
)
|
||||
|
||||
return migration_1
|
||||
@@ -0,0 +1,37 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration10Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._update_error_cols(cursor)
|
||||
|
||||
def _update_error_cols(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
- Adds `error_type` and `error_message` columns to the session queue table.
|
||||
- Renames the `error` column to `error_traceback`.
|
||||
"""
|
||||
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN error_type TEXT;")
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN error_message TEXT;")
|
||||
cursor.execute("ALTER TABLE session_queue RENAME COLUMN error TO error_traceback;")
|
||||
|
||||
|
||||
def build_migration_10() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 9 to 10.
|
||||
|
||||
This migration does the following:
|
||||
- Adds `error_type` and `error_message` columns to the session queue table.
|
||||
- Renames the `error` column to `error_traceback`.
|
||||
"""
|
||||
migration_10 = Migration(
|
||||
id="migration_10",
|
||||
depends_on="migration_9",
|
||||
from_version=9,
|
||||
to_version=10,
|
||||
callback=Migration10Callback(),
|
||||
)
|
||||
|
||||
return migration_10
|
||||
@@ -0,0 +1,77 @@
|
||||
import shutil
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
LEGACY_CORE_MODELS = [
|
||||
# OpenPose
|
||||
"any/annotators/dwpose/yolox_l.onnx",
|
||||
"any/annotators/dwpose/dw-ll_ucoco_384.onnx",
|
||||
# DepthAnything
|
||||
"any/annotators/depth_anything/depth_anything_vitl14.pth",
|
||||
"any/annotators/depth_anything/depth_anything_vitb14.pth",
|
||||
"any/annotators/depth_anything/depth_anything_vits14.pth",
|
||||
# Lama inpaint
|
||||
"core/misc/lama/lama.pt",
|
||||
# RealESRGAN upscale
|
||||
"core/upscaling/realesrgan/RealESRGAN_x4plus.pth",
|
||||
"core/upscaling/realesrgan/RealESRGAN_x4plus_anime_6B.pth",
|
||||
"core/upscaling/realesrgan/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth",
|
||||
"core/upscaling/realesrgan/RealESRGAN_x2plus.pth",
|
||||
]
|
||||
|
||||
|
||||
class Migration11Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._remove_convert_cache()
|
||||
self._remove_downloaded_models()
|
||||
self._remove_unused_core_models()
|
||||
|
||||
def _remove_convert_cache(self) -> None:
|
||||
"""Rename models/.cache to models/.convert_cache."""
|
||||
self._logger.info("Removing models/.cache directory. Converted models will now be cached in .convert_cache.")
|
||||
legacy_convert_path = self._app_config.root_path / "models" / ".cache"
|
||||
shutil.rmtree(legacy_convert_path, ignore_errors=True)
|
||||
|
||||
def _remove_downloaded_models(self) -> None:
|
||||
"""Remove models from their old locations; they will re-download when needed."""
|
||||
self._logger.info(
|
||||
"Removing legacy just-in-time models. Downloaded models will now be cached in .download_cache."
|
||||
)
|
||||
for model_path in LEGACY_CORE_MODELS:
|
||||
legacy_dest_path = self._app_config.models_path / model_path
|
||||
legacy_dest_path.unlink(missing_ok=True)
|
||||
|
||||
def _remove_unused_core_models(self) -> None:
|
||||
"""Remove unused core models and their directories."""
|
||||
self._logger.info("Removing defunct core models.")
|
||||
for dir in ["face_restoration", "misc", "upscaling"]:
|
||||
path_to_remove = self._app_config.models_path / "core" / dir
|
||||
shutil.rmtree(path_to_remove, ignore_errors=True)
|
||||
shutil.rmtree(self._app_config.models_path / "any" / "annotators", ignore_errors=True)
|
||||
|
||||
|
||||
def build_migration_11(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""
|
||||
Build the migration from database version 10 to 11.
|
||||
|
||||
This migration does the following:
|
||||
- Moves "core" models previously downloaded with download_with_progress_bar() into new
|
||||
"models/.download_cache" directory.
|
||||
- Renames "models/.cache" to "models/.convert_cache".
|
||||
"""
|
||||
migration_11 = Migration(
|
||||
id="migration_11",
|
||||
depends_on="migration_10",
|
||||
from_version=10,
|
||||
to_version=11,
|
||||
callback=Migration11Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
|
||||
return migration_11
|
||||
@@ -0,0 +1,37 @@
|
||||
import shutil
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration12Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig) -> None:
|
||||
self._app_config = app_config
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._remove_model_convert_cache_dir()
|
||||
|
||||
def _remove_model_convert_cache_dir(self) -> None:
|
||||
"""
|
||||
Removes unused model convert cache directory
|
||||
"""
|
||||
convert_cache = self._app_config.convert_cache_path
|
||||
shutil.rmtree(convert_cache, ignore_errors=True)
|
||||
|
||||
|
||||
def build_migration_12(app_config: InvokeAIAppConfig) -> Migration:
|
||||
"""
|
||||
Build the migration from database version 11 to 12.
|
||||
|
||||
This migration removes the now-unused model convert cache directory.
|
||||
"""
|
||||
migration_12 = Migration(
|
||||
id="migration_12",
|
||||
depends_on="migration_11",
|
||||
from_version=11,
|
||||
to_version=12,
|
||||
callback=Migration12Callback(app_config),
|
||||
)
|
||||
|
||||
return migration_12
|
||||
@@ -0,0 +1,33 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration13Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._add_archived_col(cursor)
|
||||
|
||||
def _add_archived_col(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
- Adds `archived` columns to the board table.
|
||||
"""
|
||||
|
||||
cursor.execute("ALTER TABLE boards ADD COLUMN archived BOOLEAN DEFAULT FALSE;")
|
||||
|
||||
|
||||
def build_migration_13() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 12 to 13..
|
||||
|
||||
This migration does the following:
|
||||
- Adds `archived` columns to the board table.
|
||||
"""
|
||||
migration_13 = Migration(
|
||||
id="migration_13",
|
||||
depends_on="migration_12",
|
||||
from_version=12,
|
||||
to_version=13,
|
||||
callback=Migration13Callback(),
|
||||
)
|
||||
|
||||
return migration_13
|
||||
@@ -0,0 +1,63 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration14Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._create_style_presets(cursor)
|
||||
|
||||
def _create_style_presets(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create the table used to store style presets."""
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS style_presets (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
preset_data TEXT NOT NULL,
|
||||
type TEXT NOT NULL DEFAULT "user",
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW'))
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
# Add trigger for `updated_at`.
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS style_presets
|
||||
AFTER UPDATE
|
||||
ON style_presets FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE style_presets SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
# Add indexes for searchable fields
|
||||
indices = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_style_presets_name ON style_presets(name);",
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
|
||||
def build_migration_14() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 13 to 14..
|
||||
|
||||
This migration does the following:
|
||||
- Create the table used to store style presets.
|
||||
"""
|
||||
migration_14 = Migration(
|
||||
id="migration_14",
|
||||
depends_on="migration_13",
|
||||
from_version=13,
|
||||
to_version=14,
|
||||
callback=Migration14Callback(),
|
||||
)
|
||||
|
||||
return migration_14
|
||||
@@ -0,0 +1,36 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration15Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._add_origin_col(cursor)
|
||||
|
||||
def _add_origin_col(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
- Adds `origin` column to the session queue table.
|
||||
- Adds `destination` column to the session queue table.
|
||||
"""
|
||||
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN origin TEXT;")
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN destination TEXT;")
|
||||
|
||||
|
||||
def build_migration_15() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 14 to 15.
|
||||
|
||||
This migration does the following:
|
||||
- Adds `origin` column to the session queue table.
|
||||
- Adds `destination` column to the session queue table.
|
||||
"""
|
||||
migration_15 = Migration(
|
||||
id="migration_15",
|
||||
depends_on="migration_14",
|
||||
from_version=14,
|
||||
to_version=15,
|
||||
callback=Migration15Callback(),
|
||||
)
|
||||
|
||||
return migration_15
|
||||
@@ -0,0 +1,33 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration16Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._add_retried_from_item_id_col(cursor)
|
||||
|
||||
def _add_retried_from_item_id_col(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
- Adds `retried_from_item_id` column to the session queue table.
|
||||
"""
|
||||
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN retried_from_item_id INTEGER;")
|
||||
|
||||
|
||||
def build_migration_16() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 15 to 16.
|
||||
|
||||
This migration does the following:
|
||||
- Adds `retried_from_item_id` column to the session queue table.
|
||||
"""
|
||||
migration_16 = Migration(
|
||||
id="migration_16",
|
||||
depends_on="migration_15",
|
||||
from_version=15,
|
||||
to_version=16,
|
||||
callback=Migration16Callback(),
|
||||
)
|
||||
|
||||
return migration_16
|
||||
@@ -0,0 +1,37 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration17Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._add_workflows_tags_col(cursor)
|
||||
|
||||
def _add_workflows_tags_col(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
- Adds `tags` column to the workflow_library table. It is a generated column that extracts the tags from the
|
||||
workflow JSON.
|
||||
"""
|
||||
|
||||
cursor.execute(
|
||||
"ALTER TABLE workflow_library ADD COLUMN tags TEXT GENERATED ALWAYS AS (json_extract(workflow, '$.tags')) VIRTUAL;"
|
||||
)
|
||||
|
||||
|
||||
def build_migration_17() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 16 to 17.
|
||||
|
||||
This migration does the following:
|
||||
- Adds `tags` column to the workflow_library table. It is a generated column that extracts the tags from the
|
||||
workflow JSON.
|
||||
"""
|
||||
migration_17 = Migration(
|
||||
id="migration_17",
|
||||
depends_on="migration_16",
|
||||
from_version=16,
|
||||
to_version=17,
|
||||
callback=Migration17Callback(),
|
||||
)
|
||||
|
||||
return migration_17
|
||||
@@ -0,0 +1,49 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration18Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._make_workflow_opened_at_nullable(cursor)
|
||||
|
||||
def _make_workflow_opened_at_nullable(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
Make the `opened_at` column nullable in the `workflow_library` table. This is accomplished by:
|
||||
- Dropping the existing `idx_workflow_library_opened_at` index (must be done before dropping the column)
|
||||
- Dropping the existing `opened_at` column
|
||||
- Adding a new nullable column `opened_at` (no data migration needed, all values will be NULL)
|
||||
- Adding a new `idx_workflow_library_opened_at` index on the `opened_at` column
|
||||
"""
|
||||
# For index renaming in SQLite, we need to drop and recreate
|
||||
cursor.execute("DROP INDEX IF EXISTS idx_workflow_library_opened_at;")
|
||||
# Rename existing column to deprecated
|
||||
cursor.execute("ALTER TABLE workflow_library DROP COLUMN opened_at;")
|
||||
# Add new nullable column - all values will be NULL - no migration of data needed
|
||||
cursor.execute("ALTER TABLE workflow_library ADD COLUMN opened_at DATETIME;")
|
||||
# Create new index on the new column
|
||||
cursor.execute(
|
||||
"CREATE INDEX idx_workflow_library_opened_at ON workflow_library(opened_at);",
|
||||
)
|
||||
|
||||
|
||||
def build_migration_18() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 17 to 18.
|
||||
|
||||
This migration does the following:
|
||||
- Make the `opened_at` column nullable in the `workflow_library` table. This is accomplished by:
|
||||
- Dropping the existing `idx_workflow_library_opened_at` index (must be done before dropping the column)
|
||||
- Dropping the existing `opened_at` column
|
||||
- Adding a new nullable column `opened_at` (no data migration needed, all values will be NULL)
|
||||
- Adding a new `idx_workflow_library_opened_at` index on the `opened_at` column
|
||||
"""
|
||||
migration_18 = Migration(
|
||||
id="migration_18",
|
||||
depends_on="migration_17",
|
||||
from_version=17,
|
||||
to_version=18,
|
||||
callback=Migration18Callback(),
|
||||
)
|
||||
|
||||
return migration_18
|
||||
@@ -0,0 +1,39 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
|
||||
|
||||
|
||||
class Migration19Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig):
|
||||
self.models_path = app_config.models_path
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._populate_size(cursor)
|
||||
self._add_size_column(cursor)
|
||||
|
||||
def _add_size_column(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"ALTER TABLE models ADD COLUMN file_size INTEGER "
|
||||
"GENERATED ALWAYS as (json_extract(config, '$.file_size')) VIRTUAL NOT NULL"
|
||||
)
|
||||
|
||||
def _populate_size(self, cursor: sqlite3.Cursor) -> None:
|
||||
all_models = cursor.execute("SELECT id, path FROM models;").fetchall()
|
||||
|
||||
for model_id, model_path in all_models:
|
||||
mod = ModelOnDisk(self.models_path / model_path)
|
||||
cursor.execute(
|
||||
"UPDATE models SET config = json_set(config, '$.file_size', ?) WHERE id = ?", (mod.size(), model_id)
|
||||
)
|
||||
|
||||
|
||||
def build_migration_19(app_config: InvokeAIAppConfig) -> Migration:
|
||||
return Migration(
|
||||
id="migration_19",
|
||||
depends_on="migration_18",
|
||||
from_version=18,
|
||||
to_version=19,
|
||||
callback=Migration19Callback(app_config),
|
||||
)
|
||||
@@ -0,0 +1,168 @@
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
|
||||
from pydantic import ValidationError
|
||||
from tqdm import tqdm
|
||||
|
||||
from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase
|
||||
from invokeai.app.services.image_files.image_files_common import ImageFileNotFoundException
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
from invokeai.app.services.workflow_records.workflow_records_common import (
|
||||
UnsafeWorkflowWithVersionValidator,
|
||||
)
|
||||
|
||||
|
||||
class Migration2Callback:
|
||||
def __init__(self, image_files: ImageFileStorageBase, logger: Logger):
|
||||
self._image_files = image_files
|
||||
self._logger = logger
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor):
|
||||
self._add_images_has_workflow(cursor)
|
||||
self._add_session_queue_workflow(cursor)
|
||||
self._drop_old_workflow_tables(cursor)
|
||||
self._add_workflow_library(cursor)
|
||||
self._drop_model_manager_metadata(cursor)
|
||||
self._migrate_embedded_workflows(cursor)
|
||||
|
||||
def _add_images_has_workflow(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add the `has_workflow` column to `images` table."""
|
||||
cursor.execute("PRAGMA table_info(images)")
|
||||
columns = [column[1] for column in cursor.fetchall()]
|
||||
|
||||
if "has_workflow" not in columns:
|
||||
cursor.execute("ALTER TABLE images ADD COLUMN has_workflow BOOLEAN DEFAULT FALSE;")
|
||||
|
||||
def _add_session_queue_workflow(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add the `workflow` column to `session_queue` table."""
|
||||
|
||||
cursor.execute("PRAGMA table_info(session_queue)")
|
||||
columns = [column[1] for column in cursor.fetchall()]
|
||||
|
||||
if "workflow" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN workflow TEXT;")
|
||||
|
||||
def _drop_old_workflow_tables(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Drops the `workflows` and `workflow_images` tables."""
|
||||
cursor.execute("DROP TABLE IF EXISTS workflow_images;")
|
||||
cursor.execute("DROP TABLE IF EXISTS workflows;")
|
||||
|
||||
def _add_workflow_library(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Adds the `workflow_library` table and drops the `workflows` and `workflow_images` tables."""
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS workflow_library (
|
||||
workflow_id TEXT NOT NULL PRIMARY KEY,
|
||||
workflow TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- updated manually when retrieving workflow
|
||||
opened_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Generated columns, needed for indexing and searching
|
||||
category TEXT GENERATED ALWAYS as (json_extract(workflow, '$.meta.category')) VIRTUAL NOT NULL,
|
||||
name TEXT GENERATED ALWAYS as (json_extract(workflow, '$.name')) VIRTUAL NOT NULL,
|
||||
description TEXT GENERATED ALWAYS as (json_extract(workflow, '$.description')) VIRTUAL NOT NULL
|
||||
);
|
||||
""",
|
||||
]
|
||||
|
||||
indices = [
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_library_created_at ON workflow_library(created_at);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_library_updated_at ON workflow_library(updated_at);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_library_opened_at ON workflow_library(opened_at);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_library_category ON workflow_library(category);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_library_name ON workflow_library(name);",
|
||||
"CREATE INDEX IF NOT EXISTS idx_workflow_library_description ON workflow_library(description);",
|
||||
]
|
||||
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_workflow_library_updated_at
|
||||
AFTER UPDATE
|
||||
ON workflow_library FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE workflow_library
|
||||
SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE workflow_id = old.workflow_id;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
def _drop_model_manager_metadata(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Drops the `model_manager_metadata` table."""
|
||||
cursor.execute("DROP TABLE IF EXISTS model_manager_metadata;")
|
||||
|
||||
def _migrate_embedded_workflows(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
In the v3.5.0 release, InvokeAI changed how it handles embedded workflows. The `images` table in
|
||||
the database now has a `has_workflow` column, indicating if an image has a workflow embedded.
|
||||
|
||||
This migrate callback checks each image for the presence of an embedded workflow, then updates its entry
|
||||
in the database accordingly.
|
||||
"""
|
||||
# Get all image names
|
||||
cursor.execute("SELECT image_name FROM images")
|
||||
image_names: list[str] = [image[0] for image in cursor.fetchall()]
|
||||
total_image_names = len(image_names)
|
||||
|
||||
if not total_image_names:
|
||||
return
|
||||
|
||||
self._logger.info(f"Migrating workflows for {total_image_names} images")
|
||||
|
||||
# Migrate the images
|
||||
to_migrate: list[tuple[bool, str]] = []
|
||||
pbar = tqdm(image_names)
|
||||
for idx, image_name in enumerate(pbar):
|
||||
pbar.set_description(f"Checking image {idx + 1}/{total_image_names} for workflow")
|
||||
try:
|
||||
pil_image = self._image_files.get(image_name)
|
||||
except ImageFileNotFoundException:
|
||||
self._logger.warning(f"Image {image_name} not found, skipping")
|
||||
continue
|
||||
except Exception as e:
|
||||
self._logger.warning(f"Error while checking image {image_name}, skipping: {e}")
|
||||
continue
|
||||
if "invokeai_workflow" in pil_image.info:
|
||||
try:
|
||||
UnsafeWorkflowWithVersionValidator.validate_json(pil_image.info.get("invokeai_workflow", ""))
|
||||
except ValidationError:
|
||||
self._logger.warning(f"Image {image_name} has invalid embedded workflow, skipping")
|
||||
continue
|
||||
to_migrate.append((True, image_name))
|
||||
|
||||
self._logger.info(f"Adding {len(to_migrate)} embedded workflows to database")
|
||||
cursor.executemany("UPDATE images SET has_workflow = ? WHERE image_name = ?", to_migrate)
|
||||
|
||||
|
||||
def build_migration_2(image_files: ImageFileStorageBase, logger: Logger) -> Migration:
|
||||
"""
|
||||
Builds the migration from database version 1 to 2.
|
||||
|
||||
Introduced in v3.5.0 for the new workflow library.
|
||||
|
||||
:param image_files: The image files service, used to check for embedded workflows
|
||||
:param logger: The logger, used to log progress during embedded workflows handling
|
||||
|
||||
This migration does the following:
|
||||
- Add `has_workflow` column to `images` table
|
||||
- Add `workflow` column to `session_queue` table
|
||||
- Drop `workflows` and `workflow_images` tables
|
||||
- Add `workflow_library` table
|
||||
- Drops the `model_manager_metadata` table
|
||||
- Drops the `model_config` table, recreating it (at this point, there is no user data in this table)
|
||||
- Populates the `has_workflow` column in the `images` table (requires `image_files` & `logger` dependencies)
|
||||
"""
|
||||
migration_2 = Migration(
|
||||
id="migration_2",
|
||||
depends_on="migration_1",
|
||||
from_version=1,
|
||||
to_version=2,
|
||||
callback=Migration2Callback(image_files=image_files, logger=logger),
|
||||
)
|
||||
|
||||
return migration_2
|
||||
@@ -0,0 +1,39 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration20Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"""
|
||||
-- many-to-many relationship table for models
|
||||
CREATE TABLE IF NOT EXISTS model_relationships (
|
||||
-- model_key_1 and model_key_2 are the same as the key(primary key) in the models table
|
||||
model_key_1 TEXT NOT NULL,
|
||||
model_key_2 TEXT NOT NULL,
|
||||
created_at TEXT DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
PRIMARY KEY (model_key_1, model_key_2),
|
||||
-- model_key_1 < model_key_2, to ensure uniqueness and prevent duplicates
|
||||
FOREIGN KEY (model_key_1) REFERENCES models(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (model_key_2) REFERENCES models(id) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
-- Creates an index to keep performance equal when searching for model_key_1 or model_key_2
|
||||
CREATE INDEX IF NOT EXISTS keyx_model_relationships_model_key_2
|
||||
ON model_relationships(model_key_2)
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_20() -> Migration:
|
||||
return Migration(
|
||||
id="migration_20",
|
||||
depends_on="migration_19",
|
||||
from_version=19,
|
||||
to_version=20,
|
||||
callback=Migration20Callback(),
|
||||
)
|
||||
+53
@@ -0,0 +1,53 @@
|
||||
"""Add workflow-call relationship columns to session_queue."""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class AddWorkflowCallQueueMetadataCallback:
|
||||
"""Add durable parent/child workflow-call relationship columns to session_queue."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='session_queue';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(session_queue);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "workflow_call_id" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN workflow_call_id TEXT;")
|
||||
cursor.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_workflow_call_id ON session_queue(workflow_call_id);"
|
||||
)
|
||||
|
||||
if "parent_item_id" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN parent_item_id INTEGER;")
|
||||
cursor.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_parent_item_id ON session_queue(parent_item_id);"
|
||||
)
|
||||
|
||||
if "parent_session_id" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN parent_session_id TEXT;")
|
||||
cursor.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_parent_session_id ON session_queue(parent_session_id);"
|
||||
)
|
||||
|
||||
if "root_item_id" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN root_item_id INTEGER;")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_queue_root_item_id ON session_queue(root_item_id);")
|
||||
|
||||
if "workflow_call_depth" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN workflow_call_depth INTEGER;")
|
||||
cursor.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_session_queue_workflow_call_depth ON session_queue(workflow_call_depth);"
|
||||
)
|
||||
|
||||
|
||||
def build_migration() -> Migration:
|
||||
return Migration(
|
||||
id="2026_07_01_add_workflow_call_queue_metadata",
|
||||
depends_on="migration_33",
|
||||
callback=AddWorkflowCallQueueMetadataCallback(),
|
||||
)
|
||||
+60
@@ -0,0 +1,60 @@
|
||||
"""Add indexes supporting round-robin dequeue.
|
||||
|
||||
The round-robin dequeue in multiuser mode relies on two access shapes:
|
||||
|
||||
1. A per-user "best pending item" selection (``status = 'pending'`` partitioned by
|
||||
``user_id`` and ordered by ``priority DESC, item_id ASC``).
|
||||
2. A per-candidate "last served" lookup (``MAX(started_at)`` for a single ``user_id``),
|
||||
evaluated as a correlated subquery once per user that has pending work.
|
||||
|
||||
With only the pre-existing single-column indexes on ``status``, ``priority``, and
|
||||
``user_id``, the pending selection falls back to scanning the table and the last-served
|
||||
lookup scans all retained history. Because completed/failed/canceled history is retained
|
||||
(and ``max_queue_history`` defaults to unbounded), that cost would grow with total queue
|
||||
history rather than with the number of pending items / active users.
|
||||
|
||||
``idx_session_queue_round_robin_pending`` lets the planner satisfy the pending selection
|
||||
without touching historical rows. ``idx_session_queue_user_started_at`` turns each
|
||||
``MAX(started_at) WHERE user_id = ?`` into an indexed seek (the planner's min/max
|
||||
optimization reads the tail of the user's index range) rather than a scan, so the
|
||||
last-served lookup costs ``O(log n)`` per active user instead of scanning history.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class RoundRobinIndexesCallback:
|
||||
"""Add composite indexes matching the round-robin dequeue query shapes."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='session_queue';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
# Pending-item selection: WHERE status = 'pending', PARTITION BY user_id,
|
||||
# ORDER BY priority DESC, item_id ASC.
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE INDEX IF NOT EXISTS idx_session_queue_round_robin_pending
|
||||
ON session_queue (status, user_id, priority DESC, item_id ASC);
|
||||
"""
|
||||
)
|
||||
|
||||
# Last-served lookup: MAX(started_at) WHERE user_id = ?)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE INDEX IF NOT EXISTS idx_session_queue_user_started_at
|
||||
ON session_queue (user_id, started_at);
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration() -> Migration:
|
||||
return Migration(
|
||||
id="2026_07_03_round_robin_indexes",
|
||||
# migration_27 added the user_id column to session_queue, which both indexes cover.
|
||||
depends_on="migration_27",
|
||||
callback=RoundRobinIndexesCallback(),
|
||||
)
|
||||
@@ -0,0 +1,42 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration21Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE client_state (
|
||||
id INTEGER PRIMARY KEY CHECK(id = 1),
|
||||
data TEXT NOT NULL, -- Frontend will handle the shape of this data
|
||||
updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP)
|
||||
);
|
||||
"""
|
||||
)
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TRIGGER tg_client_state_updated_at
|
||||
AFTER UPDATE ON client_state
|
||||
FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE client_state
|
||||
SET updated_at = CURRENT_TIMESTAMP
|
||||
WHERE id = OLD.id;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_21() -> Migration:
|
||||
"""Builds the migration object for migrating from version 20 to version 21. This includes:
|
||||
- Creating the `client_state` table.
|
||||
- Adding a trigger to update the `updated_at` field on updates.
|
||||
"""
|
||||
return Migration(
|
||||
id="migration_21",
|
||||
depends_on="migration_20",
|
||||
from_version=20,
|
||||
to_version=21,
|
||||
callback=Migration21Callback(),
|
||||
)
|
||||
@@ -0,0 +1,91 @@
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration22Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
self._models_dir = app_config.models_path.resolve()
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._logger.info("Removing UNIQUE(name, base, type) constraint from models table")
|
||||
|
||||
# Step 1: Rename the existing models table
|
||||
cursor.execute("ALTER TABLE models RENAME TO models_old;")
|
||||
|
||||
# Step 2: Create the new models table without the UNIQUE(name, base, type) constraint
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE models (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
hash TEXT GENERATED ALWAYS as (json_extract(config, '$.hash')) VIRTUAL NOT NULL,
|
||||
base TEXT GENERATED ALWAYS as (json_extract(config, '$.base')) VIRTUAL NOT NULL,
|
||||
type TEXT GENERATED ALWAYS as (json_extract(config, '$.type')) VIRTUAL NOT NULL,
|
||||
path TEXT GENERATED ALWAYS as (json_extract(config, '$.path')) VIRTUAL NOT NULL,
|
||||
format TEXT GENERATED ALWAYS as (json_extract(config, '$.format')) VIRTUAL NOT NULL,
|
||||
name TEXT GENERATED ALWAYS as (json_extract(config, '$.name')) VIRTUAL NOT NULL,
|
||||
description TEXT GENERATED ALWAYS as (json_extract(config, '$.description')) VIRTUAL,
|
||||
source TEXT GENERATED ALWAYS as (json_extract(config, '$.source')) VIRTUAL NOT NULL,
|
||||
source_type TEXT GENERATED ALWAYS as (json_extract(config, '$.source_type')) VIRTUAL NOT NULL,
|
||||
source_api_response TEXT GENERATED ALWAYS as (json_extract(config, '$.source_api_response')) VIRTUAL,
|
||||
trigger_phrases TEXT GENERATED ALWAYS as (json_extract(config, '$.trigger_phrases')) VIRTUAL,
|
||||
file_size INTEGER GENERATED ALWAYS as (json_extract(config, '$.file_size')) VIRTUAL NOT NULL,
|
||||
-- Serialized JSON representation of the whole config object, which will contain additional fields from subclasses
|
||||
config TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Explicit unique constraint on path
|
||||
UNIQUE(path)
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
# Step 3: Copy all data from the old table to the new table
|
||||
# Only copy the stored columns (id, config, created_at, updated_at), not the virtual columns
|
||||
cursor.execute(
|
||||
"INSERT INTO models (id, config, created_at, updated_at) "
|
||||
"SELECT id, config, created_at, updated_at FROM models_old;"
|
||||
)
|
||||
|
||||
# Step 4: Drop the old table
|
||||
cursor.execute("DROP TABLE models_old;")
|
||||
|
||||
# Step 5: Recreate indexes
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS base_index ON models(base);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS type_index ON models(type);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS name_index ON models(name);")
|
||||
|
||||
# Step 6: Recreate the updated_at trigger
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TRIGGER models_updated_at
|
||||
AFTER UPDATE
|
||||
ON models FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE models SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_22(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""Builds the migration object for migrating from version 21 to version 22.
|
||||
|
||||
This migration:
|
||||
- Removes the UNIQUE constraint on the combination of (base, name, type) columns in the models table
|
||||
- Adds an explicit UNIQUE contraint on the path column
|
||||
"""
|
||||
|
||||
return Migration(
|
||||
id="migration_22",
|
||||
depends_on="migration_21",
|
||||
from_version=21,
|
||||
to_version=22,
|
||||
callback=Migration22Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
@@ -0,0 +1,195 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from copy import deepcopy
|
||||
from logging import Logger
|
||||
from typing import Any
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig, AnyModelConfigValidator
|
||||
from invokeai.backend.model_manager.configs.unknown import Unknown_Config
|
||||
from invokeai.backend.model_manager.taxonomy import (
|
||||
BaseModelType,
|
||||
ClipVariantType,
|
||||
FluxVariantType,
|
||||
ModelFormat,
|
||||
ModelType,
|
||||
ModelVariantType,
|
||||
SchedulerPredictionType,
|
||||
)
|
||||
|
||||
|
||||
class Migration23Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
self._models_dir = app_config.models_path.resolve()
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
# Grab all model records
|
||||
cursor.execute("SELECT id, config FROM models;")
|
||||
rows = cursor.fetchall()
|
||||
|
||||
migrated_count = 0
|
||||
fallback_count = 0
|
||||
|
||||
for model_id, config_json in rows:
|
||||
try:
|
||||
# Migrate the config JSON to the latest schema
|
||||
config_dict: dict[str, Any] = json.loads(config_json)
|
||||
migrated_config = self._parse_and_migrate_config(config_dict)
|
||||
|
||||
if isinstance(migrated_config, Unknown_Config):
|
||||
fallback_count += 1
|
||||
else:
|
||||
migrated_count += 1
|
||||
|
||||
# Write the migrated config back to the database
|
||||
cursor.execute(
|
||||
"UPDATE models SET config = ? WHERE id = ?;",
|
||||
(migrated_config.model_dump_json(), model_id),
|
||||
)
|
||||
except ValidationError as e:
|
||||
self._logger.error("Invalid config schema for model %s: %s", model_id, e)
|
||||
raise
|
||||
except json.JSONDecodeError as e:
|
||||
self._logger.error("Invalid config JSON for model %s: %s", model_id, e)
|
||||
raise
|
||||
|
||||
if migrated_count > 0 and fallback_count == 0:
|
||||
self._logger.info(f"Migration complete: {migrated_count} model configs migrated")
|
||||
elif migrated_count > 0 and fallback_count > 0:
|
||||
self._logger.warning(
|
||||
f"Migration complete: {migrated_count} model configs migrated, "
|
||||
f"{fallback_count} model configs could not be migrated and were saved as unknown models",
|
||||
)
|
||||
elif migrated_count == 0 and fallback_count > 0:
|
||||
self._logger.warning(
|
||||
f"Migration complete: all {fallback_count} model configs could not be migrated and were saved as unknown models",
|
||||
)
|
||||
else:
|
||||
self._logger.info("Migration complete: no model configs needed migration")
|
||||
|
||||
def _parse_and_migrate_config(self, config_dict: dict[str, Any]) -> AnyModelConfig:
|
||||
# In v6.9.0 we made some improvements to the model taxonomy and the model config schemas. There are a changes
|
||||
# we need to make to old configs to bring them up to date.
|
||||
|
||||
type = config_dict.get("type")
|
||||
format = config_dict.get("format")
|
||||
base = config_dict.get("base")
|
||||
|
||||
if base == BaseModelType.Flux.value and type == ModelType.Main.value:
|
||||
# Prior to v6.9.0, we used an awkward combination of `config_path` and `variant` to distinguish between FLUX
|
||||
# variants.
|
||||
#
|
||||
# `config_path` was set to one of:
|
||||
# - flux-dev
|
||||
# - flux-dev-fill
|
||||
# - flux-schnell
|
||||
#
|
||||
# `variant` was set to ModelVariantType.Inpaint for FLUX Fill models and ModelVariantType.Normal for all other FLUX
|
||||
# models.
|
||||
#
|
||||
# We now use the `variant` field to directly represent the FLUX variant type, and `config_path` is no longer used.
|
||||
|
||||
# Extract and remove `config_path` if present.
|
||||
config_path = config_dict.pop("config_path", None)
|
||||
|
||||
match config_path:
|
||||
case "flux-dev":
|
||||
config_dict["variant"] = FluxVariantType.Dev.value
|
||||
case "flux-dev-fill":
|
||||
config_dict["variant"] = FluxVariantType.DevFill.value
|
||||
case "flux-schnell":
|
||||
config_dict["variant"] = FluxVariantType.Schnell.value
|
||||
case _:
|
||||
# Unknown config_path - default to Dev variant
|
||||
config_dict["variant"] = FluxVariantType.Dev.value
|
||||
|
||||
if (
|
||||
base
|
||||
in {
|
||||
BaseModelType.StableDiffusion1.value,
|
||||
BaseModelType.StableDiffusion2.value,
|
||||
BaseModelType.StableDiffusionXL.value,
|
||||
BaseModelType.StableDiffusionXLRefiner.value,
|
||||
}
|
||||
and type == ModelType.Main.value
|
||||
):
|
||||
# Prior to v6.9.0, the prediction_type field was optional and would default to Epsilon if not present.
|
||||
# We now make it explicit and always present. Use the existing value if present, otherwise default to
|
||||
# Epsilon, matching the probe logic.
|
||||
#
|
||||
# It's only on SD1.x, SD2.x, and SDXL main models.
|
||||
config_dict["prediction_type"] = config_dict.get("prediction_type", SchedulerPredictionType.Epsilon.value)
|
||||
|
||||
# Prior to v6.9.0, the variant field was optional and would default to Normal if not present.
|
||||
# We now make it explicit and always present. Use the existing value if present, otherwise default to
|
||||
# Normal. It's only on SD main models.
|
||||
config_dict["variant"] = config_dict.get("variant", ModelVariantType.Normal.value)
|
||||
|
||||
if base == BaseModelType.Flux.value and type == ModelType.LoRA.value and format == ModelFormat.Diffusers.value:
|
||||
# Prior to v6.9.0, we used the Diffusers format for FLUX LoRA models that used the diffusers _key_
|
||||
# structure. This was misleading, as everywhere else in the application, we used the Diffusers format
|
||||
# to indicate that the model files were in the Diffusers _file_ format (i.e. a directory containing
|
||||
# the weights and config files).
|
||||
#
|
||||
# At runtime, we check the LoRA's state dict directly to determine the key structure, so we do not need
|
||||
# to rely on the format field for this purpose. As of v6.9.0, we always use the LyCORIS format for single-
|
||||
# file LoRAs, regardless of the key structure.
|
||||
#
|
||||
# This change allows LoRA model identification to not need a special case for FLUX LoRAs in the diffusers
|
||||
# key format.
|
||||
config_dict["format"] = ModelFormat.LyCORIS.value
|
||||
|
||||
if type == ModelType.CLIPVision.value:
|
||||
# Prior to v6.9.0, some CLIP Vision models were associated with a specific base model architecture:
|
||||
# - CLIP-ViT-bigG-14-laion2B-39B-b160k is the image encoder for SDXL IP Adapter and was associated with SDXL
|
||||
# - CLIP-ViT-H-14-laion2B-s32B-b79K is the image encoder for SD1.5 IP Adapter and was associated with SD1.5
|
||||
#
|
||||
# While this made some sense at the time, it is more correct and flexible to treat CLIP Vision models
|
||||
# as independent of any specific base model architecture.
|
||||
config_dict["base"] = BaseModelType.Any.value
|
||||
|
||||
if type == ModelType.CLIPEmbed.value:
|
||||
# Prior to v6.9.0, some CLIP Embed models did not have a variant set. The default was the L variant.
|
||||
# We now make it explicit and always present. Use the existing value if present, otherwise default to
|
||||
# L variant. Also, treat CLIP Embed models as independent of any specific base model architecture.
|
||||
config_dict["base"] = BaseModelType.Any.value
|
||||
config_dict["variant"] = config_dict.get("variant", ClipVariantType.L.value)
|
||||
|
||||
try:
|
||||
migrated_config = AnyModelConfigValidator.validate_python(config_dict)
|
||||
# This could be a ValidationError or any other error that occurs during validation. A failure to generate a
|
||||
# union discriminator could raise a ValueError, for example. Who knows what else could fail - catch all.
|
||||
except Exception as e:
|
||||
self._logger.error("Failed to validate migrated config, attempting to save as unknown model: %s", e)
|
||||
cloned_config_dict = deepcopy(config_dict)
|
||||
cloned_config_dict.pop("base", None)
|
||||
cloned_config_dict.pop("type", None)
|
||||
cloned_config_dict.pop("format", None)
|
||||
|
||||
migrated_config = Unknown_Config(
|
||||
**cloned_config_dict,
|
||||
base=BaseModelType.Unknown,
|
||||
type=ModelType.Unknown,
|
||||
format=ModelFormat.Unknown,
|
||||
)
|
||||
return migrated_config
|
||||
|
||||
|
||||
def build_migration_23(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""Builds the migration object for migrating from version 22 to version 23.
|
||||
|
||||
This migration updates model configurations to the latest config schemas for v6.9.0.
|
||||
"""
|
||||
|
||||
return Migration(
|
||||
id="migration_23",
|
||||
depends_on="migration_22",
|
||||
from_version=22,
|
||||
to_version=23,
|
||||
callback=Migration23Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
@@ -0,0 +1,242 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfigValidator
|
||||
|
||||
|
||||
class NormalizeResult(NamedTuple):
|
||||
new_relative_path: str | None
|
||||
rollback_ops: list[tuple[Path, Path]]
|
||||
|
||||
|
||||
class Migration24Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
self._models_dir = app_config.models_path.resolve()
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
# Grab all model records
|
||||
cursor.execute("SELECT id, config FROM models;")
|
||||
rows = cursor.fetchall()
|
||||
|
||||
for model_id, config_json in rows:
|
||||
try:
|
||||
config = AnyModelConfigValidator.validate_json(config_json)
|
||||
except ValidationError:
|
||||
# This could happen if the config schema changed in a way that makes old configs invalid. Unlikely
|
||||
# for users, more likely for devs testing out migration paths.
|
||||
self._logger.warning("Skipping model %s: invalid config schema", model_id)
|
||||
continue
|
||||
except json.JSONDecodeError:
|
||||
# This should never happen, as we use pydantic to serialize the config to JSON.
|
||||
self._logger.warning("Skipping model %s: invalid config JSON", model_id)
|
||||
continue
|
||||
|
||||
# We'll use a savepoint so we can roll back the database update if something goes wrong, and a simple
|
||||
# rollback of file operations if needed.
|
||||
cursor.execute("SAVEPOINT migrate_model")
|
||||
try:
|
||||
new_relative_path, rollback_ops = self._normalize_model_storage(
|
||||
key=config.key,
|
||||
path_value=config.path,
|
||||
)
|
||||
except Exception as err:
|
||||
self._logger.error("Error normalizing model %s: %s", config.key, err)
|
||||
cursor.execute("ROLLBACK TO SAVEPOINT migrate_model")
|
||||
cursor.execute("RELEASE SAVEPOINT migrate_model")
|
||||
continue
|
||||
|
||||
if new_relative_path is None:
|
||||
cursor.execute("RELEASE SAVEPOINT migrate_model")
|
||||
continue
|
||||
|
||||
config.path = new_relative_path
|
||||
try:
|
||||
cursor.execute(
|
||||
"UPDATE models SET config = ? WHERE id = ?;",
|
||||
(config.model_dump_json(), model_id),
|
||||
)
|
||||
except Exception as err:
|
||||
self._logger.error("Database update failed for model %s: %s", config.key, err)
|
||||
cursor.execute("ROLLBACK TO SAVEPOINT migrate_model")
|
||||
cursor.execute("RELEASE SAVEPOINT migrate_model")
|
||||
self._rollback_file_ops(rollback_ops)
|
||||
continue
|
||||
|
||||
cursor.execute("RELEASE SAVEPOINT migrate_model")
|
||||
|
||||
self._prune_empty_directories()
|
||||
|
||||
def _normalize_model_storage(self, key: str, path_value: str) -> NormalizeResult:
|
||||
models_dir = self._models_dir
|
||||
stored_path = Path(path_value)
|
||||
|
||||
relative_path: Path | None
|
||||
if stored_path.is_absolute():
|
||||
# If the stored path is absolute, we need to check if it's inside the models directory, which means it is
|
||||
# an Invoke-managed model. If it's outside, it is user-managed we leave it alone.
|
||||
try:
|
||||
relative_path = stored_path.resolve().relative_to(models_dir)
|
||||
except ValueError:
|
||||
self._logger.info("Leaving user-managed model %s at %s", key, stored_path)
|
||||
return NormalizeResult(new_relative_path=None, rollback_ops=[])
|
||||
else:
|
||||
# Relative paths are always relative to the models directory and thus Invoke-managed.
|
||||
relative_path = stored_path
|
||||
|
||||
# If the relative path is empty, assume something is wrong. Warn and skip.
|
||||
if not relative_path.parts:
|
||||
self._logger.warning("Skipping model %s: empty relative path", key)
|
||||
return NormalizeResult(new_relative_path=None, rollback_ops=[])
|
||||
|
||||
# Sanity check: the path is relative. It should be present in the models directory.
|
||||
absolute_path = (models_dir / relative_path).resolve()
|
||||
if not absolute_path.exists():
|
||||
self._logger.warning(
|
||||
"Skipping model %s: expected model files at %s but nothing was found",
|
||||
key,
|
||||
absolute_path,
|
||||
)
|
||||
return NormalizeResult(new_relative_path=None, rollback_ops=[])
|
||||
|
||||
if relative_path.parts[0] == key:
|
||||
# Already normalized. Still ensure the stored path is relative.
|
||||
normalized_path = relative_path.as_posix()
|
||||
# If the stored path is already the normalized path, no change is needed.
|
||||
new_relative_path = normalized_path if stored_path.as_posix() != normalized_path else None
|
||||
return NormalizeResult(new_relative_path=new_relative_path, rollback_ops=[])
|
||||
|
||||
# We'll store the file operations we perform so we can roll them back if needed.
|
||||
rollback_ops: list[tuple[Path, Path]] = []
|
||||
|
||||
# Destination directory is models_dir/<key> - a flat directory structure.
|
||||
destination_dir = models_dir / key
|
||||
|
||||
try:
|
||||
if absolute_path.is_file():
|
||||
destination_dir.mkdir(parents=True, exist_ok=True)
|
||||
dest_file = destination_dir / absolute_path.name
|
||||
# This really shouldn't happen.
|
||||
if dest_file.exists():
|
||||
self._logger.warning(
|
||||
"Destination for model %s already exists at %s; skipping move",
|
||||
key,
|
||||
dest_file,
|
||||
)
|
||||
return NormalizeResult(new_relative_path=None, rollback_ops=[])
|
||||
|
||||
self._logger.info("Moving model file %s -> %s", absolute_path, dest_file)
|
||||
|
||||
# `Path.rename()` effectively moves the file or directory.
|
||||
absolute_path.rename(dest_file)
|
||||
rollback_ops.append((dest_file, absolute_path))
|
||||
|
||||
return NormalizeResult(
|
||||
new_relative_path=(Path(key) / dest_file.name).as_posix(),
|
||||
rollback_ops=rollback_ops,
|
||||
)
|
||||
|
||||
if absolute_path.is_dir():
|
||||
dest_path = destination_dir
|
||||
# This really shouldn't happen.
|
||||
if dest_path.exists():
|
||||
self._logger.warning(
|
||||
"Destination directory %s already exists for model %s; skipping",
|
||||
dest_path,
|
||||
key,
|
||||
)
|
||||
return NormalizeResult(new_relative_path=None, rollback_ops=[])
|
||||
|
||||
self._logger.info("Moving model directory %s -> %s", absolute_path, dest_path)
|
||||
|
||||
# `Path.rename()` effectively moves the file or directory.
|
||||
absolute_path.rename(dest_path)
|
||||
rollback_ops.append((dest_path, absolute_path))
|
||||
|
||||
return NormalizeResult(
|
||||
new_relative_path=Path(key).as_posix(),
|
||||
rollback_ops=rollback_ops,
|
||||
)
|
||||
|
||||
# Maybe a broken symlink or something else weird?
|
||||
self._logger.warning("Skipping model %s: path %s is neither a file nor directory", key, absolute_path)
|
||||
return NormalizeResult(new_relative_path=None, rollback_ops=[])
|
||||
except Exception:
|
||||
self._rollback_file_ops(rollback_ops)
|
||||
raise
|
||||
|
||||
def _rollback_file_ops(self, rollback_ops: list[tuple[Path, Path]]) -> None:
|
||||
# This is a super-simple rollback that just reverses the move operations we performed.
|
||||
for source, destination in reversed(rollback_ops):
|
||||
try:
|
||||
if source.exists():
|
||||
source.rename(destination)
|
||||
except Exception as err:
|
||||
self._logger.error("Failed to rollback move %s -> %s: %s", source, destination, err)
|
||||
|
||||
def _prune_empty_directories(self) -> None:
|
||||
# These directories are system directories we want to keep even if empty. Technically, the app should not
|
||||
# have any problems if these are removed, creating them as needed, but it's cleaner to just leave them alone.
|
||||
keep_names = {"model_images", ".download_cache"}
|
||||
keep_dirs = {self._models_dir / name for name in keep_names}
|
||||
removed_dirs: set[Path] = set()
|
||||
|
||||
# Walk the models directory tree from the bottom up, removing empty directories. We sort by path length
|
||||
# descending to ensure we visit children before parents.
|
||||
for directory in sorted(self._models_dir.rglob("*"), key=lambda p: len(p.parts), reverse=True):
|
||||
if not directory.is_dir():
|
||||
continue
|
||||
if directory == self._models_dir:
|
||||
continue
|
||||
if any(directory == keep or keep in directory.parents for keep in keep_dirs):
|
||||
continue
|
||||
|
||||
try:
|
||||
next(directory.iterdir())
|
||||
except StopIteration:
|
||||
try:
|
||||
directory.rmdir()
|
||||
removed_dirs.add(directory)
|
||||
self._logger.debug("Removed empty directory %s", directory)
|
||||
except OSError:
|
||||
# Directory not empty (or some other error) - bail out.
|
||||
self._logger.warning("Failed to prune directory %s - not empty?", directory)
|
||||
continue
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
self._logger.info("Pruned %d empty directories under %s", len(removed_dirs), self._models_dir)
|
||||
|
||||
|
||||
def build_migration_24(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""Builds the migration object for migrating from version 23 to version 24.
|
||||
|
||||
This migration normalizes on-disk model storage so that each model lives within
|
||||
a directory named by its key inside the Invoke-managed models directory, and
|
||||
updates database records to reference the new relative paths.
|
||||
|
||||
This migration behaves a bit differently than others. Because it involves FS operations, if we rolled the
|
||||
DB back on any failure, we could leave the FS out of sync with the DB. Instead, we use savepoints
|
||||
to roll back individual model updates on failure, and we roll back any FS operations we performed
|
||||
for that model.
|
||||
|
||||
If a model cannot be migrated for any reason (invalid config, missing files, FS errors, DB errors), we log a
|
||||
warning and skip it, leaving it in its original state and location. The model will still work, but it will be in
|
||||
the "wrong" location on disk.
|
||||
"""
|
||||
|
||||
return Migration(
|
||||
id="migration_24",
|
||||
depends_on="migration_23",
|
||||
from_version=23,
|
||||
to_version=24,
|
||||
callback=Migration24Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
from typing import Any
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
from invokeai.backend.model_manager.taxonomy import ModelType, Qwen3VariantType
|
||||
|
||||
|
||||
class Migration25Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute("SELECT id, config FROM models;")
|
||||
rows = cursor.fetchall()
|
||||
|
||||
migrated_count = 0
|
||||
|
||||
for model_id, config_json in rows:
|
||||
try:
|
||||
config_dict: dict[str, Any] = json.loads(config_json)
|
||||
|
||||
if config_dict.get("type") != ModelType.Qwen3Encoder.value:
|
||||
continue
|
||||
|
||||
if "variant" in config_dict:
|
||||
continue
|
||||
|
||||
config_dict["variant"] = Qwen3VariantType.Qwen3_4B.value
|
||||
|
||||
cursor.execute(
|
||||
"UPDATE models SET config = ? WHERE id = ?;",
|
||||
(json.dumps(config_dict), model_id),
|
||||
)
|
||||
migrated_count += 1
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
self._logger.error("Invalid config JSON for model %s: %s", model_id, e)
|
||||
raise
|
||||
|
||||
if migrated_count > 0:
|
||||
self._logger.info(f"Migration complete: {migrated_count} Qwen3 encoder configs updated with variant field")
|
||||
else:
|
||||
self._logger.info("Migration complete: no Qwen3 encoder configs needed migration")
|
||||
|
||||
|
||||
def build_migration_25(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""Builds the migration object for migrating from version 24 to version 25.
|
||||
|
||||
This migration adds the variant field to existing Qwen3 encoder models.
|
||||
Models installed before the variant field was added will default to Qwen3_4B (for Z-Image compatibility).
|
||||
"""
|
||||
|
||||
return Migration(
|
||||
id="migration_25",
|
||||
depends_on="migration_24",
|
||||
from_version=24,
|
||||
to_version=25,
|
||||
callback=Migration25Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
@@ -0,0 +1,117 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, ZImageVariantType
|
||||
|
||||
|
||||
class Migration26Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
|
||||
def _detect_variant_from_scheduler(self, model_path: Path) -> ZImageVariantType:
|
||||
"""Detect Z-Image variant from scheduler config for Diffusers models.
|
||||
|
||||
Z-Image variants are distinguished by the scheduler shift value:
|
||||
- Turbo (distilled): shift = 3.0
|
||||
- Base (undistilled): shift = 6.0
|
||||
"""
|
||||
scheduler_config_path = model_path / "scheduler" / "scheduler_config.json"
|
||||
|
||||
if not scheduler_config_path.exists():
|
||||
return ZImageVariantType.Turbo
|
||||
|
||||
try:
|
||||
with open(scheduler_config_path, "r", encoding="utf-8") as f:
|
||||
scheduler_config = json.load(f)
|
||||
|
||||
shift = scheduler_config.get("shift", 3.0)
|
||||
|
||||
# ZBase (undistilled) uses shift = 6.0, Turbo uses shift = 3.0
|
||||
if shift >= 5.0:
|
||||
return ZImageVariantType.ZBase
|
||||
else:
|
||||
return ZImageVariantType.Turbo
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
self._logger.warning(f"Could not read scheduler config: {e}, defaulting to Turbo")
|
||||
return ZImageVariantType.Turbo
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute("SELECT id, config FROM models;")
|
||||
rows = cursor.fetchall()
|
||||
|
||||
migrated_turbo = 0
|
||||
migrated_base = 0
|
||||
|
||||
for model_id, config_json in rows:
|
||||
try:
|
||||
config_dict: dict[str, Any] = json.loads(config_json)
|
||||
|
||||
# Only migrate Z-Image main models
|
||||
if config_dict.get("base") != BaseModelType.ZImage.value:
|
||||
continue
|
||||
|
||||
if config_dict.get("type") != ModelType.Main.value:
|
||||
continue
|
||||
|
||||
# Skip if variant already set
|
||||
if "variant" in config_dict:
|
||||
continue
|
||||
|
||||
# Determine variant based on format
|
||||
model_format = config_dict.get("format")
|
||||
model_path = config_dict.get("path")
|
||||
|
||||
if model_format == ModelFormat.Diffusers.value and model_path:
|
||||
# For Diffusers models, detect from scheduler config
|
||||
variant = self._detect_variant_from_scheduler(Path(model_path))
|
||||
else:
|
||||
# For Checkpoint/GGUF, default to Turbo (Base only available as Diffusers)
|
||||
variant = ZImageVariantType.Turbo
|
||||
|
||||
config_dict["variant"] = variant.value
|
||||
|
||||
cursor.execute(
|
||||
"UPDATE models SET config = ? WHERE id = ?;",
|
||||
(json.dumps(config_dict), model_id),
|
||||
)
|
||||
|
||||
if variant == ZImageVariantType.ZBase:
|
||||
migrated_base += 1
|
||||
else:
|
||||
migrated_turbo += 1
|
||||
|
||||
except json.JSONDecodeError as e:
|
||||
self._logger.error("Invalid config JSON for model %s: %s", model_id, e)
|
||||
raise
|
||||
|
||||
total = migrated_turbo + migrated_base
|
||||
if total > 0:
|
||||
self._logger.info(
|
||||
f"Migration complete: {total} Z-Image model configs updated "
|
||||
f"({migrated_turbo} Turbo, {migrated_base} Base)"
|
||||
)
|
||||
else:
|
||||
self._logger.info("Migration complete: no Z-Image model configs needed migration")
|
||||
|
||||
|
||||
def build_migration_26(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""Builds the migration object for migrating from version 25 to version 26.
|
||||
|
||||
This migration adds the variant field to existing Z-Image main models.
|
||||
Models installed before the variant field was added will default to Turbo
|
||||
(the only variant available before Z-Image Base support was added).
|
||||
"""
|
||||
|
||||
return Migration(
|
||||
id="migration_26",
|
||||
depends_on="migration_25",
|
||||
from_version=25,
|
||||
to_version=26,
|
||||
callback=Migration26Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
@@ -0,0 +1,368 @@
|
||||
"""Migration 27: Add multi-user support, per-user client state, and app settings.
|
||||
|
||||
This migration adds the database schema for multi-user support, including:
|
||||
- users table for user accounts
|
||||
- user_sessions table for session management
|
||||
- user_invitations table for invitation system
|
||||
- shared_boards table for board sharing
|
||||
- Adding user_id columns to existing tables for data ownership
|
||||
- Restructuring client_state table to support per-user storage
|
||||
- app_settings table for storing JWT secret and other app-level settings
|
||||
"""
|
||||
|
||||
import json
|
||||
import secrets
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration27Callback:
|
||||
"""Migration to add multi-user support, per-user client state, and app settings."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._create_users_table(cursor)
|
||||
self._create_user_sessions_table(cursor)
|
||||
self._create_user_invitations_table(cursor)
|
||||
self._create_shared_boards_table(cursor)
|
||||
self._update_boards_table(cursor)
|
||||
self._update_images_table(cursor)
|
||||
self._update_workflows_table(cursor)
|
||||
self._update_session_queue_table(cursor)
|
||||
self._update_style_presets_table(cursor)
|
||||
self._create_system_user(cursor)
|
||||
self._update_client_state_table(cursor)
|
||||
self._create_app_settings_table(cursor)
|
||||
self._generate_jwt_secret(cursor)
|
||||
|
||||
def _create_users_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create users table."""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
user_id TEXT NOT NULL PRIMARY KEY,
|
||||
email TEXT NOT NULL UNIQUE,
|
||||
display_name TEXT,
|
||||
password_hash TEXT NOT NULL,
|
||||
is_admin BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
is_active BOOLEAN NOT NULL DEFAULT TRUE,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
last_login_at DATETIME
|
||||
);
|
||||
""")
|
||||
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_email ON users(email);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_is_admin ON users(is_admin);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_users_is_active ON users(is_active);")
|
||||
|
||||
cursor.execute("""
|
||||
CREATE TRIGGER IF NOT EXISTS tg_users_updated_at
|
||||
AFTER UPDATE ON users FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE users SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE user_id = old.user_id;
|
||||
END;
|
||||
""")
|
||||
|
||||
def _create_user_sessions_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create user_sessions table for session management."""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS user_sessions (
|
||||
session_id TEXT NOT NULL PRIMARY KEY,
|
||||
user_id TEXT NOT NULL,
|
||||
token_hash TEXT NOT NULL,
|
||||
expires_at DATETIME NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
last_activity_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
||||
);
|
||||
""")
|
||||
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_sessions_user_id ON user_sessions(user_id);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_sessions_token_hash ON user_sessions(token_hash);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_sessions_expires_at ON user_sessions(expires_at);")
|
||||
|
||||
def _create_user_invitations_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create user_invitations table for invitation system."""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS user_invitations (
|
||||
invitation_id TEXT NOT NULL PRIMARY KEY,
|
||||
email TEXT NOT NULL,
|
||||
invited_by TEXT NOT NULL,
|
||||
invitation_code TEXT NOT NULL UNIQUE,
|
||||
is_admin BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
expires_at DATETIME NOT NULL,
|
||||
used_at DATETIME,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
FOREIGN KEY (invited_by) REFERENCES users(user_id) ON DELETE CASCADE
|
||||
);
|
||||
""")
|
||||
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_invitations_email ON user_invitations(email);")
|
||||
cursor.execute(
|
||||
"CREATE INDEX IF NOT EXISTS idx_user_invitations_invitation_code ON user_invitations(invitation_code);"
|
||||
)
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_invitations_expires_at ON user_invitations(expires_at);")
|
||||
|
||||
def _create_shared_boards_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create shared_boards table for board sharing."""
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS shared_boards (
|
||||
board_id TEXT NOT NULL,
|
||||
user_id TEXT NOT NULL,
|
||||
can_edit BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
shared_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
PRIMARY KEY (board_id, user_id),
|
||||
FOREIGN KEY (board_id) REFERENCES boards(board_id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
||||
);
|
||||
""")
|
||||
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_shared_boards_user_id ON shared_boards(user_id);")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_shared_boards_board_id ON shared_boards(board_id);")
|
||||
|
||||
def _update_boards_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add user_id and is_public columns to boards table."""
|
||||
# Check if boards table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='boards';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
# Check if user_id column exists
|
||||
cursor.execute("PRAGMA table_info(boards);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "user_id" not in columns:
|
||||
cursor.execute("ALTER TABLE boards ADD COLUMN user_id TEXT DEFAULT 'system';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_boards_user_id ON boards(user_id);")
|
||||
|
||||
if "is_public" not in columns:
|
||||
cursor.execute("ALTER TABLE boards ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT FALSE;")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_boards_is_public ON boards(is_public);")
|
||||
|
||||
def _update_images_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add user_id column to images table."""
|
||||
# Check if images table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='images';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(images);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "user_id" not in columns:
|
||||
cursor.execute("ALTER TABLE images ADD COLUMN user_id TEXT DEFAULT 'system';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_images_user_id ON images(user_id);")
|
||||
|
||||
def _update_workflows_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add user_id and is_public columns to workflows table."""
|
||||
# Check if workflows table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='workflows';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(workflows);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "user_id" not in columns:
|
||||
cursor.execute("ALTER TABLE workflows ADD COLUMN user_id TEXT DEFAULT 'system';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_workflows_user_id ON workflows(user_id);")
|
||||
|
||||
if "is_public" not in columns:
|
||||
cursor.execute("ALTER TABLE workflows ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT FALSE;")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_workflows_is_public ON workflows(is_public);")
|
||||
|
||||
def _update_session_queue_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add user_id column to session_queue table."""
|
||||
# Check if session_queue table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='session_queue';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(session_queue);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "user_id" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN user_id TEXT DEFAULT 'system';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_queue_user_id ON session_queue(user_id);")
|
||||
|
||||
def _update_style_presets_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add user_id and is_public columns to style_presets table."""
|
||||
# Check if style_presets table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='style_presets';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(style_presets);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "user_id" not in columns:
|
||||
cursor.execute("ALTER TABLE style_presets ADD COLUMN user_id TEXT DEFAULT 'system';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_style_presets_user_id ON style_presets(user_id);")
|
||||
|
||||
if "is_public" not in columns:
|
||||
cursor.execute("ALTER TABLE style_presets ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT FALSE;")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_style_presets_is_public ON style_presets(is_public);")
|
||||
|
||||
def _create_system_user(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create system user for backward compatibility.
|
||||
|
||||
The system user is NOT an admin - it's just used to own existing data
|
||||
from before multi-user support was added. Real admin users should be
|
||||
created through the /auth/setup endpoint.
|
||||
"""
|
||||
cursor.execute("""
|
||||
INSERT OR IGNORE INTO users (user_id, email, display_name, password_hash, is_admin, is_active)
|
||||
VALUES ('system', 'system@system.invokeai', 'System', '', FALSE, TRUE);
|
||||
""")
|
||||
|
||||
def _update_client_state_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Restructure client_state table to support per-user storage."""
|
||||
# Check if client_state table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='client_state';")
|
||||
if cursor.fetchone() is None:
|
||||
# Table doesn't exist, create it with the new schema
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE client_state (
|
||||
user_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP),
|
||||
PRIMARY KEY (user_id, key),
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
)
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_client_state_user_id ON client_state(user_id);")
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TRIGGER tg_client_state_updated_at
|
||||
AFTER UPDATE ON client_state
|
||||
FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE client_state
|
||||
SET updated_at = CURRENT_TIMESTAMP
|
||||
WHERE user_id = OLD.user_id AND key = OLD.key;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
return
|
||||
|
||||
# Table exists with old schema - migrate it
|
||||
# Get existing data if the data column is present (it may be absent if an older
|
||||
# version of migration 21 was deployed without the column)
|
||||
cursor.execute("PRAGMA table_info(client_state);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
existing_data = {}
|
||||
if "data" in columns:
|
||||
cursor.execute("SELECT data FROM client_state WHERE id = 1;")
|
||||
row = cursor.fetchone()
|
||||
if row is not None:
|
||||
try:
|
||||
existing_data = json.loads(row[0])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# If data is corrupt, just start fresh
|
||||
pass
|
||||
|
||||
# Drop the old table
|
||||
cursor.execute("DROP TABLE IF EXISTS client_state;")
|
||||
|
||||
# Create new table with per-user schema
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE client_state (
|
||||
user_id TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
updated_at DATETIME NOT NULL DEFAULT (CURRENT_TIMESTAMP),
|
||||
PRIMARY KEY (user_id, key),
|
||||
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_client_state_user_id ON client_state(user_id);")
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TRIGGER tg_client_state_updated_at
|
||||
AFTER UPDATE ON client_state
|
||||
FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE client_state
|
||||
SET updated_at = CURRENT_TIMESTAMP
|
||||
WHERE user_id = OLD.user_id AND key = OLD.key;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
# Migrate existing data to 'system' user
|
||||
for key, value in existing_data.items():
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO client_state (user_id, key, value)
|
||||
VALUES ('system', ?, ?);
|
||||
""",
|
||||
(key, value),
|
||||
)
|
||||
|
||||
def _create_app_settings_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create app_settings table for storing application-level configuration."""
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS app_settings (
|
||||
key TEXT NOT NULL PRIMARY KEY,
|
||||
value TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW'))
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
CREATE TRIGGER IF NOT EXISTS tg_app_settings_updated_at
|
||||
AFTER UPDATE ON app_settings
|
||||
FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE app_settings SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE key = OLD.key;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
def _generate_jwt_secret(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Generate and store a cryptographically secure JWT secret key.
|
||||
|
||||
The secret is a 64-character hexadecimal string (256 bits of entropy),
|
||||
which is suitable for HS256 JWT signing.
|
||||
"""
|
||||
# Check if JWT secret already exists
|
||||
cursor.execute("SELECT value FROM app_settings WHERE key = 'jwt_secret';")
|
||||
existing_secret = cursor.fetchone()
|
||||
|
||||
if existing_secret is None:
|
||||
# Generate a new cryptographically secure secret (256 bits)
|
||||
jwt_secret = secrets.token_hex(32) # 32 bytes = 256 bits = 64 hex characters
|
||||
|
||||
# Store in database
|
||||
cursor.execute(
|
||||
"INSERT INTO app_settings (key, value) VALUES ('jwt_secret', ?);",
|
||||
(jwt_secret,),
|
||||
)
|
||||
|
||||
|
||||
def build_migration_27() -> Migration:
|
||||
"""Builds the migration object for migrating from version 26 to version 27.
|
||||
|
||||
This migration adds multi-user support, per-user client state, and app settings
|
||||
(including a JWT secret) to the database schema.
|
||||
"""
|
||||
return Migration(
|
||||
id="migration_27",
|
||||
depends_on="migration_26",
|
||||
from_version=26,
|
||||
to_version=27,
|
||||
callback=Migration27Callback(),
|
||||
)
|
||||
@@ -0,0 +1,50 @@
|
||||
"""Migration 28: Add per-user workflow isolation columns to workflow_library.
|
||||
|
||||
This migration adds the database columns required for multiuser workflow isolation
|
||||
to the workflow_library table:
|
||||
- user_id: the owner of the workflow (defaults to 'system' for existing workflows)
|
||||
- is_public: whether the workflow is shared with all users
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration28Callback:
|
||||
"""Migration to add user_id and is_public to the workflow_library table."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._update_workflow_library_table(cursor)
|
||||
|
||||
def _update_workflow_library_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add user_id and is_public columns to workflow_library table."""
|
||||
cursor.execute("PRAGMA table_info(workflow_library);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "user_id" not in columns:
|
||||
cursor.execute("ALTER TABLE workflow_library ADD COLUMN user_id TEXT DEFAULT 'system';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_workflow_library_user_id ON workflow_library(user_id);")
|
||||
|
||||
if "is_public" not in columns:
|
||||
cursor.execute("ALTER TABLE workflow_library ADD COLUMN is_public BOOLEAN NOT NULL DEFAULT FALSE;")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_workflow_library_is_public ON workflow_library(is_public);")
|
||||
cursor.execute(
|
||||
"UPDATE workflow_library SET is_public = TRUE WHERE user_id = 'system';"
|
||||
) # one-time fix for legacy workflows
|
||||
|
||||
|
||||
def build_migration_28() -> Migration:
|
||||
"""Builds the migration object for migrating from version 27 to version 28.
|
||||
|
||||
This migration adds per-user workflow isolation to the workflow_library table:
|
||||
- user_id column: identifies the owner of each workflow
|
||||
- is_public column: controls whether a workflow is shared with all users
|
||||
"""
|
||||
return Migration(
|
||||
id="migration_28",
|
||||
depends_on="migration_27",
|
||||
from_version=27,
|
||||
to_version=28,
|
||||
callback=Migration28Callback(),
|
||||
)
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Migration 29: Add board_visibility column to boards table.
|
||||
|
||||
This migration adds a board_visibility column to the boards table to support
|
||||
three visibility levels:
|
||||
- 'private': only the board owner (and admins) can view/modify
|
||||
- 'shared': all users can view, but only the owner (and admins) can modify
|
||||
- 'public': all users can view; only the owner (and admins) can modify the
|
||||
board structure (rename/archive/delete)
|
||||
|
||||
Existing boards with is_public = 1 are migrated to 'public'.
|
||||
All other existing boards default to 'private'.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration29Callback:
|
||||
"""Migration to add board_visibility column to the boards table."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._update_boards_table(cursor)
|
||||
|
||||
def _update_boards_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add board_visibility column to boards table."""
|
||||
# Check if boards table exists
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='boards';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(boards);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "board_visibility" not in columns:
|
||||
cursor.execute("ALTER TABLE boards ADD COLUMN board_visibility TEXT NOT NULL DEFAULT 'private';")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_boards_board_visibility ON boards(board_visibility);")
|
||||
# Migrate existing is_public = 1 boards to 'public'
|
||||
if "is_public" in columns:
|
||||
cursor.execute("UPDATE boards SET board_visibility = 'public' WHERE is_public = 1;")
|
||||
|
||||
|
||||
def build_migration_29() -> Migration:
|
||||
"""Builds the migration object for migrating from version 28 to version 29.
|
||||
|
||||
This migration adds the board_visibility column to the boards table,
|
||||
supporting 'private', 'shared', and 'public' visibility levels.
|
||||
"""
|
||||
return Migration(
|
||||
id="migration_29",
|
||||
depends_on="migration_28",
|
||||
from_version=28,
|
||||
to_version=29,
|
||||
callback=Migration29Callback(),
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
import sqlite3
|
||||
from logging import Logger
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration3Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._drop_model_manager_metadata(cursor)
|
||||
self._recreate_model_config(cursor)
|
||||
|
||||
def _drop_model_manager_metadata(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Drops the `model_manager_metadata` table."""
|
||||
cursor.execute("DROP TABLE IF EXISTS model_manager_metadata;")
|
||||
|
||||
def _recreate_model_config(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
Drops the `model_config` table, recreating it.
|
||||
|
||||
In 3.4.0, this table used explicit columns but was changed to use json_extract 3.5.0.
|
||||
|
||||
Because this table is not used in production, we are able to simply drop it and recreate it.
|
||||
"""
|
||||
|
||||
cursor.execute("DROP TABLE IF EXISTS model_config;")
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS model_config (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
-- The next 3 fields are enums in python, unrestricted string here
|
||||
base TEXT GENERATED ALWAYS as (json_extract(config, '$.base')) VIRTUAL NOT NULL,
|
||||
type TEXT GENERATED ALWAYS as (json_extract(config, '$.type')) VIRTUAL NOT NULL,
|
||||
name TEXT GENERATED ALWAYS as (json_extract(config, '$.name')) VIRTUAL NOT NULL,
|
||||
path TEXT GENERATED ALWAYS as (json_extract(config, '$.path')) VIRTUAL NOT NULL,
|
||||
format TEXT GENERATED ALWAYS as (json_extract(config, '$.format')) VIRTUAL NOT NULL,
|
||||
original_hash TEXT, -- could be null
|
||||
-- Serialized JSON representation of the whole config object,
|
||||
-- which will contain additional fields from subclasses
|
||||
config TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- unique constraint on combo of name, base and type
|
||||
UNIQUE(name, base, type)
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_3(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
|
||||
"""
|
||||
Build the migration from database version 2 to 3.
|
||||
|
||||
This migration does the following:
|
||||
- Drops the `model_config` table, recreating it
|
||||
- Migrates data from `models.yaml` into the `model_config` table
|
||||
"""
|
||||
migration_3 = Migration(
|
||||
id="migration_3",
|
||||
depends_on="migration_2",
|
||||
from_version=2,
|
||||
to_version=3,
|
||||
callback=Migration3Callback(app_config=app_config, logger=logger),
|
||||
)
|
||||
|
||||
return migration_3
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Migration 30: Add per-item queue status sequencing.
|
||||
|
||||
This migration adds a `status_sequence` column to `session_queue` so queue item
|
||||
status updates can be ordered across asynchronous event and snapshot channels.
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration30Callback:
|
||||
"""Add a per-queue-item status sequence for cross-channel ordering."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='session_queue';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(session_queue);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "status_sequence" not in columns:
|
||||
cursor.execute("ALTER TABLE session_queue ADD COLUMN status_sequence INTEGER DEFAULT 0;")
|
||||
cursor.execute("UPDATE session_queue SET status_sequence = 0 WHERE status_sequence IS NULL;")
|
||||
|
||||
|
||||
def build_migration_30() -> Migration:
|
||||
return Migration(
|
||||
id="migration_30",
|
||||
depends_on="migration_29",
|
||||
from_version=29,
|
||||
to_version=30,
|
||||
callback=Migration30Callback(),
|
||||
)
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Migration 31: Add image_subfolder column to images table.
|
||||
|
||||
This migration adds an image_subfolder column to the images table to support
|
||||
configurable image subfolder strategies (flat, date, type, hash).
|
||||
Existing images get an empty string (flat/root directory).
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration31Callback:
|
||||
"""Migration to add image_subfolder column to images table."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._add_image_subfolder_column(cursor)
|
||||
|
||||
def _add_image_subfolder_column(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Add image_subfolder column to images table."""
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='images';")
|
||||
if cursor.fetchone() is None:
|
||||
return
|
||||
|
||||
cursor.execute("PRAGMA table_info(images);")
|
||||
columns = [row[1] for row in cursor.fetchall()]
|
||||
|
||||
if "image_subfolder" not in columns:
|
||||
cursor.execute("ALTER TABLE images ADD COLUMN image_subfolder TEXT NOT NULL DEFAULT '';")
|
||||
|
||||
|
||||
def build_migration_31() -> Migration:
|
||||
"""Builds the migration object for migrating from version 30 to version 31.
|
||||
|
||||
This migration adds an image_subfolder column to the images table.
|
||||
"""
|
||||
return Migration(
|
||||
id="migration_31",
|
||||
depends_on="migration_30",
|
||||
from_version=30,
|
||||
to_version=31,
|
||||
callback=Migration31Callback(),
|
||||
)
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Migration 32: Repair model_relationships foreign keys.
|
||||
|
||||
Migration 22 rebuilt the `models` table by renaming it to `models_old`, creating a
|
||||
fresh `models` table, copying the data over, and dropping `models_old`. Because
|
||||
modern SQLite (with `legacy_alter_table` off) rewrites foreign-key references in
|
||||
*other* tables when a table is renamed, the foreign keys in `model_relationships`
|
||||
were silently repointed at `models_old` -- which was then dropped.
|
||||
|
||||
This left the related-models links referencing a table that no longer exists,
|
||||
breaking `ON DELETE CASCADE` and foreign-key integrity for related models.
|
||||
|
||||
This migration rebuilds `model_relationships` so its foreign keys reference
|
||||
`models(id)` again, preserving existing links and dropping any orphaned rows whose
|
||||
model keys no longer exist (those would violate the restored foreign keys).
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration32Callback:
|
||||
"""Migration to repair the broken foreign keys on the model_relationships table."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._repair_model_relationships_fks(cursor)
|
||||
|
||||
def _repair_model_relationships_fks(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='model_relationships';")
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
# Table does not exist (fresh db will create it correctly), nothing to repair.
|
||||
return
|
||||
|
||||
existing_sql: str = row[0]
|
||||
if "models_old" not in existing_sql:
|
||||
# Foreign keys already point at the correct table, nothing to repair.
|
||||
return
|
||||
|
||||
# Rebuild the table with the correct foreign keys referencing models(id).
|
||||
cursor.execute("ALTER TABLE model_relationships RENAME TO model_relationships_old;")
|
||||
cursor.execute(
|
||||
"""
|
||||
-- many-to-many relationship table for models
|
||||
CREATE TABLE model_relationships (
|
||||
-- model_key_1 and model_key_2 are the same as the key(primary key) in the models table
|
||||
model_key_1 TEXT NOT NULL,
|
||||
model_key_2 TEXT NOT NULL,
|
||||
created_at TEXT DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
PRIMARY KEY (model_key_1, model_key_2),
|
||||
-- model_key_1 < model_key_2, to ensure uniqueness and prevent duplicates
|
||||
FOREIGN KEY (model_key_1) REFERENCES models(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (model_key_2) REFERENCES models(id) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
# Copy over the existing links, dropping any orphaned rows whose model keys no
|
||||
# longer exist -- these would violate the restored foreign keys.
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO model_relationships (model_key_1, model_key_2, created_at)
|
||||
SELECT model_key_1, model_key_2, created_at
|
||||
FROM model_relationships_old
|
||||
WHERE model_key_1 IN (SELECT id FROM models)
|
||||
AND model_key_2 IN (SELECT id FROM models);
|
||||
"""
|
||||
)
|
||||
|
||||
# Drop the old table first so its index name is freed before we recreate it.
|
||||
cursor.execute("DROP TABLE model_relationships_old;")
|
||||
cursor.execute(
|
||||
"""
|
||||
-- Creates an index to keep performance equal when searching for model_key_1 or model_key_2
|
||||
CREATE INDEX IF NOT EXISTS keyx_model_relationships_model_key_2
|
||||
ON model_relationships(model_key_2);
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_32() -> Migration:
|
||||
"""Builds the migration object for migrating from version 31 to version 32.
|
||||
|
||||
This migration repairs the foreign keys on the model_relationships table, which were
|
||||
broken by migration 22 rebuilding the models table.
|
||||
"""
|
||||
return Migration(
|
||||
id="migration_32",
|
||||
depends_on="migration_31",
|
||||
from_version=31,
|
||||
to_version=32,
|
||||
callback=Migration32Callback(),
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration33Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS image_subfolder_move_jobs (
|
||||
id INTEGER PRIMARY KEY,
|
||||
state TEXT NOT NULL CHECK (
|
||||
state IN ('planned', 'moving', 'moved', 'committed', 'error')
|
||||
),
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
error_message TEXT
|
||||
);
|
||||
"""
|
||||
)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS image_subfolder_move_items (
|
||||
job_id INTEGER NOT NULL REFERENCES image_subfolder_move_jobs(id),
|
||||
image_name TEXT NOT NULL REFERENCES images(image_name),
|
||||
old_subfolder TEXT NOT NULL,
|
||||
new_subfolder TEXT NOT NULL,
|
||||
is_intermediate BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
old_path TEXT,
|
||||
new_path TEXT,
|
||||
old_thumbnail_path TEXT,
|
||||
new_thumbnail_path TEXT,
|
||||
state TEXT NOT NULL CHECK (
|
||||
state IN ('planned', 'moved', 'committed', 'error')
|
||||
),
|
||||
error_message TEXT,
|
||||
PRIMARY KEY (job_id, image_name)
|
||||
);
|
||||
"""
|
||||
)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE INDEX IF NOT EXISTS idx_image_subfolder_move_items_job_state
|
||||
ON image_subfolder_move_items(job_id, state);
|
||||
"""
|
||||
)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE INDEX IF NOT EXISTS idx_image_subfolder_move_items_image_name
|
||||
ON image_subfolder_move_items(image_name);
|
||||
"""
|
||||
)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS tg_image_subfolder_move_jobs_updated_at
|
||||
AFTER UPDATE
|
||||
ON image_subfolder_move_jobs FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE image_subfolder_move_jobs
|
||||
SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_33() -> Migration:
|
||||
return Migration(
|
||||
from_version=32,
|
||||
to_version=33,
|
||||
callback=Migration33Callback(),
|
||||
)
|
||||
@@ -0,0 +1,85 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration4Callback:
|
||||
"""Callback to do step 4 of migration."""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None: # noqa D102
|
||||
self._create_model_metadata(cursor)
|
||||
self._create_model_tags(cursor)
|
||||
self._create_tags(cursor)
|
||||
self._create_triggers(cursor)
|
||||
|
||||
def _create_model_metadata(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Create the table used to store model metadata downloaded from remote sources."""
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS model_metadata (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
name TEXT GENERATED ALWAYS AS (json_extract(metadata, '$.name')) VIRTUAL NOT NULL,
|
||||
author TEXT GENERATED ALWAYS AS (json_extract(metadata, '$.author')) VIRTUAL NOT NULL,
|
||||
-- Serialized JSON representation of the whole metadata object,
|
||||
-- which will contain additional fields from subclasses
|
||||
metadata TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
FOREIGN KEY(id) REFERENCES model_config(id) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def _create_model_tags(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS model_tags (
|
||||
model_id TEXT NOT NULL,
|
||||
tag_id INTEGER NOT NULL,
|
||||
FOREIGN KEY(model_id) REFERENCES model_config(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY(tag_id) REFERENCES tags(tag_id) ON DELETE CASCADE,
|
||||
UNIQUE(model_id,tag_id)
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def _create_tags(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS tags (
|
||||
tag_id INTEGER NOT NULL PRIMARY KEY,
|
||||
tag_text TEXT NOT NULL UNIQUE
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def _create_triggers(self, cursor: sqlite3.Cursor) -> None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS model_metadata_updated_at
|
||||
AFTER UPDATE
|
||||
ON model_metadata FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE model_metadata SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_4() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 3 to 4.
|
||||
|
||||
Adds the tables needed to store model metadata and tags.
|
||||
"""
|
||||
migration_4 = Migration(
|
||||
id="migration_4",
|
||||
depends_on="migration_3",
|
||||
from_version=3,
|
||||
to_version=4,
|
||||
callback=Migration4Callback(),
|
||||
)
|
||||
|
||||
return migration_4
|
||||
@@ -0,0 +1,36 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration5Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._drop_graph_executions(cursor)
|
||||
|
||||
def _drop_graph_executions(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Drops the `graph_executions` table."""
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DROP TABLE IF EXISTS graph_executions;
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_5() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 4 to 5.
|
||||
|
||||
Introduced in v3.6.3, this migration:
|
||||
- Drops the `graph_executions` table. We are able to do this because we are moving the graph storage
|
||||
to be purely in-memory.
|
||||
"""
|
||||
migration_5 = Migration(
|
||||
id="migration_5",
|
||||
depends_on="migration_4",
|
||||
from_version=4,
|
||||
to_version=5,
|
||||
callback=Migration5Callback(),
|
||||
)
|
||||
|
||||
return migration_5
|
||||
@@ -0,0 +1,64 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration6Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._recreate_model_triggers(cursor)
|
||||
self._delete_ip_adapters(cursor)
|
||||
|
||||
def _recreate_model_triggers(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
Adds the timestamp trigger to the model_config table.
|
||||
|
||||
This trigger was inadvertently dropped in earlier migration scripts.
|
||||
"""
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS model_config_updated_at
|
||||
AFTER UPDATE
|
||||
ON model_config FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE model_config SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
)
|
||||
|
||||
def _delete_ip_adapters(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""
|
||||
Delete all the IP adapters.
|
||||
|
||||
The model manager will automatically find and re-add them after the migration
|
||||
is done. This allows the manager to add the correct image encoder to their
|
||||
configuration records.
|
||||
"""
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DELETE FROM model_config
|
||||
WHERE type='ip_adapter';
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def build_migration_6() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 5 to 6.
|
||||
|
||||
This migration does the following:
|
||||
- Adds the model_config_updated_at trigger if it does not exist
|
||||
- Delete all ip_adapter models so that the model prober can find and
|
||||
update with the correct image processor model.
|
||||
"""
|
||||
migration_6 = Migration(
|
||||
id="migration_6",
|
||||
depends_on="migration_5",
|
||||
from_version=5,
|
||||
to_version=6,
|
||||
callback=Migration6Callback(),
|
||||
)
|
||||
|
||||
return migration_6
|
||||
@@ -0,0 +1,90 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration7Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._create_models_table(cursor)
|
||||
self._drop_old_models_tables(cursor)
|
||||
|
||||
def _drop_old_models_tables(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Drops the old model_records, model_metadata, model_tags and tags tables."""
|
||||
|
||||
tables = ["model_config", "model_metadata", "model_tags", "tags"]
|
||||
|
||||
for table in tables:
|
||||
cursor.execute(f"DROP TABLE IF EXISTS {table};")
|
||||
|
||||
def _create_models_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the v4.0.0 models table."""
|
||||
|
||||
tables = [
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS models (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
hash TEXT GENERATED ALWAYS as (json_extract(config, '$.hash')) VIRTUAL NOT NULL,
|
||||
base TEXT GENERATED ALWAYS as (json_extract(config, '$.base')) VIRTUAL NOT NULL,
|
||||
type TEXT GENERATED ALWAYS as (json_extract(config, '$.type')) VIRTUAL NOT NULL,
|
||||
path TEXT GENERATED ALWAYS as (json_extract(config, '$.path')) VIRTUAL NOT NULL,
|
||||
format TEXT GENERATED ALWAYS as (json_extract(config, '$.format')) VIRTUAL NOT NULL,
|
||||
name TEXT GENERATED ALWAYS as (json_extract(config, '$.name')) VIRTUAL NOT NULL,
|
||||
description TEXT GENERATED ALWAYS as (json_extract(config, '$.description')) VIRTUAL,
|
||||
source TEXT GENERATED ALWAYS as (json_extract(config, '$.source')) VIRTUAL NOT NULL,
|
||||
source_type TEXT GENERATED ALWAYS as (json_extract(config, '$.source_type')) VIRTUAL NOT NULL,
|
||||
source_api_response TEXT GENERATED ALWAYS as (json_extract(config, '$.source_api_response')) VIRTUAL,
|
||||
trigger_phrases TEXT GENERATED ALWAYS as (json_extract(config, '$.trigger_phrases')) VIRTUAL,
|
||||
-- Serialized JSON representation of the whole config object, which will contain additional fields from subclasses
|
||||
config TEXT NOT NULL,
|
||||
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- Updated via trigger
|
||||
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
|
||||
-- unique constraint on combo of name, base and type
|
||||
UNIQUE(name, base, type)
|
||||
);
|
||||
"""
|
||||
]
|
||||
|
||||
# Add trigger for `updated_at`.
|
||||
triggers = [
|
||||
"""--sql
|
||||
CREATE TRIGGER IF NOT EXISTS models_updated_at
|
||||
AFTER UPDATE
|
||||
ON models FOR EACH ROW
|
||||
BEGIN
|
||||
UPDATE models SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
|
||||
WHERE id = old.id;
|
||||
END;
|
||||
"""
|
||||
]
|
||||
|
||||
# Add indexes for searchable fields
|
||||
indices = [
|
||||
"CREATE INDEX IF NOT EXISTS base_index ON models(base);",
|
||||
"CREATE INDEX IF NOT EXISTS type_index ON models(type);",
|
||||
"CREATE INDEX IF NOT EXISTS name_index ON models(name);",
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS path_index ON models(path);",
|
||||
]
|
||||
|
||||
for stmt in tables + indices + triggers:
|
||||
cursor.execute(stmt)
|
||||
|
||||
|
||||
def build_migration_7() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 6 to 7.
|
||||
|
||||
This migration does the following:
|
||||
- Adds the new models table
|
||||
- Drops the old model_records, model_metadata, model_tags and tags tables.
|
||||
- TODO(MM2): Migrates model names and descriptions from `models.yaml` to the new table (?).
|
||||
"""
|
||||
migration_7 = Migration(
|
||||
id="migration_7",
|
||||
depends_on="migration_6",
|
||||
from_version=6,
|
||||
to_version=7,
|
||||
callback=Migration7Callback(),
|
||||
)
|
||||
|
||||
return migration_7
|
||||
@@ -0,0 +1,93 @@
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration8Callback:
|
||||
def __init__(self, app_config: InvokeAIAppConfig) -> None:
|
||||
self._app_config = app_config
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._drop_model_config_table(cursor)
|
||||
self._migrate_abs_models_to_rel(cursor)
|
||||
|
||||
def _drop_model_config_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Drops the old model_config table. This was missed in a previous migration."""
|
||||
|
||||
cursor.execute("DROP TABLE IF EXISTS model_config;")
|
||||
|
||||
def _migrate_abs_models_to_rel(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Check all model paths & legacy config paths to determine if they are inside Invoke-managed directories. If
|
||||
they are, update the paths to be relative to the managed directories.
|
||||
|
||||
This migration is a no-op for normal users (their paths will already be relative), but is necessary for users
|
||||
who have been testing the RCs with their live databases. The paths were made absolute in the initial RC, but this
|
||||
change was reverted. To smooth over the revert for our tests, we can migrate the paths back to relative.
|
||||
"""
|
||||
|
||||
models_path = self._app_config.models_path
|
||||
legacy_conf_path = self._app_config.legacy_conf_path
|
||||
legacy_conf_dir = self._app_config.legacy_conf_dir
|
||||
|
||||
stmt = """---sql
|
||||
SELECT
|
||||
id,
|
||||
path,
|
||||
json_extract(config, '$.config_path') as config_path
|
||||
FROM models;
|
||||
"""
|
||||
|
||||
all_models = cursor.execute(stmt).fetchall()
|
||||
|
||||
for model_id, model_path, model_config_path in all_models:
|
||||
# If the model path is inside the models directory, update it to be relative to the models directory.
|
||||
if Path(model_path).is_relative_to(models_path):
|
||||
new_path = Path(model_path).relative_to(models_path)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE models
|
||||
SET config = json_set(config, '$.path', ?)
|
||||
WHERE id = ?;
|
||||
""",
|
||||
(str(new_path), model_id),
|
||||
)
|
||||
# If the model has a legacy config path and it is inside the legacy conf directory, update it to be
|
||||
# relative to the legacy conf directory. This also fixes up cases in which the config path was
|
||||
# incorrectly relativized to the root directory. It will now be relativized to the legacy conf directory.
|
||||
if model_config_path:
|
||||
if Path(model_config_path).is_relative_to(legacy_conf_path):
|
||||
new_config_path = Path(model_config_path).relative_to(legacy_conf_path)
|
||||
elif Path(model_config_path).is_relative_to(legacy_conf_dir):
|
||||
new_config_path = Path(*Path(model_config_path).parts[1:])
|
||||
else:
|
||||
new_config_path = None
|
||||
if new_config_path:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE models
|
||||
SET config = json_set(config, '$.config_path', ?)
|
||||
WHERE id = ?;
|
||||
""",
|
||||
(str(new_config_path), model_id),
|
||||
)
|
||||
|
||||
|
||||
def build_migration_8(app_config: InvokeAIAppConfig) -> Migration:
|
||||
"""
|
||||
Build the migration from database version 7 to 8.
|
||||
|
||||
This migration does the following:
|
||||
- Removes the `model_config` table.
|
||||
- Migrates absolute model & legacy config paths to be relative to the models directory.
|
||||
"""
|
||||
migration_8 = Migration(
|
||||
id="migration_8",
|
||||
depends_on="migration_7",
|
||||
from_version=7,
|
||||
to_version=8,
|
||||
callback=Migration8Callback(app_config),
|
||||
)
|
||||
|
||||
return migration_8
|
||||
@@ -0,0 +1,31 @@
|
||||
import sqlite3
|
||||
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration
|
||||
|
||||
|
||||
class Migration9Callback:
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._empty_session_queue(cursor)
|
||||
|
||||
def _empty_session_queue(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Empties the session queue. This is done to prevent any lingering session queue items from causing pydantic errors due to changed schemas."""
|
||||
|
||||
cursor.execute("DELETE FROM session_queue;")
|
||||
|
||||
|
||||
def build_migration_9() -> Migration:
|
||||
"""
|
||||
Build the migration from database version 8 to 9.
|
||||
|
||||
This migration does the following:
|
||||
- Empties the session queue. This is done to prevent any lingering session queue items from causing pydantic errors due to changed schemas.
|
||||
"""
|
||||
migration_9 = Migration(
|
||||
id="migration_9",
|
||||
depends_on="migration_8",
|
||||
from_version=8,
|
||||
to_version=9,
|
||||
callback=Migration9Callback(),
|
||||
)
|
||||
|
||||
return migration_9
|
||||
@@ -0,0 +1,284 @@
|
||||
import sqlite3
|
||||
from typing import Optional, Protocol, runtime_checkable
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class MigrateCallback(Protocol):
|
||||
"""
|
||||
A callback that performs a migration.
|
||||
|
||||
Migrate callbacks are provided an open cursor to the database. They should not commit their
|
||||
transaction; this is handled by the migrator.
|
||||
|
||||
If the callback needs to access additional dependencies, will be provided to the callback at runtime.
|
||||
|
||||
See :class:`Migration` for an example.
|
||||
"""
|
||||
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None: ...
|
||||
|
||||
|
||||
class MigrationError(RuntimeError):
|
||||
"""Raised when a migration fails."""
|
||||
|
||||
|
||||
class MigrationVersionError(ValueError):
|
||||
"""Raised when a migration version is invalid."""
|
||||
|
||||
|
||||
class Migration(BaseModel):
|
||||
"""
|
||||
Represents a migration for a SQLite database.
|
||||
|
||||
:param from_version: The legacy database version on which this migration may be run
|
||||
:param to_version: The legacy database version that results from this migration
|
||||
:param id: The stable migration ID. Legacy migrations default to ``migration_{to_version}``.
|
||||
:param depends_on: The stable ID of the migration that must run first.
|
||||
:param migrate_callback: The callback to run to perform the migration
|
||||
|
||||
Migrations are executed according to their stable ID dependencies. Existing legacy migrations also keep
|
||||
``from_version`` and ``to_version`` so older numeric migration state can be mapped to applied migration IDs.
|
||||
New graph-only migrations may omit legacy versions, but must provide an explicit ``id``.
|
||||
|
||||
Migration callbacks will be provided an open cursor to the database. They should not commit their
|
||||
transaction; this is handled by the migrator.
|
||||
|
||||
It is suggested to use a class to define the migration callback and a builder function to create
|
||||
the :class:`Migration`. This allows the callback to be provided with additional dependencies and
|
||||
keeps things tidy, as all migration logic is self-contained.
|
||||
|
||||
Example:
|
||||
```py
|
||||
# Define the migration callback class
|
||||
class Migration1Callback:
|
||||
# This migration needs a logger, so we define a class that accepts a logger in its constructor.
|
||||
def __init__(self, image_files: ImageFileStorageBase) -> None:
|
||||
self._image_files = ImageFileStorageBase
|
||||
|
||||
# This dunder method allows the instance of the class to be called like a function.
|
||||
def __call__(self, cursor: sqlite3.Cursor) -> None:
|
||||
self._add_with_banana_column(cursor)
|
||||
self._do_something_with_images(cursor)
|
||||
|
||||
def _add_with_banana_column(self, cursor: sqlite3.Cursor) -> None:
|
||||
\"""Adds the with_banana column to the sushi table.\"""
|
||||
# Execute SQL using the cursor, taking care to *not commit* a transaction
|
||||
cursor.execute('ALTER TABLE sushi ADD COLUMN with_banana BOOLEAN DEFAULT TRUE;')
|
||||
|
||||
def _do_something_with_images(self, cursor: sqlite3.Cursor) -> None:
|
||||
\"""Does something with the image files service.\"""
|
||||
self._image_files.get(...)
|
||||
|
||||
# Define the migration builder function. This function creates an instance of the migration callback
|
||||
# class and returns a Migration.
|
||||
def build_migration_1(image_files: ImageFileStorageBase) -> Migration:
|
||||
\"""Builds the migration from database version 0 to 1.
|
||||
Requires the image files service to...
|
||||
\"""
|
||||
|
||||
migration_1 = Migration(
|
||||
from_version=0,
|
||||
to_version=1,
|
||||
migrate_callback=Migration1Callback(image_files=image_files),
|
||||
)
|
||||
|
||||
return migration_1
|
||||
|
||||
# Register the migration after all dependencies have been initialized
|
||||
db = SqliteDatabase(db_path, logger)
|
||||
migrator = SqliteMigrator(db)
|
||||
migrator.register_migration(build_migration_1(image_files))
|
||||
migrator.run_migrations()
|
||||
```
|
||||
"""
|
||||
|
||||
from_version: Optional[int] = Field(
|
||||
default=None, ge=0, strict=True, description="The database version on which this migration may be run"
|
||||
)
|
||||
to_version: Optional[int] = Field(
|
||||
default=None, ge=1, strict=True, description="The database version that results from this migration"
|
||||
)
|
||||
id: Optional[str] = Field(default=None, description="Stable migration ID")
|
||||
depends_on: Optional[str] = Field(default=None, description="Stable ID of the migration dependency")
|
||||
callback: MigrateCallback = Field(description="The callback to run to perform the migration")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_versions_and_ids(self) -> "Migration":
|
||||
"""Validates legacy versions and derives stable IDs for legacy migrations."""
|
||||
has_from_version = self.from_version is not None
|
||||
has_to_version = self.to_version is not None
|
||||
if has_from_version != has_to_version:
|
||||
raise MigrationVersionError("from_version and to_version must both be provided for legacy migrations")
|
||||
if self.from_version is not None and self.to_version is not None and self.to_version != self.from_version + 1:
|
||||
raise MigrationVersionError("to_version must be one greater than from_version")
|
||||
if self.id is None and self.to_version is not None:
|
||||
self.id = f"migration_{self.to_version}"
|
||||
if self.id is None:
|
||||
raise MigrationVersionError("id is required for graph-only migrations")
|
||||
if self.depends_on is None and self.from_version is not None and self.from_version > 0:
|
||||
self.depends_on = f"migration_{self.from_version}"
|
||||
if self.depends_on == self.id:
|
||||
raise MigrationVersionError("migration cannot depend on itself")
|
||||
return self
|
||||
|
||||
def __hash__(self) -> int:
|
||||
# Callables are not hashable, so we need to implement our own __hash__ function to use this class in a set.
|
||||
if self.from_version is not None and self.to_version is not None:
|
||||
return hash((self.from_version, self.to_version))
|
||||
return hash(self.id)
|
||||
|
||||
@property
|
||||
def sort_key(self) -> tuple[int, int, str]:
|
||||
"""Deterministic sort key for runnable migrations."""
|
||||
if self.to_version is None:
|
||||
return (1, 0, self.id or "")
|
||||
return (0, self.to_version, self.id or "")
|
||||
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
|
||||
class MigrationSet:
|
||||
"""
|
||||
A set of Migrations. Performs validation during migration registration and provides utility methods.
|
||||
|
||||
Migrations should be registered with `register()`. Once all are registered, `validate_dependency_graph()`
|
||||
should be called to ensure that dependencies are complete and acyclic. `validate_migration_chain()` is retained for
|
||||
legacy chain validation tests and compatibility checks.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._migrations: set[Migration] = set()
|
||||
|
||||
def register(self, migration: Migration) -> None:
|
||||
"""Registers a migration."""
|
||||
migration_from_already_registered = migration.from_version is not None and any(
|
||||
m.from_version == migration.from_version for m in self._migrations if m.from_version is not None
|
||||
)
|
||||
migration_to_already_registered = migration.to_version is not None and any(
|
||||
m.to_version == migration.to_version for m in self._migrations if m.to_version is not None
|
||||
)
|
||||
if migration_from_already_registered or migration_to_already_registered:
|
||||
raise MigrationVersionError("Migration with from_version or to_version already registered")
|
||||
migration_id_already_registered = any(m.id == migration.id for m in self._migrations)
|
||||
if migration_id_already_registered:
|
||||
raise MigrationVersionError("Migration with id already registered")
|
||||
self._migrations.add(migration)
|
||||
|
||||
def get(self, from_version: int) -> Optional[Migration]:
|
||||
"""Gets the migration that may be run on the given database version."""
|
||||
# register() ensures that there is only one migration with a given from_version, so this is safe.
|
||||
return next((m for m in self._migrations if m.from_version == from_version), None)
|
||||
|
||||
def validate_migration_chain(self) -> None:
|
||||
"""
|
||||
Validates that the migrations form a single chain of migrations from version 0 to the latest version,
|
||||
Raises a MigrationError if there is a problem.
|
||||
"""
|
||||
if self.count == 0:
|
||||
return
|
||||
if self.latest_version == 0:
|
||||
return
|
||||
next_migration = self.get(from_version=0)
|
||||
if next_migration is None:
|
||||
raise MigrationError("Migration chain is fragmented")
|
||||
touched_count = 1
|
||||
while next_migration is not None:
|
||||
next_migration = self.get(next_migration.to_version)
|
||||
if next_migration is not None:
|
||||
touched_count += 1
|
||||
if touched_count != self.count:
|
||||
raise MigrationError("Migration chain is fragmented")
|
||||
|
||||
def validate_dependency_graph(self) -> None:
|
||||
"""Validates migration ID dependencies."""
|
||||
migration_ids = {migration.id for migration in self._migrations}
|
||||
migrations_by_id = self.migrations_by_id
|
||||
for migration in self._migrations:
|
||||
if migration.depends_on is not None and migration.depends_on not in migration_ids:
|
||||
raise MigrationError(
|
||||
f"Migration '{migration.id}' depends on unknown migration '{migration.depends_on}'"
|
||||
)
|
||||
if migration.to_version is not None and migration.depends_on is not None:
|
||||
dependency = migrations_by_id[migration.depends_on]
|
||||
if dependency.to_version is None:
|
||||
raise MigrationError(
|
||||
f"Legacy migration '{migration.id}' cannot depend on graph-only migration '{dependency.id}'"
|
||||
)
|
||||
|
||||
visiting: set[str] = set()
|
||||
visited: set[str] = set()
|
||||
|
||||
def visit(migration: Migration) -> None:
|
||||
migration_id = migration.id
|
||||
if migration_id is None:
|
||||
raise MigrationError("Migration is missing id")
|
||||
if migration_id in visited:
|
||||
return
|
||||
if migration_id in visiting:
|
||||
raise MigrationError("Migration dependency graph contains a cycle")
|
||||
visiting.add(migration_id)
|
||||
if migration.depends_on is not None:
|
||||
visit(migrations_by_id[migration.depends_on])
|
||||
visiting.remove(migration_id)
|
||||
visited.add(migration_id)
|
||||
|
||||
for migration in self._migrations:
|
||||
visit(migration)
|
||||
|
||||
def get_migration_plan(self, applied_migration_ids: set[str]) -> list[Migration]:
|
||||
"""Gets a deterministic migration plan from the set of applied migration IDs."""
|
||||
self.validate_dependency_graph()
|
||||
known_migration_ids = set(self.migrations_by_id)
|
||||
unknown_applied_ids = applied_migration_ids - known_migration_ids
|
||||
if unknown_applied_ids:
|
||||
unknown_ids = ", ".join(sorted(unknown_applied_ids))
|
||||
raise MigrationError(f"Database contains unknown applied migration IDs: {unknown_ids}")
|
||||
|
||||
plan: list[Migration] = []
|
||||
planned_or_applied_ids = set(applied_migration_ids)
|
||||
remaining = {
|
||||
migration.id: migration for migration in self._migrations if migration.id not in applied_migration_ids
|
||||
}
|
||||
|
||||
while remaining:
|
||||
runnable = sorted(
|
||||
(
|
||||
migration
|
||||
for migration in remaining.values()
|
||||
if migration.depends_on is None or migration.depends_on in planned_or_applied_ids
|
||||
),
|
||||
key=lambda migration: migration.sort_key,
|
||||
)
|
||||
if not runnable:
|
||||
raise MigrationError("Migration dependency graph cannot be resolved")
|
||||
migration = runnable[0]
|
||||
plan.append(migration)
|
||||
planned_or_applied_ids.add(migration.id or "")
|
||||
del remaining[migration.id]
|
||||
return plan
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
"""The count of registered migrations."""
|
||||
return len(self._migrations)
|
||||
|
||||
@property
|
||||
def latest_version(self) -> int:
|
||||
"""Gets latest to_version among registered migrations. Returns 0 if there are no migrations registered."""
|
||||
if self.count == 0:
|
||||
return 0
|
||||
legacy_migrations = [migration for migration in self._migrations if migration.to_version is not None]
|
||||
if len(legacy_migrations) == 0:
|
||||
return 0
|
||||
latest_version = sorted(legacy_migrations, key=lambda m: m.to_version or 0)[-1].to_version
|
||||
return latest_version or 0
|
||||
|
||||
@property
|
||||
def migrations(self) -> tuple[Migration, ...]:
|
||||
return tuple(sorted(self._migrations, key=lambda migration: migration.sort_key))
|
||||
|
||||
@property
|
||||
def migrations_by_id(self) -> dict[str, Migration]:
|
||||
return {migration.id or "": migration for migration in self._migrations}
|
||||
@@ -0,0 +1,302 @@
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration, MigrationError, MigrationSet
|
||||
|
||||
|
||||
class SqliteMigrator:
|
||||
"""
|
||||
Manages migrations for a SQLite database.
|
||||
|
||||
:param db: The instance of :class:`SqliteDatabase` to migrate.
|
||||
|
||||
Migrations should be registered with :meth:`register_migration`, either directly or via the migration loader.
|
||||
They are planned by stable migration ID dependencies and recorded in the ``applied_migrations`` table.
|
||||
Legacy numeric versions are still written for migrations that define ``to_version``.
|
||||
|
||||
Each migration is run in a transaction. If a migration fails, the transaction is rolled back.
|
||||
|
||||
Example Usage:
|
||||
```py
|
||||
db = SqliteDatabase(db_path="my_db.db", logger=logger)
|
||||
migrator = SqliteMigrator(db=db)
|
||||
migrator.register_migration(build_migration_1())
|
||||
migrator.register_migration(build_migration_2())
|
||||
migrator.run_migrations()
|
||||
```
|
||||
"""
|
||||
|
||||
backup_path: Optional[Path] = None
|
||||
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
self._db = db
|
||||
self._logger = db._logger
|
||||
self._migration_set = MigrationSet()
|
||||
self._backup_path: Optional[Path] = None
|
||||
|
||||
def register_migration(self, migration: Migration) -> None:
|
||||
"""Registers a migration."""
|
||||
self._migration_set.register(migration)
|
||||
self._logger.debug(f"Registered migration {migration.from_version} -> {migration.to_version}")
|
||||
|
||||
def run_migrations(self) -> bool:
|
||||
"""Migrates the database to the latest version."""
|
||||
# This throws if there is a problem.
|
||||
self._migration_set.validate_dependency_graph()
|
||||
cursor = self._db._conn.cursor()
|
||||
self._validate_existing_applied_migrations(cursor=cursor)
|
||||
self._create_migrations_table(cursor=cursor)
|
||||
self._validate_existing_legacy_migrations(cursor=cursor)
|
||||
if self._needs_applied_migrations_bootstrap(cursor=cursor):
|
||||
self._backup_db()
|
||||
self._create_applied_migrations_table(cursor=cursor)
|
||||
self._validate_existing_applied_legacy_migrations(cursor=cursor)
|
||||
self._bootstrap_applied_migrations_from_legacy_versions(cursor=cursor)
|
||||
|
||||
if self._migration_set.count == 0:
|
||||
self._logger.debug("No migrations registered")
|
||||
return False
|
||||
|
||||
applied_migration_ids = self._get_applied_migration_ids(cursor=cursor)
|
||||
migration_plan = self._migration_set.get_migration_plan(applied_migration_ids=applied_migration_ids)
|
||||
if len(migration_plan) == 0:
|
||||
self._logger.debug("Database is up to date, no migrations to run")
|
||||
return False
|
||||
|
||||
self._logger.info("Database update needed")
|
||||
|
||||
self._backup_db()
|
||||
|
||||
for migration in migration_plan:
|
||||
self._run_migration(migration)
|
||||
self._logger.info("Database updated successfully")
|
||||
return True
|
||||
|
||||
def _backup_db(self) -> None:
|
||||
"""Makes a backup of the db if it is a file db and a backup has not already been made."""
|
||||
if self._backup_path is not None:
|
||||
return
|
||||
if self._db._db_path is not None:
|
||||
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
self._backup_path = self._db._db_path.parent / f"{self._db._db_path.stem}_backup_{timestamp}.db"
|
||||
self._logger.info(f"Backing up database to {str(self._backup_path)}")
|
||||
with closing(sqlite3.connect(self._backup_path)) as backup_conn:
|
||||
self._db._conn.backup(backup_conn)
|
||||
else:
|
||||
self._logger.info("Using in-memory database, no backup needed")
|
||||
|
||||
def _run_migration(self, migration: Migration) -> None:
|
||||
"""Runs a single migration."""
|
||||
try:
|
||||
# Using sqlite3.Connection as a context manager commits a the transaction on exit, or rolls it back if an
|
||||
# exception is raised.
|
||||
with self._db._conn as conn:
|
||||
cursor = conn.cursor()
|
||||
self._create_applied_migrations_table(cursor)
|
||||
if migration.from_version is not None and self._get_current_version(cursor) != migration.from_version:
|
||||
raise MigrationError(
|
||||
f"Database is at version {self._get_current_version(cursor)}, expected {migration.from_version}"
|
||||
)
|
||||
self._logger.debug(f"Running migration '{migration.id}'")
|
||||
|
||||
# Run the actual migration
|
||||
migration.callback(cursor)
|
||||
|
||||
if migration.to_version is not None:
|
||||
cursor.execute("INSERT INTO migrations (version) VALUES (?);", (migration.to_version,))
|
||||
cursor.execute(
|
||||
"INSERT INTO applied_migrations (migration_id, legacy_version) VALUES (?, ?);",
|
||||
(migration.id, migration.to_version),
|
||||
)
|
||||
|
||||
self._logger.debug(f"Successfully ran migration '{migration.id}'")
|
||||
# We want to catch *any* error, mirroring the behaviour of the sqlite3 module.
|
||||
except Exception as e:
|
||||
# The connection context manager has already rolled back the migration, so we don't need to do anything.
|
||||
msg = f"Error running migration '{migration.id}': {e}"
|
||||
self._logger.error(msg)
|
||||
raise MigrationError(msg) from e
|
||||
|
||||
def _create_migrations_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the migrations table for the database, if one does not already exist."""
|
||||
try:
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='migrations';")
|
||||
if cursor.fetchone() is not None:
|
||||
return
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE migrations (
|
||||
version INTEGER PRIMARY KEY,
|
||||
migrated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW'))
|
||||
);
|
||||
"""
|
||||
)
|
||||
cursor.execute("INSERT INTO migrations (version) VALUES (0);")
|
||||
cursor.connection.commit()
|
||||
self._logger.debug("Created migrations table")
|
||||
except sqlite3.Error as e:
|
||||
msg = f"Problem creating migrations table: {e}"
|
||||
self._logger.error(msg)
|
||||
cursor.connection.rollback()
|
||||
raise MigrationError(msg) from e
|
||||
|
||||
def _create_applied_migrations_table(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Creates the applied migrations table for stable migration IDs."""
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
CREATE TABLE IF NOT EXISTS applied_migrations (
|
||||
migration_id TEXT PRIMARY KEY,
|
||||
legacy_version INTEGER UNIQUE,
|
||||
migrated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW'))
|
||||
);
|
||||
"""
|
||||
)
|
||||
except sqlite3.Error as e:
|
||||
msg = f"Problem creating applied_migrations table: {e}"
|
||||
self._logger.error(msg)
|
||||
cursor.connection.rollback()
|
||||
raise MigrationError(msg) from e
|
||||
|
||||
def _bootstrap_applied_migrations_from_legacy_versions(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Backfills applied migration IDs from legacy numeric migration rows."""
|
||||
try:
|
||||
cursor.execute("SELECT version FROM migrations WHERE version > 0 ORDER BY version;")
|
||||
legacy_versions = [row[0] for row in cursor.fetchall()]
|
||||
registered_migration_ids = self._migration_set.migrations_by_id
|
||||
for legacy_version in legacy_versions:
|
||||
migration_id = f"migration_{legacy_version}"
|
||||
if migration_id not in registered_migration_ids:
|
||||
cursor.connection.rollback()
|
||||
raise MigrationError(f"Database contains unknown legacy migration version: {legacy_version}")
|
||||
cursor.execute(
|
||||
"SELECT legacy_version FROM applied_migrations WHERE migration_id = ?;",
|
||||
(migration_id,),
|
||||
)
|
||||
migration_row = cursor.fetchone()
|
||||
if migration_row is not None and migration_row[0] != legacy_version:
|
||||
cursor.connection.rollback()
|
||||
raise MigrationError(
|
||||
"Database contains inconsistent applied migration state: "
|
||||
f"{migration_id} is recorded with legacy version {migration_row[0]}, "
|
||||
f"expected {legacy_version}"
|
||||
)
|
||||
cursor.execute(
|
||||
"SELECT migration_id FROM applied_migrations WHERE legacy_version = ?;",
|
||||
(legacy_version,),
|
||||
)
|
||||
legacy_row = cursor.fetchone()
|
||||
if legacy_row is not None and legacy_row[0] != migration_id:
|
||||
cursor.connection.rollback()
|
||||
raise MigrationError(
|
||||
"Database contains inconsistent applied migration state: "
|
||||
f"legacy version {legacy_version} is recorded for {legacy_row[0]}, expected {migration_id}"
|
||||
)
|
||||
cursor.execute(
|
||||
"INSERT OR IGNORE INTO applied_migrations (migration_id, legacy_version) VALUES (?, ?);",
|
||||
(migration_id, legacy_version),
|
||||
)
|
||||
cursor.connection.commit()
|
||||
except sqlite3.Error as e:
|
||||
msg = f"Problem bootstrapping applied migrations: {e}"
|
||||
self._logger.error(msg)
|
||||
cursor.connection.rollback()
|
||||
raise MigrationError(msg) from e
|
||||
|
||||
def _validate_existing_applied_migrations(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Validates existing applied migration IDs before creating or mutating migrator metadata."""
|
||||
applied_migration_ids = self._get_applied_migration_ids(cursor=cursor)
|
||||
if len(applied_migration_ids) == 0:
|
||||
return
|
||||
known_migration_ids = set(self._migration_set.migrations_by_id)
|
||||
unknown_applied_ids = applied_migration_ids - known_migration_ids
|
||||
if unknown_applied_ids:
|
||||
unknown_ids = ", ".join(sorted(unknown_applied_ids))
|
||||
raise MigrationError(f"Database contains unknown applied migration IDs: {unknown_ids}")
|
||||
|
||||
def _validate_existing_legacy_migrations(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Validates existing legacy migration versions before creating applied migration metadata."""
|
||||
try:
|
||||
cursor.execute("SELECT version FROM migrations WHERE version > 0 ORDER BY version;")
|
||||
except sqlite3.OperationalError as e:
|
||||
if "no such table" in str(e):
|
||||
return
|
||||
raise
|
||||
|
||||
registered_migration_ids = self._migration_set.migrations_by_id
|
||||
for row in cursor.fetchall():
|
||||
legacy_version = row[0]
|
||||
migration_id = f"migration_{legacy_version}"
|
||||
if migration_id not in registered_migration_ids:
|
||||
raise MigrationError(f"Database contains unknown legacy migration version: {legacy_version}")
|
||||
|
||||
def _validate_existing_applied_legacy_migrations(self, cursor: sqlite3.Cursor) -> None:
|
||||
"""Validates applied IDs for legacy migrations against legacy numeric rows."""
|
||||
registered_migrations = self._migration_set.migrations_by_id
|
||||
cursor.execute("SELECT migration_id, legacy_version FROM applied_migrations;")
|
||||
applied_rows = cursor.fetchall()
|
||||
|
||||
cursor.execute("SELECT version FROM migrations WHERE version > 0;")
|
||||
legacy_versions = {row[0] for row in cursor.fetchall()}
|
||||
|
||||
for row in applied_rows:
|
||||
migration_id = row[0]
|
||||
legacy_version = row[1]
|
||||
migration = registered_migrations[migration_id]
|
||||
if migration.to_version is None:
|
||||
continue
|
||||
if legacy_version != migration.to_version:
|
||||
raise MigrationError(
|
||||
"Database contains inconsistent applied migration state: "
|
||||
f"{migration_id} is recorded with legacy version {legacy_version}, "
|
||||
f"expected {migration.to_version}"
|
||||
)
|
||||
if legacy_version not in legacy_versions:
|
||||
raise MigrationError(
|
||||
"Database contains inconsistent applied migration state: "
|
||||
f"{migration_id} is applied, but legacy version {legacy_version} is missing"
|
||||
)
|
||||
|
||||
def _needs_applied_migrations_bootstrap(self, cursor: sqlite3.Cursor) -> bool:
|
||||
"""Checks whether legacy numeric rows need to be written to applied_migrations."""
|
||||
cursor.execute("SELECT version FROM migrations WHERE version > 0;")
|
||||
legacy_versions = {row[0] for row in cursor.fetchall()}
|
||||
if len(legacy_versions) == 0:
|
||||
return False
|
||||
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='applied_migrations';")
|
||||
if cursor.fetchone() is None:
|
||||
return True
|
||||
|
||||
cursor.execute("SELECT legacy_version FROM applied_migrations WHERE legacy_version IS NOT NULL;")
|
||||
applied_legacy_versions = {row[0] for row in cursor.fetchall()}
|
||||
return not legacy_versions.issubset(applied_legacy_versions)
|
||||
|
||||
@classmethod
|
||||
def _get_applied_migration_ids(cls, cursor: sqlite3.Cursor) -> set[str]:
|
||||
"""Gets applied stable migration IDs."""
|
||||
try:
|
||||
cursor.execute("SELECT migration_id FROM applied_migrations;")
|
||||
return {row[0] for row in cursor.fetchall()}
|
||||
except sqlite3.OperationalError as e:
|
||||
if "no such table" in str(e):
|
||||
return set()
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def _get_current_version(cls, cursor: sqlite3.Cursor) -> int:
|
||||
"""Gets the current version of the database, or 0 if the migrations table does not exist."""
|
||||
try:
|
||||
cursor.execute("SELECT MAX(version) FROM migrations;")
|
||||
version: int = cursor.fetchone()[0]
|
||||
if version is None:
|
||||
return 0
|
||||
return version
|
||||
except sqlite3.OperationalError as e:
|
||||
if "no such table" in str(e):
|
||||
return 0
|
||||
raise
|
||||
@@ -0,0 +1,229 @@
|
||||
from collections.abc import Mapping
|
||||
from copy import deepcopy
|
||||
from typing import Any, get_args, get_origin
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic_core import PydanticUndefined
|
||||
|
||||
from invokeai.app.invocations.baseinvocation import InvocationRegistry
|
||||
from invokeai.app.invocations.call_saved_workflow import parse_call_saved_workflow_dynamic_input
|
||||
from invokeai.app.invocations.fields import ImageField
|
||||
from invokeai.app.services.session_processor.workflow_call_batch import build_child_workflow_sessions
|
||||
from invokeai.app.services.session_queue.session_queue_common import TooManySessionsError
|
||||
from invokeai.app.services.shared.graph import Graph, GraphExecutionState, WorkflowCallFrame
|
||||
from invokeai.app.services.shared.workflow_call_compatibility_common import (
|
||||
WorkflowCallCompatibility,
|
||||
WorkflowCallCompatibilityReason,
|
||||
)
|
||||
from invokeai.app.services.shared.workflow_graph_builder import (
|
||||
InvalidWorkflowInputError,
|
||||
UnsupportedWorkflowNodeError,
|
||||
get_exposed_workflow_input_names,
|
||||
)
|
||||
|
||||
|
||||
def _count_workflow_return_nodes(workflow: dict[str, Any]) -> int:
|
||||
workflow_return_count = 0
|
||||
for node in workflow.get("nodes", []):
|
||||
if not isinstance(node, dict) or node.get("type") != "invocation":
|
||||
continue
|
||||
data = node.get("data")
|
||||
if isinstance(data, dict) and data.get("type") == "workflow_return":
|
||||
workflow_return_count += 1
|
||||
return workflow_return_count
|
||||
|
||||
|
||||
def _is_mapping(value: Any) -> bool:
|
||||
return isinstance(value, Mapping)
|
||||
|
||||
|
||||
def _build_placeholder_model(annotation: type[BaseModel]) -> Any:
|
||||
values: dict[str, Any] = {}
|
||||
for field_name, field_info in annotation.model_fields.items():
|
||||
if field_info.default is not PydanticUndefined:
|
||||
values[field_name] = deepcopy(field_info.default)
|
||||
continue
|
||||
if field_info.default_factory is not None:
|
||||
values[field_name] = field_info.default_factory()
|
||||
continue
|
||||
placeholder = _get_placeholder_for_annotation(field_info.annotation)
|
||||
if placeholder is None:
|
||||
return None
|
||||
values[field_name] = placeholder
|
||||
return annotation.model_construct(**values)
|
||||
|
||||
|
||||
def _get_placeholder_for_annotation(annotation: Any) -> Any:
|
||||
origin = get_origin(annotation)
|
||||
if origin is not None:
|
||||
if origin is list:
|
||||
return []
|
||||
if origin is dict:
|
||||
return {}
|
||||
if origin is tuple:
|
||||
return []
|
||||
if origin is set:
|
||||
return []
|
||||
args = [arg for arg in get_args(annotation) if arg is not type(None)]
|
||||
if args:
|
||||
return _get_placeholder_for_annotation(args[0])
|
||||
return None
|
||||
|
||||
if annotation is Any:
|
||||
return {}
|
||||
if annotation is str:
|
||||
return ""
|
||||
if annotation is int:
|
||||
return 0
|
||||
if annotation is float:
|
||||
return 0.0
|
||||
if annotation is bool:
|
||||
return False
|
||||
if annotation is ImageField:
|
||||
return ImageField(image_name="compatibility-placeholder")
|
||||
if isinstance(annotation, type) and issubclass(annotation, BaseModel):
|
||||
return _build_placeholder_model(annotation)
|
||||
return None
|
||||
|
||||
|
||||
def _build_compatibility_workflow_inputs(workflow: dict[str, Any]) -> dict[str, Any]:
|
||||
workflow_inputs: dict[str, Any] = {}
|
||||
workflow_nodes = workflow.get("nodes", [])
|
||||
if not isinstance(workflow_nodes, list):
|
||||
return workflow_inputs
|
||||
|
||||
nodes_by_id = {
|
||||
node.get("id"): node
|
||||
for node in workflow_nodes
|
||||
if _is_mapping(node) and isinstance(node.get("id"), str) and _is_mapping(node.get("data"))
|
||||
}
|
||||
|
||||
for input_name in get_exposed_workflow_input_names(workflow):
|
||||
node_id, field_name = parse_call_saved_workflow_dynamic_input(input_name)
|
||||
node = nodes_by_id.get(node_id)
|
||||
if not _is_mapping(node):
|
||||
continue
|
||||
node_data = node.get("data")
|
||||
if not _is_mapping(node_data):
|
||||
continue
|
||||
node_type = node_data.get("type")
|
||||
if not isinstance(node_type, str):
|
||||
continue
|
||||
invocation_class = InvocationRegistry.get_invocation_for_type(node_type)
|
||||
if invocation_class is None:
|
||||
continue
|
||||
field_info = invocation_class.model_fields.get(field_name)
|
||||
if field_info is None:
|
||||
continue
|
||||
if field_info.default is not PydanticUndefined:
|
||||
workflow_inputs[input_name] = deepcopy(field_info.default)
|
||||
continue
|
||||
if field_info.default_factory is not None:
|
||||
workflow_inputs[input_name] = field_info.default_factory()
|
||||
continue
|
||||
placeholder = _get_placeholder_for_annotation(field_info.annotation)
|
||||
if placeholder is not None:
|
||||
workflow_inputs[input_name] = placeholder
|
||||
|
||||
return workflow_inputs
|
||||
|
||||
|
||||
def _is_unsupported_batch_input_message(message: str) -> bool:
|
||||
return any(
|
||||
marker in message
|
||||
for marker in (
|
||||
"batch child workflow",
|
||||
"batch group",
|
||||
"batch input",
|
||||
"batch inputs",
|
||||
"batch node",
|
||||
"batch-special child workflow nodes",
|
||||
"connected batch",
|
||||
"generator-backed batch",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_workflow_call_compatibility(
|
||||
*,
|
||||
workflow: dict[str, Any],
|
||||
workflow_id: str,
|
||||
services: Any,
|
||||
user_id: str | None,
|
||||
maximum_children: int,
|
||||
resolve_generator_items: bool = True,
|
||||
) -> WorkflowCallCompatibility:
|
||||
workflow_return_count = _count_workflow_return_nodes(workflow)
|
||||
if workflow_return_count == 0:
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=WorkflowCallCompatibilityReason.MissingWorkflowReturn,
|
||||
message="The workflow must contain exactly one workflow_return node.",
|
||||
)
|
||||
if workflow_return_count > 1:
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=WorkflowCallCompatibilityReason.MultipleWorkflowReturn,
|
||||
message="The workflow must not contain more than one workflow_return node.",
|
||||
)
|
||||
|
||||
try:
|
||||
workflow_inputs = _build_compatibility_workflow_inputs(workflow)
|
||||
build_child_workflow_sessions(
|
||||
parent_session=GraphExecutionState(graph=Graph()),
|
||||
workflow=workflow,
|
||||
workflow_inputs=workflow_inputs,
|
||||
call_frame=WorkflowCallFrame(
|
||||
prepared_call_node_id="compatibility-call",
|
||||
source_call_node_id="compatibility-call",
|
||||
workflow_id=workflow_id,
|
||||
depth=1,
|
||||
),
|
||||
maximum_children=maximum_children,
|
||||
services=services,
|
||||
user_id=user_id,
|
||||
resolve_generator_items=resolve_generator_items,
|
||||
)
|
||||
except InvalidWorkflowInputError as e:
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=WorkflowCallCompatibilityReason.InvalidInputs,
|
||||
message=str(e),
|
||||
)
|
||||
except UnsupportedWorkflowNodeError as e:
|
||||
message = str(e)
|
||||
reason = WorkflowCallCompatibilityReason.UnsupportedNode
|
||||
if _is_unsupported_batch_input_message(message):
|
||||
reason = WorkflowCallCompatibilityReason.UnsupportedBatchInput
|
||||
elif "exactly one workflow_return" in message and workflow_return_count == 0:
|
||||
reason = WorkflowCallCompatibilityReason.MissingWorkflowReturn
|
||||
elif "exactly one workflow_return" in message:
|
||||
reason = WorkflowCallCompatibilityReason.InvalidGraph
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=reason,
|
||||
message=message,
|
||||
)
|
||||
except TooManySessionsError as e:
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=WorkflowCallCompatibilityReason.ExceedsCapacity,
|
||||
message=str(e),
|
||||
)
|
||||
except ValueError as e:
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=WorkflowCallCompatibilityReason.InvalidGraph,
|
||||
message=str(e),
|
||||
)
|
||||
except Exception as e:
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=False,
|
||||
reason=WorkflowCallCompatibilityReason.Unknown,
|
||||
message=str(e),
|
||||
)
|
||||
|
||||
return WorkflowCallCompatibility(
|
||||
is_callable=True,
|
||||
reason=WorkflowCallCompatibilityReason.Ok,
|
||||
)
|
||||
@@ -0,0 +1,21 @@
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class WorkflowCallCompatibilityReason(str, Enum):
|
||||
Ok = "ok"
|
||||
MissingWorkflowReturn = "missing_workflow_return"
|
||||
MultipleWorkflowReturn = "multiple_workflow_return"
|
||||
UnsupportedNode = "unsupported_node"
|
||||
UnsupportedBatchInput = "unsupported_batch_input"
|
||||
InvalidGraph = "invalid_graph"
|
||||
InvalidInputs = "invalid_inputs"
|
||||
ExceedsCapacity = "exceeds_capacity"
|
||||
Unknown = "unknown"
|
||||
|
||||
|
||||
class WorkflowCallCompatibility(BaseModel):
|
||||
is_callable: bool = Field(description="Whether the workflow can currently be executed by call_saved_workflow.")
|
||||
reason: WorkflowCallCompatibilityReason = Field(description="Structured compatibility result.")
|
||||
message: str | None = Field(default=None, description="Human-readable compatibility detail when unavailable.")
|
||||
@@ -0,0 +1,375 @@
|
||||
from collections.abc import Mapping, MutableMapping, Sequence
|
||||
from typing import Any
|
||||
|
||||
from invokeai.app.invocations.baseinvocation import Classification, InvocationRegistry
|
||||
from invokeai.app.invocations.call_saved_workflow import (
|
||||
CALL_SAVED_WORKFLOW_DYNAMIC_FIELD_PREFIX,
|
||||
parse_call_saved_workflow_dynamic_input,
|
||||
)
|
||||
from invokeai.app.services.shared.graph import Edge, EdgeConnection, Graph
|
||||
|
||||
CONNECTOR_INPUT_HANDLE = "in"
|
||||
CONNECTOR_OUTPUT_HANDLE = "out"
|
||||
|
||||
|
||||
class UnsupportedWorkflowNodeError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
class InvalidWorkflowInputError(ValueError):
|
||||
pass
|
||||
|
||||
|
||||
def _is_mapping(value: Any) -> bool:
|
||||
return isinstance(value, Mapping)
|
||||
|
||||
|
||||
def _is_invocation_node(node: Any) -> bool:
|
||||
return _is_mapping(node) and node.get("type") == "invocation" and _is_mapping(node.get("data"))
|
||||
|
||||
|
||||
def _is_connector_node(node: Any) -> bool:
|
||||
return _is_mapping(node) and node.get("type") == "connector"
|
||||
|
||||
|
||||
def _build_dynamic_input_name(node_id: str, field_name: str) -> str:
|
||||
return f"{CALL_SAVED_WORKFLOW_DYNAMIC_FIELD_PREFIX}{node_id}::{field_name}"
|
||||
|
||||
|
||||
def _get_form_elements(workflow: Mapping[str, Any]) -> tuple[Mapping[str, Any], str | None]:
|
||||
form = workflow.get("form")
|
||||
if not _is_mapping(form):
|
||||
return {}, None
|
||||
|
||||
elements = form.get("elements")
|
||||
root_element_id = form.get("rootElementId")
|
||||
if not _is_mapping(elements) or not isinstance(root_element_id, str):
|
||||
return {}, None
|
||||
|
||||
return elements, root_element_id
|
||||
|
||||
|
||||
def _collect_exposed_inputs_from_form(workflow: Mapping[str, Any]) -> set[str]:
|
||||
elements, root_element_id = _get_form_elements(workflow)
|
||||
if not elements or root_element_id is None:
|
||||
return set()
|
||||
|
||||
exposed_inputs: set[str] = set()
|
||||
stack = [root_element_id]
|
||||
visited: set[str] = set()
|
||||
|
||||
while stack:
|
||||
element_id = stack.pop()
|
||||
if element_id in visited:
|
||||
continue
|
||||
visited.add(element_id)
|
||||
|
||||
element = elements.get(element_id)
|
||||
if not _is_mapping(element):
|
||||
continue
|
||||
|
||||
if element.get("type") == "node-field":
|
||||
data = element.get("data")
|
||||
if _is_mapping(data):
|
||||
field_identifier = data.get("fieldIdentifier")
|
||||
if _is_mapping(field_identifier):
|
||||
node_id = field_identifier.get("nodeId")
|
||||
field_name = field_identifier.get("fieldName")
|
||||
if isinstance(node_id, str) and isinstance(field_name, str):
|
||||
exposed_inputs.add(_build_dynamic_input_name(node_id, field_name))
|
||||
|
||||
data = element.get("data")
|
||||
if _is_mapping(data):
|
||||
children = data.get("children")
|
||||
if isinstance(children, Sequence):
|
||||
for child_id in reversed(children):
|
||||
if isinstance(child_id, str):
|
||||
stack.append(child_id)
|
||||
|
||||
return exposed_inputs
|
||||
|
||||
|
||||
def get_exposed_workflow_input_names(workflow: Mapping[str, Any]) -> set[str]:
|
||||
exposed_inputs = _collect_exposed_inputs_from_form(workflow)
|
||||
if exposed_inputs:
|
||||
return exposed_inputs
|
||||
|
||||
workflow_exposed_fields = workflow.get("exposedFields", [])
|
||||
if not isinstance(workflow_exposed_fields, Sequence):
|
||||
return set()
|
||||
|
||||
fallback_inputs: set[str] = set()
|
||||
for field in workflow_exposed_fields:
|
||||
if not _is_mapping(field):
|
||||
continue
|
||||
node_id = field.get("nodeId")
|
||||
field_name = field.get("fieldName")
|
||||
if isinstance(node_id, str) and isinstance(field_name, str):
|
||||
fallback_inputs.add(_build_dynamic_input_name(node_id, field_name))
|
||||
|
||||
return fallback_inputs
|
||||
|
||||
|
||||
def apply_workflow_inputs_to_workflow(workflow: MutableMapping[str, Any], workflow_inputs: Mapping[str, Any]) -> None:
|
||||
if not workflow_inputs:
|
||||
return
|
||||
|
||||
allowed_inputs = get_exposed_workflow_input_names(workflow)
|
||||
for input_name, value in workflow_inputs.items():
|
||||
if input_name not in allowed_inputs:
|
||||
raise InvalidWorkflowInputError(
|
||||
f"call_saved_workflow input '{input_name}' is not exposed by the selected workflow"
|
||||
)
|
||||
|
||||
node_id, field_name = parse_call_saved_workflow_dynamic_input(input_name)
|
||||
workflow_nodes = workflow.get("nodes", [])
|
||||
if not isinstance(workflow_nodes, list):
|
||||
raise InvalidWorkflowInputError(
|
||||
f"call_saved_workflow input '{input_name}' targets missing child workflow node '{node_id}'"
|
||||
)
|
||||
matching_node = next(
|
||||
(
|
||||
node
|
||||
for node in workflow_nodes
|
||||
if _is_mapping(node)
|
||||
and _is_mapping(node.get("data"))
|
||||
and node.get("id") == node_id
|
||||
and node["data"].get("id") == node_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if matching_node is None:
|
||||
raise InvalidWorkflowInputError(
|
||||
f"call_saved_workflow input '{input_name}' targets missing child workflow node '{node_id}'"
|
||||
)
|
||||
matching_node_data = matching_node["data"]
|
||||
node_type = matching_node_data.get("type")
|
||||
if not isinstance(node_type, str):
|
||||
raise InvalidWorkflowInputError(
|
||||
f"call_saved_workflow input '{input_name}' targets missing child workflow node '{node_id}'"
|
||||
)
|
||||
invocation_class = InvocationRegistry.get_invocation_for_type(node_type)
|
||||
if invocation_class is None or field_name not in invocation_class.model_fields:
|
||||
raise InvalidWorkflowInputError(
|
||||
f"call_saved_workflow input '{input_name}' targets missing child workflow field '{field_name}'"
|
||||
)
|
||||
inputs = matching_node_data.setdefault("inputs", {})
|
||||
if not _is_mapping(inputs):
|
||||
raise InvalidWorkflowInputError(
|
||||
f"call_saved_workflow input '{input_name}' targets invalid child workflow inputs on '{node_id}'"
|
||||
)
|
||||
inputs[field_name] = {"value": value}
|
||||
|
||||
|
||||
def apply_workflow_inputs_to_graph(
|
||||
graph: Graph, workflow: Mapping[str, Any], workflow_inputs: Mapping[str, Any]
|
||||
) -> None:
|
||||
if not workflow_inputs:
|
||||
return
|
||||
|
||||
mutable_workflow = dict(workflow)
|
||||
apply_workflow_inputs_to_workflow(mutable_workflow, workflow_inputs)
|
||||
for input_name, value in workflow_inputs.items():
|
||||
node_id, field_name = parse_call_saved_workflow_dynamic_input(input_name)
|
||||
node = graph.nodes.get(node_id)
|
||||
if node is None:
|
||||
continue
|
||||
setattr(node, field_name, value)
|
||||
|
||||
|
||||
def _raise_if_unsupported_invocation_type(node_type: str, node_id: str) -> None:
|
||||
invocation_class = InvocationRegistry.get_invocation_for_type(node_type)
|
||||
if invocation_class is None:
|
||||
return
|
||||
|
||||
if (
|
||||
invocation_class.UIConfig.category == "batch"
|
||||
and invocation_class.UIConfig.classification == Classification.Special
|
||||
and not node_type.endswith("_generator")
|
||||
):
|
||||
raise UnsupportedWorkflowNodeError(
|
||||
f"call_saved_workflow does not yet support batch-special child workflow nodes such as "
|
||||
f"'{node_type}' (node '{node_id}')"
|
||||
)
|
||||
|
||||
|
||||
def _validate_callable_workflow_nodes(workflow_nodes: Sequence[Any]) -> None:
|
||||
workflow_return_node_ids: list[str] = []
|
||||
|
||||
for node in workflow_nodes:
|
||||
if not _is_invocation_node(node):
|
||||
continue
|
||||
|
||||
data = node["data"]
|
||||
node_id = data.get("id")
|
||||
node_type = data.get("type")
|
||||
if not isinstance(node_id, str) or not isinstance(node_type, str):
|
||||
continue
|
||||
|
||||
_raise_if_unsupported_invocation_type(node_type, node_id)
|
||||
|
||||
if node_type == "workflow_return":
|
||||
workflow_return_node_ids.append(node_id)
|
||||
|
||||
if len(workflow_return_node_ids) != 1:
|
||||
raise UnsupportedWorkflowNodeError(
|
||||
"call_saved_workflow requires the selected workflow to contain exactly one workflow_return node"
|
||||
)
|
||||
|
||||
|
||||
def _get_default_edges(workflow_edges: Sequence[Any]) -> list[Mapping[str, Any]]:
|
||||
return [edge for edge in workflow_edges if _is_mapping(edge) and edge.get("type") == "default"]
|
||||
|
||||
|
||||
def _get_connector_input_edge(
|
||||
connector_id: str, workflow_edges: Sequence[Mapping[str, Any]]
|
||||
) -> Mapping[str, Any] | None:
|
||||
return next(
|
||||
(
|
||||
edge
|
||||
for edge in workflow_edges
|
||||
if edge.get("target") == connector_id and edge.get("targetHandle") == CONNECTOR_INPUT_HANDLE
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_connector_source(
|
||||
connector_id: str, workflow_nodes: dict[str, Mapping[str, Any]], workflow_edges: Sequence[Mapping[str, Any]]
|
||||
) -> tuple[str, str] | None:
|
||||
visited: set[str] = set()
|
||||
|
||||
def resolve(node_id: str) -> tuple[str, str] | None:
|
||||
if node_id in visited:
|
||||
return None
|
||||
visited.add(node_id)
|
||||
|
||||
incoming_edge = _get_connector_input_edge(node_id, workflow_edges)
|
||||
if incoming_edge is None:
|
||||
return None
|
||||
|
||||
source_id = incoming_edge.get("source")
|
||||
source_handle = incoming_edge.get("sourceHandle")
|
||||
if not isinstance(source_id, str) or not isinstance(source_handle, str):
|
||||
return None
|
||||
|
||||
source_node = workflow_nodes.get(source_id)
|
||||
if source_node is None:
|
||||
return None
|
||||
|
||||
if _is_invocation_node(source_node):
|
||||
return (source_id, source_handle)
|
||||
|
||||
if _is_connector_node(source_node):
|
||||
return resolve(source_id)
|
||||
|
||||
return None
|
||||
|
||||
return resolve(connector_id)
|
||||
|
||||
|
||||
def build_graph_from_workflow(workflow: Mapping[str, Any]) -> Graph:
|
||||
workflow_nodes_raw = workflow.get("nodes", [])
|
||||
workflow_edges_raw = workflow.get("edges", [])
|
||||
_validate_callable_workflow_nodes(workflow_nodes_raw if isinstance(workflow_nodes_raw, Sequence) else [])
|
||||
|
||||
workflow_nodes = {
|
||||
node["id"]: node for node in workflow_nodes_raw if _is_mapping(node) and isinstance(node.get("id"), str)
|
||||
}
|
||||
default_edges = _get_default_edges(workflow_edges_raw if isinstance(workflow_edges_raw, Sequence) else [])
|
||||
|
||||
parsed_nodes: dict[str, dict[str, Any]] = {}
|
||||
for node in workflow_nodes.values():
|
||||
if not _is_invocation_node(node):
|
||||
continue
|
||||
|
||||
data = node["data"]
|
||||
node_id = data.get("id")
|
||||
node_type = data.get("type")
|
||||
if not isinstance(node_id, str) or not isinstance(node_type, str):
|
||||
continue
|
||||
|
||||
graph_node: dict[str, Any] = {
|
||||
"id": node_id,
|
||||
"type": node_type,
|
||||
"use_cache": data.get("useCache", False),
|
||||
"is_intermediate": data.get("isIntermediate", False),
|
||||
}
|
||||
|
||||
inputs = data.get("inputs", {})
|
||||
if _is_mapping(inputs):
|
||||
for field_name, field_value in inputs.items():
|
||||
if not isinstance(field_name, str) or not _is_mapping(field_value):
|
||||
continue
|
||||
# Saved workflows may include input metadata for unfilled optional fields without a "value".
|
||||
# Omit those fields so invocation defaults are applied instead of forcing None.
|
||||
if "value" in field_value:
|
||||
# The frontend board picker may persist sentinel strings; graph nodes model both
|
||||
# automatic and no-board selection as the board field default, None.
|
||||
if field_name == "board" and field_value["value"] in ("auto", "none"):
|
||||
continue
|
||||
graph_node[field_name] = field_value["value"]
|
||||
|
||||
parsed_nodes[node_id] = graph_node
|
||||
|
||||
parsed_edges: list[dict[str, dict[str, str]]] = []
|
||||
seen_edges: set[tuple[str, str, str, str]] = set()
|
||||
|
||||
for edge in default_edges:
|
||||
source_id = edge.get("source")
|
||||
target_id = edge.get("target")
|
||||
source_handle = edge.get("sourceHandle")
|
||||
target_handle = edge.get("targetHandle")
|
||||
if not all(isinstance(v, str) for v in (source_id, target_id, source_handle, target_handle)):
|
||||
continue
|
||||
|
||||
target_node = workflow_nodes.get(target_id)
|
||||
if not _is_invocation_node(target_node):
|
||||
continue
|
||||
|
||||
source_node = workflow_nodes.get(source_id)
|
||||
resolved_source: tuple[str, str] | None = None
|
||||
if _is_invocation_node(source_node):
|
||||
resolved_source = (source_id, source_handle)
|
||||
elif _is_connector_node(source_node):
|
||||
resolved_source = _resolve_connector_source(source_id, workflow_nodes, default_edges)
|
||||
|
||||
if resolved_source is None:
|
||||
continue
|
||||
|
||||
resolved_source_id, resolved_source_handle = resolved_source
|
||||
edge_key = (resolved_source_id, resolved_source_handle, target_id, target_handle)
|
||||
if edge_key in seen_edges:
|
||||
continue
|
||||
seen_edges.add(edge_key)
|
||||
|
||||
parsed_edges.append(
|
||||
{
|
||||
"source": {
|
||||
"node_id": resolved_source_id,
|
||||
"field": resolved_source_handle,
|
||||
},
|
||||
"destination": {
|
||||
"node_id": target_id,
|
||||
"field": target_handle,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
for edge in parsed_edges:
|
||||
destination_node_id = edge["destination"]["node_id"]
|
||||
destination_field = edge["destination"]["field"]
|
||||
parsed_nodes[destination_node_id].pop(destination_field, None)
|
||||
|
||||
return Graph.model_validate(
|
||||
{
|
||||
"nodes": parsed_nodes,
|
||||
"edges": [
|
||||
Edge(
|
||||
source=EdgeConnection(**edge["source"]),
|
||||
destination=EdgeConnection(**edge["destination"]),
|
||||
)
|
||||
for edge in parsed_edges
|
||||
],
|
||||
}
|
||||
)
|
||||
Reference in New Issue
Block a user