addded bria nodes for bria3.1 and bria3.2

Ubuntu · Ubuntu · commit 9e5e1ec0daf2 · 2025-07-09T18:38:09.000Z
diff --git a/invokeai/nodes/__init__.py b/invokeai/nodes/__init__.py
@@ -0,0 +1 @@
+from .bria_nodes import *
diff --git a/invokeai/nodes/bria_nodes/bria_decoder.py b/invokeai/nodes/bria_nodes/bria_decoder.py
@@ -0,0 +1,46 @@
+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="Bria Decoder",
+    tags=["image", "bria"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class BriaDecoderInvocation(BaseInvocation):
+    latents: LatentsField = InputField(
+        description=FieldDescriptions.latents,
+        input=Input.Connection,
+    )
+    vae: VAEField = InputField(
+        description=FieldDescriptions.vae,
+        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)
diff --git a/invokeai/nodes/bria_nodes/bria_denoiser.py b/invokeai/nodes/bria_nodes/bria_denoiser.py
@@ -0,0 +1,133 @@
+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.services.shared.invocation_context import InvocationContext
+from invokeai.invocation_api import BaseInvocation, Classification, InputField, invocation, invocation_output
+
+from invokeai.backend.bria.pipeline import get_original_sigmas, retrieve_timesteps
+from invokeai.backend.bria.transformer_bria import BriaTransformer2DModel
+
+@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):
+    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",
+    )
+    latents: LatentsField = InputField(
+        description="Latents to denoise",
+        input=Input.Connection,
+        title="Latents",
+    )
+    latent_image_ids: LatentsField = InputField(
+        description="Latent Image IDs to denoise",
+        input=Input.Connection,
+        title="Latent Image IDs",
+    )
+    pos_embeds: LatentsField = InputField(
+        description="Positive Prompt Embeds",
+        input=Input.Connection,
+        title="Positive Prompt Embeds",
+    )
+    neg_embeds: LatentsField = InputField(
+        description="Negative Prompt Embeds",
+        input=Input.Connection,
+        title="Negative Prompt Embeds",
+    )
+    text_ids: LatentsField = InputField(
+        description="Text IDs",
+        input=Input.Connection,
+        title="Text IDs",
+    )
+
+    @torch.no_grad()
+    def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
+        latents = context.tensors.load(self.latents.latents_name)
+        pos_embeds = context.tensors.load(self.pos_embeds.latents_name)
+        neg_embeds = context.tensors.load(self.neg_embeds.latents_name)
+        text_ids = context.tensors.load(self.text_ids.latents_name)
+        latent_image_ids = context.tensors.load(self.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,
+        ):
+            assert isinstance(transformer, BriaTransformer2DModel)
+            assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler)
+            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)
+
+            for t in timesteps:
+                # Prepare model input efficiently
+                if self.guidance_scale > 1:
+                    latent_model_input = torch.cat([latents] * 2)
+                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)
+
+                noise_pred = transformer(
+                        latent_model_input,
+                        encoder_hidden_states=prompt_embeds,
+                        timestep=timestep_tensor,
+                        img_ids=latent_image_ids,
+                        txt_ids=text_ids,
+                        guidance=None,
+                        return_dict=False,
+                    )[0]
+
+                if self.guidance_scale > 1:
+                    noise_uncond, noise_text = noise_pred.chunk(2)
+                    noise_pred = noise_uncond + self.guidance_scale * (noise_text - noise_uncond)
+
+                # Convert timestep for scheduler
+                t_step = float(t.item()) if isinstance(t, torch.Tensor) else float(t)
+                
+                # Use scheduler step with proper dtypes
+                latents = scheduler.step(noise_pred, t_step, latents, return_dict=False)[0]
+
+        assert isinstance(latents, torch.Tensor)
+        saved_input_latents_tensor = context.tensors.save(latents)
+        latents_output = LatentsField(latents_name=saved_input_latents_tensor)
+        return BriaDenoiseInvocationOutput(latents=latents_output)
diff --git a/invokeai/nodes/bria_nodes/bria_latent_sampler.py b/invokeai/nodes/bria_nodes/bria_latent_sampler.py
@@ -0,0 +1,79 @@
+import torch
+
+from invokeai.app.invocations.fields import Input, InputField
+from invokeai.app.invocations.model import TransformerField
+from invokeai.app.invocations.primitives import (
+    BaseInvocationOutput,
+    FieldDescriptions,
+    Input,
+    LatentsField,
+    OutputField,
+)
+from invokeai.backend.model_manager.config import MainDiffusersConfig
+from invokeai.invocation_api import (
+    BaseInvocation,
+    Classification,
+    InputField,
+    InvocationContext,
+    invocation,
+    invocation_output,
+)
+
+
+@invocation_output("bria_latent_sampler_output")
+class BriaLatentSamplerInvocationOutput(BaseInvocationOutput):
+    """Base class for nodes that output a CogView text conditioning tensor."""
+
+    latents: LatentsField = OutputField(description=FieldDescriptions.cond)
+    latent_image_ids: LatentsField = OutputField(description=FieldDescriptions.cond)
+
+
+@invocation(
+    "bria_latent_sampler",
+    title="Latent Sampler - Bria",
+    tags=["image", "bria"],
+    category="image",
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class BriaLatentSamplerInvocation(BaseInvocation):
+    seed: int = InputField(
+        default=42,
+        title="Seed",
+        description="The seed to use for the latent sampler",
+    )
+    transformer: TransformerField = InputField(
+        description="Bria model (Transformer) to load",
+        input=Input.Connection,
+        title="Transformer",
+    )
+
+    def invoke(self, context: InvocationContext) -> BriaLatentSamplerInvocationOutput:
+        device = torch.device("cuda")
+        transformer_config = context.models.get_config(self.transformer.transformer)
+        if not isinstance(transformer_config, MainDiffusersConfig):
+            raise ValueError("Transformer config is not a MainDiffusersConfig")
+        # TODO: get latent channels from transformer config
+        latent_channels = 16
+        latent_height, latent_width = 128, 128
+        shrunk = latent_channels // 4
+        gen = torch.Generator(device=device).manual_seed(self.seed)
+
+        noise4d = torch.randn((1, shrunk, latent_height, latent_width), device=device, generator=gen)
+        latents = noise4d.view(1, shrunk, latent_height // 2, 2, latent_width // 2, 2).permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(1, (latent_height // 2) * (latent_width // 2), shrunk * 4)
+
+        latent_image_ids = torch.zeros((latent_height // 2, latent_width // 2, 3), device=device, dtype=torch.long)
+        latent_image_ids[..., 1] = torch.arange(latent_height // 2, device=device)[:, None]
+        latent_image_ids[..., 2] = torch.arange(latent_width // 2, device=device)[None, :]
+        latent_image_ids = latent_image_ids.view(-1, 3)
+
+        saved_latents_tensor = context.tensors.save(latents)
+        saved_latent_image_ids_tensor = context.tensors.save(latent_image_ids)
+        latents_output = LatentsField(latents_name=saved_latents_tensor)
+        latent_image_ids_output = LatentsField(latents_name=saved_latent_image_ids_tensor)
+
+        return BriaLatentSamplerInvocationOutput(
+            latents=latents_output,
+            latent_image_ids=latent_image_ids_output,
+        )
diff --git a/invokeai/nodes/bria_nodes/bria_model_loader.py b/invokeai/nodes/bria_nodes/bria_model_loader.py
@@ -0,0 +1,60 @@
+from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField, UIType
+from invokeai.app.invocations.model import (
+    ModelIdentifierField,
+    SubModelType,
+    T5EncoderField,
+    TransformerField,
+    VAEField,
+)
+from invokeai.invocation_api import (
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    InputField,
+    InvocationContext,
+    OutputField,
+    invocation,
+    invocation_output,
+)
+
+
+@invocation_output("bria_model_loader_output")
+class BriaModelLoaderOutput(BaseInvocationOutput):
+    """Bria base model loader output"""
+
+    transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer")
+    t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder")
+    vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE")
+
+
+@invocation(
+    "bria_model_loader",
+    title="Main Model - Bria",
+    tags=["model", "bria"],
+    version="1.0.0",
+    classification=Classification.Prototype,
+)
+class BriaModelLoaderInvocation(BaseInvocation):
+    """Loads a bria base model, outputting its submodels."""
+
+    model: ModelIdentifierField = InputField(
+        description="Bria model (Transformer) to load",
+        ui_type=UIType.BriaMainModel,
+        input=Input.Direct,
+    )
+
+    def invoke(self, context: InvocationContext) -> BriaModelLoaderOutput:
+        for key in [self.model.key]:
+            if not context.models.exists(key):
+                raise ValueError(f"Unknown model: {key}")
+
+        transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer})
+        text_encoder = self.model.model_copy(update={"submodel_type": SubModelType.TextEncoder})
+        tokenizer = self.model.model_copy(update={"submodel_type": SubModelType.Tokenizer})
+        vae = self.model.model_copy(update={"submodel_type": SubModelType.VAE})
+
+        return BriaModelLoaderOutput(
+            transformer=TransformerField(transformer=transformer, loras=[]),
+            t5_encoder=T5EncoderField(tokenizer=tokenizer, text_encoder=text_encoder, loras=[]),
+            vae=VAEField(vae=vae),
+        )
diff --git a/invokeai/nodes/bria_nodes/bria_text_encoder.py b/invokeai/nodes/bria_nodes/bria_text_encoder.py
