chore: import upstream snapshot with attribution

invokeai/backend/model_manager/load/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development Team
+"""
+Init file for the model loader.
+"""
+
+from importlib import import_module
+from pathlib import Path
+
+from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig, ModelLoaderBase
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry, ModelLoaderRegistryBase
+
+# This registers the subclasses that implement loaders of specific model types
+loaders = [x.stem for x in Path(Path(__file__).parent, "model_loaders").glob("*.py") if x.stem != "__init__"]
+for module in loaders:
+    import_module(f"{__package__}.model_loaders.{module}")
+
+__all__ = [
+    "LoadedModel",
+    "LoadedModelWithoutConfig",
+    "ModelCache",
+    "ModelLoaderBase",
+    "ModelLoader",
+    "ModelLoaderRegistryBase",
+    "ModelLoaderRegistry",
+]

invokeai/backend/model_manager/load/load_base.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""
+Base class for model loading in InvokeAI.
+"""
+
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from logging import Logger
+from pathlib import Path
+from typing import Any, Dict, Generator, Optional, Tuple
+
+import torch
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
+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.model_cache import ModelCache
+from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
+
+
+class LoadedModelWithoutConfig:
+    """Context manager object that mediates transfer from RAM<->VRAM.
+
+    This is a context manager object that has two distinct APIs:
+
+    1. Older API (deprecated):
+    Use the LoadedModel object directly as a context manager.  It will move the model into VRAM (on CUDA devices), and
+    return the model in a form suitable for passing to torch.
+    Example:
+    ```
+    loaded_model_= loader.get_model_by_key('f13dd932', SubModelType('vae'))
+    with loaded_model as vae:
+      image = vae.decode(latents)[0]
+    ```
+
+    2. Newer API (recommended):
+    Call the LoadedModel's `model_on_device()` method in a context. It returns a tuple consisting of a copy of the
+    model's state dict in CPU RAM followed by a copy of the model in VRAM. The state dict is provided to allow LoRAs and
+    other model patchers to return the model to its unpatched state without expensive copy and restore operations.
+
+    Example:
+    ```
+    loaded_model_= loader.get_model_by_key('f13dd932', SubModelType('vae'))
+    with loaded_model.model_on_device() as (state_dict, vae):
+        image = vae.decode(latents)[0]
+    ```
+
+    The state_dict should be treated as a read-only object and never modified. Also be aware that some loadable models
+    do not have a state_dict, in which case this value will be None.
+    """
+
+    def __init__(self, cache_record: CacheRecord, cache: ModelCache):
+        self._cache_record = cache_record
+        self._cache = cache
+
+    def __enter__(self) -> AnyModel:
+        self._cache.lock(self._cache_record, None)
+        try:
+            self.repair_required_tensors_on_device()
+            return self.model
+        except Exception:
+            self._cache.unlock(self._cache_record)
+            raise
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        self._cache.unlock(self._cache_record)
+
+    @contextmanager
+    def model_on_device(
+        self, working_mem_bytes: Optional[int] = None
+    ) -> Generator[Tuple[Optional[Dict[str, torch.Tensor]], AnyModel], None, None]:
+        """Return a tuple consisting of the model's state dict (if it exists) and the locked model on execution device.
+
+        :param working_mem_bytes: The amount of working memory to keep available on the compute device when loading the
+            model.
+        """
+        self._cache.lock(self._cache_record, working_mem_bytes)
+        try:
+            self.repair_required_tensors_on_device()
+            yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model)
+        finally:
+            self._cache.unlock(self._cache_record)
+
+    @property
+    def model(self) -> AnyModel:
+        """Return the model without locking it."""
+        return self._cache_record.cached_model.model
+
+    def repair_required_tensors_on_device(self) -> int:
+        """Repair required tensors that should be resident on the cached model's execution device."""
+        cached_model = self._cache_record.cached_model
+        if not isinstance(cached_model, CachedModelWithPartialLoad):
+            return 0
+        return cached_model.repair_required_tensors_on_compute_device()
+
+
+class LoadedModel(LoadedModelWithoutConfig):
+    """Context manager object that mediates transfer from RAM<->VRAM."""
+
+    def __init__(self, config: Optional[AnyModelConfig], cache_record: CacheRecord, cache: ModelCache):
+        super().__init__(cache_record=cache_record, cache=cache)
+        self.config = config
+
+
+class ModelLoaderBase(ABC):
+    """Abstract base class for loading models into RAM/VRAM."""
+
+    @abstractmethod
+    def __init__(
+        self,
+        app_config: InvokeAIAppConfig,
+        logger: Logger,
+        ram_cache: ModelCache,
+    ):
+        """Initialize the loader."""
+        pass
+
+    @abstractmethod
+    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+        """
+        Return a model given its confguration.
+
+        Given a model identified in the model configuration backend,
+        return a ModelInfo object that can be used to retrieve the model.
+
+        :param model_config: Model configuration, as returned by ModelConfigRecordStore
+        :param submodel_type: an ModelType enum indicating the portion of
+               the model to retrieve (e.g. ModelType.Vae)
+        """
+        pass
+
+    @abstractmethod
+    def get_size_fs(
+        self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None
+    ) -> int:
+        """Return size in bytes of the model, calculated before loading."""
+        pass
+
+    @property
+    @abstractmethod
+    def ram_cache(self) -> ModelCache:
+        """Return the ram cache associated with this loader."""
+        pass

invokeai/backend/model_manager/load/load_default.py

@@ -0,0 +1,312 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Default implementation of model loading in InvokeAI."""
+
+import re
+from logging import Logger
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.model_manager.configs.base import Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase
+from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache, get_model_cache_key
+from invokeai.backend.model_manager.load.model_util import calc_model_size_by_fs
+from invokeai.backend.model_manager.load.optimizations import skip_torch_weight_init
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    SubModelType,
+)
+from invokeai.backend.util.devices import TorchDevice
+
+# Layer classes that benefit from FP8 storage. Mirrors diffusers'
+# `_GO_LC_SUPPORTED_PYTORCH_LAYERS` so the plain-nn.Module fallback path makes the same
+# precision/quality trade-offs as the ModelMixin path. Notably excludes norm and embedding
+# wrapper modules — those are handled by their direct param types (Embedding is included
+# but pos_embed/patch_embed are filtered by `_FP8_DEFAULT_SKIP_PATTERNS`).
+_FP8_SUPPORTED_PYTORCH_LAYERS: tuple[type[torch.nn.Module], ...] = (
+    torch.nn.Linear,
+    torch.nn.Conv1d,
+    torch.nn.Conv2d,
+    torch.nn.Conv3d,
+    torch.nn.ConvTranspose1d,
+    torch.nn.ConvTranspose2d,
+    torch.nn.ConvTranspose3d,
+    torch.nn.Embedding,
+)
+
+# Module-path regexes (matched against `named_modules()` dotted paths) for precision-sensitive
+# layers that should never be cast to FP8. Mirrors diffusers' `DEFAULT_SKIP_MODULES_PATTERN`
+# — without these, FLUX RMSNorm.scale and similar tiny learned scalars get crushed to FP8 and
+# inference quality degrades. Includes anything named `norm`, position/patch embeddings, and
+# the in/out projection of transformer blocks.
+_FP8_DEFAULT_SKIP_PATTERNS: tuple[str, ...] = (
+    "pos_embed",
+    "patch_embed",
+    "norm",
+    r"^proj_in$",
+    r"^proj_out$",
+)
+
+
+# TO DO: The loader is not thread safe!
+class ModelLoader(ModelLoaderBase):
+    """Default implementation of ModelLoaderBase."""
+
+    def __init__(
+        self,
+        app_config: InvokeAIAppConfig,
+        logger: Logger,
+        ram_cache: ModelCache,
+    ):
+        """Initialize the loader."""
+        self._app_config = app_config
+        self._logger = logger
+        self._ram_cache = ram_cache
+        self._torch_dtype = TorchDevice.choose_torch_dtype()
+        self._torch_device = TorchDevice.choose_torch_device()
+
+    def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
+        """
+        Return a model given its configuration.
+
+        Given a model's configuration as returned by the ModelRecordConfigStore service,
+        return a LoadedModel object that can be used for inference.
+
+        :param model config: Configuration record for this model
+        :param submodel_type: an ModelType enum indicating the portion of
+               the model to retrieve (e.g. ModelType.Vae)
+        """
+        model_path = self._get_model_path(model_config)
+
+        if not model_path.exists():
+            raise FileNotFoundError(f"Files for model '{model_config.name}' not found at {model_path}")
+
+        with skip_torch_weight_init():
+            cache_record = self._load_and_cache(model_config, submodel_type)
+        return LoadedModel(config=model_config, cache_record=cache_record, cache=self._ram_cache)
+
+    @property
+    def ram_cache(self) -> ModelCache:
+        """Return the ram cache associated with this loader."""
+        return self._ram_cache
+
+    def _get_model_path(self, config: AnyModelConfig) -> Path:
+        model_base = self._app_config.models_path
+        return (model_base / config.path).resolve()
+
+    def _get_execution_device(
+        self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> Optional[torch.device]:
+        """Determine the execution device for a model based on its configuration.
+
+        CPU-only execution is only applied to text encoder submodels to save VRAM while keeping
+        the denoiser on GPU for performance. Conditioning tensors are moved to GPU after encoding.
+
+        Returns:
+            torch.device("cpu") if the model should run on CPU only, None otherwise (use cache default).
+        """
+        # Check if this is a text encoder submodel of a main model with cpu_only setting
+        if hasattr(config, "default_settings") and config.default_settings is not None:
+            if hasattr(config.default_settings, "cpu_only") and config.default_settings.cpu_only is True:
+                # Only apply CPU execution to text encoder submodels
+                if submodel_type in [SubModelType.TextEncoder, SubModelType.TextEncoder2, SubModelType.TextEncoder3]:
+                    return torch.device("cpu")
+
+        # Check if this is a standalone text encoder config with cpu_only field (T5Encoder, Qwen3Encoder, etc.)
+        if hasattr(config, "cpu_only") and config.cpu_only is True:
+            return torch.device("cpu")
+
+        return None
+
+    def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> CacheRecord:
+        stats_name = ":".join([config.base, config.type, config.name, (submodel_type or "")])
+        try:
+            return self._ram_cache.get(key=get_model_cache_key(config.key, submodel_type), stats_name=stats_name)
+        except IndexError:
+            pass
+
+        config.path = str(self._get_model_path(config))
+        self._ram_cache.make_room(self.get_size_fs(config, Path(config.path), submodel_type))
+        loaded_model = self._load_model(config, submodel_type)
+
+        # Determine execution device from model config, considering submodel type
+        execution_device = self._get_execution_device(config, submodel_type)
+
+        self._ram_cache.put(
+            get_model_cache_key(config.key, submodel_type),
+            model=loaded_model,
+            execution_device=execution_device,
+        )
+
+        return self._ram_cache.get(key=get_model_cache_key(config.key, submodel_type), stats_name=stats_name)
+
+    def get_size_fs(
+        self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None
+    ) -> int:
+        """Get the size of the model on disk."""
+        return calc_model_size_by_fs(
+            model_path=model_path,
+            subfolder=submodel_type.value if submodel_type else None,
+            variant=config.repo_variant if isinstance(config, Diffusers_Config_Base) else None,
+        )
+
+    def _should_use_fp8(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> bool:
+        """Check if FP8 layerwise casting should be applied to a model."""
+        # FP8 storage only works on CUDA
+        if self._torch_device.type != "cuda":
+            return False
+
+        # Z-Image has dtype mismatch issues with diffusers' layerwise casting
+        # (skipped modules produce bf16, hooked modules expect fp16).
+        from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
+
+        if hasattr(config, "base") and config.base == BaseModelType.ZImage:
+            return False
+
+        # VAEs are excluded — fp8 storage causes noticeable quality degradation in decode.
+        if hasattr(config, "type") and config.type == ModelType.VAE:
+            return False
+
+        # LoRAs (including ControlLoRA) are excluded — they are not run as a standalone forward pass,
+        # they are patched into a base model, so the layerwise-casting hooks would never fire. The
+        # toggle is also hidden in the UI for ControlLoRA; this guard handles legacy persisted values.
+        if hasattr(config, "type") and config.type in (ModelType.LoRA, ModelType.ControlLoRa):
+            return False
+
+        # Don't apply FP8 to text encoders, tokenizers, schedulers, VAEs, etc.
+        _excluded_submodel_types = {
+            SubModelType.TextEncoder,
+            SubModelType.TextEncoder2,
+            SubModelType.TextEncoder3,
+            SubModelType.Tokenizer,
+            SubModelType.Tokenizer2,
+            SubModelType.Tokenizer3,
+            SubModelType.Scheduler,
+            SubModelType.SafetyChecker,
+            SubModelType.VAE,
+            SubModelType.VAEDecoder,
+            SubModelType.VAEEncoder,
+        }
+        if submodel_type in _excluded_submodel_types:
+            return False
+
+        # Check default_settings.fp8_storage (Main models, ControlNet)
+        if hasattr(config, "default_settings") and config.default_settings is not None:
+            if hasattr(config.default_settings, "fp8_storage") and config.default_settings.fp8_storage is True:
+                return True
+
+        return False
+
+    def _apply_fp8_layerwise_casting(
+        self, model: AnyModel, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None
+    ) -> AnyModel:
+        """Apply FP8 layerwise casting to a model if enabled in its config."""
+        if not self._should_use_fp8(config, submodel_type):
+            return model
+
+        storage_dtype = torch.float8_e4m3fn
+        compute_dtype = self._torch_dtype
+
+        # Detect the model's current dtype to use as compute dtype, since models
+        # (e.g. Flux) may require a specific dtype (bf16) that differs from the global torch dtype (fp16).
+        if isinstance(model, torch.nn.Module):
+            first_param = next(model.parameters(), None)
+            if first_param is not None:
+                compute_dtype = first_param.dtype
+
+        # We use our own hook-based path for every nn.Module — including diffusers ModelMixin —
+        # rather than `model.enable_layerwise_casting()`. Diffusers' LayerwiseCastingHook installs
+        # an instance-level `forward` attribute that captures the original `Linear.forward` in a
+        # closure. `ModelCache.put()` later runs `apply_custom_layers_to_model`, which constructs a
+        # new `CustomLinear` sharing the original Linear's `__dict__` — so the diffusers wrapper
+        # carries over and routes calls back to the captured original forward, silently bypassing
+        # `CustomLinear.forward` and its `cast_to_device` autocast. With partial loading (e.g. FLUX.2
+        # Klein 9B) some weights stay on CPU, the diffusers pre_forward only casts dtype, and
+        # `F.linear` then sees input on cuda and weight on cpu. Our `register_forward_pre_hook` /
+        # `register_forward_hook` path fires around `nn.Module._call_impl` without replacing
+        # `forward`, so `CustomLinear.forward` is still reached.
+        if isinstance(model, torch.nn.Module):
+            self._apply_fp8_to_nn_module(model, storage_dtype=storage_dtype, compute_dtype=compute_dtype)
+        else:
+            return model
+
+        param_bytes = sum(p.nelement() * p.element_size() for p in model.parameters())
+        self._logger.info(
+            f"FP8 layerwise casting enabled for {config.name} "
+            f"(storage=float8_e4m3fn, compute={compute_dtype}, "
+            f"param_size={param_bytes / (1024**2):.0f}MB)"
+        )
+        return model
+
+    @staticmethod
+    def _apply_fp8_to_nn_module(model: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype) -> None:
+        """Apply FP8 layerwise casting to a plain nn.Module.
+
+        Mirrors diffusers' `apply_layerwise_casting` semantics: only the layer classes in
+        `_FP8_SUPPORTED_PYTORCH_LAYERS` are cast, and modules whose dotted path matches any of
+        `_FP8_DEFAULT_SKIP_PATTERNS` (norm, pos_embed, patch_embed, proj_in/out) are skipped.
+        Without the skip list, precision-sensitive tiny learned scalars (e.g. FLUX RMSNorm.scale)
+        get crushed to FP8 and quality degrades noticeably.
+        """
+        for module_name, module in model.named_modules():
+            if not isinstance(module, _FP8_SUPPORTED_PYTORCH_LAYERS):
+                continue
+            if any(re.search(pattern, module_name) for pattern in _FP8_DEFAULT_SKIP_PATTERNS):
+                continue
+            params = list(module.parameters(recurse=False))
+            if not params:
+                continue
+
+            for param in params:
+                param.data = param.data.to(storage_dtype)
+
+            ModelLoader._wrap_forward_with_fp8_cast(module, storage_dtype, compute_dtype)
+
+    @staticmethod
+    def _wrap_forward_with_fp8_cast(
+        module: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype
+    ) -> None:
+        """Register pre/post forward hooks that cast params to compute dtype on entry and back
+        to storage dtype on exit.
+
+        We use hooks (rather than overriding `module.forward`) for two reasons:
+
+        1. **Correct dispatch after `apply_custom_layers_to_model`.** `ModelCache.put()` calls
+           `apply_custom_layers_to_model`, which creates a NEW `CustomLinear` instance and
+           shares the original `Linear.__dict__` (see `wrap_custom_layer`). Anything stored in
+           that dict — including an instance-level `forward` attribute — gets carried over to
+           the new object. An overridden `forward` would close over the OLD instance, so calls
+           to the new `CustomLinear` would silently route to `Linear.forward(old_instance, ...)`
+           and bypass the LoRA-patch-aware branch in `CustomLinear.forward`. Hooks, by contrast,
+           live in `_forward_hooks` / `_forward_pre_hooks` and are dispatched by
+           `nn.Module.__call__` with the *actual* called instance — so they run on the new
+           `CustomLinear` and the class's `forward` is still resolved normally.
+
+        2. **Exception safety.** `register_forward_hook(..., always_call=True)` fires the
+           post-hook even when `forward` raises. The plain pre-hook/post-hook pair without
+           `always_call` would leave params in compute dtype on exception, defeating FP8
+           storage savings and making cache size accounting stale.
+        """
+
+        def pre_hook(mod: torch.nn.Module, _args: object) -> None:
+            for p in mod.parameters(recurse=False):
+                p.data = p.data.to(compute_dtype)
+
+        def post_hook(mod: torch.nn.Module, _args: object, _output: object) -> None:
+            for p in mod.parameters(recurse=False):
+                p.data = p.data.to(storage_dtype)
+
+        module.register_forward_pre_hook(pre_hook)
+        module.register_forward_hook(post_hook, always_call=True)
+
+    # This needs to be implemented in the subclass
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        raise NotImplementedError

invokeai/backend/model_manager/load/memory_snapshot.py

@@ -0,0 +1,100 @@
+import gc
+from typing import Optional
+
+import psutil
+import torch
+from typing_extensions import Self
+
+from invokeai.backend.model_manager.util.libc_util import LibcUtil, Struct_mallinfo2
+
+GB = 2**30  # 1 GB
+
+
+class MemorySnapshot:
+    """A snapshot of RAM and VRAM usage. All values are in bytes."""
+
+    def __init__(self, process_ram: int, vram: Optional[int], malloc_info: Optional[Struct_mallinfo2]):
+        """Initialize a MemorySnapshot.
+
+        Most of the time, `MemorySnapshot` will be constructed with `MemorySnapshot.capture()`.
+
+        Args:
+            process_ram (int): CPU RAM used by the current process.
+            vram (Optional[int]): VRAM used by torch.
+            malloc_info (Optional[Struct_mallinfo2]): Malloc info obtained from LibcUtil.
+        """
+        self.process_ram = process_ram
+        self.vram = vram
+        self.malloc_info = malloc_info
+
+    @classmethod
+    def capture(cls, run_garbage_collector: bool = True) -> Self:
+        """Capture and return a MemorySnapshot.
+
+        Note: This function has significant overhead, particularly if `run_garbage_collector == True`.
+
+        Args:
+            run_garbage_collector (bool, optional): If true, gc.collect() will be run before checking the process RAM
+                usage. Defaults to True.
+
+        Returns:
+            MemorySnapshot
+        """
+        if run_garbage_collector:
+            gc.collect()
+
+        # According to the psutil docs (https://psutil.readthedocs.io/en/latest/#psutil.Process.memory_info), rss is
+        # supported on all platforms.
+        process_ram = psutil.Process().memory_info().rss
+
+        if torch.cuda.is_available():
+            vram = torch.cuda.memory_allocated()
+        else:
+            # TODO: We could add support for mps.current_allocated_memory() as well. Leaving out for now until we have
+            # time to test it properly.
+            vram = None
+
+        try:
+            malloc_info = LibcUtil().mallinfo2()
+        except (OSError, AttributeError):
+            # OSError: This is expected in environments that do not have the 'libc.so.6' shared library.
+            # AttributeError: This is expected in environments that have `libc.so.6` but do not have the `mallinfo2` (e.g. glibc < 2.33)
+            # TODO: Does `mallinfo` work?
+            malloc_info = None
+
+        return cls(process_ram, vram, malloc_info)
+
+
+def get_pretty_snapshot_diff(snapshot_1: Optional[MemorySnapshot], snapshot_2: Optional[MemorySnapshot]) -> str:
+    """Get a pretty string describing the difference between two `MemorySnapshot`s."""
+
+    def get_msg_line(prefix: str, val1: int, val2: int) -> str:
+        diff = val2 - val1
+        return f"{prefix: <30} ({(diff / GB):+5.3f}): {(val1 / GB):5.3f}GB -> {(val2 / GB):5.3f}GB\n"
+
+    msg = ""
+
+    if snapshot_1 is None or snapshot_2 is None:
+        return msg
+
+    msg += get_msg_line("Process RAM", snapshot_1.process_ram, snapshot_2.process_ram)
+
+    if snapshot_1.malloc_info is not None and snapshot_2.malloc_info is not None:
+        msg += get_msg_line("libc mmap allocated", snapshot_1.malloc_info.hblkhd, snapshot_2.malloc_info.hblkhd)
+
+        msg += get_msg_line("libc arena used", snapshot_1.malloc_info.uordblks, snapshot_2.malloc_info.uordblks)
+
+        msg += get_msg_line("libc arena free", snapshot_1.malloc_info.fordblks, snapshot_2.malloc_info.fordblks)
+
+        libc_total_allocated_1 = snapshot_1.malloc_info.arena + snapshot_1.malloc_info.hblkhd
+        libc_total_allocated_2 = snapshot_2.malloc_info.arena + snapshot_2.malloc_info.hblkhd
+        msg += get_msg_line("libc total allocated", libc_total_allocated_1, libc_total_allocated_2)
+
+        libc_total_used_1 = snapshot_1.malloc_info.uordblks + snapshot_1.malloc_info.hblkhd
+        libc_total_used_2 = snapshot_2.malloc_info.uordblks + snapshot_2.malloc_info.hblkhd
+        msg += get_msg_line("libc total used", libc_total_used_1, libc_total_used_2)
+
+    if snapshot_1.vram is not None and snapshot_2.vram is not None:
+        msg += get_msg_line("VRAM", snapshot_1.vram, snapshot_2.vram)
+
+    return msg

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")

invokeai/backend/model_manager/load/model_loader_registry.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Development team
+"""
+This module implements a system in which model loaders register the
+type, base and format of models that they know how to load.
+
+Use like this:
+
+  cls, model_config, submodel_type = ModelLoaderRegistry.get_implementation(model_config, submodel_type)  # type: ignore
+  loaded_model = cls(
+       app_config=app_config,
+       logger=logger,
+       ram_cache=ram_cache,
+       convert_cache=convert_cache
+    ).load_model(model_config, submodel_type)
+
+"""
+
+from abc import ABC, abstractmethod
+from typing import Callable, Dict, Optional, Tuple, Type, TypeVar
+
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load import ModelLoaderBase
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType
+
+
+class ModelLoaderRegistryBase(ABC):
+    """This class allows model loaders to register their type, base and format."""
+
+    @classmethod
+    @abstractmethod
+    def register(
+        cls, type: ModelType, format: ModelFormat, base: BaseModelType = BaseModelType.Any
+    ) -> Callable[[Type[ModelLoaderBase]], Type[ModelLoaderBase]]:
+        """Define a decorator which registers the subclass of loader."""
+
+    @classmethod
+    @abstractmethod
+    def get_implementation(
+        cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]
+    ) -> Tuple[Type[ModelLoaderBase], Config_Base, Optional[SubModelType]]:
+        """
+        Get subclass of ModelLoaderBase registered to handle base and type.
+
+        Parameters:
+        :param config: Model configuration record, as returned by ModelRecordService
+        :param submodel_type: Submodel to fetch (main models only)
+        :return: tuple(loader_class, model_config, submodel_type)
+
+        Note that the returned model config may be different from one what passed
+        in, in the event that a submodel type is provided.
+        """
+
+
+TModelLoader = TypeVar("TModelLoader", bound=ModelLoaderBase)
+
+
+class ModelLoaderRegistry(ModelLoaderRegistryBase):
+    """
+    This class allows model loaders to register their type, base and format.
+    """
+
+    _registry: Dict[str, Type[ModelLoaderBase]] = {}
+
+    @classmethod
+    def register(
+        cls, type: ModelType, format: ModelFormat, base: BaseModelType = BaseModelType.Any
+    ) -> Callable[[Type[TModelLoader]], Type[TModelLoader]]:
+        """Define a decorator which registers the subclass of loader."""
+
+        def decorator(subclass: Type[TModelLoader]) -> Type[TModelLoader]:
+            key = cls._to_registry_key(base, type, format)
+            if key in cls._registry:
+                raise Exception(
+                    f"{subclass.__name__} is trying to register as a loader for {base}/{type}/{format}, but this type of model has already been registered by {cls._registry[key].__name__}"
+                )
+            cls._registry[key] = subclass
+            return subclass
+
+        return decorator
+
+    @classmethod
+    def get_implementation(
+        cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]
+    ) -> Tuple[Type[ModelLoaderBase], Config_Base, Optional[SubModelType]]:
+        """Get subclass of ModelLoaderBase registered to handle base and type."""
+
+        key1 = cls._to_registry_key(config.base, config.type, config.format)  # for a specific base type
+        key2 = cls._to_registry_key(BaseModelType.Any, config.type, config.format)  # with wildcard Any
+        implementation = cls._registry.get(key1) or cls._registry.get(key2)
+        if not implementation:
+            raise NotImplementedError(
+                f"No subclass of LoadedModel is registered for base={config.base}, type={config.type}, format={config.format}"
+            )
+        return implementation, config, submodel_type
+
+    @staticmethod
+    def _to_registry_key(base: BaseModelType, type: ModelType, format: ModelFormat) -> str:
+        return "-".join([base.value, type.value, format.value])

invokeai/backend/model_manager/load/model_loaders/__init__.py

@@ -0,0 +1,3 @@
+"""
+Init file for model_loaders.
+"""

invokeai/backend/model_manager/load/model_loaders/anima.py

@@ -0,0 +1,212 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for Anima model loading in InvokeAI."""
+
+from pathlib import Path
+from typing import Optional
+
+import accelerate
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base
+from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_Anima_Config
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.configs.main import Main_Checkpoint_Anima_Config
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.logging import InvokeAILogger
+
+logger = InvokeAILogger.get_logger(__name__)
+
+
+def _strip_anima_bundle_prefix(sd: dict) -> dict:
+    """Strip the transformer-key prefix from an Anima single-file checkpoint.
+
+    Handles both packaging formats:
+      - Official format: keys prefixed with `net.` (e.g. `net.blocks.0...`)
+      - ComfyUI bundled format: transformer keys prefixed with `model.diffusion_model.`
+        alongside `first_stage_model.*` (VAE) and `cond_stage_model.*` (text encoder).
+
+    Only keys under the detected prefix are kept; unrelated keys from bundled
+    checkpoints (VAE, text encoder) are dropped. If no known prefix is present, the
+    state dict is returned unchanged.
+    """
+    prefix_to_strip = None
+    for prefix in ["model.diffusion_model.", "net."]:
+        if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)):
+            prefix_to_strip = prefix
+            break
+
+    if prefix_to_strip is None:
+        return sd
+
+    stripped_sd: dict = {}
+    for key, value in sd.items():
+        if isinstance(key, str) and key.startswith(prefix_to_strip):
+            stripped_sd[key[len(prefix_to_strip) :]] = value
+        # Skip non-transformer keys from bundled checkpoints (VAE, text encoder)
+    return stripped_sd
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Anima, type=ModelType.Main, format=ModelFormat.Checkpoint)
+class AnimaCheckpointModel(ModelLoader):
+    """Class to load Anima transformer models from single-file checkpoints.
+
+    The Anima checkpoint contains both the MiniTrainDIT backbone and the LLM Adapter
+    under a shared `net.` prefix. The loader strips this prefix and instantiates
+    the AnimaTransformer model with the correct architecture parameters.
+    """
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, Checkpoint_Config_Base):
+            raise ValueError("Only CheckpointConfigBase models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Transformer:
+                return self._load_from_singlefile(config)
+
+        raise ValueError(
+            f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    def _load_from_singlefile(
+        self,
+        config: AnyModelConfig,
+    ) -> AnyModel:
+        from safetensors.torch import load_file
+
+        from invokeai.backend.anima.anima_transformer import AnimaTransformer
+
+        if not isinstance(config, Main_Checkpoint_Anima_Config):
+            raise TypeError(
+                f"Expected Main_Checkpoint_Anima_Config, got {type(config).__name__}. "
+                "Model configuration type mismatch."
+            )
+        model_path = Path(config.path)
+
+        # Load the state dict from safetensors
+        sd = load_file(model_path)
+
+        # Strip the transformer-key prefix (`net.` or bundled `model.diffusion_model.`).
+        sd = _strip_anima_bundle_prefix(sd)
+
+        # Create an empty AnimaTransformer with Anima's default architecture parameters
+        with accelerate.init_empty_weights():
+            model = AnimaTransformer(
+                max_img_h=240,
+                max_img_w=240,
+                max_frames=1,
+                in_channels=16,
+                out_channels=16,
+                patch_spatial=2,
+                patch_temporal=1,
+                concat_padding_mask=True,
+                model_channels=2048,
+                num_blocks=28,
+                num_heads=16,
+                mlp_ratio=4.0,
+                crossattn_emb_channels=1024,
+                pos_emb_cls="rope3d",
+                # Anima reuses the Cosmos-Predict2 2B Text2Image DiT, which trains with
+                # rope_scale=(t=1.0, h=4.0, w=4.0). The NTK-scaled spatial RoPE base is
+                # mandatory; omitting it (theta=10000 on all axes) shifts every step's
+                # velocity ~7% off and compounds into degraded images. Matches diffusers
+                # CosmosTransformer3DModel rope_scale via *_extrapolation_ratio.
+                rope_h_extrapolation_ratio=4.0,
+                rope_w_extrapolation_ratio=4.0,
+                rope_t_extrapolation_ratio=1.0,
+                use_adaln_lora=True,
+                adaln_lora_dim=256,
+                extra_per_block_abs_pos_emb=False,
+                image_model="anima",
+            )
+
+        # Determine safe dtype
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_anima_inference_dtype(target_device)
+
+        # Handle memory management
+        new_sd_size = sum(ten.nelement() * model_dtype.itemsize for ten in sd.values())
+        self._ram_cache.make_room(new_sd_size)
+
+        # Convert to target dtype (skip non-float tensors like embedding indices)
+        for k in sd.keys():
+            if sd[k].is_floating_point():
+                sd[k] = sd[k].to(model_dtype)
+
+        # Filter out tensors that are regenerated at runtime and therefore not part of the
+        # in-memory module state. Some community-trained checkpoints (e.g. animaCatTower_v10)
+        # serialize derived pos_embedder buffers/cached tensors that the official model
+        # registers as non-persistent (or recomputes locally).
+        runtime_only_suffixes = (
+            ".inv_freq",
+            "pos_embedder.dim_spatial_range",
+            "pos_embedder.dim_temporal_range",
+            "pos_embedder.seq",
+        )
+        keys_to_remove = [k for k in sd.keys() if k.endswith(runtime_only_suffixes)]
+        for k in keys_to_remove:
+            del sd[k]
+
+        load_result = model.load_state_dict(sd, assign=True, strict=False)
+        if load_result.unexpected_keys:
+            raise RuntimeError(
+                f"Checkpoint contains {len(load_result.unexpected_keys)} unexpected keys. "
+                f"This may indicate a corrupted or incompatible checkpoint. "
+                f"First 5 unexpected keys: {load_result.unexpected_keys[:5]}"
+            )
+        if load_result.missing_keys:
+            logger.warning(
+                f"Checkpoint is missing {len(load_result.missing_keys)} keys "
+                f"(expected for inv_freq buffers). First 5: {load_result.missing_keys[:5]}"
+            )
+        return model
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Anima, type=ModelType.ControlNet, format=ModelFormat.Checkpoint)
+class AnimaControlNetLLLiteModel(ModelLoader):
+    """Class to load Anima ControlNet-LLLite adapter models from safetensors checkpoints.
+
+    LLLite adapters are standalone files holding a shared conditioning trunk
+    (lllite_conditioning1) plus tiny per-Linear modules (lllite_dit_blocks_*).
+    Hyperparameters are stored in the safetensors metadata (`lllite.*` keys) with
+    state-dict-shape fallbacks.
+    """
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        from safetensors import safe_open
+        from safetensors.torch import load_file
+
+        from invokeai.backend.anima.control_net_lllite import AnimaControlNetLLLite
+
+        if not isinstance(config, ControlNet_Checkpoint_Anima_Config):
+            raise ValueError("Only ControlNet_Checkpoint_Anima_Config models are supported here.")
+
+        # ControlNet type models don't use submodel_type - load the adapter directly
+        model_path = Path(config.path)
+
+        sd = load_file(model_path)
+        with safe_open(model_path, framework="pt", device="cpu") as f:
+            metadata = f.metadata()
+
+        model = AnimaControlNetLLLite.from_state_dict(sd, metadata)
+
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_anima_inference_dtype(target_device)
+        model.to(dtype=model_dtype)
+
+        return model

invokeai/backend/model_manager/load/model_loaders/clip_vision.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+from typing import Optional
+
+from transformers import CLIPVisionModelWithProjection
+
+from invokeai.backend.model_manager.configs.base import Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.CLIPVision, format=ModelFormat.Diffusers)
+class ClipVisionLoader(ModelLoader):
+    """Class to load CLIPVision models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, Diffusers_Config_Base):
+            raise ValueError("Only DiffusersConfigBase models are currently supported here.")
+
+        if submodel_type is not None:
+            raise Exception("There are no submodels in CLIP Vision models.")
+
+        model_path = Path(config.path)
+
+        model = CLIPVisionModelWithProjection.from_pretrained(
+            model_path, torch_dtype=self._torch_dtype, local_files_only=True
+        )
+        assert isinstance(model, CLIPVisionModelWithProjection)
+
+        return model

invokeai/backend/model_manager/load/model_loaders/cogview4.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.CogView4, type=ModelType.Main, format=ModelFormat.Diffusers)
+class CogView4DiffusersModel(GenericDiffusersLoader):
+    """Class to load CogView4 main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if isinstance(config, Checkpoint_Config_Base):
+            raise NotImplementedError("CheckpointConfigBase is not implemented for CogView4 models.")
+
+        if submodel_type is None:
+            raise Exception("A submodel type must be provided when loading main pipelines.")
+
+        model_path = Path(config.path)
+        load_class = self.get_hf_load_class(model_path, submodel_type)
+        repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None
+        variant = repo_variant.value if repo_variant else None
+        model_path = model_path / submodel_type.value
+
+        # We force bfloat16 for CogView4 models. It produces black images with float16. I haven't tracked down
+        # specifically which model(s) is/are responsible.
+        dtype = torch.bfloat16
+        try:
+            result: AnyModel = load_class.from_pretrained(
+                model_path,
+                torch_dtype=dtype,
+                variant=variant,
+                local_files_only=True,
+            )
+        except OSError as e:
+            if variant and "no file named" in str(
+                e
+            ):  # try without the variant, just in case user's preferences changed
+                result = load_class.from_pretrained(model_path, torch_dtype=dtype, local_files_only=True)
+            else:
+                raise e
+
+        result = self._apply_fp8_layerwise_casting(result, config, submodel_type)
+        return result

invokeai/backend/model_manager/load/model_loaders/controlnet.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for ControlNet model loading in InvokeAI."""
+
+from typing import Optional
+
+from diffusers import ControlNetModel
+
+from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+
+
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusion1, type=ModelType.ControlNet, format=ModelFormat.Diffusers
+)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusion1, type=ModelType.ControlNet, format=ModelFormat.Checkpoint
+)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusion2, type=ModelType.ControlNet, format=ModelFormat.Diffusers
+)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusion2, type=ModelType.ControlNet, format=ModelFormat.Checkpoint
+)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusionXL, type=ModelType.ControlNet, format=ModelFormat.Diffusers
+)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusionXL, type=ModelType.ControlNet, format=ModelFormat.Checkpoint
+)
+class ControlNetLoader(GenericDiffusersLoader):
+    """Class to load ControlNet models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if isinstance(config, ControlNet_Checkpoint_Config_Base):
+            result = ControlNetModel.from_single_file(
+                config.path,
+                torch_dtype=self._torch_dtype,
+            )
+            result = self._apply_fp8_layerwise_casting(result, config, submodel_type)
+            return result
+        else:
+            return super()._load_model(config, submodel_type)

invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py

@@ -0,0 +1,326 @@
+"""Pure state-dict conversion helpers for FLUX.2 single-file checkpoints.
+
+These functions translate the BFL (black-forest-labs) key layout used by FLUX.2
+single-file transformer and VAE checkpoints into the diffusers layout that the
+`Flux2Transformer2DModel` / `AutoencoderKLFlux2` architectures expect.
+
+They are intentionally free of any file/model-loading side effects so the key
+remapping can be unit-tested against a synthetic state dict (see
+`tests/backend/model_manager/load/`). Both the checkpoint and GGUF FLUX.2 loaders
+delegate here; GGUF quantized tensors are dequantized only where a fused weight
+must be split (diffusers uses separate Q/K/V projections).
+
+Based on the diffusers `convert_flux2_to_diffusers.py` key mappings.
+"""
+
+import re
+
+import torch
+
+
+def _flux2_chunk_tensor(tensor, chunks: int):
+    """Chunk a tensor along dim 0, dequantizing GGUF tensors first.
+
+    diffusers uses separate Q/K/V projections, so a fused GGUF weight cannot stay
+    quantized through the split.
+    """
+    if hasattr(tensor, "get_dequantized_tensor"):
+        tensor = tensor.get_dequantized_tensor()
+    return tensor.chunk(chunks, dim=0)
+
+
+def _flux2_malformed_for_chunk(tensor, chunks: int) -> bool:
+    """Return True if a plain tensor cannot be evenly split into `chunks` along dim 0.
+
+    GGUF quantized tensors are always considered well-formed here (their logical
+    shape is only known after dequantization, matching the original GGUF loader,
+    which split them unconditionally).
+    """
+    if hasattr(tensor, "get_dequantized_tensor"):
+        return False
+    return tensor.dim() < 1 or tensor.shape[0] % chunks != 0
+
+
+def _flux2_swap_scale_shift(weight):
+    """Swap the scale/shift halves of an AdaLayerNorm weight (BFL vs diffusers order)."""
+    if hasattr(weight, "get_dequantized_tensor"):
+        weight = weight.get_dequantized_tensor()
+    elif weight.dim() < 1 or weight.shape[0] % 2 != 0:
+        # Defensive: leave malformed plain tensors untouched.
+        return weight
+    shift, scale = weight.chunk(2, dim=0)
+    return torch.cat([scale, shift], dim=0)
+
+
+def _convert_flux2_double_block_key(key: str, tensor, converted: dict) -> str | None:
+    """Convert a `double_blocks.X.*` key to `transformer_blocks.X.*` format.
+
+    Returns the new key, or None if the key was consumed by writing directly into
+    `converted` (fused QKV split into separate projections).
+    """
+    parts = key.split(".")
+    block_idx = parts[1]
+    rest = ".".join(parts[2:])
+
+    prefix = f"transformer_blocks.{block_idx}"
+
+    # Attention QKV: BFL uses fused qkv, diffusers uses separate Q/K/V.
+    if "img_attn.qkv.weight" in rest:
+        if _flux2_malformed_for_chunk(tensor, 3):
+            return key
+        q, k, v = _flux2_chunk_tensor(tensor, 3)
+        converted[f"{prefix}.attn.to_q.weight"] = q
+        converted[f"{prefix}.attn.to_k.weight"] = k
+        converted[f"{prefix}.attn.to_v.weight"] = v
+        return None
+    elif "txt_attn.qkv.weight" in rest:
+        if _flux2_malformed_for_chunk(tensor, 3):
+            return key
+        q, k, v = _flux2_chunk_tensor(tensor, 3)
+        converted[f"{prefix}.attn.add_q_proj.weight"] = q
+        converted[f"{prefix}.attn.add_k_proj.weight"] = k
+        converted[f"{prefix}.attn.add_v_proj.weight"] = v
+        return None
+
+    # Attention output projection
+    if "img_attn.proj.weight" in rest:
+        return f"{prefix}.attn.to_out.0.weight"
+    elif "txt_attn.proj.weight" in rest:
+        return f"{prefix}.attn.to_add_out.weight"
+
+    # Attention norms
+    if "img_attn.norm.query_norm.scale" in rest or "img_attn.norm.query_norm.weight" in rest:
+        return f"{prefix}.attn.norm_q.weight"
+    elif "img_attn.norm.key_norm.scale" in rest or "img_attn.norm.key_norm.weight" in rest:
+        return f"{prefix}.attn.norm_k.weight"
+    elif "txt_attn.norm.query_norm.scale" in rest or "txt_attn.norm.query_norm.weight" in rest:
+        return f"{prefix}.attn.norm_added_q.weight"
+    elif "txt_attn.norm.key_norm.scale" in rest or "txt_attn.norm.key_norm.weight" in rest:
+        return f"{prefix}.attn.norm_added_k.weight"
+
+    # MLP layers
+    if "img_mlp.0.weight" in rest:
+        return f"{prefix}.ff.linear_in.weight"
+    elif "img_mlp.2.weight" in rest:
+        return f"{prefix}.ff.linear_out.weight"
+    elif "txt_mlp.0.weight" in rest:
+        return f"{prefix}.ff_context.linear_in.weight"
+    elif "txt_mlp.2.weight" in rest:
+        return f"{prefix}.ff_context.linear_out.weight"
+
+    return key
+
+
+def _convert_flux2_single_block_key(key: str, tensor, converted: dict) -> str | None:
+    """Convert a `single_blocks.X.*` key to `single_transformer_blocks.X.*` format."""
+    parts = key.split(".")
+    block_idx = parts[1]
+    rest = ".".join(parts[2:])
+
+    prefix = f"single_transformer_blocks.{block_idx}"
+
+    # linear1 is the fused QKV+MLP projection
+    if "linear1.weight" in rest:
+        return f"{prefix}.attn.to_qkv_mlp_proj.weight"
+    elif "linear2.weight" in rest:
+        return f"{prefix}.attn.to_out.weight"
+
+    # Norms
+    if "norm.query_norm.scale" in rest or "norm.query_norm.weight" in rest:
+        return f"{prefix}.attn.norm_q.weight"
+    elif "norm.key_norm.scale" in rest or "norm.key_norm.weight" in rest:
+        return f"{prefix}.attn.norm_k.weight"
+
+    return key
+
+
+def convert_flux2_bfl_to_diffusers(sd: dict) -> dict:
+    """Convert a FLUX.2 transformer BFL-format state dict to diffusers format."""
+    converted: dict = {}
+
+    # Basic key renames
+    key_renames = {
+        "img_in.weight": "x_embedder.weight",
+        "txt_in.weight": "context_embedder.weight",
+        "time_in.in_layer.weight": "time_guidance_embed.timestep_embedder.linear_1.weight",
+        "time_in.out_layer.weight": "time_guidance_embed.timestep_embedder.linear_2.weight",
+        "guidance_in.in_layer.weight": "time_guidance_embed.guidance_embedder.linear_1.weight",
+        "guidance_in.out_layer.weight": "time_guidance_embed.guidance_embedder.linear_2.weight",
+        "double_stream_modulation_img.lin.weight": "double_stream_modulation_img.linear.weight",
+        "double_stream_modulation_txt.lin.weight": "double_stream_modulation_txt.linear.weight",
+        "single_stream_modulation.lin.weight": "single_stream_modulation.linear.weight",
+        "final_layer.linear.weight": "proj_out.weight",
+        "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight",
+    }
+
+    for old_key, tensor in sd.items():
+        new_key = old_key
+
+        # Apply basic renames
+        if old_key in key_renames:
+            new_key = key_renames[old_key]
+            # Apply scale-shift swap for adaLN modulation weights (BFL and diffusers use
+            # different parameter ordering for AdaLayerNorm).
+            if old_key == "final_layer.adaLN_modulation.1.weight":
+                tensor = _flux2_swap_scale_shift(tensor)
+            converted[new_key] = tensor
+            continue
+
+        # Convert double_blocks.X.* to transformer_blocks.X.*
+        if old_key.startswith("double_blocks."):
+            new_key = _convert_flux2_double_block_key(old_key, tensor, converted)
+            if new_key is None:
+                continue  # Key was handled specially
+        # Convert single_blocks.X.* to single_transformer_blocks.X.*
+        elif old_key.startswith("single_blocks."):
+            new_key = _convert_flux2_single_block_key(old_key, tensor, converted)
+            if new_key is None:
+                continue  # Key was handled specially
+
+        if new_key != old_key or new_key not in converted:
+            converted[new_key] = tensor
+
+    return converted
+
+
+def _convert_flux2_vae_mid_attention_key(rest: str, tensor, block: str):
+    """Map a `{enc,dec}.mid.attn_1.*` key to the diffusers mid_block attention layout.
+
+    BFL uses Conv2d (shape `[out, in, 1, 1]`), diffusers uses Linear (`[out, in]`), so
+    weight tensors are squeezed. Returns `(new_key, tensor)`.
+    """
+    attn_prefix = f"{block}.mid_block.attentions.0"
+    if rest.startswith("q."):
+        new_key = f"{attn_prefix}.to_q.{rest[2:]}"
+    elif rest.startswith("k."):
+        new_key = f"{attn_prefix}.to_k.{rest[2:]}"
+    elif rest.startswith("v."):
+        new_key = f"{attn_prefix}.to_v.{rest[2:]}"
+    elif rest.startswith("proj_out."):
+        new_key = f"{attn_prefix}.to_out.0.{rest[9:]}"
+    elif rest.startswith("norm."):
+        return f"{attn_prefix}.group_norm.{rest[5:]}", tensor
+    else:
+        return f"{attn_prefix}.{rest}", tensor
+
+    if rest.endswith(".weight") and hasattr(tensor, "dim") and tensor.dim() == 4:
+        tensor = tensor.squeeze(-1).squeeze(-1)
+    return new_key, tensor
+
+
+def convert_flux2_vae_bfl_to_diffusers(sd: dict) -> dict:
+    """Convert a FLUX.2 VAE BFL-format state dict to diffusers format.
+
+    Key differences:
+    - encoder.down.X.block.Y -> encoder.down_blocks.X.resnets.Y
+    - encoder.down.X.downsample.conv -> encoder.down_blocks.X.downsamplers.0.conv
+    - encoder.mid.block_1/2 -> encoder.mid_block.resnets.0/1
+    - encoder.mid.attn_1.q/k/v -> encoder.mid_block.attentions.0.to_q/k/v
+    - encoder.norm_out -> encoder.conv_norm_out
+    - encoder.quant_conv -> quant_conv (top-level)
+    - decoder.up.X -> decoder.up_blocks.(num_blocks-1-X) (reversed order!)
+    - decoder.post_quant_conv -> post_quant_conv (top-level)
+    - *.nin_shortcut -> *.conv_shortcut
+    """
+    converted: dict = {}
+    num_up_blocks = 4  # Standard VAE has 4 up blocks
+
+    for old_key, tensor in sd.items():
+        new_key = old_key
+
+        # Encoder down blocks: encoder.down.X.block.Y -> encoder.down_blocks.X.resnets.Y
+        match = re.match(r"encoder\.down\.(\d+)\.block\.(\d+)\.(.*)", old_key)
+        if match:
+            block_idx, resnet_idx, rest = match.groups()
+            rest = rest.replace("nin_shortcut", "conv_shortcut")
+            new_key = f"encoder.down_blocks.{block_idx}.resnets.{resnet_idx}.{rest}"
+            converted[new_key] = tensor
+            continue
+
+        # Encoder downsamplers: encoder.down.X.downsample.conv -> encoder.down_blocks.X.downsamplers.0.conv
+        match = re.match(r"encoder\.down\.(\d+)\.downsample\.conv\.(.*)", old_key)
+        if match:
+            block_idx, rest = match.groups()
+            new_key = f"encoder.down_blocks.{block_idx}.downsamplers.0.conv.{rest}"
+            converted[new_key] = tensor
+            continue
+
+        # Encoder mid block resnets: encoder.mid.block_1/2 -> encoder.mid_block.resnets.0/1
+        match = re.match(r"encoder\.mid\.block_(\d+)\.(.*)", old_key)
+        if match:
+            block_num, rest = match.groups()
+            resnet_idx = int(block_num) - 1  # block_1 -> resnets.0, block_2 -> resnets.1
+            new_key = f"encoder.mid_block.resnets.{resnet_idx}.{rest}"
+            converted[new_key] = tensor
+            continue
+
+        # Encoder mid block attention: encoder.mid.attn_1.* -> encoder.mid_block.attentions.0.*
+        match = re.match(r"encoder\.mid\.attn_1\.(.*)", old_key)
+        if match:
+            new_key, tensor = _convert_flux2_vae_mid_attention_key(match.group(1), tensor, "encoder")
+            converted[new_key] = tensor
+            continue
+
+        # Encoder norm_out -> conv_norm_out
+        if old_key.startswith("encoder.norm_out."):
+            new_key = old_key.replace("encoder.norm_out.", "encoder.conv_norm_out.")
+            converted[new_key] = tensor
+            continue
+
+        # Encoder quant_conv -> quant_conv (move to top level)
+        if old_key.startswith("encoder.quant_conv."):
+            new_key = old_key.replace("encoder.quant_conv.", "quant_conv.")
+            converted[new_key] = tensor
+            continue
+
+        # Decoder up blocks (reversed order!): decoder.up.X -> decoder.up_blocks.(num_blocks-1-X)
+        match = re.match(r"decoder\.up\.(\d+)\.block\.(\d+)\.(.*)", old_key)
+        if match:
+            block_idx, resnet_idx, rest = match.groups()
+            new_block_idx = num_up_blocks - 1 - int(block_idx)
+            rest = rest.replace("nin_shortcut", "conv_shortcut")
+            new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.{rest}"
+            converted[new_key] = tensor
+            continue
+
+        # Decoder upsamplers (reversed order!)
+        match = re.match(r"decoder\.up\.(\d+)\.upsample\.conv\.(.*)", old_key)
+        if match:
+            block_idx, rest = match.groups()
+            new_block_idx = num_up_blocks - 1 - int(block_idx)
+            new_key = f"decoder.up_blocks.{new_block_idx}.upsamplers.0.conv.{rest}"
+            converted[new_key] = tensor
+            continue
+
+        # Decoder mid block resnets: decoder.mid.block_1/2 -> decoder.mid_block.resnets.0/1
+        match = re.match(r"decoder\.mid\.block_(\d+)\.(.*)", old_key)
+        if match:
+            block_num, rest = match.groups()
+            resnet_idx = int(block_num) - 1
+            new_key = f"decoder.mid_block.resnets.{resnet_idx}.{rest}"
+            converted[new_key] = tensor
+            continue
+
+        # Decoder mid block attention: decoder.mid.attn_1.* -> decoder.mid_block.attentions.0.*
+        match = re.match(r"decoder\.mid\.attn_1\.(.*)", old_key)
+        if match:
+            new_key, tensor = _convert_flux2_vae_mid_attention_key(match.group(1), tensor, "decoder")
+            converted[new_key] = tensor
+            continue
+
+        # Decoder norm_out -> conv_norm_out
+        if old_key.startswith("decoder.norm_out."):
+            new_key = old_key.replace("decoder.norm_out.", "decoder.conv_norm_out.")
+            converted[new_key] = tensor
+            continue
+
+        # Decoder post_quant_conv -> post_quant_conv (move to top level)
+        if old_key.startswith("decoder.post_quant_conv."):
+            new_key = old_key.replace("decoder.post_quant_conv.", "post_quant_conv.")
+            converted[new_key] = tensor
+            continue
+
+        # Keep other keys as-is (like encoder.conv_in, decoder.conv_in, decoder.conv_out, bn.*)
+        converted[new_key] = tensor
+
+    return converted

invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py

@@ -0,0 +1,104 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for simple diffusers model loading in InvokeAI."""
+
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+from invokeai.backend.model_manager.configs.base import Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.T2IAdapter, format=ModelFormat.Diffusers)
+class GenericDiffusersLoader(ModelLoader):
+    """Class to load simple diffusers models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        model_path = Path(config.path)
+        model_class = self.get_hf_load_class(model_path)
+        if submodel_type is not None:
+            raise Exception(f"There are no submodels in models of type {model_class}")
+        repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None
+        variant = repo_variant.value if repo_variant else None
+        try:
+            result: AnyModel = model_class.from_pretrained(
+                model_path, torch_dtype=self._torch_dtype, variant=variant, local_files_only=True
+            )
+        except OSError as e:
+            if variant and "no file named" in str(
+                e
+            ):  # try without the variant, just in case user's preferences changed
+                result = model_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, local_files_only=True)
+            else:
+                raise e
+        result = self._apply_fp8_layerwise_casting(result, config, submodel_type)
+        return result
+
+    # TO DO: Add exception handling
+    def get_hf_load_class(self, model_path: Path, submodel_type: Optional[SubModelType] = None) -> ModelMixin:
+        """Given the model path and submodel, returns the diffusers ModelMixin subclass needed to load."""
+        result = None
+        if submodel_type:
+            try:
+                config = self._load_diffusers_config(model_path, config_name="model_index.json")
+                module, class_name = config[submodel_type.value]
+                result = self._hf_definition_to_type(module=module, class_name=class_name)
+            except KeyError as e:
+                raise ValueError(f'The "{submodel_type}" submodel is not available for this model.') from e
+        else:
+            try:
+                config = self._load_diffusers_config(model_path, config_name="config.json")
+                if class_name := config.get("_class_name"):
+                    result = self._hf_definition_to_type(module="diffusers", class_name=class_name)
+                elif class_name := config.get("architectures"):
+                    result = self._hf_definition_to_type(module="transformers", class_name=class_name[0])
+                else:
+                    raise RuntimeError("Unable to decipher Load Class based on given config.json")
+            except KeyError as e:
+                raise ValueError("An expected config.json file is missing from this model.") from e
+        assert result is not None
+        return result
+
+    # TO DO: Add exception handling
+    def _hf_definition_to_type(self, module: str, class_name: str) -> ModelMixin:  # fix with correct type
+        if module in [
+            "diffusers",
+            "transformers",
+            "invokeai.backend.quantization.fast_quantized_transformers_model",
+            "invokeai.backend.quantization.fast_quantized_diffusion_model",
+        ]:
+            res_type = sys.modules[module]
+        else:
+            res_type = sys.modules["diffusers"].pipelines
+        result: ModelMixin = getattr(res_type, class_name)
+        return result
+
+    def _load_diffusers_config(self, model_path: Path, config_name: str = "config.json") -> dict[str, Any]:
+        return ConfigLoader.load_config(model_path, config_name=config_name)
+
+
+class ConfigLoader(ConfigMixin):
+    """Subclass of ConfigMixin for loading diffusers configuration files."""
+
+    @classmethod
+    def load_config(cls, *args: Any, **kwargs: Any) -> dict[str, Any]:  # pyright: ignore [reportIncompatibleMethodOverride]
+        """Load a diffusrs ConfigMixin configuration."""
+        cls.config_name = kwargs.pop("config_name")
+        # TODO(psyche): the types on this diffusers method are not correct
+        return super().load_config(*args, **kwargs)  # type: ignore

invokeai/backend/model_manager/load/model_loaders/ip_adapter.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for IP Adapter model loading in InvokeAI."""
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from invokeai.backend.ip_adapter.ip_adapter import build_ip_adapter
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load import ModelLoader, ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
+from invokeai.backend.raw_model import RawModel
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.IPAdapter, format=ModelFormat.InvokeAI)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.IPAdapter, format=ModelFormat.Checkpoint)
+class IPAdapterInvokeAILoader(ModelLoader):
+    """Class to load IP Adapter diffusers models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("There are no submodels in an IP-Adapter model.")
+        model_path = Path(config.path)
+        model: RawModel = build_ip_adapter(
+            ip_adapter_ckpt_path=model_path,
+            device=torch.device("cpu"),
+            dtype=self._torch_dtype,
+        )
+        return model

invokeai/backend/model_manager/load/model_loaders/llava_onevision.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+from typing import Optional
+
+from transformers import LlavaOnevisionForConditionalGeneration
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LlavaOnevision, format=ModelFormat.Diffusers)
+class LlavaOnevisionModelLoader(ModelLoader):
+    """Class for loading LLaVA Onevision VLLM models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("Unexpected submodel requested for LLaVA OneVision model.")
+
+        model_path = Path(config.path)
+        model = LlavaOnevisionForConditionalGeneration.from_pretrained(
+            model_path, local_files_only=True, torch_dtype=self._torch_dtype
+        )
+        assert isinstance(model, LlavaOnevisionForConditionalGeneration)
+        return model

invokeai/backend/model_manager/load/model_loaders/lora.py

@@ -0,0 +1,198 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for LoRA model loading in InvokeAI."""
+
+from logging import Logger
+from pathlib import Path
+from typing import Optional
+
+import torch
+from safetensors.torch import load_file
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.omi.omi import convert_from_omi
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+from invokeai.backend.patches.lora_conversions.anima_lora_conversion_utils import lora_model_from_anima_state_dict
+from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_aitoolkit_format,
+    lora_model_from_flux_aitoolkit_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_bfl_peft_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_bfl_peft_format,
+    lora_model_from_flux2_bfl_peft_state_dict,
+    lora_model_from_flux_bfl_peft_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import (
+    is_state_dict_likely_flux_control,
+    lora_model_from_flux_control_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import (
+    is_state_dict_flux2_diffusers_format,
+    is_state_dict_likely_in_flux_diffusers_format,
+    lora_model_from_flux2_diffusers_state_dict,
+    lora_model_from_flux_diffusers_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_kohya_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_kohya_format,
+    lora_model_from_flux_kohya_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_onetrainer_bfl_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_onetrainer_bfl_format,
+    lora_model_from_flux_onetrainer_bfl_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_onetrainer_format,
+    lora_model_from_flux_onetrainer_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.flux_xlabs_lora_conversion_utils import (
+    is_state_dict_likely_in_flux_xlabs_format,
+    lora_model_from_flux_xlabs_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.peft_adapter_utils import normalize_peft_adapter_names
+from invokeai.backend.patches.lora_conversions.qwen_image_lora_conversion_utils import (
+    lora_model_from_qwen_image_state_dict,
+)
+from invokeai.backend.patches.lora_conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict
+from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format
+from invokeai.backend.patches.lora_conversions.z_image_lora_conversion_utils import lora_model_from_z_image_state_dict
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.LoRA, format=ModelFormat.OMI)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.LoRA, format=ModelFormat.OMI)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LoRA, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.LoRA, format=ModelFormat.LyCORIS)
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.ControlLoRa, format=ModelFormat.LyCORIS)
+@ModelLoaderRegistry.register(base=BaseModelType.Flux, type=ModelType.ControlLoRa, format=ModelFormat.Diffusers)
+class LoRALoader(ModelLoader):
+    """Class to load LoRA models."""
+
+    # We cheat a little bit to get access to the model base
+    def __init__(
+        self,
+        app_config: InvokeAIAppConfig,
+        logger: Logger,
+        ram_cache: ModelCache,
+    ):
+        """Initialize the loader."""
+        super().__init__(app_config, logger, ram_cache)
+        self._model_base: Optional[BaseModelType] = None
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("There are no submodels in a LoRA model.")
+        model_path = Path(config.path)
+        assert self._model_base is not None
+
+        # Load the state dict from the model file.
+        if model_path.suffix == ".safetensors":
+            state_dict = load_file(model_path.absolute().as_posix(), device="cpu")
+        else:
+            state_dict = torch.load(model_path, map_location="cpu")
+
+        # Strip 'bundle_emb' keys - these are unused and currently cause downstream errors.
+        # To revisit later to determine if they're needed/useful.
+        state_dict = {k: v for k, v in state_dict.items() if not k.startswith("bundle_emb")}
+
+        # Normalize PEFT named-adapter keys (e.g. `lora_A.default.weight` → `lora_A.weight`)
+        # so the downstream format detectors and converters see canonical PEFT keys.
+        state_dict = normalize_peft_adapter_names(state_dict)
+
+        # At the time of writing, we support the OMI standard for base models Flux and SDXL
+        if config.format == ModelFormat.OMI and self._model_base in [
+            BaseModelType.StableDiffusionXL,
+            BaseModelType.Flux,
+        ]:
+            state_dict = convert_from_omi(state_dict, config.base)  # type: ignore
+
+        # Apply state_dict key conversions, if necessary.
+        if self._model_base == BaseModelType.StableDiffusionXL:
+            state_dict = convert_sdxl_keys_to_diffusers_format(state_dict)
+            model = lora_model_from_sd_state_dict(state_dict=state_dict)
+        elif self._model_base in (BaseModelType.Flux, BaseModelType.Flux2):
+            if config.format is ModelFormat.OMI:
+                # HACK(ryand): We set alpha=None for diffusers PEFT format models. These models are typically
+                # distributed as a single file without the associated metadata containing the alpha value. We chose
+                # alpha=None, because this is treated as alpha=rank internally in `LoRALayerBase.scale()`. alpha=rank
+                # is a popular choice. For example, in the diffusers training scripts:
+                # https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_flux.py#L1194
+                #
+                # We assume the same for LyCORIS models in diffusers key format.
+                model = lora_model_from_flux_diffusers_state_dict(state_dict=state_dict, alpha=None)
+            elif config.format is ModelFormat.LyCORIS:
+                if is_state_dict_likely_in_flux_diffusers_format(state_dict=state_dict):
+                    if is_state_dict_flux2_diffusers_format(state_dict=state_dict):
+                        # Flux2 Klein native diffusers naming (to_qkv_mlp_proj, ff.linear_in, etc.)
+                        model = lora_model_from_flux2_diffusers_state_dict(state_dict=state_dict, alpha=None)
+                    else:
+                        # Flux.1 diffusers naming (to_q/to_k/to_v, ff.net.0.proj, etc.)
+                        model = lora_model_from_flux_diffusers_state_dict(state_dict=state_dict, alpha=None)
+                elif is_state_dict_likely_in_flux_kohya_format(state_dict=state_dict):
+                    model = lora_model_from_flux_kohya_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_in_flux_onetrainer_bfl_format(state_dict=state_dict):
+                    model = lora_model_from_flux_onetrainer_bfl_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_in_flux_onetrainer_format(state_dict=state_dict):
+                    model = lora_model_from_flux_onetrainer_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_flux_control(state_dict=state_dict):
+                    model = lora_model_from_flux_control_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_in_flux_aitoolkit_format(state_dict=state_dict):
+                    model = lora_model_from_flux_aitoolkit_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_in_flux_xlabs_format(state_dict=state_dict):
+                    model = lora_model_from_flux_xlabs_state_dict(state_dict=state_dict)
+                elif is_state_dict_likely_in_flux_bfl_peft_format(state_dict=state_dict):
+                    if self._model_base == BaseModelType.Flux2:
+                        # FLUX.2 Klein uses Flux2Transformer2DModel (diffusers naming),
+                        # so we need to convert BFL keys to diffusers naming.
+                        model = lora_model_from_flux2_bfl_peft_state_dict(state_dict=state_dict, alpha=None)
+                    else:
+                        # FLUX.1 uses BFL Flux class, so BFL keys work directly.
+                        model = lora_model_from_flux_bfl_peft_state_dict(state_dict=state_dict, alpha=None)
+                else:
+                    raise ValueError("LoRA model is in unsupported FLUX format")
+            else:
+                raise ValueError(f"LoRA model is in unsupported FLUX format: {config.format}")
+        elif self._model_base in [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]:
+            # Currently, we don't apply any conversions for SD1 and SD2 LoRA models.
+            model = lora_model_from_sd_state_dict(state_dict=state_dict)
+        elif self._model_base == BaseModelType.ZImage:
+            # Z-Image LoRAs use diffusers PEFT format with transformer and/or Qwen3 encoder layers.
+            # We set alpha=None to use rank as alpha (common default).
+            model = lora_model_from_z_image_state_dict(state_dict=state_dict, alpha=None)
+        elif self._model_base == BaseModelType.QwenImage:
+            model = lora_model_from_qwen_image_state_dict(state_dict=state_dict, alpha=None)
+        elif self._model_base == BaseModelType.Anima:
+            # Anima LoRAs use Kohya-style or diffusers PEFT format targeting Cosmos DiT blocks.
+            model = lora_model_from_anima_state_dict(state_dict=state_dict, alpha=None)
+        else:
+            raise ValueError(f"Unsupported LoRA base model: {self._model_base}")
+
+        model.to(dtype=self._torch_dtype)
+        return model
+
+    def _get_model_path(self, config: AnyModelConfig) -> Path:
+        # cheating a little - we remember this variable for using in the subsequent call to _load_model()
+        self._model_base = config.base
+
+        model_base_path = self._app_config.models_path
+        model_path = model_base_path / config.path
+
+        if config.format == ModelFormat.Diffusers:
+            for ext in ["safetensors", "bin"]:  # return path to the safetensors file inside the folder
+                path = model_base_path / config.path / f"pytorch_lora_weights.{ext}"
+                if path.exists():
+                    model_path = path
+                    break
+
+        return model_path.resolve()

invokeai/backend/model_manager/load/model_loaders/onnx.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for Onnx model loading in InvokeAI."""
+
+# This should work the same as Stable Diffusion pipelines
+from pathlib import Path
+from typing import Optional
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.ONNX, format=ModelFormat.ONNX)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.ONNX, format=ModelFormat.Olive)
+class OnnyxDiffusersModel(GenericDiffusersLoader):
+    """Class to load onnx models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not submodel_type is not None:
+            raise Exception("A submodel type must be provided when loading onnx pipelines.")
+        model_path = Path(config.path)
+        load_class = self.get_hf_load_class(model_path, submodel_type)
+        repo_variant = getattr(config, "repo_variant", None)
+        variant = repo_variant.value if repo_variant else None
+        model_path = model_path / submodel_type.value
+        result: AnyModel = load_class.from_pretrained(
+            model_path,
+            torch_dtype=self._torch_dtype,
+            variant=variant,
+            local_files_only=True,
+        )
+        return result

invokeai/backend/model_manager/load/model_loaders/qwen_image.py

@@ -0,0 +1,556 @@
+from pathlib import Path
+from typing import Optional
+
+import accelerate
+import torch
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.configs.main import (
+    Main_Checkpoint_QwenImage_Config,
+    Main_GGUF_QwenImage_Config,
+)
+from invokeai.backend.model_manager.configs.qwen_vl_encoder import (
+    QwenVLEncoder_Checkpoint_Config,
+    QwenVLEncoder_Diffusers_Config,
+)
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    QwenImageVariantType,
+    SubModelType,
+)
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
+from invokeai.backend.util.devices import TorchDevice
+
+
+def _strip_comfyui_prefix(sd: dict) -> dict:
+    """Strip ComfyUI-style `model.diffusion_model.` / `diffusion_model.` prefixes from keys."""
+    prefix_to_strip = None
+    for prefix in ["model.diffusion_model.", "diffusion_model."]:
+        if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)):
+            prefix_to_strip = prefix
+            break
+    if prefix_to_strip is None:
+        return sd
+    stripped: dict = {}
+    for key, value in sd.items():
+        if isinstance(key, str) and key.startswith(prefix_to_strip):
+            stripped[key[len(prefix_to_strip) :]] = value
+        else:
+            stripped[key] = value
+    return stripped
+
+
+def _dequantize_comfyui_fp8(sd: dict, compute_dtype: torch.dtype) -> int:
+    """Dequantize ComfyUI-style fp8_scaled weights in-place. Returns count of dequantized tensors.
+
+    Weights are dequantized directly to `compute_dtype` (typically bf16) instead of via a
+    full-precision float32 intermediate. The previous float32 path materialised a complete
+    4-byte/param copy of the model before a separate downcast pass, spiking peak RAM to ~2x the
+    final bf16 size (~80GB for the 20B Qwen-Image transformer). Multiplying in the target dtype
+    keeps the dict at the bf16 model size plus a single transient tensor. fp8 has only 3 mantissa
+    bits and bf16 shares float32's exponent range, so the bf16 multiply loses no meaningful
+    precision here.
+
+    Two key naming schemes are in the wild:
+      - `<path>.weight` + `<path>.weight_scale`  (FLUX, Z-Image style)
+      - `<path>.weight` + `<path>.scale_weight`  (Qwen2.5-VL fp8_scaled style, also
+        emits `<path>.scale_input` for activation scaling that we discard).
+    """
+    scale_suffixes = (".weight_scale", ".scale_weight")
+    weight_scale_keys = [k for k in sd.keys() if isinstance(k, str) and k.endswith(scale_suffixes)]
+    count = 0
+    for scale_key in weight_scale_keys:
+        for suffix in scale_suffixes:
+            if scale_key.endswith(suffix):
+                weight_key = scale_key[: -len(suffix)] + ".weight"
+                break
+        if weight_key not in sd:
+            continue
+        weight = sd[weight_key].to(compute_dtype)
+        scale = sd[scale_key].to(compute_dtype)
+        if scale.shape != weight.shape and scale.numel() > 1:
+            for dim in range(len(weight.shape)):
+                if dim < len(scale.shape) and scale.shape[dim] != weight.shape[dim]:
+                    block_size = weight.shape[dim] // scale.shape[dim]
+                    if block_size > 1:
+                        scale = scale.repeat_interleave(block_size, dim=dim)
+        sd[weight_key] = weight * scale
+        count += 1
+    return count
+
+
+def _remap_qwen_vl_checkpoint_keys(sd: dict) -> dict:
+    """Remap legacy ComfyUI Qwen2.5-VL single-file keys to the transformers layout.
+
+    ComfyUI single-file checkpoints use the legacy Qwen2.5-VL key layout
+    (`visual.X`, `model.X`); transformers ≥4.50 expects `model.visual.X` and
+    `model.language_model.X`. This applies the same conversion mapping that
+    `Qwen2_5_VLForConditionalGeneration.from_pretrained` would, since
+    `load_state_dict` does not.
+
+    transformers ≤4.x exposed this as `_checkpoint_conversion_mapping`, but 5.x
+    dropped it (returns `{}`), so we fall back to the legacy mapping ourselves. The
+    negative lookahead keeps already-converted keys untouched, so the remap is safe
+    (and idempotent) for both legacy and new-layout single-file checkpoints.
+    """
+    import re
+
+    from transformers import Qwen2_5_VLForConditionalGeneration
+
+    key_mapping = Qwen2_5_VLForConditionalGeneration._checkpoint_conversion_mapping or {
+        r"^visual": "model.visual",
+        r"^model(?!\.(language_model|visual))": "model.language_model",
+    }
+    if not key_mapping:
+        return sd
+
+    remapped_sd: dict = {}
+    for old_key, tensor in sd.items():
+        new_key = old_key
+        if isinstance(old_key, str):
+            for pattern, replacement in key_mapping.items():
+                new_key, n_replace = re.subn(pattern, replacement, new_key)
+                if n_replace > 0:
+                    break
+        remapped_sd[new_key] = tensor
+    return remapped_sd
+
+
+def _strip_quantization_metadata(sd: dict) -> None:
+    """Strip ComfyUI fp8 quantization metadata keys in-place."""
+    keys_to_drop = [
+        k
+        for k in sd.keys()
+        if isinstance(k, str)
+        and (
+            k.endswith(".weight_scale")
+            or k.endswith(".scale_weight")
+            or k.endswith(".scale_input")
+            or "comfy_quant" in k
+            or k == "scaled_fp8"
+        )
+    ]
+    for k in keys_to_drop:
+        del sd[k]
+
+
+def _build_qwen_image_transformer_config(sd: dict, is_edit: bool) -> dict:
+    """Auto-detect Qwen Image transformer architecture parameters from the state dict.
+
+    Works for both GGUF (GGMLTensor) and plain safetensors (torch.Tensor) state dicts.
+    Mutates nothing.
+    """
+    from diffusers import QwenImageTransformer2DModel
+
+    def _shape(t):
+        return t.tensor_shape if isinstance(t, GGMLTensor) else t.shape
+
+    num_layers = 0
+    for key in sd.keys():
+        if isinstance(key, str) and key.startswith("transformer_blocks."):
+            parts = key.split(".")
+            if len(parts) >= 2:
+                try:
+                    num_layers = max(num_layers, int(parts[1]) + 1)
+                except ValueError:
+                    pass
+
+    num_attention_heads = 24
+    attention_head_dim = 128
+    in_channels = 64
+
+    if "img_in.weight" in sd:
+        shape = _shape(sd["img_in.weight"])
+        hidden_dim = shape[0]
+        in_channels = shape[1]
+        num_attention_heads = hidden_dim // attention_head_dim
+
+    joint_attention_dim = 3584
+    if "txt_in.weight" in sd:
+        joint_attention_dim = _shape(sd["txt_in.weight"])[1]
+
+    model_config: dict = {
+        "patch_size": 2,
+        "in_channels": in_channels,
+        "out_channels": 16,
+        "num_layers": num_layers if num_layers > 0 else 60,
+        "attention_head_dim": attention_head_dim,
+        "num_attention_heads": num_attention_heads,
+        "joint_attention_dim": joint_attention_dim,
+        "guidance_embeds": False,
+        "axes_dims_rope": (16, 56, 56),
+    }
+
+    # zero_cond_t enables dual modulation for noisy vs reference patches in edit-variant
+    # models. Setting it on txt2img models produces garbage. Requires diffusers 0.37+.
+    import inspect
+
+    if is_edit and "zero_cond_t" in inspect.signature(QwenImageTransformer2DModel.__init__).parameters:
+        model_config["zero_cond_t"] = True
+
+    return model_config
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.QwenImage, type=ModelType.Main, format=ModelFormat.Diffusers)
+class QwenImageDiffusersModel(GenericDiffusersLoader):
+    """Class to load Qwen Image Edit main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if isinstance(config, Checkpoint_Config_Base):
+            raise NotImplementedError("CheckpointConfigBase is not implemented for Qwen Image Edit models.")
+
+        if submodel_type is None:
+            raise Exception("A submodel type must be provided when loading main pipelines.")
+
+        model_path = Path(config.path)
+        load_class = self.get_hf_load_class(model_path, submodel_type)
+        repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None
+        variant = repo_variant.value if repo_variant else None
+        model_path = model_path / submodel_type.value
+
+        # We force bfloat16 for Qwen Image Edit models.
+        # Use `dtype` (newer) with fallback to `torch_dtype` (older diffusers).
+        dtype_kwarg = {"dtype": torch.bfloat16}
+        try:
+            result: AnyModel = load_class.from_pretrained(
+                model_path,
+                **dtype_kwarg,
+                variant=variant,
+                local_files_only=True,
+            )
+        except TypeError:
+            # Older diffusers uses torch_dtype instead of dtype
+            dtype_kwarg = {"torch_dtype": torch.bfloat16}
+            result = load_class.from_pretrained(
+                model_path,
+                **dtype_kwarg,
+                variant=variant,
+                local_files_only=True,
+            )
+        except OSError as e:
+            if variant and "no file named" in str(e):
+                result = load_class.from_pretrained(model_path, **dtype_kwarg, local_files_only=True)
+            else:
+                raise e
+
+        result = self._apply_fp8_layerwise_casting(result, config, submodel_type)
+        return result
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.QwenImage, type=ModelType.Main, format=ModelFormat.GGUFQuantized)
+class QwenImageGGUFCheckpointModel(ModelLoader):
+    """Class to load GGUF-quantized Qwen Image Edit transformer models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, Checkpoint_Config_Base):
+            raise ValueError("Only CheckpointConfigBase models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Transformer:
+                return self._load_from_singlefile(config)
+
+        raise ValueError(
+            f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel:
+        from diffusers import QwenImageTransformer2DModel
+
+        if not isinstance(config, Main_GGUF_QwenImage_Config):
+            raise TypeError(f"Expected Main_GGUF_QwenImage_Config, got {type(config).__name__}.")
+        model_path = Path(config.path)
+
+        target_device = TorchDevice.choose_torch_device()
+        compute_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+
+        sd = gguf_sd_loader(model_path, compute_dtype=compute_dtype)
+        sd = _strip_comfyui_prefix(sd)
+
+        is_edit = getattr(config, "variant", None) == QwenImageVariantType.Edit
+        model_config = _build_qwen_image_transformer_config(sd, is_edit=is_edit)
+
+        with accelerate.init_empty_weights():
+            model = QwenImageTransformer2DModel(**model_config)
+
+        model.load_state_dict(sd, strict=False, assign=True)
+        return model
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.QwenImage, type=ModelType.Main, format=ModelFormat.Checkpoint)
+class QwenImageCheckpointModel(ModelLoader):
+    """Loads Qwen Image transformer models from single-file safetensors checkpoints
+    (e.g. ComfyUI fp8_scaled, plain bf16/fp16). Dequantizes ComfyUI fp8 scaling to
+    bf16 at load time; the `default_settings.fp8_storage` toggle then optionally
+    re-casts to fp8 for VRAM savings."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, Checkpoint_Config_Base):
+            raise ValueError("Only CheckpointConfigBase models are currently supported here.")
+
+        match submodel_type:
+            case SubModelType.Transformer:
+                model = self._load_from_singlefile(config)
+                return self._apply_fp8_layerwise_casting(model, config, submodel_type)
+
+        raise ValueError(
+            f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel:
+        from diffusers import QwenImageTransformer2DModel
+        from safetensors.torch import load_file
+
+        from invokeai.backend.util.logging import InvokeAILogger
+
+        logger = InvokeAILogger.get_logger(self.__class__.__name__)
+
+        if not isinstance(config, Main_Checkpoint_QwenImage_Config):
+            raise TypeError(f"Expected Main_Checkpoint_QwenImage_Config, got {type(config).__name__}.")
+        model_path = Path(config.path)
+
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+
+        sd = load_file(str(model_path))
+        sd = _strip_comfyui_prefix(sd)
+
+        dequantized = _dequantize_comfyui_fp8(sd, model_dtype)
+        if dequantized > 0:
+            logger.info(f"Dequantized {dequantized} ComfyUI-quantized weights")
+        _strip_quantization_metadata(sd)
+
+        is_edit = getattr(config, "variant", None) == QwenImageVariantType.Edit
+        model_config = _build_qwen_image_transformer_config(sd, is_edit=is_edit)
+
+        with accelerate.init_empty_weights():
+            model = QwenImageTransformer2DModel(**model_config)
+
+        # Dequantized fp8 weights are already at model_dtype; this only casts any remaining
+        # non-quantized float weights (e.g. a plain fp16/fp32 checkpoint) to the compute dtype
+        # so the cache reservation below is sized from the actual post-cast tensors.
+        for k in list(sd.keys()):
+            if sd[k].is_floating_point():
+                sd[k] = sd[k].to(model_dtype)
+
+        new_sd_size = sum(t.nelement() * t.element_size() for t in sd.values())
+        self._ram_cache.make_room(new_sd_size)
+
+        model.load_state_dict(sd, strict=False, assign=True)
+        return model
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.QwenVLEncoder, format=ModelFormat.QwenVLEncoder)
+class QwenVLEncoderLoader(ModelLoader):
+    """Loads a standalone Qwen2.5-VL encoder (text_encoder/ + tokenizer/ + processor/)."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, QwenVLEncoder_Diffusers_Config):
+            raise TypeError(f"Expected QwenVLEncoder_Diffusers_Config, got {type(config).__name__}.")
+
+        from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration
+
+        model_path = Path(config.path)
+
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+
+        match submodel_type:
+            case SubModelType.Tokenizer:
+                tokenizer_path = model_path / "tokenizer"
+                return AutoTokenizer.from_pretrained(str(tokenizer_path), local_files_only=True)
+            case SubModelType.TextEncoder:
+                encoder_path = model_path / "text_encoder"
+                return Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                    str(encoder_path),
+                    torch_dtype=model_dtype,
+                    low_cpu_mem_usage=True,
+                    local_files_only=True,
+                )
+
+        raise ValueError(
+            f"Only Tokenizer and TextEncoder submodels are supported. "
+            f"Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.QwenVLEncoder, format=ModelFormat.Checkpoint)
+class QwenVLEncoderCheckpointLoader(ModelLoader):
+    """Loads a single-file Qwen2.5-VL encoder checkpoint (e.g. ComfyUI fp8_scaled).
+
+    The checkpoint bundles the language model and the visual tower into one
+    safetensors file. Tokenizer + processor are pulled from HuggingFace
+    (`Qwen/Qwen2.5-VL-7B-Instruct`) on first use, with offline cache fallback.
+    """
+
+    DEFAULT_HF_REPO = "Qwen/Qwen2.5-VL-7B-Instruct"
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if not isinstance(config, QwenVLEncoder_Checkpoint_Config):
+            raise TypeError(f"Expected QwenVLEncoder_Checkpoint_Config, got {type(config).__name__}.")
+
+        match submodel_type:
+            case SubModelType.Tokenizer:
+                return self._load_tokenizer_with_offline_fallback()
+            case SubModelType.TextEncoder:
+                return self._load_text_encoder_from_singlefile(config)
+
+        raise ValueError(
+            f"Only Tokenizer and TextEncoder submodels are supported. "
+            f"Received: {submodel_type.value if submodel_type else 'None'}"
+        )
+
+    def _load_tokenizer_with_offline_fallback(self) -> AnyModel:
+        from transformers import AutoTokenizer
+
+        from invokeai.backend.util.logging import InvokeAILogger
+
+        logger = InvokeAILogger.get_logger(self.__class__.__name__)
+
+        try:
+            return AutoTokenizer.from_pretrained(self.DEFAULT_HF_REPO, local_files_only=True)
+        except OSError:
+            logger.info(
+                f"Tokenizer for single-file Qwen VL encoder not found in HuggingFace cache; "
+                f"downloading from {self.DEFAULT_HF_REPO} (one-time, requires network access)."
+            )
+            try:
+                return AutoTokenizer.from_pretrained(self.DEFAULT_HF_REPO)
+            except OSError as e:
+                raise RuntimeError(
+                    f"Failed to load Qwen VL tokenizer. Single-file Qwen VL encoder checkpoints do not "
+                    f"include the tokenizer; it must be downloaded from HuggingFace ({self.DEFAULT_HF_REPO}) "
+                    f"on first use. Either restore network access, or install the encoder in the "
+                    f"diffusers folder layout (text_encoder/ + tokenizer/) instead. Original error: {e}"
+                ) from e
+
+    def _load_text_encoder_from_singlefile(self, config: QwenVLEncoder_Checkpoint_Config) -> AnyModel:
+        from safetensors.torch import load_file
+        from transformers import AutoConfig, Qwen2_5_VLForConditionalGeneration
+
+        from invokeai.backend.util.logging import InvokeAILogger
+
+        logger = InvokeAILogger.get_logger(self.__class__.__name__)
+
+        model_path = Path(config.path)
+
+        target_device = TorchDevice.choose_torch_device()
+        model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device)
+
+        sd = load_file(str(model_path))
+
+        # Dequantize ComfyUI-style fp8 weights, then strip the now-unused quantization
+        # metadata (`scale_input` is the activation scale ComfyUI's fp8 matmul kernels
+        # use at runtime — we run the encoder in bf16 after dequantization).
+        dequantized_count = _dequantize_comfyui_fp8(sd, model_dtype)
+        if dequantized_count > 0:
+            logger.info(f"Dequantized {dequantized_count} ComfyUI-quantized weights")
+        _strip_quantization_metadata(sd)
+
+        # ComfyUI single-file checkpoints use the legacy Qwen2.5-VL key layout
+        # (`visual.X`, `model.X`); remap to the `model.visual.X` / `model.language_model.X`
+        # layout transformers expects. See `_remap_qwen_vl_checkpoint_keys` for details.
+        sd = _remap_qwen_vl_checkpoint_keys(sd)
+
+        # Cast to compute dtype (skip integer/index tensors)
+        for k in list(sd.keys()):
+            if sd[k].is_floating_point():
+                sd[k] = sd[k].to(model_dtype)
+
+        # Fetch the architecture config from HuggingFace (small, ~5KB).
+        # Offline fallback: tries cache first, downloads only if missing.
+        try:
+            qwen_config = AutoConfig.from_pretrained(self.DEFAULT_HF_REPO, local_files_only=True)
+        except OSError:
+            logger.info(
+                f"Architecture config for single-file Qwen VL encoder not found in HuggingFace cache; "
+                f"downloading from {self.DEFAULT_HF_REPO} (one-time, ~5KB, requires network access)."
+            )
+            try:
+                qwen_config = AutoConfig.from_pretrained(self.DEFAULT_HF_REPO)
+            except OSError as e:
+                raise RuntimeError(
+                    f"Failed to load Qwen VL architecture config. Single-file Qwen VL encoder checkpoints "
+                    f"do not include the model config; it must be downloaded from HuggingFace "
+                    f"({self.DEFAULT_HF_REPO}) on first use. Either restore network access, or install the "
+                    f"encoder in the diffusers folder layout (text_encoder/config.json + tokenizer/) "
+                    f"instead. Original error: {e}"
+                ) from e
+        qwen_config.torch_dtype = model_dtype
+
+        new_sd_size = sum(t.nelement() * t.element_size() for t in sd.values())
+        self._ram_cache.make_room(new_sd_size)
+
+        with accelerate.init_empty_weights():
+            model = Qwen2_5_VLForConditionalGeneration(qwen_config)
+
+        # Load weights; allow missing keys for tied lm_head and re-initialised buffers.
+        load_result = model.load_state_dict(sd, strict=False, assign=True)
+        if load_result.unexpected_keys:
+            logger.warning(
+                f"{len(load_result.unexpected_keys)} unexpected keys in checkpoint, "
+                f"first 5: {load_result.unexpected_keys[:5]}"
+            )
+
+        # Tie lm_head ↔ embed_tokens if config requires it and lm_head wasn't loaded
+        if getattr(qwen_config, "tie_word_embeddings", False):
+            try:
+                if hasattr(model, "lm_head") and model.lm_head.weight.is_meta:
+                    model.lm_head.weight = model.model.embed_tokens.weight
+                else:
+                    model.tie_weights()
+            except AttributeError:
+                model.tie_weights()
+
+        # Re-initialise any leftover meta buffers (RoPE inv_freq etc.)
+        for name, buffer in list(model.named_buffers()):
+            if not buffer.is_meta:
+                continue
+            parts = name.rsplit(".", 1)
+            if len(parts) == 2:
+                parent = model.get_submodule(parts[0])
+                buffer_name = parts[1]
+            else:
+                parent = model
+                buffer_name = name
+            # Replace meta buffer with a real (zero) tensor of the same shape; the model
+            # will recompute or refill these as needed at first forward pass.
+            try:
+                shape = buffer.shape
+                parent.register_buffer(buffer_name, torch.zeros(shape, dtype=model_dtype), persistent=False)
+            except Exception:
+                logger.warning(f"Could not re-initialise meta buffer {name}")
+
+        meta_params = [name for name, p in model.named_parameters() if p.is_meta]
+        if meta_params:
+            raise RuntimeError(f"Failed to load all parameters from checkpoint. Meta tensors remain: {meta_params[:5]}")
+
+        model.eval()
+        return model

invokeai/backend/model_manager/load/model_loaders/sig_lip.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+from typing import Optional
+
+from transformers import SiglipVisionModel
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.SigLIP, format=ModelFormat.Diffusers)
+class SigLIPModelLoader(ModelLoader):
+    """Class for loading SigLIP models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("Unexpected submodel requested for LLaVA OneVision model.")
+
+        model_path = Path(config.path)
+        model = SiglipVisionModel.from_pretrained(model_path, local_files_only=True, torch_dtype=self._torch_dtype)
+        return model

invokeai/backend/model_manager/load/model_loaders/spandrel_image_to_image.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
+from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
+
+
+@ModelLoaderRegistry.register(
+    base=BaseModelType.Any, type=ModelType.SpandrelImageToImage, format=ModelFormat.Checkpoint
+)
+class SpandrelImageToImageModelLoader(ModelLoader):
+    """Class for loading Spandrel Image-to-Image models (i.e. models wrapped by spandrel.ImageModelDescriptor)."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("Unexpected submodel requested for Spandrel model.")
+
+        model_path = Path(config.path)
+        model = SpandrelImageToImageModel.load_from_file(model_path)
+
+        torch_dtype = self._torch_dtype
+        if not model.supports_dtype(torch_dtype):
+            self._logger.warning(
+                f"The configured dtype ('{self._torch_dtype}') is not supported by the {model.get_model_type_name()} "
+                "model. Falling back to 'float32'."
+            )
+            torch_dtype = torch.float32
+        model.to(dtype=torch_dtype)
+
+        return model

invokeai/backend/model_manager/load/model_loaders/stable_diffusion.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for StableDiffusion model loading in InvokeAI."""
+
+from pathlib import Path
+from typing import Optional
+
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import StableDiffusionXLPipeline
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl_inpaint import (
+    StableDiffusionXLInpaintPipeline,
+)
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.configs.main import (
+    Main_Checkpoint_SD1_Config,
+    Main_Checkpoint_SD2_Config,
+    Main_Checkpoint_SDXL_Config,
+    Main_Checkpoint_SDXLRefiner_Config,
+    Main_Diffusers_SD1_Config,
+    Main_Diffusers_SD2_Config,
+    Main_Diffusers_SDXL_Config,
+    Main_Diffusers_SDXLRefiner_Config,
+)
+from invokeai.backend.model_manager.load.model_cache.model_cache import get_model_cache_key
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    ModelVariantType,
+    SubModelType,
+)
+from invokeai.backend.util.silence_warnings import SilenceWarnings
+
+VARIANT_TO_IN_CHANNEL_MAP = {
+    ModelVariantType.Normal: 4,
+    ModelVariantType.Depth: 5,
+    ModelVariantType.Inpaint: 9,
+}
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion2, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusionXLRefiner, type=ModelType.Main, format=ModelFormat.Diffusers
+)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion3, type=ModelType.Main, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion1, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusion2, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(base=BaseModelType.StableDiffusionXL, type=ModelType.Main, format=ModelFormat.Checkpoint)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.StableDiffusionXLRefiner, type=ModelType.Main, format=ModelFormat.Checkpoint
+)
+class StableDiffusionDiffusersModel(GenericDiffusersLoader):
+    """Class to load main models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if isinstance(config, Checkpoint_Config_Base):
+            return self._load_from_singlefile(config, submodel_type)
+
+        if submodel_type is None:
+            raise Exception("A submodel type must be provided when loading main pipelines.")
+
+        model_path = Path(config.path)
+        load_class = self.get_hf_load_class(model_path, submodel_type)
+        repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None
+        variant = repo_variant.value if repo_variant else None
+        model_path = model_path / submodel_type.value
+        try:
+            result: AnyModel = load_class.from_pretrained(
+                model_path,
+                torch_dtype=self._torch_dtype,
+                variant=variant,
+                local_files_only=True,
+            )
+        except OSError as e:
+            if variant and "no file named" in str(
+                e
+            ):  # try without the variant, just in case user's preferences changed
+                result = load_class.from_pretrained(model_path, torch_dtype=self._torch_dtype, local_files_only=True)
+            else:
+                raise e
+
+        result = self._apply_fp8_layerwise_casting(result, config, submodel_type)
+        return result
+
+    def _load_from_singlefile(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        load_classes = {
+            BaseModelType.StableDiffusion1: {
+                ModelVariantType.Normal: StableDiffusionPipeline,
+                ModelVariantType.Inpaint: StableDiffusionInpaintPipeline,
+            },
+            BaseModelType.StableDiffusion2: {
+                ModelVariantType.Normal: StableDiffusionPipeline,
+                ModelVariantType.Inpaint: StableDiffusionInpaintPipeline,
+            },
+            BaseModelType.StableDiffusionXL: {
+                ModelVariantType.Normal: StableDiffusionXLPipeline,
+                ModelVariantType.Inpaint: StableDiffusionXLInpaintPipeline,
+            },
+            BaseModelType.StableDiffusionXLRefiner: {
+                ModelVariantType.Normal: StableDiffusionXLPipeline,
+            },
+        }
+        assert isinstance(
+            config,
+            (
+                Main_Diffusers_SD1_Config,
+                Main_Diffusers_SD2_Config,
+                Main_Diffusers_SDXL_Config,
+                Main_Diffusers_SDXLRefiner_Config,
+                Main_Checkpoint_SD1_Config,
+                Main_Checkpoint_SD2_Config,
+                Main_Checkpoint_SDXL_Config,
+                Main_Checkpoint_SDXLRefiner_Config,
+            ),
+        )
+        try:
+            load_class = load_classes[config.base][config.variant]
+        except KeyError as e:
+            raise Exception(f"No diffusers pipeline known for base={config.base}, variant={config.variant}") from e
+
+        # Without SilenceWarnings we get log messages like this:
+        # site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.
+        # warnings.warn(
+        # Some weights of the model checkpoint were not used when initializing CLIPTextModel:
+        # ['text_model.embeddings.position_ids']
+        # Some weights of the model checkpoint were not used when initializing CLIPTextModelWithProjection:
+        # ['text_model.embeddings.position_ids']
+
+        with SilenceWarnings():
+            pipeline = load_class.from_single_file(config.path, torch_dtype=self._torch_dtype)
+
+        if not submodel_type:
+            return pipeline
+
+        # Proactively load the various submodels into the RAM cache so that we don't have to re-load
+        # the entire pipeline every time a new submodel is needed.
+        for subtype in SubModelType:
+            if subtype == submodel_type:
+                continue
+            if submodel := getattr(pipeline, subtype.value, None):
+                self._apply_fp8_layerwise_casting(submodel, config, subtype)
+                self._ram_cache.put(get_model_cache_key(config.key, subtype), model=submodel)
+        result = getattr(pipeline, submodel_type.value)
+        result = self._apply_fp8_layerwise_casting(result, config, submodel_type)
+        return result

invokeai/backend/model_manager/load/model_loaders/text_llm.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+from typing import Optional
+
+import torch
+from transformers import AutoModelForCausalLM
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.TextLLM, format=ModelFormat.Diffusers)
+class TextLLMModelLoader(ModelLoader):
+    """Class for loading text causal language models (Llama, Phi, Qwen, Mistral, etc.)."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("Unexpected submodel requested for TextLLM model.")
+
+        # Use float32 for CPU-only models since CPU fp16 is emulated and slow.
+        dtype = self._torch_dtype
+        if getattr(config, "cpu_only", False) is True:
+            dtype = torch.float32
+
+        model_path = Path(config.path)
+        model = AutoModelForCausalLM.from_pretrained(model_path, local_files_only=True, torch_dtype=dtype)
+        return model

invokeai/backend/model_manager/load/model_loaders/textual_inversion.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for TI model loading in InvokeAI."""
+
+from pathlib import Path
+from typing import Optional
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+from invokeai.backend.textual_inversion import TextualInversionModelRaw
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.TextualInversion, format=ModelFormat.EmbeddingFile)
+@ModelLoaderRegistry.register(
+    base=BaseModelType.Any, type=ModelType.TextualInversion, format=ModelFormat.EmbeddingFolder
+)
+class TextualInversionLoader(ModelLoader):
+    """Class to load TI models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if submodel_type is not None:
+            raise ValueError("There are no submodels in a TI model.")
+        model = TextualInversionModelRaw.from_checkpoint(
+            file_path=config.path,
+            dtype=self._torch_dtype,
+        )
+        return model
+
+    # override
+    def _get_model_path(self, config: AnyModelConfig) -> Path:
+        model_path = self._app_config.models_path / config.path
+
+        if config.format == ModelFormat.EmbeddingFolder:
+            path = model_path / "learned_embeds.bin"
+        else:
+            path = model_path
+
+        if not path.exists():
+            raise OSError(f"The embedding file at {path} was not found")
+
+        return path

invokeai/backend/model_manager/load/model_loaders/vae.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team
+"""Class for VAE model loading in InvokeAI."""
+
+from typing import Optional
+
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.configs.vae import (
+    VAE_Checkpoint_Anima_Config,
+    VAE_Checkpoint_Config_Base,
+    VAE_Checkpoint_QwenImage_Config,
+)
+from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
+from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
+from invokeai.backend.model_manager.taxonomy import (
+    AnyModel,
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+    SubModelType,
+)
+
+
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.VAE, format=ModelFormat.Diffusers)
+@ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.VAE, format=ModelFormat.Checkpoint)
+class VAELoader(GenericDiffusersLoader):
+    """Class to load VAE models."""
+
+    def _load_model(
+        self,
+        config: AnyModelConfig,
+        submodel_type: Optional[SubModelType] = None,
+    ) -> AnyModel:
+        if isinstance(config, VAE_Checkpoint_Anima_Config):
+            from diffusers.models.autoencoders import AutoencoderKLWan
+
+            return AutoencoderKLWan.from_single_file(
+                config.path,
+                torch_dtype=self._torch_dtype,
+            )
+        elif isinstance(config, VAE_Checkpoint_QwenImage_Config):
+            return self._load_qwen_image_vae(config)
+        elif isinstance(config, VAE_Checkpoint_Config_Base):
+            return AutoencoderKL.from_single_file(
+                config.path,
+                torch_dtype=self._torch_dtype,
+            )
+        else:
+            return super()._load_model(config, submodel_type)
+
+    def _load_qwen_image_vae(self, config: VAE_Checkpoint_QwenImage_Config) -> AnyModel:
+        """Load a Qwen Image VAE from a single safetensors file.
+
+        The Qwen Image VAE checkpoint is expected to be in the diffusers state-dict
+        layout (i.e. the same keys as `vae/diffusion_pytorch_model.safetensors` from
+        the Qwen-Image repo). `AutoencoderKLQwenImage` does not register a single-file
+        conversion in diffusers, so we instantiate the model with default config and
+        load the state dict directly.
+        """
+        import accelerate
+        from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage
+        from safetensors.torch import load_file
+
+        sd = load_file(config.path)
+
+        if self._torch_dtype is not None:
+            for k in list(sd.keys()):
+                if sd[k].is_floating_point():
+                    sd[k] = sd[k].to(self._torch_dtype)
+
+        new_sd_size = sum(t.nelement() * t.element_size() for t in sd.values())
+        self._ram_cache.make_room(new_sd_size)
+
+        with accelerate.init_empty_weights():
+            model = AutoencoderKLQwenImage()
+
+        model.load_state_dict(sd, strict=True, assign=True)
+        model.eval()
+        return model

invokeai/backend/model_manager/load/model_util.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2024 The InvokeAI Development Team
+"""Various utility functions needed by the loader and caching system."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Optional
+
+import onnxruntime as ort
+import torch
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from transformers import CLIPTokenizer, PreTrainedTokenizerBase, T5Tokenizer
+
+from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline
+from invokeai.backend.image_util.grounding_dino.grounding_dino_pipeline import GroundingDinoPipeline
+from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
+from invokeai.backend.ip_adapter.ip_adapter import IPAdapter
+from invokeai.backend.model_manager.taxonomy import AnyModel
+from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
+from invokeai.backend.textual_inversion import TextualInversionModelRaw
+from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+
+
+def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
+    """Get size of a model in memory in bytes."""
+    # TODO(ryand): We should create a CacheableModel interface for all models, and move the size calculations down to
+    # the models themselves.
+    if isinstance(model, DiffusionPipeline):
+        return _calc_pipeline_by_data(model)
+    elif isinstance(model, torch.nn.Module):
+        return calc_module_size(model)
+    elif isinstance(model, IAIOnnxRuntimeModel):
+        return _calc_onnx_model_by_data(model)
+    elif isinstance(model, SchedulerMixin):
+        return 0
+    elif isinstance(model, CLIPTokenizer):
+        # TODO(ryand): Accurately calculate the tokenizer's size. It's small enough that it shouldn't matter for now.
+        return 0
+    elif isinstance(
+        model,
+        (
+            TextualInversionModelRaw,
+            IPAdapter,
+            ModelPatchRaw,
+            SpandrelImageToImageModel,
+            GroundingDinoPipeline,
+            SegmentAnythingPipeline,
+            DepthAnythingPipeline,
+        ),
+    ):
+        return model.calc_size()
+    elif isinstance(model, ort.InferenceSession):
+        if model._model_bytes is not None:
+            # If the model is already loaded, return the size of the model bytes
+            return len(model._model_bytes)
+        elif model._model_path is not None:
+            # If the model is not loaded, return the size of the model path
+            return calc_model_size_by_fs(Path(model._model_path))
+        else:
+            # If neither is available, return 0
+            return 0
+    elif isinstance(
+        model,
+        T5Tokenizer,
+    ):
+        # HACK(ryand): len(model) just returns the vocabulary size, so this is blatantly wrong. It should be small
+        # relative to the text encoder that it's used with, so shouldn't matter too much, but we should fix this at some
+        # point.
+        return len(model)
+    elif isinstance(model, PreTrainedTokenizerBase):
+        # Catch-all for other tokenizer types (e.g., Qwen2Tokenizer, Qwen3Tokenizer).
+        # Tokenizers are small relative to models, so returning 0 is acceptable.
+        return 0
+    else:
+        # TODO(ryand): Promote this from a log to an exception once we are confident that we are handling all of the
+        # supported model types.
+        logger.warning(
+            f"Failed to calculate model size for unexpected model type: {type(model)}. The model will be treated as "
+            "having size 0."
+        )
+        return 0
+
+
+def _calc_pipeline_by_data(pipeline: DiffusionPipeline) -> int:
+    res = 0
+    assert hasattr(pipeline, "components")
+    for submodel_key in pipeline.components.keys():
+        submodel = getattr(pipeline, submodel_key)
+        if submodel is not None and isinstance(submodel, torch.nn.Module):
+            res += calc_module_size(submodel)
+    return res
+
+
+def calc_module_size(model: torch.nn.Module) -> int:
+    """Calculate the size (in bytes) of a torch.nn.Module."""
+    mem_params = sum([calc_tensor_size(param) for param in model.parameters()])
+    mem_bufs = sum([calc_tensor_size(buf) for buf in model.buffers()])
+    return mem_params + mem_bufs
+
+
+def _calc_onnx_model_by_data(model: IAIOnnxRuntimeModel) -> int:
+    tensor_size = model.tensors.size() * 2  # The session doubles this
+    mem = tensor_size  # in bytes
+    return mem
+
+
+def calc_model_size_by_fs(model_path: Path, subfolder: Optional[str] = None, variant: Optional[str] = None) -> int:
+    """Estimate the size of a model on disk in bytes."""
+    if model_path.is_file():
+        return model_path.stat().st_size
+
+    if subfolder is not None:
+        model_path = model_path / subfolder
+
+    # this can happen when, for example, the safety checker is not downloaded.
+    if not model_path.exists():
+        return 0
+
+    all_files = [f for f in model_path.iterdir() if (model_path / f).is_file()]
+
+    fp16_files = {f for f in all_files if ".fp16." in f.name or ".fp16-" in f.name}
+    bit8_files = {f for f in all_files if ".8bit." in f.name or ".8bit-" in f.name}
+    other_files = set(all_files) - fp16_files - bit8_files
+
+    if not variant:  # ModelRepoVariant.DEFAULT evaluates to empty string for compatability with HF
+        files = other_files
+    elif variant == "fp16":
+        files = fp16_files
+    elif variant == "8bit":
+        files = bit8_files
+    else:
+        raise NotImplementedError(f"Unknown variant: {variant}")
+
+    # try read from index if exists
+    index_postfix = ".index.json"
+    if variant is not None:
+        index_postfix = f".index.{variant}.json"
+
+    for file in files:
+        if not file.name.endswith(index_postfix):
+            continue
+        try:
+            with open(model_path / file, "r") as f:
+                index_data = json.loads(f.read())
+            return int(index_data["metadata"]["total_size"])
+        except Exception:
+            pass
+
+    # calculate files size if there is no index file
+    formats = [
+        (".safetensors",),  # safetensors
+        (".bin",),  # torch
+        (".onnx", ".pb"),  # onnx
+        (".msgpack",),  # flax
+        (".ckpt",),  # tf
+        (".h5",),  # tf2
+        (".gguf",),  # gguf quantized
+    ]
+
+    for file_format in formats:
+        model_files = [f for f in files if f.suffix in file_format]
+        if len(model_files) == 0:
+            continue
+
+        model_size = 0
+        for model_file in model_files:
+            file_stats = (model_path / model_file).stat()
+            model_size += file_stats.st_size
+        return model_size
+
+    return 0  # scheduler/feature_extractor/tokenizer - models without loading to gpu

invokeai/backend/model_manager/load/optimizations.py

@@ -0,0 +1,31 @@
+from contextlib import contextmanager
+from typing import Any, Generator
+
+import torch
+
+
+def _no_op(*args: Any, **kwargs: Any) -> None:
+    pass
+
+
+@contextmanager
+def skip_torch_weight_init() -> Generator[None, None, None]:
+    """Monkey patch several of the common torch layers (torch.nn.Linear, torch.nn.Conv1d, etc.) to skip weight initialization.
+
+    By default, `torch.nn.Linear` and `torch.nn.ConvNd` layers initialize their weights (according to a particular
+    distribution) when __init__ is called. This weight initialization step can take a significant amount of time, and is
+    completely unnecessary if the intent is to load checkpoint weights from disk for the layer. This context manager
+    monkey-patches common torch layers to skip the weight initialization step.
+    """
+    torch_modules = [torch.nn.Linear, torch.nn.modules.conv._ConvNd, torch.nn.Embedding]
+    saved_functions = [hasattr(m, "reset_parameters") and m.reset_parameters for m in torch_modules]
+
+    try:
+        for torch_module in torch_modules:
+            assert hasattr(torch_module, "reset_parameters")
+            torch_module.reset_parameters = _no_op
+        yield None
+    finally:
+        for torch_module, saved_function in zip(torch_modules, saved_functions, strict=True):
+            assert hasattr(torch_module, "reset_parameters")
+            torch_module.reset_parameters = saved_function