chore: import upstream snapshot with attribution

invokeai/backend/model_manager/load/model_cache/cache_record.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+
+from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_only_full_load import (
+    CachedModelOnlyFullLoad,
+)
+from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import (
+    CachedModelWithPartialLoad,
+)
+
+
+@dataclass
+class CacheRecord:
+    """A class that represents a model in the model cache."""
+
+    # Cache key.
+    key: str
+    # Model in memory.
+    cached_model: CachedModelWithPartialLoad | CachedModelOnlyFullLoad
+    _locks: int = 0
+    # Set by ModelCache.drop_model() when the entry was locked at invalidation time.
+    # ModelCache.unlock() evicts the entry as soon as the last lock releases so a setting
+    # change (e.g. fp8_storage toggled during an in-flight generation) takes effect on the
+    # next load instead of silently being ignored.
+    is_stale: bool = False
+
+    def lock(self) -> None:
+        """Lock this record."""
+        self._locks += 1
+
+    def unlock(self) -> None:
+        """Unlock this record."""
+        self._locks -= 1
+        assert self._locks >= 0
+
+    @property
+    def is_locked(self) -> bool:
+        """Return true if record is locked."""
+        return self._locks > 0

invokeai/backend/model_manager/load/model_cache/cache_stats.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass, field
+from typing import Dict
+
+
+@dataclass
+class CacheStats(object):
+    """Collect statistics on cache performance."""
+
+    hits: int = 0  # cache hits
+    misses: int = 0  # cache misses
+    high_watermark: int = 0  # amount of cache used
+    in_cache: int = 0  # number of models in cache
+    cleared: int = 0  # number of models cleared to make space
+    cache_size: int = 0  # total size of cache
+    loaded_model_sizes: Dict[str, int] = field(default_factory=dict)

invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py

@@ -0,0 +1,121 @@
+from typing import Any
+
+import torch
+
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+class CachedModelOnlyFullLoad:
+    """A wrapper around a PyTorch model to handle full loads and unloads between the CPU and the compute device.
+    Note: "VRAM" is used throughout this class to refer to the memory on the compute device. It could be CUDA memory,
+    MPS memory, etc.
+    """
+
+    def __init__(
+        self, model: torch.nn.Module | Any, compute_device: torch.device, total_bytes: int, keep_ram_copy: bool = False
+    ):
+        """Initialize a CachedModelOnlyFullLoad.
+        Args:
+            model (torch.nn.Module | Any): The model to wrap. Should be on the CPU.
+            compute_device (torch.device): The compute device to move the model to.
+            total_bytes (int): The total size (in bytes) of all the weights in the model.
+            keep_ram_copy (bool): Whether to keep a read-only copy of the model's state dict in RAM. Keeping a RAM copy
+                increases RAM usage, but speeds up model offload from VRAM and LoRA patching (assuming there is
+                sufficient RAM).
+        """
+        # model is often a torch.nn.Module, but could be any model type. Throughout this class, we handle both cases.
+        self._model = model
+        self._compute_device = compute_device
+        self._offload_device = torch.device("cpu")
+
+        # A CPU read-only copy of the model's state dict.
+        self._cpu_state_dict: dict[str, torch.Tensor] | None = None
+        if isinstance(model, torch.nn.Module) and keep_ram_copy:
+            self._cpu_state_dict = model.state_dict()
+
+        self._total_bytes = total_bytes
+        self._is_in_vram = False
+
+    @property
+    def model(self) -> torch.nn.Module:
+        return self._model
+
+    def get_cpu_state_dict(self) -> dict[str, torch.Tensor] | None:
+        """Get a read-only copy of the model's state dict in RAM."""
+        # TODO(ryand): Document this better.
+        return self._cpu_state_dict
+
+    def total_bytes(self) -> int:
+        """Get the total size (in bytes) of all the weights in the model."""
+        return self._total_bytes
+
+    def cur_vram_bytes(self) -> int:
+        """Get the size (in bytes) of the weights that are currently in VRAM."""
+        if self._is_in_vram:
+            return self._total_bytes
+        else:
+            return 0
+
+    def is_in_vram(self) -> bool:
+        """Return true if the model is currently in VRAM."""
+        return self._is_in_vram
+
+    @property
+    def compute_device(self) -> torch.device:
+        """Return the compute device for this model."""
+        return self._compute_device
+
+    def full_load_to_vram(self) -> int:
+        """Load all weights into VRAM (if supported by the model).
+        Returns:
+            The number of bytes loaded into VRAM.
+        """
+        if self._is_in_vram:
+            # Already in VRAM.
+            return 0
+
+        if not hasattr(self._model, "to"):
+            # Model doesn't support moving to a device.
+            return 0
+
+        if self._cpu_state_dict is not None:
+            new_state_dict: dict[str, torch.Tensor] = {}
+            for k, v in self._cpu_state_dict.items():
+                new_state_dict[k] = v.to(self._compute_device, copy=True)
+            self._model.load_state_dict(new_state_dict, assign=True)
+
+        check_for_gguf = hasattr(self._model, "state_dict") and self._model.state_dict().get("img_in.weight")
+        if isinstance(check_for_gguf, GGMLTensor):
+            old_value = torch.__future__.get_overwrite_module_params_on_conversion()
+            torch.__future__.set_overwrite_module_params_on_conversion(True)
+            self._model.to(self._compute_device)
+            torch.__future__.set_overwrite_module_params_on_conversion(old_value)
+        else:
+            self._model.to(self._compute_device)
+
+        self._is_in_vram = True
+        return self._total_bytes
+
+    def full_unload_from_vram(self) -> int:
+        """Unload all weights from VRAM.
+        Returns:
+            The number of bytes unloaded from VRAM.
+        """
+        if not self._is_in_vram:
+            # Already in RAM.
+            return 0
+
+        if self._cpu_state_dict is not None:
+            self._model.load_state_dict(self._cpu_state_dict, assign=True)
+
+        check_for_gguf = hasattr(self._model, "state_dict") and self._model.state_dict().get("img_in.weight")
+        if isinstance(check_for_gguf, GGMLTensor):
+            old_value = torch.__future__.get_overwrite_module_params_on_conversion()
+            torch.__future__.set_overwrite_module_params_on_conversion(True)
+            self._model.to(self._offload_device)
+            torch.__future__.set_overwrite_module_params_on_conversion(old_value)
+        else:
+            self._model.to(self._offload_device)
+
+        self._is_in_vram = False
+        return self._total_bytes

invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py

@@ -0,0 +1,365 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+from invokeai.backend.util.logging import InvokeAILogger
+
+
+class CachedModelWithPartialLoad:
+    """A wrapper around a PyTorch model to handle partial loads and unloads between the CPU and the compute device.
+
+    Note: "VRAM" is used throughout this class to refer to the memory on the compute device. It could be CUDA memory,
+    MPS memory, etc.
+    """
+
+    def __init__(self, model: torch.nn.Module, compute_device: torch.device, keep_ram_copy: bool = False):
+        self._model = model
+        self._compute_device = compute_device
+
+        model_state_dict = model.state_dict()
+        # A CPU read-only copy of the model's state dict. Used for faster model unloads from VRAM, and to speed up LoRA
+        # patching. Set to `None` if keep_ram_copy is False.
+        self._cpu_state_dict: dict[str, torch.Tensor] | None = model_state_dict if keep_ram_copy else None
+
+        # A dictionary of the size of each tensor in the state dict.
+        # HACK(ryand): We use this dictionary any time we are doing byte tracking calculations. We do this for
+        # consistency in case the application code has modified the model's size (e.g. by casting to a different
+        # precision). Of course, this means that we are making model cache load/unload decisions based on model size
+        # data that may not be fully accurate.
+        self._state_dict_bytes = {k: calc_tensor_size(v) for k, v in model_state_dict.items()}
+
+        self._total_bytes = sum(self._state_dict_bytes.values())
+        self._cur_vram_bytes: int | None = None
+
+        self._modules_that_support_autocast = self._find_modules_that_support_autocast()
+        self._keys_in_modules_that_do_not_support_autocast = self._find_keys_in_modules_that_do_not_support_autocast(
+            model_state_dict
+        )
+        self._state_dict_keys_by_module_prefix = self._group_state_dict_keys_by_module_prefix(model_state_dict)
+
+    def _find_modules_that_support_autocast(self) -> dict[str, torch.nn.Module]:
+        """Find all modules that support autocasting."""
+        return {n: m for n, m in self._model.named_modules() if isinstance(m, CustomModuleMixin)}  # type: ignore
+
+    def _find_keys_in_modules_that_do_not_support_autocast(self, state_dict: dict[str, torch.Tensor]) -> set[str]:
+        keys_in_modules_that_do_not_support_autocast: set[str] = set()
+        for key in state_dict.keys():
+            for module_name in self._modules_that_support_autocast.keys():
+                if key.startswith(module_name):
+                    break
+            else:
+                keys_in_modules_that_do_not_support_autocast.add(key)
+        return keys_in_modules_that_do_not_support_autocast
+
+    def _group_state_dict_keys_by_module_prefix(self, state_dict: dict[str, torch.Tensor]) -> dict[str, list[str]]:
+        """A helper function that groups state dict keys by module prefix.
+
+        Example:
+        ```
+        state_dict = {
+            "weight": ...,
+            "module.submodule.weight": ...,
+            "module.submodule.bias": ...,
+            "module.other_submodule.weight": ...,
+            "module.other_submodule.bias": ...,
+        }
+
+        output = group_state_dict_keys_by_module_prefix(state_dict)
+
+        # The output will be:
+        output = {
+            "": [
+                "weight",
+            ],
+            "module.submodule": [
+                "module.submodule.weight",
+                "module.submodule.bias",
+            ],
+            "module.other_submodule": [
+                "module.other_submodule.weight",
+                "module.other_submodule.bias",
+            ],
+        }
+        ```
+        """
+        state_dict_keys_by_module_prefix: dict[str, list[str]] = {}
+        for key in state_dict.keys():
+            split = key.rsplit(".", 1)
+            # `split` will have length 1 if the root module has parameters.
+            module_name = split[0] if len(split) > 1 else ""
+            if module_name not in state_dict_keys_by_module_prefix:
+                state_dict_keys_by_module_prefix[module_name] = []
+            state_dict_keys_by_module_prefix[module_name].append(key)
+        return state_dict_keys_by_module_prefix
+
+    def _move_non_persistent_buffers_to_device(self, device: torch.device):
+        """Move the non-persistent buffers to the target device. These buffers are not included in the state dict,
+        so we need to move them manually.
+        """
+        # HACK(ryand): Typically, non-persistent buffers are moved when calling module.to(device). We don't move entire
+        # modules, because we manage the devices of individual tensors using the state dict. Since non-persistent
+        # buffers are not included in the state dict, we need to handle them manually. The only way to do this is by
+        # using private torch.nn.Module attributes.
+        for module in self._model.modules():
+            for name, buffer in module.named_buffers():
+                if name in module._non_persistent_buffers_set:
+                    module._buffers[name] = buffer.to(device, copy=True)
+
+    def _set_autocast_enabled_in_all_modules(self, enabled: bool):
+        """Set autocast_enabled flag in all modules that support device autocasting."""
+        for module in self._modules_that_support_autocast.values():
+            module.set_device_autocasting_enabled(enabled)
+
+    @property
+    def model(self) -> torch.nn.Module:
+        return self._model
+
+    def get_cpu_state_dict(self) -> dict[str, torch.Tensor] | None:
+        """Get a read-only copy of the model's state dict in RAM."""
+        # TODO(ryand): Document this better.
+        return self._cpu_state_dict
+
+    def total_bytes(self) -> int:
+        """Get the total size (in bytes) of all the weights in the model."""
+        return self._total_bytes
+
+    def cur_vram_bytes(self) -> int:
+        """Get the size (in bytes) of the weights that are currently in VRAM."""
+        if self._cur_vram_bytes is None:
+            cur_state_dict = self._model.state_dict()
+            self._cur_vram_bytes = sum(
+                self._state_dict_bytes[k]
+                for k, v in cur_state_dict.items()
+                if v.device.type == self._compute_device.type
+            )
+        return self._cur_vram_bytes
+
+    @property
+    def compute_device(self) -> torch.device:
+        """Return the compute device for this model."""
+        return self._compute_device
+
+    def full_load_to_vram(self) -> int:
+        """Load all weights into VRAM."""
+        return self.partial_load_to_vram(self.total_bytes())
+
+    def full_unload_from_vram(self) -> int:
+        """Unload all weights from VRAM."""
+        return self.partial_unload_from_vram(self.total_bytes())
+
+    @torch.no_grad()
+    def repair_required_tensors_on_compute_device(self) -> int:
+        """Repair required non-autocast tensors that were left off the compute device.
+
+        This can happen if an interrupted run leaves the model in a partially inconsistent state. Any repaired device
+        movement invalidates the cached VRAM accounting.
+        """
+        cur_state_dict = self._model.state_dict()
+        keys_to_repair = {
+            key
+            for key in self._keys_in_modules_that_do_not_support_autocast
+            if cur_state_dict[key].device.type != self._compute_device.type
+        }
+        if len(keys_to_repair) == 0:
+            return 0
+
+        self._load_state_dict_with_device_conversion(cur_state_dict, keys_to_repair, self._compute_device)
+        self._move_non_persistent_buffers_to_device(self._compute_device)
+        self._cur_vram_bytes = None
+        return len(keys_to_repair)
+
+    def _load_state_dict_with_device_conversion(
+        self, state_dict: dict[str, torch.Tensor], keys_to_convert: set[str], target_device: torch.device
+    ):
+        if self._cpu_state_dict is not None:
+            # Run the fast version.
+            self._load_state_dict_with_fast_device_conversion(
+                state_dict=state_dict,
+                keys_to_convert=keys_to_convert,
+                target_device=target_device,
+                cpu_state_dict=self._cpu_state_dict,
+            )
+        else:
+            # Run the low-virtual-memory version.
+            self._load_state_dict_with_jit_device_conversion(
+                state_dict=state_dict,
+                keys_to_convert=keys_to_convert,
+                target_device=target_device,
+            )
+
+    def _load_state_dict_with_jit_device_conversion(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        keys_to_convert: set[str],
+        target_device: torch.device,
+    ):
+        """A custom state dict loading implementation with good peak memory properties.
+
+        This implementation has the important property that it copies parameters to the target device one module at a time
+        rather than applying all of the device conversions and then calling load_state_dict(). This is done to minimize the
+        peak virtual memory usage. Specifically, we want to avoid a case where we hold references to all of the CPU weights
+        and CUDA weights simultaneously, because Windows will reserve virtual memory for both.
+        """
+        for module_name, module in self._model.named_modules():
+            module_keys = self._state_dict_keys_by_module_prefix.get(module_name, [])
+            # Calculate the length of the module name prefix.
+            prefix_len = len(module_name)
+            if prefix_len > 0:
+                prefix_len += 1
+
+            module_state_dict = {}
+            for key in module_keys:
+                if key in keys_to_convert:
+                    # It is important that we overwrite `state_dict[key]` to avoid keeping two copies of the same
+                    # parameter.
+                    state_dict[key] = state_dict[key].to(target_device)
+                # Note that we keep parameters that have not been moved to a new device in case the module implements
+                # weird custom state dict loading logic that requires all parameters to be present.
+                module_state_dict[key[prefix_len:]] = state_dict[key]
+
+            if len(module_state_dict) > 0:
+                # We set strict=False, because if `module` has both parameters and child modules, then we are loading a
+                # state dict that only contains the parameters of `module` (not its children).
+                # We assume that it is rare for non-leaf modules to have parameters. Calling load_state_dict() on non-leaf
+                # modules will recurse through all of the children, so is a bit wasteful.
+                incompatible_keys = module.load_state_dict(module_state_dict, strict=False, assign=True)
+                # Missing keys are ok, unexpected keys are not.
+                assert len(incompatible_keys.unexpected_keys) == 0
+
+    def _load_state_dict_with_fast_device_conversion(
+        self,
+        state_dict: dict[str, torch.Tensor],
+        keys_to_convert: set[str],
+        target_device: torch.device,
+        cpu_state_dict: dict[str, torch.Tensor],
+    ):
+        """Convert parameters to the target device and load them into the model. Leverages the `cpu_state_dict` to speed
+        up transfers of weights to the CPU.
+        """
+        for key in keys_to_convert:
+            if target_device.type == "cpu":
+                state_dict[key] = cpu_state_dict[key]
+            else:
+                state_dict[key] = state_dict[key].to(target_device)
+
+        self._model.load_state_dict(state_dict, assign=True)
+
+    @torch.no_grad()
+    def partial_load_to_vram(self, vram_bytes_to_load: int) -> int:
+        """Load more weights into VRAM without exceeding vram_bytes_to_load.
+
+        Returns:
+            The number of bytes loaded into VRAM.
+        """
+        # TODO(ryand): Handle the case where an exception is thrown while loading or unloading weights. At the very
+        # least, we should reset self._cur_vram_bytes to None.
+
+        vram_bytes_loaded = 0
+
+        cur_state_dict = self._model.state_dict()
+
+        # Identify the keys that will be loaded into VRAM.
+        keys_to_load: set[str] = set()
+
+        # First, process the keys that *must* be loaded into VRAM.
+        for key in self._keys_in_modules_that_do_not_support_autocast:
+            param = cur_state_dict[key]
+            if param.device.type == self._compute_device.type:
+                continue
+
+            keys_to_load.add(key)
+            param_size = self._state_dict_bytes[key]
+            vram_bytes_loaded += param_size
+
+        if vram_bytes_loaded > vram_bytes_to_load:
+            logger = InvokeAILogger.get_logger()
+            logger.warning(
+                f"Loading {vram_bytes_loaded / 2**20} MB into VRAM, but only {vram_bytes_to_load / 2**20} MB were "
+                "requested. This is the minimum set of weights in VRAM required to run the model."
+            )
+
+        # Next, process the keys that can optionally be loaded into VRAM.
+        fully_loaded = True
+        for key, param in cur_state_dict.items():
+            # Skip the keys that have already been processed above.
+            if key in keys_to_load:
+                continue
+
+            if param.device.type == self._compute_device.type:
+                continue
+
+            param_size = self._state_dict_bytes[key]
+            if vram_bytes_loaded + param_size > vram_bytes_to_load:
+                # TODO(ryand): Should we just break here? If we couldn't fit this parameter into VRAM, is it really
+                # worth continuing to search for a smaller parameter that would fit?
+                fully_loaded = False
+                continue
+
+            keys_to_load.add(key)
+            vram_bytes_loaded += param_size
+
+        if len(keys_to_load) > 0:
+            # We load the entire state dict, not just the parameters that changed, in case there are modules that
+            # override _load_from_state_dict() and do some funky stuff that requires the entire state dict.
+            # Alternatively, in the future, grouping parameters by module could probably solve this problem.
+            self._load_state_dict_with_device_conversion(cur_state_dict, keys_to_load, self._compute_device)
+
+        if self._cur_vram_bytes is not None:
+            self._cur_vram_bytes += vram_bytes_loaded
+
+        if fully_loaded:
+            self._set_autocast_enabled_in_all_modules(False)
+        else:
+            self._set_autocast_enabled_in_all_modules(True)
+
+        # Move all non-persistent buffers to the compute device. These are a weird edge case and do not participate in
+        # the vram_bytes_loaded tracking.
+        self._move_non_persistent_buffers_to_device(self._compute_device)
+
+        return vram_bytes_loaded
+
+    @torch.no_grad()
+    def partial_unload_from_vram(self, vram_bytes_to_free: int, keep_required_weights_in_vram: bool = False) -> int:
+        """Unload weights from VRAM until vram_bytes_to_free bytes are freed. Or the entire model is unloaded.
+
+        :param keep_required_weights_in_vram: If True, any weights that must be kept in VRAM to run the model will be
+            kept in VRAM.
+
+        Returns:
+            The number of bytes unloaded from VRAM.
+        """
+        vram_bytes_freed = 0
+        required_weights_in_vram = 0
+
+        offload_device = "cpu"
+        cur_state_dict = self._model.state_dict()
+
+        # Identify the keys that will be offloaded to CPU.
+        keys_to_offload: set[str] = set()
+
+        for key, param in cur_state_dict.items():
+            if vram_bytes_freed >= vram_bytes_to_free:
+                break
+
+            if param.device.type == offload_device:
+                continue
+
+            if keep_required_weights_in_vram and key in self._keys_in_modules_that_do_not_support_autocast:
+                required_weights_in_vram += self._state_dict_bytes[key]
+                continue
+
+            keys_to_offload.add(key)
+            vram_bytes_freed += self._state_dict_bytes[key]
+
+        if len(keys_to_offload) > 0:
+            self._load_state_dict_with_device_conversion(cur_state_dict, keys_to_offload, torch.device("cpu"))
+
+        if self._cur_vram_bytes is not None:
+            self._cur_vram_bytes -= vram_bytes_freed
+
+        # We may have gone from a fully-loaded model to a partially-loaded model, so we need to reapply the custom
+        # layers.
+        self._set_autocast_enabled_in_all_modules(True)
+        return vram_bytes_freed

invokeai/backend/model_manager/load/model_cache/dev_utils.py

@@ -0,0 +1,33 @@
+from contextlib import contextmanager
+
+import torch
+
+from invokeai.backend.util.logging import InvokeAILogger
+
+
+@contextmanager
+def log_operation_vram_usage(operation_name: str):
+    """A helper function for tuning working memory requirements for memory-intensive ops.
+
+    Sample usage:
+
+    ```python
+    with log_operation_vram_usage("some_operation"):
+        some_operation()
+    ```
+    """
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+    max_allocated_before = torch.cuda.max_memory_allocated()
+    max_reserved_before = torch.cuda.max_memory_reserved()
+    try:
+        yield
+    finally:
+        torch.cuda.synchronize()
+        max_allocated_after = torch.cuda.max_memory_allocated()
+        max_reserved_after = torch.cuda.max_memory_reserved()
+        logger = InvokeAILogger.get_logger()
+        logger.info(
+            f">>>{operation_name} Peak VRAM allocated: {(max_allocated_after - max_allocated_before) / 2**20} MB, "
+            f"Peak VRAM reserved: {(max_reserved_after - max_reserved_before) / 2**20} MB"
+        )

invokeai/backend/model_manager/load/model_cache/model_cache.py

@@ -0,0 +1,931 @@
+import gc
+import logging
+import threading
+import time
+from dataclasses import dataclass
+from functools import wraps
+from logging import Logger
+from typing import Any, Callable, Dict, List, Optional, Protocol
+
+import psutil
+import torch
+
+from invokeai.backend.model_manager.load.memory_snapshot import MemorySnapshot
+from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
+from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
+from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_only_full_load import (
+    CachedModelOnlyFullLoad,
+)
+from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import (
+    CachedModelWithPartialLoad,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
+    apply_custom_layers_to_model,
+)
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
+from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter
+
+# Size of a GB in bytes.
+GB = 2**30
+
+# Size of a MB in bytes.
+MB = 2**20
+
+
+# TODO(ryand): Where should this go? The ModelCache shouldn't be concerned with submodels.
+def get_model_cache_key(model_key: str, submodel_type: Optional[SubModelType] = None) -> str:
+    """Get the cache key for a model based on the optional submodel type."""
+    if submodel_type:
+        return f"{model_key}:{submodel_type.value}"
+    else:
+        return model_key
+
+
+def synchronized(method: Callable[..., Any]) -> Callable[..., Any]:
+    """A decorator that applies the class's self._lock to the method."""
+
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        with self._lock:  # Automatically acquire and release the lock
+            return method(self, *args, **kwargs)
+
+    return wrapper
+
+
+def record_activity(method: Callable[..., Any]) -> Callable[..., Any]:
+    """A decorator that records activity after a method completes successfully.
+
+    Note: This decorator should be applied to methods that already hold self._lock.
+    """
+
+    @wraps(method)
+    def wrapper(self, *args, **kwargs):
+        result = method(self, *args, **kwargs)
+        self._record_activity()
+        return result
+
+    return wrapper
+
+
+@dataclass
+class CacheEntrySnapshot:
+    cache_key: str
+    total_bytes: int
+    current_vram_bytes: int
+
+
+class CacheMissCallback(Protocol):
+    def __call__(
+        self,
+        model_key: str,
+        cache_snapshot: dict[str, CacheEntrySnapshot],
+    ) -> None: ...
+
+
+class CacheHitCallback(Protocol):
+    def __call__(
+        self,
+        model_key: str,
+        cache_snapshot: dict[str, CacheEntrySnapshot],
+    ) -> None: ...
+
+
+class CacheModelsClearedCallback(Protocol):
+    def __call__(
+        self,
+        models_cleared: int,
+        bytes_requested: int,
+        bytes_freed: int,
+        cache_snapshot: dict[str, CacheEntrySnapshot],
+    ) -> None: ...
+
+
+class ModelCache:
+    """A cache for managing models in memory.
+
+    The cache is based on two levels of model storage:
+    - execution_device: The device where most models are executed (typically "cuda", "mps", or "cpu").
+    - storage_device: The device where models are offloaded when not in active use (typically "cpu").
+
+    The model cache is based on the following assumptions:
+    - storage_device_mem_size > execution_device_mem_size
+    - disk_to_storage_device_transfer_time >> storage_device_to_execution_device_transfer_time
+
+    A copy of all models in the cache is always kept on the storage_device. A subset of the models also have a copy on
+    the execution_device.
+
+    Models are moved between the storage_device and the execution_device as necessary. Cache size limits are enforced
+    on both the storage_device and the execution_device. The execution_device cache uses a smallest-first offload
+    policy. The storage_device cache uses a least-recently-used (LRU) offload policy.
+
+    Note: Neither of these offload policies has really been compared against alternatives. It's likely that different
+    policies would be better, although the optimal policies are likely heavily dependent on usage patterns and HW
+    configuration.
+
+    The cache returns context manager generators designed to load the model into the execution device (often GPU) within
+    the context, and unload outside the context.
+
+    Example usage:
+    ```
+    cache = ModelCache(max_cache_size=7.5, max_vram_cache_size=6.0)
+    with cache.get_model('runwayml/stable-diffusion-1-5') as SD1:
+        do_something_on_gpu(SD1)
+    ```
+    """
+
+    def __init__(
+        self,
+        execution_device_working_mem_gb: float,
+        enable_partial_loading: bool,
+        keep_ram_copy_of_weights: bool,
+        max_ram_cache_size_gb: float | None = None,
+        max_vram_cache_size_gb: float | None = None,
+        execution_device: torch.device | str = "cuda",
+        storage_device: torch.device | str = "cpu",
+        log_memory_usage: bool = False,
+        logger: Optional[Logger] = None,
+        keep_alive_minutes: float = 0,
+    ):
+        """Initialize the model RAM cache.
+
+        :param execution_device_working_mem_gb: The amount of working memory to keep on the GPU (in GB) i.e. non-model
+            VRAM.
+        :param enable_partial_loading: Whether to enable partial loading of models.
+        :param max_ram_cache_size_gb: The maximum amount of CPU RAM to use for model caching in GB. This parameter is
+            kept to maintain compatibility with previous versions of the model cache, but should be deprecated in the
+            future. If set, this parameter overrides the default cache size logic.
+        :param max_vram_cache_size_gb: The amount of VRAM to use for model caching in GB. This parameter is kept to
+            maintain compatibility with previous versions of the model cache, but should be deprecated in the future.
+            If set, this parameter overrides the default cache size logic.
+        :param execution_device: Torch device to load active model into [torch.device('cuda')]
+        :param storage_device: Torch device to save inactive model in [torch.device('cpu')]
+        :param log_memory_usage: If True, a memory snapshot will be captured before and after every model cache
+            operation, and the result will be logged (at debug level). There is a time cost to capturing the memory
+            snapshots, so it is recommended to disable this feature unless you are actively inspecting the model cache's
+            behaviour.
+        :param logger: InvokeAILogger to use (otherwise creates one)
+        :param keep_alive_minutes: How long to keep models in cache after last use (in minutes). 0 means keep indefinitely.
+        """
+        self._enable_partial_loading = enable_partial_loading
+        self._keep_ram_copy_of_weights = keep_ram_copy_of_weights
+        self._execution_device_working_mem_gb = execution_device_working_mem_gb
+        self._execution_device: torch.device = torch.device(execution_device)
+        self._storage_device: torch.device = torch.device(storage_device)
+
+        self._max_ram_cache_size_gb = max_ram_cache_size_gb
+        self._max_vram_cache_size_gb = max_vram_cache_size_gb
+
+        self._logger = PrefixedLoggerAdapter(
+            logger or InvokeAILogger.get_logger(self.__class__.__name__), "MODEL CACHE"
+        )
+        self._log_memory_usage = log_memory_usage
+        self._stats: Optional[CacheStats] = None
+
+        self._cached_models: Dict[str, CacheRecord] = {}
+        self._cache_stack: List[str] = []
+
+        self._ram_cache_size_bytes = self._calc_ram_available_to_model_cache()
+
+        # A lock applied to all public method calls to make the ModelCache thread-safe.
+        # At the time of writing, the ModelCache should only be accessed from two threads:
+        # - The graph execution thread
+        # - Requests to empty the cache from a separate thread
+        self._lock = threading.RLock()
+
+        self._on_cache_hit_callbacks: set[CacheHitCallback] = set()
+        self._on_cache_miss_callbacks: set[CacheMissCallback] = set()
+        self._on_cache_models_cleared_callbacks: set[CacheModelsClearedCallback] = set()
+
+        # Keep-alive timeout support
+        self._keep_alive_minutes = keep_alive_minutes
+        self._last_activity_time: Optional[float] = None
+        self._timeout_timer: Optional[threading.Timer] = None
+        self._shutdown_event = threading.Event()
+
+    def on_cache_hit(self, cb: CacheHitCallback) -> Callable[[], None]:
+        self._on_cache_hit_callbacks.add(cb)
+
+        def unsubscribe() -> None:
+            self._on_cache_hit_callbacks.discard(cb)
+
+        return unsubscribe
+
+    def on_cache_miss(self, cb: CacheMissCallback) -> Callable[[], None]:
+        self._on_cache_miss_callbacks.add(cb)
+
+        def unsubscribe() -> None:
+            self._on_cache_miss_callbacks.discard(cb)
+
+        return unsubscribe
+
+    def on_cache_models_cleared(self, cb: CacheModelsClearedCallback) -> Callable[[], None]:
+        self._on_cache_models_cleared_callbacks.add(cb)
+
+        def unsubscribe() -> None:
+            self._on_cache_models_cleared_callbacks.discard(cb)
+
+        return unsubscribe
+
+    @property
+    @synchronized
+    def stats(self) -> Optional[CacheStats]:
+        """Return collected CacheStats object."""
+        return self._stats
+
+    @stats.setter
+    @synchronized
+    def stats(self, stats: CacheStats) -> None:
+        """Set the CacheStats object for collecting cache statistics."""
+        self._stats = stats
+        # Populate the cache size in the stats object when it's set
+        if self._stats is not None:
+            self._stats.cache_size = self._ram_cache_size_bytes
+
+    def _record_activity(self) -> None:
+        """Record model activity and reset the timeout timer if configured.
+
+        Note: This method should only be called when self._lock is already held.
+        """
+        if self._keep_alive_minutes <= 0:
+            return
+
+        self._last_activity_time = time.time()
+
+        # Cancel any existing timer
+        if self._timeout_timer is not None:
+            self._timeout_timer.cancel()
+
+        # Start a new timer
+        timeout_seconds = self._keep_alive_minutes * 60
+        self._timeout_timer = threading.Timer(timeout_seconds, self._on_timeout)
+        # Set as daemon so it doesn't prevent application shutdown
+        self._timeout_timer.daemon = True
+        self._timeout_timer.start()
+        self._logger.debug(f"Model cache activity recorded. Timeout set to {self._keep_alive_minutes} minutes.")
+
+    @synchronized
+    @record_activity
+    def _on_timeout(self) -> None:
+        """Called when the keep-alive timeout expires. Clears the model cache."""
+        if self._shutdown_event.is_set():
+            return
+
+        # Double-check if there has been activity since the timer was set
+        # This handles the race condition where activity occurred just before the timer fired
+        if self._last_activity_time is not None and self._keep_alive_minutes > 0:
+            elapsed_minutes = (time.time() - self._last_activity_time) / 60
+            if elapsed_minutes < self._keep_alive_minutes:
+                # Activity occurred, don't clear cache
+                self._logger.debug(
+                    f"Model cache timeout fired but activity detected {elapsed_minutes:.2f} minutes ago. "
+                    f"Skipping cache clear."
+                )
+                return
+
+        # Check if there are any unlocked models that can be cleared
+        unlocked_models = [key for key, entry in self._cached_models.items() if not entry.is_locked]
+
+        if len(unlocked_models) > 0:
+            self._logger.info(
+                f"Model cache keep-alive timeout of {self._keep_alive_minutes} minutes expired. "
+                f"Clearing {len(unlocked_models)} unlocked model(s) from cache."
+            )
+            # Clear the cache by requesting a very large amount of space.
+            # This is the same logic used by the "Clear Model Cache" button.
+            # Using 1000 GB ensures all unlocked models are removed.
+            self._make_room_internal(1000 * GB)
+        elif len(self._cached_models) > 0:
+            # All models are locked, don't log at info level
+            self._logger.debug(
+                f"Model cache timeout fired but all {len(self._cached_models)} model(s) are locked. "
+                f"Skipping cache clear."
+            )
+        else:
+            self._logger.debug("Model cache timeout fired but cache is already empty.")
+
+    @synchronized
+    def shutdown(self) -> None:
+        """Shutdown the model cache, cancelling any pending timers."""
+        self._shutdown_event.set()
+        if self._timeout_timer is not None:
+            self._timeout_timer.cancel()
+            self._timeout_timer = None
+
+    @synchronized
+    @record_activity
+    def put(self, key: str, model: AnyModel, execution_device: Optional[torch.device] = None) -> None:
+        """Add a model to the cache.
+
+        Args:
+            key: Cache key for the model
+            model: The model to cache
+            execution_device: Optional device to use for this specific model. If None, uses the cache's default
+                execution_device. Use torch.device("cpu") to force a model to run on CPU.
+        """
+        if key in self._cached_models:
+            self._logger.debug(
+                f"Attempted to add model {key} ({model.__class__.__name__}), but it already exists in the cache. No action necessary."
+            )
+            return
+
+        size = calc_model_size_by_data(self._logger, model)
+        self._make_room_internal(size)
+
+        # Inject custom modules into the model.
+        if isinstance(model, torch.nn.Module):
+            apply_custom_layers_to_model(model)
+
+        # Use the provided execution device, or fall back to the cache's default
+        effective_execution_device = execution_device if execution_device is not None else self._execution_device
+
+        # Partial loading only makes sense on CUDA.
+        # - When running on CPU, there is no 'loading' to do.
+        # - When running on MPS, memory is shared with the CPU, so the default OS memory management already handles this
+        #   well.
+        running_with_cuda = effective_execution_device.type == "cuda"
+
+        # Wrap model.
+        if isinstance(model, torch.nn.Module) and running_with_cuda and self._enable_partial_loading:
+            wrapped_model = CachedModelWithPartialLoad(
+                model, effective_execution_device, keep_ram_copy=self._keep_ram_copy_of_weights
+            )
+        else:
+            wrapped_model = CachedModelOnlyFullLoad(
+                model, effective_execution_device, size, keep_ram_copy=self._keep_ram_copy_of_weights
+            )
+
+        cache_record = CacheRecord(key=key, cached_model=wrapped_model)
+        self._cached_models[key] = cache_record
+        self._cache_stack.append(key)
+        self._logger.debug(
+            f"Added model {key} (Type: {model.__class__.__name__}, Wrap mode: {wrapped_model.__class__.__name__}, Model size: {size / MB:.2f}MB)"
+        )
+
+    @synchronized
+    def _get_cache_snapshot(self) -> dict[str, CacheEntrySnapshot]:
+        overview: dict[str, CacheEntrySnapshot] = {}
+        for cache_key, cache_entry in self._cached_models.items():
+            total_bytes = cache_entry.cached_model.total_bytes()
+            current_vram_bytes = cache_entry.cached_model.cur_vram_bytes()
+            overview[cache_key] = CacheEntrySnapshot(
+                cache_key=cache_key,
+                total_bytes=total_bytes,
+                current_vram_bytes=current_vram_bytes,
+            )
+
+        return overview
+
+    @synchronized
+    @record_activity
+    def get(self, key: str, stats_name: Optional[str] = None) -> CacheRecord:
+        """Retrieve a model from the cache.
+
+        :param key: Model key
+        :param stats_name: A human-readable id for the model for the purposes of stats reporting.
+
+        Raises IndexError if the model is not in the cache.
+        """
+        if key in self._cached_models:
+            if self.stats:
+                self.stats.hits += 1
+        else:
+            for cb in self._on_cache_miss_callbacks:
+                cb(model_key=key, cache_snapshot=self._get_cache_snapshot())
+            if self.stats:
+                self.stats.misses += 1
+            self._logger.debug(f"Cache miss: {key}")
+            raise IndexError(f"The model with key {key} is not in the cache.")
+
+        cache_entry = self._cached_models[key]
+
+        # more stats
+        if self.stats:
+            stats_name = stats_name or key
+            self.stats.high_watermark = max(self.stats.high_watermark, self._get_ram_in_use())
+            self.stats.in_cache = len(self._cached_models)
+            self.stats.loaded_model_sizes[stats_name] = max(
+                self.stats.loaded_model_sizes.get(stats_name, 0), cache_entry.cached_model.total_bytes()
+            )
+
+        # This moves the entry to the top (right end) of the stack.
+        self._cache_stack = [k for k in self._cache_stack if k != key]
+        self._cache_stack.append(key)
+
+        self._logger.debug(f"Cache hit: {key} (Type: {cache_entry.cached_model.model.__class__.__name__})")
+        for cb in self._on_cache_hit_callbacks:
+            cb(model_key=key, cache_snapshot=self._get_cache_snapshot())
+
+        return cache_entry
+
+    @synchronized
+    @record_activity
+    def lock(self, cache_entry: CacheRecord, working_mem_bytes: Optional[int]) -> None:
+        """Lock a model for use and move it into VRAM."""
+        if cache_entry.key not in self._cached_models:
+            self._logger.info(
+                f"Locking model cache entry {cache_entry.key} "
+                f"(Type: {cache_entry.cached_model.model.__class__.__name__}), but it has already been dropped from "
+                "the RAM cache. This is a sign that the model loading order is non-optimal in the invocation code "
+                "(See https://github.com/invoke-ai/InvokeAI/issues/7513)."
+            )
+        # cache_entry = self._cached_models[key]
+        cache_entry.lock()
+
+        self._logger.debug(
+            f"Locking model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})"
+        )
+
+        # Check if the model's specific compute_device is CPU, not just the cache's default execution_device
+        model_compute_device = cache_entry.cached_model.compute_device
+        if model_compute_device.type == "cpu":
+            # Models configured for CPU execution don't need to be loaded into VRAM
+            self._logger.debug(f"Model {cache_entry.key} is configured for CPU execution, skipping VRAM load")
+            return
+
+        try:
+            self._load_locked_model(cache_entry, working_mem_bytes)
+            self._logger.debug(
+                f"Finished locking model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})"
+            )
+        except torch.cuda.OutOfMemoryError:
+            self._logger.warning("Insufficient GPU memory to load model. Aborting")
+            cache_entry.unlock()
+            raise
+        except Exception:
+            cache_entry.unlock()
+            raise
+
+        self._log_cache_state()
+
+    @synchronized
+    @record_activity
+    def unlock(self, cache_entry: CacheRecord) -> None:
+        """Unlock a model."""
+        if cache_entry.key not in self._cached_models:
+            self._logger.info(
+                f"Unlocking model cache entry {cache_entry.key} "
+                f"(Type: {cache_entry.cached_model.model.__class__.__name__}), but it has already been dropped from "
+                "the RAM cache. This is a sign that the model loading order is non-optimal in the invocation code "
+                "(See https://github.com/invoke-ai/InvokeAI/issues/7513)."
+            )
+        # cache_entry = self._cached_models[key]
+        cache_entry.unlock()
+        self._logger.debug(
+            f"Unlocked model {cache_entry.key} (Type: {cache_entry.cached_model.model.__class__.__name__})"
+        )
+
+        # If `drop_model()` marked this entry stale (e.g. settings changed while a generation
+        # was using it), evict now so the next load rebuilds with the new settings rather than
+        # silently reusing the pre-change cached module.
+        if cache_entry.is_stale and not cache_entry.is_locked and cache_entry.key in self._cached_models:
+            bytes_freed = cache_entry.cached_model.total_bytes()
+            self._delete_cache_entry(cache_entry)
+            if self.stats:
+                self.stats.cleared = (self.stats.cleared or 0) + 1
+            snapshot = self._get_cache_snapshot()
+            for cb in self._on_cache_models_cleared_callbacks:
+                cb(
+                    models_cleared=1,
+                    bytes_requested=0,
+                    bytes_freed=bytes_freed,
+                    cache_snapshot=snapshot,
+                )
+            gc.collect()
+            TorchDevice.empty_cache()
+            self._logger.debug(f"Evicted stale cache entry {cache_entry.key} after unlock.")
+
+    def _load_locked_model(self, cache_entry: CacheRecord, working_mem_bytes: Optional[int] = None) -> None:
+        """Helper function for self.lock(). Loads a locked model into VRAM."""
+        start_time = time.time()
+
+        # Calculate model_vram_needed, the amount of additional VRAM that will be used if we fully load the model into
+        # VRAM.
+        model_cur_vram_bytes = cache_entry.cached_model.cur_vram_bytes()
+        model_total_bytes = cache_entry.cached_model.total_bytes()
+        model_vram_needed = model_total_bytes - model_cur_vram_bytes
+
+        vram_available = self._get_vram_available(working_mem_bytes)
+        self._logger.debug(
+            f"Before unloading: {self._get_vram_state_str(model_cur_vram_bytes, model_total_bytes, vram_available)}"
+        )
+
+        # Make room for the model in VRAM.
+        # 1. If the model can fit entirely in VRAM, then make enough room for it to be loaded fully.
+        # 2. If the model can't fit fully into VRAM, then unload all other models and load as much of the model as
+        #    possible.
+        vram_bytes_freed = self._offload_unlocked_models(model_vram_needed, working_mem_bytes)
+        self._logger.debug(f"Unloaded models (if necessary): vram_bytes_freed={(vram_bytes_freed / MB):.2f}MB")
+
+        # Check the updated vram_available after offloading.
+        vram_available = self._get_vram_available(working_mem_bytes)
+        self._logger.debug(
+            f"After unloading: {self._get_vram_state_str(model_cur_vram_bytes, model_total_bytes, vram_available)}"
+        )
+
+        if vram_available < 0:
+            # There is insufficient VRAM available. As a last resort, try to unload the model being locked from VRAM,
+            # as it may still be loaded from a previous use.
+            vram_bytes_freed_from_own_model = self._move_model_to_ram(cache_entry, -vram_available)
+            vram_available = self._get_vram_available(working_mem_bytes)
+            self._logger.debug(
+                f"Unloaded {vram_bytes_freed_from_own_model / MB:.2f}MB from the model being locked ({cache_entry.key})."
+            )
+
+        # Move as much of the model as possible into VRAM.
+        # For testing, only allow 10% of the model to be loaded into VRAM.
+        # vram_available = int(model_vram_needed * 0.1)
+        # We add 1 MB to the available VRAM to account for small errors in memory tracking (e.g. off-by-one). A fully
+        # loaded model is much faster than a 95% loaded model.
+        model_bytes_loaded = self._move_model_to_vram(cache_entry, vram_available + MB)
+
+        model_cur_vram_bytes = cache_entry.cached_model.cur_vram_bytes()
+        vram_available = self._get_vram_available(working_mem_bytes)
+        loaded_percent = model_cur_vram_bytes / model_total_bytes if model_total_bytes > 0 else 0
+        # Use the model's actual compute_device for logging, not the cache's default
+        model_device = cache_entry.cached_model.compute_device
+        self._logger.info(
+            f"Loaded model '{cache_entry.key}' ({cache_entry.cached_model.model.__class__.__name__}) onto "
+            f"{model_device.type} device in {(time.time() - start_time):.2f}s. "
+            f"Total model size: {model_total_bytes / MB:.2f}MB, "
+            f"VRAM: {model_cur_vram_bytes / MB:.2f}MB ({loaded_percent:.1%})"
+        )
+        self._logger.debug(
+            f"Loaded model onto execution device: model_bytes_loaded={(model_bytes_loaded / MB):.2f}MB, "
+        )
+        self._logger.debug(
+            f"After loading: {self._get_vram_state_str(model_cur_vram_bytes, model_total_bytes, vram_available)}"
+        )
+
+    def _move_model_to_vram(self, cache_entry: CacheRecord, vram_available: int) -> int:
+        try:
+            if isinstance(cache_entry.cached_model, CachedModelWithPartialLoad):
+                return cache_entry.cached_model.partial_load_to_vram(vram_available)
+            elif isinstance(cache_entry.cached_model, CachedModelOnlyFullLoad):  # type: ignore
+                # Partial load is not supported, so we have not choice but to try and fit it all into VRAM.
+                return cache_entry.cached_model.full_load_to_vram()
+            else:
+                raise ValueError(f"Unsupported cached model type: {type(cache_entry.cached_model)}")
+        except Exception as e:
+            if isinstance(e, torch.cuda.OutOfMemoryError):
+                self._logger.warning("Insufficient GPU memory to load model. Aborting")
+            # If an exception occurs, the model could be left in a bad state, so we delete it from the cache entirely.
+            self._delete_cache_entry(cache_entry)
+            raise
+
+    def _move_model_to_ram(self, cache_entry: CacheRecord, vram_bytes_to_free: int) -> int:
+        try:
+            if isinstance(cache_entry.cached_model, CachedModelWithPartialLoad):
+                return cache_entry.cached_model.partial_unload_from_vram(
+                    vram_bytes_to_free, keep_required_weights_in_vram=cache_entry.is_locked
+                )
+            elif isinstance(cache_entry.cached_model, CachedModelOnlyFullLoad):  # type: ignore
+                return cache_entry.cached_model.full_unload_from_vram()
+            else:
+                raise ValueError(f"Unsupported cached model type: {type(cache_entry.cached_model)}")
+        except Exception:
+            # If an exception occurs, the model could be left in a bad state, so we delete it from the cache entirely.
+            self._delete_cache_entry(cache_entry)
+            raise
+
+    def _get_vram_available(self, working_mem_bytes: Optional[int]) -> int:
+        """Calculate the amount of additional VRAM available for the cache to use (takes into account the working
+        memory).
+        """
+        # If self._max_vram_cache_size_gb is set, then it overrides the default logic.
+        if self._max_vram_cache_size_gb is not None:
+            vram_total_available_to_cache = int(self._max_vram_cache_size_gb * GB)
+            return vram_total_available_to_cache - self._get_vram_in_use()
+
+        working_mem_bytes_default = int(self._execution_device_working_mem_gb * GB)
+        working_mem_bytes = max(working_mem_bytes or working_mem_bytes_default, working_mem_bytes_default)
+
+        if self._execution_device.type == "cuda":
+            # TODO(ryand): It is debatable whether we should use memory_reserved() or memory_allocated() here.
+            # memory_reserved() includes memory reserved by the torch CUDA memory allocator that may or may not be
+            # re-used for future allocations. For now, we use memory_allocated() to be conservative.
+            # vram_reserved = torch.cuda.memory_reserved(self._execution_device)
+            vram_allocated = torch.cuda.memory_allocated(self._execution_device)
+            vram_free, _vram_total = torch.cuda.mem_get_info(self._execution_device)
+            vram_available_to_process = vram_free + vram_allocated
+        elif self._execution_device.type == "mps":
+            vram_reserved = torch.mps.driver_allocated_memory()
+            # TODO(ryand): Is it accurate that MPS shares memory with the CPU?
+            vram_free = psutil.virtual_memory().available
+            vram_available_to_process = vram_free + vram_reserved
+        else:
+            raise ValueError(f"Unsupported execution device: {self._execution_device.type}")
+
+        vram_total_available_to_cache = vram_available_to_process - working_mem_bytes
+        vram_cur_available_to_cache = vram_total_available_to_cache - self._get_vram_in_use()
+        return vram_cur_available_to_cache
+
+    def _get_vram_in_use(self) -> int:
+        """Get the amount of VRAM currently in use by the cache."""
+        if self._execution_device.type == "cuda":
+            return torch.cuda.memory_allocated()
+        elif self._execution_device.type == "mps":
+            return torch.mps.current_allocated_memory()
+        else:
+            raise ValueError(f"Unsupported execution device type: {self._execution_device.type}")
+        # Alternative definition of VRAM in use:
+        # return sum(ce.cached_model.cur_vram_bytes() for ce in self._cached_models.values())
+
+    def _calc_ram_available_to_model_cache(self) -> int:
+        """Calculate the amount of RAM available for the cache to use."""
+        # If self._max_ram_cache_size_gb is set, then it overrides the default logic.
+        if self._max_ram_cache_size_gb is not None:
+            self._logger.info(f"Using user-defined RAM cache size: {self._max_ram_cache_size_gb} GB.")
+            return int(self._max_ram_cache_size_gb * GB)
+
+        # Heuristics for dynamically calculating the RAM cache size, **in order of increasing priority**:
+        # 1. As an initial default, use 50% of the total RAM for InvokeAI.
+        #   - Assume a 2GB baseline for InvokeAI's non-model RAM usage, and use the rest of the RAM for the model cache.
+        # 2. On a system with a lot of RAM, users probably don't want InvokeAI to eat up too much RAM.
+        #    There are diminishing returns to storing more and more models. So, we apply an upper bound. (Keep in mind
+        #    that most OSes have some amount of disk caching, which we still benefit from if there is excess memory,
+        #    even if we drop models from the cache.)
+        #    - On systems without a CUDA device, the upper bound is 32GB.
+        #    - On systems with a CUDA device, the upper bound is 1x the amount of VRAM (less the working memory).
+        # 3. Absolute minimum of 4GB.
+
+        # NOTE(ryand): We explored dynamically adjusting the RAM cache size based on memory pressure (using psutil), but
+        # decided against it for now, for the following reasons:
+        # - It was surprisingly difficult to get memory metrics with consistent definitions across OSes. (If you go
+        #   down this path again, don't underestimate the amount of complexity here and be sure to test rigorously on all
+        #   OSes.)
+        # - Making the RAM cache size dynamic opens the door for performance regressions that are hard to diagnose and
+        #   hard for users to understand. It is better for users to see that their RAM is maxed out, and then override
+        #   the default value if desired.
+
+        # Lookup the total VRAM size for the CUDA execution device.
+        total_cuda_vram_bytes: int | None = None
+        if self._execution_device.type == "cuda":
+            _, total_cuda_vram_bytes = torch.cuda.mem_get_info(self._execution_device)
+
+        # Apply heuristic 1.
+        # ------------------
+        heuristics_applied = [1]
+        total_system_ram_bytes = psutil.virtual_memory().total
+        # Assumed baseline RAM used by InvokeAI for non-model stuff.
+        baseline_ram_used_by_invokeai = 2 * GB
+        ram_available_to_model_cache = int(total_system_ram_bytes * 0.5 - baseline_ram_used_by_invokeai)
+
+        # Apply heuristic 2.
+        # ------------------
+        max_ram_cache_size_bytes = 32 * GB
+        if total_cuda_vram_bytes is not None:
+            if self._max_vram_cache_size_gb is not None:
+                max_ram_cache_size_bytes = int(self._max_vram_cache_size_gb * GB)
+            else:
+                max_ram_cache_size_bytes = total_cuda_vram_bytes - int(self._execution_device_working_mem_gb * GB)
+        if ram_available_to_model_cache > max_ram_cache_size_bytes:
+            heuristics_applied.append(2)
+            ram_available_to_model_cache = max_ram_cache_size_bytes
+
+        # Apply heuristic 3.
+        # ------------------
+        if ram_available_to_model_cache < 4 * GB:
+            heuristics_applied.append(3)
+            ram_available_to_model_cache = 4 * GB
+
+        self._logger.info(
+            f"Calculated model RAM cache size: {ram_available_to_model_cache / MB:.2f} MB. Heuristics applied: {heuristics_applied}."
+        )
+        return ram_available_to_model_cache
+
+    def _get_ram_in_use(self) -> int:
+        """Get the amount of RAM currently in use."""
+        return sum(ce.cached_model.total_bytes() for ce in self._cached_models.values())
+
+    def _get_ram_available(self) -> int:
+        """Get the amount of RAM available for the cache to use."""
+        return self._ram_cache_size_bytes - self._get_ram_in_use()
+
+    def _capture_memory_snapshot(self) -> Optional[MemorySnapshot]:
+        if self._log_memory_usage:
+            return MemorySnapshot.capture()
+        return None
+
+    def _get_vram_state_str(self, model_cur_vram_bytes: int, model_total_bytes: int, vram_available: int) -> str:
+        """Helper function for preparing a VRAM state log string."""
+        model_cur_vram_bytes_percent = model_cur_vram_bytes / model_total_bytes if model_total_bytes > 0 else 0
+        return (
+            f"model_total={model_total_bytes / MB:.0f} MB, "
+            + f"model_vram={model_cur_vram_bytes / MB:.0f} MB ({model_cur_vram_bytes_percent:.1%} %), "
+            # + f"vram_total={int(self._max_vram_cache_size * GB)/MB:.0f} MB, "
+            + f"vram_available={(vram_available / MB):.0f} MB, "
+        )
+
+    def _offload_unlocked_models(self, vram_bytes_required: int, working_mem_bytes: Optional[int] = None) -> int:
+        """Offload models from the execution_device until vram_bytes_required bytes are available, or all models are
+        offloaded. Of course, locked models are not offloaded.
+
+        Returns:
+            int: The number of bytes freed based on believed model sizes. The actual change in VRAM may be different.
+        """
+        self._logger.debug(
+            f"Offloading unlocked models with goal of making room for {vram_bytes_required / MB:.2f}MB of VRAM."
+        )
+        vram_bytes_freed = 0
+        # TODO(ryand): Give more thought to the offloading policy used here.
+        cache_entries_increasing_size = sorted(self._cached_models.values(), key=lambda x: x.cached_model.total_bytes())
+        for cache_entry in cache_entries_increasing_size:
+            # We do not fully trust the count of bytes freed, so we check again on each iteration.
+            vram_available = self._get_vram_available(working_mem_bytes)
+            vram_bytes_to_free = vram_bytes_required - vram_available
+            if vram_bytes_to_free <= 0:
+                break
+            if cache_entry.is_locked:
+                # TODO(ryand): In the future, we may want to partially unload locked models, but this requires careful
+                # handling of model patches (e.g. LoRA).
+                continue
+            cache_entry_bytes_freed = self._move_model_to_ram(cache_entry, vram_bytes_to_free)
+            if cache_entry_bytes_freed > 0:
+                self._logger.debug(
+                    f"Unloaded {cache_entry.key} from VRAM to free {(cache_entry_bytes_freed / MB):.0f} MB."
+                )
+            vram_bytes_freed += cache_entry_bytes_freed
+
+        TorchDevice.empty_cache()
+        return vram_bytes_freed
+
+    def _log_cache_state(self, title: str = "Model cache state:", include_entry_details: bool = True):
+        if self._logger.getEffectiveLevel() > logging.DEBUG:
+            # Short circuit if the logger is not set to debug. Some of the data lookups could take a non-negligible
+            # amount of time.
+            return
+
+        log = f"{title}\n"
+
+        log_format = "  {:<30} Limit: {:>7.1f} MB, Used: {:>7.1f} MB ({:>5.1%}), Available: {:>7.1f} MB ({:>5.1%})\n"
+
+        ram_in_use_bytes = self._get_ram_in_use()
+        ram_available_bytes = self._get_ram_available()
+        ram_size_bytes = ram_in_use_bytes + ram_available_bytes
+        ram_in_use_bytes_percent = ram_in_use_bytes / ram_size_bytes if ram_size_bytes > 0 else 0
+        ram_available_bytes_percent = ram_available_bytes / ram_size_bytes if ram_size_bytes > 0 else 0
+        log += log_format.format(
+            f"Storage Device ({self._storage_device.type})",
+            ram_size_bytes / MB,
+            ram_in_use_bytes / MB,
+            ram_in_use_bytes_percent,
+            ram_available_bytes / MB,
+            ram_available_bytes_percent,
+        )
+
+        if self._execution_device.type != "cpu":
+            vram_in_use_bytes = self._get_vram_in_use()
+            vram_available_bytes = self._get_vram_available(None)
+            vram_size_bytes = vram_in_use_bytes + vram_available_bytes
+            vram_in_use_bytes_percent = vram_in_use_bytes / vram_size_bytes if vram_size_bytes > 0 else 0
+            vram_available_bytes_percent = vram_available_bytes / vram_size_bytes if vram_size_bytes > 0 else 0
+            log += log_format.format(
+                f"Compute Device ({self._execution_device.type})",
+                vram_size_bytes / MB,
+                vram_in_use_bytes / MB,
+                vram_in_use_bytes_percent,
+                vram_available_bytes / MB,
+                vram_available_bytes_percent,
+            )
+
+        if torch.cuda.is_available():
+            log += "  {:<30} {:.1f} MB\n".format("CUDA Memory Allocated:", torch.cuda.memory_allocated() / MB)
+        log += "  {:<30} {}\n".format("Total models:", len(self._cached_models))
+
+        if include_entry_details and len(self._cached_models) > 0:
+            log += "  Models:\n"
+            log_format = (
+                "    {:<80} total={:>7.1f} MB, vram={:>7.1f} MB ({:>5.1%}), ram={:>7.1f} MB ({:>5.1%}), locked={}\n"
+            )
+            for cache_record in self._cached_models.values():
+                total_bytes = cache_record.cached_model.total_bytes()
+                cur_vram_bytes = cache_record.cached_model.cur_vram_bytes()
+                cur_vram_bytes_percent = cur_vram_bytes / total_bytes if total_bytes > 0 else 0
+                cur_ram_bytes = total_bytes - cur_vram_bytes
+                cur_ram_bytes_percent = cur_ram_bytes / total_bytes if total_bytes > 0 else 0
+
+                log += log_format.format(
+                    f"{cache_record.key} ({cache_record.cached_model.model.__class__.__name__}):",
+                    total_bytes / MB,
+                    cur_vram_bytes / MB,
+                    cur_vram_bytes_percent,
+                    cur_ram_bytes / MB,
+                    cur_ram_bytes_percent,
+                    cache_record.is_locked,
+                )
+
+        self._logger.debug(log)
+
+    @synchronized
+    def make_room(self, bytes_needed: int) -> None:
+        """Make enough room in the cache to accommodate a new model of indicated size.
+
+        Note: This function deletes all of the cache's internal references to a model in order to free it. If there are
+        external references to the model, there's nothing that the cache can do about it, and those models will not be
+        garbage-collected.
+        """
+        self._make_room_internal(bytes_needed)
+
+    def _make_room_internal(self, bytes_needed: int) -> None:
+        """Internal implementation of make_room(). Assumes the lock is already held."""
+        self._logger.debug(f"Making room for {bytes_needed / MB:.2f}MB of RAM.")
+        self._log_cache_state(title="Before dropping models:")
+
+        ram_bytes_available = self._get_ram_available()
+        ram_bytes_to_free = max(0, bytes_needed - ram_bytes_available)
+
+        ram_bytes_freed = 0
+        pos = 0
+        models_cleared = 0
+        while ram_bytes_freed < ram_bytes_to_free and pos < len(self._cache_stack):
+            model_key = self._cache_stack[pos]
+            cache_entry = self._cached_models[model_key]
+
+            if not cache_entry.is_locked:
+                ram_bytes_freed += cache_entry.cached_model.total_bytes()
+                self._logger.debug(
+                    f"Dropping {model_key} from RAM cache to free {(cache_entry.cached_model.total_bytes() / MB):.2f}MB."
+                )
+                self._delete_cache_entry(cache_entry)
+                del cache_entry
+                models_cleared += 1
+            else:
+                pos += 1
+
+        if models_cleared > 0:
+            # There would likely be some 'garbage' to be collected regardless of whether a model was cleared or not, but
+            # there is a significant time cost to calling `gc.collect()`, so we want to use it sparingly. (The time cost
+            # is high even if no garbage gets collected.)
+            #
+            # Calling gc.collect(...) when a model is cleared seems like a good middle-ground:
+            # - If models had to be cleared, it's a signal that we are close to our memory limit.
+            # - If models were cleared, there's a good chance that there's a significant amount of garbage to be
+            #   collected.
+            #
+            # Keep in mind that gc is only responsible for handling reference cycles. Most objects should be cleaned up
+            # immediately when their reference count hits 0.
+            if self.stats:
+                self.stats.cleared = models_cleared
+            for cb in self._on_cache_models_cleared_callbacks:
+                cb(
+                    models_cleared=models_cleared,
+                    bytes_requested=bytes_needed,
+                    bytes_freed=ram_bytes_freed,
+                    cache_snapshot=self._get_cache_snapshot(),
+                )
+            gc.collect()
+
+        TorchDevice.empty_cache()
+        self._logger.debug(f"Dropped {models_cleared} models to free {ram_bytes_freed / MB:.2f}MB of RAM.")
+        self._log_cache_state(title="After dropping models:")
+
+    def _delete_cache_entry(self, cache_entry: CacheRecord) -> None:
+        """Delete cache_entry from the cache if it exists. No exception is thrown if it doesn't exist."""
+        self._cache_stack = [key for key in self._cache_stack if key != cache_entry.key]
+        self._cached_models.pop(cache_entry.key, None)
+
+    @synchronized
+    def drop_model(self, model_key: str) -> int:
+        """Drop all cache entries belonging to a model so the next load rebuilds them.
+
+        Cache keys are `<model_key>` or `<model_key>:<submodel>` (see `get_model_cache_key`),
+        so a single model may have multiple entries. Locked entries are marked `is_stale` and
+        evicted by `unlock()` as soon as the last lock releases — without that, a setting
+        toggled during an in-flight generation would survive on the locked entry and quietly
+        get reused by the next generation.
+
+        Returns the number of entries immediately dropped (locked entries that are only marked
+        stale do not count).
+        """
+        prefix = f"{model_key}:"
+        matching: list[CacheRecord] = [
+            entry for key, entry in self._cached_models.items() if key == model_key or key.startswith(prefix)
+        ]
+
+        dropped: list[CacheRecord] = []
+        bytes_freed = 0
+        for entry in matching:
+            if entry.is_locked:
+                entry.is_stale = True
+                continue
+            bytes_freed += entry.cached_model.total_bytes()
+            self._delete_cache_entry(entry)
+            dropped.append(entry)
+
+        if dropped:
+            if self.stats:
+                self.stats.cleared = len(dropped)
+            snapshot = self._get_cache_snapshot()
+            for cb in self._on_cache_models_cleared_callbacks:
+                cb(
+                    models_cleared=len(dropped),
+                    bytes_requested=0,
+                    bytes_freed=bytes_freed,
+                    cache_snapshot=snapshot,
+                )
+            gc.collect()
+            TorchDevice.empty_cache()
+        return len(dropped)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/cast_to_device.py

@@ -0,0 +1,15 @@
+from typing import TypeVar
+
+import torch
+
+T = TypeVar("T", torch.Tensor, None, torch.Tensor | None)
+
+
+def cast_to_device(t: T, to_device: torch.device) -> T:
+    """Helper function to cast an optional tensor to a target device."""
+    if t is None:
+        return t
+
+    if t.device.type != to_device.type:
+        return t.to(to_device)
+    return t

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/README.md

@@ -0,0 +1,8 @@
+
+This directory contains custom implementations of common torch.nn.Module classes that add support for:
+- Streaming weights to the execution device
+- Applying sidecar patches at execution time (e.g. sidecar LoRA layers)
+
+Each custom class sub-classes the original module type that is is replacing, so the following properties are preserved:
+- `isinstance(m, torch.nn.OrginalModule)` should still work.
+- Patching the weights directly (e.g. for LoRA) should still work. (Of course, this is not possible for quantized layers, hence the sidecar support.)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_conv1d.py

@@ -0,0 +1,43 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.utils import (
+    add_nullable_tensors,
+)
+
+
+class CustomConv1d(torch.nn.Conv1d, CustomModuleMixin):
+    def _autocast_forward_with_patches(self, input: torch.Tensor) -> torch.Tensor:
+        weight = cast_to_device(self.weight, input.device)
+        bias = cast_to_device(self.bias, input.device)
+
+        # Prepare the original parameters for the patch aggregation.
+        orig_params = {"weight": weight, "bias": bias}
+        # Filter out None values.
+        orig_params = {k: v for k, v in orig_params.items() if v is not None}
+
+        aggregated_param_residuals = self._aggregate_patch_parameters(
+            patches_and_weights=self._patches_and_weights,
+            orig_params=orig_params,
+            device=input.device,
+        )
+
+        weight = add_nullable_tensors(weight, aggregated_param_residuals.get("weight", None))
+        bias = add_nullable_tensors(bias, aggregated_param_residuals.get("bias", None))
+        return self._conv_forward(input, weight, bias)
+
+    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight = cast_to_device(self.weight, input.device)
+        bias = cast_to_device(self.bias, input.device)
+        return self._conv_forward(input, weight, bias)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            return self._autocast_forward_with_patches(input)
+        elif self._device_autocasting_enabled:
+            return self._autocast_forward(input)
+        else:
+            return super().forward(input)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_conv2d.py

@@ -0,0 +1,74 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.utils import (
+    add_nullable_tensors,
+)
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+class CustomConv2d(torch.nn.Conv2d, CustomModuleMixin):
+    def _cast_tensor_for_input(self, tensor: torch.Tensor | None, input: torch.Tensor) -> torch.Tensor | None:
+        tensor = cast_to_device(tensor, input.device)
+        if (
+            tensor is not None
+            and input.is_floating_point()
+            and tensor.is_floating_point()
+            and not isinstance(tensor, GGMLTensor)
+            and tensor.dtype != input.dtype
+        ):
+            tensor = tensor.to(dtype=input.dtype)
+        return tensor
+
+    def _autocast_forward_with_patches(self, input: torch.Tensor) -> torch.Tensor:
+        weight = self._cast_tensor_for_input(self.weight, input)
+        bias = self._cast_tensor_for_input(self.bias, input)
+
+        # Prepare the original parameters for the patch aggregation.
+        orig_params = {"weight": weight, "bias": bias}
+        # Filter out None values.
+        orig_params = {k: v for k, v in orig_params.items() if v is not None}
+
+        aggregated_param_residuals = self._aggregate_patch_parameters(
+            patches_and_weights=self._patches_and_weights,
+            orig_params=orig_params,
+            device=input.device,
+        )
+
+        residual_weight = self._cast_tensor_for_input(aggregated_param_residuals.get("weight", None), input)
+        residual_bias = self._cast_tensor_for_input(aggregated_param_residuals.get("bias", None), input)
+        weight = add_nullable_tensors(weight, residual_weight)
+        bias = add_nullable_tensors(bias, residual_bias)
+        return self._conv_forward(input, weight, bias)
+
+    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight = self._cast_tensor_for_input(self.weight, input)
+        bias = self._cast_tensor_for_input(self.bias, input)
+        return self._conv_forward(input, weight, bias)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            return self._autocast_forward_with_patches(input)
+        elif self._device_autocasting_enabled:
+            return self._autocast_forward(input)
+        elif input.is_floating_point() and (
+            (
+                self.weight.is_floating_point()
+                and not isinstance(self.weight, GGMLTensor)
+                and self.weight.dtype != input.dtype
+            )
+            or (
+                self.bias is not None
+                and self.bias.is_floating_point()
+                and not isinstance(self.bias, GGMLTensor)
+                and self.bias.dtype != input.dtype
+            )
+        ):
+            weight = self._cast_tensor_for_input(self.weight, input)
+            bias = self._cast_tensor_for_input(self.bias, input)
+            return self._conv_forward(input, weight, bias)
+        else:
+            return super().forward(input)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_diffusers_rms_norm.py

@@ -0,0 +1,40 @@
+import torch
+from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+
+
+class CustomDiffusersRMSNorm(DiffusersRMSNorm, CustomModuleMixin):
+    """Custom wrapper for diffusers RMSNorm that supports device autocasting for partial model loading."""
+
+    def _autocast_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        weight = cast_to_device(self.weight, hidden_states.device) if self.weight is not None else None
+        bias = cast_to_device(self.bias, hidden_states.device) if self.bias is not None else None
+
+        input_dtype = hidden_states.dtype
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+
+        if weight is not None:
+            # convert into half-precision if necessary
+            if weight.dtype in [torch.float16, torch.bfloat16]:
+                hidden_states = hidden_states.to(weight.dtype)
+            hidden_states = hidden_states * weight
+            if bias is not None:
+                hidden_states = hidden_states + bias
+        else:
+            hidden_states = hidden_states.to(input_dtype)
+
+        return hidden_states
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            raise RuntimeError("DiffusersRMSNorm layers do not support patches")
+
+        if self._device_autocasting_enabled:
+            return self._autocast_forward(hidden_states)
+        else:
+            return super().forward(hidden_states)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_embedding.py

@@ -0,0 +1,29 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+
+
+class CustomEmbedding(torch.nn.Embedding, CustomModuleMixin):
+    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight = cast_to_device(self.weight, input.device)
+        return torch.nn.functional.embedding(
+            input,
+            weight,
+            self.padding_idx,
+            self.max_norm,
+            self.norm_type,
+            self.scale_grad_by_freq,
+            self.sparse,
+        )
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            raise RuntimeError("Embedding layers do not support patches")
+
+        if self._device_autocasting_enabled:
+            return self._autocast_forward(input)
+        else:
+            return super().forward(input)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_flux_rms_norm.py

@@ -0,0 +1,36 @@
+import torch
+
+from invokeai.backend.flux.modules.layers import RMSNorm
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.patches.layers.set_parameter_layer import SetParameterLayer
+
+
+class CustomFluxRMSNorm(RMSNorm, CustomModuleMixin):
+    def _autocast_forward_with_patches(self, x: torch.Tensor) -> torch.Tensor:
+        # Currently, CustomFluxRMSNorm layers only support patching with a single SetParameterLayer.
+        assert len(self._patches_and_weights) == 1
+        patch, _patch_weight = self._patches_and_weights[0]
+        assert isinstance(patch, SetParameterLayer)
+        assert patch.param_name == "scale"
+
+        scale = cast_to_device(patch.weight, x.device)
+
+        # Apply the patch.
+        # NOTE(ryand): Currently, we ignore the patch weight when running as a sidecar. It's not clear how this should
+        # be handled.
+        return torch.nn.functional.rms_norm(x, scale.shape, scale, eps=1e-6)
+
+    def _autocast_forward(self, x: torch.Tensor) -> torch.Tensor:
+        scale = cast_to_device(self.scale, x.device)
+        return torch.nn.functional.rms_norm(x, scale.shape, scale, eps=1e-6)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            return self._autocast_forward_with_patches(x)
+        elif self._device_autocasting_enabled:
+            return self._autocast_forward(x)
+        else:
+            return super().forward(x)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_group_norm.py

@@ -0,0 +1,22 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+
+
+class CustomGroupNorm(torch.nn.GroupNorm, CustomModuleMixin):
+    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight = cast_to_device(self.weight, input.device)
+        bias = cast_to_device(self.bias, input.device)
+        return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            raise RuntimeError("GroupNorm layers do not support patches")
+
+        if self._device_autocasting_enabled:
+            return self._autocast_forward(input)
+        else:
+            return super().forward(input)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_invoke_linear_8_bit_lt.py

@@ -0,0 +1,66 @@
+import bitsandbytes as bnb
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_linear import (
+    autocast_linear_forward_sidecar_patches,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.patches.layers.param_shape_utils import get_param_shape
+from invokeai.backend.quantization.bnb_llm_int8 import InvokeLinear8bitLt
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+class CustomInvokeLinear8bitLt(InvokeLinear8bitLt, CustomModuleMixin):
+    def _cast_tensor_for_input(self, tensor: torch.Tensor | None, input: torch.Tensor) -> torch.Tensor | None:
+        tensor = cast_to_device(tensor, input.device)
+        if (
+            tensor is not None
+            and input.is_floating_point()
+            and tensor.is_floating_point()
+            and not isinstance(tensor, GGMLTensor)
+            and tensor.dtype != input.dtype
+        ):
+            tensor = tensor.to(dtype=input.dtype)
+        return tensor
+
+    def _cast_weight_bias_for_input(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # See the matching method on CustomInvokeLinearNF4 for the rationale. Int8Params doesn't have
+        # the same packed-shape problem as Params4bit, but we still substitute a meta tensor so that
+        # patches don't accidentally read the quantized weight values.
+        weight = torch.empty(get_param_shape(self.weight), device="meta")
+        bias = self._cast_tensor_for_input(self.bias, input)
+        return weight, bias
+
+    def _autocast_forward_with_patches(self, x: torch.Tensor) -> torch.Tensor:
+        return autocast_linear_forward_sidecar_patches(self, x, self._patches_and_weights)
+
+    def _autocast_forward(self, x: torch.Tensor) -> torch.Tensor:
+        matmul_state = bnb.MatmulLtState()
+        matmul_state.threshold = self.state.threshold
+        matmul_state.has_fp16_weights = self.state.has_fp16_weights
+        matmul_state.use_pool = self.state.use_pool
+        matmul_state.is_training = self.training
+        # The underlying InvokeInt8Params weight must already be quantized.
+        assert self.weight.CB is not None
+        matmul_state.CB = cast_to_device(self.weight.CB, x.device)
+        matmul_state.SCB = cast_to_device(self.weight.SCB, x.device)
+
+        # weights are cast automatically as Int8Params, but the bias has to be cast manually.
+        if self.bias is not None and self.bias.dtype != x.dtype:
+            self.bias.data = self.bias.data.to(x.dtype)
+
+        # NOTE(ryand): The second parameter should not be needed at all given our expected inference configuration, but
+        # it's dtype field must be accessible, even though it's not used. We pass in self.weight even though it could be
+        # on the wrong device.
+        return bnb.matmul(x, self.weight, bias=cast_to_device(self.bias, x.device), state=matmul_state)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            return self._autocast_forward_with_patches(x)
+        elif self._device_autocasting_enabled:
+            return self._autocast_forward(x)
+        else:
+            return super().forward(x)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_invoke_linear_nf4.py

@@ -0,0 +1,93 @@
+import copy
+
+import bitsandbytes as bnb
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_linear import (
+    autocast_linear_forward_sidecar_patches,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.patches.layers.param_shape_utils import get_param_shape
+from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+class CustomInvokeLinearNF4(InvokeLinearNF4, CustomModuleMixin):
+    def _cast_tensor_for_input(self, tensor: torch.Tensor | None, input: torch.Tensor) -> torch.Tensor | None:
+        tensor = cast_to_device(tensor, input.device)
+        if (
+            tensor is not None
+            and input.is_floating_point()
+            and tensor.is_floating_point()
+            and not isinstance(tensor, GGMLTensor)
+            and tensor.dtype != input.dtype
+        ):
+            tensor = tensor.to(dtype=input.dtype)
+        return tensor
+
+    def _cast_weight_bias_for_input(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # The NF4 weight is a Params4bit whose .shape reports the *packed-byte* layout, not the logical
+        # (out_features, in_features) shape. We hand patches a meta-device tensor with the correct
+        # logical shape so that shape-only patches (LoRA, LoHA, MergedLayerPatch over LoRA, ...) work.
+        # Patches that read the original weight values (e.g. SetParameterLayer, DoRA) are not supported
+        # on NF4-quantized modules.
+        weight = torch.empty(get_param_shape(self.weight), device="meta")
+        bias = self._cast_tensor_for_input(self.bias, input)
+        return weight, bias
+
+    def _autocast_forward_with_patches(self, x: torch.Tensor) -> torch.Tensor:
+        return autocast_linear_forward_sidecar_patches(self, x, self._patches_and_weights)
+
+    def _autocast_forward(self, x: torch.Tensor) -> torch.Tensor:
+        bnb.nn.modules.fix_4bit_weight_quant_state_from_module(self)
+
+        # weights are cast automatically as Int8Params, but the bias has to be cast manually
+        if self.bias is not None and self.bias.dtype != x.dtype:
+            self.bias.data = self.bias.data.to(x.dtype)
+
+        if not self.compute_type_is_set:
+            self.set_compute_type(x)
+            self.compute_type_is_set = True
+
+        inp_dtype = x.dtype
+        if self.compute_dtype is not None:
+            x = x.to(self.compute_dtype)
+
+        bias = None if self.bias is None else self.bias.to(self.compute_dtype)
+
+        # HACK(ryand): Casting self.weight to the device also casts the self.weight.quant_state in-place (i.e. it
+        # does not follow the tensor semantics of returning a new copy when converting to a different device). This
+        # means that quant_state elements that started on the CPU would be left on the GPU, which we don't want. To
+        # avoid this side effect we make a shallow copy of the original quant_state so that we can restore it. Fixing
+        # this properly would require more invasive changes to the bitsandbytes library.
+
+        # Make a shallow copy of the quant_state so that we can undo the in-place modification that occurs when casting
+        # to a new device.
+        weight_was_offloaded = self.weight.device.type != x.device.type
+        old_quant_state = copy.copy(self.weight.quant_state)
+        weight = cast_to_device(self.weight, x.device)
+        self.weight.quant_state = old_quant_state
+
+        # For some reason, the quant_state.to(...) implementation fails to cast the quant_state.code field. We do this
+        # manually here.
+        weight.quant_state.code = cast_to_device(weight.quant_state.code, x.device)
+
+        bias = cast_to_device(self.bias, x.device)
+        if weight_was_offloaded and x.numel() == x.shape[-1]:
+            # bitsandbytes routes single-vector inputs through gemv_4bit, which can fail with CPU-stored,
+            # device-autocasted Params4bit weights on some CUDA/bnb combinations. Use the same dequantized
+            # matmul path that bnb.matmul_4bit uses for batched inputs.
+            dequantized_weight = bnb.functional.dequantize_4bit(weight, weight.quant_state).to(x.dtype)
+            return torch.nn.functional.linear(x, dequantized_weight, bias).to(inp_dtype)
+        return bnb.matmul_4bit(x, weight.t(), bias=bias, quant_state=weight.quant_state).to(inp_dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            return self._autocast_forward_with_patches(x)
+        elif self._device_autocasting_enabled:
+            return self._autocast_forward(x)
+        else:
+            return super().forward(x)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_layer_norm.py

@@ -0,0 +1,25 @@
+import torch
+import torch.nn.functional as F
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+
+
+class CustomLayerNorm(torch.nn.LayerNorm, CustomModuleMixin):
+    """Custom wrapper for torch.nn.LayerNorm that supports device autocasting for partial model loading."""
+
+    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight = cast_to_device(self.weight, input.device) if self.weight is not None else None
+        bias = cast_to_device(self.bias, input.device) if self.bias is not None else None
+        return F.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            raise RuntimeError("LayerNorm layers do not support patches")
+
+        if self._device_autocasting_enabled:
+            return self._autocast_forward(input)
+        else:
+            return super().forward(input)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py

@@ -0,0 +1,121 @@
+import copy
+
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+def linear_lora_forward(input: torch.Tensor, lora_layer: LoRALayer, lora_weight: float) -> torch.Tensor:
+    """An optimized implementation of the residual calculation for a sidecar linear LoRALayer."""
+    # up matrix and down matrix have different ranks so we can't simply multiply them
+    if lora_layer.up.shape[1] != lora_layer.down.shape[0]:
+        x = torch.nn.functional.linear(input, lora_layer.get_weight(lora_weight), bias=lora_layer.bias)
+        x *= lora_weight * lora_layer.scale()
+        return x
+
+    x = torch.nn.functional.linear(input, lora_layer.down)
+    if lora_layer.mid is not None:
+        x = torch.nn.functional.linear(x, lora_layer.mid)
+    x = torch.nn.functional.linear(x, lora_layer.up, bias=lora_layer.bias)
+    x *= lora_weight * lora_layer.scale()
+    return x
+
+
+def autocast_linear_forward_sidecar_patches(
+    orig_module: torch.nn.Linear, input: torch.Tensor, patches_and_weights: list[tuple[BaseLayerPatch, float]]
+) -> torch.Tensor:
+    """A function that runs a linear layer (quantized or non-quantized) with sidecar patches for a linear layer.
+    Compatible with both quantized and non-quantized Linear layers.
+    """
+    # First, apply the original linear layer.
+    # NOTE: We slice the input to match the original weight shape in order to work with FluxControlLoRAs, which
+    # change the linear layer's in_features.
+    orig_input = input
+    input = orig_input[..., : orig_module.in_features]
+    output = orig_module._autocast_forward(input)
+
+    # Then, apply layers for which we have optimized implementations.
+    unprocessed_patches_and_weights: list[tuple[BaseLayerPatch, float]] = []
+    for patch, patch_weight in patches_and_weights:
+        # Shallow copy the patch so that we can cast it to the target device without modifying the original patch.
+        patch = copy.copy(patch)
+        patch.to(input.device)
+
+        if isinstance(patch, FluxControlLoRALayer):
+            # Note that we use the original input here, not the sliced input.
+            output += linear_lora_forward(orig_input, patch, patch_weight)
+        elif isinstance(patch, LoRALayer):
+            output += linear_lora_forward(input, patch, patch_weight)
+        else:
+            unprocessed_patches_and_weights.append((patch, patch_weight))
+
+    # Finally, apply any remaining patches.
+    if len(unprocessed_patches_and_weights) > 0:
+        weight, bias = orig_module._cast_weight_bias_for_input(input)
+        # Prepare the original parameters for the patch aggregation.
+        orig_params = {"weight": weight, "bias": bias}
+        # Filter out None values.
+        orig_params = {k: v for k, v in orig_params.items() if v is not None}
+
+        aggregated_param_residuals = orig_module._aggregate_patch_parameters(
+            unprocessed_patches_and_weights, orig_params=orig_params, device=input.device
+        )
+        residual_weight = orig_module._cast_tensor_for_input(aggregated_param_residuals["weight"], input)
+        residual_bias = orig_module._cast_tensor_for_input(aggregated_param_residuals.get("bias", None), input)
+        assert residual_weight is not None
+        output += torch.nn.functional.linear(input, residual_weight, residual_bias)
+
+    return output
+
+
+class CustomLinear(torch.nn.Linear, CustomModuleMixin):
+    def _cast_tensor_for_input(self, tensor: torch.Tensor | None, input: torch.Tensor) -> torch.Tensor | None:
+        tensor = cast_to_device(tensor, input.device)
+        if (
+            tensor is not None
+            and input.is_floating_point()
+            and tensor.is_floating_point()
+            and not isinstance(tensor, GGMLTensor)
+            and tensor.dtype != input.dtype
+        ):
+            tensor = tensor.to(dtype=input.dtype)
+        return tensor
+
+    def _cast_weight_bias_for_input(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
+        weight = self._cast_tensor_for_input(self.weight, input)
+        bias = self._cast_tensor_for_input(self.bias, input)
+        assert weight is not None
+        return weight, bias
+
+    def _autocast_forward_with_patches(self, input: torch.Tensor) -> torch.Tensor:
+        return autocast_linear_forward_sidecar_patches(self, input, self._patches_and_weights)
+
+    def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor:
+        weight, bias = self._cast_weight_bias_for_input(input)
+        return torch.nn.functional.linear(input, weight, bias)
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if len(self._patches_and_weights) > 0:
+            return self._autocast_forward_with_patches(input)
+        elif self._device_autocasting_enabled:
+            return self._autocast_forward(input)
+        elif input.is_floating_point() and (
+            (self.weight.is_floating_point() and self.weight.dtype != input.dtype)
+            or (
+                self.bias is not None
+                and self.bias.is_floating_point()
+                and not isinstance(self.bias, GGMLTensor)
+                and self.bias.dtype != input.dtype
+            )
+        ):
+            weight, bias = self._cast_weight_bias_for_input(input)
+            return torch.nn.functional.linear(input, weight, bias)
+        else:
+            return super().forward(input)

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_module_mixin.py

@@ -0,0 +1,82 @@
+import copy
+
+import torch
+
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.param_shape_utils import get_param_shape
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+class CustomModuleMixin:
+    """A mixin class for custom modules that enables device autocasting of module parameters."""
+
+    def __init__(self):
+        self._device_autocasting_enabled = False
+        self._patches_and_weights: list[tuple[BaseLayerPatch, float]] = []
+
+    def set_device_autocasting_enabled(self, enabled: bool):
+        """Pass True to enable autocasting of module parameters to the same device as the input tensor. Pass False to
+        disable autocasting, which results in slightly faster execution speed when we know that device autocasting is
+        not needed.
+        """
+        self._device_autocasting_enabled = enabled
+
+    def is_device_autocasting_enabled(self) -> bool:
+        """Check if device autocasting is enabled for the module."""
+        return self._device_autocasting_enabled
+
+    def add_patch(self, patch: BaseLayerPatch, patch_weight: float):
+        """Add a patch to the module."""
+        self._patches_and_weights.append((patch, patch_weight))
+
+    def clear_patches(self):
+        """Clear all patches from the module."""
+        self._patches_and_weights = []
+
+    def get_num_patches(self) -> int:
+        """Get the number of patches in the module."""
+        return len(self._patches_and_weights)
+
+    def _aggregate_patch_parameters(
+        self,
+        patches_and_weights: list[tuple[BaseLayerPatch, float]],
+        orig_params: dict[str, torch.Tensor],
+        device: torch.device | None = None,
+    ):
+        """Helper function that aggregates the parameters from all patches into a single dict."""
+        # HACK(ryand): If the original parameters are in a quantized format whose weights can't be accessed, we replace
+        # them with dummy tensors on the 'meta' device. This allows patch layers to access the shapes of the original
+        # parameters. But, of course, any sub-layers that need to access the actual values of the parameters will fail.
+        for param_name in orig_params.keys():
+            param = orig_params[param_name]
+            if isinstance(param, torch.nn.Parameter) and type(param.data) is torch.Tensor:
+                pass
+            elif type(param) is torch.Tensor:
+                # Plain tensor (e.g. after cast_to_device moved a Parameter to another device).
+                pass
+            elif type(param) is GGMLTensor:
+                # Move to device and dequantize here. Doing it in the patch layer can result in redundant casts /
+                # dequantizations.
+                orig_params[param_name] = param.to(device=device).get_dequantized_tensor()
+            else:
+                orig_params[param_name] = torch.empty(get_param_shape(param), device="meta")
+
+        params: dict[str, torch.Tensor] = {}
+
+        for patch, patch_weight in patches_and_weights:
+            if device is not None:
+                # Shallow copy the patch so that we can cast it to the target device without modifying the original patch.
+                patch = copy.copy(patch)
+                patch.to(device)
+
+            # TODO(ryand): `self` could be a quantized module. Depending on what the patch is doing with the original
+            # parameters, this might fail or return incorrect results.
+            layer_params = patch.get_parameters(orig_params, weight=patch_weight)
+
+            for param_name, param_weight in layer_params.items():
+                if param_name not in params:
+                    params[param_name] = param_weight
+                else:
+                    params[param_name] += param_weight
+
+        return params

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/utils.py

@@ -0,0 +1,30 @@
+from typing import overload
+
+import torch
+
+
+@overload
+def add_nullable_tensors(a: None, b: None) -> None: ...
+
+
+@overload
+def add_nullable_tensors(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ...
+
+
+@overload
+def add_nullable_tensors(a: torch.Tensor, b: None) -> torch.Tensor: ...
+
+
+@overload
+def add_nullable_tensors(a: None, b: torch.Tensor) -> torch.Tensor: ...
+
+
+def add_nullable_tensors(a: torch.Tensor | None, b: torch.Tensor | None) -> torch.Tensor | None:
+    if a is None and b is None:
+        return None
+    elif a is None:
+        return b
+    elif b is None:
+        return a
+    else:
+        return a + b

invokeai/backend/model_manager/load/model_cache/torch_module_autocast/torch_module_autocast.py

@@ -0,0 +1,114 @@
+from typing import TypeVar
+
+import torch
+from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm
+
+from invokeai.backend.flux.modules.layers import RMSNorm as FluxRMSNorm
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_conv1d import (
+    CustomConv1d,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_conv2d import (
+    CustomConv2d,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_diffusers_rms_norm import (
+    CustomDiffusersRMSNorm,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_embedding import (
+    CustomEmbedding,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_flux_rms_norm import (
+    CustomFluxRMSNorm,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_group_norm import (
+    CustomGroupNorm,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_layer_norm import (
+    CustomLayerNorm,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_linear import (
+    CustomLinear,
+)
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import (
+    CustomModuleMixin,
+)
+
+AUTOCAST_MODULE_TYPE_MAPPING: dict[type[torch.nn.Module], type[torch.nn.Module]] = {
+    torch.nn.Linear: CustomLinear,
+    torch.nn.Conv1d: CustomConv1d,
+    torch.nn.Conv2d: CustomConv2d,
+    torch.nn.GroupNorm: CustomGroupNorm,
+    torch.nn.Embedding: CustomEmbedding,
+    torch.nn.LayerNorm: CustomLayerNorm,
+    FluxRMSNorm: CustomFluxRMSNorm,
+    DiffusersRMSNorm: CustomDiffusersRMSNorm,
+}
+
+try:
+    # These dependencies are not expected to be present on MacOS.
+    from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_8_bit_lt import (
+        CustomInvokeLinear8bitLt,
+    )
+    from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_nf4 import (
+        CustomInvokeLinearNF4,
+    )
+    from invokeai.backend.quantization.bnb_llm_int8 import InvokeLinear8bitLt
+    from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
+
+    AUTOCAST_MODULE_TYPE_MAPPING[InvokeLinear8bitLt] = CustomInvokeLinear8bitLt
+    AUTOCAST_MODULE_TYPE_MAPPING[InvokeLinearNF4] = CustomInvokeLinearNF4
+except ImportError:
+    pass
+
+
+AUTOCAST_MODULE_TYPE_MAPPING_INVERSE = {v: k for k, v in AUTOCAST_MODULE_TYPE_MAPPING.items()}
+
+
+T = TypeVar("T", bound=torch.nn.Module)
+
+
+def wrap_custom_layer(module_to_wrap: torch.nn.Module, custom_layer_type: type[T]) -> T:
+    # HACK(ryand): We use custom initialization logic so that we can initialize a new custom layer instance from an
+    # existing layer instance without calling __init__() on the original layer class. We achieve this by copying
+    # the attributes from the original layer instance to the new instance.
+    custom_layer = custom_layer_type.__new__(custom_layer_type)
+    # Note that we share the __dict__.
+    # TODO(ryand): In the future, we may want to do a shallow copy of the __dict__.
+    custom_layer.__dict__ = module_to_wrap.__dict__
+
+    # Initialize the CustomModuleMixin fields.
+    CustomModuleMixin.__init__(custom_layer)  # type: ignore
+    return custom_layer
+
+
+def unwrap_custom_layer(custom_layer: torch.nn.Module, original_layer_type: type[torch.nn.Module]):
+    # HACK(ryand): We use custom initialization logic so that we can initialize a new custom layer instance from an
+    # existing layer instance without calling __init__() on the original layer class. We achieve this by copying
+    # the attributes from the original layer instance to the new instance.
+    original_layer = original_layer_type.__new__(original_layer_type)
+    # Note that we share the __dict__.
+    # TODO(ryand): In the future, we may want to do a shallow copy of the __dict__ and strip out the CustomModuleMixin
+    # fields.
+    original_layer.__dict__ = custom_layer.__dict__
+    return original_layer
+
+
+def apply_custom_layers_to_model(module: torch.nn.Module, device_autocasting_enabled: bool = False):
+    for name, submodule in module.named_children():
+        override_type = AUTOCAST_MODULE_TYPE_MAPPING.get(type(submodule), None)
+        if override_type is not None:
+            custom_layer = wrap_custom_layer(submodule, override_type)
+            # TODO(ryand): In the future, we should manage this flag on a per-module basis.
+            custom_layer.set_device_autocasting_enabled(device_autocasting_enabled)
+            setattr(module, name, custom_layer)
+        else:
+            # Recursively apply to submodules
+            apply_custom_layers_to_model(submodule, device_autocasting_enabled)
+
+
+def remove_custom_layers_from_model(module: torch.nn.Module):
+    for name, submodule in module.named_children():
+        override_type = AUTOCAST_MODULE_TYPE_MAPPING_INVERSE.get(type(submodule), None)
+        if override_type is not None:
+            setattr(module, name, unwrap_custom_layer(submodule, override_type))
+        else:
+            remove_custom_layers_from_model(submodule)

invokeai/backend/model_manager/load/model_cache/utils.py

@@ -0,0 +1,20 @@
+import itertools
+
+import torch
+
+
+def get_effective_device(model: torch.nn.Module) -> torch.device:
+    """A utility to infer the 'effective' device of a model.
+
+    This utility handles the case where a model is partially loaded onto the GPU, so is safer than just calling:
+    `next(iter(model.parameters())).device`.
+
+    In the worst case, this utility has to check all model parameters, so if you already know the intended model device,
+    then it is better to avoid calling this function.
+    """
+    # If all parameters are on the CPU, return the CPU device. Otherwise, return the first non-CPU device.
+    for p in itertools.chain(model.parameters(), model.buffers()):
+        if p.device.type != "cpu":
+            return p.device
+
+    return torch.device("cpu")