`longcat_audio_dit` model/pipeline review

[hlky]

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:

diffusers/src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py

Lines 280 to 302 in 0f1abc4

	prompt_embeds, prompt_embeds_len = self.encode_prompt(normalized_prompts, device)
	duration_tensor = torch.full((batch_size,), duration, device=device, dtype=torch.long)
	mask = _lens_to_mask(duration_tensor)
	text_mask = _lens_to_mask(prompt_embeds_len, length=prompt_embeds.shape[1])
	if negative_prompt is None:
	negative_prompt_embeds = torch.zeros_like(prompt_embeds)
	negative_prompt_embeds_len = prompt_embeds_len
	negative_prompt_embeds_mask = text_mask
	else:
	if isinstance(negative_prompt, str):
	negative_prompt = [negative_prompt] * batch_size
	else:
	negative_prompt = list(negative_prompt)
	negative_prompt_embeds, negative_prompt_embeds_len = self.encode_prompt(negative_prompt, device)
	negative_prompt_embeds_mask = _lens_to_mask(
	negative_prompt_embeds_len, length=negative_prompt_embeds.shape[1]
	)
	latent_cond = torch.zeros(batch_size, duration, self.latent_dim, device=device, dtype=prompt_embeds.dtype)
	latents = self.prepare_latents(
	batch_size, duration, device, prompt_embeds.dtype, generator=generator, latents=latents
	)

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:

diffusers/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py

Lines 692 to 707 in 0f1abc4

	text_audio_duration_embeds.dtype,
	device,
	generator,
	latents,
	initial_audio_waveforms,
	num_waveforms_per_prompt,
	audio_channels=self.vae.config.audio_channels,
	)
	# 6. Prepare extra step kwargs
	extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
	# 7. Prepare rotary positional embedding
	rotary_embedding = get_1d_rotary_pos_embed(
	self.rotary_embed_dim,
	latents.shape[2] + audio_duration_embeds.shape[1],

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:

diffusers/src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py

Lines 136 to 153 in 0f1abc4

	def encode_prompt(self, prompt: str | list[str], device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
	if isinstance(prompt, str):
	prompt = [prompt]
	model_max_length = getattr(self.tokenizer, "model_max_length", 512)
	if not isinstance(model_max_length, int) or model_max_length <= 0 or model_max_length > 32768:
	model_max_length = 512
	text_inputs = self.tokenizer(
	prompt,
	padding="longest",
	truncation=True,
	max_length=model_max_length,
	return_tensors="pt",
	)
	input_ids = text_inputs.input_ids.to(device)
	attention_mask = text_inputs.attention_mask.to(device)
	with torch.no_grad():
	output = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
	prompt_embeds = output.last_hidden_state

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:

diffusers/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py

Lines 202 to 207 in 0f1abc4

	c_mults = [1] + (c_mults or [1, 2, 4, 8, 16])
	strides = list(strides or [2] * (len(c_mults) - 1))
	if len(strides) < len(c_mults) - 1:
	strides.extend([strides[-1] if strides else 2] * (len(c_mults) - 1 - len(strides)))
	else:
	strides = strides[: len(c_mults) - 1]

diffusers/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py

Lines 251 to 256 in 0f1abc4

	c_mults = [1] + (c_mults or [1, 2, 4, 8, 16])
	strides = list(strides or [2] * (len(c_mults) - 1))
	if len(strides) < len(c_mults) - 1:
	strides.extend([strides[-1] if strides else 2] * (len(c_mults) - 1 - len(strides)))
	else:
	strides = strides[: len(c_mults) - 1]

diffusers/src/diffusers/pipelines/longcat_audio_dit/pipeline_longcat_audio_dit.py

Lines 121 to 123 in 0f1abc4

	self.sample_rate = getattr(vae.config, "sample_rate", 24000)
	self.vae_scale_factor = getattr(vae.config, "downsampling_ratio", 2048)
	self.latent_dim = getattr(transformer.config, "latent_dim", 64)

diffusers/tests/pipelines/longcat_audio_dit/test_longcat_audio_dit.py

Lines 63 to 71 in 0f1abc4

	vae = LongCatAudioDiTVae(
	in_channels=1,
	channels=16,
	c_mults=[1, 2],
	strides=[2],
	latent_dim=8,
	encoder_latent_dim=16,
	downsampling_ratio=2,
	sample_rate=24000,

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:

diffusers/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py

Lines 984 to 991 in 0f1abc4

	height = int(audio_length_in_s / vocoder_upsample_factor)
	original_waveform_length = int(audio_length_in_s * self.vocoder.config.sampling_rate)
	if height % self.vae_scale_factor != 0:
	height = int(np.ceil(height / self.vae_scale_factor)) * self.vae_scale_factor
	logger.info(
	f"Audio length in seconds {audio_length_in_s} is increased to {height * vocoder_upsample_factor} "
	f"so that it can be handled by the model. It will be cut to {audio_length_in_s} after the "

diffusers/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py

Lines 1106 to 1108 in 0f1abc4

	audio = self.mel_spectrogram_to_waveform(mel_spectrogram)
	audio = audio[:, :original_waveform_length]

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:

diffusers/src/diffusers/models/transformers/transformer_longcat_audio_dit.py

Lines 230 to 255 in 0f1abc4

	class AudioDiTAttention(nn.Module, AttentionModuleMixin):
	def __init__(
	self,
	q_dim: int,
	kv_dim: int | None,
	heads: int,
	dim_head: int,
	dropout: float = 0.0,
	bias: bool = True,
	qk_norm: bool = False,
	eps: float = 1e-6,
	processor: AttentionModuleMixin | None = None,
	):
	super().__init__()
	kv_dim = q_dim if kv_dim is None else kv_dim
	self.heads = heads
	self.inner_dim = dim_head * heads
	self.to_q = nn.Linear(q_dim, self.inner_dim, bias=bias)
	self.to_k = nn.Linear(kv_dim, self.inner_dim, bias=bias)
	self.to_v = nn.Linear(kv_dim, self.inner_dim, bias=bias)
	self.qk_norm = qk_norm
	if qk_norm:
	self.q_norm = RMSNorm(self.inner_dim, eps=eps)
	self.k_norm = RMSNorm(self.inner_dim, eps=eps)
	self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, q_dim, bias=bias), nn.Dropout(dropout)])
	self.set_processor(processor or AudioDiTSelfAttnProcessor())

diffusers/src/diffusers/models/transformers/transformer_longcat_audio_dit.py

Lines 455 to 457 in 0f1abc4

	class LongCatAudioDiTTransformer(ModelMixin, ConfigMixin):
	_supports_gradient_checkpointing = False
	_repeated_blocks = ["AudioDiTBlock"]

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:

diffusers/src/diffusers/models/transformers/transformer_longcat_image.py

Lines 136 to 163 in 0f1abc4

	_default_processor_cls = LongCatImageAttnProcessor
	_available_processors = [
	LongCatImageAttnProcessor,
	]
	def __init__(
	self,
	query_dim: int,
	heads: int = 8,
	dim_head: int = 64,
	dropout: float = 0.0,
	bias: bool = False,
	added_kv_proj_dim: int | None = None,
	added_proj_bias: bool | None = True,
	out_bias: bool = True,
	eps: float = 1e-5,
	out_dim: int = None,
	context_pre_only: bool | None = None,
	pre_only: bool = False,
	elementwise_affine: bool = True,
	processor=None,
	):
	super().__init__()
	self.head_dim = dim_head
	self.inner_dim = out_dim if out_dim is not None else dim_head * heads
	self.query_dim = query_dim
	self.use_bias = bias

diffusers/src/diffusers/models/transformers/transformer_longcat_image.py

Lines 397 to 403 in 0f1abc4

	class LongCatImageTransformer2DModel(
	ModelMixin,
	ConfigMixin,
	PeftAdapterMixin,
	FromOriginalModelMixin,
	CacheMixin,
	AttentionMixin,

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:

diffusers/tests/pipelines/longcat_audio_dit/test_longcat_audio_dit.py

Lines 189 to 210 in 0f1abc4

	class LongCatAudioDiTPipelineSlowTests(unittest.TestCase):
	pipeline_class = LongCatAudioDiTPipeline
	def test_longcat_audio_pipeline_from_pretrained_real_local_weights(self):
	model_path = Path(
	os.getenv("LONGCAT_AUDIO_DIT_MODEL_PATH", "/data/models/meituan-longcat/LongCat-AudioDiT-1B")
	)
	tokenizer_path_env = os.getenv("LONGCAT_AUDIO_DIT_TOKENIZER_PATH")
	if tokenizer_path_env is None:
	raise unittest.SkipTest("LONGCAT_AUDIO_DIT_TOKENIZER_PATH is not set")
	tokenizer_path = Path(tokenizer_path_env)
	if not model_path.exists():
	raise unittest.SkipTest(f"LongCat-AudioDiT model path not found: {model_path}")
	if not tokenizer_path.exists():
	raise unittest.SkipTest(f"LongCat-AudioDiT tokenizer path not found: {tokenizer_path}")
	pipe = LongCatAudioDiTPipeline.from_pretrained(
	model_path,
	tokenizer=tokenizer_path,
	torch_dtype=torch.float16,
	local_files_only=True,

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:

diffusers/tests/pipelines/audioldm2/test_audioldm2.py

Line 594 in 0f1abc4

audioldm_pipe = AudioLDM2Pipeline.from_pretrained("cvssp/audioldm2")

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,
)