From bca3e27c96b942db49ccab8ddf824e7a54d43ed1 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Sat, 29 Nov 2025 10:42:38 +0530 Subject: [PATCH] up --- .../test_models_transformer_z_image.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 4fcb7728df..309a646de5 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -21,7 +21,7 @@ import torch from diffusers import ZImageTransformer2DModel -from ...testing_utils import torch_device +from ...testing_utils import is_flaky, torch_device from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin @@ -139,6 +139,22 @@ class ZImageTransformerTests(ModelTesterMixin, unittest.TestCase): def test_group_offloading_with_disk(self): super().test_group_offloading_with_disk() + @is_flaky(max_attempts=10) + def test_auto_model(self): + super().test_auto_model() + + @is_flaky(max_attempts=10) + def test_determinism(self): + super().test_determinism() + + @is_flaky(max_attempts=10) + def test_from_save_pretrained(self): + super().test_from_save_pretrained() + + @is_flaky(max_attempts=10) + def test_from_save_pretrained_variant(self): + super().test_from_save_pretrained_variant() + class Flux2TransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): model_class = ZImageTransformer2DModel