stable_audio model/pipeline review

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py#L485-L487
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py#L670-L680

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/stable_audio_transformer.py#L206-L208
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/api/pipelines/stable_audio.md#L66-L72

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L565-L566
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_wan.py#L546-L548

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/stable_audio_transformer.py#L24-L24
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/stable_audio_transformer.py#L105-L121
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_processor.py#L2991-L3103

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L75-L125

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py#L439-L445

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/stable_audio/test_stable_audio.py#L413-L423
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/stable_audio/test_stable_audio.py#L426-L478

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: https://github.com/huggingface/diffusers/issues/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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/transformers/test_models_transformer_longcat_audio_dit.py#L84-L101

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/docs/source/en/api/pipelines/stable_audio.md#L33-L37
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py#L490-L764

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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/audioldm2/pipeline_audioldm2.py#L707-L727

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

Bắt đầu với src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py, src/diffusers/models/transformers/stable_audio_transformer.py và bộ xử lý attention dùng chung được tham chiếu trong issue. Tái hiện các trường hợp audio ban đầu theo batch và dtype, sau đó kiểm tra các bài kiểm thử pipeline Stable Audio hiện có và các tiền lệ trong bài kiểm thử model. Hoàn tất khi các hành vi được báo cáo đã được bao phủ hoặc sửa, các kỳ vọng về device-map và backend attention đã được xác minh, đồng thời tài liệu khớp với pipeline đã triển khai.