Expose max_area for condition images in Flux2 pipelines

The condition/reference image downscale threshold was hardcoded to
1024*1024 in Flux2Pipeline, Flux2KleinPipeline and Flux2KleinKVPipeline,
silently downscaling any reference image above ~1MP. Expose it as a
max_area __call__ argument (default unchanged), mirroring the existing
max_area parameter of FluxKontextPipeline.

Add fast tests for the new argument on the three pipelines (mirroring
the FluxKontext max_area test style) and a nightly integration test
consuming a 2048x2048 condition image at full resolution.

Author: sqhuang

src/diffusers/pipelines/flux2/pipeline_flux2.py

@@ -765,6 +765,7 @@ def __call__(
         max_sequence_length: int = 512,
         text_encoder_out_layers: tuple[int] = (10, 20, 30),
         caption_upsample_temperature: float = None,
+        max_area: int = 1024**2,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -832,6 +833,9 @@ def __call__(
             caption_upsample_temperature (`float`):
                 When specified, we will try to perform caption upsampling for potentially improved outputs. We
                 recommend setting it to 0.15 if caption upsampling is to be performed.
+            max_area (`int`, defaults to `1024 ** 2`):
+                The maximum area (in pixels) allowed for each condition image. Condition images whose area exceeds
+                this value are downscaled to fit it while preserving their aspect ratio.
 
         Examples:
 
@@ -891,8 +895,8 @@ def __call__(
             condition_images = []
             for img in image:
                 image_width, image_height = img.size
-                if image_width * image_height > 1024 * 1024:
-                    img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
+                if image_width * image_height > max_area:
+                    img = self.image_processor._resize_to_target_area(img, max_area)
                     image_width, image_height = img.size
 
                 multiple_of = self.vae_scale_factor * 2

src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

@@ -632,6 +632,7 @@ def __call__(
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
         max_sequence_length: int = 512,
         text_encoder_out_layers: tuple[int] = (9, 18, 27),
+        max_area: int = 1024**2,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -700,6 +701,9 @@ def __call__(
             max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
             text_encoder_out_layers (`tuple[int]`):
                 Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
+            max_area (`int`, defaults to `1024 ** 2`):
+                The maximum area (in pixels) allowed for each condition image. Condition images whose area exceeds
+                this value are downscaled to fit it while preserving their aspect ratio.
 
         Examples:
 
@@ -769,8 +773,8 @@ def __call__(
             condition_images = []
             for img in image:
                 image_width, image_height = img.size
-                if image_width * image_height > 1024 * 1024:
-                    img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
+                if image_width * image_height > max_area:
+                    img = self.image_processor._resize_to_target_area(img, max_area)
                     image_width, image_height = img.size
 
                 multiple_of = self.vae_scale_factor * 2

src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py

@@ -627,6 +627,7 @@ def __call__(
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
         max_sequence_length: int = 512,
         text_encoder_out_layers: tuple[int] = (9, 18, 27),
+        max_area: int = 1024**2,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -668,6 +669,9 @@ def __call__(
                 Maximum sequence length for the prompt.
             text_encoder_out_layers (`tuple[int]`):
                 Layer indices for text encoder hidden state extraction.
+            max_area (`int`, defaults to `1024 ** 2`):
+                The maximum area (in pixels) allowed for each condition image. Condition images whose area exceeds
+                this value are downscaled to fit it while preserving their aspect ratio.
 
         Examples:
 
@@ -720,8 +724,8 @@ def __call__(
             condition_images = []
             for img in image:
                 image_width, image_height = img.size
-                if image_width * image_height > 1024 * 1024:
-                    img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
+                if image_width * image_height > max_area:
+                    img = self.image_processor._resize_to_target_area(img, max_area)
                     image_width, image_height = img.size
 
                 multiple_of = self.vae_scale_factor * 2

tests/pipelines/flux2/test_pipeline_flux2.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import torch
+from PIL import Image
 from transformers import AutoProcessor, Mistral3Config, Mistral3ForConditionalGeneration
 
 from diffusers import (
@@ -186,3 +187,15 @@ def test_flux_image_output_shape(self):
                 (expected_height, expected_width),
                 f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}",
             )
+
+    def test_image_input_max_area(self):
+        # `max_area` (previously hardcoded to 1024**2) is the condition-image downscale threshold:
+        # condition images whose area exceeds it are downscaled while preserving aspect ratio.
+        pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        inputs = self.get_dummy_inputs(torch_device)
+        height, width = inputs["height"], inputs["width"]
+
+        inputs.update({"image": Image.new("RGB", (128, 128)), "max_area": 64 * 64})
+        image = pipe(**inputs).images[0]
+        output_height, output_width, _ = image.shape
+        self.assertEqual((output_height, output_width), (height, width))

tests/pipelines/flux2/test_pipeline_flux2_klein.py

@@ -17,6 +17,8 @@
 from ...testing_utils import (
     backend_empty_cache,
     backend_synchronize,
+    nightly,
+    require_big_accelerator,
     require_torch_neuron,
     torch_device,
 )
@@ -183,6 +185,18 @@ def test_image_input(self):
         # fmt: on
         assert np.allclose(expected_slice, generated_slice, atol=1e-4, rtol=1e-4)
 
+    def test_image_input_max_area(self):
+        # `max_area` (previously hardcoded to 1024**2) is the condition-image downscale threshold:
+        # condition images whose area exceeds it are downscaled while preserving aspect ratio.
+        pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        inputs = self.get_dummy_inputs(torch_device)
+        height, width = inputs["height"], inputs["width"]
+
+        inputs.update({"image": Image.new("RGB", (128, 128)), "max_area": 64 * 64})
+        image = pipe(**inputs).images[0]
+        output_height, output_width, _ = image.shape
+        self.assertEqual((output_height, output_width), (height, width))
+
     @unittest.skip("Needs to be revisited")
     def test_encode_prompt_works_in_isolation(self):
         pass
@@ -279,3 +293,49 @@ def test_flux2_klein_neuron_compile_128(self):
             (image >= 0.0).all() and (image <= 1.0).all(),
             "Output pixel values outside [0, 1]",
         )
+
+
+@nightly
+@require_big_accelerator
+class Flux2KleinPipelineConditionImageSlowTests(unittest.TestCase):
+    ckpt_id = "black-forest-labs/FLUX.2-klein-4B"
+    prompt = "A small cactus with a happy face in the Sahara desert."
+
+    def setUp(self):
+        super().setUp()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def tearDown(self):
+        super().tearDown()
+        gc.collect()
+        backend_empty_cache(torch_device)
+
+    def test_flux2_klein_2048_condition_image(self):
+        # A 2048x2048 condition image used to be silently downscaled to fit the hardcoded
+        # 1024**2 threshold; passing max_area=2048**2 lets the pipeline consume it at full
+        # resolution.
+        pipe = Flux2KleinPipeline.from_pretrained(self.ckpt_id, torch_dtype=torch.bfloat16)
+        pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        generator = torch.Generator("cpu").manual_seed(0)
+        condition_image = Image.new("RGB", (2048, 2048), (128, 128, 128))
+        image = pipe(
+            prompt=self.prompt,
+            image=condition_image,
+            height=512,
+            width=512,
+            num_inference_steps=4,
+            guidance_scale=1.0,
+            generator=generator,
+            max_area=2048 * 2048,
+            output_type="np",
+        ).images
+
+        self.assertEqual(image.shape, (1, 512, 512, 3))
+        self.assertFalse(np.isnan(image).any(), "Output contains NaN values")
+        self.assertTrue(
+            (image >= 0.0).all() and (image <= 1.0).all(),
+            "Output pixel values outside [0, 1]",
+        )

tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py

@@ -167,6 +167,19 @@ def test_without_image(self):
         image = pipe(**inputs).images
         self.assertEqual(image.shape, (1, 8, 8, 3))
 
+    def test_image_input_max_area(self):
+        # `max_area` (previously hardcoded to 1024**2) is the condition-image downscale threshold:
+        # condition images whose area exceeds it are downscaled while preserving aspect ratio.
+        pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
+        inputs = self.get_dummy_inputs(torch_device)
+        height, width = inputs["height"], inputs["width"]
+
+        # the dummy 64x64 condition image exceeds max_area -> downscale path
+        inputs["max_area"] = 32 * 32
+        image = pipe(**inputs).images[0]
+        output_height, output_width, _ = image.shape
+        self.assertEqual((output_height, output_width), (height, width))
+
     @unittest.skip("Needs to be revisited")
     def test_encode_prompt_works_in_isolation(self):
         pass