ace++ and uno framework integration

Author: Attashe

invokeai/app/invocations/fields.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Callable, Optional, Tuple
+from typing import Any, Callable, Optional, Tuple, List
 
 from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, model_validator
 from pydantic.fields import _Unset
@@ -280,6 +280,15 @@ class FluxReduxConditioningField(BaseModel):
     )
 
 
+class FluxUnoReferenceField(BaseModel):
+    """A FLUX Uno image list primitive value"""
+
+    image_names: List[str] = 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

@@ -7,6 +7,7 @@
 import torch
 import torchvision.transforms as tv_transforms
 from PIL import Image
+import torchvision.transforms.functional as TVF
 from torchvision.transforms.functional import resize as tv_resize
 from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
 
@@ -17,6 +18,7 @@
     FluxConditioningField,
     FluxFillConditioningField,
     FluxReduxConditioningField,
+    FluxUnoReferenceField,
     ImageField,
     Input,
     InputField,
@@ -27,6 +29,7 @@
 from invokeai.app.invocations.flux_controlnet import FluxControlNetField
 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
@@ -42,6 +45,7 @@
 from invokeai.backend.flux.sampling_utils import (
     clip_timestep_schedule_fractional,
     generate_img_ids,
+    prepare_multi_ip,
     get_noise,
     get_schedule,
     pack,
@@ -109,6 +113,11 @@ class FluxDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard):
         description="FLUX Redux conditioning tensor.",
         input=Input.Connection,
     )
+    uno_reference: FluxUnoReferenceField | None = InputField(
+        default=None,
+        description="FLUX Redux conditioning tensor.",
+        input=Input.Connection,
+    )
     fill_conditioning: FluxFillConditioningField | None = InputField(
         default=None,
         description="FLUX Fill conditioning.",
@@ -284,6 +293,15 @@ def _run_diffusion(
 
         img_ids = generate_img_ids(h=latent_h, w=latent_w, batch_size=b, device=x.device, dtype=x.dtype)
 
+        is_flux_uno = self.uno_reference is not None
+        if is_flux_uno:
+            # Encode reference images and prepare position ids
+            uno_ref_imgs = self._prep_uno_reference_imgs(context)
+            uno_ref_imgs, uno_ref_ids = prepare_multi_ip(x, uno_ref_imgs)
+        else:
+            uno_ref_imgs = None
+            uno_ref_ids = None
+
         # Pack all latent tensors.
         init_latents = pack(init_latents) if init_latents is not None else None
         inpaint_mask = pack(inpaint_mask) if inpaint_mask is not None else None
@@ -391,6 +409,8 @@ def _run_diffusion(
                 pos_ip_adapter_extensions=pos_ip_adapter_extensions,
                 neg_ip_adapter_extensions=neg_ip_adapter_extensions,
                 img_cond=img_cond,
+                uno_ref_imgs=uno_ref_imgs,
+                uno_ref_ids=uno_ref_ids,
             )
 
         x = unpack(x.float(), self.height, self.width)
@@ -657,6 +677,30 @@ 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'
+        vae_info = context.models.load(self.controlnet_vae.vae)
+        
+        assert self.uno_reference is not None, "Needs reference images for UNO"
+
+        ref_img_names: list[str] = self.uno_reference.image_names
+        ref_latents: list[torch.Tensor] = []
+        
+        # TODO: Maybe move reference side to UNO Node
+        ref_long_side = 512 if len(ref_img_names) <= 1 else 320
+        
+        for img_name in ref_img_names:
+            image_pil = context.images.get_pil(img_name)
+            image_pil = image_pil.convert("RGB")  # To correct resizing
+            image_pil = preprocess_ref(image_pil, ref_long_side)  # resize and crop
+            
+            image_tensor = (TVF.to_tensor(image_pil) * 2.0 - 1.0).unsqueeze(0).float()
+            ref_latent = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=image_tensor)
+            ref_latents.append(ref_latent)
+        
+        return ref_latents
 
     def _prep_structural_control_img_cond(self, context: InvocationContext) -> torch.Tensor | None:
         if self.control_lora is None:
@@ -714,6 +758,7 @@ def _prep_flux_fill_img_cond(
         cond_img = context.images.get_pil(self.fill_conditioning.image.image_name, mode="RGB")
         cond_img = cond_img.resize((self.width, self.height), Image.Resampling.BICUBIC)
         cond_img = np.array(cond_img)
+
         cond_img = torch.from_numpy(cond_img).float() / 127.5 - 1.0
         cond_img = einops.rearrange(cond_img, "h w c -> 1 c h w")
         cond_img = cond_img.to(device=device, dtype=dtype)

invokeai/app/invocations/flux_uno.py

@@ -0,0 +1,93 @@
+from typing import Literal, Optional
+
+from PIL import Image
+import torchvision.transforms.functional as TVF
+
+from invokeai.app.invocations.baseinvocation import (
+    BaseModel,
+    BaseInvocation,
+    BaseInvocationOutput,
+    Classification,
+    invocation,
+    invocation_output,
+)
+from invokeai.app.invocations.fields import (
+    InputField,
+    OutputField,
+    FluxUnoReferenceField
+)
+from invokeai.app.invocations.primitives import ImageField
+from invokeai.app.services.shared.invocation_context import InvocationContext
+
+
+def preprocess_ref(raw_image: Image.Image, long_size: int = 512) -> Image.Image:
+    """Resize and center crop reference image
+    Code from https://github.com/bytedance/UNO/blob/main/uno/flux/pipeline.py
+    """
+    # Get the width and height of the original image
+    image_w, image_h = raw_image.size
+
+    # Calculate the long and short sides
+    if image_w >= image_h:
+        new_w = long_size
+        new_h = int((long_size / image_w) * image_h)
+    else:
+        new_h = long_size
+        new_w = int((long_size / image_h) * image_w)
+
+    # Scale proportionally to the new width and height
+    raw_image = raw_image.resize((new_w, new_h), resample=Image.LANCZOS)
+    target_w = new_w // 16 * 16
+    target_h = new_h // 16 * 16
+
+    # Calculate the starting coordinates of the clipping to achieve center clipping
+    left = (new_w - target_w) // 2
+    top = (new_h - target_h) // 2
+    right = left + target_w
+    bottom = top + target_h
+
+    # Center crop
+    raw_image = raw_image.crop((left, top, right, bottom))
+
+    # Convert to RGB mode
+    raw_image = raw_image.convert("RGB")
+    return raw_image
+
+
+@invocation_output("flux_uno_output")
+class FluxUnoOutput(BaseInvocationOutput):
+    """The conditioning output of a FLUX Redux invocation."""
+
+    uno_refs: FluxUnoReferenceField = OutputField(
+        description="Reference images container", title="Reference images"
+    )
+
+# TODO(attashe): adjust tags and category
+@invocation(
+    "flux_uno",
+    title="FLUX UNO",
+    tags=["ip_adapter", "control"],
+    category="ip_adapter",
+    version="2.1.0",
+    classification=Classification.Beta,
+)
+class FluxReduxInvocation(BaseInvocation):
+    """Runs a FLUX Redux model to generate a conditioning tensor."""
+
+    image: ImageField = InputField(description="The FLUX Redux image prompt.")
+    image2: Optional[ImageField] = InputField(default=None, description="2nd reference")
+    image3: Optional[ImageField] = InputField(default=None, description="3rd reference")
+    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)
+        )