chore: ruff

Author: psychedelicious

invokeai/app/invocations/fields.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Callable, Optional, Tuple, List
+from typing import Any, Callable, Optional, Tuple
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, model_validator
 from pydantic.fields import _Unset
@@ -286,12 +286,13 @@ class FluxReduxConditioningField(BaseModel):
 class FluxUnoReferenceField(BaseModel):
     """A FLUX Uno image list primitive value"""
 
-    image_names: List[str] = Field(
+    image_names: list[str] | None = Field(
         default=None,
         description="The name of the image associated with this conditioning tensor. This is used to store the image "
         "in the context.",
     )
 
+
 class FluxFillConditioningField(BaseModel):
     """A FLUX Fill conditioning field."""
 

invokeai/app/invocations/flux_denoise.py

@@ -6,8 +6,8 @@
 import numpy.typing as npt
 import torch
 import torchvision.transforms as tv_transforms
-from PIL import Image
 import torchvision.transforms.functional as TVF
+from PIL import Image
 from torchvision.transforms.functional import resize as tv_resize
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 
@@ -27,9 +27,9 @@
     WithMetadata,
 )
 from invokeai.app.invocations.flux_controlnet import FluxControlNetField
+from invokeai.app.invocations.flux_uno import preprocess_ref
 from invokeai.app.invocations.flux_vae_encode import FluxVaeEncodeInvocation
 from invokeai.app.invocations.ip_adapter import IPAdapterField
-from invokeai.app.invocations.flux_uno import preprocess_ref
 from invokeai.app.invocations.model import ControlLoRAField, LoRAField, TransformerField, VAEField
 from invokeai.app.invocations.primitives import LatentsOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
@@ -45,10 +45,10 @@
 from invokeai.backend.flux.sampling_utils import (
     clip_timestep_schedule_fractional,
     generate_img_ids,
-    prepare_multi_ip,
     get_noise,
     get_schedule,
     pack,
+    prepare_multi_ip,
     unpack,
 )
 from invokeai.backend.flux.text_conditioning import FluxReduxConditioning, FluxTextConditioning
@@ -677,7 +677,7 @@ def _prep_controlnet_extensions(
                 raise ValueError(f"Unsupported ControlNet model type: {type(model)}")
 
         return controlnet_extensions
-    
+
     def _prep_uno_reference_imgs(self, context: InvocationContext) -> list[torch.Tensor]:
         # Load the conditioning image and resize it to the target image size.
         assert self.controlnet_vae is not None, 'Controlnet Vae must be set for UNO encoding'

invokeai/app/invocations/flux_uno.py

@@ -9,11 +9,7 @@
     invocation,
     invocation_output,
 )
-from invokeai.app.invocations.fields import (
-    InputField,
-    OutputField,
-    FluxUnoReferenceField
-)
+from invokeai.app.invocations.fields import FluxUnoReferenceField, InputField, OutputField
 from invokeai.app.invocations.primitives import ImageField
 from invokeai.app.services.shared.invocation_context import InvocationContext
 
@@ -56,9 +52,7 @@ def preprocess_ref(raw_image: Image.Image, long_size: int = 512) -> Image.Image:
 class FluxUnoOutput(BaseInvocationOutput):
     """The conditioning output of a FLUX Redux invocation."""
 
-    uno_refs: FluxUnoReferenceField = OutputField(
-        description="Reference images container", title="Reference images"
-    )
+    uno_refs: FluxUnoReferenceField = OutputField(description="Reference images container", title="Reference images")
 
 
 @invocation(
@@ -78,14 +72,10 @@ class FluxUnoInvocation(BaseInvocation):
     image4: Optional[ImageField] = InputField(default=None, description="4th reference")
 
     def invoke(self, context: InvocationContext) -> FluxUnoOutput:
-
         images: list[str] = []
         for image in [self.image, self.image2, self.image3, self.image4]:
             if image is not None:
                 image_pil = context.images.get_pil(image.image_name)
                 images.append(context.images.save(image=image_pil).image_name)
-        
-        return FluxUnoOutput(
-            uno_refs=FluxUnoReferenceField(
-                image_names=images)
-        )
+
+        return FluxUnoOutput(uno_refs=FluxUnoReferenceField(image_names=images))

invokeai/app/invocations/image_context_utils.py

@@ -19,42 +19,39 @@
 - Gaussian Blur Mask
 - Image Concatenate
 """
+
 import math
-from PIL import Image
-from typing import Optional, List, Dict, Any, Tuple, Literal
+from typing import List, Literal, Optional
 
-import cv2
 import numpy as np
 import torch
 import torchvision.transforms as T
+from PIL import Image
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output
 from invokeai.app.invocations.fields import (
-    Field,
-    Input,
     ImageField,
+    Input,
     InputField,
     OutputField,
     TensorField,
     WithBoard,
-    WithMetadata
+    WithMetadata,
 )
-from invokeai.app.invocations.primitives import ImageOutput, MaskOutput
+from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
-from invokeai.backend.image_util.util import cv2_to_pil, pil_to_cv2
-
 
 DIRECTION_OPTIONS = Literal["right", "left", "down", "up"]
 
 
-def concat_images(image1: Image.Image, image2: Image.Image,
-                  direction: str = "right",
-                  match_image_size=True) -> Image.Image:
+def concat_images(
+    image1: Image.Image, image2: Image.Image, direction: str = "right", match_image_size=True
+) -> Image.Image:
     """Concatenate two images either horizontally or vertically."""
     # Ensure that image modes are same
     if image1.mode != image2.mode:
         image2 = image2.convert(image1.mode)
-    
+
     if direction == "right" or direction == "left":
         if direction == "left":
             image1, image2 = image2, image1
@@ -73,7 +70,7 @@ def concat_images(image1: Image.Image, image2: Image.Image,
         new_image.paste(image2, (0, image1.height))
     else:
         raise ValueError("Mode must be either 'horizontal' or 'vertical'.")
-    
+
     return new_image
 
 
@@ -89,7 +86,9 @@ class ConcatImagesInvocation(BaseInvocation, WithMetadata, WithBoard):
 
     image1: ImageField = InputField(description="The first image to process")
     image2: ImageField = InputField(description="The second image to process")
-    mode: DIRECTION_OPTIONS = InputField(default="horizontal", description="Mode of concatenation: 'horizontal' or 'vertical'")
+    mode: DIRECTION_OPTIONS = InputField(
+        default="horizontal", description="Mode of concatenation: 'horizontal' or 'vertical'"
+    )
 
     def invoke(self, context: InvocationContext) -> ImageOutput:
         image1 = context.images.get_pil(self.image1.image_name)
@@ -103,9 +102,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 class InpaintCropOutput(BaseInvocationOutput):
     """The output of Inpain Crop Invocation."""
 
-    image_crop: ImageField = OutputField(
-        description="Cropped part of image", title="Conditioning"
-    )
+    image_crop: ImageField = OutputField(description="Cropped part of image", title="Conditioning")
     stitcher: List[int] = OutputField(description="Parameter for stitching image after inpainting")
 
 
@@ -117,17 +114,17 @@ class InpaintCropOutput(BaseInvocationOutput):
 )
 class InpaintCropInvocation(BaseInvocation, WithMetadata, WithBoard):
     "Crop from image masked area with resize and expand options"
-    
+
     image: ImageField = InputField(description="The source image")
     mask: TensorField = InputField(description="Inpaint mask")
-    
+
     def invoke(self, context: InvocationContext) -> ImageOutput:
         image = context.images.get_pil(self.image.image_name, "RGB")
         mask = context.tensors.load(self.mask.tensor_name)
-        
+
         # TODO: Finish InpaintCrop implementation
         image_crop = Image.new("RGB", (256, 256))
-        
+
         image_dto = context.images.save(image=image_crop)
         return ImageOutput.build(image_dto)
 
@@ -136,9 +133,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 class ACEppProcessorOutput(BaseInvocationOutput):
     """The conditioning output of a FLUX Fill invocation."""
 
-    image: ImageField = OutputField(
-        description="Concatted image", title="Image"
-    )
+    image: ImageField = OutputField(description="Concatted image", title="Image")
     mask: TensorField = OutputField(description="Inpaint mask")
     crop_pad: int = OutputField(description="Padding to crop result")
     crop_width: int = OutputField(description="Width of output area")
@@ -155,15 +150,15 @@ class ACEppProcessor(BaseInvocation):
     reference_image: ImageField = InputField(description="Reference Image")
     edit_image: Optional[ImageField] = InputField(description="Edit Image", default=None, input=Input.Connection)
     edit_mask: Optional[TensorField] = InputField(description="Edit Mask", default=None, input=Input.Connection)
-    
+
     width: int = InputField(default=512, gt=0, description="The width of the crop rectangle")
     height: int = InputField(default=512, gt=0, description="The height of the crop rectangle")
-    
+
     max_seq_len: int = InputField(default=4096, gt=2048, le=5120, description="The height of the crop rectangle")
-    
+
     def image_check(self, image_pil: Image.Image) -> torch.Tensor:
         max_aspect_ratio = 4
-        
+
         image = self.transform_pil_tensor(image_pil)
         image = image.unsqueeze(0)
         # preprocess
@@ -173,20 +168,18 @@ def image_check(self, image_pil: Image.Image) -> torch.Tensor:
         elif W / H > max_aspect_ratio:
             image[0] = T.CenterCrop([H, int(max_aspect_ratio * H)])(image[0])
         return image[0]
-    
+
     def transform_pil_tensor(self, pil_image: Image.Image) -> torch.Tensor:
-        transform = T.Compose([
-            T.ToTensor()
-        ])
+        transform = T.Compose([T.ToTensor()])
         tensor_image: torch.Tensor = transform(pil_image)
         return tensor_image
-    
+
     def invoke(self, context: InvocationContext) -> ACEppProcessorOutput:
         d = 16  # Flux pixels per patch rate
-        
+
         image_pil = context.images.get_pil(self.reference_image.image_name, "RGB")
         image = self.image_check(image_pil) - 0.5
-        
+
         if self.edit_image is None:
             edit_image = torch.zeros((3, self.height, self.width))
             edit_mask = torch.ones((1, self.height, self.width))
@@ -199,23 +192,21 @@ def invoke(self, context: InvocationContext) -> ACEppProcessorOutput:
                 edit_mask = torch.ones((eH, eW))
             else:
                 edit_mask = context.tensors.load(self.edit_mask.tensor_name)
-        
+
         out_H, out_W = edit_image.shape[-2:]
-            
+
         _, H, W = image.shape
         _, eH, eW = edit_image.shape
-        
+
         # align height with edit_image
         scale = eH / H
         tH, tW = eH, int(W * scale)
-        
-        reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)(
-                image)
+
+        reference_image = T.Resize((tH, tW), interpolation=T.InterpolationMode.BILINEAR, antialias=True)(image)
         edit_image = torch.cat([reference_image, edit_image], dim=-1)
-        edit_mask = torch.cat([torch.zeros((1, reference_image.shape[1], reference_image.shape[2])), edit_mask],
-                                dim=-1)
+        edit_mask = torch.cat([torch.zeros((1, reference_image.shape[1], reference_image.shape[2])), edit_mask], dim=-1)
         slice_w = reference_image.shape[-1]
-    
+
         H, W = edit_image.shape[-2:]
         scale = min(1.0, math.sqrt(self.max_seq_len * 2 / ((H / d) * (W / d))))
         rH = int(H * scale) // d * d
@@ -235,7 +226,7 @@ def invoke(self, context: InvocationContext) -> ACEppProcessorOutput:
         # Convert to torch.bool
         edit_mask = edit_mask > 0.5
         image_out = Image.fromarray((edit_image[0].numpy() * 255).astype(np.uint8))
-        
+
         image_dto = context.images.save(image=image_out)
         mask_name = context.tensors.save(edit_mask)
         return ACEppProcessorOutput(

invokeai/backend/flux/model.py

@@ -123,10 +123,10 @@ def forward(
         img_end = img.shape[1]  # length of original image vector
         if uno_ref_imgs is not None and uno_ref_ids is not None:
             img_in = [img] + [self.img_in(ref) for ref in uno_ref_imgs]
-            img_ids = [ids] + [ref_ids for ref_ids in uno_ref_ids]
-            img = torch.cat(img_in, dim=1)  
+            img_ids = [ids] + uno_ref_ids
+            img = torch.cat(img_in, dim=1)
             ids = torch.cat(img_ids, dim=1)
-            
+
         pe = self.pe_embedder(ids)
 
         # Validate double_block_residuals shape.