Skip to content

stable_audio_conditioner

Stable Audio Open 1.0 conditioner.

T5-base text encoder + two NumberConditioners (seconds_start, seconds_total), wrapped by StableAudioMultiConditioner which produces the cross-attention and global-conditioning tensors the DiT expects.

Classes

fastvideo.models.encoders.stable_audio_conditioner.NumberConditioner

NumberConditioner(output_dim: int, min_val: float = 0, max_val: float = 1)

Bases: _Conditioner

Float-valued conditioner with min/max clamping + NumberEmbedder.

Source code in fastvideo/models/encoders/stable_audio_conditioner.py
def __init__(self, output_dim: int, min_val: float = 0, max_val: float = 1) -> None:
    super().__init__(output_dim, output_dim)
    self.min_val = min_val
    self.max_val = max_val
    self.embedder = NumberEmbedder(features=output_dim)

fastvideo.models.encoders.stable_audio_conditioner.StableAudioMultiConditioner

StableAudioMultiConditioner(config: StableAudioConditionerConfig | None = None)

Bases: Module

SA-Open-1.0 conditioner: T5 prompt + duration NumberConditioners.

All hardcoded constants (cond_dim, sub-conditioner ids, T5 model name + max_length, NumberConditioner ranges) live on StableAudioConditionerConfig — see fastvideo/configs/models/encoders/stable_audio_conditioner.py.

Source code in fastvideo/models/encoders/stable_audio_conditioner.py
def __init__(self, config: StableAudioConditionerConfig | None = None) -> None:
    super().__init__()
    self.config = config or StableAudioConditionerConfig()
    arch = self.config.arch_config
    # Build sub-conditioners from the `configs` list (mirrors
    # upstream's `MultiConditioner` factory).
    sub: dict[str, nn.Module] = {}
    for spec in arch.configs:
        sid = spec["id"]
        stype = spec["type"]
        scfg = spec["config"]
        if stype == "t5":
            sub[sid] = T5Conditioner(output_dim=arch.cond_dim,
                                     t5_model_name=scfg["t5_model_name"],
                                     max_length=scfg["max_length"],
                                     dtype=arch.t5_dtype)
        elif stype == "number":
            sub[sid] = NumberConditioner(output_dim=arch.cond_dim,
                                         min_val=scfg["min_val"], max_val=scfg["max_val"])
        else:
            raise ValueError(f"Unknown sub-conditioner type {stype!r} for id {sid!r}.")
    self.conditioners = nn.ModuleDict(sub)
    self.cross_attention_cond_ids = tuple(arch.cross_attention_cond_ids)
    self.global_cond_ids = tuple(arch.global_cond_ids)

Methods:

fastvideo.models.encoders.stable_audio_conditioner.StableAudioMultiConditioner.from_official_state_dict classmethod
from_official_state_dict(state_dict: dict[str, Tensor], prefix: str = 'conditioner.') -> 'StableAudioMultiConditioner'

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

Source code in fastvideo/models/encoders/stable_audio_conditioner.py
@classmethod
def from_official_state_dict(cls, state_dict: dict[str, torch.Tensor],
                             prefix: str = "conditioner.") -> "StableAudioMultiConditioner":
    """Load NumberConditioner weights from a raw `stable_audio_tools`
    monolithic state dict. Kept for tests / older checkpoints;
    production loads go through the standard `ConditionerLoader`
    against the converted Diffusers repo.
    """
    mc = cls()
    own_state = mc.state_dict()
    loaded: dict[str, torch.Tensor] = {}
    for k, v in state_dict.items():
        if not k.startswith(prefix):
            continue
        stripped = k[len(prefix):]
        if stripped in own_state:
            loaded[stripped] = v
    # T5 keys are intentionally absent from the checkpoint.
    missing = [k for k in own_state.keys() if k not in loaded
               and not k.startswith("conditioners.prompt.")]
    unexpected = [k for k in loaded.keys() if k not in own_state]
    if missing or unexpected:
        raise RuntimeError(
            f"StableAudioMultiConditioner load mismatch — missing={missing[:5]} unexpected={unexpected[:5]}"
        )
    mc.load_state_dict(loaded, strict=False)
    return mc
fastvideo.models.encoders.stable_audio_conditioner.StableAudioMultiConditioner.get_conditioning_inputs
get_conditioning_inputs(cond: dict[str, tuple[Tensor, Tensor]]) -> tuple[Tensor, Tensor, Tensor]

Pack conditioner outputs into the (cross_attn_cond, cross_attn_mask, global_embed) triple the DiT consumes. Order is driven by cross_attention_cond_ids / global_cond_ids from the config — SA-1.0 uses three sub-conditioners (prompt + seconds_start + seconds_total); SA-small uses two (prompt + seconds_total).

Source code in fastvideo/models/encoders/stable_audio_conditioner.py
def get_conditioning_inputs(
    self, cond: dict[str, tuple[torch.Tensor, torch.Tensor]]
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Pack conditioner outputs into the (cross_attn_cond,
    cross_attn_mask, global_embed) triple the DiT consumes. Order
    is driven by `cross_attention_cond_ids` / `global_cond_ids`
    from the config — SA-1.0 uses three sub-conditioners
    (prompt + seconds_start + seconds_total); SA-small uses two
    (prompt + seconds_total).
    """
    x_embs = [cond[i][0] for i in self.cross_attention_cond_ids]
    x_masks = [cond[i][1] for i in self.cross_attention_cond_ids]
    cross_attn_cond = torch.cat(x_embs, dim=1)
    cross_attn_mask = torch.cat(x_masks, dim=1)
    global_embed = torch.cat([cond[i][0][:, 0] for i in self.global_cond_ids], dim=-1)
    return cross_attn_cond, cross_attn_mask, global_embed

fastvideo.models.encoders.stable_audio_conditioner.T5Conditioner

T5Conditioner(output_dim: int, t5_model_name: str = 't5-base', max_length: int = 128, dtype: str = 'float16')

Bases: _Conditioner

T5 text conditioner. Pads to model_max_length (=128 for the SA repo's tokenizer, NOT the standard 512) and emits a masked last-hidden-state.

Source code in fastvideo/models/encoders/stable_audio_conditioner.py
def __init__(self, output_dim: int, t5_model_name: str = "t5-base",
             max_length: int = 128, dtype: str = "float16") -> None:
    super().__init__(self.T5_MODEL_DIMS[t5_model_name], output_dim, project_out=False)
    from transformers import AutoTokenizer, T5EncoderModel
    self.max_length = max_length
    self.tokenizer = AutoTokenizer.from_pretrained(t5_model_name)
    # T5 loaded directly in fp16 (config-driven) to match official
    # `stable_audio_tools/models/conditioners.py:334`. Registered as
    # a normal submodule so `.to(device)` / `torch.compile` track it;
    # `from_official_state_dict` filters `conditioners.prompt.*` from
    # the missing-key check (T5 weights are absent from the SA
    # checkpoint by design).
    # Explicit lookup so a typo (e.g. "fp16" instead of "float16") errors
    # at load time rather than silently falling back to a wrong dtype.
    torch_dtype = getattr(torch, dtype)
    if not isinstance(torch_dtype, torch.dtype):
        raise ValueError(f"T5Conditioner dtype={dtype!r} is not a torch.dtype.")
    self._t5_dtype = torch_dtype
    self.model = (T5EncoderModel.from_pretrained(t5_model_name).eval().requires_grad_(False).to(torch_dtype))