Skip to content

stable_audio

Stable Audio Open 1.0 DiT.

Continuous transformer with rotary self-attention, GQA cross-attention, and prepend global conditioning. 24 layers, embed_dim=1536, head_dim=64.

Classes

fastvideo.models.dits.stable_audio.FourierFeatures 

FourierFeatures(in_features: int, out_features: int, std: float = 1.0)

Bases: Module

Random-Fourier learned-frequency timestep encoder.

Source code in fastvideo/models/dits/stable_audio.py 
34
35
36
37
def __init__(self, in_features: int, out_features: int, std: float = 1.0) -> None:
    super().__init__()
    assert out_features % 2 == 0
    self.weight = nn.Parameter(torch.randn([out_features // 2, in_features]) * std)

fastvideo.models.dits.stable_audio.StableAudioDiT 

StableAudioDiT(config: StableAudioConfig | None = None, hf_config: dict[str, Any] | None = None)

Bases: BaseDiT

Stable Audio Open 1.0 diffusion transformer.

Source code in fastvideo/models/dits/stable_audio.py 
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
def __init__(self, config: StableAudioConfig | None = None,
             hf_config: dict[str, Any] | None = None) -> None:
    if config is None:
        config = StableAudioConfig()
    super().__init__(config=config, hf_config=hf_config or {})
    arch = config.arch_config
    self.hidden_size = arch.hidden_size
    self.num_attention_heads = arch.num_attention_heads
    self.num_channels_latents = arch.num_channels_latents
    io_channels = arch.io_channels
    embed_dim = arch.embed_dim
    depth = arch.depth
    num_heads = arch.num_attention_heads
    cond_token_dim = arch.cond_token_dim
    global_cond_dim = arch.global_cond_dim
    project_cond_tokens = arch.project_cond_tokens
    project_global_cond = arch.project_global_cond
    qk_norm = arch.qk_norm

    self.cond_token_dim = cond_token_dim
    timestep_features_dim = 256
    self.timestep_features = FourierFeatures(1, timestep_features_dim)
    self.to_timestep_embed = nn.Sequential(
        ReplicatedLinear(timestep_features_dim, embed_dim, bias=True),
        nn.SiLU(),
        ReplicatedLinear(embed_dim, embed_dim, bias=True),
    )
    self.diffusion_objective = "v"

    cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim
    self.to_cond_embed = nn.Sequential(
        ReplicatedLinear(cond_token_dim, cond_embed_dim, bias=False),
        nn.SiLU(),
        ReplicatedLinear(cond_embed_dim, cond_embed_dim, bias=False),
    )

    global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
    self.to_global_embed = nn.Sequential(
        ReplicatedLinear(global_cond_dim, global_embed_dim, bias=False),
        nn.SiLU(),
        ReplicatedLinear(global_embed_dim, global_embed_dim, bias=False),
    )

    self.transformer = ContinuousTransformer(
        dim=embed_dim, depth=depth, dim_heads=embed_dim // num_heads, dim_in=io_channels,
        dim_out=io_channels, cross_attend=True, cond_token_dim=cond_embed_dim,
        qk_norm=qk_norm,
    )

    self.preprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
    nn.init.zeros_(self.preprocess_conv.weight)
    self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
    nn.init.zeros_(self.postprocess_conv.weight)

    self.io_channels = io_channels
    self.embed_dim = embed_dim
    self.depth = depth
    self.num_heads = num_heads
    self.__post_init__()

Methods:

fastvideo.models.dits.stable_audio.StableAudioDiT.forward 
forward(x: Tensor, t: Tensor, *, cross_attn_cond: Tensor, global_embed: Tensor) -> Tensor

Forward over a single batch. CFG batching is the caller's job.

Source code in fastvideo/models/dits/stable_audio.py 
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
def forward(self, x: torch.Tensor, t: torch.Tensor, *, cross_attn_cond: torch.Tensor,
            global_embed: torch.Tensor) -> torch.Tensor:
    """Forward over a single batch. CFG batching is the caller's job."""
    model_dtype = next(self.parameters()).dtype
    x = x.to(model_dtype)
    t = t.to(model_dtype)
    cross_attn_cond = cross_attn_cond.to(model_dtype)
    global_embed = global_embed.to(model_dtype)

    cross_attn_cond = self._seq_apply(self.to_cond_embed, cross_attn_cond)
    global_embed = self._seq_apply(self.to_global_embed, global_embed)
    timestep_embed = self._seq_apply(self.to_timestep_embed, self.timestep_features(t[:, None]))
    global_embed = global_embed + timestep_embed
    prepend_inputs = global_embed.unsqueeze(1)

    x = self.preprocess_conv(x) + x
    x = rearrange(x, "b c t -> b t c")
    out = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond)
    out = rearrange(out, "b t c -> b c t")[:, :, prepend_inputs.shape[1]:]
    return self.postprocess_conv(out) + out
fastvideo.models.dits.stable_audio.StableAudioDiT.from_official_state_dict classmethod 
from_official_state_dict(state_dict: dict[str, Tensor], prefix: str = 'model.model.') -> 'StableAudioDiT'

Load from a raw stable_audio_tools monolithic state dict. Kept for tests / older checkpoints; production loads go through the standard TransformerLoader against the converted Diffusers repo.

Source code in fastvideo/models/dits/stable_audio.py 
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
@classmethod
def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor],
                             prefix: str = "model.model.") -> "StableAudioDiT":
    """Load from a raw `stable_audio_tools` monolithic state dict.
    Kept for tests / older checkpoints; production loads go through
    the standard `TransformerLoader` against the converted Diffusers
    repo.
    """
    model = cls()
    mapping_fn = get_param_names_mapping(model.config.arch_config.param_names_mapping)
    remapped: dict[str, torch.Tensor] = {}
    for k, v in state_dict.items():
        if not k.startswith(prefix):
            continue
        new_key, _, _ = mapping_fn(k)
        remapped[new_key] = v
    missing, unexpected = model.load_state_dict(remapped, strict=True)
    if missing or unexpected:
        raise RuntimeError(
            f"StableAudioDiT load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}")
    return model

Functions: