replace load_and_cache_model() with load_remote_model() and load_local_odel()

Author: Lincoln Stein

docs/contributing/MODEL_MANAGER.md

@@ -1585,9 +1585,9 @@ Within invocations, the following methods are available from the
 
 ### context.download_and_cache_model(source) -> Path
 
-This method accepts a `source` of a model, downloads and caches it
-locally, and returns a Path to the local model. The source can be a
-local file or directory, a URL, or a HuggingFace repo_id.
+This method accepts a `source` of a remote model, downloads and caches
+it locally, and then returns a Path to the local model. The source can
+be a direct download URL or a HuggingFace repo_id.
 
 In the case of HuggingFace repo_id, the following variants are
 recognized:
@@ -1602,16 +1602,34 @@ directory using this syntax:
 
 * stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors
 
-### context.load_and_cache_model(source, [loader]) -> LoadedModel
+### context.load_local_model(model_path, [loader]) -> LoadedModel
 
-This method takes a model source, downloads it, caches it, and then
-loads it into the RAM cache for use in inference. The optional loader
-is a Callable that accepts a Path to the object, and returns a
-`Dict[str, torch.Tensor]`. If no loader is provided, then the method
-will use `torch.load()` for a .ckpt or .bin checkpoint file,
-`safetensors.torch.load_file()` for a safetensors checkpoint file, or
-`*.from_pretrained()` for a directory that looks like a
-diffusers directory.
+This method loads a local model from the indicated path, returning a
+`LoadedModel`. The optional loader is a Callable that accepts a Path
+to the object, and returns a `AnyModel` object. If no loader is
+provided, then the method will use `torch.load()` for a .ckpt or .bin
+checkpoint file, `safetensors.torch.load_file()` for a safetensors
+checkpoint file, or `cls.from_pretrained()` for a directory that looks
+like a diffusers directory.
+
+### context.load_remote_model(source, [loader]) -> LoadedModel
+
+This method accepts a `source` of a remote model, downloads and caches
+it locally, loads it, and returns a `LoadedModel`. The source can be a
+direct download URL or a HuggingFace repo_id.
+
+In the case of HuggingFace repo_id, the following variants are
+recognized:
+
+* stabilityai/stable-diffusion-v4           -- default model
+* stabilityai/stable-diffusion-v4:fp16      -- fp16 variant
+* stabilityai/stable-diffusion-v4:fp16:vae  -- the fp16 vae subfolder
+* stabilityai/stable-diffusion-v4:onnx:vae  -- the onnx variant vae subfolder
+
+You can also point at an arbitrary individual file within a repo_id
+directory using this syntax:
+
+* stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors
 
 
 

invokeai/app/invocations/controlnet_image_processors.py

@@ -611,7 +611,7 @@ def loader(model_path: Path):
                 model_path, model_size=self.model_size, device=TorchDevice.choose_torch_device()
             )
 
-        with self._context.models.load_and_cache_model(
+        with self._context.models.load_remote_model(
             source=DEPTH_ANYTHING_MODELS[self.model_size], loader=loader
         ) as model:
             depth_anything_detector = DepthAnythingDetector(model, TorchDevice.choose_torch_device())

invokeai/app/invocations/infill.py

@@ -134,7 +134,7 @@ class LaMaInfillInvocation(InfillImageProcessorInvocation):
     """Infills transparent areas of an image using the LaMa model"""
 
     def infill(self, image: Image.Image):
-        with self._context.models.load_and_cache_model(
+        with self._context.models.load_remote_model(
             source="https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
             loader=LaMA.load_jit_model,
         ) as model:

invokeai/app/invocations/upscale.py

@@ -91,7 +91,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             context.logger.error(msg)
             raise ValueError(msg)
 
-        loadnet = context.models.load_and_cache_model(
+        loadnet = context.models.load_remote_model(
             source=ESRGAN_MODEL_URLS[self.model_name],
         )
 

invokeai/app/services/model_install/model_install_base.py

@@ -5,6 +5,8 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
+from pydantic.networks import AnyHttpUrl
+
 from invokeai.app.services.config import InvokeAIAppConfig
 from invokeai.app.services.download import DownloadQueueServiceBase
 from invokeai.app.services.events.events_base import EventServiceBase
@@ -241,7 +243,7 @@ def sync_model_path(self, key: str) -> AnyModelConfig:
         """
 
     @abstractmethod
-    def download_and_cache_model(self, source: str) -> Path:
+    def download_and_cache_model(self, source: str | AnyHttpUrl) -> Path:
         """
         Download the model file located at source to the models cache and return its Path.
 

invokeai/app/services/model_install/model_install_default.py

@@ -15,6 +15,7 @@
 import yaml
 from huggingface_hub import HfFolder
 from pydantic.networks import AnyHttpUrl
+from pydantic_core import Url
 from requests import Session
 
 from invokeai.app.services.config import InvokeAIAppConfig
@@ -374,7 +375,7 @@ def _download_cache_path(cls, source: Union[str, AnyHttpUrl], app_config: Invoke
 
     def download_and_cache_model(
         self,
-        source: str,
+        source: str | AnyHttpUrl,
     ) -> Path:
         """Download the model file located at source to the models cache and return its Path."""
         model_path = self._download_cache_path(str(source), self._app_config)
@@ -388,7 +389,7 @@ def download_and_cache_model(
                 return contents[0]
 
         model_path.mkdir(parents=True, exist_ok=True)
-        model_source = self._guess_source(source)
+        model_source = self._guess_source(str(source))
         remote_files, _ = self._remote_files_from_source(model_source)
         job = self._multifile_download(
             dest=model_path,
@@ -447,7 +448,7 @@ def _guess_source(self, source: str) -> ModelSource:
             )
         elif re.match(r"^https?://[^/]+", source):
             source_obj = URLModelSource(
-                url=AnyHttpUrl(source),
+                url=Url(source),
             )
         else:
             raise ValueError(f"Unsupported model source: '{source}'")

invokeai/app/services/model_load/model_load_base.py

@@ -3,9 +3,7 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Callable, Dict, Optional
-
-from torch import Tensor
+from typing import Callable, Optional
 
 from invokeai.backend.model_manager import AnyModel, AnyModelConfig, SubModelType
 from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig
@@ -37,7 +35,7 @@ def convert_cache(self) -> ModelConvertCacheBase:
 
     @abstractmethod
     def load_model_from_path(
-        self, model_path: Path, loader: Optional[Callable[[Path], Dict[str, Tensor]]] = None
+        self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
     ) -> LoadedModelWithoutConfig:
         """
         Load the model file or directory located at the indicated Path.

invokeai/app/services/model_load/model_load_default.py

@@ -2,11 +2,10 @@
 """Implementation of model loader service."""
 
 from pathlib import Path
-from typing import Callable, Dict, Optional, Type
+from typing import Callable, Optional, Type
 
 from picklescan.scanner import scan_file_path
 from safetensors.torch import load_file as safetensors_load_file
-from torch import Tensor
 from torch import load as torch_load
 
 from invokeai.app.services.config import InvokeAIAppConfig
@@ -86,7 +85,7 @@ def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubMo
         return loaded_model
 
     def load_model_from_path(
-        self, model_path: Path, loader: Optional[Callable[[Path], Dict[str, Tensor] | AnyModel]] = None
+        self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
     ) -> LoadedModelWithoutConfig:
         cache_key = str(model_path)
         ram_cache = self.ram_cache
@@ -95,11 +94,11 @@ def load_model_from_path(
         except IndexError:
             pass
 
-        def torch_load_file(checkpoint: Path) -> Dict[str, Tensor]:
+        def torch_load_file(checkpoint: Path) -> AnyModel:
             scan_result = scan_file_path(checkpoint)
             if scan_result.infected_files != 0:
                 raise Exception("The model at {checkpoint} is potentially infected by malware. Aborting load.")
-            result: Dict[str, Tensor] = torch_load(checkpoint, map_location="cpu")
+            result = torch_load(checkpoint, map_location="cpu")
             return result
 
         def diffusers_load_directory(directory: Path) -> AnyModel:
@@ -109,18 +108,16 @@ def diffusers_load_directory(directory: Path) -> AnyModel:
                 ram_cache=self._ram_cache,
                 convert_cache=self.convert_cache,
             ).get_hf_load_class(directory)
-            result: AnyModel = load_class.from_pretrained(model_path, torch_dtype=TorchDevice.choose_torch_dtype())
-            return result
-
-        if loader is None:
-            loader = (
-                diffusers_load_directory
-                if model_path.is_dir()
-                else torch_load_file
-                if model_path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin"))
-                else lambda path: safetensors_load_file(path, device="cpu")
-            )
-
+            return load_class.from_pretrained(model_path, torch_dtype=TorchDevice.choose_torch_dtype())
+
+        loader = loader or (
+            diffusers_load_directory
+            if model_path.is_dir()
+            else torch_load_file
+            if model_path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin"))
+            else lambda path: safetensors_load_file(path, device="cpu")
+        )
+        assert loader is not None
         raw_model = loader(model_path)
         ram_cache.put(key=cache_key, model=raw_model)
         return LoadedModelWithoutConfig(_locker=ram_cache.get(key=cache_key))

invokeai/app/services/shared/invocation_context.py

@@ -15,7 +15,14 @@
 from invokeai.app.services.invocation_services import InvocationServices
 from invokeai.app.services.model_records.model_records_base import UnknownModelException
 from invokeai.app.util.step_callback import stable_diffusion_step_callback
-from invokeai.backend.model_manager.config import AnyModelConfig, BaseModelType, ModelFormat, ModelType, SubModelType
+from invokeai.backend.model_manager.config import (
+    AnyModel,
+    AnyModelConfig,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
 from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
@@ -449,21 +456,42 @@ def download_and_cache_model(
         installed, the cached path will be returned. Otherwise it will be downloaded.
 
         Args:
-            source: A model path, URL or repo_id.
+            source: A URL that points to the model, or a huggingface repo_id.
 
         Returns:
             Path to the downloaded model
         """
-
         return self._services.model_manager.install.download_and_cache_model(source=source)
 
-    def load_and_cache_model(
+    def load_local_model(
         self,
-        source: Path | str | AnyHttpUrl,
-        loader: Optional[Callable[[Path], dict[str, Tensor]]] = None,
+        model_path: Path,
+        loader: Optional[Callable[[Path], AnyModel]] = None,
     ) -> LoadedModelWithoutConfig:
         """
-        Download, cache, and load the model file located at the indicated URL.
+        Load the model file located at the indicated path
+
+        If a loader callable is provided, it will be invoked to load the model. Otherwise,
+        `safetensors.torch.load_file()` or `torch.load()` will be called to load the model.
+
+        Be aware that the LoadedModelWithoutConfig object has no `config` attribute
+
+        Args:
+            path: A model Path
+            loader: A Callable that expects a Path and returns a dict[str|int, Any]
+
+        Returns:
+            A LoadedModelWithoutConfig object.
+        """
+        return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
+
+    def load_remote_model(
+        self,
+        source: str | AnyHttpUrl,
+        loader: Optional[Callable[[Path], AnyModel]] = None,
+    ) -> LoadedModelWithoutConfig:
+        """
+        Download, cache, and load the model file located at the indicated URL or repo_id.
 
         If the model is already downloaded, it will be loaded from the cache.
 
@@ -473,18 +501,14 @@ def load_and_cache_model(
         Be aware that the LoadedModelWithoutConfig object has no `config` attribute
 
         Args:
-            source: A model Path, URL, or repoid.
+            source: A URL or huggingface repoid.
             loader: A Callable that expects a Path and returns a dict[str|int, Any]
 
         Returns:
             A LoadedModelWithoutConfig object.
         """
-
-        if isinstance(source, Path):
-            return self._services.model_manager.load.load_model_from_path(model_path=source, loader=loader)
-        else:
-            model_path = self._services.model_manager.install.download_and_cache_model(source=str(source))
-            return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
+        model_path = self._services.model_manager.install.download_and_cache_model(source=str(source))
+        return self._services.model_manager.load.load_model_from_path(model_path=model_path, loader=loader)
 
 
 class ConfigInterface(InvocationContextInterface):

invokeai/app/services/shared/sqlite_migrator/migrations/migration_11.py

@@ -59,14 +59,12 @@ def _remove_unused_core_models(self) -> None:
 
 def build_migration_11(app_config: InvokeAIAppConfig, logger: Logger) -> Migration:
     """
-    Build the migration from database version 9 to 10.
+    Build the migration from database version 10 to 11.
 
     This migration does the following:
     - Moves "core" models previously downloaded with download_with_progress_bar() into new
       "models/.download_cache" directory.
     - Renames "models/.cache" to "models/.convert_cache".
-    - Adds `error_type` and `error_message` columns to the session queue table.
-    - Renames the `error` column to `error_traceback`.
     """
     migration_11 = Migration(
         from_version=10,