From 1e4cf09f6ead523a961ab10ea4a1e86897345cdf Mon Sep 17 00:00:00 2001 From: Vladimir Babin Date: Tue, 13 Oct 2026 08:37:21 +0300 Subject: [PATCH 1/2] Do not identify FLUX ControlNet checkpoints as FLUX main models --- .../backend/model_manager/configs/main.py | 7 ++ .../test_flux_controlnet_not_main_config.py | 72 +++++++++++++++++++ 2 files changed, 79 insertions(+) create mode 100644 tests/backend/model_manager/configs/test_flux_controlnet_not_main_config.py diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 052bbc970ca..694eae1533e 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -6,6 +6,10 @@ import torch from pydantic import BaseModel, Field +from invokeai.backend.flux.controlnet.state_dict_utils import ( + is_state_dict_instantx_controlnet, + is_state_dict_xlabs_controlnet, +) from invokeai.backend.ltx2 import checkpoint_layout as ltx2_layout from invokeai.backend.model_manager.checkpoint_prefix import COMFYUI_KEY_PREFIXES from invokeai.backend.model_manager.configs.base import ( @@ -590,6 +594,9 @@ def _validate_is_flux(cls, mod: ModelOnDisk) -> None: ): raise NotAMatchError("state dict does not look like a FLUX checkpoint") + if is_state_dict_xlabs_controlnet(state_dict) or is_state_dict_instantx_controlnet(state_dict): + raise NotAMatchError("state dict is a FLUX ControlNet, not a main model") + # Exclude FLUX.2 models - they have their own config class if _is_flux2_model(state_dict): raise NotAMatchError("model is a FLUX.2 model, not FLUX.1") diff --git a/tests/backend/model_manager/configs/test_flux_controlnet_not_main_config.py b/tests/backend/model_manager/configs/test_flux_controlnet_not_main_config.py new file mode 100644 index 00000000000..4cc12001dfe --- /dev/null +++ b/tests/backend/model_manager/configs/test_flux_controlnet_not_main_config.py @@ -0,0 +1,72 @@ +"""Regression test for XLabs FLUX ControlNet checkpoints being identified as FLUX main models. + +An XLabs FLUX ControlNet checkpoint (e.g. flux-canny-controlnet-v3.safetensors) contains two FLUX +double blocks, so it carries the `double_blocks.0.img_attn.norm.key_norm.scale` key that +`Main_Checkpoint_FLUX_Config` keys off. It also matches `is_state_dict_xlabs_controlnet`. Both config +classes used to match, and the factory tie-break ranks main models first, so the file was installed as +a main model and could not be used as a ControlNet (invoke-ai/InvokeAI#9832). +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_FLUX_Config +from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError +from invokeai.backend.model_manager.configs.main import Main_Checkpoint_FLUX_Config +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType + +_OVERRIDE_FIELDS: dict[str, object] = { + "hash": "blake3:fakehash", + "path": "/fake/models/flux-canny-controlnet-v3.safetensors", + "file_size": 1000, + "name": "flux-canny-controlnet-v3", + "description": "test", + "source": "test", + "source_type": "path", + "key": "test-key", +} + +# Keys mirroring an XLabs FLUX ControlNet checkpoint: the controlnet-specific keys that +# is_state_dict_xlabs_controlnet looks for, plus the two FLUX double blocks (and img_in / guidance_in) +# the file actually contains. There are no single_blocks or final_layer keys. +XLABS_CONTROLNET_KEYS = { + "controlnet_blocks.0.weight": SimpleNamespace(shape=(3072, 3072)), + "controlnet_blocks.0.bias": SimpleNamespace(shape=(3072,)), + "input_hint_block.0.weight": SimpleNamespace(shape=(16, 3, 3, 3)), + "input_hint_block.0.bias": SimpleNamespace(shape=(16,)), + "pos_embed_input.weight": SimpleNamespace(shape=(3072, 64)), + "pos_embed_input.bias": SimpleNamespace(shape=(3072,)), + "double_blocks.0.img_attn.norm.key_norm.scale": SimpleNamespace(shape=(128,)), + "double_blocks.0.img_attn.qkv.weight": SimpleNamespace(shape=(9216, 3072)), + "img_in.weight": SimpleNamespace(shape=(3072, 64)), + "guidance_in.out_layer.weight": SimpleNamespace(shape=(3072, 3072)), + "time_in.in_layer.weight": SimpleNamespace(shape=(3072, 256)), + "txt_in.weight": SimpleNamespace(shape=(3072, 4096)), + "vector_in.in_layer.weight": SimpleNamespace(shape=(3072, 768)), +} + + +def _make_mod(state_dict: dict[str, object]) -> MagicMock: + mod = MagicMock() + mod.load_state_dict.return_value = state_dict + mod.metadata.return_value = {} + return mod + + +def test_xlabs_controlnet_does_not_match_flux_main(): + mod = _make_mod(dict(XLABS_CONTROLNET_KEYS)) + + with pytest.raises(NotAMatchError, match="FLUX ControlNet"): + Main_Checkpoint_FLUX_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) + + +def test_xlabs_controlnet_matches_flux_controlnet(): + mod = _make_mod(dict(XLABS_CONTROLNET_KEYS)) + + config = ControlNet_Checkpoint_FLUX_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) + + assert config.base is BaseModelType.Flux + assert config.type is ModelType.ControlNet + assert config.format is ModelFormat.Checkpoint From fabf8ec9315316a1e8df2e1f2ad77b7c6d614015 Mon Sep 17 00:00:00 2001 From: Vladimir Babin Date: Wed, 14 Oct 2026 14:00:05 +0300 Subject: [PATCH 2/2] fix(model manager): drop the unreachable InstantX check from FLUX main validation InstantX ControlNets use diffusers keys and have no BFL double_blocks key, so the key check above already rejects them. --- invokeai/backend/model_manager/configs/main.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 694eae1533e..823c91261ad 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -7,7 +7,6 @@ from pydantic import BaseModel, Field from invokeai.backend.flux.controlnet.state_dict_utils import ( - is_state_dict_instantx_controlnet, is_state_dict_xlabs_controlnet, ) from invokeai.backend.ltx2 import checkpoint_layout as ltx2_layout @@ -594,7 +593,7 @@ def _validate_is_flux(cls, mod: ModelOnDisk) -> None: ): raise NotAMatchError("state dict does not look like a FLUX checkpoint") - if is_state_dict_xlabs_controlnet(state_dict) or is_state_dict_instantx_controlnet(state_dict): + if is_state_dict_xlabs_controlnet(state_dict): raise NotAMatchError("state dict is a FLUX ControlNet, not a main model") # Exclude FLUX.2 models - they have their own config class