diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py index 9a85d31eec..ec38e0f1b6 100644 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ b/invokeai/app/services/shared/sqlite/sqlite_util.py @@ -24,6 +24,7 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_18 import from invokeai.app.services.shared.sqlite_migrator.migrations.migration_19 import build_migration_19 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_20 import build_migration_20 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_21 import build_migration_21 +from invokeai.app.services.shared.sqlite_migrator.migrations.migration_22 import build_migration_22 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator @@ -65,6 +66,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto migrator.register_migration(build_migration_19(app_config=config)) migrator.register_migration(build_migration_20()) migrator.register_migration(build_migration_21()) + migrator.register_migration(build_migration_22(app_config=config, logger=logger)) migrator.run_migrations() return db diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_22.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_22.py new file mode 100644 index 0000000000..c79b58bf2a --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_22.py @@ -0,0 +1,237 @@ +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.config import AnyModelConfig, AnyModelConfigValidator + + +class NormalizeResult(NamedTuple): + new_relative_path: str | None + rollback_ops: list[tuple[Path, Path]] + + +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: + # Grab all model records + cursor.execute("SELECT id, config FROM models;") + rows = cursor.fetchall() + + for model_id, config_json in rows: + try: + # Get the model config as a pydantic object + config = self._load_model_config(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/ - 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 _load_model_config(self, config_json: str) -> AnyModelConfig: + # The typing of the validator says it returns Unknown, but it's really a AnyModelConfig. This utility function + # just makes that clear. + return AnyModelConfigValidator.validate_json(config_json) + + +def build_migration_22(app_config: InvokeAIAppConfig, logger: Logger) -> Migration: + """Builds the migration object for migrating from version 21 to version 22. + + 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. + """ + + return Migration( + from_version=21, + to_version=22, + callback=Migration22Callback(app_config=app_config, logger=logger), + )