github/InvokeAI
1
1
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI.git synced 2026-04-23 03:00:31 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat(mm): add migration to flat model storage

Browse Source
This commit is contained in:
psychedelicious
2025-09-19 00:47:07 +10:00
parent d67d652dcf
commit ebd8e7dc04
2 changed files with 239 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
invokeai/app/services/shared/sqlite/sqlite_util.py
View File

@@ -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

237
invokeai/app/services/shared/sqlite_migrator/migrations/migration_22.py Normal file
View File

@@ -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/<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 _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),
)