Flux2 pipelines: expose the hardcoded 1MP condition-image downscale threshold (max_area)

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

@@ -1,4 +1,5 @@
 import torch
+from PIL import Image
 from transformers import AutoProcessor, Mistral3Config, Mistral3ForConditionalGeneration
 
 from diffusers import (
@@ -184,6 +185,18 @@ def test_flux_image_output_shape(self):
                 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.get_pipeline().to(torch_device)
+        inputs = self.get_dummy_inputs()
+        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
+        assert (output_height, output_width) == (height, width)
+
 
 class TestFlux2PipelineMemory(Flux2PipelineTesterConfig, MemoryTesterMixin):
     """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux2 pipeline."""

tests/pipelines/flux2/test_pipeline_flux2_klein.py

@@ -18,6 +18,8 @@
     assert_tensors_close,
     backend_empty_cache,
     backend_synchronize,
+    nightly,
+    require_big_accelerator,
     require_torch_neuron,
     torch_device,
 )
@@ -189,6 +191,18 @@ def test_image_input(self):
         # fmt: on
         assert_tensors_close(generated_slice, expected_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.get_pipeline().to(torch_device)
+        inputs = self.get_dummy_inputs()
+        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
+        assert (output_height, output_width) == (height, width)
+
     @pytest.mark.skip("Needs to be revisited")
     def test_encode_prompt_works_in_isolation(self):
         pass
@@ -283,3 +297,44 @@ def test_flux2_klein_neuron_compile_128(self):
         assert image.shape == (1, 128, 128, 3)
         assert not np.isnan(image).any(), "Output contains NaN values"
         assert (image >= 0.0).all() and (image <= 1.0).all(), "Output pixel values outside [0, 1]"
+
+
+@nightly
+@require_big_accelerator
+class TestFlux2KleinPipelineConditionImageSlow:
+    ckpt_id = "black-forest-labs/FLUX.2-klein-4B"
+    prompt = "A small cactus with a happy face in the Sahara desert."
+
+    @pytest.fixture(autouse=True)
+    def cleanup(self):
+        gc.collect()
+        backend_empty_cache(torch_device)
+        yield
+        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
+
+        assert image.shape == (1, 512, 512, 3)
+        assert not np.isnan(image).any(), "Output contains NaN values"
+        assert (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
         assert image.shape == (1, *self.output_shape)
 
+    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.get_pipeline().to(torch_device)
+        inputs = self.get_dummy_inputs()
+        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
+        assert (output_height, output_width) == (height, width)
+
     @pytest.mark.skip("Needs to be revisited")
     def test_encode_prompt_works_in_isolation(self):
         pass