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.
fastvideo.models.encoders.stable_audio_conditioner.StableAudioConditionerConfig
StableAudioConditionerConfig(arch_config: ArchConfig = StableAudioConditionerArchConfig(), prefix: str = 'stable_audio_conditioner', quant_config: QuantizationConfig | None = None, lora_config: Any | None = None, *, _resolved_attention_backend: AttentionBackendEnum | None = None)fastvideo.models.encoders.stable_audio_conditioner.NumberEmbedder
NumberEmbedder(features: int, dim: int = 256)Source
class NumberEmbedder(nn.Module):
def __init__(self, features: int, dim: int = 256) -> None:
super().__init__()
self.features = features
self.embedding = _time_positional_embedding(dim=dim, out_features=features)
def forward(self, x: torch.Tensor | list[float]) -> torch.Tensor:
if not torch.is_tensor(x):
device = next(self.embedding.parameters()).device
x = torch.tensor(x, device=device)
shape = x.shape
x = rearrange(x, "... -> (...)")
out = self.embedding(x)
return out.view(*shape, self.features)</code></pre></details>
fastvideo.models.encoders.stable_audio_conditioner.NumberEmbedder.forward
forward(x: torch.Tensor | list[float]) -> torch.Tensor
Source
def forward(self, x: torch.Tensor | list[float]) -> torch.Tensor:
if not torch.is_tensor(x):
device = next(self.embedding.parameters()).device
x = torch.tensor(x, device=device)
shape = x.shape
x = rearrange(x, "... -> (...)")
out = self.embedding(x)
return out.view(*shape, self.features)
fastvideo.models.encoders.stable_audio_conditioner.T5Conditioner
T5Conditioner(output_dim: int, t5_model_name: str = 't5-base', max_length: int = 128, dtype: str = 'float16')
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
class T5Conditioner(_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.
"""
T5_MODEL_DIMS = {"t5-base": 768}
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))
def forward(self, texts: list[str], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]:
encoded = self.tokenizer(texts,
truncation=True,
max_length=self.max_length,
padding="max_length",
return_tensors="pt")
input_ids = encoded["input_ids"].to(device)
attention_mask = encoded["attention_mask"].to(device).to(torch.bool)
# Mirror official's `autocast(fp16)` wrap on T5 forward.
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self._t5_dtype):
embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"]
embeddings = self.proj_out(embeddings) * attention_mask.unsqueeze(-1).float()
return embeddings, attention_mask</code></pre></details>
fastvideo.models.encoders.stable_audio_conditioner.T5Conditioner.forward
forward(texts: list[str], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(self, texts: list[str], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]:
encoded = self.tokenizer(texts,
truncation=True,
max_length=self.max_length,
padding="max_length",
return_tensors="pt")
input_ids = encoded["input_ids"].to(device)
attention_mask = encoded["attention_mask"].to(device).to(torch.bool)
# Mirror official's `autocast(fp16)` wrap on T5 forward.
with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self._t5_dtype):
embeddings = self.model(input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"]
embeddings = self.proj_out(embeddings) * attention_mask.unsqueeze(-1).float()
return embeddings, attention_mask
fastvideo.models.encoders.stable_audio_conditioner.NumberConditioner
NumberConditioner(output_dim: int, min_val: float = 0, max_val: float = 1)
Float-valued conditioner with min/max clamping + NumberEmbedder.
Source
class NumberConditioner(_Conditioner):
"""Float-valued conditioner with min/max clamping + NumberEmbedder."""
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)
def forward(self, floats: list[float], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]:
floats = [float(x) for x in floats]
floats_t = torch.tensor(floats, device=device).clamp(self.min_val, self.max_val)
normalized = (floats_t - self.min_val) / (self.max_val - self.min_val)
emb_dtype = next(self.embedder.parameters()).dtype
normalized = normalized.to(emb_dtype)
float_embeds = self.embedder(normalized).unsqueeze(1)
return float_embeds, torch.ones(float_embeds.shape[0], 1, device=device)</code></pre></details>
fastvideo.models.encoders.stable_audio_conditioner.NumberConditioner.forward
forward(floats: list[float], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]
Source
def forward(self, floats: list[float], device: torch.device | str) -> tuple[torch.Tensor, torch.Tensor]:
floats = [float(x) for x in floats]
floats_t = torch.tensor(floats, device=device).clamp(self.min_val, self.max_val)
normalized = (floats_t - self.min_val) / (self.max_val - self.min_val)
emb_dtype = next(self.embedder.parameters()).dtype
normalized = normalized.to(emb_dtype)
float_embeds = self.embedder(normalized).unsqueeze(1)
return float_embeds, torch.ones(float_embeds.shape[0], 1, device=device)
fastvideo.models.encoders.stable_audio_conditioner.StableAudioMultiConditioner
StableAudioMultiConditioner(config: StableAudioConditionerConfig | None = None)
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
class StableAudioMultiConditioner(nn.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`.
"""
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)
def forward(self, batch_metadata: list[dict],
device: torch.device | str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
out: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
for key, conditioner in self.conditioners.items():
inputs = [x[key] for x in batch_metadata]
out[key] = conditioner(inputs, device)
return out
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
@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</code></pre></details>
fastvideo.models.encoders.stable_audio_conditioner.StableAudioMultiConditioner.forward
forward(batch_metadata: list[dict], device: torch.device | str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]
Source
def forward(self, batch_metadata: list[dict],
device: torch.device | str) -> dict[str, tuple[torch.Tensor, torch.Tensor]]:
out: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
for key, conditioner in self.conditioners.items():
inputs = [x[key] for x in batch_metadata]
out[key] = conditioner(inputs, device)
return out
fastvideo.models.encoders.stable_audio_conditioner.StableAudioMultiConditioner.get_conditioning_inputs
get_conditioning_inputs(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).
Source
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.StableAudioMultiConditioner.from_official_state_dict
from_official_state_dict(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.
Source
@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