cr fixes 3

Author: ilanbria

invokeai/app/invocations/bria_denoiser.py

@@ -5,7 +5,6 @@
 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
@@ -72,11 +71,6 @@ class BriaDenoiseInvocation(BaseInvocation):
         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,
@@ -87,6 +81,16 @@ class BriaDenoiseInvocation(BaseInvocation):
         input=Input.Connection,
         title="Negative Prompt Embeds",
     )
+    latents: LatentsField = InputField(
+        description="Latent noise with latent image ids to denoise",
+        input=Input.Connection,
+        title="Latent Noise",
+    )
+    latent_image_ids: LatentsField = InputField(
+        description="Latent image ids to denoise",
+        input=Input.Connection,
+        title="Latent Image IDs",
+    )
     control: BriaControlNetField | list[BriaControlNetField] | None = InputField(
         description="ControlNet",
         input=Input.Connection,
@@ -96,10 +100,10 @@ class BriaDenoiseInvocation(BaseInvocation):
 
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
-        latents = context.tensors.load(self.latent_noise.latents.latents_name)
+        latents = context.tensors.load(self.latents.latents_name)
+        latent_image_ids = context.tensors.load(self.latent_image_ids.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

invokeai/app/invocations/bria_latent_noise.py

@@ -18,14 +18,11 @@
 )
 
 
-class BriaLatentNoiseOutput(BaseModel):
-    latents: LatentsField
-    latent_image_ids: LatentsField
-
 @invocation_output("bria_latent_noise_output")
 class BriaLatentNoiseInvocationOutput(BaseInvocationOutput):
     """Base class for nodes that output Bria latent tensors."""
-    latent_noise: BriaLatentNoiseOutput = OutputField(description="The latent noise, containing latents and latent image ids.")
+    latents: LatentsField = OutputField(description="The latent noise")
+    latent_image_ids: LatentsField = OutputField(description="The latent image ids.")
     height: int = OutputField(description="The height of the output image")
     width: int = OutputField(description="The width of the output image")
 
@@ -86,10 +83,8 @@ def invoke(self, context: InvocationContext) -> BriaLatentNoiseInvocationOutput:
         latent_image_ids_output = LatentsField(latents_name=saved_latent_image_ids_tensor)
 
         return BriaLatentNoiseInvocationOutput(
-            latent_noise=BriaLatentNoiseOutput(
-                latents=latents_output,
-                latent_image_ids=latent_image_ids_output,
-            ),
+            latents=latents_output,
+            latent_image_ids=latent_image_ids_output,
             height=self.height,
             width=self.width,
         )

invokeai/app/invocations/bria_decoder.py renamed to invokeai/app/invocations/bria_latents_to_image.py

@@ -9,14 +9,14 @@
 
 
 @invocation(
-    "bria_decoder",
-    title="Decoder - Bria",
+    "bria_latents_to_image",
+    title="Latents to Image - Bria",
     tags=["image", "bria"],
     category="image",
     version="1.0.0",
     classification=Classification.Prototype,
 )
-class BriaDecoderInvocation(BaseInvocation):
+class BriaLatentsToImageInvocation(BaseInvocation):
     """
     Decode Bria latents to an image.
     """

invokeai/backend/bria/pipeline_bria_controlnet.py

@@ -259,10 +259,6 @@ def __call__(
             prompt (`str` or `List[str]`, *optional*):
                 The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
                 instead.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
             num_inference_steps (`int`, *optional*, defaults to 50):
                 The number of denoising steps. More denoising steps usually lead to a higher quality image at the
                 expense of slower inference.
@@ -323,8 +319,6 @@ def __call__(
             `tuple`. When returning a tuple, the first element is a list with the generated images.
         """
 
-        height = height or self.default_sample_size * self.vae_scale_factor
-        width = width or self.default_sample_size * self.vae_scale_factor
         control_guidance_start, control_guidance_end = self.get_control_start_end(
             control_guidance_start=control_guidance_start, control_guidance_end=control_guidance_end
         )
@@ -335,8 +329,8 @@ def __call__(
         )
         self.check_inputs(
             prompt,
-            height,
-            width,
+            height=height,
+            width=width,
             negative_prompt=negative_prompt,
             prompt_embeds=prompt_embeds,
             negative_prompt_embeds=negative_prompt_embeds,
@@ -517,7 +511,7 @@ def __call__(
                             order=1,
                             total_steps=num_inference_steps,
                             timestep=int(t),
-                            latents=latents.view(1, 64, 64, 4, 2, 2).permute(0, 3, 1, 4, 2, 5).reshape(1, 4, 128, 128),
+                            latents=self._unpack_latents(latents, height, width, self.vae_scale_factor),
                         ),
                     )