Feat -- Stable Audio 3 by buffett0323 · Pull Request #14119 · huggingface/diffusers

buffett0323 · 2026-07-03T23:05:16Z

What does this PR do?

Adds Stable Audio 3 Medium to Diffusers.

Before submitting

Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
- Did you read the Coding with AI agents guide?
- Did you self-review the diff against .ai/review-rules.md?
Did you read the contributor guideline?
Did you read our philosophy doc? (important for complex PRs)
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case. Support for Stable Audio 3 Medium #13793
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?
Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

…n mask The diffusers Stable Audio 3 pipeline produced noise instead of music because the cross-attention conditioning was built differently from the reference: - The reference text conditioner uses padding_mode="learned": padded text positions (~245 of 256 for a short prompt) are filled with a trained `padding_embedding`, and the DiT attends to *all* positions (its cross-attention mask is intentionally disabled). Our pipeline instead zeroed padded positions and masked them out, wiping ~95% of the conditioning signal. Changes: - transformer_stable_audio3: add learned `prompt_padding_embedding`; in forward, replace masked text positions with it (in cond_token_dim space, before to_cond_embed) then attend to the full context, matching the reference. - pipeline_stable_audio_3: stop zeroing padded positions (the DiT now handles them); default `silence_padding_duration` to 6.0 (reference headroom default). - convert_..._to_diffusers: convert `conditioner.conditioners.prompt.padding_embedding` into the DiT. - scheduling_ping_pong: pin schedule endpoints (sigmas[0]=1.0, sigmas[-1]=0.0) to match the reference LogSNRShift endpoint preservation. - run_..._inference: coerce float16->float32 on CPU. - tests: update transformer state-dict expectations (523 tensors, add prompt_padding_embedding) and ping-pong sigma endpoint assertions. Verified against the reference with identical noise + identical conditioning: full 8-step ping-pong trajectories agree to 4e-6 (final latent). Removed the dev-only parity scripts (verify_*_dit_parity, verify_*_vae_parity) from the tree. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

sergereview

🤗 Serge says:

Adds Stable Audio 3 end-to-end (SAME VAE, DiT, duration embedder, two schedulers, text-to-audio + inpaint pipelines, conversion script, docs, and tests). The modeling code is clean and well-documented, but there are a few correctness bugs, repo-convention violations, and dead-code issues to resolve before merging.

Correctness

StableAudio3InpaintPipeline: passing mask_start_seconds/mask_end_seconds as ints (e.g. mask_start_seconds=2) crashes with a TypeError — the code only checks isinstance(..., float) and then calls list(2).
StableAudio3InpaintPipeline.__call__ hardcodes num_inference_steps: int = 8, while the base StableAudio3Pipeline deliberately defaults to None and falls back to the scheduler config (8 for ping-pong, 100 for Euler — the base pipeline even has a test asserting this). The inpaint default will silently under-sample the non-distilled Euler checkpoint.
The inpaint pipeline silently falls back to plain text-to-audio when no audio/mask is given. This contradicts "do not guess user intent and silently correct behavior" — it should raise instead of degrading behind a warning.
encode_duration accepts per-prompt duration lists, but check_inputs (duration <= 0) and the latent-length computation (duration + silence_padding_duration) assume a scalar — a list either crashes or silently misbehaves. Remove the list path or make it consistent.

Packaging / imports

src/diffusers/pipelines/stable_audio_3/__init__.py uses plain eager imports instead of the lazy-module + OptionalDependencyNotAvailable pattern used by every sibling package (e.g. stable_audio/__init__.py). Since pipeline_stable_audio_3.py imports transformers at module level, importing this subpackage without transformers raises instead of surfacing the dummy objects.

Pipeline API

StableAudio3Pipeline.__call__ uses the deprecated callback/callback_steps pattern; new pipelines should implement callback_on_step_end.
EXAMPLE_DOC_STRING is defined but __call__ lacks the @replace_example_docstring decorator, so it's dead.

Duplication

PingPongScheduler and StableAudio3EulerScheduler share identical set_timesteps, add_noise, scale_model_input, __len__, and near-identical docstrings. Link one to the other with # Copied from headers (and make fix-copies) so they don't drift.

Dead / stray code

docs/source/en/api/pipelines/autodoc.md is a stray two-line artifact not referenced in _toctree.yml — remove it.
Conversion script: _check_shapes is never called; convert_dit returns a skipped list that is never appended to yet reported in the summary print; redundant import math inside convert_duration_embedder; the sys.path.insert(0, repo_root / "src") hack is unnecessary (other conversion scripts assume an installed diffusers). Also, the written model_index.json hard-codes _diffusers_version: "0.0.0.dev0" instead of the actual version.
AutoencoderSAME.encode_audio/decode_audio chunked helpers aren't used by any pipeline — only by their own tests. Wire them into the pipelines or drop them.
_SoftNormBottleneck.encode carries a training-time running-std update ported from the research repo; this is inference-only code and should be deleted.
scripts/run_stable_audio_3_inference.py is a one-off smoke test duplicating the docs example; there's no precedent for run_*_inference.py scripts in scripts/, and the conversion script already does a round-trip sanity check — drop it.

Tests

TestStableAudio3DiTProductionStructure instantiates the full production config (embed_dim=1536, depth=24 → ~1.4B params, ~5.5 GB fp32) in the fast CPU suite. Mark it @slow or use init_empty_weights/meta device for the structural check.
test_speed_vs_reference in the SAME autoencoder tests imports the external stable_audio_3 research package, asserts nothing, and just prints a benchmark table — this is tooling, not a test, and should be removed.
The model tests (AutoencoderSAME, StableAudio3DiTModel) are hand-rolled unittest classes; deriving from ModelTesterMixin (and SchedulerCommonTest for the schedulers) would get save/load, dtype, and offloading coverage for free.
The pipeline fast tests using PipelineTesterMixin look good, and the dummy-object/__init__ registrations are consistent.

serge v0.1.0 · model: claude-fable-5 · 12 LLM turns · 17 tool calls · 450.2s · 881430 in / 34077 out tokens

sergereview · 2026-07-06T03:03:32Z

@@ -0,0 +1,3 @@
+from .modeling_stable_audio_3 import StableAudio3DurationEmbedder


Every other pipeline subpackage (e.g. stable_audio/__init__.py) uses the lazy _LazyModule structure with an OptionalDependencyNotAvailable guard for transformers. Because pipeline_stable_audio_3.py does from transformers import GemmaTokenizer, ... at module level, importing diffusers.pipelines.stable_audio_3 in an environment without transformers raises an ImportError instead of resolving to the dummy objects. Please follow the same guarded lazy-import pattern as the sibling packages.

solved, thanks!

sergereview · 2026-07-06T03:03:32Z

@@ -0,0 +1,2 @@
+# autodoc


This looks like an accidentally committed artifact — it contains only # autodoc and is not referenced in _toctree.yml. Please remove the file.

removed, thanks!

sergereview · 2026-07-06T03:03:32Z

+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def _check_shapes(converted_sd: dict, model_sd: dict, label: str) -> tuple[int, int]:


_check_shapes is defined but never called anywhere in the script — dead code. Per the repo's coding rules (no unused code paths), either use it in the conversion flow or delete it.

sergereview · 2026-07-06T03:03:32Z

+    Returns (state_dict, list_of_skipped_keys).
+    """
+    out = {}
+    skipped = []


skipped is initialized here and returned, but nothing in convert_dit ever appends to it, so the "{len(dit_skipped)} skipped/no-op" message in convert() always prints 0 and is misleading. Either track keys that don't match any rename (as convert_vae does via _convert_trb_block) or drop the list and the tuple return.

sergereview · 2026-07-06T03:03:32Z

+
+def convert_duration_embedder(ref_sd: dict, min_freq: float = 0.5, max_freq: float = 10000.0) -> dict:
+    """Build diffusers StableAudio3DurationEmbedder state dict."""
+    import math


math is already imported at module top (line 54); this local re-import is redundant.

Suggested change

import math

sergereview · 2026-07-06T03:03:32Z

+        mask: Optional[torch.Tensor] = None,
+        mask_start_seconds: Optional[Union[float, List[float]]] = None,
+        mask_end_seconds: Optional[Union[float, List[float]]] = None,
+        num_inference_steps: int = 8,


This hardcoded default of 8 is inconsistent with the base StableAudio3Pipeline.__call__, which defaults num_inference_steps=None and falls back to self.scheduler.config.num_inference_steps (8 for PingPongScheduler, 100 for StableAudio3EulerScheduler). There is even a base-pipeline test (test_stable_audio_3_default_steps_follow_scheduler) asserting that behavior. With the Euler base checkpoint, this default will silently run only 8 steps. Use Optional[int] = None and the same scheduler-config fallback here.

nice, solved!

sergereview · 2026-07-06T03:03:32Z

+        Returns:
+            [`~pipelines.AudioPipelineOutput`] with ``.audios``.
+        """
+        if audio is None and mask is None and mask_start_seconds is None:


Per the repo rules ("Do not guess user intent and silently correct behavior... raise a concise error for unsupported cases rather than adding complex fallback logic"), an inpaint pipeline called without any inpainting inputs should raise a ValueError, not silently degrade to plain text-to-audio behind a warning. Users who want text-to-audio should use StableAudio3Pipeline directly. Also, the warning text says "unconditional" generation, which is inaccurate — it is still text-conditioned.

sergereview · 2026-07-06T03:03:33Z

+        self.sigmas: Optional[torch.Tensor] = None
+        self.timesteps: Optional[torch.Tensor] = None
+
+    def set_timesteps(


set_timesteps, add_noise, scale_model_input, and __len__ are byte-for-byte identical to PingPongScheduler (only step differs). Diffusers keeps this kind of duplication in sync via # Copied from diffusers.schedulers.scheduling_ping_pong.PingPongScheduler.set_timesteps headers + make fix-copies — please add them so the two schedules can't drift apart.

sergereview · 2026-07-06T03:03:33Z

+    exactly, so converted SA3 Medium weights load with ``strict=True``."""
+
+    def test_state_dict_matches_checkpoint(self):
+        model = StableAudio3DiTModel(**PROD_CFG)


This instantiates the full production config (embed_dim=1536, depth=24 → roughly 1.4B parameters, ~5.5 GB in fp32) inside the fast CPU test suite, and all of those weights get randomly initialized just to inspect the state-dict keys/shapes. This will be very slow and memory-heavy in CI. Either mark the test @slow, or build the model on the meta device (e.g. accelerate.init_empty_weights()) since only names and shapes are checked.

sergereview · 2026-07-06T03:03:33Z

+    """
+
+    @slow
+    def test_speed_vs_reference(self):


This is a benchmark, not a test: it imports the external stable_audio_3 research package, asserts nothing, and only prints timing tables (with a bare except Exception: pass at the end). Per the repo guidelines, research-repo comparison scaffolding should not be carried into the library — please remove this test (and the now-unused SMALL_CFG alias at line 72).

buffett0323 · 2026-07-07T06:06:53Z

Hi, I just updated the code, fixing serge review's suggestion, thanks!

sdevil7th · 2026-07-07T09:02:29Z

Hi @buffett0323 , thanks for working on Stable Audio 3 support. The PR already covers a lot of important ground- conversion, SAME autoencoder, transformer, schedulers, text-to-audio, inpainting, docs, and tests!

One workflow gap I noticed is native audio-to-audio variation. Stability’s reference library exposes this via init_audio + init_noise_level:

model.generate(
    init_audio=init_audio,
    init_noise_level=0.5,
    prompt="...",
    duration=30,
)

In this PR, I see text-to-audio and inpainting via audio + mask controls, but not an equivalent init-audio/noise-strength path. Inpainting can approximate some variation-like behavior, but it is not the same workflow: true audio-to-audio variation noises/denoises the whole source according to init_noise_level.

Suggested improvement: either add native audio-to-audio support, e.g. StableAudio3AudioToAudioPipeline or explicit init_audio / init_noise_level args, or document clearly that audio-to-audio variation is not exposed yet so users do not confuse it with inpainting.

One related base-checkpoint question: the PR docs mention stable-audio-3-medium-base, while Stability’s reference docs treat cfg_scale / negative_prompt as base-model-only controls. The current pipeline does not expose guidance_scale / negative_prompt. If base checkpoints are in scope for this PR, should that be supported or explicitly documented as out of scope?

Also, while testing the converter, I noticed two implementation issues:

Scheduler selection appears to miss one valid SA3 Medium config shape. The converter currently reads model.diffusion_objective / top-level diffusion_objective, but my model_config.json has it under model.diffusion.diffusion_objective with value rf_denoiser. In that case the converter silently falls back to StableAudio3EulerScheduler with 100 steps instead of PingPongScheduler with 8 steps.
The converter uses Unicode characters in console output (──, →, ✓, etc.). On Windows environments using a non-UTF-8 code page, this can raise UnicodeEncodeError during conversion. It may be safer for conversion scripts to either force UTF-8 output or keep logs ASCII-safe.

dg845

Thanks for the PR! Left an initial design review :).

dg845 · 2026-07-09T03:54:32Z

+# ──────────────────────────────────────────────────────────────────────────────
+
+
+class _DynamicTanh(nn.Module):


Suggested change

class _DynamicTanh(nn.Module):

class DynamicTanh(nn.Module):

nit: I don't think we need to prefix all of the helper module names with _, as the modeling code for other autoencoders doesn't do this. For example, see AutoencoderOobleck:

diffusers/src/diffusers/models/autoencoders/autoencoder_oobleck.py

Line 30 in 208704a

class Snake1d(nn.Module):

dg845 · 2026-07-09T03:56:48Z

+    − Attn(Q2,K2,V) — so that attention patterns common to both heads cancel out, improving focus.
+    """
+
+    def __init__(self, dim: int, dim_heads: int = 128, use_differential: bool = True, qk_norm_eps: float = 1e-3):


The qk_norm_eps argument here doesn't appear to be used, is this intentional? If it isn't needed we should remove it.

dg845 · 2026-07-09T04:00:22Z

+# ──────────────────────────────────────────────────────────────────────────────
+
+
+class _Attention(nn.Module):


Would it be possible to refactor _Attention to use the diffusers attention module (AttentionModuleMixin) + attention processor pattern? For example, this is how Flux 2 implements the attention module:

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 494 to 496 in 208704a

class Flux2Attention(torch.nn.Module, AttentionModuleMixin):

_default_processor_cls = Flux2AttnProcessor

_available_processors = [Flux2AttnProcessor, Flux2KVAttnProcessor]

and an attention processor:

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 326 to 328 in 208704a

class Flux2AttnProcessor:

_attention_backend = None

_parallel_config = None

dg845 · 2026-07-09T04:05:10Z

+            k1, k2 = self.k_norm(k1), self.k_norm(k2)
+            q1, q2 = _apply_rotary(q1, freqs), _apply_rotary(q2, freqs)
+            k1, k2 = _apply_rotary(k1, freqs), _apply_rotary(k2, freqs)
+            out = F.scaled_dot_product_attention(q1, k1, v, attn_mask=attn_mask) - F.scaled_dot_product_attention(


Could we use dispatch_attention_fn instead of F.scaled_dot_product_attention here? This would allow us to support different attention backends such as Flash Attention (with torch native SDPA still supported as the default "native" backend). For example, this is what Flux 2 does:

diffusers/src/diffusers/models/transformers/transformer_flux2.py

Lines 369 to 376 in 208704a

hidden_states = dispatch_attention_fn(

query,

key,

value,

attn_mask=attention_mask,

backend=self._attention_backend,

parallel_config=self._parallel_config,

)

For reference, you can also look at the attention backend docs here.

dg845 · 2026-07-09T04:06:55Z

+        self.sampling_rate = sampling_rate
+        self.latent_dim = latent_dim


Suggested change

self.sampling_rate = sampling_rate

self.latent_dim = latent_dim

I don't think we need to define these attributes as they are already registered to the config and available via e.g. self.config.sampling_rate.

dg845 · 2026-07-09T04:42:52Z

+        if output_type == "latent":
+            return AudioPipelineOutput(audios=latents)


nit: I think we should also support return_dict=False in the output_type="latent" case here. For example, we could do something like

if output_type == "latent": audio = latents else: audio = self.vae.decode(latents).sample ... ... if not return_dict: return (audio,) return AudioPipelineOutput(audios=audio)

dg845 · 2026-07-09T04:49:01Z

+from .pipeline_stable_audio_3 import StableAudio3Pipeline
+
+
+class StableAudio3InpaintPipeline(StableAudio3Pipeline):


Suggested change

class StableAudio3InpaintPipeline(StableAudio3Pipeline):

class StableAudio3InpaintPipeline(DiffusionPipeline):

StableAudio3InpaintPipeline should not inherit from StableAudio3Pipeline but rather should inherit directly from DiffusionPipeline. We can copy over shared methods such as encode_prompt, encode_duration, etc. with the # Copied from mechanism to sync the implementations:

class StableAudio3InpaintPipeline(StableAudio3Pipeline): ... # Copied from diffusers.pipelines.stable_audio_3.pipeline_stable_audio3.StableAudio3Pipeline.encode_prompt def encode_prompt( self, prompt: Optional[Union[str, List[str]]], ..., ): ... ...

dg845 · 2026-07-09T04:49:50Z

+    @torch.no_grad()
+    def __call__(


nit: could we also have an example for the inpainting pipeline?

dg845 · 2026-07-09T04:52:55Z

+    pred_original_sample: Optional[torch.Tensor] = None
+
+
+class StableAudio3EulerScheduler(SchedulerMixin, ConfigMixin):


Can we reuse the existing FlowMatchEulerDiscreteScheduler in place of StableAudio3EulerScheduler? I think supplying the log-SNR sigma schedule as a custom sigmas schedule to FlowMatchEulerDiscreteScheduler.set_timesteps could potentially be equivalent.

dg845 · 2026-07-09T04:56:10Z

+    pred_original_sample: Optional[torch.Tensor] = None
+
+
+class PingPongScheduler(SchedulerMixin, ConfigMixin):


Similarly, could we also reuse FlowMatchEulerDiscreteScheduler in place of PingPongScheduler? I think using FlowMatchEulerDiscreteScheduler with stochastic_sampling=True might be equivalent as it does a similar denoise + re-noise to the next noise level operation like PingPongScheduler does:

diffusers/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py

Lines 506 to 509 in 208704a

if self.config.stochastic_sampling:

x0 = sample - current_sigma * model_output

noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)

prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise

github-actions Bot added documentation Improvements or additions to documentation models tests utils pipelines schedulers fixes-issue size/L PR with diff > 200 LOC labels Jul 3, 2026

buffett0323 and others added 8 commits July 4, 2026 21:15

Stable Audio 3 First Version Commit, require model dict loading on GPUs

bf4d378

Fixing bugs for sa3 VAE

3acd64d

Make quality fix

a366039

Fix the unit test failure

20d1733

Style fix

1e1d382

Autopipeline available

daa4a01

Adding document description for testing examples.

1a6af0d

buffett0323 force-pushed the feat/stable-audio-medium branch from 7f9ee7e to 1a6af0d Compare July 5, 2026 04:15

sayakpaul requested a review from dg845 July 6, 2026 02:54

sergereview Bot reviewed Jul 6, 2026

View reviewed changes

Solving fix for serge review bot suggestion

7620761

dg845 reviewed Jul 9, 2026

View reviewed changes

		@@ -0,0 +1,3 @@
		from .modeling_stable_audio_3 import StableAudio3DurationEmbedder

		# ──────────────────────────────────────────────────────────────────────────────


		def _check_shapes(converted_sd: dict, model_sd: dict, label: str) -> tuple[int, int]:

		# ──────────────────────────────────────────────────────────────────────────────


		class _DynamicTanh(nn.Module):

		# ──────────────────────────────────────────────────────────────────────────────


		class _Attention(nn.Module):

	class Flux2Attention(torch.nn.Module, AttentionModuleMixin):
	_default_processor_cls = Flux2AttnProcessor
	_available_processors = [Flux2AttnProcessor, Flux2KVAttnProcessor]

	class Flux2AttnProcessor:
	_attention_backend = None
	_parallel_config = None

	hidden_states = dispatch_attention_fn(
	query,
	key,
	value,
	attn_mask=attention_mask,
	backend=self._attention_backend,
	parallel_config=self._parallel_config,
	)

		self.sampling_rate = sampling_rate
		self.latent_dim = latent_dim

		if output_type == "latent":
		return AudioPipelineOutput(audios=latents)

		from .pipeline_stable_audio_3 import StableAudio3Pipeline


		class StableAudio3InpaintPipeline(StableAudio3Pipeline):

	class StableAudio3InpaintPipeline(StableAudio3Pipeline):
	class StableAudio3InpaintPipeline(DiffusionPipeline):

		pred_original_sample: Optional[torch.Tensor] = None


		class StableAudio3EulerScheduler(SchedulerMixin, ConfigMixin):

		pred_original_sample: Optional[torch.Tensor] = None


		class PingPongScheduler(SchedulerMixin, ConfigMixin):

	if self.config.stochastic_sampling:
	x0 = sample - current_sigma * model_output
	noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
	prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise

Uh oh!

Conversation

buffett0323 commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

sergereview Bot left a comment

Choose a reason for hiding this comment

Correctness

Packaging / imports

Pipeline API

Duplication

Dead / stray code

Tests

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sergereview Bot Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

buffett0323 commented Jul 7, 2026

Uh oh!

sdevil7th commented Jul 7, 2026

Uh oh!

dg845 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

buffett0323 commented Jul 3, 2026 •

edited

Loading