chore: import upstream snapshot with attribution

tests/backend/model_manager/configs/test_anima_controlnet_config.py

@@ -0,0 +1,193 @@
+"""Tests for Anima ControlNet-LLLite config probing.
+
+Anima LLLite adapters (v2 named-key format) are identified by the presence of both the shared
+conditioning trunk (`lllite_conditioning1.*`) and per-module weights (`lllite_dit_blocks_*`).
+SDXL ControlNet-LLLite models (`lllite_unet_*`) and Z-Image Control adapters
+(`control_layers.*` etc.) must not match.
+"""
+
+import os
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from invokeai.backend.model_manager.configs.controlnet import (
+    ControlNet_Checkpoint_Anima_Config,
+    _has_anima_lllite_keys,
+)
+from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
+
+# Opt-in test against the real adapter weights (anima-lllite-inpainting-v2.safetensors).
+_REAL_WEIGHTS_ENV_VAR = "ANIMA_LLLITE_WEIGHTS_PATH"
+REAL_WEIGHTS_PATH = Path(os.environ[_REAL_WEIGHTS_ENV_VAR]) if _REAL_WEIGHTS_ENV_VAR in os.environ else None
+
+_OVERRIDE_FIELDS: dict[str, object] = {
+    "hash": "blake3:fakehash",
+    "path": "/fake/models/anima-lllite.safetensors",
+    "file_size": 1000,
+    "name": "anima-lllite",
+    "description": "test",
+    "source": "test",
+    "source_type": "path",
+    "key": "test-key",
+}
+
+ANIMA_LLLITE_KEYS = [
+    "lllite_conditioning1.conv1.weight",
+    "lllite_conditioning1.conv1.bias",
+    "lllite_conditioning1.proj.weight",
+    "lllite_conditioning1.out_norm.weight",
+    "lllite_dit_blocks_0_self_attn_q_proj.down.weight",
+    "lllite_dit_blocks_0_self_attn_q_proj.depth_embed",
+    "lllite_dit_blocks_27_mlp_layer1.up.weight",
+]
+
+SDXL_LLLITE_KEYS = [
+    "lllite_unet_input_blocks_4_1_transformer_blocks_0_attn1_to_q.cond_emb.weight",
+    "lllite_unet_input_blocks_4_1_transformer_blocks_0_attn1_to_q.down.0.weight",
+    "lllite_unet_middle_block_1_transformer_blocks_0_attn1_to_q.up.0.weight",
+    "lllite_unet_output_blocks_0_1_transformer_blocks_0_attn1_to_q.mid.0.weight",
+]
+
+Z_IMAGE_CONTROL_KEYS = [
+    "control_layers.0.attention.qkv.weight",
+    "control_layers.1.attention.out.weight",
+    "control_all_x_embedder.2-1.weight",
+    "control_noise_refiner.0.attention.qkv.weight",
+]
+
+
+def _make_state_dict(keys: list[str], conv1_in_channels: int = 3) -> dict[str, object]:
+    sd: dict[str, object] = dict.fromkeys(keys)
+    if "lllite_conditioning1.conv1.weight" in sd:
+        sd["lllite_conditioning1.conv1.weight"] = SimpleNamespace(shape=(160, conv1_in_channels, 4, 4))
+    return sd
+
+
+def _make_mod(state_dict: dict[str, object], metadata: dict[str, str] | None = None) -> MagicMock:
+    mod = MagicMock()
+    mod.load_state_dict.return_value = state_dict
+    mod.metadata.return_value = metadata or {}
+    return mod
+
+
+class TestHasAnimaLLLiteKeys:
+    """Tests for the _has_anima_lllite_keys heuristic used during model identification."""
+
+    def test_anima_lllite_keys(self):
+        assert _has_anima_lllite_keys(_make_state_dict(ANIMA_LLLITE_KEYS)) is True
+
+    def test_trunk_only_does_not_match(self):
+        sd = _make_state_dict([k for k in ANIMA_LLLITE_KEYS if k.startswith("lllite_conditioning1.")])
+        assert _has_anima_lllite_keys(sd) is False
+
+    def test_modules_only_does_not_match(self):
+        sd = _make_state_dict([k for k in ANIMA_LLLITE_KEYS if k.startswith("lllite_dit_blocks_")])
+        assert _has_anima_lllite_keys(sd) is False
+
+    def test_sdxl_lllite_does_not_match(self):
+        assert _has_anima_lllite_keys(_make_state_dict(SDXL_LLLITE_KEYS)) is False
+
+    def test_z_image_control_does_not_match(self):
+        assert _has_anima_lllite_keys(_make_state_dict(Z_IMAGE_CONTROL_KEYS)) is False
+
+    def test_empty_state_dict(self):
+        assert _has_anima_lllite_keys({}) is False
+
+
+class TestAnimaControlNetConfigProbe:
+    """Tests for ControlNet_Checkpoint_Anima_Config.from_model_on_disk."""
+
+    def test_matches_anima_lllite(self):
+        mod = _make_mod(_make_state_dict(ANIMA_LLLITE_KEYS))
+
+        config = ControlNet_Checkpoint_Anima_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+        assert config.base is BaseModelType.Anima
+        assert config.type is ModelType.ControlNet
+        assert config.format is ModelFormat.Checkpoint
+
+    @pytest.mark.parametrize(
+        "keys",
+        [SDXL_LLLITE_KEYS, Z_IMAGE_CONTROL_KEYS, []],
+        ids=["sdxl_lllite", "z_image_control", "empty"],
+    )
+    def test_rejects_non_anima_lllite(self, keys: list[str]):
+        mod = _make_mod(_make_state_dict(keys))
+
+        with pytest.raises(NotAMatchError, match="does not look like an Anima ControlNet-LLLite"):
+            ControlNet_Checkpoint_Anima_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+
+class TestCondInChannels:
+    """Tests for the cond_in_channels field (metadata, conv1-shape fallback, old-JSON rehydration)."""
+
+    def test_populated_from_metadata(self):
+        mod = _make_mod(
+            _make_state_dict(ANIMA_LLLITE_KEYS, conv1_in_channels=3),
+            metadata={"lllite.cond_in_channels": "4"},
+        )
+
+        config = ControlNet_Checkpoint_Anima_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+        assert config.cond_in_channels == 4
+
+    def test_fallback_to_conv1_shape_when_metadata_absent(self):
+        mod = _make_mod(_make_state_dict(ANIMA_LLLITE_KEYS, conv1_in_channels=4))
+
+        config = ControlNet_Checkpoint_Anima_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+        assert config.cond_in_channels == 4
+
+    def test_old_json_rehydrates_with_none(self):
+        """Configs stored before the field existed must still validate, with cond_in_channels=None."""
+        mod = _make_mod(_make_state_dict(ANIMA_LLLITE_KEYS))
+        config = ControlNet_Checkpoint_Anima_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+        assert config.cond_in_channels == 3
+
+        stored = config.model_dump(mode="json")
+        del stored["cond_in_channels"]
+        rehydrated = ControlNet_Checkpoint_Anima_Config.model_validate(stored)
+
+        assert rehydrated.cond_in_channels is None
+
+    def test_override_skips_probe(self):
+        """An explicit override must win without probing the state dict (which may be unprobeable)."""
+        keys = [k for k in ANIMA_LLLITE_KEYS if k != "lllite_conditioning1.conv1.weight"]
+        mod = _make_mod(_make_state_dict(keys))
+
+        config = ControlNet_Checkpoint_Anima_Config.from_model_on_disk(
+            mod, dict(_OVERRIDE_FIELDS) | {"cond_in_channels": 4}
+        )
+
+        assert config.cond_in_channels == 4
+
+    def test_missing_conv1_without_metadata_is_not_a_match(self):
+        """A malformed file with LLLite keys but no conv1.weight must raise NotAMatchError, not KeyError."""
+        keys = [k for k in ANIMA_LLLITE_KEYS if k != "lllite_conditioning1.conv1.weight"]
+        mod = _make_mod(_make_state_dict(keys))
+
+        with pytest.raises(NotAMatchError, match="no lllite_conditioning1.conv1.weight"):
+            ControlNet_Checkpoint_Anima_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+
+@pytest.mark.skipif(
+    REAL_WEIGHTS_PATH is None,
+    reason=f"set {_REAL_WEIGHTS_ENV_VAR} to the real LLLite weights file to run",
+)
+def test_real_file_classifies_as_anima_controlnet():
+    """The real adapter file must classify uniquely as an Anima ControlNet via the full factory."""
+    from invokeai.backend.model_manager.configs.factory import ModelConfigFactory
+
+    assert REAL_WEIGHTS_PATH is not None
+    # A stale path must fail loudly, not silently skip.
+    assert REAL_WEIGHTS_PATH.is_file(), f"{_REAL_WEIGHTS_ENV_VAR} points to a missing file: {REAL_WEIGHTS_PATH}"
+    result = ModelConfigFactory.from_model_on_disk(REAL_WEIGHTS_PATH, allow_unknown=False)
+
+    assert result.config is not None
+    assert isinstance(result.config, ControlNet_Checkpoint_Anima_Config)
+    assert result.match_count == 1
+    assert result.config.cond_in_channels == 4

tests/backend/model_manager/configs/test_anima_model_identification.py

@@ -0,0 +1,162 @@
+import pytest
+
+from invokeai.backend.model_manager.configs.main import _has_anima_keys
+
+
+def _make_state_dict(prefixes: list[str], keys: list[str]) -> dict[str, object]:
+    """Build a minimal fake state dict with the given prefixes applied to the given keys."""
+    return {f"{prefix}{key}": None for prefix in prefixes for key in keys}
+
+
+# Minimal keys that satisfy both llm_adapter and cosmos DiT requirements
+ANIMA_LLM_ADAPTER_KEYS = ["llm_adapter.blocks.0.cross_attn.k_norm.weight"]
+ANIMA_COSMOS_DIT_KEYS = [
+    "blocks.0.adaln_modulation_cross_attn.1.weight",
+    "t_embedder.1.linear_1.weight",
+    "x_embedder.proj.1.weight",
+    "final_layer.adaln_modulation.1.weight",
+]
+
+
+class TestHasAnimaKeys:
+    """Tests for _has_anima_keys heuristic used during model identification."""
+
+    def test_bare_keys(self):
+        """Bare keys (no prefix) should be recognized."""
+        sd = _make_state_dict([""], ANIMA_LLM_ADAPTER_KEYS + ANIMA_COSMOS_DIT_KEYS)
+        assert _has_anima_keys(sd) is True
+
+    def test_net_prefix(self):
+        """Official format with `net.` prefix should be recognized."""
+        sd = _make_state_dict(["net."], ANIMA_LLM_ADAPTER_KEYS + ANIMA_COSMOS_DIT_KEYS)
+        assert _has_anima_keys(sd) is True
+
+    def test_comfyui_bundled_prefix(self):
+        """ComfyUI bundled format with `model.diffusion_model.` prefix should be recognized."""
+        sd = _make_state_dict(["model.diffusion_model."], ANIMA_LLM_ADAPTER_KEYS + ANIMA_COSMOS_DIT_KEYS)
+        assert _has_anima_keys(sd) is True
+
+    def test_comfyui_bundled_with_extra_keys(self):
+        """Bundled checkpoint with VAE and text encoder keys should still be recognized."""
+        sd = _make_state_dict(["model.diffusion_model."], ANIMA_LLM_ADAPTER_KEYS + ANIMA_COSMOS_DIT_KEYS)
+        # Add bundled VAE and text encoder keys (should not interfere)
+        sd["first_stage_model.conv1.weight"] = None
+        sd["first_stage_model.encoder.downsamples.0.weight"] = None
+        sd["cond_stage_model.qwen3_06b.transformer.model.embed_tokens.weight"] = None
+        assert _has_anima_keys(sd) is True
+
+    def test_missing_llm_adapter_keys(self):
+        """Should not match if llm_adapter keys are absent."""
+        sd = _make_state_dict([""], ANIMA_COSMOS_DIT_KEYS)
+        assert _has_anima_keys(sd) is False
+
+    def test_missing_cosmos_dit_keys(self):
+        """Should not match if Cosmos DiT keys are absent."""
+        sd = _make_state_dict([""], ANIMA_LLM_ADAPTER_KEYS)
+        assert _has_anima_keys(sd) is False
+
+    def test_empty_state_dict(self):
+        """Empty state dict should not match."""
+        assert _has_anima_keys({}) is False
+
+    def test_unrelated_keys(self):
+        """State dict with unrelated keys should not match."""
+        sd = {
+            "model.diffusion_model.input_blocks.0.0.weight": None,
+            "model.diffusion_model.output_blocks.0.0.weight": None,
+            "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight": None,
+        }
+        assert _has_anima_keys(sd) is False
+
+    @pytest.mark.parametrize(
+        "prefix",
+        ["", "net.", "model.diffusion_model."],
+    )
+    def test_all_prefixes_parametrized(self, prefix: str):
+        """All supported prefix formats should be recognized."""
+        sd = _make_state_dict([prefix], ANIMA_LLM_ADAPTER_KEYS + ANIMA_COSMOS_DIT_KEYS)
+        assert _has_anima_keys(sd) is True
+
+
+class TestAnimaDoesNotConflictWithOtherModels:
+    """Verify that _has_anima_keys does not false-positive on similar model architectures."""
+
+    def test_flux_bundled_checkpoint(self):
+        """FLUX bundled checkpoints use double_blocks/single_blocks, not blocks — should not match."""
+        sd = {
+            "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale": None,
+            "model.diffusion_model.double_blocks.0.img_attn.proj.weight": None,
+            "model.diffusion_model.single_blocks.0.linear1.weight": None,
+            "model.diffusion_model.context_embedder.weight": None,
+            "model.diffusion_model.img_in.weight": None,
+        }
+        assert _has_anima_keys(sd) is False
+
+    def test_sd1_bundled_checkpoint(self):
+        """SD1/SD2/SDXL bundled checkpoints use input_blocks/output_blocks — should not match."""
+        sd = {
+            "model.diffusion_model.input_blocks.0.0.weight": None,
+            "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight": None,
+            "model.diffusion_model.output_blocks.0.0.weight": None,
+            "model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_q.weight": None,
+            "first_stage_model.encoder.down.0.block.0.conv1.weight": None,
+            "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight": None,
+        }
+        assert _has_anima_keys(sd) is False
+
+    def test_raw_cosmos_dit_without_llm_adapter(self):
+        """A raw Cosmos Predict2 DiT (without Anima's LLM adapter) should not match."""
+        sd = {
+            "blocks.0.adaln_modulation_cross_attn.1.weight": None,
+            "blocks.0.self_attn.q_proj.weight": None,
+            "t_embedder.1.linear_1.weight": None,
+            "x_embedder.proj.1.weight": None,
+            "final_layer.adaln_modulation.1.weight": None,
+        }
+        assert _has_anima_keys(sd) is False
+
+    def test_z_image_checkpoint(self):
+        """Z-Image uses blocks.* but with cap_embedder/context_refiner — should not match."""
+        sd = {
+            "model.diffusion_model.blocks.0.attn.to_q.weight": None,
+            "model.diffusion_model.blocks.0.attn.to_k.weight": None,
+            "model.diffusion_model.cap_embedder.0.weight": None,
+            "model.diffusion_model.context_refiner.blocks.0.weight": None,
+            "model.diffusion_model.t_embedder.mlp.0.weight": None,
+            "model.diffusion_model.x_embedder.proj.weight": None,
+        }
+        # Z-Image has blocks/t_embedder/x_embedder but NOT llm_adapter
+        assert _has_anima_keys(sd) is False
+
+    def test_qwen_image_checkpoint(self):
+        """QwenImage uses txt_in/txt_norm/img_in — should not match."""
+        sd = {
+            "txt_in.weight": None,
+            "txt_norm.weight": None,
+            "img_in.weight": None,
+            "double_blocks.0.img_attn.proj.weight": None,
+            "single_blocks.0.linear1.weight": None,
+        }
+        assert _has_anima_keys(sd) is False
+
+    def test_flux_lora_does_not_match(self):
+        """FLUX LoRA weights should not match as Anima."""
+        sd = {
+            "double_blocks.0.img_attn.proj.lora_down.weight": None,
+            "double_blocks.0.img_attn.proj.lora_up.weight": None,
+            "single_blocks.0.linear1.lora_down.weight": None,
+        }
+        assert _has_anima_keys(sd) is False
+
+    def test_cosmos_dit_bundled_without_llm_adapter(self):
+        """Bundled Cosmos DiT (model.diffusion_model. prefix) but no llm_adapter — should not match."""
+        sd = {
+            "model.diffusion_model.blocks.0.self_attn.q_proj.weight": None,
+            "model.diffusion_model.t_embedder.1.linear_1.weight": None,
+            "model.diffusion_model.x_embedder.proj.1.weight": None,
+            "model.diffusion_model.final_layer.adaln_modulation.1.weight": None,
+            "first_stage_model.encoder.downsamples.0.weight": None,
+            "cond_stage_model.transformer.model.embed_tokens.weight": None,
+        }
+        # Has all the Cosmos DiT keys but missing llm_adapter — not Anima
+        assert _has_anima_keys(sd) is False

tests/backend/model_manager/configs/test_double_variant_regression.py

@@ -0,0 +1,115 @@
+"""Regression tests for the double-variant kwarg bug.
+
+When override_fields contains a field (variant, repo_variant, prediction_type, etc.)
+that is also computed and passed as an explicit kwarg to cls(), using .get() instead
+of .pop() causes TypeError("got multiple values for keyword argument ...").
+
+These tests verify that .pop() is used consistently, so override values don't conflict
+with explicitly computed values.
+"""
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from invokeai.backend.model_manager.taxonomy import QwenImageVariantType
+
+# Required fields for the Pydantic config model
+_REQUIRED_FIELDS = {
+    "hash": "blake3:fakehash",
+    "path": "/fake/models/test-model",
+    "file_size": 1000,
+    "name": "test-model",
+    "description": "test",
+    "source": "test",
+    "source_type": "path",
+    "key": "test-key",
+}
+
+
+def _make_mock_dir(dirname: str = "test-model") -> MagicMock:
+    """Create a mock ModelOnDisk for a Diffusers directory."""
+    mod = MagicMock()
+    mod.path = Path(f"/fake/models/{dirname}")
+    return mod
+
+
+class TestDoubleVariantRegression:
+    """Verify that override_fields with variant/repo_variant don't cause double-kwarg errors."""
+
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_class_name")
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_dir")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_qwen_image_diffusers_with_variant_in_overrides(self, _rfo, _rid, _rfc):
+        """Installing a Qwen Image Edit Diffusers model with variant in override_fields should not crash."""
+        from invokeai.backend.model_manager.configs.main import Main_Diffusers_QwenImage_Config
+
+        mod = _make_mock_dir("Qwen-Image-Edit-2511")
+
+        # Simulate what happens when a starter model provides variant
+        overrides = {
+            **_REQUIRED_FIELDS,
+            "variant": QwenImageVariantType.Edit,
+        }
+
+        from invokeai.backend.model_manager.configs.base import ModelRepoVariant
+
+        with patch.object(
+            Main_Diffusers_QwenImage_Config, "_get_repo_variant_or_raise", return_value=ModelRepoVariant("")
+        ):
+            with patch.object(
+                Main_Diffusers_QwenImage_Config,
+                "_get_qwen_image_variant",
+                return_value=QwenImageVariantType.Edit,
+            ):
+                # This would previously raise: TypeError("got multiple values for keyword argument 'variant'")
+                config = Main_Diffusers_QwenImage_Config.from_model_on_disk(mod, overrides)
+
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_class_name")
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_dir")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_qwen_image_diffusers_override_variant_takes_precedence(self, _rfo, _rid, _rfc):
+        """An explicit variant override should take precedence over auto-detection."""
+        from invokeai.backend.model_manager.configs.base import ModelRepoVariant
+        from invokeai.backend.model_manager.configs.main import Main_Diffusers_QwenImage_Config
+
+        mod = _make_mock_dir("Qwen-Image-2512")
+
+        overrides = {
+            **_REQUIRED_FIELDS,
+            "variant": QwenImageVariantType.Edit,  # explicitly override to Edit
+        }
+
+        with patch.object(
+            Main_Diffusers_QwenImage_Config, "_get_repo_variant_or_raise", return_value=ModelRepoVariant("")
+        ):
+            with patch.object(
+                Main_Diffusers_QwenImage_Config,
+                "_get_qwen_image_variant",
+                return_value=QwenImageVariantType.Generate,  # auto-detect says Generate
+            ):
+                config = Main_Diffusers_QwenImage_Config.from_model_on_disk(mod, overrides)
+
+        # Override should win over auto-detection
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_qwen_image_gguf_with_variant_in_overrides(self, _rfo, _rif, _hgt, _hqk):
+        """Installing a Qwen Image Edit GGUF with variant in override_fields should not crash."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        mod = MagicMock()
+        mod.path = Path("/fake/models/qwen-image-edit-2511-Q4_K_M.gguf")
+        mod.load_state_dict.return_value = {}
+
+        overrides = {
+            **_REQUIRED_FIELDS,
+            "variant": QwenImageVariantType.Edit,
+        }
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(mod, overrides)
+        assert config.variant == QwenImageVariantType.Edit

tests/backend/model_manager/configs/test_lora_default_settings.py

@@ -0,0 +1,112 @@
+import pytest
+from pydantic import ValidationError
+
+from invokeai.backend.model_manager.configs.lora import LoraModelDefaultSettings
+
+
+def test_accepts_none_for_all_fields() -> None:
+    settings = LoraModelDefaultSettings()
+    assert settings.weight is None
+    assert settings.weight_min is None
+    assert settings.weight_max is None
+
+
+def test_accepts_both_bounds_with_min_less_than_max() -> None:
+    settings = LoraModelDefaultSettings(weight_min=-2.0, weight_max=3.0)
+    assert settings.weight_min == -2.0
+    assert settings.weight_max == 3.0
+
+
+def test_rejects_both_bounds_with_min_greater_than_or_equal_to_max() -> None:
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight_min=2.0, weight_max=2.0)
+
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight_min=3.0, weight_max=1.0)
+
+
+def test_accepts_only_weight_min_within_default_range() -> None:
+    # Default max is 2.0; a weight_min of 1.0 leaves a valid effective range [1.0, 2.0].
+    settings = LoraModelDefaultSettings(weight_min=1.0)
+    assert settings.weight_min == 1.0
+    assert settings.weight_max is None
+
+
+def test_rejects_only_weight_min_above_default_max() -> None:
+    # Reproduces the partial-bound bug: saving only weight_min=3 used to be accepted but
+    # rendered as a min>max slider in LoRACard. The effective max defaults to 2.0.
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight_min=3.0)
+
+
+def test_rejects_only_weight_min_equal_to_default_max() -> None:
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight_min=2.0)
+
+
+def test_accepts_only_weight_max_within_default_range() -> None:
+    # Default min is -1.0; a weight_max of 0.5 leaves a valid effective range [-1.0, 0.5].
+    settings = LoraModelDefaultSettings(weight_max=0.5)
+    assert settings.weight_max == 0.5
+    assert settings.weight_min is None
+
+
+def test_rejects_only_weight_max_below_default_min() -> None:
+    # Reproduces the symmetric partial-bound bug for weight_max <= -1.
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight_max=-2.0)
+
+
+def test_rejects_only_weight_max_equal_to_default_min() -> None:
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight_max=-1.0)
+
+
+def test_accepts_weight_within_effective_range_with_explicit_bounds() -> None:
+    settings = LoraModelDefaultSettings(weight=0.5, weight_min=-1.0, weight_max=2.0)
+    assert settings.weight == 0.5
+
+
+def test_accepts_weight_within_default_effective_range() -> None:
+    # With no bounds set the effective range is [-1.0, 2.0].
+    settings = LoraModelDefaultSettings(weight=1.5)
+    assert settings.weight == 1.5
+
+
+def test_rejects_weight_above_explicit_effective_range() -> None:
+    # Reproduces the out-of-range-weight bug: weight=10 with bounds [-1, 2] used to be
+    # accepted but the slider in LoRACard could only show [-1, 2].
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight=10.0, weight_min=-1.0, weight_max=2.0)
+
+
+def test_rejects_weight_below_explicit_effective_range() -> None:
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight=-5.0, weight_min=-1.0, weight_max=2.0)
+
+
+def test_rejects_weight_above_default_effective_range() -> None:
+    # No bounds set; default max is 2.0.
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight=2.5)
+
+
+def test_rejects_weight_outside_partial_bound() -> None:
+    # Only weight_min set; weight must respect effective range [weight_min, default_max].
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight=0.0, weight_min=0.5)
+
+    # Only weight_max set; weight must respect effective range [default_min, weight_max].
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(weight=1.0, weight_max=0.5)
+
+
+def test_accepts_weight_at_effective_bounds() -> None:
+    # Inclusive bounds.
+    LoraModelDefaultSettings(weight=-1.0, weight_min=-1.0, weight_max=2.0)
+    LoraModelDefaultSettings(weight=2.0, weight_min=-1.0, weight_max=2.0)
+
+
+def test_rejects_extra_fields() -> None:
+    with pytest.raises(ValidationError):
+        LoraModelDefaultSettings(extra_field=True)  # type: ignore[call-arg]

tests/backend/model_manager/configs/test_qwen3_encoder_config.py

@@ -0,0 +1,81 @@
+"""Regression tests for Qwen3 Encoder config probing.
+
+See https://github.com/invoke-ai/InvokeAI/issues/9090
+
+`Qwen2.5-1.5B-Instruct` (a standalone causal LM) was being misidentified as a
+`Qwen3Encoder` because the diffusers-style config check matched any directory with
+`config.json` at the root and a Qwen* class name. A complete causal LM also bundles
+tokenizer files at the root, while standalone text_encoder downloads do not — we
+use that to disambiguate.
+"""
+
+import json
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import MagicMock
+
+import pytest
+
+from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
+from invokeai.backend.model_manager.configs.qwen3_encoder import Qwen3Encoder_Qwen3Encoder_Config
+
+_OVERRIDE_FIELDS: dict[str, object] = {
+    "hash": "blake3:fakehash",
+    "path": "/fake/models/test-model",
+    "file_size": 1000,
+    "name": "test-model",
+    "description": "test",
+    "source": "test",
+    "source_type": "path",
+    "key": "test-key",
+}
+
+
+def _write_config(path: Path, hidden_size: int = 2560, architecture: str = "Qwen2ForCausalLM") -> None:
+    path.write_text(json.dumps({"architectures": [architecture], "hidden_size": hidden_size}))
+
+
+@pytest.mark.parametrize("tokenizer_file", ["tokenizer.json", "tokenizer.model", "tokenizer_config.json"])
+def test_complete_causal_lm_is_rejected(tokenizer_file: str) -> None:
+    """A directory with config.json + tokenizer files at root is a TextLLM, not a Qwen3 encoder."""
+    with TemporaryDirectory() as tmpdir:
+        root = Path(tmpdir)
+        _write_config(root / "config.json")
+        (root / tokenizer_file).write_text("{}")
+
+        mod = MagicMock()
+        mod.path = root
+
+        with pytest.raises(NotAMatchError, match="complete causal LM"):
+            Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+
+def test_standalone_text_encoder_subfolder_still_matches() -> None:
+    """A standalone text_encoder download (config.json at root, no tokenizer files) should still match."""
+    with TemporaryDirectory() as tmpdir:
+        root = Path(tmpdir)
+        _write_config(root / "config.json")
+
+        mod = MagicMock()
+        mod.path = root
+
+        config = Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+        assert config.type.value == "qwen3_encoder"
+
+
+def test_nested_text_encoder_with_root_tokenizer_still_matches() -> None:
+    """A model with text_encoder/config.json should match even if tokenizer files exist at root.
+
+    The tokenizer-at-root heuristic only applies to the standalone (root-level config.json) case.
+    """
+    with TemporaryDirectory() as tmpdir:
+        root = Path(tmpdir)
+        (root / "text_encoder").mkdir()
+        _write_config(root / "text_encoder" / "config.json")
+        (root / "tokenizer.json").write_text("{}")
+
+        mod = MagicMock()
+        mod.path = root
+
+        config = Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+        assert config.type.value == "qwen3_encoder"

tests/backend/model_manager/configs/test_qwen_image_checkpoint_variant_detection.py

@@ -0,0 +1,187 @@
+"""Tests for Qwen Image single-file checkpoint variant detection.
+
+Mirrors `test_qwen_image_gguf_variant_detection.py`. The Checkpoint and GGUF
+configs share the same variant inference (`_infer_qwen_image_variant`):
+
+1. Explicit `variant` in override_fields wins.
+2. Presence of the `__index_timestep_zero__` tensor → Edit.
+3. Filename heuristic: "edit" substring in the stem → Edit.
+4. Otherwise default to Generate.
+
+Also tests that the Checkpoint config rejects GGUF state dicts (and vice
+versa), so the two configs don't both match the same file.
+"""
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
+from invokeai.backend.model_manager.taxonomy import QwenImageVariantType
+
+_REQUIRED_FIELDS = {
+    "hash": "blake3:fakehash",
+    "path": "/fake/models/test.safetensors",
+    "file_size": 1000,
+    "name": "test-model",
+    "description": "test",
+    "source": "test",
+    "source_type": "path",
+    "key": "test-key",
+}
+
+
+class TestCheckpointQwenImageVariantDetection:
+    def _make_mock_mod(self, filename: str) -> MagicMock:
+        mod = MagicMock()
+        mod.path = Path(f"/fake/models/{filename}")
+        return mod
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_edit_in_filename_sets_edit_variant(self, _rfo, _rif, _hgt, _hqk):
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-edit-2511-fp8.safetensors")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_edit_case_insensitive(self, _rfo, _rif, _hgt, _hqk):
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("Qwen-Image-EDIT-2511-fp8.safetensors")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_no_marker_no_edit_in_filename_defaults_to_generate(self, _rfo, _rif, _hgt, _hqk):
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-2512-bf16.safetensors")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Generate
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_marker_tensor_sets_edit_variant(self, _rfo, _rif, _hgt, _hqk):
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("some-arbitrary-name.safetensors")
+        mod.load_state_dict.return_value = {"__index_timestep_zero__": object()}
+
+        config = Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_explicit_variant_override_not_overwritten(self, _rfo, _rif, _hgt, _hqk):
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-edit-2511-fp8.safetensors")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_Checkpoint_QwenImage_Config.from_model_on_disk(
+            mod, {**_REQUIRED_FIELDS, "variant": QwenImageVariantType.Generate}
+        )
+        assert config.variant == QwenImageVariantType.Generate
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_rejects_gguf_state_dict(self, _rfo, _rif, _hgt, _hqk):
+        """Checkpoint config must NOT match files that look GGUF-quantized."""
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-edit-2511-Q4_K_M.gguf")
+        mod.load_state_dict.return_value = {}
+
+        with pytest.raises(NotAMatchError):
+            Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=False)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_rejects_non_qwen_state_dict(self, _rfo, _rif, _hgt, _hqk):
+        """Checkpoint config must NOT match files whose state dict isn't Qwen Image."""
+        from invokeai.backend.model_manager.configs.main import Main_Checkpoint_QwenImage_Config
+
+        mod = self._make_mock_mod("not-a-qwen-model.safetensors")
+        mod.load_state_dict.return_value = {}
+
+        with pytest.raises(NotAMatchError):
+            Main_Checkpoint_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+
+
+class TestHasQwenImageKeys:
+    """Detection must agree with the loader, which strips ComfyUI prefixes before loading."""
+
+    def test_bare_keys_detected(self):
+        from invokeai.backend.model_manager.configs.main import _has_qwen_image_keys
+
+        sd = {"txt_in.weight": 1, "txt_norm.weight": 1, "img_in.weight": 1}
+        assert _has_qwen_image_keys(sd)
+
+    @pytest.mark.parametrize("prefix", ["model.diffusion_model.", "diffusion_model."])
+    def test_comfyui_prefixed_keys_detected(self, prefix: str):
+        """A ComfyUI checkpoint with prefixed keys must still be identified so it reaches the loader."""
+        from invokeai.backend.model_manager.configs.main import _has_qwen_image_keys
+
+        sd = {f"{prefix}txt_in.weight": 1, f"{prefix}txt_norm.weight": 1, f"{prefix}img_in.weight": 1}
+        assert _has_qwen_image_keys(sd)
+
+    def test_flux_rejected(self):
+        from invokeai.backend.model_manager.configs.main import _has_qwen_image_keys
+
+        sd = {"txt_in.weight": 1, "txt_norm.weight": 1, "img_in.weight": 1, "context_embedder.weight": 1}
+        assert not _has_qwen_image_keys(sd)
+
+    def test_prefixed_marker_sets_edit_variant(self):
+        """The Edit marker tensor may also carry a ComfyUI prefix."""
+        from invokeai.backend.model_manager.configs.main import _infer_qwen_image_variant
+        from invokeai.backend.model_manager.taxonomy import QwenImageVariantType
+
+        sd = {"model.diffusion_model.__index_timestep_zero__": object()}
+        assert _infer_qwen_image_variant(sd, Path("/fake/plain-name.safetensors")) == QwenImageVariantType.Edit
+
+
+class TestEditTokenHeuristic:
+    """The filename "edit" heuristic must match the token, not any substring."""
+
+    @pytest.mark.parametrize(
+        "stem",
+        ["qwen-image-edit-2511", "qwen_image_edit_2509", "Qwen-Image-EDIT", "model.edit"],
+    )
+    def test_edit_token_matches(self, stem: str):
+        from invokeai.backend.model_manager.configs.main import _infer_qwen_image_variant
+        from invokeai.backend.model_manager.taxonomy import QwenImageVariantType
+
+        assert _infer_qwen_image_variant({}, Path(f"/fake/{stem}.safetensors")) == QwenImageVariantType.Edit
+
+    @pytest.mark.parametrize("stem", ["credited-model", "edited-final", "unedited", "qwen-image"])
+    def test_edit_substring_does_not_false_positive(self, stem: str):
+        from invokeai.backend.model_manager.configs.main import _infer_qwen_image_variant
+        from invokeai.backend.model_manager.taxonomy import QwenImageVariantType
+
+        assert _infer_qwen_image_variant({}, Path(f"/fake/{stem}.safetensors")) == QwenImageVariantType.Generate

tests/backend/model_manager/configs/test_qwen_image_gguf_variant_detection.py

@@ -0,0 +1,122 @@
+"""Tests for GGUF Qwen Image variant detection.
+
+Detection precedence:
+1. Explicit `variant` in override_fields wins.
+2. Presence of the `__index_timestep_zero__` tensor in the state dict marks an Edit model.
+3. Otherwise fall back to a filename heuristic ("edit" in the stem → Edit).
+4. Otherwise default to Generate.
+"""
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from invokeai.backend.model_manager.taxonomy import QwenImageVariantType
+
+# Required fields for the Pydantic config model
+_REQUIRED_FIELDS = {
+    "hash": "blake3:fakehash",
+    "path": "/fake/models/test.gguf",
+    "file_size": 1000,
+    "name": "test-model",
+    "description": "test",
+    "source": "test",
+    "source_type": "path",
+    "key": "test-key",
+}
+
+
+class TestGGUFQwenImageVariantDetection:
+    """Test that GGUF Qwen Image models infer the edit variant from filename."""
+
+    def _make_mock_mod(self, filename: str) -> MagicMock:
+        """Create a mock ModelOnDisk with the given filename."""
+        mod = MagicMock()
+        mod.path = Path(f"/fake/models/{filename}")
+        return mod
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_edit_in_filename_sets_edit_variant(self, _rfo, _rif, _hgt, _hqk):
+        """A GGUF file with 'edit' in the name should be tagged as edit variant."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-edit-2511-Q4_K_M.gguf")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_edit_case_insensitive(self, _rfo, _rif, _hgt, _hqk):
+        """The 'edit' check should be case-insensitive."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        mod = self._make_mock_mod("Qwen-Image-EDIT-2511-Q8_0.gguf")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_no_marker_no_edit_in_filename_defaults_to_generate(self, _rfo, _rif, _hgt, _hqk):
+        """A GGUF file without the marker tensor or 'edit' in the name should default to Generate."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-2512-Q4_K_M.gguf")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Generate
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_marker_tensor_sets_edit_variant(self, _rfo, _rif, _hgt, _hqk):
+        """Presence of `__index_timestep_zero__` in the state dict should set the Edit variant."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        # Filename has no "edit" marker, but the tensor is present
+        mod = self._make_mock_mod("some-arbitrary-name.gguf")
+        mod.load_state_dict.return_value = {"__index_timestep_zero__": object()}
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_marker_tensor_takes_precedence_over_filename(self, _rfo, _rif, _hgt, _hqk):
+        """The marker tensor wins even when the filename has no 'edit' substring."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-2512-Q4_K_M.gguf")
+        mod.load_state_dict.return_value = {"__index_timestep_zero__": object()}
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS})
+        assert config.variant == QwenImageVariantType.Edit
+
+    @patch("invokeai.backend.model_manager.configs.main._has_qwen_image_keys", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main._has_ggml_tensors", return_value=True)
+    @patch("invokeai.backend.model_manager.configs.main.raise_if_not_file")
+    @patch("invokeai.backend.model_manager.configs.main.raise_for_override_fields")
+    def test_explicit_variant_override_not_overwritten(self, _rfo, _rif, _hgt, _hqk):
+        """An explicit variant in override_fields should not be overwritten by filename heuristic."""
+        from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+
+        mod = self._make_mock_mod("qwen-image-edit-2511-Q4_K_M.gguf")
+        mod.load_state_dict.return_value = {}
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(
+            mod, {**_REQUIRED_FIELDS, "variant": QwenImageVariantType.Generate}
+        )
+        assert config.variant == QwenImageVariantType.Generate

tests/backend/model_manager/configs/test_qwen_image_main_config.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import MagicMock
+
+import gguf
+import pytest
+import torch
+
+from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config
+from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
+
+
+def _build_ggml_tensor() -> GGMLTensor:
+    return GGMLTensor(
+        data=torch.zeros((1,), dtype=torch.uint8),
+        ggml_quantization_type=gguf.GGMLQuantizationType.Q4_0,
+        tensor_shape=torch.Size([1, 1]),
+        compute_dtype=torch.float32,
+    )
+
+
+@pytest.mark.parametrize("is_edit_model", [True, False])
+def test_qwen_gguf_config_sets_a_variant_for_imported_models(is_edit_model: bool) -> None:
+    with TemporaryDirectory() as tmpdir:
+        model_path = Path(tmpdir) / ("qwen-image-edit.gguf" if is_edit_model else "qwen-image.gguf")
+        model_name = "Qwen Image Edit GGUF" if is_edit_model else "Qwen Image GGUF"
+        model_path.touch()
+
+        mod = MagicMock()
+        mod.path = model_path
+        mod.load_state_dict.return_value = {
+            "txt_in.weight": _build_ggml_tensor(),
+            "txt_norm.weight": _build_ggml_tensor(),
+            "img_in.weight": _build_ggml_tensor(),
+        }
+
+        config = Main_GGUF_QwenImage_Config.from_model_on_disk(
+            mod,
+            {
+                "hash": "test-hash",
+                "path": str(model_path),
+                "file_size": model_path.stat().st_size,
+                "name": model_name,
+                "source": str(model_path),
+                "source_type": "path",
+            },
+        )
+
+    if is_edit_model:
+        assert config.variant == "edit"
+    else:
+        assert config.variant == "generate"

tests/backend/model_manager/configs/test_qwen_vl_encoder_config.py

@@ -0,0 +1,101 @@
+"""Tests for Qwen VL encoder config identification.
+
+The single-file checkpoint identifier reads only the safetensors key index
+instead of loading the full tensor data — a 7GB fp8 encoder otherwise pins
+~7GB of RAM during model scan.
+"""
+
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError
+from invokeai.backend.model_manager.configs.qwen_vl_encoder import (
+    QwenVLEncoder_Checkpoint_Config,
+    _has_qwen_vl_keys,
+    _read_safetensors_keys,
+)
+
+_OVERRIDE_FIELDS: dict[str, object] = {
+    "hash": "blake3:fakehash",
+    "path": "/fake/models/test-model.safetensors",
+    "file_size": 1000,
+    "name": "test-model",
+    "description": "test",
+    "source": "test",
+    "source_type": "path",
+    "key": "test-key",
+}
+
+
+def _write_safetensors(path: Path, keys: list[str]) -> None:
+    """Write a safetensors file with tiny placeholder tensors for the given keys."""
+    sd = {k: torch.zeros(1, dtype=torch.float32) for k in keys}
+    save_file(sd, str(path))
+
+
+def test_has_qwen_vl_keys_accepts_lm_plus_visual() -> None:
+    assert _has_qwen_vl_keys(["model.embed_tokens.weight", "visual.patch_embed.proj.weight"])
+    assert _has_qwen_vl_keys(["model.layers.0.self_attn.q_proj.weight", "visual.blocks.0.norm1.weight"])
+
+
+def test_has_qwen_vl_keys_rejects_lm_only() -> None:
+    """Text-only Qwen3/Qwen2 encoders have LM keys but no visual tower — must not match."""
+    assert not _has_qwen_vl_keys(["model.embed_tokens.weight", "model.layers.0.self_attn.q_proj.weight"])
+
+
+def test_has_qwen_vl_keys_rejects_empty() -> None:
+    assert not _has_qwen_vl_keys([])
+
+
+def test_read_safetensors_keys_returns_keys_without_loading_tensors() -> None:
+    with TemporaryDirectory() as tmpdir:
+        path = Path(tmpdir) / "tiny.safetensors"
+        _write_safetensors(path, ["model.embed_tokens.weight", "visual.patch_embed.proj.weight"])
+
+        keys = _read_safetensors_keys(path)
+
+        assert set(keys) == {"model.embed_tokens.weight", "visual.patch_embed.proj.weight"}
+
+
+def test_checkpoint_config_matches_qwen_vl_safetensors() -> None:
+    with TemporaryDirectory() as tmpdir:
+        path = Path(tmpdir) / "qwen_vl.safetensors"
+        _write_safetensors(path, ["model.embed_tokens.weight", "visual.patch_embed.proj.weight"])
+
+        mod = MagicMock()
+        mod.path = path
+
+        config = QwenVLEncoder_Checkpoint_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+        assert config.type.value == "qwen_vl_encoder"
+        assert config.format.value == "checkpoint"
+
+
+def test_checkpoint_config_rejects_lm_only_safetensors() -> None:
+    """A text-only LM checkpoint must not be identified as a Qwen VL encoder."""
+    with TemporaryDirectory() as tmpdir:
+        path = Path(tmpdir) / "text_only.safetensors"
+        _write_safetensors(path, ["model.embed_tokens.weight", "model.layers.0.self_attn.q_proj.weight"])
+
+        mod = MagicMock()
+        mod.path = path
+
+        with pytest.raises(NotAMatchError, match="does not look like a Qwen2.5-VL"):
+            QwenVLEncoder_Checkpoint_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))
+
+
+def test_checkpoint_config_rejects_non_safetensors_extension() -> None:
+    """Bin/ckpt/pt files should be rejected cheaply without attempting to read the header."""
+    with TemporaryDirectory() as tmpdir:
+        path = Path(tmpdir) / "weights.bin"
+        path.write_bytes(b"not a safetensors file")
+
+        mod = MagicMock()
+        mod.path = path
+
+        with pytest.raises(NotAMatchError, match="expected a .safetensors file"):
+            QwenVLEncoder_Checkpoint_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS))

tests/backend/model_manager/data/invokeai_root/README

@@ -0,0 +1 @@
+This is an empty invokeai root that is used as a template for model manager tests.

tests/backend/model_manager/data/invokeai_root/configs/stable-diffusion/v1-inference.yaml

@@ -0,0 +1,79 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: invokeai.backend.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+
+    scheduler_config: # 10000 warmup steps
+      target: invokeai.backend.stable_diffusion.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    personalization_config:
+      target: invokeai.backend.stable_diffusion.embedding_manager.EmbeddingManager
+      params:
+        placeholder_strings: ["*"]
+        initializer_words: ['sculpture']
+        per_image_tokens: false
+        num_vectors_per_token: 1
+        progressive_words: False
+
+    unet_config:
+      target: invokeai.backend.stable_diffusion.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: invokeai.backend.stable_diffusion.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: invokeai.backend.stable_diffusion.encoders.modules.WeightedFrozenCLIPEmbedder

tests/backend/model_manager/data/invokeai_root/databases/README

@@ -0,0 +1 @@
+This is a template empty invokeai root directory used to test model management.

tests/backend/model_manager/data/invokeai_root/models/README

@@ -0,0 +1 @@
+This is a template empty invokeai root directory used to test model management.

tests/backend/model_manager/data/test_files/test-diffusers-main/model_index.json

@@ -0,0 +1,34 @@
+{
+  "_class_name": "StableDiffusionXLPipeline",
+  "_diffusers_version": "0.23.0",
+  "_name_or_path": "stabilityai/sdxl-turbo",
+  "force_zeros_for_empty_prompt": true,
+  "scheduler": [
+    "diffusers",
+    "EulerAncestralDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "CLIPTextModel"
+  ],
+  "text_encoder_2": [
+    "transformers",
+    "CLIPTextModelWithProjection"
+  ],
+  "tokenizer": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "tokenizer_2": [
+    "transformers",
+    "CLIPTokenizer"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DConditionModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/scheduler/scheduler_config.json

@@ -0,0 +1,17 @@
+{
+  "_class_name": "EulerAncestralDiscreteScheduler",
+  "_diffusers_version": "0.23.0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "interpolation_type": "linear",
+  "num_train_timesteps": 1000,
+  "prediction_type": "epsilon",
+  "sample_max_value": 1.0,
+  "set_alpha_to_one": false,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "timestep_spacing": "trailing",
+  "trained_betas": null
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/text_encoder/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "/home/lstein/.cache/huggingface/hub/models--stabilityai--sdxl-turbo/snapshots/fbda35297a8280789ffe2e25206800702fa5c4c1/text_encoder",
+  "architectures": [
+    "CLIPTextModel"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 768,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "projection_dim": 768,
+  "torch_dtype": "float16",
+  "transformers_version": "4.35.0",
+  "vocab_size": 49408
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/text_encoder_2/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "/home/lstein/.cache/huggingface/hub/models--stabilityai--sdxl-turbo/snapshots/fbda35297a8280789ffe2e25206800702fa5c4c1/text_encoder_2",
+  "architectures": [
+    "CLIPTextModelWithProjection"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 0,
+  "dropout": 0.0,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_size": 1280,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 77,
+  "model_type": "clip_text_model",
+  "num_attention_heads": 20,
+  "num_hidden_layers": 32,
+  "pad_token_id": 1,
+  "projection_dim": 1280,
+  "torch_dtype": "float16",
+  "transformers_version": "4.35.0",
+  "vocab_size": 49408
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/tokenizer/tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/tokenizer_2/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "!",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/tokenizer_2/tokenizer_config.json

@@ -0,0 +1,38 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "!",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49406": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": true,
+  "do_lower_case": true,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "!",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<|endoftext|>"
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/unet/config.json

@@ -0,0 +1,73 @@
+{
+  "_class_name": "UNet2DConditionModel",
+  "_diffusers_version": "0.23.0",
+  "_name_or_path": "/home/lstein/.cache/huggingface/hub/models--stabilityai--sdxl-turbo/snapshots/fbda35297a8280789ffe2e25206800702fa5c4c1/unet",
+  "act_fn": "silu",
+  "addition_embed_type": "text_time",
+  "addition_embed_type_num_heads": 64,
+  "addition_time_embed_dim": 256,
+  "attention_head_dim": [
+    5,
+    10,
+    20
+  ],
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "cross_attention_dim": 2048,
+  "cross_attention_norm": null,
+  "down_block_types": [
+    "DownBlock2D",
+    "CrossAttnDownBlock2D",
+    "CrossAttnDownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "dropout": 0.0,
+  "dual_cross_attention": false,
+  "encoder_hid_dim": null,
+  "encoder_hid_dim_type": null,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_only_cross_attention": null,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DCrossAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "only_cross_attention": false,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": 2816,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": [
+    1,
+    2,
+    10
+  ],
+  "up_block_types": [
+    "CrossAttnUpBlock2D",
+    "CrossAttnUpBlock2D",
+    "UpBlock2D"
+  ],
+  "upcast_attention": null,
+  "use_linear_projection": true
+}

tests/backend/model_manager/data/test_files/test-diffusers-main/vae/config.json

@@ -0,0 +1,32 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.23.0",
+  "_name_or_path": "/home/lstein/.cache/huggingface/hub/models--stabilityai--sdxl-turbo/snapshots/fbda35297a8280789ffe2e25206800702fa5c4c1/vae",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 1024,
+  "scaling_factor": 0.13025,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_only_full_load.py

@@ -0,0 +1,143 @@
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_only_full_load import (
+    CachedModelOnlyFullLoad,
+)
+from tests.backend.model_manager.load.model_cache.cached_model.utils import (
+    DummyModule,
+    parameterize_keep_ram_copy,
+    parameterize_mps_and_cuda,
+)
+
+
+class NonTorchModel:
+    """A model that does not sub-class torch.nn.Module."""
+
+    def __init__(self):
+        self.linear = torch.nn.Linear(10, 32)
+
+    def run_inference(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_total_bytes(device: str, keep_ram_copy: bool):
+    model = DummyModule()
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device(device), total_bytes=100, keep_ram_copy=keep_ram_copy
+    )
+    assert cached_model.total_bytes() == 100
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_is_in_vram(device: str, keep_ram_copy: bool):
+    model = DummyModule()
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device(device), total_bytes=100, keep_ram_copy=keep_ram_copy
+    )
+    assert not cached_model.is_in_vram()
+    assert cached_model.cur_vram_bytes() == 0
+
+    cached_model.full_load_to_vram()
+    assert cached_model.is_in_vram()
+    assert cached_model.cur_vram_bytes() == 100
+
+    cached_model.full_unload_from_vram()
+    assert not cached_model.is_in_vram()
+    assert cached_model.cur_vram_bytes() == 0
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_full_load_and_unload(device: str, keep_ram_copy: bool):
+    model = DummyModule()
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device(device), total_bytes=100, keep_ram_copy=keep_ram_copy
+    )
+    assert cached_model.full_load_to_vram() == 100
+    assert cached_model.is_in_vram()
+    assert all(p.device.type == device for p in cached_model.model.parameters())
+
+    assert cached_model.full_unload_from_vram() == 100
+    assert not cached_model.is_in_vram()
+    assert all(p.device.type == "cpu" for p in cached_model.model.parameters())
+
+
+@parameterize_mps_and_cuda
+def test_cached_model_get_cpu_state_dict(device: str):
+    model = DummyModule()
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device(device), total_bytes=100, keep_ram_copy=True
+    )
+    assert not cached_model.is_in_vram()
+
+    # The CPU state dict can be accessed and has the expected properties.
+    cpu_state_dict = cached_model.get_cpu_state_dict()
+    assert cpu_state_dict is not None
+    assert len(cpu_state_dict) == len(model.state_dict())
+    assert all(p.device.type == "cpu" for p in cpu_state_dict.values())
+
+    # Full load the model into VRAM.
+    cached_model.full_load_to_vram()
+    assert cached_model.is_in_vram()
+
+    # The CPU state dict is still available, and still on the CPU.
+    cpu_state_dict = cached_model.get_cpu_state_dict()
+    assert cpu_state_dict is not None
+    assert len(cpu_state_dict) == len(model.state_dict())
+    assert all(p.device.type == "cpu" for p in cpu_state_dict.values())
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_full_load_and_inference(device: str, keep_ram_copy: bool):
+    model = DummyModule()
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device(device), total_bytes=100, keep_ram_copy=keep_ram_copy
+    )
+    assert not cached_model.is_in_vram()
+
+    # Run inference on the CPU.
+    x = torch.randn(1, 10)
+    output1 = model(x)
+    assert output1.device.type == "cpu"
+
+    # Full load the model into VRAM.
+    cached_model.full_load_to_vram()
+    assert cached_model.is_in_vram()
+
+    # Run inference on the GPU.
+    output2 = model(x.to(device))
+    assert output2.device.type == device
+
+    # The outputs should be the same for both runs.
+    assert torch.allclose(output1, output2.to("cpu"))
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_non_torch_model(device: str, keep_ram_copy: bool):
+    model = NonTorchModel()
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device(device), total_bytes=100, keep_ram_copy=keep_ram_copy
+    )
+    assert not cached_model.is_in_vram()
+
+    # The model does not have a CPU state dict.
+    assert cached_model.get_cpu_state_dict() is None
+
+    # Attempting to load the model into VRAM should have no effect.
+    cached_model.full_load_to_vram()
+    assert not cached_model.is_in_vram()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Attempting to unload the model from VRAM should have no effect.
+    cached_model.full_unload_from_vram()
+    assert not cached_model.is_in_vram()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Running inference on the CPU should work.
+    output1 = model.run_inference(torch.randn(1, 10))
+    assert output1.device.type == "cpu"

tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_with_partial_load.py

@@ -0,0 +1,341 @@
+import itertools
+
+import pytest
+import torch
+
+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.util.calc_tensor_size import calc_tensor_size
+from tests.backend.model_manager.load.model_cache.cached_model.utils import (
+    DummyModule,
+    parameterize_keep_ram_copy,
+    parameterize_mps_and_cuda,
+)
+
+
+@pytest.fixture
+def model():
+    model = DummyModule()
+    apply_custom_layers_to_model(model)
+    return model
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_total_bytes(device: str, model: DummyModule, keep_ram_copy: bool):
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    linear1_numel = 10 * 32 + 32
+    linear2_numel = 32 * 64 + 64
+    buffer1_numel = 64
+    # Note that the non-persistent buffer (buffer2) is not included in .total_bytes() calculation.
+    assert cached_model.total_bytes() == (linear1_numel + linear2_numel + buffer1_numel) * 4
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_cur_vram_bytes(device: str, model: DummyModule, keep_ram_copy: bool):
+    # Model starts in CPU memory.
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Full load the model into VRAM.
+    cached_model.full_load_to_vram()
+    assert cached_model.cur_vram_bytes() > 0
+    assert cached_model.cur_vram_bytes() == cached_model.total_bytes()
+    assert all(p.device.type == device for p in model.parameters())
+    assert all(p.device.type == device for p in model.buffers())
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_partial_load(device: str, model: DummyModule, keep_ram_copy: bool):
+    # Model starts in CPU memory.
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Partially load the model into VRAM.
+    target_vram_bytes = int(model_total_bytes * 0.6)
+    loaded_bytes = cached_model.partial_load_to_vram(target_vram_bytes)
+
+    # Check that the model is partially loaded into VRAM.
+    assert loaded_bytes > 0
+    assert loaded_bytes < model_total_bytes
+    assert loaded_bytes == cached_model.cur_vram_bytes()
+    assert loaded_bytes == sum(
+        calc_tensor_size(p)
+        for n, p in itertools.chain(model.named_parameters(), model.named_buffers())
+        if p.device.type == device and n != "buffer2"
+    )
+
+    # Check that the model's modules have device autocasting enabled.
+    assert model.linear1.is_device_autocasting_enabled()
+    assert model.linear2.is_device_autocasting_enabled()
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_partial_unload(device: str, model: DummyModule, keep_ram_copy: bool):
+    # Model starts in CPU memory.
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Full load the model into VRAM.
+    cached_model.full_load_to_vram()
+    assert cached_model.cur_vram_bytes() == model_total_bytes
+
+    # Partially unload the model from VRAM.
+    bytes_to_free = int(model_total_bytes * 0.4)
+    freed_bytes = cached_model.partial_unload_from_vram(bytes_to_free)
+
+    # Check that the model is partially unloaded from VRAM.
+    assert freed_bytes >= bytes_to_free
+    assert freed_bytes < model_total_bytes
+    assert freed_bytes == model_total_bytes - cached_model.cur_vram_bytes()
+    assert freed_bytes == sum(
+        calc_tensor_size(p) for p in itertools.chain(model.parameters(), model.buffers()) if p.device.type == "cpu"
+    )
+
+    # Check that the model's modules still have device autocasting enabled.
+    assert model.linear1.is_device_autocasting_enabled()
+    assert model.linear2.is_device_autocasting_enabled()
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_partial_unload_keep_required_weights_in_vram(
+    device: str, model: DummyModule, keep_ram_copy: bool
+):
+    # Model starts in CPU memory.
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Full load the model into VRAM.
+    cached_model.full_load_to_vram()
+    assert cached_model.cur_vram_bytes() == model_total_bytes
+
+    # Partially unload the model from VRAM, but request the required weights to be kept in VRAM.
+    bytes_to_free = int(model_total_bytes)
+    freed_bytes = cached_model.partial_unload_from_vram(bytes_to_free, keep_required_weights_in_vram=True)
+
+    # Check that the model is partially unloaded from VRAM.
+    assert freed_bytes < model_total_bytes
+    assert freed_bytes == model_total_bytes - cached_model.cur_vram_bytes()
+    assert freed_bytes == sum(
+        calc_tensor_size(p) for p in itertools.chain(model.parameters(), model.buffers()) if p.device.type == "cpu"
+    )
+    # The parameters should be offloaded to the CPU, because they are in Linear layers.
+    assert all(p.device.type == "cpu" for p in model.parameters())
+    # The buffer should still be on the device, because it is in a layer that does not support autocast.
+    assert all(p.device.type == device for p in model.buffers())
+
+    # Check that the model's modules still have device autocasting enabled.
+    assert model.linear1.is_device_autocasting_enabled()
+    assert model.linear2.is_device_autocasting_enabled()
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_full_load_and_unload(device: str, model: DummyModule, keep_ram_copy: bool):
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+
+    # Model starts in CPU memory.
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Full load the model into VRAM.
+    loaded_bytes = cached_model.full_load_to_vram()
+    assert loaded_bytes > 0
+    assert loaded_bytes == model_total_bytes
+    assert loaded_bytes == cached_model.cur_vram_bytes()
+    assert all(p.device.type == device for p in itertools.chain(model.parameters(), model.buffers()))
+    assert not model.linear1.is_device_autocasting_enabled()
+    assert not model.linear2.is_device_autocasting_enabled()
+
+    # Full unload the model from VRAM.
+    unloaded_bytes = cached_model.full_unload_from_vram()
+
+    # Check that the model is fully unloaded from VRAM.
+    assert unloaded_bytes > 0
+    assert unloaded_bytes == model_total_bytes
+    assert cached_model.cur_vram_bytes() == 0
+    # Note that the non-persistent buffer (buffer2) is not required to be unloaded from VRAM.
+    assert all(
+        p.device.type == "cpu"
+        for n, p in itertools.chain(model.named_parameters(), model.named_buffers())
+        if n != "buffer2"
+    )
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_full_load_from_partial(device: str, model: DummyModule, keep_ram_copy: bool):
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+
+    # Model starts in CPU memory.
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Partially load the model into VRAM.
+    target_vram_bytes = int(model_total_bytes * 0.6)
+    loaded_bytes = cached_model.partial_load_to_vram(target_vram_bytes)
+    assert loaded_bytes > 0
+    assert loaded_bytes < model_total_bytes
+    assert loaded_bytes == cached_model.cur_vram_bytes()
+    assert model.linear1.is_device_autocasting_enabled()
+    assert model.linear2.is_device_autocasting_enabled()
+
+    # Full load the rest of the model into VRAM.
+    loaded_bytes_2 = cached_model.full_load_to_vram()
+    assert loaded_bytes_2 > 0
+    assert loaded_bytes_2 < model_total_bytes
+    assert loaded_bytes + loaded_bytes_2 == cached_model.cur_vram_bytes()
+    assert loaded_bytes + loaded_bytes_2 == model_total_bytes
+    assert all(p.device.type == device for p in itertools.chain(model.parameters(), model.buffers()))
+    assert not model.linear1.is_device_autocasting_enabled()
+    assert not model.linear2.is_device_autocasting_enabled()
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_full_unload_from_partial(device: str, model: DummyModule, keep_ram_copy: bool):
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+
+    # Model starts in CPU memory.
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Partially load the model into VRAM.
+    target_vram_bytes = int(model_total_bytes * 0.6)
+    loaded_bytes = cached_model.partial_load_to_vram(target_vram_bytes)
+    assert loaded_bytes > 0
+    assert loaded_bytes < model_total_bytes
+    assert loaded_bytes == cached_model.cur_vram_bytes()
+
+    # Full unload the model from VRAM.
+    unloaded_bytes = cached_model.full_unload_from_vram()
+    assert unloaded_bytes > 0
+    assert unloaded_bytes == loaded_bytes
+    assert cached_model.cur_vram_bytes() == 0
+    # Note that the non-persistent buffer (buffer2) is not required to be unloaded from VRAM.
+    assert all(
+        p.device.type == "cpu"
+        for n, p in itertools.chain(model.named_parameters(), model.named_buffers())
+        if n != "buffer2"
+    )
+
+
+@parameterize_mps_and_cuda
+def test_cached_model_get_cpu_state_dict(device: str, model: DummyModule):
+    cached_model = CachedModelWithPartialLoad(model=model, compute_device=torch.device(device), keep_ram_copy=True)
+
+    # Model starts in CPU memory.
+    assert cached_model.cur_vram_bytes() == 0
+
+    # The CPU state dict can be accessed and has the expected properties.
+    cpu_state_dict = cached_model.get_cpu_state_dict()
+    assert cpu_state_dict is not None
+    assert len(cpu_state_dict) == len(model.state_dict())
+    assert all(p.device.type == "cpu" for p in cpu_state_dict.values())
+
+    # Full load the model into VRAM.
+    cached_model.full_load_to_vram()
+    assert cached_model.cur_vram_bytes() == cached_model.total_bytes()
+
+    # The CPU state dict is still available, and still on the CPU.
+    cpu_state_dict = cached_model.get_cpu_state_dict()
+    assert cpu_state_dict is not None
+    assert len(cpu_state_dict) == len(model.state_dict())
+    assert all(p.device.type == "cpu" for p in cpu_state_dict.values())
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_full_load_and_inference(device: str, model: DummyModule, keep_ram_copy: bool):
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    # Model starts in CPU memory.
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Run inference on the CPU.
+    x = torch.randn(1, 10)
+    output1 = model(x)
+    assert output1.device.type == "cpu"
+
+    # Full load the model into VRAM.
+    loaded_bytes = cached_model.full_load_to_vram()
+    assert loaded_bytes > 0
+    assert loaded_bytes == model_total_bytes
+    assert loaded_bytes == cached_model.cur_vram_bytes()
+    assert all(p.device.type == device for p in itertools.chain(model.parameters(), model.buffers()))
+
+    # Run inference on the GPU.
+    output2 = model(x.to(device))
+    assert output2.device.type == device
+
+    # The outputs should be the same for both runs.
+    assert torch.allclose(output1, output2.to("cpu"))
+
+
+@parameterize_mps_and_cuda
+@parameterize_keep_ram_copy
+def test_cached_model_partial_load_and_inference(device: str, model: DummyModule, keep_ram_copy: bool):
+    # Model starts in CPU memory.
+    cached_model = CachedModelWithPartialLoad(
+        model=model, compute_device=torch.device(device), keep_ram_copy=keep_ram_copy
+    )
+    model_total_bytes = cached_model.total_bytes()
+    assert cached_model.cur_vram_bytes() == 0
+
+    # Run inference on the CPU.
+    x = torch.randn(1, 10)
+    output1 = model(x)
+    assert output1.device.type == "cpu"
+
+    # Partially load the model into VRAM.
+    target_vram_bytes = int(model_total_bytes * 0.6)
+    loaded_bytes = cached_model.partial_load_to_vram(target_vram_bytes)
+
+    # Check that the model is partially loaded into VRAM.
+    assert loaded_bytes > 0
+    assert loaded_bytes < model_total_bytes
+    assert loaded_bytes == cached_model.cur_vram_bytes()
+    assert loaded_bytes == sum(
+        calc_tensor_size(p)
+        for n, p in itertools.chain(model.named_parameters(), model.named_buffers())
+        if p.device.type == device and n != "buffer2"
+    )
+    # Check that the model's modules have device autocasting enabled.
+    assert model.linear1.is_device_autocasting_enabled()
+    assert model.linear2.is_device_autocasting_enabled()
+
+    # Run inference on the GPU.
+    output2 = model(x.to(device))
+    assert output2.device.type == device
+
+    # The output should be the same as the output from the CPU.
+    assert torch.allclose(output1, output2.to("cpu"))

tests/backend/model_manager/load/model_cache/cached_model/test_repair_required_tensors.py

@@ -0,0 +1,47 @@
+import pytest
+import torch
+
+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,
+)
+
+
+class ModelWithRequiredScale(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(4, 4)
+        self.scale = torch.nn.Parameter(torch.ones(4))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x) * self.scale
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param(
+            torch.device("cuda"), marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA device")
+        ),
+        pytest.param(
+            torch.device("mps"),
+            marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires MPS device"),
+        ),
+    ],
+)
+@pytest.mark.parametrize("keep_ram_copy", [True, False])
+@torch.no_grad()
+def test_repair_required_tensors_on_compute_device(device: torch.device, keep_ram_copy: bool):
+    model = ModelWithRequiredScale()
+    apply_custom_layers_to_model(model, device_autocasting_enabled=True)
+    cached_model = CachedModelWithPartialLoad(model=model, compute_device=device, keep_ram_copy=keep_ram_copy)
+
+    cached_model._cur_vram_bytes = 0
+    repaired_tensors = cached_model.repair_required_tensors_on_compute_device()
+
+    assert repaired_tensors == 1
+    assert cached_model._cur_vram_bytes is None
+    assert model.scale.device.type == device.type
+    assert all(param.device.type == "cpu" for param in model.linear.parameters())

tests/backend/model_manager/load/model_cache/cached_model/utils.py

@@ -0,0 +1,41 @@
+import os
+
+import pytest
+import torch
+
+
+class DummyModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear1 = torch.nn.Linear(10, 32)
+        self.linear2 = torch.nn.Linear(32, 64)
+        self.register_buffer("buffer1", torch.ones(64))
+        # Non-persistent buffers are not included in the state dict. We need to make sure that this case is handled
+        # correctly by the partial loading code.
+        self.register_buffer("buffer2", torch.ones(64), persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear1(x)
+        x = self.linear2(x)
+        x = x + self.buffer1
+        x = x + self.buffer2
+        return x
+
+
+is_github_ci = os.getenv("GITHUB_ACTIONS") == "true"
+
+parameterize_mps_and_cuda = pytest.mark.parametrize(
+    ("device"),
+    [
+        pytest.param(
+            "mps",
+            marks=pytest.mark.skipif(
+                is_github_ci or not torch.backends.mps.is_available(),
+                reason="MPS is very flaky in CI" if is_github_ci else "MPS is not available.",
+            ),
+        ),
+        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.")),
+    ],
+)
+
+parameterize_keep_ram_copy = pytest.mark.parametrize("keep_ram_copy", [True, False])

tests/backend/model_manager/load/model_cache/test_model_cache_drop_model.py

@@ -0,0 +1,224 @@
+"""Tests for `ModelCache.drop_model` — used by the model_manager API to invalidate cached
+entries when a setting that changes how a model loads (e.g. `fp8_storage`, `cpu_only`) is
+toggled. Without this, the toggle is silently a no-op until the entry is evicted by other
+means (clear cache, eviction under memory pressure, restart).
+
+Also covers:
+- Locked entries are marked stale and evicted by `unlock()` — without that, a setting toggled
+  during an in-flight generation would survive on the locked entry and silently be reused.
+- `stats.cleared` and the `cleared` callbacks fire on invalidation, mirroring the eviction
+  path through `_make_room_internal`, so observers and stats stay accurate.
+"""
+
+import logging
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+
+
+@pytest.fixture
+def mock_logger():
+    logger = MagicMock()
+    logger.getEffectiveLevel.return_value = logging.INFO
+    return logger
+
+
+@pytest.fixture
+def cache(mock_logger):
+    cache = ModelCache(
+        execution_device_working_mem_gb=1.0,
+        enable_partial_loading=False,
+        keep_ram_copy_of_weights=True,
+        execution_device="cpu",
+        storage_device="cpu",
+        logger=mock_logger,
+    )
+    yield cache
+    cache.shutdown()
+
+
+def test_drop_model_removes_all_submodel_entries(cache: ModelCache):
+    """A model with multiple submodels has multiple cache keys (`<key>` and `<key>:<submodel>`);
+    drop_model must drop them all together so the next load rebuilds with the new settings.
+    """
+    model_key = "abc123"
+    cache.put(model_key, torch.randn(4))
+    cache.put(f"{model_key}:unet", torch.randn(4))
+    cache.put(f"{model_key}:text_encoder", torch.randn(4))
+    cache.put("other_model", torch.randn(4))
+    cache.put("other_model:unet", torch.randn(4))
+
+    dropped = cache.drop_model(model_key)
+
+    assert dropped == 3
+    assert model_key not in cache._cached_models
+    assert f"{model_key}:unet" not in cache._cached_models
+    assert f"{model_key}:text_encoder" not in cache._cached_models
+    # Unrelated model is left alone.
+    assert "other_model" in cache._cached_models
+    assert "other_model:unet" in cache._cached_models
+
+
+def test_drop_model_marks_locked_entries_stale_without_evicting(cache: ModelCache):
+    """Locked entries are in active use; we must not yank them out from under inference.
+    But we also must not silently retain them after the lock releases — otherwise a setting
+    toggle that happened during inference would survive and the next generation would reuse
+    the pre-change cached module. drop_model marks locked entries `is_stale=True`; unlock
+    evicts them as soon as the last lock releases.
+    """
+    model_key = "abc123"
+    cache.put(model_key, torch.randn(4))
+    cache.put(f"{model_key}:unet", torch.randn(4))
+
+    locked_entry = cache._cached_models[f"{model_key}:unet"]
+    locked_entry.lock()
+
+    dropped = cache.drop_model(model_key)
+
+    assert dropped == 1
+    assert model_key not in cache._cached_models
+    assert f"{model_key}:unet" in cache._cached_models
+    assert locked_entry.is_stale is True
+
+
+def test_unlock_evicts_stale_entry(cache: ModelCache):
+    """The flip side of `marks_locked_entries_stale`: the next `unlock` after a stale-marking
+    invalidation must actually remove the entry.
+    """
+    model_key = "abc123"
+    cache.put(model_key, torch.randn(4))
+    entry = cache._cached_models[model_key]
+    entry.lock()
+
+    cache.drop_model(model_key)
+
+    # Entry still here while locked.
+    assert model_key in cache._cached_models
+    assert entry.is_stale is True
+
+    cache.unlock(entry)
+
+    assert model_key not in cache._cached_models
+
+
+def test_unlock_does_not_evict_non_stale_entry(cache: ModelCache):
+    """The stale-eviction path must not affect ordinary unlock — only stale-marked entries
+    should be evicted on unlock.
+    """
+    model_key = "abc123"
+    cache.put(model_key, torch.randn(4))
+    entry = cache._cached_models[model_key]
+    entry.lock()
+
+    cache.unlock(entry)
+
+    # No drop_model was called, so entry should still be there.
+    assert model_key in cache._cached_models
+
+
+def test_unlock_only_evicts_when_last_lock_releases(cache: ModelCache):
+    """If the entry is held by multiple locks (the cache supports re-entrant locking via
+    `_locks`), eviction must wait until they all release. Otherwise we'd yank the entry out
+    from under a caller that still expects it loaded.
+    """
+    model_key = "abc123"
+    cache.put(model_key, torch.randn(4))
+    entry = cache._cached_models[model_key]
+    entry.lock()
+    entry.lock()
+
+    cache.drop_model(model_key)
+    assert entry.is_stale is True
+
+    cache.unlock(entry)
+    # Still locked by one holder — must remain.
+    assert model_key in cache._cached_models
+
+    cache.unlock(entry)
+    # Now fully released — eviction happens.
+    assert model_key not in cache._cached_models
+
+
+def test_drop_model_updates_stats_and_fires_callbacks(cache: ModelCache):
+    """drop_model is a real eviction path — observers watching for cache changes (stats,
+    cleared callbacks) must see it just like the make_room eviction path. Otherwise the UI
+    cache-stats panel and any external observer would miss invalidations.
+    """
+    model_key = "abc123"
+    # Use real nn.Modules so `total_bytes()` is non-zero (raw tensors are sized as 0 by
+    # `calc_model_size_by_data` since the cache doesn't know what they are).
+    cache.put(model_key, torch.nn.Linear(4, 4))
+    cache.put(f"{model_key}:unet", torch.nn.Linear(4, 4))
+
+    cache.stats = CacheStats()
+    callback = MagicMock()
+    cache.on_cache_models_cleared(callback)
+
+    dropped = cache.drop_model(model_key)
+
+    assert dropped == 2
+    assert cache.stats.cleared == 2
+    callback.assert_called_once()
+    kwargs = callback.call_args.kwargs
+    assert kwargs["models_cleared"] == 2
+    assert kwargs["bytes_requested"] == 0  # not a make-room call
+    assert kwargs["bytes_freed"] > 0
+
+
+def test_unlock_stale_eviction_updates_stats_and_fires_callbacks(cache: ModelCache):
+    """Stale-entry eviction is also a cache change observers care about."""
+    model_key = "abc123"
+    cache.put(model_key, torch.randn(4))
+    entry = cache._cached_models[model_key]
+    entry.lock()
+
+    cache.drop_model(model_key)
+
+    cache.stats = CacheStats()
+    callback = MagicMock()
+    cache.on_cache_models_cleared(callback)
+
+    cache.unlock(entry)
+
+    assert model_key not in cache._cached_models
+    assert cache.stats.cleared == 1
+    callback.assert_called_once()
+
+
+def test_drop_model_with_no_matches_does_not_fire_callbacks(cache: ModelCache):
+    """No-op invalidations should be silent — don't spam observers with empty events."""
+    cache.put("other_model", torch.randn(4))
+    callback = MagicMock()
+    cache.on_cache_models_cleared(callback)
+
+    dropped = cache.drop_model("does_not_exist")
+
+    assert dropped == 0
+    callback.assert_not_called()
+
+
+def test_drop_model_with_no_matches_is_noop(cache: ModelCache):
+    cache.put("other_model", torch.randn(4))
+
+    dropped = cache.drop_model("does_not_exist")
+
+    assert dropped == 0
+    assert "other_model" in cache._cached_models
+
+
+def test_drop_model_does_not_match_prefix_substring(cache: ModelCache):
+    """`drop_model("abc")` must not drop `abcd` — only the exact key or `abc:<submodel>`."""
+    cache.put("abc", torch.randn(4))
+    cache.put("abcd", torch.randn(4))
+    cache.put("abc:unet", torch.randn(4))
+
+    dropped = cache.drop_model("abc")
+
+    assert dropped == 2
+    assert "abc" not in cache._cached_models
+    assert "abc:unet" not in cache._cached_models
+    assert "abcd" in cache._cached_models

tests/backend/model_manager/load/model_cache/test_model_cache_timeout.py

@@ -0,0 +1,126 @@
+"""Tests for model cache keep-alive timeout functionality."""
+
+import logging
+import time
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+
+
+@pytest.fixture
+def mock_logger():
+    """Create a mock logger."""
+    logger = MagicMock()
+    # Configure the mock to return a valid log level for getEffectiveLevel()
+    logger.getEffectiveLevel.return_value = logging.INFO
+    return logger
+
+
+@pytest.fixture
+def model_cache_with_timeout(mock_logger):
+    """Create a ModelCache instance with a short timeout for testing."""
+    cache = ModelCache(
+        execution_device_working_mem_gb=1.0,
+        enable_partial_loading=False,
+        keep_ram_copy_of_weights=True,
+        execution_device="cpu",
+        storage_device="cpu",
+        logger=mock_logger,
+        keep_alive_minutes=0.01,  # 0.6 seconds for fast testing
+    )
+    yield cache
+    cache.shutdown()
+
+
+@pytest.fixture
+def model_cache_no_timeout(mock_logger):
+    """Create a ModelCache instance without timeout (default behavior)."""
+    cache = ModelCache(
+        execution_device_working_mem_gb=1.0,
+        enable_partial_loading=False,
+        keep_ram_copy_of_weights=True,
+        execution_device="cpu",
+        storage_device="cpu",
+        logger=mock_logger,
+        keep_alive_minutes=0,  # 0 means no timeout
+    )
+    yield cache
+    cache.shutdown()
+
+
+def test_timeout_clears_cache(model_cache_with_timeout):
+    """Test that the cache is cleared after the timeout expires."""
+    cache = model_cache_with_timeout
+
+    # Add a simple tensor to the cache
+    test_tensor = torch.randn(10, 10)
+    cache.put("test_model", test_tensor)
+
+    # Verify the model is in the cache
+    assert "test_model" in cache._cached_models
+
+    # Wait for the timeout to expire (0.01 minutes = 0.6 seconds + buffer)
+    time.sleep(1.5)
+
+    # Verify the cache has been cleared
+    assert len(cache._cached_models) == 0
+
+
+def test_activity_resets_timeout(model_cache_with_timeout):
+    """Test that model activity resets the timeout."""
+    cache = model_cache_with_timeout
+
+    # Add a simple tensor to the cache
+    test_tensor = torch.randn(10, 10)
+    cache.put("test_model", test_tensor)
+
+    # Wait half the timeout
+    time.sleep(0.4)
+
+    # Access the model to reset the timeout
+    cache.get("test_model")
+
+    # Wait another half timeout (model should still be in cache)
+    time.sleep(0.4)
+
+    # Verify the model is still in the cache
+    assert "test_model" in cache._cached_models
+
+
+def test_no_timeout_keeps_models(model_cache_no_timeout):
+    """Test that models are kept indefinitely when timeout is 0."""
+    cache = model_cache_no_timeout
+
+    # Add a simple tensor to the cache
+    test_tensor = torch.randn(10, 10)
+    cache.put("test_model", test_tensor)
+
+    # Verify the model is in the cache
+    assert "test_model" in cache._cached_models
+
+    # Wait longer than what would be a timeout
+    time.sleep(1.0)
+
+    # Verify the model is still in the cache
+    assert "test_model" in cache._cached_models
+
+
+def test_shutdown_cancels_timer(model_cache_with_timeout):
+    """Test that shutdown properly cancels the timeout timer."""
+    cache = model_cache_with_timeout
+
+    # Add a model to start the timer
+    test_tensor = torch.randn(10, 10)
+    cache.put("test_model", test_tensor)
+
+    # Shutdown the cache
+    cache.shutdown()
+
+    # Wait for what would be the timeout
+    time.sleep(1.0)
+
+    # The model should still be in the cache since shutdown was called
+    assert "test_model" in cache._cached_models

tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_all_custom_modules.py

@@ -0,0 +1,763 @@
+import copy
+from collections.abc import Callable
+
+import gguf
+import pytest
+import torch
+
+from invokeai.backend.flux.modules.layers import RMSNorm
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
+    AUTOCAST_MODULE_TYPE_MAPPING,
+    AUTOCAST_MODULE_TYPE_MAPPING_INVERSE,
+    unwrap_custom_layer,
+    wrap_custom_layer,
+)
+from invokeai.backend.patches.layer_patcher import LayerPatcher
+from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch
+from invokeai.backend.patches.layers.dora_layer import DoRALayer
+from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer
+from invokeai.backend.patches.layers.lokr_layer import LoKRLayer
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
+from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch, Range
+from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage
+from tests.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.test_custom_invoke_linear_8_bit_lt import (
+    build_linear_8bit_lt_layer,
+)
+from tests.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.test_custom_invoke_linear_nf4 import (
+    build_linear_nf4_layer,
+)
+from tests.backend.quantization.gguf.test_ggml_tensor import quantize_tensor
+
+
+def build_linear_layer_with_ggml_quantized_tensor(orig_layer: torch.nn.Linear | None = None):
+    if orig_layer is None:
+        orig_layer = torch.nn.Linear(32, 64)
+
+    ggml_quantized_weight = quantize_tensor(orig_layer.weight, gguf.GGMLQuantizationType.Q8_0)
+    orig_layer.weight = torch.nn.Parameter(ggml_quantized_weight)
+    ggml_quantized_bias = quantize_tensor(orig_layer.bias, gguf.GGMLQuantizationType.Q8_0)
+    orig_layer.bias = torch.nn.Parameter(ggml_quantized_bias)
+    return orig_layer
+
+
+parameterize_all_devices = pytest.mark.parametrize(
+    ("device"),
+    [
+        pytest.param("cpu"),
+        pytest.param(
+            "mps", marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available.")
+        ),
+        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.")),
+    ],
+)
+
+parameterize_cuda_and_mps = pytest.mark.parametrize(
+    ("device"),
+    [
+        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available.")),
+        pytest.param(
+            "mps", marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS is not available.")
+        ),
+    ],
+)
+
+
+LayerUnderTest = tuple[torch.nn.Module, torch.Tensor, bool]
+
+
+@pytest.fixture(
+    params=[
+        "linear",
+        "conv1d",
+        "conv2d",
+        "group_norm",
+        "embedding",
+        "flux_rms_norm",
+        "linear_with_ggml_quantized_tensor",
+        "invoke_linear_8_bit_lt",
+        "invoke_linear_nf4",
+    ]
+)
+def layer_under_test(request: pytest.FixtureRequest) -> LayerUnderTest:
+    """A fixture that returns a tuple of (layer, input, supports_cpu_inference) for the layer under test."""
+    layer_type = request.param
+    if layer_type == "linear":
+        return (torch.nn.Linear(8, 16), torch.randn(1, 8), True)
+    elif layer_type == "conv1d":
+        return (torch.nn.Conv1d(8, 16, 3), torch.randn(1, 8, 5), True)
+    elif layer_type == "conv2d":
+        return (torch.nn.Conv2d(8, 16, 3), torch.randn(1, 8, 5, 5), True)
+    elif layer_type == "group_norm":
+        return (torch.nn.GroupNorm(2, 8), torch.randn(1, 8, 5), True)
+    elif layer_type == "embedding":
+        return (torch.nn.Embedding(4, 8), torch.tensor([0, 1], dtype=torch.long), True)
+    elif layer_type == "flux_rms_norm":
+        return (RMSNorm(8), torch.randn(1, 8), True)
+    elif layer_type == "linear_with_ggml_quantized_tensor":
+        return (build_linear_layer_with_ggml_quantized_tensor(), torch.randn(1, 32), True)
+    elif layer_type == "invoke_linear_8_bit_lt":
+        return (build_linear_8bit_lt_layer(), torch.randn(1, 32), False)
+    elif layer_type == "invoke_linear_nf4":
+        return (build_linear_nf4_layer(), torch.randn(1, 64), False)
+    else:
+        raise ValueError(f"Unsupported layer_type: {layer_type}")
+
+
+def layer_to_device_via_state_dict(layer: torch.nn.Module, device: str):
+    """A helper function to move a layer to a device by roundtripping through a state dict. This most closely matches
+    how models are moved in the app. Some of the quantization types have broken semantics around calling .to() on the
+    layer directly, so this is a workaround.
+
+    We should fix this in the future.
+    Relevant article: https://pytorch.org/tutorials/recipes/recipes/swap_tensors.html
+    """
+    state_dict = layer.state_dict()
+    state_dict = {k: v.to(device) for k, v in state_dict.items()}
+    layer.load_state_dict(state_dict, assign=True)
+
+
+def wrap_single_custom_layer(layer: torch.nn.Module):
+    custom_layer_type = AUTOCAST_MODULE_TYPE_MAPPING[type(layer)]
+    return wrap_custom_layer(layer, custom_layer_type)
+
+
+def unwrap_single_custom_layer(layer: torch.nn.Module):
+    orig_layer_type = AUTOCAST_MODULE_TYPE_MAPPING_INVERSE[type(layer)]
+    return unwrap_custom_layer(layer, orig_layer_type)
+
+
+class ZeroParamPatch(BaseLayerPatch):
+    """A minimal parameter patch that exercises the aggregated sidecar patch path."""
+
+    def get_parameters(self, orig_parameters: dict[str, torch.Tensor], weight: float) -> dict[str, torch.Tensor]:
+        return {name: torch.zeros_like(param) for name, param in orig_parameters.items()}
+
+    def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None):
+        return self
+
+    def calc_size(self) -> int:
+        return 0
+
+
+def _cpu_dtype_supported(
+    layer_factory: Callable[[], torch.nn.Module],
+    input_factory: Callable[[torch.dtype], torch.Tensor],
+    dtype: torch.dtype,
+) -> bool:
+    try:
+        layer = layer_factory().to(dtype=dtype)
+        input_tensor = input_factory(dtype)
+        with torch.no_grad():
+            _ = layer(input_tensor)
+        return True
+    except (RuntimeError, TypeError, NotImplementedError):
+        return False
+
+
+def _cpu_dtype_param(
+    dtype: torch.dtype,
+    layer_factory: Callable[[], torch.nn.Module],
+    input_factory: Callable[[torch.dtype], torch.Tensor],
+):
+    supported = _cpu_dtype_supported(layer_factory, input_factory, dtype)
+    return pytest.param(
+        dtype,
+        id=str(dtype).removeprefix("torch."),
+        marks=pytest.mark.skipif(not supported, reason=f"CPU {dtype} is not supported for this op"),
+    )
+
+
+LINEAR_CPU_MIXED_DTYPE_PARAMS = [
+    _cpu_dtype_param(torch.bfloat16, lambda: torch.nn.Linear(8, 16), lambda dtype: torch.randn(2, 8, dtype=dtype)),
+    _cpu_dtype_param(torch.float16, lambda: torch.nn.Linear(8, 16), lambda dtype: torch.randn(2, 8, dtype=dtype)),
+]
+
+
+CONV2D_CPU_MIXED_DTYPE_PARAMS = [
+    _cpu_dtype_param(
+        torch.bfloat16,
+        lambda: torch.nn.Conv2d(8, 16, 3),
+        lambda dtype: torch.randn(2, 8, 5, 5, dtype=dtype),
+    ),
+    _cpu_dtype_param(
+        torch.float16,
+        lambda: torch.nn.Conv2d(8, 16, 3),
+        lambda dtype: torch.randn(2, 8, 5, 5, dtype=dtype),
+    ),
+]
+
+
+def test_isinstance(layer_under_test: LayerUnderTest):
+    """Test that isinstance() and type() behave as expected after wrapping a layer in a custom layer."""
+    orig_layer, _, _ = layer_under_test
+    orig_type = type(orig_layer)
+
+    custom_layer = wrap_single_custom_layer(orig_layer)
+
+    assert isinstance(custom_layer, orig_type)
+    assert type(custom_layer) is not orig_type
+
+
+def test_wrap_and_unwrap(layer_under_test: LayerUnderTest):
+    """Test that wrapping and unwrapping a layer behaves as expected."""
+    orig_layer, _, _ = layer_under_test
+    orig_type = type(orig_layer)
+
+    # Wrap the original layer and assert that attributes of the custom layer can be accessed.
+    custom_layer = wrap_single_custom_layer(orig_layer)
+    custom_layer.set_device_autocasting_enabled(True)
+    assert custom_layer._device_autocasting_enabled
+
+    # Unwrap the custom layer.
+    # Assert that the methods of the wrapped layer are no longer accessible.
+    unwrapped_layer = unwrap_single_custom_layer(custom_layer)
+    with pytest.raises(AttributeError):
+        _ = unwrapped_layer.set_device_autocasting_enabled(True)
+    # For now, we have chosen to allow attributes to persist. We may revisit this in the future.
+    assert unwrapped_layer._device_autocasting_enabled
+    assert type(unwrapped_layer) is orig_type
+
+
+@parameterize_all_devices
+def test_state_dict(device: str, layer_under_test: LayerUnderTest):
+    """Test that .state_dict() behaves the same on the original layer and the wrapped layer."""
+    orig_layer, _, _ = layer_under_test
+
+    # Get the original layer on the test device.
+    orig_layer.to(device)
+    orig_state_dict = orig_layer.state_dict()
+
+    # Wrap the original layer.
+    custom_layer = copy.deepcopy(orig_layer)
+    custom_layer = wrap_single_custom_layer(custom_layer)
+
+    custom_state_dict = custom_layer.state_dict()
+
+    assert set(orig_state_dict.keys()) == set(custom_state_dict.keys())
+    for k in orig_state_dict:
+        assert orig_state_dict[k].shape == custom_state_dict[k].shape
+        assert orig_state_dict[k].dtype == custom_state_dict[k].dtype
+        assert orig_state_dict[k].device == custom_state_dict[k].device
+        assert torch.allclose(orig_state_dict[k], custom_state_dict[k])
+
+
+@parameterize_all_devices
+def test_load_state_dict(device: str, layer_under_test: LayerUnderTest):
+    """Test that .load_state_dict() behaves the same on the original layer and the wrapped layer."""
+    orig_layer, _, _ = layer_under_test
+
+    orig_layer.to(device)
+
+    custom_layer = copy.deepcopy(orig_layer)
+    custom_layer = wrap_single_custom_layer(custom_layer)
+
+    # Do a state dict roundtrip.
+    orig_state_dict = orig_layer.state_dict()
+    custom_state_dict = custom_layer.state_dict()
+
+    orig_layer.load_state_dict(custom_state_dict, assign=True)
+    custom_layer.load_state_dict(orig_state_dict, assign=True)
+
+    orig_state_dict = orig_layer.state_dict()
+    custom_state_dict = custom_layer.state_dict()
+
+    # Assert that the state dicts are the same after the roundtrip.
+    assert set(orig_state_dict.keys()) == set(custom_state_dict.keys())
+    for k in orig_state_dict:
+        assert orig_state_dict[k].shape == custom_state_dict[k].shape
+        assert orig_state_dict[k].dtype == custom_state_dict[k].dtype
+        assert orig_state_dict[k].device == custom_state_dict[k].device
+        assert torch.allclose(orig_state_dict[k], custom_state_dict[k])
+
+
+@parameterize_all_devices
+def test_inference_on_device(device: str, layer_under_test: LayerUnderTest):
+    """Test that inference behaves the same on the original layer and the wrapped layer when all weights are on the
+    device.
+    """
+    orig_layer, layer_input, supports_cpu_inference = layer_under_test
+
+    if device == "cpu" and not supports_cpu_inference:
+        pytest.skip("Layer does not support CPU inference.")
+
+    layer_to_device_via_state_dict(orig_layer, device)
+
+    custom_layer = copy.deepcopy(orig_layer)
+    custom_layer = wrap_single_custom_layer(custom_layer)
+
+    # Run inference with the original layer.
+    x = layer_input.to(device)
+    orig_output = orig_layer(x)
+
+    # Run inference with the wrapped layer.
+    custom_output = custom_layer(x)
+
+    assert torch.allclose(orig_output, custom_output)
+
+
+@parameterize_cuda_and_mps
+def test_inference_autocast_from_cpu_to_device(device: str, layer_under_test: LayerUnderTest):
+    """Test that inference behaves the same on the original layer and the wrapped layer when all weights are on the
+    device.
+    """
+    orig_layer, layer_input, supports_cpu_inference = layer_under_test
+
+    if device == "cpu" and not supports_cpu_inference:
+        pytest.skip("Layer does not support CPU inference.")
+
+    # Make sure the original layer is on the device.
+    layer_to_device_via_state_dict(orig_layer, device)
+
+    x = layer_input.to(device)
+
+    # Run inference with the original layer on the device.
+    orig_output = orig_layer(x)
+
+    # Move the original layer to the CPU.
+    layer_to_device_via_state_dict(orig_layer, "cpu")
+
+    is_nf4_layer = type(orig_layer).__name__ == "InvokeLinearNF4"
+    # Inference should fail with an input on the device. Do not probe raw NF4 here: with CPU-stored weights and a
+    # single-row CUDA input, some bitsandbytes versions hit an unsafe gemv_4bit path instead of raising safely.
+    if not is_nf4_layer:
+        with pytest.raises((RuntimeError, ValueError)):
+            _ = orig_layer(x)
+
+    # Wrap the original layer.
+    custom_layer = copy.deepcopy(orig_layer)
+    custom_layer = wrap_single_custom_layer(custom_layer)
+
+    # Inference should still fail with autocasting disabled. See the raw NF4 note above.
+    custom_layer.set_device_autocasting_enabled(False)
+    if not is_nf4_layer:
+        with pytest.raises((RuntimeError, ValueError)):
+            _ = custom_layer(x)
+
+    # Run inference with the wrapped layer on the device.
+    custom_layer.set_device_autocasting_enabled(True)
+    custom_output = custom_layer(x)
+    assert custom_output.device.type == device
+
+    if is_nf4_layer:
+        assert torch.allclose(orig_output, custom_output, atol=1e-5)
+    else:
+        assert torch.allclose(orig_output, custom_output)
+
+
+PatchUnderTest = tuple[list[tuple[BaseLayerPatch, float]], torch.Tensor]
+
+
+def _has_dora_patch(patches: list[tuple[BaseLayerPatch, float]]) -> bool:
+    return any(isinstance(patch, DoRALayer) for patch, _ in patches)
+
+
+def _is_bnb_quantized_linear(layer: torch.nn.Module) -> bool:
+    return type(layer).__name__ in {"InvokeLinear8bitLt", "InvokeLinearNF4"}
+
+
+@pytest.fixture(
+    params=[
+        "single_lora",
+        "multiple_loras",
+        "concatenated_lora",
+        "flux_control_lora",
+        "single_lokr",
+        "single_dora",
+    ]
+)
+def patch_under_test(request: pytest.FixtureRequest) -> PatchUnderTest:
+    """A fixture that returns a tuple of (patches, input) for the patch under test."""
+    layer_type = request.param
+    torch.manual_seed(0)
+
+    # The assumed in/out features of the base linear layer.
+    in_features = 32
+    out_features = 64
+
+    rank = 4
+
+    if layer_type == "single_lora":
+        lora_layer = LoRALayer(
+            up=torch.randn(out_features, rank),
+            mid=None,
+            down=torch.randn(rank, in_features),
+            alpha=1.0,
+            bias=torch.randn(out_features),
+        )
+        input = torch.randn(1, in_features)
+        return ([(lora_layer, 0.7)], input)
+    elif layer_type == "multiple_loras":
+        lora_layer = LoRALayer(
+            up=torch.randn(out_features, rank),
+            mid=None,
+            down=torch.randn(rank, in_features),
+            alpha=1.0,
+            bias=torch.randn(out_features),
+        )
+        lora_layer_2 = LoRALayer(
+            up=torch.randn(out_features, rank),
+            mid=None,
+            down=torch.randn(rank, in_features),
+            alpha=1.0,
+            bias=torch.randn(out_features),
+        )
+
+        input = torch.randn(1, in_features)
+        return ([(lora_layer, 1.0), (lora_layer_2, 0.5)], input)
+    elif layer_type == "concatenated_lora":
+        sub_layer_out_features = [16, 16, 32]
+
+        # Create a MergedLayerPatch.
+        sub_layers: list[LoRALayer] = []
+        sub_layer_ranges: list[Range] = []
+        dim_0_offset = 0
+        for out_features in sub_layer_out_features:
+            down = torch.randn(rank, in_features)
+            up = torch.randn(out_features, rank)
+            bias = torch.randn(out_features)
+            sub_layers.append(LoRALayer(up=up, mid=None, down=down, alpha=1.0, bias=bias))
+            sub_layer_ranges.append(Range(dim_0_offset, dim_0_offset + out_features))
+            dim_0_offset += out_features
+        merged_layer_patch = MergedLayerPatch(sub_layers, sub_layer_ranges)
+
+        input = torch.randn(1, in_features)
+        return ([(merged_layer_patch, 0.7)], input)
+    elif layer_type == "flux_control_lora":
+        # Create a FluxControlLoRALayer.
+        patched_in_features = 40
+        lora_layer = FluxControlLoRALayer(
+            up=torch.randn(out_features, rank),
+            mid=None,
+            down=torch.randn(rank, patched_in_features),
+            alpha=1.0,
+            bias=torch.randn(out_features),
+        )
+
+        input = torch.randn(1, patched_in_features)
+        return ([(lora_layer, 0.7)], input)
+    elif layer_type == "single_lokr":
+        lokr_layer = LoKRLayer(
+            w1=torch.randn(rank, rank),
+            w1_a=None,
+            w1_b=None,
+            w2=torch.randn(out_features // rank, in_features // rank),
+            w2_a=None,
+            w2_b=None,
+            t2=None,
+            alpha=1.0,
+            bias=torch.randn(out_features),
+        )
+        input = torch.randn(1, in_features)
+        return ([(lokr_layer, 0.7)], input)
+    elif layer_type == "single_dora":
+        # Regression coverage for #8624: DoRA + partial-loading + CPU->device autocast.
+        # Scaled down so the patched weight stays well-conditioned for allclose comparisons.
+        # dora_scale has shape (1, in_features) to broadcast against direction_norm in
+        # DoRALayer.get_weight — see dora_layer.py:74-82.
+        dora_layer = DoRALayer(
+            up=torch.randn(out_features, rank) * 0.01,
+            down=torch.randn(rank, in_features) * 0.01,
+            dora_scale=torch.ones(1, in_features),
+            alpha=1.0,
+            bias=torch.randn(out_features) * 0.01,
+        )
+        input = torch.randn(1, in_features)
+        return ([(dora_layer, 0.7)], input)
+    else:
+        raise ValueError(f"Unsupported layer_type: {layer_type}")
+
+
+@parameterize_all_devices
+def test_linear_sidecar_patches(device: str, patch_under_test: PatchUnderTest):
+    patches, input = patch_under_test
+
+    torch.manual_seed(0)
+
+    # Build the base layer under test.
+    layer = torch.nn.Linear(32, 64)
+
+    # Move the layer and input to the device.
+    layer_to_device_via_state_dict(layer, device)
+    input = input.to(torch.device(device))
+
+    # Patch the LoRA layer into the linear layer.
+    layer_patched = copy.deepcopy(layer)
+    for patch, weight in patches:
+        LayerPatcher._apply_model_layer_patch(
+            module_to_patch=layer_patched,
+            module_to_patch_key="",
+            patch=patch,
+            patch_weight=weight,
+            original_weights=OriginalWeightsStorage(),
+        )
+
+    # Wrap the original layer in a custom layer and add the patch to it as a sidecar.
+    custom_layer = wrap_single_custom_layer(layer)
+    for patch, weight in patches:
+        patch.to(torch.device(device))
+        custom_layer.add_patch(patch, weight)
+
+    # Run inference with the original layer and the patched layer and assert they are equal.
+    output_patched = layer_patched(input)
+    output_custom = custom_layer(input)
+    assert torch.allclose(output_patched, output_custom, atol=1e-6)
+
+
+@parameterize_cuda_and_mps
+def test_linear_sidecar_patches_with_autocast_from_cpu_to_device(device: str, patch_under_test: PatchUnderTest):
+    """Test that the output of a linear layer with sidecar patches is the same when the layer is on the device and
+    when the layer is on the CPU and the patches are autocasted to the device.
+    """
+    patches, input = patch_under_test
+
+    # Build the base layer under test.
+    layer = torch.nn.Linear(32, 64)
+
+    # Move the layer and input to the device.
+    layer_to_device_via_state_dict(layer, device)
+    input = input.to(torch.device(device))
+
+    # Wrap the original layer in a custom layer and add the patch to it.
+    custom_layer = wrap_single_custom_layer(layer)
+    for patch, weight in patches:
+        patch.to(torch.device(device))
+        custom_layer.add_patch(patch, weight)
+
+    # Run inference with the custom layer on the device.
+    expected_output = custom_layer(input)
+
+    # Move the custom layer to the CPU.
+    layer_to_device_via_state_dict(custom_layer, "cpu")
+
+    # Move the patches to the CPU.
+    custom_layer.clear_patches()
+    for patch, weight in patches:
+        patch.to(torch.device("cpu"))
+        custom_layer.add_patch(patch, weight)
+
+    # Run inference with an input on the device, and all layer weights on the CPU. The weights should be autocasted to
+    # the device.
+    autocast_output = custom_layer(input)
+    assert autocast_output.device.type == device
+
+    # Assert that the outputs with and without autocasting are the same.
+    assert torch.allclose(expected_output, autocast_output, atol=1e-6)
+
+
+@pytest.fixture(
+    params=[
+        "linear_ggml_quantized",
+        "invoke_linear_8_bit_lt",
+        "invoke_linear_nf4",
+    ]
+)
+def quantized_linear_layer_under_test(request: pytest.FixtureRequest):
+    in_features = 32
+    out_features = 64
+    torch.manual_seed(0)
+    layer_type = request.param
+    orig_layer = torch.nn.Linear(in_features, out_features)
+    if layer_type == "linear_ggml_quantized":
+        return orig_layer, build_linear_layer_with_ggml_quantized_tensor(orig_layer)
+    elif layer_type == "invoke_linear_8_bit_lt":
+        return orig_layer, build_linear_8bit_lt_layer(orig_layer)
+    elif layer_type == "invoke_linear_nf4":
+        return orig_layer, build_linear_nf4_layer(orig_layer)
+    else:
+        raise ValueError(f"Unsupported layer_type: {layer_type}")
+
+
+@parameterize_cuda_and_mps
+def test_quantized_linear_sidecar_patches(
+    device: str,
+    quantized_linear_layer_under_test: tuple[torch.nn.Module, torch.nn.Module],
+    patch_under_test: PatchUnderTest,
+):
+    """Test that patches can be applied to quantized linear layers and that the output is the same as when the patch is
+    applied to a non-quantized linear layer.
+    """
+    patches, input = patch_under_test
+
+    linear_layer, quantized_linear_layer = quantized_linear_layer_under_test
+    expect_dora_incompatible = _is_bnb_quantized_linear(quantized_linear_layer) and _has_dora_patch(patches)
+
+    # Move everything to the device.
+    layer_to_device_via_state_dict(linear_layer, device)
+    layer_to_device_via_state_dict(quantized_linear_layer, device)
+    input = input.to(torch.device(device))
+
+    # Wrap both layers in custom layers.
+    linear_layer_custom = wrap_single_custom_layer(linear_layer)
+    quantized_linear_layer_custom = wrap_single_custom_layer(quantized_linear_layer)
+
+    # Apply the patches to the custom layers.
+    for patch, weight in patches:
+        patch.to(torch.device(device))
+        linear_layer_custom.add_patch(patch, weight)
+        quantized_linear_layer_custom.add_patch(patch, weight)
+
+    # Run inference with the original layer and the patched layer and assert they are equal.
+    output_linear_patched = linear_layer_custom(input)
+    if expect_dora_incompatible:
+        with pytest.raises(RuntimeError, match="not compatible with DoRA patches"):
+            quantized_linear_layer_custom(input)
+        return
+
+    output_quantized_patched = quantized_linear_layer_custom(input)
+    assert torch.allclose(output_linear_patched, output_quantized_patched, rtol=0.2, atol=0.2)
+
+
+@parameterize_cuda_and_mps
+def test_quantized_linear_sidecar_patches_with_autocast_from_cpu_to_device(
+    device: str,
+    quantized_linear_layer_under_test: tuple[torch.nn.Module, torch.nn.Module],
+    patch_under_test: PatchUnderTest,
+):
+    """Test that the output of a linear layer with sidecar patches is the same when the layer is on the device and
+    when the layer is on the CPU and the patches are autocasted to the device.
+    """
+    patches, input = patch_under_test
+
+    _, quantized_linear_layer = quantized_linear_layer_under_test
+    expect_dora_incompatible = _is_bnb_quantized_linear(quantized_linear_layer) and _has_dora_patch(patches)
+
+    # Move everything to the device.
+    layer_to_device_via_state_dict(quantized_linear_layer, device)
+    input = input.to(torch.device(device))
+
+    # Wrap the quantized linear layer in a custom layer and add the patch to it.
+    quantized_linear_layer_custom = wrap_single_custom_layer(quantized_linear_layer)
+    for patch, weight in patches:
+        patch.to(torch.device(device))
+        quantized_linear_layer_custom.add_patch(patch, weight)
+
+    # Run inference with the custom layer on the device.
+    if expect_dora_incompatible:
+        with pytest.raises(RuntimeError, match="not compatible with DoRA patches"):
+            quantized_linear_layer_custom(input)
+        return
+
+    expected_output = quantized_linear_layer_custom(input)
+
+    # Move the custom layer to the CPU.
+    layer_to_device_via_state_dict(quantized_linear_layer_custom, "cpu")
+
+    # Move the patches to the CPU.
+    quantized_linear_layer_custom.clear_patches()
+    for patch, weight in patches:
+        patch.to(torch.device("cpu"))
+        quantized_linear_layer_custom.add_patch(patch, weight)
+
+    # Run inference with an input on the device, and all layer weights on the CPU. The weights should be autocasted to
+    # the device.
+    autocast_output = quantized_linear_layer_custom(input)
+    assert autocast_output.device.type == device
+
+    # Assert that the outputs with and without autocasting are the same.
+    assert torch.allclose(expected_output, autocast_output, atol=1e-6)
+
+
+@pytest.mark.parametrize("dtype", LINEAR_CPU_MIXED_DTYPE_PARAMS)
+@torch.no_grad()
+def test_linear_mixed_dtype_inference_without_patches(dtype: torch.dtype):
+    layer = wrap_single_custom_layer(torch.nn.Linear(8, 16))
+    input = torch.randn(2, 8, dtype=dtype)
+
+    output = layer(input)
+
+    assert output.dtype == input.dtype
+    assert output.shape == (2, 16)
+
+
+@pytest.mark.parametrize("dtype", LINEAR_CPU_MIXED_DTYPE_PARAMS)
+@torch.no_grad()
+def test_linear_mixed_dtype_inference_without_patches_bias_only_mismatch(dtype: torch.dtype):
+    layer = torch.nn.Linear(8, 16).to(dtype=dtype)
+    layer.bias = torch.nn.Parameter(layer.bias.detach().to(torch.float32))
+    layer = wrap_single_custom_layer(layer)
+    input = torch.randn(2, 8, dtype=dtype)
+
+    output = layer(input)
+
+    assert output.dtype == input.dtype
+    assert output.shape == (2, 16)
+
+
+@pytest.mark.parametrize("dtype", CONV2D_CPU_MIXED_DTYPE_PARAMS)
+@torch.no_grad()
+def test_conv2d_mixed_dtype_inference_without_patches(dtype: torch.dtype):
+    layer = wrap_single_custom_layer(torch.nn.Conv2d(8, 16, 3))
+    input = torch.randn(2, 8, 5, 5, dtype=dtype)
+
+    output = layer(input)
+
+    assert output.dtype == input.dtype
+    assert output.shape == (2, 16, 3, 3)
+
+
+@pytest.mark.parametrize("dtype", LINEAR_CPU_MIXED_DTYPE_PARAMS)
+@torch.no_grad()
+def test_linear_mixed_dtype_sidecar_parameter_patch(dtype: torch.dtype):
+    layer = wrap_single_custom_layer(torch.nn.Linear(8, 16))
+    layer.add_patch(ZeroParamPatch(), 1.0)
+    input = torch.randn(2, 8, dtype=dtype)
+
+    output = layer(input)
+
+    assert output.dtype == input.dtype
+    assert output.shape == (2, 16)
+
+
+@pytest.mark.parametrize("dtype", CONV2D_CPU_MIXED_DTYPE_PARAMS)
+@torch.no_grad()
+def test_conv2d_mixed_dtype_sidecar_parameter_patch(dtype: torch.dtype):
+    layer = wrap_single_custom_layer(torch.nn.Conv2d(8, 16, 3))
+    layer.add_patch(ZeroParamPatch(), 1.0)
+    input = torch.randn(2, 8, 5, 5, dtype=dtype)
+
+    output = layer(input)
+
+    assert output.dtype == input.dtype
+    assert output.shape == (2, 16, 3, 3)
+
+
+@torch.no_grad()
+def test_aggregate_patch_parameters_preserves_plain_tensor_with_dora():
+    """Regression test for #8624: when partial-loading autocasts a CPU Parameter onto the
+    compute device, cast_to_device returns a plain torch.Tensor (not a Parameter). The
+    aggregator must treat that as a real tensor and not substitute a meta-device dummy —
+    otherwise DoRA's quantization guard falsely triggers on non-quantized base models.
+
+    This test is CPU-only and simulates the hand-off by constructing a plain torch.Tensor
+    directly; the equivalent CUDA/MPS E2E flow is exercised by the "single_dora" variant
+    of test_linear_sidecar_patches_with_autocast_from_cpu_to_device.
+    """
+    layer = wrap_single_custom_layer(torch.nn.Linear(32, 64))
+
+    rank = 4
+    dora_patch = DoRALayer(
+        up=torch.randn(64, rank) * 0.01,
+        down=torch.randn(rank, 32) * 0.01,
+        dora_scale=torch.ones(1, 32),
+        alpha=1.0,
+        bias=None,
+    )
+
+    # Plain torch.Tensor — the shape _cast_weight_bias_for_input hands into
+    # _aggregate_patch_parameters after autocasting a Parameter across devices.
+    plain_weight = torch.randn(64, 32)
+    assert type(plain_weight) is torch.Tensor
+
+    orig_params = {"weight": plain_weight}
+    params = layer._aggregate_patch_parameters(
+        patches_and_weights=[(dora_patch, 1.0)],
+        orig_params=orig_params,
+        device=torch.device("cpu"),
+    )
+
+    # Pre-fix, orig_params["weight"] would have been replaced by a meta-device dummy,
+    # causing DoRALayer.get_parameters to raise "not compatible with DoRA patches".
+    assert orig_params["weight"].device.type == "cpu"
+    assert params["weight"].shape == (64, 32)
+    assert params["weight"].device.type == "cpu"
+    assert not torch.isnan(params["weight"]).any()

tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_custom_flux_rms_norm.py

@@ -0,0 +1,31 @@
+import torch
+
+from invokeai.backend.flux.modules.layers import RMSNorm
+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.torch_module_autocast import (
+    wrap_custom_layer,
+)
+from invokeai.backend.patches.layers.set_parameter_layer import SetParameterLayer
+
+
+def test_custom_flux_rms_norm_patch():
+    """Test a SetParameterLayer patch on a CustomFluxRMSNorm layer."""
+    # Create a RMSNorm layer.
+    dim = 8
+    rms_norm = RMSNorm(dim)
+
+    # Create a SetParameterLayer.
+    new_scale = torch.randn(dim)
+    set_parameter_layer = SetParameterLayer("scale", new_scale)
+
+    # Wrap the RMSNorm layer in a CustomFluxRMSNorm layer.
+    custom_flux_rms_norm = wrap_custom_layer(rms_norm, CustomFluxRMSNorm)
+    custom_flux_rms_norm.add_patch(set_parameter_layer, 1.0)
+
+    # Run the CustomFluxRMSNorm layer.
+    input = torch.randn(1, dim)
+    expected_output = torch.nn.functional.rms_norm(input, new_scale.shape, new_scale, eps=1e-6)
+    output_custom = custom_flux_rms_norm(input)
+    assert torch.allclose(output_custom, expected_output, atol=1e-6)

tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_custom_invoke_linear_8_bit_lt.py

@@ -0,0 +1,82 @@
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
+    wrap_custom_layer,
+)
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA is not available", allow_module_level=True)
+else:
+    from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_8_bit_lt import (
+        CustomInvokeLinear8bitLt,
+    )
+    from invokeai.backend.quantization.bnb_llm_int8 import InvokeLinear8bitLt
+
+
+def build_linear_8bit_lt_layer(orig_layer: torch.nn.Linear | None = None):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    torch.manual_seed(1)
+
+    if orig_layer is None:
+        orig_layer = torch.nn.Linear(32, 64)
+    orig_layer_state_dict = orig_layer.state_dict()
+
+    # Prepare a quantized InvokeLinear8bitLt layer.
+    quantized_layer = InvokeLinear8bitLt(
+        input_features=orig_layer.in_features, output_features=orig_layer.out_features, has_fp16_weights=False
+    )
+    quantized_layer.load_state_dict(orig_layer_state_dict)
+    quantized_layer.to("cuda")
+
+    # Assert that the InvokeLinear8bitLt layer is quantized.
+    assert quantized_layer.weight.CB is not None
+    assert quantized_layer.weight.SCB is not None
+    assert quantized_layer.weight.CB.dtype == torch.int8
+
+    return quantized_layer
+
+
+@pytest.fixture
+def linear_8bit_lt_layer():
+    return build_linear_8bit_lt_layer()
+
+
+def test_custom_invoke_linear_8bit_lt_all_weights_on_cuda(linear_8bit_lt_layer: InvokeLinear8bitLt):
+    """Test CustomInvokeLinear8bitLt inference with all weights on the GPU."""
+    # Run inference on the original layer.
+    x = torch.randn(1, 32).to("cuda")
+    y_quantized = linear_8bit_lt_layer(x)
+
+    # Wrap the InvokeLinear8bitLt layer in a CustomInvokeLinear8bitLt layer, and run inference on it.
+    custom_linear_8bit_lt_layer = wrap_custom_layer(linear_8bit_lt_layer, CustomInvokeLinear8bitLt)
+    y_custom = custom_linear_8bit_lt_layer(x)
+
+    # Assert that the quantized and custom layers produce the same output.
+    assert torch.allclose(y_quantized, y_custom, atol=1e-5)
+
+
+def test_custom_invoke_linear_8bit_lt_all_weights_on_cpu(linear_8bit_lt_layer: InvokeLinear8bitLt):
+    """Test CustomInvokeLinear8bitLt inference with all weights on the CPU (streaming to the GPU)."""
+    # Run inference on the original layer.
+    x = torch.randn(1, 32).to("cuda")
+    y_quantized = linear_8bit_lt_layer(x)
+
+    # Copy the state dict to the CPU and reload it.
+    state_dict = linear_8bit_lt_layer.state_dict()
+    state_dict = {k: v.to("cpu") for k, v in state_dict.items()}
+    linear_8bit_lt_layer.load_state_dict(state_dict)
+
+    # Inference of the original layer should fail.
+    with pytest.raises((RuntimeError, ValueError)):
+        linear_8bit_lt_layer(x)
+
+    # Wrap the InvokeLinear8bitLt layer in a CustomInvokeLinear8bitLt layer, and run inference on it.
+    custom_linear_8bit_lt_layer = wrap_custom_layer(linear_8bit_lt_layer, CustomInvokeLinear8bitLt)
+    custom_linear_8bit_lt_layer.set_device_autocasting_enabled(True)
+    y_custom = custom_linear_8bit_lt_layer(x)
+
+    # Assert that the quantized and custom layers produce the same output.
+    assert torch.allclose(y_quantized, y_custom, atol=1e-5)

tests/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/test_custom_invoke_linear_nf4.py

@@ -0,0 +1,108 @@
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
+    wrap_custom_layer,
+)
+
+if not torch.cuda.is_available():
+    pytest.skip("CUDA is not available", allow_module_level=True)
+else:
+    from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_invoke_linear_nf4 import (
+        CustomInvokeLinearNF4,
+    )
+    from invokeai.backend.quantization.bnb_nf4 import InvokeLinearNF4
+
+
+def build_linear_nf4_layer(orig_layer: torch.nn.Linear | None = None):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is not available")
+
+    torch.manual_seed(1)
+
+    if orig_layer is None:
+        orig_layer = torch.nn.Linear(64, 16)
+
+    orig_layer_state_dict = orig_layer.state_dict()
+
+    # Prepare a quantized InvokeLinearNF4 layer.
+    quantized_layer = InvokeLinearNF4(input_features=orig_layer.in_features, output_features=orig_layer.out_features)
+    quantized_layer.load_state_dict(orig_layer_state_dict)
+    quantized_layer.to("cuda")
+
+    # Assert that the InvokeLinearNF4 layer is quantized.
+    assert quantized_layer.weight.bnb_quantized
+
+    return quantized_layer
+
+
+@pytest.fixture
+def linear_nf4_layer():
+    return build_linear_nf4_layer()
+
+
+def test_custom_invoke_linear_nf4_all_weights_on_cuda(linear_nf4_layer: InvokeLinearNF4):
+    """Test CustomInvokeLinearNF4 inference with all weights on the GPU."""
+    # Run inference on the original layer.
+    x = torch.randn(1, 64).to("cuda")
+    y_quantized = linear_nf4_layer(x)
+
+    # Wrap the InvokeLinearNF4 layer in a CustomInvokeLinearNF4 layer, and run inference on it.
+    custom_linear_nf4_layer = wrap_custom_layer(linear_nf4_layer, CustomInvokeLinearNF4)
+    custom_linear_nf4_layer.set_device_autocasting_enabled(True)
+    y_custom = custom_linear_nf4_layer(x)
+
+    # Assert that the quantized and custom layers produce the same output.
+    assert torch.allclose(y_quantized, y_custom, atol=1e-5)
+
+
+def test_custom_invoke_linear_nf4_all_weights_on_cuda_uses_bnb_single_vector_path(
+    linear_nf4_layer: InvokeLinearNF4,
+):
+    """GPU-resident single-vector inference should keep using bnb's gemv_4bit path."""
+    x = torch.randn(1, 64).to("cuda")
+    custom_linear_nf4_layer = wrap_custom_layer(linear_nf4_layer, CustomInvokeLinearNF4)
+    custom_linear_nf4_layer.set_device_autocasting_enabled(True)
+
+    with patch("bitsandbytes.functional.dequantize_4bit") as mock_dequantize:
+        _ = custom_linear_nf4_layer(x)
+
+    mock_dequantize.assert_not_called()
+
+
+# We run with two different input dimensions, because the NF4 layer follows a different code path depending on the
+# input dimension, and this has caused issues in the past.
+@pytest.mark.parametrize("input_dim_0", [1, 2])
+def test_custom_invoke_linear_nf4_all_weights_on_cpu(linear_nf4_layer: InvokeLinearNF4, input_dim_0: int):
+    """Test CustomInvokeLinearNF4 inference with all weights on the CPU (streaming to the GPU)."""
+    # Run inference on the original layer.
+    x = torch.randn(input_dim_0, 64).to(device="cuda")
+    y_quantized = linear_nf4_layer(x)
+
+    # Copy the state dict to the CPU and reload it.
+    state_dict = linear_nf4_layer.state_dict()
+    state_dict = {k: v.to("cpu") for k, v in state_dict.items()}
+    linear_nf4_layer.load_state_dict(state_dict)
+
+    # Do not call the raw bitsandbytes NF4 layer here. With CPU-stored weights and a single-row CUDA input, some
+    # bitsandbytes versions hit an unsafe gemv_4bit path instead of raising a Python exception. The custom layer below
+    # is the behavior under test.
+
+    # Wrap the InvokeLinearNF4 layer in a CustomInvokeLinearNF4 layer, and run inference on it.
+    custom_linear_nf4_layer = wrap_custom_layer(linear_nf4_layer, CustomInvokeLinearNF4)
+    custom_linear_nf4_layer.set_device_autocasting_enabled(True)
+    y_custom = custom_linear_nf4_layer(x)
+
+    # Assert that the state dict (and the tensors that it references) are still on the CPU.
+    assert all(v.device == torch.device("cpu") for v in state_dict.values())
+
+    # Assert that the weight, bias, and quant_state are all on the CPU.
+    assert custom_linear_nf4_layer.weight.device == torch.device("cpu")
+    assert custom_linear_nf4_layer.bias.device == torch.device("cpu")
+    assert custom_linear_nf4_layer.weight.quant_state.absmax.device == torch.device("cpu")
+    assert custom_linear_nf4_layer.weight.quant_state.code.device == torch.device("cpu")
+
+    # Assert that the quantized and custom layers produce the same output.
+    assert torch.allclose(y_quantized, y_custom, atol=1e-5)

tests/backend/model_manager/load/model_cache/torch_module_autocast/test_torch_module_autocast.py

@@ -0,0 +1,132 @@
+import os
+
+import gguf
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
+    apply_custom_layers_to_model,
+    remove_custom_layers_from_model,
+)
+from tests.backend.quantization.gguf.test_ggml_tensor import quantize_tensor
+
+try:
+    from invokeai.backend.quantization.bnb_llm_int8 import InvokeLinear8bitLt, quantize_model_llm_int8
+except ImportError:
+    # This is expected to fail on MacOS
+    pass
+
+cuda_and_mps = pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param(
+            torch.device("cuda"), marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA device")
+        ),
+        pytest.param(
+            torch.device("mps"),
+            marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires MPS device"),
+        ),
+    ],
+)
+
+
+class ModelWithLinearLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(32, 64)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+@pytest.fixture(params=["none", "gguf"])
+def model(request: pytest.FixtureRequest) -> torch.nn.Module:
+    if request.param == "none":
+        return ModelWithLinearLayer()
+    elif request.param == "gguf":
+        # Initialize ModelWithLinearLayer and replace the linear layer weight with a GGML quantized weight.
+        model = ModelWithLinearLayer()
+        ggml_quantized_weight = quantize_tensor(model.linear.weight, gguf.GGMLQuantizationType.Q8_0)
+        model.linear.weight = torch.nn.Parameter(ggml_quantized_weight)
+        return model
+    else:
+        raise ValueError(f"Invalid quantization type: {request.param}")
+
+
+@cuda_and_mps
+@torch.no_grad()
+def test_torch_module_autocast_linear_layer(device: torch.device, model: torch.nn.Module):
+    # Skip this test with MPS on GitHub Actions. It fails but I haven't taken the tie to figure out why. It passes
+    # locally on MacOS.
+    if os.environ.get("GITHUB_ACTIONS") == "true" and device.type == "mps":
+        pytest.skip("This test is flaky on GitHub Actions")
+
+    # Model parameters should start off on the CPU.
+    assert all(p.device.type == "cpu" for p in model.parameters())
+
+    torch.manual_seed(0)
+
+    # Run inference on the CPU.
+    x = torch.randn(1, 32, device="cpu")
+    expected = model(x)
+    assert expected.device.type == "cpu"
+
+    # Apply the custom layers to the model.
+    apply_custom_layers_to_model(model, device_autocasting_enabled=True)
+
+    # Run the model on the device.
+    autocast_result = model(x.to(device))
+
+    # The model output should be on the device.
+    assert autocast_result.device.type == device.type
+    # The model parameters should still be on the CPU.
+    assert all(p.device.type == "cpu" for p in model.parameters())
+
+    # Remove the custom layers from the model.
+    remove_custom_layers_from_model(model)
+
+    # After removing the custom layers, the model should no longer be able to run inference on the device.
+    with pytest.raises(RuntimeError):
+        _ = model(x.to(device))
+
+    # Run inference again on the CPU.
+    after_result = model(x)
+
+    assert after_result.device.type == "cpu"
+
+    # The results from all inference runs should be the same.
+    assert torch.allclose(autocast_result.to("cpu"), expected, atol=1e-5)
+    assert torch.allclose(after_result, expected, atol=1e-5)
+
+
+@torch.no_grad()
+def test_torch_module_autocast_bnb_llm_int8_linear_layer():
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA device")
+
+    torch.manual_seed(0)
+
+    model = ModelWithLinearLayer()
+    model = quantize_model_llm_int8(model, modules_to_not_convert=set())
+    # The act of moving the model to the CUDA device will trigger quantization.
+    model.to("cuda")
+    # Confirm that the layer is quantized.
+    assert isinstance(model.linear, InvokeLinear8bitLt)
+    assert model.linear.weight.CB is not None
+    assert model.linear.weight.SCB is not None
+
+    # Run inference on the GPU.
+    x = torch.randn(1, 32)
+    expected = model(x.to("cuda"))
+    assert expected.device.type == "cuda"
+
+    # Move the model back to the CPU and add the custom layers to the model.
+    model.to("cpu")
+    apply_custom_layers_to_model(model, device_autocasting_enabled=True)
+
+    # Run inference with weights being streamed to the GPU.
+    autocast_result = model(x.to("cuda"))
+    assert autocast_result.device.type == "cuda"
+
+    # The results from all inference runs should be the same.
+    assert torch.allclose(autocast_result, expected, atol=1e-5)

tests/backend/model_manager/load/state_dicts/anima_comfyui_keys.py

@@ -0,0 +1,155 @@
+"""Representative key layout of an official Anima transformer single-file checkpoint.
+
+Captured from `anima-base-v1.0.safetensors`. The full checkpoint has 685 tensors under the
+`net.` prefix; this fixture keeps `net.blocks.0.*` plus all non-block `net.*` keys. Used to
+exercise `_strip_anima_bundle_prefix` (the `net.` -> unprefixed strip). The ComfyUI-bundled
+`model.diffusion_model.*` variant is covered by a small synthetic dict in the test.
+"""
+
+state_dict_keys: dict[str, list[int]] = {
+    "net.blocks.0.adaln_modulation_cross_attn.1.weight": [256, 2048],
+    "net.blocks.0.adaln_modulation_cross_attn.2.weight": [6144, 256],
+    "net.blocks.0.adaln_modulation_mlp.1.weight": [256, 2048],
+    "net.blocks.0.adaln_modulation_mlp.2.weight": [6144, 256],
+    "net.blocks.0.adaln_modulation_self_attn.1.weight": [256, 2048],
+    "net.blocks.0.adaln_modulation_self_attn.2.weight": [6144, 256],
+    "net.blocks.0.cross_attn.k_norm.weight": [128],
+    "net.blocks.0.cross_attn.k_proj.weight": [2048, 1024],
+    "net.blocks.0.cross_attn.output_proj.weight": [2048, 2048],
+    "net.blocks.0.cross_attn.q_norm.weight": [128],
+    "net.blocks.0.cross_attn.q_proj.weight": [2048, 2048],
+    "net.blocks.0.cross_attn.v_proj.weight": [2048, 1024],
+    "net.blocks.0.mlp.layer1.weight": [8192, 2048],
+    "net.blocks.0.mlp.layer2.weight": [2048, 8192],
+    "net.blocks.0.self_attn.k_norm.weight": [128],
+    "net.blocks.0.self_attn.k_proj.weight": [2048, 2048],
+    "net.blocks.0.self_attn.output_proj.weight": [2048, 2048],
+    "net.blocks.0.self_attn.q_norm.weight": [128],
+    "net.blocks.0.self_attn.q_proj.weight": [2048, 2048],
+    "net.blocks.0.self_attn.v_proj.weight": [2048, 2048],
+    "net.final_layer.adaln_modulation.1.weight": [256, 2048],
+    "net.final_layer.adaln_modulation.2.weight": [4096, 256],
+    "net.final_layer.linear.weight": [64, 2048],
+    "net.llm_adapter.blocks.0.cross_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.0.cross_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.cross_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.cross_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.0.cross_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.cross_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.mlp.0.bias": [4096],
+    "net.llm_adapter.blocks.0.mlp.0.weight": [4096, 1024],
+    "net.llm_adapter.blocks.0.mlp.2.bias": [1024],
+    "net.llm_adapter.blocks.0.mlp.2.weight": [1024, 4096],
+    "net.llm_adapter.blocks.0.norm_cross_attn.weight": [1024],
+    "net.llm_adapter.blocks.0.norm_mlp.weight": [1024],
+    "net.llm_adapter.blocks.0.norm_self_attn.weight": [1024],
+    "net.llm_adapter.blocks.0.self_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.0.self_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.self_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.self_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.0.self_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.0.self_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.cross_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.1.cross_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.cross_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.cross_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.1.cross_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.cross_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.mlp.0.bias": [4096],
+    "net.llm_adapter.blocks.1.mlp.0.weight": [4096, 1024],
+    "net.llm_adapter.blocks.1.mlp.2.bias": [1024],
+    "net.llm_adapter.blocks.1.mlp.2.weight": [1024, 4096],
+    "net.llm_adapter.blocks.1.norm_cross_attn.weight": [1024],
+    "net.llm_adapter.blocks.1.norm_mlp.weight": [1024],
+    "net.llm_adapter.blocks.1.norm_self_attn.weight": [1024],
+    "net.llm_adapter.blocks.1.self_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.1.self_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.self_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.self_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.1.self_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.1.self_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.cross_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.2.cross_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.cross_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.cross_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.2.cross_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.cross_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.mlp.0.bias": [4096],
+    "net.llm_adapter.blocks.2.mlp.0.weight": [4096, 1024],
+    "net.llm_adapter.blocks.2.mlp.2.bias": [1024],
+    "net.llm_adapter.blocks.2.mlp.2.weight": [1024, 4096],
+    "net.llm_adapter.blocks.2.norm_cross_attn.weight": [1024],
+    "net.llm_adapter.blocks.2.norm_mlp.weight": [1024],
+    "net.llm_adapter.blocks.2.norm_self_attn.weight": [1024],
+    "net.llm_adapter.blocks.2.self_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.2.self_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.self_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.self_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.2.self_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.2.self_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.cross_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.3.cross_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.cross_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.cross_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.3.cross_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.cross_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.mlp.0.bias": [4096],
+    "net.llm_adapter.blocks.3.mlp.0.weight": [4096, 1024],
+    "net.llm_adapter.blocks.3.mlp.2.bias": [1024],
+    "net.llm_adapter.blocks.3.mlp.2.weight": [1024, 4096],
+    "net.llm_adapter.blocks.3.norm_cross_attn.weight": [1024],
+    "net.llm_adapter.blocks.3.norm_mlp.weight": [1024],
+    "net.llm_adapter.blocks.3.norm_self_attn.weight": [1024],
+    "net.llm_adapter.blocks.3.self_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.3.self_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.self_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.self_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.3.self_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.3.self_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.cross_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.4.cross_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.cross_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.cross_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.4.cross_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.cross_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.mlp.0.bias": [4096],
+    "net.llm_adapter.blocks.4.mlp.0.weight": [4096, 1024],
+    "net.llm_adapter.blocks.4.mlp.2.bias": [1024],
+    "net.llm_adapter.blocks.4.mlp.2.weight": [1024, 4096],
+    "net.llm_adapter.blocks.4.norm_cross_attn.weight": [1024],
+    "net.llm_adapter.blocks.4.norm_mlp.weight": [1024],
+    "net.llm_adapter.blocks.4.norm_self_attn.weight": [1024],
+    "net.llm_adapter.blocks.4.self_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.4.self_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.self_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.self_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.4.self_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.4.self_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.cross_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.5.cross_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.cross_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.cross_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.5.cross_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.cross_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.mlp.0.bias": [4096],
+    "net.llm_adapter.blocks.5.mlp.0.weight": [4096, 1024],
+    "net.llm_adapter.blocks.5.mlp.2.bias": [1024],
+    "net.llm_adapter.blocks.5.mlp.2.weight": [1024, 4096],
+    "net.llm_adapter.blocks.5.norm_cross_attn.weight": [1024],
+    "net.llm_adapter.blocks.5.norm_mlp.weight": [1024],
+    "net.llm_adapter.blocks.5.norm_self_attn.weight": [1024],
+    "net.llm_adapter.blocks.5.self_attn.k_norm.weight": [64],
+    "net.llm_adapter.blocks.5.self_attn.k_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.self_attn.o_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.self_attn.q_norm.weight": [64],
+    "net.llm_adapter.blocks.5.self_attn.q_proj.weight": [1024, 1024],
+    "net.llm_adapter.blocks.5.self_attn.v_proj.weight": [1024, 1024],
+    "net.llm_adapter.embed.weight": [32128, 1024],
+    "net.llm_adapter.norm.weight": [1024],
+    "net.llm_adapter.out_proj.bias": [1024],
+    "net.llm_adapter.out_proj.weight": [1024, 1024],
+    "net.t_embedder.1.linear_1.weight": [2048, 2048],
+    "net.t_embedder.1.linear_2.weight": [6144, 2048],
+    "net.t_embedding_norm.weight": [2048],
+    "net.x_embedder.proj.1.weight": [2048, 68],
+}

tests/backend/model_manager/load/state_dicts/flux2_transformer_bfl_keys.py

@@ -0,0 +1,41 @@
+"""Representative BFL-format key layout of a FLUX.2 transformer single-file checkpoint.
+
+Captured from `flux-2-klein-9b-kv.safetensors` (FLUX.2 Klein 9B, bf16). The full
+checkpoint has 201 tensors (8 double blocks, 24 single blocks); this
+fixture keeps every top-level key plus block 0 of the double/single stacks, which is enough
+to exercise `convert_flux2_bfl_to_diffusers` (fused-QKV split, block renames, adaLN
+scale/shift swap) and validate against a single-layer `Flux2Transformer2DModel`.
+
+BFL layout uses `double_blocks.*`, `single_blocks.*`, `img_in`, `txt_in`, `time_in`,
+`*_modulation.lin`, `final_layer.*`; diffusers expects `transformer_blocks.*`,
+`single_transformer_blocks.*`, `x_embedder`, `context_embedder`, `time_guidance_embed.*`,
+`proj_out`, `norm_out`.
+"""
+
+state_dict_keys: dict[str, list[int]] = {
+    "double_blocks.0.img_attn.norm.key_norm.scale": [128],
+    "double_blocks.0.img_attn.norm.query_norm.scale": [128],
+    "double_blocks.0.img_attn.proj.weight": [4096, 4096],
+    "double_blocks.0.img_attn.qkv.weight": [12288, 4096],
+    "double_blocks.0.img_mlp.0.weight": [24576, 4096],
+    "double_blocks.0.img_mlp.2.weight": [4096, 12288],
+    "double_blocks.0.txt_attn.norm.key_norm.scale": [128],
+    "double_blocks.0.txt_attn.norm.query_norm.scale": [128],
+    "double_blocks.0.txt_attn.proj.weight": [4096, 4096],
+    "double_blocks.0.txt_attn.qkv.weight": [12288, 4096],
+    "double_blocks.0.txt_mlp.0.weight": [24576, 4096],
+    "double_blocks.0.txt_mlp.2.weight": [4096, 12288],
+    "double_stream_modulation_img.lin.weight": [24576, 4096],
+    "double_stream_modulation_txt.lin.weight": [24576, 4096],
+    "final_layer.adaLN_modulation.1.weight": [8192, 4096],
+    "final_layer.linear.weight": [128, 4096],
+    "img_in.weight": [4096, 128],
+    "single_blocks.0.linear1.weight": [36864, 4096],
+    "single_blocks.0.linear2.weight": [4096, 16384],
+    "single_blocks.0.norm.key_norm.scale": [128],
+    "single_blocks.0.norm.query_norm.scale": [128],
+    "single_stream_modulation.lin.weight": [12288, 4096],
+    "time_in.in_layer.weight": [4096, 256],
+    "time_in.out_layer.weight": [4096, 4096],
+    "txt_in.weight": [4096, 12288],
+}

tests/backend/model_manager/load/state_dicts/flux2_vae_bfl_keys.py

@@ -0,0 +1,266 @@
+"""BFL-format key layout of a FLUX.2 VAE single-file checkpoint (full, 251 keys).
+
+Captured from `flux2-vae.safetensors` (standard FLUX.2 VAE, block_out_channels=(128,256,512,512)).
+The full key set is kept (the VAE is small), so `convert_flux2_vae_bfl_to_diffusers` can be
+validated for *complete* coverage against `AutoencoderKLFlux2` — every converted key must be a
+real parameter and every parameter must be covered.
+
+BFL layout uses `encoder.down.*`, `decoder.up.*` (reversed order!), `{enc,dec}.mid.block_N`,
+`{enc,dec}.mid.attn_1.*`, `norm_out`, `encoder.quant_conv`, `decoder.post_quant_conv`;
+diffusers expects `down_blocks`/`up_blocks`/`mid_block.resnets`/`mid_block.attentions`/
+`conv_norm_out` and top-level `quant_conv`/`post_quant_conv`.
+"""
+
+state_dict_keys: dict[str, list[int]] = {
+    "bn.num_batches_tracked": [],
+    "bn.running_mean": [128],
+    "bn.running_var": [128],
+    "decoder.conv_in.bias": [512],
+    "decoder.conv_in.weight": [512, 32, 3, 3],
+    "decoder.conv_norm_out.bias": [128],
+    "decoder.conv_norm_out.weight": [128],
+    "decoder.conv_out.bias": [3],
+    "decoder.conv_out.weight": [3, 128, 3, 3],
+    "decoder.mid_block.attentions.0.group_norm.bias": [512],
+    "decoder.mid_block.attentions.0.group_norm.weight": [512],
+    "decoder.mid_block.attentions.0.to_k.bias": [512],
+    "decoder.mid_block.attentions.0.to_k.weight": [512, 512],
+    "decoder.mid_block.attentions.0.to_out.0.bias": [512],
+    "decoder.mid_block.attentions.0.to_out.0.weight": [512, 512],
+    "decoder.mid_block.attentions.0.to_q.bias": [512],
+    "decoder.mid_block.attentions.0.to_q.weight": [512, 512],
+    "decoder.mid_block.attentions.0.to_v.bias": [512],
+    "decoder.mid_block.attentions.0.to_v.weight": [512, 512],
+    "decoder.mid_block.resnets.0.conv1.bias": [512],
+    "decoder.mid_block.resnets.0.conv1.weight": [512, 512, 3, 3],
+    "decoder.mid_block.resnets.0.conv2.bias": [512],
+    "decoder.mid_block.resnets.0.conv2.weight": [512, 512, 3, 3],
+    "decoder.mid_block.resnets.0.norm1.bias": [512],
+    "decoder.mid_block.resnets.0.norm1.weight": [512],
+    "decoder.mid_block.resnets.0.norm2.bias": [512],
+    "decoder.mid_block.resnets.0.norm2.weight": [512],
+    "decoder.mid_block.resnets.1.conv1.bias": [512],
+    "decoder.mid_block.resnets.1.conv1.weight": [512, 512, 3, 3],
+    "decoder.mid_block.resnets.1.conv2.bias": [512],
+    "decoder.mid_block.resnets.1.conv2.weight": [512, 512, 3, 3],
+    "decoder.mid_block.resnets.1.norm1.bias": [512],
+    "decoder.mid_block.resnets.1.norm1.weight": [512],
+    "decoder.mid_block.resnets.1.norm2.bias": [512],
+    "decoder.mid_block.resnets.1.norm2.weight": [512],
+    "decoder.up_blocks.0.resnets.0.conv1.bias": [512],
+    "decoder.up_blocks.0.resnets.0.conv1.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.0.resnets.0.conv2.bias": [512],
+    "decoder.up_blocks.0.resnets.0.conv2.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.0.resnets.0.norm1.bias": [512],
+    "decoder.up_blocks.0.resnets.0.norm1.weight": [512],
+    "decoder.up_blocks.0.resnets.0.norm2.bias": [512],
+    "decoder.up_blocks.0.resnets.0.norm2.weight": [512],
+    "decoder.up_blocks.0.resnets.1.conv1.bias": [512],
+    "decoder.up_blocks.0.resnets.1.conv1.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.0.resnets.1.conv2.bias": [512],
+    "decoder.up_blocks.0.resnets.1.conv2.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.0.resnets.1.norm1.bias": [512],
+    "decoder.up_blocks.0.resnets.1.norm1.weight": [512],
+    "decoder.up_blocks.0.resnets.1.norm2.bias": [512],
+    "decoder.up_blocks.0.resnets.1.norm2.weight": [512],
+    "decoder.up_blocks.0.resnets.2.conv1.bias": [512],
+    "decoder.up_blocks.0.resnets.2.conv1.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.0.resnets.2.conv2.bias": [512],
+    "decoder.up_blocks.0.resnets.2.conv2.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.0.resnets.2.norm1.bias": [512],
+    "decoder.up_blocks.0.resnets.2.norm1.weight": [512],
+    "decoder.up_blocks.0.resnets.2.norm2.bias": [512],
+    "decoder.up_blocks.0.resnets.2.norm2.weight": [512],
+    "decoder.up_blocks.0.upsamplers.0.conv.bias": [512],
+    "decoder.up_blocks.0.upsamplers.0.conv.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.0.conv1.bias": [512],
+    "decoder.up_blocks.1.resnets.0.conv1.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.0.conv2.bias": [512],
+    "decoder.up_blocks.1.resnets.0.conv2.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.0.norm1.bias": [512],
+    "decoder.up_blocks.1.resnets.0.norm1.weight": [512],
+    "decoder.up_blocks.1.resnets.0.norm2.bias": [512],
+    "decoder.up_blocks.1.resnets.0.norm2.weight": [512],
+    "decoder.up_blocks.1.resnets.1.conv1.bias": [512],
+    "decoder.up_blocks.1.resnets.1.conv1.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.1.conv2.bias": [512],
+    "decoder.up_blocks.1.resnets.1.conv2.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.1.norm1.bias": [512],
+    "decoder.up_blocks.1.resnets.1.norm1.weight": [512],
+    "decoder.up_blocks.1.resnets.1.norm2.bias": [512],
+    "decoder.up_blocks.1.resnets.1.norm2.weight": [512],
+    "decoder.up_blocks.1.resnets.2.conv1.bias": [512],
+    "decoder.up_blocks.1.resnets.2.conv1.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.2.conv2.bias": [512],
+    "decoder.up_blocks.1.resnets.2.conv2.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.1.resnets.2.norm1.bias": [512],
+    "decoder.up_blocks.1.resnets.2.norm1.weight": [512],
+    "decoder.up_blocks.1.resnets.2.norm2.bias": [512],
+    "decoder.up_blocks.1.resnets.2.norm2.weight": [512],
+    "decoder.up_blocks.1.upsamplers.0.conv.bias": [512],
+    "decoder.up_blocks.1.upsamplers.0.conv.weight": [512, 512, 3, 3],
+    "decoder.up_blocks.2.resnets.0.conv1.bias": [256],
+    "decoder.up_blocks.2.resnets.0.conv1.weight": [256, 512, 3, 3],
+    "decoder.up_blocks.2.resnets.0.conv2.bias": [256],
+    "decoder.up_blocks.2.resnets.0.conv2.weight": [256, 256, 3, 3],
+    "decoder.up_blocks.2.resnets.0.conv_shortcut.bias": [256],
+    "decoder.up_blocks.2.resnets.0.conv_shortcut.weight": [256, 512, 1, 1],
+    "decoder.up_blocks.2.resnets.0.norm1.bias": [512],
+    "decoder.up_blocks.2.resnets.0.norm1.weight": [512],
+    "decoder.up_blocks.2.resnets.0.norm2.bias": [256],
+    "decoder.up_blocks.2.resnets.0.norm2.weight": [256],
+    "decoder.up_blocks.2.resnets.1.conv1.bias": [256],
+    "decoder.up_blocks.2.resnets.1.conv1.weight": [256, 256, 3, 3],
+    "decoder.up_blocks.2.resnets.1.conv2.bias": [256],
+    "decoder.up_blocks.2.resnets.1.conv2.weight": [256, 256, 3, 3],
+    "decoder.up_blocks.2.resnets.1.norm1.bias": [256],
+    "decoder.up_blocks.2.resnets.1.norm1.weight": [256],
+    "decoder.up_blocks.2.resnets.1.norm2.bias": [256],
+    "decoder.up_blocks.2.resnets.1.norm2.weight": [256],
+    "decoder.up_blocks.2.resnets.2.conv1.bias": [256],
+    "decoder.up_blocks.2.resnets.2.conv1.weight": [256, 256, 3, 3],
+    "decoder.up_blocks.2.resnets.2.conv2.bias": [256],
+    "decoder.up_blocks.2.resnets.2.conv2.weight": [256, 256, 3, 3],
+    "decoder.up_blocks.2.resnets.2.norm1.bias": [256],
+    "decoder.up_blocks.2.resnets.2.norm1.weight": [256],
+    "decoder.up_blocks.2.resnets.2.norm2.bias": [256],
+    "decoder.up_blocks.2.resnets.2.norm2.weight": [256],
+    "decoder.up_blocks.2.upsamplers.0.conv.bias": [256],
+    "decoder.up_blocks.2.upsamplers.0.conv.weight": [256, 256, 3, 3],
+    "decoder.up_blocks.3.resnets.0.conv1.bias": [128],
+    "decoder.up_blocks.3.resnets.0.conv1.weight": [128, 256, 3, 3],
+    "decoder.up_blocks.3.resnets.0.conv2.bias": [128],
+    "decoder.up_blocks.3.resnets.0.conv2.weight": [128, 128, 3, 3],
+    "decoder.up_blocks.3.resnets.0.conv_shortcut.bias": [128],
+    "decoder.up_blocks.3.resnets.0.conv_shortcut.weight": [128, 256, 1, 1],
+    "decoder.up_blocks.3.resnets.0.norm1.bias": [256],
+    "decoder.up_blocks.3.resnets.0.norm1.weight": [256],
+    "decoder.up_blocks.3.resnets.0.norm2.bias": [128],
+    "decoder.up_blocks.3.resnets.0.norm2.weight": [128],
+    "decoder.up_blocks.3.resnets.1.conv1.bias": [128],
+    "decoder.up_blocks.3.resnets.1.conv1.weight": [128, 128, 3, 3],
+    "decoder.up_blocks.3.resnets.1.conv2.bias": [128],
+    "decoder.up_blocks.3.resnets.1.conv2.weight": [128, 128, 3, 3],
+    "decoder.up_blocks.3.resnets.1.norm1.bias": [128],
+    "decoder.up_blocks.3.resnets.1.norm1.weight": [128],
+    "decoder.up_blocks.3.resnets.1.norm2.bias": [128],
+    "decoder.up_blocks.3.resnets.1.norm2.weight": [128],
+    "decoder.up_blocks.3.resnets.2.conv1.bias": [128],
+    "decoder.up_blocks.3.resnets.2.conv1.weight": [128, 128, 3, 3],
+    "decoder.up_blocks.3.resnets.2.conv2.bias": [128],
+    "decoder.up_blocks.3.resnets.2.conv2.weight": [128, 128, 3, 3],
+    "decoder.up_blocks.3.resnets.2.norm1.bias": [128],
+    "decoder.up_blocks.3.resnets.2.norm1.weight": [128],
+    "decoder.up_blocks.3.resnets.2.norm2.bias": [128],
+    "decoder.up_blocks.3.resnets.2.norm2.weight": [128],
+    "encoder.conv_in.bias": [128],
+    "encoder.conv_in.weight": [128, 3, 3, 3],
+    "encoder.conv_norm_out.bias": [512],
+    "encoder.conv_norm_out.weight": [512],
+    "encoder.conv_out.bias": [64],
+    "encoder.conv_out.weight": [64, 512, 3, 3],
+    "encoder.down_blocks.0.downsamplers.0.conv.bias": [128],
+    "encoder.down_blocks.0.downsamplers.0.conv.weight": [128, 128, 3, 3],
+    "encoder.down_blocks.0.resnets.0.conv1.bias": [128],
+    "encoder.down_blocks.0.resnets.0.conv1.weight": [128, 128, 3, 3],
+    "encoder.down_blocks.0.resnets.0.conv2.bias": [128],
+    "encoder.down_blocks.0.resnets.0.conv2.weight": [128, 128, 3, 3],
+    "encoder.down_blocks.0.resnets.0.norm1.bias": [128],
+    "encoder.down_blocks.0.resnets.0.norm1.weight": [128],
+    "encoder.down_blocks.0.resnets.0.norm2.bias": [128],
+    "encoder.down_blocks.0.resnets.0.norm2.weight": [128],
+    "encoder.down_blocks.0.resnets.1.conv1.bias": [128],
+    "encoder.down_blocks.0.resnets.1.conv1.weight": [128, 128, 3, 3],
+    "encoder.down_blocks.0.resnets.1.conv2.bias": [128],
+    "encoder.down_blocks.0.resnets.1.conv2.weight": [128, 128, 3, 3],
+    "encoder.down_blocks.0.resnets.1.norm1.bias": [128],
+    "encoder.down_blocks.0.resnets.1.norm1.weight": [128],
+    "encoder.down_blocks.0.resnets.1.norm2.bias": [128],
+    "encoder.down_blocks.0.resnets.1.norm2.weight": [128],
+    "encoder.down_blocks.1.downsamplers.0.conv.bias": [256],
+    "encoder.down_blocks.1.downsamplers.0.conv.weight": [256, 256, 3, 3],
+    "encoder.down_blocks.1.resnets.0.conv1.bias": [256],
+    "encoder.down_blocks.1.resnets.0.conv1.weight": [256, 128, 3, 3],
+    "encoder.down_blocks.1.resnets.0.conv2.bias": [256],
+    "encoder.down_blocks.1.resnets.0.conv2.weight": [256, 256, 3, 3],
+    "encoder.down_blocks.1.resnets.0.conv_shortcut.bias": [256],
+    "encoder.down_blocks.1.resnets.0.conv_shortcut.weight": [256, 128, 1, 1],
+    "encoder.down_blocks.1.resnets.0.norm1.bias": [128],
+    "encoder.down_blocks.1.resnets.0.norm1.weight": [128],
+    "encoder.down_blocks.1.resnets.0.norm2.bias": [256],
+    "encoder.down_blocks.1.resnets.0.norm2.weight": [256],
+    "encoder.down_blocks.1.resnets.1.conv1.bias": [256],
+    "encoder.down_blocks.1.resnets.1.conv1.weight": [256, 256, 3, 3],
+    "encoder.down_blocks.1.resnets.1.conv2.bias": [256],
+    "encoder.down_blocks.1.resnets.1.conv2.weight": [256, 256, 3, 3],
+    "encoder.down_blocks.1.resnets.1.norm1.bias": [256],
+    "encoder.down_blocks.1.resnets.1.norm1.weight": [256],
+    "encoder.down_blocks.1.resnets.1.norm2.bias": [256],
+    "encoder.down_blocks.1.resnets.1.norm2.weight": [256],
+    "encoder.down_blocks.2.downsamplers.0.conv.bias": [512],
+    "encoder.down_blocks.2.downsamplers.0.conv.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.2.resnets.0.conv1.bias": [512],
+    "encoder.down_blocks.2.resnets.0.conv1.weight": [512, 256, 3, 3],
+    "encoder.down_blocks.2.resnets.0.conv2.bias": [512],
+    "encoder.down_blocks.2.resnets.0.conv2.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.2.resnets.0.conv_shortcut.bias": [512],
+    "encoder.down_blocks.2.resnets.0.conv_shortcut.weight": [512, 256, 1, 1],
+    "encoder.down_blocks.2.resnets.0.norm1.bias": [256],
+    "encoder.down_blocks.2.resnets.0.norm1.weight": [256],
+    "encoder.down_blocks.2.resnets.0.norm2.bias": [512],
+    "encoder.down_blocks.2.resnets.0.norm2.weight": [512],
+    "encoder.down_blocks.2.resnets.1.conv1.bias": [512],
+    "encoder.down_blocks.2.resnets.1.conv1.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.2.resnets.1.conv2.bias": [512],
+    "encoder.down_blocks.2.resnets.1.conv2.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.2.resnets.1.norm1.bias": [512],
+    "encoder.down_blocks.2.resnets.1.norm1.weight": [512],
+    "encoder.down_blocks.2.resnets.1.norm2.bias": [512],
+    "encoder.down_blocks.2.resnets.1.norm2.weight": [512],
+    "encoder.down_blocks.3.resnets.0.conv1.bias": [512],
+    "encoder.down_blocks.3.resnets.0.conv1.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.3.resnets.0.conv2.bias": [512],
+    "encoder.down_blocks.3.resnets.0.conv2.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.3.resnets.0.norm1.bias": [512],
+    "encoder.down_blocks.3.resnets.0.norm1.weight": [512],
+    "encoder.down_blocks.3.resnets.0.norm2.bias": [512],
+    "encoder.down_blocks.3.resnets.0.norm2.weight": [512],
+    "encoder.down_blocks.3.resnets.1.conv1.bias": [512],
+    "encoder.down_blocks.3.resnets.1.conv1.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.3.resnets.1.conv2.bias": [512],
+    "encoder.down_blocks.3.resnets.1.conv2.weight": [512, 512, 3, 3],
+    "encoder.down_blocks.3.resnets.1.norm1.bias": [512],
+    "encoder.down_blocks.3.resnets.1.norm1.weight": [512],
+    "encoder.down_blocks.3.resnets.1.norm2.bias": [512],
+    "encoder.down_blocks.3.resnets.1.norm2.weight": [512],
+    "encoder.mid_block.attentions.0.group_norm.bias": [512],
+    "encoder.mid_block.attentions.0.group_norm.weight": [512],
+    "encoder.mid_block.attentions.0.to_k.bias": [512],
+    "encoder.mid_block.attentions.0.to_k.weight": [512, 512],
+    "encoder.mid_block.attentions.0.to_out.0.bias": [512],
+    "encoder.mid_block.attentions.0.to_out.0.weight": [512, 512],
+    "encoder.mid_block.attentions.0.to_q.bias": [512],
+    "encoder.mid_block.attentions.0.to_q.weight": [512, 512],
+    "encoder.mid_block.attentions.0.to_v.bias": [512],
+    "encoder.mid_block.attentions.0.to_v.weight": [512, 512],
+    "encoder.mid_block.resnets.0.conv1.bias": [512],
+    "encoder.mid_block.resnets.0.conv1.weight": [512, 512, 3, 3],
+    "encoder.mid_block.resnets.0.conv2.bias": [512],
+    "encoder.mid_block.resnets.0.conv2.weight": [512, 512, 3, 3],
+    "encoder.mid_block.resnets.0.norm1.bias": [512],
+    "encoder.mid_block.resnets.0.norm1.weight": [512],
+    "encoder.mid_block.resnets.0.norm2.bias": [512],
+    "encoder.mid_block.resnets.0.norm2.weight": [512],
+    "encoder.mid_block.resnets.1.conv1.bias": [512],
+    "encoder.mid_block.resnets.1.conv1.weight": [512, 512, 3, 3],
+    "encoder.mid_block.resnets.1.conv2.bias": [512],
+    "encoder.mid_block.resnets.1.conv2.weight": [512, 512, 3, 3],
+    "encoder.mid_block.resnets.1.norm1.bias": [512],
+    "encoder.mid_block.resnets.1.norm1.weight": [512],
+    "encoder.mid_block.resnets.1.norm2.bias": [512],
+    "encoder.mid_block.resnets.1.norm2.weight": [512],
+    "post_quant_conv.bias": [32],
+    "post_quant_conv.weight": [32, 32, 1, 1],
+    "quant_conv.bias": [64],
+    "quant_conv.weight": [64, 64, 1, 1],
+}

tests/backend/model_manager/load/state_dicts/qwen_vl_encoder_comfyui_keys.py

@@ -0,0 +1,77 @@
+"""Representative key layout of a ComfyUI single-file Qwen2.5-VL encoder checkpoint.
+
+Captured from `qwen_2.5_vl_7b_fp8_scaled.safetensors` (Qwen2.5-VL-7B, ComfyUI fp8_scaled).
+The full checkpoint has 1446 tensors (32 visual blocks, 28 language layers); this
+fixture keeps every top-level/structural key plus block 0 of each repeated stack, which is
+enough to exercise the `visual.* -> model.visual.*` / `model.* -> model.language_model.*`
+remap and the fp8 metadata stripping without shipping a ~1446-key dict.
+
+Legacy ComfyUI layout uses `visual.*`, `model.*`, `lm_head.*` (transformers >=4.50 expects
+`model.visual.*` and `model.language_model.*`). `scale_weight` / `scale_input` / `scaled_fp8`
+are ComfyUI fp8 quantization metadata.
+"""
+
+state_dict_keys: dict[str, list[int]] = {
+    "visual.blocks.0.attn.proj.bias": [1280],
+    "visual.blocks.0.attn.proj.scale_input": [],
+    "visual.blocks.0.attn.proj.scale_weight": [],
+    "visual.blocks.0.attn.proj.weight": [1280, 1280],
+    "visual.blocks.0.attn.qkv.bias": [3840],
+    "visual.blocks.0.attn.qkv.scale_input": [],
+    "visual.blocks.0.attn.qkv.scale_weight": [],
+    "visual.blocks.0.attn.qkv.weight": [3840, 1280],
+    "visual.blocks.0.mlp.down_proj.bias": [1280],
+    "visual.blocks.0.mlp.down_proj.scale_input": [],
+    "visual.blocks.0.mlp.down_proj.scale_weight": [],
+    "visual.blocks.0.mlp.down_proj.weight": [1280, 3420],
+    "visual.blocks.0.mlp.gate_proj.bias": [3420],
+    "visual.blocks.0.mlp.gate_proj.scale_input": [],
+    "visual.blocks.0.mlp.gate_proj.scale_weight": [],
+    "visual.blocks.0.mlp.gate_proj.weight": [3420, 1280],
+    "visual.blocks.0.mlp.up_proj.bias": [3420],
+    "visual.blocks.0.mlp.up_proj.scale_input": [],
+    "visual.blocks.0.mlp.up_proj.scale_weight": [],
+    "visual.blocks.0.mlp.up_proj.weight": [3420, 1280],
+    "visual.blocks.0.norm1.weight": [1280],
+    "visual.blocks.0.norm2.weight": [1280],
+    "visual.merger.ln_q.weight": [1280],
+    "visual.merger.mlp.0.bias": [5120],
+    "visual.merger.mlp.0.scale_input": [],
+    "visual.merger.mlp.0.scale_weight": [],
+    "visual.merger.mlp.0.weight": [5120, 5120],
+    "visual.merger.mlp.2.bias": [3584],
+    "visual.merger.mlp.2.scale_input": [],
+    "visual.merger.mlp.2.scale_weight": [],
+    "visual.merger.mlp.2.weight": [3584, 5120],
+    "visual.patch_embed.proj.weight": [1280, 3, 2, 14, 14],
+    "model.embed_tokens.weight": [152064, 3584],
+    "model.layers.0.input_layernorm.weight": [3584],
+    "model.layers.0.mlp.down_proj.scale_input": [],
+    "model.layers.0.mlp.down_proj.scale_weight": [],
+    "model.layers.0.mlp.down_proj.weight": [3584, 18944],
+    "model.layers.0.mlp.gate_proj.scale_input": [],
+    "model.layers.0.mlp.gate_proj.scale_weight": [],
+    "model.layers.0.mlp.gate_proj.weight": [18944, 3584],
+    "model.layers.0.mlp.up_proj.scale_input": [],
+    "model.layers.0.mlp.up_proj.scale_weight": [],
+    "model.layers.0.mlp.up_proj.weight": [18944, 3584],
+    "model.layers.0.post_attention_layernorm.weight": [3584],
+    "model.layers.0.self_attn.k_proj.bias": [512],
+    "model.layers.0.self_attn.k_proj.scale_input": [],
+    "model.layers.0.self_attn.k_proj.scale_weight": [],
+    "model.layers.0.self_attn.k_proj.weight": [512, 3584],
+    "model.layers.0.self_attn.o_proj.scale_input": [],
+    "model.layers.0.self_attn.o_proj.scale_weight": [],
+    "model.layers.0.self_attn.o_proj.weight": [3584, 3584],
+    "model.layers.0.self_attn.q_proj.bias": [3584],
+    "model.layers.0.self_attn.q_proj.scale_input": [],
+    "model.layers.0.self_attn.q_proj.scale_weight": [],
+    "model.layers.0.self_attn.q_proj.weight": [3584, 3584],
+    "model.layers.0.self_attn.v_proj.bias": [512],
+    "model.layers.0.self_attn.v_proj.scale_input": [],
+    "model.layers.0.self_attn.v_proj.scale_weight": [],
+    "model.layers.0.self_attn.v_proj.weight": [512, 3584],
+    "model.norm.weight": [3584],
+    "lm_head.weight": [152064, 3584],
+    "scaled_fp8": [0],
+}

tests/backend/model_manager/load/state_dicts/utils.py

@@ -0,0 +1,13 @@
+"""Shared helpers for model-loader state-dict fixtures.
+
+Mirrors `tests/backend/patches/lora_conversions/lora_state_dicts/utils.py`: a fixture module
+exports `state_dict_keys: dict[str, list[int]]` (key name -> shape, captured from a real
+checkpoint) and tests expand it to a mock state dict with `keys_to_mock_state_dict()`.
+"""
+
+import torch
+
+
+def keys_to_mock_state_dict(keys: dict[str, list[int]]) -> dict[str, torch.Tensor]:
+    """Build a state dict of empty tensors from a {key: shape} mapping."""
+    return {k: torch.empty(shape) for k, shape in keys.items()}

tests/backend/model_manager/load/state_dicts/z_image_transformer_comfyui_keys.py

@@ -0,0 +1,69 @@
+"""Representative ComfyUI key layout of a Z-Image transformer single-file checkpoint.
+
+Captured from `zimageTurboBadmilk_v10.safetensors` (Z-Image Turbo) after stripping the
+`model.diffusion_model.` prefix -- this is exactly the input `_convert_z_image_gguf_to_diffusers`
+receives (the converter runs on both the checkpoint and GGUF paths after prefix stripping).
+The full transformer has 453 keys (context_refiner / noise_refiner / layers stacks); this
+fixture keeps block 0 of each stack plus all non-block keys, and adds a synthetic
+`norm_final.weight` to exercise the skip branch.
+
+Legacy layout uses fused `*.attention.qkv.*`, `*.attention.out.*`, `*.attention.q_norm/k_norm`,
+`x_embedder.*`, `final_layer.*`, `x_pad_token`, `cap_pad_token`; diffusers expects split
+`to_q/to_k/to_v`, `to_out.0`, `norm_q/norm_k`, `all_x_embedder.2-1.*`, `all_final_layer.2-1.*`.
+"""
+
+state_dict_keys: dict[str, list[int]] = {
+    "cap_embedder.0.weight": [2560],
+    "cap_embedder.1.bias": [3840],
+    "cap_embedder.1.weight": [3840, 2560],
+    "cap_pad_token": [1, 3840],
+    "context_refiner.0.attention.k_norm.weight": [128],
+    "context_refiner.0.attention.out.weight": [3840, 3840],
+    "context_refiner.0.attention.q_norm.weight": [128],
+    "context_refiner.0.attention.qkv.weight": [11520, 3840],
+    "context_refiner.0.attention_norm1.weight": [3840],
+    "context_refiner.0.attention_norm2.weight": [3840],
+    "context_refiner.0.feed_forward.w1.weight": [10240, 3840],
+    "context_refiner.0.feed_forward.w2.weight": [3840, 10240],
+    "context_refiner.0.feed_forward.w3.weight": [10240, 3840],
+    "context_refiner.0.ffn_norm1.weight": [3840],
+    "context_refiner.0.ffn_norm2.weight": [3840],
+    "final_layer.adaLN_modulation.1.bias": [3840],
+    "final_layer.adaLN_modulation.1.weight": [3840, 256],
+    "final_layer.linear.bias": [64],
+    "final_layer.linear.weight": [64, 3840],
+    "layers.0.adaLN_modulation.0.bias": [15360],
+    "layers.0.adaLN_modulation.0.weight": [15360, 256],
+    "layers.0.attention.k_norm.weight": [128],
+    "layers.0.attention.out.weight": [3840, 3840],
+    "layers.0.attention.q_norm.weight": [128],
+    "layers.0.attention.qkv.weight": [11520, 3840],
+    "layers.0.attention_norm1.weight": [3840],
+    "layers.0.attention_norm2.weight": [3840],
+    "layers.0.feed_forward.w1.weight": [10240, 3840],
+    "layers.0.feed_forward.w2.weight": [3840, 10240],
+    "layers.0.feed_forward.w3.weight": [10240, 3840],
+    "layers.0.ffn_norm1.weight": [3840],
+    "layers.0.ffn_norm2.weight": [3840],
+    "noise_refiner.0.adaLN_modulation.0.bias": [15360],
+    "noise_refiner.0.adaLN_modulation.0.weight": [15360, 256],
+    "noise_refiner.0.attention.k_norm.weight": [128],
+    "noise_refiner.0.attention.out.weight": [3840, 3840],
+    "noise_refiner.0.attention.q_norm.weight": [128],
+    "noise_refiner.0.attention.qkv.weight": [11520, 3840],
+    "noise_refiner.0.attention_norm1.weight": [3840],
+    "noise_refiner.0.attention_norm2.weight": [3840],
+    "noise_refiner.0.feed_forward.w1.weight": [10240, 3840],
+    "noise_refiner.0.feed_forward.w2.weight": [3840, 10240],
+    "noise_refiner.0.feed_forward.w3.weight": [10240, 3840],
+    "noise_refiner.0.ffn_norm1.weight": [3840],
+    "noise_refiner.0.ffn_norm2.weight": [3840],
+    "t_embedder.mlp.0.bias": [1024],
+    "t_embedder.mlp.0.weight": [1024, 256],
+    "t_embedder.mlp.2.bias": [256],
+    "t_embedder.mlp.2.weight": [256, 1024],
+    "x_embedder.bias": [3840],
+    "x_embedder.weight": [3840, 64],
+    "x_pad_token": [1, 3840],
+    "norm_final.weight": [2304],
+}

tests/backend/model_manager/load/test_anima_state_dict_utils.py

@@ -0,0 +1,38 @@
+"""Unit tests for the Anima single-file prefix-stripping helper."""
+
+import torch
+
+from invokeai.backend.model_manager.load.model_loaders.anima import _strip_anima_bundle_prefix
+from tests.backend.model_manager.load.state_dicts.anima_comfyui_keys import state_dict_keys as anima_keys
+from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict
+
+
+class TestStripAnimaBundlePrefix:
+    def test_official_net_prefix_is_stripped(self):
+        sd = keys_to_mock_state_dict(anima_keys)
+        assert all(k.startswith("net.") for k in sd)
+
+        out = _strip_anima_bundle_prefix(sd)
+
+        assert len(out) == len(sd)
+        assert not any(k.startswith("net.") for k in out)
+        # Every key had exactly its `net.` prefix removed.
+        assert {"net." + k for k in out} == set(sd.keys())
+
+    def test_comfyui_bundle_keeps_only_transformer_keys(self):
+        # ComfyUI bundles the transformer under `model.diffusion_model.` alongside the VAE and
+        # text encoder, which must be dropped.
+        sd = {
+            "model.diffusion_model.blocks.0.attn.qkv.weight": torch.empty(1),
+            "model.diffusion_model.final_layer.weight": torch.empty(1),
+            "first_stage_model.encoder.conv_in.weight": torch.empty(1),
+            "cond_stage_model.transformer.embeddings.weight": torch.empty(1),
+        }
+
+        out = _strip_anima_bundle_prefix(sd)
+
+        assert set(out.keys()) == {"blocks.0.attn.qkv.weight", "final_layer.weight"}
+
+    def test_no_known_prefix_is_a_noop(self):
+        sd = {"blocks.0.attn.qkv.weight": torch.empty(1)}
+        assert _strip_anima_bundle_prefix(sd) is sd

tests/backend/model_manager/load/test_flux2_state_dict_utils.py

@@ -0,0 +1,101 @@
+"""Unit tests for the FLUX.2 BFL->diffusers state-dict converters.
+
+Fixtures are captured from real single-file checkpoints (see the fixture module docstrings).
+The meta-device tests instantiate the actual diffusers architectures with `init_empty_weights`
+(no real weights, no GPU) and assert that every converted key is a real parameter -- the same
+kind of check that would have caught the Qwen VL remap regression.
+"""
+
+import accelerate
+import torch
+
+from invokeai.backend.model_manager.load.model_loaders.flux2_state_dict_utils import (
+    _flux2_swap_scale_shift,
+    convert_flux2_bfl_to_diffusers,
+    convert_flux2_vae_bfl_to_diffusers,
+)
+from tests.backend.model_manager.load.state_dicts.flux2_transformer_bfl_keys import (
+    state_dict_keys as flux2_transformer_keys,
+)
+from tests.backend.model_manager.load.state_dicts.flux2_vae_bfl_keys import (
+    state_dict_keys as flux2_vae_keys,
+)
+from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict
+
+
+class TestConvertFlux2Transformer:
+    def test_fused_qkv_is_split_and_blocks_renamed(self):
+        sd = keys_to_mock_state_dict(flux2_transformer_keys)
+
+        converted = convert_flux2_bfl_to_diffusers(sd)
+
+        # Fused img/txt QKV are split into separate projections.
+        assert "transformer_blocks.0.attn.to_q.weight" in converted
+        assert "transformer_blocks.0.attn.to_k.weight" in converted
+        assert "transformer_blocks.0.attn.to_v.weight" in converted
+        assert "transformer_blocks.0.attn.add_q_proj.weight" in converted
+        # No fused/BFL-named keys remain.
+        assert not any("img_attn.qkv" in k or "double_blocks." in k or "single_blocks." in k for k in converted)
+        # Top-level renames.
+        assert "x_embedder.weight" in converted
+        assert "context_embedder.weight" in converted
+        assert "proj_out.weight" in converted
+
+    def test_converted_keys_are_all_real_transformer_params(self):
+        """Meta-device coverage: every converted key must exist in Flux2Transformer2DModel."""
+        from diffusers import Flux2Transformer2DModel
+
+        converted = convert_flux2_bfl_to_diffusers(keys_to_mock_state_dict(flux2_transformer_keys))
+
+        # The fixture keeps block 0 of each stack -> a single-layer model covers it.
+        with accelerate.init_empty_weights():
+            model = Flux2Transformer2DModel(num_layers=1, num_single_layers=1)
+        params = set(model.state_dict().keys())
+
+        unmatched = sorted(k for k in converted if k not in params)
+        assert not unmatched, f"converted keys with no matching model parameter: {unmatched}"
+
+
+class TestConvertFlux2Vae:
+    def test_full_bijective_coverage_against_arch(self):
+        """The full VAE fixture must convert to exactly the AutoencoderKLFlux2 parameter set."""
+        from diffusers import AutoencoderKLFlux2
+
+        converted = convert_flux2_vae_bfl_to_diffusers(keys_to_mock_state_dict(flux2_vae_keys))
+
+        with accelerate.init_empty_weights():
+            vae = AutoencoderKLFlux2(block_out_channels=(128, 256, 512, 512))
+        params = set(vae.state_dict().keys())
+
+        unmatched = sorted(k for k in converted if k not in params)
+        missing = sorted(k for k in params if k not in converted)
+        assert not unmatched, f"converted keys with no matching VAE parameter: {unmatched}"
+        assert not missing, f"VAE parameters not covered by the converted checkpoint: {missing}"
+
+    def test_up_block_order_is_reversed(self):
+        # BFL decoder.up.X maps to diffusers up_blocks.(3 - X).
+        sd = {
+            "decoder.up.0.block.0.norm1.weight": torch.empty(1),
+            "decoder.up.3.block.0.norm1.weight": torch.empty(1),
+        }
+        converted = convert_flux2_vae_bfl_to_diffusers(sd)
+        assert "decoder.up_blocks.3.resnets.0.norm1.weight" in converted
+        assert "decoder.up_blocks.0.resnets.0.norm1.weight" in converted
+
+    def test_mid_attention_conv_weights_are_squeezed_to_linear(self):
+        # BFL stores mid attention as Conv2d [out, in, 1, 1]; diffusers uses Linear [out, in].
+        sd = {"encoder.mid.attn_1.q.weight": torch.empty(8, 8, 1, 1)}
+        converted = convert_flux2_vae_bfl_to_diffusers(sd)
+        assert converted["encoder.mid_block.attentions.0.to_q.weight"].shape == (8, 8)
+
+
+class TestSwapScaleShift:
+    def test_swaps_the_two_halves(self):
+        # First half = shift, second half = scale; diffusers wants them swapped.
+        weight = torch.cat([torch.zeros(2), torch.ones(2)])  # [shift=0, scale=1]
+        swapped = _flux2_swap_scale_shift(weight)
+        assert torch.allclose(swapped, torch.cat([torch.ones(2), torch.zeros(2)]))
+
+    def test_leaves_malformed_tensor_untouched(self):
+        weight = torch.ones(3)  # odd length -> cannot be split
+        assert torch.allclose(_flux2_swap_scale_shift(weight), weight)

tests/backend/model_manager/load/test_load_default_fp8.py

@@ -0,0 +1,371 @@
+"""Tests for `ModelLoader` FP8 helpers.
+
+Covers:
+- `_should_use_fp8` excludes ControlLoRA (the LoRA loader never runs the layerwise
+  casting helper, and a LoRA isn't a standalone forward module — so a persisted
+  `fp8_storage=true` must be a no-op).
+- `_wrap_forward_with_fp8_cast` uses pre/post hooks with `always_call=True`, so it is
+  exception-safe AND survives `apply_custom_layers_to_model`'s instance swap. Without
+  hooks, an instance-level `forward` override would be carried into the new CustomLinear
+  via the shared `__dict__` and silently bypass `CustomLinear.forward` — breaking LoRA
+  patch dispatch for FP8 checkpoint models.
+- `_apply_fp8_to_nn_module` skips precision-sensitive layers (norm, pos_embed, etc.)
+  so FLUX RMSNorm.scale and friends aren't crushed to FP8.
+"""
+
+from logging import getLogger
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.load_default import ModelLoader
+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.torch_module_autocast import (
+    apply_custom_layers_to_model,
+)
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
+
+
+def _make_loader(device: str = "cuda") -> ModelLoader:
+    """Build a ModelLoader without going through dependency injection.
+
+    `_should_use_fp8` and `_wrap_forward_with_fp8_cast` only depend on `_torch_device`,
+    so we instantiate via __new__ and set the minimum state directly.
+    """
+    loader = ModelLoader.__new__(ModelLoader)
+    loader._torch_device = torch.device(device)
+    loader._torch_dtype = torch.float16
+    loader._logger = getLogger("test")
+    return loader
+
+
+def _make_config(model_type: ModelType, fp8: bool, base: BaseModelType = BaseModelType.Flux):
+    return SimpleNamespace(
+        type=model_type,
+        base=base,
+        name="test",
+        default_settings=SimpleNamespace(fp8_storage=fp8),
+    )
+
+
+def test_should_use_fp8_excludes_control_lora():
+    """ControlLoRA gets the FP8 toggle in the UI history but the LoRA loader never applies
+    layerwise casting (the model isn't run as a standalone forward pass — it patches into a
+    base model). The loader must silently ignore a persisted `fp8_storage=true` to avoid
+    misleading users who toggled it under a prior version.
+    """
+    loader = _make_loader(device="cuda")
+    with patch("torch.cuda.is_available", return_value=True):
+        assert loader._should_use_fp8(_make_config(ModelType.ControlLoRa, fp8=True)) is False
+
+
+def test_should_use_fp8_excludes_lora():
+    loader = _make_loader(device="cuda")
+    assert loader._should_use_fp8(_make_config(ModelType.LoRA, fp8=True)) is False
+
+
+def test_should_use_fp8_returns_true_for_main_with_fp8():
+    loader = _make_loader(device="cuda")
+    assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=True)) is True
+
+
+def test_should_use_fp8_returns_false_for_main_without_fp8():
+    loader = _make_loader(device="cuda")
+    assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=False)) is False
+
+
+def test_should_use_fp8_returns_false_on_cpu():
+    loader = _make_loader(device="cpu")
+    assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=True)) is False
+
+
+class _RaisingModule(torch.nn.Module):
+    """A module whose forward unconditionally raises — used to test that the FP8 wrapper's
+    storage-dtype cleanup runs even when forward fails."""
+
+    def __init__(self):
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.zeros(4))
+        self.bias = torch.nn.Parameter(torch.zeros(4))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        raise RuntimeError("boom")
+
+
+def _fp8_supported() -> bool:
+    return hasattr(torch, "float8_e4m3fn")
+
+
+@pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available")
+def test_wrap_forward_restores_storage_dtype_on_exception():
+    """When forward raises, params must be returned to storage dtype. Otherwise FP8 storage
+    savings silently revert to fp16/bf16 and the cache's size accounting becomes stale.
+    """
+    storage_dtype = torch.float8_e4m3fn
+    compute_dtype = torch.bfloat16
+
+    module = _RaisingModule()
+    for p in module.parameters(recurse=False):
+        p.data = p.data.to(storage_dtype)
+
+    ModelLoader._wrap_forward_with_fp8_cast(module, storage_dtype, compute_dtype)
+
+    # Sanity: params start in storage dtype.
+    assert module.weight.dtype == storage_dtype
+    assert module.bias.dtype == storage_dtype
+
+    with pytest.raises(RuntimeError, match="boom"):
+        module(torch.zeros(4, dtype=compute_dtype))
+
+    # Critical assertion: cleanup ran despite the exception.
+    assert module.weight.dtype == storage_dtype
+    assert module.bias.dtype == storage_dtype
+
+
+@pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available")
+def test_wrap_forward_casts_to_compute_then_back_on_success():
+    """Happy-path sanity check: params are in compute dtype during forward, storage dtype after."""
+    storage_dtype = torch.float8_e4m3fn
+    compute_dtype = torch.bfloat16
+
+    seen_dtypes: list[torch.dtype] = []
+
+    class _CaptureModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.zeros(4))
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            seen_dtypes.append(self.weight.dtype)
+            return x + self.weight
+
+    module = _CaptureModule()
+    for p in module.parameters(recurse=False):
+        p.data = p.data.to(storage_dtype)
+
+    ModelLoader._wrap_forward_with_fp8_cast(module, storage_dtype, compute_dtype)
+
+    module(torch.zeros(4, dtype=compute_dtype))
+
+    assert seen_dtypes == [compute_dtype]
+    assert module.weight.dtype == storage_dtype
+
+
+def test_apply_fp8_to_nn_module_uses_wrapper():
+    """`_apply_fp8_to_nn_module` should delegate per-module wrapping to
+    `_wrap_forward_with_fp8_cast`, which encapsulates the hook registration.
+    """
+    module = torch.nn.Linear(4, 4)
+    with patch.object(ModelLoader, "_wrap_forward_with_fp8_cast") as mock_wrap:
+        ModelLoader._apply_fp8_to_nn_module(module, torch.float16, torch.float32)
+    mock_wrap.assert_called_once_with(module, torch.float16, torch.float32)
+
+
+def test_apply_fp8_to_nn_module_skips_norm_modules():
+    """Modules whose path matches `norm` must not be cast — diffusers' `enable_layerwise_casting`
+    does the same. FLUX RMSNorm.scale is the canonical example: a tiny learned scalar that
+    breaks badly in FP8.
+    """
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.norm1 = torch.nn.LayerNorm(4)
+            self.linear = torch.nn.Linear(4, 4)
+
+    storage_dtype = torch.float16
+    compute_dtype = torch.float32
+    model = _Model()
+    for p in model.parameters():
+        p.data = p.data.to(compute_dtype)
+
+    ModelLoader._apply_fp8_to_nn_module(model, storage_dtype, compute_dtype)
+
+    # Linear params get cast to storage dtype.
+    assert model.linear.weight.dtype == storage_dtype
+    # Norm params stay in compute dtype — they must not be cast.
+    assert model.norm1.weight.dtype == compute_dtype
+    assert model.norm1.bias.dtype == compute_dtype
+
+
+def test_apply_fp8_to_nn_module_skips_pos_embed_and_proj_in_out():
+    """Position embeddings and the in/out projection of transformer blocks are also on the
+    diffusers default skip list — they're precision-sensitive.
+    """
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.pos_embed = torch.nn.Linear(4, 4)
+            self.proj_in = torch.nn.Linear(4, 4)
+            self.proj_out = torch.nn.Linear(4, 4)
+            self.attn = torch.nn.Linear(4, 4)
+
+    storage_dtype = torch.float16
+    compute_dtype = torch.float32
+    model = _Model()
+    for p in model.parameters():
+        p.data = p.data.to(compute_dtype)
+
+    ModelLoader._apply_fp8_to_nn_module(model, storage_dtype, compute_dtype)
+
+    assert model.attn.weight.dtype == storage_dtype
+    assert model.pos_embed.weight.dtype == compute_dtype
+    assert model.proj_in.weight.dtype == compute_dtype
+    assert model.proj_out.weight.dtype == compute_dtype
+
+
+def test_apply_fp8_to_nn_module_skips_unsupported_layer_types():
+    """Only the layer classes in `_FP8_SUPPORTED_PYTORCH_LAYERS` are cast — matches diffusers'
+    behavior. A custom RMSNorm-style module with a raw Parameter must be left alone, otherwise
+    its learned scalar gets clobbered.
+    """
+
+    class _ScaleModule(torch.nn.Module):
+        """Mimics FLUX RMSNorm — a tiny learned scalar that must not be cast to FP8."""
+
+        def __init__(self):
+            super().__init__()
+            self.scale = torch.nn.Parameter(torch.ones(4))
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x * self.scale
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.rms = _ScaleModule()
+            self.linear = torch.nn.Linear(4, 4)
+
+    storage_dtype = torch.float16
+    compute_dtype = torch.float32
+    model = _Model()
+    for p in model.parameters():
+        p.data = p.data.to(compute_dtype)
+
+    ModelLoader._apply_fp8_to_nn_module(model, storage_dtype, compute_dtype)
+
+    assert model.linear.weight.dtype == storage_dtype
+    # Critical: the RMS-style scalar lives on a custom module type, not in the supported list.
+    assert model.rms.scale.dtype == compute_dtype
+
+
+def test_wrap_forward_reaches_custom_linear_after_apply_custom_layers():
+    """Production order: `_load_model` applies FP8 wrapping, THEN `ModelCache.put()` calls
+    `apply_custom_layers_to_model` which constructs a NEW `CustomLinear` object via
+    `CustomLinear.__new__` and points its `__dict__` at the original `Linear.__dict__`
+    (see `wrap_custom_layer`). The new object is installed on the parent in place of the
+    original Linear.
+
+    An instance-level `forward` override would be carried into the new CustomLinear via the
+    shared dict but would close over the OLD Linear instance — so calls to the new
+    CustomLinear would silently route to `Linear.forward(old_instance, ...)` and bypass
+    `CustomLinear.forward`, where LoRA/ControlLoRA patches are applied. This is the bug a
+    reviewer reproduced on a fresh worktree.
+
+    Hooks fix this because `nn.Module._call_impl` dispatches them with the *actual* called
+    instance, and `self.forward(...)` is resolved by normal class lookup — reaching
+    `CustomLinear.forward`. This test exercises the production wrapping path (real
+    `apply_custom_layers_to_model`) and asserts CustomLinear.forward is reached by attaching
+    a sentinel patch list and observing that the patch-aware branch runs.
+    """
+
+    class Parent(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.child = torch.nn.Linear(4, 4, bias=False)
+
+    parent = Parent()
+    original_linear = parent.child
+
+    ModelLoader._wrap_forward_with_fp8_cast(original_linear, torch.float16, torch.float32)
+
+    apply_custom_layers_to_model(parent)
+    new_child = parent.child
+
+    # Sanity: production wrapping replaced the child with a NEW CustomLinear instance.
+    assert isinstance(new_child, CustomLinear)
+    assert new_child is not original_linear
+
+    # Attach a sentinel patch so CustomLinear.forward routes through the LoRA-aware branch
+    # (see custom_linear.py: `if len(self._patches_and_weights) > 0`). If that branch fires,
+    # our FP8 wrapping is correctly dispatched through CustomLinear.forward.
+    patch_was_invoked = {"hit": False}
+
+    class _SentinelPatch:
+        def __init__(self):
+            self.hit = patch_was_invoked
+
+        def __call__(self, *_args, **_kwargs):  # not actually called
+            pass
+
+    # Patch the CustomLinear's patch-handling branch to record that it was reached.
+    original_patch_branch = CustomLinear._autocast_forward_with_patches
+
+    def tracked_patch_branch(self, input):
+        patch_was_invoked["hit"] = True
+        # Return a same-shape tensor so the outer caller doesn't choke.
+        return torch.zeros_like(input @ self.weight.t())
+
+    new_child._patches_and_weights = [(_SentinelPatch(), 1.0)]
+    try:
+        CustomLinear._autocast_forward_with_patches = tracked_patch_branch
+        _ = new_child(torch.zeros(1, 4, dtype=torch.float32))
+    finally:
+        CustomLinear._autocast_forward_with_patches = original_patch_branch
+        new_child._patches_and_weights = []
+
+    assert patch_was_invoked["hit"] is True, (
+        "FP8-wrapped forward did not reach CustomLinear.forward — LoRA/ControlLoRA patches "
+        "would be silently bypassed on FP8 checkpoint models."
+    )
+
+
+def test_apply_fp8_layerwise_casting_uses_hook_path_for_model_mixin():
+    """Regression test for the FLUX.2 Klein 9B partial-load device-mismatch crash.
+
+    Diffusers' `enable_layerwise_casting()` registers a `LayerwiseCastingHook` whose
+    `pre_forward` only casts dtype (not device) and whose hook system replaces
+    `Linear.forward` with a wrapper that calls the *original* `Linear.forward` captured
+    before the hook was installed. `ModelCache.put()` later wraps Linear as CustomLinear
+    sharing `__dict__`, so the diffusers wrapper is carried into the new CustomLinear and
+    routes calls to the captured original Linear.forward — bypassing
+    `CustomLinear.forward`'s `cast_to_device`. On partial load (some weights on CPU,
+    input on cuda), this raises a device-mismatch error.
+
+    The fix routes ModelMixin through `_apply_fp8_to_nn_module` (hook-based,
+    `forward`-preserving). This test asserts that path is taken even when the model
+    inherits from ModelMixin.
+    """
+    from diffusers.models.modeling_utils import ModelMixin
+
+    class _FakeModelMixin(ModelMixin):
+        # ModelMixin requires a config_name class attribute and a config dict for serialization.
+        # We never serialize, so we only need to satisfy isinstance() checks.
+        config_name = "config.json"
+
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(4, 4, bias=False)
+
+        def forward(self, x):
+            return self.linear(x)
+
+    loader = _make_loader(device="cuda")
+    config = _make_config(ModelType.Main, fp8=True)
+
+    model = _FakeModelMixin()
+
+    with (
+        patch.object(ModelLoader, "_should_use_fp8", return_value=True),
+        patch.object(ModelLoader, "_apply_fp8_to_nn_module") as mock_to_nn,
+        patch.object(_FakeModelMixin, "enable_layerwise_casting") as mock_enable,
+    ):
+        loader._apply_fp8_layerwise_casting(model, config)
+
+    mock_to_nn.assert_called_once()
+    mock_enable.assert_not_called()

tests/backend/model_manager/load/test_loaded_model.py

@@ -0,0 +1,82 @@
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig
+from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord
+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,
+)
+
+
+class ModelWithRequiredScale(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = torch.nn.Linear(4, 4)
+        self.scale = torch.nn.Parameter(torch.ones(4))
+
+
+class FakeCache:
+    def __init__(self):
+        self.lock_calls = 0
+        self.unlock_calls = 0
+
+    def lock(self, cache_record: CacheRecord, working_mem_bytes: int | None) -> None:
+        del cache_record, working_mem_bytes
+        self.lock_calls += 1
+
+    def unlock(self, cache_record: CacheRecord) -> None:
+        del cache_record
+        self.unlock_calls += 1
+
+
+def test_model_on_device_repairs_required_tensors_for_partial_models():
+    model = ModelWithRequiredScale()
+    apply_custom_layers_to_model(model, device_autocasting_enabled=True)
+    cached_model = CachedModelWithPartialLoad(model=model, compute_device=torch.device("meta"), keep_ram_copy=False)
+    loaded_model = LoadedModelWithoutConfig(
+        cache_record=CacheRecord(key="test", cached_model=cached_model), cache=FakeCache()
+    )
+
+    with loaded_model.model_on_device():
+        assert model.scale.device.type == "meta"
+        assert all(param.device.type == "cpu" for param in model.linear.parameters())
+
+
+def test_model_on_device_leaves_full_load_models_unchanged():
+    model = torch.nn.Linear(4, 4)
+    cached_model = CachedModelOnlyFullLoad(
+        model=model, compute_device=torch.device("meta"), total_bytes=1, keep_ram_copy=False
+    )
+    loaded_model = LoadedModelWithoutConfig(
+        cache_record=CacheRecord(key="test", cached_model=cached_model), cache=FakeCache()
+    )
+
+    with loaded_model.model_on_device() as (_, returned_model):
+        assert returned_model is model
+        assert all(param.device.type == "cpu" for param in model.parameters())
+
+
+def test_enter_unlocks_if_repair_raises():
+    class BrokenCachedModel(CachedModelWithPartialLoad):
+        def repair_required_tensors_on_compute_device(self) -> int:
+            raise RuntimeError("repair failed")
+
+    model = ModelWithRequiredScale()
+    apply_custom_layers_to_model(model, device_autocasting_enabled=True)
+    cached_model = BrokenCachedModel(model=model, compute_device=torch.device("meta"), keep_ram_copy=False)
+    fake_cache = FakeCache()
+    loaded_model = LoadedModelWithoutConfig(
+        cache_record=CacheRecord(key="test", cached_model=cached_model), cache=fake_cache
+    )
+
+    with pytest.raises(RuntimeError, match="repair failed"):
+        loaded_model.__enter__()
+
+    assert fake_cache.lock_calls == 1
+    assert fake_cache.unlock_calls == 1

tests/backend/model_manager/load/test_qwen_image_state_dict_utils.py

@@ -0,0 +1,169 @@
+"""Unit tests for the pure state-dict helpers in the Qwen-Image / Qwen-VL loader.
+
+These freeze the checkpoint key-surgery that the loaders perform before instantiating a model,
+so a regression like the transformers-5.x one (where `_checkpoint_conversion_mapping` became
+`{}` and the `visual.* -> model.visual.*` remap was silently skipped) fails here instead of at
+the user's first load.
+"""
+
+import torch
+
+from invokeai.backend.model_manager.load.model_loaders.qwen_image import (
+    _build_qwen_image_transformer_config,
+    _dequantize_comfyui_fp8,
+    _remap_qwen_vl_checkpoint_keys,
+    _strip_comfyui_prefix,
+    _strip_quantization_metadata,
+)
+from tests.backend.model_manager.load.state_dicts.qwen_vl_encoder_comfyui_keys import (
+    state_dict_keys as qwen_vl_keys,
+)
+from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict
+
+# Prefixes the Qwen2.5-VL architecture (transformers >=4.50) actually exposes. Frozen here on
+# purpose: if a remap regresses, converted keys stop matching this set.
+_VALID_QWEN_VL_PREFIXES = ("model.visual.", "model.language_model.", "lm_head")
+
+
+class TestRemapQwenVlCheckpointKeys:
+    def test_every_key_maps_to_the_transformers_layout(self):
+        """Every legacy ComfyUI key must land under `model.visual.*` / `model.language_model.*`."""
+        sd = keys_to_mock_state_dict(qwen_vl_keys)
+        # The loader strips fp8 metadata (`scaled_fp8` etc.) before remapping; mirror that order.
+        _strip_quantization_metadata(sd)
+
+        remapped = _remap_qwen_vl_checkpoint_keys(sd)
+
+        assert len(remapped) == len(sd)
+        for key in remapped:
+            assert key.startswith(_VALID_QWEN_VL_PREFIXES), f"key not remapped to a known layout: {key}"
+        # No key may survive in the legacy layout.
+        assert not any(k.startswith("visual.") for k in remapped)
+        assert not any(k.startswith(("model.layers.", "model.embed_tokens", "model.norm")) for k in remapped)
+
+    def test_specific_keys_from_the_bug_report(self):
+        """The exact keys that failed to load in the original bug report are remapped."""
+        sd = {
+            "visual.blocks.0.attn.qkv.weight": torch.empty(1),
+            "visual.patch_embed.proj.weight": torch.empty(1),
+            "model.layers.0.self_attn.q_proj.weight": torch.empty(1),
+            "model.embed_tokens.weight": torch.empty(1),
+            "lm_head.weight": torch.empty(1),
+        }
+
+        remapped = _remap_qwen_vl_checkpoint_keys(sd)
+
+        assert "model.visual.blocks.0.attn.qkv.weight" in remapped
+        assert "model.visual.patch_embed.proj.weight" in remapped
+        assert "model.language_model.layers.0.self_attn.q_proj.weight" in remapped
+        assert "model.language_model.embed_tokens.weight" in remapped
+        assert "lm_head.weight" in remapped  # unchanged
+
+    def test_idempotent_on_already_converted_layout(self):
+        """Re-running the remap on new-layout keys must not double-prefix them."""
+        sd = keys_to_mock_state_dict(qwen_vl_keys)
+
+        once = _remap_qwen_vl_checkpoint_keys(sd)
+        twice = _remap_qwen_vl_checkpoint_keys(once)
+
+        assert set(once.keys()) == set(twice.keys())
+
+    def test_fallback_when_transformers_mapping_is_empty(self, monkeypatch):
+        """Even if transformers stops providing `_checkpoint_conversion_mapping`, the remap fires.
+
+        transformers 5.x returns `{}` here; forcing that value pins the fallback that fixes the
+        original bug.
+        """
+        from transformers import Qwen2_5_VLForConditionalGeneration
+
+        monkeypatch.setattr(Qwen2_5_VLForConditionalGeneration, "_checkpoint_conversion_mapping", {})
+
+        remapped = _remap_qwen_vl_checkpoint_keys(
+            {
+                "visual.blocks.0.norm1.weight": torch.empty(1),
+                "model.layers.0.input_layernorm.weight": torch.empty(1),
+            }
+        )
+
+        assert "model.visual.blocks.0.norm1.weight" in remapped
+        assert "model.language_model.layers.0.input_layernorm.weight" in remapped
+
+
+class TestStripQuantizationMetadata:
+    def test_drops_fp8_metadata_keeps_weights(self):
+        sd = keys_to_mock_state_dict(qwen_vl_keys)
+        # The captured checkpoint is fp8_scaled, so it really does ship this metadata.
+        assert any(k.endswith((".scale_weight", ".scale_input")) or k == "scaled_fp8" for k in sd)
+        n_weights_before = sum(1 for k in sd if k.endswith(".weight"))
+
+        _strip_quantization_metadata(sd)
+
+        assert not any(
+            k.endswith((".scale_weight", ".scale_input")) or "comfy_quant" in k or k == "scaled_fp8" for k in sd
+        )
+        # Real weights are untouched.
+        assert sum(1 for k in sd if k.endswith(".weight")) == n_weights_before
+
+
+class TestDequantizeComfyuiFp8:
+    def test_scalar_scale(self):
+        sd = {
+            "l.weight": torch.full((2, 2), 2.0),
+            "l.scale_weight": torch.tensor(3.0),
+            "l.scale_input": torch.tensor(9.0),  # activation scale, must be ignored
+        }
+
+        count = _dequantize_comfyui_fp8(sd, torch.float32)
+
+        assert count == 1
+        assert torch.allclose(sd["l.weight"], torch.full((2, 2), 6.0))
+
+    def test_block_wise_scale_is_broadcast(self):
+        # Per-block scale [2, 1] must be repeat_interleaved up to the weight shape [4, 2].
+        sd = {
+            "l.weight": torch.ones(4, 2),
+            "l.weight_scale": torch.tensor([[10.0], [20.0]]),
+        }
+
+        count = _dequantize_comfyui_fp8(sd, torch.float32)
+
+        assert count == 1
+        expected = torch.tensor([[10.0, 10.0], [10.0, 10.0], [20.0, 20.0], [20.0, 20.0]])
+        assert torch.allclose(sd["l.weight"], expected)
+
+
+class TestStripComfyuiPrefix:
+    def test_strips_diffusion_model_prefix(self):
+        sd = {
+            "model.diffusion_model.transformer_blocks.0.img_mod.1.weight": torch.empty(1),
+            "model.diffusion_model.img_in.weight": torch.empty(1),
+        }
+        out = _strip_comfyui_prefix(sd)
+        assert set(out.keys()) == {"transformer_blocks.0.img_mod.1.weight", "img_in.weight"}
+
+    def test_no_prefix_is_a_noop(self):
+        sd = {"transformer_blocks.0.x": torch.empty(1)}
+        assert _strip_comfyui_prefix(sd) is sd
+
+
+class TestBuildQwenImageTransformerConfig:
+    def test_infers_layer_count_and_dims_from_shapes(self):
+        # torch-order (logical) shapes, as the GGMLTensor.tensor_shape / safetensors path exposes.
+        sd = {
+            "img_in.weight": torch.empty(3072, 64),
+            "txt_in.weight": torch.empty(3072, 3584),
+            "transformer_blocks.0.img_mod.1.weight": torch.empty(1),
+            "transformer_blocks.1.img_mod.1.weight": torch.empty(1),
+            "transformer_blocks.5.img_mod.1.weight": torch.empty(1),
+        }
+
+        cfg = _build_qwen_image_transformer_config(sd, is_edit=False)
+
+        assert cfg["num_layers"] == 6  # max block index (5) + 1
+        assert cfg["in_channels"] == 64
+        assert cfg["num_attention_heads"] == 24  # 3072 // 128
+        assert cfg["joint_attention_dim"] == 3584
+
+    def test_empty_state_dict_falls_back_to_defaults(self):
+        cfg = _build_qwen_image_transformer_config({}, is_edit=False)
+        assert cfg["num_layers"] == 60

tests/backend/model_manager/load/test_z_image_state_dict_utils.py

@@ -0,0 +1,66 @@
+"""Unit tests for the Z-Image GGUF/ComfyUI -> diffusers state-dict converter."""
+
+import torch
+
+from invokeai.backend.model_manager.load.model_loaders.z_image import _convert_z_image_gguf_to_diffusers
+from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict
+from tests.backend.model_manager.load.state_dicts.z_image_transformer_comfyui_keys import (
+    state_dict_keys as z_image_keys,
+)
+
+
+class TestConvertZImageGgufToDiffusers:
+    def test_fused_qkv_split(self):
+        sd = keys_to_mock_state_dict(z_image_keys)
+        n_qkv = sum(1 for k in sd if k.endswith(".attention.qkv.weight"))
+        assert n_qkv > 0
+
+        out = _convert_z_image_gguf_to_diffusers(sd)
+
+        # Each fused qkv weight becomes three separate projections.
+        assert sum(1 for k in out if k.endswith(".attention.to_q.weight")) == n_qkv
+        assert sum(1 for k in out if k.endswith(".attention.to_k.weight")) == n_qkv
+        assert sum(1 for k in out if k.endswith(".attention.to_v.weight")) == n_qkv
+        assert not any(".attention.qkv." in k for k in out)
+
+    def test_key_renames(self):
+        out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys))
+        # q_norm/k_norm -> norm_q/norm_k, attention.out -> attention.to_out.0
+        assert any(k.endswith(".attention.norm_q.weight") for k in out)
+        assert any(k.endswith(".attention.norm_k.weight") for k in out)
+        assert any(k.endswith(".attention.to_out.0.weight") for k in out)
+        assert not any(".q_norm." in k or ".k_norm." in k for k in out)
+        assert not any(".attention.out." in k for k in out)
+
+    def test_embedder_and_final_layer_renamed(self):
+        out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys))
+        assert any(k.startswith("all_x_embedder.2-1.") for k in out)
+        assert any(k.startswith("all_final_layer.2-1.") for k in out)
+        assert not any(k.startswith("x_embedder.") or k.startswith("final_layer.") for k in out)
+
+    def test_norm_final_is_dropped(self):
+        # The diffusers model uses a non-learnable final LayerNorm, so norm_final.* is skipped.
+        assert any(k.startswith("norm_final.") for k in z_image_keys)
+        out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys))
+        assert not any(k.startswith("norm_final.") for k in out)
+
+    def test_pad_tokens_are_2d_after_conversion(self):
+        # The diffusers model expects a leading batch dim on the pad tokens. The checkpoint
+        # already stores them 2D; GGUF ships them 1D (see the reshape test below).
+        out = _convert_z_image_gguf_to_diffusers(keys_to_mock_state_dict(z_image_keys))
+        for pad in ("x_pad_token", "cap_pad_token"):
+            assert out[pad].dim() == 2
+            assert out[pad].shape[0] == 1
+
+    def test_1d_pad_token_gains_batch_dim(self):
+        # GGUF stores pad tokens as [dim]; they must be reshaped to [1, dim].
+        out = _convert_z_image_gguf_to_diffusers({"x_pad_token": torch.arange(4.0)})
+        assert out["x_pad_token"].shape == (1, 4)
+
+    def test_qkv_split_preserves_values(self):
+        # A [6, 2] fused qkv splits into three [2, 2] chunks in order q, k, v.
+        qkv = torch.arange(12, dtype=torch.float32).reshape(6, 2)
+        out = _convert_z_image_gguf_to_diffusers({"blk.attention.qkv.weight": qkv})
+        assert torch.allclose(out["blk.attention.to_q.weight"], qkv[0:2])
+        assert torch.allclose(out["blk.attention.to_k.weight"], qkv[2:4])
+        assert torch.allclose(out["blk.attention.to_v.weight"], qkv[4:6])

tests/backend/model_manager/model_loading/test_model_load.py

@@ -0,0 +1,25 @@
+"""
+Test model loading
+"""
+
+from pathlib import Path
+
+from invokeai.app.services.model_manager import ModelManagerServiceBase
+from invokeai.backend.textual_inversion import TextualInversionModelRaw
+
+
+def test_loading(mm2_model_manager: ModelManagerServiceBase, embedding_file: Path):
+    store = mm2_model_manager.store
+    matches = store.search_by_attr(model_name="test_embedding")
+    assert len(matches) == 0
+    key = mm2_model_manager.install.register_path(embedding_file)
+    loaded_model = mm2_model_manager.load.load_model(store.get_model(key))
+    assert loaded_model is not None
+    assert loaded_model.config.key == key
+    with loaded_model as model:
+        assert isinstance(model, TextualInversionModelRaw)
+
+    config = mm2_model_manager.store.get_model(key)
+    loaded_model_2 = mm2_model_manager.load.load_model(config)
+
+    assert loaded_model.config.key == loaded_model_2.config.key

tests/backend/model_manager/model_manager_fixtures.py

@@ -0,0 +1,359 @@
+# Fixtures to support testing of the model_manager v2 installer, metadata and record store
+
+import os
+import shutil
+from pathlib import Path
+
+import pytest
+from requests.sessions import Session
+from requests_testadapter import TestAdapter, TestSession
+
+from invokeai.app.services.config import InvokeAIAppConfig
+from invokeai.app.services.download import DownloadQueueService, DownloadQueueServiceBase
+from invokeai.app.services.model_install import ModelInstallService, ModelInstallServiceBase
+from invokeai.app.services.model_load import ModelLoadService, ModelLoadServiceBase
+from invokeai.app.services.model_manager import ModelManagerService, ModelManagerServiceBase
+from invokeai.app.services.model_records import ModelRecordServiceBase, ModelRecordServiceSQL
+from invokeai.backend.model_manager.configs.lora import LoRA_Diffusers_SD1_Config, LoRA_Diffusers_SDXL_Config
+from invokeai.backend.model_manager.configs.main import Main_Checkpoint_SD1_Config, Main_Diffusers_SDXL_Config
+from invokeai.backend.model_manager.configs.vae import VAE_Diffusers_SD1_Config
+from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
+from invokeai.backend.model_manager.taxonomy import (
+    BaseModelType,
+    ModelFormat,
+    ModelRepoVariant,
+    ModelSourceType,
+    ModelType,
+    ModelVariantType,
+    SchedulerPredictionType,
+)
+from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.logging import InvokeAILogger
+from tests.backend.model_manager.model_metadata.metadata_examples import (
+    HFTestLoraMetadata,
+    RepoCivitaiModelMetadata1,
+    RepoCivitaiVersionMetadata1,
+    RepoHFMetadata1,
+    RepoHFMetadata1_nofp16,
+    RepoHFModelJson1,
+)
+from tests.fixtures.sqlite_database import create_mock_sqlite_database
+from tests.test_nodes import TestEventService
+
+
+# Create a temporary directory using the contents of `./data/invokeai_root` as the template
+@pytest.fixture
+def mm2_root_dir(tmp_path_factory) -> Path:
+    root_template = Path(__file__).resolve().parent / "data" / "invokeai_root"
+    temp_dir: Path = tmp_path_factory.mktemp("data") / "invokeai_root"
+    shutil.copytree(root_template, temp_dir)
+    return temp_dir
+
+
+@pytest.fixture
+def mm2_model_files(tmp_path_factory) -> Path:
+    root_template = Path(__file__).resolve().parent / "data" / "test_files"
+    temp_dir: Path = tmp_path_factory.mktemp("data") / "test_files"
+    shutil.copytree(root_template, temp_dir)
+    return temp_dir
+
+
+@pytest.fixture
+def embedding_file(mm2_model_files: Path) -> Path:
+    return mm2_model_files / "test_embedding.safetensors"
+
+
+# Can be used to test diffusers model directory loading, but
+# the test file adds ~10MB of space.
+# @pytest.fixture
+# def vae_directory(mm2_model_files: Path) -> Path:
+#     return mm2_model_files / "taesdxl"
+
+
+@pytest.fixture
+def diffusers_dir(mm2_model_files: Path) -> Path:
+    return mm2_model_files / "test-diffusers-main"
+
+
+@pytest.fixture
+def mm2_app_config(mm2_root_dir: Path) -> InvokeAIAppConfig:
+    app_config = InvokeAIAppConfig(models_dir=mm2_root_dir / "models", log_level="info", allow_unknown_models=False)
+    app_config._root = mm2_root_dir
+    return app_config
+
+
+@pytest.fixture
+def mm2_download_queue(mm2_session: Session) -> DownloadQueueServiceBase:
+    download_queue = DownloadQueueService(requests_session=mm2_session)
+    download_queue.start()
+    yield download_queue
+    download_queue.stop()
+
+
+@pytest.fixture
+def mm2_loader(mm2_app_config: InvokeAIAppConfig) -> ModelLoadServiceBase:
+    ram_cache = ModelCache(
+        execution_device_working_mem_gb=mm2_app_config.device_working_mem_gb,
+        enable_partial_loading=mm2_app_config.enable_partial_loading,
+        keep_ram_copy_of_weights=mm2_app_config.keep_ram_copy_of_weights,
+        max_ram_cache_size_gb=mm2_app_config.max_cache_ram_gb,
+        max_vram_cache_size_gb=mm2_app_config.max_cache_vram_gb,
+        execution_device=TorchDevice.choose_torch_device(),
+        logger=InvokeAILogger.get_logger(),
+    )
+    return ModelLoadService(
+        app_config=mm2_app_config,
+        ram_cache=ram_cache,
+    )
+
+
+@pytest.fixture
+def mm2_installer(
+    mm2_app_config: InvokeAIAppConfig,
+    mm2_download_queue: DownloadQueueServiceBase,
+    mm2_session: Session,
+) -> ModelInstallServiceBase:
+    logger = InvokeAILogger.get_logger()
+    db = create_mock_sqlite_database(mm2_app_config, logger)
+    events = TestEventService()
+    store = ModelRecordServiceSQL(db, logger)
+
+    installer = ModelInstallService(
+        app_config=mm2_app_config,
+        record_store=store,
+        download_queue=mm2_download_queue,
+        event_bus=events,
+        session=mm2_session,
+    )
+    installer.start()
+    yield installer
+    installer.stop()
+
+
+@pytest.fixture
+def mm2_record_store(mm2_app_config: InvokeAIAppConfig) -> ModelRecordServiceBase:
+    logger = InvokeAILogger.get_logger(config=mm2_app_config)
+    db = create_mock_sqlite_database(mm2_app_config, logger)
+    store = ModelRecordServiceSQL(db, logger)
+    # add five simple config records to the database
+    config1 = VAE_Diffusers_SD1_Config(
+        key="test_config_1",
+        path="/tmp/foo1",
+        format=ModelFormat.Diffusers,
+        name="test2",
+        base=BaseModelType.StableDiffusion1,
+        type=ModelType.VAE,
+        hash="111222333444",
+        file_size=4096,
+        source="stabilityai/sdxl-vae",
+        source_type=ModelSourceType.HFRepoID,
+        repo_variant=ModelRepoVariant.Default,
+    )
+    config2 = Main_Checkpoint_SD1_Config(
+        key="test_config_2",
+        path="/tmp/foo2.ckpt",
+        name="model1",
+        format=ModelFormat.Checkpoint,
+        base=BaseModelType.StableDiffusion1,
+        type=ModelType.Main,
+        config_path="/tmp/foo.yaml",
+        variant=ModelVariantType.Normal,
+        hash="111222333444",
+        file_size=8192,
+        source="https://civitai.com/models/206883/split",
+        source_type=ModelSourceType.Url,
+        prediction_type=SchedulerPredictionType.Epsilon,
+    )
+    config3 = Main_Diffusers_SDXL_Config(
+        key="test_config_3",
+        path="/tmp/foo3",
+        format=ModelFormat.Diffusers,
+        name="test3",
+        base=BaseModelType.StableDiffusionXL,
+        type=ModelType.Main,
+        hash="111222333444",
+        file_size=8193,
+        source="author3/model3",
+        description="This is test 3",
+        source_type=ModelSourceType.HFRepoID,
+        variant=ModelVariantType.Normal,
+        prediction_type=SchedulerPredictionType.Epsilon,
+        repo_variant=ModelRepoVariant.Default,
+    )
+    config4 = LoRA_Diffusers_SDXL_Config(
+        key="test_config_4",
+        path="/tmp/foo4",
+        format=ModelFormat.Diffusers,
+        name="test4",
+        base=BaseModelType.StableDiffusionXL,
+        type=ModelType.LoRA,
+        hash="111222333444",
+        file_size=5000,
+        source="author4/model4",
+        source_type=ModelSourceType.HFRepoID,
+    )
+    config5 = LoRA_Diffusers_SD1_Config(
+        key="test_config_5",
+        path="/tmp/foo5",
+        format=ModelFormat.Diffusers,
+        name="test5",
+        base=BaseModelType.StableDiffusion1,
+        type=ModelType.LoRA,
+        hash="111222333444",
+        file_size=5001,
+        source="author4/model5",
+        source_type=ModelSourceType.HFRepoID,
+    )
+    store.add_model(config1)
+    store.add_model(config2)
+    store.add_model(config3)
+    store.add_model(config4)
+    store.add_model(config5)
+    return store
+
+
+@pytest.fixture
+def mm2_model_manager(
+    mm2_record_store: ModelRecordServiceBase, mm2_installer: ModelInstallServiceBase, mm2_loader: ModelLoadServiceBase
+) -> ModelManagerServiceBase:
+    return ModelManagerService(store=mm2_record_store, install=mm2_installer, load=mm2_loader)
+
+
+@pytest.fixture
+def mm2_session(embedding_file: Path, diffusers_dir: Path) -> Session:
+    """This fixtures defines a series of mock URLs for testing download and installation."""
+    sess: Session = TestSession()
+    sess.mount(
+        "https://test.com/missing_model.safetensors",
+        TestAdapter(
+            b"missing",
+            status=404,
+        ),
+    )
+    sess.mount(
+        "https://huggingface.co/api/models/stabilityai/sdxl-turbo",
+        TestAdapter(
+            RepoHFMetadata1,
+            headers={"Content-Type": "application/json; charset=utf-8", "Content-Length": len(RepoHFMetadata1)},
+        ),
+    )
+    sess.mount(
+        "https://huggingface.co/api/models/stabilityai/sdxl-turbo-nofp16",
+        TestAdapter(
+            RepoHFMetadata1_nofp16,
+            headers={"Content-Type": "application/json; charset=utf-8", "Content-Length": len(RepoHFMetadata1_nofp16)},
+        ),
+    )
+    sess.mount(
+        "https://civitai.com/api/v1/model-versions/242807",
+        TestAdapter(
+            RepoCivitaiVersionMetadata1,
+            headers={
+                "Content-Length": len(RepoCivitaiVersionMetadata1),
+            },
+        ),
+    )
+    sess.mount(
+        "https://civitai.com/api/v1/models/215485",
+        TestAdapter(
+            RepoCivitaiModelMetadata1,
+            headers={
+                "Content-Length": len(RepoCivitaiModelMetadata1),
+            },
+        ),
+    )
+    sess.mount(
+        "https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/model_index.json",
+        TestAdapter(
+            RepoHFModelJson1,
+            headers={
+                "Content-Length": len(RepoHFModelJson1),
+            },
+        ),
+    )
+    with open(embedding_file, "rb") as f:
+        data = f.read()  # file is small - just 15K
+    sess.mount(
+        "https://www.test.foo/download/test_embedding.safetensors",
+        TestAdapter(data, headers={"Content-Type": "application/octet-stream", "Content-Length": len(data)}),
+    )
+    sess.mount(
+        "https://huggingface.co/api/models/stabilityai/sdxl-turbo",
+        TestAdapter(
+            RepoHFMetadata1,
+            headers={"Content-Type": "application/json; charset=utf-8", "Content-Length": len(RepoHFMetadata1)},
+        ),
+    )
+    sess.mount(
+        "https://huggingface.co/api/models/InvokeAI-test/textual_inversion_tests?blobs=True",
+        TestAdapter(
+            HFTestLoraMetadata,
+            headers={"Content-Type": "application/json; charset=utf-8", "Content-Length": len(HFTestLoraMetadata)},
+        ),
+    )
+    sess.mount(
+        "https://huggingface.co/InvokeAI-test/textual_inversion_tests/resolve/main/learned_embeds-steps-1000.safetensors",
+        TestAdapter(
+            data,
+            headers={"Content-Type": "application/json; charset=utf-8", "Content-Length": len(data)},
+        ),
+    )
+    for root, _, files in os.walk(diffusers_dir):
+        for name in files:
+            path = Path(root, name)
+            url_base = path.relative_to(diffusers_dir).as_posix()
+            url = f"https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/{url_base}"
+            with open(path, "rb") as f:
+                data = f.read()
+            sess.mount(
+                url,
+                TestAdapter(
+                    data,
+                    headers={
+                        "Content-Type": "application/json; charset=utf-8",
+                        "Content-Length": len(data),
+                    },
+                ),
+            )
+
+    for i in ["12345", "9999", "54321"]:
+        content = (
+            b"I am a safetensors file " + bytearray(i, "utf-8") + bytearray(32_000)
+        )  # for pause tests, must make content large
+        sess.mount(
+            f"http://www.civitai.com/models/{i}",
+            TestAdapter(
+                content,
+                headers={
+                    "Content-Length": len(content),
+                    "Content-Disposition": f'filename="mock{i}.safetensors"',
+                },
+            ),
+        )
+
+    sess.mount(
+        "http://www.huggingface.co/foo.txt",
+        TestAdapter(
+            content,
+            headers={
+                "Content-Length": len(content),
+                "Content-Disposition": 'filename="foo.safetensors"',
+            },
+        ),
+    )
+
+    # here are some malformed URLs to test
+    # missing the content length
+    sess.mount(
+        "http://www.civitai.com/models/missing",
+        TestAdapter(
+            b"Missing content length",
+            headers={
+                "Content-Disposition": 'filename="missing.txt"',
+            },
+        ),
+    )
+    # not found test
+    sess.mount("http://www.civitai.com/models/broken", TestAdapter(b"Not found", status=404))
+
+    return sess

tests/backend/model_manager/test_external_api_config.py

@@ -0,0 +1,54 @@
+import pytest
+from pydantic import ValidationError
+
+from invokeai.backend.model_manager.configs.external_api import (
+    ExternalApiModelConfig,
+    ExternalApiModelDefaultSettings,
+    ExternalImageSize,
+    ExternalModelCapabilities,
+)
+
+
+def test_external_api_model_config_defaults() -> None:
+    capabilities = ExternalModelCapabilities(modes=["txt2img"], supports_seed=True)
+
+    config = ExternalApiModelConfig(
+        name="Test External",
+        provider_id="openai",
+        provider_model_id="gpt-image-1",
+        capabilities=capabilities,
+    )
+
+    assert config.path == "external://openai/gpt-image-1"
+    assert config.source == "external://openai/gpt-image-1"
+    assert config.hash == "external:openai:gpt-image-1"
+    assert config.file_size == 0
+    assert config.default_settings is None
+    assert config.capabilities.supports_seed is True
+
+
+def test_external_api_model_capabilities_allows_aspect_ratio_sizes() -> None:
+    capabilities = ExternalModelCapabilities(
+        modes=["txt2img"],
+        allowed_aspect_ratios=["1:1"],
+        aspect_ratio_sizes={"1:1": ExternalImageSize(width=1024, height=1024)},
+    )
+
+    assert capabilities.aspect_ratio_sizes is not None
+    assert capabilities.aspect_ratio_sizes["1:1"].width == 1024
+
+
+def test_external_api_model_config_rejects_extra_fields() -> None:
+    with pytest.raises(ValidationError):
+        ExternalModelCapabilities(modes=["txt2img"], supports_seed=True, extra_field=True)  # type: ignore
+
+    with pytest.raises(ValidationError):
+        ExternalApiModelDefaultSettings(width=512, extra_field=True)  # type: ignore
+
+
+def test_external_api_model_config_validates_limits() -> None:
+    with pytest.raises(ValidationError):
+        ExternalModelCapabilities(modes=["txt2img"], max_images_per_request=0)
+
+    with pytest.raises(ValidationError):
+        ExternalApiModelDefaultSettings(width=0)

tests/backend/model_manager/test_libc_util.py

@@ -0,0 +1,27 @@
+import pytest
+
+from invokeai.backend.model_manager.util.libc_util import LibcUtil, Struct_mallinfo2
+
+
+def test_libc_util_mallinfo2():
+    """Smoke test of LibcUtil().mallinfo2()."""
+    try:
+        libc = LibcUtil()
+    except OSError:
+        # TODO: Set the expected result preemptively based on the system properties.
+        pytest.xfail("libc shared library is not available on this system.")
+
+    try:
+        info = libc.mallinfo2()
+    except AttributeError:
+        pytest.xfail("`mallinfo2` is not available on this system, likely due to glibc < 2.33.")
+
+    assert info.arena > 0
+
+
+def test_struct_mallinfo2_to_str():
+    """Smoke test of Struct_mallinfo2.__str__()."""
+    info = Struct_mallinfo2()
+    info_str = str(info)
+
+    assert len(info_str) > 0

tests/backend/model_manager/test_memory_snapshot.py

@@ -0,0 +1,39 @@
+import pytest
+
+from invokeai.backend.model_manager.load.memory_snapshot import MemorySnapshot, get_pretty_snapshot_diff
+from invokeai.backend.model_manager.util.libc_util import Struct_mallinfo2
+
+
+def test_memory_snapshot_capture():
+    """Smoke test of MemorySnapshot.capture()."""
+    snapshot = MemorySnapshot.capture()
+
+    # We just check process_ram, because it is the only field that should be supported on all platforms.
+    assert snapshot.process_ram > 0
+
+
+snapshots = [
+    MemorySnapshot(process_ram=1, vram=2, malloc_info=Struct_mallinfo2()),
+    MemorySnapshot(process_ram=1, vram=2, malloc_info=None),
+    MemorySnapshot(process_ram=1, vram=None, malloc_info=Struct_mallinfo2()),
+    MemorySnapshot(process_ram=1, vram=None, malloc_info=None),
+    None,
+]
+
+
+@pytest.mark.parametrize("snapshot_1", snapshots)
+@pytest.mark.parametrize("snapshot_2", snapshots)
+def test_get_pretty_snapshot_diff(snapshot_1, snapshot_2):
+    """Test that get_pretty_snapshot_diff() works with various combinations of missing MemorySnapshot fields."""
+    msg = get_pretty_snapshot_diff(snapshot_1, snapshot_2)
+    print(msg)
+
+    expected_lines = 0
+    if snapshot_1 is not None and snapshot_2 is not None:
+        expected_lines += 1
+        if snapshot_1.vram is not None and snapshot_2.vram is not None:
+            expected_lines += 1
+        if snapshot_1.malloc_info is not None and snapshot_2.malloc_info is not None:
+            expected_lines += 5
+
+    assert len(msg.splitlines()) == expected_lines

tests/backend/model_manager/test_model_load_optimization.py

@@ -0,0 +1,73 @@
+import pytest
+import torch
+
+from invokeai.backend.model_manager.load.optimizations import _no_op, skip_torch_weight_init
+
+
+@pytest.mark.parametrize(
+    ["torch_module", "layer_args"],
+    [
+        (torch.nn.Linear, {"in_features": 10, "out_features": 20}),
+        (torch.nn.Conv1d, {"in_channels": 10, "out_channels": 20, "kernel_size": 3}),
+        (torch.nn.Conv2d, {"in_channels": 10, "out_channels": 20, "kernel_size": 3}),
+        (torch.nn.Conv3d, {"in_channels": 10, "out_channels": 20, "kernel_size": 3}),
+        (torch.nn.Embedding, {"num_embeddings": 10, "embedding_dim": 10}),
+    ],
+)
+def test_skip_torch_weight_init_linear(torch_module, layer_args):
+    """Test the interactions between `skip_torch_weight_init()` and various torch modules."""
+    seed = 123
+
+    # Initialize a torch layer *before* applying `skip_torch_weight_init()`.
+    reset_params_fn_before = torch_module.reset_parameters
+    torch.manual_seed(seed)
+    layer_before = torch_module(**layer_args)
+
+    # Initialize a torch layer while `skip_torch_weight_init()` is applied.
+    with skip_torch_weight_init():
+        reset_params_fn_during = torch_module.reset_parameters
+        torch.manual_seed(123)
+        layer_during = torch_module(**layer_args)
+
+    # Initialize a torch layer *after* applying `skip_torch_weight_init()`.
+    reset_params_fn_after = torch_module.reset_parameters
+    torch.manual_seed(123)
+    layer_after = torch_module(**layer_args)
+
+    # Check that reset_parameters is skipped while `skip_torch_weight_init()` is active.
+    assert reset_params_fn_during == _no_op
+    assert not torch.allclose(layer_before.weight, layer_during.weight)
+    if hasattr(layer_before, "bias"):
+        assert not torch.allclose(layer_before.bias, layer_during.bias)
+
+    # Check that the original behavior is restored after `skip_torch_weight_init()` ends.
+    assert reset_params_fn_before is reset_params_fn_after
+    assert torch.allclose(layer_before.weight, layer_after.weight)
+    if hasattr(layer_before, "bias"):
+        assert torch.allclose(layer_before.bias, layer_after.bias)
+
+
+def test_skip_torch_weight_init_restores_base_class_behavior():
+    """Test that `skip_torch_weight_init()` correctly restores the original behavior of torch.nn.Conv*d modules. This
+    test was created to catch a previous bug where `reset_parameters` was being copied from the base `_ConvNd` class to
+    its child classes (like `Conv1d`).
+    """
+    with skip_torch_weight_init():
+        # There is no need to do anything while the context manager is applied, we're just testing that the original
+        # behavior is restored correctly.
+        pass
+
+    # Mock the behavior of another library that monkey patches `torch.nn.modules.conv._ConvNd.reset_parameters` and
+    # expects it to affect all of the sub-classes (e.g. `torch.nn.Conv1D`, `torch.nn.Conv2D`, etc.).
+    called_monkey_patched_fn = False
+
+    def monkey_patched_fn(*args, **kwargs):
+        nonlocal called_monkey_patched_fn
+        called_monkey_patched_fn = True
+
+    saved_fn = torch.nn.modules.conv._ConvNd.reset_parameters
+    torch.nn.modules.conv._ConvNd.reset_parameters = monkey_patched_fn
+    _ = torch.nn.Conv1d(10, 20, 3)
+    torch.nn.modules.conv._ConvNd.reset_parameters = saved_fn
+
+    assert called_monkey_patched_fn

tests/backend/model_manager/util/test_hf_model_select.py

@@ -0,0 +1,405 @@
+from pathlib import Path
+from typing import List
+
+import pytest
+
+from invokeai.backend.model_manager.taxonomy import ModelRepoVariant
+from invokeai.backend.model_manager.util.select_hf_files import filter_files
+
+
+# This is the full list of model paths returned by the HF API for sdxl-base
+@pytest.fixture
+def sdxl_base_files() -> List[Path]:
+    return [
+        Path(x)
+        for x in [
+            ".gitattributes",
+            "01.png",
+            "LICENSE.md",
+            "README.md",
+            "comparison.png",
+            "model_index.json",
+            "pipeline.png",
+            "scheduler/scheduler_config.json",
+            "sd_xl_base_1.0.safetensors",
+            "sd_xl_base_1.0_0.9vae.safetensors",
+            "sd_xl_offset_example-lora_1.0.safetensors",
+            "text_encoder/config.json",
+            "text_encoder/flax_model.msgpack",
+            "text_encoder/model.fp16.safetensors",
+            "text_encoder/model.onnx",
+            "text_encoder/model.safetensors",
+            "text_encoder/openvino_model.bin",
+            "text_encoder/openvino_model.xml",
+            "text_encoder_2/config.json",
+            "text_encoder_2/flax_model.msgpack",
+            "text_encoder_2/model.fp16.safetensors",
+            "text_encoder_2/model.onnx",
+            "text_encoder_2/model.onnx_data",
+            "text_encoder_2/model.safetensors",
+            "text_encoder_2/openvino_model.bin",
+            "text_encoder_2/openvino_model.xml",
+            "tokenizer/merges.txt",
+            "tokenizer/special_tokens_map.json",
+            "tokenizer/tokenizer_config.json",
+            "tokenizer/vocab.json",
+            "tokenizer_2/merges.txt",
+            "tokenizer_2/special_tokens_map.json",
+            "tokenizer_2/tokenizer_config.json",
+            "tokenizer_2/vocab.json",
+            "unet/config.json",
+            "unet/diffusion_flax_model.msgpack",
+            "unet/diffusion_pytorch_model.fp16.safetensors",
+            "unet/diffusion_pytorch_model.safetensors",
+            "unet/model.onnx",
+            "unet/model.onnx_data",
+            "unet/openvino_model.bin",
+            "unet/openvino_model.xml",
+            "vae/config.json",
+            "vae/diffusion_flax_model.msgpack",
+            "vae/diffusion_pytorch_model.fp16.safetensors",
+            "vae/diffusion_pytorch_model.safetensors",
+            "vae_1_0/config.json",
+            "vae_1_0/diffusion_pytorch_model.fp16.safetensors",
+            "vae_1_0/diffusion_pytorch_model.safetensors",
+            "vae_decoder/config.json",
+            "vae_decoder/model.onnx",
+            "vae_decoder/openvino_model.bin",
+            "vae_decoder/openvino_model.xml",
+            "vae_encoder/config.json",
+            "vae_encoder/model.onnx",
+            "vae_encoder/openvino_model.bin",
+            "vae_encoder/openvino_model.xml",
+        ]
+    ]
+
+
+# This are what we expect to get when various diffusers variants are requested
+@pytest.mark.parametrize(
+    "variant,expected_list",
+    [
+        (
+            None,
+            [
+                "model_index.json",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/model.safetensors",
+                "text_encoder_2/config.json",
+                "text_encoder_2/model.safetensors",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "tokenizer_2/merges.txt",
+                "tokenizer_2/special_tokens_map.json",
+                "tokenizer_2/tokenizer_config.json",
+                "tokenizer_2/vocab.json",
+                "unet/config.json",
+                "unet/diffusion_pytorch_model.safetensors",
+                "vae/config.json",
+                "vae/diffusion_pytorch_model.safetensors",
+                "vae_1_0/config.json",
+                "vae_1_0/diffusion_pytorch_model.safetensors",
+            ],
+        ),
+        (
+            ModelRepoVariant.Default,
+            [
+                "model_index.json",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/model.safetensors",
+                "text_encoder_2/config.json",
+                "text_encoder_2/model.safetensors",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "tokenizer_2/merges.txt",
+                "tokenizer_2/special_tokens_map.json",
+                "tokenizer_2/tokenizer_config.json",
+                "tokenizer_2/vocab.json",
+                "unet/config.json",
+                "unet/diffusion_pytorch_model.safetensors",
+                "vae/config.json",
+                "vae/diffusion_pytorch_model.safetensors",
+                "vae_1_0/config.json",
+                "vae_1_0/diffusion_pytorch_model.safetensors",
+            ],
+        ),
+        (
+            ModelRepoVariant.OpenVINO,
+            [
+                "model_index.json",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/openvino_model.bin",
+                "text_encoder/openvino_model.xml",
+                "text_encoder_2/config.json",
+                "text_encoder_2/openvino_model.bin",
+                "text_encoder_2/openvino_model.xml",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "tokenizer_2/merges.txt",
+                "tokenizer_2/special_tokens_map.json",
+                "tokenizer_2/tokenizer_config.json",
+                "tokenizer_2/vocab.json",
+                "unet/config.json",
+                "unet/openvino_model.bin",
+                "unet/openvino_model.xml",
+                "vae_decoder/config.json",
+                "vae_decoder/openvino_model.bin",
+                "vae_decoder/openvino_model.xml",
+                "vae_encoder/config.json",
+                "vae_encoder/openvino_model.bin",
+                "vae_encoder/openvino_model.xml",
+            ],
+        ),
+        (
+            ModelRepoVariant.FP16,
+            [
+                "model_index.json",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/model.fp16.safetensors",
+                "text_encoder_2/config.json",
+                "text_encoder_2/model.fp16.safetensors",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "tokenizer_2/merges.txt",
+                "tokenizer_2/special_tokens_map.json",
+                "tokenizer_2/tokenizer_config.json",
+                "tokenizer_2/vocab.json",
+                "unet/config.json",
+                "unet/diffusion_pytorch_model.fp16.safetensors",
+                "vae/config.json",
+                "vae/diffusion_pytorch_model.fp16.safetensors",
+                "vae_1_0/config.json",
+                "vae_1_0/diffusion_pytorch_model.fp16.safetensors",
+            ],
+        ),
+        (
+            ModelRepoVariant.ONNX,
+            [
+                "model_index.json",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/model.onnx",
+                "text_encoder_2/config.json",
+                "text_encoder_2/model.onnx",
+                "text_encoder_2/model.onnx_data",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "tokenizer_2/merges.txt",
+                "tokenizer_2/special_tokens_map.json",
+                "tokenizer_2/tokenizer_config.json",
+                "tokenizer_2/vocab.json",
+                "unet/config.json",
+                "unet/model.onnx",
+                "unet/model.onnx_data",
+                "vae_decoder/config.json",
+                "vae_decoder/model.onnx",
+                "vae_encoder/config.json",
+                "vae_encoder/model.onnx",
+            ],
+        ),
+        (
+            ModelRepoVariant.Flax,
+            [
+                "model_index.json",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/flax_model.msgpack",
+                "text_encoder_2/config.json",
+                "text_encoder_2/flax_model.msgpack",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "tokenizer_2/merges.txt",
+                "tokenizer_2/special_tokens_map.json",
+                "tokenizer_2/tokenizer_config.json",
+                "tokenizer_2/vocab.json",
+                "unet/config.json",
+                "unet/diffusion_flax_model.msgpack",
+                "vae/config.json",
+                "vae/diffusion_flax_model.msgpack",
+            ],
+        ),
+    ],
+)
+def test_select(sdxl_base_files: List[Path], variant: ModelRepoVariant, expected_list: List[str]) -> None:
+    print(f"testing variant {variant}")
+    filtered_files = filter_files(sdxl_base_files, variant)
+    assert set(filtered_files) == {Path(x) for x in expected_list}
+
+
+@pytest.fixture
+def sd15_test_files() -> list[Path]:
+    return [
+        Path(f)
+        for f in [
+            "feature_extractor/preprocessor_config.json",
+            "safety_checker/config.json",
+            "safety_checker/model.fp16.safetensors",
+            "safety_checker/model.safetensors",
+            "safety_checker/pytorch_model.bin",
+            "safety_checker/pytorch_model.fp16.bin",
+            "scheduler/scheduler_config.json",
+            "text_encoder/config.json",
+            "text_encoder/model.fp16.safetensors",
+            "text_encoder/model.safetensors",
+            "text_encoder/pytorch_model.bin",
+            "text_encoder/pytorch_model.fp16.bin",
+            "tokenizer/merges.txt",
+            "tokenizer/special_tokens_map.json",
+            "tokenizer/tokenizer_config.json",
+            "tokenizer/vocab.json",
+            "unet/config.json",
+            "unet/diffusion_pytorch_model.bin",
+            "unet/diffusion_pytorch_model.fp16.bin",
+            "unet/diffusion_pytorch_model.fp16.safetensors",
+            "unet/diffusion_pytorch_model.non_ema.bin",
+            "unet/diffusion_pytorch_model.non_ema.safetensors",
+            "unet/diffusion_pytorch_model.safetensors",
+            "vae/config.json",
+            "vae/diffusion_pytorch_model.bin",
+            "vae/diffusion_pytorch_model.fp16.bin",
+            "vae/diffusion_pytorch_model.fp16.safetensors",
+            "vae/diffusion_pytorch_model.safetensors",
+        ]
+    ]
+
+
+@pytest.mark.parametrize(
+    "variant,expected_files",
+    [
+        (
+            ModelRepoVariant.FP16,
+            [
+                "feature_extractor/preprocessor_config.json",
+                "safety_checker/config.json",
+                "safety_checker/model.fp16.safetensors",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/model.fp16.safetensors",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "unet/config.json",
+                "unet/diffusion_pytorch_model.fp16.safetensors",
+                "vae/config.json",
+                "vae/diffusion_pytorch_model.fp16.safetensors",
+            ],
+        ),
+        (
+            ModelRepoVariant.FP32,
+            [
+                "feature_extractor/preprocessor_config.json",
+                "safety_checker/config.json",
+                "safety_checker/model.safetensors",
+                "scheduler/scheduler_config.json",
+                "text_encoder/config.json",
+                "text_encoder/model.safetensors",
+                "tokenizer/merges.txt",
+                "tokenizer/special_tokens_map.json",
+                "tokenizer/tokenizer_config.json",
+                "tokenizer/vocab.json",
+                "unet/config.json",
+                "unet/diffusion_pytorch_model.safetensors",
+                "vae/config.json",
+                "vae/diffusion_pytorch_model.safetensors",
+            ],
+        ),
+    ],
+)
+def test_select_multiple_weights(
+    sd15_test_files: list[Path], variant: ModelRepoVariant, expected_files: list[str]
+) -> None:
+    filtered_files = filter_files(sd15_test_files, variant)
+    assert set(filtered_files) == {Path(f) for f in expected_files}
+
+
+@pytest.fixture
+def flux_schnell_test_files() -> list[Path]:
+    return [
+        Path(f)
+        for f in [
+            "FLUX.1-schnell/.gitattributes",
+            "FLUX.1-schnell/README.md",
+            "FLUX.1-schnell/ae.safetensors",
+            "FLUX.1-schnell/flux1-schnell.safetensors",
+            "FLUX.1-schnell/model_index.json",
+            "FLUX.1-schnell/scheduler/scheduler_config.json",
+            "FLUX.1-schnell/schnell_grid.jpeg",
+            "FLUX.1-schnell/text_encoder/config.json",
+            "FLUX.1-schnell/text_encoder/model.safetensors",
+            "FLUX.1-schnell/text_encoder_2/config.json",
+            "FLUX.1-schnell/text_encoder_2/model-00001-of-00002.safetensors",
+            "FLUX.1-schnell/text_encoder_2/model-00002-of-00002.safetensors",
+            "FLUX.1-schnell/text_encoder_2/model.safetensors.index.json",
+            "FLUX.1-schnell/tokenizer/merges.txt",
+            "FLUX.1-schnell/tokenizer/special_tokens_map.json",
+            "FLUX.1-schnell/tokenizer/tokenizer_config.json",
+            "FLUX.1-schnell/tokenizer/vocab.json",
+            "FLUX.1-schnell/tokenizer_2/special_tokens_map.json",
+            "FLUX.1-schnell/tokenizer_2/spiece.model",
+            "FLUX.1-schnell/tokenizer_2/tokenizer.json",
+            "FLUX.1-schnell/tokenizer_2/tokenizer_config.json",
+            "FLUX.1-schnell/transformer/config.json",
+            "FLUX.1-schnell/transformer/diffusion_pytorch_model-00001-of-00003.safetensors",
+            "FLUX.1-schnell/transformer/diffusion_pytorch_model-00002-of-00003.safetensors",
+            "FLUX.1-schnell/transformer/diffusion_pytorch_model-00003-of-00003.safetensors",
+            "FLUX.1-schnell/transformer/diffusion_pytorch_model.safetensors.index.json",
+            "FLUX.1-schnell/vae/config.json",
+            "FLUX.1-schnell/vae/diffusion_pytorch_model.safetensors",
+        ]
+    ]
+
+
+@pytest.mark.parametrize(
+    ["variant", "expected_files"],
+    [
+        (
+            ModelRepoVariant.Default,
+            [
+                "FLUX.1-schnell/model_index.json",
+                "FLUX.1-schnell/scheduler/scheduler_config.json",
+                "FLUX.1-schnell/text_encoder/config.json",
+                "FLUX.1-schnell/text_encoder/model.safetensors",
+                "FLUX.1-schnell/text_encoder_2/config.json",
+                "FLUX.1-schnell/text_encoder_2/model-00001-of-00002.safetensors",
+                "FLUX.1-schnell/text_encoder_2/model-00002-of-00002.safetensors",
+                "FLUX.1-schnell/text_encoder_2/model.safetensors.index.json",
+                "FLUX.1-schnell/tokenizer/merges.txt",
+                "FLUX.1-schnell/tokenizer/special_tokens_map.json",
+                "FLUX.1-schnell/tokenizer/tokenizer_config.json",
+                "FLUX.1-schnell/tokenizer/vocab.json",
+                "FLUX.1-schnell/tokenizer_2/special_tokens_map.json",
+                "FLUX.1-schnell/tokenizer_2/spiece.model",
+                "FLUX.1-schnell/tokenizer_2/tokenizer.json",
+                "FLUX.1-schnell/tokenizer_2/tokenizer_config.json",
+                "FLUX.1-schnell/transformer/config.json",
+                "FLUX.1-schnell/transformer/diffusion_pytorch_model-00001-of-00003.safetensors",
+                "FLUX.1-schnell/transformer/diffusion_pytorch_model-00002-of-00003.safetensors",
+                "FLUX.1-schnell/transformer/diffusion_pytorch_model-00003-of-00003.safetensors",
+                "FLUX.1-schnell/transformer/diffusion_pytorch_model.safetensors.index.json",
+                "FLUX.1-schnell/vae/config.json",
+                "FLUX.1-schnell/vae/diffusion_pytorch_model.safetensors",
+            ],
+        ),
+    ],
+)
+def test_select_flux_schnell_files(
+    flux_schnell_test_files: list[Path], variant: ModelRepoVariant, expected_files: list[str]
+) -> None:
+    filtered_files = filter_files(flux_schnell_test_files, variant)
+    assert set(filtered_files) == {Path(f) for f in expected_files}