ruff fix

Author: Ilan Tchenak

invokeai/app/invocations/bria_controlnet.py

@@ -1,31 +1,41 @@
-from invokeai.backend.bria.controlnet_bria import BRIA_CONTROL_MODES
+import cv2
+import numpy as np
+from PIL import Image
 from pydantic import BaseModel, Field
-from invokeai.invocation_api import ImageOutput, Classification
+
 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.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
-import numpy as np
-import cv2
-from PIL import Image
-
+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.backend.bria.controlnet_aux.open_pose import OpenposeDetector, Body, Hand, Face
+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"""
@@ -49,12 +59,8 @@ class BriaControlNetInvocation(BaseInvocation, WithMetadata, WithBoard):
     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"
-    )
+    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))
@@ -70,7 +76,7 @@ def invoke(self, context: InvocationContext) -> BriaControlNetOutput:
             control_image = convert_to_grayscale(image_in)
         elif self.control_mode == "tile":
             control_image = tile(16, image_in)
-            
+
         control_image = resize_img(control_image)
         image_dto = context.images.save(image=control_image)
         image_output = ImageOutput.build(image_dto)
@@ -99,6 +105,7 @@ def invoke(self, context: InvocationContext) -> BriaControlNetOutput:
     1.7708333333333333: {"width": 1360, "height": 768},
 }
 
+
 def extract_depth(image: Image.Image, context: InvocationContext):
     loaded_model = context.models.load_remote_model(DEPTH_SMALL_V2_URL, DepthAnythingPipeline.load_model)
 
@@ -107,6 +114,7 @@ def extract_depth(image: Image.Image, context: InvocationContext):
         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)
@@ -115,10 +123,10 @@ def extract_openpose(image: Image.Image, context: InvocationContext):
     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)
@@ -130,13 +138,17 @@ def extract_canny(input_image):
 
 
 def convert_to_grayscale(image):
-    gray_image = image.convert('L').convert('RGB')
+    gray_image = image.convert("L").convert("RGB")
     return gray_image
 
+
 def tile(downscale_factor, input_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)
+    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_ratio = control_image.width / control_image.height
     ratio = min(RATIO_CONFIGS_1024.keys(), key=lambda k: abs(k - image_ratio))

invokeai/app/invocations/bria_decoder.py

@@ -30,15 +30,15 @@ class BriaDecoderInvocation(BaseInvocation):
     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 / 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)

invokeai/app/invocations/bria_denoiser.py

@@ -1,19 +1,20 @@
 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.backend.bria.pipeline_bria_controlnet import BriaControlNetPipeline
-from invokeai.app.invocations.bria_controlnet import BriaControlNetField
 
 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.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.invocation_api import BaseInvocation, Classification, invocation, invocation_output
+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.invocation_api import BaseInvocation, Classification, invocation, invocation_output
+
 
 @invocation_output("bria_denoise_output")
 class BriaDenoiseInvocationOutput(BaseInvocationOutput):
@@ -80,7 +81,7 @@ class BriaDenoiseInvocation(BaseInvocation):
         description="ControlNet",
         input=Input.Connection,
         title="ControlNet",
-        default = None,
+        default=None,
     )
 
     @torch.no_grad()
@@ -106,7 +107,7 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
             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))
+            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:
@@ -134,7 +135,7 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
                 width=1024,
                 height=1024,
                 controlnet_conditioning_scale=control_scales,
-                num_inference_steps=self.num_steps, 
+                num_inference_steps=self.num_steps,
                 max_sequence_length=128,
                 guidance_scale=self.guidance_scale,
                 latents=latents,
@@ -150,36 +151,30 @@ def invoke(self, context: InvocationContext) -> BriaDenoiseInvocationOutput:
         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
+        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_modes.append(BriaControlModes[controlnet.mode].value)   
+                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:
-                    raise FileNotFoundError(f"Control image {controlnet.image.image_name} not found. Make sure not to delete the preprocessed image before finishing the pipeline.")
+                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, 
+            control_images=control_images,
+            control_modes=control_modes,
             width=width,
             height=height,
-            device=device, 
-            )
+            device=device,
+        )
         return control_model, tensored_control_images, tensored_control_modes, control_scales
-        

invokeai/app/invocations/bria_latent_sampler.py

@@ -1,19 +1,16 @@
 import torch
 
-from invokeai.app.invocations.fields import Input, InputField
+from invokeai.app.invocations.fields import Input, InputField, OutputField
 from invokeai.app.invocations.model import TransformerField
 from invokeai.app.invocations.primitives import (
     BaseInvocationOutput,
     FieldDescriptions,
-    Input,
     LatentsField,
-    OutputField,
 )
 from invokeai.backend.bria.pipeline_bria_controlnet import prepare_latents
 from invokeai.invocation_api import (
     BaseInvocation,
     Classification,
-    InputField,
     InvocationContext,
     invocation,
     invocation_output,
@@ -56,7 +53,7 @@ def invoke(self, context: InvocationContext) -> BriaLatentSamplerInvocationOutpu
 
         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,
@@ -66,7 +63,7 @@ def invoke(self, context: InvocationContext) -> BriaLatentSamplerInvocationOutpu
             dtype=dtype,
             device=device,
             generator=generator,
-            )
+        )
 
         saved_latents_tensor = context.tensors.save(latents)
         saved_latent_image_ids_tensor = context.tensors.save(latent_image_ids)

invokeai/app/invocations/bria_model_loader.py

@@ -10,9 +10,7 @@
     BaseInvocation,
     BaseInvocationOutput,
     Classification,
-    InputField,
     InvocationContext,
-    OutputField,
     invocation,
     invocation_output,
 )

invokeai/app/invocations/bria_text_encoder.py

@@ -19,8 +19,6 @@
     invocation_output,
 )
 
-from invokeai.backend.bria.bria_utils import get_t5_prompt_embeds, is_ng_none
-
 
 @invocation_output("bria_text_encoder_output")
 class BriaTextEncoderInvocationOutput(BaseInvocationOutput):
@@ -70,7 +68,7 @@ def invoke(self, context: InvocationContext) -> BriaTextEncoderInvocationOutput:
         ):
             assert isinstance(tokenizer, T5TokenizerFast)
             assert isinstance(text_encoder, T5EncoderModel)
-            
+
         (prompt_embeds, negative_prompt_embeds, text_ids) = encode_prompt(
             prompt=self.prompt,
             tokenizer=tokenizer,
@@ -81,7 +79,7 @@ def invoke(self, context: InvocationContext) -> BriaTextEncoderInvocationOutput:
             max_sequence_length=self.max_length,
             lora_scale=1.0,
         )
-        
+
         saved_pos_tensor = context.tensors.save(prompt_embeds)
         saved_neg_tensor = context.tensors.save(negative_prompt_embeds)
         saved_text_ids_tensor = context.tensors.save(text_ids)

invokeai/backend/bria/bria_utils.py

@@ -87,7 +87,7 @@ def is_ng_none(negative_prompt):
         negative_prompt is None
         or negative_prompt == ""
         or (isinstance(negative_prompt, list) and negative_prompt[0] is None)
-        or (type(negative_prompt) == list and negative_prompt[0] == "")
+        or (isinstance(negative_prompt, list) and negative_prompt[0] == "")
     )
 
 

invokeai/backend/bria/controlnet_aux/__init__.py

@@ -1,5 +1,6 @@
 __version__ = "0.0.9"
 
-from .canny import CannyDetector
-from .open_pose import OpenposeDetector
+from invokeai.backend.bria.controlnet_aux.canny import CannyDetector as CannyDetector
+from invokeai.backend.bria.controlnet_aux.open_pose import OpenposeDetector as OpenposeDetector
 
+__all__ = ["CannyDetector", "OpenposeDetector"]

invokeai/backend/bria/controlnet_aux/canny/__init__.py

@@ -1,15 +1,27 @@
 import warnings
+
 import cv2
 import numpy as np
 from PIL import Image
-from ..util import HWC3, resize_image
+
+from invokeai.backend.bria.controlnet_aux.util import HWC3, resize_image
+
 
 class CannyDetector:
-    def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, image_resolution=512, output_type=None, **kwargs):
+    def __call__(
+        self,
+        input_image=None,
+        low_threshold=100,
+        high_threshold=200,
+        detect_resolution=512,
+        image_resolution=512,
+        output_type=None,
+        **kwargs,
+    ):
         if "img" in kwargs:
-            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning, stacklevel=2)
             input_image = kwargs.pop("img")
-        
+
         if input_image is None:
             raise ValueError("input_image must be defined.")
 
@@ -18,19 +30,19 @@ def __call__(self, input_image=None, low_threshold=100, high_threshold=200, dete
             output_type = output_type or "pil"
         else:
             output_type = output_type or "np"
-        
+
         input_image = HWC3(input_image)
         input_image = resize_image(input_image, detect_resolution)
 
         detected_map = cv2.Canny(input_image, low_threshold, high_threshold)
-        detected_map = HWC3(detected_map)      
-         
+        detected_map = HWC3(detected_map)
+
         img = resize_image(input_image, image_resolution)
         H, W, C = img.shape
 
         detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
-        
+
         if output_type == "pil":
             detected_map = Image.fromarray(detected_map)
-            
+
         return detected_map