chore: import upstream snapshot with attribution

invokeai/backend/model_manager/README.md

@@ -0,0 +1,212 @@
+# Model Management System
+
+This document describes Invoke's model management system and common tasks for extending model support.
+
+## Overview
+
+The model management system handles the full lifecycle of models: identification, loading, and running. The system is extensible and supports multiple model architectures, formats, and quantization schemes.
+
+### Three Major Subsystems
+
+1. **Model Identification** (`configs/`): Determines model type, architecture, format, and metadata when users install models.
+2. **Model Loading** (`load/`): Loads models from disk into memory for inference.
+3. **Model Running**: Executes inference on loaded models. Implementation is scattered across the codebase, typically in architecture-specific inference code adjacent to `model_manager/`. The inference code is run in nodes in the graph execution system.
+
+## Core Concepts
+
+### Model Taxonomy
+
+The `taxonomy.py` module defines the type system for models:
+
+- `ModelType`: The kind of model (e.g., `Main`, `LoRA`, `ControlNet`, `VAE`).
+- `ModelFormat`: Storage format - may imply a quantization or some other quality (e.g., `Diffusers`, `Checkpoint`, `LyCORIS`, `BnbQuantizednf4b`).
+- `BaseModelType`: Associated pipeline architecture (e.g., `StableDiffusion1`, `StableDiffusionXL`, `Flux`). Models without an associated base use `Any` (e.g., `CLIPVision` is its own thing).
+- `ModelVariantType`, `FluxVariantType`, `ClipVariantType`: Architecture-specific variants.
+
+These enums form a discriminated union that uniquely identifies each model configuration class.
+
+### Model "Configs"
+
+Model configs are Pydantic models that describe a model on disk. They include the model taxonomy, path, and any metadata needed for loading or running the model.
+
+Model configs are stored in the database.
+
+### Model Identification
+
+When a user installs a model, the system attempts to identify it by trying each registered config class until one matches.
+
+**Config Classes** (`configs/`):
+
+- All config classes inherit from `Config_Base`, either directly or indirectly via some intermediary class (e.g., `Diffusers_Config_Base`, `Checkpoint_Config_Base`, or something narrower).
+- Each config class represents a specific, unique combination of `type`, `format`, `base`, and optional `variant`.
+- Config classes must implement `from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict) -> Self`. This method inspects the model on disk and raises `NotAMatchError` if the model doesn't match the config class, or returns an instance of the config class if it does.
+  - `ModelOnDisk` is a helper class that abstracts the model weights. It should be the entrypoint for inspecting the model (e.g., loading state dicts).
+- Override fields allow users to provide hints (e.g., when differentiating between SD1/SD2/SDXL VAEs with identical structures).
+
+**Identification Process**:
+
+1. `ModelConfigFactory.from_model_on_disk()` is called with a path to the model.
+2. The factory iterates through all registered config classes, calling `from_model_on_disk()` on each.
+3. Each config class inspects the model (state dict keys, tensor shapes, config files, etc.).
+4. If a match is found, the config instance is returned. If multiple matches are found, they are prioritized (e.g., main models over LoRAs).
+5. If no match is found, an `Unknown_Config` is returned as a fallback.
+
+**Utilities** (`identification_utils.py`):
+
+- `NotAMatchError`: Exception raised when a model doesn't match a config class.
+- `get_config_dict_or_raise()`: Load JSON config files from diffusers/transformers models.
+- `raise_for_class_name()`: Validate class names in config files.
+- `raise_for_override_fields()`: Validate user-provided override fields against the config schema.
+- `state_dict_has_any_keys_*()`: Helpers for inspecting state dict keys.
+
+### Model Loading
+
+Model loaders handle instantiating models from disk into memory.
+
+**Loader Classes** (`load/model_loaders/`):
+
+- Loaders register themselves with a decorator `@ModelLoaderRegistry.register(base=..., type=..., format=...)`. The `type`, `format` and `base` indicate which configs classes the loader can handle.
+- Each loader implements `_load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType]) -> AnyModel`.
+- Loaders are responsible for:
+  - Loading model weights from the config's path.
+  - Instantiating the correct model class (often using diffusers, transformers, or custom implementations).
+  - Returning the in-memory model representation.
+
+**Model Cache** (`load/model_cache/`):
+
+> This system typically does not require changes to support new model types, but it is important to understand how it works.
+
+- Manages models in memory with RAM and VRAM limits.
+- Handles moving models between CPU (storage device) and GPU (execution device).
+- Implements LRU eviction for RAM and smallest-first offload for VRAM.
+- Supports partial loading for large models on CUDA.
+- Thread-safe with locks on all public methods.
+
+**Loading Process**:
+
+1. The appropriate loader is selected based on the model config's `base`, `type`, and `format` attributes.
+2. The loader's `_load_model()` method is called with the model config.
+3. The loaded model is added to the model cache via `ModelCache.put()`.
+4. When needed, the model is moved into VRAM via `ModelCache.get()` and `ModelCache.lock()`.
+
+### Model Running
+
+Model running is architecture-specific and typically implemented in folders adjacent to `model_manager/`.
+
+Inference code doesn't necessarily follow any specific pattern, and doesn't interact directly with the model management system except to receive model configs and loaded models.
+
+At a high level, when a node needs to run a model, it will:
+
+- Receive a model identifier as an input or constant. This is typically the model's database ID (aka the `key`).
+- The node will use the `InvocationContext` API to load the model. The request is dispatched to the model manager which will load the model and return the a model loader with a context manager that yields the in-memory model, mediating VRAM/RAM management as needed.
+- The node will run inference using the loaded model using whatever patterns or libraries it needs.
+
+## Common Tasks
+
+### Task 1: Improving Identification for a Supported Model Type
+
+When identification fails or produces incorrect results for a model that should be supported, you may need to refine the identification logic.
+
+**Steps**:
+
+1. Obtain the failing model file or directory.
+2. Create a test case for it, following the instructions in `tests/model_identification/README.md`.
+3. Review the relevant config class in `configs/` (e.g., `configs/lora.py` for LoRA models).
+4. Examine the `from_model_on_disk()` method for some existing models to understand the patterns for identification logic.
+5. Inspect the failing model's files and structure:
+   - For checkpoint files: Load the state dict and examine keys and tensor shapes.
+   - For diffusers models: Examine the config files and directory structure.
+6. Update the identification logic to handle the new model variant. Common approaches:
+   - Check for specific state dict keys or key patterns.
+   - Inspect tensor shapes (e.g., `state_dict[key].shape`).
+   - Parse config files for class names or configuration values.
+   - Use helper functions from `identification_utils.py`.
+7. Run the test suite to verify the new logic works and doesn't break existing tests: `pytest tests/model_identification/test_identification.py`.
+   - Make sure you have installed the test dependencies (e.g. `uv pip install -e ".[dev,test]"`).
+   - If the model type is complex or has multiple variants, consider adding more test cases to cover edge cases.
+8. If, after successfully adding identification support for the model, it still doesn't work, you may need to update loading and/or inference code as well.
+
+**Key Files**:
+
+- Config class: `configs/<model_type>.py`
+- Identification utilities: `configs/identification_utils.py`
+- Taxonomy: `taxonomy.py`
+- Test README: `tests/model_identification/README.md`
+
+### Task 2: Adding Support for a New Model Type
+
+Adding a new model type requires implementing identification and loading logic. Inference and new nodes ("invocations") may be required if the model type doesn't fit into existing architectures or nodes.
+
+**Steps**:
+
+#### 1. Define Taxonomy
+
+- Add a new `ModelType` enum value in `taxonomy.py` if needed.
+- Determine the appropriate `BaseModelType` (or use `Any` if not architecture-specific).
+- Add a new `ModelFormat` if the model uses a unique storage format.
+
+You may need to add other attributes, depending on the model.
+
+#### 2. Implement Config Class
+
+- Create a new config file in `configs/` (e.g., `configs/new_model.py`).
+- Define a config class inheriting from `Config_Base` and appropriate format base class:
+  - `Diffusers_Config_Base` for diffusers-style models.
+  - `Checkpoint_Config_Base` for single-file checkpoint models.
+- Define `type`, `format`, and `base` as `Literal` fields with defaults. Remember, these must uniquely identify the config class.
+- Implement `from_model_on_disk()`:
+  - Validate the model is the correct format (file vs directory).
+  - Inspect state dict keys, tensor shapes, or config files.
+  - Raise `NotAMatchError` if the model doesn't match.
+  - Extract any additional metadata needed (e.g., variant, prediction type).
+  - Return an instance of the config class.
+- Register the config in `configs/factory.py`:
+  - Add the config class to the `AnyModelConfig` union.
+  - Add an `Annotated[YourConfig, YourConfig.get_tag()]` entry.
+
+#### 3. Implement Loader Class
+
+- Create a new loader file in `load/model_loaders/` (e.g., `load/model_loaders/new_model.py`).
+- Define a loader class inheriting from `ModelLoader`.
+- Decorate with `@ModelLoaderRegistry.register(base=..., type=..., format=...)`.
+- Implement `_load_model()`:
+  - Load model weights from `config.path`.
+  - Instantiate the model using the appropriate library (diffusers, transformers, or custom).
+  - Handle `submodel_type` if the model has submodels (e.g., text encoders, VAE).
+  - Return the in-memory model representation.
+
+#### 4. Add Tests
+
+Follow the instructions in `tests/model_identification/README.md`.
+
+#### 5. Implement Inference and Nodes (if needed)
+
+- If the model type requires new inference logic, implement it in an appropriate location.
+- Create nodes for the model if it doesn't fit into existing nodes. Search for subclasses of `BaseInvocation` for many examples.
+
+### 6. Frontend Support
+
+#### Workflows tab
+
+Typically, you will not need to do anything for the model to work in the Workflow Editor. When you define the node's model field, you can provide constraints for what type of models are selectable. The UI will automatically filter the list of models based on the model taxonomy.
+
+For example, this field definition in a node will allow users to select only "main" (pipeline) Stable Diffusion 1.x or 2.x models:
+
+```py
+model: ModelIdentifierField = InputField(
+    ui_model_base=[BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2],
+    ui_model_type=ModelType.Main,
+)
+```
+
+This same pattern works for any combination of `type`, `base`, `format`, and `variant`.
+
+#### Canvas / Generate tabs
+
+The Canvas and Generate tabs use graphs internally, but they don't expose the full graph editor UI. Instead, they provide a simplified interface for common tasks.
+
+They use "graph builder" functions, which take the user's selected settings and build a graph behind the scenes. We have one graph builder for each model architecture.
+
+Updating or adding a graph builder can be a bit complex, and you'd likely need to update other UI components and state management to support the new model type.
+
+The SDXL graph builder is a good example: `invokeai/frontend/web/src/features/nodes/util/graph/generation/buildSDXLGraph.ts`

invokeai/backend/model_manager/configs/base.py

@@ -0,0 +1,257 @@
+from abc import ABC, abstractmethod
+from enum import Enum
+from inspect import isabstract
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Literal,
+    Self,
+    Type,
+)
+
+from pydantic import BaseModel, ConfigDict, Field, Tag, field_validator
+from pydantic_core import PydanticUndefined
+
+from invokeai.app.util.misc import uuid_string
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    AnyVariant,
+    BaseModelType,
+    ModelFormat,
+    ModelRepoVariant,
+    ModelSourceType,
+    ModelType,
+)
+
+if TYPE_CHECKING:
+    pass
+
+
+class Config_Base(ABC, BaseModel):
+    """
+    Abstract base class for model configurations. A model config describes a specific combination of model base, type and
+    format, along with other metadata about the model. For example, a Stable Diffusion 1.x main model in checkpoint format
+    would have base=sd-1, type=main, format=checkpoint.
+
+    To create a new config type, inherit from this class and implement its interface:
+    - Define method 'from_model_on_disk' that returns an instance of the class or raises NotAMatch. This method will be
+        called during model installation to determine the correct config class for a model.
+    - Define fields 'type', 'base' and 'format' as pydantic fields. These should be Literals with a single value. A
+        default must be provided for each of these fields.
+
+    If multiple combinations of base, type and format need to be supported, create a separate subclass for each.
+
+    See MinimalConfigExample in test_model_probe.py for an example implementation.
+    """
+
+    # These fields are common to all model configs.
+
+    key: str = Field(
+        default_factory=uuid_string,
+        description="A unique key for this model.",
+    )
+    hash: str = Field(
+        description="The hash of the model file(s).",
+    )
+    path: str = Field(
+        description="Path to the model on the filesystem. Relative paths are relative to the Invoke root directory.",
+    )
+    file_size: int = Field(
+        description="The size of the model in bytes.",
+    )
+    name: str = Field(
+        description="Name of the model.",
+    )
+    description: str | None = Field(
+        default=None,
+        description="Model description",
+    )
+    source: str = Field(
+        description="The original source of the model (path, URL or repo_id).",
+    )
+    source_type: ModelSourceType = Field(
+        description="The type of source",
+    )
+    source_api_response: str | None = Field(
+        default=None,
+        description="The original API response from the source, as stringified JSON.",
+    )
+    source_url: str | None = Field(
+        default=None,
+        description="Optional URL for the model (e.g. download page or model page).",
+    )
+
+    @field_validator("source_url", mode="before")
+    @classmethod
+    def validate_source_url(cls, v: Any) -> str | None:
+        if v is None or v == "":
+            return None
+        if not isinstance(v, str):
+            raise ValueError("source_url must be a string")
+        if not v.startswith(("https://", "http://")):
+            raise ValueError("source_url must be an http or https URL")
+        return v
+
+    cover_image: str | None = Field(
+        default=None,
+        description="Url for image to preview model",
+    )
+
+    CONFIG_CLASSES: ClassVar[set[Type["Config_Base"]]] = set()
+    """Set of all non-abstract subclasses of Config_Base, for use during model probing. In other words, this is the set
+    of all known model config types."""
+
+    model_config = ConfigDict(
+        validate_assignment=True,
+        json_schema_serialization_defaults_required=True,
+        json_schema_mode_override="serialization",
+    )
+
+    @classmethod
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # Register non-abstract subclasses so we can iterate over them later during model probing. Note that
+        # isabstract() will return False if the class does not have any abstract methods, even if it inherits from ABC.
+        # We must check for ABC lest we unintentionally register some abstract model config classes.
+        if not isabstract(cls) and ABC not in cls.__bases__:
+            cls.CONFIG_CLASSES.add(cls)
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs):
+        # Ensure that model configs define 'base', 'type' and 'format' fields and provide defaults for them. Each
+        # subclass is expected to represent a single combination of base, type and format.
+        #
+        # This pydantic dunder method is called after the pydantic model for a class is created. The normal
+        # __init_subclass__ is too early to do this check.
+        for name in ("type", "base", "format"):
+            if name not in cls.model_fields:
+                raise NotImplementedError(f"{cls.__name__} must define a '{name}' field")
+            if cls.model_fields[name].default is PydanticUndefined:
+                raise NotImplementedError(f"{cls.__name__} must define a default for the '{name}' field")
+
+    @classmethod
+    def get_tag(cls) -> Tag:
+        """Constructs a pydantic discriminated union tag for this model config class. When a config is deserialized,
+        pydantic uses the tag to determine which subclass to instantiate.
+
+        The tag is a dot-separated string of the type, format, base and variant (if applicable).
+        """
+        tag_strings: list[str] = []
+        for name in ("type", "format", "base", "variant"):
+            if field := cls.model_fields.get(name):
+                # The check in __pydantic_init_subclass__ ensures that type, format and base are always present with
+                # defaults. variant does not require a default, but if it has one, we need to add it to the tag. We can
+                # check for the presence of a default by seeing if it's not PydanticUndefined, a sentinel value used by
+                # pydantic to indicate that no default was provided.
+                if field.default is not PydanticUndefined and field.default is not None:
+                    # We expect each of these fields has an Enum for its default; we want the value of the enum.
+                    tag_strings.append(field.default.value)
+        return Tag(".".join(tag_strings))
+
+    @staticmethod
+    def get_model_discriminator_value(v: Any) -> str:
+        """Computes the discriminator value for a model config discriminated union."""
+        # This is called by pydantic during deserialization and serialization to determine which model the data
+        # represents. It can get either a dict (during deserialization) or an instance of a Config_Base subclass
+        # (during serialization).
+        #
+        # See: https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions-with-callable-discriminator
+        if isinstance(v, Config_Base):
+            # We have an instance of a ModelConfigBase subclass - use its tag directly.
+            return v.get_tag().tag
+        if isinstance(v, dict):
+            # We have a dict - attempt to compute a tag from its fields.
+            tag_strings: list[str] = []
+            if type_ := v.get("type"):
+                if isinstance(type_, Enum):
+                    type_ = str(type_.value)
+                elif not isinstance(type_, str):
+                    raise ValueError("Model config dict 'type' field must be a string or Enum")
+                tag_strings.append(type_)
+
+            if format_ := v.get("format"):
+                if isinstance(format_, Enum):
+                    format_ = str(format_.value)
+                elif not isinstance(format_, str):
+                    raise ValueError("Model config dict 'format' field must be a string or Enum")
+                tag_strings.append(format_)
+
+            if base_ := v.get("base"):
+                if isinstance(base_, Enum):
+                    base_ = str(base_.value)
+                elif not isinstance(base_, str):
+                    raise ValueError("Model config dict 'base' field must be a string or Enum")
+                tag_strings.append(base_)
+
+            # Special case: CLIP Embed models also need the variant to distinguish them.
+            if (
+                type_ == ModelType.CLIPEmbed.value
+                and format_ == ModelFormat.Diffusers.value
+                and base_ == BaseModelType.Any.value
+            ):
+                if variant_ := v.get("variant"):
+                    if isinstance(variant_, Enum):
+                        variant_ = variant_.value
+                    elif not isinstance(variant_, str):
+                        raise ValueError("Model config dict 'variant' field must be a string or Enum")
+                    tag_strings.append(variant_)
+                else:
+                    raise ValueError("CLIP Embed model config dict must include a 'variant' field")
+
+            return ".".join(tag_strings)
+        else:
+            raise ValueError(
+                "Model config discriminator value must be computed from a dict or ModelConfigBase instance"
+            )
+
+    @classmethod
+    @abstractmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        """Given the model on disk and any override fields, attempt to construct an instance of this config class.
+
+        This method serves to identify whether the model on disk matches this config class, and if so, to extract any
+        additional metadata needed to instantiate the config.
+
+        Implementations should raise a NotAMatchError if the model does not match this config class."""
+        raise NotImplementedError(f"from_model_on_disk not implemented for {cls.__name__}")
+
+
+class Checkpoint_Config_Base(ABC, BaseModel):
+    """Base class for checkpoint-style models."""
+
+    config_path: str | None = Field(
+        description="Path to the config for this model, if any.",
+        default=None,
+    )
+
+
+class Diffusers_Config_Base(ABC, BaseModel):
+    """Base class for diffusers-style models."""
+
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    repo_variant: ModelRepoVariant = Field(ModelRepoVariant.Default)
+
+    @classmethod
+    def _get_repo_variant_or_raise(cls, mod: ModelOnDisk) -> ModelRepoVariant:
+        # get all files ending in .bin or .safetensors
+        weight_files = list(mod.path.glob("**/*.safetensors"))
+        weight_files.extend(list(mod.path.glob("**/*.bin")))
+        for x in weight_files:
+            if ".fp16" in x.suffixes:
+                return ModelRepoVariant.FP16
+            if "openvino_model" in x.name:
+                return ModelRepoVariant.OpenVINO
+            if "flax_model" in x.name:
+                return ModelRepoVariant.Flax
+            if x.suffix == ".onnx":
+                return ModelRepoVariant.ONNX
+        return ModelRepoVariant.Default
+
+
+class SubmodelDefinition(BaseModel):
+    path_or_prefix: str
+    model_type: ModelType
+    variant: AnyVariant | None = None
+
+    model_config = ConfigDict(protected_namespaces=())

invokeai/backend/model_manager/configs/clip_embed.py

@@ -0,0 +1,92 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    get_config_dict_or_raise,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ClipVariantType,
+    ModelFormat,
+    ModelType,
+)
+
+
+def get_clip_variant_type_from_config(config: dict[str, Any]) -> ClipVariantType | None:
+    try:
+        hidden_size = config.get("hidden_size")
+        match hidden_size:
+            case 1280:
+                return ClipVariantType.G
+            case 768:
+                return ClipVariantType.L
+            case _:
+                return None
+    except Exception:
+        return None
+
+
+class CLIPEmbed_Diffusers_Config_Base(Diffusers_Config_Base):
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.CLIPEmbed] = Field(default=ModelType.CLIPEmbed)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            {
+                mod.path / "config.json",
+                mod.path / "text_encoder" / "config.json",
+            },
+            {
+                "CLIPModel",
+                "CLIPTextModel",
+                "CLIPTextModelWithProjection",
+            },
+        )
+
+        cls._validate_variant(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_variant(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model variant does not match this config class."""
+        expected_variant = cls.model_fields["variant"].default
+        config = get_config_dict_or_raise(
+            {
+                mod.path / "config.json",
+                mod.path / "text_encoder" / "config.json",
+            },
+        )
+        recognized_variant = get_clip_variant_type_from_config(config)
+
+        if recognized_variant is None:
+            raise NotAMatchError("unable to determine CLIP variant from config")
+
+        if expected_variant is not recognized_variant:
+            raise NotAMatchError(f"variant is {recognized_variant}, not {expected_variant}")
+
+
+class CLIPEmbed_Diffusers_G_Config(CLIPEmbed_Diffusers_Config_Base, Config_Base):
+    variant: Literal[ClipVariantType.G] = Field(default=ClipVariantType.G)
+
+
+class CLIPEmbed_Diffusers_L_Config(CLIPEmbed_Diffusers_Config_Base, Config_Base):
+    variant: Literal[ClipVariantType.L] = Field(default=ClipVariantType.L)

invokeai/backend/model_manager/configs/clip_vision.py

@@ -0,0 +1,58 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    get_class_name_from_config_dict_or_raise,
+    get_config_dict_or_raise,
+    raise_for_override_fields,
+    raise_if_not_dir,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+
+class CLIPVision_Diffusers_Config(Diffusers_Config_Base, Config_Base):
+    """Model config for CLIPVision."""
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.CLIPVision] = Field(default=ModelType.CLIPVision)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls.raise_if_config_doesnt_look_like_clip_vision(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def raise_if_config_doesnt_look_like_clip_vision(cls, mod: ModelOnDisk) -> None:
+        config_dict = get_config_dict_or_raise(mod.path / "config.json")
+        class_name = get_class_name_from_config_dict_or_raise(config_dict)
+
+        if class_name == "CLIPVisionModelWithProjection":
+            looks_like_clip_vision = True
+        elif class_name == "CLIPModel" and "vision_config" in config_dict:
+            looks_like_clip_vision = True
+        else:
+            looks_like_clip_vision = False
+
+        if not looks_like_clip_vision:
+            raise NotAMatchError(
+                f"config class name is {class_name}, not CLIPVisionModelWithProjection or CLIPModel with vision_config"
+            )

invokeai/backend/model_manager/configs/controlnet.py

@@ -0,0 +1,342 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import Any
+
+from invokeai.backend.flux.controlnet.state_dict_utils import (
+    is_state_dict_instantx_controlnet,
+    is_state_dict_xlabs_controlnet,
+)
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    common_config_paths,
+    get_config_dict_or_raise,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    raise_if_not_file,
+    state_dict_has_any_keys_starting_with,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+MODEL_NAME_TO_PREPROCESSOR = {
+    "canny": "canny_image_processor",
+    "mlsd": "mlsd_image_processor",
+    "depth": "depth_anything_image_processor",
+    "bae": "normalbae_image_processor",
+    "normal": "normalbae_image_processor",
+    "sketch": "pidi_image_processor",
+    "scribble": "lineart_image_processor",
+    "lineart anime": "lineart_anime_image_processor",
+    "lineart_anime": "lineart_anime_image_processor",
+    "lineart": "lineart_image_processor",
+    "soft": "hed_image_processor",
+    "softedge": "hed_image_processor",
+    "hed": "hed_image_processor",
+    "shuffle": "content_shuffle_image_processor",
+    "pose": "dw_openpose_image_processor",
+    "mediapipe": "mediapipe_face_processor",
+    "pidi": "pidi_image_processor",
+    "zoe": "zoe_depth_image_processor",
+    "color": "color_map_image_processor",
+}
+
+
+class ControlAdapterDefaultSettings(BaseModel):
+    # This could be narrowed to controlnet processor nodes, but they change. Leaving this a string is safer.
+    preprocessor: str | None
+    fp8_storage: bool | None = Field(
+        default=None,
+        description="Store weights in FP8 to reduce VRAM usage (~50% savings). Weights are cast to compute dtype during inference.",
+    )
+    model_config = ConfigDict(extra="forbid")
+
+    @classmethod
+    def from_model_name(cls, model_name: str) -> Self:
+        for k, v in MODEL_NAME_TO_PREPROCESSOR.items():
+            model_name_lower = model_name.lower()
+            if k in model_name_lower:
+                return cls(preprocessor=v)
+        return cls(preprocessor=None)
+
+
+class ControlNet_Diffusers_Config_Base(Diffusers_Config_Base):
+    """Model config for ControlNet models (diffusers version)."""
+
+    type: Literal[ModelType.ControlNet] = Field(default=ModelType.ControlNet)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    default_settings: ControlAdapterDefaultSettings | None = Field(None)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "ControlNetModel",
+                "FluxControlNetModel",
+            },
+        )
+
+        cls._validate_base(mod)
+
+        repo_variant = {"repo_variant": override_fields.get("repo_variant", cls._get_repo_variant_or_raise(mod))}
+        args = override_fields | repo_variant
+        return cls(**args)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
+        config_dict = get_config_dict_or_raise(common_config_paths(mod.path))
+
+        if config_dict.get("_class_name") == "FluxControlNetModel":
+            return BaseModelType.Flux
+
+        dimension = config_dict.get("cross_attention_dim")
+
+        match dimension:
+            case 768:
+                return BaseModelType.StableDiffusion1
+            case 1024:
+                # No obvious way to distinguish between sd2-base and sd2-768, but we don't really differentiate them
+                # anyway.
+                return BaseModelType.StableDiffusion2
+            case 2048:
+                return BaseModelType.StableDiffusionXL
+            case _:
+                raise NotAMatchError(f"unrecognized cross_attention_dim {dimension}")
+
+
+class ControlNet_Diffusers_SD1_Config(ControlNet_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class ControlNet_Diffusers_SD2_Config(ControlNet_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class ControlNet_Diffusers_SDXL_Config(ControlNet_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class ControlNet_Diffusers_FLUX_Config(ControlNet_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.Flux] = Field(default=BaseModelType.Flux)
+
+
+class ControlNet_Checkpoint_Config_Base(Checkpoint_Config_Base):
+    """Model config for ControlNet models (diffusers version)."""
+
+    type: Literal[ModelType.ControlNet] = Field(default=ModelType.ControlNet)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    default_settings: ControlAdapterDefaultSettings | None = Field(None)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_controlnet(mod)
+
+        cls._validate_base(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _validate_looks_like_controlnet(cls, mod: ModelOnDisk) -> None:
+        if not state_dict_has_any_keys_starting_with(
+            mod.load_state_dict(),
+            {
+                "controlnet",
+                "control_model",
+                "input_blocks",
+                # XLabs FLUX ControlNet models have keys starting with "controlnet_blocks."
+                # For example: https://huggingface.co/XLabs-AI/flux-controlnet-collections/blob/86ab1e915a389d5857135c00e0d350e9e38a9048/flux-canny-controlnet_v2.safetensors
+                # TODO(ryand): This is very fragile. XLabs FLUX ControlNet models also contain keys starting with
+                # "double_blocks.", which we check for above. But, I'm afraid to modify this logic because it is so
+                # delicate.
+                "controlnet_blocks",
+            },
+        ):
+            raise NotAMatchError("state dict does not look like a ControlNet checkpoint")
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
+        state_dict = mod.load_state_dict()
+
+        if is_state_dict_xlabs_controlnet(state_dict) or is_state_dict_instantx_controlnet(state_dict):
+            # TODO(ryand): Should I distinguish between XLabs, InstantX and other ControlNet models by implementing
+            # get_format()?
+            return BaseModelType.Flux
+
+        for key in (
+            "control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight",
+            "controlnet_mid_block.bias",
+            "input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight",
+            "down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight",
+        ):
+            if key not in state_dict:
+                continue
+            width = state_dict[key].shape[-1]
+            match width:
+                case 768:
+                    return BaseModelType.StableDiffusion1
+                case 1024:
+                    return BaseModelType.StableDiffusion2
+                case 2048:
+                    return BaseModelType.StableDiffusionXL
+                case 1280:
+                    return BaseModelType.StableDiffusionXL
+                case _:
+                    pass
+
+        raise NotAMatchError("unable to determine base type from state dict")
+
+
+class ControlNet_Checkpoint_SD1_Config(ControlNet_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class ControlNet_Checkpoint_SD2_Config(ControlNet_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class ControlNet_Checkpoint_SDXL_Config(ControlNet_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class ControlNet_Checkpoint_FLUX_Config(ControlNet_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.Flux] = Field(default=BaseModelType.Flux)
+
+
+def _has_z_image_control_keys(state_dict: dict) -> bool:
+    """Check if state dict contains Z-Image Control specific keys."""
+    z_image_control_keys = {"control_layers", "control_all_x_embedder", "control_noise_refiner"}
+    for key in state_dict.keys():
+        if isinstance(key, str):
+            prefix = key.split(".")[0]
+            if prefix in z_image_control_keys:
+                return True
+    return False
+
+
+class ControlNet_Checkpoint_ZImage_Config(Checkpoint_Config_Base, Config_Base):
+    """Model config for Z-Image Control adapter models (Safetensors checkpoint).
+
+    Z-Image Control models are standalone adapters containing only the control layers
+    (control_layers, control_all_x_embedder, control_noise_refiner) that extend
+    the base Z-Image transformer with spatial conditioning capabilities.
+
+    Supports: Canny, HED, Depth, Pose, MLSD.
+    Recommended control_context_scale: 0.65-0.80.
+    """
+
+    type: Literal[ModelType.ControlNet] = Field(default=ModelType.ControlNet)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.ZImage] = Field(default=BaseModelType.ZImage)
+    default_settings: ControlAdapterDefaultSettings | None = Field(None)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_z_image_control(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_looks_like_z_image_control(cls, mod: ModelOnDisk) -> None:
+        state_dict = mod.load_state_dict()
+        if not _has_z_image_control_keys(state_dict):
+            raise NotAMatchError("state dict does not look like a Z-Image Control model")
+
+
+def _has_anima_lllite_keys(state_dict: dict) -> bool:
+    """Check if state dict contains Anima ControlNet-LLLite specific keys.
+
+    Anima LLLite adapters (v2 named-key format) have a shared conditioning trunk under
+    `lllite_conditioning1.*` and per-module weights under `lllite_dit_blocks_*`. SDXL
+    ControlNet-LLLite models use `lllite_unet_*` keys instead and do not match.
+    """
+    return state_dict_has_any_keys_starting_with(
+        state_dict, "lllite_conditioning1."
+    ) and state_dict_has_any_keys_starting_with(state_dict, "lllite_dit_blocks_")
+
+
+class ControlNet_Checkpoint_Anima_Config(Checkpoint_Config_Base, Config_Base):
+    """Model config for Anima ControlNet-LLLite adapter models (Safetensors checkpoint).
+
+    Anima LLLite adapters are standalone adapters consisting of a shared conditioning trunk
+    (lllite_conditioning1) that encodes a conditioning image, plus tiny per-Linear modules
+    (lllite_dit_blocks_*) that perturb the inputs of target Linears in the Anima DiT.
+    """
+
+    type: Literal[ModelType.ControlNet] = Field(default=ModelType.ControlNet)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.Anima] = Field(default=BaseModelType.Anima)
+    default_settings: ControlAdapterDefaultSettings | None = Field(None)
+    cond_in_channels: int | None = Field(
+        default=None,
+        description="Number of conditioning image channels (3 = RGB control image, 4 = RGB + inpaint mask). None for "
+        "models installed before this field was recorded.",
+    )
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_anima_lllite(mod)
+
+        args = dict(override_fields)
+        if "cond_in_channels" not in args:
+            args["cond_in_channels"] = cls._get_cond_in_channels(mod)
+        return cls(**args)
+
+    @classmethod
+    def _get_cond_in_channels(cls, mod: ModelOnDisk) -> int:
+        # Mirrors AnimaControlNetLLLite.from_state_dict: prefer the saved `lllite.*` hyperparam, falling back to
+        # the conv1 weight shape (ch_half, cond_in_channels, 4, 4).
+        meta_value = mod.metadata().get("lllite.cond_in_channels")
+        if meta_value is not None:
+            return int(meta_value)
+        conv1_weight = mod.load_state_dict().get("lllite_conditioning1.conv1.weight")
+        if conv1_weight is None:
+            raise NotAMatchError("state dict has Anima ControlNet-LLLite keys but no lllite_conditioning1.conv1.weight")
+        return int(conv1_weight.shape[1])
+
+    @classmethod
+    def _validate_looks_like_anima_lllite(cls, mod: ModelOnDisk) -> None:
+        state_dict = mod.load_state_dict()
+        if not _has_anima_lllite_keys(state_dict):
+            raise NotAMatchError("state dict does not look like an Anima ControlNet-LLLite model")

invokeai/backend/model_manager/configs/external_api.py

@@ -0,0 +1,113 @@
+from typing import Literal, Self
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelSourceType, ModelType
+
+ExternalGenerationMode = Literal["txt2img", "img2img", "inpaint"]
+ExternalMaskFormat = Literal["alpha", "binary", "none"]
+ExternalPanelControlName = Literal["reference_images", "dimensions", "seed"]
+
+
+class ExternalImageSize(BaseModel):
+    width: int = Field(gt=0)
+    height: int = Field(gt=0)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ExternalResolutionPreset(BaseModel):
+    label: str = Field(min_length=1, description="Display label, e.g. '1:1 (1K)'")
+    aspect_ratio: str = Field(min_length=1, description="Aspect ratio string, e.g. '1:1'")
+    image_size: str = Field(min_length=1, description="Image size preset, e.g. '1K'")
+    width: int = Field(gt=0)
+    height: int = Field(gt=0)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ExternalModelCapabilities(BaseModel):
+    modes: list[ExternalGenerationMode] = Field(default_factory=lambda: ["txt2img"])
+    supports_reference_images: bool = Field(default=False)
+    supports_negative_prompt: bool = Field(default=True)
+    supports_seed: bool = Field(default=False)
+    supports_guidance: bool = Field(default=False)
+    supports_steps: bool = Field(default=False)
+    max_images_per_request: int | None = Field(default=None, gt=0)
+    max_image_size: ExternalImageSize | None = Field(default=None)
+    allowed_aspect_ratios: list[str] | None = Field(default=None)
+    aspect_ratio_sizes: dict[str, ExternalImageSize] | None = Field(default=None)
+    resolution_presets: list[ExternalResolutionPreset] | None = Field(default=None)
+    max_reference_images: int | None = Field(default=None, gt=0)
+    mask_format: ExternalMaskFormat = Field(default="none")
+    input_image_required_for: list[ExternalGenerationMode] | None = Field(default=None)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ExternalApiModelDefaultSettings(BaseModel):
+    width: int | None = Field(default=None, gt=0)
+    height: int | None = Field(default=None, gt=0)
+    num_images: int | None = Field(default=None, gt=0)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ExternalModelPanelControl(BaseModel):
+    name: ExternalPanelControlName
+    slider_min: float | None = Field(default=None)
+    slider_max: float | None = Field(default=None)
+    number_input_min: float | None = Field(default=None)
+    number_input_max: float | None = Field(default=None)
+    fine_step: float | None = Field(default=None)
+    coarse_step: float | None = Field(default=None)
+    marks: list[float] | None = Field(default=None)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ExternalModelPanelSchema(BaseModel):
+    prompts: list[ExternalModelPanelControl] = Field(default_factory=list)
+    image: list[ExternalModelPanelControl] = Field(default_factory=list)
+    generation: list[ExternalModelPanelControl] = Field(default_factory=list)
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class ExternalApiModelConfig(Config_Base):
+    base: Literal[BaseModelType.External] = Field(default=BaseModelType.External)
+    type: Literal[ModelType.ExternalImageGenerator] = Field(default=ModelType.ExternalImageGenerator)
+    format: Literal[ModelFormat.ExternalApi] = Field(default=ModelFormat.ExternalApi)
+
+    provider_id: str = Field(min_length=1, description="External provider ID")
+    provider_model_id: str = Field(min_length=1, description="Provider-specific model ID")
+    capabilities: ExternalModelCapabilities = Field(description="Provider capability matrix")
+    default_settings: ExternalApiModelDefaultSettings | None = Field(default=None)
+    panel_schema: ExternalModelPanelSchema | None = Field(default=None)
+    tags: list[str] | None = Field(default=None)
+    is_default: bool = Field(default=False)
+
+    source_type: ModelSourceType = Field(default=ModelSourceType.External)
+    path: str = Field(default="")
+    source: str = Field(default="")
+    hash: str = Field(default="")
+    file_size: int = Field(default=0, ge=0)
+
+    model_config = ConfigDict(extra="forbid")
+
+    @model_validator(mode="after")
+    def _populate_external_fields(self) -> "ExternalApiModelConfig":
+        if not self.path:
+            self.path = f"external://{self.provider_id}/{self.provider_model_id}"
+        if not self.source:
+            self.source = self.path
+        if not self.hash:
+            self.hash = f"external:{self.provider_id}:{self.provider_model_id}"
+        return self
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, object]) -> Self:
+        raise NotAMatchError("external API models are not probed from disk")

invokeai/backend/model_manager/configs/factory.py

@@ -0,0 +1,587 @@
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import (
+    Union,
+)
+
+from pydantic import Discriminator, TypeAdapter, ValidationError
+from typing_extensions import Annotated, Any
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.app.util.misc import uuid_string
+from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.clip_embed import CLIPEmbed_Diffusers_G_Config, CLIPEmbed_Diffusers_L_Config
+from invokeai.backend.model_manager.configs.clip_vision import CLIPVision_Diffusers_Config
+from invokeai.backend.model_manager.configs.controlnet import (
+    ControlAdapterDefaultSettings,
+    ControlNet_Checkpoint_Anima_Config,
+    ControlNet_Checkpoint_FLUX_Config,
+    ControlNet_Checkpoint_SD1_Config,
+    ControlNet_Checkpoint_SD2_Config,
+    ControlNet_Checkpoint_SDXL_Config,
+    ControlNet_Checkpoint_ZImage_Config,
+    ControlNet_Diffusers_FLUX_Config,
+    ControlNet_Diffusers_SD1_Config,
+    ControlNet_Diffusers_SD2_Config,
+    ControlNet_Diffusers_SDXL_Config,
+)
+from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig
+from invokeai.backend.model_manager.configs.flux_redux import FLUXRedux_Checkpoint_Config
+from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
+from invokeai.backend.model_manager.configs.ip_adapter import (
+    IPAdapter_Checkpoint_FLUX_Config,
+    IPAdapter_Checkpoint_SD1_Config,
+    IPAdapter_Checkpoint_SD2_Config,
+    IPAdapter_Checkpoint_SDXL_Config,
+    IPAdapter_InvokeAI_SD1_Config,
+    IPAdapter_InvokeAI_SD2_Config,
+    IPAdapter_InvokeAI_SDXL_Config,
+)
+from invokeai.backend.model_manager.configs.llava_onevision import LlavaOnevision_Diffusers_Config
+from invokeai.backend.model_manager.configs.lora import (
+    ControlLoRA_LyCORIS_FLUX_Config,
+    LoRA_Diffusers_Flux2_Config,
+    LoRA_Diffusers_FLUX_Config,
+    LoRA_Diffusers_SD1_Config,
+    LoRA_Diffusers_SD2_Config,
+    LoRA_Diffusers_SDXL_Config,
+    LoRA_Diffusers_ZImage_Config,
+    LoRA_LyCORIS_Anima_Config,
+    LoRA_LyCORIS_Flux2_Config,
+    LoRA_LyCORIS_FLUX_Config,
+    LoRA_LyCORIS_QwenImage_Config,
+    LoRA_LyCORIS_SD1_Config,
+    LoRA_LyCORIS_SD2_Config,
+    LoRA_LyCORIS_SDXL_Config,
+    LoRA_LyCORIS_ZImage_Config,
+    LoRA_OMI_FLUX_Config,
+    LoRA_OMI_SDXL_Config,
+    LoraModelDefaultSettings,
+)
+from invokeai.backend.model_manager.configs.main import (
+    Main_BnBNF4_FLUX_Config,
+    Main_Checkpoint_Anima_Config,
+    Main_Checkpoint_Flux2_Config,
+    Main_Checkpoint_FLUX_Config,
+    Main_Checkpoint_QwenImage_Config,
+    Main_Checkpoint_SD1_Config,
+    Main_Checkpoint_SD2_Config,
+    Main_Checkpoint_SDXL_Config,
+    Main_Checkpoint_SDXLRefiner_Config,
+    Main_Checkpoint_ZImage_Config,
+    Main_Diffusers_CogView4_Config,
+    Main_Diffusers_Flux2_Config,
+    Main_Diffusers_FLUX_Config,
+    Main_Diffusers_QwenImage_Config,
+    Main_Diffusers_SD1_Config,
+    Main_Diffusers_SD2_Config,
+    Main_Diffusers_SD3_Config,
+    Main_Diffusers_SDXL_Config,
+    Main_Diffusers_SDXLRefiner_Config,
+    Main_Diffusers_ZImage_Config,
+    Main_GGUF_Flux2_Config,
+    Main_GGUF_FLUX_Config,
+    Main_GGUF_QwenImage_Config,
+    Main_GGUF_ZImage_Config,
+    MainModelDefaultSettings,
+)
+from invokeai.backend.model_manager.configs.qwen3_encoder import (
+    Qwen3Encoder_Checkpoint_Config,
+    Qwen3Encoder_GGUF_Config,
+    Qwen3Encoder_Qwen3Encoder_Config,
+)
+from invokeai.backend.model_manager.configs.qwen_vl_encoder import (
+    QwenVLEncoder_Checkpoint_Config,
+    QwenVLEncoder_Diffusers_Config,
+)
+from invokeai.backend.model_manager.configs.siglip import SigLIP_Diffusers_Config
+from invokeai.backend.model_manager.configs.spandrel import Spandrel_Checkpoint_Config
+from invokeai.backend.model_manager.configs.t2i_adapter import (
+    T2IAdapter_Diffusers_SD1_Config,
+    T2IAdapter_Diffusers_SDXL_Config,
+)
+from invokeai.backend.model_manager.configs.t5_encoder import T5Encoder_BnBLLMint8_Config, T5Encoder_T5Encoder_Config
+from invokeai.backend.model_manager.configs.text_llm import TextLLM_Diffusers_Config
+from invokeai.backend.model_manager.configs.textual_inversion import (
+    TI_File_SD1_Config,
+    TI_File_SD2_Config,
+    TI_File_SDXL_Config,
+    TI_Folder_SD1_Config,
+    TI_Folder_SD2_Config,
+    TI_Folder_SDXL_Config,
+)
+from invokeai.backend.model_manager.configs.unknown import Unknown_Config
+from invokeai.backend.model_manager.configs.vae import (
+    VAE_Checkpoint_Anima_Config,
+    VAE_Checkpoint_Flux2_Config,
+    VAE_Checkpoint_FLUX_Config,
+    VAE_Checkpoint_QwenImage_Config,
+    VAE_Checkpoint_SD1_Config,
+    VAE_Checkpoint_SD2_Config,
+    VAE_Checkpoint_SDXL_Config,
+    VAE_Diffusers_Flux2_Config,
+    VAE_Diffusers_SD1_Config,
+    VAE_Diffusers_SDXL_Config,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelSourceType,
+    ModelType,
+    variant_type_adapter,
+)
+
+logger = logging.getLogger(__name__)
+app_config = get_config()
+
+# Known model file extensions for sanity checking
+_MODEL_EXTENSIONS = {
+    ".safetensors",
+    ".ckpt",
+    ".pt",
+    ".pth",
+    ".bin",
+    ".gguf",
+    ".onnx",
+}
+
+# Known config file names for diffusers/transformers models
+_CONFIG_FILES = {
+    "model_index.json",
+    "config.json",
+}
+
+# Maximum number of files in a directory to be considered a model
+_MAX_FILES_IN_MODEL_DIR = 50
+
+# Maximum depth to search for model files in directories
+_MAX_SEARCH_DEPTH = 2
+
+
+# The types are listed explicitly because IDEs/LSPs can't identify the correct types
+# when AnyModelConfig is constructed dynamically using ModelConfigBase.all_config_classes
+AnyModelConfig = Annotated[
+    Union[
+        # Main (Pipeline) - diffusers format
+        Annotated[Main_Diffusers_SD1_Config, Main_Diffusers_SD1_Config.get_tag()],
+        Annotated[Main_Diffusers_SD2_Config, Main_Diffusers_SD2_Config.get_tag()],
+        Annotated[Main_Diffusers_SDXL_Config, Main_Diffusers_SDXL_Config.get_tag()],
+        Annotated[Main_Diffusers_SDXLRefiner_Config, Main_Diffusers_SDXLRefiner_Config.get_tag()],
+        Annotated[Main_Diffusers_SD3_Config, Main_Diffusers_SD3_Config.get_tag()],
+        Annotated[Main_Diffusers_FLUX_Config, Main_Diffusers_FLUX_Config.get_tag()],
+        Annotated[Main_Diffusers_Flux2_Config, Main_Diffusers_Flux2_Config.get_tag()],
+        Annotated[Main_Diffusers_CogView4_Config, Main_Diffusers_CogView4_Config.get_tag()],
+        Annotated[Main_Diffusers_QwenImage_Config, Main_Diffusers_QwenImage_Config.get_tag()],
+        Annotated[Main_Diffusers_ZImage_Config, Main_Diffusers_ZImage_Config.get_tag()],
+        # Main (Pipeline) - checkpoint format
+        # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
+        # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
+        Annotated[Main_Checkpoint_SD1_Config, Main_Checkpoint_SD1_Config.get_tag()],
+        Annotated[Main_Checkpoint_SD2_Config, Main_Checkpoint_SD2_Config.get_tag()],
+        Annotated[Main_Checkpoint_SDXL_Config, Main_Checkpoint_SDXL_Config.get_tag()],
+        Annotated[Main_Checkpoint_SDXLRefiner_Config, Main_Checkpoint_SDXLRefiner_Config.get_tag()],
+        Annotated[Main_Checkpoint_Flux2_Config, Main_Checkpoint_Flux2_Config.get_tag()],
+        Annotated[Main_Checkpoint_FLUX_Config, Main_Checkpoint_FLUX_Config.get_tag()],
+        Annotated[Main_Checkpoint_QwenImage_Config, Main_Checkpoint_QwenImage_Config.get_tag()],
+        Annotated[Main_Checkpoint_ZImage_Config, Main_Checkpoint_ZImage_Config.get_tag()],
+        Annotated[Main_Checkpoint_Anima_Config, Main_Checkpoint_Anima_Config.get_tag()],
+        # Main (Pipeline) - quantized formats
+        # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
+        # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
+        Annotated[Main_BnBNF4_FLUX_Config, Main_BnBNF4_FLUX_Config.get_tag()],
+        Annotated[Main_GGUF_Flux2_Config, Main_GGUF_Flux2_Config.get_tag()],
+        Annotated[Main_GGUF_FLUX_Config, Main_GGUF_FLUX_Config.get_tag()],
+        Annotated[Main_GGUF_QwenImage_Config, Main_GGUF_QwenImage_Config.get_tag()],
+        Annotated[Main_GGUF_ZImage_Config, Main_GGUF_ZImage_Config.get_tag()],
+        # VAE - checkpoint format
+        Annotated[VAE_Checkpoint_SD1_Config, VAE_Checkpoint_SD1_Config.get_tag()],
+        Annotated[VAE_Checkpoint_SD2_Config, VAE_Checkpoint_SD2_Config.get_tag()],
+        Annotated[VAE_Checkpoint_SDXL_Config, VAE_Checkpoint_SDXL_Config.get_tag()],
+        Annotated[VAE_Checkpoint_FLUX_Config, VAE_Checkpoint_FLUX_Config.get_tag()],
+        Annotated[VAE_Checkpoint_Flux2_Config, VAE_Checkpoint_Flux2_Config.get_tag()],
+        Annotated[VAE_Checkpoint_QwenImage_Config, VAE_Checkpoint_QwenImage_Config.get_tag()],
+        Annotated[VAE_Checkpoint_Anima_Config, VAE_Checkpoint_Anima_Config.get_tag()],
+        # VAE - diffusers format
+        Annotated[VAE_Diffusers_SD1_Config, VAE_Diffusers_SD1_Config.get_tag()],
+        Annotated[VAE_Diffusers_SDXL_Config, VAE_Diffusers_SDXL_Config.get_tag()],
+        Annotated[VAE_Diffusers_Flux2_Config, VAE_Diffusers_Flux2_Config.get_tag()],
+        # ControlNet - checkpoint format
+        Annotated[ControlNet_Checkpoint_SD1_Config, ControlNet_Checkpoint_SD1_Config.get_tag()],
+        Annotated[ControlNet_Checkpoint_SD2_Config, ControlNet_Checkpoint_SD2_Config.get_tag()],
+        Annotated[ControlNet_Checkpoint_SDXL_Config, ControlNet_Checkpoint_SDXL_Config.get_tag()],
+        Annotated[ControlNet_Checkpoint_FLUX_Config, ControlNet_Checkpoint_FLUX_Config.get_tag()],
+        Annotated[ControlNet_Checkpoint_ZImage_Config, ControlNet_Checkpoint_ZImage_Config.get_tag()],
+        Annotated[ControlNet_Checkpoint_Anima_Config, ControlNet_Checkpoint_Anima_Config.get_tag()],
+        # ControlNet - diffusers format
+        Annotated[ControlNet_Diffusers_SD1_Config, ControlNet_Diffusers_SD1_Config.get_tag()],
+        Annotated[ControlNet_Diffusers_SD2_Config, ControlNet_Diffusers_SD2_Config.get_tag()],
+        Annotated[ControlNet_Diffusers_SDXL_Config, ControlNet_Diffusers_SDXL_Config.get_tag()],
+        Annotated[ControlNet_Diffusers_FLUX_Config, ControlNet_Diffusers_FLUX_Config.get_tag()],
+        # LoRA - LyCORIS format
+        # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
+        # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
+        Annotated[LoRA_LyCORIS_SD1_Config, LoRA_LyCORIS_SD1_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_SD2_Config, LoRA_LyCORIS_SD2_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_SDXL_Config, LoRA_LyCORIS_SDXL_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_Flux2_Config, LoRA_LyCORIS_Flux2_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_FLUX_Config, LoRA_LyCORIS_FLUX_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_ZImage_Config, LoRA_LyCORIS_ZImage_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_QwenImage_Config, LoRA_LyCORIS_QwenImage_Config.get_tag()],
+        Annotated[LoRA_LyCORIS_Anima_Config, LoRA_LyCORIS_Anima_Config.get_tag()],
+        # LoRA - OMI format
+        Annotated[LoRA_OMI_SDXL_Config, LoRA_OMI_SDXL_Config.get_tag()],
+        Annotated[LoRA_OMI_FLUX_Config, LoRA_OMI_FLUX_Config.get_tag()],
+        # LoRA - diffusers format
+        # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation
+        # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models
+        Annotated[LoRA_Diffusers_SD1_Config, LoRA_Diffusers_SD1_Config.get_tag()],
+        Annotated[LoRA_Diffusers_SD2_Config, LoRA_Diffusers_SD2_Config.get_tag()],
+        Annotated[LoRA_Diffusers_SDXL_Config, LoRA_Diffusers_SDXL_Config.get_tag()],
+        Annotated[LoRA_Diffusers_Flux2_Config, LoRA_Diffusers_Flux2_Config.get_tag()],
+        Annotated[LoRA_Diffusers_FLUX_Config, LoRA_Diffusers_FLUX_Config.get_tag()],
+        Annotated[LoRA_Diffusers_ZImage_Config, LoRA_Diffusers_ZImage_Config.get_tag()],
+        # ControlLoRA - diffusers format
+        Annotated[ControlLoRA_LyCORIS_FLUX_Config, ControlLoRA_LyCORIS_FLUX_Config.get_tag()],
+        # T5 Encoder - all formats
+        Annotated[T5Encoder_T5Encoder_Config, T5Encoder_T5Encoder_Config.get_tag()],
+        Annotated[T5Encoder_BnBLLMint8_Config, T5Encoder_BnBLLMint8_Config.get_tag()],
+        # Qwen3 Encoder
+        Annotated[Qwen3Encoder_Qwen3Encoder_Config, Qwen3Encoder_Qwen3Encoder_Config.get_tag()],
+        Annotated[Qwen3Encoder_Checkpoint_Config, Qwen3Encoder_Checkpoint_Config.get_tag()],
+        Annotated[Qwen3Encoder_GGUF_Config, Qwen3Encoder_GGUF_Config.get_tag()],
+        # Qwen VL Encoder (Qwen2.5-VL multimodal encoder for Qwen Image)
+        Annotated[QwenVLEncoder_Diffusers_Config, QwenVLEncoder_Diffusers_Config.get_tag()],
+        Annotated[QwenVLEncoder_Checkpoint_Config, QwenVLEncoder_Checkpoint_Config.get_tag()],
+        # TI - file format
+        Annotated[TI_File_SD1_Config, TI_File_SD1_Config.get_tag()],
+        Annotated[TI_File_SD2_Config, TI_File_SD2_Config.get_tag()],
+        Annotated[TI_File_SDXL_Config, TI_File_SDXL_Config.get_tag()],
+        # TI - folder format
+        Annotated[TI_Folder_SD1_Config, TI_Folder_SD1_Config.get_tag()],
+        Annotated[TI_Folder_SD2_Config, TI_Folder_SD2_Config.get_tag()],
+        Annotated[TI_Folder_SDXL_Config, TI_Folder_SDXL_Config.get_tag()],
+        # IP Adapter - InvokeAI format
+        Annotated[IPAdapter_InvokeAI_SD1_Config, IPAdapter_InvokeAI_SD1_Config.get_tag()],
+        Annotated[IPAdapter_InvokeAI_SD2_Config, IPAdapter_InvokeAI_SD2_Config.get_tag()],
+        Annotated[IPAdapter_InvokeAI_SDXL_Config, IPAdapter_InvokeAI_SDXL_Config.get_tag()],
+        # IP Adapter - checkpoint format
+        Annotated[IPAdapter_Checkpoint_SD1_Config, IPAdapter_Checkpoint_SD1_Config.get_tag()],
+        Annotated[IPAdapter_Checkpoint_SD2_Config, IPAdapter_Checkpoint_SD2_Config.get_tag()],
+        Annotated[IPAdapter_Checkpoint_SDXL_Config, IPAdapter_Checkpoint_SDXL_Config.get_tag()],
+        Annotated[IPAdapter_Checkpoint_FLUX_Config, IPAdapter_Checkpoint_FLUX_Config.get_tag()],
+        # T2I Adapter - diffusers format
+        Annotated[T2IAdapter_Diffusers_SD1_Config, T2IAdapter_Diffusers_SD1_Config.get_tag()],
+        Annotated[T2IAdapter_Diffusers_SDXL_Config, T2IAdapter_Diffusers_SDXL_Config.get_tag()],
+        # Misc models
+        Annotated[Spandrel_Checkpoint_Config, Spandrel_Checkpoint_Config.get_tag()],
+        Annotated[CLIPEmbed_Diffusers_G_Config, CLIPEmbed_Diffusers_G_Config.get_tag()],
+        Annotated[CLIPEmbed_Diffusers_L_Config, CLIPEmbed_Diffusers_L_Config.get_tag()],
+        Annotated[CLIPVision_Diffusers_Config, CLIPVision_Diffusers_Config.get_tag()],
+        Annotated[SigLIP_Diffusers_Config, SigLIP_Diffusers_Config.get_tag()],
+        Annotated[FLUXRedux_Checkpoint_Config, FLUXRedux_Checkpoint_Config.get_tag()],
+        Annotated[LlavaOnevision_Diffusers_Config, LlavaOnevision_Diffusers_Config.get_tag()],
+        Annotated[TextLLM_Diffusers_Config, TextLLM_Diffusers_Config.get_tag()],
+        Annotated[ExternalApiModelConfig, ExternalApiModelConfig.get_tag()],
+        # Unknown model (fallback)
+        Annotated[Unknown_Config, Unknown_Config.get_tag()],
+    ],
+    Discriminator(Config_Base.get_model_discriminator_value),
+]
+
+AnyModelConfigValidator = TypeAdapter[AnyModelConfig](AnyModelConfig)
+"""Pydantic TypeAdapter for the AnyModelConfig union, used for parsing and validation.
+
+If you need to parse/validate a dict or JSON into an AnyModelConfig, you should probably use
+ModelConfigFactory.from_dict or ModelConfigFactory.from_json instead as they may implement
+additional logic in the future.
+"""
+
+
+@dataclass
+class ModelClassificationResult:
+    """Result of attempting to classify a model on disk into a specific model config.
+
+    Attributes:
+        match: The best matching model config, or None if no match was found.
+        results: A mapping of model config class names to either an instance of that class (if it matched)
+            or an Exception (if it didn't match or an error occurred during matching).
+    """
+
+    config: AnyModelConfig | None
+    details: dict[str, AnyModelConfig | Exception]
+
+    @property
+    def all_matches(self) -> list[AnyModelConfig]:
+        """Returns a list of all matching model configs found."""
+        return [r for r in self.details.values() if isinstance(r, Config_Base)]
+
+    @property
+    def match_count(self) -> int:
+        """Returns the number of matching model configs found."""
+        return len(self.all_matches)
+
+
+class ModelConfigFactory:
+    @staticmethod
+    def from_dict(fields: dict[str, Any]) -> AnyModelConfig:
+        """Return the appropriate config object from raw dict values."""
+        model = AnyModelConfigValidator.validate_python(fields)
+        return model
+
+    @staticmethod
+    def from_json(json: str | bytes | bytearray) -> AnyModelConfig:
+        """Return the appropriate config object from json."""
+        model = AnyModelConfigValidator.validate_json(json)
+        return model
+
+    @staticmethod
+    def build_common_fields(
+        mod: ModelOnDisk,
+        override_fields: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Builds the common fields for all model configs.
+
+        Args:
+            mod: The model on disk to extract fields from.
+            overrides: A optional dictionary of fields to override. These fields will take precedence over the values
+                extracted from the model on disk.
+
+        - Casts string fields to their Enum types.
+        - Does not validate the fields against the model config schema.
+        """
+
+        _overrides: dict[str, Any] = override_fields or {}
+        fields: dict[str, Any] = {}
+
+        if "type" in _overrides:
+            fields["type"] = ModelType(_overrides["type"])
+
+        if "format" in _overrides:
+            fields["format"] = ModelFormat(_overrides["format"])
+
+        if "base" in _overrides:
+            fields["base"] = BaseModelType(_overrides["base"])
+
+        if "source_type" in _overrides:
+            fields["source_type"] = ModelSourceType(_overrides["source_type"])
+
+        if "variant" in _overrides:
+            fields["variant"] = variant_type_adapter.validate_strings(_overrides["variant"])
+
+        fields["path"] = mod.path.as_posix()
+        fields["source"] = _overrides.get("source") or fields["path"]
+        fields["source_type"] = _overrides.get("source_type") or ModelSourceType.Path
+        fields["name"] = _overrides.get("name") or mod.name
+        fields["hash"] = _overrides.get("hash") or mod.hash()
+        fields["key"] = _overrides.get("key") or uuid_string()
+        fields["description"] = _overrides.get("description")
+        fields["file_size"] = _overrides.get("file_size") or mod.size()
+
+        return fields
+
+    @staticmethod
+    def _validate_path_looks_like_model(path: Path) -> None:
+        """Perform basic sanity checks to ensure a path looks like a model.
+
+        This prevents wasting time trying to identify obviously non-model paths like
+        home directories or downloads folders. Raises RuntimeError if the path doesn't
+        pass basic checks.
+
+        Args:
+            path: The path to validate
+
+        Raises:
+            ValueError: If the path doesn't look like a model
+        """
+        if path.is_file():
+            # For files, just check the extension
+            if path.suffix.lower() not in _MODEL_EXTENSIONS:
+                raise ValueError(
+                    f"File extension {path.suffix} is not a recognized model format. "
+                    f"Expected one of: {', '.join(sorted(_MODEL_EXTENSIONS))}"
+                )
+        else:
+            # For directories, do a quick file count check with early exit
+            total_files = 0
+            # Ignore hidden files and directories
+            paths_to_check = (
+                p
+                for p in path.rglob("*")
+                if not p.name.startswith(".") and not any(part.startswith(".") for part in p.parts)
+            )
+            for item in paths_to_check:
+                if item.is_file():
+                    total_files += 1
+                    if total_files > _MAX_FILES_IN_MODEL_DIR:
+                        raise ValueError(
+                            f"Directory contains more than {_MAX_FILES_IN_MODEL_DIR} files. "
+                            "This looks like a general-purpose directory rather than a model. "
+                            "Please provide a path to a specific model file or model directory."
+                        )
+
+            # Check if it has config files at root (diffusers/transformers marker)
+            has_root_config = any((path / config).exists() for config in _CONFIG_FILES)
+
+            if has_root_config:
+                # Has a config file, looks like a valid model directory
+                return
+
+            # Otherwise, search for model files within depth limit
+            def find_model_files(current_path: Path, depth: int) -> bool:
+                if depth > _MAX_SEARCH_DEPTH:
+                    return False
+                try:
+                    for item in current_path.iterdir():
+                        if item.is_file() and item.suffix.lower() in _MODEL_EXTENSIONS:
+                            return True
+                        elif item.is_dir() and find_model_files(item, depth + 1):
+                            return True
+                except PermissionError:
+                    pass
+                return False
+
+            if not find_model_files(path, 0):
+                raise ValueError(
+                    f"No model files or config files found in directory {path}. "
+                    f"Expected to find model files with extensions: {', '.join(sorted(_MODEL_EXTENSIONS))} "
+                    f"or config files: {', '.join(sorted(_CONFIG_FILES))}"
+                )
+
+    @staticmethod
+    def matches_sort_key(m: AnyModelConfig) -> int:
+        """Sort key function to prioritize model config matches in case of multiple matches."""
+
+        # It is possible that we have multiple matches. We need to prioritize them.
+
+        # Known cases where multiple matches can occur:
+        # - SD main models can look like a LoRA when they have merged in LoRA weights. Prefer the main model.
+        # - SD main models in diffusers format can look like a CLIP Embed; they have a text_encoder folder with
+        #   a config.json file. Prefer the main model.
+
+        # Given the above cases, we can prioritize the matches by type. If we find more cases, we may need a more
+        # sophisticated approach.
+        match m.type:
+            case ModelType.Main:
+                return 0
+            case ModelType.LoRA:
+                return 1
+            case ModelType.CLIPEmbed:
+                return 2
+            case _:
+                return 3
+
+    @staticmethod
+    def from_model_on_disk(
+        mod: str | Path | ModelOnDisk,
+        override_fields: dict[str, Any] | None = None,
+        hash_algo: HASHING_ALGORITHMS = "blake3_single",
+        allow_unknown: bool = True,
+    ) -> ModelClassificationResult:
+        """Classify a model on disk and return the best matching model config.
+
+        Args:
+            mod: The model on disk to classify. Can be a path (str or Path) or a ModelOnDisk instance.
+            override_fields: Optional dictionary of fields to override. These fields will take precedence
+                over the values extracted from the model on disk, but this cannot force a match if the
+                model on disk doesn't actually match the config class.
+            hash_algo: The hashing algorithm to use when computing the model hash if needed.
+
+        Returns:
+            A ModelClassificationResult containing the best matching model config (or None if no match)
+            and a mapping of all attempted model config classes to either an instance of that class (if it matched)
+            or an Exception (if it didn't match or an error occurred during matching).
+
+        Raises:
+            ValueError: If the provided path doesn't look like a model.
+        """
+        if isinstance(mod, Path | str):
+            mod = ModelOnDisk(Path(mod), hash_algo)
+
+        # Perform basic sanity checks before attempting any config matching
+        # This rejects obviously non-model paths early, saving time
+        ModelConfigFactory._validate_path_looks_like_model(mod.path)
+
+        # We will always need these fields to build any model config.
+        fields = ModelConfigFactory.build_common_fields(mod, override_fields)
+
+        # Store results as a mapping of config class to either an instance of that class or an exception
+        # that was raised when trying to build it.
+        details: dict[str, AnyModelConfig | Exception] = {}
+
+        # Try to build an instance of each model config class that uses the classify API.
+        # Each class will either return an instance of itself or raise NotAMatch if it doesn't match.
+        # Other exceptions may be raised if something unexpected happens during matching or building.
+        for candidate_class in filter(lambda x: x is not Unknown_Config, Config_Base.CONFIG_CLASSES):
+            candidate_name = candidate_class.__name__
+            try:
+                # Technically, from_model_on_disk returns a Config_Base, but in practice it will always be a member of
+                # the AnyModelConfig union.
+                details[candidate_name] = candidate_class.from_model_on_disk(mod, fields)  # type: ignore
+            except NotAMatchError as e:
+                # This means the model didn't match this config class. It's not an error, just no match.
+                details[candidate_name] = e
+            except ValidationError as e:
+                # This means the model matched, but we couldn't create the pydantic model instance for the config.
+                # Maybe invalid overrides were provided?
+                details[candidate_name] = e
+            except Exception as e:
+                # Some other unexpected error occurred. Store the exception for reporting later.
+                details[candidate_name] = e
+
+        # Extract just the successful matches
+        matches = [r for r in details.values() if isinstance(r, Config_Base)]
+
+        if not matches:
+            if not allow_unknown:
+                # No matches and we are not allowed to fall back to Unknown_Config
+                return ModelClassificationResult(config=None, details=details)
+            else:
+                # Fall back to Unknown_Config
+                # This should always succeed as Unknown_Config.from_model_on_disk never raises NotAMatch
+                config = Unknown_Config.from_model_on_disk(mod, fields)
+                details[Unknown_Config.__name__] = config
+                return ModelClassificationResult(config=config, details=details)
+
+        matches.sort(key=ModelConfigFactory.matches_sort_key)
+        config = matches[0]
+
+        # Now do any post-processing needed for specific model types/bases/etc.
+        match config.type:
+            case ModelType.Main:
+                # Pass variant if available (e.g., for Flux2 models)
+                variant = getattr(config, "variant", None)
+                config.default_settings = MainModelDefaultSettings.from_base(config.base, variant)
+            case ModelType.ControlNet | ModelType.T2IAdapter | ModelType.ControlLoRa:
+                config.default_settings = ControlAdapterDefaultSettings.from_model_name(config.name)
+            case ModelType.LoRA:
+                config.default_settings = LoraModelDefaultSettings()
+            case _:
+                pass
+
+        return ModelClassificationResult(config=config, details=details)
+
+
+MODEL_NAME_TO_PREPROCESSOR = {
+    "canny": "canny_image_processor",
+    "mlsd": "mlsd_image_processor",
+    "depth": "depth_anything_image_processor",
+    "bae": "normalbae_image_processor",
+    "normal": "normalbae_image_processor",
+    "sketch": "pidi_image_processor",
+    "scribble": "lineart_image_processor",
+    "lineart anime": "lineart_anime_image_processor",
+    "lineart_anime": "lineart_anime_image_processor",
+    "lineart": "lineart_image_processor",
+    "soft": "hed_image_processor",
+    "softedge": "hed_image_processor",
+    "hed": "hed_image_processor",
+    "shuffle": "content_shuffle_image_processor",
+    "pose": "dw_openpose_image_processor",
+    "mediapipe": "mediapipe_face_processor",
+    "pidi": "pidi_image_processor",
+    "zoe": "zoe_depth_image_processor",
+    "color": "color_map_image_processor",
+}

invokeai/backend/model_manager/configs/flux_redux.py

@@ -0,0 +1,40 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.flux.redux.flux_redux_state_dict_utils import is_state_dict_likely_flux_redux
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_override_fields,
+    raise_if_not_file,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+
+class FLUXRedux_Checkpoint_Config(Config_Base):
+    """Model config for FLUX Tools Redux model."""
+
+    type: Literal[ModelType.FluxRedux] = Field(default=ModelType.FluxRedux)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.Flux] = Field(default=BaseModelType.Flux)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        if not is_state_dict_likely_flux_redux(mod.load_state_dict()):
+            raise NotAMatchError("model does not match FLUX Tools Redux heuristics")
+
+        return cls(**override_fields)

invokeai/backend/model_manager/configs/identification_utils.py

@@ -0,0 +1,206 @@
+import json
+from functools import cache
+from pathlib import Path
+
+from pydantic import BaseModel, ValidationError
+from pydantic_core import CoreSchema, SchemaValidator
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+
+
+class NotAMatchError(Exception):
+    """Exception for when a model does not match a config class.
+
+    Args:
+        reason: The reason why the model did not match.
+    """
+
+    def __init__(self, reason: str):
+        super().__init__(reason)
+
+
+def get_config_dict_or_raise(config_path: Path | set[Path]) -> dict[str, Any]:
+    """Load the diffusers/transformers model config file and return it as a dictionary. The config file is expected
+    to be in JSON format.
+
+    Args:
+        config_path: The path to the config file, or a set of paths to try.
+
+    Returns:
+        The config file as a dictionary.
+
+    Raises:
+        NotAMatch if the config file is missing or cannot be loaded.
+    """
+    paths_to_check = config_path if isinstance(config_path, set) else {config_path}
+
+    problems: dict[Path, str] = {}
+
+    for p in paths_to_check:
+        if not p.exists():
+            problems[p] = "file does not exist"
+            continue
+
+        try:
+            with open(p, "r") as file:
+                config = json.load(file)
+
+            return config
+        except Exception as e:
+            problems[p] = str(e)
+            continue
+
+    raise NotAMatchError(f"unable to load config file(s): {problems}")
+
+
+def get_class_name_from_config_dict_or_raise(config: Path | set[Path] | dict[str, Any]) -> str:
+    """Load the diffusers/transformers model config file and return the class name.
+
+    Args:
+        config_path: The path to the config file, or a set of paths to try.
+
+    Returns:
+        The class name from the config file.
+
+    Raises:
+        NotAMatch if the config file is missing or does not contain a valid class name.
+    """
+
+    if not isinstance(config, dict):
+        config = get_config_dict_or_raise(config)
+
+    try:
+        if "_class_name" in config:
+            # This is a diffusers-style config
+            config_class_name = config["_class_name"]
+        elif "architectures" in config:
+            # This is a transformers-style config
+            config_class_name = config["architectures"][0]
+        else:
+            raise ValueError("missing _class_name or architectures field")
+    except Exception as e:
+        raise NotAMatchError(f"unable to determine class name from config file: {config}") from e
+
+    if not isinstance(config_class_name, str):
+        raise NotAMatchError(f"_class_name or architectures field is not a string: {config_class_name}")
+
+    return config_class_name
+
+
+def raise_for_class_name(config: Path | set[Path] | dict[str, Any], class_name: str | set[str]) -> None:
+    """Get the class name from the config file and raise NotAMatch if it is not in the expected set.
+
+    Args:
+        config_path: The path to the config file, or a set of paths to try.
+        class_name: The expected class name, or a set of expected class names.
+
+    Raises:
+        NotAMatch if the class name is not in the expected set.
+    """
+
+    class_name = {class_name} if isinstance(class_name, str) else class_name
+
+    actual_class_name = get_class_name_from_config_dict_or_raise(config)
+    if actual_class_name not in class_name:
+        raise NotAMatchError(f"invalid class name from config: {actual_class_name}")
+
+
+def raise_for_override_fields(candidate_config_class: type[BaseModel], override_fields: dict[str, Any]) -> None:
+    """Check if the provided override fields are valid for the config class using pydantic.
+
+    For example, if the candidate config class has a field "base" of type Literal[BaseModelType.StableDiffusion1], and
+    the override fields contain "base": BaseModelType.Flux, this function will raise NotAMatch.
+
+    Internally, this function extracts the pydantic schema for each individual override field from the candidate config
+    class and validates the override value against that schema. Post-instantiation validators are not run.
+
+    Args:
+        candidate_config_class: The config class that is being tested.
+        override_fields: The override fields provided by the user.
+
+    Raises:
+        NotAMatch if any override field is invalid for the config class.
+    """
+    for field_name, override_value in override_fields.items():
+        if field_name not in candidate_config_class.model_fields:
+            raise NotAMatchError(f"unknown override field: {field_name}")
+        try:
+            PydanticFieldValidator.validate_field(candidate_config_class, field_name, override_value)
+        except ValidationError as e:
+            raise NotAMatchError(f"invalid override for field '{field_name}': {e}") from e
+
+
+def raise_if_not_file(mod: ModelOnDisk) -> None:
+    """Raise NotAMatch if the model path is not a file."""
+    if not mod.path.is_file():
+        raise NotAMatchError("model path is not a file")
+
+
+def raise_if_not_dir(mod: ModelOnDisk) -> None:
+    """Raise NotAMatch if the model path is not a directory."""
+    if not mod.path.is_dir():
+        raise NotAMatchError("model path is not a directory")
+
+
+def state_dict_has_any_keys_exact(state_dict: dict[str | int, Any], keys: str | set[str]) -> bool:
+    """Returns true if the state dict has any of the specified keys."""
+    _keys = {keys} if isinstance(keys, str) else keys
+    return any(key in state_dict for key in _keys)
+
+
+def state_dict_has_any_keys_starting_with(state_dict: dict[str | int, Any], prefixes: str | set[str]) -> bool:
+    """Returns true if the state dict has any keys starting with any of the specified prefixes."""
+    _prefixes = {prefixes} if isinstance(prefixes, str) else prefixes
+    return any(any(key.startswith(prefix) for prefix in _prefixes) for key in state_dict.keys() if isinstance(key, str))
+
+
+def state_dict_has_any_keys_ending_with(state_dict: dict[str | int, Any], suffixes: str | set[str]) -> bool:
+    """Returns true if the state dict has any keys ending with any of the specified suffixes."""
+    _suffixes = {suffixes} if isinstance(suffixes, str) else suffixes
+    return any(any(key.endswith(suffix) for suffix in _suffixes) for key in state_dict.keys() if isinstance(key, str))
+
+
+def common_config_paths(path: Path) -> set[Path]:
+    """Returns common config file paths for models stored in directories."""
+    return {path / "config.json", path / "model_index.json"}
+
+
+class PydanticFieldValidator:
+    """Utility class for validating individual fields of a Pydantic model without instantiating the whole model.
+
+    See: https://github.com/pydantic/pydantic/discussions/7367#discussioncomment-14213144
+    """
+
+    @staticmethod
+    def find_field_schema(model: type[BaseModel], field_name: str) -> CoreSchema:
+        """Find the Pydantic core schema for a specific field in a model."""
+        schema: CoreSchema = model.__pydantic_core_schema__.copy()
+        # we shallow copied, be careful not to mutate the original schema!
+
+        assert schema["type"] in ["definitions", "model"]
+
+        # find the field schema
+        field_schema = schema["schema"]  # type: ignore
+        while "fields" not in field_schema:
+            field_schema = field_schema["schema"]  # type: ignore
+
+        field_schema = field_schema["fields"][field_name]["schema"]  # type: ignore
+
+        # if the original schema is a definition schema, replace the model schema with the field schema
+        if schema["type"] == "definitions":
+            schema["schema"] = field_schema
+            return schema
+        else:
+            return field_schema
+
+    @cache
+    @staticmethod
+    def get_validator(model: type[BaseModel], field_name: str) -> SchemaValidator:
+        """Get a SchemaValidator for a specific field in a model."""
+        return SchemaValidator(PydanticFieldValidator.find_field_schema(model, field_name))
+
+    @staticmethod
+    def validate_field(model: type[BaseModel], field_name: str, value: Any) -> Any:
+        """Validate a value for a specific field in a model."""
+        return PydanticFieldValidator.get_validator(model, field_name).validate_python(value)

invokeai/backend/model_manager/configs/ip_adapter.py

@@ -0,0 +1,180 @@
+from abc import ABC
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import BaseModel, Field
+from typing_extensions import Any
+
+from invokeai.backend.flux.ip_adapter.state_dict_utils import is_state_dict_xlabs_ip_adapter
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    raise_if_not_file,
+    state_dict_has_any_keys_starting_with,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+
+class IPAdapter_Config_Base(ABC, BaseModel):
+    type: Literal[ModelType.IPAdapter] = Field(default=ModelType.IPAdapter)
+
+
+class IPAdapter_InvokeAI_Config_Base(IPAdapter_Config_Base):
+    """Model config for IP Adapter diffusers format models."""
+
+    format: Literal[ModelFormat.InvokeAI] = Field(default=ModelFormat.InvokeAI)
+
+    # TODO(ryand): Should we deprecate this field? From what I can tell, it hasn't been probed correctly for a long
+    # time. Need to go through the history to make sure I'm understanding this fully.
+    image_encoder_model_id: str = Field()
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_has_weights_file(mod)
+
+        cls._validate_has_image_encoder_metadata_file(mod)
+
+        cls._validate_base(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _validate_has_weights_file(cls, mod: ModelOnDisk) -> None:
+        weights_file = mod.path / "ip_adapter.bin"
+        if not weights_file.exists():
+            raise NotAMatchError("missing ip_adapter.bin weights file")
+
+    @classmethod
+    def _validate_has_image_encoder_metadata_file(cls, mod: ModelOnDisk) -> None:
+        image_encoder_metadata_file = mod.path / "image_encoder.txt"
+        if not image_encoder_metadata_file.exists():
+            raise NotAMatchError("missing image_encoder.txt metadata file")
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
+        state_dict = mod.load_state_dict()
+
+        try:
+            cross_attention_dim = state_dict["ip_adapter"]["1.to_k_ip.weight"].shape[-1]
+        except Exception as e:
+            raise NotAMatchError(f"unable to determine cross attention dimension: {e}") from e
+
+        match cross_attention_dim:
+            case 768:
+                return BaseModelType.StableDiffusion1
+            case 1024:
+                return BaseModelType.StableDiffusion2
+            case 2048:
+                return BaseModelType.StableDiffusionXL
+            case _:
+                raise NotAMatchError(f"unrecognized cross attention dimension {cross_attention_dim}")
+
+
+class IPAdapter_InvokeAI_SD1_Config(IPAdapter_InvokeAI_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class IPAdapter_InvokeAI_SD2_Config(IPAdapter_InvokeAI_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class IPAdapter_InvokeAI_SDXL_Config(IPAdapter_InvokeAI_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class IPAdapter_Checkpoint_Config_Base(IPAdapter_Config_Base):
+    """Model config for IP Adapter checkpoint format models."""
+
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_ip_adapter(mod)
+
+        cls._validate_base(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _validate_looks_like_ip_adapter(cls, mod: ModelOnDisk) -> None:
+        if not state_dict_has_any_keys_starting_with(
+            mod.load_state_dict(),
+            {
+                "image_proj.",
+                "ip_adapter.",
+                # XLabs FLUX IP-Adapter models have keys startinh with "ip_adapter_proj_model.".
+                "ip_adapter_proj_model.",
+            },
+        ):
+            raise NotAMatchError("model does not match Checkpoint IP Adapter heuristics")
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
+        state_dict = mod.load_state_dict()
+
+        if is_state_dict_xlabs_ip_adapter(state_dict):
+            return BaseModelType.Flux
+
+        try:
+            cross_attention_dim = state_dict["ip_adapter.1.to_k_ip.weight"].shape[-1]
+        except Exception as e:
+            raise NotAMatchError(f"unable to determine cross attention dimension: {e}") from e
+
+        match cross_attention_dim:
+            case 768:
+                return BaseModelType.StableDiffusion1
+            case 1024:
+                return BaseModelType.StableDiffusion2
+            case 2048:
+                return BaseModelType.StableDiffusionXL
+            case _:
+                raise NotAMatchError(f"unrecognized cross attention dimension {cross_attention_dim}")
+
+
+class IPAdapter_Checkpoint_SD1_Config(IPAdapter_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class IPAdapter_Checkpoint_SD2_Config(IPAdapter_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class IPAdapter_Checkpoint_SDXL_Config(IPAdapter_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class IPAdapter_Checkpoint_FLUX_Config(IPAdapter_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.Flux] = Field(default=BaseModelType.Flux)

invokeai/backend/model_manager/configs/llava_onevision.py

@@ -0,0 +1,43 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    common_config_paths,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelType,
+)
+
+
+class LlavaOnevision_Diffusers_Config(Diffusers_Config_Base, Config_Base):
+    """Model config for Llava Onevision models."""
+
+    type: Literal[ModelType.LlavaOnevision] = Field(default=ModelType.LlavaOnevision)
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "LlavaOnevisionForConditionalGeneration",
+            },
+        )
+
+        return cls(**override_fields)

invokeai/backend/model_manager/configs/qwen3_encoder.py

@@ -0,0 +1,311 @@
+import json
+from typing import Any, Literal, Optional, Self
+
+from pydantic import Field
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    raise_if_not_file,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, Qwen3VariantType
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+def _has_qwen3_keys(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict contains Qwen3 model keys.
+
+    Supports both:
+    - PyTorch/diffusers format: model.layers.0., model.embed_tokens.weight
+    - GGUF/llama.cpp format: blk.0., token_embd.weight
+    """
+    # PyTorch/diffusers format indicators
+    pytorch_indicators = ["model.layers.0.", "model.embed_tokens.weight"]
+    # GGUF/llama.cpp format indicators
+    gguf_indicators = ["blk.0.", "token_embd.weight"]
+
+    for key in state_dict.keys():
+        if isinstance(key, str):
+            # Check PyTorch format
+            for indicator in pytorch_indicators:
+                if key.startswith(indicator) or key == indicator:
+                    return True
+            # Check GGUF format
+            for indicator in gguf_indicators:
+                if key.startswith(indicator) or key == indicator:
+                    return True
+    return False
+
+
+def _has_ggml_tensors(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict contains GGML tensors (GGUF quantized)."""
+    return any(isinstance(v, GGMLTensor) for v in state_dict.values())
+
+
+def _has_qwen_vl_visual_tower(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict bundles a Qwen2.5-VL / Qwen2-VL vision tower.
+
+    Qwen-VL encoders ship the visual tower (`visual.blocks.*`, `visual.patch_embed.*`)
+    alongside the language model, whereas a text-only Qwen3 encoder never does. A Qwen-VL
+    file otherwise satisfies the Qwen3 key heuristic (it has `model.layers.*` /
+    `model.embed_tokens.weight` too), so without this check it matches *both* the Qwen3 and
+    the QwenVLEncoder configs and the tiebreak can misroute it to Qwen3. We use it to keep
+    the two mutually exclusive.
+    """
+    for key in state_dict.keys():
+        if isinstance(key, str) and (key.startswith("visual.blocks.") or key.startswith("visual.patch_embed.")):
+            return True
+    return False
+
+
+def _get_qwen3_variant_from_state_dict(state_dict: dict[str | int, Any]) -> Optional[Qwen3VariantType]:
+    """Determine Qwen3 variant (0.6B, 4B, or 8B) from state dict based on hidden_size.
+
+    The hidden_size can be determined from the embed_tokens.weight tensor shape:
+    - Qwen3 0.6B: hidden_size = 1024
+    - Qwen3 4B: hidden_size = 2560
+    - Qwen3 8B: hidden_size = 4096
+
+    For GGUF format, the key is 'token_embd.weight'.
+    For PyTorch format, the key is 'model.embed_tokens.weight'.
+    """
+    # Hidden size thresholds
+    QWEN3_06B_HIDDEN_SIZE = 1024
+    QWEN3_4B_HIDDEN_SIZE = 2560
+    QWEN3_8B_HIDDEN_SIZE = 4096
+
+    # Try to find embed_tokens weight
+    embed_key = None
+    for key in state_dict.keys():
+        if isinstance(key, str):
+            if key == "model.embed_tokens.weight" or key == "token_embd.weight":
+                embed_key = key
+                break
+
+    if embed_key is None:
+        return None
+
+    tensor = state_dict[embed_key]
+
+    # Get hidden_size from tensor shape
+    # Shape is [vocab_size, hidden_size]
+    if isinstance(tensor, GGMLTensor):
+        # GGUF tensor
+        if hasattr(tensor, "shape") and len(tensor.shape) >= 2:
+            hidden_size = tensor.shape[1]
+        else:
+            return None
+    elif hasattr(tensor, "shape"):
+        # PyTorch tensor
+        if len(tensor.shape) >= 2:
+            hidden_size = tensor.shape[1]
+        else:
+            return None
+    else:
+        return None
+
+    # Determine variant based on hidden_size
+    if hidden_size == QWEN3_06B_HIDDEN_SIZE:
+        return Qwen3VariantType.Qwen3_06B
+    elif hidden_size == QWEN3_4B_HIDDEN_SIZE:
+        return Qwen3VariantType.Qwen3_4B
+    elif hidden_size == QWEN3_8B_HIDDEN_SIZE:
+        return Qwen3VariantType.Qwen3_8B
+    else:
+        # Unknown size, default to 4B (more common)
+        return Qwen3VariantType.Qwen3_4B
+
+
+class Qwen3Encoder_Checkpoint_Config(Checkpoint_Config_Base, Config_Base):
+    """Configuration for single-file Qwen3 Encoder models (safetensors)."""
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.Qwen3Encoder] = Field(default=ModelType.Qwen3Encoder)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+    variant: Qwen3VariantType = Field(description="Qwen3 model size variant (4B or 8B)")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_qwen3_model(mod)
+
+        cls._validate_does_not_look_like_gguf_quantized(mod)
+
+        # Determine variant from state dict
+        variant = cls._get_variant_or_default(mod)
+
+        return cls(variant=variant, **override_fields)
+
+    @classmethod
+    def _get_variant_or_default(cls, mod: ModelOnDisk) -> Qwen3VariantType:
+        """Get variant from state dict, defaulting to 4B if unknown."""
+        state_dict = mod.load_state_dict()
+        variant = _get_qwen3_variant_from_state_dict(state_dict)
+        return variant if variant is not None else Qwen3VariantType.Qwen3_4B
+
+    @classmethod
+    def _validate_looks_like_qwen3_model(cls, mod: ModelOnDisk) -> None:
+        state_dict = mod.load_state_dict()
+        if not _has_qwen3_keys(state_dict):
+            raise NotAMatchError("state dict does not look like a Qwen3 model")
+        # Reject Qwen2.5-VL / Qwen2-VL encoders: they carry a visual tower and must be
+        # classified as QwenVLEncoder (text-only Qwen3 encoders never have one).
+        if _has_qwen_vl_visual_tower(state_dict):
+            raise NotAMatchError(
+                "state dict bundles a Qwen-VL visual tower; this is a Qwen-VL encoder, not a text-only Qwen3 encoder"
+            )
+
+    @classmethod
+    def _validate_does_not_look_like_gguf_quantized(cls, mod: ModelOnDisk) -> None:
+        has_ggml = _has_ggml_tensors(mod.load_state_dict())
+        if has_ggml:
+            raise NotAMatchError("state dict looks like GGUF quantized")
+
+
+class Qwen3Encoder_Qwen3Encoder_Config(Config_Base):
+    """Configuration for Qwen3 Encoder models in a diffusers-like format.
+
+    The model weights are expected to be in a folder called text_encoder inside the model directory,
+    compatible with Qwen2VLForConditionalGeneration or similar architectures used by Z-Image.
+    """
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.Qwen3Encoder] = Field(default=ModelType.Qwen3Encoder)
+    format: Literal[ModelFormat.Qwen3Encoder] = Field(default=ModelFormat.Qwen3Encoder)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+    variant: Qwen3VariantType = Field(description="Qwen3 model size variant (4B or 8B)")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        # Exclude full pipeline models - these should be matched as main models, not just Qwen3 encoders.
+        # Full pipelines have model_index.json at root (diffusers format) or a transformer subfolder.
+        model_index_path = mod.path / "model_index.json"
+        transformer_path = mod.path / "transformer"
+        if model_index_path.exists() or transformer_path.exists():
+            raise NotAMatchError(
+                "directory looks like a full diffusers pipeline (has model_index.json or transformer folder), "
+                "not a standalone Qwen3 encoder"
+            )
+
+        # Check for text_encoder config - support both:
+        # 1. Full model structure: model_root/text_encoder/config.json
+        # 2. Standalone text_encoder download: model_root/config.json (when text_encoder subfolder is downloaded separately)
+        config_path_nested = mod.path / "text_encoder" / "config.json"
+        config_path_direct = mod.path / "config.json"
+
+        if config_path_nested.exists():
+            expected_config_path = config_path_nested
+        elif config_path_direct.exists():
+            # Standalone text_encoder downloads do not bundle tokenizer files. If we see tokenizer files at the
+            # root next to config.json, this is a complete causal LM (TextLLM), not a Qwen3 encoder subfolder.
+            tokenizer_files = ("tokenizer.json", "tokenizer.model", "tokenizer_config.json")
+            if any((mod.path / f).exists() for f in tokenizer_files):
+                raise NotAMatchError(
+                    "directory looks like a complete causal LM (config.json and tokenizer files at root), "
+                    "not a standalone Qwen3 encoder"
+                )
+            expected_config_path = config_path_direct
+        else:
+            raise NotAMatchError(
+                f"unable to load config file(s): {{PosixPath('{config_path_nested}'): 'file does not exist'}}"
+            )
+
+        # Qwen3 uses Qwen2VLForConditionalGeneration or similar
+        raise_for_class_name(
+            expected_config_path,
+            {
+                "Qwen2VLForConditionalGeneration",
+                "Qwen2ForCausalLM",
+                "Qwen3ForCausalLM",
+            },
+        )
+
+        # Determine variant from config.json hidden_size
+        variant = cls._get_variant_from_config(expected_config_path)
+
+        return cls(variant=variant, **override_fields)
+
+    @classmethod
+    def _get_variant_from_config(cls, config_path) -> Qwen3VariantType:
+        """Get variant from config.json based on hidden_size."""
+        QWEN3_06B_HIDDEN_SIZE = 1024
+        QWEN3_4B_HIDDEN_SIZE = 2560
+        QWEN3_8B_HIDDEN_SIZE = 4096
+
+        try:
+            with open(config_path, "r", encoding="utf-8") as f:
+                config = json.load(f)
+            hidden_size = config.get("hidden_size")
+            if hidden_size == QWEN3_8B_HIDDEN_SIZE:
+                return Qwen3VariantType.Qwen3_8B
+            elif hidden_size == QWEN3_4B_HIDDEN_SIZE:
+                return Qwen3VariantType.Qwen3_4B
+            elif hidden_size == QWEN3_06B_HIDDEN_SIZE:
+                return Qwen3VariantType.Qwen3_06B
+            else:
+                # Default to 4B for unknown sizes
+                return Qwen3VariantType.Qwen3_4B
+        except (json.JSONDecodeError, OSError):
+            return Qwen3VariantType.Qwen3_4B
+
+
+class Qwen3Encoder_GGUF_Config(Checkpoint_Config_Base, Config_Base):
+    """Configuration for GGUF-quantized Qwen3 Encoder models."""
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.Qwen3Encoder] = Field(default=ModelType.Qwen3Encoder)
+    format: Literal[ModelFormat.GGUFQuantized] = Field(default=ModelFormat.GGUFQuantized)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+    variant: Qwen3VariantType = Field(description="Qwen3 model size variant (4B or 8B)")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_qwen3_model(mod)
+
+        cls._validate_looks_like_gguf_quantized(mod)
+
+        # Determine variant from state dict
+        variant = cls._get_variant_or_default(mod)
+
+        return cls(variant=variant, **override_fields)
+
+    @classmethod
+    def _get_variant_or_default(cls, mod: ModelOnDisk) -> Qwen3VariantType:
+        """Get variant from state dict, defaulting to 4B if unknown."""
+        state_dict = mod.load_state_dict()
+        variant = _get_qwen3_variant_from_state_dict(state_dict)
+        return variant if variant is not None else Qwen3VariantType.Qwen3_4B
+
+    @classmethod
+    def _validate_looks_like_qwen3_model(cls, mod: ModelOnDisk) -> None:
+        state_dict = mod.load_state_dict()
+        if not _has_qwen3_keys(state_dict):
+            raise NotAMatchError("state dict does not look like a Qwen3 model")
+        # Reject Qwen2.5-VL / Qwen2-VL encoders: they carry a visual tower and must be
+        # classified as QwenVLEncoder (text-only Qwen3 encoders never have one).
+        if _has_qwen_vl_visual_tower(state_dict):
+            raise NotAMatchError(
+                "state dict bundles a Qwen-VL visual tower; this is a Qwen-VL encoder, not a text-only Qwen3 encoder"
+            )
+
+    @classmethod
+    def _validate_looks_like_gguf_quantized(cls, mod: ModelOnDisk) -> None:
+        has_ggml = _has_ggml_tensors(mod.load_state_dict())
+        if not has_ggml:
+            raise NotAMatchError("state dict does not look like GGUF quantized")

invokeai/backend/model_manager/configs/qwen_vl_encoder.py

@@ -0,0 +1,154 @@
+import json
+from pathlib import Path
+from typing import Any, Iterable, Literal, Self
+
+from pydantic import Field
+from safetensors import safe_open
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    raise_if_not_file,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
+
+_RECOGNIZED_TEXT_ENCODER_CLASSES = {
+    "Qwen2_5_VLForConditionalGeneration",
+    "Qwen2VLForConditionalGeneration",
+}
+
+
+def _has_qwen_vl_keys(keys: Iterable[str]) -> bool:
+    """A Qwen2.5-VL/Qwen2-VL checkpoint must have both LM weights and a visual
+    tower — that's what distinguishes it from text-only Qwen3/Qwen2 encoders."""
+    has_lm = False
+    has_vision = False
+    for k in keys:
+        if not isinstance(k, str):
+            continue
+        if not has_lm and (k == "model.embed_tokens.weight" or k.startswith("model.layers.")):
+            has_lm = True
+        if not has_vision and (k.startswith("visual.patch_embed.") or k.startswith("visual.blocks.")):
+            has_vision = True
+        if has_lm and has_vision:
+            return True
+    return False
+
+
+def _read_safetensors_keys(path: Path) -> list[str]:
+    """Read only the key index from a safetensors file without loading tensor data.
+
+    Avoids holding multi-GB encoder weights in RAM just to classify the file.
+    """
+    with safe_open(str(path), framework="pt", device="cpu") as f:
+        return list(f.keys())
+
+
+class QwenVLEncoder_Diffusers_Config(Config_Base):
+    """Configuration for standalone Qwen2.5-VL encoder models in diffusers-style folder layout.
+
+    Expected structure:
+        <model_root>/
+            text_encoder/
+                config.json (with `_class_name` or `architectures` listing
+                             `Qwen2_5_VLForConditionalGeneration`)
+                model.safetensors
+            tokenizer/
+                tokenizer_config.json
+                ...
+            processor/                  (optional, for vision preprocessing)
+                preprocessor_config.json
+
+    This lets users avoid downloading the full ~40 GB Qwen Image diffusers pipeline
+    when they only need the Qwen2.5-VL encoder for use with a GGUF transformer.
+    """
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.QwenVLEncoder] = Field(default=ModelType.QwenVLEncoder)
+    format: Literal[ModelFormat.QwenVLEncoder] = Field(default=ModelFormat.QwenVLEncoder)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        # Reject anything that looks like a full pipeline (those are matched as Main models).
+        if (mod.path / "model_index.json").exists() or (mod.path / "transformer").exists():
+            raise NotAMatchError(
+                "directory looks like a full diffusers pipeline (has model_index.json or transformer folder), "
+                "not a standalone Qwen VL encoder"
+            )
+
+        text_encoder_dir = mod.path / "text_encoder"
+        tokenizer_dir = mod.path / "tokenizer"
+
+        if not text_encoder_dir.is_dir():
+            raise NotAMatchError("missing text_encoder/ subfolder")
+        if not tokenizer_dir.is_dir():
+            raise NotAMatchError("missing tokenizer/ subfolder")
+
+        config_path = text_encoder_dir / "config.json"
+        if not config_path.is_file():
+            raise NotAMatchError(f"missing {config_path}")
+
+        try:
+            with open(config_path, "r", encoding="utf-8") as f:
+                cfg = json.load(f)
+        except (OSError, json.JSONDecodeError) as e:
+            raise NotAMatchError(f"could not read text_encoder/config.json: {e}") from e
+
+        class_name = cfg.get("_class_name")
+        architectures = cfg.get("architectures") or []
+        candidates = {class_name, *architectures} - {None}
+
+        if not candidates & _RECOGNIZED_TEXT_ENCODER_CLASSES:
+            raise NotAMatchError(
+                f"text_encoder class is {sorted(candidates) or 'unknown'}, "
+                f"expected one of {sorted(_RECOGNIZED_TEXT_ENCODER_CLASSES)}"
+            )
+
+        return cls(**override_fields)
+
+
+class QwenVLEncoder_Checkpoint_Config(Checkpoint_Config_Base, Config_Base):
+    """Configuration for single-file Qwen2.5-VL encoder checkpoints (safetensors).
+
+    This matches ComfyUI-style consolidated single-file encoders such as
+    `qwen_2.5_vl_7b_fp8_scaled.safetensors`, which bundle the language model
+    and the visual tower into one file (typically with FP8 + per-tensor
+    `weight_scale` ComfyUI quantization).
+
+    The matching tokenizer + processor are pulled from HuggingFace
+    (`Qwen/Qwen2.5-VL-7B-Instruct`) on first use and cached for offline use.
+    """
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.QwenVLEncoder] = Field(default=ModelType.QwenVLEncoder)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        # Only safetensors checkpoints are supported as single-file Qwen VL encoders.
+        # Reject other extensions cheaply before attempting to read keys.
+        if mod.path.suffix != ".safetensors":
+            raise NotAMatchError(f"expected a .safetensors file, got {mod.path.suffix or '(no suffix)'}")
+
+        # Read only the key index — a 7GB fp8 encoder weighs ~7GB on disk, but we
+        # only need the key names to classify it, not the tensor data.
+        try:
+            keys = _read_safetensors_keys(mod.path)
+        except Exception as e:
+            raise NotAMatchError(f"could not read safetensors header: {e}") from e
+
+        if not _has_qwen_vl_keys(keys):
+            raise NotAMatchError("state dict does not look like a Qwen2.5-VL/Qwen2-VL checkpoint")
+
+        return cls(**override_fields)

invokeai/backend/model_manager/configs/siglip.py

@@ -0,0 +1,45 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    common_config_paths,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+
+class SigLIP_Diffusers_Config(Diffusers_Config_Base, Config_Base):
+    """Model config for SigLIP."""
+
+    type: Literal[ModelType.SigLIP] = Field(default=ModelType.SigLIP)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "SiglipModel",
+            },
+        )
+
+        return cls(**override_fields)

invokeai/backend/model_manager/configs/spandrel.py

@@ -0,0 +1,54 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_override_fields,
+    raise_if_not_file,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+from invokeai.backend.spandrel_image_to_image_model import SpandrelImageToImageModel
+
+
+class Spandrel_Checkpoint_Config(Config_Base):
+    """Model config for Spandrel Image to Image models."""
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.SpandrelImageToImage] = Field(default=ModelType.SpandrelImageToImage)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_spandrel_loads_model(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_spandrel_loads_model(cls, mod: ModelOnDisk) -> None:
+        try:
+            # It would be nice to avoid having to load the Spandrel model from disk here. A couple of options were
+            # explored to avoid this:
+            # 1. Call `SpandrelImageToImageModel.load_from_state_dict(ckpt)`, where `ckpt` is a state_dict on the meta
+            #    device. Unfortunately, some Spandrel models perform operations during initialization that are not
+            #    supported on meta tensors.
+            # 2. Spandrel has internal logic to determine a model's type from its state_dict before loading the model.
+            #    This logic is not exposed in spandrel's public API. We could copy the logic here, but then we have to
+            #    maintain it, and the risk of false positive detections is higher.
+            SpandrelImageToImageModel.load_from_file(mod.path)
+        except Exception as e:
+            raise NotAMatchError("model does not match SpandrelImageToImage heuristics") from e

invokeai/backend/model_manager/configs/t2i_adapter.py

@@ -0,0 +1,79 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.controlnet import ControlAdapterDefaultSettings
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    common_config_paths,
+    get_config_dict_or_raise,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+
+class T2IAdapter_Diffusers_Config_Base(Diffusers_Config_Base):
+    """Model config for T2I."""
+
+    type: Literal[ModelType.T2IAdapter] = Field(default=ModelType.T2IAdapter)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    default_settings: ControlAdapterDefaultSettings | None = Field(None)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "T2IAdapter",
+            },
+        )
+
+        cls._validate_base(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
+        config_dict = get_config_dict_or_raise(common_config_paths(mod.path))
+
+        adapter_type = config_dict.get("adapter_type")
+
+        match adapter_type:
+            case "full_adapter_xl":
+                return BaseModelType.StableDiffusionXL
+            case "full_adapter" | "light_adapter":
+                return BaseModelType.StableDiffusion1
+            case _:
+                raise NotAMatchError(f"unrecognized adapter_type '{adapter_type}'")
+
+
+class T2IAdapter_Diffusers_SD1_Config(T2IAdapter_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class T2IAdapter_Diffusers_SDXL_Config(T2IAdapter_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)

invokeai/backend/model_manager/configs/t5_encoder.py

@@ -0,0 +1,82 @@
+from typing import Any, Literal, Self
+
+from pydantic import Field
+
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    state_dict_has_any_keys_ending_with,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
+
+
+class T5Encoder_T5Encoder_Config(Config_Base):
+    """Configuration for T5 Encoder models in a bespoke, diffusers-like format. The model weights are expected to be in
+    a folder called text_encoder_2 inside the model directory, with a config file named model.safetensors.index.json."""
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.T5Encoder] = Field(default=ModelType.T5Encoder)
+    format: Literal[ModelFormat.T5Encoder] = Field(default=ModelFormat.T5Encoder)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        expected_config_path = mod.path / "text_encoder_2" / "config.json"
+        expected_class_name = "T5EncoderModel"
+        raise_for_class_name(expected_config_path, expected_class_name)
+
+        cls.raise_if_doesnt_have_unquantized_config_file(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def raise_if_doesnt_have_unquantized_config_file(cls, mod: ModelOnDisk) -> None:
+        has_unquantized_config = (mod.path / "text_encoder_2" / "model.safetensors.index.json").exists()
+
+        if not has_unquantized_config:
+            raise NotAMatchError("missing text_encoder_2/model.safetensors.index.json")
+
+
+class T5Encoder_BnBLLMint8_Config(Config_Base):
+    """Configuration for T5 Encoder models quantized by bitsandbytes' LLM.int8."""
+
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    type: Literal[ModelType.T5Encoder] = Field(default=ModelType.T5Encoder)
+    format: Literal[ModelFormat.BnbQuantizedLlmInt8b] = Field(default=ModelFormat.BnbQuantizedLlmInt8b)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        expected_config_path = mod.path / "text_encoder_2" / "config.json"
+        expected_class_name = "T5EncoderModel"
+        raise_for_class_name(expected_config_path, expected_class_name)
+
+        cls.raise_if_filename_doesnt_look_like_bnb_quantized(mod)
+
+        cls.raise_if_state_dict_doesnt_look_like_bnb_quantized(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def raise_if_filename_doesnt_look_like_bnb_quantized(cls, mod: ModelOnDisk) -> None:
+        filename_looks_like_bnb = any(x for x in mod.weight_files() if "llm_int8" in x.as_posix())
+        if not filename_looks_like_bnb:
+            raise NotAMatchError("filename does not look like bnb quantized llm_int8")
+
+    @classmethod
+    def raise_if_state_dict_doesnt_look_like_bnb_quantized(cls, mod: ModelOnDisk) -> None:
+        has_scb_key_suffix = state_dict_has_any_keys_ending_with(mod.load_state_dict(), "SCB")
+        if not has_scb_key_suffix:
+            raise NotAMatchError("state dict does not look like bnb quantized llm_int8")

invokeai/backend/model_manager/configs/text_llm.py

@@ -0,0 +1,52 @@
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    common_config_paths,
+    get_class_name_from_config_dict_or_raise,
+    raise_for_override_fields,
+    raise_if_not_dir,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelType,
+)
+
+
+class TextLLM_Diffusers_Config(Diffusers_Config_Base, Config_Base):
+    """Model config for text-only causal language models (e.g. Llama, Phi, Qwen, Mistral)."""
+
+    type: Literal[ModelType.TextLLM] = Field(default=ModelType.TextLLM)
+    base: Literal[BaseModelType.Any] = Field(default=BaseModelType.Any)
+    cpu_only: bool | None = Field(default=None, description="Whether this model should run on CPU only")
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        # Check that the model's architecture is a causal language model.
+        # This covers LlamaForCausalLM, PhiForCausalLM, Phi3ForCausalLM, Qwen2ForCausalLM,
+        # MistralForCausalLM, GemmaForCausalLM, GPTNeoXForCausalLM, etc.
+        class_name = get_class_name_from_config_dict_or_raise(common_config_paths(mod.path))
+        if not class_name.endswith("ForCausalLM"):
+            raise NotAMatchError(f"model architecture '{class_name}' is not a causal language model")
+
+        # Verify tokenizer files exist to avoid runtime failures
+        tokenizer_files = {"tokenizer.json", "tokenizer.model", "tokenizer_config.json"}
+        if not any((mod.path / f).exists() for f in tokenizer_files):
+            raise NotAMatchError(
+                f"no tokenizer files found in '{mod.path}' "
+                f"(expected at least one of: {', '.join(sorted(tokenizer_files))})"
+            )
+
+        return cls(**override_fields)

invokeai/backend/model_manager/configs/textual_inversion.py

@@ -0,0 +1,156 @@
+from abc import ABC
+from pathlib import Path
+from typing import (
+    Literal,
+    Self,
+)
+
+import torch
+from pydantic import BaseModel, Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    raise_if_not_file,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+
+class TI_Config_Base(ABC, BaseModel):
+    type: Literal[ModelType.TextualInversion] = Field(default=ModelType.TextualInversion)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk, path: Path | None = None) -> None:
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod, path)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _file_looks_like_embedding(cls, mod: ModelOnDisk, path: Path | None = None) -> bool:
+        try:
+            p = path or mod.path
+
+            if not p.exists():
+                return False
+
+            if p.is_dir():
+                return False
+
+            if p.name in [f"learned_embeds.{s}" for s in mod.weight_files()]:
+                return True
+
+            state_dict = mod.load_state_dict(p)
+
+            # Heuristic: textual inversion embeddings have these keys
+            if any(key in {"string_to_param", "emb_params", "clip_g"} for key in state_dict.keys()):
+                return True
+
+            # Heuristic: small state dict with all tensor values
+            if (len(state_dict)) < 10 and all(isinstance(v, torch.Tensor) for v in state_dict.values()):
+                return True
+
+            return False
+        except Exception:
+            return False
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk, path: Path | None = None) -> BaseModelType:
+        p = path or mod.path
+
+        try:
+            state_dict = mod.load_state_dict(p)
+        except Exception as e:
+            raise NotAMatchError(f"unable to load state dict from {p}: {e}") from e
+
+        try:
+            if "string_to_token" in state_dict:
+                token_dim = list(state_dict["string_to_param"].values())[0].shape[-1]
+            elif "emb_params" in state_dict:
+                token_dim = state_dict["emb_params"].shape[-1]
+            elif "clip_g" in state_dict:
+                token_dim = state_dict["clip_g"].shape[-1]
+            else:
+                token_dim = list(state_dict.values())[0].shape[0]
+        except Exception as e:
+            raise NotAMatchError(f"unable to determine token dimension from state dict in {p}: {e}") from e
+
+        match token_dim:
+            case 768:
+                return BaseModelType.StableDiffusion1
+            case 1024:
+                return BaseModelType.StableDiffusion2
+            case 1280:
+                return BaseModelType.StableDiffusionXL
+            case _:
+                raise NotAMatchError(f"unrecognized token dimension {token_dim}")
+
+
+class TI_File_Config_Base(TI_Config_Base):
+    """Model config for textual inversion embeddings."""
+
+    format: Literal[ModelFormat.EmbeddingFile] = Field(default=ModelFormat.EmbeddingFile)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        if not cls._file_looks_like_embedding(mod):
+            raise NotAMatchError("model does not look like a textual inversion embedding file")
+
+        cls._validate_base(mod)
+
+        return cls(**override_fields)
+
+
+class TI_File_SD1_Config(TI_File_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class TI_File_SD2_Config(TI_File_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class TI_File_SDXL_Config(TI_File_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class TI_Folder_Config_Base(TI_Config_Base):
+    """Model config for textual inversion embeddings."""
+
+    format: Literal[ModelFormat.EmbeddingFolder] = Field(default=ModelFormat.EmbeddingFolder)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        for p in mod.weight_files():
+            if cls._file_looks_like_embedding(mod, p):
+                cls._validate_base(mod, p)
+                return cls(**override_fields)
+
+        raise NotAMatchError("model does not look like a textual inversion embedding folder")
+
+
+class TI_Folder_SD1_Config(TI_Folder_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class TI_Folder_SD2_Config(TI_Folder_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class TI_Folder_SDXL_Config(TI_Folder_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)

invokeai/backend/model_manager/configs/unknown.py

@@ -0,0 +1,44 @@
+from copy import deepcopy
+from typing import Any, Literal, Self
+
+from pydantic import Field
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.model_manager.configs.base import Config_Base
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+app_config = get_config()
+
+
+class Unknown_Config(Config_Base):
+    """Model config for unknown models, used as a fallback when we cannot positively identify a model."""
+
+    base: Literal[BaseModelType.Unknown] = Field(default=BaseModelType.Unknown)
+    type: Literal[ModelType.Unknown] = Field(default=ModelType.Unknown)
+    format: Literal[ModelFormat.Unknown] = Field(default=ModelFormat.Unknown)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        """Create an Unknown_Config for models that couldn't be positively identified.
+
+        Note: Basic path validation (file extensions, directory structure) is already
+        performed by ModelConfigFactory before this method is called.
+        """
+
+        cloned_override_fields = deepcopy(override_fields)
+        cloned_override_fields.pop("base", None)
+        cloned_override_fields.pop("type", None)
+        cloned_override_fields.pop("format", None)
+
+        return cls(
+            **cloned_override_fields,
+            # Override the type/format/base to ensure it's marked as unknown.
+            base=BaseModelType.Unknown,
+            type=ModelType.Unknown,
+            format=ModelFormat.Unknown,
+        )

invokeai/backend/model_manager/configs/vae.py

@@ -0,0 +1,345 @@
+import re
+from typing import (
+    Literal,
+    Self,
+)
+
+from pydantic import Field
+from typing_extensions import Any
+
+from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base, Diffusers_Config_Base
+from invokeai.backend.model_manager.configs.identification_utils import (
+    NotAMatchError,
+    common_config_paths,
+    get_config_dict_or_raise,
+    raise_for_class_name,
+    raise_for_override_fields,
+    raise_if_not_dir,
+    raise_if_not_file,
+    state_dict_has_any_keys_starting_with,
+)
+from invokeai.backend.model_manager.model_on_disk import ModelOnDisk
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelType,
+)
+
+REGEX_TO_BASE: dict[str, BaseModelType] = {
+    r"xl": BaseModelType.StableDiffusionXL,
+    r"sd2": BaseModelType.StableDiffusion2,
+    r"vae": BaseModelType.StableDiffusion1,
+    r"FLUX.1-schnell_ae": BaseModelType.Flux,
+}
+
+
+def _is_qwen_image_vae(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict is a Qwen Image VAE (AutoencoderKLQwenImage).
+
+    Qwen Image VAE can be identified by:
+    1. Diffusers-format encoder/decoder keys (`encoder.conv_in`, `decoder.conv_in`)
+    2. 5-dimensional convolution weights (3D causal convolutions vs. standard 2D conv in SD/SDXL/FLUX VAEs)
+    3. 16-dimensional latent space (z_dim=16)
+    """
+    decoder_conv_in_key = "decoder.conv_in.weight"
+    if decoder_conv_in_key not in state_dict:
+        return False
+    weight = state_dict[decoder_conv_in_key]
+    shape = getattr(weight, "shape", None)
+    if shape is None or len(shape) != 5:
+        return False
+    # z_dim is the input channel dim of decoder.conv_in
+    return shape[1] == 16
+
+
+def _is_flux2_vae(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict is a FLUX.2 VAE (AutoencoderKLFlux2).
+
+    FLUX.2 VAE can be identified by:
+    1. Batch Normalization layers (bn.running_mean, bn.running_var) - unique to FLUX.2
+    2. 32-dimensional latent space (decoder.conv_in has 32 input channels)
+
+    FLUX.1 VAE has 16-dimensional latent space and no BatchNorm layers.
+    """
+    # Check for BN layer which is unique to FLUX.2 VAE
+    has_bn = "bn.running_mean" in state_dict or "bn.running_var" in state_dict
+
+    # Check for 32-channel latent space (FLUX.2 has 32, FLUX.1 has 16)
+    decoder_conv_in_key = "decoder.conv_in.weight"
+    has_32_latent_channels = decoder_conv_in_key in state_dict and state_dict[decoder_conv_in_key].shape[1] == 32
+
+    return has_bn or has_32_latent_channels
+
+
+class VAE_Checkpoint_Config_Base(Checkpoint_Config_Base):
+    """Model config for standalone VAE models."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_vae(mod)
+
+        cls._validate_base(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _validate_looks_like_vae(cls, mod: ModelOnDisk) -> None:
+        state_dict = mod.load_state_dict()
+        if not state_dict_has_any_keys_starting_with(
+            state_dict,
+            {
+                "encoder.conv_in",
+                "decoder.conv_in",
+            },
+        ):
+            raise NotAMatchError("model does not match Checkpoint VAE heuristics")
+
+        # Exclude FLUX.2 VAEs - they have their own config class
+        if _is_flux2_vae(state_dict):
+            raise NotAMatchError("model is a FLUX.2 VAE, not a standard VAE")
+
+        # Exclude Qwen Image VAEs - they have their own config class
+        if _is_qwen_image_vae(state_dict):
+            raise NotAMatchError("model is a Qwen Image VAE, not a standard VAE")
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType:
+        # First, try to identify by latent space dimensions (most reliable)
+        state_dict = mod.load_state_dict()
+        decoder_conv_in_key = "decoder.conv_in.weight"
+        if decoder_conv_in_key in state_dict:
+            latent_channels = state_dict[decoder_conv_in_key].shape[1]
+            if latent_channels == 16:
+                # Flux1 VAE has 16-dimensional latent space
+                return BaseModelType.Flux
+            elif latent_channels == 4:
+                # SD/SDXL VAE has 4-dimensional latent space
+                # Try to distinguish SD1/SD2/SDXL by name, fallback to SD1
+                for regexp, base in REGEX_TO_BASE.items():
+                    if re.search(regexp, mod.path.name, re.IGNORECASE):
+                        return base
+                # Default to SD1 if we can't determine from name
+                return BaseModelType.StableDiffusion1
+
+        # Fallback: guess based on name
+        for regexp, base in REGEX_TO_BASE.items():
+            if re.search(regexp, mod.path.name, re.IGNORECASE):
+                return base
+
+        raise NotAMatchError("cannot determine base type")
+
+
+class VAE_Checkpoint_SD1_Config(VAE_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class VAE_Checkpoint_SD2_Config(VAE_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion2] = Field(default=BaseModelType.StableDiffusion2)
+
+
+class VAE_Checkpoint_SDXL_Config(VAE_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class VAE_Checkpoint_FLUX_Config(VAE_Checkpoint_Config_Base, Config_Base):
+    base: Literal[BaseModelType.Flux] = Field(default=BaseModelType.Flux)
+
+
+class VAE_Checkpoint_Flux2_Config(Checkpoint_Config_Base, Config_Base):
+    """Model config for FLUX.2 VAE checkpoint models (AutoencoderKLFlux2)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.Flux2] = Field(default=BaseModelType.Flux2)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        cls._validate_looks_like_vae(mod)
+
+        cls._validate_is_flux2_vae(mod)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_looks_like_vae(cls, mod: ModelOnDisk) -> None:
+        if not state_dict_has_any_keys_starting_with(
+            mod.load_state_dict(),
+            {
+                "encoder.conv_in",
+                "decoder.conv_in",
+            },
+        ):
+            raise NotAMatchError("model does not match Checkpoint VAE heuristics")
+
+    @classmethod
+    def _validate_is_flux2_vae(cls, mod: ModelOnDisk) -> None:
+        """Validate that this is a FLUX.2 VAE, not FLUX.1."""
+        state_dict = mod.load_state_dict()
+        if not _is_flux2_vae(state_dict):
+            raise NotAMatchError("state dict does not look like a FLUX.2 VAE")
+
+
+class VAE_Checkpoint_QwenImage_Config(Checkpoint_Config_Base, Config_Base):
+    """Model config for Qwen Image VAE checkpoint models (AutoencoderKLQwenImage)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.QwenImage] = Field(default=BaseModelType.QwenImage)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        state_dict = mod.load_state_dict()
+        if not _is_qwen_image_vae(state_dict):
+            raise NotAMatchError("state dict does not look like a Qwen Image VAE")
+
+        return cls(**override_fields)
+
+
+def _has_anima_vae_keys(state_dict: dict[str | int, Any]) -> bool:
+    """Check if state dict looks like an Anima QwenImage VAE (AutoencoderKLQwenImage).
+
+    The Anima VAE has a distinctive structure with:
+    - encoder.downsamples.* (instead of encoder.down_blocks)
+    - decoder.upsamples.* (instead of decoder.up_blocks)
+    - decoder.head.* / decoder.middle.*
+    - Top-level conv1/conv2 weights
+    """
+    required_prefixes = {
+        "encoder.downsamples.",
+        "decoder.upsamples.",
+        "decoder.middle.",
+    }
+    return all(any(str(k).startswith(prefix) for k in state_dict) for prefix in required_prefixes)
+
+
+class VAE_Checkpoint_Anima_Config(Checkpoint_Config_Base, Config_Base):
+    """Model config for Anima QwenImage VAE checkpoint models (AutoencoderKLQwenImage)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint)
+    base: Literal[BaseModelType.Anima] = Field(default=BaseModelType.Anima)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_file(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        state_dict = mod.load_state_dict()
+        if not _has_anima_vae_keys(state_dict):
+            raise NotAMatchError("state dict does not look like an Anima QwenImage VAE")
+
+        return cls(**override_fields)
+
+
+class VAE_Diffusers_Config_Base(Diffusers_Config_Base):
+    """Model config for standalone VAE models (diffusers version)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "AutoencoderKL",
+                "AutoencoderTiny",
+            },
+        )
+
+        # Unfortunately it is difficult to distinguish SD1 and SDXL VAEs by config alone, so we may need to
+        # guess based on name if the config is inconclusive.
+        override_name = override_fields.get("name")
+        cls._validate_base(mod, override_name)
+
+        return cls(**override_fields)
+
+    @classmethod
+    def _validate_base(cls, mod: ModelOnDisk, override_name: str | None = None) -> None:
+        """Raise `NotAMatch` if the model base does not match this config class."""
+        expected_base = cls.model_fields["base"].default
+        recognized_base = cls._get_base_or_raise(mod, override_name)
+        if expected_base is not recognized_base:
+            raise NotAMatchError(f"base is {recognized_base}, not {expected_base}")
+
+    @classmethod
+    def _config_looks_like_sdxl(cls, config: dict[str, Any]) -> bool:
+        # Heuristic: These config values that distinguish Stability's SD 1.x VAE from their SDXL VAE.
+        return config.get("scaling_factor", 0) == 0.13025 and config.get("sample_size") in [512, 1024]
+
+    @classmethod
+    def _name_looks_like_sdxl(cls, mod: ModelOnDisk, override_name: str | None = None) -> bool:
+        # Heuristic: SD and SDXL VAE are the same shape (3-channel RGB to 4-channel float scaled down
+        # by a factor of 8), so we can't necessarily tell them apart by config hyperparameters. Best
+        # we can do is guess based on name.
+        return bool(re.search(r"xl\b", override_name or mod.path.name, re.IGNORECASE))
+
+    @classmethod
+    def _get_base_or_raise(cls, mod: ModelOnDisk, override_name: str | None = None) -> BaseModelType:
+        config_dict = get_config_dict_or_raise(common_config_paths(mod.path))
+        if cls._config_looks_like_sdxl(config_dict):
+            return BaseModelType.StableDiffusionXL
+        elif cls._name_looks_like_sdxl(mod, override_name):
+            return BaseModelType.StableDiffusionXL
+        else:
+            # TODO(psyche): Figure out how to positively identify SD1 here, and raise if we can't. Until then, YOLO.
+            return BaseModelType.StableDiffusion1
+
+
+class VAE_Diffusers_SD1_Config(VAE_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusion1] = Field(default=BaseModelType.StableDiffusion1)
+
+
+class VAE_Diffusers_SDXL_Config(VAE_Diffusers_Config_Base, Config_Base):
+    base: Literal[BaseModelType.StableDiffusionXL] = Field(default=BaseModelType.StableDiffusionXL)
+
+
+class VAE_Diffusers_Flux2_Config(Diffusers_Config_Base, Config_Base):
+    """Model config for FLUX.2 VAE models in diffusers format (AutoencoderKLFlux2)."""
+
+    type: Literal[ModelType.VAE] = Field(default=ModelType.VAE)
+    format: Literal[ModelFormat.Diffusers] = Field(default=ModelFormat.Diffusers)
+    base: Literal[BaseModelType.Flux2] = Field(default=BaseModelType.Flux2)
+
+    @classmethod
+    def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self:
+        raise_if_not_dir(mod)
+
+        raise_for_override_fields(cls, override_fields)
+
+        raise_for_class_name(
+            common_config_paths(mod.path),
+            {
+                "AutoencoderKLFlux2",
+            },
+        )
+
+        return cls(**override_fields)

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

invokeai/backend/model_manager/metadata/__init__.py

@@ -0,0 +1,40 @@
+"""
+Initialization file for invokeai.backend.model_manager.metadata
+
+Usage:
+
+from invokeai.backend.model_manager.metadata import(
+   AnyModelRepoMetadata,
+   CommercialUsage,
+   LicenseRestrictions,
+   HuggingFaceMetadata,
+)
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+data = HuggingFaceMetadataFetch().from_id("<REPO_ID>")
+assert isinstance(data, HuggingFaceMetadata)
+"""
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch, ModelMetadataFetchBase
+from invokeai.backend.model_manager.metadata.metadata_base import (
+    AnyModelRepoMetadata,
+    AnyModelRepoMetadataValidator,
+    BaseMetadata,
+    HuggingFaceMetadata,
+    ModelMetadataWithFiles,
+    RemoteModelFile,
+    UnknownMetadataException,
+)
+
+__all__ = [
+    "AnyModelRepoMetadata",
+    "AnyModelRepoMetadataValidator",
+    "HuggingFaceMetadata",
+    "HuggingFaceMetadataFetch",
+    "ModelMetadataFetchBase",
+    "BaseMetadata",
+    "ModelMetadataWithFiles",
+    "RemoteModelFile",
+    "UnknownMetadataException",
+]

invokeai/backend/model_manager/metadata/fetch/__init__.py

@@ -0,0 +1,16 @@
+"""
+Initialization file for invokeai.backend.model_manager.metadata.fetch
+
+Usage:
+from invokeai.backend.model_manager.metadata.fetch import (
+    HuggingFaceMetadataFetch,
+)
+
+data = HuggingFaceMetadataFetch().from_id("<repo_id>")
+assert isinstance(data, HuggingFaceMetadata)
+"""
+
+from invokeai.backend.model_manager.metadata.fetch.fetch_base import ModelMetadataFetchBase
+from invokeai.backend.model_manager.metadata.fetch.huggingface import HuggingFaceMetadataFetch
+
+__all__ = ["ModelMetadataFetchBase", "HuggingFaceMetadataFetch"]

invokeai/backend/model_manager/metadata/fetch/fetch_base.py

@@ -0,0 +1,68 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""
+This module is the base class for subclasses that fetch metadata from model repositories
+
+Usage:
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+data = HuggingFaceMetadataFetch().from_id("<REPO_ID>")
+assert isinstance(data, HuggingFaceMetadata)
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from pydantic.networks import AnyHttpUrl
+from requests.sessions import Session
+
+from invokeai.backend.model_manager.metadata.metadata_base import (
+    AnyModelRepoMetadata,
+    AnyModelRepoMetadataValidator,
+    BaseMetadata,
+)
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+
+
+class ModelMetadataFetchBase(ABC):
+    """Fetch metadata from remote generative model repositories."""
+
+    @abstractmethod
+    def __init__(self, session: Optional[Session] = None):
+        """
+        Initialize the fetcher with an optional requests.sessions.Session object.
+
+        By providing a configurable Session object, we can support unit tests on
+        this module without an internet connection.
+        """
+        pass
+
+    @abstractmethod
+    def from_url(self, url: AnyHttpUrl) -> AnyModelRepoMetadata:
+        """
+        Given a URL to a model repository, return a ModelMetadata object.
+
+        This method will raise a `UnknownMetadataException`
+        in the event that the requested model metadata is not found at the provided location.
+        """
+        pass
+
+    @abstractmethod
+    def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyModelRepoMetadata:
+        """
+        Given an ID for a model, return a ModelMetadata object.
+
+        :param id: An ID.
+        :param variant: A model variant from the ModelRepoVariant enum.
+
+        This method will raise a `UnknownMetadataException`
+        in the event that the requested model's metadata is not found at the provided id.
+        """
+        pass
+
+    @classmethod
+    def from_json(cls, json: str) -> AnyModelRepoMetadata:
+        """Given the JSON representation of the metadata, return the corresponding Pydantic object."""
+        metadata: BaseMetadata = AnyModelRepoMetadataValidator.validate_json(json)  # type: ignore
+        return metadata

invokeai/backend/model_manager/metadata/fetch/huggingface.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""
+This module fetches model metadata objects from the HuggingFace model repository,
+using either a `repo_id` or the model page URL.
+
+Usage:
+
+from invokeai.backend.model_manager.metadata.fetch import HuggingFaceMetadataFetch
+
+fetcher = HuggingFaceMetadataFetch()
+metadata = fetcher.from_url("https://huggingface.co/stabilityai/sdxl-turbo")
+print(metadata.tags)
+"""
+
+import json
+import re
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Optional
+
+import requests
+from huggingface_hub import HfApi, hf_hub_url
+from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundError
+from pydantic.networks import AnyHttpUrl
+from requests.sessions import Session
+
+from invokeai.backend.model_manager.metadata.fetch.fetch_base import ModelMetadataFetchBase
+from invokeai.backend.model_manager.metadata.metadata_base import (
+    AnyModelRepoMetadata,
+    HuggingFaceMetadata,
+    RemoteModelFile,
+    UnknownMetadataException,
+)
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+
+HF_MODEL_RE = r"https?://huggingface.co/([\w\-.]+/[\w\-.]+)"
+
+
+class HuggingFaceMetadataFetch(ModelMetadataFetchBase):
+    """Fetch model metadata from HuggingFace."""
+
+    def __init__(self, session: Optional[Session] = None):
+        """
+        Initialize the fetcher with an optional requests.sessions.Session object.
+
+        By providing a configurable Session object, we can support unit tests on
+        this module without an internet connection.
+        """
+        self._requests = session or requests.Session()
+        self._has_custom_session = session is not None
+
+    @classmethod
+    def from_json(cls, json: str) -> HuggingFaceMetadata:
+        """Given the JSON representation of the metadata, return the corresponding Pydantic object."""
+        metadata = HuggingFaceMetadata.model_validate_json(json)
+        return metadata
+
+    def _model_info_via_session(self, repo_id: str, variant: Optional[ModelRepoVariant] = None) -> SimpleNamespace:
+        """Fetch model info using the injected requests session (for testing/custom backends)."""
+        params = {"blobs": "true"}
+        url = f"https://huggingface.co/api/models/{repo_id}"
+        if variant is not None:
+            url += f"/revision/{variant}"
+        resp = self._requests.get(url, params=params)
+        if resp.status_code == 404:
+            error_code = resp.headers.get("X-Error-Code", "")
+            if error_code == "RevisionNotFound" or (variant is not None):
+                raise RevisionNotFoundError(f"Revision '{variant}' not found for repo '{repo_id}'.")
+            raise RepositoryNotFoundError(f"Repository '{repo_id}' not found.")
+        resp.raise_for_status()
+        data = resp.json()
+        # Convert siblings dicts to SimpleNamespace objects matching HfApi.model_info() shape
+        siblings = []
+        for s in data.get("siblings", []):
+            siblings.append(
+                SimpleNamespace(
+                    rfilename=s.get("rfilename"),
+                    size=s.get("size") or (s.get("lfs", {}) or {}).get("size"),
+                    lfs=s.get("lfs"),
+                )
+            )
+        return SimpleNamespace(id=data["id"], siblings=siblings)
+
+    def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyModelRepoMetadata:
+        """Return a HuggingFaceMetadata object given the model's repo_id."""
+        # Little loop which tries fetching a revision corresponding to the selected variant.
+        # If not available, then set variant to None and get the default.
+        # If this too fails, raise exception.
+
+        model_info = None
+
+        # Handling for our special syntax - we only want the base HF `org/repo` here.
+        repo_id = id.split("::")[0] or id
+        while not model_info:
+            try:
+                # Use the injected session when provided (supports testing with mock adapters).
+                # Otherwise use HfApi which uses httpx internally.
+                if self._has_custom_session:
+                    model_info = self._model_info_via_session(repo_id, variant)
+                else:
+                    model_info = HfApi().model_info(repo_id=repo_id, files_metadata=True, revision=variant)
+            except RepositoryNotFoundError as excp:
+                raise UnknownMetadataException(f"'{repo_id}' not found. See trace for details.") from excp
+            except RevisionNotFoundError:
+                if variant is None:
+                    raise
+                else:
+                    variant = None
+
+        files: list[RemoteModelFile] = []
+
+        _, name = repo_id.split("/")
+
+        for s in model_info.siblings or []:
+            assert s.rfilename is not None
+            assert s.size is not None
+            files.append(
+                RemoteModelFile(
+                    url=hf_hub_url(repo_id, s.rfilename, revision=variant or "main"),
+                    path=Path(name, s.rfilename),
+                    size=s.size,
+                    sha256=s.lfs.get("sha256") if s.lfs else None,
+                )
+            )
+
+        # diffusers models have a `model_index.json` or `config.json` file
+        is_diffusers = any(str(f.url).endswith(("model_index.json", "config.json")) for f in files)
+
+        # These URLs will be exposed to the user - I think these are the only file types we fully support
+        ckpt_urls = (
+            None
+            if is_diffusers
+            else [
+                f.url
+                for f in files
+                if str(f.url).endswith(
+                    (
+                        ".safetensors",
+                        ".bin",
+                        ".pth",
+                        ".pt",
+                        ".ckpt",
+                    )
+                )
+            ]
+        )
+
+        return HuggingFaceMetadata(
+            id=model_info.id,
+            name=name,
+            files=files,
+            api_response=json.dumps(model_info.__dict__, default=str),
+            is_diffusers=is_diffusers,
+            ckpt_urls=ckpt_urls,
+        )
+
+    def from_url(self, url: AnyHttpUrl) -> AnyModelRepoMetadata:
+        """
+        Return a HuggingFaceMetadata object given the model's web page URL.
+
+        In the case of an invalid or missing URL, raises a ModelNotFound exception.
+        """
+        if match := re.match(HF_MODEL_RE, str(url), re.IGNORECASE):
+            repo_id = match.group(1)
+            return self.from_id(repo_id)
+        else:
+            raise UnknownMetadataException(f"'{url}' does not look like a HuggingFace model page")

invokeai/backend/model_manager/metadata/metadata_base.py

@@ -0,0 +1,136 @@
+# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
+
+"""This module defines core text-to-image model metadata fields.
+
+Metadata comprises any descriptive information that is not essential
+for getting the model to run. For example "author" is metadata, while
+"type", "base" and "format" are not. The latter fields are part of the
+model's config, as defined in invokeai.backend.model_manager.config.
+
+Note that the "name" and "description" are also present in `config`
+records. This is intentional. The config record fields are intended to
+be editable by the user as a form of customization. The metadata
+versions of these fields are intended to be kept in sync with the
+remote repo.
+"""
+
+from pathlib import Path
+from typing import List, Literal, Optional, Union
+
+from huggingface_hub import hf_hub_url
+from pydantic import BaseModel, Field, TypeAdapter
+from pydantic.networks import AnyHttpUrl
+from requests.sessions import Session
+from typing_extensions import Annotated
+
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+from invokeai.backend.model_manager.util.select_hf_files import filter_files
+
+
+class UnknownMetadataException(Exception):
+    """Raised when no metadata is available for a model."""
+
+
+class RemoteModelFile(BaseModel):
+    """Information about a downloadable file that forms part of a model."""
+
+    url: AnyHttpUrl = Field(description="The url to download this model file")
+    path: Path = Field(description="The path to the file, relative to the model root")
+    size: Optional[int] = Field(description="The size of this file, in bytes", default=0)
+    sha256: Optional[str] = Field(description="SHA256 hash of this model (not always available)", default=None)
+
+    def __hash__(self) -> int:
+        return hash(str(self))
+
+
+class ModelMetadataBase(BaseModel):
+    """Base class for model metadata information."""
+
+    name: str = Field(description="model's name")
+
+
+class BaseMetadata(ModelMetadataBase):
+    """Adds typing data for discriminated union."""
+
+    type: Literal["basemetadata"] = "basemetadata"
+
+
+class ModelMetadataWithFiles(ModelMetadataBase):
+    """Base class for metadata that contains a list of downloadable model file(s)."""
+
+    files: List[RemoteModelFile] = Field(description="model files and their sizes", default_factory=list)
+
+    def download_urls(
+        self,
+        variant: Optional[ModelRepoVariant] = None,
+        subfolder: Optional[Path] = None,
+        session: Optional[Session] = None,
+    ) -> List[RemoteModelFile]:
+        """
+        Return a list of URLs needed to download the model.
+
+        :param variant: Return files needed to reconstruct the indicated variant (e.g. ModelRepoVariant('fp16'))
+        :param subfolder: Return files in the designated subfolder only
+        :param session: A request.Session object for offline testing
+
+        Note that the "variant" and "subfolder" concepts currently only apply to HuggingFace.
+        However Civitai does have fields for the precision and format of its models, and may
+        provide variant selection criteria in the future.
+        """
+        return self.files
+
+
+class HuggingFaceMetadata(ModelMetadataWithFiles):
+    """Extended metadata fields provided by HuggingFace."""
+
+    type: Literal["huggingface"] = "huggingface"
+    id: str = Field(description="The HF model id")
+    api_response: Optional[str] = Field(description="Response from the HF API as stringified JSON", default=None)
+    is_diffusers: bool = Field(description="Whether the metadata is for a Diffusers format model", default=False)
+    ckpt_urls: Optional[List[AnyHttpUrl]] = Field(
+        description="URLs for all checkpoint format models in the metadata", default=None
+    )
+
+    def download_urls(
+        self,
+        variant: Optional[ModelRepoVariant] = None,
+        subfolder: Optional[Path] = None,
+        subfolders: Optional[List[Path]] = None,
+        session: Optional[Session] = None,
+    ) -> List[RemoteModelFile]:
+        """
+        Return list of downloadable files, filtering by variant and subfolder(s), if any.
+
+        :param variant: Return model files needed to reconstruct the indicated variant
+        :param subfolder: Return model files from the designated subfolder only (deprecated, use subfolders)
+        :param subfolders: Return model files from the designated subfolders
+        :param session: A request.Session object used for internet-free testing
+
+        Note that there is special variant-filtering behavior here:
+        When the fp16 variant is requested and not available, the
+        full-precision model is returned.
+        """
+        session = session or Session()
+
+        paths = filter_files([x.path for x in self.files], variant, subfolder, subfolders)  #  all files in the model
+
+        # Determine prefix for model_index.json check - only applies for single subfolder
+        prefix = ""
+        if subfolder and not subfolders:
+            prefix = f"{subfolder}/"
+
+        # the next step reads model_index.json to determine which subdirectories belong
+        # to the model (only for single subfolder case)
+        if Path(f"{prefix}model_index.json") in paths:
+            url = hf_hub_url(self.id, filename="model_index.json", subfolder=str(subfolder) if subfolder else None)
+            resp = session.get(url)
+            resp.raise_for_status()
+            submodels = resp.json()
+            paths = [Path(subfolder or "", x) for x in paths if Path(x).parent.as_posix() in submodels]
+            paths.insert(0, Path(f"{prefix}model_index.json"))
+
+        return [x for x in self.files if x.path in paths]
+
+
+AnyModelRepoMetadata = Annotated[Union[BaseMetadata, HuggingFaceMetadata], Field(discriminator="type")]
+AnyModelRepoMetadataValidator = TypeAdapter(AnyModelRepoMetadata)

invokeai/backend/model_manager/model_on_disk.py

@@ -0,0 +1,144 @@
+from pathlib import Path
+from typing import Any, Optional, TypeAlias
+
+import safetensors.torch
+import torch
+from picklescan.scanner import scan_file_path
+from safetensors import safe_open
+
+from invokeai.app.services.config.config_default import get_config
+from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader
+from invokeai.backend.util.logging import InvokeAILogger
+from invokeai.backend.util.silence_warnings import SilenceWarnings
+
+StateDict: TypeAlias = dict[str | int, Any]  # When are the keys int?
+
+logger = InvokeAILogger.get_logger()
+
+
+class ModelOnDisk:
+    """A utility class representing a model stored on disk."""
+
+    def __init__(self, path: Path, hash_algo: HASHING_ALGORITHMS = "blake3_single"):
+        self.path = path
+        if self.path.suffix in {".safetensors", ".bin", ".pt", ".ckpt"}:
+            self.name = path.stem
+        else:
+            self.name = path.name
+        self.hash_algo = hash_algo
+        # Having a cache helps users of ModelOnDisk (i.e. configs) to save state
+        # This prevents redundant computations during matching and parsing
+        self._state_dict_cache: dict[Path, Any] = {}
+        self._metadata_cache: dict[Path, Any] = {}
+
+    def hash(self) -> str:
+        return ModelHash(algorithm=self.hash_algo).hash(self.path)
+
+    def size(self) -> int:
+        if self.path.is_file():
+            return self.path.stat().st_size
+        return sum(file.stat().st_size for file in self.path.rglob("*"))
+
+    def weight_files(self) -> set[Path]:
+        if self.path.is_file():
+            return {self.path}
+        extensions = {".safetensors", ".pt", ".pth", ".ckpt", ".bin", ".gguf"}
+        return {f for f in self.path.rglob("*") if f.suffix in extensions and f.is_file()}
+
+    def metadata(self, path: Optional[Path] = None) -> dict[str, str]:
+        path = path or self.path
+        if path in self._metadata_cache:
+            return self._metadata_cache[path]
+        try:
+            with safe_open(self.path, framework="pt", device="cpu") as f:
+                metadata = f.metadata()
+                assert isinstance(metadata, dict)
+        except Exception:
+            metadata = {}
+
+        self._metadata_cache[path] = metadata
+        return metadata
+
+    def repo_variant(self) -> Optional[ModelRepoVariant]:
+        if self.path.is_file():
+            return None
+
+        weight_files = list(self.path.glob("**/*.safetensors"))
+        weight_files.extend(list(self.path.glob("**/*.bin")))
+        for x in weight_files:
+            if ".fp16" in x.suffixes:
+                return ModelRepoVariant.FP16
+            if "openvino_model" in x.name:
+                return ModelRepoVariant.OpenVINO
+            if "flax_model" in x.name:
+                return ModelRepoVariant.Flax
+            if x.suffix == ".onnx":
+                return ModelRepoVariant.ONNX
+        return ModelRepoVariant.Default
+
+    def load_state_dict(self, path: Optional[Path] = None) -> StateDict:
+        if path in self._state_dict_cache:
+            return self._state_dict_cache[path]
+
+        path = self.resolve_weight_file(path)
+
+        if path in self._state_dict_cache:
+            return self._state_dict_cache[path]
+
+        with SilenceWarnings():
+            if path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin")):
+                scan_result = scan_file_path(path)
+                if scan_result.infected_files != 0:
+                    if get_config().unsafe_disable_picklescan:
+                        logger.warning(
+                            f"The model {path.stem} is potentially infected by malware, but picklescan is disabled. "
+                            "Proceeding with caution."
+                        )
+                    else:
+                        raise RuntimeError(
+                            f"The model {path.stem} is potentially infected by malware. Aborting import."
+                        )
+                if scan_result.scan_err:
+                    if get_config().unsafe_disable_picklescan:
+                        logger.warning(
+                            f"Error scanning the model at {path.stem} for malware, but picklescan is disabled. "
+                            "Proceeding with caution."
+                        )
+                    else:
+                        raise RuntimeError(f"Error scanning the model at {path.stem} for malware. Aborting import.")
+                checkpoint = torch.load(path, map_location="cpu")
+                assert isinstance(checkpoint, dict)
+            elif path.suffix.endswith(".gguf"):
+                checkpoint = gguf_sd_loader(path, compute_dtype=torch.float32)
+            elif path.suffix.endswith(".safetensors"):
+                checkpoint = safetensors.torch.load_file(path)
+            else:
+                raise ValueError(f"Unrecognized model extension: {path.suffix}")
+
+        state_dict = checkpoint.get("state_dict", checkpoint)
+
+        # Normalize PEFT named-adapter keys (e.g. `lora_A.default.weight` → `lora_A.weight`).
+        # Pattern is LoRA-specific, so this is a no-op for non-LoRA state dicts.
+        from invokeai.backend.patches.lora_conversions.peft_adapter_utils import normalize_peft_adapter_names
+
+        state_dict = normalize_peft_adapter_names(state_dict)
+
+        self._state_dict_cache[path] = state_dict
+        return state_dict
+
+    def resolve_weight_file(self, path: Optional[Path] = None) -> Path:
+        if not path:
+            weight_files = list(self.weight_files())
+            match weight_files:
+                case []:
+                    raise ValueError("No weight files found for this model")
+                case [p]:
+                    return p
+                case ps if len(ps) >= 2:
+                    raise ValueError(
+                        f"Multiple weight files found for this model: {ps}. "
+                        f"Please specify the intended file using the 'path' argument"
+                    )
+        return path

invokeai/backend/model_manager/omi/__init__.py

@@ -0,0 +1,7 @@
+from invokeai.backend.model_manager.omi.omi import convert_from_omi
+from invokeai.backend.model_manager.omi.vendor.model_spec.architecture import (
+    flux_dev_1_lora,
+    stable_diffusion_xl_1_lora,
+)
+
+__all__ = ["flux_dev_1_lora", "stable_diffusion_xl_1_lora", "convert_from_omi"]

invokeai/backend/model_manager/omi/omi.py

@@ -0,0 +1,21 @@
+from invokeai.backend.model_manager.model_on_disk import StateDict
+from invokeai.backend.model_manager.omi.vendor.convert.lora import (
+    convert_flux_lora as omi_flux,
+)
+from invokeai.backend.model_manager.omi.vendor.convert.lora import (
+    convert_lora_util as lora_util,
+)
+from invokeai.backend.model_manager.omi.vendor.convert.lora import (
+    convert_sdxl_lora as omi_sdxl,
+)
+from invokeai.backend.model_manager.taxonomy import BaseModelType
+
+
+def convert_from_omi(weights_sd: StateDict, base: BaseModelType):
+    keyset = {
+        BaseModelType.Flux: omi_flux.convert_flux_lora_key_sets(),
+        BaseModelType.StableDiffusionXL: omi_sdxl.convert_sdxl_lora_key_sets(),
+    }[base]
+    source = "omi"
+    target = "legacy_diffusers"
+    return lora_util.__convert(weights_sd, keyset, source, target)  # type: ignore

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_clip.py

@@ -0,0 +1,20 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+
+
+def map_clip(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("text_projection", "text_projection", parent=key_prefix)]
+
+    for k in map_prefix_range("text_model.encoder.layers", "text_model.encoder.layers", parent=key_prefix):
+        keys += [LoraConversionKeySet("mlp.fc1", "mlp.fc1", parent=k)]
+        keys += [LoraConversionKeySet("mlp.fc2", "mlp.fc2", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.k_proj", "self_attn.k_proj", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.out_proj", "self_attn.out_proj", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.q_proj", "self_attn.q_proj", parent=k)]
+        keys += [LoraConversionKeySet("self_attn.v_proj", "self_attn.v_proj", parent=k)]
+
+    return keys

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_flux_lora.py

@@ -0,0 +1,84 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_clip import map_clip
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_t5 import map_t5
+
+
+def __map_double_transformer_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("img_attn.qkv.0", "attn.to_q", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_attn.qkv.1", "attn.to_k", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_attn.qkv.2", "attn.to_v", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("txt_attn.qkv.0", "attn.add_q_proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_attn.qkv.1", "attn.add_k_proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_attn.qkv.2", "attn.add_v_proj", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("img_attn.proj", "attn.to_out.0", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_mlp.0", "ff.net.0.proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_mlp.2", "ff.net.2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_mod.lin", "norm1.linear", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("txt_attn.proj", "attn.to_add_out", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_mlp.0", "ff_context.net.0.proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_mlp.2", "ff_context.net.2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("txt_mod.lin", "norm1_context.linear", parent=key_prefix)]
+
+    return keys
+
+
+def __map_single_transformer_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("linear1.0", "attn.to_q", parent=key_prefix)]
+    keys += [LoraConversionKeySet("linear1.1", "attn.to_k", parent=key_prefix)]
+    keys += [LoraConversionKeySet("linear1.2", "attn.to_v", parent=key_prefix)]
+    keys += [LoraConversionKeySet("linear1.3", "proj_mlp", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("linear2", "proj_out", parent=key_prefix)]
+    keys += [LoraConversionKeySet("modulation.lin", "norm.linear", parent=key_prefix)]
+
+    return keys
+
+
+def __map_transformer(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("txt_in", "context_embedder", parent=key_prefix)]
+    keys += [
+        LoraConversionKeySet("final_layer.adaLN_modulation.1", "norm_out.linear", parent=key_prefix, swap_chunks=True)
+    ]
+    keys += [LoraConversionKeySet("final_layer.linear", "proj_out", parent=key_prefix)]
+    keys += [
+        LoraConversionKeySet("guidance_in.in_layer", "time_text_embed.guidance_embedder.linear_1", parent=key_prefix)
+    ]
+    keys += [
+        LoraConversionKeySet("guidance_in.out_layer", "time_text_embed.guidance_embedder.linear_2", parent=key_prefix)
+    ]
+    keys += [LoraConversionKeySet("vector_in.in_layer", "time_text_embed.text_embedder.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("vector_in.out_layer", "time_text_embed.text_embedder.linear_2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("time_in.in_layer", "time_text_embed.timestep_embedder.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("time_in.out_layer", "time_text_embed.timestep_embedder.linear_2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("img_in.proj", "x_embedder", parent=key_prefix)]
+
+    for k in map_prefix_range("double_blocks", "transformer_blocks", parent=key_prefix):
+        keys += __map_double_transformer_block(k)
+
+    for k in map_prefix_range("single_blocks", "single_transformer_blocks", parent=key_prefix):
+        keys += __map_single_transformer_block(k)
+
+    return keys
+
+
+def convert_flux_lora_key_sets() -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("bundle_emb", "bundle_emb")]
+    keys += __map_transformer(LoraConversionKeySet("transformer", "lora_transformer"))
+    keys += map_clip(LoraConversionKeySet("clip_l", "lora_te1"))
+    keys += map_t5(LoraConversionKeySet("t5", "lora_te2"))
+
+    return keys

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_lora_util.py

@@ -0,0 +1,217 @@
+import torch
+from torch import Tensor
+from typing_extensions import Self
+
+
+class LoraConversionKeySet:
+    def __init__(
+        self,
+        omi_prefix: str,
+        diffusers_prefix: str,
+        legacy_diffusers_prefix: str | None = None,
+        parent: Self | None = None,
+        swap_chunks: bool = False,
+        filter_is_last: bool | None = None,
+        next_omi_prefix: str | None = None,
+        next_diffusers_prefix: str | None = None,
+    ):
+        if parent is not None:
+            self.omi_prefix = combine(parent.omi_prefix, omi_prefix)
+            self.diffusers_prefix = combine(parent.diffusers_prefix, diffusers_prefix)
+        else:
+            self.omi_prefix = omi_prefix
+            self.diffusers_prefix = diffusers_prefix
+
+        if legacy_diffusers_prefix is None:
+            self.legacy_diffusers_prefix = self.diffusers_prefix.replace(".", "_")
+        elif parent is not None:
+            self.legacy_diffusers_prefix = combine(parent.legacy_diffusers_prefix, legacy_diffusers_prefix).replace(
+                ".", "_"
+            )
+        else:
+            self.legacy_diffusers_prefix = legacy_diffusers_prefix
+
+        self.parent = parent
+        self.swap_chunks = swap_chunks
+        self.filter_is_last = filter_is_last
+        self.prefix = parent
+
+        if next_omi_prefix is None and parent is not None:
+            self.next_omi_prefix = parent.next_omi_prefix
+            self.next_diffusers_prefix = parent.next_diffusers_prefix
+            self.next_legacy_diffusers_prefix = parent.next_legacy_diffusers_prefix
+        elif next_omi_prefix is not None and parent is not None:
+            self.next_omi_prefix = combine(parent.omi_prefix, next_omi_prefix)
+            self.next_diffusers_prefix = combine(parent.diffusers_prefix, next_diffusers_prefix)
+            self.next_legacy_diffusers_prefix = combine(parent.legacy_diffusers_prefix, next_diffusers_prefix).replace(
+                ".", "_"
+            )
+        elif next_omi_prefix is not None and parent is None:
+            self.next_omi_prefix = next_omi_prefix
+            self.next_diffusers_prefix = next_diffusers_prefix
+            self.next_legacy_diffusers_prefix = next_diffusers_prefix.replace(".", "_")
+        else:
+            self.next_omi_prefix = None
+            self.next_diffusers_prefix = None
+            self.next_legacy_diffusers_prefix = None
+
+    def __get_omi(self, in_prefix: str, key: str) -> str:
+        return self.omi_prefix + key.removeprefix(in_prefix)
+
+    def __get_diffusers(self, in_prefix: str, key: str) -> str:
+        return self.diffusers_prefix + key.removeprefix(in_prefix)
+
+    def __get_legacy_diffusers(self, in_prefix: str, key: str) -> str:
+        key = self.legacy_diffusers_prefix + key.removeprefix(in_prefix)
+
+        suffix = key[key.rfind(".") :]
+        if suffix not in [".alpha", ".dora_scale"]:  # some keys only have a single . in the suffix
+            suffix = key[key.removesuffix(suffix).rfind(".") :]
+        key = key.removesuffix(suffix)
+
+        return key.replace(".", "_") + suffix
+
+    def get_key(self, in_prefix: str, key: str, target: str) -> str:
+        if target == "omi":
+            return self.__get_omi(in_prefix, key)
+        elif target == "diffusers":
+            return self.__get_diffusers(in_prefix, key)
+        elif target == "legacy_diffusers":
+            return self.__get_legacy_diffusers(in_prefix, key)
+        return key
+
+    def __str__(self) -> str:
+        return f"omi: {self.omi_prefix}, diffusers: {self.diffusers_prefix}, legacy: {self.legacy_diffusers_prefix}"
+
+
+def combine(left: str, right: str) -> str:
+    left = left.rstrip(".")
+    right = right.lstrip(".")
+    if left == "" or left is None:
+        return right
+    elif right == "" or right is None:
+        return left
+    else:
+        return left + "." + right
+
+
+def map_prefix_range(
+    omi_prefix: str,
+    diffusers_prefix: str,
+    parent: LoraConversionKeySet,
+) -> list[LoraConversionKeySet]:
+    # 100 should be a safe upper bound. increase if it's not enough in the future
+    return [
+        LoraConversionKeySet(
+            omi_prefix=f"{omi_prefix}.{i}",
+            diffusers_prefix=f"{diffusers_prefix}.{i}",
+            parent=parent,
+            next_omi_prefix=f"{omi_prefix}.{i + 1}",
+            next_diffusers_prefix=f"{diffusers_prefix}.{i + 1}",
+        )
+        for i in range(100)
+    ]
+
+
+def __convert(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+    source: str,
+    target: str,
+) -> dict[str, Tensor]:
+    out_states = {}
+
+    if source == target:
+        return dict(state_dict)
+
+    # TODO: maybe replace with a non O(n^2) algorithm
+    for key, tensor in state_dict.items():
+        for key_set in key_sets:
+            in_prefix = ""
+
+            if source == "omi":
+                in_prefix = key_set.omi_prefix
+            elif source == "diffusers":
+                in_prefix = key_set.diffusers_prefix
+            elif source == "legacy_diffusers":
+                in_prefix = key_set.legacy_diffusers_prefix
+
+            if not key.startswith(in_prefix):
+                continue
+
+            if key_set.filter_is_last is not None:
+                next_prefix = None
+                if source == "omi":
+                    next_prefix = key_set.next_omi_prefix
+                elif source == "diffusers":
+                    next_prefix = key_set.next_diffusers_prefix
+                elif source == "legacy_diffusers":
+                    next_prefix = key_set.next_legacy_diffusers_prefix
+
+                is_last = not any(k.startswith(next_prefix) for k in state_dict)
+                if key_set.filter_is_last != is_last:
+                    continue
+
+            name = key_set.get_key(in_prefix, key, target)
+
+            can_swap_chunks = target == "omi" or source == "omi"
+            if key_set.swap_chunks and name.endswith(".lora_up.weight") and can_swap_chunks:
+                chunk_0, chunk_1 = tensor.chunk(2, dim=0)
+                tensor = torch.cat([chunk_1, chunk_0], dim=0)
+
+            out_states[name] = tensor
+
+            break  # only map the first matching key set
+
+    return out_states
+
+
+def __detect_source(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> str:
+    omi_count = 0
+    diffusers_count = 0
+    legacy_diffusers_count = 0
+
+    for key in state_dict:
+        for key_set in key_sets:
+            if key.startswith(key_set.omi_prefix):
+                omi_count += 1
+            if key.startswith(key_set.diffusers_prefix):
+                diffusers_count += 1
+            if key.startswith(key_set.legacy_diffusers_prefix):
+                legacy_diffusers_count += 1
+
+    if omi_count > diffusers_count and omi_count > legacy_diffusers_count:
+        return "omi"
+    if diffusers_count > omi_count and diffusers_count > legacy_diffusers_count:
+        return "diffusers"
+    if legacy_diffusers_count > omi_count and legacy_diffusers_count > diffusers_count:
+        return "legacy_diffusers"
+
+    return ""
+
+
+def convert_to_omi(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> dict[str, Tensor]:
+    source = __detect_source(state_dict, key_sets)
+    return __convert(state_dict, key_sets, source, "omi")
+
+
+def convert_to_diffusers(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> dict[str, Tensor]:
+    source = __detect_source(state_dict, key_sets)
+    return __convert(state_dict, key_sets, source, "diffusers")
+
+
+def convert_to_legacy_diffusers(
+    state_dict: dict[str, Tensor],
+    key_sets: list[LoraConversionKeySet],
+) -> dict[str, Tensor]:
+    source = __detect_source(state_dict, key_sets)
+    return __convert(state_dict, key_sets, source, "legacy_diffusers")

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_sdxl_lora.py

@@ -0,0 +1,125 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_clip import map_clip
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+
+
+def __map_unet_resnet_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("emb_layers.1", "time_emb_proj", parent=key_prefix)]
+    keys += [LoraConversionKeySet("in_layers.2", "conv1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("out_layers.3", "conv2", parent=key_prefix)]
+    keys += [LoraConversionKeySet("skip_connection", "conv_shortcut", parent=key_prefix)]
+
+    return keys
+
+
+def __map_unet_attention_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("proj_in", "proj_in", parent=key_prefix)]
+    keys += [LoraConversionKeySet("proj_out", "proj_out", parent=key_prefix)]
+    for k in map_prefix_range("transformer_blocks", "transformer_blocks", parent=key_prefix):
+        keys += [LoraConversionKeySet("attn1.to_q", "attn1.to_q", parent=k)]
+        keys += [LoraConversionKeySet("attn1.to_k", "attn1.to_k", parent=k)]
+        keys += [LoraConversionKeySet("attn1.to_v", "attn1.to_v", parent=k)]
+        keys += [LoraConversionKeySet("attn1.to_out.0", "attn1.to_out.0", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_q", "attn2.to_q", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_k", "attn2.to_k", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_v", "attn2.to_v", parent=k)]
+        keys += [LoraConversionKeySet("attn2.to_out.0", "attn2.to_out.0", parent=k)]
+        keys += [LoraConversionKeySet("ff.net.0.proj", "ff.net.0.proj", parent=k)]
+        keys += [LoraConversionKeySet("ff.net.2", "ff.net.2", parent=k)]
+
+    return keys
+
+
+def __map_unet_down_blocks(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("1.0", "0.resnets.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("2.0", "0.resnets.1", parent=key_prefix))
+    keys += [LoraConversionKeySet("3.0.op", "0.downsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("4.0", "1.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("4.1", "1.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("5.0", "1.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("5.1", "1.attentions.1", parent=key_prefix))
+    keys += [LoraConversionKeySet("6.0.op", "1.downsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("7.0", "2.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("7.1", "2.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("8.0", "2.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("8.1", "2.attentions.1", parent=key_prefix))
+
+    return keys
+
+
+def __map_unet_mid_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("0", "resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("1", "attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("2", "resnets.1", parent=key_prefix))
+
+    return keys
+
+
+def __map_unet_up_block(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("0.0", "0.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("0.1", "0.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("1.0", "0.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("1.1", "0.attentions.1", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("2.0", "0.resnets.2", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("2.1", "0.attentions.2", parent=key_prefix))
+    keys += [LoraConversionKeySet("2.2.conv", "0.upsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("3.0", "1.resnets.0", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("3.1", "1.attentions.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("4.0", "1.resnets.1", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("4.1", "1.attentions.1", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("5.0", "1.resnets.2", parent=key_prefix))
+    keys += __map_unet_attention_block(LoraConversionKeySet("5.1", "1.attentions.2", parent=key_prefix))
+    keys += [LoraConversionKeySet("5.2.conv", "1.upsamplers.0.conv", parent=key_prefix)]
+
+    keys += __map_unet_resnet_block(LoraConversionKeySet("6.0", "2.resnets.0", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("7.0", "2.resnets.1", parent=key_prefix))
+    keys += __map_unet_resnet_block(LoraConversionKeySet("8.0", "2.resnets.2", parent=key_prefix))
+
+    return keys
+
+
+def __map_unet(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("input_blocks.0.0", "conv_in", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("time_embed.0", "time_embedding.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("time_embed.2", "time_embedding.linear_2", parent=key_prefix)]
+
+    keys += [LoraConversionKeySet("label_emb.0.0", "add_embedding.linear_1", parent=key_prefix)]
+    keys += [LoraConversionKeySet("label_emb.0.2", "add_embedding.linear_2", parent=key_prefix)]
+
+    keys += __map_unet_down_blocks(LoraConversionKeySet("input_blocks", "down_blocks", parent=key_prefix))
+    keys += __map_unet_mid_block(LoraConversionKeySet("middle_block", "mid_block", parent=key_prefix))
+    keys += __map_unet_up_block(LoraConversionKeySet("output_blocks", "up_blocks", parent=key_prefix))
+
+    keys += [LoraConversionKeySet("out.0", "conv_norm_out", parent=key_prefix)]
+    keys += [LoraConversionKeySet("out.2", "conv_out", parent=key_prefix)]
+
+    return keys
+
+
+def convert_sdxl_lora_key_sets() -> list[LoraConversionKeySet]:
+    keys = []
+
+    keys += [LoraConversionKeySet("bundle_emb", "bundle_emb")]
+    keys += __map_unet(LoraConversionKeySet("unet", "lora_unet"))
+    keys += map_clip(LoraConversionKeySet("clip_l", "lora_te1"))
+    keys += map_clip(LoraConversionKeySet("clip_g", "lora_te2"))
+
+    return keys

invokeai/backend/model_manager/omi/vendor/convert/lora/convert_t5.py

@@ -0,0 +1,19 @@
+from invokeai.backend.model_manager.omi.vendor.convert.lora.convert_lora_util import (
+    LoraConversionKeySet,
+    map_prefix_range,
+)
+
+
+def map_t5(key_prefix: LoraConversionKeySet) -> list[LoraConversionKeySet]:
+    keys = []
+
+    for k in map_prefix_range("encoder.block", "encoder.block", parent=key_prefix):
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.k", "layer.0.SelfAttention.k", parent=k)]
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.o", "layer.0.SelfAttention.o", parent=k)]
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.q", "layer.0.SelfAttention.q", parent=k)]
+        keys += [LoraConversionKeySet("layer.0.SelfAttention.v", "layer.0.SelfAttention.v", parent=k)]
+        keys += [LoraConversionKeySet("layer.1.DenseReluDense.wi_0", "layer.1.DenseReluDense.wi_0", parent=k)]
+        keys += [LoraConversionKeySet("layer.1.DenseReluDense.wi_1", "layer.1.DenseReluDense.wi_1", parent=k)]
+        keys += [LoraConversionKeySet("layer.1.DenseReluDense.wo", "layer.1.DenseReluDense.wo", parent=k)]
+
+    return keys

invokeai/backend/model_manager/omi/vendor/model_spec/architecture.py

@@ -0,0 +1,31 @@
+stable_diffusion_1_lora = "stable-diffusion-v1/lora"
+stable_diffusion_1_inpainting_lora = "stable-diffusion-v1-inpainting/lora"
+
+stable_diffusion_2_512_lora = "stable-diffusion-v2-512/lora"
+stable_diffusion_2_768_v_lora = "stable-diffusion-v2-768-v/lora"
+stable_diffusion_2_depth_lora = "stable-diffusion-v2-depth/lora"
+stable_diffusion_2_inpainting_lora = "stable-diffusion-v2-inpainting/lora"
+
+stable_diffusion_3_medium_lora = "stable-diffusion-v3-medium/lora"
+stable_diffusion_35_medium_lora = "stable-diffusion-v3.5-medium/lora"
+stable_diffusion_35_large_lora = "stable-diffusion-v3.5-large/lora"
+
+stable_diffusion_xl_1_lora = "stable-diffusion-xl-v1-base/lora"
+stable_diffusion_xl_1_inpainting_lora = "stable-diffusion-xl-v1-base-inpainting/lora"
+
+wuerstchen_2_lora = "wuerstchen-v2-prior/lora"
+stable_cascade_1_stage_a_lora = "stable-cascade-v1-stage-a/lora"
+stable_cascade_1_stage_b_lora = "stable-cascade-v1-stage-b/lora"
+stable_cascade_1_stage_c_lora = "stable-cascade-v1-stage-c/lora"
+
+pixart_alpha_lora = "pixart-alpha/lora"
+pixart_sigma_lora = "pixart-sigma/lora"
+
+flux_dev_1_lora = "Flux.1-dev/lora"
+flux_fill_dev_1_lora = "Flux.1-fill-dev/lora"
+
+sana_lora = "sana/lora"
+
+hunyuan_video_lora = "hunyuan-video/lora"
+
+hi_dream_i1_lora = "hidream-i1/lora"

invokeai/backend/model_manager/search.py

@@ -0,0 +1,142 @@
+# Copyright 2023, Lincoln D. Stein and the InvokeAI Team
+"""
+Abstract base class and implementation for recursive directory search for models.
+
+Example usage:
+```
+    from invokeai.backend.model_manager import ModelSearch, ModelProbe
+
+    def find_main_models(model: Path) -> bool:
+        info = ModelProbe.probe(model)
+        if info.model_type == 'main' and info.base_type == 'sd-1':
+            return True
+        else:
+            return False
+
+    search = ModelSearch(on_model_found=report_it)
+    found = search.search('/tmp/models')
+    print(found)   #  list of matching model paths
+    print(search.stats)  #  search stats
+```
+"""
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional
+
+from invokeai.backend.util.logging import InvokeAILogger
+
+
+@dataclass
+class SearchStats:
+    """Statistics about the search.
+
+    Attributes:
+        items_scanned: number of items scanned
+        models_found: number of models found
+        models_filtered: number of models that passed the filter
+    """
+
+    items_scanned = 0
+    models_found = 0
+    models_filtered = 0
+
+
+class ModelSearch:
+    """Searches a directory tree for models, using a callback to filter the results.
+
+    Usage:
+        search = ModelSearch()
+        search.on_model_found = lambda path : 'anime' in path.as_posix()
+        found = search.search(Path('/tmp/models1'))
+    """
+
+    def __init__(
+        self,
+        on_search_started: Optional[Callable[[Path], None]] = None,
+        on_model_found: Optional[Callable[[Path], bool]] = None,
+        on_search_completed: Optional[Callable[[set[Path]], None]] = None,
+    ) -> None:
+        """Create a new ModelSearch object.
+
+        Args:
+            on_search_started: callback to be invoked when the search starts
+            on_model_found: callback to be invoked when a model is found. The callback should return True if the model
+                should be included in the results.
+            on_search_completed: callback to be invoked when the search is completed
+        """
+        self.stats = SearchStats()
+        self.logger = InvokeAILogger.get_logger()
+        self.on_search_started = on_search_started
+        self.on_model_found = on_model_found
+        self.on_search_completed = on_search_completed
+        self.models_found: set[Path] = set()
+
+    def search_started(self) -> None:
+        self.models_found = set()
+        if self.on_search_started:
+            self.on_search_started(self._directory)
+
+    def model_found(self, model: Path) -> None:
+        self.stats.models_found += 1
+        if self.on_model_found is None or self.on_model_found(model):
+            self.stats.models_filtered += 1
+            self.models_found.add(model)
+
+    def search_completed(self) -> None:
+        if self.on_search_completed is not None:
+            self.on_search_completed(self.models_found)
+
+    def search(self, directory: Path) -> set[Path]:
+        self._directory = Path(directory)
+        self._directory = self._directory.resolve()
+        self.stats = SearchStats()  # zero out
+        self.search_started()  # This will initialize _models_found to empty
+        self._walk_directory(self._directory)
+        self.search_completed()
+        return self.models_found
+
+    def _walk_directory(self, path: Path, max_depth: int = 20) -> None:
+        """Recursively walk the directory tree, looking for models."""
+        absolute_path = Path(path)
+        if (
+            len(absolute_path.parts) - len(self._directory.parts) > max_depth
+            or not absolute_path.exists()
+            or absolute_path.parent in self.models_found
+        ):
+            return
+        entries = os.scandir(absolute_path.as_posix())
+        entries = [entry for entry in entries if not entry.name.startswith(".")]
+        dirs = [entry for entry in entries if entry.is_dir()]
+        file_names = [entry.name for entry in entries if entry.is_file()]
+        if any(
+            x in file_names
+            for x in [
+                "config.json",
+                "model_index.json",
+                "learned_embeds.bin",
+                "pytorch_lora_weights.bin",
+                "image_encoder.txt",
+            ]
+        ):
+            try:
+                self.model_found(absolute_path)
+                return
+            except KeyboardInterrupt:
+                raise
+            except Exception as e:
+                self.logger.warning(str(e))
+                return
+
+        for n in file_names:
+            if n.endswith((".ckpt", ".bin", ".pth", ".safetensors", ".pt", ".gguf")):
+                try:
+                    self.model_found(absolute_path / n)
+                except KeyboardInterrupt:
+                    raise
+                except Exception as e:
+                    self.logger.warning(str(e))
+
+        for d in dirs:
+            self._walk_directory(absolute_path / d)

invokeai/backend/model_manager/single_file_config_files.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass
+
+from invokeai.backend.model_manager.configs.factory import AnyModelConfig
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelType,
+    ModelVariantType,
+    SchedulerPredictionType,
+)
+
+
+@dataclass(frozen=True)
+class LegacyConfigKey:
+    type: ModelType
+    base: BaseModelType
+    variant: ModelVariantType | None = None
+    pred: SchedulerPredictionType | None = None
+
+    @classmethod
+    def from_model_config(cls, config: AnyModelConfig) -> "LegacyConfigKey":
+        variant = getattr(config, "variant", None)
+        pred = getattr(config, "prediction_type", None)
+        return cls(type=config.type, base=config.base, variant=variant, pred=pred)
+
+
+LEGACY_CONFIG_MAP: dict[LegacyConfigKey, str] = {
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion1,
+        ModelVariantType.Normal,
+        SchedulerPredictionType.Epsilon,
+    ): "stable-diffusion/v1-inference.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion1,
+        ModelVariantType.Normal,
+        SchedulerPredictionType.VPrediction,
+    ): "stable-diffusion/v1-inference-v.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion1,
+        ModelVariantType.Inpaint,
+    ): "stable-diffusion/v1-inpainting-inference.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion2,
+        ModelVariantType.Normal,
+        SchedulerPredictionType.Epsilon,
+    ): "stable-diffusion/v2-inference.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion2,
+        ModelVariantType.Normal,
+        SchedulerPredictionType.VPrediction,
+    ): "stable-diffusion/v2-inference-v.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion2,
+        ModelVariantType.Inpaint,
+        SchedulerPredictionType.Epsilon,
+    ): "stable-diffusion/v2-inpainting-inference.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion2,
+        ModelVariantType.Inpaint,
+        SchedulerPredictionType.VPrediction,
+    ): "stable-diffusion/v2-inpainting-inference-v.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusion2,
+        ModelVariantType.Depth,
+    ): "stable-diffusion/v2-midas-inference.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusionXL,
+        ModelVariantType.Normal,
+    ): "stable-diffusion/sd_xl_base.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusionXL,
+        ModelVariantType.Inpaint,
+    ): "stable-diffusion/sd_xl_inpaint.yaml",
+    LegacyConfigKey(
+        ModelType.Main,
+        BaseModelType.StableDiffusionXLRefiner,
+        ModelVariantType.Normal,
+    ): "stable-diffusion/sd_xl_refiner.yaml",
+    LegacyConfigKey(ModelType.ControlNet, BaseModelType.StableDiffusion1): "controlnet/cldm_v15.yaml",
+    LegacyConfigKey(ModelType.ControlNet, BaseModelType.StableDiffusion2): "controlnet/cldm_v21.yaml",
+    LegacyConfigKey(ModelType.VAE, BaseModelType.StableDiffusion1): "stable-diffusion/v1-inference.yaml",
+    LegacyConfigKey(ModelType.VAE, BaseModelType.StableDiffusion2): "stable-diffusion/v2-inference.yaml",
+    LegacyConfigKey(ModelType.VAE, BaseModelType.StableDiffusionXL): "stable-diffusion/sd_xl_base.yaml",
+}

invokeai/backend/model_manager/taxonomy.py

@@ -0,0 +1,269 @@
+from enum import Enum
+from typing import Dict, TypeAlias, Union
+
+import onnxruntime as ort
+import torch
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from pydantic import TypeAdapter
+
+from invokeai.backend.raw_model import RawModel
+
+# ModelMixin is the base class for all diffusers and transformers models
+# RawModel is the InvokeAI wrapper class for ip_adapters, loras, textual_inversion and onnx runtime
+AnyModel: TypeAlias = Union[
+    ModelMixin,
+    RawModel,
+    torch.nn.Module,
+    Dict[str, torch.Tensor],
+    DiffusionPipeline,
+    ort.InferenceSession,
+]
+"""Type alias for any kind of runtime, in-memory model representation. For example, a torch module or diffusers pipeline."""
+
+
+class BaseModelType(str, Enum):
+    """An enumeration of base model architectures. For example, Stable Diffusion 1.x, Stable Diffusion 2.x, FLUX, etc.
+
+    Every model config must have a base architecture type.
+
+    Not all models are associated with a base architecture. For example, CLIP models are their own thing, not related
+    to any particular model architecture. To simplify internal APIs and make it easier to work with models, we use a
+    fallback/null value `BaseModelType.Any` for these models, instead of making the model base optional."""
+
+    Any = "any"
+    """`Any` is essentially a fallback/null value for models with no base architecture association.
+    For example, CLIP models are not related to Stable Diffusion, FLUX, or any other model arch."""
+    StableDiffusion1 = "sd-1"
+    """Indicates the model is associated with the Stable Diffusion 1.x model architecture, including 1.4 and 1.5."""
+    StableDiffusion2 = "sd-2"
+    """Indicates the model is associated with the Stable Diffusion 2.x model architecture, including 2.0 and 2.1."""
+    StableDiffusion3 = "sd-3"
+    """Indicates the model is associated with the Stable Diffusion 3.5 model architecture."""
+    StableDiffusionXL = "sdxl"
+    """Indicates the model is associated with the Stable Diffusion XL model architecture."""
+    StableDiffusionXLRefiner = "sdxl-refiner"
+    """Indicates the model is associated with the Stable Diffusion XL Refiner model architecture."""
+    Flux = "flux"
+    """Indicates the model is associated with FLUX.1 model architecture, including FLUX Dev, Schnell and Fill."""
+    Flux2 = "flux2"
+    """Indicates the model is associated with FLUX.2 model architecture, including FLUX2 Klein."""
+    CogView4 = "cogview4"
+    """Indicates the model is associated with CogView 4 model architecture."""
+    ZImage = "z-image"
+    """Indicates the model is associated with Z-Image model architecture, including Z-Image-Turbo."""
+    External = "external"
+    """Indicates the model is hosted by an external provider."""
+    QwenImage = "qwen-image"
+    """Indicates the model is associated with Qwen Image Edit 2511 model architecture."""
+    Anima = "anima"
+    """Indicates the model is associated with Anima model architecture (Cosmos Predict2 DiT + LLM Adapter)."""
+    Unknown = "unknown"
+    """Indicates the model's base architecture is unknown."""
+
+
+class ModelType(str, Enum):
+    """Model type."""
+
+    ONNX = "onnx"
+    Main = "main"
+    VAE = "vae"
+    LoRA = "lora"
+    ControlLoRa = "control_lora"
+    ControlNet = "controlnet"  # used by model_probe
+    TextualInversion = "embedding"
+    IPAdapter = "ip_adapter"
+    CLIPVision = "clip_vision"
+    CLIPEmbed = "clip_embed"
+    T2IAdapter = "t2i_adapter"
+    T5Encoder = "t5_encoder"
+    Qwen3Encoder = "qwen3_encoder"
+    QwenVLEncoder = "qwen_vl_encoder"
+    SpandrelImageToImage = "spandrel_image_to_image"
+    SigLIP = "siglip"
+    FluxRedux = "flux_redux"
+    LlavaOnevision = "llava_onevision"
+    TextLLM = "text_llm"
+    ExternalImageGenerator = "external_image_generator"
+    Unknown = "unknown"
+
+
+class SubModelType(str, Enum):
+    """Submodel type."""
+
+    UNet = "unet"
+    Transformer = "transformer"
+    TextEncoder = "text_encoder"
+    TextEncoder2 = "text_encoder_2"
+    TextEncoder3 = "text_encoder_3"
+    Tokenizer = "tokenizer"
+    Tokenizer2 = "tokenizer_2"
+    Tokenizer3 = "tokenizer_3"
+    VAE = "vae"
+    VAEDecoder = "vae_decoder"
+    VAEEncoder = "vae_encoder"
+    Scheduler = "scheduler"
+    SafetyChecker = "safety_checker"
+
+
+class ClipVariantType(str, Enum):
+    """Variant type."""
+
+    L = "large"
+    G = "gigantic"
+
+
+class ModelVariantType(str, Enum):
+    """Variant type."""
+
+    Normal = "normal"
+    Inpaint = "inpaint"
+    Depth = "depth"
+
+
+class FluxVariantType(str, Enum):
+    """FLUX.1 model variants."""
+
+    Schnell = "schnell"
+    Dev = "dev"
+    DevFill = "dev_fill"
+
+
+class Flux2VariantType(str, Enum):
+    """FLUX.2 model variants."""
+
+    Klein4B = "klein_4b"
+    """Flux2 Klein 4B variant using Qwen3 4B text encoder (distilled)."""
+
+    Klein4BBase = "klein_4b_base"
+    """Flux2 Klein 4B Base variant - undistilled foundation model using Qwen3 4B text encoder."""
+
+    Klein9B = "klein_9b"
+    """Flux2 Klein 9B variant using Qwen3 8B text encoder (distilled)."""
+
+    Klein9BBase = "klein_9b_base"
+    """Flux2 Klein 9B Base variant - undistilled foundation model using Qwen3 8B text encoder."""
+
+
+class ZImageVariantType(str, Enum):
+    """Z-Image model variants."""
+
+    Turbo = "turbo"
+    """Z-Image Turbo - distilled model optimized for 8 steps, no CFG support."""
+
+    ZBase = "zbase"
+    """Z-Image Base - undistilled foundation model with full CFG and negative prompt support."""
+
+
+class QwenImageVariantType(str, Enum):
+    """Qwen Image model variants."""
+
+    Generate = "generate"
+    """Qwen Image - text-to-image generation model."""
+
+    Edit = "edit"
+    """Qwen Image Edit - image editing model with reference image support."""
+
+
+class Qwen3VariantType(str, Enum):
+    """Qwen3 text encoder variants based on model size."""
+
+    Qwen3_4B = "qwen3_4b"
+    """Qwen3 4B text encoder (hidden_size=2560). Used by FLUX.2 Klein 4B and Z-Image."""
+
+    Qwen3_8B = "qwen3_8b"
+    """Qwen3 8B text encoder (hidden_size=4096). Used by FLUX.2 Klein 9B."""
+
+    Qwen3_06B = "qwen3_06b"
+    """Qwen3 0.6B text encoder (hidden_size=1024). Used by Anima."""
+
+
+class ModelFormat(str, Enum):
+    """Storage format of model."""
+
+    OMI = "omi"
+    Diffusers = "diffusers"
+    Checkpoint = "checkpoint"
+    LyCORIS = "lycoris"
+    ONNX = "onnx"
+    Olive = "olive"
+    EmbeddingFile = "embedding_file"
+    EmbeddingFolder = "embedding_folder"
+    InvokeAI = "invokeai"
+    T5Encoder = "t5_encoder"
+    Qwen3Encoder = "qwen3_encoder"
+    QwenVLEncoder = "qwen_vl_encoder"
+    BnbQuantizedLlmInt8b = "bnb_quantized_int8b"
+    BnbQuantizednf4b = "bnb_quantized_nf4b"
+    GGUFQuantized = "gguf_quantized"
+    ExternalApi = "external_api"
+    Unknown = "unknown"
+
+
+class SchedulerPredictionType(str, Enum):
+    """Scheduler prediction type."""
+
+    Epsilon = "epsilon"
+    VPrediction = "v_prediction"
+    Sample = "sample"
+
+
+class ModelRepoVariant(str, Enum):
+    """Various hugging face variants on the diffusers format."""
+
+    Default = ""  # model files without "fp16" or other qualifier
+    FP16 = "fp16"
+    FP32 = "fp32"
+    ONNX = "onnx"
+    OpenVINO = "openvino"
+    Flax = "flax"
+
+
+class ModelSourceType(str, Enum):
+    """Model source type."""
+
+    Path = "path"
+    Url = "url"
+    HFRepoID = "hf_repo_id"
+    External = "external"
+
+
+class FluxLoRAFormat(str, Enum):
+    """Flux LoRA formats."""
+
+    Diffusers = "flux.diffusers"
+    Kohya = "flux.kohya"
+    OneTrainer = "flux.onetrainer"
+    Control = "flux.control"
+    AIToolkit = "flux.aitoolkit"
+    XLabs = "flux.xlabs"
+    BflPeft = "flux.bfl_peft"
+    OneTrainerBfl = "flux.onetrainer_bfl"
+
+
+AnyVariant: TypeAlias = Union[
+    ModelVariantType,
+    ClipVariantType,
+    FluxVariantType,
+    Flux2VariantType,
+    ZImageVariantType,
+    QwenImageVariantType,
+    Qwen3VariantType,
+]
+variant_type_adapter = TypeAdapter[
+    ModelVariantType
+    | ClipVariantType
+    | FluxVariantType
+    | Flux2VariantType
+    | ZImageVariantType
+    | QwenImageVariantType
+    | Qwen3VariantType
+](
+    ModelVariantType
+    | ClipVariantType
+    | FluxVariantType
+    | Flux2VariantType
+    | ZImageVariantType
+    | QwenImageVariantType
+    | Qwen3VariantType
+)

invokeai/backend/model_manager/util/libc_util.py

@@ -0,0 +1,78 @@
+import ctypes
+
+
+class Struct_mallinfo2(ctypes.Structure):
+    """A ctypes Structure that matches the libc mallinfo2 struct.
+
+    Docs:
+    - https://man7.org/linux/man-pages/man3/mallinfo.3.html
+    - https://www.gnu.org/software/libc/manual/html_node/Statistics-of-Malloc.html
+
+    struct mallinfo2 {
+        size_t arena;     /* Non-mmapped space allocated (bytes) */
+        size_t ordblks;   /* Number of free chunks */
+        size_t smblks;    /* Number of free fastbin blocks */
+        size_t hblks;     /* Number of mmapped regions */
+        size_t hblkhd;    /* Space allocated in mmapped regions (bytes) */
+        size_t usmblks;   /* See below */
+        size_t fsmblks;   /* Space in freed fastbin blocks (bytes) */
+        size_t uordblks;  /* Total allocated space (bytes) */
+        size_t fordblks;  /* Total free space (bytes) */
+        size_t keepcost;  /* Top-most, releasable space (bytes) */
+    };
+    """
+
+    _fields_ = [
+        ("arena", ctypes.c_size_t),
+        ("ordblks", ctypes.c_size_t),
+        ("smblks", ctypes.c_size_t),
+        ("hblks", ctypes.c_size_t),
+        ("hblkhd", ctypes.c_size_t),
+        ("usmblks", ctypes.c_size_t),
+        ("fsmblks", ctypes.c_size_t),
+        ("uordblks", ctypes.c_size_t),
+        ("fordblks", ctypes.c_size_t),
+        ("keepcost", ctypes.c_size_t),
+    ]
+
+    def __str__(self) -> str:
+        s = ""
+        s += (
+            f"{'arena': <10}= {(self.arena / 2**30):15.5f}   # Non-mmapped space allocated (GB) (uordblks + fordblks)\n"
+        )
+        s += f"{'ordblks': <10}= {(self.ordblks): >15}   # Number of free chunks\n"
+        s += f"{'smblks': <10}= {(self.smblks): >15}   # Number of free fastbin blocks \n"
+        s += f"{'hblks': <10}= {(self.hblks): >15}   # Number of mmapped regions \n"
+        s += f"{'hblkhd': <10}= {(self.hblkhd / 2**30):15.5f}   # Space allocated in mmapped regions (GB)\n"
+        s += f"{'usmblks': <10}= {(self.usmblks): >15}   # Unused\n"
+        s += f"{'fsmblks': <10}= {(self.fsmblks / 2**30):15.5f}   # Space in freed fastbin blocks (GB)\n"
+        s += (
+            f"{'uordblks': <10}= {(self.uordblks / 2**30):15.5f}   # Space used by in-use allocations (non-mmapped)"
+            " (GB)\n"
+        )
+        s += f"{'fordblks': <10}= {(self.fordblks / 2**30):15.5f}   # Space in free blocks (non-mmapped) (GB)\n"
+        s += f"{'keepcost': <10}= {(self.keepcost / 2**30):15.5f}   # Top-most, releasable space (GB)\n"
+        return s
+
+
+class LibcUtil:
+    """A utility class for interacting with the C Standard Library (`libc`) via ctypes.
+
+    Note that this class will raise on __init__() if 'libc.so.6' can't be found. Take care to handle environments where
+    this shared library is not available.
+
+    TODO: Improve cross-OS compatibility of this class.
+    """
+
+    def __init__(self) -> None:
+        self._libc = ctypes.cdll.LoadLibrary("libc.so.6")
+
+    def mallinfo2(self) -> Struct_mallinfo2:
+        """Calls `libc` `mallinfo2`.
+
+        Docs: https://man7.org/linux/man-pages/man3/mallinfo.3.html
+        """
+        mallinfo2 = self._libc.mallinfo2
+        mallinfo2.restype = Struct_mallinfo2
+        result: Struct_mallinfo2 = mallinfo2()
+        return result