feat(mm): normalized model storage

invokeai/app/api/routers/model_manager.py

@@ -297,10 +297,8 @@ async def update_model_record(
     """Update a model's config."""
     logger = ApiDependencies.invoker.services.logger
     record_store = ApiDependencies.invoker.services.model_manager.store
-    installer = ApiDependencies.invoker.services.model_manager.install
     try:
-        record_store.update_model(key, changes=changes)
-        config = installer.sync_model_path(key)
+        config = record_store.update_model(key, changes=changes)
         config = add_cover_image_to_model_config(config, ApiDependencies)
         logger.info(f"Updated model: {key}")
     except UnknownModelException as e:

invokeai/app/services/model_install/model_install_base.py

@@ -12,7 +12,6 @@
 from invokeai.app.services.invoker import Invoker
 from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
 from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceBase
-from invokeai.backend.model_manager import AnyModelConfig
 
 if TYPE_CHECKING:
     from invokeai.app.services.events.events_base import EventServiceBase
@@ -231,19 +230,6 @@ def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]:
         will block indefinitely until the installs complete.
         """
 
-    @abstractmethod
-    def sync_model_path(self, key: str) -> AnyModelConfig:
-        """
-        Move model into the location indicated by its basetype, type and name.
-
-        Call this after updating a model's attributes in order to move
-        the model's path into the location indicated by its basetype, type and
-        name. Applies only to models whose paths are within the root `models_dir`
-        directory.
-
-        May raise an UnknownModelException.
-        """
-
     @abstractmethod
     def download_and_cache_model(self, source: str | AnyHttpUrl) -> Path:
         """

invokeai/app/services/model_install/model_install_default.py

@@ -180,28 +180,32 @@ def install_path(
         self,
         model_path: Union[Path, str],
         config: Optional[ModelRecordChanges] = None,
-    ) -> str:  # noqa D102
+    ) -> str:
         model_path = Path(model_path)
         config = config or ModelRecordChanges()
         info: AnyModelConfig = self._probe(Path(model_path), config)  # type: ignore
 
-        if preferred_name := config.name:
-            if Path(model_path).is_file():
-                # Careful! Don't use pathlib.Path(...).with_suffix - it can will strip everything after the first dot.
-                preferred_name = f"{preferred_name}{model_path.suffix}"
-
-        dest_path = (
-            self.app_config.models_path / info.base.value / info.type.value / (preferred_name or model_path.name)
-        )
+        dest_dir = self.app_config.models_path / info.key
         try:
-            new_path = self._move_model(model_path, dest_path)
-        except FileExistsError as excp:
+            if dest_dir.exists():
+                raise FileExistsError(
+                    f"Cannot install model {model_path.name} to {dest_dir}: destination already exists"
+                )
+            dest_dir.mkdir(parents=True)
+            dest_path = dest_dir / model_path.name if model_path.is_file() else dest_dir
+            if model_path.is_file():
+                move(model_path, dest_path)
+            elif model_path.is_dir():
+                # Move the contents of the directory, not the directory itself
+                for item in model_path.iterdir():
+                    move(item, dest_dir / item.name)
+        except FileExistsError as e:
             raise DuplicateModelException(
-                f"A model named {model_path.name} is already installed at {dest_path.as_posix()}"
-            ) from excp
+                f"A model named {model_path.name} is already installed at {dest_dir.as_posix()}"
+            ) from e
 
         return self._register(
-            new_path,
+            dest_path,
             config,
             info,
         )
@@ -589,49 +593,6 @@ def on_model_found(model_path: Path) -> bool:
         found_models = search.search(self._app_config.models_path)
         self._logger.info(f"{len(found_models)} new models registered")
 
-    def sync_model_path(self, key: str) -> AnyModelConfig:
-        """
-        Move model into the location indicated by its basetype, type and name.
-
-        Call this after updating a model's attributes in order to move
-        the model's path into the location indicated by its basetype, type and
-        name. Applies only to models whose paths are within the root `models_dir`
-        directory.
-
-        May raise an UnknownModelException.
-        """
-        model = self.record_store.get_model(key)
-        models_dir = self.app_config.models_path
-        old_path = self.app_config.models_path / model.path
-
-        if not old_path.is_relative_to(models_dir):
-            # The model is not in the models directory - we don't need to move it.
-            return model
-
-        new_path = models_dir / model.base.value / model.type.value / old_path.name
-
-        if old_path == new_path or new_path.exists() and old_path == new_path.resolve():
-            return model
-
-        self._logger.info(f"Moving {model.name} to {new_path}.")
-        new_path = self._move_model(old_path, new_path)
-        model.path = new_path.relative_to(models_dir).as_posix()
-        self.record_store.update_model(key, ModelRecordChanges(path=model.path))
-        return model
-
-    def _move_model(self, old_path: Path, new_path: Path) -> Path:
-        if old_path == new_path:
-            return old_path
-
-        if new_path.exists():
-            raise FileExistsError(f"Cannot move {old_path} to {new_path}: destination already exists")
-
-        new_path.parent.mkdir(parents=True, exist_ok=True)
-
-        move(old_path, new_path)
-
-        return new_path
-
     def _probe(self, model_path: Path, config: Optional[ModelRecordChanges] = None):
         config = config or ModelRecordChanges()
         hash_algo = self._app_config.hashing_algorithm

invokeai/app/services/shared/sqlite/sqlite_util.py

@@ -24,6 +24,7 @@
 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

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/<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),
+    )