feat(mm): support bria-3 controlnets

Author: psychedelicious

invokeai/backend/model_manager/config.py

@@ -547,6 +547,34 @@ def parse(cls, mod: ModelOnDisk) -> dict[str, Any]:
     def get_tag(cls) -> Tag:
         return Tag(f"{ModelType.Main.value}.{ModelFormat.Diffusers.value}.{BaseModelType.Bria.value}")
 
+class BriaControlNetDiffusersConfig(DiffusersConfigBase, ControlAdapterConfigBase, ModelConfigBase):
+    """Model config for Bria/Diffusers ControlNet models."""
+
+    type: Literal[ModelType.ControlNet] = ModelType.ControlNet
+    format: Literal[ModelFormat.Diffusers] = ModelFormat.Diffusers
+    base: Literal[BaseModelType.Bria] = BaseModelType.Bria
+
+    @classmethod
+    def matches(cls, mod: ModelOnDisk) -> bool:
+        if mod.path.is_file():
+            return False
+
+        config_path = mod.path / "config.json"
+        if config_path.exists():
+            with open(config_path) as file:
+                transformer_conf = json.load(file)
+            if transformer_conf["_class_name"] == "BriaTransformer2DModel":
+                return True
+
+        return False
+
+    @classmethod
+    def parse(cls, mod: ModelOnDisk) -> dict[str, Any]:
+        return {}
+
+    @classmethod
+    def get_tag(cls) -> Tag:
+        return Tag(f"{ModelType.ControlNet.value}.{ModelFormat.Diffusers.value}.{BaseModelType.Bria.value}")
 
 
 class IPAdapterConfigBase(ABC, BaseModel):
@@ -732,6 +760,7 @@ def get_model_discriminator_value(v: Any) -> str:
         Annotated[ControlLoRADiffusersConfig, ControlLoRADiffusersConfig.get_tag()],
         Annotated[LoRADiffusersConfig, LoRADiffusersConfig.get_tag()],
         Annotated[BriaDiffusersConfig, BriaDiffusersConfig.get_tag()],
+        Annotated[BriaControlNetDiffusersConfig, BriaControlNetDiffusersConfig.get_tag()],
         Annotated[T5EncoderConfig, T5EncoderConfig.get_tag()],
         Annotated[T5EncoderBnbQuantizedLlmInt8bConfig, T5EncoderBnbQuantizedLlmInt8bConfig.get_tag()],
         Annotated[TextualInversionFileConfig, TextualInversionFileConfig.get_tag()],

invokeai/backend/model_manager/legacy_probe.py

@@ -126,8 +126,6 @@ class ModelProbe(object):
     }
 
     CLASS2TYPE = {
-        "BriaPipeline": ModelType.Main,
-        "BriaTransformer2DModel": ModelType.ControlNet,
         "FluxPipeline": ModelType.Main,
         "StableDiffusionPipeline": ModelType.Main,
         "StableDiffusionInpaintPipeline": ModelType.Main,
@@ -876,8 +874,6 @@ def get_base_type(self) -> BaseModelType:
                 return BaseModelType.StableDiffusion3
             elif transformer_conf["_class_name"] == "CogView4Transformer2DModel":
                 return BaseModelType.CogView4
-            elif transformer_conf["_class_name"] == "BriaTransformer2DModel":
-                return BaseModelType.Bria
             else:
                 raise InvalidModelConfigException(f"Unknown base model for {self.model_path}")
 
@@ -1027,9 +1023,6 @@ def get_base_type(self) -> BaseModelType:
         if config.get("_class_name", None) == "FluxControlNetModel":
             return BaseModelType.Flux
 
-        if config.get("_class_name", None) == "BriaTransformer2DModel":
-            return BaseModelType.Bria
-
         # no obvious way to distinguish between sd2-base and sd2-768
         dimension = config["cross_attention_dim"]
         if dimension == 768: