chore: import upstream snapshot with attribution

tests/backend/anima/test_control_net_lllite.py

@@ -0,0 +1,764 @@
+"""Tests for the Anima ControlNet-LLLite adapter — construction from a saved
+state dict, exact-passthrough guarantees, forward-swap binding/restore,
+multi-adapter composition, and the conditioning image preprocessing helpers."""
+
+import os
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from invokeai.app.invocations.anima_denoise import AnimaDenoiseInvocation
+from invokeai.app.invocations.anima_lllite import AnimaLLLiteField
+from invokeai.app.invocations.model import ModelIdentifierField
+from invokeai.backend.anima.control_net_lllite import (
+    AnimaControlNetLLLite,
+    build_inpaint_cond_image,
+    prepare_cond_image,
+    prepare_mask,
+    target_cond_hw,
+)
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
+
+# Opt-in test against the real adapter weights (anima-lllite-inpainting-v2.safetensors).
+_REAL_WEIGHTS_ENV_VAR = "ANIMA_LLLITE_WEIGHTS_PATH"
+REAL_WEIGHTS_PATH = Path(os.environ[_REAL_WEIGHTS_ENV_VAR]) if _REAL_WEIGHTS_ENV_VAR in os.environ else None
+
+COND_DIM = 16
+COND_EMB_DIM = 8
+MLP_DIM = 8
+IN_DIM = 16
+N_BLOCKS = 2
+KINDS = ("self_attn_q_proj", "self_attn_k_proj", "self_attn_v_proj", "mlp_layer1")
+
+
+def make_synthetic_state_dict(seed: int = 0, mlp_dim: int = MLP_DIM) -> dict[str, torch.Tensor]:
+    """Saved-format (v2 named-key) state dict for a tiny 2-block adapter with
+    4-channel (inpaint) conditioning and one trunk resblock."""
+    g = torch.Generator().manual_seed(seed)
+
+    def t(*shape: int) -> torch.Tensor:
+        return torch.randn(*shape, generator=g)
+
+    ch_half = COND_DIM // 2
+    sd = {
+        "lllite_conditioning1.conv1.weight": t(ch_half, 4, 4, 4),
+        "lllite_conditioning1.conv1.bias": t(ch_half),
+        "lllite_conditioning1.norm1.weight": t(ch_half),
+        "lllite_conditioning1.norm1.bias": t(ch_half),
+        "lllite_conditioning1.conv2.weight": t(ch_half, ch_half, 3, 3),
+        "lllite_conditioning1.conv2.bias": t(ch_half),
+        "lllite_conditioning1.norm2.weight": t(ch_half),
+        "lllite_conditioning1.norm2.bias": t(ch_half),
+        "lllite_conditioning1.conv3.weight": t(COND_DIM, ch_half, 4, 4),
+        "lllite_conditioning1.conv3.bias": t(COND_DIM),
+        "lllite_conditioning1.norm3.weight": t(COND_DIM),
+        "lllite_conditioning1.norm3.bias": t(COND_DIM),
+        "lllite_conditioning1.resblocks.0.norm1.weight": t(COND_DIM),
+        "lllite_conditioning1.resblocks.0.norm1.bias": t(COND_DIM),
+        "lllite_conditioning1.resblocks.0.conv1.weight": t(COND_DIM, COND_DIM, 3, 3),
+        "lllite_conditioning1.resblocks.0.conv1.bias": t(COND_DIM),
+        "lllite_conditioning1.resblocks.0.norm2.weight": t(COND_DIM),
+        "lllite_conditioning1.resblocks.0.norm2.bias": t(COND_DIM),
+        "lllite_conditioning1.resblocks.0.conv2.weight": t(COND_DIM, COND_DIM, 3, 3),
+        "lllite_conditioning1.resblocks.0.conv2.bias": t(COND_DIM),
+        "lllite_conditioning1.proj.weight": t(COND_EMB_DIM, COND_DIM, 1, 1),
+        "lllite_conditioning1.proj.bias": t(COND_EMB_DIM),
+        "lllite_conditioning1.out_norm.weight": t(COND_EMB_DIM),
+        "lllite_conditioning1.out_norm.bias": t(COND_EMB_DIM),
+    }
+    for i in range(N_BLOCKS):
+        for kind in KINDS:
+            p = f"lllite_dit_blocks_{i}_{kind}"
+            sd[f"{p}.down.weight"] = t(mlp_dim, IN_DIM)
+            sd[f"{p}.down.bias"] = t(mlp_dim)
+            sd[f"{p}.mid.weight"] = t(mlp_dim, mlp_dim + COND_EMB_DIM)
+            sd[f"{p}.mid.bias"] = t(mlp_dim)
+            sd[f"{p}.cond_to_film.weight"] = t(2 * mlp_dim, COND_EMB_DIM)
+            sd[f"{p}.cond_to_film.bias"] = t(2 * mlp_dim)
+            sd[f"{p}.up.weight"] = t(IN_DIM, mlp_dim)
+            sd[f"{p}.up.bias"] = t(IN_DIM)
+            sd[f"{p}.depth_embed"] = t(COND_EMB_DIM)
+    return sd
+
+
+class FakeAttention(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+
+
+class FakeMlp(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.layer1 = nn.Linear(dim, dim * 2, bias=False)
+
+
+class FakeBlock(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.self_attn = FakeAttention(dim)
+        self.cross_attn = FakeAttention(dim)
+        self.mlp = FakeMlp(dim)
+
+
+class FakeTransformer(nn.Module):
+    def __init__(self, dim: int, n_blocks: int):
+        super().__init__()
+        self.blocks = nn.ModuleList([FakeBlock(dim) for _ in range(n_blocks)])
+
+
+def make_model_and_transformer(dim: int = IN_DIM) -> tuple[AnimaControlNetLLLite, FakeTransformer]:
+    model = AnimaControlNetLLLite.from_state_dict(make_synthetic_state_dict(), None)
+    torch.manual_seed(123)
+    transformer = FakeTransformer(dim, N_BLOCKS)
+    return model, transformer
+
+
+def matching_cond_image() -> torch.Tensor:
+    """4ch cond image sized for latent 4x4 -> trunk tokens S = 2*2 = 4."""
+    return torch.randn(1, 4, 32, 32, generator=torch.Generator().manual_seed(7)).clamp(-1, 1)
+
+
+def plain_linear(linear: nn.Linear, x: torch.Tensor) -> torch.Tensor:
+    return F.linear(x, linear.weight, linear.bias)
+
+
+# ----------------------------------------------------------------------------
+# Preprocessing helpers
+# ----------------------------------------------------------------------------
+
+
+def test_target_cond_hw_even_latent() -> None:
+    assert target_cond_hw(128, 128) == (1024, 1024)
+    assert target_cond_hw(64, 96) == (512, 768)
+
+
+def test_target_cond_hw_odd_latent_pads_to_patch_multiple() -> None:
+    assert target_cond_hw(129, 129) == (1040, 1040)
+    assert target_cond_hw(129, 64) == (1040, 512)
+    assert target_cond_hw(5, 5, patch_spatial=4) == (64, 64)
+
+
+def test_prepare_cond_image_no_resize_is_exact_rescale() -> None:
+    rgb = torch.rand(1, 3, 16, 16, generator=torch.Generator().manual_seed(0))
+    out = prepare_cond_image(rgb, latent_h=2, latent_w=2)
+    assert torch.equal(out, rgb * 2.0 - 1.0)
+
+
+def test_prepare_cond_image_resizes_takes_first_frame_and_stays_in_range() -> None:
+    rgb = torch.rand(2, 3, 100, 100, generator=torch.Generator().manual_seed(1))
+    out = prepare_cond_image(rgb, latent_h=9, latent_w=8)
+    assert out.shape == (1, 3, 80, 64)
+    assert out.min().item() >= -1.0
+    assert out.max().item() <= 1.0
+
+
+def test_prepare_cond_image_rejects_bad_shape() -> None:
+    with pytest.raises(ValueError, match="Unexpected cond image shape"):
+        prepare_cond_image(torch.rand(1, 4, 16, 16), latent_h=2, latent_w=2)
+
+
+def test_prepare_mask_binarizes_at_half() -> None:
+    mask = torch.full((1, 1, 16, 16), 0.49)
+    mask[:, :, 8:, :] = 0.5
+    out = prepare_mask(mask, latent_h=2, latent_w=2)
+    assert torch.equal(out[:, :, :8, :], torch.zeros(1, 1, 8, 16))
+    assert torch.equal(out[:, :, 8:, :], torch.ones(1, 1, 8, 16))
+
+
+def test_prepare_mask_accepts_3d_and_resizes_nearest() -> None:
+    mask = torch.zeros(1, 8, 8)
+    mask[:, :4, :] = 1.0
+    out = prepare_mask(mask, latent_h=4, latent_w=4)
+    assert out.shape == (1, 1, 32, 32)
+    assert set(out.unique().tolist()) <= {0.0, 1.0}
+    assert torch.equal(out[:, :, :16, :], torch.ones(1, 1, 16, 32))
+    assert torch.equal(out[:, :, 16:, :], torch.zeros(1, 1, 16, 32))
+
+
+def test_prepare_mask_rejects_bad_shape() -> None:
+    with pytest.raises(ValueError, match="Unexpected mask shape"):
+        prepare_mask(torch.rand(1, 3, 16, 16), latent_h=2, latent_w=2)
+
+
+def test_build_inpaint_cond_image_polarity_and_range() -> None:
+    rgb = torch.full((1, 3, 4, 4), 0.5)
+    mask = torch.zeros(1, 1, 4, 4)
+    mask[:, :, :2, :] = 1.0  # white = inpaint area
+
+    out = build_inpaint_cond_image(rgb, mask, masked_input=True)
+    assert out.shape == (1, 4, 4, 4)
+    # RGB zeroed under the inpaint area, untouched elsewhere.
+    assert torch.equal(out[:, :3, :2, :], torch.zeros(1, 3, 2, 4))
+    assert torch.equal(out[:, :3, 2:, :], rgb[:, :, 2:, :])
+    # Mask channel rescaled to [-1, 1]: +1 = inpaint, -1 = keep.
+    assert torch.equal(out[:, 3:, :2, :], torch.ones(1, 1, 2, 4))
+    assert torch.equal(out[:, 3:, 2:, :], -torch.ones(1, 1, 2, 4))
+
+
+def test_build_inpaint_cond_image_without_masked_input_keeps_rgb() -> None:
+    rgb = torch.full((1, 3, 4, 4), 0.5)
+    mask = torch.ones(1, 1, 4, 4)
+    out = build_inpaint_cond_image(rgb, mask, masked_input=False)
+    assert torch.equal(out[:, :3], rgb)
+    assert torch.equal(out[:, 3:], torch.ones(1, 1, 4, 4))
+
+
+# ----------------------------------------------------------------------------
+# Construction from state dict
+# ----------------------------------------------------------------------------
+
+
+def test_from_state_dict_synthetic_shape_fallbacks() -> None:
+    sd = make_synthetic_state_dict()
+    model = AnimaControlNetLLLite.from_state_dict(sd, None)
+
+    assert len(model.lllite_modules) == N_BLOCKS * len(KINDS)
+    assert model.cond_in_channels == 4
+    assert model.cond_emb_dim == COND_EMB_DIM
+    assert model.mlp_dim == MLP_DIM
+    assert model.cond_dim == COND_DIM
+    assert model.cond_resblocks == 1
+    assert model.use_aspp is False
+    assert model.inpaint_masked_input is False  # metadata-only, defaults False
+
+    # Weights actually landed where they belong.
+    assert torch.equal(model.conditioning1.conv1.weight, sd["lllite_conditioning1.conv1.weight"])
+    by_name = {m.lllite_name: m for m in model.lllite_modules}
+    m0 = by_name["lllite_dit_blocks_0_self_attn_q_proj"]
+    assert torch.equal(m0.depth_embed, sd["lllite_dit_blocks_0_self_attn_q_proj.depth_embed"])
+    assert torch.equal(m0.up.weight, sd["lllite_dit_blocks_0_self_attn_q_proj.up.weight"])
+    m1 = by_name["lllite_dit_blocks_1_mlp_layer1"]
+    assert torch.equal(m1.down.weight, sd["lllite_dit_blocks_1_mlp_layer1.down.weight"])
+
+    assert not any(p.requires_grad for p in model.parameters())
+    assert not model.training
+
+
+def test_from_state_dict_metadata_wins() -> None:
+    metadata = {
+        "lllite.cond_emb_dim": str(COND_EMB_DIM),
+        "lllite.mlp_dim": str(MLP_DIM),
+        "lllite.cond_dim": str(COND_DIM),
+        "lllite.cond_resblocks": "1",
+        "lllite.use_aspp": "false",
+        "lllite.cond_in_channels": "4",
+        "lllite.inpaint_masked_input": "true",
+    }
+    model = AnimaControlNetLLLite.from_state_dict(make_synthetic_state_dict(), metadata)
+    assert model.inpaint_masked_input is True
+    assert model.cond_in_channels == 4
+
+
+def test_from_state_dict_rejects_legacy_format() -> None:
+    sd = make_synthetic_state_dict()
+    sd["lllite_modules.0.down.weight"] = torch.zeros(MLP_DIM, IN_DIM)
+    with pytest.raises(ValueError, match="legacy"):
+        AnimaControlNetLLLite.from_state_dict(sd, None)
+
+
+def test_from_state_dict_strict_on_unknown_keys() -> None:
+    sd = make_synthetic_state_dict()
+    sd["some_unrelated_key"] = torch.zeros(1)
+    with pytest.raises(RuntimeError, match="some_unrelated_key"):
+        AnimaControlNetLLLite.from_state_dict(sd, None)
+
+    sd = make_synthetic_state_dict()
+    sd["lllite_dit_blocks_0_self_attn_q_proj.extra.weight"] = torch.zeros(1)
+    with pytest.raises(RuntimeError, match="extra"):
+        AnimaControlNetLLLite.from_state_dict(sd, None)
+
+
+def test_from_state_dict_strict_on_missing_keys() -> None:
+    sd = make_synthetic_state_dict()
+    del sd["lllite_dit_blocks_0_self_attn_q_proj.depth_embed"]
+    with pytest.raises(RuntimeError, match="depth_embed"):
+        AnimaControlNetLLLite.from_state_dict(sd, None)
+
+
+def test_from_state_dict_requires_modules() -> None:
+    sd = {k: v for k, v in make_synthetic_state_dict().items() if k.startswith("lllite_conditioning1.")}
+    with pytest.raises(ValueError, match="no LLLite modules"):
+        AnimaControlNetLLLite.from_state_dict(sd, None)
+
+
+def test_from_state_dict_missing_down_weight_raises_value_error() -> None:
+    sd = make_synthetic_state_dict()
+    del sd["lllite_dit_blocks_0_self_attn_q_proj.down.weight"]
+    with pytest.raises(ValueError, match="missing key 'lllite_dit_blocks_0_self_attn_q_proj.down.weight'"):
+        AnimaControlNetLLLite.from_state_dict(sd, None)
+
+
+# ----------------------------------------------------------------------------
+# Binding / restore
+# ----------------------------------------------------------------------------
+
+
+def test_apply_to_swaps_forward_and_restore_is_bit_exact() -> None:
+    model, transformer = make_model_and_transformer()
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(2))
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    expected = plain_linear(q_proj, x)
+
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    model.set_cond_image(matching_cond_image())
+    assert not torch.equal(q_proj(x), expected)
+
+    model.restore()
+    assert torch.equal(q_proj(x), expected)
+
+
+def test_apply_to_is_idempotent() -> None:
+    model, transformer = make_model_and_transformer()
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(3))
+    v_proj = transformer.blocks[1].self_attn.v_proj
+    expected = plain_linear(v_proj, x)
+
+    model.apply_to(transformer)
+    model.apply_to(transformer)  # must not double-wrap
+    model.restore()
+    assert torch.equal(v_proj(x), expected)
+
+
+def test_restore_without_apply_is_safe() -> None:
+    model, _ = make_model_and_transformer()
+    model.restore()
+    model.restore()
+
+
+def test_restore_does_not_pin_instance_level_forward() -> None:
+    # An instance-level `forward` left behind after restore() would silently
+    # bypass class-level forward swaps that share the module __dict__ (see
+    # wrap_custom_layer notes in model_manager/load/load_default.py).
+    model, transformer = make_model_and_transformer()
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    assert "forward" not in q_proj.__dict__
+
+    model.apply_to(transformer)
+    assert "forward" in q_proj.__dict__
+    model.restore()
+    assert "forward" not in q_proj.__dict__
+
+
+def test_restore_preserves_preexisting_instance_level_forward() -> None:
+    model, transformer = make_model_and_transformer()
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    sentinel = q_proj.forward
+    q_proj.forward = sentinel  # pre-existing instance-level forward
+
+    model.apply_to(transformer)
+    model.restore()
+    assert q_proj.__dict__.get("forward") is sentinel
+
+
+def test_apply_to_rejects_in_features_mismatch() -> None:
+    model, _ = make_model_and_transformer()
+    transformer = FakeTransformer(IN_DIM * 2, N_BLOCKS)
+    with pytest.raises(ValueError, match="in_features"):
+        model.apply_to(transformer)
+
+
+def test_apply_to_rejects_missing_block() -> None:
+    model, _ = make_model_and_transformer()
+    transformer = FakeTransformer(IN_DIM, 1)
+    with pytest.raises(ValueError, match="block"):
+        model.apply_to(transformer)
+
+
+# ----------------------------------------------------------------------------
+# Forward: passthrough guarantees and active path
+# ----------------------------------------------------------------------------
+
+
+def test_passthrough_multiplier_zero_is_bit_exact() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_cond_image(matching_cond_image())
+    model.set_multiplier(0.0)
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(4))
+    for block in transformer.blocks:
+        for linear in (block.self_attn.q_proj, block.self_attn.k_proj, block.self_attn.v_proj):
+            assert torch.equal(linear(x), plain_linear(linear, x))
+
+
+def test_passthrough_no_cond_is_bit_exact() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(5))
+    assert torch.equal(q_proj(x), plain_linear(q_proj, x))
+
+    model.set_cond_image(matching_cond_image())
+    assert not torch.equal(q_proj(x), plain_linear(q_proj, x))
+    model.set_cond_image(None)  # clearing re-enables passthrough
+    assert torch.equal(q_proj(x), plain_linear(q_proj, x))
+
+
+def test_passthrough_on_seq_len_mismatch_is_bit_exact() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    # Cond image for latent 6x6 -> S = 3*3 = 9, but x has S = 4.
+    model.set_cond_image(torch.randn(1, 4, 48, 48, generator=torch.Generator().manual_seed(6)))
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(7))
+    assert torch.equal(q_proj(x), plain_linear(q_proj, x))
+
+
+def test_passthrough_on_non_divisible_batch_is_bit_exact() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    model.set_cond_image(matching_cond_image().repeat(2, 1, 1, 1))  # cond batch 2
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    x = torch.randn(3, 4, IN_DIM, generator=torch.Generator().manual_seed(8))  # 3 % 2 != 0
+    assert torch.equal(q_proj(x), plain_linear(q_proj, x))
+
+
+def test_cfg_batch_broadcast_matches_single_sample() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    model.set_cond_image(matching_cond_image())  # cond batch 1
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    xa = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(9))
+    xb = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(10))
+    y_batched = q_proj(torch.cat([xa, xb], dim=0))
+    assert not torch.equal(y_batched, plain_linear(q_proj, torch.cat([xa, xb], dim=0)))
+    assert torch.allclose(y_batched[0:1], q_proj(xa), rtol=1e-5, atol=1e-6)
+    assert torch.allclose(y_batched[1:2], q_proj(xb), rtol=1e-5, atol=1e-6)
+
+
+def test_5d_mlp_input_matches_flattened_3d_path() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    model.set_cond_image(matching_cond_image())
+    layer1 = transformer.blocks[0].mlp.layer1
+    x5 = torch.randn(1, 1, 2, 2, IN_DIM, generator=torch.Generator().manual_seed(11))
+    y5 = layer1(x5)
+    assert y5.shape == (1, 1, 2, 2, IN_DIM * 2)
+    y3 = layer1(x5.reshape(1, 4, IN_DIM))
+    assert torch.equal(y5, y3.reshape(1, 1, 2, 2, -1))
+    assert not torch.equal(y5, plain_linear(layer1, x5))
+
+
+def test_5d_passthrough_keeps_shape() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    # Seq mismatch: trunk S = 9, x flattened S = 4 -> identity fallback.
+    model.set_cond_image(torch.randn(1, 4, 48, 48, generator=torch.Generator().manual_seed(12)))
+    layer1 = transformer.blocks[0].mlp.layer1
+    x5 = torch.randn(1, 1, 2, 2, IN_DIM, generator=torch.Generator().manual_seed(13))
+    assert torch.equal(layer1(x5), plain_linear(layer1, x5))
+
+
+def test_forward_casts_mismatched_input_dtype() -> None:
+    model, transformer = make_model_and_transformer()
+    model.apply_to(transformer)
+    model.set_multiplier(1.0)
+    model.set_cond_image(matching_cond_image())
+    transformer.to(torch.bfloat16)
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(14)).to(torch.bfloat16)
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    y = q_proj(x)
+    assert y.dtype == torch.bfloat16
+    assert not torch.equal(y, plain_linear(q_proj, x))
+
+
+# ----------------------------------------------------------------------------
+# Multi-adapter composition (denoise applies in list order, restores LIFO)
+# ----------------------------------------------------------------------------
+
+
+class _IdentityCapture(nn.Linear):
+    """nn.Linear whose forward returns its input unchanged — binding a LLLite
+    module to it reads out the perturbed input the module passes to the
+    forward it wrapped."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+
+def make_two_adapters() -> tuple[AnimaControlNetLLLite, AnimaControlNetLLLite, FakeTransformer]:
+    """Two adapters with different mlp dims (8 and 16) plus one fake tree."""
+    model_a = AnimaControlNetLLLite.from_state_dict(make_synthetic_state_dict(seed=0, mlp_dim=8), None)
+    model_b = AnimaControlNetLLLite.from_state_dict(make_synthetic_state_dict(seed=1, mlp_dim=16), None)
+    torch.manual_seed(123)
+    transformer = FakeTransformer(IN_DIM, N_BLOCKS)
+    return model_a, model_b, transformer
+
+
+def cond_image_for(seed: int) -> torch.Tensor:
+    return torch.randn(1, 4, 32, 32, generator=torch.Generator().manual_seed(seed)).clamp(-1, 1)
+
+
+def test_two_adapters_chain_with_second_wrapping_first() -> None:
+    model_a, model_b, transformer = make_two_adapters()
+    assert model_a.mlp_dim != model_b.mlp_dim
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(20))
+
+    model_a.set_multiplier(1.0)
+    model_a.set_cond_image(cond_image_for(30))
+    model_b.set_multiplier(1.0)
+    model_b.set_cond_image(cond_image_for(31))
+
+    # Capture x + delta_B(x): the input B hands to the forward it wrapped.
+    capture = _IdentityCapture(IN_DIM, IN_DIM, bias=False)
+    module_b = next(m for m in model_b.lllite_modules if m.lllite_name == "lllite_dit_blocks_0_self_attn_q_proj")
+    module_b.bind(capture)
+    try:
+        x_plus_delta_b = capture(x)
+    finally:
+        module_b.unbind()
+    assert not torch.equal(x_plus_delta_b, x)
+
+    model_a.apply_to(transformer)
+    expected = q_proj(x_plus_delta_b)  # A's wrapped forward on B's perturbed input
+
+    model_b.apply_to(transformer)  # B wraps A
+    assert torch.equal(q_proj(x), expected)
+
+    model_b.restore()
+    model_a.restore()
+
+
+def test_two_adapters_reverse_restore_returns_pristine_dispatch() -> None:
+    model_a, model_b, transformer = make_two_adapters()
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(21))
+
+    targets = [
+        linear
+        for block in transformer.blocks
+        for linear in (block.self_attn.q_proj, block.self_attn.k_proj, block.self_attn.v_proj, block.mlp.layer1)
+    ]
+    pre_outputs = [linear(x) for linear in targets]
+
+    for model, seed in ((model_a, 32), (model_b, 33)):
+        model.apply_to(transformer)
+        model.set_multiplier(1.0)
+        model.set_cond_image(cond_image_for(seed))
+    assert all("forward" in linear.__dict__ for linear in targets)
+
+    # LIFO: the last adapter applied is restored first.
+    model_b.restore()
+    model_a.restore()
+
+    for linear, expected in zip(targets, pre_outputs, strict=True):
+        assert "forward" not in linear.__dict__
+        assert torch.equal(linear(x), expected)
+
+
+def test_each_adapter_multiplier_gates_only_its_own_contribution() -> None:
+    model_a, model_b, transformer = make_two_adapters()
+    q_proj = transformer.blocks[0].self_attn.q_proj
+    x = torch.randn(1, 4, IN_DIM, generator=torch.Generator().manual_seed(22))
+    y_plain = plain_linear(q_proj, x)
+    cond_a = cond_image_for(34)
+    cond_b = cond_image_for(35)
+
+    # Single-adapter baselines.
+    model_a.apply_to(transformer)
+    model_a.set_multiplier(1.0)
+    model_a.set_cond_image(cond_a)
+    y_a_only = q_proj(x)
+    model_a.restore()
+
+    model_b.apply_to(transformer)
+    model_b.set_multiplier(1.0)
+    model_b.set_cond_image(cond_b)
+    y_b_only = q_proj(x)
+    model_b.restore()
+
+    assert not torch.equal(y_a_only, y_plain)
+    assert not torch.equal(y_b_only, y_plain)
+    assert not torch.equal(y_a_only, y_b_only)
+
+    # Both applied: zeroing one adapter's multiplier removes exactly its contribution.
+    model_a.apply_to(transformer)
+    model_b.apply_to(transformer)
+    model_a.set_cond_image(cond_a)
+    model_b.set_cond_image(cond_b)
+
+    model_a.set_multiplier(1.0)
+    model_b.set_multiplier(0.0)
+    assert torch.equal(q_proj(x), y_a_only)
+
+    model_a.set_multiplier(0.0)
+    model_b.set_multiplier(1.0)
+    assert torch.equal(q_proj(x), y_b_only)
+
+    model_a.set_multiplier(0.0)
+    model_b.set_multiplier(0.0)
+    assert torch.equal(q_proj(x), y_plain)
+
+    model_a.set_multiplier(1.0)
+    model_b.set_multiplier(1.0)
+    y_both = q_proj(x)
+    assert not torch.equal(y_both, y_a_only)
+    assert not torch.equal(y_both, y_b_only)
+
+    model_b.restore()
+    model_a.restore()
+
+
+# ----------------------------------------------------------------------------
+# Denoise input normalization (duplicate models share one cached instance)
+# ----------------------------------------------------------------------------
+
+
+def _lllite_field(key: str) -> AnimaLLLiteField:
+    return AnimaLLLiteField(
+        image_name="cond_image",
+        control_model=ModelIdentifierField(
+            key=key,
+            hash="blake3:0000",
+            name=f"adapter-{key}",
+            base=BaseModelType.Anima,
+            type=ModelType.ControlNet,
+        ),
+    )
+
+
+def test_normalize_control_lllite_to_list() -> None:
+    assert AnimaDenoiseInvocation._normalize_control_lllite(None) == []
+    single = _lllite_field("key-a")
+    assert AnimaDenoiseInvocation._normalize_control_lllite(single) == [single]
+    pair = [_lllite_field("key-a"), _lllite_field("key-b")]
+    assert AnimaDenoiseInvocation._normalize_control_lllite(pair) == pair
+
+
+def test_normalize_control_lllite_sorts_by_model_key() -> None:
+    """Apply order must be deterministic: collect-node input order follows random node ids."""
+    field_a = _lllite_field("key-a")
+    field_b = _lllite_field("key-b")
+    assert AnimaDenoiseInvocation._normalize_control_lllite([field_b, field_a]) == [field_a, field_b]
+
+
+def test_normalize_control_lllite_rejects_duplicate_model_keys() -> None:
+    pair = [_lllite_field("key-a"), _lllite_field("key-a")]
+    with pytest.raises(ValueError, match="used by more than one"):
+        AnimaDenoiseInvocation._normalize_control_lllite(pair)
+
+
+# ----------------------------------------------------------------------------
+# Step-range gate (_get_lllite_multiplier)
+# ----------------------------------------------------------------------------
+
+
+def _stepped_field(weight: float, begin_step_percent: float, end_step_percent: float) -> AnimaLLLiteField:
+    return AnimaLLLiteField(
+        image_name="cond_image",
+        control_model=ModelIdentifierField(
+            key="key-a",
+            hash="blake3:0000",
+            name="adapter",
+            base=BaseModelType.Anima,
+            type=ModelType.ControlNet,
+        ),
+        weight=weight,
+        begin_step_percent=begin_step_percent,
+        end_step_percent=end_step_percent,
+    )
+
+
+def test_lllite_multiplier_full_range_returns_weight() -> None:
+    """A 0%..100% window returns the field weight at every step, boundaries included."""
+    field = _stepped_field(weight=0.7, begin_step_percent=0.0, end_step_percent=1.0)
+    for step_index in (0, 1, 10, 19, 20):
+        assert AnimaDenoiseInvocation._get_lllite_multiplier(field, step_index, total_steps=20) == 0.7
+
+
+def test_lllite_multiplier_zero_outside_window() -> None:
+    """A 25%..75% window over 20 steps gates to first_step=5, last_step=15 (both inclusive)."""
+    field = _stepped_field(weight=0.5, begin_step_percent=0.25, end_step_percent=0.75)
+    # Below the window.
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 4, total_steps=20) == 0.0
+    # First applied step (floor(0.25 * 20) = 5) is inclusive.
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 5, total_steps=20) == 0.5
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 10, total_steps=20) == 0.5
+    # Last applied step (ceil(0.75 * 20) = 15) is inclusive.
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 15, total_steps=20) == 0.5
+    # Above the window.
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 16, total_steps=20) == 0.0
+
+
+def test_lllite_multiplier_floor_ceil_rounding() -> None:
+    """A non-integer step boundary floors the start and ceils the end (widening, not truncating)."""
+    # 0.34 * 10 = 3.4 -> first_step = floor(3.4) = 3, last_step = ceil(3.4) = 4.
+    field = _stepped_field(weight=1.0, begin_step_percent=0.34, end_step_percent=0.34)
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 2, total_steps=10) == 0.0
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 3, total_steps=10) == 1.0
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 4, total_steps=10) == 1.0
+    assert AnimaDenoiseInvocation._get_lllite_multiplier(field, 5, total_steps=10) == 0.0
+
+
+# ----------------------------------------------------------------------------
+# Real weight file (optional; skipped when the file is not present)
+# ----------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(
+    REAL_WEIGHTS_PATH is None,
+    reason=f"set {_REAL_WEIGHTS_ENV_VAR} to the real LLLite weights file to run",
+)
+def test_from_state_dict_real_file() -> None:
+    from safetensors import safe_open
+    from safetensors.torch import load_file
+
+    assert REAL_WEIGHTS_PATH is not None
+    # A stale path must fail loudly, not silently skip.
+    assert REAL_WEIGHTS_PATH.is_file(), f"{_REAL_WEIGHTS_ENV_VAR} points to a missing file: {REAL_WEIGHTS_PATH}"
+    sd = load_file(str(REAL_WEIGHTS_PATH))
+    with safe_open(str(REAL_WEIGHTS_PATH), framework="pt") as f:
+        metadata = f.metadata()
+
+    model = AnimaControlNetLLLite.from_state_dict(sd, metadata)
+
+    assert len(model.lllite_modules) == 112
+    assert model.cond_in_channels == 4
+    assert model.inpaint_masked_input is True
+    assert model.cond_emb_dim == 64
+    assert model.mlp_dim == 64
+    assert model.cond_dim == 128
+    assert model.cond_resblocks == 4
+    assert model.use_aspp is False
+
+    for m in model.lllite_modules:
+        assert m.in_dim == 2048
+        assert m.down.weight.shape == (64, 2048)
+        assert m.mid.weight.shape == (64, 128)
+        assert m.cond_to_film.weight.shape == (128, 64)
+        assert m.up.weight.shape == (2048, 64)
+        assert m.depth_embed.shape == (64,)
+
+    # Strict loading consumed every saved key (1056 = 48 trunk + 112 * 9).
+    assert len(model.state_dict()) == len(sd) == 1056
+
+    # Forward smoke test on one module bound to a 2048-wide Linear.
+    model = model.to(torch.float32)
+    linear = nn.Linear(2048, 2048)
+    module = model.lllite_modules[0]
+    module.bind(linear)
+    try:
+        model.set_multiplier(1.0)
+        # Cond image for latent 4x4 -> S = 2*2 = 4 trunk tokens.
+        model.set_cond_image(torch.randn(1, 4, 32, 32, generator=torch.Generator().manual_seed(15)))
+        assert module.cond_emb is not None
+        assert module.cond_emb.shape == (1, 4, 64)
+        x = torch.randn(1, 4, 2048, generator=torch.Generator().manual_seed(16))
+        y = linear(x)
+        assert y.shape == (1, 4, 2048)
+        assert torch.isfinite(y).all()
+        assert not torch.equal(y, plain_linear(linear, x))
+    finally:
+        module.unbind()

tests/backend/anima/test_scheduler_driver.py

@@ -0,0 +1,195 @@
+"""Tests for AnimaSchedulerDriver — the helper that hides per-scheduler API quirks
+(sigmas= vs num_inference_steps=, Heun's doubled timestep array, set_begin_index)
+behind a uniform iteration interface."""
+
+import inspect
+
+import pytest
+import torch
+
+from invokeai.app.invocations.anima_denoise import loglinear_timestep_shift
+from invokeai.backend.anima.scheduler_driver import AnimaSchedulerDriver
+from invokeai.backend.flux.schedulers import ANIMA_SCHEDULER_MAP, ANIMA_SHIFT
+
+
+def _anima_sigmas(num_steps: int) -> list[float]:
+    return [loglinear_timestep_shift(ANIMA_SHIFT, 1.0 - i / num_steps) for i in range(num_steps + 1)]
+
+
+@pytest.mark.parametrize("scheduler_name", ["heun", "dpmpp_2m", "dpmpp_2m_sde", "er_sde"])
+def test_driver_full_schedule_iteration_count(scheduler_name: str) -> None:
+    """For full schedules (no clipping), the driver yields enough iterations to
+    cover one full denoise. Heun yields 2N-1 iterations for N user steps."""
+    num_steps = 8
+    sigmas = _anima_sigmas(num_steps)
+    driver = AnimaSchedulerDriver(
+        scheduler_name=scheduler_name,
+        sigmas=sigmas,
+        steps=num_steps,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        device=torch.device("cpu"),
+        seed=0,
+    )
+    iterations = list(driver.iterations())
+    if driver.is_heun:
+        assert len(iterations) == 2 * num_steps - 1
+    else:
+        assert len(iterations) == num_steps
+
+
+@pytest.mark.parametrize("scheduler_name", ["dpmpp_2m", "dpmpp_2m_sde", "er_sde"])
+def test_driver_single_step_schedulers_complete_user_step_every_iteration(scheduler_name: str) -> None:
+    """Non-Heun schedulers report completes_user_step on every iteration."""
+    num_steps = 6
+    driver = AnimaSchedulerDriver(
+        scheduler_name=scheduler_name,
+        sigmas=_anima_sigmas(num_steps),
+        steps=num_steps,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        device=torch.device("cpu"),
+        seed=0,
+    )
+    user_step_count = sum(1 for it in driver.iterations() if it.completes_user_step)
+    assert user_step_count == num_steps
+
+
+def test_driver_heun_completes_user_step_on_second_order_and_terminal() -> None:
+    """Heun yields one completion per user step: each pair's 2nd-order half plus
+    the unpaired terminal 1st-order step (sigma_prev==0)."""
+    num_steps = 4
+    driver = AnimaSchedulerDriver(
+        scheduler_name="heun",
+        sigmas=_anima_sigmas(num_steps),
+        steps=num_steps,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        device=torch.device("cpu"),
+        seed=0,
+    )
+    # state_in_first_order only toggles once scheduler.step runs, so drive a fake
+    # step per iteration to mirror production behaviour.
+    completes_flags = []
+    for it in driver.iterations():
+        completes_flags.append(it.completes_user_step)
+        driver.scheduler.step(
+            model_output=torch.zeros((1, 1, 1, 4, 4)),
+            timestep=it.sched_timestep,
+            sample=torch.zeros((1, 1, 1, 4, 4)),
+        )
+    # N=4 → 7 iterations: indices 1, 3, 5 (SO halves) + 6 (terminal FO) = 4 completions.
+    assert sum(completes_flags) == num_steps
+    assert completes_flags[-1] is True, "terminal Heun first-order step must complete its user step"
+
+
+@pytest.mark.parametrize(
+    ("denoising_start", "denoising_end"),
+    [(0.0, 1.0), (0.2, 1.0), (0.0, 0.8), (0.2, 0.8), (0.5, 0.75)],
+)
+def test_driver_dpmpp_clipped_schedule_starts_at_correct_sigma(denoising_start: float, denoising_end: float) -> None:
+    """DPM++ doesn't accept sigmas= on diffusers 0.35.1; the driver's set_begin_index
+    fallback must expose a first iteration whose sigma matches the clipped Anima reference.
+
+    DPM++ constructs its internal flow schedule via ``np.linspace(1, T, N+1)[:-1]`` rather
+    than the closed-form Anima loglinear shift, so the leading-edge sigma is offset by up
+    to ~2e-3 from the Anima reference. That offset is a property of the scheduler family,
+    not the driver — same offset exists in the pre-refactor code path.
+    """
+    num_steps = 30
+    full_sigmas = _anima_sigmas(num_steps)
+    k_start = int(denoising_start * num_steps)
+    expected_first_sigma = full_sigmas[k_start]
+
+    cls, _ = ANIMA_SCHEDULER_MAP["dpmpp_2m"]
+    accepts_sigmas = "sigmas" in inspect.signature(cls(num_train_timesteps=1000).set_timesteps).parameters
+
+    driver = AnimaSchedulerDriver(
+        scheduler_name="dpmpp_2m",
+        sigmas=full_sigmas[k_start : int(denoising_end * num_steps) + 1] if accepts_sigmas else full_sigmas,
+        steps=num_steps,
+        denoising_start=denoising_start,
+        denoising_end=denoising_end,
+        device=torch.device("cpu"),
+        seed=0,
+    )
+    first_iter = next(driver.iterations())
+    assert abs(first_iter.sigma_curr - expected_first_sigma) < 2e-3
+
+
+@pytest.mark.parametrize(
+    ("denoising_start", "denoising_end", "steps"),
+    [(0.2, 0.8, 30), (0.0, 0.8, 30), (0.2, 1.0, 30), (0.5, 0.75, 20)],
+)
+def test_driver_heun_clipped_schedule_iteration_count(denoising_start: float, denoising_end: float, steps: int) -> None:
+    """Heun clipped schedule: iteration count is 2*(k_end-k_start), clamped so
+    denoising_end=1.0 doesn't run past the 2N-1 array."""
+    full_sigmas = _anima_sigmas(steps)
+    k_start = int(denoising_start * steps)
+    k_end = int(denoising_end * steps)
+
+    driver = AnimaSchedulerDriver(
+        scheduler_name="heun",
+        sigmas=full_sigmas,
+        steps=steps,
+        denoising_start=denoising_start,
+        denoising_end=denoising_end,
+        device=torch.device("cpu"),
+        seed=0,
+    )
+
+    # If Heun's set_timesteps accepts sigmas=, the driver will pass the full schedule directly
+    # and yield 2*steps-1 iterations regardless of clipping. The set_begin_index path applies
+    # only when sigmas= is unsupported.
+    accepts_sigmas = "sigmas" in inspect.signature(driver.scheduler.set_timesteps).parameters
+    if accepts_sigmas:
+        # Driver took the sigma-passing path; sigmas were not pre-clipped here, so the count
+        # reflects the full schedule.
+        assert driver.num_iterations == 2 * steps - 1
+        return
+
+    expected = min(2 * (k_end - k_start), len(driver.scheduler.timesteps) - driver.begin_index)
+    assert driver.num_iterations == expected
+    assert driver.begin_index == 2 * k_start
+
+
+def test_driver_terminal_sigma_prev_is_zero() -> None:
+    """The last iteration's sigma_prev must be 0.0 (terminal noise level)."""
+    driver = AnimaSchedulerDriver(
+        scheduler_name="dpmpp_2m",
+        sigmas=_anima_sigmas(8),
+        steps=8,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        device=torch.device("cpu"),
+        seed=0,
+    )
+    last_iter = list(driver.iterations())[-1]
+    assert last_iter.sigma_prev == 0.0
+
+
+def test_driver_seed_determinism() -> None:
+    """Same seed → identical step_generator state → reproducible SDE noise."""
+    sigmas = _anima_sigmas(8)
+    driver_a = AnimaSchedulerDriver(
+        scheduler_name="er_sde",
+        sigmas=sigmas,
+        steps=8,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        device=torch.device("cpu"),
+        seed=42,
+    )
+    driver_b = AnimaSchedulerDriver(
+        scheduler_name="er_sde",
+        sigmas=sigmas,
+        steps=8,
+        denoising_start=0.0,
+        denoising_end=1.0,
+        device=torch.device("cpu"),
+        seed=42,
+    )
+    # Same seed → same first random draw.
+    a = torch.randn((1, 4), generator=driver_a.step_generator)
+    b = torch.randn((1, 4), generator=driver_b.step_generator)
+    assert torch.equal(a, b)

tests/backend/anima/test_t5_tokenizer.py

@@ -0,0 +1,38 @@
+"""Tests for the bundled T5-XXL tokenizer used by Anima.
+
+Anima feeds T5-XXL token IDs into the LLM Adapter's learned embedding table
+(nn.Embedding(32128, 1024)). The tokenizer is vendored in the package so users
+do not need to install a 9GB T5-XXL encoder just to obtain a ~2MB tokenizer.
+"""
+
+from invokeai.backend.anima.t5_tokenizer import ANIMA_T5_VOCAB_SIZE, load_bundled_t5_tokenizer
+
+
+def test_bundled_tokenizer_is_fast() -> None:
+    tokenizer = load_bundled_t5_tokenizer()
+    assert tokenizer.is_fast
+
+
+def test_bundled_tokenizer_known_ids() -> None:
+    tokenizer = load_bundled_t5_tokenizer()
+    ids = tokenizer("a cat sitting on a mat", truncation=True, max_length=512).input_ids
+    assert ids == [3, 9, 1712, 3823, 30, 3, 9, 6928, 1]
+
+
+def test_bundled_tokenizer_appends_eos() -> None:
+    tokenizer = load_bundled_t5_tokenizer()
+    assert tokenizer("", truncation=True, max_length=512).input_ids == [1]
+
+
+def test_bundled_tokenizer_ids_within_adapter_embedding() -> None:
+    tokenizer = load_bundled_t5_tokenizer()
+    ids = tokenizer(
+        "a very long and unusual prompt with rare tokens: zxqwv 12345",
+        truncation=True,
+        max_length=512,
+    ).input_ids
+    assert all(0 <= i < ANIMA_T5_VOCAB_SIZE for i in ids)
+
+
+def test_bundled_tokenizer_is_cached() -> None:
+    assert load_bundled_t5_tokenizer() is load_bundled_t5_tokenizer()

tests/backend/flux/controlnet/instantx_flux_controlnet_state_dict.py

@@ -0,0 +1,374 @@
+# State dict keys and shapes for an InstantX FLUX ControlNet Union model. Intended to be used for unit tests.
+# These keys were extracted from:
+# https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/blob/4f32d6f2b220f8873d49bb8acc073e1df180c994/diffusion_pytorch_model.safetensors
+instantx_sd_shapes = {
+    "context_embedder.bias": [3072],
+    "context_embedder.weight": [3072, 4096],
+    "controlnet_blocks.0.bias": [3072],
+    "controlnet_blocks.0.weight": [3072, 3072],
+    "controlnet_blocks.1.bias": [3072],
+    "controlnet_blocks.1.weight": [3072, 3072],
+    "controlnet_blocks.2.bias": [3072],
+    "controlnet_blocks.2.weight": [3072, 3072],
+    "controlnet_blocks.3.bias": [3072],
+    "controlnet_blocks.3.weight": [3072, 3072],
+    "controlnet_blocks.4.bias": [3072],
+    "controlnet_blocks.4.weight": [3072, 3072],
+    "controlnet_mode_embedder.weight": [10, 3072],
+    "controlnet_single_blocks.0.bias": [3072],
+    "controlnet_single_blocks.0.weight": [3072, 3072],
+    "controlnet_single_blocks.1.bias": [3072],
+    "controlnet_single_blocks.1.weight": [3072, 3072],
+    "controlnet_single_blocks.2.bias": [3072],
+    "controlnet_single_blocks.2.weight": [3072, 3072],
+    "controlnet_single_blocks.3.bias": [3072],
+    "controlnet_single_blocks.3.weight": [3072, 3072],
+    "controlnet_single_blocks.4.bias": [3072],
+    "controlnet_single_blocks.4.weight": [3072, 3072],
+    "controlnet_single_blocks.5.bias": [3072],
+    "controlnet_single_blocks.5.weight": [3072, 3072],
+    "controlnet_single_blocks.6.bias": [3072],
+    "controlnet_single_blocks.6.weight": [3072, 3072],
+    "controlnet_single_blocks.7.bias": [3072],
+    "controlnet_single_blocks.7.weight": [3072, 3072],
+    "controlnet_single_blocks.8.bias": [3072],
+    "controlnet_single_blocks.8.weight": [3072, 3072],
+    "controlnet_single_blocks.9.bias": [3072],
+    "controlnet_single_blocks.9.weight": [3072, 3072],
+    "controlnet_x_embedder.bias": [3072],
+    "controlnet_x_embedder.weight": [3072, 64],
+    "single_transformer_blocks.0.attn.norm_k.weight": [128],
+    "single_transformer_blocks.0.attn.norm_q.weight": [128],
+    "single_transformer_blocks.0.attn.to_k.bias": [3072],
+    "single_transformer_blocks.0.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.0.attn.to_q.bias": [3072],
+    "single_transformer_blocks.0.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.0.attn.to_v.bias": [3072],
+    "single_transformer_blocks.0.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.0.norm.linear.bias": [9216],
+    "single_transformer_blocks.0.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.0.proj_mlp.bias": [12288],
+    "single_transformer_blocks.0.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.0.proj_out.bias": [3072],
+    "single_transformer_blocks.0.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.1.attn.norm_k.weight": [128],
+    "single_transformer_blocks.1.attn.norm_q.weight": [128],
+    "single_transformer_blocks.1.attn.to_k.bias": [3072],
+    "single_transformer_blocks.1.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.1.attn.to_q.bias": [3072],
+    "single_transformer_blocks.1.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.1.attn.to_v.bias": [3072],
+    "single_transformer_blocks.1.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.1.norm.linear.bias": [9216],
+    "single_transformer_blocks.1.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.1.proj_mlp.bias": [12288],
+    "single_transformer_blocks.1.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.1.proj_out.bias": [3072],
+    "single_transformer_blocks.1.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.2.attn.norm_k.weight": [128],
+    "single_transformer_blocks.2.attn.norm_q.weight": [128],
+    "single_transformer_blocks.2.attn.to_k.bias": [3072],
+    "single_transformer_blocks.2.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.2.attn.to_q.bias": [3072],
+    "single_transformer_blocks.2.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.2.attn.to_v.bias": [3072],
+    "single_transformer_blocks.2.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.2.norm.linear.bias": [9216],
+    "single_transformer_blocks.2.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.2.proj_mlp.bias": [12288],
+    "single_transformer_blocks.2.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.2.proj_out.bias": [3072],
+    "single_transformer_blocks.2.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.3.attn.norm_k.weight": [128],
+    "single_transformer_blocks.3.attn.norm_q.weight": [128],
+    "single_transformer_blocks.3.attn.to_k.bias": [3072],
+    "single_transformer_blocks.3.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.3.attn.to_q.bias": [3072],
+    "single_transformer_blocks.3.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.3.attn.to_v.bias": [3072],
+    "single_transformer_blocks.3.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.3.norm.linear.bias": [9216],
+    "single_transformer_blocks.3.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.3.proj_mlp.bias": [12288],
+    "single_transformer_blocks.3.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.3.proj_out.bias": [3072],
+    "single_transformer_blocks.3.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.4.attn.norm_k.weight": [128],
+    "single_transformer_blocks.4.attn.norm_q.weight": [128],
+    "single_transformer_blocks.4.attn.to_k.bias": [3072],
+    "single_transformer_blocks.4.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.4.attn.to_q.bias": [3072],
+    "single_transformer_blocks.4.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.4.attn.to_v.bias": [3072],
+    "single_transformer_blocks.4.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.4.norm.linear.bias": [9216],
+    "single_transformer_blocks.4.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.4.proj_mlp.bias": [12288],
+    "single_transformer_blocks.4.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.4.proj_out.bias": [3072],
+    "single_transformer_blocks.4.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.5.attn.norm_k.weight": [128],
+    "single_transformer_blocks.5.attn.norm_q.weight": [128],
+    "single_transformer_blocks.5.attn.to_k.bias": [3072],
+    "single_transformer_blocks.5.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.5.attn.to_q.bias": [3072],
+    "single_transformer_blocks.5.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.5.attn.to_v.bias": [3072],
+    "single_transformer_blocks.5.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.5.norm.linear.bias": [9216],
+    "single_transformer_blocks.5.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.5.proj_mlp.bias": [12288],
+    "single_transformer_blocks.5.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.5.proj_out.bias": [3072],
+    "single_transformer_blocks.5.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.6.attn.norm_k.weight": [128],
+    "single_transformer_blocks.6.attn.norm_q.weight": [128],
+    "single_transformer_blocks.6.attn.to_k.bias": [3072],
+    "single_transformer_blocks.6.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.6.attn.to_q.bias": [3072],
+    "single_transformer_blocks.6.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.6.attn.to_v.bias": [3072],
+    "single_transformer_blocks.6.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.6.norm.linear.bias": [9216],
+    "single_transformer_blocks.6.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.6.proj_mlp.bias": [12288],
+    "single_transformer_blocks.6.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.6.proj_out.bias": [3072],
+    "single_transformer_blocks.6.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.7.attn.norm_k.weight": [128],
+    "single_transformer_blocks.7.attn.norm_q.weight": [128],
+    "single_transformer_blocks.7.attn.to_k.bias": [3072],
+    "single_transformer_blocks.7.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.7.attn.to_q.bias": [3072],
+    "single_transformer_blocks.7.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.7.attn.to_v.bias": [3072],
+    "single_transformer_blocks.7.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.7.norm.linear.bias": [9216],
+    "single_transformer_blocks.7.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.7.proj_mlp.bias": [12288],
+    "single_transformer_blocks.7.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.7.proj_out.bias": [3072],
+    "single_transformer_blocks.7.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.8.attn.norm_k.weight": [128],
+    "single_transformer_blocks.8.attn.norm_q.weight": [128],
+    "single_transformer_blocks.8.attn.to_k.bias": [3072],
+    "single_transformer_blocks.8.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.8.attn.to_q.bias": [3072],
+    "single_transformer_blocks.8.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.8.attn.to_v.bias": [3072],
+    "single_transformer_blocks.8.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.8.norm.linear.bias": [9216],
+    "single_transformer_blocks.8.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.8.proj_mlp.bias": [12288],
+    "single_transformer_blocks.8.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.8.proj_out.bias": [3072],
+    "single_transformer_blocks.8.proj_out.weight": [3072, 15360],
+    "single_transformer_blocks.9.attn.norm_k.weight": [128],
+    "single_transformer_blocks.9.attn.norm_q.weight": [128],
+    "single_transformer_blocks.9.attn.to_k.bias": [3072],
+    "single_transformer_blocks.9.attn.to_k.weight": [3072, 3072],
+    "single_transformer_blocks.9.attn.to_q.bias": [3072],
+    "single_transformer_blocks.9.attn.to_q.weight": [3072, 3072],
+    "single_transformer_blocks.9.attn.to_v.bias": [3072],
+    "single_transformer_blocks.9.attn.to_v.weight": [3072, 3072],
+    "single_transformer_blocks.9.norm.linear.bias": [9216],
+    "single_transformer_blocks.9.norm.linear.weight": [9216, 3072],
+    "single_transformer_blocks.9.proj_mlp.bias": [12288],
+    "single_transformer_blocks.9.proj_mlp.weight": [12288, 3072],
+    "single_transformer_blocks.9.proj_out.bias": [3072],
+    "single_transformer_blocks.9.proj_out.weight": [3072, 15360],
+    "time_text_embed.guidance_embedder.linear_1.bias": [3072],
+    "time_text_embed.guidance_embedder.linear_1.weight": [3072, 256],
+    "time_text_embed.guidance_embedder.linear_2.bias": [3072],
+    "time_text_embed.guidance_embedder.linear_2.weight": [3072, 3072],
+    "time_text_embed.text_embedder.linear_1.bias": [3072],
+    "time_text_embed.text_embedder.linear_1.weight": [3072, 768],
+    "time_text_embed.text_embedder.linear_2.bias": [3072],
+    "time_text_embed.text_embedder.linear_2.weight": [3072, 3072],
+    "time_text_embed.timestep_embedder.linear_1.bias": [3072],
+    "time_text_embed.timestep_embedder.linear_1.weight": [3072, 256],
+    "time_text_embed.timestep_embedder.linear_2.bias": [3072],
+    "time_text_embed.timestep_embedder.linear_2.weight": [3072, 3072],
+    "transformer_blocks.0.attn.add_k_proj.bias": [3072],
+    "transformer_blocks.0.attn.add_k_proj.weight": [3072, 3072],
+    "transformer_blocks.0.attn.add_q_proj.bias": [3072],
+    "transformer_blocks.0.attn.add_q_proj.weight": [3072, 3072],
+    "transformer_blocks.0.attn.add_v_proj.bias": [3072],
+    "transformer_blocks.0.attn.add_v_proj.weight": [3072, 3072],
+    "transformer_blocks.0.attn.norm_added_k.weight": [128],
+    "transformer_blocks.0.attn.norm_added_q.weight": [128],
+    "transformer_blocks.0.attn.norm_k.weight": [128],
+    "transformer_blocks.0.attn.norm_q.weight": [128],
+    "transformer_blocks.0.attn.to_add_out.bias": [3072],
+    "transformer_blocks.0.attn.to_add_out.weight": [3072, 3072],
+    "transformer_blocks.0.attn.to_k.bias": [3072],
+    "transformer_blocks.0.attn.to_k.weight": [3072, 3072],
+    "transformer_blocks.0.attn.to_out.0.bias": [3072],
+    "transformer_blocks.0.attn.to_out.0.weight": [3072, 3072],
+    "transformer_blocks.0.attn.to_q.bias": [3072],
+    "transformer_blocks.0.attn.to_q.weight": [3072, 3072],
+    "transformer_blocks.0.attn.to_v.bias": [3072],
+    "transformer_blocks.0.attn.to_v.weight": [3072, 3072],
+    "transformer_blocks.0.ff.net.0.proj.bias": [12288],
+    "transformer_blocks.0.ff.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.0.ff.net.2.bias": [3072],
+    "transformer_blocks.0.ff.net.2.weight": [3072, 12288],
+    "transformer_blocks.0.ff_context.net.0.proj.bias": [12288],
+    "transformer_blocks.0.ff_context.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.0.ff_context.net.2.bias": [3072],
+    "transformer_blocks.0.ff_context.net.2.weight": [3072, 12288],
+    "transformer_blocks.0.norm1.linear.bias": [18432],
+    "transformer_blocks.0.norm1.linear.weight": [18432, 3072],
+    "transformer_blocks.0.norm1_context.linear.bias": [18432],
+    "transformer_blocks.0.norm1_context.linear.weight": [18432, 3072],
+    "transformer_blocks.1.attn.add_k_proj.bias": [3072],
+    "transformer_blocks.1.attn.add_k_proj.weight": [3072, 3072],
+    "transformer_blocks.1.attn.add_q_proj.bias": [3072],
+    "transformer_blocks.1.attn.add_q_proj.weight": [3072, 3072],
+    "transformer_blocks.1.attn.add_v_proj.bias": [3072],
+    "transformer_blocks.1.attn.add_v_proj.weight": [3072, 3072],
+    "transformer_blocks.1.attn.norm_added_k.weight": [128],
+    "transformer_blocks.1.attn.norm_added_q.weight": [128],
+    "transformer_blocks.1.attn.norm_k.weight": [128],
+    "transformer_blocks.1.attn.norm_q.weight": [128],
+    "transformer_blocks.1.attn.to_add_out.bias": [3072],
+    "transformer_blocks.1.attn.to_add_out.weight": [3072, 3072],
+    "transformer_blocks.1.attn.to_k.bias": [3072],
+    "transformer_blocks.1.attn.to_k.weight": [3072, 3072],
+    "transformer_blocks.1.attn.to_out.0.bias": [3072],
+    "transformer_blocks.1.attn.to_out.0.weight": [3072, 3072],
+    "transformer_blocks.1.attn.to_q.bias": [3072],
+    "transformer_blocks.1.attn.to_q.weight": [3072, 3072],
+    "transformer_blocks.1.attn.to_v.bias": [3072],
+    "transformer_blocks.1.attn.to_v.weight": [3072, 3072],
+    "transformer_blocks.1.ff.net.0.proj.bias": [12288],
+    "transformer_blocks.1.ff.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.1.ff.net.2.bias": [3072],
+    "transformer_blocks.1.ff.net.2.weight": [3072, 12288],
+    "transformer_blocks.1.ff_context.net.0.proj.bias": [12288],
+    "transformer_blocks.1.ff_context.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.1.ff_context.net.2.bias": [3072],
+    "transformer_blocks.1.ff_context.net.2.weight": [3072, 12288],
+    "transformer_blocks.1.norm1.linear.bias": [18432],
+    "transformer_blocks.1.norm1.linear.weight": [18432, 3072],
+    "transformer_blocks.1.norm1_context.linear.bias": [18432],
+    "transformer_blocks.1.norm1_context.linear.weight": [18432, 3072],
+    "transformer_blocks.2.attn.add_k_proj.bias": [3072],
+    "transformer_blocks.2.attn.add_k_proj.weight": [3072, 3072],
+    "transformer_blocks.2.attn.add_q_proj.bias": [3072],
+    "transformer_blocks.2.attn.add_q_proj.weight": [3072, 3072],
+    "transformer_blocks.2.attn.add_v_proj.bias": [3072],
+    "transformer_blocks.2.attn.add_v_proj.weight": [3072, 3072],
+    "transformer_blocks.2.attn.norm_added_k.weight": [128],
+    "transformer_blocks.2.attn.norm_added_q.weight": [128],
+    "transformer_blocks.2.attn.norm_k.weight": [128],
+    "transformer_blocks.2.attn.norm_q.weight": [128],
+    "transformer_blocks.2.attn.to_add_out.bias": [3072],
+    "transformer_blocks.2.attn.to_add_out.weight": [3072, 3072],
+    "transformer_blocks.2.attn.to_k.bias": [3072],
+    "transformer_blocks.2.attn.to_k.weight": [3072, 3072],
+    "transformer_blocks.2.attn.to_out.0.bias": [3072],
+    "transformer_blocks.2.attn.to_out.0.weight": [3072, 3072],
+    "transformer_blocks.2.attn.to_q.bias": [3072],
+    "transformer_blocks.2.attn.to_q.weight": [3072, 3072],
+    "transformer_blocks.2.attn.to_v.bias": [3072],
+    "transformer_blocks.2.attn.to_v.weight": [3072, 3072],
+    "transformer_blocks.2.ff.net.0.proj.bias": [12288],
+    "transformer_blocks.2.ff.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.2.ff.net.2.bias": [3072],
+    "transformer_blocks.2.ff.net.2.weight": [3072, 12288],
+    "transformer_blocks.2.ff_context.net.0.proj.bias": [12288],
+    "transformer_blocks.2.ff_context.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.2.ff_context.net.2.bias": [3072],
+    "transformer_blocks.2.ff_context.net.2.weight": [3072, 12288],
+    "transformer_blocks.2.norm1.linear.bias": [18432],
+    "transformer_blocks.2.norm1.linear.weight": [18432, 3072],
+    "transformer_blocks.2.norm1_context.linear.bias": [18432],
+    "transformer_blocks.2.norm1_context.linear.weight": [18432, 3072],
+    "transformer_blocks.3.attn.add_k_proj.bias": [3072],
+    "transformer_blocks.3.attn.add_k_proj.weight": [3072, 3072],
+    "transformer_blocks.3.attn.add_q_proj.bias": [3072],
+    "transformer_blocks.3.attn.add_q_proj.weight": [3072, 3072],
+    "transformer_blocks.3.attn.add_v_proj.bias": [3072],
+    "transformer_blocks.3.attn.add_v_proj.weight": [3072, 3072],
+    "transformer_blocks.3.attn.norm_added_k.weight": [128],
+    "transformer_blocks.3.attn.norm_added_q.weight": [128],
+    "transformer_blocks.3.attn.norm_k.weight": [128],
+    "transformer_blocks.3.attn.norm_q.weight": [128],
+    "transformer_blocks.3.attn.to_add_out.bias": [3072],
+    "transformer_blocks.3.attn.to_add_out.weight": [3072, 3072],
+    "transformer_blocks.3.attn.to_k.bias": [3072],
+    "transformer_blocks.3.attn.to_k.weight": [3072, 3072],
+    "transformer_blocks.3.attn.to_out.0.bias": [3072],
+    "transformer_blocks.3.attn.to_out.0.weight": [3072, 3072],
+    "transformer_blocks.3.attn.to_q.bias": [3072],
+    "transformer_blocks.3.attn.to_q.weight": [3072, 3072],
+    "transformer_blocks.3.attn.to_v.bias": [3072],
+    "transformer_blocks.3.attn.to_v.weight": [3072, 3072],
+    "transformer_blocks.3.ff.net.0.proj.bias": [12288],
+    "transformer_blocks.3.ff.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.3.ff.net.2.bias": [3072],
+    "transformer_blocks.3.ff.net.2.weight": [3072, 12288],
+    "transformer_blocks.3.ff_context.net.0.proj.bias": [12288],
+    "transformer_blocks.3.ff_context.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.3.ff_context.net.2.bias": [3072],
+    "transformer_blocks.3.ff_context.net.2.weight": [3072, 12288],
+    "transformer_blocks.3.norm1.linear.bias": [18432],
+    "transformer_blocks.3.norm1.linear.weight": [18432, 3072],
+    "transformer_blocks.3.norm1_context.linear.bias": [18432],
+    "transformer_blocks.3.norm1_context.linear.weight": [18432, 3072],
+    "transformer_blocks.4.attn.add_k_proj.bias": [3072],
+    "transformer_blocks.4.attn.add_k_proj.weight": [3072, 3072],
+    "transformer_blocks.4.attn.add_q_proj.bias": [3072],
+    "transformer_blocks.4.attn.add_q_proj.weight": [3072, 3072],
+    "transformer_blocks.4.attn.add_v_proj.bias": [3072],
+    "transformer_blocks.4.attn.add_v_proj.weight": [3072, 3072],
+    "transformer_blocks.4.attn.norm_added_k.weight": [128],
+    "transformer_blocks.4.attn.norm_added_q.weight": [128],
+    "transformer_blocks.4.attn.norm_k.weight": [128],
+    "transformer_blocks.4.attn.norm_q.weight": [128],
+    "transformer_blocks.4.attn.to_add_out.bias": [3072],
+    "transformer_blocks.4.attn.to_add_out.weight": [3072, 3072],
+    "transformer_blocks.4.attn.to_k.bias": [3072],
+    "transformer_blocks.4.attn.to_k.weight": [3072, 3072],
+    "transformer_blocks.4.attn.to_out.0.bias": [3072],
+    "transformer_blocks.4.attn.to_out.0.weight": [3072, 3072],
+    "transformer_blocks.4.attn.to_q.bias": [3072],
+    "transformer_blocks.4.attn.to_q.weight": [3072, 3072],
+    "transformer_blocks.4.attn.to_v.bias": [3072],
+    "transformer_blocks.4.attn.to_v.weight": [3072, 3072],
+    "transformer_blocks.4.ff.net.0.proj.bias": [12288],
+    "transformer_blocks.4.ff.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.4.ff.net.2.bias": [3072],
+    "transformer_blocks.4.ff.net.2.weight": [3072, 12288],
+    "transformer_blocks.4.ff_context.net.0.proj.bias": [12288],
+    "transformer_blocks.4.ff_context.net.0.proj.weight": [12288, 3072],
+    "transformer_blocks.4.ff_context.net.2.bias": [3072],
+    "transformer_blocks.4.ff_context.net.2.weight": [3072, 12288],
+    "transformer_blocks.4.norm1.linear.bias": [18432],
+    "transformer_blocks.4.norm1.linear.weight": [18432, 3072],
+    "transformer_blocks.4.norm1_context.linear.bias": [18432],
+    "transformer_blocks.4.norm1_context.linear.weight": [18432, 3072],
+    "x_embedder.bias": [3072],
+    "x_embedder.weight": [3072, 64],
+}
+
+
+# InstantX FLUX ControlNet config for unit tests.
+# Copied from https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union/blob/main/config.json
+instantx_config = {
+    "_class_name": "FluxControlNetModel",
+    "_diffusers_version": "0.30.0.dev0",
+    "_name_or_path": "/mnt/wangqixun/",
+    "attention_head_dim": 128,
+    "axes_dims_rope": [16, 56, 56],
+    "guidance_embeds": True,
+    "in_channels": 64,
+    "joint_attention_dim": 4096,
+    "num_attention_heads": 24,
+    "num_layers": 5,
+    "num_mode": 10,
+    "num_single_layers": 10,
+    "patch_size": 1,
+    "pooled_projection_dim": 768,
+}

tests/backend/flux/controlnet/test_state_dict_utils.py

@@ -0,0 +1,108 @@
+import sys
+
+import pytest
+import torch
+
+from invokeai.backend.flux.controlnet.instantx_controlnet_flux import InstantXControlNetFlux
+from invokeai.backend.flux.controlnet.state_dict_utils import (
+    convert_diffusers_instantx_state_dict_to_bfl_format,
+    infer_flux_params_from_state_dict,
+    infer_instantx_num_control_modes_from_state_dict,
+    is_state_dict_instantx_controlnet,
+    is_state_dict_xlabs_controlnet,
+)
+from tests.backend.flux.controlnet.instantx_flux_controlnet_state_dict import instantx_config, instantx_sd_shapes
+from tests.backend.flux.controlnet.xlabs_flux_controlnet_state_dict import xlabs_sd_shapes
+
+
+@pytest.mark.parametrize(
+    ["sd_shapes", "expected"],
+    [
+        (xlabs_sd_shapes, True),
+        (instantx_sd_shapes, False),
+        (["foo"], False),
+    ],
+)
+def test_is_state_dict_xlabs_controlnet(sd_shapes: dict[str, list[int]], expected: bool):
+    sd = dict.fromkeys(sd_shapes)
+    assert is_state_dict_xlabs_controlnet(sd) == expected
+
+
+@pytest.mark.parametrize(
+    ["sd_keys", "expected"],
+    [
+        (instantx_sd_shapes, True),
+        (xlabs_sd_shapes, False),
+        (["foo"], False),
+    ],
+)
+def test_is_state_dict_instantx_controlnet(sd_keys: list[str], expected: bool):
+    sd = dict.fromkeys(sd_keys)
+    assert is_state_dict_instantx_controlnet(sd) == expected
+
+
+def test_convert_diffusers_instantx_state_dict_to_bfl_format():
+    """Smoke test convert_diffusers_instantx_state_dict_to_bfl_format() to ensure that it handles all of the keys."""
+    sd = {k: torch.zeros(1) for k in instantx_sd_shapes}
+    bfl_sd = convert_diffusers_instantx_state_dict_to_bfl_format(sd)
+    assert bfl_sd is not None
+
+
+# TODO(ryand): Figure out why some tests in this file are failing on the MacOS CI runners. It seems to be related to
+# using the meta device. I can't reproduce the issue on my local MacOS system.
+
+
+@pytest.mark.skipif(sys.platform == "darwin", reason="Skipping on macOS")
+def test_infer_flux_params_from_state_dict():
+    # Construct a dummy state_dict with tensors of the correct shape on the meta device.
+    with torch.device("meta"):
+        sd = {k: torch.zeros(v) for k, v in instantx_sd_shapes.items()}
+
+    sd = convert_diffusers_instantx_state_dict_to_bfl_format(sd)
+    flux_params = infer_flux_params_from_state_dict(sd)
+
+    assert flux_params.in_channels == instantx_config["in_channels"]
+    assert flux_params.vec_in_dim == instantx_config["pooled_projection_dim"]
+    assert flux_params.context_in_dim == instantx_config["joint_attention_dim"]
+    assert flux_params.hidden_size // flux_params.num_heads == instantx_config["attention_head_dim"]
+    assert flux_params.num_heads == instantx_config["num_attention_heads"]
+    assert flux_params.mlp_ratio == 4
+    assert flux_params.depth == instantx_config["num_layers"]
+    assert flux_params.depth_single_blocks == instantx_config["num_single_layers"]
+    assert flux_params.axes_dim == instantx_config["axes_dims_rope"]
+    assert flux_params.theta == 10000
+    assert flux_params.qkv_bias
+    assert flux_params.guidance_embed == instantx_config["guidance_embeds"]
+
+
+@pytest.mark.skipif(sys.platform == "darwin", reason="Skipping on macOS")
+def test_infer_instantx_num_control_modes_from_state_dict():
+    # Construct a dummy state_dict with tensors of the correct shape on the meta device.
+    with torch.device("meta"):
+        sd = {k: torch.zeros(v) for k, v in instantx_sd_shapes.items()}
+
+    sd = convert_diffusers_instantx_state_dict_to_bfl_format(sd)
+    num_control_modes = infer_instantx_num_control_modes_from_state_dict(sd)
+
+    assert num_control_modes == instantx_config["num_mode"]
+
+
+@pytest.mark.skipif(sys.platform == "darwin", reason="Skipping on macOS")
+def test_load_instantx_from_state_dict():
+    # Construct a dummy state_dict with tensors of the correct shape on the meta device.
+    with torch.device("meta"):
+        sd = {k: torch.zeros(v) for k, v in instantx_sd_shapes.items()}
+
+    sd = convert_diffusers_instantx_state_dict_to_bfl_format(sd)
+    flux_params = infer_flux_params_from_state_dict(sd)
+    num_control_modes = infer_instantx_num_control_modes_from_state_dict(sd)
+
+    with torch.device("meta"):
+        model = InstantXControlNetFlux(flux_params, num_control_modes)
+
+    model_sd = model.state_dict()
+
+    assert set(model_sd.keys()) == set(sd.keys())
+    for key, tensor in model_sd.items():
+        assert isinstance(tensor, torch.Tensor)
+        assert tensor.shape == sd[key].shape

tests/backend/flux/controlnet/xlabs_flux_controlnet_state_dict.py

@@ -0,0 +1,91 @@
+# State dict keys and shapes for an XLabs FLUX ControlNet model. Intended to be used for unit tests.
+# These keys were extracted from:
+# https://huggingface.co/XLabs-AI/flux-controlnet-collections/blob/86ab1e915a389d5857135c00e0d350e9e38a9048/flux-canny-controlnet_v2.safetensors
+xlabs_sd_shapes = {
+    "controlnet_blocks.0.bias": [3072],
+    "controlnet_blocks.0.weight": [3072, 3072],
+    "controlnet_blocks.1.bias": [3072],
+    "controlnet_blocks.1.weight": [3072, 3072],
+    "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.bias": [3072],
+    "double_blocks.0.img_attn.proj.weight": [3072, 3072],
+    "double_blocks.0.img_attn.qkv.bias": [9216],
+    "double_blocks.0.img_attn.qkv.weight": [9216, 3072],
+    "double_blocks.0.img_mlp.0.bias": [12288],
+    "double_blocks.0.img_mlp.0.weight": [12288, 3072],
+    "double_blocks.0.img_mlp.2.bias": [3072],
+    "double_blocks.0.img_mlp.2.weight": [3072, 12288],
+    "double_blocks.0.img_mod.lin.bias": [18432],
+    "double_blocks.0.img_mod.lin.weight": [18432, 3072],
+    "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.bias": [3072],
+    "double_blocks.0.txt_attn.proj.weight": [3072, 3072],
+    "double_blocks.0.txt_attn.qkv.bias": [9216],
+    "double_blocks.0.txt_attn.qkv.weight": [9216, 3072],
+    "double_blocks.0.txt_mlp.0.bias": [12288],
+    "double_blocks.0.txt_mlp.0.weight": [12288, 3072],
+    "double_blocks.0.txt_mlp.2.bias": [3072],
+    "double_blocks.0.txt_mlp.2.weight": [3072, 12288],
+    "double_blocks.0.txt_mod.lin.bias": [18432],
+    "double_blocks.0.txt_mod.lin.weight": [18432, 3072],
+    "double_blocks.1.img_attn.norm.key_norm.scale": [128],
+    "double_blocks.1.img_attn.norm.query_norm.scale": [128],
+    "double_blocks.1.img_attn.proj.bias": [3072],
+    "double_blocks.1.img_attn.proj.weight": [3072, 3072],
+    "double_blocks.1.img_attn.qkv.bias": [9216],
+    "double_blocks.1.img_attn.qkv.weight": [9216, 3072],
+    "double_blocks.1.img_mlp.0.bias": [12288],
+    "double_blocks.1.img_mlp.0.weight": [12288, 3072],
+    "double_blocks.1.img_mlp.2.bias": [3072],
+    "double_blocks.1.img_mlp.2.weight": [3072, 12288],
+    "double_blocks.1.img_mod.lin.bias": [18432],
+    "double_blocks.1.img_mod.lin.weight": [18432, 3072],
+    "double_blocks.1.txt_attn.norm.key_norm.scale": [128],
+    "double_blocks.1.txt_attn.norm.query_norm.scale": [128],
+    "double_blocks.1.txt_attn.proj.bias": [3072],
+    "double_blocks.1.txt_attn.proj.weight": [3072, 3072],
+    "double_blocks.1.txt_attn.qkv.bias": [9216],
+    "double_blocks.1.txt_attn.qkv.weight": [9216, 3072],
+    "double_blocks.1.txt_mlp.0.bias": [12288],
+    "double_blocks.1.txt_mlp.0.weight": [12288, 3072],
+    "double_blocks.1.txt_mlp.2.bias": [3072],
+    "double_blocks.1.txt_mlp.2.weight": [3072, 12288],
+    "double_blocks.1.txt_mod.lin.bias": [18432],
+    "double_blocks.1.txt_mod.lin.weight": [18432, 3072],
+    "guidance_in.in_layer.bias": [3072],
+    "guidance_in.in_layer.weight": [3072, 256],
+    "guidance_in.out_layer.bias": [3072],
+    "guidance_in.out_layer.weight": [3072, 3072],
+    "img_in.bias": [3072],
+    "img_in.weight": [3072, 64],
+    "input_hint_block.0.bias": [16],
+    "input_hint_block.0.weight": [16, 3, 3, 3],
+    "input_hint_block.10.bias": [16],
+    "input_hint_block.10.weight": [16, 16, 3, 3],
+    "input_hint_block.12.bias": [16],
+    "input_hint_block.12.weight": [16, 16, 3, 3],
+    "input_hint_block.14.bias": [16],
+    "input_hint_block.14.weight": [16, 16, 3, 3],
+    "input_hint_block.2.bias": [16],
+    "input_hint_block.2.weight": [16, 16, 3, 3],
+    "input_hint_block.4.bias": [16],
+    "input_hint_block.4.weight": [16, 16, 3, 3],
+    "input_hint_block.6.bias": [16],
+    "input_hint_block.6.weight": [16, 16, 3, 3],
+    "input_hint_block.8.bias": [16],
+    "input_hint_block.8.weight": [16, 16, 3, 3],
+    "pos_embed_input.bias": [3072],
+    "pos_embed_input.weight": [3072, 64],
+    "time_in.in_layer.bias": [3072],
+    "time_in.in_layer.weight": [3072, 256],
+    "time_in.out_layer.bias": [3072],
+    "time_in.out_layer.weight": [3072, 3072],
+    "txt_in.bias": [3072],
+    "txt_in.weight": [3072, 4096],
+    "vector_in.in_layer.bias": [3072],
+    "vector_in.in_layer.weight": [3072, 768],
+    "vector_in.out_layer.bias": [3072],
+    "vector_in.out_layer.weight": [3072, 3072],
+}

tests/backend/flux/dype/test_dype.py

@@ -0,0 +1,499 @@
+"""Tests for DyPE (Dynamic Position Extrapolation) module."""
+
+import torch
+
+from invokeai.backend.flux.dype.base import (
+    DyPEConfig,
+    compute_vision_yarn_freqs,
+    get_timestep_kappa,
+)
+from invokeai.backend.flux.dype.embed import DyPEEmbedND
+from invokeai.backend.flux.dype.presets import (
+    DYPE_PRESET_4K,
+    DYPE_PRESET_AREA,
+    DYPE_PRESET_AUTO,
+    DYPE_PRESET_MANUAL,
+    DYPE_PRESET_OFF,
+    DYPE_PRESETS,
+    get_dype_config_for_area,
+    get_dype_config_for_resolution,
+    get_dype_config_from_preset,
+)
+from invokeai.backend.flux.dype.rope import rope_dype
+from invokeai.backend.flux.extensions.dype_extension import DyPEExtension
+
+
+class TestDyPEConfig:
+    """Tests for DyPEConfig dataclass."""
+
+    def test_default_values(self):
+        config = DyPEConfig()
+        assert config.enable_dype is True
+        assert config.base_resolution == 1024
+        assert config.dype_scale == 2.0
+        assert config.dype_exponent == 2.0
+        assert config.dype_start_sigma == 1.0
+
+    def test_custom_values(self):
+        config = DyPEConfig(
+            enable_dype=False,
+            base_resolution=512,
+            dype_scale=4.0,
+            dype_exponent=3.0,
+            dype_start_sigma=0.5,
+        )
+        assert config.enable_dype is False
+        assert config.base_resolution == 512
+        assert config.dype_scale == 4.0
+
+
+class TestDyPEExtension:
+    """Tests for DyPE extension helpers."""
+
+    def test_resolve_step_sigma_prefers_scheduler_sigmas_tensor(self):
+        sigma = DyPEExtension.resolve_step_sigma(
+            fallback_sigma=0.42,
+            step_index=1,
+            scheduler_sigmas=torch.tensor([1.0, 0.75, 0.5]),
+        )
+        assert sigma == 0.75
+
+    def test_resolve_step_sigma_falls_back_without_scheduler_sigmas(self):
+        sigma = DyPEExtension.resolve_step_sigma(
+            fallback_sigma=0.42,
+            step_index=1,
+            scheduler_sigmas=None,
+        )
+        assert sigma == 0.42
+
+
+class TestKappa:
+    """Tests for the DyPE timestep scheduler."""
+
+    def test_get_timestep_kappa_clamps_to_zero_without_scale(self):
+        assert (
+            get_timestep_kappa(
+                current_sigma=0.5,
+                dype_scale=0.0,
+                dype_exponent=2.0,
+                dype_start_sigma=1.0,
+            )
+            == 0.0
+        )
+
+    def test_get_timestep_kappa_is_stronger_early(self):
+        early_kappa = get_timestep_kappa(
+            current_sigma=1.0,
+            dype_scale=2.0,
+            dype_exponent=2.0,
+            dype_start_sigma=1.0,
+        )
+        late_kappa = get_timestep_kappa(
+            current_sigma=0.1,
+            dype_scale=2.0,
+            dype_exponent=2.0,
+            dype_start_sigma=1.0,
+        )
+
+        assert early_kappa == 2.0
+        assert late_kappa < early_kappa
+
+    def test_get_timestep_kappa_clamps_above_start_sigma(self):
+        kappa = get_timestep_kappa(
+            current_sigma=2.0,
+            dype_scale=2.0,
+            dype_exponent=2.0,
+            dype_start_sigma=1.0,
+        )
+        assert kappa == 2.0
+
+
+class TestRopeDype:
+    """Tests for DyPE-enhanced RoPE function."""
+
+    def test_rope_dype_shape(self):
+        """Test that rope_dype returns correct shape."""
+        pos = torch.zeros(1, 64)
+        dim = 64
+        theta = 10000
+
+        config = DyPEConfig()
+        result = rope_dype(
+            pos=pos,
+            dim=dim,
+            theta=theta,
+            current_sigma=0.5,
+            target_height=2048,
+            target_width=2048,
+            dype_config=config,
+        )
+
+        # Shape should be (batch, seq_len, dim/2, 2, 2)
+        assert result.shape == (1, 64, dim // 2, 2, 2)
+
+    def test_rope_dype_no_scaling(self):
+        """When target is same as base, output should match base rope."""
+        pos = torch.arange(16).unsqueeze(0).float()
+        dim = 32
+        theta = 10000
+
+        config = DyPEConfig(base_resolution=1024)
+
+        # No scaling needed
+        result_no_scale = rope_dype(
+            pos=pos,
+            dim=dim,
+            theta=theta,
+            current_sigma=0.5,
+            target_height=1024,
+            target_width=1024,
+            dype_config=config,
+        )
+
+        # With scaling
+        result_with_scale = rope_dype(
+            pos=pos,
+            dim=dim,
+            theta=theta,
+            current_sigma=0.5,
+            target_height=2048,
+            target_width=2048,
+            dype_config=config,
+        )
+
+        # Results should be different when scaling is applied
+        assert not torch.allclose(result_no_scale, result_with_scale)
+
+    def test_rope_dype_late_stage_moves_toward_base_rope(self):
+        """Late-stage DyPE should be closer to base RoPE than early-stage DyPE."""
+        pos = torch.arange(16).unsqueeze(0).float()
+        dim = 32
+        theta = 10000
+
+        config = DyPEConfig(base_resolution=1024)
+
+        base_result = rope_dype(
+            pos=pos,
+            dim=dim,
+            theta=theta,
+            current_sigma=1.0,
+            target_height=1024,
+            target_width=1024,
+            dype_config=config,
+        )
+        early_result = rope_dype(
+            pos=pos,
+            dim=dim,
+            theta=theta,
+            current_sigma=1.0,
+            target_height=2048,
+            target_width=2048,
+            dype_config=config,
+        )
+        late_result = rope_dype(
+            pos=pos,
+            dim=dim,
+            theta=theta,
+            current_sigma=0.05,
+            target_height=2048,
+            target_width=2048,
+            dype_config=config,
+        )
+
+        early_delta = torch.mean(torch.abs(early_result - base_result))
+        late_delta = torch.mean(torch.abs(late_result - base_result))
+
+        assert late_delta < early_delta
+
+
+class TestDyPEEmbedND:
+    """Tests for DyPEEmbedND module."""
+
+    def test_init(self):
+        """Test DyPEEmbedND initialization."""
+        config = DyPEConfig()
+        embedder = DyPEEmbedND(
+            dim=128,
+            theta=10000,
+            axes_dim=[16, 56, 56],
+            dype_config=config,
+        )
+
+        assert embedder.dim == 128
+        assert embedder.theta == 10000
+        assert embedder.axes_dim == [16, 56, 56]
+
+    def test_set_step_state(self):
+        """Test step state update."""
+        config = DyPEConfig()
+        embedder = DyPEEmbedND(
+            dim=128,
+            theta=10000,
+            axes_dim=[16, 56, 56],
+            dype_config=config,
+        )
+
+        embedder.set_step_state(sigma=0.5, height=2048, width=2048)
+
+        assert embedder._current_sigma == 0.5
+        assert embedder._target_height == 2048
+        assert embedder._target_width == 2048
+
+    def test_forward_shape(self):
+        """Test forward pass output shape."""
+        config = DyPEConfig()
+        embedder = DyPEEmbedND(
+            dim=128,
+            theta=10000,
+            axes_dim=[16, 56, 56],
+            dype_config=config,
+        )
+
+        # Create input ids tensor (batch=1, seq_len=64, n_axes=3)
+        ids = torch.zeros(1, 64, 3)
+
+        result = embedder(ids)
+
+        # Output should have shape (batch, 1, seq_len, dim)
+        # Actually the shape is (batch, 1, seq_len, dim/2, 2, 2) based on rope output
+        assert result.dim() == 6
+        assert result.shape[0] == 1  # batch
+        assert result.shape[1] == 1  # unsqueeze
+        assert result.shape[2] == 64  # seq_len
+
+
+class TestDyPEPresets:
+    """Tests for DyPE preset configurations."""
+
+    def test_preset_4k_exists(self):
+        """Test that 4K preset is defined."""
+        assert DYPE_PRESET_4K in DYPE_PRESETS
+
+    def test_get_dype_config_for_resolution_below_threshold(self):
+        """When resolution is below threshold, should return None."""
+        config = get_dype_config_for_resolution(
+            width=1024,
+            height=1024,
+            activation_threshold=1536,
+        )
+        assert config is None
+
+        config = get_dype_config_for_resolution(
+            width=1536,
+            height=1024,
+            activation_threshold=1536,
+        )
+        assert config is None
+
+    def test_get_dype_config_for_resolution_above_threshold(self):
+        """When resolution is above threshold, should return config."""
+        config = get_dype_config_for_resolution(
+            width=2048,
+            height=2048,
+            activation_threshold=1536,
+        )
+        assert config is not None
+        assert config.enable_dype is True
+
+    def test_get_dype_config_for_resolution_dynamic_scale(self):
+        """Higher resolution should result in higher dype_scale."""
+        config_2k = get_dype_config_for_resolution(
+            width=2048,
+            height=2048,
+            base_resolution=1024,
+            activation_threshold=1536,
+        )
+        config_4k = get_dype_config_for_resolution(
+            width=4096,
+            height=4096,
+            base_resolution=1024,
+            activation_threshold=1536,
+        )
+
+        assert config_2k is not None
+        assert config_4k is not None
+        assert config_4k.dype_scale > config_2k.dype_scale
+
+    def test_get_dype_config_for_area_below_threshold(self):
+        """When area is below threshold area, should return None."""
+        config = get_dype_config_for_area(
+            width=1024,
+            height=1024,
+        )
+        assert config is None
+
+    def test_get_dype_config_for_area_above_threshold(self):
+        """When area is above threshold area, should return config."""
+        config = get_dype_config_for_area(
+            width=2048,
+            height=1536,
+            base_resolution=1024,
+        )
+        assert config is not None
+        assert config.enable_dype is True
+
+    def test_get_dype_config_for_area_penalizes_extreme_aspect_ratios(self):
+        balanced_extreme = get_dype_config_for_area(
+            width=2304,
+            height=1152,
+            base_resolution=1024,
+        )
+        extreme = get_dype_config_for_area(
+            width=2304,
+            height=960,
+            base_resolution=1024,
+        )
+        balanced_same_area = get_dype_config_for_area(
+            width=2048,
+            height=1080,
+            base_resolution=1024,
+        )
+
+        assert balanced_extreme is not None
+        assert extreme is not None
+        assert balanced_same_area is not None
+        assert extreme.dype_scale < balanced_extreme.dype_scale
+        assert extreme.dype_scale < balanced_same_area.dype_scale
+
+    def test_get_dype_config_for_area_is_closer_to_auto_strength(self):
+        area = get_dype_config_for_area(
+            width=1728,
+            height=1152,
+            base_resolution=1024,
+        )
+        auto = get_dype_config_for_resolution(
+            width=1728,
+            height=1152,
+            base_resolution=1024,
+            activation_threshold=1536,
+        )
+
+        assert area is not None
+        assert auto is not None
+        assert area.dype_scale > auto.dype_scale * 0.9
+        assert area.dype_scale < auto.dype_scale * 1.1
+
+    def test_get_dype_config_for_area_uses_higher_exponent_than_old_curve(self):
+        config = get_dype_config_for_area(
+            width=1536,
+            height=1024,
+            base_resolution=1024,
+        )
+
+        assert config is not None
+        assert 1.25 <= config.dype_exponent <= 2.0
+
+    def test_get_dype_config_from_preset_area(self):
+        """Preset AREA should use area-based config."""
+        config = get_dype_config_from_preset(
+            preset=DYPE_PRESET_AREA,
+            width=2048,
+            height=1536,
+        )
+        assert config is not None
+        assert config.enable_dype is True
+
+    def test_get_dype_config_from_preset_off(self):
+        """Preset OFF should return None."""
+        config = get_dype_config_from_preset(
+            preset=DYPE_PRESET_OFF,
+            width=2048,
+            height=2048,
+        )
+        assert config is None
+
+    def test_get_dype_config_from_preset_auto(self):
+        """Preset AUTO should use resolution-based config."""
+        config = get_dype_config_from_preset(
+            preset=DYPE_PRESET_AUTO,
+            width=2048,
+            height=2048,
+        )
+        assert config is not None
+        assert config.enable_dype is True
+
+    def test_get_dype_config_from_preset_4k(self):
+        """Preset 4K should use 4K settings."""
+        config = get_dype_config_from_preset(
+            preset=DYPE_PRESET_4K,
+            width=3840,
+            height=2160,
+        )
+        assert config is not None
+        assert config.enable_dype is True
+
+    def test_get_dype_config_from_preset_manual_custom_overrides(self):
+        """Custom scale/exponent should override defaults only with 'manual' preset."""
+        config = get_dype_config_from_preset(
+            preset=DYPE_PRESET_MANUAL,
+            width=2048,
+            height=2048,
+            custom_scale=5.0,
+            custom_exponent=10.0,
+        )
+        assert config is not None
+        assert config.dype_scale == 5.0
+        assert config.dype_exponent == 10.0
+
+    def test_get_dype_config_from_preset_4k_ignores_custom(self):
+        """4K preset should ignore custom scale/exponent values."""
+        config = get_dype_config_from_preset(
+            preset=DYPE_PRESET_4K,
+            width=3840,
+            height=2160,
+            custom_scale=5.0,
+            custom_exponent=10.0,
+        )
+        assert config is not None
+        # Custom values should be ignored - preset values used instead
+        assert config.dype_scale == 2.0  # 4K preset default
+        assert config.dype_exponent == 2.0  # 4K preset default
+
+
+class TestFrequencyComputation:
+    """Tests for frequency computation functions."""
+
+    def test_compute_vision_yarn_freqs_shape(self):
+        """Test vision_yarn frequency computation shape."""
+        pos = torch.arange(16).unsqueeze(0).float()
+        config = DyPEConfig()
+
+        cos, sin = compute_vision_yarn_freqs(
+            pos=pos,
+            dim=32,
+            theta=10000,
+            scale_h=2.0,
+            scale_w=2.0,
+            current_sigma=0.5,
+            dype_config=config,
+        )
+
+        assert cos.shape == sin.shape
+        assert cos.shape[0] == 1  # batch
+        assert cos.shape[1] == 16  # seq_len
+
+    def test_compute_vision_yarn_freqs_reverts_to_base_rope_at_zero_sigma(self):
+        pos = torch.arange(16).unsqueeze(0).float()
+        config = DyPEConfig()
+
+        dy_cos, dy_sin = compute_vision_yarn_freqs(
+            pos=pos,
+            dim=32,
+            theta=10000,
+            scale_h=2.0,
+            scale_w=2.0,
+            current_sigma=0.0,
+            dype_config=config,
+        )
+        base_cos, base_sin = compute_vision_yarn_freqs(
+            pos=pos,
+            dim=32,
+            theta=10000,
+            scale_h=1.0,
+            scale_w=1.0,
+            current_sigma=0.0,
+            dype_config=config,
+        )
+
+        assert torch.allclose(dy_cos, base_cos)
+        assert torch.allclose(dy_sin, base_sin)

tests/backend/flux/ip_adapter/test_xlabs_ip_adapter_flux.py

@@ -0,0 +1,78 @@
+import sys
+
+import accelerate
+import pytest
+import torch
+
+from invokeai.backend.flux.ip_adapter.state_dict_utils import (
+    infer_xlabs_ip_adapter_params_from_state_dict,
+    is_state_dict_xlabs_ip_adapter,
+)
+from invokeai.backend.flux.ip_adapter.xlabs_ip_adapter_flux import (
+    XlabsIpAdapterFlux,
+    XlabsIpAdapterParams,
+)
+from tests.backend.flux.ip_adapter.xlabs_flux_ip_adapter_state_dict import xlabs_flux_ip_adapter_sd_shapes
+from tests.backend.flux.ip_adapter.xlabs_flux_ip_adapter_v2_state_dict import xlabs_flux_ip_adapter_v2_sd_shapes
+
+
+@pytest.mark.parametrize("sd_shapes", [xlabs_flux_ip_adapter_sd_shapes, xlabs_flux_ip_adapter_v2_sd_shapes])
+def test_is_state_dict_xlabs_ip_adapter(sd_shapes: dict[str, list[int]]):
+    # Construct a dummy state_dict.
+    sd = dict.fromkeys(sd_shapes)
+
+    assert is_state_dict_xlabs_ip_adapter(sd)
+
+
+@pytest.mark.skipif(sys.platform == "darwin", reason="Skipping on macOS")
+@pytest.mark.parametrize(
+    ["sd_shapes", "expected_params"],
+    [
+        (
+            xlabs_flux_ip_adapter_sd_shapes,
+            XlabsIpAdapterParams(
+                num_double_blocks=19,
+                context_dim=4096,
+                hidden_dim=3072,
+                clip_embeddings_dim=768,
+                clip_extra_context_tokens=4,
+            ),
+        ),
+        (
+            xlabs_flux_ip_adapter_v2_sd_shapes,
+            XlabsIpAdapterParams(
+                num_double_blocks=19,
+                context_dim=4096,
+                hidden_dim=3072,
+                clip_embeddings_dim=768,
+                clip_extra_context_tokens=16,
+            ),
+        ),
+    ],
+)
+def test_infer_xlabs_ip_adapter_params_from_state_dict(
+    sd_shapes: dict[str, list[int]], expected_params: XlabsIpAdapterParams
+):
+    # Construct a dummy state_dict with tensors of the correct shape on the meta device.
+    with torch.device("meta"):
+        sd = {k: torch.zeros(v) for k, v in sd_shapes.items()}
+
+    params = infer_xlabs_ip_adapter_params_from_state_dict(sd)
+    assert params == expected_params
+
+
+@pytest.mark.skipif(sys.platform == "darwin", reason="Skipping on macOS")
+@pytest.mark.parametrize("sd_shapes", [xlabs_flux_ip_adapter_sd_shapes, xlabs_flux_ip_adapter_v2_sd_shapes])
+def test_initialize_xlabs_ip_adapter_flux_from_state_dict(sd_shapes: dict[str, list[int]]):
+    # Construct a dummy state_dict with tensors of the correct shape on the meta device.
+    with torch.device("meta"):
+        sd = {k: torch.zeros(v) for k, v in sd_shapes.items()}
+
+    # Initialize the XLabs IP-Adapter from the state_dict.
+    params = infer_xlabs_ip_adapter_params_from_state_dict(sd)
+
+    with accelerate.init_empty_weights():
+        model = XlabsIpAdapterFlux(params=params)
+
+    # Smoke test state_dict loading.
+    model.load_xlabs_state_dict(sd)

tests/backend/flux/ip_adapter/xlabs_flux_ip_adapter_state_dict.py

@@ -0,0 +1,85 @@
+# State dict keys and shapes for an XLabs FLUX IP-Adapter model. Intended to be used for unit tests.
+# These keys were extracted from:
+# https://huggingface.co/XLabs-AI/flux-ip-adapter/resolve/main/ip_adapter.safetensors
+xlabs_flux_ip_adapter_sd_shapes = {
+    "double_blocks.0.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.0.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.0.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.1.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.1.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.1.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.1.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.10.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.10.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.10.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.10.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.11.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.11.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.11.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.11.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.12.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.12.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.12.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.12.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.13.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.13.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.13.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.13.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.14.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.14.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.14.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.14.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.15.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.15.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.15.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.15.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.16.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.16.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.16.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.16.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.17.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.17.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.17.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.17.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.18.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.18.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.18.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.18.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.2.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.2.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.2.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.2.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.3.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.3.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.3.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.3.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.4.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.4.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.4.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.4.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.5.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.5.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.5.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.5.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.6.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.6.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.6.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.6.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.7.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.7.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.7.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.7.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.8.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.8.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.8.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.8.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.9.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.9.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.9.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.9.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "ip_adapter_proj_model.norm.bias": [4096],
+    "ip_adapter_proj_model.norm.weight": [4096],
+    "ip_adapter_proj_model.proj.bias": [16384],
+    "ip_adapter_proj_model.proj.weight": [16384, 768],
+}

tests/backend/flux/ip_adapter/xlabs_flux_ip_adapter_v2_state_dict.py

@@ -0,0 +1,85 @@
+# State dict keys and shapes for an XLabs FLUX IP-Adapter V2 model. Intended to be used for unit tests.
+# These keys were extracted from:
+# https://huggingface.co/XLabs-AI/flux-ip-adapter-v2/blob/main/ip_adapter.safetensors
+xlabs_flux_ip_adapter_v2_sd_shapes = {
+    "double_blocks.0.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.0.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.0.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.0.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.1.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.1.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.1.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.1.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.10.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.10.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.10.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.10.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.11.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.11.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.11.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.11.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.12.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.12.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.12.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.12.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.13.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.13.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.13.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.13.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.14.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.14.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.14.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.14.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.15.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.15.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.15.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.15.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.16.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.16.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.16.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.16.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.17.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.17.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.17.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.17.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.18.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.18.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.18.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.18.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.2.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.2.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.2.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.2.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.3.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.3.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.3.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.3.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.4.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.4.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.4.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.4.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.5.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.5.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.5.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.5.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.6.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.6.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.6.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.6.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.7.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.7.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.7.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.7.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.8.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.8.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.8.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.8.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "double_blocks.9.processor.ip_adapter_double_stream_k_proj.bias": [3072],
+    "double_blocks.9.processor.ip_adapter_double_stream_k_proj.weight": [3072, 4096],
+    "double_blocks.9.processor.ip_adapter_double_stream_v_proj.bias": [3072],
+    "double_blocks.9.processor.ip_adapter_double_stream_v_proj.weight": [3072, 4096],
+    "ip_adapter_proj_model.norm.bias": [4096],
+    "ip_adapter_proj_model.norm.weight": [4096],
+    "ip_adapter_proj_model.proj.bias": [65536],
+    "ip_adapter_proj_model.proj.weight": [65536, 768],
+}

tests/backend/flux/modules/test_conditioner.py

@@ -0,0 +1,45 @@
+import torch
+
+from invokeai.backend.flux.modules.conditioner import HFEncoder
+
+
+class FakeTokenizer:
+    def __call__(
+        self,
+        text,
+        truncation,
+        max_length,
+        return_length,
+        return_overflowing_tokens,
+        padding,
+        return_tensors,
+    ):
+        del text, truncation, max_length, return_length, return_overflowing_tokens, padding, return_tensors
+        return {"input_ids": torch.tensor([[1, 2, 3]], dtype=torch.long)}
+
+
+class FakeEncoderOutput(dict):
+    pass
+
+
+class FakePartiallyLoadedEncoder(torch.nn.Module):
+    def __init__(self, effective_device: torch.device):
+        super().__init__()
+        self.register_parameter("cpu_param", torch.nn.Parameter(torch.ones(1)))
+        self.register_buffer("active_buffer", torch.ones(1, device=effective_device))
+        self.forward_input_device: torch.device | None = None
+
+    def forward(self, input_ids: torch.Tensor, attention_mask=None, output_hidden_states: bool = False):
+        del attention_mask, output_hidden_states
+        self.forward_input_device = input_ids.device
+        return FakeEncoderOutput(pooler_output=torch.ones((1, 4), dtype=torch.float32))
+
+
+def test_hf_encoder_uses_effective_device_for_partially_loaded_models():
+    effective_device = torch.device("meta")
+    encoder = FakePartiallyLoadedEncoder(effective_device=effective_device)
+    hf_encoder = HFEncoder(encoder=encoder, tokenizer=FakeTokenizer(), is_clip=True, max_length=77)
+
+    hf_encoder(["test prompt"])
+
+    assert encoder.forward_input_device == effective_device

tests/backend/flux/redux/test_flux_redux_state_dict_utils.py

@@ -0,0 +1,28 @@
+# The state dict keys and shapes for a FLUX Redux model.
+# Model source: https://huggingface.co/black-forest-labs/FLUX.1-Redux-dev/blob/1282f955f706b5240161278f2ef261d2a29ad649/flux1-redux-dev.safetensors
+# The keys and shapes were extracted with extract_sd_keys_and_shapes.py.
+import torch
+
+from invokeai.backend.flux.redux.flux_redux_state_dict_utils import is_state_dict_likely_flux_redux
+from tests.backend.patches.lora_conversions.lora_state_dicts.utils import keys_to_mock_state_dict
+
+flux_redux_keys_and_shapes = {
+    "redux_down.bias": [4096],
+    "redux_down.weight": [4096, 12288],
+    "redux_up.bias": [12288],
+    "redux_up.weight": [12288, 1152],
+}
+
+
+def test_is_state_dict_likely_flux_redux_true():
+    # Expand flux_redux_keys_and_shapes to a mock state dict.
+    sd = keys_to_mock_state_dict(flux_redux_keys_and_shapes)
+    assert is_state_dict_likely_flux_redux(sd)
+
+
+def test_is_state_dict_likely_flux_redux_extra_key():
+    # Expand flux_redux_keys_and_shapes to a mock state dict.
+    sd = keys_to_mock_state_dict(flux_redux_keys_and_shapes)
+    # Add an extra key to the state dict.
+    sd["extra_key"] = torch.rand(1)
+    assert not is_state_dict_likely_flux_redux(sd)

tests/backend/flux/test_anima_schedulers.py

@@ -0,0 +1,319 @@
+"""Tests for Anima scheduler registry."""
+
+import typing
+
+import pytest
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from invokeai.backend.flux.schedulers import (
+    ANIMA_SCHEDULER_LABELS,
+    ANIMA_SCHEDULER_MAP,
+    ANIMA_SCHEDULER_NAME_VALUES,
+)
+
+
+def test_anima_scheduler_map_entries_are_class_kwargs_tuples():
+    """Every entry must be (SchedulerClass, kwargs_dict)."""
+    for name, entry in ANIMA_SCHEDULER_MAP.items():
+        assert isinstance(entry, tuple), f"{name} is not a tuple"
+        assert len(entry) == 2, f"{name} tuple has wrong arity"
+        cls, kwargs = entry
+        assert isinstance(cls, type) and issubclass(cls, SchedulerMixin), (
+            f"{name} first element is not a SchedulerMixin subclass"
+        )
+        assert isinstance(kwargs, dict), f"{name} second element is not a dict"
+
+
+def test_anima_scheduler_map_entries_can_be_constructed():
+    """Every entry must construct cleanly by splatting its kwargs."""
+    for name, (cls, kwargs) in ANIMA_SCHEDULER_MAP.items():
+        scheduler = cls(num_train_timesteps=1000, **kwargs)
+        assert isinstance(scheduler, SchedulerMixin), f"{name} did not produce a SchedulerMixin"
+
+
+def test_anima_scheduler_labels_cover_every_map_key():
+    for name in ANIMA_SCHEDULER_MAP.keys():
+        assert name in ANIMA_SCHEDULER_LABELS, f"{name} has no label"
+
+
+def test_anima_scheduler_map_includes_new_dpmpp_entries():
+    assert "dpmpp_2m" in ANIMA_SCHEDULER_MAP
+    assert "dpmpp_2m_sde" in ANIMA_SCHEDULER_MAP
+
+
+def test_anima_dpmpp_2m_uses_flow_prediction():
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+
+    cls, kwargs = ANIMA_SCHEDULER_MAP["dpmpp_2m"]
+    assert kwargs["prediction_type"] == "flow_prediction"
+    assert kwargs["use_flow_sigmas"] is True
+    assert kwargs["flow_shift"] == ANIMA_SHIFT
+    assert kwargs["solver_order"] == 2
+    assert "algorithm_type" not in kwargs  # deterministic, default algorithm
+
+
+def test_anima_dpmpp_2m_sde_uses_sde_algorithm():
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+
+    cls, kwargs = ANIMA_SCHEDULER_MAP["dpmpp_2m_sde"]
+    assert kwargs["prediction_type"] == "flow_prediction"
+    assert kwargs["use_flow_sigmas"] is True
+    assert kwargs["flow_shift"] == ANIMA_SHIFT
+    assert kwargs["algorithm_type"] == "sde-dpmsolver++"
+    assert kwargs["solver_order"] == 2
+
+
+def test_anima_dpmpp_2m_produces_anima_compatible_sigma_schedule():
+    """The DPM++ 2M scheduler, when run through the same dispatch logic as
+    anima_denoise._run_diffusion, must produce a sigma schedule equivalent
+    to Anima's reference schedule (loglinear_timestep_shift with shift=3.0).
+
+    On diffusers 0.35.1, DPMSolverMultistepScheduler.set_timesteps does not
+    accept `sigmas=`, so the runtime falls back to num_inference_steps and
+    relies on the scheduler's internal flow_shift=3.0 to compute equivalent
+    sigmas. This test verifies that equivalence end-to-end.
+    """
+    import inspect
+
+    from invokeai.app.invocations.anima_denoise import loglinear_timestep_shift
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+
+    num_steps = 10
+
+    # Reference: Anima's own pre-shifted sigma schedule.
+    anima_sigmas = [loglinear_timestep_shift(ANIMA_SHIFT, 1.0 - i / num_steps) for i in range(num_steps + 1)]
+
+    cls, kwargs = ANIMA_SCHEDULER_MAP["dpmpp_2m"]
+    scheduler = cls(num_train_timesteps=1000, **kwargs)
+
+    # Mirror anima_denoise.py:502-506 dispatch.
+    sig = inspect.signature(scheduler.set_timesteps)
+    if "sigmas" in sig.parameters:
+        scheduler.set_timesteps(sigmas=anima_sigmas, device="cpu")
+    else:
+        scheduler.set_timesteps(num_inference_steps=num_steps, device="cpu")
+
+    diffusers_sigmas = [float(s) for s in scheduler.sigmas[: len(anima_sigmas)]]
+    max_diff = max(abs(a - b) for a, b in zip(anima_sigmas, diffusers_sigmas, strict=True))
+    assert max_diff < 1e-3, f"DPM++ 2M sigma schedule diverges from Anima reference (max abs diff = {max_diff:.6f})"
+
+
+def test_anima_dpmpp_2m_with_denoising_start_honors_clipped_schedule():
+    """DPM++ img2img: the set_begin_index path must start at the correct sigma.
+
+    When DPMSolverMultistepScheduler doesn't accept sigmas=, anima_denoise falls back to
+    set_timesteps(num_inference_steps=full_steps) + set_begin_index(start_idx).  The
+    effective first sigma must match the clipped Anima reference schedule within 1e-3.
+    """
+    import inspect
+
+    from invokeai.app.invocations.anima_denoise import loglinear_timestep_shift
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+
+    num_steps = 30
+    denoising_start = 0.5
+    start_idx = int(denoising_start * num_steps)  # mirrors anima_denoise clipping math
+
+    full_sigmas = [loglinear_timestep_shift(ANIMA_SHIFT, 1.0 - i / num_steps) for i in range(num_steps + 1)]
+    expected_first_sigma = full_sigmas[start_idx]
+
+    cls, kwargs = ANIMA_SCHEDULER_MAP["dpmpp_2m"]
+    scheduler = cls(num_train_timesteps=1000, **kwargs)
+    sig = inspect.signature(scheduler.set_timesteps)
+
+    if "sigmas" in sig.parameters:
+        # Future diffusers: sigmas= supported, clipped schedule passed directly.
+        scheduler.set_timesteps(sigmas=full_sigmas[start_idx:], device="cpu")
+        actual_first_sigma = float(scheduler.sigmas[0])
+    else:
+        # Current diffusers: use set_begin_index on the full schedule.
+        scheduler.set_timesteps(num_inference_steps=num_steps, device="cpu")
+        scheduler.set_begin_index(start_idx)
+        actual_first_sigma = float(scheduler.sigmas[start_idx])
+
+    assert abs(actual_first_sigma - expected_first_sigma) < 1e-3, (
+        f"DPM++ first sigma with denoising_start=0.5: got {actual_first_sigma:.6f}, expected {expected_first_sigma:.6f}"
+    )
+
+
+def test_anima_set_begin_index_path_step_count_with_denoising_end():
+    """set_begin_index fallback must honour denoising_end, not just denoising_start.
+
+    Regression test: the old formula (len(timesteps) - begin_index) ignored denoising_end
+    and ran past it. For steps=30, denoising_start=0.2, denoising_end=0.8 the correct
+    step count is 18, not 24.
+    """
+    import inspect
+
+    from invokeai.app.invocations.anima_denoise import loglinear_timestep_shift
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+
+    num_steps = 30
+    denoising_start = 0.2
+    denoising_end = 0.8
+
+    # Reference step count from the Euler path (clipped sigmas).
+    full_sigmas = [loglinear_timestep_shift(ANIMA_SHIFT, 1.0 - i / num_steps) for i in range(num_steps + 1)]
+    total_sigmas = len(full_sigmas)
+    start_idx = int(denoising_start * (total_sigmas - 1))
+    end_idx = int(denoising_end * (total_sigmas - 1)) + 1
+    expected_steps = (end_idx - start_idx) - 1  # 18
+
+    cls, kwargs = ANIMA_SCHEDULER_MAP["dpmpp_2m"]
+    scheduler = cls(num_train_timesteps=1000, **kwargs)
+    sig = inspect.signature(scheduler.set_timesteps)
+
+    scheduler_begin_index = int(denoising_start * num_steps)
+    if "sigmas" in sig.parameters:
+        clipped = full_sigmas[start_idx:end_idx]
+        scheduler.set_timesteps(sigmas=clipped, device="cpu")
+        num_scheduler_steps = len(scheduler.timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps=num_steps, device="cpu")
+        scheduler.set_begin_index(scheduler_begin_index)
+        num_scheduler_steps = int(denoising_end * num_steps) - scheduler_begin_index
+
+    assert num_scheduler_steps == expected_steps, (
+        f"DPM++ scheduler step count with denoising_start={denoising_start}, "
+        f"denoising_end={denoising_end}: got {num_scheduler_steps}, expected {expected_steps}"
+    )
+
+
+@pytest.mark.parametrize(
+    ["denoising_start", "denoising_end", "steps"],
+    [
+        (0.2, 0.8, 30),  # mid-range: 18 logical steps → 36 doubled calls
+        (0.0, 0.8, 30),  # start-only clip: 24 logical steps → 48 doubled calls
+        (0.2, 1.0, 30),  # end=1.0: clamp kicks in (last step first-order only)
+        (0.5, 0.75, 20),  # different step count
+    ],
+)
+def test_anima_heun_set_begin_index_path_begin_index_and_step_count(
+    denoising_start: float, denoising_end: float, steps: int
+):
+    """Heun img2img: set_begin_index must use doubled-array index and step count must
+    account for Heun's 2N-1 timestep structure.
+
+    Logical step k maps to doubled-array begin index 2k.  For a range [k_start, k_end)
+    the total calls is 2*(k_end-k_start), clamped to len(timesteps)-begin_index so that
+    denoising_end=1.0 correctly gets the 2N-1 (not 2N) count.
+    """
+    from diffusers import FlowMatchHeunDiscreteScheduler
+
+    k_start = int(denoising_start * steps)
+    k_end = int(denoising_end * steps)
+
+    scheduler = FlowMatchHeunDiscreteScheduler(num_train_timesteps=1000, shift=1.0)
+    scheduler.set_timesteps(num_inference_steps=steps, device="cpu")
+
+    expected_begin_index = 2 * k_start
+    expected_steps = min(2 * (k_end - k_start), len(scheduler.timesteps) - expected_begin_index)
+
+    # Verify the doubled structure: len(timesteps) == 2*steps - 1
+    assert len(scheduler.timesteps) == 2 * steps - 1, (
+        f"Heun timesteps length: expected {2 * steps - 1}, got {len(scheduler.timesteps)}"
+    )
+
+    # The fixed code's begin index must map logical step to doubled-array space.
+    assert expected_begin_index == 2 * k_start
+
+    # For mid-range (denoising_end < 1): all steps in range have first + second order.
+    if denoising_end < 1.0:
+        assert expected_steps == 2 * (k_end - k_start), (
+            f"mid-range step count: expected {2 * (k_end - k_start)}, got {expected_steps}"
+        )
+
+    # For denoising_end=1.0: last step is first-order only → clamped to 2N-1-begin.
+    if denoising_end == 1.0 and k_start > 0:
+        full_from_begin = len(scheduler.timesteps) - expected_begin_index
+        assert expected_steps == full_from_begin
+
+    # Bounds check: begin_index + num_steps must not exceed len(timesteps).
+    assert expected_begin_index + expected_steps <= len(scheduler.timesteps), (
+        f"step range [{expected_begin_index}, {expected_begin_index + expected_steps}) "
+        f"exceeds timesteps length {len(scheduler.timesteps)}"
+    )
+
+    # Sigma sanity: the sigma at the doubled begin index must equal the sigma at logical k_start.
+    # sigmas has 2N entries; sigmas[2k] == s_k for all k.
+    assert len(scheduler.sigmas) == 2 * steps
+    sigma_at_begin = scheduler.sigmas[expected_begin_index].item()
+    sigma_at_logical_k = scheduler.sigmas[2 * k_start].item()
+    assert abs(sigma_at_begin - sigma_at_logical_k) < 1e-6
+
+
+def test_anima_literal_covers_every_map_key():
+    """Catch the silent failure mode where a new entry lands in the map but
+    the Literal isn't updated — Pydantic validation would still accept it
+    via runtime introspection but type-check tooling would not."""
+    literal_values = set(typing.get_args(ANIMA_SCHEDULER_NAME_VALUES))
+    for name in ANIMA_SCHEDULER_MAP:
+        assert name in literal_values, f"{name} is in the map but missing from the Literal"
+
+
+def test_anima_scheduler_literal_includes_er_sde():
+    """er_sde must appear in the literal type, have a label, and be
+    registered in ANIMA_SCHEDULER_MAP for dispatch through the universal
+    scheduler path."""
+    literal_args = typing.get_args(ANIMA_SCHEDULER_NAME_VALUES)
+    assert "er_sde" in literal_args
+    assert "er_sde" in ANIMA_SCHEDULER_LABELS
+    assert ANIMA_SCHEDULER_LABELS["er_sde"] == "ER-SDE"
+    assert "er_sde" in ANIMA_SCHEDULER_MAP
+
+
+def test_anima_heun_uses_anima_shift_for_internal_schedule():
+    """Heun does NOT accept set_timesteps(sigmas=...) so it always builds its own internal
+    schedule. With shift=1.0 (the previous setting), that schedule was linear and gave the
+    wrong noise levels for img2img — Heun's sigmas[2*k_start] would be ~0.48 when Anima's
+    reference at user step k_start (denoising_start=0.5) is ~0.75. The model would receive
+    a timestep matching neither the latents nor its training distribution.
+
+    Fix: give Heun shift=ANIMA_SHIFT so its internal schedule approximates Anima's reference.
+    """
+    from invokeai.app.invocations.anima_denoise import loglinear_timestep_shift
+
+    cls, kwargs = ANIMA_SCHEDULER_MAP["heun"]
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+
+    assert kwargs["shift"] == ANIMA_SHIFT, (
+        f"Heun must use shift={ANIMA_SHIFT} (Anima's loglinear shift) since it doesn't accept "
+        f"sigmas=; got shift={kwargs['shift']}"
+    )
+
+    # Verify the schedule approximates Anima's reference for the bulk of user steps.
+    # The two formulas diverge at the tail (Heun uses linspace(1, T, N+1), Anima uses
+    # 1 - i/N), so we tolerate up to 5% absolute. The previous shift=1.0 bug gave 25-40%+
+    # divergence at mid-schedule, so any reasonable tolerance catches that regression.
+    steps = 30
+    anima_ref = [loglinear_timestep_shift(ANIMA_SHIFT, 1.0 - i / steps) for i in range(steps + 1)]
+    s = cls(num_train_timesteps=1000, **kwargs)
+    s.set_timesteps(num_inference_steps=steps, device="cpu")
+
+    # Heun's sigmas array has 2*N entries; sigmas[2*k] is the noise level at user step k.
+    for k in (0, 5, 10, 15):
+        heun_sigma = s.sigmas[2 * k].item()
+        ref = anima_ref[k]
+        assert abs(heun_sigma - ref) < 0.05, (
+            f"Heun internal sigma at user step {k} ({heun_sigma:.4f}) diverges from "
+            f"Anima reference ({ref:.4f}) by more than 5% — likely a shift kwarg regression"
+        )
+
+
+def test_anima_scheduler_map_er_sde_entry():
+    """ANIMA_SCHEDULER_MAP['er_sde'] must map to ERSDEScheduler with rectified-flow kwargs.
+
+    This is the wiring that lets Anima dispatch er_sde through the universal scheduler
+    path (replacing the legacy elif is_er_sde: branch in anima_denoise.py).
+    """
+    from invokeai.backend.flux.schedulers import ANIMA_SHIFT
+    from invokeai.backend.rectified_flow.er_sde_scheduler import ERSDEScheduler
+
+    assert "er_sde" in ANIMA_SCHEDULER_MAP, "er_sde must be in ANIMA_SCHEDULER_MAP"
+    cls, kwargs = ANIMA_SCHEDULER_MAP["er_sde"]
+    assert cls is ERSDEScheduler
+    assert kwargs["use_flow_sigmas"] is True
+    assert kwargs["prediction_type"] == "flow_prediction"
+    assert kwargs["solver_order"] == 3
+    assert kwargs["stochastic"] is True
+    assert kwargs["flow_shift"] == ANIMA_SHIFT

tests/backend/flux/test_denoise.py

@@ -0,0 +1,249 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from invokeai.backend.flux.denoise import denoise
+from invokeai.backend.flux.schedulers import FLUX_SCHEDULER_MAP
+
+
+class _FakeFluxModel:
+    def __call__(
+        self,
+        img: torch.Tensor,
+        img_ids: torch.Tensor,
+        txt: torch.Tensor,
+        txt_ids: torch.Tensor,
+        y: torch.Tensor,
+        timesteps: torch.Tensor,
+        guidance: torch.Tensor,
+        timestep_index: int,
+        total_num_timesteps: int,
+        controlnet_double_block_residuals: list[torch.Tensor] | None,
+        controlnet_single_block_residuals: list[torch.Tensor] | None,
+        ip_adapter_extensions: list[object],
+        regional_prompting_extension: object,
+    ) -> torch.Tensor:
+        return torch.zeros_like(img)
+
+
+class _FakeDyPEExtension:
+    def __init__(self) -> None:
+        self.sigmas: list[float] = []
+
+    def patch_model(self, model: object) -> tuple[object, None]:
+        return object(), None
+
+    def update_step_state(self, embedder: object, sigma: float) -> None:
+        self.sigmas.append(sigma)
+
+
+class _FakeScheduler:
+    def __init__(self) -> None:
+        self.config = SimpleNamespace(num_train_timesteps=1000)
+        self.timesteps = torch.tensor([], dtype=torch.float32)
+        self.sigmas = torch.tensor([], dtype=torch.float32)
+
+    def set_timesteps(self, sigmas: list[float], device: torch.device) -> None:
+        del device
+        self.sigmas = torch.tensor(sigmas, dtype=torch.float32)
+        self.timesteps = torch.tensor([900.0, 400.0], dtype=torch.float32)
+
+    def step(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor) -> SimpleNamespace:
+        del model_output, timestep
+        return SimpleNamespace(prev_sample=sample)
+
+
+class _FakeHeunScheduler:
+    def __init__(self) -> None:
+        self.config = SimpleNamespace(num_train_timesteps=1000)
+        self.timesteps = torch.tensor([], dtype=torch.float32)
+        self.sigmas = torch.tensor([], dtype=torch.float32)
+        self.state_in_first_order = True
+        self._step_index = 0
+
+    def set_timesteps(self, sigmas: list[float], device: torch.device) -> None:
+        del device
+        # Duplicate each user-facing step to mimic a second-order scheduler.
+        self.sigmas = torch.tensor([1.0, 1.0, 0.25, 0.25, 0.0], dtype=torch.float32)
+        self.timesteps = torch.tensor([900.0, 850.0, 400.0, 350.0], dtype=torch.float32)
+        self._step_index = 0
+        self.state_in_first_order = True
+
+    def step(self, model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor) -> SimpleNamespace:
+        del model_output, timestep
+        self._step_index += 1
+        self.state_in_first_order = self._step_index % 2 == 0
+        return SimpleNamespace(prev_sample=sample)
+
+
+class _FakePbar:
+    def update(self, value: int) -> None:
+        del value
+
+    def close(self) -> None:
+        return None
+
+
+def _fake_tqdm(iterable=None, **kwargs):
+    del kwargs
+    if iterable is None:
+        return _FakePbar()
+    return iterable
+
+
+def _build_regional_prompting_extension(batch_size: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        regional_text_conditioning=SimpleNamespace(
+            t5_embeddings=torch.zeros(batch_size, 1, 4),
+            t5_txt_ids=torch.zeros(batch_size, 1, 3),
+            clip_embeddings=torch.zeros(batch_size, 4),
+        )
+    )
+
+
+def test_denoise_euler_path_updates_dype_with_sigma(monkeypatch):
+    monkeypatch.setattr("invokeai.backend.flux.denoise.tqdm", _fake_tqdm)
+
+    model = _FakeFluxModel()
+    dype_extension = _FakeDyPEExtension()
+    img = torch.zeros(1, 2, 4)
+    img_ids = torch.zeros(1, 2, 3)
+    regional_prompting_extension = _build_regional_prompting_extension(batch_size=1)
+    callback_steps: list[int] = []
+
+    result = denoise(
+        model=model,
+        img=img,
+        img_ids=img_ids,
+        pos_regional_prompting_extension=regional_prompting_extension,
+        neg_regional_prompting_extension=None,
+        timesteps=[1.0, 0.5, 0.0],
+        step_callback=lambda state: callback_steps.append(state.step),
+        guidance=1.0,
+        cfg_scale=[1.0, 1.0],
+        inpaint_extension=None,
+        controlnet_extensions=[],
+        pos_ip_adapter_extensions=[],
+        neg_ip_adapter_extensions=[],
+        img_cond=None,
+        img_cond_seq=None,
+        img_cond_seq_ids=None,
+        dype_extension=dype_extension,
+        scheduler=None,
+    )
+
+    assert torch.equal(result, img)
+    assert dype_extension.sigmas == [1.0, 0.5]
+    assert callback_steps == [1, 2]
+
+
+def test_denoise_scheduler_path_prefers_scheduler_sigmas_for_dype(monkeypatch):
+    monkeypatch.setattr("invokeai.backend.flux.denoise.tqdm", _fake_tqdm)
+
+    model = _FakeFluxModel()
+    scheduler = _FakeScheduler()
+    dype_extension = _FakeDyPEExtension()
+    img = torch.zeros(1, 2, 4)
+    img_ids = torch.zeros(1, 2, 3)
+    regional_prompting_extension = _build_regional_prompting_extension(batch_size=1)
+
+    denoise(
+        model=model,
+        img=img,
+        img_ids=img_ids,
+        pos_regional_prompting_extension=regional_prompting_extension,
+        neg_regional_prompting_extension=None,
+        timesteps=[1.0, 0.25, 0.0],
+        step_callback=lambda state: None,
+        guidance=1.0,
+        cfg_scale=[1.0, 1.0],
+        inpaint_extension=None,
+        controlnet_extensions=[],
+        pos_ip_adapter_extensions=[],
+        neg_ip_adapter_extensions=[],
+        img_cond=None,
+        img_cond_seq=None,
+        img_cond_seq_ids=None,
+        dype_extension=dype_extension,
+        scheduler=scheduler,
+    )
+
+    # Scheduler timesteps normalize to [0.9, 0.4], so this asserts the scheduler
+    # sigma sequence is what DyPE actually consumes.
+    assert dype_extension.sigmas == [1.0, 0.25]
+
+
+def test_denoise_heun_scheduler_path_uses_internal_scheduler_sigmas(monkeypatch):
+    monkeypatch.setattr("invokeai.backend.flux.denoise.tqdm", _fake_tqdm)
+
+    model = _FakeFluxModel()
+    scheduler = _FakeHeunScheduler()
+    dype_extension = _FakeDyPEExtension()
+    img = torch.zeros(1, 2, 4)
+    img_ids = torch.zeros(1, 2, 3)
+    regional_prompting_extension = _build_regional_prompting_extension(batch_size=1)
+    callback_steps: list[int] = []
+
+    denoise(
+        model=model,
+        img=img,
+        img_ids=img_ids,
+        pos_regional_prompting_extension=regional_prompting_extension,
+        neg_regional_prompting_extension=None,
+        timesteps=[1.0, 0.25, 0.0],
+        step_callback=lambda state: callback_steps.append(state.step),
+        guidance=1.0,
+        cfg_scale=[1.0, 1.0],
+        inpaint_extension=None,
+        controlnet_extensions=[],
+        pos_ip_adapter_extensions=[],
+        neg_ip_adapter_extensions=[],
+        img_cond=None,
+        img_cond_seq=None,
+        img_cond_seq_ids=None,
+        dype_extension=dype_extension,
+        scheduler=scheduler,
+    )
+
+    assert dype_extension.sigmas == [1.0, 1.0, 0.25, 0.25]
+    assert callback_steps == [1, 2]
+
+
+@pytest.mark.parametrize("scheduler_name", sorted(FLUX_SCHEDULER_MAP))
+def test_denoise_real_flux_schedulers_update_dype_from_internal_sigma_schedule(monkeypatch, scheduler_name):
+    monkeypatch.setattr("invokeai.backend.flux.denoise.tqdm", _fake_tqdm)
+
+    model = _FakeFluxModel()
+    scheduler = FLUX_SCHEDULER_MAP[scheduler_name](num_train_timesteps=1000)
+    dype_extension = _FakeDyPEExtension()
+    img = torch.zeros(1, 2, 4)
+    img_ids = torch.zeros(1, 2, 3)
+    regional_prompting_extension = _build_regional_prompting_extension(batch_size=1)
+    callback_steps: list[int] = []
+
+    denoise(
+        model=model,
+        img=img,
+        img_ids=img_ids,
+        pos_regional_prompting_extension=regional_prompting_extension,
+        neg_regional_prompting_extension=None,
+        timesteps=[1.0, 0.25, 0.0],
+        step_callback=lambda state: callback_steps.append(state.step),
+        guidance=1.0,
+        cfg_scale=[1.0, 1.0],
+        inpaint_extension=None,
+        controlnet_extensions=[],
+        pos_ip_adapter_extensions=[],
+        neg_ip_adapter_extensions=[],
+        img_cond=None,
+        img_cond_seq=None,
+        img_cond_seq_ids=None,
+        dype_extension=dype_extension,
+        scheduler=scheduler,
+    )
+
+    assert dype_extension.sigmas
+    expected_sigmas = [float(sigma) for sigma in scheduler.sigmas[: len(dype_extension.sigmas)]]
+    assert dype_extension.sigmas == expected_sigmas
+    assert callback_steps

tests/backend/flux/test_sampling_utils.py

@@ -0,0 +1,76 @@
+import pytest
+import torch
+
+from invokeai.backend.flux.sampling_utils import clip_timestep_schedule, clip_timestep_schedule_fractional
+
+
+def float_lists_almost_equal(list1: list[float], list2: list[float], tol: float = 1e-6) -> bool:
+    return all(abs(a - b) < tol for a, b in zip(list1, list2, strict=True))
+
+
+@pytest.mark.parametrize(
+    ["denoising_start", "denoising_end", "expected_timesteps", "raises"],
+    [
+        (0.0, 1.0, [1.0, 0.75, 0.5, 0.25, 0.0], False),  # Default case.
+        (-0.1, 1.0, [], True),  # Negative denoising_start should raise.
+        (0.0, 1.1, [], True),  # denoising_end > 1 should raise.
+        (0.5, 0.0, [], True),  # denoising_start > denoising_end should raise.
+        (0.0, 0.0, [1.0], False),  # denoising_end == 0.
+        (1.0, 1.0, [0.0], False),  # denoising_start == 1.
+        (0.2, 0.8, [1.0, 0.75, 0.5, 0.25], False),  # Middle of the schedule.
+        # If we denoise from 0.0 to x, then from x to 1.0, it is important that denoise_end = x and denoise_start = x
+        # map to the same timestep. We test this first when x is equal to a timestep, then when it falls between two
+        # timesteps.
+        # x = 0.5
+        (0.0, 0.5, [1.0, 0.75, 0.5], False),
+        (0.5, 1.0, [0.5, 0.25, 0.0], False),
+        # x = 0.3
+        (0.0, 0.3, [1.0, 0.75], False),
+        (0.3, 1.0, [0.75, 0.5, 0.25, 0.0], False),
+    ],
+)
+def test_clip_timestep_schedule(
+    denoising_start: float, denoising_end: float, expected_timesteps: list[float], raises: bool
+):
+    timesteps = torch.linspace(1, 0, 5).tolist()
+    if raises:
+        with pytest.raises(AssertionError):
+            clip_timestep_schedule(timesteps, denoising_start, denoising_end)
+    else:
+        assert float_lists_almost_equal(
+            clip_timestep_schedule(timesteps, denoising_start, denoising_end), expected_timesteps
+        )
+
+
+@pytest.mark.parametrize(
+    ["denoising_start", "denoising_end", "expected_timesteps", "raises"],
+    [
+        (0.0, 1.0, [1.0, 0.75, 0.5, 0.25, 0.0], False),  # Default case.
+        (-0.1, 1.0, [], True),  # Negative denoising_start should raise.
+        (0.0, 1.1, [], True),  # denoising_end > 1 should raise.
+        (0.5, 0.0, [], True),  # denoising_start > denoising_end should raise.
+        (0.0, 0.0, [1.0], False),  # denoising_end == 0.
+        (1.0, 1.0, [0.0], False),  # denoising_start == 1.
+        (0.2, 0.8, [0.8, 0.75, 0.5, 0.25, 0.2], False),  # Middle of the schedule.
+        # If we denoise from 0.0 to x, then from x to 1.0, it is important that denoise_end = x and denoise_start = x
+        # map to the same timestep. We test this first when x is equal to a timestep, then when it falls between two
+        # timesteps.
+        # x = 0.5
+        (0.0, 0.5, [1.0, 0.75, 0.5], False),
+        (0.5, 1.0, [0.5, 0.25, 0.0], False),
+        # x = 0.3
+        (0.0, 0.3, [1.0, 0.75, 0.7], False),
+        (0.3, 1.0, [0.7, 0.5, 0.25, 0.0], False),
+    ],
+)
+def test_clip_timestep_schedule_fractional(
+    denoising_start: float, denoising_end: float, expected_timesteps: list[float], raises: bool
+):
+    timesteps = torch.linspace(1, 0, 5).tolist()
+    if raises:
+        with pytest.raises(AssertionError):
+            clip_timestep_schedule_fractional(timesteps, denoising_start, denoising_end)
+    else:
+        assert float_lists_almost_equal(
+            clip_timestep_schedule_fractional(timesteps, denoising_start, denoising_end), expected_timesteps
+        )

tests/backend/image_util/test_color_conversion.py

@@ -0,0 +1,620 @@
+import pytest
+import torch
+
+from invokeai.backend.image_util import color_conversion
+from invokeai.invocation_api import (
+    hsl_from_linear_srgb,
+    hsl_from_srgb,
+    lab_from_linear_srgb,
+    lab_from_srgb,
+    lab_from_xyz,
+    linear_srgb_from_hsl,
+    linear_srgb_from_lab,
+    linear_srgb_from_oklab,
+    linear_srgb_from_oklch,
+    linear_srgb_from_srgb,
+    linear_srgb_from_xyz,
+    okhsl_from_srgb,
+    okhsv_from_srgb,
+    oklab_from_linear_srgb,
+    oklab_from_oklch,
+    oklab_from_srgb,
+    oklab_from_xyz,
+    oklch_from_linear_srgb,
+    oklch_from_oklab,
+    oklch_from_srgb,
+    oklch_from_xyz,
+    srgb_from_hsl,
+    srgb_from_lab,
+    srgb_from_linear_srgb,
+    srgb_from_okhsl,
+    srgb_from_okhsv,
+    srgb_from_oklab,
+    srgb_from_oklch,
+    srgb_from_xyz,
+    xyz_d50_to_d65,
+    xyz_d65_to_d50,
+    xyz_from_lab,
+    xyz_from_linear_srgb,
+    xyz_from_oklab,
+    xyz_from_oklch,
+    xyz_from_srgb,
+)
+
+
+def test_srgb_oklab_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [1.0, 0.1]],
+            [[0.0, 1.0], [0.0, 0.6]],
+            [[0.0, 1.0], [0.0, 0.9]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = srgb_from_linear_srgb(linear_srgb_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(srgb))))
+
+    assert torch.allclose(round_tripped, srgb, atol=1e-5)
+
+
+def test_oklab_from_srgb_matches_explicit_conversion_path() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [1.0, 0.1]],
+            [[0.0, 1.0], [0.0, 0.6]],
+            [[0.0, 1.0], [0.0, 0.9]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = oklab_from_srgb(srgb)
+    via_linear_srgb = oklab_from_linear_srgb(linear_srgb_from_srgb(srgb))
+
+    assert torch.allclose(direct, via_linear_srgb, atol=1e-6)
+
+
+def test_srgb_from_oklab_matches_explicit_conversion_path() -> None:
+    oklab = torch.tensor(
+        [
+            [[0.6, 0.4]],
+            [[0.2, -0.1]],
+            [[0.1, 0.05]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = srgb_from_oklab(oklab)
+    via_linear_srgb = srgb_from_linear_srgb(linear_srgb_from_oklab(oklab))
+
+    assert torch.allclose(direct, via_linear_srgb, atol=1e-6)
+
+
+def test_oklab_from_pure_srgb_red_matches_reference_value() -> None:
+    srgb_red = torch.tensor([[[1.0]], [[0.0]], [[0.0]]], dtype=torch.float32)
+
+    oklab_red = oklab_from_linear_srgb(linear_srgb_from_srgb(srgb_red))
+
+    assert torch.allclose(
+        oklab_red[:, 0, 0],
+        torch.tensor([0.62795536, 0.22486306, 0.1258463], dtype=torch.float32),
+        atol=1e-6,
+    )
+
+
+def test_oklab_oklch_round_trip() -> None:
+    oklab = torch.tensor(
+        [
+            [[0.6, 0.4]],
+            [[0.2, -0.1]],
+            [[0.1, 0.05]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = oklab_from_oklch(oklch_from_oklab(oklab))
+
+    assert torch.allclose(round_tripped, oklab, atol=1e-6)
+
+
+def test_oklch_from_linear_srgb_matches_explicit_conversion_path() -> None:
+    linear_srgb = torch.tensor(
+        [
+            [[0.1, 0.9]],
+            [[0.4, 0.2]],
+            [[0.7, 0.3]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = oklch_from_linear_srgb(linear_srgb)
+    via_oklab = oklch_from_oklab(oklab_from_linear_srgb(linear_srgb))
+
+    assert torch.allclose(direct, via_oklab, atol=1e-6)
+
+
+def test_oklch_from_srgb_and_back_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.2, 0.9]],
+            [[0.4, 0.3]],
+            [[0.8, 0.1]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct_round_trip = srgb_from_oklch(oklch_from_srgb(srgb))
+    explicit_round_trip = srgb_from_linear_srgb(
+        linear_srgb_from_oklch(oklch_from_oklab(oklab_from_linear_srgb(linear_srgb_from_srgb(srgb))))
+    )
+
+    assert torch.allclose(direct_round_trip, srgb, atol=1e-5)
+    assert torch.allclose(explicit_round_trip, srgb, atol=1e-5)
+
+
+def test_linear_srgb_from_oklch_matches_oklab_path() -> None:
+    oklch = torch.tensor(
+        [
+            [[0.7, 0.5]],
+            [[0.12, 0.04]],
+            [[30.0, 210.0]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = linear_srgb_from_oklch(oklch)
+    via_oklab = linear_srgb_from_oklab(oklab_from_oklch(oklch))
+
+    assert torch.allclose(direct, via_oklab, atol=1e-6)
+    assert direct.shape == (3, 1, 2)
+
+
+def test_hsl_srgb_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.9]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = srgb_from_hsl(hsl_from_srgb(srgb))
+
+    assert torch.allclose(round_tripped, srgb, atol=1e-5)
+
+
+def test_hsl_hue_is_expressed_in_degrees() -> None:
+    srgb = torch.tensor(
+        [
+            [[1.0, 0.0, 0.0]],
+            [[0.0, 1.0, 0.0]],
+            [[0.0, 0.0, 1.0]],
+        ],
+        dtype=torch.float32,
+    )
+
+    hsl = hsl_from_srgb(srgb)
+
+    assert torch.allclose(hsl[0, 0, :], torch.tensor([0.0, 120.0, 240.0]), atol=1e-3)
+    assert torch.allclose(hsl[1, 0, :], torch.tensor([1.0, 1.0, 1.0]), atol=1e-6)
+    assert torch.allclose(hsl[2, 0, :], torch.tensor([0.5, 0.5, 0.5]), atol=1e-6)
+
+
+def test_hsl_from_grayscale_has_zero_saturation() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.1, 0.8]],
+            [[0.1, 0.8]],
+            [[0.1, 0.8]],
+        ],
+        dtype=torch.float32,
+    )
+
+    hsl = hsl_from_srgb(srgb)
+
+    assert torch.allclose(hsl[1, ...], torch.zeros_like(hsl[1, ...]), atol=1e-6)
+
+
+def test_hsl_from_linear_srgb_matches_explicit_conversion_path() -> None:
+    linear_srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.8]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = hsl_from_linear_srgb(linear_srgb)
+    via_srgb = hsl_from_srgb(srgb_from_linear_srgb(linear_srgb))
+
+    assert torch.allclose(direct, via_srgb, atol=1e-6)
+
+
+def test_linear_srgb_from_hsl_matches_explicit_conversion_path() -> None:
+    hsl = torch.tensor(
+        [
+            [[0.0, 216.0], [90.0, 324.0]],
+            [[1.0, 0.25], [0.75, 0.1]],
+            [[0.5, 0.4], [0.2, 0.8]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = linear_srgb_from_hsl(hsl)
+    via_srgb = linear_srgb_from_srgb(srgb_from_hsl(hsl))
+
+    assert torch.allclose(direct, via_srgb, atol=1e-6)
+
+
+def test_srgb_from_hsl_wraps_degree_hue_values() -> None:
+    hsl = torch.tensor(
+        [
+            [[360.0, -120.0]],
+            [[1.0, 1.0]],
+            [[0.5, 0.5]],
+        ],
+        dtype=torch.float32,
+    )
+
+    rgb = srgb_from_hsl(hsl)
+
+    assert torch.allclose(rgb[:, 0, 0], torch.tensor([1.0, 0.0, 0.0]), atol=1e-5)
+    assert torch.allclose(rgb[:, 0, 1], torch.tensor([0.0, 0.0, 1.0]), atol=1e-5)
+
+
+def test_okhsl_srgb_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.05, 0.95], [0.2, 0.8]],
+            [[0.4, 0.2], [0.6, 0.1]],
+            [[0.9, 0.05], [0.3, 0.7]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = srgb_from_okhsl(okhsl_from_srgb(srgb))
+
+    assert torch.allclose(round_tripped, srgb, atol=5e-4)
+
+
+def test_okhsl_and_okhsv_hue_are_expressed_in_degrees() -> None:
+    srgb = torch.tensor(
+        [
+            [[1.0, 0.0, 0.0]],
+            [[0.0, 1.0, 0.0]],
+            [[0.0, 0.0, 1.0]],
+        ],
+        dtype=torch.float32,
+    )
+
+    okhsl = okhsl_from_srgb(srgb)
+    okhsv = okhsv_from_srgb(srgb)
+
+    assert torch.allclose(okhsl[0, 0, :], torch.tensor([29.2473, 142.4848, 264.0487]), atol=2e-2)
+    assert torch.allclose(okhsv[0, 0, :], torch.tensor([29.2473, 142.4848, 264.0487]), atol=2e-2)
+    assert torch.allclose(okhsl[1, 0, :], torch.tensor([1.0, 1.0, 1.0]), atol=1e-6)
+    assert torch.allclose(okhsv[1, 0, :], torch.tensor([1.0, 1.0, 1.0]), atol=1e-6)
+
+
+def test_okhsl_and_okhsv_srgb_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.05, 0.95], [0.2, 0.8]],
+            [[0.4, 0.2], [0.6, 0.1]],
+            [[0.9, 0.05], [0.3, 0.7]],
+        ],
+        dtype=torch.float32,
+    )
+
+    okhsl_round_tripped = srgb_from_okhsl(okhsl_from_srgb(srgb))
+    okhsv_round_tripped = srgb_from_okhsv(okhsv_from_srgb(srgb))
+
+    assert torch.allclose(okhsl_round_tripped, srgb, atol=5e-4)
+    assert torch.allclose(okhsv_round_tripped, srgb, atol=5e-4)
+
+
+def test_okhsl_and_okhsv_outputs_keep_hue_in_degrees_and_other_channels_in_unit_range() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0]],
+            [[1.0, 0.0]],
+            [[0.5, 0.25]],
+        ],
+        dtype=torch.float32,
+    )
+
+    okhsl = okhsl_from_srgb(srgb)
+    okhsv = okhsv_from_srgb(srgb)
+
+    assert torch.all(okhsl[0, ...] >= 0.0)
+    assert torch.all(okhsl[0, ...] <= 360.0)
+    assert torch.all(okhsl[1:, ...] >= 0.0)
+    assert torch.all(okhsl[1:, ...] <= 1.0)
+    assert torch.all(okhsv[0, ...] >= 0.0)
+    assert torch.all(okhsv[0, ...] <= 360.0)
+    assert torch.all(okhsv[1:, ...] >= 0.0)
+    assert torch.all(okhsv[1:, ...] <= 1.0)
+
+
+def test_okhsl_and_okhsv_wrap_degree_hue_values() -> None:
+    okhsl = torch.tensor([[[389.2473]], [[1.0]], [[0.5681]]], dtype=torch.float32)
+    okhsl_wrapped = torch.tensor([[[29.2473]], [[1.0]], [[0.5681]]], dtype=torch.float32)
+    okhsv = torch.tensor([[[389.2473]], [[1.0]], [[1.0]]], dtype=torch.float32)
+    okhsv_wrapped = torch.tensor([[[29.2473]], [[1.0]], [[1.0]]], dtype=torch.float32)
+
+    assert torch.allclose(srgb_from_okhsl(okhsl), srgb_from_okhsl(okhsl_wrapped), atol=1e-5)
+    assert torch.allclose(srgb_from_okhsv(okhsv), srgb_from_okhsv(okhsv_wrapped), atol=1e-5)
+
+
+def test_linear_srgb_xyz_round_trip() -> None:
+    linear_srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.8]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = linear_srgb_from_xyz(xyz_from_linear_srgb(linear_srgb))
+
+    assert torch.allclose(round_tripped, linear_srgb, atol=5e-5)
+
+
+def test_oklab_from_xyz_matches_explicit_conversion_path() -> None:
+    xyz = torch.tensor(
+        [
+            [[0.4124, 0.9505]],
+            [[0.2126, 1.0]],
+            [[0.0193, 1.0888]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = oklab_from_xyz(xyz)
+    via_linear_srgb = oklab_from_linear_srgb(linear_srgb_from_xyz(xyz))
+
+    assert torch.allclose(direct, via_linear_srgb, atol=1e-6)
+
+
+def test_xyz_from_oklab_matches_explicit_conversion_path() -> None:
+    oklab = torch.tensor(
+        [
+            [[0.6, 0.4]],
+            [[0.2, -0.1]],
+            [[0.1, 0.05]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = xyz_from_oklab(oklab)
+    via_linear_srgb = xyz_from_linear_srgb(linear_srgb_from_oklab(oklab))
+
+    assert torch.allclose(direct, via_linear_srgb, atol=1e-6)
+
+
+def test_oklch_from_xyz_matches_explicit_conversion_path() -> None:
+    xyz = torch.tensor(
+        [
+            [[0.4124, 0.9505]],
+            [[0.2126, 1.0]],
+            [[0.0193, 1.0888]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = oklch_from_xyz(xyz)
+    via_oklab = oklch_from_oklab(oklab_from_xyz(xyz))
+
+    assert torch.allclose(direct, via_oklab, atol=1e-6)
+
+
+def test_xyz_from_oklch_matches_explicit_conversion_path() -> None:
+    oklch = torch.tensor(
+        [
+            [[0.7, 0.5]],
+            [[0.12, 0.04]],
+            [[30.0, 210.0]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = xyz_from_oklch(oklch)
+    via_oklab = xyz_from_oklab(oklab_from_oklch(oklch))
+
+    assert torch.allclose(direct, via_oklab, atol=1e-6)
+
+
+def test_lab_from_linear_srgb_matches_explicit_conversion_path() -> None:
+    linear_srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.8]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = lab_from_linear_srgb(linear_srgb)
+    via_xyz = lab_from_xyz(xyz_from_linear_srgb(linear_srgb))
+
+    assert torch.allclose(direct, via_xyz, atol=1e-6)
+
+
+def test_linear_srgb_from_lab_matches_explicit_conversion_path() -> None:
+    lab = torch.tensor(
+        [
+            [[0.0, 100.0], [50.0, 75.0]],
+            [[0.0, 0.0], [10.0, -20.0]],
+            [[0.0, 0.0], [-5.0, 30.0]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = linear_srgb_from_lab(lab)
+    via_xyz = linear_srgb_from_xyz(xyz_from_lab(lab))
+
+    assert torch.allclose(direct, via_xyz, atol=1e-6)
+
+
+def test_srgb_xyz_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.8]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = srgb_from_xyz(xyz_from_srgb(srgb))
+
+    assert torch.allclose(round_tripped, srgb, atol=5e-4)
+
+
+def test_lab_from_srgb_and_back_round_trip() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.8]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = srgb_from_lab(lab_from_srgb(srgb))
+
+    assert torch.allclose(round_tripped, srgb, atol=5e-4)
+
+
+def test_xyz_lab_round_trip_for_d65_and_d50() -> None:
+    xyz = torch.tensor(
+        [
+            [[0.4124, 0.9505]],
+            [[0.2126, 1.0]],
+            [[0.0193, 1.0888]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped_d65 = xyz_from_lab(lab_from_xyz(xyz, reference_illuminant="D65"), reference_illuminant="D65")
+    round_tripped_d50 = xyz_from_lab(lab_from_xyz(xyz, reference_illuminant="D50"), reference_illuminant="D50")
+
+    assert torch.allclose(round_tripped_d65, xyz, atol=1e-4)
+    assert torch.allclose(round_tripped_d50, xyz, atol=1e-4)
+
+
+def test_xyz_d65_to_d50_maps_reference_white() -> None:
+    xyz_d65 = torch.tensor([[[0.950489]], [[1.0]], [[1.088840]]], dtype=torch.float32)
+
+    xyz_d50 = xyz_d65_to_d50(xyz_d65)
+
+    assert torch.allclose(xyz_d50[:, 0, 0], torch.tensor([0.964212, 1.0, 0.825188]), atol=5e-4)
+
+
+def test_xyz_d50_to_d65_round_trip() -> None:
+    xyz_d65 = torch.tensor(
+        [
+            [[0.4124, 0.9505]],
+            [[0.2126, 1.0]],
+            [[0.0193, 1.0888]],
+        ],
+        dtype=torch.float32,
+    )
+
+    round_tripped = xyz_d50_to_d65(xyz_d65_to_d50(xyz_d65))
+
+    assert torch.allclose(round_tripped, xyz_d65, atol=1e-5)
+
+
+def test_lab_from_srgb_d50_matches_adapted_xyz_path() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.0, 1.0], [0.25, 0.8]],
+            [[0.2, 0.8], [0.75, 0.1]],
+            [[1.0, 0.1], [0.5, 0.4]],
+        ],
+        dtype=torch.float32,
+    )
+
+    direct = lab_from_srgb(srgb, reference_illuminant="D50")
+    via_xyz = lab_from_xyz(xyz_d65_to_d50(xyz_from_srgb(srgb)), reference_illuminant="D50")
+
+    assert torch.allclose(direct, via_xyz, atol=1e-4)
+
+
+def test_lab_from_xyz_matches_reference_white_and_black() -> None:
+    xyz = torch.tensor(
+        [
+            [[0.0, 0.950489]],
+            [[0.0, 1.0]],
+            [[0.0, 1.088840]],
+        ],
+        dtype=torch.float32,
+    )
+
+    lab = lab_from_xyz(xyz, reference_illuminant="D65")
+
+    assert torch.allclose(lab[:, 0, 0], torch.tensor([0.0, 0.0, 0.0]), atol=1e-4)
+    assert torch.allclose(lab[:, 0, 1], torch.tensor([100.0, 0.0, 0.0]), atol=1e-3)
+
+
+def test_invalid_tensor_shape_raises_value_error() -> None:
+    with pytest.raises(ValueError, match="3xHxW"):
+        oklab_from_srgb(torch.zeros((2, 2), dtype=torch.float32))
+
+
+def test_invalid_reference_illuminant_raises_value_error() -> None:
+    xyz = torch.ones((3, 1, 1), dtype=torch.float32)
+
+    with pytest.raises(ValueError, match="Unsupported reference_illuminant"):
+        lab_from_xyz(xyz, reference_illuminant="E")
+
+
+def test_okhsl_from_srgb_forwards_steps_parameters(monkeypatch: pytest.MonkeyPatch) -> None:
+    recorded: dict[str, int] = {}
+
+    def fake_get_cs_tensor(
+        l_tensor: torch.Tensor, units_ab_tensor: torch.Tensor, steps: int = 1, steps_outer: int = 1
+    ) -> torch.Tensor:
+        recorded["steps"] = steps
+        recorded["steps_outer"] = steps_outer
+        return torch.ones((3, *l_tensor.shape), dtype=l_tensor.dtype, device=l_tensor.device)
+
+    monkeypatch.setattr(color_conversion, "_get_cs_tensor", fake_get_cs_tensor)
+
+    srgb = torch.tensor(
+        [
+            [[0.2, 0.9]],
+            [[0.4, 0.3]],
+            [[0.8, 0.1]],
+        ],
+        dtype=torch.float32,
+    )
+
+    color_conversion.okhsl_from_srgb(srgb, steps=3, steps_outer=4)
+
+    assert recorded == {"steps": 3, "steps_outer": 4}
+
+
+def test_public_hsl_okhsl_okhsv_conversions_preserve_dtype_and_device() -> None:
+    srgb = torch.tensor(
+        [
+            [[0.05, 0.95], [0.2, 0.8]],
+            [[0.4, 0.2], [0.6, 0.1]],
+            [[0.9, 0.05], [0.3, 0.7]],
+        ],
+        dtype=torch.float64,
+    )
+
+    hsl = hsl_from_srgb(srgb)
+    okhsl = okhsl_from_srgb(srgb)
+    okhsv = okhsv_from_srgb(srgb)
+    srgb_from_plain_hsl = srgb_from_hsl(hsl)
+    srgb_from_perceptual_hsl = srgb_from_okhsl(okhsl)
+    srgb_from_perceptual_hsv = srgb_from_okhsv(okhsv)
+
+    for output in (hsl, okhsl, okhsv, srgb_from_plain_hsl, srgb_from_perceptual_hsl, srgb_from_perceptual_hsv):
+        assert output.dtype == srgb.dtype
+        assert output.device == srgb.device

tests/backend/image_util/test_vendor_mutable_defaults.py

@@ -0,0 +1,106 @@
+"""Tests for the mutable default argument fix in imwatermark/vendor.py
+and the bare except fix in sqlite_database.py."""
+
+from logging import Logger
+from unittest import mock
+
+import pytest
+
+from invokeai.backend.image_util.imwatermark.vendor import EmbedMaxDct, WatermarkEncoder
+
+
+class TestSetByBitsNoSharedState:
+    """set_by_bits() used to have bits=[] as a default arg.
+    If it were still mutable, successive calls without an explicit arg
+    would accumulate state. After the fix (bits=None), each call gets
+    a fresh list."""
+
+    def test_set_by_bits_default_is_independent(self):
+        enc1 = WatermarkEncoder()
+        enc1.set_by_bits()
+        assert enc1._watermarks == []
+        assert enc1._wmLen == 0
+
+        enc2 = WatermarkEncoder()
+        enc2.set_by_bits()
+        assert enc2._watermarks == []
+        assert enc2._wmLen == 0
+
+    def test_set_by_bits_with_explicit_arg(self):
+        enc = WatermarkEncoder()
+        enc.set_by_bits([1, 0, 1])
+        assert enc._watermarks == [1, 0, 1]
+        assert enc._wmLen == 3
+        assert enc._wmType == "bits"
+
+
+class TestEmbedMaxDctNoSharedState:
+    """EmbedMaxDct.__init__ used to have watermarks=[] and scales=[0,36,36].
+    After the fix (both default to None), each instance gets its own list."""
+
+    def test_default_watermarks_independent(self):
+        e1 = EmbedMaxDct()
+        e1._watermarks.append(999)
+
+        e2 = EmbedMaxDct()
+        assert 999 not in e2._watermarks
+        assert e2._watermarks == []
+
+    def test_default_scales_independent(self):
+        e1 = EmbedMaxDct()
+        e1._scales.append(72)
+
+        e2 = EmbedMaxDct()
+        assert e2._scales == [0, 36, 36]
+
+    def test_explicit_args_still_work(self):
+        wm = [1, 0, 1, 1]
+        sc = [0, 50, 50]
+        e = EmbedMaxDct(watermarks=wm, wmLen=4, scales=sc, block=8)
+        assert e._watermarks == wm
+        assert e._wmLen == 4
+        assert e._scales == sc
+        assert e._block == 8
+
+
+class TestTransactionExceptException:
+    """The transaction() context manager used to have a bare `except:`.
+    After the fix it uses `except Exception:`, so BaseException subclasses
+    like KeyboardInterrupt and SystemExit should propagate instead of
+    being silently caught and rolled back."""
+
+    @staticmethod
+    def _make_db():
+        """Create a minimal SqliteDatabase-like object with transaction()."""
+        # Import here so the test stays focused; we just need the real class.
+        from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
+
+        logger = mock.MagicMock(spec=Logger)
+        db = SqliteDatabase(db_path=None, logger=logger, verbose=False)
+        return db
+
+    def test_regular_exception_rolls_back(self):
+        db = self._make_db()
+
+        # create a table first in a successful transaction
+        with db.transaction() as cursor:
+            cursor.execute("CREATE TABLE t (id INTEGER)")
+
+        # now try to insert and fail — the insert should be rolled back
+        with pytest.raises(ValueError):
+            with db.transaction() as cursor:
+                cursor.execute("INSERT INTO t VALUES (42)")
+                raise ValueError("boom")
+
+        # the row should not exist after rollback
+        with db.transaction() as cursor:
+            cursor.execute("SELECT * FROM t")
+            assert cursor.fetchone() is None
+
+    def test_keyboard_interrupt_propagates(self):
+        with pytest.raises(KeyboardInterrupt):
+            raise KeyboardInterrupt()
+
+    def test_system_exit_propagates(self):
+        with pytest.raises(SystemExit):
+            raise SystemExit(1)

tests/backend/ip_adapter/test_ip_adapter.py

@@ -0,0 +1,84 @@
+import pytest
+import torch
+
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, SubModelType
+from invokeai.backend.stable_diffusion.diffusion.unet_attention_patcher import UNetAttentionPatcher
+from invokeai.backend.util.test_utils import install_and_load_model
+
+
+def build_dummy_sd15_unet_input(torch_device):
+    batch_size = 1
+    num_channels = 4
+    sizes = (32, 32)
+
+    noise = torch.randn((batch_size, num_channels) + sizes).to(torch_device)
+    time_step = torch.tensor([10]).to(torch_device)
+    encoder_hidden_states = torch.randn((batch_size, 77, 768)).to(torch_device)
+
+    return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}
+
+
+@pytest.mark.parametrize(
+    "model_params",
+    [
+        # SD1.5, IPAdapter
+        {
+            "ip_adapter_model_id": "InvokeAI/ip_adapter_sd15",
+            "ip_adapter_model_name": "ip_adapter_sd15",
+            "base_model": BaseModelType.StableDiffusion1,
+            "unet_model_id": "runwayml/stable-diffusion-v1-5",
+            "unet_model_name": "stable-diffusion-v1-5",
+        },
+        # SD1.5, IPAdapterPlus
+        {
+            "ip_adapter_model_id": "InvokeAI/ip_adapter_plus_sd15",
+            "ip_adapter_model_name": "ip_adapter_plus_sd15",
+            "base_model": BaseModelType.StableDiffusion1,
+            "unet_model_id": "runwayml/stable-diffusion-v1-5",
+            "unet_model_name": "stable-diffusion-v1-5",
+        },
+        # SD1.5, IPAdapterFull
+        {
+            "ip_adapter_model_id": "InvokeAI/ip-adapter-full-face_sd15",
+            "ip_adapter_model_name": "ip-adapter-full-face_sd15",
+            "base_model": BaseModelType.StableDiffusion1,
+            "unet_model_id": "runwayml/stable-diffusion-v1-5",
+            "unet_model_name": "stable-diffusion-v1-5",
+        },
+    ],
+)
+@pytest.mark.slow
+def test_ip_adapter_unet_patch(model_params, model_installer, torch_device):
+    """Smoke test that IP-Adapter weights can be loaded and used to patch a UNet."""
+    ip_adapter_info = install_and_load_model(
+        model_installer=model_installer,
+        model_path_id_or_url=model_params["ip_adapter_model_id"],
+        model_name=model_params["ip_adapter_model_name"],
+        base_model=model_params["base_model"],
+        model_type=ModelType.IPAdapter,
+    )
+
+    unet_info = install_and_load_model(
+        model_installer=model_installer,
+        model_path_id_or_url=model_params["unet_model_id"],
+        model_name=model_params["unet_model_name"],
+        base_model=model_params["base_model"],
+        model_type=ModelType.Main,
+        submodel_type=SubModelType.UNet,
+    )
+
+    dummy_unet_input = build_dummy_sd15_unet_input(torch_device)
+
+    with torch.no_grad(), ip_adapter_info as ip_adapter, unet_info as unet:
+        ip_adapter.to(torch_device, dtype=torch.float32)
+        unet.to(torch_device, dtype=torch.float32)
+
+        # ip_embeds shape: (batch_size, num_ip_images, seq_len, ip_image_embedding_len)
+        ip_embeds = torch.randn((1, 3, 4, 768)).to(torch_device)
+
+        cross_attention_kwargs = {"ip_adapter_image_prompt_embeds": [ip_embeds]}
+        ip_adapter_unet_patcher = UNetAttentionPatcher([ip_adapter])
+        with ip_adapter_unet_patcher.apply_ip_adapter_attention(unet):
+            output = unet(**dummy_unet_input, cross_attention_kwargs=cross_attention_kwargs).sample
+
+    assert output.shape == dummy_unet_input["sample"].shape

tests/backend/model_manager/configs/test_anima_controlnet_config.py

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

tests/backend/model_manager/configs/test_anima_model_identification.py

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

tests/backend/model_manager/configs/test_double_variant_regression.py

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

tests/backend/model_manager/configs/test_lora_default_settings.py

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

tests/backend/model_manager/configs/test_qwen3_encoder_config.py

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

tests/backend/model_manager/configs/test_qwen_image_checkpoint_variant_detection.py

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

tests/backend/model_manager/configs/test_qwen_image_gguf_variant_detection.py

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

tests/backend/model_manager/configs/test_qwen_image_main_config.py

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

tests/backend/model_manager/configs/test_qwen_vl_encoder_config.py

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

tests/backend/model_manager/data/invokeai_root/README

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

tests/backend/model_manager/load/test_anima_state_dict_utils.py

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

tests/backend/model_manager/load/test_flux2_state_dict_utils.py

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

tests/backend/model_manager/load/test_load_default_fp8.py

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

tests/backend/model_manager/load/test_loaded_model.py

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

tests/backend/model_manager/load/test_qwen_image_state_dict_utils.py

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

tests/backend/model_manager/load/test_z_image_state_dict_utils.py

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

tests/backend/model_manager/model_loading/test_model_load.py

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

tests/backend/model_manager/model_manager_fixtures.py

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

tests/backend/model_manager/test_external_api_config.py

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

tests/backend/model_manager/test_libc_util.py

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

tests/backend/model_manager/test_memory_snapshot.py

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

tests/backend/model_manager/test_model_load_optimization.py

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

tests/backend/model_manager/util/test_hf_model_select.py

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

tests/backend/patches/layers/test_flux_control_lora_layer.py

@@ -0,0 +1,25 @@
+import torch
+
+from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer
+
+
+def test_flux_control_lora_layer_get_parameters():
+    """Test getting weight and bias parameters from FluxControlLoRALayer."""
+    small_in_features = 4
+    big_in_features = 8
+    out_features = 16
+    rank = 4
+    alpha = 16.0
+    layer = FluxControlLoRALayer(
+        up=torch.ones(out_features, rank), mid=None, down=torch.ones(rank, big_in_features), alpha=alpha, bias=None
+    )
+
+    # Create mock original module
+    orig_module = torch.nn.Linear(small_in_features, out_features)
+
+    # Test that get_parameters() behaves as expected in spite of the difference in in_features shapes.
+    params = layer.get_parameters(dict(orig_module.named_parameters(recurse=False)), weight=1.0)
+    assert "weight" in params
+    assert params["weight"].shape == (out_features, big_in_features)
+    assert params["weight"].allclose(torch.ones(out_features, big_in_features) * alpha)
+    assert "bias" not in params  # No bias in this case

tests/backend/patches/layers/test_lora_layer.py

@@ -0,0 +1,114 @@
+import logging
+
+import pytest
+import torch
+
+from invokeai.backend.patches.layers.lora_layer import LoRALayer
+
+
+def test_lora_layer_init_from_state_dict():
+    """Test initializing a LoRALayer from state dict values."""
+    # Create mock state dict values
+    in_features = 8
+    out_features = 16
+    rank = 4
+    alpha = 16.0
+    values = {
+        "lora_up.weight": torch.ones(out_features, rank),
+        "lora_down.weight": torch.ones(rank, in_features),
+        "alpha": torch.tensor(alpha),
+    }
+    layer = LoRALayer.from_state_dict_values(values)
+
+    assert layer.up.shape == (out_features, rank)
+    assert layer.down.shape == (rank, in_features)
+    assert layer._alpha == alpha
+    assert layer.bias is None
+
+
+def test_lora_layer_init_from_state_dict_with_unhandled_keys_logs_warning(caplog: pytest.LogCaptureFixture):
+    """Test initializing a LoRALayer from state dict values with an unhandled key."""
+    in_features = 8
+    out_features = 16
+    rank = 4
+    alpha = 16.0
+    values = {
+        "lora_up.weight": torch.ones(out_features, rank),
+        "lora_down.weight": torch.ones(rank, in_features),
+        "alpha": torch.tensor(alpha),
+        "unhandled_key": torch.randn(4, 4),
+    }
+
+    with caplog.at_level(logging.WARNING):
+        _ = LoRALayer.from_state_dict_values(values)
+
+    assert (
+        "Unexpected keys found in LoRA/LyCORIS layer, model might work incorrectly! Unexpected keys: {'unhandled_key'}"
+        in caplog.text
+    )
+
+
+@pytest.mark.parametrize(
+    ["device"],
+    [
+        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA device")),
+        pytest.param(
+            "mps", marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires MPS device")
+        ),
+    ],
+)
+def test_lora_layer_to(device: str):
+    in_features = 8
+    out_features = 16
+    rank = 4
+    alpha = 16.0
+    values = {
+        "lora_up.weight": torch.ones(out_features, rank),
+        "lora_down.weight": torch.ones(rank, in_features),
+        "alpha": torch.tensor(alpha),
+    }
+    layer = LoRALayer.from_state_dict_values(values)
+
+    # Layer is initialized on the CPU.
+    assert layer.up.device.type == "cpu"
+    assert layer.down.device.type == "cpu"
+
+    # Test moving to device.
+    layer.to(device=torch.device(device))
+    assert layer.up.device.type == device
+    assert layer.down.device.type == device
+
+
+def test_lora_layer_calc_size():
+    """Test calculating memory size of LoRALayer tensors."""
+    # Initialize weights with random shapes.
+    up = torch.randn(1, 2)
+    mid = torch.randn(2, 3)
+    down = torch.randn(3, 4)
+    bias = torch.randn(5)
+    layer = LoRALayer(up=up, mid=mid, down=down, alpha=8.0, bias=bias)
+
+    assert layer.calc_size() == sum(tensor.numel() * tensor.element_size() for tensor in [up, mid, down, bias])
+
+
+def test_lora_layer_get_parameters():
+    """Test getting weight and bias parameters from LoRALayer."""
+    in_features = 8
+    out_features = 16
+    rank = 4
+    alpha = 16.0
+    values = {
+        "lora_up.weight": torch.ones(out_features, rank),
+        "lora_down.weight": torch.ones(rank, in_features),
+        "alpha": torch.tensor(alpha),
+    }
+    layer = LoRALayer.from_state_dict_values(values)
+
+    # Create mock original module
+    orig_module = torch.nn.Linear(in_features, out_features)
+
+    params = layer.get_parameters(dict(orig_module.named_parameters(recurse=False)), weight=1.0)
+    assert "weight" in params
+    assert params["weight"].shape == orig_module.weight.shape
+    assert params["weight"].allclose(torch.ones(out_features, in_features) * alpha)
+    assert "bias" not in params  # No bias in this case

tests/backend/patches/layers/test_set_parameter_layer.py

@@ -0,0 +1,49 @@
+import pytest
+import torch
+
+from invokeai.backend.patches.layers.set_parameter_layer import SetParameterLayer
+
+
+def test_set_parameter_layer_get_parameters():
+    orig_module = torch.nn.Linear(4, 8)
+
+    target_weight = torch.randn(8, 4)
+    layer = SetParameterLayer(param_name="weight", weight=target_weight)
+
+    params = layer.get_parameters(dict(orig_module.named_parameters(recurse=False)), weight=1.0)
+    assert len(params) == 1
+    new_weight = orig_module.weight + params["weight"]
+    assert torch.allclose(new_weight, target_weight)
+
+
+@pytest.mark.parametrize(
+    ["device"],
+    [
+        pytest.param("cuda", marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA device")),
+        pytest.param(
+            "mps", marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires MPS device")
+        ),
+    ],
+)
+def test_set_parameter_layer_to(device: str):
+    """Test moving SetParameterLayer to different device/dtype."""
+
+    target_weight = torch.randn(8, 4)
+    layer = SetParameterLayer(param_name="weight", weight=target_weight)
+
+    # SetParameterLayer should be initialized on CPU.
+    assert layer.weight.device.type == "cpu"  # type: ignore
+
+    # Move to device.
+    layer.to(device=torch.device(device))
+    assert layer.weight.device.type == device  # type: ignore
+
+
+def test_set_parameter_layer_calc_size():
+    """Test calculating parameter size of SetParameterLayer"""
+    param = torch.randn(4, 8)
+    layer = SetParameterLayer(param_name="weight", weight=param)
+
+    # Size should be number of elements * bytes per element
+    expected_size = param.nelement() * param.element_size()
+    assert layer.calc_size() == expected_size

tests/backend/patches/lora_conversions/lora_state_dicts/anima_lora_kohya_format.py

@@ -0,0 +1,42 @@
+# A sample state dict in the Kohya Anima LoRA format.
+# These keys are based on Anima LoRAs targeting the Cosmos Predict2 DiT transformer.
+# Keys follow the pattern: lora_unet_blocks_{N}_{component}.{suffix}
+state_dict_keys: dict[str, list[int]] = {
+    # Block 0 - cross attention
+    "lora_unet_blocks_0_cross_attn_k_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_cross_attn_k_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_cross_attn_k_proj.alpha": [],
+    "lora_unet_blocks_0_cross_attn_q_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_cross_attn_q_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_cross_attn_q_proj.alpha": [],
+    "lora_unet_blocks_0_cross_attn_v_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_cross_attn_v_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_cross_attn_v_proj.alpha": [],
+    "lora_unet_blocks_0_cross_attn_output_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_cross_attn_output_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_cross_attn_output_proj.alpha": [],
+    # Block 0 - self attention
+    "lora_unet_blocks_0_self_attn_k_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_self_attn_k_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_self_attn_k_proj.alpha": [],
+    "lora_unet_blocks_0_self_attn_q_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_self_attn_q_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_self_attn_q_proj.alpha": [],
+    "lora_unet_blocks_0_self_attn_v_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_self_attn_v_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_self_attn_v_proj.alpha": [],
+    "lora_unet_blocks_0_self_attn_output_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_self_attn_output_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_self_attn_output_proj.alpha": [],
+    # Block 0 - MLP
+    "lora_unet_blocks_0_mlp_layer1.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_mlp_layer1.lora_up.weight": [8192, 8],
+    "lora_unet_blocks_0_mlp_layer1.alpha": [],
+    "lora_unet_blocks_0_mlp_layer2.lora_down.weight": [8, 8192],
+    "lora_unet_blocks_0_mlp_layer2.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_mlp_layer2.alpha": [],
+    # Block 0 - adaln modulation
+    "lora_unet_blocks_0_adaln_modulation_cross_attn_1.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_adaln_modulation_cross_attn_1.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_adaln_modulation_cross_attn_1.alpha": [],
+}

tests/backend/patches/lora_conversions/lora_state_dicts/anima_lora_kohya_with_te_format.py

@@ -0,0 +1,34 @@
+# A sample state dict in the Kohya Anima LoRA format with Qwen3 text encoder layers.
+# Contains both lora_unet_ (transformer) and lora_te_ (Qwen3 encoder) keys.
+state_dict_keys: dict[str, list[int]] = {
+    # Transformer block 0 - cross attention
+    "lora_unet_blocks_0_cross_attn_k_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_cross_attn_k_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_cross_attn_k_proj.alpha": [],
+    "lora_unet_blocks_0_cross_attn_q_proj.lora_down.weight": [8, 2048],
+    "lora_unet_blocks_0_cross_attn_q_proj.lora_up.weight": [2048, 8],
+    "lora_unet_blocks_0_cross_attn_q_proj.alpha": [],
+    # Qwen3 text encoder layer 0 - self attention
+    "lora_te_layers_0_self_attn_q_proj.lora_down.weight": [8, 1024],
+    "lora_te_layers_0_self_attn_q_proj.lora_up.weight": [1024, 8],
+    "lora_te_layers_0_self_attn_q_proj.alpha": [],
+    "lora_te_layers_0_self_attn_k_proj.lora_down.weight": [8, 1024],
+    "lora_te_layers_0_self_attn_k_proj.lora_up.weight": [1024, 8],
+    "lora_te_layers_0_self_attn_k_proj.alpha": [],
+    "lora_te_layers_0_self_attn_v_proj.lora_down.weight": [8, 1024],
+    "lora_te_layers_0_self_attn_v_proj.lora_up.weight": [1024, 8],
+    "lora_te_layers_0_self_attn_v_proj.alpha": [],
+    "lora_te_layers_0_self_attn_o_proj.lora_down.weight": [8, 1024],
+    "lora_te_layers_0_self_attn_o_proj.lora_up.weight": [1024, 8],
+    "lora_te_layers_0_self_attn_o_proj.alpha": [],
+    # Qwen3 text encoder layer 0 - MLP
+    "lora_te_layers_0_mlp_gate_proj.lora_down.weight": [8, 1024],
+    "lora_te_layers_0_mlp_gate_proj.lora_up.weight": [2816, 8],
+    "lora_te_layers_0_mlp_gate_proj.alpha": [],
+    "lora_te_layers_0_mlp_down_proj.lora_down.weight": [8, 2816],
+    "lora_te_layers_0_mlp_down_proj.lora_up.weight": [1024, 8],
+    "lora_te_layers_0_mlp_down_proj.alpha": [],
+    "lora_te_layers_0_mlp_up_proj.lora_down.weight": [8, 1024],
+    "lora_te_layers_0_mlp_up_proj.lora_up.weight": [2816, 8],
+    "lora_te_layers_0_mlp_up_proj.alpha": [],
+}

tests/backend/patches/lora_conversions/lora_state_dicts/anima_lora_lokr_format.py

@@ -0,0 +1,18 @@
+# A sample state dict in the LoKR Anima LoRA format (with DoRA).
+# Some Anima LoRAs use LoKR weights (lokr_w1/lokr_w2) combined with DoRA (dora_scale).
+# The dora_scale should be stripped from LoKR layers during conversion.
+state_dict_keys: dict[str, list[int]] = {
+    # Block 0 - cross attention with LoKR + DoRA
+    "diffusion_model.blocks.0.cross_attn.k_proj.lokr_w1": [2048, 8],
+    "diffusion_model.blocks.0.cross_attn.k_proj.lokr_w2": [8, 2048],
+    "diffusion_model.blocks.0.cross_attn.k_proj.alpha": [],
+    "diffusion_model.blocks.0.cross_attn.k_proj.dora_scale": [2048],
+    "diffusion_model.blocks.0.cross_attn.q_proj.lokr_w1": [2048, 8],
+    "diffusion_model.blocks.0.cross_attn.q_proj.lokr_w2": [8, 2048],
+    "diffusion_model.blocks.0.cross_attn.q_proj.alpha": [],
+    "diffusion_model.blocks.0.cross_attn.q_proj.dora_scale": [2048],
+    # Block 0 - self attention with LoKR (no DoRA)
+    "diffusion_model.blocks.0.self_attn.k_proj.lokr_w1": [2048, 8],
+    "diffusion_model.blocks.0.self_attn.k_proj.lokr_w2": [8, 2048],
+    "diffusion_model.blocks.0.self_attn.k_proj.alpha": [],
+}

tests/backend/patches/lora_conversions/lora_state_dicts/anima_lora_peft_format.py

@@ -0,0 +1,19 @@
+# A sample state dict in the diffusers PEFT Anima LoRA format.
+# Keys follow the pattern: diffusion_model.blocks.{N}.{component}.lora_{A|B}.weight
+state_dict_keys: dict[str, list[int]] = {
+    # Block 0 - cross attention
+    "diffusion_model.blocks.0.cross_attn.k_proj.lora_A.weight": [8, 2048],
+    "diffusion_model.blocks.0.cross_attn.k_proj.lora_B.weight": [2048, 8],
+    "diffusion_model.blocks.0.cross_attn.q_proj.lora_A.weight": [8, 2048],
+    "diffusion_model.blocks.0.cross_attn.q_proj.lora_B.weight": [2048, 8],
+    "diffusion_model.blocks.0.cross_attn.v_proj.lora_A.weight": [8, 2048],
+    "diffusion_model.blocks.0.cross_attn.v_proj.lora_B.weight": [2048, 8],
+    # Block 0 - self attention
+    "diffusion_model.blocks.0.self_attn.k_proj.lora_A.weight": [8, 2048],
+    "diffusion_model.blocks.0.self_attn.k_proj.lora_B.weight": [2048, 8],
+    "diffusion_model.blocks.0.self_attn.q_proj.lora_A.weight": [8, 2048],
+    "diffusion_model.blocks.0.self_attn.q_proj.lora_B.weight": [2048, 8],
+    # Block 0 - MLP
+    "diffusion_model.blocks.0.mlp.layer1.lora_A.weight": [8, 2048],
+    "diffusion_model.blocks.0.mlp.layer1.lora_B.weight": [8192, 8],
+}

tests/backend/patches/lora_conversions/lora_state_dicts/flux_lokr_bfl_format.py

@@ -0,0 +1,22 @@
+# A sample state dict in the BFL LOKR format (FLUX.1 hidden_size=3072).
+# These keys represent a LOKR model using BFL internal key names with 'diffusion_model.' prefix.
+state_dict_keys = {
+    "diffusion_model.double_blocks.0.img_attn.proj.lokr_w1": [32, 96],
+    "diffusion_model.double_blocks.0.img_attn.proj.lokr_w2": [32, 32],
+    "diffusion_model.double_blocks.0.img_attn.proj.alpha": [],
+    "diffusion_model.double_blocks.0.img_attn.qkv.lokr_w1": [32, 96],
+    "diffusion_model.double_blocks.0.img_attn.qkv.lokr_w2": [32, 288],
+    "diffusion_model.double_blocks.0.img_attn.qkv.alpha": [],
+    "diffusion_model.double_blocks.0.img_mlp.0.lokr_w1": [32, 96],
+    "diffusion_model.double_blocks.0.img_mlp.0.lokr_w2": [32, 128],
+    "diffusion_model.double_blocks.0.img_mlp.0.alpha": [],
+    "diffusion_model.double_blocks.0.img_mlp.2.lokr_w1": [32, 128],
+    "diffusion_model.double_blocks.0.img_mlp.2.lokr_w2": [32, 96],
+    "diffusion_model.double_blocks.0.img_mlp.2.alpha": [],
+    "diffusion_model.single_blocks.0.linear1.lokr_w1": [32, 128],
+    "diffusion_model.single_blocks.0.linear1.lokr_w2": [32, 128],
+    "diffusion_model.single_blocks.0.linear1.alpha": [],
+    "diffusion_model.single_blocks.0.linear2.lokr_w1": [32, 64],
+    "diffusion_model.single_blocks.0.linear2.lokr_w2": [32, 48],
+    "diffusion_model.single_blocks.0.linear2.alpha": [],
+}

tests/backend/patches/lora_conversions/lora_state_dicts/flux_lora_aitoolkit_format.py

@@ -0,0 +1,458 @@
+state_dict_keys = {
+    "diffusion_model.double_blocks.0.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.0.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.0.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.0.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.0.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.0.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.0.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.0.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.0.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.0.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.0.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.0.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.0.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.0.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.0.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.0.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.1.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.1.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.1.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.1.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.1.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.1.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.1.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.1.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.1.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.1.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.1.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.1.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.1.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.1.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.1.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.1.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.10.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.10.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.10.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.10.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.10.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.10.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.10.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.10.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.10.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.10.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.10.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.10.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.10.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.10.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.10.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.10.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.11.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.11.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.11.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.11.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.11.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.11.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.11.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.11.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.11.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.11.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.11.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.11.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.11.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.11.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.11.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.11.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.12.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.12.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.12.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.12.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.12.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.12.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.12.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.12.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.12.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.12.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.12.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.12.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.12.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.12.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.12.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.12.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.13.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.13.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.13.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.13.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.13.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.13.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.13.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.13.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.13.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.13.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.13.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.13.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.13.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.13.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.13.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.13.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.14.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.14.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.14.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.14.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.14.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.14.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.14.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.14.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.14.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.14.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.14.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.14.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.14.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.14.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.14.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.14.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.15.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.15.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.15.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.15.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.15.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.15.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.15.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.15.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.15.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.15.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.15.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.15.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.15.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.15.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.15.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.15.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.16.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.16.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.16.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.16.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.16.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.16.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.16.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.16.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.16.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.16.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.16.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.16.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.16.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.16.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.16.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.16.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.17.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.17.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.17.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.17.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.17.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.17.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.17.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.17.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.17.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.17.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.17.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.17.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.17.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.17.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.17.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.17.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.18.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.18.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.18.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.18.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.18.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.18.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.18.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.18.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.18.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.18.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.18.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.18.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.18.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.18.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.18.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.18.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.2.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.2.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.2.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.2.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.2.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.2.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.2.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.2.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.2.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.2.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.2.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.2.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.2.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.2.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.2.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.2.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.3.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.3.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.3.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.3.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.3.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.3.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.3.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.3.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.3.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.3.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.3.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.3.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.3.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.3.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.3.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.3.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.4.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.4.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.4.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.4.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.4.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.4.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.4.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.4.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.4.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.4.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.4.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.4.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.4.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.4.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.4.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.4.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.5.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.5.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.5.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.5.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.5.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.5.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.5.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.5.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.5.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.5.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.5.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.5.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.5.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.5.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.5.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.5.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.6.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.6.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.6.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.6.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.6.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.6.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.6.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.6.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.6.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.6.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.6.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.6.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.6.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.6.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.6.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.6.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.7.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.7.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.7.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.7.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.7.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.7.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.7.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.7.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.7.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.7.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.7.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.7.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.7.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.7.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.7.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.7.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.8.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.8.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.8.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.8.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.8.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.8.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.8.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.8.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.8.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.8.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.8.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.8.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.8.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.8.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.8.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.8.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.9.img_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.9.img_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.9.img_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.9.img_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.9.img_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.9.img_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.9.img_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.9.img_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.9.txt_attn.proj.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.9.txt_attn.proj.lora_B.weight": [3072, 16],
+    "diffusion_model.double_blocks.9.txt_attn.qkv.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.9.txt_attn.qkv.lora_B.weight": [9216, 16],
+    "diffusion_model.double_blocks.9.txt_mlp.0.lora_A.weight": [16, 3072],
+    "diffusion_model.double_blocks.9.txt_mlp.0.lora_B.weight": [12288, 16],
+    "diffusion_model.double_blocks.9.txt_mlp.2.lora_A.weight": [16, 12288],
+    "diffusion_model.double_blocks.9.txt_mlp.2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.0.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.0.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.0.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.0.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.1.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.1.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.1.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.1.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.10.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.10.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.10.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.10.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.11.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.11.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.11.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.11.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.12.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.12.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.12.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.12.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.13.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.13.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.13.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.13.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.14.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.14.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.14.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.14.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.15.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.15.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.15.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.15.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.16.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.16.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.16.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.16.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.17.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.17.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.17.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.17.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.18.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.18.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.18.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.18.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.19.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.19.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.19.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.19.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.2.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.2.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.2.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.2.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.20.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.20.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.20.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.20.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.21.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.21.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.21.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.21.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.22.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.22.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.22.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.22.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.23.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.23.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.23.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.23.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.24.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.24.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.24.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.24.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.25.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.25.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.25.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.25.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.26.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.26.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.26.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.26.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.27.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.27.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.27.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.27.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.28.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.28.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.28.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.28.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.29.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.29.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.29.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.29.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.3.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.3.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.3.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.3.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.30.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.30.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.30.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.30.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.31.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.31.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.31.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.31.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.32.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.32.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.32.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.32.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.33.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.33.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.33.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.33.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.34.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.34.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.34.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.34.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.35.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.35.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.35.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.35.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.36.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.36.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.36.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.36.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.37.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.37.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.37.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.37.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.4.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.4.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.4.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.4.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.5.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.5.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.5.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.5.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.6.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.6.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.6.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.6.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.7.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.7.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.7.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.7.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.8.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.8.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.8.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.8.linear2.lora_B.weight": [3072, 16],
+    "diffusion_model.single_blocks.9.linear1.lora_A.weight": [16, 3072],
+    "diffusion_model.single_blocks.9.linear1.lora_B.weight": [21504, 16],
+    "diffusion_model.single_blocks.9.linear2.lora_A.weight": [16, 15360],
+    "diffusion_model.single_blocks.9.linear2.lora_B.weight": [3072, 16],
+}

tests/backend/patches/lora_conversions/lora_state_dicts/flux_lora_diffusers_base_model_format.py

@@ -0,0 +1,766 @@
+# A sample state dict in the Diffusers FLUX LoRA format with base_model.model prefix.
+# These keys are based on the LoRA model in peft_adapter_model.safetensors
+state_dict_keys = {
+    "base_model.model.proj_out.lora_A.weight": [4, 3072],
+    "base_model.model.proj_out.lora_B.weight": [64, 4],
+    "base_model.model.single_transformer_blocks.0.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.0.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.0.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.0.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.0.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.0.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.0.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.0.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.0.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.0.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.1.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.1.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.1.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.1.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.1.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.1.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.1.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.1.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.1.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.1.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.10.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.10.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.10.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.10.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.10.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.10.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.10.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.10.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.10.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.10.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.11.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.11.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.11.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.11.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.11.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.11.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.11.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.11.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.11.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.11.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.12.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.12.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.12.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.12.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.12.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.12.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.12.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.12.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.12.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.12.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.13.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.13.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.13.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.13.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.13.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.13.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.13.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.13.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.13.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.13.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.14.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.14.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.14.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.14.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.14.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.14.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.14.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.14.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.14.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.14.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.15.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.15.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.15.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.15.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.15.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.15.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.15.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.15.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.15.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.15.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.16.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.16.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.16.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.16.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.16.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.16.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.16.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.16.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.16.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.16.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.17.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.17.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.17.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.17.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.17.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.17.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.17.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.17.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.17.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.17.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.18.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.18.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.18.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.18.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.18.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.18.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.18.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.18.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.18.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.18.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.19.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.19.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.19.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.19.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.19.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.19.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.19.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.19.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.19.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.19.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.2.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.2.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.2.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.2.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.2.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.2.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.2.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.2.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.2.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.2.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.20.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.20.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.20.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.20.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.20.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.20.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.20.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.20.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.20.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.20.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.21.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.21.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.21.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.21.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.21.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.21.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.21.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.21.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.21.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.21.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.22.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.22.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.22.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.22.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.22.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.22.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.22.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.22.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.22.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.22.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.23.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.23.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.23.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.23.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.23.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.23.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.23.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.23.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.23.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.23.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.24.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.24.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.24.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.24.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.24.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.24.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.24.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.24.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.24.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.24.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.25.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.25.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.25.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.25.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.25.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.25.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.25.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.25.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.25.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.25.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.26.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.26.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.26.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.26.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.26.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.26.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.26.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.26.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.26.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.26.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.27.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.27.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.27.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.27.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.27.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.27.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.27.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.27.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.27.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.27.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.28.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.28.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.28.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.28.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.28.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.28.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.28.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.28.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.28.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.28.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.29.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.29.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.29.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.29.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.29.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.29.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.29.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.29.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.29.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.29.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.3.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.3.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.3.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.3.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.3.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.3.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.3.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.3.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.3.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.3.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.30.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.30.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.30.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.30.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.30.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.30.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.30.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.30.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.30.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.30.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.31.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.31.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.31.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.31.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.31.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.31.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.31.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.31.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.31.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.31.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.32.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.32.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.32.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.32.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.32.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.32.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.32.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.32.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.32.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.32.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.33.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.33.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.33.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.33.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.33.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.33.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.33.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.33.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.33.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.33.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.34.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.34.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.34.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.34.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.34.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.34.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.34.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.34.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.34.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.34.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.35.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.35.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.35.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.35.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.35.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.35.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.35.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.35.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.35.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.35.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.36.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.36.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.36.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.36.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.36.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.36.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.36.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.36.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.36.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.36.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.37.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.37.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.37.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.37.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.37.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.37.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.37.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.37.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.37.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.37.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.4.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.4.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.4.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.4.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.4.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.4.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.4.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.4.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.4.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.4.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.5.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.5.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.5.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.5.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.5.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.5.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.5.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.5.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.5.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.5.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.6.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.6.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.6.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.6.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.6.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.6.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.6.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.6.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.6.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.6.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.7.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.7.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.7.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.7.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.7.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.7.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.7.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.7.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.7.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.7.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.8.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.8.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.8.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.8.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.8.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.8.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.8.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.8.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.8.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.8.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.9.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.9.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.9.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.9.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.9.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.9.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.single_transformer_blocks.9.proj_mlp.lora_A.weight": [4, 3072],
+    "base_model.model.single_transformer_blocks.9.proj_mlp.lora_B.weight": [12288, 4],
+    "base_model.model.single_transformer_blocks.9.proj_out.lora_A.weight": [4, 15360],
+    "base_model.model.single_transformer_blocks.9.proj_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.0.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.0.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.0.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.1.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.1.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.1.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.1.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.10.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.10.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.10.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.10.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.11.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.11.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.11.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.11.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.12.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.12.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.12.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.12.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.13.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.13.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.13.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.13.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.14.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.14.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.14.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.14.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.15.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.15.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.15.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.15.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.16.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.16.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.16.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.16.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.17.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.17.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.17.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.17.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.18.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.18.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.18.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.18.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.2.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.2.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.2.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.2.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.3.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.3.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.3.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.3.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.4.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.4.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.4.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.4.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.5.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.5.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.5.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.5.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.6.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.6.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.6.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.6.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.7.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.7.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.7.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.7.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.8.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.8.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.8.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.8.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.9.attn.add_k_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.add_k_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.add_q_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.add_q_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.add_v_proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.add_v_proj.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.to_add_out.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.to_add_out.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.to_k.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.to_k.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.to_out.0.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.to_out.0.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.to_q.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.to_q.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.attn.to_v.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.attn.to_v.lora_B.weight": [3072, 4],
+    "base_model.model.transformer_blocks.9.ff.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.ff.net.0.proj.lora_B.weight": [12288, 4],
+    "base_model.model.transformer_blocks.9.ff_context.net.0.proj.lora_A.weight": [4, 3072],
+    "base_model.model.transformer_blocks.9.ff_context.net.0.proj.lora_B.weight": [12288, 4],
+}