tests: flux variant probe testing

invokeai/backend/model_manager/legacy_probe.py

@@ -569,6 +569,8 @@ def get_variant_type(self) -> ModelVariantType:
 
             if in_channels is None:
                 # If we cannot find the in_channels, we assume that this is a normal variant. Log a warning.
+                # If this occurs, we should add a test case for the affected model here:
+                # tests/backend/flux/test_flux_state_dict_utils.py
                 logger.warning(
                     f"{self.model_path} does not have img_in.weight or model.diffusion_model.img_in.weight key. Assuming normal variant."
                 )

tests/backend/flux/test_flux_state_dict_utils.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+
+import pytest
+
+from invokeai.backend.flux.flux_state_dict_utils import get_flux_in_channels_from_state_dict
+from invokeai.backend.model_manager.config import ModelOnDisk
+
+test_cases = [
+    # Unquantized
+    ("FLUX Dev.safetensors", 64),
+    ("FLUX Schnell.safetensors", 64),
+    ("FLUX Fill.safetensors", 384),
+    # BNB-NF4 quantized
+    ("FLUX Dev (Quantized).safetensors", 1),  # BNB-NF4
+    ("FLUX Schnell (Quantized).safetensors", 1),  # BNB-NF4
+    # GGUF quantized FLUX Fill
+    ("flux1-fill-dev-Q8_0.gguf", 384),
+    # Fine-tune w/ "model.diffusion_model.img_in.weight" instead of "img_in.weight"
+    ("midjourneyReplica_flux1Dev.safetensors", 64),
+    # Not a FLUX model, testing fallback case
+    ("Noodles Style.safetensors", None),
+]
+
+
+@pytest.mark.parametrize("model_file_name,expected_in_channels", test_cases)
+def test_get_flux_in_channels_from_state_dict(model_file_name: str, expected_in_channels: int, override_model_loading):
+    model_path = Path(f"tests/test_model_probe/stripped_models/{model_file_name}")
+
+    mod = ModelOnDisk(model_path)
+
+    state_dict = mod.load_state_dict()
+
+    in_channels = get_flux_in_channels_from_state_dict(state_dict)
+
+    assert in_channels == expected_in_channels

tests/conftest.py

@@ -94,6 +94,9 @@ def override_model_loading(monkeypatch):
     monkeypatch.setattr(safetensors.torch, "load", load_stripped_model)
     monkeypatch.setattr(safetensors.torch, "load_file", load_stripped_model)
     monkeypatch.setattr(gguf_loaders, "gguf_sd_loader", load_stripped_model)
+    monkeypatch.setattr("invokeai.backend.model_manager.config.gguf_sd_loader", load_stripped_model)
+    monkeypatch.setattr("invokeai.backend.model_manager.util.model_util.gguf_sd_loader", load_stripped_model)
+    monkeypatch.setattr("invokeai.backend.model_manager.legacy_probe.gguf_sd_loader", load_stripped_model)
 
     def fake_scan(*args, **kwargs):
         return SimpleNamespace(infected_files=0, scan_err=None)