chore: import upstream snapshot with attribution

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