Merge branch 'main' into refactor-caching-tests

src/diffusers/quantizers/quantization_config.py

@@ -36,7 +36,7 @@ from typing import Any, Callable
 
 from packaging import version
 
-from ..utils import deprecate, is_torch_available, is_torchao_available, is_torchao_version, logging
+from ..utils import is_torch_available, is_torchao_available, is_torchao_version, logging
 
 
 if is_torch_available():
@@ -844,8 +844,6 @@ class QuantoConfig(QuantizationConfigMixin):
         modules_to_not_convert: list[str] | None = None,
         **kwargs,
     ):
-        deprecation_message = "`QuantoConfig` is deprecated and will be removed in version 1.0.0."
-        deprecate("QuantoConfig", "1.0.0", deprecation_message)
         self.quant_method = QuantizationMethod.QUANTO
         self.weights_dtype = weights_dtype
         self.modules_to_not_convert = modules_to_not_convert

src/diffusers/quantizers/quanto/quanto_quantizer.py

@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, Any
 from diffusers.utils.import_utils import is_optimum_quanto_version
 
 from ...utils import (
-    deprecate,
     get_module_from_name,
     is_accelerate_available,
     is_accelerate_version,
@@ -43,9 +42,6 @@ class QuantoQuantizer(DiffusersQuantizer):
         super().__init__(quantization_config, **kwargs)
 
     def validate_environment(self, *args, **kwargs):
-        deprecation_message = "The Quanto quantizer is deprecated and will be removed in version 1.0.0."
-        deprecate("QuantoQuantizer", "1.0.0", deprecation_message)
-
         if not is_optimum_quanto_available():
             raise ImportError(
                 "Loading an optimum-quanto quantized model requires optimum-quanto library (`pip install optimum-quanto`)"

tests/hooks/test_mag_cache.py

@@ -1,244 +0,0 @@
-# Copyright 2025 HuggingFace Inc.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import unittest
-
-import numpy as np
-import torch
-
-from diffusers import MagCacheConfig, apply_mag_cache
-from diffusers.hooks._helpers import TransformerBlockMetadata, TransformerBlockRegistry
-from diffusers.models import ModelMixin
-from diffusers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-
-
-class DummyBlock(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
-        # Output is double input
-        # This ensures Residual = 2*Input - Input = Input
-        return hidden_states * 2.0
-
-
-class DummyTransformer(ModelMixin):
-    def __init__(self):
-        super().__init__()
-        self.transformer_blocks = torch.nn.ModuleList([DummyBlock(), DummyBlock()])
-
-    def forward(self, hidden_states, encoder_hidden_states=None):
-        for block in self.transformer_blocks:
-            hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states)
-        return hidden_states
-
-
-class TupleOutputBlock(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
-        # Returns a tuple
-        return hidden_states * 2.0, encoder_hidden_states
-
-
-class TupleTransformer(ModelMixin):
-    def __init__(self):
-        super().__init__()
-        self.transformer_blocks = torch.nn.ModuleList([TupleOutputBlock()])
-
-    def forward(self, hidden_states, encoder_hidden_states=None):
-        for block in self.transformer_blocks:
-            # Emulate Flux-like behavior
-            output = block(hidden_states, encoder_hidden_states=encoder_hidden_states)
-            hidden_states = output[0]
-            encoder_hidden_states = output[1]
-        return hidden_states, encoder_hidden_states
-
-
-class MagCacheTests(unittest.TestCase):
-    def setUp(self):
-        # Register standard dummy block
-        TransformerBlockRegistry.register(
-            DummyBlock,
-            TransformerBlockMetadata(return_hidden_states_index=None, return_encoder_hidden_states_index=None),
-        )
-        # Register tuple block (Flux style)
-        TransformerBlockRegistry.register(
-            TupleOutputBlock,
-            TransformerBlockMetadata(return_hidden_states_index=0, return_encoder_hidden_states_index=1),
-        )
-
-    def _set_context(self, model, context_name):
-        """Helper to set context on all hooks in the model."""
-        for module in model.modules():
-            if hasattr(module, "_diffusers_hook"):
-                module._diffusers_hook._set_context(context_name)
-
-    def _get_calibration_data(self, model):
-        for module in model.modules():
-            if hasattr(module, "_diffusers_hook"):
-                hook = module._diffusers_hook.get_hook("mag_cache_block_hook")
-                if hook:
-                    return hook.state_manager.get_state().calibration_ratios
-        return []
-
-    def test_mag_cache_validation(self):
-        """Test that missing mag_ratios raises ValueError."""
-        with self.assertRaises(ValueError):
-            MagCacheConfig(num_inference_steps=10, calibrate=False)
-
-    def test_mag_cache_skipping_logic(self):
-        """
-        Tests that MagCache correctly calculates residuals and skips blocks when conditions are met.
-        """
-        model = DummyTransformer()
-
-        # Dummy ratios: [1.0, 1.0] implies 0 accumulated error if we skip
-        ratios = np.array([1.0, 1.0])
-
-        config = MagCacheConfig(
-            threshold=100.0,
-            num_inference_steps=2,
-            retention_ratio=0.0,  # Enable immediate skipping
-            max_skip_steps=5,
-            mag_ratios=ratios,
-        )
-
-        apply_mag_cache(model, config)
-        self._set_context(model, "test_context")
-
-        # Step 0: Input 10.0 -> Output 40.0 (2 blocks * 2x each)
-        # HeadInput=10. Output=40. Residual=30.
-        input_t0 = torch.tensor([[[10.0]]])
-        output_t0 = model(input_t0)
-        self.assertTrue(torch.allclose(output_t0, torch.tensor([[[40.0]]])), "Step 0 failed")
-
-        # Step 1: Input 11.0.
-        # If Skipped: Output = Input(11) + Residual(30) = 41.0
-        # If Computed: Output = 11 * 4 = 44.0
-        input_t1 = torch.tensor([[[11.0]]])
-        output_t1 = model(input_t1)
-
-        self.assertTrue(
-            torch.allclose(output_t1, torch.tensor([[[41.0]]])), f"Expected Skip (41.0), got {output_t1.item()}"
-        )
-
-    def test_mag_cache_retention(self):
-        """Test that retention_ratio prevents skipping even if error is low."""
-        model = DummyTransformer()
-        # Ratios that imply 0 error, so it *would* skip if retention allowed it
-        ratios = np.array([1.0, 1.0])
-
-        config = MagCacheConfig(
-            threshold=100.0,
-            num_inference_steps=2,
-            retention_ratio=1.0,  # Force retention for ALL steps
-            mag_ratios=ratios,
-        )
-
-        apply_mag_cache(model, config)
-        self._set_context(model, "test_context")
-
-        # Step 0
-        model(torch.tensor([[[10.0]]]))
-
-        # Step 1: Should COMPUTE (44.0) not SKIP (41.0) because of retention
-        input_t1 = torch.tensor([[[11.0]]])
-        output_t1 = model(input_t1)
-
-        self.assertTrue(
-            torch.allclose(output_t1, torch.tensor([[[44.0]]])),
-            f"Expected Compute (44.0) due to retention, got {output_t1.item()}",
-        )
-
-    def test_mag_cache_tuple_outputs(self):
-        """Test compatibility with models returning (hidden, encoder_hidden) like Flux."""
-        model = TupleTransformer()
-        ratios = np.array([1.0, 1.0])
-
-        config = MagCacheConfig(threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=ratios)
-
-        apply_mag_cache(model, config)
-        self._set_context(model, "test_context")
-
-        # Step 0: Compute. Input 10.0 -> Output 20.0 (1 block * 2x)
-        # Residual = 10.0
-        input_t0 = torch.tensor([[[10.0]]])
-        enc_t0 = torch.tensor([[[1.0]]])
-        out_0, _ = model(input_t0, encoder_hidden_states=enc_t0)
-        self.assertTrue(torch.allclose(out_0, torch.tensor([[[20.0]]])))
-
-        # Step 1: Skip. Input 11.0.
-        # Skipped Output = 11 + 10 = 21.0
-        input_t1 = torch.tensor([[[11.0]]])
-        out_1, _ = model(input_t1, encoder_hidden_states=enc_t0)
-
-        self.assertTrue(
-            torch.allclose(out_1, torch.tensor([[[21.0]]])), f"Tuple skip failed. Expected 21.0, got {out_1.item()}"
-        )
-
-    def test_mag_cache_reset(self):
-        """Test that state resets correctly after num_inference_steps."""
-        model = DummyTransformer()
-        config = MagCacheConfig(
-            threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=np.array([1.0, 1.0])
-        )
-        apply_mag_cache(model, config)
-        self._set_context(model, "test_context")
-
-        input_t = torch.ones(1, 1, 1)
-
-        model(input_t)  # Step 0
-        model(input_t)  # Step 1 (Skipped)
-
-        # Step 2 (Reset -> Step 0) -> Should Compute
-        # Input 2.0 -> Output 8.0
-        input_t2 = torch.tensor([[[2.0]]])
-        output_t2 = model(input_t2)
-
-        self.assertTrue(torch.allclose(output_t2, torch.tensor([[[8.0]]])), "State did not reset correctly")
-
-    def test_mag_cache_calibration(self):
-        """Test that calibration mode records ratios."""
-        model = DummyTransformer()
-        config = MagCacheConfig(num_inference_steps=2, calibrate=True)
-        apply_mag_cache(model, config)
-        self._set_context(model, "test_context")
-
-        # Step 0
-        # HeadInput = 10. Output = 40. Residual = 30.
-        # Ratio 0 is placeholder 1.0
-        model(torch.tensor([[[10.0]]]))
-
-        # Check intermediate state
-        ratios = self._get_calibration_data(model)
-        self.assertEqual(len(ratios), 1)
-        self.assertEqual(ratios[0], 1.0)
-
-        # Step 1
-        # HeadInput = 10. Output = 40. Residual = 30.
-        # PrevResidual = 30. CurrResidual = 30.
-        # Ratio = 30/30 = 1.0
-        model(torch.tensor([[[10.0]]]))
-
-        # Verify it computes fully (no skip)
-        # If it skipped, output would be 41.0. It should be 40.0
-        # Actually in test setup, input is same (10.0) so output 40.0.
-        # Let's ensure list is empty after reset (end of step 1)
-        ratios_after = self._get_calibration_data(model)
-        self.assertEqual(ratios_after, [])

tests/models/testing_utils/__init__.py

@@ -5,8 +5,12 @@ from .cache import (
     FasterCacheTesterMixin,
     FirstBlockCacheConfigMixin,
     FirstBlockCacheTesterMixin,
+    MagCacheConfigMixin,
+    MagCacheTesterMixin,
     PyramidAttentionBroadcastConfigMixin,
     PyramidAttentionBroadcastTesterMixin,
+    TaylorSeerCacheConfigMixin,
+    TaylorSeerCacheTesterMixin,
 )
 from .common import BaseModelTesterConfig, ModelTesterMixin
 from .compile import TorchCompileTesterMixin
@@ -50,6 +54,8 @@ __all__ = [
     "FasterCacheTesterMixin",
     "FirstBlockCacheConfigMixin",
     "FirstBlockCacheTesterMixin",
+    "MagCacheConfigMixin",
+    "MagCacheTesterMixin",
     "GGUFCompileTesterMixin",
     "GGUFConfigMixin",
     "GGUFTesterMixin",
@@ -65,6 +71,8 @@ __all__ = [
     "ModelTesterMixin",
     "PyramidAttentionBroadcastConfigMixin",
     "PyramidAttentionBroadcastTesterMixin",
+    "TaylorSeerCacheConfigMixin",
+    "TaylorSeerCacheTesterMixin",
     "QuantizationCompileTesterMixin",
     "QuantizationTesterMixin",
     "QuantoCompileTesterMixin",

tests/models/testing_utils/cache.py

@@ -18,10 +18,18 @@ import gc
 import pytest
 import torch
 
-from diffusers.hooks import FasterCacheConfig, FirstBlockCacheConfig, PyramidAttentionBroadcastConfig
+from diffusers.hooks import (
+    FasterCacheConfig,
+    FirstBlockCacheConfig,
+    MagCacheConfig,
+    PyramidAttentionBroadcastConfig,
+    TaylorSeerCacheConfig,
+)
 from diffusers.hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK
 from diffusers.hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK
+from diffusers.hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK
 from diffusers.hooks.pyramid_attention_broadcast import _PYRAMID_ATTENTION_BROADCAST_HOOK
+from diffusers.hooks.taylorseer_cache import _TAYLORSEER_CACHE_HOOK
 from diffusers.models.cache_utils import CacheMixin
 
 from ...testing_utils import assert_tensors_close, backend_empty_cache, is_cache, torch_device
@@ -554,3 +562,192 @@ class FasterCacheTesterMixin(FasterCacheConfigMixin, CacheTesterMixin):
     @require_cache_mixin
     def test_faster_cache_reset_stateful_cache(self):
         self._test_reset_stateful_cache()
+
+
+@is_cache
+class MagCacheConfigMixin:
+    """
+    Base mixin providing MagCache config.
+
+    Expected class attributes:
+        - model_class: The model class to test (must use CacheMixin)
+    """
+
+    # Default MagCache config - can be overridden by subclasses.
+    # Uses neutral ratios [1.0, 1.0] and a high threshold so the second
+    # inference step is always skipped, which is required by _test_cache_inference.
+    MAG_CACHE_CONFIG = {
+        "num_inference_steps": 2,
+        "retention_ratio": 0.0,
+        "threshold": 100.0,
+        "mag_ratios": [1.0, 1.0],
+    }
+
+    def _get_cache_config(self):
+        return MagCacheConfig(**self.MAG_CACHE_CONFIG)
+
+    def _get_hook_names(self):
+        return [_MAG_CACHE_LEADER_BLOCK_HOOK, _MAG_CACHE_BLOCK_HOOK]
+
+
+@is_cache
+class MagCacheTesterMixin(MagCacheConfigMixin, CacheTesterMixin):
+    """
+    Mixin class for testing MagCache on models.
+
+    Expected class attributes:
+        - model_class: The model class to test (must use CacheMixin)
+
+    Expected methods to be implemented by subclasses:
+        - get_init_dict(): Returns dict of arguments to initialize the model
+        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
+
+    Pytest mark: cache
+        Use `pytest -m "not cache"` to skip these tests
+    """
+
+    @require_cache_mixin
+    def test_mag_cache_enable_disable_state(self):
+        self._test_cache_enable_disable_state()
+
+    @require_cache_mixin
+    def test_mag_cache_double_enable_raises_error(self):
+        self._test_cache_double_enable_raises_error()
+
+    @require_cache_mixin
+    def test_mag_cache_hooks_registered(self):
+        self._test_cache_hooks_registered()
+
+    @require_cache_mixin
+    def test_mag_cache_inference(self):
+        self._test_cache_inference()
+
+    @require_cache_mixin
+    def test_mag_cache_context_manager(self):
+        self._test_cache_context_manager()
+
+    @require_cache_mixin
+    def test_mag_cache_reset_stateful_cache(self):
+        self._test_reset_stateful_cache()
+
+
+@is_cache
+class TaylorSeerCacheConfigMixin:
+    """
+    Base mixin providing TaylorSeerCache config.
+
+    Expected class attributes:
+        - model_class: The model class to test (must use CacheMixin)
+    """
+
+    # Default TaylorSeerCache config - can be overridden by subclasses.
+    # Uses a low cache_interval and disable_cache_before_step=0 so the second
+    # inference step is always predicted, which is required by _test_cache_inference.
+    TAYLORSEER_CACHE_CONFIG = {
+        "cache_interval": 3,
+        "disable_cache_before_step": 1,
+        "max_order": 1,
+    }
+
+    def _get_cache_config(self):
+        return TaylorSeerCacheConfig(**self.TAYLORSEER_CACHE_CONFIG)
+
+    def _get_hook_names(self):
+        return [_TAYLORSEER_CACHE_HOOK]
+
+
+@is_cache
+class TaylorSeerCacheTesterMixin(TaylorSeerCacheConfigMixin, CacheTesterMixin):
+    """
+    Mixin class for testing TaylorSeerCache on models.
+
+    Expected class attributes:
+        - model_class: The model class to test (must use CacheMixin)
+
+    Expected methods to be implemented by subclasses:
+        - get_init_dict(): Returns dict of arguments to initialize the model
+        - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass
+
+    Pytest mark: cache
+        Use `pytest -m "not cache"` to skip these tests
+    """
+
+    @torch.no_grad()
+    def _test_cache_inference(self):
+        """Test that model can run inference with TaylorSeer cache enabled (requires cache_context)."""
+        init_dict = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
+        model = self.model_class(**init_dict).to(torch_device)
+        model.eval()
+
+        config = self._get_cache_config()
+        model.enable_cache(config)
+
+        # TaylorSeer requires cache_context to be set for inference
+        with model.cache_context("taylorseer_test"):
+            # First pass populates the cache
+            _ = model(**inputs_dict, return_dict=False)[0]
+
+            # Create modified inputs for second pass
+            inputs_dict_step2 = inputs_dict.copy()
+            if self.cache_input_key in inputs_dict_step2:
+                inputs_dict_step2[self.cache_input_key] = inputs_dict_step2[self.cache_input_key] + torch.randn_like(
+                    inputs_dict_step2[self.cache_input_key]
+                )
+
+            # Second pass - TaylorSeer should use cached Taylor series predictions
+            output_with_cache = model(**inputs_dict_step2, return_dict=False)[0]
+
+        assert output_with_cache is not None, "Model output should not be None with cache enabled."
+        assert not torch.isnan(output_with_cache).any(), "Model output contains NaN with cache enabled."
+
+        # Run same inputs without cache to compare
+        model.disable_cache()
+        output_without_cache = model(**inputs_dict_step2, return_dict=False)[0]
+
+        # Cached output should be different from non-cached output (due to approximation)
+        assert not torch.allclose(output_without_cache, output_with_cache, atol=1e-5), (
+            "Cached output should be different from non-cached output due to cache approximation."
+        )
+
+    @torch.no_grad()
+    def _test_reset_stateful_cache(self):
+        """Test that _reset_stateful_cache resets the TaylorSeer cache state (requires cache_context)."""
+        init_dict = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
+        model = self.model_class(**init_dict).to(torch_device)
+        model.eval()
+
+        config = self._get_cache_config()
+        model.enable_cache(config)
+
+        with model.cache_context("taylorseer_test"):
+            _ = model(**inputs_dict, return_dict=False)[0]
+
+        model._reset_stateful_cache()
+
+        model.disable_cache()
+
+    @require_cache_mixin
+    def test_taylorseer_cache_enable_disable_state(self):
+        self._test_cache_enable_disable_state()
+
+    @require_cache_mixin
+    def test_taylorseer_cache_double_enable_raises_error(self):
+        self._test_cache_double_enable_raises_error()
+
+    @require_cache_mixin
+    def test_taylorseer_cache_hooks_registered(self):
+        self._test_cache_hooks_registered()
+
+    @require_cache_mixin
+    def test_taylorseer_cache_inference(self):
+        self._test_cache_inference()
+
+    @require_cache_mixin
+    def test_taylorseer_cache_context_manager(self):
+        self._test_cache_context_manager()
+
+    @require_cache_mixin
+    def test_taylorseer_cache_reset_stateful_cache(self):
+        self._test_reset_stateful_cache()

tests/models/transformers/test_models_transformer_flux.py

@@ -37,6 +37,7 @@ from ..testing_utils import (
     IPAdapterTesterMixin,
     LoraHotSwappingForModelTesterMixin,
     LoraTesterMixin,
+    MagCacheTesterMixin,
     MemoryTesterMixin,
     ModelOptCompileTesterMixin,
     ModelOptTesterMixin,
@@ -45,6 +46,7 @@ from ..testing_utils import (
     QuantoCompileTesterMixin,
     QuantoTesterMixin,
     SingleFileTesterMixin,
+    TaylorSeerCacheTesterMixin,
     TorchAoCompileTesterMixin,
     TorchAoTesterMixin,
     TorchCompileTesterMixin,
@@ -430,3 +432,11 @@ class TestFluxTransformerFasterCache(FluxTransformerTesterConfig, FasterCacheTes
         "tensor_format": "BCHW",
         "is_guidance_distilled": True,
     }
+
+
+class TestFluxTransformerMagCache(FluxTransformerTesterConfig, MagCacheTesterMixin):
+    """MagCache tests for Flux Transformer."""
+
+
+class TestFluxTransformerTaylorSeerCache(FluxTransformerTesterConfig, TaylorSeerCacheTesterMixin):
+    """TaylorSeerCache tests for Flux Transformer."""

tests/models/transformers/test_models_transformer_ltx.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,47 +13,59 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import unittest
+
 import torch
 
 from diffusers import LTXVideoTransformer3DModel
-from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import (
-    BaseModelTesterConfig,
-    MemoryTesterMixin,
-    ModelTesterMixin,
-    TorchCompileTesterMixin,
-    TrainingTesterMixin,
-)
+from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin
 
 
 enable_full_determinism()
 
 
-class LTXTransformerTesterConfig(BaseModelTesterConfig):
-    @property
-    def model_class(self):
-        return LTXVideoTransformer3DModel
+class LTXTransformerTests(ModelTesterMixin, unittest.TestCase):
+    model_class = LTXVideoTransformer3DModel
+    main_input_name = "hidden_states"
+    uses_custom_attn_processor = True
 
     @property
-    def output_shape(self) -> tuple[int, int]:
-        return (512, 4)
+    def dummy_input(self):
+        batch_size = 2
+        num_channels = 4
+        num_frames = 2
+        height = 16
+        width = 16
+        embedding_dim = 16
+        sequence_length = 16
 
-    @property
-    def input_shape(self) -> tuple[int, int]:
-        return (512, 4)
+        hidden_states = torch.randn((batch_size, num_frames * height * width, num_channels)).to(torch_device)
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        encoder_attention_mask = torch.ones((batch_size, sequence_length)).bool().to(torch_device)
+        timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
 
-    @property
-    def main_input_name(self) -> str:
-        return "hidden_states"
-
-    @property
-    def generator(self):
-        return torch.Generator("cpu").manual_seed(0)
-
-    def get_init_dict(self):
         return {
+            "hidden_states": hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "timestep": timestep,
+            "encoder_attention_mask": encoder_attention_mask,
+            "num_frames": num_frames,
+            "height": height,
+            "width": width,
+        }
+
+    @property
+    def input_shape(self):
+        return (512, 4)
+
+    @property
+    def output_shape(self):
+        return (512, 4)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
             "in_channels": 4,
             "out_channels": 4,
             "num_attention_heads": 2,
@@ -62,57 +75,16 @@ class LTXTransformerTesterConfig(BaseModelTesterConfig):
             "qk_norm": "rms_norm_across_heads",
             "caption_channels": 16,
         }
-
-    def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
-        batch_size = 2
-        num_channels = 4
-        num_frames = 2
-        height = 16
-        width = 16
-        embedding_dim = 16
-        sequence_length = 16
-
-        return {
-            "hidden_states": randn_tensor(
-                (batch_size, num_frames * height * width, num_channels),
-                generator=self.generator,
-                device=torch_device,
-            ),
-            "encoder_hidden_states": randn_tensor(
-                (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
-            ),
-            "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
-            "encoder_attention_mask": torch.ones((batch_size, sequence_length)).bool().to(torch_device),
-            "num_frames": num_frames,
-            "height": height,
-            "width": width,
-        }
-
-
-class TestLTXTransformer(LTXTransformerTesterConfig, ModelTesterMixin):
-    """Core model tests for LTX Video Transformer."""
-
-
-class TestLTXTransformerMemory(LTXTransformerTesterConfig, MemoryTesterMixin):
-    """Memory optimization tests for LTX Video Transformer."""
-
-
-class TestLTXTransformerTraining(LTXTransformerTesterConfig, TrainingTesterMixin):
-    """Training tests for LTX Video Transformer."""
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
 
     def test_gradient_checkpointing_is_applied(self):
-        super().test_gradient_checkpointing_is_applied(expected_set={"LTXVideoTransformer3DModel"})
+        expected_set = {"LTXVideoTransformer3DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
 
-class TestLTXTransformerCompile(LTXTransformerTesterConfig, TorchCompileTesterMixin):
-    """Torch compile tests for LTX Video Transformer."""
+class LTXTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = LTXVideoTransformer3DModel
 
-
-# TODO: Add pretrained_model_name_or_path once a tiny LTX model is available on the Hub
-# class TestLTXTransformerBitsAndBytes(LTXTransformerTesterConfig, BitsAndBytesTesterMixin):
-#     """BitsAndBytes quantization tests for LTX Video Transformer."""
-
-
-# TODO: Add pretrained_model_name_or_path once a tiny LTX model is available on the Hub
-# class TestLTXTransformerTorchAo(LTXTransformerTesterConfig, TorchAoTesterMixin):
-#     """TorchAo quantization tests for LTX Video Transformer."""
+    def prepare_init_args_and_inputs_for_common(self):
+        return LTXTransformerTests().prepare_init_args_and_inputs_for_common()

tests/models/transformers/test_models_transformer_ltx2.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 HuggingFace Inc.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,49 +13,77 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
+import unittest
+
 import torch
 
 from diffusers import LTX2VideoTransformer3DModel
-from diffusers.utils.torch_utils import randn_tensor
 
 from ...testing_utils import enable_full_determinism, torch_device
-from ..testing_utils import (
-    AttentionTesterMixin,
-    BaseModelTesterConfig,
-    MemoryTesterMixin,
-    ModelTesterMixin,
-    TorchCompileTesterMixin,
-    TrainingTesterMixin,
-)
+from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin
 
 
 enable_full_determinism()
 
 
-class LTX2TransformerTesterConfig(BaseModelTesterConfig):
-    @property
-    def model_class(self):
-        return LTX2VideoTransformer3DModel
+class LTX2TransformerTests(ModelTesterMixin, unittest.TestCase):
+    model_class = LTX2VideoTransformer3DModel
+    main_input_name = "hidden_states"
+    uses_custom_attn_processor = True
 
     @property
-    def output_shape(self) -> tuple[int, int]:
-        return (512, 4)
+    def dummy_input(self):
+        # Common
+        batch_size = 2
 
-    @property
-    def input_shape(self) -> tuple[int, int]:
-        return (512, 4)
+        # Video
+        num_frames = 2
+        num_channels = 4
+        height = 16
+        width = 16
 
-    @property
-    def main_input_name(self) -> str:
-        return "hidden_states"
+        # Audio
+        audio_num_frames = 9
+        audio_num_channels = 2
+        num_mel_bins = 2
 
-    @property
-    def generator(self):
-        return torch.Generator("cpu").manual_seed(0)
+        # Text
+        embedding_dim = 16
+        sequence_length = 16
+
+        hidden_states = torch.randn((batch_size, num_frames * height * width, num_channels)).to(torch_device)
+        audio_hidden_states = torch.randn((batch_size, audio_num_frames, audio_num_channels * num_mel_bins)).to(
+            torch_device
+        )
+        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        audio_encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
+        encoder_attention_mask = torch.ones((batch_size, sequence_length)).bool().to(torch_device)
+        timestep = torch.rand((batch_size,)).to(torch_device) * 1000
 
-    def get_init_dict(self):
         return {
+            "hidden_states": hidden_states,
+            "audio_hidden_states": audio_hidden_states,
+            "encoder_hidden_states": encoder_hidden_states,
+            "audio_encoder_hidden_states": audio_encoder_hidden_states,
+            "timestep": timestep,
+            "encoder_attention_mask": encoder_attention_mask,
+            "num_frames": num_frames,
+            "height": height,
+            "width": width,
+            "audio_num_frames": audio_num_frames,
+            "fps": 25.0,
+        }
+
+    @property
+    def input_shape(self):
+        return (512, 4)
+
+    @property
+    def output_shape(self):
+        return (512, 4)
+
+    def prepare_init_args_and_inputs_for_common(self):
+        init_dict = {
             "in_channels": 4,
             "out_channels": 4,
             "patch_size": 1,
@@ -72,80 +101,122 @@ class LTX2TransformerTesterConfig(BaseModelTesterConfig):
             "caption_channels": 16,
             "rope_double_precision": False,
         }
-
-    def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
-        batch_size = 2
-        num_frames = 2
-        num_channels = 4
-        height = 16
-        width = 16
-        audio_num_frames = 9
-        audio_num_channels = 2
-        num_mel_bins = 2
-        embedding_dim = 16
-        sequence_length = 16
-
-        return {
-            "hidden_states": randn_tensor(
-                (batch_size, num_frames * height * width, num_channels),
-                generator=self.generator,
-                device=torch_device,
-            ),
-            "audio_hidden_states": randn_tensor(
-                (batch_size, audio_num_frames, audio_num_channels * num_mel_bins),
-                generator=self.generator,
-                device=torch_device,
-            ),
-            "encoder_hidden_states": randn_tensor(
-                (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
-            ),
-            "audio_encoder_hidden_states": randn_tensor(
-                (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
-            ),
-            "timestep": (randn_tensor((batch_size,), generator=self.generator, device=torch_device).abs() * 1000),
-            "encoder_attention_mask": torch.ones((batch_size, sequence_length)).bool().to(torch_device),
-            "num_frames": num_frames,
-            "height": height,
-            "width": width,
-            "audio_num_frames": audio_num_frames,
-            "fps": 25.0,
-        }
-
-
-class TestLTX2Transformer(LTX2TransformerTesterConfig, ModelTesterMixin):
-    """Core model tests for LTX2 Video Transformer."""
-
-
-class TestLTX2TransformerMemory(LTX2TransformerTesterConfig, MemoryTesterMixin):
-    """Memory optimization tests for LTX2 Video Transformer."""
-
-
-class TestLTX2TransformerTraining(LTX2TransformerTesterConfig, TrainingTesterMixin):
-    """Training tests for LTX2 Video Transformer."""
+        inputs_dict = self.dummy_input
+        return init_dict, inputs_dict
 
     def test_gradient_checkpointing_is_applied(self):
-        super().test_gradient_checkpointing_is_applied(expected_set={"LTX2VideoTransformer3DModel"})
+        expected_set = {"LTX2VideoTransformer3DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+    # def test_ltx2_consistency(self, seed=0, dtype=torch.float32):
+    #     torch.manual_seed(seed)
+    #     init_dict, _ = self.prepare_init_args_and_inputs_for_common()
+
+    #     # Calculate dummy inputs in a custom manner to ensure compatibility with original code
+    #     batch_size = 2
+    #     num_frames = 9
+    #     latent_frames = 2
+    #     text_embedding_dim = 16
+    #     text_seq_len = 16
+    #     fps = 25.0
+    #     sampling_rate = 16000.0
+    #     hop_length = 160.0
+
+    #     sigma = torch.rand((1,), generator=torch.manual_seed(seed), dtype=dtype, device="cpu") * 1000
+    #     timestep = (sigma * torch.ones((batch_size,), dtype=dtype, device="cpu")).to(device=torch_device)
+
+    #     num_channels = 4
+    #     latent_height = 4
+    #     latent_width = 4
+    #     hidden_states = torch.randn(
+    #         (batch_size, num_channels, latent_frames, latent_height, latent_width),
+    #         generator=torch.manual_seed(seed),
+    #         dtype=dtype,
+    #         device="cpu",
+    #     )
+    #     # Patchify video latents (with patch_size (1, 1, 1))
+    #     hidden_states = hidden_states.reshape(batch_size, -1, latent_frames, 1, latent_height, 1, latent_width, 1)
+    #     hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+    #     encoder_hidden_states = torch.randn(
+    #         (batch_size, text_seq_len, text_embedding_dim),
+    #         generator=torch.manual_seed(seed),
+    #         dtype=dtype,
+    #         device="cpu",
+    #     )
+
+    #     audio_num_channels = 2
+    #     num_mel_bins = 2
+    #     latent_length = int((sampling_rate / hop_length / 4) * (num_frames / fps))
+    #     audio_hidden_states = torch.randn(
+    #         (batch_size, audio_num_channels, latent_length, num_mel_bins),
+    #         generator=torch.manual_seed(seed),
+    #         dtype=dtype,
+    #         device="cpu",
+    #     )
+    #     # Patchify audio latents
+    #     audio_hidden_states = audio_hidden_states.transpose(1, 2).flatten(2, 3)
+    #     audio_encoder_hidden_states = torch.randn(
+    #         (batch_size, text_seq_len, text_embedding_dim),
+    #         generator=torch.manual_seed(seed),
+    #         dtype=dtype,
+    #         device="cpu",
+    #     )
+
+    #     inputs_dict = {
+    #         "hidden_states": hidden_states.to(device=torch_device),
+    #         "audio_hidden_states": audio_hidden_states.to(device=torch_device),
+    #         "encoder_hidden_states": encoder_hidden_states.to(device=torch_device),
+    #         "audio_encoder_hidden_states": audio_encoder_hidden_states.to(device=torch_device),
+    #         "timestep": timestep,
+    #         "num_frames": latent_frames,
+    #         "height": latent_height,
+    #         "width": latent_width,
+    #         "audio_num_frames": num_frames,
+    #         "fps": 25.0,
+    #     }
+
+    #     model = self.model_class.from_pretrained(
+    #         "diffusers-internal-dev/dummy-ltx2",
+    #         subfolder="transformer",
+    #         device_map="cpu",
+    #     )
+    #     # torch.manual_seed(seed)
+    #     # model = self.model_class(**init_dict)
+    #     model.to(torch_device)
+    #     model.eval()
+
+    #     with attention_backend("native"):
+    #         with torch.no_grad():
+    #             output = model(**inputs_dict)
+
+    #             video_output, audio_output = output.to_tuple()
+
+    #     self.assertIsNotNone(video_output)
+    #     self.assertIsNotNone(audio_output)
+
+    #     # input & output have to have the same shape
+    #     video_expected_shape = (batch_size, latent_frames * latent_height * latent_width, num_channels)
+    #     self.assertEqual(video_output.shape, video_expected_shape, "Video input and output shapes do not match")
+    #     audio_expected_shape = (batch_size, latent_length, audio_num_channels * num_mel_bins)
+    #     self.assertEqual(audio_output.shape, audio_expected_shape, "Audio input and output shapes do not match")
+
+    #     # Check against expected slice
+    #     # fmt: off
+    #     video_expected_slice = torch.tensor([0.4783, 1.6954, -1.2092, 0.1762, 0.7801, 1.2025, -1.4525, -0.2721, 0.3354, 1.9144, -1.5546, 0.0831, 0.4391, 1.7012, -1.7373, -0.2676])
+    #     audio_expected_slice = torch.tensor([-0.4236, 0.4750, 0.3901, -0.4339, -0.2782, 0.4357, 0.4526, -0.3927, -0.0980, 0.4870, 0.3964, -0.3169, -0.3974, 0.4408, 0.3809, -0.4692])
+    #     # fmt: on
+
+    #     video_output_flat = video_output.cpu().flatten().float()
+    #     video_generated_slice = torch.cat([video_output_flat[:8], video_output_flat[-8:]])
+    #     self.assertTrue(torch.allclose(video_generated_slice, video_expected_slice, atol=1e-4))
+
+    #     audio_output_flat = audio_output.cpu().flatten().float()
+    #     audio_generated_slice = torch.cat([audio_output_flat[:8], audio_output_flat[-8:]])
+    #     self.assertTrue(torch.allclose(audio_generated_slice, audio_expected_slice, atol=1e-4))
 
 
-class TestLTX2TransformerAttention(LTX2TransformerTesterConfig, AttentionTesterMixin):
-    """Attention processor tests for LTX2 Video Transformer."""
+class LTX2TransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase):
+    model_class = LTX2VideoTransformer3DModel
 
-    @pytest.mark.skip(
-        "LTX2Attention does not set is_cross_attention, so fuse_projections tries to fuse Q+K+V together even for cross-attention modules with different input dimensions."
-    )
-    def test_fuse_unfuse_qkv_projections(self, atol=1e-3, rtol=0):
-        pass
-
-
-class TestLTX2TransformerCompile(LTX2TransformerTesterConfig, TorchCompileTesterMixin):
-    """Torch compile tests for LTX2 Video Transformer."""
-
-
-# TODO: Add pretrained_model_name_or_path once a tiny LTX2 model is available on the Hub
-# class TestLTX2TransformerBitsAndBytes(LTX2TransformerTesterConfig, BitsAndBytesTesterMixin):
-#     """BitsAndBytes quantization tests for LTX2 Video Transformer."""
-
-
-# TODO: Add pretrained_model_name_or_path once a tiny LTX2 model is available on the Hub
-# class TestLTX2TransformerTorchAo(LTX2TransformerTesterConfig, TorchAoTesterMixin):
-#     """TorchAo quantization tests for LTX2 Video Transformer."""
+    def prepare_init_args_and_inputs_for_common(self):
+        return LTX2TransformerTests().prepare_init_args_and_inputs_for_common()