diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2.py b/src/diffusers/pipelines/flux2/pipeline_flux2.py index b1645b4ae244..990c2e82c35d 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2.py +++ b/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 diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py index d768e6127f26..1ec883694047 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein.py +++ b/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 diff --git a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py b/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py index 78ed42f20afb..ba1d9e191943 100644 --- a/src/diffusers/pipelines/flux2/pipeline_flux2_klein_kv.py +++ b/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 diff --git a/tests/pipelines/flux2/test_pipeline_flux2.py b/tests/pipelines/flux2/test_pipeline_flux2.py index c065ff3ed191..1f2e99253bec 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2.py +++ b/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)) diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index 6db70c6367ab..4291b03ea861 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/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]", + ) diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py b/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py index 4f77579af6d6..553b5e9c1176 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py +++ b/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