audioldm2 model/pipeline review

`audioldm2` model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against .ai/review-rules.md and all present referenced rule files. AGENTS.md was referenced by the rules but is not present in this checkout.

Duplicate search: checked GitHub issues/PRs for audioldm2, AudioLDM2ProjectionModel, AudioLDM2UNet2DConditionModel, pipeline_audioldm2, modeling_audioldm2, cross_attention_kwargs, gradient-checkpointing masks, scoring, and cross_attention_dim IndexError. I did not find likely duplicates for the issues below. Existing #12630 / PR #13111 cover a different GPT2Model AttributeError.

Issue 1: Projection mask fallback crashes when only one mask is provided

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L143-L153

Problem:
AudioLDM2ProjectionModel.forward() calls new_ones((hidden_states[:2])) and new_ones((hidden_states_1[:2])). Those are tensor slices, not shape tuples, so direct projection-model use crashes when one encoder mask is provided and the other is omitted. The first branch is also placed after concatenating hidden_states, so even changing it to hidden_states.shape[:2] there would create the wrong sequence length.

Impact:
Users providing one precomputed attention mask cannot use the projection model directly, and pipeline paths that mix precomputed embeddings/masks are fragile.

Reproduction:

import torch
from diffusers import AudioLDM2ProjectionModel

model = AudioLDM2ProjectionModel(text_encoder_dim=3, text_encoder_1_dim=4, langauge_model_dim=5)
h0 = torch.randn(2, 1, 3)
h1 = torch.randn(2, 7, 4)
mask1 = torch.ones(2, 7, dtype=torch.long)

model(hidden_states=h0, hidden_states_1=h1, attention_mask=None, attention_mask_1=mask1)
# TypeError: new_ones(): argument 'size' must be tuple of ints

Relevant precedent:
Use tensor .shape[:2] for mask creation, as the pipeline already does for default prompt masks.

Suggested fix:

# before concatenating hidden_states and hidden_states_1
if attention_mask is None and attention_mask_1 is not None:
    attention_mask = attention_mask_1.new_ones(hidden_states.shape[:2])
elif attention_mask is not None and attention_mask_1 is None:
    attention_mask_1 = attention_mask.new_ones(hidden_states_1.shape[:2])

hidden_states = torch.cat([hidden_states, hidden_states_1], dim=1)

Issue 2: Omitted second encoder states ignore the first encoder mask

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1039-L1044
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1200-L1205
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1352-L1357

Problem:
The blocks first replace encoder_hidden_states_1=None with encoder_hidden_states, then decide whether to fallback encoder_attention_mask_1 based on the already-mutated encoder_hidden_states_1. As a result, the second cross-attention stream uses the first hidden states but drops the first mask.

Impact:
Calling AudioLDM2UNet2DConditionModel with only encoder_hidden_states and encoder_attention_mask produces different results than explicitly passing the same states/mask as stream 1.

Reproduction:

import torch
from diffusers import AudioLDM2UNet2DConditionModel

torch.manual_seed(0)
model = AudioLDM2UNet2DConditionModel(
    sample_size=8, in_channels=4, out_channels=4, block_out_channels=(8,),
    layers_per_block=1, norm_num_groups=1,
    down_block_types=("CrossAttnDownBlock2D",),
    up_block_types=("CrossAttnUpBlock2D",),
    cross_attention_dim=((None, 8, 8),),
    attention_head_dim=1,
).eval()

sample = torch.randn(1, 4, 8, 8)
encoder = torch.randn(1, 5, 8)
mask = torch.tensor([[1, 1, 1, 0, 0]])

with torch.no_grad():
    implicit = model(sample, 1, encoder_hidden_states=encoder, encoder_attention_mask=mask).sample
    explicit = model(
        sample, 1, encoder_hidden_states=encoder, encoder_attention_mask=mask,
        encoder_hidden_states_1=encoder, encoder_attention_mask_1=mask,
    ).sample

print((implicit - explicit).abs().max())
# tensor(0.1561...)

Relevant precedent:
The intended fallback is visible from the code itself: stream 1 defaults to stream 0 when omitted.

Suggested fix:

if encoder_hidden_states_1 is None:
    encoder_hidden_states_1 = encoder_hidden_states
    if encoder_attention_mask_1 is None:
        encoder_attention_mask_1 = encoder_attention_mask

Issue 3: `cross_attention_kwargs` is accepted but ignored

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py#L1069-L1077
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1081-L1087
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1241-L1247
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1399-L1405

Problem:
AudioLDM2Pipeline.__call__() exposes cross_attention_kwargs, but does not pass it to the UNet. The UNet block normal paths also do not pass cross_attention_kwargs to Transformer2DModel.

Impact:
Custom attention processors and LoRA-style attention kwargs silently do nothing in normal inference.

Reproduction:

import torch
from diffusers import AudioLDM2UNet2DConditionModel
from diffusers.models.attention_processor import AttnProcessor

class MarkerProcessor(AttnProcessor):
    def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, temb=None, marker=None, **kwargs):
        if marker:
            raise RuntimeError("cross_attention_kwargs propagated")
        return super().__call__(attn, hidden_states, encoder_hidden_states, attention_mask, temb, **kwargs)

model = AudioLDM2UNet2DConditionModel(
    sample_size=8, in_channels=4, out_channels=4, block_out_channels=(8,),
    layers_per_block=1, norm_num_groups=1,
    down_block_types=("CrossAttnDownBlock2D",),
    up_block_types=("CrossAttnUpBlock2D",),
    cross_attention_dim=(8,),
    attention_head_dim=1,
).eval()
model.set_attn_processor(MarkerProcessor())

with torch.no_grad():
    model(torch.randn(1, 4, 8, 8), 1, encoder_hidden_states=torch.randn(1, 5, 8), cross_attention_kwargs={"marker": True})

print("marker was ignored")

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_blocks.py#L1257-L1277

Suggested fix:

# pipeline denoise call
noise_pred = self.unet(
    latent_model_input,
    t,
    encoder_hidden_states=generated_prompt_embeds,
    encoder_hidden_states_1=prompt_embeds,
    encoder_attention_mask_1=attention_mask,
    cross_attention_kwargs=cross_attention_kwargs,
    return_dict=False,
)[0]

# all AudioLDM2 cross-attn block normal paths
hidden_states = self.attentions[i * num_attention_per_layer + idx](
    hidden_states,
    encoder_hidden_states=forward_encoder_hidden_states,
    cross_attention_kwargs=cross_attention_kwargs,
    attention_mask=attention_mask,
    encoder_attention_mask=forward_encoder_attention_mask,
    return_dict=False,
)[0]

Issue 4: Gradient checkpointing misroutes `Transformer2DModel` arguments

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1059-L1068
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1219-L1228
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py#L1377-L1386

Problem:
The gradient-checkpointing branch calls Transformer2DModel positionally using a stale signature. With the current signature, cross_attention_kwargs, attention_mask, and encoder_attention_mask are shifted into the wrong parameters.

Impact:
_supports_gradient_checkpointing=True is advertised, but training with gradient checkpointing and encoder masks can crash or use the wrong masks.

Reproduction:

import torch
from diffusers import AudioLDM2UNet2DConditionModel

model = AudioLDM2UNet2DConditionModel(
    sample_size=8, in_channels=4, out_channels=4, block_out_channels=(8,),
    layers_per_block=1, norm_num_groups=1,
    down_block_types=("CrossAttnDownBlock2D",),
    up_block_types=("CrossAttnUpBlock2D",),
    cross_attention_dim=(8,),
    attention_head_dim=1,
)
model.enable_gradient_checkpointing()

sample = torch.randn(1, 4, 8, 8, requires_grad=True)
encoder = torch.randn(1, 5, 8, requires_grad=True)
mask = torch.tensor([[1, 1, 1, 0, 0]])

model(sample, 1, encoder_hidden_states=encoder, encoder_attention_mask=mask)
# RuntimeError: The size of tensor a (64) must match the size of tensor b (69)...

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/unets/unet_2d_blocks.py#L1257-L1277

Suggested fix:
Do not checkpoint-call Transformer2DModel positionally. Follow the regular UNet block pattern: checkpoint the ResNet, then call the attention module with keyword arguments so the nested Transformer2DModel handles its own checkpointing.

hidden_states = self._gradient_checkpointing_func(self.resnets[i], hidden_states, temb)
hidden_states = self.attentions[i * num_attention_per_layer + idx](
    hidden_states,
    encoder_hidden_states=forward_encoder_hidden_states,
    cross_attention_kwargs=cross_attention_kwargs,
    attention_mask=attention_mask,
    encoder_attention_mask=forward_encoder_attention_mask,
    return_dict=False,
)[0]

Issue 5: Automatic scoring ranks across the full batch, not per prompt

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py#L724-L728

Problem:
score_waveforms() sorts each prompt against every generated waveform, then index-selects globally. For batched prompts with num_waveforms_per_prompt > 1, a prompt can select waveforms generated for another prompt.

Impact:
The returned audio order can mix prompts, so users may receive a waveform for the wrong prompt after automatic scoring.

Reproduction:

import torch
from types import SimpleNamespace
from diffusers import AudioLDM2Pipeline

class Inputs(dict):
    def to(self, device): return self

class Tokenizer:
    def __call__(self, text, return_tensors=None, padding=None):
        return Inputs(input_ids=torch.ones(len(text), 2, dtype=torch.long))

class FeatureExtractor:
    sampling_rate = 16000
    def __call__(self, audio, return_tensors=None, sampling_rate=None):
        return SimpleNamespace(input_features=torch.zeros(len(audio), 1, 4))

class TextEncoder:
    def __call__(self, **inputs):
        return SimpleNamespace(logits_per_text=torch.tensor([[0.1, 0.2, 0.9, 0.8], [0.7, 0.6, 0.5, 0.4]]))

pipe = AudioLDM2Pipeline.__new__(AudioLDM2Pipeline)
pipe.tokenizer = Tokenizer()
pipe.feature_extractor = FeatureExtractor()
pipe.text_encoder = TextEncoder()
pipe.vocoder = SimpleNamespace(config=SimpleNamespace(sampling_rate=16000))

audio = torch.arange(16, dtype=torch.float32).view(4, 4)
ranked = pipe.score_waveforms(["prompt a", "prompt b"], audio, num_waveforms_per_prompt=2, device="cpu", dtype=torch.float32)
print(ranked[:, 0].tolist())
# [8.0, 12.0, 0.0, 4.0] selects rows 2/3 for prompt 0 and rows 0/1 for prompt 1

Relevant precedent:
No good in-repo precedent found; MusicLDMPipeline copies this method and appears to carry the same behavior.

Suggested fix:

batch_size = len(text) if isinstance(text, list) else 1
selected = []
for i in range(batch_size):
    start = i * num_waveforms_per_prompt
    end = start + num_waveforms_per_prompt
    local_scores = logits_per_text[i, start:end]
    local_indices = torch.argsort(local_scores, descending=True) + start
    selected.append(local_indices)

indices = torch.cat(selected)
audio = torch.index_select(audio, 0, indices.cpu())

Issue 6: Tuple `cross_attention_dim` length is not validated

Problem:
The constructor validates mismatched cross_attention_dim only when it is a list, but the public type allows tuples. A short tuple falls through and later raises an IndexError.

Impact:
Invalid configs fail with an opaque internal error instead of the intended config validation error. This is especially confusing for from_config() / custom checkpoint users.

Reproduction:

from diffusers import AudioLDM2UNet2DConditionModel

AudioLDM2UNet2DConditionModel(
    block_out_channels=(8, 16),
    layers_per_block=1,
    norm_num_groups=1,
    down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),
    up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),
    cross_attention_dim=(8,),
)
# IndexError: tuple index out of range

Relevant precedent:
The same constructor already validates tuple-like attention_head_dim and layers_per_block.

Suggested fix:

if not isinstance(cross_attention_dim, int) and len(cross_attention_dim) != len(down_block_types):
    raise ValueError(
        f"Must provide the same number of `cross_attention_dim` as `down_block_types`. "
        f"`cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
    )

if not isinstance(transformer_layers_per_block, int) and len(transformer_layers_per_block) != len(down_block_types):
    raise ValueError(
        f"Must provide the same number of `transformer_layers_per_block` as `down_block_types`. "
        f"`transformer_layers_per_block`: {transformer_layers_per_block}. `down_block_types`: {down_block_types}."
    )

Coverage status

Fast and slow AudioLDM2 tests exist:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/audioldm2/test_audioldm2.py#L62-L650

Slow coverage is present for cvssp/audioldm2, LMS, cvssp/audioldm2-large, and anhnct/audioldm2_gigaspeech. Several offload / isolation tests remain skipped:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/audioldm2/test_audioldm2.py#L533-L550

Local pytest collection could not complete in .venv because the installed torch build is missing torch._C._distributed_c10d; the standalone repro snippets above were run successfully under .venv.

Description