Text-to-Speech
Safetensors
English
steertts
steer-tts
tts
speech-synthesis
english
controllable-tts
steerable-tts
prosody-control
expressive-tts
style-control
inline-markup
word-emphasis
zero-shot-voice-cloning
voice-cloning
interpretable
non-autoregressive
fastpitch
mamba
mamba-3
state-space-model
hybrid-attention
flow-matching
vocos
small-model
quantized
int8
on-device
streaming
non-commercial
Eval Results (legacy)
| """SteerTTS acoustic model (inference path). | |
| text (ARPAbet) -> encoder -> per-phone duration / pitch / energy predictors -> length regulator -> decoder -> | |
| 100-band log-mel (Vocos-compatible). Every block is conditioned through adaptive LayerNorm on | |
| g = MLP([speaker vector, style tag, control z-vector * mask, mask]). Encoder and decoder blocks are "hybrid": | |
| multi-head self-attention followed by a bidirectional Mamba-3 branch. | |
| This is the inference subset of the training code (recipes/steer_tts/steer/model.py): the same modules and | |
| operations in the same order, so a released checkpoint gives bit-identical mels. Training-only parts (the | |
| alignment module, losses, torch.compile paths) are left out; the released weights do not contain the aligner. | |
| """ | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from .config import CONTROL_DIMS, N_MELS | |
| from .mamba3 import Mamba3Layer, scan_forward_ssd | |
| from .text import N_BASE, PAD | |
| N_CTRL = len(CONTROL_DIMS) | |
| def seq_mask(lens, max_len=None): | |
| max_len = max_len or int(lens.max()) | |
| return torch.arange(max_len, device=lens.device)[None] < lens[:, None] | |
| class SinusoidalPos(nn.Module): | |
| def __init__(self, d): | |
| super().__init__() | |
| self.d = d | |
| def forward(self, T, device): | |
| pos = torch.arange(T, device=device, dtype=torch.float32)[:, None] | |
| div = torch.exp(torch.arange(0, self.d, 2, device=device).float() * (-math.log(10000.0) / self.d)) | |
| pe = torch.zeros(T, self.d, device=device) | |
| pe[:, 0::2] = torch.sin(pos * div) | |
| pe[:, 1::2] = torch.cos(pos * div) | |
| return pe | |
| class AdaLN(nn.Module): | |
| """LayerNorm whose scale/shift are predicted from the conditioning vector g.""" | |
| def __init__(self, d, g_dim): | |
| super().__init__() | |
| self.ln = nn.LayerNorm(d, elementwise_affine=False) | |
| self.proj = nn.Linear(g_dim, 2 * d) | |
| def forward(self, x, g): # x [B,T,d], g [B,g] | |
| s, b = self.proj(g)[:, None].chunk(2, -1) | |
| return self.ln(x) * (1 + s) + b | |
| class BiMamba3(nn.Module): | |
| """Bidirectional Mamba-3 token mixer: a forward and a backward Mamba3Layer, summed. The backward pass reverses | |
| each sequence within its own length, so padding never enters the state. Runs the SSD path in fp32.""" | |
| def __init__(self, d, d_state=32, head_dim=32, rank=1, chunk=64): | |
| super().__init__() | |
| self.fwd = Mamba3Layer(d, d_state=d_state, rank=rank, head_dim=head_dim) | |
| self.bwd = Mamba3Layer(d, d_state=d_state, rank=rank, head_dim=head_dim) | |
| self.chunk = chunk | |
| def forward(self, x, mask): # x [B,T,d], mask [B,T] True=valid | |
| T = x.shape[1] | |
| t = torch.arange(T, device=x.device)[None] | |
| lens = mask.sum(1, keepdim=True) | |
| rev = torch.where(t < lens, lens - 1 - t, t)[..., None].expand(-1, -1, x.shape[2]) # an involution | |
| with torch.autocast(x.device.type, enabled=False): | |
| x = (x * mask[..., None]).float() | |
| y = scan_forward_ssd(self.fwd, x, self.chunk, vectorized=True) | |
| y = y + torch.gather(scan_forward_ssd(self.bwd, torch.gather(x, 1, rev), self.chunk, vectorized=True), 1, rev) | |
| return y * mask[..., None] | |
| class FFTBlock(nn.Module): | |
| """Pre-norm block: token mixer + conv FFN, all norms adaptive on g. | |
| mixer: "attn" (self-attention), "mamba3" (bidirectional Mamba-3), "hybrid" (attention, then a Mamba-3 branch).""" | |
| def __init__(self, d, heads, ffn, k, g_dim, p=0.1, mixer="attn", ssm_state=32): | |
| super().__init__() | |
| assert mixer in ("attn", "mamba3", "hybrid"), mixer | |
| self.mixer = mixer | |
| self.n1 = AdaLN(d, g_dim) | |
| if mixer == "mamba3": | |
| self.mix = BiMamba3(d, d_state=ssm_state) | |
| else: | |
| self.att = nn.MultiheadAttention(d, heads, dropout=p, batch_first=True) | |
| if mixer == "hybrid": | |
| self.n_ssm = AdaLN(d, g_dim) | |
| self.ssm = BiMamba3(d, d_state=ssm_state) | |
| self.n2 = AdaLN(d, g_dim) | |
| self.c1 = nn.Conv1d(d, ffn, k, padding=k // 2) | |
| self.c2 = nn.Conv1d(ffn, d, 1) | |
| self.drop = nn.Dropout(p) | |
| def forward(self, x, mask, g): # mask [B,T] True=valid | |
| h = self.n1(x, g) | |
| if self.mixer == "mamba3": | |
| h = self.mix(h, mask) | |
| else: | |
| h, _ = self.att(h, h, h, key_padding_mask=~mask, need_weights=False) | |
| x = x + self.drop(h) | |
| if self.mixer == "hybrid": | |
| x = x + self.drop(self.ssm(self.n_ssm(x, g), mask)) | |
| h = self.n2(x, g).transpose(1, 2) | |
| h = self.c2(F.gelu(self.c1(h))).transpose(1, 2) | |
| x = x + self.drop(h) | |
| return x * mask[..., None] | |
| class Stack(nn.Module): | |
| def __init__(self, n, d, heads, ffn, k, g_dim, p, mixer="attn", ssm_state=32): | |
| super().__init__() | |
| self.pos = SinusoidalPos(d) | |
| self.layers = nn.ModuleList([FFTBlock(d, heads, ffn, k, g_dim, p, mixer, ssm_state) for _ in range(n)]) | |
| self.out_norm = AdaLN(d, g_dim) | |
| def forward(self, x, mask, g): | |
| x = x + self.pos(x.shape[1], x.device)[None] | |
| for l in self.layers: | |
| x = l(x, mask, g) | |
| return self.out_norm(x, g) * mask[..., None] | |
| class VariancePredictor(nn.Module): | |
| def __init__(self, d, g_dim, k=3, p=0.15): | |
| super().__init__() | |
| self.c1 = nn.Conv1d(d, d, k, padding=k // 2) | |
| self.n1 = AdaLN(d, g_dim) | |
| self.c2 = nn.Conv1d(d, d, k, padding=k // 2) | |
| self.n2 = AdaLN(d, g_dim) | |
| self.out = nn.Linear(d, 1) | |
| self.drop = nn.Dropout(p) | |
| def forward(self, x, mask, g): | |
| h = F.relu(self.c1(x.transpose(1, 2))).transpose(1, 2) | |
| h = self.drop(self.n1(h, g)) | |
| h = F.relu(self.c2(h.transpose(1, 2))).transpose(1, 2) | |
| h = self.drop(self.n2(h, g)) | |
| return self.out(h).squeeze(-1) * mask | |
| def length_regulate(x, dur, max_len=None): | |
| """x [B,N,d], dur [B,N] int -> [B,T,d], mel_lens.""" | |
| lens = dur.sum(1) | |
| T = int(max_len or lens.max()) | |
| idx = torch.cumsum(dur, 1) | |
| frames = torch.arange(T, device=x.device)[None] | |
| tok = (frames[:, :, None] >= idx[:, None, :]).sum(-1).clamp(max=x.shape[1] - 1) | |
| out = torch.gather(x, 1, tok[..., None].expand(-1, -1, x.shape[2])) | |
| return out * seq_mask(lens, T)[..., None], lens | |
| class SteerTTSModel(nn.Module): | |
| """Arguments mirror the training class; `arch` in config.json holds the non-default ones.""" | |
| def __init__(self, n_speakers, n_tags, d=192, heads=2, ffn=384, k=3, enc_layers=4, dec_layers=4, g_dim=128, | |
| spk_dim=64, tag_dim=32, p=0.1, enc_mixer="attn", dec_mixer="attn", ssm_state=32, spk_enc_dim=0, | |
| n_symbols=None, n_env=0, ctrl_dims=None, no_mask=False, ctrl_norm="spk", level_norm=True): | |
| super().__init__() | |
| assert ctrl_dims is None and not no_mask and ctrl_norm == "spk", "ablation checkpoints are not supported" | |
| n_symbols = n_symbols or N_BASE | |
| self.emb = nn.Embedding(n_symbols, d, padding_idx=PAD) | |
| self.spk = nn.Embedding(n_speakers + 1, spk_dim) # training-speaker table (unused when a speaker vector is given) | |
| self.spk_proj = nn.Linear(spk_enc_dim, spk_dim) if spk_enc_dim else None # speaker-verification vector -> slot | |
| self.tag = nn.Embedding(n_tags, tag_dim) # index 0 = 'none' | |
| self.cond = nn.Sequential(nn.Linear(spk_dim + tag_dim + 2 * N_CTRL, 256), nn.SiLU(), nn.Linear(256, g_dim)) | |
| # acoustic-environment input of the training code; present in the weights but never trained to do anything | |
| # useful in this release, so the pipeline does not expose it | |
| self.env_proj = nn.Sequential(nn.Linear(n_env, 128), nn.SiLU(), nn.Linear(128, g_dim)) if n_env else None | |
| self.encoder = Stack(enc_layers, d, heads, ffn, k, g_dim, p, enc_mixer, ssm_state) | |
| self.decoder = Stack(dec_layers, d, heads, ffn, k, g_dim, p, dec_mixer, ssm_state) | |
| self.dur_pred = VariancePredictor(d, g_dim) | |
| self.pitch_pred = VariancePredictor(d, g_dim) | |
| self.energy_pred = VariancePredictor(d, g_dim) | |
| self.pitch_emb = nn.Conv1d(1, d, 3, padding=1) | |
| self.energy_emb = nn.Conv1d(1, d, 3, padding=1) | |
| self.to_mel = nn.Linear(d, N_MELS) | |
| def condition(self, spk, tag, ctrl, cmask, spk_vec=None, env=None): | |
| """ctrl [B,12] z-values, cmask [B,12] 1 = this dim is requested, 0 = unspecified (the model picks it).""" | |
| c = torch.nan_to_num(ctrl) * cmask | |
| s = self.spk_proj(spk_vec) if (self.spk_proj is not None and spk_vec is not None) else self.spk(spk) | |
| g = self.cond(torch.cat([s, self.tag(tag), c, cmask], -1)) | |
| if self.env_proj is not None and env is not None: | |
| g = g + self.env_proj(env) | |
| return g | |
| def encode(self, ids, text_lens, g): | |
| tmask = seq_mask(text_lens, ids.shape[1]) | |
| e = self.emb(ids) | |
| return e, self.encoder(e, tmask, g), tmask | |
| def infer(self, ids, spk, tag, ctrl, cmask, pace=1.0, word_idx=None, word_shift=None, spk_vec=None, env=None, | |
| guidance=1.0, base=None): | |
| """ids [1,N]. word_shift: {word_index: dict(pitch=st/PITCH_SCALE, energy=dB/10, dur=factor)}. | |
| guidance s != 1 with base=(tag0, ctrl0, cmask0): out = out(base) + s * (out(request) - out(base)) on every | |
| stage (durations first, so both passes share one timeline).""" | |
| g = self.condition(spk, tag, ctrl, cmask, spk_vec, env) | |
| g0 = self.condition(spk, *base, spk_vec, env) if (base is not None and guidance != 1.0) else None | |
| text_lens = torch.tensor([ids.shape[1]], device=ids.device) | |
| _, h, tmask = self.encode(ids, text_lens, g) | |
| logdur, pitch, energy = self.dur_pred(h, tmask, g), self.pitch_pred(h, tmask, g), self.energy_pred(h, tmask, g) | |
| if g0 is not None: | |
| _, h0, _ = self.encode(ids, text_lens, g0) | |
| mix = lambda a, a0: a0 + guidance * (a - a0) | |
| logdur = mix(logdur, self.dur_pred(h0, tmask, g0)) | |
| pitch = mix(pitch, self.pitch_pred(h0, tmask, g0)) | |
| energy = mix(energy, self.energy_pred(h0, tmask, g0)) | |
| dur = (torch.exp(logdur) - 1).clamp_min(0) / pace | |
| if word_shift and word_idx is not None: | |
| for w, sh in word_shift.items(): | |
| sel = word_idx[0] == w | |
| pitch[0, sel] += sh.get("pitch", 0.0) | |
| energy[0, sel] += sh.get("energy", 0.0) | |
| dur[0, sel] *= sh.get("dur", 1.0) | |
| dur = torch.round(dur).long().clamp_min(0) | |
| dur[:, 0] = dur[:, 0].clamp_min(1) | |
| pe = self.pitch_emb(pitch[:, None]).transpose(1, 2) + self.energy_emb(energy[:, None]).transpose(1, 2) | |
| x, lens = length_regulate(h + pe, dur) | |
| mel = self.to_mel(self.decoder(x, seq_mask(lens, x.shape[1]), g)).transpose(1, 2) | |
| if g0 is not None: | |
| x0, _ = length_regulate(h0 + pe, dur) | |
| mel0 = self.to_mel(self.decoder(x0, seq_mask(lens, x0.shape[1]), g0)).transpose(1, 2) | |
| mel = mel0 + guidance * (mel - mel0) | |
| return mel, dict(dur=dur, pitch=pitch, energy=energy, g=g) | |