cr fixes 1

Author: ilanbria

invokeai/app/invocations/bria_controlnet.py

@@ -1,5 +1,3 @@
-import cv2
-import numpy as np
 from PIL import Image
 from pydantic import BaseModel, Field
 
@@ -20,9 +18,7 @@
 )
 from invokeai.app.invocations.model import ModelIdentifierField
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.bria.controlnet_aux.open_pose import Body, Face, Hand, OpenposeDetector
 from invokeai.backend.bria.controlnet_bria import BRIA_CONTROL_MODES
-from invokeai.backend.image_util.depth_anything.depth_anything_pipeline import DepthAnythingPipeline
 from invokeai.invocation_api import Classification, ImageOutput
 
 DEPTH_SMALL_V2_URL = "depth-anything/Depth-Anything-V2-Small-hf"
@@ -41,7 +37,6 @@ class BriaControlNetOutput(BaseInvocationOutput):
     """Bria ControlNet info"""
 
     control: BriaControlNetField = OutputField(description=FieldDescriptions.control)
-    preprocessed_images: ImageField = OutputField(description="The preprocessed control image")
 
 
 @invocation(
@@ -64,24 +59,18 @@ class BriaControlNetInvocation(BaseInvocation, WithMetadata, WithBoard):
 
     def invoke(self, context: InvocationContext) -> BriaControlNetOutput:
         image_in = resize_img(context.images.get_pil(self.control_image.image_name))
-        if self.control_mode == "canny":
-            control_image = extract_canny(image_in)
-        elif self.control_mode == "depth":
-            control_image = extract_depth(image_in, context)
-        elif self.control_mode == "pose":
-            control_image = extract_openpose(image_in, context)
-        elif self.control_mode == "colorgrid":
+        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)
-        image_output = ImageOutput.build(image_dto)
         return BriaControlNetOutput(
-            preprocessed_images=image_output.image,
             control=BriaControlNetField(
                 image=ImageField(image_name=image_dto.image_name),
                 model=self.control_model,
@@ -106,50 +95,20 @@ def invoke(self, context: InvocationContext) -> BriaControlNetOutput:
 }
 
 
-def extract_depth(image: Image.Image, context: InvocationContext):
-    loaded_model = context.models.load_remote_model(DEPTH_SMALL_V2_URL, DepthAnythingPipeline.load_model)
 
-    with loaded_model as depth_anything_detector:
-        assert isinstance(depth_anything_detector, DepthAnythingPipeline)
-        depth_map = depth_anything_detector.generate_depth(image)
-    return depth_map
-
-
-def extract_openpose(image: Image.Image, context: InvocationContext):
-    body_model = context.models.load_remote_model(f"{HF_LLLYASVIEL}body_pose_model.pth", Body)
-    hand_model = context.models.load_remote_model(f"{HF_LLLYASVIEL}hand_pose_model.pth", Hand)
-    face_model = context.models.load_remote_model(f"{HF_LLLYASVIEL}facenet.pth", Face)
-
-    with body_model as body_model, hand_model as hand_model, face_model as face_model:
-        open_pose_model = OpenposeDetector(body_model, hand_model, face_model)
-        processed_image_open_pose = open_pose_model(image, hand_and_face=True)
-
-    processed_image_open_pose = processed_image_open_pose.resize(image.size)
-    return processed_image_open_pose
-
-
-def extract_canny(input_image):
-    image = np.array(input_image)
-    image = cv2.Canny(image, 100, 200)
-    image = image[:, :, None]
-    image = np.concatenate([image, image, image], axis=2)
-    canny_image = Image.fromarray(image)
-    return canny_image
-
-
-def convert_to_grayscale(image):
+def convert_to_grayscale(image: Image.Image) -> Image.Image:
     gray_image = image.convert("L").convert("RGB")
     return gray_image
 
 
-def tile(downscale_factor, input_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):
+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"]

invokeai/app/invocations/bria_decoder.py

@@ -17,6 +17,10 @@
     classification=Classification.Prototype,
 )
 class BriaDecoderInvocation(BaseInvocation):
+    """
+    Decode Bria latents to an image.
+    """
+
     vae: VAEField = InputField(
         description=FieldDescriptions.vae,
         input=Input.Connection,

invokeai/app/invocations/bria_denoiser.py

@@ -1,18 +1,21 @@
-from typing import List, Tuple
+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.fields import Input, InputField, LatentsField, OutputField
+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
 
 
@@ -30,6 +33,11 @@ class BriaDenoiseInvocationOutput(BaseInvocationOutput):
     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"
     )
@@ -52,31 +60,31 @@ class BriaDenoiseInvocation(BaseInvocation):
         input=Input.Connection,
         title="VAE",
     )
-    latents: LatentsField = InputField(
-        description="Latents to denoise",
-        input=Input.Connection,
-        title="Latents",
+    height: int = InputField(
+        default=1024,
+        title="Height",
+        description="The height of the output image",
     )
-    latent_image_ids: LatentsField = InputField(
-        description="Latent Image IDs to denoise",
+    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 Image IDs",
+        title="Latent Noise",
     )
-    pos_embeds: LatentsField = InputField(
+    pos_embeds: FluxConditioningField = InputField(
         description="Positive Prompt Embeds",
         input=Input.Connection,
         title="Positive Prompt Embeds",
     )
-    neg_embeds: LatentsField = InputField(
+    neg_embeds: FluxConditioningField = 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",
-    )
     control: BriaControlNetField | list[BriaControlNetField] | None = InputField(
         description="ControlNet",
         input=Input.Connection,
@@ -86,11 +94,10 @@ class BriaDenoiseInvocation(BaseInvocation):
 
     @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)
+        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
@@ -114,11 +121,12 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
                 control_model, control_images, control_modes, control_scales = self._prepare_multi_control(
                     context=context,
                     vae=vae,
-                    width=1024,
-                    height=1024,
+                    width=self.width,
+                    height=self.height,
                     device=vae.device,
                 )
 
+
             pipeline = BriaControlNetPipeline(
                 transformer=transformer,
                 scheduler=scheduler,
@@ -129,31 +137,32 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
             )
             pipeline.to(device=transformer.device, dtype=transformer.dtype)
 
-            latents = pipeline(
+            output_latents = pipeline(
                 control_image=control_images,
                 control_mode=control_modes,
-                width=1024,
-                height=1024,
+                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,
-                text_ids=text_ids,
                 prompt_embeds=pos_embeds,
                 negative_prompt_embeds=neg_embeds,
                 output_type="latent",
+                step_callback=_build_step_callback(context),
             )[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)
+            
+
+        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[torch.Tensor], List[float]]:
+    ) -> 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:
@@ -178,3 +187,11 @@ def _prepare_multi_control(
             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

invokeai/app/invocations/bria_latent_sampler.py renamed to invokeai/app/invocations/bria_latent_noise.py

@@ -1,4 +1,5 @@
 import torch
+from pydantic import BaseModel, Field 
 
 from invokeai.app.invocations.fields import Input, InputField, OutputField
 from invokeai.app.invocations.model import TransformerField
@@ -17,23 +18,28 @@
 )
 
 
-@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)
+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.")
+    height: int = OutputField(description="The height of the output image", default=1024)
+    width: int = OutputField(description="The width of the output image", default=1024)
 
 @invocation(
-    "bria_latent_sampler",
-    title="Latent Sampler - Bria",
+    "bria_latent_noise",
+    title="Latent Noise - Bria",
     tags=["image", "bria"],
     category="image",
     version="1.0.0",
     classification=Classification.Prototype,
 )
-class BriaLatentSamplerInvocation(BaseInvocation):
+class BriaLatentNoiseInvocation(BaseInvocation):
+    """ Generate latent noise for Bria. """
+
     seed: int = InputField(
         default=42,
         title="Seed",
@@ -44,22 +50,31 @@ class BriaLatentSamplerInvocation(BaseInvocation):
         input=Input.Connection,
         title="Transformer",
     )
+    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",
+    )
 
     @torch.no_grad()
-    def invoke(self, context: InvocationContext) -> BriaLatentSamplerInvocationOutput:
+    def invoke(self, context: InvocationContext) -> BriaLatentNoiseInvocationOutput:
         with context.models.load(self.transformer.transformer) as transformer:
             device = transformer.device
             dtype = transformer.dtype
 
-        height, width = 1024, 1024
         generator = torch.Generator(device=device).manual_seed(self.seed)
 
         num_channels_latents = 4
         latents, latent_image_ids = prepare_latents(
             batch_size=1,
             num_channels_latents=num_channels_latents,
-            height=height,
-            width=width,
+            height=self.height,
+            width=self.width,
             dtype=dtype,
             device=device,
             generator=generator,
@@ -70,7 +85,11 @@ def invoke(self, context: InvocationContext) -> BriaLatentSamplerInvocationOutpu
         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,
+        return BriaLatentNoiseInvocationOutput(
+            latent_noise=BriaLatentNoiseOutput(
+                latents=latents_output,
+                latent_image_ids=latent_image_ids_output,
+            ),
+            height=self.height,
+            width=self.width,
         )