Feat -- Stable Audio 3 by buffett0323 · Pull Request #14119 · huggingface/diffusers · GitHub
Skip to content

Feat -- Stable Audio 3 - #14119

Merged
dg845 merged 19 commits into
huggingface:mainfrom
buffett0323:feat/stable-audio-medium
Aug 19, 2026
Merged

dg845 merged 19 commits into
huggingface:mainfrom
buffett0323:feat/stable-audio-medium

Conversation

@buffett0323

@buffett0323 buffett0323 commented Jul 3, 2026 •

Copy link
Copy Markdown
Contributor

What does this PR do?

Adds Stable Audio 3 Medium to Diffusers.

Fixes #13793

Before submitting

Who can review?

@yiyixuxu @DN6 @sayakpaul

@github-actions 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
buffett0323 force-pushed the feat/stable-audio-medium branch from 7f9ee7e to 1a6af0d Compare July 5, 2026 04:15
@sayakpaul
sayakpaul requested a review from dg845 July 6, 2026 02:54

@sergereview sergereview Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤗 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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

solved, thanks!

Comment thread docs/source/en/api/pipelines/autodoc.md Outdated
@@ -0,0 +1,2 @@
# autodoc

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed, thanks!

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


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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_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.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Suggested change
import math

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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

self.sigmas: Optional[torch.Tensor] = None
self.timesteps: Optional[torch.Tensor] = None

def set_timesteps(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

exactly, so converted SA3 Medium weights load with ``strict=True``."""

def test_state_dict_matches_checkpoint(self):
model = StableAudio3DiTModel(**PROD_CFG)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"""

@slow
def test_speed_vs_reference(self):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor Author

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

@sdevil7th

Copy link
Copy Markdown

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:

  1. 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.

  2. 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 dg845 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3.py Outdated
Comment thread src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3_inpaint.py Outdated
pred_original_sample: Optional[torch.Tensor] = None


class StableAudio3EulerScheduler(SchedulerMixin, ConfigMixin):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

pred_original_sample: Optional[torch.Tensor] = None


class PingPongScheduler(SchedulerMixin, ConfigMixin):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

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

@buffett0323
buffett0323 force-pushed the feat/stable-audio-medium branch from 0c1a9ba to f327a24 Compare July 12, 2026 08:08
@buffett0323

Copy link
Copy Markdown
Contributor Author

@dg845 Hi, thanks for your time and suggestion. I've revised the code. If there's anything else needed, I would love to help, thanks again!

@buffett0323

Copy link
Copy Markdown
Contributor Author

@dg845 Sorry to bother you again, I saw there was a CI/CD bug after you merged from main and I just solved it. Can you take a look and do it again? Thanks for your help and time!

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@buffett0323

buffett0323 commented Jul 21, 2026 •

Copy link
Copy Markdown
Contributor Author

@dg845 Doc build failures are resolved, and I believe the only remaining failure is unrelated to this PR. Thanks again for your time and support!

@dg845 dg845 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for iterating! Left some follow up comments :).

Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py
Comment thread src/diffusers/models/autoencoders/autoencoder_same.py Outdated
Comment on lines +124 to +125
sigmas[0] = 1.0 # sigma_max: start from pure noise
sigmas[-1] = 0.0 # end fully denoised

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is overwriting sigmas[0] to be 1.0 and sigmas[-1] to be 0.0 intentional here? This seems a little weird as we are effectively overwriting logsnr_min and logsnr_max to $-\infty$ and $\infty$, respectively (although they still determine the spacing of the schedule).

sigmas = torch.sigmoid(-logsnr) # (N+1,), decreasing
sigmas[0] = 1.0 # sigma_max: start from pure noise
sigmas[-1] = 0.0 # end fully denoised
sigmas = sigmas * sigma_max

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Similarly, is scaling the sigmas schedule by sigma_max intentional? This also seems a little weird since the resulting sigmas schedule is in general no longer uniform in log-SNR space.

Comment thread tests/models/autoencoders/test_models_autoencoder_same.py Outdated
Comment thread tests/models/transformers/test_models_transformer_stable_audio3.py Outdated
Comment thread scripts/convert_stable_audio_3_to_diffusers.py Outdated
buffett0323 and others added 8 commits July 22, 2026 20:22
…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>
@buffett0323
buffett0323 force-pushed the feat/stable-audio-medium branch from c8784ae to 874d993 Compare July 23, 2026 03:32
@buffett0323

Copy link
Copy Markdown
Contributor Author

@dg845 Hi, I just updated the code, thanks again for your time and effort!

@buffett0323

Copy link
Copy Markdown
Contributor Author

@dg845 Sorry to bother you again, gentle reminder for this PR whenever you get a chance! Let me know if you have any questions or feedback. Thanks again for your help!!

@dg845 dg845 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for the delay and thanks for your patience! I think this PR is close to merge, left a few small comments :).

Comment thread docs/source/en/api/models/autoencoder_same.md Outdated
Comment thread scripts/convert_stable_audio_3_to_diffusers.py Outdated
Comment thread scripts/convert_stable_audio_3_to_diffusers.py Outdated
Comment thread src/diffusers/pipelines/stable_audio_3/pipeline_stable_audio_3.py Outdated
Comment on lines +359 to +360
# The trailing `sinusoidal_blocks` layers use sin gating (matches the reference).
sinusoidal=(transformer_depth - i) <= sinusoidal_blocks,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# The trailing `sinusoidal_blocks` layers use sin gating (matches the reference).
sinusoidal=(transformer_depth - i) <= sinusoidal_blocks,
# The trailing `sinusoidal_blocks` layers use sin gating (matches the reference).
sinusoidal=(transformer_depth - i) < sinusoidal_blocks,

Sorry, I think my previous suggestion was wrong: using <= results in samples which are pure noise, and reverting back to using < produces good samples again. Can you confirm this against the reference implementation and update the comment to clarify which is correct?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@buffett0323

Copy link
Copy Markdown
Contributor Author

Fixed, thanks again for your time @dg845 , please take a look!

@dg845 dg845 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks again for your patience and work on this PR! Can you run make fix-copies to fix the CI failure? This should propagate the prepare_latents change in #14119 (comment) to the inpainting and audio-to-audio pipeline copies.

dim_heads=min(dim_heads, tdim),
use_differential=use_differential,
ff_mult=ff_mult,
# The trailing `sinusoidal_blocks` layers use sin gating (matches the reference).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
# The trailing `sinusoidal_blocks` layers use sin gating (matches the reference).
# Sin-gate the trailing `sinusoidal_blocks - 1` layers ("<", not "<="), which matches the
# `stable-audio-tools` SAME reference (`stable_audio_tools/models/autoencoders.py`).

nit: update the comment to clarify the reference behavior.

use_differential: Whether to use differential attention.
ff_mult: Feed-forward expansion factor.
sliding_window: Sliding-window half-width in latents (band half-width is ``sliding_window * (stride + 1)``).
sinusoidal_blocks: Number of trailing transformer layers that use ``sin`` FFN gating instead of SiLU.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
sinusoidal_blocks: Number of trailing transformer layers that use ``sin`` FFN gating instead of SiLU.
sinusoidal_blocks (`int`, defaults to `0`):
The `sinusoidal_blocks - 1` trailing layers will use `sin` FFN gating instead of SiLU gating. This matches
the reference `stable-audio-tools` SAME transformer block `sinusoidal_blocks` semantics.

nit: update the docstring to clarify the sinusoidal_blocks behavior.

@buffett0323

Copy link
Copy Markdown
Contributor Author

Fixed and uploaded, can you please help me for the CI/CD? Thanks!
@dg845

@dg845

dg845 commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

@dg845
dg845 merged commit 360bef8 into huggingface:main Aug 19, 2026
14 of 15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation fixes-issue models pipelines size/L PR with diff > 200 LOC tests utils

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support for Stable Audio 3 Medium

4 participants