feat(mm): reidentify models

Add route and model record service method to reidentify a model. This
re-probes the model files and replaces the model's config with the new
one if it does not error.

Author: psychedelicious

invokeai/app/api/routers/model_manager.py

@@ -28,7 +28,7 @@
     UnknownModelException,
 )
 from invokeai.app.util.suppress_output import SuppressOutput
-from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig, ModelConfigFactory
 from invokeai.backend.model_manager.configs.main import (
     Main_Checkpoint_SD1_Config,
     Main_Checkpoint_SD2_Config,
@@ -38,6 +38,7 @@
 from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
 from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch
 from invokeai.backend.model_manager.metadata.metadata_base import ModelMetadataWithFiles, UnknownMetadataException
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
 from invokeai.backend.model_manager.search import ModelSearch
 from invokeai.backend.model_manager.starter_models import (
     STARTER_BUNDLES,
@@ -191,6 +192,40 @@ async def get_model_record(
         raise HTTPException(status_code=404, detail=str(e))
 
 
+@model_manager_router.post(
+    "/i/{key}/reidentify",
+    operation_id="reidentify_model",
+    responses={
+        200: {
+            "description": "The model configuration was retrieved successfully",
+            "content": {"application/json": {"example": example_model_config}},
+        },
+        400: {"description": "Bad request"},
+        404: {"description": "The model could not be found"},
+    },
+)
+async def reidentify_model(
+    key: Annotated[str, Path(description="Key of the model to reidentify.")],
+) -> AnyModelConfig:
+    """Attempt to reidentify a model by re-probing its weights file."""
+    try:
+        config = ApiDependencies.invoker.services.model_manager.store.get_model(key)
+        models_path = ApiDependencies.invoker.services.configuration.models_path
+        if pathlib.Path(config.path).is_relative_to(models_path):
+            model_path = pathlib.Path(config.path)
+        else:
+            model_path = models_path / config.path
+        mod = ModelOnDisk(model_path)
+        result = ModelConfigFactory.from_model_on_disk(mod)
+        if result.config is None:
+            raise InvalidModelException("Unable to identify model format")
+        result.config.key = config.key  # retain the same key
+        new_config = ApiDependencies.invoker.services.model_manager.store.replace_model(config.key, result.config)
+        return new_config
+    except UnknownModelException as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+
 class FoundModel(BaseModel):
     path: str = Field(description="Path to the model")
     is_installed: bool = Field(description="Whether or not the model is already installed")

invokeai/app/services/model_records/model_records_base.py

@@ -138,6 +138,18 @@ def update_model(self, key: str, changes: ModelRecordChanges, allow_class_change
         """
         pass
 
+    @abstractmethod
+    def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig:
+        """
+        Replace the model record entirely, returning the new record.
+
+        This is used when we re-identify a model and have a new config object.
+
+        :param key: Unique key for the model to be updated.
+        :param new_config: The new model config to write.
+        """
+        pass
+
     @abstractmethod
     def get_model(self, key: str) -> AnyModelConfig:
         """

invokeai/app/services/model_records/model_records_sql.py

@@ -179,6 +179,23 @@ def update_model(self, key: str, changes: ModelRecordChanges, allow_class_change
 
         return self.get_model(key)
 
+    def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig:
+        if key != new_config.key:
+            raise ValueError("key does not match new_config.key")
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """--sql
+                UPDATE models
+                SET
+                    config=?
+                WHERE id=?;
+                """,
+                (new_config.model_dump_json(), key),
+            )
+            if cursor.rowcount == 0:
+                raise UnknownModelException("model not found")
+        return self.get_model(key)
+
     def get_model(self, key: str) -> AnyModelConfig:
         """
         Retrieve the ModelConfigBase instance for the indicated model.