`longcat_audio_dit` model/pipeline review

longcat_audio_dit model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules. Reviewed target pipeline/model files, public imports/lazy loading, docs, fast/slow tests, dtype/device handling, attention processor behavior, scheduler usage, serialization/loading paths, and offload coverage.

Duplicate search: searched GitHub Issues/PRs for LongCatAudioDiT, LongCat-AudioDiT, affected classes, and the failure modes below. No duplicates found for these findings. Related open PR found: #13525, but it only covers negative prompt normalization.

Local test note: targeted pytest collection was blocked in this .venv by ModuleNotFoundError: torch._C._distributed_c10d; direct .venv Python repros below were run instead.

Issue 1: Pipeline uses text-encoder dtype for transformer latents

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py#L280-L302

Problem:
LongCatAudioDiTPipeline derives latent, latent condition, and timestep dtype from prompt_embeds.dtype. If the text encoder stays float32 while the transformer is loaded in bfloat16/float16, the first transformer linear layer receives float activations with bf16/fp16 weights and fails.

Impact:
Mixed-dtype component loading is a normal diffusers use case. This makes the pipeline fragile when users keep the text encoder in fp32 or override component dtypes.

Reproduction:

import torch
from types import SimpleNamespace
from diffusers import LongCatAudioDiTPipeline, LongCatAudioDiTTransformer, LongCatAudioDiTVae

class Tok:
    model_max_length = 4
    def __call__(self, prompt, **kwargs):
        b = len(prompt)
        return SimpleNamespace(input_ids=torch.ones(b, 4, dtype=torch.long), attention_mask=torch.ones(b, 4))

class Text(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = torch.nn.Embedding(8, 32)
    def forward(self, input_ids, attention_mask, output_hidden_states=False):
        h = self.emb(input_ids)  # float32
        return SimpleNamespace(last_hidden_state=h, hidden_states=(h,))

pipe = LongCatAudioDiTPipeline(
    vae=LongCatAudioDiTVae(in_channels=1, channels=4, c_mults=[1], strides=[2], latent_dim=8, encoder_latent_dim=16),
    text_encoder=Text(),
    tokenizer=Tok(),
    transformer=LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False).to(torch.bfloat16),
)
pipe.set_progress_bar_config(disable=True)
pipe("x", audio_duration_s=0.1, num_inference_steps=1, guidance_scale=1.0, output_type="latent")

Relevant precedent:
Stable Audio prepares latents from the denoiser/latent dtype rather than text embedding dtype:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py#L692-L707

Suggested fix:

latent_dtype = self.transformer.dtype

prompt_embeds = prompt_embeds.to(dtype=latent_dtype)
negative_prompt_embeds = negative_prompt_embeds.to(dtype=latent_dtype)

latent_cond = torch.zeros(batch_size, duration, self.latent_dim, device=device, dtype=latent_dtype)
latents = self.prepare_latents(batch_size, duration, device, latent_dtype, generator=generator, latents=latents)

curr_t = (t / self.scheduler.config.num_train_timesteps).expand(batch_size).to(dtype=latent_dtype)

Issue 2: encode_prompt disables gradients internally

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py#L136-L153

Problem:
encode_prompt wraps the text encoder call in torch.no_grad(). The pipeline rules state __call__ should own inference no-grad; helper methods should remain usable with gradients for training, prompt optimization, and direct embedding workflows.

Impact:
Calling pipe.encode_prompt(...) directly under torch.enable_grad() still returns detached embeddings.

Reproduction:

import torch
from types import SimpleNamespace
from diffusers import LongCatAudioDiTPipeline, LongCatAudioDiTTransformer, LongCatAudioDiTVae

class Tok:
    model_max_length = 4
    def __call__(self, prompt, **kwargs):
        return SimpleNamespace(input_ids=torch.ones(1, 4, dtype=torch.long), attention_mask=torch.ones(1, 4))

class Text(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = torch.nn.Embedding(8, 32)
    def forward(self, input_ids, attention_mask, output_hidden_states=False):
        h = self.emb(input_ids)
        return SimpleNamespace(last_hidden_state=h, hidden_states=(h,))

pipe = LongCatAudioDiTPipeline(
    vae=LongCatAudioDiTVae(in_channels=1, channels=4, c_mults=[1], strides=[2], latent_dim=8, encoder_latent_dim=16),
    text_encoder=Text(),
    tokenizer=Tok(),
    transformer=LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False),
)
with torch.enable_grad():
    embeds, _ = pipe.encode_prompt("x", torch.device("cpu"))
print(embeds.requires_grad, embeds.grad_fn)  # False, None

Relevant precedent:
The pipeline rule explicitly calls this out in .ai/pipelines.md.

Suggested fix:

# Remove the inner no_grad block.
output = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)

Issue 3: VAE downsampling ratio can disagree with actual strides

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py#L202-L207
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py#L251-L256
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py#L121-L123
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/longcat_audio_dit/test_longcat_audio_dit.py#L63-L71

Problem:
The pipeline trusts vae.config.downsampling_ratio, but the VAE’s real temporal scale is the product of normalized strides. The fast test fixture already sets downsampling_ratio=2 while its VAE decodes 10 latent frames to 40 samples, i.e. an actual ratio of 4.

Impact:
Duration calculation can be wrong for custom configs and tests can pass while exercising the wrong waveform length.

Reproduction:

import torch
from diffusers import LongCatAudioDiTVae

vae = LongCatAudioDiTVae(
    in_channels=1, channels=16, c_mults=[1, 2], strides=[2],
    latent_dim=8, encoder_latent_dim=16, downsampling_ratio=2,
)
decoded = vae.decode(torch.zeros(1, 8, 10)).sample
print("configured ratio:", vae.config.downsampling_ratio)
print("actual ratio:", decoded.shape[-1] // 10)  # 4

Relevant precedent:
AudioLDM2 rounds internal lengths to the VAE scale and cuts back to the requested waveform length:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py#L984-L991
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py#L1106-L1108

Suggested fix:
Validate or derive the ratio from normalized strides, and update the fast fixture to match. For defaults, use the published checkpoint stride pattern [2, 4, 4, 8, 8] if downsampling_ratio=2048 remains the default.

actual_downsampling_ratio = math.prod(strides)
if downsampling_ratio != actual_downsampling_ratio:
    raise ValueError(
        f"`downsampling_ratio` ({downsampling_ratio}) must match product(strides) ({actual_downsampling_ratio})."
    )

Issue 4: Attention module is only partially wired into diffusers attention APIs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_longcat_audio_dit.py#L230-L255
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_longcat_audio_dit.py#L455-L457

Problem:
AudioDiTAttention inherits AttentionModuleMixin but does not define _default_processor_cls, _available_processors, or self.use_bias. Its inherited fuse_projections() crashes. The transformer also does not inherit AttentionMixin, so model-level attn_processors / set_attn_processor APIs are absent and attention tests mostly skip.

Impact:
Users cannot manage attention processors through standard model APIs, and direct QKV fusion on the attention module raises.

Reproduction:

from diffusers import LongCatAudioDiTTransformer

model = LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False)
print(hasattr(model, "attn_processors"))  # False

try:
    model.blocks[0].self_attn.fuse_projections()
except Exception as e:
    print(type(e).__name__, e)  # AttributeError: no attribute 'use_bias'

Relevant precedent:
LongCat Image wires the same mixin pattern fully:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_longcat_image.py#L136-L163
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_longcat_image.py#L397-L403

Suggested fix:

from ..attention import AttentionMixin, AttentionModuleMixin

class AudioDiTAttention(nn.Module, AttentionModuleMixin):
    _default_processor_cls = AudioDiTSelfAttnProcessor
    _available_processors = [AudioDiTSelfAttnProcessor, AudioDiTCrossAttnProcessor]
    _supports_qkv_fusion = False

    def __init__(..., bias: bool = True, ...):
        ...
        self.use_bias = bias
        self.set_processor(processor or self._default_processor_cls())

class LongCatAudioDiTTransformer(ModelMixin, AttentionMixin, ConfigMixin):
    ...

Issue 5: Slow test passes a tokenizer path as a component override

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/longcat_audio_dit/test_longcat_audio_dit.py#L189-L210

Problem:
The slow test calls LongCatAudioDiTPipeline.from_pretrained(..., tokenizer=tokenizer_path). In from_pretrained, a component kwarg is treated as an already-instantiated component, not as a path to load. This raises a type error when the env vars are set. The slow test exists, but it is not effective as written.

Impact:
Slow coverage will skip by default and fail when configured, so the real checkpoint path is not covered.

Reproduction:

import tempfile
from pathlib import Path
from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
from diffusers import LongCatAudioDiTPipeline, LongCatAudioDiTTransformer, LongCatAudioDiTVae

tok = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
pipe = LongCatAudioDiTPipeline(
    vae=LongCatAudioDiTVae(in_channels=1, channels=4, c_mults=[1], strides=[2], latent_dim=8, encoder_latent_dim=16),
    text_encoder=UMT5EncoderModel(UMT5Config(d_model=32, num_layers=1, num_heads=4, d_ff=64, vocab_size=tok.vocab_size)),
    tokenizer=tok,
    transformer=LongCatAudioDiTTransformer(dit_dim=64, dit_depth=1, dit_heads=4, dit_text_dim=32, latent_dim=8, text_conv=False),
)
with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as d:
    pipe.save_pretrained(d)
    LongCatAudioDiTPipeline.from_pretrained(d, tokenizer=Path(d) / "tokenizer", local_files_only=True)

Relevant precedent:
Other audio slow tests load public repos directly rather than passing component paths as overrides:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/audioldm2/test_audioldm2.py#L594

Suggested fix:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
pipe = LongCatAudioDiTPipeline.from_pretrained(
    model_path,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    local_files_only=True,
)

Start with the affected pipeline, VAE, transformer, and test files named in the issue, then run the targeted longcat_audio_dit tests; collection is currently blocked by the reported torch._C._distributed_c10d error. Use the supplied reproductions and compare the Stable Audio, AudioLDM2, and LongCat Image precedents. Done means the five reported behaviors are corrected and fast and slow coverage exercises the intended paths.