Merge branch 'main' into sd3-test-refactor

.github/workflows/claude_review.yml

@@ -32,9 +32,6 @@ jobs:
       )
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-        with:
-          fetch-depth: 1
       - uses: anthropics/claude-code-action@v1
         with:
           anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}

src/diffusers/pipelines/pipeline_utils.py

@@ -764,7 +764,7 @@ class DiffusionPipeline(ConfigMixin, PushToHubMixin):
         token = kwargs.pop("token", None)
         revision = kwargs.pop("revision", None)
         from_flax = kwargs.pop("from_flax", False)
-        torch_dtype = kwargs.pop("torch_dtype", torch.float32)
+        torch_dtype = kwargs.pop("torch_dtype", None)
         custom_pipeline = kwargs.pop("custom_pipeline", None)
         custom_revision = kwargs.pop("custom_revision", None)
         provider = kwargs.pop("provider", None)

src/diffusers/quantizers/quantization_config.py

@@ -470,8 +470,8 @@ class TorchAoConfig(QuantizationConfigMixin):
         self.post_init()
 
     def post_init(self):
-        if is_torchao_version("<", "0.15.0"):
-            raise ValueError("TorchAoConfig requires torchao >= 0.15.0. Please upgrade with `pip install -U torchao`.")
+        if is_torchao_version("<=", "0.9.0"):
+            raise ValueError("TorchAoConfig requires torchao > 0.9.0. Please upgrade with `pip install -U torchao`.")
 
         from torchao.quantization.quant_api import AOBaseConfig
 
@@ -495,8 +495,8 @@ class TorchAoConfig(QuantizationConfigMixin):
     @classmethod
     def from_dict(cls, config_dict, return_unused_kwargs=False, **kwargs):
         """Create configuration from a dictionary."""
-        if not is_torchao_version(">=", "0.15.0"):
-            raise NotImplementedError("TorchAoConfig requires torchao >= 0.15.0 for construction from dict")
+        if not is_torchao_version(">", "0.9.0"):
+            raise NotImplementedError("TorchAoConfig requires torchao > 0.9.0 for construction from dict")
         config_dict = config_dict.copy()
         quant_type = config_dict.pop("quant_type")
 

src/diffusers/quantizers/torchao/torchao_quantizer.py

@@ -113,7 +113,7 @@ if (
     is_torch_available()
     and is_torch_version(">=", "2.6.0")
     and is_torchao_available()
-    and is_torchao_version(">=", "0.15.0")
+    and is_torchao_version(">=", "0.7.0")
 ):
     _update_torch_safe_globals()
 
@@ -168,10 +168,10 @@ class TorchAoHfQuantizer(DiffusersQuantizer):
             raise ImportError(
                 "Loading a TorchAO quantized model requires the torchao library. Please install with `pip install torchao`"
             )
-        torchao_version = version.parse(importlib.metadata.version("torchao"))
-        if torchao_version < version.parse("0.15.0"):
+        torchao_version = version.parse(importlib.metadata.version("torch"))
+        if torchao_version < version.parse("0.7.0"):
             raise RuntimeError(
-                f"The minimum required version of `torchao` is 0.15.0, but the current version is {torchao_version}. Please upgrade with `pip install -U torchao`."
+                f"The minimum required version of `torchao` is 0.7.0, but the current version is {torchao_version}. Please upgrade with `pip install -U torchao`."
             )
 
         self.offload = False

tests/models/transformers/test_models_transformer_sd3.py

@@ -13,58 +13,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
 import torch
 
 from diffusers import SD3Transformer2DModel
-from diffusers.utils.import_utils import is_xformers_available
+from diffusers.utils.torch_utils import randn_tensor
 
-from ...testing_utils import (
-    enable_full_determinism,
-    torch_device,
+from ...testing_utils import enable_full_determinism, torch_device
+from ..testing_utils import (
+    BaseModelTesterConfig,
+    BitsAndBytesTesterMixin,
+    ModelTesterMixin,
+    TorchAoTesterMixin,
+    TorchCompileTesterMixin,
+    TrainingTesterMixin,
 )
-from ..test_modeling_common import ModelTesterMixin
 
 
 enable_full_determinism()
 
 
-class SD3TransformerTests(ModelTesterMixin, unittest.TestCase):
-    model_class = SD3Transformer2DModel
-    main_input_name = "hidden_states"
-    model_split_percents = [0.8, 0.8, 0.9]
+# ======================== SD3 Transformer ========================
+
+
+class SD3TransformerTesterConfig(BaseModelTesterConfig):
+    @property
+    def model_class(self):
+        return SD3Transformer2DModel
 
     @property
-    def dummy_input(self):
-        batch_size = 2
-        num_channels = 4
-        height = width = embedding_dim = 32
-        pooled_embedding_dim = embedding_dim * 2
-        sequence_length = 154
+    def pretrained_model_name_or_path(self):
+        return "hf-internal-testing/tiny-sd3-pipe"
 
-        hidden_states = torch.randn((batch_size, num_channels, height, width)).to(torch_device)
-        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
-        pooled_prompt_embeds = torch.randn((batch_size, pooled_embedding_dim)).to(torch_device)
-        timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
+    @property
+    def pretrained_model_kwargs(self):
+        return {"subfolder": "transformer"}
 
+    @property
+    def main_input_name(self) -> str:
+        return "hidden_states"
+
+    @property
+    def model_split_percents(self) -> list:
+        return [0.8, 0.8, 0.9]
+
+    @property
+    def output_shape(self) -> tuple:
+        return (4, 32, 32)
+
+    @property
+    def input_shape(self) -> tuple:
+        return (4, 32, 32)
+
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self) -> dict:
         return {
-            "hidden_states": hidden_states,
-            "encoder_hidden_states": encoder_hidden_states,
-            "pooled_projections": pooled_prompt_embeds,
-            "timestep": timestep,
-        }
-
-    @property
-    def input_shape(self):
-        return (4, 32, 32)
-
-    @property
-    def output_shape(self):
-        return (4, 32, 32)
-
-    def prepare_init_args_and_inputs_for_common(self):
-        init_dict = {
             "sample_size": 32,
             "patch_size": 1,
             "in_channels": 4,
@@ -79,67 +84,79 @@ class SD3TransformerTests(ModelTesterMixin, unittest.TestCase):
             "dual_attention_layers": (),
             "qk_norm": None,
         }
-        inputs_dict = self.dummy_input
-        return init_dict, inputs_dict
 
-    @unittest.skipIf(
-        torch_device != "cuda" or not is_xformers_available(),
-        reason="XFormers attention is only available with CUDA and `xformers` installed",
-    )
-    def test_xformers_enable_works(self):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
-
-        model.enable_xformers_memory_efficient_attention()
-
-        assert model.transformer_blocks[0].attn.processor.__class__.__name__ == "XFormersJointAttnProcessor", (
-            "xformers is not enabled"
-        )
-
-    @unittest.skip("SD3Transformer2DModel uses a dedicated attention processor. This test doesn't apply")
-    def test_set_attn_processor_for_determinism(self):
-        pass
-
-    def test_gradient_checkpointing_is_applied(self):
-        expected_set = {"SD3Transformer2DModel"}
-        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
-
-
-class SD35TransformerTests(ModelTesterMixin, unittest.TestCase):
-    model_class = SD3Transformer2DModel
-    main_input_name = "hidden_states"
-    model_split_percents = [0.8, 0.8, 0.9]
-
-    @property
-    def dummy_input(self):
-        batch_size = 2
+    def get_dummy_inputs(self, batch_size: int = 2) -> dict[str, torch.Tensor]:
         num_channels = 4
         height = width = embedding_dim = 32
         pooled_embedding_dim = embedding_dim * 2
         sequence_length = 154
 
-        hidden_states = torch.randn((batch_size, num_channels, height, width)).to(torch_device)
-        encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device)
-        pooled_prompt_embeds = torch.randn((batch_size, pooled_embedding_dim)).to(torch_device)
-        timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device)
-
         return {
-            "hidden_states": hidden_states,
-            "encoder_hidden_states": encoder_hidden_states,
-            "pooled_projections": pooled_prompt_embeds,
-            "timestep": timestep,
+            "hidden_states": randn_tensor(
+                (batch_size, num_channels, height, width), generator=self.generator, device=torch_device
+            ),
+            "encoder_hidden_states": randn_tensor(
+                (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+            ),
+            "pooled_projections": randn_tensor(
+                (batch_size, pooled_embedding_dim), generator=self.generator, device=torch_device
+            ),
+            "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
         }
 
+
+class TestSD3Transformer(SD3TransformerTesterConfig, ModelTesterMixin):
+    pass
+
+
+class TestSD3TransformerTraining(SD3TransformerTesterConfig, TrainingTesterMixin):
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"SD3Transformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class TestSD3TransformerCompile(SD3TransformerTesterConfig, TorchCompileTesterMixin):
+    pass
+
+
+# ======================== SD3.5 Transformer ========================
+
+
+class SD35TransformerTesterConfig(BaseModelTesterConfig):
     @property
-    def input_shape(self):
+    def model_class(self):
+        return SD3Transformer2DModel
+
+    @property
+    def pretrained_model_name_or_path(self):
+        return "hf-internal-testing/tiny-sd35-pipe"
+
+    @property
+    def pretrained_model_kwargs(self):
+        return {"subfolder": "transformer"}
+
+    @property
+    def main_input_name(self) -> str:
+        return "hidden_states"
+
+    @property
+    def model_split_percents(self) -> list:
+        return [0.8, 0.8, 0.9]
+
+    @property
+    def output_shape(self) -> tuple:
         return (4, 32, 32)
 
     @property
-    def output_shape(self):
+    def input_shape(self) -> tuple:
         return (4, 32, 32)
 
-    def prepare_init_args_and_inputs_for_common(self):
-        init_dict = {
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self) -> dict:
+        return {
             "sample_size": 32,
             "patch_size": 1,
             "in_channels": 4,
@@ -154,47 +171,56 @@ class SD35TransformerTests(ModelTesterMixin, unittest.TestCase):
             "dual_attention_layers": (0,),
             "qk_norm": "rms_norm",
         }
-        inputs_dict = self.dummy_input
-        return init_dict, inputs_dict
 
-    @unittest.skipIf(
-        torch_device != "cuda" or not is_xformers_available(),
-        reason="XFormers attention is only available with CUDA and `xformers` installed",
-    )
-    def test_xformers_enable_works(self):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
-        model = self.model_class(**init_dict)
+    def get_dummy_inputs(self, batch_size: int = 2) -> dict[str, torch.Tensor]:
+        num_channels = 4
+        height = width = embedding_dim = 32
+        pooled_embedding_dim = embedding_dim * 2
+        sequence_length = 154
 
-        model.enable_xformers_memory_efficient_attention()
+        return {
+            "hidden_states": randn_tensor(
+                (batch_size, num_channels, height, width), generator=self.generator, device=torch_device
+            ),
+            "encoder_hidden_states": randn_tensor(
+                (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device
+            ),
+            "pooled_projections": randn_tensor(
+                (batch_size, pooled_embedding_dim), generator=self.generator, device=torch_device
+            ),
+            "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device),
+        }
 
-        assert model.transformer_blocks[0].attn.processor.__class__.__name__ == "XFormersJointAttnProcessor", (
-            "xformers is not enabled"
-        )
-
-    @unittest.skip("SD3Transformer2DModel uses a dedicated attention processor. This test doesn't apply")
-    def test_set_attn_processor_for_determinism(self):
-        pass
-
-    def test_gradient_checkpointing_is_applied(self):
-        expected_set = {"SD3Transformer2DModel"}
-        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
 
+class TestSD35Transformer(SD35TransformerTesterConfig, ModelTesterMixin):
     def test_skip_layers(self):
-        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        init_dict = self.get_init_dict()
+        inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**init_dict).to(torch_device)
 
-        # Forward pass without skipping layers
         output_full = model(**inputs_dict).sample
 
-        # Forward pass with skipping layers 0 (since there's only one layer in this test setup)
         inputs_dict_with_skip = inputs_dict.copy()
         inputs_dict_with_skip["skip_layers"] = [0]
         output_skip = model(**inputs_dict_with_skip).sample
 
-        # Check that the outputs are different
-        self.assertFalse(
-            torch.allclose(output_full, output_skip, atol=1e-5), "Outputs should differ when layers are skipped"
-        )
+        assert not torch.allclose(output_full, output_skip, atol=1e-5), "Outputs should differ when layers are skipped"
+        assert output_full.shape == output_skip.shape, "Outputs should have the same shape"
 
-        # Check that the outputs have the same shape
-        self.assertEqual(output_full.shape, output_skip.shape, "Outputs should have the same shape")
+
+class TestSD35TransformerTraining(SD35TransformerTesterConfig, TrainingTesterMixin):
+    def test_gradient_checkpointing_is_applied(self):
+        expected_set = {"SD3Transformer2DModel"}
+        super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
+
+
+class TestSD35TransformerCompile(SD35TransformerTesterConfig, TorchCompileTesterMixin):
+    pass
+
+
+class TestSD35TransformerBitsAndBytes(SD35TransformerTesterConfig, BitsAndBytesTesterMixin):
+    """BitsAndBytes quantization tests for SD3.5 Transformer."""
+
+
+class TestSD35TransformerTorchAo(SD35TransformerTesterConfig, TorchAoTesterMixin):
+    """TorchAO quantization tests for SD3.5 Transformer."""

tests/quantization/torchao/test_torchao.py

@@ -14,11 +14,13 @@
 # limitations under the License.
 
 import gc
+import importlib.metadata
 import tempfile
 import unittest
 from typing import List
 
 import numpy as np
+from packaging import version
 from parameterized import parameterized
 from transformers import AutoTokenizer, CLIPTextModel, CLIPTokenizer, T5EncoderModel
 
@@ -80,17 +82,18 @@ if is_torchao_available():
         Float8WeightOnlyConfig,
         Int4WeightOnlyConfig,
         Int8DynamicActivationInt8WeightConfig,
-        Int8DynamicActivationIntxWeightConfig,
         Int8WeightOnlyConfig,
-        IntxWeightOnlyConfig,
     )
     from torchao.quantization.linear_activation_quantized_tensor import LinearActivationQuantizedTensor
     from torchao.utils import get_model_size_in_bytes
 
+    if version.parse(importlib.metadata.version("torchao")) >= version.Version("0.10.0"):
+        from torchao.quantization import Int8DynamicActivationIntxWeightConfig, IntxWeightOnlyConfig
+
 
 @require_torch
 @require_torch_accelerator
-@require_torchao_version_greater_or_equal("0.15.0")
+@require_torchao_version_greater_or_equal("0.14.0")
 class TorchAoConfigTest(unittest.TestCase):
     def test_to_dict(self):
         """
@@ -125,7 +128,7 @@ class TorchAoConfigTest(unittest.TestCase):
 # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners
 @require_torch
 @require_torch_accelerator
-@require_torchao_version_greater_or_equal("0.15.0")
+@require_torchao_version_greater_or_equal("0.14.0")
 class TorchAoTest(unittest.TestCase):
     def tearDown(self):
         gc.collect()
@@ -524,7 +527,7 @@ class TorchAoTest(unittest.TestCase):
         inputs = self.get_dummy_inputs(torch_device)
         _ = pipe(**inputs)
 
-    @require_torchao_version_greater_or_equal("0.15.0")
+    @require_torchao_version_greater_or_equal("0.9.0")
     def test_aobase_config(self):
         quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
         components = self.get_dummy_components(quantization_config)
@@ -537,7 +540,7 @@ class TorchAoTest(unittest.TestCase):
 # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners
 @require_torch
 @require_torch_accelerator
-@require_torchao_version_greater_or_equal("0.15.0")
+@require_torchao_version_greater_or_equal("0.14.0")
 class TorchAoSerializationTest(unittest.TestCase):
     model_name = "hf-internal-testing/tiny-flux-pipe"
 
@@ -647,7 +650,7 @@ class TorchAoSerializationTest(unittest.TestCase):
         self._check_serialization_expected_slice(quant_type, expected_slice, device)
 
 
-@require_torchao_version_greater_or_equal("0.15.0")
+@require_torchao_version_greater_or_equal("0.14.0")
 class TorchAoCompileTest(QuantCompileTests, unittest.TestCase):
     @property
     def quantization_config(self):
@@ -693,7 +696,7 @@ class TorchAoCompileTest(QuantCompileTests, unittest.TestCase):
 # Slices for these tests have been obtained on our aws-g6e-xlarge-plus runners
 @require_torch
 @require_torch_accelerator
-@require_torchao_version_greater_or_equal("0.15.0")
+@require_torchao_version_greater_or_equal("0.14.0")
 @slow
 @nightly
 class SlowTorchAoTests(unittest.TestCase):
@@ -851,7 +854,7 @@ class SlowTorchAoTests(unittest.TestCase):
 
 @require_torch
 @require_torch_accelerator
-@require_torchao_version_greater_or_equal("0.15.0")
+@require_torchao_version_greater_or_equal("0.14.0")
 @slow
 @nightly
 class SlowTorchAoPreserializedModelTests(unittest.TestCase):