upd

Signed-off-by: Lancer <maruixiang6688@gmail.com>

Author: RuixiangMa

src/diffusers/models/transformers/transformer_longcat_audio_dit.py

@@ -188,8 +188,8 @@ def __call__(
         self,
         attn: "AudioDiTAttention",
         hidden_states: torch.Tensor,
-        mask: torch.BoolTensor | None = None,
-        rope: tuple | None = None,
+        attention_mask: torch.BoolTensor | None = None,
+        audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
     ) -> torch.Tensor:
         batch_size = hidden_states.shape[0]
         query = attn.to_q(hidden_states)
@@ -205,20 +205,20 @@ def __call__(
         key = key.view(batch_size, -1, attn.heads, head_dim)
         value = value.view(batch_size, -1, attn.heads, head_dim)
 
-        if rope is not None:
-            query = _apply_rotary_emb(query, rope)
-            key = _apply_rotary_emb(key, rope)
+        if audio_rotary_emb is not None:
+            query = _apply_rotary_emb(query, audio_rotary_emb)
+            key = _apply_rotary_emb(key, audio_rotary_emb)
 
         hidden_states = dispatch_attention_fn(
             query,
             key,
             value,
-            attn_mask=mask,
+            attn_mask=attention_mask,
             backend=self._attention_backend,
             parallel_config=self._parallel_config,
         )
-        if mask is not None:
-            hidden_states = hidden_states * mask[:, :, None, None].to(hidden_states.dtype)
+        if attention_mask is not None:
+            hidden_states = hidden_states * attention_mask[:, :, None, None].to(hidden_states.dtype)
 
         hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
         hidden_states = attn.to_out[0](hidden_states)
@@ -261,11 +261,14 @@ def forward(
         attention_mask: torch.BoolTensor | None = None,
         audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         prompt_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-        mask: torch.BoolTensor | None = None,
-        rope: tuple | None = None,
     ) -> torch.Tensor:
         if encoder_hidden_states is None:
-            return self.processor(self, hidden_states, mask=mask, rope=rope)
+            return self.processor(
+                self,
+                hidden_states,
+                attention_mask=attention_mask,
+                audio_rotary_emb=audio_rotary_emb,
+            )
         return self.processor(
             self,
             hidden_states,
@@ -419,7 +422,11 @@ def forward(
 
         norm_hidden_states = F.layer_norm(hidden_states.float(), (hidden_states.shape[-1],), eps=1e-6).type_as(hidden_states)
         norm_hidden_states = norm_hidden_states * (1 + scale_sa[:, None]) + shift_sa[:, None]
-        attn_output = self.self_attn(norm_hidden_states, mask=mask, rope=rope)
+        attn_output = self.self_attn(
+            norm_hidden_states,
+            attention_mask=mask,
+            audio_rotary_emb=rope,
+        )
         hidden_states = hidden_states + gate_sa.unsqueeze(1) * attn_output
 
         if self.use_cross_attn:

src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py

@@ -171,8 +171,7 @@ def __init__(
             transformer=transformer,
         )
         self.sample_rate = getattr(vae.config, "sample_rate", 24000)
-        self.latent_hop = getattr(vae.config, "downsampling_ratio", 2048)
-        self.vae_scale_factor = self.latent_hop
+        self.vae_scale_factor = getattr(vae.config, "downsampling_ratio", 2048)
         self.latent_dim = getattr(transformer.config, "latent_dim", 64)
         self.max_wav_duration = 30.0
         self.text_norm_feat = True
@@ -321,8 +320,7 @@ def from_pretrained(
 
         pipe = cls(vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer)
         pipe.sample_rate = config.get("sampling_rate", pipe.sample_rate)
-        pipe.latent_hop = config.get("latent_hop", pipe.latent_hop)
-        pipe.vae_scale_factor = pipe.latent_hop
+        pipe.vae_scale_factor = config.get("vae_scale_factor", config.get("latent_hop", pipe.vae_scale_factor))
         pipe.max_wav_duration = config.get("max_wav_duration", pipe.max_wav_duration)
         pipe.text_norm_feat = config.get("text_norm_feat", pipe.text_norm_feat)
         pipe.text_add_embed = config.get("text_add_embed", pipe.text_add_embed)

tests/models/transformers/test_models_transformer_longcat_audio_dit.py

@@ -1,54 +1,125 @@
+# coding=utf-8
+# 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 pytest
 import torch
 
 from diffusers import LongCatAudioDiTTransformer
+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,
+)
+
+
+enable_full_determinism()
+
+
+class LongCatAudioDiTTransformerTesterConfig(BaseModelTesterConfig):
+    @property
+    def main_input_name(self) -> str:
+        return "hidden_states"
+
+    @property
+    def model_class(self):
+        return LongCatAudioDiTTransformer
+
+    @property
+    def output_shape(self) -> tuple[int, ...]:
+        return (16, 8)
+
+    @property
+    def generator(self):
+        return torch.Generator("cpu").manual_seed(0)
+
+    def get_init_dict(self) -> dict[str, int | bool | float | str]:
+        return {
+            "dit_dim": 64,
+            "dit_depth": 2,
+            "dit_heads": 4,
+            "dit_text_dim": 32,
+            "latent_dim": 8,
+            "text_conv": False,
+        }
+
+    def get_dummy_inputs(self) -> dict[str, torch.Tensor]:
+        batch_size = 1
+        sequence_length = 16
+        encoder_sequence_length = 10
+        latent_dim = 8
+        text_dim = 32
+
+        return {
+            "hidden_states": randn_tensor(
+                (batch_size, sequence_length, latent_dim), generator=self.generator, device=torch_device
+            ),
+            "encoder_hidden_states": randn_tensor(
+                (batch_size, encoder_sequence_length, text_dim), generator=self.generator, device=torch_device
+            ),
+            "encoder_attention_mask": torch.ones(
+                batch_size, encoder_sequence_length, dtype=torch.bool, device=torch_device
+            ),
+            "attention_mask": torch.ones(batch_size, sequence_length, dtype=torch.bool, device=torch_device),
+            "timestep": torch.ones(batch_size, device=torch_device),
+        }
+
+
+class TestLongCatAudioDiTTransformer(LongCatAudioDiTTransformerTesterConfig, ModelTesterMixin):
+    pass
+
+
+class TestLongCatAudioDiTTransformerMemory(LongCatAudioDiTTransformerTesterConfig, MemoryTesterMixin):
+    def test_layerwise_casting_memory(self):
+        pytest.skip("LongCatAudioDiTTransformer does not support standard layerwise casting memory tests yet.")
+
+    def test_layerwise_casting_training(self):
+        pytest.skip("LongCatAudioDiTTransformer does not support standard layerwise casting training tests yet.")
+
+    def test_group_offloading_with_layerwise_casting(self, *args, **kwargs):
+        pytest.skip("LongCatAudioDiTTransformer does not support combined group offloading and layerwise casting tests yet.")
+
+
+class TestLongCatAudioDiTTransformerCompile(LongCatAudioDiTTransformerTesterConfig, TorchCompileTesterMixin):
+    def test_torch_compile_repeated_blocks(self):
+        pytest.skip("LongCatAudioDiTTransformer does not define repeated blocks for regional compilation.")
+
+
+class TestLongCatAudioDiTTransformerAttention(LongCatAudioDiTTransformerTesterConfig, AttentionTesterMixin):
+    pass
+
+
+def test_longcat_audio_attention_uses_standard_self_attn_kwargs():
+    from diffusers.models.transformers.transformer_longcat_audio_dit import AudioDiTAttention
+
+    attn = AudioDiTAttention(q_dim=4, kv_dim=None, heads=1, dim_head=4, dropout=0.0, bias=False)
+
+    eye = torch.eye(4)
+    with torch.no_grad():
+        attn.to_q.weight.copy_(eye)
+        attn.to_k.weight.copy_(eye)
+        attn.to_v.weight.copy_(eye)
+        attn.to_out[0].weight.copy_(eye)
+
+    hidden_states = torch.tensor([[[1.0, 0.0, 0.0, 0.0], [0.5, 0.5, 0.5, 0.5]]])
+    attention_mask = torch.tensor([[True, False]])
 
+    output = attn(hidden_states=hidden_states, attention_mask=attention_mask)
 
-def test_longcat_audio_transformer_forward_shape():
-    model = LongCatAudioDiTTransformer(
-        dit_dim=64,
-        dit_depth=2,
-        dit_heads=4,
-        dit_text_dim=32,
-        latent_dim=8,
-        text_conv=False,
-    )
-    hidden_states = torch.randn(2, 16, 8)
-    encoder_hidden_states = torch.randn(2, 10, 32)
-    encoder_attention_mask = torch.ones(2, 10, dtype=torch.bool)
-    timestep = torch.tensor([1.0, 1.0])
-
-    output = model(
-        hidden_states=hidden_states,
-        encoder_hidden_states=encoder_hidden_states,
-        encoder_attention_mask=encoder_attention_mask,
-        timestep=timestep,
-    )
-
-    assert output.sample.shape == hidden_states.shape
-
-
-def test_longcat_audio_transformer_masked_forward():
-    model = LongCatAudioDiTTransformer(
-        dit_dim=64,
-        dit_depth=2,
-        dit_heads=4,
-        dit_text_dim=32,
-        latent_dim=8,
-        text_conv=False,
-    )
-    hidden_states = torch.randn(2, 16, 8)
-    encoder_hidden_states = torch.randn(2, 10, 32)
-    encoder_attention_mask = torch.tensor([[1] * 10, [1] * 6 + [0] * 4], dtype=torch.bool)
-    attention_mask = torch.tensor([[1] * 16, [1] * 9 + [0] * 7], dtype=torch.bool)
-    timestep = torch.tensor([1.0, 1.0])
-
-    output = model(
-        hidden_states=hidden_states,
-        encoder_hidden_states=encoder_hidden_states,
-        encoder_attention_mask=encoder_attention_mask,
-        timestep=timestep,
-        attention_mask=attention_mask,
-    )
-
-    assert output.sample.shape == hidden_states.shape
-    assert torch.all(output.sample[1, 9:] == 0)
+    assert torch.allclose(output[:, 1], torch.zeros_like(output[:, 1]))