Config for the Stable Audio Open 1.0 DiT.

Note: the SA pipeline bypasses the standard ComposedPipelineBase component loader because the published HF repo ships a single monolithic model.safetensors (no Diffusers-style model_index.json or per-subfolder layout). The arch fields and param_names_mapping here document the architecture and key remap so the same conventions used by the rest of the DiT family apply (FSDP shard conditions, supported attention backends, future loader integrations) — they are not currently consumed by fastvideo/models/loader/fsdp_load.py for SA.

`fastvideo.configs.models.dits.stable_audio.DiTArchConfig`

DiTArchConfig(stacked_params_mapping: list[tuple[str, str, str]] = list(), _fsdp_shard_conditions: list = list(), _compile_conditions: list = list(), param_names_mapping: dict = dict(), reverse_param_names_mapping: dict = dict(), lora_param_names_mapping: dict = dict(), cast_prompt_embeds_to_dit_dtype: bool = False, _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.SAGE_ATTN_THREE, AttentionBackendEnum.ATTN_QAT_INFER, AttentionBackendEnum.ATTN_QAT_TRAIN, AttentionBackendEnum.SLA_ATTN, AttentionBackendEnum.SAGE_SLA_ATTN), hidden_size: int = 0, num_attention_heads: int = 0, num_channels_latents: int = 0, in_channels: int = 0, out_channels: int = 0, exclude_lora_layers: list[str] = list(), boundary_ratio: float | None = None)

`fastvideo.configs.models.dits.stable_audio.DiTConfig`

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

`fastvideo.configs.models.dits.stable_audio.DiTConfig.add_cli_args`

add_cli_args(parser: Any, prefix: str = 'dit-config') -> Any

Add CLI arguments for DiTConfig fields

`fastvideo.configs.models.dits.stable_audio.AttentionBackendEnum`

AttentionBackendEnum()

`fastvideo.configs.models.dits.stable_audio.StableAudioArchConfig`

StableAudioArchConfig(stacked_params_mapping: list[tuple[str, str, str]] = list(), _fsdp_shard_conditions: list = (lambda: [_is_transformer_layer])(), _compile_conditions: list = list(), param_names_mapping: dict = (lambda: {'^model\\.model\\.(.*?)\\.gamma$': '\\1.weight', '^model\\.model\\.(.*?)\\.beta$': '\\1.bias', '^model\\.model\\.(.*)$': '\\1'})(), reverse_param_names_mapping: dict = dict(), lora_param_names_mapping: dict = dict(), cast_prompt_embeds_to_dit_dtype: bool = False, _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA), hidden_size: int = 0, num_attention_heads: int = 24, num_channels_latents: int = 0, in_channels: int = 0, out_channels: int = 0, exclude_lora_layers: list[str] = list(), boundary_ratio: float | None = None, io_channels: int = 64, embed_dim: int = 1536, depth: int = 24, cond_token_dim: int = 768, global_cond_dim: int = 1536, project_cond_tokens: bool = False, project_global_cond: bool = True, qk_norm: str | None = None)

Source

@dataclass
class StableAudioArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_transformer_layer])

    # SA's checkpoint is `stable_audio_tools` raw format (not Diffusers),
    # so the only remaps are: strip the `model.model.` host-pipeline
    # prefix, and rename `nn.LayerNorm`'s `gamma`/`beta` to torch's
    # canonical `weight`/`bias`. Linear / cross-attention naming already
    # matches FastVideo's conventions, so no further remap is needed.
    param_names_mapping: dict = field(
        default_factory=lambda: {
            r"^model\.model\.(.*?)\.gamma$": r"\1.weight",
            r"^model\.model\.(.*?)\.beta$": r"\1.bias",
            r"^model\.model\.(.*)$": r"\1",
        })

    # SA only supports backends compatible with single-GPU LocalAttention.
    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.TORCH_SDPA,
    )

    # Architecture constants (from the published `model_config.json` for
    # `stabilityai/stable-audio-open-1.0`).
    io_channels: int = 64
    embed_dim: int = 1536
    depth: int = 24
    num_attention_heads: int = 24
    cond_token_dim: int = 768
    global_cond_dim: int = 1536
    project_cond_tokens: bool = False
    project_global_cond: bool = True
    # Set to "ln" to wrap attention Q/K in LayerNorm (used by
    # `stable-audio-open-small`; absent in the 1.0 base).
    qk_norm: str | None = None

    def __post_init__(self) -> None:
        super().__post_init__()
        self.hidden_size = self.embed_dim
        self.in_channels = self.io_channels
        self.out_channels = self.io_channels
        self.num_channels_latents = self.io_channels
        self.attention_head_dim = self.embed_dim // self.num_attention_heads

`fastvideo.configs.models.dits.stable_audio.StableAudioConfig`

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

Source

@dataclass
class StableAudioConfig(DiTConfig):
    arch_config: DiTArchConfig = field(default_factory=StableAudioArchConfig)

    prefix: str = "StableAudio"