huggingface · sayakpaul · Jun 17, 2026 · Jul 3, 2026 · Jul 3, 2026 · Jul 3, 2026
diff --git a/tests/pipelines/flux/test_pipeline_flux.py b/tests/pipelines/flux/test_pipeline_flux.py
@@ -1,7 +1,7 @@
 import gc
-import unittest
 
 import numpy as np
+import pytest
 import torch
 from huggingface_hub import hf_hub_download
 from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel
@@ -23,45 +23,28 @@
     slow,
     torch_device,
 )
-from ..test_pipelines_common import (
+from ..testing_utils import (
+    BasePipelineTesterConfig,
     FasterCacheTesterMixin,
     FirstBlockCacheTesterMixin,
     FluxIPAdapterTesterMixin,
     MagCacheTesterMixin,
+    MemoryTesterMixin,
     PipelineTesterMixin,
     PyramidAttentionBroadcastTesterMixin,
     TaylorSeerCacheTesterMixin,
     check_qkv_fused_layers_exist,
 )
 
 
-class FluxPipelineFastTests(
-    PipelineTesterMixin,
-    FluxIPAdapterTesterMixin,
-    PyramidAttentionBroadcastTesterMixin,
-    FasterCacheTesterMixin,
-    FirstBlockCacheTesterMixin,
-    TaylorSeerCacheTesterMixin,
-    MagCacheTesterMixin,
-    unittest.TestCase,
-):
+class FluxPipelineTesterConfig(BasePipelineTesterConfig):
     pipeline_class = FluxPipeline
     params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"])
     batch_params = frozenset(["prompt"])
 
-    # there is no xformers processor for Flux
-    test_xformers_attention = False
     test_layerwise_casting = True
     test_group_offloading = True
 
-    faster_cache_config = FasterCacheConfig(
-        spatial_attention_block_skip_range=2,
-        spatial_attention_timestep_skip_range=(-1, 901),
-        unconditional_batch_skip_range=2,
-        attention_weight_callback=lambda _: 0.5,
-        is_guidance_distilled=True,
-    )
-
     def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1):
         torch.manual_seed(0)
         transformer = FluxTransformer2DModel(
@@ -146,6 +129,8 @@ def get_dummy_inputs(self, device, seed=0):
         }
         return inputs
 
+
+class TestFluxPipeline(FluxPipelineTesterConfig, PipelineTesterMixin):
     def test_flux_different_prompts(self):
         pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
 
@@ -160,7 +145,7 @@ def test_flux_different_prompts(self):
 
         # Outputs should be different here
         # For some reasons, they don't show large differences
-        self.assertGreater(max_diff, 1e-6, "Outputs should be different for different prompts.")
+        assert max_diff > 1e-6, "Outputs should be different for different prompts."
 
     def test_fused_qkv_projections(self):
         device = "cpu"  # ensure determinism for the device-dependent torch.Generator
@@ -176,9 +161,8 @@ def test_fused_qkv_projections(self):
         # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added
         # to the pipeline level.
         pipe.transformer.fuse_qkv_projections()
-        self.assertTrue(
-            check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]),
-            ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."),
+        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(device)
@@ -190,17 +174,14 @@ def test_fused_qkv_projections(self):
         image = pipe(**inputs).images
         image_slice_disabled = image[0, -3:, -3:, -1]
 
-        self.assertTrue(
-            np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3),
-            ("Fusion of QKV projections shouldn't affect the outputs."),
+        assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), (
+            "Fusion of QKV projections shouldn't affect the outputs."
         )
-        self.assertTrue(
-            np.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 np.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."
         )
-        self.assertTrue(
-            np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2),
-            ("Original outputs should match 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."
         )
 
     def test_flux_image_output_shape(self):
@@ -215,10 +196,8 @@ def test_flux_image_output_shape(self):
             inputs.update({"height": height, "width": width})
             image = pipe(**inputs).images[0]
             output_height, output_width, _ = image.shape
-            self.assertEqual(
-                (output_height, output_width),
-                (expected_height, expected_width),
-                f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}",
+            assert (output_height, output_width) == (expected_height, expected_width), (
+                f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}"
             )
 
     def test_flux_true_cfg(self):
@@ -230,8 +209,8 @@ def test_flux_true_cfg(self):
         inputs["negative_prompt"] = "bad quality"
         inputs["true_cfg_scale"] = 2.0
         true_cfg_out = pipe(**inputs, generator=torch.manual_seed(0)).images[0]
-        self.assertFalse(
-            np.allclose(no_true_cfg_out, true_cfg_out), "Outputs should be different when true_cfg_scale is set."
+        assert not np.allclose(no_true_cfg_out, true_cfg_out), (
+            "Outputs should be different when true_cfg_scale is set."
         )
 
     def test_flux_negative_embeds_shape_check(self):
@@ -248,25 +227,60 @@ def test_flux_negative_embeds_shape_check(self):
             "output_type": "latent",
         }
 
-        with self.assertRaisesRegex(ValueError, "must have the same shape when passed directly"):
+        with pytest.raises(ValueError, match="must have the same shape when passed directly"):
             pipe(**base_inputs, true_cfg_scale=2.0, generator=torch.manual_seed(0))
 
         pipe(**base_inputs, true_cfg_scale=1.0, generator=torch.manual_seed(0))
 
 
+class TestFluxPipelineIPAdapter(FluxPipelineTesterConfig, FluxIPAdapterTesterMixin):
+    """IP-Adapter tests for the Flux pipeline."""
+
+
+class TestFluxPipelineMemory(FluxPipelineTesterConfig, MemoryTesterMixin):
+    """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux pipeline."""
+
+
+class TestFluxPipelinePyramidAttentionBroadcast(FluxPipelineTesterConfig, PyramidAttentionBroadcastTesterMixin):
+    """Pyramid Attention Broadcast cache tests for the Flux pipeline."""
+
+
+class TestFluxPipelineFasterCache(FluxPipelineTesterConfig, FasterCacheTesterMixin):
+    """FasterCache tests for the Flux pipeline."""
+
+    # Flux is guidance-distilled, so the FasterCache tester must skip the low/high-frequency-delta state checks.
+    faster_cache_config = FasterCacheConfig(
+        spatial_attention_block_skip_range=2,
+        spatial_attention_timestep_skip_range=(-1, 901),
+        unconditional_batch_skip_range=2,
+        attention_weight_callback=lambda _: 0.5,
+        is_guidance_distilled=True,
+    )
+
+
+class TestFluxPipelineFirstBlockCache(FluxPipelineTesterConfig, FirstBlockCacheTesterMixin):
+    """First Block Cache tests for the Flux pipeline."""
+
+
+class TestFluxPipelineTaylorSeerCache(FluxPipelineTesterConfig, TaylorSeerCacheTesterMixin):
+    """TaylorSeer cache tests for the Flux pipeline."""
+
+
+class TestFluxPipelineMagCache(FluxPipelineTesterConfig, MagCacheTesterMixin):
+    """MagCache tests for the Flux pipeline."""
+
+
 @nightly
 @require_big_accelerator
-class FluxPipelineSlowTests(unittest.TestCase):
+class TestFluxPipelineSlow:
     pipeline_class = FluxPipeline
     repo_id = "black-forest-labs/FLUX.1-schnell"
 
-    def setUp(self):
-        super().setUp()
+    @pytest.fixture(autouse=True)
+    def cleanup(self):
         gc.collect()
         backend_empty_cache(torch_device)
-
-    def tearDown(self):
-        super().tearDown()
+        yield
         gc.collect()
         backend_empty_cache(torch_device)
 
@@ -312,27 +326,23 @@ def test_flux_inference(self):
         # fmt: on
 
         max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten())
-        self.assertLess(
-            max_diff, 1e-4, f"Image slice is different from expected slice: {image_slice} != {expected_slice}"
-        )
+        assert max_diff < 1e-4, f"Image slice is different from expected slice: {image_slice} != {expected_slice}"
 
 
 @slow
 @require_big_accelerator
-class FluxIPAdapterPipelineSlowTests(unittest.TestCase):
+class TestFluxIPAdapterPipelineSlow:
     pipeline_class = FluxPipeline
     repo_id = "black-forest-labs/FLUX.1-dev"
     image_encoder_pretrained_model_name_or_path = "openai/clip-vit-large-patch14"
     weight_name = "ip_adapter.safetensors"
     ip_adapter_repo_id = "XLabs-AI/flux-ip-adapter"
 
-    def setUp(self):
-        super().setUp()
+    @pytest.fixture(autouse=True)
+    def cleanup(self):
         gc.collect()
         backend_empty_cache(torch_device)
-
-    def tearDown(self):
-        super().tearDown()
+        yield
         gc.collect()
         backend_empty_cache(torch_device)
 
@@ -392,6 +402,4 @@ def test_flux_ip_adapter_inference(self):
         # fmt: on
 
         max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten())
-        self.assertLess(
-            max_diff, 1e-4, f"Image slice is different from expected slice: {image_slice} != {expected_slice}"
-        )
+        assert max_diff < 1e-4, f"Image slice is different from expected slice: {image_slice} != {expected_slice}"