diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py index 83a42917a7..43c3edd5f5 100644 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ b/invokeai/app/services/shared/sqlite/sqlite_util.py @@ -5,6 +5,7 @@ from invokeai.app.services.image_files.image_files_base import ImageFileStorageB from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase from invokeai.app.services.shared.sqlite_migrator.migrations.migration_1 import build_migration_1 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_2 import build_migration_2 +from invokeai.app.services.shared.sqlite_migrator.migrations.migration_3 import build_migration_3 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator @@ -27,6 +28,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto migrator = SqliteMigrator(db=db) migrator.register_migration(build_migration_1()) migrator.register_migration(build_migration_2(image_files=image_files, logger=logger)) + migrator.register_migration(build_migration_3()) migrator.run_migrations() return db diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py index 9b9dedcc58..99922e2fc1 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py @@ -11,6 +11,8 @@ from invokeai.app.services.workflow_records.workflow_records_common import ( UnsafeWorkflowWithVersionValidator, ) +from .util.migrate_yaml_config_1 import MigrateModelYamlToDb1 + class Migration2Callback: def __init__(self, image_files: ImageFileStorageBase, logger: Logger): @@ -24,6 +26,7 @@ class Migration2Callback: self._add_workflow_library(cursor) self._drop_model_manager_metadata(cursor) self._recreate_model_config(cursor) + self._migrate_model_config_records(cursor) self._migrate_embedded_workflows(cursor) def _add_images_has_workflow(self, cursor: sqlite3.Cursor) -> None: @@ -131,6 +134,11 @@ class Migration2Callback: """ ) + def _migrate_model_config_records(self, cursor: sqlite3.Cursor) -> None: + """After updating the model config table, we repopulate it.""" + model_record_migrator = MigrateModelYamlToDb1(cursor) + model_record_migrator.migrate() + 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 diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_3.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_3.py new file mode 100644 index 0000000000..2ffef13dd4 --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_3.py @@ -0,0 +1,75 @@ +import sqlite3 + +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration + +from .util.migrate_yaml_config_1 import MigrateModelYamlToDb1 + + +class Migration3Callback: + def __init__(self) -> None: + pass + + def __call__(self, cursor: sqlite3.Cursor) -> None: + self._drop_model_manager_metadata(cursor) + self._recreate_model_config(cursor) + self._migrate_model_config_records(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 _migrate_model_config_records(self, cursor: sqlite3.Cursor) -> None: + """After updating the model config table, we repopulate it.""" + model_record_migrator = MigrateModelYamlToDb1(cursor) + model_record_migrator.migrate() + + +def build_migration_3() -> 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( + from_version=2, + to_version=3, + callback=Migration3Callback(), + ) + + return migration_3 diff --git a/invokeai/backend/model_manager/migrate_to_db.py b/invokeai/app/services/shared/sqlite_migrator/migrations/util/migrate_yaml_config_1.py similarity index 56% rename from invokeai/backend/model_manager/migrate_to_db.py rename to invokeai/app/services/shared/sqlite_migrator/migrations/util/migrate_yaml_config_1.py index e68a2eab36..f2476ed0f6 100644 --- a/invokeai/backend/model_manager/migrate_to_db.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/util/migrate_yaml_config_1.py @@ -1,8 +1,12 @@ # Copyright (c) 2023 Lincoln D. Stein """Migrate from the InvokeAI v2 models.yaml format to the v3 sqlite format.""" +import json +import sqlite3 from hashlib import sha1 from logging import Logger +from pathlib import Path +from typing import Optional from omegaconf import DictConfig, OmegaConf from pydantic import TypeAdapter @@ -10,13 +14,12 @@ from pydantic import TypeAdapter from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.model_records import ( DuplicateModelException, - ModelRecordServiceSQL, UnknownModelException, ) -from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase from invokeai.backend.model_manager.config import ( AnyModelConfig, BaseModelType, + ModelConfigFactory, ModelType, ) from invokeai.backend.model_manager.hash import FastModelHash @@ -25,9 +28,9 @@ from invokeai.backend.util.logging import InvokeAILogger ModelsValidator = TypeAdapter(AnyModelConfig) -class MigrateModelYamlToDb: +class MigrateModelYamlToDb1: """ - Migrate the InvokeAI models.yaml format (VERSION 3.0.0) to SQL3 database format (VERSION 3.2.0) + Migrate the InvokeAI models.yaml format (VERSION 3.0.0) to SQL3 database format (VERSION 3.5.0). The class has one externally useful method, migrate(), which scans the currently models.yaml file and imports all its entries into invokeai.db. @@ -41,17 +44,13 @@ class MigrateModelYamlToDb: config: InvokeAIAppConfig logger: Logger + cursor: sqlite3.Cursor - def __init__(self) -> None: + def __init__(self, cursor: sqlite3.Cursor = None) -> None: self.config = InvokeAIAppConfig.get_config() self.config.parse_args() self.logger = InvokeAILogger.get_logger() - - def get_db(self) -> ModelRecordServiceSQL: - """Fetch the sqlite3 database for this installation.""" - db_path = None if self.config.use_memory_db else self.config.db_path - db = SqliteDatabase(db_path=db_path, logger=self.logger, verbose=self.config.log_sql) - return ModelRecordServiceSQL(db) + self.cursor = cursor def get_yaml(self) -> DictConfig: """Fetch the models.yaml DictConfig for this installation.""" @@ -62,8 +61,10 @@ class MigrateModelYamlToDb: def migrate(self) -> None: """Do the migration from models.yaml to invokeai.db.""" - db = self.get_db() - yaml = self.get_yaml() + try: + yaml = self.get_yaml() + except OSError: + return for model_key, stanza in yaml.items(): if model_key == "__metadata__": @@ -86,22 +87,62 @@ class MigrateModelYamlToDb: new_config: AnyModelConfig = ModelsValidator.validate_python(stanza) # type: ignore # see https://github.com/pydantic/pydantic/discussions/7094 try: - if original_record := db.search_by_path(stanza.path): - key = original_record[0].key + if original_record := self._search_by_path(stanza.path): + key = original_record.key self.logger.info(f"Updating model {model_name} with information from models.yaml using key {key}") - db.update_model(key, new_config) + self._update_model(key, new_config) else: self.logger.info(f"Adding model {model_name} with key {model_key}") - db.add_model(new_key, new_config) + self._add_model(new_key, new_config) except DuplicateModelException: self.logger.warning(f"Model {model_name} is already in the database") except UnknownModelException: self.logger.warning(f"Model at {stanza.path} could not be found in database") + def _search_by_path(self, path: Path) -> Optional[AnyModelConfig]: + self.cursor.execute( + """--sql + SELECT config FROM model_config + WHERE path=?; + """, + (str(path),), + ) + results = [ModelConfigFactory.make_config(json.loads(x[0])) for x in self.cursor.fetchall()] + return results[0] if results else None -def main(): - MigrateModelYamlToDb().migrate() + def _update_model(self, key: str, config: AnyModelConfig) -> None: + record = ModelConfigFactory.make_config(config, key=key) # ensure it is a valid config obect + json_serialized = record.model_dump_json() # and turn it into a json string. + self.cursor.execute( + """--sql + UPDATE model_config + SET + config=? + WHERE id=?; + """, + (json_serialized, key), + ) + if self.cursor.rowcount == 0: + raise UnknownModelException("model not found") - -if __name__ == "__main__": - main() + def _add_model(self, key: str, config: AnyModelConfig) -> None: + record = ModelConfigFactory.make_config(config, key=key) # ensure it is a valid config obect. + json_serialized = record.model_dump_json() # and turn it into a json string. + try: + self.cursor.execute( + """--sql + INSERT INTO model_config ( + id, + original_hash, + config + ) + VALUES (?,?,?); + """, + ( + key, + record.original_hash, + json_serialized, + ), + ) + except sqlite3.IntegrityError as exc: + raise DuplicateModelException(f"{record.name}: model is already in database") from exc diff --git a/pyproject.toml b/pyproject.toml index 0b8d258e7d..98018dc7cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,7 +138,6 @@ dependencies = [ "invokeai-node-web" = "invokeai.app.api_app:invoke_api" "invokeai-import-images" = "invokeai.frontend.install.import_images:main" "invokeai-db-maintenance" = "invokeai.backend.util.db_maintenance:main" -"invokeai-migrate-models-to-db" = "invokeai.backend.model_manager.migrate_to_db:main" [project.urls] "Homepage" = "https://invoke-ai.github.io/InvokeAI/"