Added support for bria's controlnet

Ilan Tchenak · Ubuntu · commit 1076d8eb067f · 2025-07-10T20:15:14.000Z
diff --git a/invokeai/backend/bria/controlnet_bria.py b/invokeai/backend/bria/controlnet_bria.py
diff --git a/invokeai/backend/bria/controlnet_utils.py b/invokeai/backend/bria/controlnet_utils.py
@@ -0,0 +1,69 @@
+from typing import List, Tuple
+from PIL import Image
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+
+from diffusers.image_processor import VaeImageProcessor
+
+import torch
+
+
+
+@torch.no_grad()
+def prepare_control_images(
+    vae: AutoencoderKL,
+    control_images: list[Image.Image],
+    control_modes: list[int],
+    width: int,
+    height: int,
+    device: torch.device,
+) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+    
+    tensored_control_images = []
+    tensored_control_modes = []
+    for idx, control_image_ in enumerate(control_images):
+        tensored_control_image = _prepare_image(
+            image=control_image_,
+            width=width,
+            height=height,
+            device=device,
+            dtype=vae.dtype,
+        )
+        height, width = tensored_control_image.shape[-2:]
+
+        # vae encode
+        tensored_control_image = vae.encode(tensored_control_image).latent_dist.sample()
+        tensored_control_image = (tensored_control_image) * 16
+
+        # pack
+        height_control_image, width_control_image = tensored_control_image.shape[2:]
+        tensored_control_image = _pack_latents(
+            tensored_control_image,
+            height_control_image,
+            width_control_image,
+        )
+        tensored_control_images.append(tensored_control_image)
+        tensored_control_modes.append(torch.tensor(control_modes[idx]).expand(
+            tensored_control_image.shape[0]).to(device, dtype=torch.long))
+
+    return tensored_control_images, tensored_control_modes
+
+def _prepare_image(
+    image: Image.Image,
+    width: int,
+    height: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    image = image.convert("RGB")
+    image = VaeImageProcessor(vae_scale_factor=16).preprocess(image, height=height, width=width)
+    image = image.repeat_interleave(1, dim=0)
+    image = image.to(device=device, dtype=dtype)
+    return image
+
+def _pack_latents(latents, height, width):
+    latents = latents.view(1, 4, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    latents = latents.reshape(1, (height // 2) * (width // 2), 16)
+
+    return latents
+
diff --git a/invokeai/backend/model_manager/legacy_probe.py b/invokeai/backend/model_manager/legacy_probe.py
@@ -126,7 +126,7 @@ class ModelProbe(object):
 
     CLASS2TYPE = {
         "BriaPipeline": ModelType.Main,
-        "BriaControlNetModel": ModelType.ControlNet,
+        "BriaTransformer2DModel": ModelType.ControlNet,
         "FluxPipeline": ModelType.Main,
         "StableDiffusionPipeline": ModelType.Main,
         "StableDiffusionInpaintPipeline": ModelType.Main,
@@ -1014,7 +1014,7 @@ def get_base_type(self) -> BaseModelType:
         if config.get("_class_name", None) == "FluxControlNetModel":
             return BaseModelType.Flux
 
-        if config.get("_class_name", None) == "BriaControlNetModel":
+        if config.get("_class_name", None) == "BriaTransformer2DModel":
             return BaseModelType.Bria
 
         # no obvious way to distinguish between sd2-base and sd2-768
diff --git a/invokeai/backend/model_manager/load/model_loaders/bria.py b/invokeai/backend/model_manager/load/model_loaders/bria.py
@@ -31,14 +31,11 @@ def _load_model(
         if isinstance(config, ControlNetCheckpointConfig):
             raise NotImplementedError("CheckpointConfigBase is not implemented for Bria models.")
 
-        if submodel_type is None:
-            raise Exception("A submodel type must be provided when loading control net pipelines.")
-
         model_path = Path(config.path)
-        load_class = self.get_hf_load_class(model_path, submodel_type)
+        load_class = self.get_hf_load_class(model_path)
         repo_variant = config.repo_variant if isinstance(config, ControlNetDiffusersConfig) else None
         variant = repo_variant.value if repo_variant else None
-        model_path = model_path / submodel_type.value
+        model_path = model_path
 
         dtype = self._torch_dtype
 
diff --git a/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py b/invokeai/backend/model_manager/load/model_loaders/generic_diffusers.py
@@ -84,6 +84,9 @@ def _hf_definition_to_type(self, module: str, class_name: str) -> ModelMixin:  #
         ]:
             if module == "transformer_bria":
                 module = "invokeai.backend.bria.transformer_bria"
+            elif class_name == "BriaTransformer2DModel":
+                class_name = "BriaControlNetModel"
+                module = "invokeai.backend.bria.controlnet_bria"
             res_type = sys.modules[module]
         else:
             res_type = sys.modules["diffusers"].pipelines
diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/BriaControlNetModelFieldInputComponent.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/BriaControlNetModelFieldInputComponent.tsx
@@ -6,8 +6,8 @@ import type {
   BriaControlNetModelFieldInputTemplate,
 } from 'features/nodes/types/field';
 import { memo, useCallback } from 'react';
-import { useBriaModels } from 'services/api/hooks/modelsByType';
-import type { MainModelConfig } from 'services/api/types';
+import { useBriaControlNetModels } from 'services/api/hooks/modelsByType';
+import type { ControlNetModelConfig } from 'services/api/types';
 
 import type { FieldComponentProps } from './types';
 
@@ -16,9 +16,9 @@ type Props = FieldComponentProps<BriaControlNetModelFieldInputInstance, BriaCont
 const BriaControlNetModelFieldInputComponent = (props: Props) => {
   const { nodeId, field } = props;
   const dispatch = useAppDispatch();
-  const [modelConfigs, { isLoading }] = useBriaModels();
+  const [modelConfigs, { isLoading }] = useBriaControlNetModels();
   const onChange = useCallback(
-    (value: MainModelConfig | null) => {
+    (value: ControlNetModelConfig | null) => {
       if (!value) {
         return;
       }
diff --git a/invokeai/nodes/bria_nodes/bria_controlnet.py b/invokeai/nodes/bria_nodes/bria_controlnet.py
@@ -0,0 +1,57 @@
+from invokeai.backend.bria.controlnet_bria import BRIA_CONTROL_MODES
+from pydantic import BaseModel, Field
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import FieldDescriptions, ImageField, InputField, OutputField, UIType
+from invokeai.app.invocations.model import ModelIdentifierField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+class BriaControlNetField(BaseModel):
+    image: ImageField = Field(description="The control image")
+    model: ModelIdentifierField = Field(description="The ControlNet model to use")
+    mode: BRIA_CONTROL_MODES = Field(description="The mode of the ControlNet")
+    conditioning_scale: float = Field(description="The weight given to the ControlNet")
+
+@invocation_output("flux_controlnet_output")
+class BriaControlNetOutput(BaseInvocationOutput):
+    """FLUX ControlNet info"""
+
+    control: BriaControlNetField = OutputField(description=FieldDescriptions.control)
+
+
+@invocation(
+    "bria_controlnet",
+    title="Bria ControlNet",
+    tags=["controlnet", "bria"],
+    category="controlnet",
+    version="1.0.0",
+)
+class BriaControlNetInvocation(BaseInvocation):
+    """Collect Bria ControlNet info to pass to denoiser node."""
+
+    control_image: ImageField = InputField(description="The control image")
+    control_model: ModelIdentifierField = InputField(
+        description=FieldDescriptions.controlnet_model, ui_type=UIType.BriaControlNetModel
+    )
+    control_mode: BRIA_CONTROL_MODES = InputField(
+        default="depth", description="The mode of the ControlNet"
+    )
+    control_weight: float = InputField(
+        default=1.0, ge=-1, le=2, description="The weight given to the ControlNet"
+    )
+
+    def invoke(self, context: InvocationContext) -> BriaControlNetOutput:
+        return BriaControlNetOutput(
+            control=BriaControlNetField(
+                image=self.control_image,
+                model=self.control_model,
+                mode=self.control_mode,
+                conditioning_scale=self.control_weight,
+            ),
+        )
diff --git a/invokeai/nodes/bria_nodes/bria_denoiser.py b/invokeai/nodes/bria_nodes/bria_denoiser.py
@@ -1,16 +1,15 @@
+from typing import List, Tuple
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from invokeai.backend.bria.controlnet_bria import BriaControlModes, BriaMultiControlNetModel
+from invokeai.backend.bria.controlnet_utils import prepare_control_images
+from invokeai.nodes.bria_nodes.bria_controlnet import BriaControlNetField
+
 import torch
 from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 
-from invokeai.app.invocations.fields import Input, InputField
-from invokeai.app.invocations.model import SubModelType, TransformerField
-from invokeai.app.invocations.primitives import (
-    BaseInvocationOutput,
-    FieldDescriptions,
-    Input,
-    InputField,
-    LatentsField,
-    OutputField,
-)
+from invokeai.app.invocations.fields import Input, InputField, LatentsField, OutputField
+from invokeai.app.invocations.model import SubModelType, TransformerField, VAEField
+from invokeai.app.invocations.primitives import BaseInvocationOutput, FieldDescriptions
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.invocation_api import BaseInvocation, Classification, InputField, invocation, invocation_output
 
@@ -43,6 +42,11 @@ class BriaDenoiseInvocation(BaseInvocation):
         input=Input.Connection,
         title="Transformer",
     )
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+        title="VAE",
+    )
     latents: LatentsField = InputField(
         description="Latents to denoise",
         input=Input.Connection,
@@ -68,6 +72,12 @@ class BriaDenoiseInvocation(BaseInvocation):
         input=Input.Connection,
         title="Text IDs",
     )
+    control: BriaControlNetField | list[BriaControlNetField] | None = InputField(
+        description="ControlNet",
+        input=Input.Connection,
+        title="ControlNet",
+        default = None,
+    )
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
@@ -83,16 +93,28 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
         with (
             context.models.load(self.transformer.transformer) as transformer,
             context.models.load(scheduler_identifier) as scheduler,
+            context.models.load(self.vae.vae) as vae,
         ):
             assert isinstance(transformer, BriaTransformer2DModel)
             assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler)
+            assert isinstance(vae, AutoencoderKL)
             dtype = transformer.dtype
             device = transformer.device
             latents, pos_embeds, neg_embeds = map(lambda x: x.to(device, dtype), (latents, pos_embeds, neg_embeds))
             prompt_embeds = torch.cat([neg_embeds, pos_embeds]) if self.guidance_scale > 1 else pos_embeds
 
             sigmas = get_original_sigmas(1000, self.num_steps)
             timesteps, _ = retrieve_timesteps(scheduler, self.num_steps, device, None, sigmas, mu=0.0)
+            width, height = 1024, 1024
+            if self.control is not None:
+                control_model, control_images, control_modes, control_scales = self._prepare_multi_control(
+                        context=context,
+                        vae=vae,
+                        width=width,
+                        height=height,
+                        device=device,
+                        
+                    )
 
             for t in timesteps:
                 # Prepare model input efficiently
@@ -101,11 +123,21 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
                 else:
                     latent_model_input = latents
                 
-                # Prepare timestep tensor efficiently
-                if isinstance(t, torch.Tensor):
-                    timestep_tensor = t.expand(latent_model_input.shape[0])
-                else:
-                    timestep_tensor = torch.tensor([t] * latent_model_input.shape[0], device=device, dtype=torch.float32)
+                timestep_tensor = t.expand(latent_model_input.shape[0])
+
+                controlnet_block_samples, controlnet_single_block_samples = None, None
+                if self.control is not None:
+                    controlnet_block_samples, controlnet_single_block_samples = control_model(
+                        hidden_states=latents,
+                        controlnet_cond=control_images, # type: ignore
+                        controlnet_mode=control_modes, # type: ignore
+                        conditioning_scale=control_scales, # type: ignore
+                        timestep=timestep_tensor,
+                        encoder_hidden_states=prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_image_ids,
+                        return_dict=False,
+                    )
 
                 noise_pred = transformer(
                         latent_model_input,
@@ -115,6 +147,8 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
                         txt_ids=text_ids,
                         guidance=None,
                         return_dict=False,
+                        controlnet_block_samples=controlnet_block_samples,
+                        controlnet_single_block_samples=controlnet_single_block_samples,
                     )[0]
 
                 if self.guidance_scale > 1:
@@ -131,3 +165,35 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
         saved_input_latents_tensor = context.tensors.save(latents)
         latents_output = LatentsField(latents_name=saved_input_latents_tensor)
         return BriaDenoiseInvocationOutput(latents=latents_output)
+
+
+
+    def _prepare_multi_control(
+        self,
+        context: InvocationContext,
+        vae: AutoencoderKL,
+        width: int,
+        height: int,
+        device: torch.device
+    ) -> Tuple[BriaMultiControlNetModel, List[torch.Tensor], List[torch.Tensor], List[float]]:
+
+        control = self.control if isinstance(self.control, list) else [self.control]
+        control_images, control_models, control_modes, control_scales = [], [], [], []
+        for controlnet in control:
+            if controlnet is not None:
+                control_models.append(context.models.load(controlnet.model).model)
+                control_images.append(context.images.get_pil(controlnet.image.image_name))
+                control_modes.append(BriaControlModes[controlnet.mode].value)   
+                control_scales.append(controlnet.conditioning_scale)
+        
+        control_model = BriaMultiControlNetModel(control_models).to(device)
+        tensored_control_images, tensored_control_modes = prepare_control_images(
+            vae=vae,
+            control_images=control_images, 
+            control_modes=control_modes, 
+            width=width,
+            height=height,
+            device=device, 
+            )
+        return control_model, tensored_control_images, tensored_control_modes, control_scales
+        
