diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 8be893cc8213..35197ccddd0b 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -36,6 +36,7 @@ torch_device, ) from ..testing_utils import ( + AttentionTesterMixin, BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, @@ -203,6 +204,10 @@ class TestAutoencoderKLMemory(AutoencoderKLTesterConfig, MemoryTesterMixin): """Memory optimization tests for AutoencoderKL.""" +class TestAutoencoderKLAttention(AutoencoderKLTesterConfig, AttentionTesterMixin): + """Attention processor tests for AutoencoderKL.""" + + class TestAutoencoderKLSlicingTiling(AutoencoderKLTesterConfig, AutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderKL.""" diff --git a/tests/models/testing_utils/attention.py b/tests/models/testing_utils/attention.py index f31323d1bf52..62aae39130c0 100644 --- a/tests/models/testing_utils/attention.py +++ b/tests/models/testing_utils/attention.py @@ -15,13 +15,14 @@ import gc import logging +from typing import NamedTuple import pytest import torch -from diffusers.models.attention import AttentionModuleMixin +from diffusers.models.attention import AttentionMixin, AttentionModuleMixin from diffusers.models.attention_dispatch import AttentionBackendName, _AttentionBackendRegistry, attention_backend -from diffusers.models.attention_processor import AttnProcessor +from diffusers.models.attention_processor import Attention, AttnProcessor from diffusers.utils import is_kernels_available, is_torch_version from ...testing_utils import assert_tensors_close, backend_empty_cache, is_attention, is_torch_compile, torch_device @@ -123,6 +124,40 @@ def _skip_if_backend_requires_nondeterminism(backend): ) +class _FusionPath(NamedTuple): + """How one of the two `fuse_qkv_projections()` implementations behaves. + + `attention_cls` is the attention base class that implementation walks: a model built on the *other* base class + gets fused right past, which is why `test_fuse_unfuse_qkv_projections` checks it before anything else. + `swaps_in_fused_processors` says what unfusing is expected to undo. + """ + + attention_cls: type + swaps_in_fused_processors: bool + + +# `AttentionMixin.fuse_qkv_projections()`, inherited by Flux, Flux2, Chroma, Wan, ... It walks `AttentionModuleMixin` +# modules and leaves the processors alone, so `unfuse_qkv_projections()` deletes the fused layers again. +_SHARED_FUSION = _FusionPath(attention_cls=AttentionModuleMixin, swaps_in_fused_processors=False) + +# The per-model implementations on UNets, SD3, PixArt, AuraFlow, CogVideoX, AutoencoderKL, ... They walk `Attention` +# modules, swap in a dedicated `Fused*` processor and stash the originals in `original_attn_processors`, so +# `unfuse_qkv_projections()` only puts those processors back - the fused layers stay on the modules. +_LEGACY_FUSION = _FusionPath(attention_cls=Attention, swaps_in_fused_processors=True) + + +def _fusion_path(model): + """Which of the two `fuse_qkv_projections()` implementations this model inherits.""" + if type(model).fuse_qkv_projections is AttentionMixin.fuse_qkv_projections: + return _SHARED_FUSION + return _LEGACY_FUSION + + +def _fused_layer_name(module): + """Name of the layer `fuse_projections()` creates: cross-attention fuses K/V only, self-attention fuses Q/K/V.""" + return "to_kv" if getattr(module, "is_cross_attention", False) else "to_qkv" + + @is_attention class AttentionTesterMixin: """ @@ -152,7 +187,7 @@ def teardown_method(self): backend_empty_cache(torch_device) @torch.no_grad() - def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0): + def test_fuse_unfuse_qkv_projections(self, request, atol=1e-3, rtol=0): init_dict = self.get_init_dict() inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict) @@ -162,46 +197,100 @@ def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0): if not hasattr(model, "fuse_qkv_projections"): pytest.skip("Model does not support QKV projection fusion.") + fusion = _fusion_path(model) + + attention_modules = [ + module for module in model.modules() if isinstance(module, (AttentionModuleMixin, Attention)) + ] + if not attention_modules: + pytest.skip("Model has no attention modules to fuse.") + + walked_modules = [module for module in attention_modules if isinstance(module, fusion.attention_cls)] + stranded_modules = [module for module in attention_modules if not isinstance(module, fusion.attention_cls)] + + if not walked_modules: + # Every attention module in this model is of the base class the model's `fuse_qkv_projections()` does not + # walk, so fusing silently does nothing. Marked xfail rather than skipped so that fixing the model shows + # up as an XPASS instead of staying quietly green. + request.node.add_marker( + pytest.mark.xfail( + reason=( + f"{type(model).__name__}.fuse_qkv_projections() only walks " + f"`{fusion.attention_cls.__name__}` modules, but every attention module in this model is " + f"a `{type(stranded_modules[0]).__name__}` instance, so nothing is ever fused." + ), + strict=True, + ) + ) + + assert walked_modules, ( + f"{type(model).__name__}.fuse_qkv_projections() does not reach any of this model's attention modules." + ) + + fusable_modules = [module for module in walked_modules if getattr(module, "_supports_qkv_fusion", True)] + if not fusable_modules: + pytest.skip("Model's attention modules do not support QKV projection fusion.") + output_before_fusion = model(**inputs_dict, return_dict=False)[0] + processors_before_fusion = model.attn_processors model.fuse_qkv_projections() - has_fused_projections = False - for module in model.modules(): - if isinstance(module, AttentionModuleMixin): - if hasattr(module, "to_qkv") or hasattr(module, "to_kv"): - has_fused_projections = True - assert module.fused_projections, "fused_projections flag should be True" - break - - if has_fused_projections: - output_after_fusion = model(**inputs_dict, return_dict=False)[0] - - assert_tensors_close( - output_before_fusion, - output_after_fusion, - atol=atol, - rtol=rtol, - msg="Output should not change after fusing projections", + for module in fusable_modules: + assert module.fused_projections, "fused_projections flag should be True" + layer_name = _fused_layer_name(module) + assert getattr(module, layer_name, None) is not None, ( + f"{type(module).__name__} should expose a fused `{layer_name}` layer after fusing." ) - model.unfuse_qkv_projections() + processors_after_fusion = model.attn_processors + assert len(processors_after_fusion) == len(processors_before_fusion), ( + "Fusing projections should not change the number of attention processors." + ) + if fusion.swaps_in_fused_processors: + for name, processor in processors_after_fusion.items(): + assert type(processor).__name__.startswith("Fused"), ( + f"Processor {name} should be a fused processor after fusing, got {type(processor).__name__}." + ) - for module in model.modules(): - if isinstance(module, AttentionModuleMixin): - assert not hasattr(module, "to_qkv"), "to_qkv should be removed after unfusing" - assert not hasattr(module, "to_kv"), "to_kv should be removed after unfusing" - assert not module.fused_projections, "fused_projections flag should be False" + output_after_fusion = model(**inputs_dict, return_dict=False)[0] - output_after_unfusion = model(**inputs_dict, return_dict=False)[0] + assert_tensors_close( + output_before_fusion, + output_after_fusion, + atol=atol, + rtol=rtol, + msg="Output should not change after fusing projections", + ) + + model.unfuse_qkv_projections() - assert_tensors_close( - output_before_fusion, - output_after_unfusion, - atol=atol, - rtol=rtol, - msg="Output should match original after unfusing projections", + if fusion.swaps_in_fused_processors: + # This path only restores the original processors; the fused layers themselves are left in place. + processors_after_unfusion = model.attn_processors + assert len(processors_after_unfusion) == len(processors_before_fusion), ( + "Unfusing projections should not change the number of attention processors." ) + for name, processor in processors_after_unfusion.items(): + assert type(processor) is type(processors_before_fusion[name]), ( + f"Processor {name} should be restored to {type(processors_before_fusion[name]).__name__} " + f"after unfusing, got {type(processor).__name__}." + ) + else: + for module in fusable_modules: + assert not hasattr(module, "to_qkv"), "to_qkv should be removed after unfusing" + assert not hasattr(module, "to_kv"), "to_kv should be removed after unfusing" + assert not module.fused_projections, "fused_projections flag should be False" + + output_after_unfusion = model(**inputs_dict, return_dict=False)[0] + + assert_tensors_close( + output_before_fusion, + output_after_unfusion, + atol=atol, + rtol=rtol, + msg="Output should match original after unfusing projections", + ) def test_get_set_processor(self): init_dict = self.get_init_dict() diff --git a/tests/models/transformers/test_models_transformer_chroma.py b/tests/models/transformers/test_models_transformer_chroma.py index dc300fbbe716..cec055306a52 100644 --- a/tests/models/transformers/test_models_transformer_chroma.py +++ b/tests/models/transformers/test_models_transformer_chroma.py @@ -20,6 +20,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( + AttentionTesterMixin, BaseModelTesterConfig, LoraHotSwappingForModelTesterMixin, LoraTesterMixin, @@ -127,6 +128,10 @@ def test_deprecated_inputs_img_txt_ids_3d(self): ) +class TestChromaTransformerAttention(ChromaTransformerTesterConfig, AttentionTesterMixin): + """Attention processor tests for Chroma Transformer.""" + + class TestChromaTransformerTraining(ChromaTransformerTesterConfig, TrainingTesterMixin): def test_gradient_checkpointing_is_applied(self): expected_set = {"ChromaTransformer2DModel"} diff --git a/tests/models/transformers/test_models_transformer_hunyuan_dit.py b/tests/models/transformers/test_models_transformer_hunyuan_dit.py index eb8728976153..4b47e5f9cef8 100644 --- a/tests/models/transformers/test_models_transformer_hunyuan_dit.py +++ b/tests/models/transformers/test_models_transformer_hunyuan_dit.py @@ -20,6 +20,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( + AttentionTesterMixin, BaseModelTesterConfig, ModelTesterMixin, TrainingTesterMixin, @@ -125,6 +126,10 @@ def test_output(self, base_model_output): super().test_output(base_model_output, expected_output_shape=(batch_size,) + self.output_shape) +class TestHunyuanDiTAttention(HunyuanDiTTesterConfig, AttentionTesterMixin): + """Attention processor tests for HunyuanDiT.""" + + class TestHunyuanDiTTraining(HunyuanDiTTesterConfig, TrainingTesterMixin): def test_gradient_checkpointing_is_applied(self): expected_set = {"HunyuanDiT2DModel"} diff --git a/tests/models/transformers/test_models_transformer_ltx2.py b/tests/models/transformers/test_models_transformer_ltx2.py index 0131d372341b..516c478df499 100644 --- a/tests/models/transformers/test_models_transformer_ltx2.py +++ b/tests/models/transformers/test_models_transformer_ltx2.py @@ -173,11 +173,16 @@ def test_gradient_checkpointing_is_applied(self): class TestLTX2TransformerAttention(LTX2TransformerTesterConfig, AttentionTesterMixin): """Attention processor tests for LTX2 Video Transformer.""" - @pytest.mark.skip( - "LTX2Attention does not set is_cross_attention, so fuse_projections tries to fuse Q+K+V together even for cross-attention modules with different input dimensions." + @pytest.mark.xfail( + reason=( + "LTX2Attention does not set is_cross_attention, so fuse_projections tries to fuse Q+K+V together even " + "for cross-attention modules with different input dimensions." + ), + raises=RuntimeError, + strict=True, ) - def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0): - pass + def test_fuse_unfuse_qkv_projections(self, request, atol=1e-3, rtol=0): + super().test_fuse_unfuse_qkv_projections(request, atol=atol, rtol=rtol) class TestLTX2TransformerCompile(LTX2TransformerTesterConfig, TorchCompileTesterMixin): diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index 6294bf80635a..d272934d86df 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import SD3Transformer2DModel @@ -20,6 +21,7 @@ from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import ( + AttentionTesterMixin, BaseModelTesterConfig, BitsAndBytesTesterMixin, LoraTesterMixin, @@ -117,6 +119,10 @@ def test_gradient_checkpointing_is_applied(self): super().test_gradient_checkpointing_is_applied(expected_set=expected_set) +class TestSD3TransformerAttention(SD3TransformerTesterConfig, AttentionTesterMixin): + """Attention processor tests for SD3 Transformer.""" + + class TestSD3TransformerCompile(SD3TransformerTesterConfig, TorchCompileTesterMixin): pass @@ -216,6 +222,22 @@ def test_gradient_checkpointing_is_applied(self): super().test_gradient_checkpointing_is_applied(expected_set=expected_set) +class TestSD35TransformerAttention(SD35TransformerTesterConfig, AttentionTesterMixin): + """Attention processor tests for SD3.5 Transformer.""" + + @pytest.mark.xfail( + reason=( + "fuse_qkv_projections() sets FusedJointAttnProcessor2_0 on every attention module, including the " + "self-attention-only `attn2` of the dual-attention layers, which is then called without " + "`encoder_hidden_states` and raises." + ), + raises=AttributeError, + strict=True, + ) + def test_fuse_unfuse_qkv_projections(self, request, atol=1e-3, rtol=0): + super().test_fuse_unfuse_qkv_projections(request, atol=atol, rtol=rtol) + + class TestSD35TransformerCompile(SD35TransformerTesterConfig, TorchCompileTesterMixin): pass diff --git a/tests/pipelines/aura_flow/test_pipeline_aura_flow.py b/tests/pipelines/aura_flow/test_pipeline_aura_flow.py index a45f65ebb89b..be54cadc29c7 100644 --- a/tests/pipelines/aura_flow/test_pipeline_aura_flow.py +++ b/tests/pipelines/aura_flow/test_pipeline_aura_flow.py @@ -3,15 +3,12 @@ from diffusers import AuraFlowPipeline, AuraFlowTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler -from ...testing_utils import assert_tensors_close from ..testing_utils import ( BasePipelineTesterConfig, LoraMemoryTesterMixin, LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -83,52 +80,6 @@ def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff= # AuraFlow pads the prompt embeddings to a common length, so batched and single runs diverge slightly more. super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) - def test_fused_qkv_projections(self): - # Run on CPU to keep the device-dependent `torch.Generator` deterministic. - pipe = self.get_pipeline() - - image = self.run_pipe(pipe) - original_image_slice = image[0, -1, -3:, -3:] - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - image = self.run_pipe(pipe) - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - image = self.run_pipe(pipe) - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - class TestAuraFlowPipelineMemory(AuraFlowPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the AuraFlow pipeline.""" diff --git a/tests/pipelines/chroma/test_pipeline_chroma.py b/tests/pipelines/chroma/test_pipeline_chroma.py index ecaf056bf716..e59b6b5f194f 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma.py +++ b/tests/pipelines/chroma/test_pipeline_chroma.py @@ -3,13 +3,12 @@ from diffusers import AutoencoderKL, ChromaPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import torch_device from ..flux.testing_utils import FluxIPAdapterTesterMixin from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -105,52 +104,6 @@ def test_chroma_different_prompts(self): # For some reasons, they don't show large differences assert max_diff > 1e-6, "Outputs should be different for different prompts." - def test_fused_qkv_projections(self): - # Run on CPU to keep the seeded generator deterministic across the three forward passes. - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - original_image_slice = image[0, -1, -3:, -3:] - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - 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() - image = pipe(**inputs).images - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - def test_chroma_image_output_shape(self): pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py index 224d2bd6c98f..76d62181f5b4 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py +++ b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py @@ -5,13 +5,12 @@ from diffusers import AutoencoderKL, ChromaImg2ImgPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler -from ...testing_utils import assert_tensors_close, floats_tensor, torch_device +from ...testing_utils import floats_tensor, torch_device from ..flux.testing_utils import FluxIPAdapterTesterMixin from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -110,52 +109,6 @@ def test_chroma_different_prompts(self): # For some reasons, they don't show large differences assert max_diff > 1e-6, "Outputs should be different for different prompts." - def test_fused_qkv_projections(self): - # Run on CPU to keep the seeded generator deterministic across the three forward passes. - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - original_image_slice = image[0, -1, -3:, -3:] - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - 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() - image = pipe(**inputs).images - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - def test_chroma_image_output_shape(self): pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py index 82db4dba98a6..4838fa464ad3 100644 --- a/tests/pipelines/cogvideo/test_cogvideox.py +++ b/tests/pipelines/cogvideo/test_cogvideox.py @@ -37,8 +37,6 @@ MemoryTesterMixin, PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -173,52 +171,6 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): "VAE tiling should not affect the inference results" ) - def test_fused_qkv_projections(self): - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames # [B, F, C, H, W] - original_image_slice = frames[0, -2:, -1, -3:, -3:] - - pipe.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_fused = frames[0, -2:, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_disabled = frames[0, -2:, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - class TestCogVideoXPipelineMemory(CogVideoXPipelineTesterConfig, MemoryTesterMixin): pass diff --git a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py index 9aa8d1e34c7e..a71a583552fe 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_fun_control.py +++ b/tests/pipelines/cogvideo/test_cogvideox_fun_control.py @@ -24,8 +24,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -166,52 +164,6 @@ def test_vae_tiling(self, expected_diff_max: float = 0.5): "VAE tiling should not affect the inference results" ) - def test_fused_qkv_projections(self): - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames # [B, F, C, H, W] - original_image_slice = frames[0, -2:, -1, -3:, -3:] - - pipe.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_fused = frames[0, -2:, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_disabled = frames[0, -2:, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - class TestCogVideoXFunControlPipelineMemory(CogVideoXFunControlPipelineTesterConfig, MemoryTesterMixin): pass diff --git a/tests/pipelines/cogvideo/test_cogvideox_image2video.py b/tests/pipelines/cogvideo/test_cogvideox_image2video.py index b4b9f20c5421..815819d90884 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_image2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_image2video.py @@ -34,8 +34,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -204,52 +202,6 @@ def test_vae_tiling(self, expected_diff_max: float = 0.3): "VAE tiling should not affect the inference results" ) - def test_fused_qkv_projections(self): - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames # [B, F, C, H, W] - original_image_slice = frames[0, -2:, -1, -3:, -3:] - - pipe.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_fused = frames[0, -2:, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_disabled = frames[0, -2:, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - class TestCogVideoXImageToVideoPipelineMemory(CogVideoXImageToVideoPipelineTesterConfig, MemoryTesterMixin): pass diff --git a/tests/pipelines/cogvideo/test_cogvideox_video2video.py b/tests/pipelines/cogvideo/test_cogvideox_video2video.py index 9b6cb76e4e6d..b86f4cd7e1bf 100644 --- a/tests/pipelines/cogvideo/test_cogvideox_video2video.py +++ b/tests/pipelines/cogvideo/test_cogvideox_video2video.py @@ -24,8 +24,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -171,52 +169,6 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): "VAE tiling should not affect the inference results" ) - def test_fused_qkv_projections(self): - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames # [B, F, C, H, W] - original_image_slice = frames[0, -2:, -1, -3:, -3:] - - pipe.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_fused = frames[0, -2:, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - frames = pipe(**inputs).frames - image_slice_disabled = frames[0, -2:, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - class TestCogVideoXVideoToVideoPipelineMemory(CogVideoXVideoToVideoPipelineTesterConfig, MemoryTesterMixin): pass diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py index 165e7bda2dc9..a4e7ba2e9a5e 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py @@ -10,12 +10,11 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -158,50 +157,6 @@ def test_flux_controlnet_different_prompts(self): assert max_diff > 1e-6, "Outputs should be different for different prompts." - def test_fused_qkv_projections(self): - # Run on CPU to keep the seeded generator deterministic across the three forward passes. - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - original_image_slice = image[0, -1, -3:, -3:] - - pipe.transformer.fuse_qkv_projections() - 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() - image = pipe(**inputs).images - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - 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.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/flux/test_pipeline_flux_control.py b/tests/pipelines/flux/test_pipeline_flux_control.py index 70f558b16cb2..1925628c6764 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control.py +++ b/tests/pipelines/flux/test_pipeline_flux_control.py @@ -33,7 +33,6 @@ LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -144,45 +143,6 @@ def test_flux_different_prompts(self): # For some reasons, they don't show large differences assert max_diff > 1e-6, "Outputs should be different for different prompts." - def test_fused_qkv_projections(self): - pipe = self.get_pipeline().to(torch_device) - - image_slice = self.run_pipe(pipe)[0, -1, -3:, -3:] - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - 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." - ) - - image_slice_fused = self.run_pipe(pipe)[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - image_slice_disabled = self.run_pipe(pipe)[0, -1, -3:, -3:] - - assert_tensors_close( - image_slice_fused, - image_slice, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - 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_tensors_close( - image_slice_disabled, - 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.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py index d4e7018bc8a6..d621db62c762 100644 --- a/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py +++ b/tests/pipelines/flux/test_pipeline_flux_control_inpaint.py @@ -9,12 +9,11 @@ FluxTransformer2DModel, ) -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -116,49 +115,6 @@ def get_dummy_inputs(self): class TestFluxControlInpaintPipeline(FluxControlInpaintPipelineTesterConfig, PipelineTesterMixin): - def test_fused_qkv_projections(self): - # Run on CPU to ensure determinism for the device-dependent torch.Generator. - pipe = self.get_pipeline() - - 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() - 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() - image_slice_fused = pipe(**inputs).images[0, -3:, -3:, -1] - - pipe.transformer.unfuse_qkv_projections() - 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_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_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.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/flux2/test_pipeline_flux2.py b/tests/pipelines/flux2/test_pipeline_flux2.py index ac72d843cd05..aba18b2aaf89 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2.py +++ b/tests/pipelines/flux2/test_pipeline_flux2.py @@ -8,14 +8,13 @@ Flux2Transformer2DModel, ) -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BasePipelineTesterConfig, LoraMemoryTesterMixin, LoraTesterMixin, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -125,51 +124,6 @@ def get_dummy_inputs(self): class TestFlux2Pipeline(Flux2PipelineTesterConfig, PipelineTesterMixin): - def test_fused_qkv_projections(self): - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - original_image_slice = image[0, -1, -3:, -3:] - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - 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() - image = pipe(**inputs).images - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - 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() diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein.py b/tests/pipelines/flux2/test_pipeline_flux2_klein.py index 0d7139b21e16..0fa833aca0bc 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein.py @@ -25,7 +25,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -113,49 +112,6 @@ def get_dummy_inputs(self): class TestFlux2KleinPipeline(Flux2KleinPipelineTesterConfig, PipelineTesterMixin): - def test_fused_qkv_projections(self): - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - original_image_slice = image[0, -1, -3:, -3:] - - pipe.transformer.fuse_qkv_projections() - 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() - image = pipe(**inputs).images - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - def test_image_output_shape(self): pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py b/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py index 29113510c1c8..a54f9dc568dc 100644 --- a/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py +++ b/tests/pipelines/flux2/test_pipeline_flux2_klein_kv.py @@ -10,12 +10,11 @@ Flux2Transformer2DModel, ) -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fused_layers_exist, ) @@ -100,50 +99,6 @@ def get_dummy_inputs(self): class TestFlux2KleinKVPipeline(Flux2KleinKVPipelineTesterConfig, PipelineTesterMixin): - def test_fused_qkv_projections(self): - # Run on CPU to keep the slice comparisons deterministic. - pipe = self.get_pipeline() - - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - original_image_slice = image[0, -1, -3:, -3:] - - pipe.transformer.fuse_qkv_projections() - 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() - image = pipe(**inputs).images - image_slice_fused = image[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs() - image = pipe(**inputs).images - image_slice_disabled = image[0, -1, -3:, -3:] - - assert_tensors_close( - original_image_slice, - image_slice_fused, - atol=1e-3, - rtol=1e-3, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_slice_fused, - image_slice_disabled, - atol=1e-3, - rtol=1e-3, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - original_image_slice, - image_slice_disabled, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - def test_image_output_shape(self): pipe = self.get_pipeline().to(torch_device) inputs = self.get_dummy_inputs() diff --git a/tests/pipelines/hunyuandit/test_hunyuan_dit.py b/tests/pipelines/hunyuandit/test_hunyuan_dit.py index 80cd76a2bcf8..f518e9802e2c 100644 --- a/tests/pipelines/hunyuandit/test_hunyuan_dit.py +++ b/tests/pipelines/hunyuandit/test_hunyuan_dit.py @@ -36,8 +36,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -149,50 +147,6 @@ def test_feed_forward_chunking(self): image_chunking, image_no_chunking, atol=1e-4, msg="Feed forward chunking should not affect the outputs." ) - def test_fused_qkv_projections(self): - # Run on CPU to ensure determinism for the device-dependent torch.Generator. - pipe = self.get_pipeline() - - original_image = pipe(**self.get_dummy_inputs(), return_dict=False)[0] - - pipe.transformer.fuse_qkv_projections() - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - image_fused = pipe(**self.get_dummy_inputs(), return_dict=False)[0] - - pipe.transformer.unfuse_qkv_projections() - image_disabled = pipe(**self.get_dummy_inputs(), return_dict=False)[0] - - assert_tensors_close( - image_fused, - original_image, - atol=1e-2, - rtol=1e-2, - msg="Fusion of QKV projections shouldn't affect the outputs.", - ) - assert_tensors_close( - image_disabled, - image_fused, - atol=1e-2, - rtol=1e-2, - msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", - ) - assert_tensors_close( - image_disabled, - original_image, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - @pytest.mark.skip( "Test not supported as `encode_prompt` is called two times separately which deivates from about 99% of the pipelines we have." ) diff --git a/tests/pipelines/pag/test_pag_hunyuan_dit.py b/tests/pipelines/pag/test_pag_hunyuan_dit.py index 41d2f68aa5c7..7a600bedac75 100644 --- a/tests/pipelines/pag/test_pag_hunyuan_dit.py +++ b/tests/pipelines/pag/test_pag_hunyuan_dit.py @@ -128,36 +128,6 @@ def test_feed_forward_chunking(self): output_chunking, output_no_chunking, atol=1e-4, msg="Forward chunking changed the output." ) - def test_fused_qkv_projections(self): - # Run on CPU to keep the device-dependent `torch.Generator` deterministic. - pipe = self.get_pipeline() - - original_output = self.run_pipe(pipe) - - pipe.transformer.fuse_qkv_projections() - output_fused = self.run_pipe(pipe) - - pipe.transformer.unfuse_qkv_projections() - output_disabled = self.run_pipe(pipe) - - assert_tensors_close( - output_fused, original_output, atol=1e-2, rtol=1e-2, msg="Fusion of QKV projections changed the outputs." - ) - assert_tensors_close( - output_disabled, - output_fused, - atol=1e-2, - rtol=1e-2, - msg="Outputs changed after the fused QKV projections were disabled.", - ) - assert_tensors_close( - output_disabled, - original_output, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - def test_pag_applied_layers(self): pipe = self.get_pipeline() diff --git a/tests/pipelines/pag/test_pag_sd3.py b/tests/pipelines/pag/test_pag_sd3.py index 6958985c954a..a140d689c5e6 100644 --- a/tests/pipelines/pag/test_pag_sd3.py +++ b/tests/pipelines/pag/test_pag_sd3.py @@ -16,12 +16,10 @@ StableDiffusion3Pipeline, ) -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import torch_device from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) from .testing_utils import PAGPipelineTesterMixin @@ -152,45 +150,6 @@ def test_stable_diffusion_3_different_negative_prompts(self): # Outputs should be different here assert (output_same_prompt - output_different_prompts).abs().max() > 1e-2 - def test_fused_qkv_projections(self): - # Run on CPU to keep the device-dependent `torch.Generator` deterministic. - pipe = self.get_pipeline() - - original_output = self.run_pipe(pipe) - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - output_fused = self.run_pipe(pipe) - - pipe.transformer.unfuse_qkv_projections() - output_disabled = self.run_pipe(pipe) - - assert_tensors_close( - output_fused, original_output, atol=1e-3, rtol=1e-3, msg="Fusion of QKV projections changed the outputs." - ) - assert_tensors_close( - output_disabled, - output_fused, - atol=1e-3, - rtol=1e-3, - msg="Outputs changed after the fused QKV projections were disabled.", - ) - assert_tensors_close( - output_disabled, - original_output, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - def test_pag_applied_layers(self): pipe = self.get_pipeline() diff --git a/tests/pipelines/pixart_sigma/test_pixart.py b/tests/pipelines/pixart_sigma/test_pixart.py index e65867dc11bd..0f9b33db7417 100644 --- a/tests/pipelines/pixart_sigma/test_pixart.py +++ b/tests/pipelines/pixart_sigma/test_pixart.py @@ -42,8 +42,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -197,45 +195,6 @@ def test_save_load_optional_components(self): def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=1e-3) - def test_fused_qkv_projections(self): - # Run on CPU to keep the device-dependent `torch.Generator` deterministic. - pipe = self.get_pipeline() - - original_output = self.run_pipe(pipe) - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - output_fused = self.run_pipe(pipe) - - pipe.transformer.unfuse_qkv_projections() - output_disabled = self.run_pipe(pipe) - - assert_tensors_close( - output_fused, original_output, atol=1e-3, rtol=1e-3, msg="Fusion of QKV projections changed the outputs." - ) - assert_tensors_close( - output_disabled, - output_fused, - atol=1e-3, - rtol=1e-3, - msg="Outputs changed after the fused QKV projections were disabled.", - ) - assert_tensors_close( - output_disabled, - original_output, - atol=1e-2, - rtol=1e-2, - msg="Original outputs should match when fused QKV projections are disabled.", - ) - class TestPixArtSigmaPipelineMemory(PixArtSigmaPipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the PixArt-sigma pipeline.""" diff --git a/tests/pipelines/stable_diffusion/test_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_stable_diffusion.py index f647f009ec58..e0142d612925 100644 --- a/tests/pipelines/stable_diffusion/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion/test_stable_diffusion.py @@ -529,31 +529,6 @@ def test_freeu_disabled(self, base_pipe_output): "Disabling of FreeU should lead to results similar to the default pipeline results." ) - def test_fused_qkv_projections(self, base_pipe_output): - # The unfused reference is the class-cached `base_pipe_output`, so the runs below reseed the global RNG to - # reproduce it and the only remaining difference comes from the projection fusion itself. - sd_pipe = self.get_pipeline().to(torch_device) - - original_image_slice = base_pipe_output[0, -1, -3:, -3:] - - sd_pipe.fuse_qkv_projections() - torch.manual_seed(0) - image_slice_fused = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - sd_pipe.unfuse_qkv_projections() - torch.manual_seed(0) - image_slice_disabled = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - assert torch.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." - ) - assert torch.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." - ) - assert torch.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." - ) - def test_pipeline_interrupt(self): sd_pipe = self.get_pipeline().to(torch_device) diff --git a/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py b/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py index 2acfd8520484..014165c8d601 100644 --- a/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py +++ b/tests/pipelines/stable_diffusion_2/test_stable_diffusion.py @@ -319,31 +319,6 @@ def test_freeu_disabled(self, base_pipe_output): "Disabling of FreeU should lead to results similar to the default pipeline results." ) - def test_fused_qkv_projections(self, base_pipe_output): - # The unfused reference is the class-cached `base_pipe_output`, so the runs below reseed the global RNG to - # reproduce it and the only remaining difference comes from the projection fusion itself. - sd_pipe = self.get_pipeline().to(torch_device) - - original_image_slice = base_pipe_output[0, -1, -3:, -3:] - - sd_pipe.fuse_qkv_projections() - torch.manual_seed(0) - image_slice_fused = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - sd_pipe.unfuse_qkv_projections() - torch.manual_seed(0) - image_slice_disabled = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - assert torch.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." - ) - assert torch.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." - ) - assert torch.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." - ) - class TestStableDiffusion2PipelineMemory(StableDiffusion2PipelineTesterConfig, MemoryTesterMixin): """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD2 pipeline.""" diff --git a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py index 4b8952188efa..833bd01b492d 100644 --- a/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py +++ b/tests/pipelines/stable_diffusion_3/test_pipeline_stable_diffusion_3.py @@ -26,8 +26,6 @@ BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -147,40 +145,6 @@ def test_inference(self): assert_tensors_close(generated_slice, expected_slice, atol=1e-3, msg="Output does not match expected slice.") - def test_fused_qkv_projections(self, base_pipe_output): - # The unfused reference is the class-cached `base_pipe_output`, so the runs below reseed the global RNG to - # reproduce it and the only remaining difference comes from the projection fusion itself. - pipe = self.get_pipeline().to(torch_device) - - original_image_slice = base_pipe_output[0, -1, -3:, -3:] - - # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added - # to the pipeline level. - pipe.transformer.fuse_qkv_projections() - assert check_qkv_fusion_processors_exist(pipe.transformer), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length( - pipe.transformer, pipe.transformer.original_attn_processors - ), "Something wrong with the attention processors concerning the fused QKV projections." - - torch.manual_seed(0) - image_slice_fused = pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - pipe.transformer.unfuse_qkv_projections() - torch.manual_seed(0) - image_slice_disabled = pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - assert torch.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." - ) - assert torch.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 torch.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." - ) - def test_skip_guidance_layers(self): pipe = self.get_pipeline().to(torch_device) diff --git a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py index b4dd632c0128..8728dd2a3bdc 100644 --- a/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py +++ b/tests/pipelines/stable_diffusion_xl/test_stable_diffusion_xl.py @@ -70,8 +70,6 @@ MemoryTesterMixin, PipelineTesterMixin, UNetLoraTesterMixin, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, ) @@ -346,43 +344,6 @@ def test_freeu_disabled(self, base_pipe_output): "Disabling of FreeU should lead to results similar to the default pipeline results." ) - def test_fused_qkv_projections(self, base_pipe_output): - # The unfused reference is the class-cached `base_pipe_output`, so the runs below reseed the global RNG to - # reproduce it and the only remaining difference comes from the projection fusion itself. - sd_pipe = self.get_pipeline().to(torch_device) - - original_image_slice = base_pipe_output[0, -1, -3:, -3:] - - sd_pipe.fuse_qkv_projections() - for component in sd_pipe.components.values(): - if ( - isinstance(component, torch.nn.Module) - and getattr(component, "original_attn_processors", None) is not None - ): - assert check_qkv_fusion_processors_exist(component), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length(component, component.original_attn_processors), ( - "Something wrong with the attention processors concerning the fused QKV projections." - ) - - torch.manual_seed(0) - image_slice_fused = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - sd_pipe.unfuse_qkv_projections() - torch.manual_seed(0) - image_slice_disabled = sd_pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] - - assert torch.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." - ) - assert torch.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." - ) - assert torch.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." - ) - def test_stable_diffusion_two_xl_mixture_of_denoiser_fast(self): components = self.get_dummy_components() pipe_1 = self.get_pipeline(**components).to(torch_device) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index 106ba55cf149..82f51595912d 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -35,7 +35,6 @@ from diffusers.hooks.taylorseer_cache import TaylorSeerCacheConfig from diffusers.image_processor import VaeImageProcessor from diffusers.loaders import FluxIPAdapterMixin, IPAdapterMixin -from diffusers.models.attention import AttentionModuleMixin from diffusers.models.attention_processor import AttnProcessor from diffusers.models.controlnets.controlnet_xs import UNetControlNetXSModel from diffusers.models.unets.unet_3d_condition import UNet3DConditionModel @@ -83,31 +82,6 @@ def check_same_shape(tensor_list): return all(shape == shapes[0] for shape in shapes[1:]) -def check_qkv_fusion_matches_attn_procs_length(model, original_attn_processors): - current_attn_processors = model.attn_processors - return len(current_attn_processors) == len(original_attn_processors) - - -def check_qkv_fusion_processors_exist(model): - current_attn_processors = model.attn_processors - proc_names = [v.__class__.__name__ for _, v in current_attn_processors.items()] - return all(p.startswith("Fused") for p in proc_names) - - -def check_qkv_fused_layers_exist(model, layer_names): - is_fused_submodules = [] - for submodule in model.modules(): - if not isinstance(submodule, AttentionModuleMixin) or not submodule._supports_qkv_fusion: - continue - is_fused_attribute_set = submodule.fused_projections - is_fused_layer = True - for layer in layer_names: - is_fused_layer = is_fused_layer and getattr(submodule, layer, None) is not None - is_fused = is_fused_attribute_set and is_fused_layer - is_fused_submodules.append(is_fused) - return all(is_fused_submodules) - - class SDFunctionTesterMixin: """ This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes. @@ -210,53 +184,6 @@ def test_freeu(self): f"Disabling of FreeU should lead to results similar to the default pipeline results but Max Abs Error={np.abs(output_no_freeu - output).max()}." ) - 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) - - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image = pipe(**inputs)[0] - original_image_slice = image[0, -3:, -3:, -1] - - pipe.fuse_qkv_projections() - for _, component in pipe.components.items(): - if ( - isinstance(component, nn.Module) - and hasattr(component, "original_attn_processors") - and component.original_attn_processors is not None - ): - assert check_qkv_fusion_processors_exist(component), ( - "Something wrong with the fused attention processors. Expected all the attention processors to be fused." - ) - assert check_qkv_fusion_matches_attn_procs_length(component, component.original_attn_processors), ( - "Something wrong with the attention processors concerning the fused QKV projections." - ) - - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image_fused = pipe(**inputs)[0] - image_slice_fused = image_fused[0, -3:, -3:, -1] - - pipe.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) - inputs["return_dict"] = False - image_disabled = pipe(**inputs)[0] - image_slice_disabled = image_disabled[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-2, rtol=1e-2), ( - "Fusion of QKV projections shouldn't affect the outputs." - ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "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." - ) - class IPAdapterTesterMixin: """ diff --git a/tests/pipelines/testing_utils/__init__.py b/tests/pipelines/testing_utils/__init__.py index 9b756ec64693..467065d52d7f 100644 --- a/tests/pipelines/testing_utils/__init__.py +++ b/tests/pipelines/testing_utils/__init__.py @@ -6,7 +6,7 @@ PyramidAttentionBroadcastTesterMixin, TaylorSeerCacheTesterMixin, ) -from .common import BasePipelineTesterConfig, PipelineTesterMixin +from .common import BasePipelineTesterConfig, PipelineTesterMixin, check_same_shape from .from_pipe import FromPipeTesterMixin from .ip_adapter import IPAdapterTesterMixin from .lora import LoraMemoryTesterMixin, LoraTesterMixin, UNetLoraTesterMixin @@ -16,12 +16,6 @@ MemoryTesterMixin, PipelineOffloadTesterMixin, ) -from .utils import ( - check_qkv_fused_layers_exist, - check_qkv_fusion_matches_attn_procs_length, - check_qkv_fusion_processors_exist, - check_same_shape, -) __all__ = [ @@ -42,8 +36,5 @@ "FirstBlockCacheTesterMixin", "TaylorSeerCacheTesterMixin", "MagCacheTesterMixin", - "check_qkv_fused_layers_exist", - "check_qkv_fusion_matches_attn_procs_length", - "check_qkv_fusion_processors_exist", "check_same_shape", ] diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 5cf57eee6ead..9f761a61d958 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -37,6 +37,11 @@ ) +def check_same_shape(tensor_list): + shapes = [tensor.shape for tensor in tensor_list] + return all(shape == shapes[0] for shape in shapes[1:]) + + def cast_module_to_dtype(module, dtype): """Cast `module` to `dtype` in place, keeping its `_keep_in_fp32_modules` submodules in float32. diff --git a/tests/pipelines/testing_utils/utils.py b/tests/pipelines/testing_utils/utils.py deleted file mode 100644 index a20c0ca9becb..000000000000 --- a/tests/pipelines/testing_utils/utils.py +++ /dev/null @@ -1,54 +0,0 @@ -# coding=utf-8 -# Copyright 2025 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. - -from diffusers.models.attention import AttentionModuleMixin - - -""" -TODO (related to methods like `check_qkv_fusion_matches_attn_procs_length()`): -After https://github.com/huggingface/diffusers/pull/14113 is merged, move those -checks out of pipeline-level testing and ensure that these are sufficiently -tested in model-level tests. -""" - - -def check_same_shape(tensor_list): - shapes = [tensor.shape for tensor in tensor_list] - return all(shape == shapes[0] for shape in shapes[1:]) - - -def check_qkv_fusion_matches_attn_procs_length(model, original_attn_processors): - current_attn_processors = model.attn_processors - return len(current_attn_processors) == len(original_attn_processors) - - -def check_qkv_fusion_processors_exist(model): - current_attn_processors = model.attn_processors - proc_names = [v.__class__.__name__ for _, v in current_attn_processors.items()] - return all(p.startswith("Fused") for p in proc_names) - - -def check_qkv_fused_layers_exist(model, layer_names): - is_fused_submodules = [] - for submodule in model.modules(): - if not isinstance(submodule, AttentionModuleMixin) or not submodule._supports_qkv_fusion: - continue - is_fused_attribute_set = submodule.fused_projections - is_fused_layer = True - for layer in layer_names: - is_fused_layer = is_fused_layer and getattr(submodule, layer, None) is not None - is_fused = is_fused_attribute_set and is_fused_layer - is_fused_submodules.append(is_fused) - return all(is_fused_submodules)