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:
```python
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:
```python
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:
```python
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:
```python
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:
```python
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:
```python
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:
```python
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:
```python
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:
```python
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.

Start with src/diffusers/pipelines/stable_audio/pipeline_stable_audio.py, src/diffusers/models/transformers/stable_audio_transformer.py, and the shared attention processor referenced in the issue. Reproduce the batched initial-audio and dtype cases, then inspect the existing Stable Audio pipeline tests and model-test precedents. Done means the reported behaviors are covered or corrected, device-map and attention-backend expectations are verified, and the documentation matches the implemented pipeline.

Scritto dal modello di indicizzazione a partire dal testo della issue.