diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit.py b/tests/pipelines/qwenimage/test_qwenimage_edit.py index 8d7f8bc914..8f767aff8d 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit.py @@ -24,7 +24,7 @@ from diffusers import ( QwenImageTransformer2DModel, ) -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import assert_tensors_close, require_accelerator_memory, torch_device from ..testing_utils import ( BasePipelineTesterConfig, MemoryTesterMixin, @@ -170,6 +170,15 @@ class TestQwenImageEditPipeline(QwenImageEditPipelineTesterConfig, PipelineTeste "VAE tiling should not affect the inference results." ) + # The condition image is always resized to ~1 megapixel before it is encoded, so the VAE attention runs over + # 65k tokens no matter how small the dummy inputs are. On GPUs without bf16 SDPA support that attention falls + # back to the math backend, which materializes the full attention matrix (~16GB). + @require_accelerator_memory(24) + @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU") + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=str) + def test_half_precision_inference_no_nan(self, dtype): + super().test_half_precision_inference_no_nan(dtype) + @pytest.mark.xfail(condition=True, reason="Preconfigured embeddings need to be revisited.", strict=True) def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol, rtol) @@ -200,4 +209,7 @@ class TestQwenImageEditPipeline(QwenImageEditPipelineTesterConfig, PipelineTeste class TestQwenImageEditPipelineMemory(QwenImageEditPipelineTesterConfig, MemoryTesterMixin): - pass + # Runs the pipeline in bf16, see the note on `test_half_precision_inference_no_nan` above. + @require_accelerator_memory(24) + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() diff --git a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py index 2eef47169f..8977eb6d82 100644 --- a/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py +++ b/tests/pipelines/qwenimage/test_qwenimage_edit_plus.py @@ -24,7 +24,7 @@ from diffusers import ( QwenImageTransformer2DModel, ) -from ...testing_utils import assert_tensors_close, torch_device +from ...testing_utils import assert_tensors_close, require_accelerator_memory, torch_device from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin @@ -164,6 +164,15 @@ class TestQwenImageEditPlusPipeline(QwenImageEditPlusPipelineTesterConfig, Pipel "VAE tiling should not affect the inference results." ) + # The condition image is always resized to ~1 megapixel before it is encoded, so the VAE attention runs over + # 65k tokens no matter how small the dummy inputs are. On GPUs without bf16 SDPA support that attention falls + # back to the math backend, which materializes the full attention matrix (~16GB). + @require_accelerator_memory(24) + @pytest.mark.skipif(torch_device not in ["cuda", "xpu"], reason="half-precision inference requires CUDA or XPU") + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=str) + def test_half_precision_inference_no_nan(self, dtype): + super().test_half_precision_inference_no_nan(dtype) + @pytest.mark.xfail(condition=True, reason="Preconfigured embeddings need to be revisited.", strict=True) def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4): super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol, rtol) @@ -206,4 +215,7 @@ class TestQwenImageEditPlusPipeline(QwenImageEditPlusPipelineTesterConfig, Pipel class TestQwenImageEditPlusPipelineMemory(QwenImageEditPlusPipelineTesterConfig, MemoryTesterMixin): - pass + # Runs the pipeline in bf16, see the note on `test_half_precision_inference_no_nan` above. + @require_accelerator_memory(24) + def test_layerwise_casting_inference(self): + super().test_layerwise_casting_inference() diff --git a/tests/testing_utils.py b/tests/testing_utils.py index be4e07d2ea..a7bd1ac08a 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -660,6 +660,28 @@ def require_big_accelerator(test_case): )(test_case) +def require_accelerator_memory(min_memory_gb: int): + """ + Decorator marking a test that needs at least `min_memory_gb` GB of memory on the accelerator it runs on. Tests + running on CPU are not affected. + """ + + def decorator(test_case): + if torch_device == "cuda" and torch.cuda.is_available(): + total_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3) + elif torch_device == "xpu" and torch.xpu.is_available(): + total_memory = torch.xpu.get_device_properties(0).total_memory / (1024**3) + else: + return test_case + + return pytest.mark.skipif( + total_memory < min_memory_gb, + reason=f"test requires an accelerator with at least {min_memory_gb} GB memory", + )(test_case) + + return decorator + + def require_torch_accelerator_with_training(test_case): """Decorator marking a test that requires an accelerator with support for training.""" return pytest.mark.skipif(