tests: fix qwen tests from getting oom'd in our CI. (#14474)

Author: sayakpaul

tests/pipelines/qwenimage/test_qwenimage_edit.py

@@ -24,7 +24,7 @@
     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 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
             "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 @@ def test_true_cfg_without_negative_prompt_embeds_mask(self):
 
 
 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()

tests/pipelines/qwenimage/test_qwenimage_edit_plus.py

@@ -24,7 +24,7 @@
     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 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2):
             "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 @@ def test_true_cfg_without_negative_prompt_embeds_mask(self):
 
 
 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()

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(