Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

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(&#91;f.cos(), f.sin()&#93;, dim=-1)</code></pre></details>

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) -&gt; None:
    super().__init__()
    inv_freq = 1.0 / (base**(torch.arange(0, dim, 2).float() / dim))
    self.register_buffer(&quot;inv_freq&quot;, inv_freq)
    self.register_buffer(&quot;scale&quot;, 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(&quot;i , j -&gt; i j&quot;, t, self.inv_freq)
    freqs = torch.cat((freqs, freqs), dim=-1)
    return freqs, 1.0</code></pre></details>

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) -&gt; 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) -&gt; torch.Tensor:
    for mod in self.ff:
        if isinstance(mod, ReplicatedLinear):
            x, _ = mod(x)
        else:
            x = mod(x)
    return x</code></pre></details>

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) -&gt; 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 = &quot;ln&quot;` 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 == &quot;ln&quot;:
        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&quot;Unsupported qk_norm=&#123;qk_norm!r&#125;; expected &#x27;ln&#x27; or None.&quot;)

    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&#91;torch.Tensor, float&#93; | None = None) -&gt; 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 &#91;batch, seq_len, num_heads, head_dim&#93;.
    q = rearrange(q, &quot;b n (h d) -&gt; b n h d&quot;, h=h)
    k = rearrange(k, &quot;b n (h d) -&gt; b n h d&quot;, h=kv_h)
    v = rearrange(v, &quot;b n (h d) -&gt; b n h d&quot;, 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 &lt; head_dim) with halves-swap, so
        # apply outside LocalAttention. q,k come in as &#91;B, S, H, D&#93;;
        # transpose to &#91;B, H, S, D&#93; for the helper.
        q_t = q.transpose(1, 2)
        k_t = k.transpose(1, 2)
        if q_t.shape&#91;-2&#93; &gt;= k_t.shape&#91;-2&#93;:
            ratio = q_t.shape&#91;-2&#93; / k_t.shape&#91;-2&#93;
            q_freqs, k_freqs = freqs, ratio * freqs
        else:
            ratio = k_t.shape&#91;-2&#93; / q_t.shape&#91;-2&#93;
            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, &quot;b n h d -&gt; b n (h d)&quot;)
    out, _ = self.to_out(out)
    return out</code></pre></details>

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 &lt; head_dim) with halves-swap, so
    # apply outside LocalAttention. q,k come in as &#91;B, S, H, D&#93;;
    # transpose to &#91;B, H, S, D&#93; for the helper.
    q_t = q.transpose(1, 2)
    k_t = k.transpose(1, 2)
    if q_t.shape&#91;-2&#93; &gt;= k_t.shape&#91;-2&#93;:
        ratio = q_t.shape&#91;-2&#93; / k_t.shape&#91;-2&#93;
        q_freqs, k_freqs = freqs, ratio * freqs
    else:
        ratio = k_t.shape&#91;-2&#93; / q_t.shape&#91;-2&#93;
        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, &quot;b n h d -&gt; b n (h d)&quot;)
out, _ = self.to_out(out)
return out</code></pre></details>

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) -&gt; 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&#91;torch.Tensor, float&#93; | None = None) -&gt; 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</code></pre></details>

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) -&gt; 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(&#91;
        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)
    &#93;)

def forward(self,
            x: torch.Tensor,
            prepend_embeds: torch.Tensor | None = None,
            context: torch.Tensor | None = None) -&gt; torch.Tensor:
    if isinstance(self.project_in, ReplicatedLinear):
        x, _ = self.project_in(x)
    if prepend_embeds is not None:
        assert prepend_embeds.shape&#91;-1&#93; == x.shape&#91;-1&#93;
        x = torch.cat((prepend_embeds, x), dim=-2)
    rotary = self.rotary_pos_emb.forward_from_seq_len(x.shape&#91;1&#93;)
    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</code></pre></details>

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 = &#123;&#125;

def __init__(self, config: StableAudioConfig | None = None, hf_config: dict&#91;str, Any&#93; | None = None) -&gt; None:
    if config is None:
        config = StableAudioConfig()
    super().__init__(config=config, hf_config=hf_config or &#123;&#125;)
    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 = &quot;v&quot;

    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) -&gt; 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) -&gt; torch.Tensor:
    &quot;&quot;&quot;Forward over a single batch. CFG batching is the caller&#x27;s job.&quot;&quot;&quot;
    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&#91;:, None&#93;))
    global_embed = global_embed + timestep_embed
    prepend_inputs = global_embed.unsqueeze(1)

    x = self.preprocess_conv(x) + x
    x = rearrange(x, &quot;b c t -&gt; b t c&quot;)
    out = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond)
    out = rearrange(out, &quot;b t c -&gt; b c t&quot;)&#91;:, :, prepend_inputs.shape&#91;1&#93;:&#93;
    return self.postprocess_conv(out) + out

@classmethod
def from_official_state_dict(cls,
                             state_dict: dict&#91;str, torch.Tensor&#93;,
                             prefix: str = &quot;model.model.&quot;) -&gt; &quot;StableAudioDiT&quot;:
    &quot;&quot;&quot;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.
    &quot;&quot;&quot;
    model = cls()
    mapping_fn = get_param_names_mapping(model.config.arch_config.param_names_mapping)
    remapped: dict&#91;str, torch.Tensor&#93; = &#123;&#125;
    for k, v in state_dict.items():
        if not k.startswith(prefix):
            continue
        new_key, _, _ = mapping_fn(k)
        remapped&#91;new_key&#93; = v
    missing, unexpected = model.load_state_dict(remapped, strict=True)
    if missing or unexpected:
        raise RuntimeError(f&quot;StableAudioDiT load mismatch — missing=&#123;missing&#91;:5&#93;&#125; unexpected=&#123;unexpected&#91;:5&#93;&#125;&quot;)
    return model</code></pre></details>

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&#91;:, None&#93;))
global_embed = global_embed + timestep_embed
prepend_inputs = global_embed.unsqueeze(1)

x = self.preprocess_conv(x) + x
x = rearrange(x, &quot;b c t -&gt; b t c&quot;)
out = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond)
out = rearrange(out, &quot;b t c -&gt; b c t&quot;)&#91;:, :, prepend_inputs.shape&#91;1&#93;:&#93;
return self.postprocess_conv(out) + out</code></pre></details>

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

Was this page helpful?