Integrate Bria 3.1/3.2 Models and ControlNet Pipelines into InvokeAI

invokeai/app/invocations/bria_controlnet.py

@@ -0,0 +1,117 @@
+from PIL import Image
+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,
+    WithBoard,
+    WithMetadata,
+)
+from invokeai.app.invocations.model import ModelIdentifierField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.bria.controlnet_bria import BRIA_CONTROL_MODES
+from invokeai.invocation_api import Classification, ImageOutput
+
+DEPTH_SMALL_V2_URL = "depth-anything/Depth-Anything-V2-Small-hf"
+HF_LLLYASVIEL = "https://huggingface.co/lllyasviel/Annotators/resolve/main/"
+
+
+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("bria_controlnet_output")
+class BriaControlNetOutput(BaseInvocationOutput):
+    """Bria ControlNet info"""
+
+    control: BriaControlNetField = OutputField(description=FieldDescriptions.control)
+
+
+@invocation(
+    "bria_controlnet",
+    title="ControlNet - Bria",
+    tags=["controlnet", "bria"],
+    category="controlnet",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class BriaControlNetInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """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:
+        image_in = resize_img(context.images.get_pil(self.control_image.image_name))
+        if self.control_mode == "colorgrid":
+            control_image = tile(64, image_in)
+        elif self.control_mode == "recolor":
+            control_image = convert_to_grayscale(image_in)
+        elif self.control_mode == "tile":
+            control_image = tile(16, image_in)
+        else:
+            control_image = image_in
+
+        control_image = resize_img(control_image)
+        image_dto = context.images.save(image=control_image)
+        return BriaControlNetOutput(
+            control=BriaControlNetField(
+                image=ImageField(image_name=image_dto.image_name),
+                model=self.control_model,
+                mode=self.control_mode,
+                conditioning_scale=self.control_weight,
+            ),
+        )
+
+
+RATIO_CONFIGS_1024 = {
+    0.6666666666666666: {"width": 832, "height": 1248},
+    0.7432432432432432: {"width": 880, "height": 1184},
+    0.8028169014084507: {"width": 912, "height": 1136},
+    1.0: {"width": 1024, "height": 1024},
+    1.2456140350877194: {"width": 1136, "height": 912},
+    1.3454545454545455: {"width": 1184, "height": 880},
+    1.4339622641509433: {"width": 1216, "height": 848},
+    1.5: {"width": 1248, "height": 832},
+    1.5490196078431373: {"width": 1264, "height": 816},
+    1.62: {"width": 1296, "height": 800},
+    1.7708333333333333: {"width": 1360, "height": 768},
+}
+
+
+
+def convert_to_grayscale(image: Image.Image) -> Image.Image:
+    gray_image = image.convert("L").convert("RGB")
+    return gray_image
+
+
+def tile(downscale_factor: int, input_image: Image.Image) -> Image.Image:
+    control_image = input_image.resize(
+        (input_image.size[0] // downscale_factor, input_image.size[1] // downscale_factor)
+    ).resize(input_image.size, Image.Resampling.NEAREST)
+    return control_image
+
+
+def resize_img(control_image: Image.Image) -> Image.Image:
+    image_ratio = control_image.width / control_image.height
+    ratio = min(RATIO_CONFIGS_1024.keys(), key=lambda k: abs(k - image_ratio))
+    to_height = RATIO_CONFIGS_1024[ratio]["height"]
+    to_width = RATIO_CONFIGS_1024[ratio]["width"]
+    resized_image = control_image.resize((to_width, to_height), resample=Image.Resampling.LANCZOS)
+    return resized_image

invokeai/app/invocations/bria_decoder.py

@@ -0,0 +1,50 @@
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from PIL import Image
+
+from invokeai.app.invocations.model import VAEField
+from invokeai.app.invocations.primitives import FieldDescriptions, Input, InputField, LatentsField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.invocation_api import BaseInvocation, Classification, ImageOutput, invocation
+
+
+@invocation(
+    "bria_decoder",
+    title="Decoder - Bria",
+    tags=["image", "bria"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class BriaDecoderInvocation(BaseInvocation):
+    """
+    Decode Bria latents to an image.
+    """
+
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+    )
+    latents: LatentsField = InputField(
+        description=FieldDescriptions.latents,
+        input=Input.Connection,
+    )
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        latents = context.tensors.load(self.latents.latents_name)
+        latents = latents.view(1, 64, 64, 4, 2, 2).permute(0, 3, 1, 4, 2, 5).reshape(1, 4, 128, 128)
+
+        with context.models.load(self.vae.vae) as vae:
+            assert isinstance(vae, AutoencoderKL)
+            latents = latents / vae.config.scaling_factor
+            latents = latents.to(device=vae.device, dtype=vae.dtype)
+
+            decoded_output = vae.decode(latents)
+            image = decoded_output.sample
+
+        # Convert to numpy with proper gradient handling
+        image = ((image.clamp(-1, 1) + 1) / 2 * 255).cpu().detach().permute(0, 2, 3, 1).numpy().astype("uint8")[0]
+        img = Image.fromarray(image)
+        image_dto = context.images.save(image=img)
+        return ImageOutput.build(image_dto)

invokeai/app/invocations/bria_denoiser.py

@@ -0,0 +1,197 @@
+from typing import Callable, List, Tuple
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+
+from invokeai.app.invocations.bria_controlnet import BriaControlNetField
+from invokeai.app.invocations.bria_latent_noise import BriaLatentNoiseOutput
+from invokeai.app.invocations.fields import FluxConditioningField, Input, InputField, LatentsField, OutputField
+from invokeai.app.invocations.model import SubModelType, T5EncoderField, TransformerField, VAEField
+from invokeai.app.invocations.primitives import BaseInvocationOutput, FieldDescriptions
+from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.bria.controlnet_bria import BriaControlModes, BriaMultiControlNetModel
+from invokeai.backend.bria.controlnet_utils import prepare_control_images
+from invokeai.backend.bria.pipeline_bria_controlnet import BriaControlNetPipeline
+from invokeai.backend.bria.transformer_bria import BriaTransformer2DModel
+from invokeai.backend.model_manager.taxonomy import BaseModelType
+from invokeai.backend.stable_diffusion.extensions.preview import PipelineIntermediateState
+from invokeai.invocation_api import BaseInvocation, Classification, invocation, invocation_output
+
+
+@invocation_output("bria_denoise_output")
+class BriaDenoiseInvocationOutput(BaseInvocationOutput):
+    latents: LatentsField = OutputField(description=FieldDescriptions.latents)
+
+
+@invocation(
+    "bria_denoise",
+    title="Denoise - Bria",
+    tags=["image", "bria"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class BriaDenoiseInvocation(BaseInvocation):
+
+    """
+    Denoise Bria latents using a Bria Pipeline.
+    """
+
+    num_steps: int = InputField(
+        default=30, title="Number of Steps", description="The number of steps to use for the denoiser"
+    )
+    guidance_scale: float = InputField(
+        default=5.0, title="Guidance Scale", description="The guidance scale to use for the denoiser"
+    )
+
+    transformer: TransformerField = InputField(
+        description="Bria model (Transformer) to load",
+        input=Input.Connection,
+        title="Transformer",
+    )
+    t5_encoder: T5EncoderField = InputField(
+        title="T5Encoder",
+        description=FieldDescriptions.t5_encoder,
+        input=Input.Connection,
+    )
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        input=Input.Connection,
+        title="VAE",
+    )
+    height: int = InputField(
+        default=1024,
+        title="Height",
+        description="The height of the output image",
+    )
+    width: int = InputField(
+        default=1024,
+        title="Width",
+        description="The width of the output image",
+    )
+    latent_noise: BriaLatentNoiseOutput = InputField(
+        description="Latent noise to denoise",
+        input=Input.Connection,
+        title="Latent Noise",
+    )
+    pos_embeds: FluxConditioningField = InputField(
+        description="Positive Prompt Embeds",
+        input=Input.Connection,
+        title="Positive Prompt Embeds",
+    )
+    neg_embeds: FluxConditioningField = InputField(
+        description="Negative Prompt Embeds",
+        input=Input.Connection,
+        title="Negative Prompt Embeds",
+    )
+    control: BriaControlNetField | list[BriaControlNetField] | None = InputField(
+        description="ControlNet",
+        input=Input.Connection,
+        title="ControlNet",
+        default=None,
+    )
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
+        latents = context.tensors.load(self.latent_noise.latents.latents_name)
+        pos_embeds = context.tensors.load(self.pos_embeds.conditioning_name)
+        neg_embeds = context.tensors.load(self.neg_embeds.conditioning_name)
+        latent_image_ids = context.tensors.load(self.latent_noise.latent_image_ids.latents_name)
+        scheduler_identifier = self.transformer.transformer.model_copy(update={"submodel_type": SubModelType.Scheduler})
+
+        device = None
+        dtype = None
+        with (
+            context.models.load(self.transformer.transformer) as transformer,
+            context.models.load(scheduler_identifier) as scheduler,
+            context.models.load(self.vae.vae) as vae,
+            context.models.load(self.t5_encoder.text_encoder) as t5_encoder,
+            context.models.load(self.t5_encoder.tokenizer) as t5_tokenizer,
+        ):
+            assert isinstance(transformer, BriaTransformer2DModel)
+            assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler)
+            assert isinstance(vae, AutoencoderKL)
+            dtype = transformer.dtype
+            device = transformer.device
+            latents, pos_embeds, neg_embeds = (x.to(device, dtype) for x in (latents, pos_embeds, neg_embeds))
+
+            control_model, control_images, control_modes, control_scales = None, None, None, None
+            if self.control is not None:
+                control_model, control_images, control_modes, control_scales = self._prepare_multi_control(
+                    context=context,
+                    vae=vae,
+                    width=self.width,
+                    height=self.height,
+                    device=vae.device,
+                )
+
+
+            pipeline = BriaControlNetPipeline(
+                transformer=transformer,
+                scheduler=scheduler,
+                vae=vae,
+                text_encoder=t5_encoder,
+                tokenizer=t5_tokenizer,
+                controlnet=control_model,
+            )
+            pipeline.to(device=transformer.device, dtype=transformer.dtype)
+
+            output_latents = pipeline(
+                control_image=control_images,
+                control_mode=control_modes,
+                width=self.width,
+                height=self.height,
+                controlnet_conditioning_scale=control_scales,
+                num_inference_steps=self.num_steps,
+                max_sequence_length=128,
+                guidance_scale=self.guidance_scale,
+                latents=latents,
+                latent_image_ids=latent_image_ids,
+                prompt_embeds=pos_embeds,
+                negative_prompt_embeds=neg_embeds,
+                output_type="latent",
+                step_callback=_build_step_callback(context),
+            )[0]
+
+
+
+        assert isinstance(output_latents, torch.Tensor)
+        saved_input_latents_tensor = context.tensors.save(output_latents)
+        return BriaDenoiseInvocationOutput(latents=LatentsField(latents_name=saved_input_latents_tensor))
+
+    def _prepare_multi_control(
+        self, context: InvocationContext, vae: AutoencoderKL, width: int, height: int, device: torch.device
+    ) -> Tuple[BriaMultiControlNetModel, List[torch.Tensor], List[int], 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_modes.append(BriaControlModes[controlnet.mode].value)
+                control_scales.append(controlnet.conditioning_scale)
+                try:
+                    control_images.append(context.images.get_pil(controlnet.image.image_name))
+                except Exception:
+                    raise FileNotFoundError(
+                        f"Control image {controlnet.image.image_name} not found. Make sure not to delete the preprocessed image before finishing the pipeline."
+                    )
+
+        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
+
+
+def _build_step_callback(context: InvocationContext) -> Callable[[PipelineIntermediateState], None]:
+    def step_callback(state: PipelineIntermediateState) -> None:
+        return 
+        context.util.sd_step_callback(state, BaseModelType.Bria)
+
+    return step_callback