diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index 2489604274b4..c334feeefee9 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -2,10 +2,12 @@ import unittest import numpy as np +import pytest import torch from diffusers import DiffusionPipeline from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor +from diffusers.utils.testing_utils import torch_device class AttnAddedKVProcessorTests(unittest.TestCase): @@ -79,6 +81,11 @@ def test_only_cross_attention(self): class DeprecatedAttentionBlockTests(unittest.TestCase): + @pytest.mark.xfail( + condition=torch.device(torch_device).type == "cuda", + reason="Test currently fails on our GPU CI because of `disfile`.", + strict=True, + ) def test_conversion_when_using_device_map(self): pipe = DiffusionPipeline.from_pretrained( "hf-internal-testing/tiny-stable-diffusion-torch", safety_checker=None diff --git a/tests/models/transformers/test_models_transformer_mochi.py b/tests/models/transformers/test_models_transformer_mochi.py index fc1412c7cd31..d284ab942949 100644 --- a/tests/models/transformers/test_models_transformer_mochi.py +++ b/tests/models/transformers/test_models_transformer_mochi.py @@ -30,6 +30,8 @@ class MochiTransformerTests(ModelTesterMixin, unittest.TestCase): model_class = MochiTransformer3DModel main_input_name = "hidden_states" uses_custom_attn_processor = True + # Overriding it because of the transformer size. + model_split_percents = [0.7, 0.6, 0.6] @property def dummy_input(self): diff --git a/tests/models/transformers/test_models_transformer_sana.py b/tests/models/transformers/test_models_transformer_sana.py index 0222bef4c7c3..83db153dadea 100644 --- a/tests/models/transformers/test_models_transformer_sana.py +++ b/tests/models/transformers/test_models_transformer_sana.py @@ -14,6 +14,7 @@ import unittest +import pytest import torch from diffusers import SanaTransformer2DModel @@ -80,3 +81,27 @@ def prepare_init_args_and_inputs_for_common(self): def test_gradient_checkpointing_is_applied(self): expected_set = {"SanaTransformer2DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + @pytest.mark.xfail( + condition=torch.device(torch_device).type == "cuda", + reason="Test currently fails.", + strict=True, + ) + def test_cpu_offload(self): + return super().test_cpu_offload() + + @pytest.mark.xfail( + condition=torch.device(torch_device).type == "cuda", + reason="Test currently fails.", + strict=True, + ) + def test_disk_offload_with_safetensors(self): + return super().test_disk_offload_with_safetensors() + + @pytest.mark.xfail( + condition=torch.device(torch_device).type == "cuda", + reason="Test currently fails.", + strict=True, + ) + def test_disk_offload_without_safetensors(self): + return super().test_disk_offload_without_safetensors()