--- title: 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.

fastvideo.models.dits.stable_audio.LocalAttention

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

fastvideo.models.dits.stable_audio.LocalAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor

Query tensor of shape [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor of shape [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor

torch.Tensor: Output tensor after local attention

fastvideo.models.dits.stable_audio.StableAudioConfig

StableAudioConfig(arch_config: DiTArchConfig = StableAudioArchConfig(), prefix: str = 'StableAudio', quant_config: QuantizationConfig | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)

fastvideo.models.dits.stable_audio.FP32LayerNorm

FP32LayerNorm()

fastvideo.models.dits.stable_audio.FP32LayerNorm.forward

forward(inputs: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.stable_audio.ReplicatedLinear

ReplicatedLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Replicated linear layer.

Parameters

input_size: int

input dimension of the linear layer.

output_size: int

output dimension of the linear layer.

bias: bool

If true, add bias.

skip_bias_add: bool

If true, skip adding bias but instead return it.

params_dtype: torch.dtype | None

Data type for the parameters.

quant_config: QuantizationConfig | None

Quantization configure.

prefix: str

The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

fastvideo.models.dits.stable_audio.ReplicatedLinear.weight_loader

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

fastvideo.models.dits.stable_audio.ReplicatedLinear.forward

forward(x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

fastvideo.models.dits.stable_audio.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.stable_audio.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

fastvideo.models.dits.stable_audio.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.stable_audio.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

fastvideo.models.dits.stable_audio.BaseDiT

BaseDiT(config: DiTConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.stable_audio.BaseDiT.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, guidance = None, **kwargs) -> torch.Tensor

fastvideo.models.dits.stable_audio.BaseDiT.preprocess_loaded_weight

preprocess_loaded_weight(name: str, tensor: torch.Tensor) -> torch.Tensor

Adapt a checkpoint tensor before sharding; default models are unchanged.

fastvideo.models.dits.stable_audio.BaseDiT.inference_config_overrides

inference_config_overrides() -> dict[str, Any]

Extra config.json keys a deployable checkpoint of this model needs.

fastvideo.models.dits.stable_audio.BaseDiT.inference_export_expected_shape

inference_export_expected_shape(output_key: str, base_shape: tuple[int, ...]) -> tuple[int, ...]

The exported shape of output_key given the base checkpoint's shape.

fastvideo.models.dits.stable_audio.get_param_names_mapping

get_param_names_mapping(mapping_dict: dict[str, str]) -> Callable[[str], tuple[str, Any, Any]]

Creates a mapping function that transforms parameter names using regex patterns.

Parameters

mapping_dict: Dict[str, str]

Dictionary mapping regex patterns to replacement patterns

param_name: str

The parameter name to be transformed

Returns

Callable[[str], tuple[str, Any, Any]]

Callable[[str], str]: A function that maps parameter names from source to target format

fastvideo.models.dits.stable_audio.FourierFeatures

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

Random-Fourier learned-frequency timestep encoder.

Source
class FourierFeatures(nn.Module):
    """Random-Fourier learned-frequency timestep encoder."""

    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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        f = 2 * math.pi * x @ self.weight.T
        return torch.cat([f.cos(), f.sin()], dim=-1)

fastvideo.models.dits.stable_audio.FourierFeatures.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    f = 2 * math.pi * x @ self.weight.T
    return torch.cat([f.cos(), f.sin()], dim=-1)

fastvideo.models.dits.stable_audio.RotaryEmbedding

RotaryEmbedding(dim: int, base: float = 10000.0)
Source
class RotaryEmbedding(nn.Module):

    def __init__(self, dim: int, base: float = 10000.0) -> None:
        super().__init__()
        inv_freq = 1.0 / (base**(torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.register_buffer("scale", None)

    def forward_from_seq_len(self, seq_len: int):
        t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32)
        freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
        freqs = torch.cat((freqs, freqs), dim=-1)
        return freqs, 1.0

fastvideo.models.dits.stable_audio.RotaryEmbedding.forward_from_seq_len

forward_from_seq_len(seq_len: int)
Source
def forward_from_seq_len(self, seq_len: int):
    t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32)
    freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
    freqs = torch.cat((freqs, freqs), dim=-1)
    return freqs, 1.0

fastvideo.models.dits.stable_audio.FeedForward

FeedForward(dim: int, mult: int = 4, zero_init_output: bool = True)
Source
class FeedForward(nn.Module):
    # Sequential layout `(GLU, Identity, Linear, Identity)` keeps the
    # checkpoint keys at indices 0 and 2.

    def __init__(self, dim: int, mult: int = 4, zero_init_output: bool = True) -> None:
        super().__init__()
        inner_dim = int(dim * mult)
        linear_in = _GLU(dim, inner_dim, nn.SiLU())
        linear_out = ReplicatedLinear(inner_dim, dim, bias=True)
        if zero_init_output:
            nn.init.zeros_(linear_out.weight)
            nn.init.zeros_(linear_out.bias)
        self.ff = nn.Sequential(linear_in, nn.Identity(), linear_out, nn.Identity())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for mod in self.ff:
            if isinstance(mod, ReplicatedLinear):
                x, _ = mod(x)
            else:
                x = mod(x)
        return x

fastvideo.models.dits.stable_audio.FeedForward.forward

forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    for mod in self.ff:
        if isinstance(mod, ReplicatedLinear):
            x, _ = mod(x)
        else:
            x = mod(x)
    return x

fastvideo.models.dits.stable_audio.Attention

Attention(dim: int, dim_heads: int = 64, dim_context: int | None = None, zero_init_output: bool = True, qk_norm: str | None = None)
Source
class Attention(nn.Module):

    def __init__(self,
                 dim: int,
                 dim_heads: int = 64,
                 dim_context: int | None = None,
                 zero_init_output: bool = True,
                 qk_norm: str | None = None) -> None:
        super().__init__()
        self.dim = dim
        self.dim_heads = dim_heads
        dim_kv = dim_context if dim_context is not None else dim
        self.num_heads = dim // dim_heads
        self.kv_heads = dim_kv // dim_heads
        if dim_context is not None:
            self.to_q = ReplicatedLinear(dim, dim, bias=False)
            self.to_kv = ReplicatedLinear(dim_kv, dim_kv * 2, bias=False)
        else:
            self.to_qkv = ReplicatedLinear(dim, dim * 3, bias=False)
        self.to_out = ReplicatedLinear(dim, dim, bias=False)
        if zero_init_output:
            nn.init.zeros_(self.to_out.weight)

        # `stable-audio-open-small` wraps Q/K in LayerNorm before attn
        # (`attn_kwargs.qk_norm = "ln"` in its `model_config.json`); the
        # 1.0 base does not. Names match upstream (`q_norm`/`k_norm`)
        # so the converted state dict loads strict.
        if qk_norm == "ln":
            self.q_norm = nn.LayerNorm(dim_heads)
            self.k_norm = nn.LayerNorm(dim_heads)
        elif qk_norm is None:
            self.q_norm = nn.Identity()
            self.k_norm = nn.Identity()
        else:
            raise ValueError(f"Unsupported qk_norm={qk_norm!r}; expected 'ln' or None.")

        self.attn = LocalAttention(num_heads=self.num_heads,
                                   head_size=dim_heads,
                                   num_kv_heads=self.kv_heads,
                                   causal=False,
                                   supported_attention_backends=_SUPPORTED_BACKENDS)

    def forward(self,
                x: torch.Tensor,
                context: torch.Tensor | None = None,
                rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor:
        h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
        kv_input = context if has_context else x
        if has_context:
            q, _ = self.to_q(x)
            kv, _ = self.to_kv(kv_input)
            k, v = kv.chunk(2, dim=-1)
        else:
            qkv, _ = self.to_qkv(x)
            q, k, v = qkv.chunk(3, dim=-1)
        # LocalAttention expects [batch, seq_len, num_heads, head_dim].
        q = rearrange(q, "b n (h d) -> b n h d", h=h)
        k = rearrange(k, "b n (h d) -> b n h d", h=kv_h)
        v = rearrange(v, "b n (h d) -> b n h d", h=kv_h)
        q = self.q_norm(q)
        k = self.k_norm(k)

        if rotary_pos_emb is not None:
            freqs, _ = rotary_pos_emb
            v_dtype = v.dtype
            # Partial rotary (rot_dim < head_dim) with halves-swap, so
            # apply outside LocalAttention. q,k come in as [B, S, H, D];
            # transpose to [B, H, S, D] for the helper.
            q_t = q.transpose(1, 2)
            k_t = k.transpose(1, 2)
            if q_t.shape[-2] >= k_t.shape[-2]:
                ratio = q_t.shape[-2] / k_t.shape[-2]
                q_freqs, k_freqs = freqs, ratio * freqs
            else:
                ratio = k_t.shape[-2] / q_t.shape[-2]
                q_freqs, k_freqs = ratio * freqs, freqs
            q = _apply_rotary_pos_emb(q_t, q_freqs).to(v_dtype).transpose(1, 2)
            k = _apply_rotary_pos_emb(k_t, k_freqs).to(v_dtype).transpose(1, 2)

        out = self.attn(q, k, v)
        out = rearrange(out, "b n h d -> b n (h d)")
        out, _ = self.to_out(out)
        return out

fastvideo.models.dits.stable_audio.Attention.forward

forward(x: torch.Tensor, context: torch.Tensor | None = None, rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor
Source
def forward(self,
            x: torch.Tensor,
            context: torch.Tensor | None = None,
            rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor:
    h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
    kv_input = context if has_context else x
    if has_context:
        q, _ = self.to_q(x)
        kv, _ = self.to_kv(kv_input)
        k, v = kv.chunk(2, dim=-1)
    else:
        qkv, _ = self.to_qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
    # LocalAttention expects [batch, seq_len, num_heads, head_dim].
    q = rearrange(q, "b n (h d) -> b n h d", h=h)
    k = rearrange(k, "b n (h d) -> b n h d", h=kv_h)
    v = rearrange(v, "b n (h d) -> b n h d", h=kv_h)
    q = self.q_norm(q)
    k = self.k_norm(k)

    if rotary_pos_emb is not None:
        freqs, _ = rotary_pos_emb
        v_dtype = v.dtype
        # Partial rotary (rot_dim < head_dim) with halves-swap, so
        # apply outside LocalAttention. q,k come in as [B, S, H, D];
        # transpose to [B, H, S, D] for the helper.
        q_t = q.transpose(1, 2)
        k_t = k.transpose(1, 2)
        if q_t.shape[-2] >= k_t.shape[-2]:
            ratio = q_t.shape[-2] / k_t.shape[-2]
            q_freqs, k_freqs = freqs, ratio * freqs
        else:
            ratio = k_t.shape[-2] / q_t.shape[-2]
            q_freqs, k_freqs = ratio * freqs, freqs
        q = _apply_rotary_pos_emb(q_t, q_freqs).to(v_dtype).transpose(1, 2)
        k = _apply_rotary_pos_emb(k_t, k_freqs).to(v_dtype).transpose(1, 2)

    out = self.attn(q, k, v)
    out = rearrange(out, "b n h d -> b n (h d)")
    out, _ = self.to_out(out)
    return out

fastvideo.models.dits.stable_audio.TransformerBlock

TransformerBlock(dim: int, dim_heads: int = 64, cross_attend: bool = False, dim_context: int | None = None, zero_init_branch_outputs: bool = True, qk_norm: str | None = None)
Source
class TransformerBlock(nn.Module):

    def __init__(self,
                 dim: int,
                 dim_heads: int = 64,
                 cross_attend: bool = False,
                 dim_context: int | None = None,
                 zero_init_branch_outputs: bool = True,
                 qk_norm: str | None = None) -> None:
        super().__init__()
        self.dim = dim
        self.dim_heads = min(dim_heads, dim)
        self.cross_attend = cross_attend
        self.pre_norm = FP32LayerNorm(dim, elementwise_affine=True)
        self.self_attn = Attention(dim,
                                   dim_heads=self.dim_heads,
                                   zero_init_output=zero_init_branch_outputs,
                                   qk_norm=qk_norm)
        if cross_attend:
            self.cross_attend_norm = FP32LayerNorm(dim, elementwise_affine=True)
            self.cross_attn = Attention(dim,
                                        dim_heads=self.dim_heads,
                                        dim_context=dim_context,
                                        zero_init_output=zero_init_branch_outputs,
                                        qk_norm=qk_norm)
        self.ff_norm = FP32LayerNorm(dim, elementwise_affine=True)
        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs)

    def forward(self,
                x: torch.Tensor,
                context: torch.Tensor | None = None,
                rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor:
        x = x + self.self_attn(self.pre_norm(x), rotary_pos_emb=rotary_pos_emb)
        if context is not None and self.cross_attend:
            x = x + self.cross_attn(self.cross_attend_norm(x), context=context)
        x = x + self.ff(self.ff_norm(x))
        return x

fastvideo.models.dits.stable_audio.TransformerBlock.forward

forward(x: torch.Tensor, context: torch.Tensor | None = None, rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor
Source
def forward(self,
            x: torch.Tensor,
            context: torch.Tensor | None = None,
            rotary_pos_emb: tuple[torch.Tensor, float] | None = None) -> torch.Tensor:
    x = x + self.self_attn(self.pre_norm(x), rotary_pos_emb=rotary_pos_emb)
    if context is not None and self.cross_attend:
        x = x + self.cross_attn(self.cross_attend_norm(x), context=context)
    x = x + self.ff(self.ff_norm(x))
    return x

fastvideo.models.dits.stable_audio.ContinuousTransformer

ContinuousTransformer(dim: int, depth: int, *, dim_heads: int = 64, dim_in: int | None = None, dim_out: int | None = None, cross_attend: bool = False, cond_token_dim: int | None = None, zero_init_branch_outputs: bool = True, qk_norm: str | None = None)
Source
class ContinuousTransformer(nn.Module):

    def __init__(self,
                 dim: int,
                 depth: int,
                 *,
                 dim_heads: int = 64,
                 dim_in: int | None = None,
                 dim_out: int | None = None,
                 cross_attend: bool = False,
                 cond_token_dim: int | None = None,
                 zero_init_branch_outputs: bool = True,
                 qk_norm: str | None = None) -> None:
        super().__init__()
        self.dim = dim
        self.depth = depth
        self.project_in = (ReplicatedLinear(dim_in, dim, bias=False) if dim_in is not None else nn.Identity())
        self.project_out = (ReplicatedLinear(dim, dim_out, bias=False) if dim_out is not None else nn.Identity())
        self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32))
        self.layers = nn.ModuleList([
            TransformerBlock(dim,
                             dim_heads=dim_heads,
                             cross_attend=cross_attend,
                             dim_context=cond_token_dim,
                             zero_init_branch_outputs=zero_init_branch_outputs,
                             qk_norm=qk_norm) for _ in range(depth)
        ])

    def forward(self,
                x: torch.Tensor,
                prepend_embeds: torch.Tensor | None = None,
                context: torch.Tensor | None = None) -> torch.Tensor:
        if isinstance(self.project_in, ReplicatedLinear):
            x, _ = self.project_in(x)
        if prepend_embeds is not None:
            assert prepend_embeds.shape[-1] == x.shape[-1]
            x = torch.cat((prepend_embeds, x), dim=-2)
        rotary = self.rotary_pos_emb.forward_from_seq_len(x.shape[1])
        for layer in self.layers:
            x = layer(x, context=context, rotary_pos_emb=rotary)
        if isinstance(self.project_out, ReplicatedLinear):
            x, _ = self.project_out(x)
        return x

fastvideo.models.dits.stable_audio.ContinuousTransformer.forward

forward(x: torch.Tensor, prepend_embeds: torch.Tensor | None = None, context: torch.Tensor | None = None) -> torch.Tensor
Source
def forward(self,
            x: torch.Tensor,
            prepend_embeds: torch.Tensor | None = None,
            context: torch.Tensor | None = None) -> torch.Tensor:
    if isinstance(self.project_in, ReplicatedLinear):
        x, _ = self.project_in(x)
    if prepend_embeds is not None:
        assert prepend_embeds.shape[-1] == x.shape[-1]
        x = torch.cat((prepend_embeds, x), dim=-2)
    rotary = self.rotary_pos_emb.forward_from_seq_len(x.shape[1])
    for layer in self.layers:
        x = layer(x, context=context, rotary_pos_emb=rotary)
    if isinstance(self.project_out, ReplicatedLinear):
        x, _ = self.project_out(x)
    return x

fastvideo.models.dits.stable_audio.StableAudioDiT

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

Stable Audio Open 1.0 diffusion transformer.

Source
class StableAudioDiT(BaseDiT):
    """Stable Audio Open 1.0 diffusion transformer."""

    _fsdp_shard_conditions = _DEFAULT_CONFIG.arch_config._fsdp_shard_conditions
    _compile_conditions = _DEFAULT_CONFIG.arch_config._compile_conditions
    param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
    reverse_param_names_mapping: dict = {}

    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__()

    @staticmethod
    def _seq_apply(seq: nn.Sequential, x: torch.Tensor) -> torch.Tensor:
        for mod in seq:
            if isinstance(mod, ReplicatedLinear):
                x, _ = mod(x)
            else:
                x = mod(x)
        return x

    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

    @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

fastvideo.models.dits.stable_audio.StableAudioDiT.forward

forward(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.

Source
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

from_official_state_dict(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.

Source
@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