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 73c480a82809..2507ab4d6f4e 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2.py +++ b/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.""" diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index 0d7139b21e16..288b8ca1f0ff 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/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]" diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py b/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py index 29113510c1c8..7f048d03a06f 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 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