up

Author: sayakpaul

src/diffusers/pipelines/flux2/image_processor.py

@@ -96,21 +96,24 @@ def check_image_input(
             )
 
         return image
-
+    
     @staticmethod
-    def _resize_to_target_area(
-        image: PIL.Image.Image, target_area: int = 1024 * 1024, return_if_small_image: bool = False
-    ) -> PIL.Image.Image:
+    def _resize_to_target_area(image: PIL.Image.Image, target_area: int = 1024 * 1024) -> PIL.Image.Image:
         image_width, image_height = image.size
-        pixel_count = image_width * image_height
-        if return_if_small_image and pixel_count <= target_area:
-            return image
 
-        scale = math.sqrt(target_area / pixel_count)
+        scale = math.sqrt(target_area / (image_width * image_height))
         width = int(image_width * scale)
         height = int(image_height * scale)
 
         return image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+    
+    @staticmethod
+    def _resize_if_exceeds_area(image, target_area=1024 * 1024) -> PIL.Image.Image:
+        image_width, image_height = image.size
+        pixel_count = image_width * image_height
+        if pixel_count <= target_area:
+            return image
+        return Flux2ImageProcessor._resize_to_target_area(image, target_area)
 
     def _resize_and_crop(
         self,

src/diffusers/pipelines/flux2/pipeline_flux2.py

@@ -57,6 +57,7 @@
         ```
 """
 
+UPSAMPLING_MAX_IMAGE_SIZE = 768**2
 
 # Adapted from
 # https://github.com/black-forest-labs/flux2/blob/5a5d316b1b42f6b59a8c9194b77c8256be848432/src/flux2/text_encoder.py#L68
@@ -95,8 +96,6 @@ def format_input(
         ]
     else:
         assert len(images) == len(prompts), "Number of images must match number of prompts"
-        images = _validate_and_process_images(images)
-
         messages = [
             [
                 {
@@ -149,7 +148,7 @@ def _validate_and_process_images(
     # cap the pixels
     images = [
         [
-            image_processor._resize_to_target_area(img_i, upsampling_max_image_size, return_if_small_image=True)
+            image_processor._resize_if_exceeds_area(img_i, upsampling_max_image_size)
             for img_i in img_i
         ]
         for img_i in images
@@ -301,7 +300,7 @@ def __init__(
         self.system_message = SYSTEM_MESSAGE
         self.system_message_upsampling_t2i = SYSTEM_MESSAGE_UPSAMPLING_T2I
         self.system_message_upsampling_i2i = SYSTEM_MESSAGE_UPSAMPLING_I2I
-        self.upsampling_max_image_size = 768**2
+        self.upsampling_max_image_size = UPSAMPLING_MAX_IMAGE_SIZE
 
     @staticmethod
     def _get_mistral_3_small_prompt_embeds(
@@ -525,6 +524,10 @@ def upsample_prompt(
         else:
             system_message = SYSTEM_MESSAGE_UPSAMPLING_I2I
 
+        # Validate and process the input images
+        if images:
+            images = _validate_and_process_images(images, self.image_processor, self.upsampling_max_image_size)
+        
         # Format input messages
         messages_batch = format_input(prompts=prompt, system_message=system_message, images=images)