sopro 1.0.0__py3-none-any.whl

sopro/nn/__init__.py

@@ -0,0 +1,20 @@
+from .blocks import GLU, AttentiveStatsPool, DepthwiseConv1d, RMSNorm, SSMLiteBlock
+from .embeddings import CodebookEmbedding, SinusoidalPositionalEmbedding, TextEmbedding
+from .speaker import SpeakerFiLM, Token2SV
+from .xattn import RefXAttn, RefXAttnBlock, TextXAttnBlock
+
+__all__ = [
+    "GLU",
+    "RMSNorm",
+    "DepthwiseConv1d",
+    "SSMLiteBlock",
+    "AttentiveStatsPool",
+    "SinusoidalPositionalEmbedding",
+    "TextEmbedding",
+    "CodebookEmbedding",
+    "Token2SV",
+    "SpeakerFiLM",
+    "RefXAttn",
+    "RefXAttnBlock",
+    "TextXAttnBlock",
+]

sopro/nn/blocks.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class GLU(nn.Module):
+    def __init__(self, d: int):
+        super().__init__()
+        self.pro = nn.Linear(d, 2 * d)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a, b = self.pro(x).chunk(2, dim=-1)
+        return a * torch.sigmoid(b)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x32 = x.float()
+        var = x32.pow(2).mean(dim=-1, keepdim=True)
+        y32 = x32 * torch.rsqrt(var + self.eps)
+        y32 = y32 * self.weight.float()
+        return y32.to(dtype=x.dtype)
+
+
+class DepthwiseConv1d(nn.Module):
+    def __init__(
+        self, d: int, kernel_size: int = 7, causal: bool = False, dilation: int = 1
+    ):
+        super().__init__()
+        self.causal = causal
+        self.dilation = dilation
+        self.kernel_size = kernel_size
+        self.dw = nn.Conv1d(d, d, kernel_size, groups=d, padding=0, dilation=dilation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        xt = x.transpose(1, 2)
+        if self.causal:
+            pad_left = (self.kernel_size - 1) * self.dilation
+            xt = F.pad(xt, (pad_left, 0))
+        else:
+            total = (self.kernel_size - 1) * self.dilation
+            left = total // 2
+            right = total - left
+            xt = F.pad(xt, (left, right))
+        y = self.dw(xt)
+        return y.transpose(1, 2)
+
+
+class SSMLiteBlock(nn.Module):
+    def __init__(
+        self,
+        d: int,
+        dropout: float = 0.05,
+        causal: bool = False,
+        kernel_size: int = 7,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        self.norm = RMSNorm(d)
+        self.glu = GLU(d)
+        self.dw = DepthwiseConv1d(
+            d, kernel_size=kernel_size, causal=causal, dilation=dilation
+        )
+        self.ff = nn.Sequential(
+            RMSNorm(d),
+            nn.Linear(d, 4 * d),
+            nn.GELU(),
+            nn.Linear(4 * d, d),
+        )
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.glu(self.norm(x))
+        h = self.dw(h)
+        x = x + self.drop(h)
+        x = x + self.drop(self.ff(x))
+        return x
+
+
+class AttentiveStatsPool(nn.Module):
+    def __init__(self, d: int):
+        super().__init__()
+        self.attn = nn.Sequential(
+            nn.Linear(d, d),
+            nn.Tanh(),
+            nn.Linear(d, 1),
+        )
+
+    def forward(
+        self, h: torch.Tensor, lengths: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        B, T, D = h.shape
+        logits = self.attn(h).squeeze(-1)
+
+        if lengths is not None:
+            mask = torch.arange(T, device=h.device)[None, :] < lengths[:, None]
+            logits = logits.masked_fill(~mask, -1e9)
+
+        w = torch.softmax(logits, dim=1).unsqueeze(-1)
+        mu = (h * w).sum(dim=1)
+        var = (w * (h - mu.unsqueeze(1)).pow(2)).sum(dim=1).clamp_min(1e-6)
+        std = torch.sqrt(var)
+        return torch.cat([mu, std], dim=-1)

sopro/nn/embeddings.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import math
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SinusoidalPositionalEmbedding(nn.Module):
+    def __init__(self, d_model: int, max_len: int = 10000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, d_model, 2, dtype=torch.float32)
+            * (-math.log(10000.0) / d_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe, persistent=False)
+
+    def forward(self, positions: torch.Tensor) -> torch.Tensor:
+        return self.pe.index_select(0, positions.long())
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, vocab_size: int, d_model: int):
+        super().__init__()
+        self.emb = nn.Embedding(vocab_size, d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.emb(x)
+
+
+class CodebookEmbedding(nn.Module):
+    def __init__(
+        self, num_codebooks: int, vocab_size: int, d_model: int, use_bos: bool = True
+    ):
+        super().__init__()
+        self.Q = int(num_codebooks)
+        self.V = int(vocab_size)
+        self.D = int(d_model)
+        self.use_bos = bool(use_bos)
+
+        table_size = self.Q * self.V + (1 if self.use_bos else 0)
+        self.emb = nn.Embedding(table_size, d_model)
+        self.bos_id = (self.Q * self.V) if self.use_bos else None
+
+    def _indices_for(self, tokens: torch.Tensor, cb_index: int) -> torch.Tensor:
+        return cb_index * self.V + tokens
+
+    def embed_tokens(self, tokens: torch.Tensor, cb_index: int) -> torch.Tensor:
+        return self.emb(self._indices_for(tokens, cb_index))
+
+    def embed_shift_by_k(
+        self, tokens: torch.Tensor, cb_index: int, k: int
+    ) -> torch.Tensor:
+        idx = self._indices_for(tokens, cb_index)
+        B, T = idx.shape
+        if (not self.use_bos) or (self.bos_id is None) or k <= 0:
+            pad_tok = idx[:, :1]
+        else:
+            pad_tok = torch.full(
+                (B, 1), self.bos_id, dtype=torch.long, device=idx.device
+            )
+
+        if k >= T:
+            idx_shift = pad_tok.expand(-1, T)
+        else:
+            pad = pad_tok.expand(-1, k)
+            idx_shift = torch.cat([pad, idx[:, :-k]], dim=1)
+
+        return self.emb(idx_shift)
+
+    def sum_embed_subset(
+        self,
+        tokens_subset: Optional[torch.Tensor],
+        cb_indices: Optional[List[int]],
+        keep_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if tokens_subset is None or cb_indices is None or len(cb_indices) == 0:
+            return 0.0
+
+        B, T, K = tokens_subset.shape
+        idx_list = []
+        for k, cb in enumerate(cb_indices):
+            idx_list.append(self._indices_for(tokens_subset[..., k], cb).unsqueeze(2))
+        idx = torch.cat(idx_list, dim=2)
+        emb = self.emb(idx)
+
+        if keep_mask is not None:
+            emb = emb * keep_mask.unsqueeze(-1).to(emb.dtype)
+
+        return emb.sum(dim=2)

sopro/nn/speaker.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .blocks import AttentiveStatsPool, DepthwiseConv1d
+
+
+class Token2SV(nn.Module):
+    def __init__(
+        self, Q: int, V: int, d: int = 192, out_dim: int = 256, dropout: float = 0.05
+    ):
+        super().__init__()
+        self.Q, self.V = int(Q), int(V)
+        self.emb = nn.Embedding(self.Q * self.V, d)
+
+        initial_weights = torch.linspace(1.0, 0.1, self.Q)
+        self.cb_weights = nn.Parameter(initial_weights)
+
+        self.enc = nn.Sequential(
+            DepthwiseConv1d(d, 7, causal=False),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            DepthwiseConv1d(d, 7, causal=False),
+            nn.GELU(),
+        )
+
+        self.pool = AttentiveStatsPool(d)
+        self.proj = nn.Linear(2 * d, out_dim)
+
+    def _get_mixed_embedding(self, embed_btqd: torch.Tensor) -> torch.Tensor:
+        w = F.softmax(self.cb_weights, dim=0).view(1, 1, self.Q, 1)
+        return (embed_btqd * w).sum(dim=2)
+
+    def forward(
+        self, tokens_btq: torch.Tensor, lengths: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        B, T, Q = tokens_btq.shape
+        q_idx = torch.arange(Q, device=tokens_btq.device, dtype=torch.long).view(
+            1, 1, Q
+        )
+        idx = q_idx * self.V + tokens_btq.long()
+        raw_emb = self.emb(idx)
+
+        if self.training:
+            keep_prob = 0.95
+            mask = torch.rand(B, T, device=tokens_btq.device) < keep_prob
+            bad = mask.sum(dim=1) == 0
+            if bad.any():
+                bad_idx = bad.nonzero(as_tuple=False).squeeze(1)
+                rand_pos = torch.randint(
+                    0, T, (bad_idx.numel(),), device=tokens_btq.device
+                )
+                mask[bad_idx, rand_pos] = True
+            raw_emb = raw_emb * mask.float().unsqueeze(-1).unsqueeze(-1)
+
+        x = self._get_mixed_embedding(raw_emb)
+        h = self.enc(x)
+        pooled = self.pool(h, lengths=lengths)
+        e = self.proj(pooled)
+        return F.normalize(e, dim=-1, eps=1e-6)
+
+
+class SpeakerFiLM(nn.Module):
+    def __init__(self, d_model: int, sv_dim: int):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(sv_dim, d_model),
+            nn.GELU(),
+            nn.Linear(d_model, 2 * d_model),
+        )
+        self.norm = nn.LayerNorm(d_model)
+        nn.init.zeros_(self.mlp[-1].weight)
+        nn.init.zeros_(self.mlp[-1].bias)
+
+    def forward(
+        self, base_bt_d: torch.Tensor, spk_b_d: torch.Tensor, strength: float = 1.0
+    ) -> torch.Tensor:
+        B, T, D = base_bt_d.shape
+        film = self.mlp(spk_b_d)
+        gamma, beta = film.chunk(2, dim=-1)
+        gamma = gamma.unsqueeze(1).expand(B, T, D)
+        beta = beta.unsqueeze(1).expand(B, T, D)
+        x = self.norm(base_bt_d)
+        return x * (1 + strength * torch.tanh(gamma)) + strength * torch.tanh(beta)

sopro/nn/xattn.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .blocks import RMSNorm
+
+
+def rms(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
+    return torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)
+
+
+class RefXAttnBlock(nn.Module):
+    def __init__(self, d_model: int, heads: int = 2, dropout: float = 0.0):
+        super().__init__()
+        self.nq = RMSNorm(d_model)
+        self.nkv = RMSNorm(d_model)
+        self.attn = nn.MultiheadAttention(
+            d_model, heads, batch_first=True, dropout=dropout
+        )
+        self.gate = nn.Parameter(torch.tensor(0.5))
+        self.gmax = 0.35
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        ref: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        q = self.nq(x)
+        kv = self.nkv(ref.float())
+
+        with torch.autocast(device_type=q.device.type, enabled=False):
+            a, _ = self.attn(
+                q.float(), kv, kv, key_padding_mask=key_padding_mask, need_weights=False
+            )
+
+        a = torch.nan_to_num(a, nan=0.0, posinf=0.0, neginf=0.0)
+
+        rms_x = rms(x.float())
+        rms_a = rms(a)
+        scale = rms_x / rms_a
+        a = (a * scale).to(x.dtype)
+
+        gate_eff = (self.gmax * torch.tanh(self.gate)).to(x.dtype)
+        return x + gate_eff * a
+
+
+class RefXAttn(nn.Module):
+    def __init__(
+        self, d_model: int, heads: int = 2, layers: int = 3, dropout: float = 0.0
+    ):
+        super().__init__()
+        self.blocks = nn.ModuleList(
+            [RefXAttnBlock(d_model, heads, dropout) for _ in range(layers)]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        ref: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        for blk in self.blocks:
+            x = blk(x, ref, key_padding_mask)
+        return x
+
+
+class TextXAttnBlock(nn.Module):
+    def __init__(self, d_model: int, heads: int = 4, dropout: float = 0.0):
+        super().__init__()
+        self.nq = RMSNorm(d_model)
+        self.nkv = RMSNorm(d_model)
+        self.attn = nn.MultiheadAttention(
+            d_model, num_heads=heads, dropout=dropout, batch_first=True
+        )
+        self.gate = nn.Parameter(torch.tensor(0.0))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        q = self.nq(x)
+        kv = self.nkv(context)
+        with torch.autocast(device_type=q.device.type, enabled=False):
+            out, _ = self.attn(
+                q.float(),
+                kv.float(),
+                kv.float(),
+                key_padding_mask=key_padding_mask,
+                need_weights=False,
+            )
+        out = torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0).to(x.dtype)
+        return x + torch.tanh(self.gate) * out

sopro/sampling.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import List
+
+import torch
+
+
+def center_crop_tokens(ref_tq: torch.Tensor, win_frames: int) -> torch.Tensor:
+    T = int(ref_tq.size(0))
+    if T <= win_frames:
+        return ref_tq
+    s = (T - win_frames) // 2
+    return ref_tq[s : s + win_frames]
+
+
+def repeated_tail(hist: List[int], max_n: int = 16) -> bool:
+    L = len(hist)
+    for n in range(3, min(max_n, L // 2) + 1):
+        if hist[-n:] == hist[-2 * n : -n]:
+            return True
+    return False
+
+
+def sample_token(
+    logits_1x1v: torch.Tensor,
+    history: List[int],
+    top_p: float = 0.9,
+    top_k: int = 0,
+    temperature: float = 1.0,
+    repetition_penalty: float = 1.0,
+    eps: float = 1e-12,
+) -> int:
+    x = logits_1x1v
+
+    x = torch.nan_to_num(x, nan=-1e9, posinf=1e9, neginf=-1e9)
+
+    if temperature and temperature != 1.0:
+        x = x / float(temperature)
+
+    if repetition_penalty != 1.0 and len(history) > 0:
+        context = history[-50:]
+        ids = torch.tensor(list(set(context)), device=x.device, dtype=torch.long)
+        if ids.numel() > 0:
+            vals = x[0, 0, ids]
+            neg = vals < 0
+            vals = torch.where(
+                neg, vals * repetition_penalty, vals / repetition_penalty
+            )
+            x = x.clone()
+            x[0, 0, ids] = vals
+
+    probs = torch.softmax(x, dim=-1).view(1, -1)
+    probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+
+    V = int(probs.size(-1))
+    if top_k and top_k > 0:
+        k = min(int(top_k), V)
+        val, idx = torch.topk(probs, k, dim=-1)
+        newp = torch.zeros_like(probs)
+        newp.scatter_(1, idx, val)
+        probs = newp
+
+        s = probs.sum(dim=-1, keepdim=True)
+        if float(s.item()) <= eps:
+            return int(torch.argmax(x[0, 0]).item())
+        probs = probs / s
+
+    if top_p is not None and top_p < 1.0:
+        sorted_probs, sorted_idx = torch.sort(probs, descending=True, dim=-1)
+        cum = torch.cumsum(sorted_probs, dim=-1)
+
+        remove = cum > float(top_p)
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = False
+
+        sorted_probs = sorted_probs.masked_fill(remove, 0.0)
+
+        s = sorted_probs.sum(dim=-1, keepdim=True)
+        if float(s.item()) <= eps:
+            return int(torch.argmax(x[0, 0]).item())
+        sorted_probs = sorted_probs / s
+
+        j = torch.multinomial(sorted_probs, 1).item()
+        token = int(sorted_idx[0, j].item())
+
+        return token
+
+    s = probs.sum(dim=-1, keepdim=True)
+    if float(s.item()) <= eps:
+        return int(torch.argmax(x[0, 0]).item())
+    probs = probs / s
+
+    return int(torch.multinomial(probs, 1).item())
+
+
+def rf_ar(ar_kernel: int, dilations: Tuple[int, ...]) -> int:
+    return 1 + (ar_kernel - 1) * int(sum(dilations))
+
+
+def rf_nar(n_layers_nar: int, kernel_size: int = 7, dilation: int = 1) -> int:
+    return 1 + (kernel_size - 1) * int(n_layers_nar) * int(dilation)

sopro/streaming.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Iterator, List, Optional, Tuple
+
+import torch
+
+from .codec.mimi import MimiDecodeState, MimiStreamDecoder
+from .model import SoproTTS, SoproTTSModel
+
+
+@dataclass
+class StreamConfig:
+    chunk_frames: int = 16
+    nar_context_frames: Optional[int] = None
+    cond_chunk_size: int = 32
+
+
+class SoproTTSStreamer:
+    def __init__(self, tts: SoproTTS, cfg: Optional[StreamConfig] = None):
+        self.tts = tts
+        self.cfg = cfg or StreamConfig()
+        self.mimi_stream = MimiStreamDecoder(tts.codec)
+
+    @torch.inference_mode()
+    def stream(
+        self,
+        text: str,
+        *,
+        ref_audio_path: Optional[str] = None,
+        ref_tokens_tq: Optional[torch.Tensor] = None,
+        max_frames: int = 400,
+        top_p: float = 0.9,
+        temperature: float = 1.05,
+        anti_loop: bool = True,
+        use_prefix: bool = True,
+        prefix_sec_fixed: Optional[float] = None,
+        style_strength: Optional[float] = None,
+        ref_seconds: Optional[float] = None,
+        chunk_frames: Optional[int] = None,
+        nar_context_frames: Optional[int] = None,
+        cond_chunk_size: Optional[int] = None,
+        use_stop_head: Optional[bool] = None,
+        stop_patience: Optional[int] = None,
+        stop_threshold: Optional[float] = None,
+        min_gen_frames: Optional[int] = None,
+    ) -> Iterator[torch.Tensor]:
+        model: SoproTTSModel = self.tts.model
+        device = self.tts.device
+
+        text_ids = self.tts.encode_text(text)
+        ref = self.tts.encode_reference(
+            ref_audio_path=ref_audio_path,
+            ref_tokens_tq=ref_tokens_tq,
+            ref_seconds=ref_seconds,
+        )
+
+        prep = model.prepare_conditioning_lazy(
+            text_ids,
+            ref,
+            max_frames=max_frames,
+            device=device,
+            style_strength=float(
+                style_strength
+                if style_strength is not None
+                else self.tts.cfg.style_strength
+            ),
+        )
+
+        cf = int(chunk_frames if chunk_frames is not None else self.cfg.chunk_frames)
+        cond_cs = int(
+            cond_chunk_size if cond_chunk_size is not None else self.cfg.cond_chunk_size
+        )
+
+        nar_ctx = (
+            nar_context_frames
+            if nar_context_frames is not None
+            else self.cfg.nar_context_frames
+        )
+        if nar_ctx is None:
+            nar_ctx = int(model.rf_nar())
+        nar_ctx = int(nar_ctx)
+
+        hist_A: List[int] = []
+
+        frames_emitted = 0
+
+        mimi_state = MimiDecodeState()
+
+        def refine_and_emit(end: int) -> Optional[torch.Tensor]:
+            nonlocal frames_emitted, mimi_state
+
+            new_start = frames_emitted
+            if end <= new_start:
+                return None
+
+            win_start = max(0, new_start - nar_ctx)
+            win_end = end
+
+            cond_win = prep["cond_all"][:, win_start:win_end, :]
+            tokens_A_win = torch.as_tensor(
+                hist_A[win_start:win_end], device=device, dtype=torch.long
+            ).unsqueeze(0)
+
+            tokens_win_tq = model.nar_refine(cond_win, tokens_A_win).squeeze(0)
+
+            tail_i = new_start - win_start
+            emit_tokens = tokens_win_tq[tail_i:, :]
+
+            wav_chunk, mimi_state = self.mimi_stream.decode_step(
+                emit_tokens, mimi_state
+            )
+            frames_emitted = end
+            return wav_chunk if wav_chunk.numel() > 0 else None
+
+        for _t, rvq1_id, _p_stop in model.ar_stream(
+            prep,
+            max_frames=max_frames,
+            top_p=top_p,
+            temperature=temperature,
+            anti_loop=anti_loop,
+            use_prefix=use_prefix,
+            prefix_sec_fixed=prefix_sec_fixed,
+            cond_chunk_size=cond_cs,
+            use_stop_head=use_stop_head,
+            stop_patience=stop_patience,
+            stop_threshold=stop_threshold,
+            min_gen_frames=min_gen_frames,
+        ):
+            hist_A.append(int(rvq1_id))
+            T = len(hist_A)
+
+            is_boundary = (T % cf) == 0
+            if not is_boundary and T < max_frames:
+                continue
+
+            wav = refine_and_emit(T)
+            if wav is not None:
+                yield wav
+
+        T_final = len(hist_A)
+        if frames_emitted < T_final:
+            wav = refine_and_emit(T_final)
+            if wav is not None:
+                yield wav
+
+
+def stream(
+    tts: SoproTTS,
+    text: str,
+    *,
+    ref_audio_path: Optional[str] = None,
+    ref_tokens_tq: Optional[torch.Tensor] = None,
+    chunk_frames: int = 6,
+    **kwargs,
+) -> Iterator[torch.Tensor]:
+    streamer = SoproTTSStreamer(tts, StreamConfig(chunk_frames=chunk_frames))
+    return streamer.stream(
+        text,
+        ref_audio_path=ref_audio_path,
+        ref_tokens_tq=ref_tokens_tq,
+        chunk_frames=chunk_frames,
+        **kwargs,
+    )