From d0f90c81b90aa69144f8f9f42d9abda0dacc513c Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 21 Aug 2026 14:24:16 +0000 Subject: [PATCH] migrate remaining flux pipeline tests to use new mixins --- tests/pipelines/flux/test_pipeline_flux.py | 132 +------------- .../flux/test_pipeline_flux_control.py | 118 +++++++------ .../test_pipeline_flux_control_img2img.py | 59 +++---- .../test_pipeline_flux_control_inpaint.py | 106 ++++++------ .../pipelines/flux/test_pipeline_flux_fill.py | 69 ++++---- .../flux/test_pipeline_flux_img2img.py | 77 +++++---- .../flux/test_pipeline_flux_inpaint.py | 70 ++++---- .../flux/test_pipeline_flux_kontext.py | 106 ++++++------ .../test_pipeline_flux_kontext_inpaint.py | 108 ++++++------ .../flux/test_pipeline_flux_redux.py | 12 +- tests/pipelines/flux/testing_utils.py | 161 ++++++++++++++++++ 11 files changed, 548 insertions(+), 470 deletions(-) create mode 100644 tests/pipelines/flux/testing_utils.py diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py index 11dd28244bca..70cd88a48c9b 100644 --- a/tests/pipelines/flux/test_pipeline_flux.py +++ b/tests/pipelines/flux/test_pipeline_flux.py @@ -1,7 +1,5 @@ import gc -import inspect import os -from typing import Any import numpy as np import pytest @@ -16,16 +14,13 @@ FluxPipeline, FluxTransformer2DModel, ) -from diffusers.loaders import FluxIPAdapterMixin from diffusers.utils.import_utils import is_peft_available from ...models.testing_utils.lora import check_if_lora_correctly_set -from ...models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict from ...testing_utils import ( Expectations, assert_tensors_close, backend_empty_cache, - is_ip_adapter, nightly, numpy_cosine_similarity_distance, require_big_accelerator, @@ -44,6 +39,7 @@ PyramidAttentionBroadcastTesterMixin, TaylorSeerCacheTesterMixin, ) +from .testing_utils import FluxIPAdapterTesterMixin if is_peft_available(): @@ -206,133 +202,9 @@ def test_flux_negative_embeds_shape_check(self): pipe(**base_inputs, true_cfg_scale=1.0, generator=torch.manual_seed(0)) -@is_ip_adapter -class TestFluxPipelineIPAdapter(FluxPipelineTesterConfig): +class TestFluxPipelineIPAdapter(FluxPipelineTesterConfig, FluxIPAdapterTesterMixin): """IP-Adapter tests for the Flux pipeline.""" - def test_pipeline_signature(self): - parameters = inspect.signature(self.pipeline_class.__call__).parameters - - assert issubclass(self.pipeline_class, FluxIPAdapterMixin) - assert "ip_adapter_image" in parameters, ( - "`ip_adapter_image` argument must be supported by the `__call__` method" - ) - assert "ip_adapter_image_embeds" in parameters, ( - "`ip_adapter_image_embeds` argument must be supported by the `__call__` method" - ) - - def _get_dummy_image_embeds(self, image_embed_dim: int = 768): - return torch.randn((1, 1, image_embed_dim), device=torch_device) - - def _modify_inputs_for_ip_adapter_test(self, inputs: dict[str, Any]): - inputs["negative_prompt"] = "" - if "true_cfg_scale" in inspect.signature(self.pipeline_class.__call__).parameters: - inputs["true_cfg_scale"] = 4.0 - # Request torch outputs so comparisons run on torch tensors directly (see `BasePipelineTesterConfig`). - inputs["output_type"] = "pt" - inputs["return_dict"] = False - return inputs - - def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): - r"""Tests for IP-Adapter. - - The following scenarios are tested: - - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter. - - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. - """ - # Raising the tolerance for this test when it's run on a CPU because we compare against static slices and - # that can be shaky (with a VVVV low probability). - expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff - - components = self.get_dummy_components() - pipe = self.pipeline_class(**components).to(torch_device) - pipe.set_progress_bar_config(disable=None) - image_embed_dim = ( - pipe.transformer.config.pooled_projection_dim - if hasattr(pipe.transformer.config, "pooled_projection_dim") - else 768 - ) - - # forward pass without ip adapter - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - if expected_pipe_slice is None: - output_without_adapter = pipe(**inputs)[0] - else: - output_without_adapter = expected_pipe_slice - - # 1. Single IP-Adapter test cases - adapter_state_dict = create_flux_ip_adapter_state_dict(pipe.transformer) - # Load through the pipeline's public IP-Adapter API. `image_encoder_pretrained_model_name_or_path=None` - # skips fetching a CLIP image encoder since we feed pre-computed `ip_adapter_image_embeds` directly. - pipe.load_ip_adapter(adapter_state_dict, weight_name="", image_encoder_pretrained_model_name_or_path=None) - - # forward pass with single ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - pipe.set_ip_adapter_scale(0.0) - output_without_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with single ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] - pipe.set_ip_adapter_scale(42.0) - output_with_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() - - assert_tensors_close( - output_without_adapter_scale, - output_without_adapter, - atol=expected_max_diff, - msg="Output without ip-adapter must be same as normal inference", - ) - max_diff_with_adapter_scale = (output_with_adapter_scale - output_without_adapter).abs().max() - assert max_diff_with_adapter_scale > 1e-2, "Output with ip-adapter must be different from normal inference" - - # 2. Multi IP-Adapter test cases - adapter_state_dict_1 = create_flux_ip_adapter_state_dict(pipe.transformer) - adapter_state_dict_2 = create_flux_ip_adapter_state_dict(pipe.transformer) - pipe.load_ip_adapter( - [adapter_state_dict_1, adapter_state_dict_2], - weight_name=["", ""], - image_encoder_pretrained_model_name_or_path=None, - ) - - # forward pass with multi ip adapter, but scale=0 which should have no effect - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - pipe.set_ip_adapter_scale([0.0, 0.0]) - output_without_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - # forward pass with multi ip adapter, but with scale of adapter weights - inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) - inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 - pipe.set_ip_adapter_scale([42.0, 42.0]) - output_with_multi_adapter_scale = pipe(**inputs)[0] - if expected_pipe_slice is not None: - output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() - - assert_tensors_close( - output_without_multi_adapter_scale, - output_without_adapter, - atol=expected_max_diff, - msg="Output without multi-ip-adapter must be same as normal inference", - ) - max_diff_with_multi_adapter_scale = (output_with_multi_adapter_scale - output_without_adapter).abs().max() - assert max_diff_with_multi_adapter_scale > 1e-2, ( - "Output with multi-ip-adapter scale must be different from normal inference" - ) - class TestFluxPipelineMemory(FluxPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_control.py b/tests/pipelines/flux/test_pipeline_flux_control.py index 44efca9b9f0e..6d07a91becd3 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control.py +++ b/tests/pipelines/flux/test_pipeline_flux_control.py @@ -1,25 +1,25 @@ -import unittest - -import numpy as np import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxControlPipeline, FluxTransformer2DModel -from ...testing_utils import torch_device -from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + check_qkv_fused_layers_exist, +) -class FluxControlPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class FluxControlPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxControlPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -86,84 +86,88 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - + def get_dummy_inputs(self): control_image = Image.new("RGB", (16, 16), 0) inputs = { "prompt": "A painting of a squirrel eating a burger", "control_image": control_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 8, "width": 8, "max_sequence_length": 48, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestFluxControlPipeline(FluxControlPipelineTesterConfig, PipelineTesterMixin): def test_flux_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU to ensure determinism for the device-dependent torch.Generator. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + inputs = self.get_dummy_inputs() + original_image_slice = pipe(**inputs).images[0, -3:, -3:, -1] # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. pipe.transformer.fuse_qkv_projections() - self.assertTrue( - check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), - ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."), + assert check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), ( + "Something wrong with the fused attention layers. Expected all the attention projections to be fused." ) - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + inputs = self.get_dummy_inputs() + image_slice_fused = pipe(**inputs).images[0, -3:, -3:, -1] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + inputs = self.get_dummy_inputs() + image_slice_disabled = pipe(**inputs).images[0, -3:, -3:, -1] + + assert_tensors_close( + image_slice_fused, + original_image_slice, + atol=1e-3, + rtol=1e-3, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_slice_disabled, + image_slice_fused, + atol=1e-3, + rtol=1e-3, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." + assert_tensors_close( + image_slice_disabled, + original_image_slice, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -172,5 +176,11 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxControlPipelineMemory(FluxControlPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux control pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py index 0f0bc0934115..516c6b92b504 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_img2img.py @@ -1,6 +1,3 @@ -import unittest - -import numpy as np import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel @@ -12,18 +9,17 @@ FluxTransformer2DModel, ) -from ...testing_utils import enable_full_determinism, torch_device -from ..test_pipelines_common import PipelineTesterMixin - +from ...testing_utils import torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -enable_full_determinism() - -class FluxControlImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class FluxControlImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxControlImg2ImgPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - test_xformers_attention = False + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -90,12 +86,7 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - + def get_dummy_inputs(self): image = Image.new("RGB", (16, 16), 0) control_image = Image.new("RGB", (16, 16), 0) @@ -103,36 +94,40 @@ def get_dummy_inputs(self, device, seed=0): "prompt": "A painting of a squirrel eating a burger", "image": image, "control_image": control_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 8, "width": 8, "max_sequence_length": 48, "strength": 0.8, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestFluxControlImg2ImgPipeline(FluxControlImg2ImgPipelineTesterConfig, PipelineTesterMixin): def test_flux_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -141,5 +136,11 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxControlImg2ImgPipelineMemory(FluxControlImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux control img2img pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py index ae2b6b829e54..d4e7018bc8a6 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py @@ -1,6 +1,3 @@ -import unittest - -import numpy as np import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel @@ -12,19 +9,22 @@ FluxTransformer2DModel, ) -from ...testing_utils import ( - torch_device, +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + check_qkv_fused_layers_exist, ) -from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist -class FluxControlInpaintPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class FluxControlInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxControlInpaintPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - - # there is no xformers processor for Flux - test_xformers_attention = False + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -91,12 +91,7 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - + def get_dummy_inputs(self): image = Image.new("RGB", (8, 8), 0) control_image = Image.new("RGB", (8, 8), 0) mask_image = Image.new("RGB", (8, 8), 255) @@ -104,7 +99,7 @@ def get_dummy_inputs(self, device, seed=0): inputs = { "prompt": "A painting of a squirrel eating a burger", "control_image": control_image, - "generator": generator, + "generator": self.get_generator(0), "image": image, "mask_image": mask_image, "strength": 0.8, @@ -113,51 +108,60 @@ def get_dummy_inputs(self, device, seed=0): "height": 8, "width": 8, "max_sequence_length": 48, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestFluxControlInpaintPipeline(FluxControlInpaintPipelineTesterConfig, PipelineTesterMixin): def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU to ensure determinism for the device-dependent torch.Generator. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + inputs = self.get_dummy_inputs() + original_image_slice = pipe(**inputs).images[0, -3:, -3:, -1] # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. pipe.transformer.fuse_qkv_projections() - self.assertTrue( - check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), - ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."), + assert check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), ( + "Something wrong with the fused attention layers. Expected all the attention projections to be fused." ) - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + inputs = self.get_dummy_inputs() + image_slice_fused = pipe(**inputs).images[0, -3:, -3:, -1] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + inputs = self.get_dummy_inputs() + image_slice_disabled = pipe(**inputs).images[0, -3:, -3:, -1] + + assert_tensors_close( + image_slice_fused, + original_image_slice, + atol=1e-3, + rtol=1e-3, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_slice_disabled, + image_slice_fused, + atol=1e-3, + rtol=1e-3, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." + assert_tensors_close( + image_slice_disabled, + original_image_slice, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -166,5 +170,11 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxControlInpaintPipelineMemory(FluxControlInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux control inpaint pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_fill.py b/tests/pipelines/flux/test_pipeline_flux_fill.py index 42cd1efad495..4dc3aa5f51d6 100644 --- a/tests/pipelines/flux/test_pipeline_flux_fill.py +++ b/tests/pipelines/flux/test_pipeline_flux_fill.py @@ -1,30 +1,21 @@ import random -import unittest -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxFillPipeline, FluxTransformer2DModel -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_pipelines_common import PipelineTesterMixin +from ...testing_utils import floats_tensor, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -enable_full_determinism() - - -class FluxFillPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class FluxFillPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxFillPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -91,47 +82,47 @@ def get_dummy_components(self): "vae": vae, } - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - mask_image = torch.ones((1, 1, 32, 32)).to(device) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) + mask_image = torch.ones((1, 1, 32, 32)).to(torch_device) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": image, "mask_image": mask_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 32, "width": 32, "max_sequence_length": 48, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestFluxFillPipeline(FluxFillPipelineTesterConfig, PipelineTesterMixin): def test_flux_fill_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -140,8 +131,14 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=1e-3) + super().test_inference_batch_single_identical(expected_max_diff=1e-3) + + +class TestFluxFillPipelineMemory(FluxFillPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux fill pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_img2img.py b/tests/pipelines/flux/test_pipeline_flux_img2img.py index 4b85243938ec..93736df73ff5 100644 --- a/tests/pipelines/flux/test_pipeline_flux_img2img.py +++ b/tests/pipelines/flux/test_pipeline_flux_img2img.py @@ -1,28 +1,22 @@ import random -import unittest -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxImg2ImgPipeline, FluxTransformer2DModel -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_pipelines_common import FluxIPAdapterTesterMixin, PipelineTesterMixin +from ...testing_utils import floats_tensor, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from .testing_utils import FluxIPAdapterTesterMixin -enable_full_determinism() - - -class FluxImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMixin, FluxIPAdapterTesterMixin): +class FluxImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxImg2ImgPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - test_xformers_attention = False + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self): torch.manual_seed(0) @@ -90,46 +84,46 @@ def get_dummy_components(self): "feature_extractor": None, } - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 8, "width": 8, "max_sequence_length": 48, "strength": 0.8, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestFluxImg2ImgPipeline(FluxImg2ImgPipelineTesterConfig, PipelineTesterMixin): def test_flux_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_flux_true_cfg_with_negative_embeds(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("generator") prompt = inputs.pop("prompt") @@ -150,14 +144,13 @@ def test_flux_true_cfg_with_negative_embeds(self): cfg_off = pipe(**inputs, generator=torch.manual_seed(0)).images[0] inputs["true_cfg_scale"] = 2.0 cfg_on = pipe(**inputs, generator=torch.manual_seed(0)).images[0] - self.assertFalse( - np.allclose(cfg_off, cfg_on), - "Precomputed negative embeds should enable true CFG when negative_prompt is None.", + assert not torch.allclose(cfg_off, cfg_on), ( + "Precomputed negative embeds should enable true CFG when negative_prompt is None." ) def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -166,5 +159,15 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxImg2ImgPipelineIPAdapter(FluxImg2ImgPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Flux img2img pipeline.""" + + +class TestFluxImg2ImgPipelineMemory(FluxImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux img2img pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_inpaint.py index 14edb9e441b5..87e778083134 100644 --- a/tests/pipelines/flux/test_pipeline_flux_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_inpaint.py @@ -1,28 +1,22 @@ import random -import unittest -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxInpaintPipeline, FluxTransformer2DModel -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, -) -from ..test_pipelines_common import FluxIPAdapterTesterMixin, PipelineTesterMixin +from ...testing_utils import floats_tensor, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from .testing_utils import FluxIPAdapterTesterMixin -enable_full_determinism() - - -class FluxInpaintPipelineFastTests(unittest.TestCase, PipelineTesterMixin, FluxIPAdapterTesterMixin): +class FluxInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxInpaintPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - test_xformers_attention = False + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -90,48 +84,48 @@ def get_dummy_components(self): "feature_extractor": None, } - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - mask_image = torch.ones((1, 1, 32, 32)).to(device) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) + mask_image = torch.ones((1, 1, 32, 32)).to(torch_device) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": image, "mask_image": mask_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 32, "width": 32, "max_sequence_length": 48, "strength": 0.8, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs + +class TestFluxInpaintPipeline(FluxInpaintPipelineTesterConfig, PipelineTesterMixin): def test_flux_inpaint_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -140,5 +134,15 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxInpaintPipelineIPAdapter(FluxInpaintPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Flux inpaint pipeline.""" + + +class TestFluxInpaintPipelineMemory(FluxInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux inpaint pipeline.""" diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext.py b/tests/pipelines/flux/test_pipeline_flux_kontext.py index 1c018f14b522..1fb5a52f7c84 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext.py @@ -1,52 +1,32 @@ -import unittest - -import numpy as np import PIL.Image import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, - FasterCacheConfig, FlowMatchEulerDiscreteScheduler, FluxKontextPipeline, FluxTransformer2DModel, ) from ...testing_utils import torch_device -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, FasterCacheTesterMixin, - FluxIPAdapterTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, ) +from .testing_utils import FluxIPAdapterTesterMixin -class FluxKontextPipelineFastTests( - unittest.TestCase, - PipelineTesterMixin, - FluxIPAdapterTesterMixin, - PyramidAttentionBroadcastTesterMixin, - FasterCacheTesterMixin, -): +class FluxKontextPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxKontextPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( ["image", "prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) - batch_params = frozenset(["image", "prompt"]) - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - - faster_cache_config = FasterCacheConfig( - spatial_attention_block_skip_range=2, - spatial_attention_timestep_skip_range=(-1, 901), - unconditional_batch_skip_range=2, - attention_weight_callback=lambda _: 0.5, - is_guidance_distilled=True, - ) + batch_input_params = frozenset(["image", "prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) @@ -114,47 +94,46 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "feature_extractor": None, } - def get_dummy_inputs(self, device, seed=0): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - + def get_dummy_inputs(self): image = PIL.Image.new("RGB", (32, 32), 0) inputs = { "image": image, "prompt": "A painting of a squirrel eating a burger", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 8, "width": 8, "max_area": 8 * 8, "max_sequence_length": 48, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "_auto_resize": False, } return inputs + +class TestFluxKontextPipeline(FluxKontextPipelineTesterConfig, PipelineTesterMixin): def test_flux_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -163,16 +142,47 @@ def test_flux_image_output_shape(self): inputs.update({"height": height, "width": width, "max_area": height * width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) def test_flux_true_cfg(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("generator") no_true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] inputs["negative_prompt"] = "bad quality" inputs["true_cfg_scale"] = 2.0 true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] - assert not np.allclose(no_true_cfg_out, true_cfg_out) + assert not torch.allclose(no_true_cfg_out, true_cfg_out), ( + "Outputs should be different when true_cfg_scale is set." + ) + + +class TestFluxKontextPipelineIPAdapter(FluxKontextPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Flux Kontext pipeline.""" + + +class TestFluxKontextPipelineMemory(FluxKontextPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux Kontext pipeline.""" + + +class TestFluxKontextPipelinePyramidAttentionBroadcast( + FluxKontextPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + """Pyramid Attention Broadcast cache tests for the Flux Kontext pipeline.""" + + +class TestFluxKontextPipelineFasterCache(FluxKontextPipelineTesterConfig, FasterCacheTesterMixin): + """FasterCache tests for the Flux Kontext pipeline.""" + + # Flux is guidance-distilled, so the FasterCache tester must skip the low/high-frequency-delta state checks. + FASTER_CACHE_CONFIG = { + "spatial_attention_block_skip_range": 2, + "spatial_attention_timestep_skip_range": (-1, 901), + "unconditional_batch_skip_range": 2, + "attention_weight_callback": lambda _: 0.5, + "is_guidance_distilled": True, + } diff --git a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py index b5f8570ebd1a..cfae5b4a2a38 100644 --- a/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_kontext_inpaint.py @@ -1,52 +1,33 @@ import random -import unittest -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel from diffusers import ( AutoencoderKL, - FasterCacheConfig, FlowMatchEulerDiscreteScheduler, FluxKontextInpaintPipeline, FluxTransformer2DModel, ) from ...testing_utils import floats_tensor, torch_device -from ..test_pipelines_common import ( +from ..testing_utils import ( + BasePipelineTesterConfig, FasterCacheTesterMixin, - FluxIPAdapterTesterMixin, + MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, ) +from .testing_utils import FluxIPAdapterTesterMixin -class FluxKontextInpaintPipelineFastTests( - unittest.TestCase, - PipelineTesterMixin, - FluxIPAdapterTesterMixin, - PyramidAttentionBroadcastTesterMixin, - FasterCacheTesterMixin, -): +class FluxKontextInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxKontextInpaintPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( ["image", "prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] ) - batch_params = frozenset(["image", "prompt"]) - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True - - faster_cache_config = FasterCacheConfig( - spatial_attention_block_skip_range=2, - spatial_attention_timestep_skip_range=(-1, 901), - unconditional_batch_skip_range=2, - attention_weight_callback=lambda _: 0.5, - is_guidance_distilled=True, - ) + batch_input_params = frozenset(["image", "prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) @@ -114,49 +95,49 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "feature_extractor": None, } - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - mask_image = torch.ones((1, 1, 32, 32)).to(device) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) + mask_image = torch.ones((1, 1, 32, 32)).to(torch_device) inputs = { "prompt": "A painting of a squirrel eating a burger", "image": image, "mask_image": mask_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 32, "width": 32, "max_sequence_length": 48, "strength": 0.8, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", "_auto_resize": False, } return inputs + +class TestFluxKontextInpaintPipeline(FluxKontextInpaintPipelineTesterConfig, PipelineTesterMixin): def test_flux_inpaint_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 56)] for height, width in height_width_pairs: @@ -176,16 +157,47 @@ def test_flux_image_output_shape(self): } ) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) def test_flux_true_cfg(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() inputs.pop("generator") no_true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] inputs["negative_prompt"] = "bad quality" inputs["true_cfg_scale"] = 2.0 true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0] - assert not np.allclose(no_true_cfg_out, true_cfg_out) + assert not torch.allclose(no_true_cfg_out, true_cfg_out), ( + "Outputs should be different when true_cfg_scale is set." + ) + + +class TestFluxKontextInpaintPipelineIPAdapter(FluxKontextInpaintPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Flux Kontext inpaint pipeline.""" + + +class TestFluxKontextInpaintPipelineMemory(FluxKontextInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux Kontext inpaint pipeline.""" + + +class TestFluxKontextInpaintPipelinePyramidAttentionBroadcast( + FluxKontextInpaintPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin +): + """Pyramid Attention Broadcast cache tests for the Flux Kontext inpaint pipeline.""" + + +class TestFluxKontextInpaintPipelineFasterCache(FluxKontextInpaintPipelineTesterConfig, FasterCacheTesterMixin): + """FasterCache tests for the Flux Kontext inpaint pipeline.""" + + # Flux is guidance-distilled, so the FasterCache tester must skip the low/high-frequency-delta state checks. + FASTER_CACHE_CONFIG = { + "spatial_attention_block_skip_range": 2, + "spatial_attention_timestep_skip_range": (-1, 901), + "unconditional_batch_skip_range": 2, + "attention_weight_callback": lambda _: 0.5, + "is_guidance_distilled": True, + } diff --git a/tests/pipelines/flux/test_pipeline_flux_redux.py b/tests/pipelines/flux/test_pipeline_flux_redux.py index bbeee28e6a62..bb50bc08f009 100644 --- a/tests/pipelines/flux/test_pipeline_flux_redux.py +++ b/tests/pipelines/flux/test_pipeline_flux_redux.py @@ -1,7 +1,7 @@ import gc -import unittest import numpy as np +import pytest import torch from diffusers import FluxPipeline, FluxPriorReduxPipeline @@ -19,19 +19,17 @@ @slow @require_big_accelerator -class FluxReduxSlowTests(unittest.TestCase): +class TestFluxReduxSlow: pipeline_class = FluxPriorReduxPipeline repo_id = "black-forest-labs/FLUX.1-Redux-dev" base_pipeline_class = FluxPipeline base_repo_id = "black-forest-labs/FLUX.1-schnell" - def setUp(self): - super().setUp() + @pytest.fixture(autouse=True) + def cleanup(self): gc.collect() backend_empty_cache(torch_device) - - def tearDown(self): - super().tearDown() + yield gc.collect() backend_empty_cache(torch_device) diff --git a/tests/pipelines/flux/testing_utils.py b/tests/pipelines/flux/testing_utils.py new file mode 100644 index 000000000000..0168a23a78c2 --- /dev/null +++ b/tests/pipelines/flux/testing_utils.py @@ -0,0 +1,161 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any + +import torch + +from diffusers.loaders import FluxIPAdapterMixin + +from ...testing_utils import assert_tensors_close, is_ip_adapter, torch_device +from ..testing_utils.common import BasePipelineOutputMixin + + +@is_ip_adapter +class FluxIPAdapterTesterMixin(BasePipelineOutputMixin): + """IP-Adapter tests shared by the Flux pipelines in this directory. + + Flux has its own IP-Adapter API (`FluxIPAdapterMixin`, image embeddings sized after the transformer's + `pooled_projection_dim`), so it doesn't reuse the Stable Diffusion `IPAdapterTesterMixin`. Compose it with a + `BasePipelineTesterConfig` subclass in its own test class, separate from the `PipelineTesterMixin` one. + """ + + def test_pipeline_signature(self): + parameters = inspect.signature(self.pipeline_class.__call__).parameters + + assert issubclass(self.pipeline_class, FluxIPAdapterMixin) + assert "ip_adapter_image" in parameters, ( + "`ip_adapter_image` argument must be supported by the `__call__` method" + ) + assert "ip_adapter_image_embeds" in parameters, ( + "`ip_adapter_image_embeds` argument must be supported by the `__call__` method" + ) + + def _get_dummy_image_embeds(self, image_embed_dim: int = 768): + return torch.randn((1, 1, image_embed_dim), device=torch_device) + + def _modify_inputs_for_ip_adapter_test(self, inputs: dict[str, Any]): + inputs["negative_prompt"] = "" + if "true_cfg_scale" in inspect.signature(self.pipeline_class.__call__).parameters: + inputs["true_cfg_scale"] = 4.0 + # Request torch outputs so comparisons run on torch tensors directly (see `BasePipelineTesterConfig`). + inputs["output_type"] = "pt" + inputs["return_dict"] = False + return inputs + + def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=None): + r"""Tests for IP-Adapter. + + The following scenarios are tested: + - Single IP-Adapter with scale=0 should produce same output as no IP-Adapter. + - Multi IP-Adapter with scale=0 should produce same output as no IP-Adapter. + - Single IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. + - Multi IP-Adapter with scale!=0 should produce different output compared to no IP-Adapter. + """ + # The state dict builder is imported here rather than at module scope: it lives in a model test module + # that calls `enable_full_determinism()` on import, which would otherwise flip that global for every test + # collected alongside this directory. + from ...models.transformers.test_models_transformer_flux import create_flux_ip_adapter_state_dict + + # Raising the tolerance for this test when it's run on a CPU because we compare against static slices and + # that can be shaky (with a VVVV low probability). + expected_max_diff = 9e-4 if torch_device == "cpu" else expected_max_diff + + components = self.get_dummy_components() + pipe = self.get_pipeline(**components).to(torch_device) + image_embed_dim = ( + pipe.transformer.config.pooled_projection_dim + if hasattr(pipe.transformer.config, "pooled_projection_dim") + else 768 + ) + + # forward pass without ip adapter + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + if expected_pipe_slice is None: + output_without_adapter = pipe(**inputs)[0] + else: + output_without_adapter = expected_pipe_slice + + # 1. Single IP-Adapter test cases + adapter_state_dict = create_flux_ip_adapter_state_dict(pipe.transformer) + # Load through the pipeline's public IP-Adapter API. `image_encoder_pretrained_model_name_or_path=None` + # skips fetching a CLIP image encoder since we feed pre-computed `ip_adapter_image_embeds` directly. + pipe.load_ip_adapter(adapter_state_dict, weight_name="", image_encoder_pretrained_model_name_or_path=None) + + # forward pass with single ip adapter, but scale=0 which should have no effect + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + pipe.set_ip_adapter_scale(0.0) + output_without_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_without_adapter_scale = output_without_adapter_scale[0, -3:, -3:, -1].flatten() + + # forward pass with single ip adapter, but with scale of adapter weights + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] + pipe.set_ip_adapter_scale(42.0) + output_with_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_with_adapter_scale = output_with_adapter_scale[0, -3:, -3:, -1].flatten() + + assert_tensors_close( + output_without_adapter_scale, + output_without_adapter, + atol=expected_max_diff, + msg="Output without ip-adapter must be same as normal inference", + ) + max_diff_with_adapter_scale = (output_with_adapter_scale - output_without_adapter).abs().max() + assert max_diff_with_adapter_scale > 1e-2, "Output with ip-adapter must be different from normal inference" + + # 2. Multi IP-Adapter test cases + adapter_state_dict_1 = create_flux_ip_adapter_state_dict(pipe.transformer) + adapter_state_dict_2 = create_flux_ip_adapter_state_dict(pipe.transformer) + pipe.load_ip_adapter( + [adapter_state_dict_1, adapter_state_dict_2], + weight_name=["", ""], + image_encoder_pretrained_model_name_or_path=None, + ) + + # forward pass with multi ip adapter, but scale=0 which should have no effect + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + pipe.set_ip_adapter_scale([0.0, 0.0]) + output_without_multi_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_without_multi_adapter_scale = output_without_multi_adapter_scale[0, -3:, -3:, -1].flatten() + + # forward pass with multi ip adapter, but with scale of adapter weights + inputs = self._modify_inputs_for_ip_adapter_test(self.get_dummy_inputs()) + inputs["ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + inputs["negative_ip_adapter_image_embeds"] = [self._get_dummy_image_embeds(image_embed_dim)] * 2 + pipe.set_ip_adapter_scale([42.0, 42.0]) + output_with_multi_adapter_scale = pipe(**inputs)[0] + if expected_pipe_slice is not None: + output_with_multi_adapter_scale = output_with_multi_adapter_scale[0, -3:, -3:, -1].flatten() + + assert_tensors_close( + output_without_multi_adapter_scale, + output_without_adapter, + atol=expected_max_diff, + msg="Output without multi-ip-adapter must be same as normal inference", + ) + max_diff_with_multi_adapter_scale = (output_with_multi_adapter_scale - output_without_adapter).abs().max() + assert max_diff_with_multi_adapter_scale > 1e-2, ( + "Output with multi-ip-adapter scale must be different from normal inference" + )