stable_audio model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Files/categories reviewed: target pipeline/model files, public lazy imports, top-level exports, config/loading/device-map behavior, dtype/device handling, offload-related tests, attention processor behavior, docs, examples, fast/nightly/slow test coverage.

Verification note: attempted .venv\Scripts\python.exe -m pytest tests/pipelines/stable_audio/test_stable_audio.py -q, but local test collection fails before Stable Audio tests run because this Windows torch build lacks torch._C._distributed_c10d while importing FSDP. Narrow reproduction snippets below were checked with .venv.

Duplicate-search status: searched GitHub Issues/PRs for stable_audio, StableAudioPipeline, StableAudioDiTModel device_map, StableAudioAttnProcessor2_0 set_attention_backend, and initial_audio_waveforms num_waveforms_per_prompt. Found related but not exact duplicates: #10861 for initial-audio scaling and #8989 for sequential offload testing. No exact duplicate found for the batch-order, _no_split_modules, attention-backend, dtype, or docs findings.

Issue 1: Batched initial audio is paired with the wrong prompt when generating multiple waveforms

Affected code:

Problem:
text_audio_duration_embeds is expanded per prompt as [prompt0, prompt0, prompt1, prompt1], but encoded initial audio is expanded with encoded_audio.repeat((num_waveforms_per_prompt, 1, 1)), producing [audio0, audio1, audio0, audio1]. For batched audio-to-audio with num_waveforms_per_prompt > 1, prompts and initial audio become misaligned.

Impact:
Users requesting multiple variations per prompt with batched initial_audio_waveforms condition some generations on another prompt's audio. Existing tests only assert output shape, so this does not get caught.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import StableAudioPipeline

class DummyLatentDist:
    def __init__(self, sample):
        self._sample = sample
    def sample(self, generator=None):
        return self._sample

class DummyVAE:
    hop_length = 1
    def encode(self, audio):
        return SimpleNamespace(latent_dist=DummyLatentDist(audio[:, :1, :]))

pipe = StableAudioPipeline.__new__(StableAudioPipeline)
pipe.scheduler = SimpleNamespace(init_noise_sigma=0.0)
pipe.transformer = SimpleNamespace(config=SimpleNamespace(sample_size=2))
pipe.vae = DummyVAE()

initial_audio = torch.tensor([[[10.0, 10.0]], [[20.0, 20.0]]])
latents = StableAudioPipeline.prepare_latents(
    pipe, batch_size=4, num_channels_vae=1, sample_size=2,
    dtype=torch.float32, device=torch.device("cpu"),
    generator=torch.Generator().manual_seed(0),
    initial_audio_waveforms=initial_audio,
    num_waveforms_per_prompt=2,
    audio_channels=1,
)
print(latents[:, 0, 0].tolist())  # [10.0, 20.0, 10.0, 20.0], expected [10.0, 10.0, 20.0, 20.0]

Relevant precedent:
repeat_interleave(..., dim=0) is the common pattern for per-prompt expansion, e.g. qwenimage modular inputs.

Suggested fix:

encoded_audio = encoded_audio.repeat_interleave(num_waveforms_per_prompt, dim=0)

Issue 2: StableAudioDiTModel cannot be loaded with device_map, despite docs using device_map="balanced"

Affected code:

Problem:
StableAudioDiTModel sets _supports_gradient_checkpointing = True but does not define _no_split_modules. Diffusers model loading raises for device_map="balanced"/"auto" unless _no_split_modules is implemented. The Stable Audio quantization docs currently show StableAudioPipeline.from_pretrained(..., device_map="balanced"), which is not supported by the transformer class.

Impact:
The documented quantized loading path is broken for the Stable Audio transformer, and users cannot use Diffusers device-map placement for the model.

Reproduction:

from diffusers import StableAudioDiTModel

model = StableAudioDiTModel(
    sample_size=4, in_channels=3, num_layers=1,
    attention_head_dim=4, num_attention_heads=2,
    num_key_value_attention_heads=2, out_channels=3,
    cross_attention_dim=4, time_proj_dim=8,
    global_states_input_dim=8, cross_attention_input_dim=4,
)

try:
    print(model._get_no_split_modules("balanced"))
except Exception as e:
    print(type(e).__name__, str(e).splitlines()[0])
# ValueError StableAudioDiTModel does not support `device_map='balanced'`.

Relevant precedent:

Suggested fix:

class StableAudioDiTModel(ModelMixin, AttentionMixin, ConfigMixin):
    _supports_gradient_checkpointing = True
    _no_split_modules = ["StableAudioDiTBlock"]

Issue 3: Stable Audio attention ignores set_attention_backend

Affected code:

Problem:
StableAudioAttnProcessor2_0 lives in the shared attention processor file, has no _attention_backend / _parallel_config fields, and calls F.scaled_dot_product_attention directly. ModelMixin.set_attention_backend() only updates processors with _attention_backend, so Stable Audio processors remain unchanged.

Impact:
Users cannot select Flash/Sage/Flex/native backend behavior for Stable Audio even though StableAudioDiTModel inherits AttentionMixin. This also leaves Stable Audio outside the newer attention-dispatch and context-parallel patterns.

Reproduction:

from diffusers import StableAudioDiTModel

model = StableAudioDiTModel(
    sample_size=4, in_channels=3, num_layers=1,
    attention_head_dim=4, num_attention_heads=2,
    num_key_value_attention_heads=2, out_channels=3,
    cross_attention_dim=4, time_proj_dim=8,
    global_states_input_dim=8, cross_attention_input_dim=4,
)
model.set_attention_backend("native")
print([(type(p).__name__, hasattr(p, "_attention_backend")) for p in model.attn_processors.values()])
# [('StableAudioAttnProcessor2_0', False), ('StableAudioAttnProcessor2_0', False)]

Relevant precedent:

Suggested fix:
Refactor the Stable Audio attention processor to the model-file attention pattern: define processor state fields, use dispatch_attention_fn, and keep Q/K/V in the (batch, sequence, heads, head_dim) layout expected by the dispatcher. This is a moderate refactor because the current implementation uses (batch, heads, sequence, head_dim) around RoPE.

Issue 4: User-provided latents keep their original dtype

Affected code:

Problem:
When latents is supplied, prepare_latents() only does latents.to(device). It does not cast to the dtype selected for the pipeline call. With a half-precision Stable Audio transformer, float32 user latents are forwarded into half-precision Conv/Linear layers.

Impact:
Mixed-precision calls can fail at runtime or run with an unintended latent dtype. This is especially relevant because the slow test path supplies precomputed latents.

Reproduction:

from types import SimpleNamespace
import torch
from diffusers import StableAudioPipeline

pipe = StableAudioPipeline.__new__(StableAudioPipeline)
pipe.scheduler = SimpleNamespace(init_noise_sigma=1.0)

latents = torch.randn(1, 3, 4, dtype=torch.float32)
out = StableAudioPipeline.prepare_latents(
    pipe, batch_size=1, num_channels_vae=3, sample_size=4,
    dtype=torch.float16, device=torch.device("cpu"),
    generator=None, latents=latents,
)
print(out.dtype)  # torch.float32, expected torch.float16

Relevant precedent:
Newer pipelines often recast latents to the active latent/model dtype before denoising, for example QwenImage and Flux paths recast around denoising.

Suggested fix:

if latents is None:
    latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
    latents = latents.to(device=device, dtype=dtype)

Issue 5: Stable Audio has no model-level tests and no @slow tests

Affected code:

Problem:
The family has fast pipeline tests and a @nightly integration test, but no tests/models/transformers/test_models_stable_audio*.py coverage and no @slow Stable Audio test. The pipeline also skips sequential offload tests and encode-prompt isolation. The sequential offload skip is already related to open issue #8989: #8989

Impact:
Model serialization/loading, attention backend behavior, _no_split_modules/device-map support, compile behavior, and model-level attention masks are not covered by the standard model test mixins. Missing slow coverage also means the non-nightly slow suite does not exercise the published checkpoint.

Reproduction:

from pathlib import Path

model_tests = list(Path("tests/models/transformers").glob("*stable*audio*.py"))
pipeline_test = Path("tests/pipelines/stable_audio/test_stable_audio.py").read_text()

print(model_tests)          # []
print("@slow" in pipeline_test)    # False
print("@nightly" in pipeline_test) # True

Relevant precedent:

Suggested fix:
Add a StableAudioDiTModel model tester using ModelTesterMixin, AttentionTesterMixin, and compile/memory coverage where supported. Add or mark a published-checkpoint pipeline test with @slow so Stable Audio is covered outside nightly-only CI. Keep #8989 referenced for sequential offload until that behavior is fixed or explicitly unsupported.

Issue 6: Stable Audio docs claim waveform scoring that the pipeline does not implement

Affected code:

Problem:
The docs say num_waveforms_per_prompt > 1 performs automatic scoring and ranks outputs by prompt similarity. StableAudioPipeline has no scoring component or score_waveforms() path; it simply returns generated audio in batch order. This looks copied from AudioLDM2/MusicLDM behavior.

Impact:
Users are told generated Stable Audio waveforms are ranked when they are not.

Reproduction:

import inspect
from diffusers import StableAudioPipeline

print(hasattr(StableAudioPipeline, "score_waveforms"))                  # False
print("score_waveforms" in inspect.getsource(StableAudioPipeline.__call__))  # False

Relevant precedent:
AudioLDM2 implements waveform scoring:

Suggested fix:
Remove the scoring/ranking sentence from the Stable Audio docs, or implement an actual scoring component before documenting ranking behavior.