sopro 1.0.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

sopro/model.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 import os
+from dataclasses import dataclass
 from typing import Dict, Iterator, List, Optional, Tuple
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 from sopro.config import SoproTTSConfig
 from sopro.hub import (
@@ -18,152 +18,36 @@ from .audio import save_audio
 from .codec.mimi import MimiCodec
 from .constants import TARGET_SR
 from .nn import (
+    ARRVQ1Generator,
     CodebookEmbedding,
-    RefXAttn,
+    RefXAttnStack,
     RMSNorm,
     SinusoidalPositionalEmbedding,
     SpeakerFiLM,
     SSMLiteBlock,
-    TextEmbedding,
-    TextXAttnBlock,
+    TextEncoder,
     Token2SV,
 )
-from .sampling import center_crop_tokens, repeated_tail
+from .nn.nar import NARSinglePass
+from .sampling import repeated_tail
 from .sampling import rf_ar as rf_ar_fn
 from .sampling import rf_nar as rf_nar_fn
 from .sampling import sample_token
 from .tokenizer import TextTokenizer
 
 
-class TextEncoder(nn.Module):
-    def __init__(
-        self,
-        cfg: SoproTTSConfig,
-        d_model: int,
-        n_layers: int,
-        tokenizer: TextTokenizer,
-    ):
-        super().__init__()
-        self.tok = tokenizer
-        self.embed = TextEmbedding(self.tok.vocab_size, d_model)
-        self.layers = nn.ModuleList(
-            [SSMLiteBlock(d_model, cfg.dropout, causal=False) for _ in range(n_layers)]
-        )
-        self.pos = SinusoidalPositionalEmbedding(d_model, max_len=cfg.max_text_len + 8)
-        self.norm = RMSNorm(d_model)
-
-    def forward(
-        self, text_ids: torch.Tensor, mask: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        x = self.embed(text_ids)
-        L = x.size(1)
-        pos = self.pos(torch.arange(L, device=x.device))
-        x = x + pos.unsqueeze(0)
-
-        x = x * mask.unsqueeze(-1).float()
-        for layer in self.layers:
-            x = layer(x)
-        x = self.norm(x)
-
-        mask_f = mask.float().unsqueeze(-1)
-        pooled = (x * mask_f).sum(dim=1) / (mask_f.sum(dim=1) + 1e-6)
-        return x, pooled
-
-
-class ARRVQ1Generator(nn.Module):
-    def __init__(self, cfg: SoproTTSConfig, d_model: int, vocab: int):
-        super().__init__()
-        ks = cfg.ar_kernel
-        dils: List[int] = []
-        while len(dils) < cfg.n_layers_ar:
-            dils.extend(list(cfg.ar_dilation_cycle))
-        dils = dils[: cfg.n_layers_ar]
-
-        self.dils = tuple(int(d) for d in dils)
-        self.blocks = nn.ModuleList(
-            [
-                SSMLiteBlock(
-                    d_model, cfg.dropout, causal=True, kernel_size=ks, dilation=d
-                )
-                for d in self.dils
-            ]
-        )
-
-        self.attn_freq = int(cfg.ar_text_attn_freq)
-        self.x_attns = nn.ModuleList()
-        for i in range(len(self.blocks)):
-            if (i + 1) % self.attn_freq == 0:
-                self.x_attns.append(
-                    TextXAttnBlock(d_model, heads=4, dropout=cfg.dropout)
-                )
-            else:
-                self.x_attns.append(nn.Identity())
-
-        self.norm = RMSNorm(d_model)
-        self.head = nn.Linear(d_model, vocab)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        text_emb: Optional[torch.Tensor] = None,
-        text_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        key_padding_mask = ~text_mask if text_mask is not None else None
-
-        if key_padding_mask is not None:
-            bad_rows = key_padding_mask.all(dim=1)
-            if bad_rows.any():
-                key_padding_mask = key_padding_mask.clone()
-                idx = torch.nonzero(bad_rows, as_tuple=False).squeeze(1)
-                key_padding_mask[idx, 0] = False
-                if text_emb is not None:
-                    text_emb = text_emb.clone()
-                    text_emb[idx, 0, :] = 0
-
-        h = x
-        for i, lyr in enumerate(self.blocks):
-            h = lyr(h)
-            if not isinstance(self.x_attns[i], nn.Identity) and text_emb is not None:
-                h = self.x_attns[i](h, text_emb, key_padding_mask=key_padding_mask)
-
-        h = self.norm(h)
-        return self.head(h)
-
-
-class StageRefiner(nn.Module):
-    def __init__(self, cfg: SoproTTSConfig, D: int, num_heads: int, codebook_size: int):
-        super().__init__()
-        self.blocks = nn.ModuleList(
-            [SSMLiteBlock(D, cfg.dropout, causal=True) for _ in range(cfg.n_layers_nar)]
-        )
-        self.norm = RMSNorm(D)
-        self.pre = nn.Linear(D, cfg.nar_head_dim)
-        self.heads = nn.ModuleList(
-            [nn.Linear(cfg.nar_head_dim, codebook_size) for _ in range(num_heads)]
-        )
-        self.mix = nn.Parameter(torch.ones(2, dtype=torch.float32))
-
-    def forward_hidden(
-        self, cond_bt_d: torch.Tensor, prev_bt_d: torch.Tensor
-    ) -> torch.Tensor:
-        w = torch.softmax(self.mix, dim=0)
-        x = w[0] * cond_bt_d + w[1] * prev_bt_d
-        for b in self.blocks:
-            x = b(x)
-        return self.norm(x)
-
-    def forward_heads(self, h: torch.Tensor) -> List[torch.Tensor]:
-        z = self.pre(h)
-        return [head(z) for head in self.heads]
-
+def _stage_range_to_indices(stage_rng: Tuple[int, int], Q: int) -> List[int]:
+    lo, hi = int(stage_rng[0]), int(stage_rng[1])
+    idxs = list(range(lo - 1, hi))
+    return [i for i in idxs if 1 <= i < Q]
 
-class StopHead(nn.Module):
-    def __init__(self, D: int):
-        super().__init__()
-        self.proj = nn.Linear(D, 1)
 
-    def forward(self, h: torch.Tensor) -> torch.Tensor:
-        return self.proj(h).squeeze(-1)
+@dataclass
+class PreparedReference:
+    ref_tokens_btq: torch.Tensor
+    sv_ref: torch.Tensor
+    ref_seq: torch.Tensor
+    ref_kv_caches: List[Dict[str, torch.Tensor]]
 
 
 class SoproTTSModel(nn.Module):
@@ -172,327 +56,165 @@ class SoproTTSModel(nn.Module):
         self.cfg = cfg
         D = int(cfg.d_model)
 
-        self.text_enc = TextEncoder(cfg, D, cfg.n_layers_text, tokenizer)
-        self.frame_pos = SinusoidalPositionalEmbedding(D, max_len=cfg.pos_emb_max + 2)
+        self.eos_id = int(cfg.codebook_size)
+
+        self.text_enc = TextEncoder(cfg, D, int(cfg.n_layers_text), tokenizer)
+        self.frame_pos = SinusoidalPositionalEmbedding(
+            D, max_len=int(cfg.pos_emb_max) + 8
+        )
 
         self.cb_embed = CodebookEmbedding(
             cfg.num_codebooks, cfg.codebook_size, D, use_bos=True
         )
-        self.rvq1_bos_id = self.cb_embed.bos_id
+
+        self.nar_prev_cb_weights = nn.Parameter(
+            torch.zeros(cfg.num_codebooks, dtype=torch.float32)
+        )
 
         self.token2sv = Token2SV(
             cfg.num_codebooks,
             cfg.codebook_size,
             d=192,
-            out_dim=cfg.sv_student_dim,
+            out_dim=int(cfg.sv_student_dim),
             dropout=cfg.dropout,
         )
-        self.spk_film = SpeakerFiLM(D, sv_dim=cfg.sv_student_dim)
-        self.cond_norm = RMSNorm(D)
+        self.spk_film = SpeakerFiLM(D, sv_dim=int(cfg.sv_student_dim))
 
-        self.ar = ARRVQ1Generator(cfg, D, cfg.codebook_size)
-        if cfg.ar_lookback > 0:
-            self.ar_hist_w = nn.Parameter(torch.zeros(cfg.ar_lookback))
-
-        def idxs(rng: Tuple[int, int]) -> List[int]:
-            lo, hi = rng
-            return list(range(lo - 1, hi))
+        self.ar = ARRVQ1Generator(cfg, D, int(cfg.codebook_size) + 1)
 
+        Q = int(cfg.num_codebooks)
         self.stage_indices: Dict[str, List[int]] = {
-            "B": idxs(cfg.stage_B),
-            "C": idxs(cfg.stage_C),
-            "D": idxs(cfg.stage_D),
-            "E": idxs(cfg.stage_E),
+            "B": _stage_range_to_indices(cfg.stage_B, Q),
+            "C": _stage_range_to_indices(cfg.stage_C, Q),
+            "D": _stage_range_to_indices(cfg.stage_D, Q),
+            "E": _stage_range_to_indices(cfg.stage_E, Q),
         }
-        self.stages = nn.ModuleDict(
-            {
-                s: StageRefiner(cfg, D, len(self.stage_indices[s]), cfg.codebook_size)
-                for s in ["B", "C", "D", "E"]
-            }
-        )
+        self.stage_order = [
+            s for s in ["B", "C", "D", "E"] if len(self.stage_indices[s]) > 0
+        ]
+
+        self.nar = NARSinglePass(cfg, D, stage_specs=self.stage_indices)
 
-        self.stop_head = StopHead(D) if cfg.use_stop_head else None
+        self.cond_norm = RMSNorm(D)
 
+        ref_enc_layers = int(getattr(cfg, "ref_enc_layers", 2))
         self.ref_enc_blocks = nn.ModuleList(
-            [SSMLiteBlock(D, cfg.dropout, causal=False) for _ in range(2)]
+            [SSMLiteBlock(D, cfg.dropout, causal=False) for _ in range(ref_enc_layers)]
         )
         self.ref_enc_norm = RMSNorm(D)
-        self.ref_xattn_stack = RefXAttn(
-            D, heads=cfg.ref_attn_heads, layers=3, dropout=cfg.dropout
+
+        self.ref_xattn = RefXAttnStack(
+            D,
+            heads=cfg.ref_xattn_heads,
+            layers=cfg.ref_xattn_layers,
+            gmax=cfg.ref_xattn_gmax,
+        )
+
+        self.register_buffer(
+            "ref_cb_weights",
+            torch.linspace(1.0, 0.1, int(cfg.num_codebooks)),
+            persistent=True,
         )
 
     def rf_ar(self) -> int:
-        return rf_ar_fn(self.cfg.ar_kernel, self.ar.dils)
+        return rf_ar_fn(
+            int(self.cfg.ar_kernel),
+            getattr(self.ar, "dils", tuple(int(x) for x in self.cfg.ar_dilation_cycle)),
+        )
 
     def rf_nar(self) -> int:
-        return rf_nar_fn(self.cfg.n_layers_nar, kernel_size=7, dilation=1)
-
-    def _pool_time(self, x: torch.Tensor, factor: int) -> torch.Tensor:
-        if factor <= 1 or x.size(1) < 2 * factor:
-            return x
-        return F.avg_pool1d(
-            x.transpose(1, 2), kernel_size=factor, stride=factor
-        ).transpose(1, 2)
-
-    def _normalize_ref_mask(
-        self, ref_mask: Optional[torch.Tensor], device: torch.device
-    ) -> Optional[torch.Tensor]:
-        if ref_mask is None:
-            return None
-        mk = ref_mask.to(device).bool()
-        if mk.ndim == 1:
-            mk = mk.unsqueeze(0)
-        return mk
-
-    def _encode_reference_seq(self, ref_tokens: torch.Tensor) -> torch.Tensor:
-        B, Tr, Q = ref_tokens.shape
-        emb_sum = 0.0
-        for q in range(Q):
-            emb_sum = emb_sum + self.cb_embed.embed_tokens(
-                ref_tokens[:, :, q], cb_index=q
-            )
-        x = emb_sum / float(Q)
-        for b in self.ref_enc_blocks:
-            x = b(x)
-        return self.ref_enc_norm(x)
+        cycle = tuple(int(x) for x in self.cfg.nar_dilation_cycle) or (1,)
+        dils: List[int] = []
+        while len(dils) < int(self.cfg.n_layers_nar):
+            dils.extend(list(cycle))
+        dils = dils[: int(self.cfg.n_layers_nar)]
+        return rf_nar_fn(int(self.cfg.nar_kernel_size), tuple(dils))
 
-    def _single_pass_ref_xattn(
-        self,
-        cond_btd: torch.Tensor,
-        ref_seq: torch.Tensor,
-        ref_mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        ref_seq_p = self._pool_time(ref_seq, 1)
-
-        key_padding_mask = None
-        if ref_mask is not None:
-            mk_bool = ref_mask.bool()
-            B, Tr = mk_bool.shape
-            pooled_len = ref_seq_p.size(1)
-            if pooled_len == Tr:
-                key_padding_mask = ~mk_bool
-            else:
-                cut = pooled_len * 2
-                mk2 = mk_bool[:, :cut].reshape(B, pooled_len, 2).any(dim=2)
-                key_padding_mask = ~mk2
+    @torch.no_grad()
+    def _encode_reference_seq(self, ref_tokens_btq: torch.Tensor) -> torch.Tensor:
+        B, Tr, Q = ref_tokens_btq.shape
 
-        return self.ref_xattn_stack(
-            cond_btd, ref_seq_p, key_padding_mask=key_padding_mask
+        w = torch.softmax(self.ref_cb_weights.float(), dim=0).to(
+            device=ref_tokens_btq.device
         )
 
-    def _base_cond_at(
-        self, t: int, txt_pool: torch.Tensor, device: torch.device
-    ) -> torch.Tensor:
-        pos = self.frame_pos(torch.tensor([t], device=device)).unsqueeze(0)
-        return txt_pool[:, None, :] + pos
-
-    def _ar_prev_from_seq(self, seq_1xT: torch.Tensor) -> torch.Tensor:
-        K = int(self.cfg.ar_lookback)
-        if K <= 0 or getattr(self, "ar_hist_w", None) is None:
-            return self.cb_embed.embed_shift_by_k(seq_1xT, cb_index=0, k=1)
-
-        ws = torch.softmax(self.ar_hist_w, dim=0)
-        acc = 0.0
-        k_max = min(K, int(seq_1xT.size(1)))
-        for k in range(1, k_max + 1):
-            acc = acc + ws[k - 1] * self.cb_embed.embed_shift_by_k(
-                seq_1xT, cb_index=0, k=k
-            )
-        return acc
+        x = 0.0
+        for q in range(Q):
+            e = self.cb_embed.embed_tokens(ref_tokens_btq[:, :, q], cb_index=q)
+            x = x + w[q].to(e.dtype) * e
 
-    @torch.no_grad()
-    def prepare_conditioning(
-        self,
-        text_ids_1d: torch.Tensor,
-        ref_tokens_tq: torch.Tensor,
-        *,
-        max_frames: int,
-        device: torch.device,
-        style_strength: float = 1.0,
-        ref_mask: Optional[torch.Tensor] = None,
-        chunk_size: Optional[int] = None,
-    ) -> Dict[str, torch.Tensor]:
-        self.eval()
+        for b in self.ref_enc_blocks:
+            x = b(x)
 
-        text_ids = text_ids_1d.to(device)
-        text_mask = torch.ones_like(text_ids, dtype=torch.bool).unsqueeze(0)
-        txt_seq, txt_pool = self.text_enc(text_ids.unsqueeze(0), text_mask)
+        return self.ref_enc_norm(x)
 
-        ref_btq = ref_tokens_tq.unsqueeze(0).to(device)
+    @torch.no_grad()
+    def prepare_reference(
+        self, ref_tokens_tq: torch.Tensor, *, device: torch.device
+    ) -> PreparedReference:
+        ref_tokens_btq = ref_tokens_tq.unsqueeze(0).to(device=device, dtype=torch.long)
+        Tr = int(ref_tokens_btq.size(1))
 
-        sv_ref = self.token2sv(ref_btq, lengths=None)
-        ref_seq = self._encode_reference_seq(ref_btq)
-        ref_mask_btr = self._normalize_ref_mask(ref_mask, device)
+        lengths = torch.tensor([Tr], device=device, dtype=torch.long)
+        sv_ref = self.token2sv(ref_tokens_btq, lengths=lengths)
 
-        T = int(max_frames)
-        if T <= 0:
-            cond_all = torch.zeros(
-                (1, 0, txt_pool.size(-1)), device=device, dtype=txt_pool.dtype
-            )
-        else:
-            pos = self.frame_pos(torch.arange(T, device=device)).unsqueeze(0)
-            base_all = txt_pool[:, None, :] + pos
-            base_all = self.spk_film(base_all, sv_ref, strength=float(style_strength))
-
-            if chunk_size is None or int(chunk_size) >= T:
-                out = self._single_pass_ref_xattn(
-                    base_all, ref_seq, ref_mask=ref_mask_btr
-                )
-                cond_all = self.cond_norm(out)
-            else:
-                cs = int(chunk_size)
-                chunks: List[torch.Tensor] = []
-                for s in range(0, T, cs):
-                    e = min(T, s + cs)
-                    q = base_all[:, s:e, :]
-                    out = self._single_pass_ref_xattn(q, ref_seq, ref_mask=ref_mask_btr)
-                    chunks.append(self.cond_norm(out))
-                cond_all = torch.cat(chunks, dim=1)
+        ref_seq = self._encode_reference_seq(ref_tokens_btq)
 
-        return {
-            "txt_seq": txt_seq,
-            "text_mask": text_mask,
-            "cond_all": cond_all,
-            "ref_btq": ref_btq,
-            "txt_pool": txt_pool,
-            "sv_ref": sv_ref,
-            "ref_seq": ref_seq,
-            "ref_mask": (
-                ref_mask_btr
-                if ref_mask_btr is not None
-                else torch.empty(0, device=device, dtype=torch.bool)
-            ),
-            "style_strength": torch.tensor(float(style_strength), device=device),
-        }
+        ref_kv_caches = self.ref_xattn.build_kv_caches(ref_seq, key_padding_mask=None)
+
+        return PreparedReference(
+            ref_tokens_btq=ref_tokens_btq,
+            sv_ref=sv_ref,
+            ref_seq=ref_seq,
+            ref_kv_caches=ref_kv_caches,
+        )
 
     @torch.no_grad()
-    def prepare_conditioning_lazy(
+    def prepare_conditioning(
         self,
         text_ids_1d: torch.Tensor,
-        ref_tokens_tq: torch.Tensor,
+        ref: PreparedReference,
         *,
         max_frames: int,
         device: torch.device,
-        style_strength: float = 1.0,
-        ref_mask: Optional[torch.Tensor] = None,
+        style_strength: float = 1.2,
     ) -> Dict[str, torch.Tensor]:
         self.eval()
+        sv_ref = ref.sv_ref.to(device)
 
         text_ids = text_ids_1d.to(device)
         text_mask = torch.ones_like(text_ids, dtype=torch.bool).unsqueeze(0)
         txt_seq, txt_pool = self.text_enc(text_ids.unsqueeze(0), text_mask)
 
-        ref_btq = ref_tokens_tq.unsqueeze(0).to(device)
-        sv_ref = self.token2sv(ref_btq, lengths=None)
-        ref_seq = self._encode_reference_seq(ref_btq)
-        ref_mask_btr = self._normalize_ref_mask(ref_mask, device)
+        if sv_ref is not None:
+            if sv_ref.dim() == 1:
+                sv_ref = sv_ref.unsqueeze(0)
+            sv_ref = sv_ref.to(device)
+        else:
+            ref_btq = ref_tokens_tq.unsqueeze(0).to(device)
+            ref_len = torch.tensor(
+                [int(ref_btq.size(1))], device=device, dtype=torch.long
+            )
+            sv_ref = self.token2sv(ref_btq, lengths=ref_len)
+
+        Tar = int(max_frames) + 1
+        pos = self.frame_pos(torch.arange(Tar, device=device)).unsqueeze(0)
+        base_ar = txt_pool[:, None, :] + pos
+        cond_ar = self.spk_film(base_ar, sv_ref, strength=float(style_strength))
 
-        D = int(txt_pool.size(-1))
-        cond_all = torch.zeros((1, 0, D), device=device, dtype=txt_pool.dtype)
+        cond_ar, _ = self.ref_xattn(
+            cond_ar, kv_caches=ref.ref_kv_caches, use_cache=True
+        )
+        cond_ar = self.cond_norm(cond_ar)
 
         return {
             "txt_seq": txt_seq,
             "text_mask": text_mask,
-            "cond_all": cond_all,
-            "ref_btq": ref_btq,
             "txt_pool": txt_pool,
             "sv_ref": sv_ref,
-            "ref_seq": ref_seq,
-            "ref_mask": (
-                ref_mask_btr
-                if ref_mask_btr is not None
-                else torch.empty(0, device=device, dtype=torch.bool)
-            ),
-            "style_strength": torch.tensor(float(style_strength), device=device),
-            "max_frames": torch.tensor(int(max_frames), device=device),
+            "cond_ar": cond_ar,
         }
 
-    @torch.no_grad()
-    def ensure_cond_upto(
-        self,
-        prep: Dict[str, torch.Tensor],
-        t_inclusive: int,
-        *,
-        chunk_size: int = 64,
-    ) -> None:
-        if t_inclusive < 0:
-            return
-
-        cond_all = prep.get("cond_all", None)
-        if cond_all is None:
-            raise KeyError("prep dict missing 'cond_all'.")
-
-        have = int(cond_all.size(1))
-        need_min = int(t_inclusive) + 1
-        if have >= need_min:
-            return
-
-        if "txt_pool" not in prep or "sv_ref" not in prep or "ref_seq" not in prep:
-            raise RuntimeError(
-                "Lazy conditioning requested but prep lacks txt_pool/sv_ref/ref_seq. "
-                "Use prepare_conditioning_lazy() or prepare_conditioning(..., chunk_size=...)."
-            )
-
-        device = cond_all.device
-        txt_pool = prep["txt_pool"]
-        sv_ref = prep["sv_ref"]
-        ref_seq = prep["ref_seq"]
-
-        style_strength = float(
-            prep.get(
-                "style_strength", torch.tensor(self.cfg.style_strength, device=device)
-            ).item()
-        )
-
-        ref_mask = prep.get("ref_mask", None)
-        if ref_mask is not None and ref_mask.numel() == 0:
-            ref_mask = None
-
-        max_frames = prep.get("max_frames", None)
-        maxT = int(max_frames.item()) if max_frames is not None else None
-
-        cs = max(1, int(chunk_size))
-
-        need = ((need_min + cs - 1) // cs) * cs
-        if maxT is not None:
-            need = min(need, maxT)
-
-        if have >= need:
-            return
-
-        new_chunks: List[torch.Tensor] = []
-        for s in range(have, need, cs):
-            e = min(need, s + cs)
-            pos = self.frame_pos(torch.arange(s, e, device=device)).unsqueeze(0)
-            base = txt_pool[:, None, :] + pos
-            base = self.spk_film(base, sv_ref, strength=style_strength)
-            out = self._single_pass_ref_xattn(base, ref_seq, ref_mask=ref_mask)
-            new_chunks.append(self.cond_norm(out))
-
-        prep["cond_all"] = torch.cat([cond_all] + new_chunks, dim=1)
-
-    def build_ar_prefix(
-        self,
-        ref_btq: torch.Tensor,
-        device: torch.device,
-        prefix_sec_fixed: Optional[float],
-        use_prefix: bool,
-    ) -> torch.Tensor:
-        if not use_prefix or ref_btq.size(1) == 0:
-            return torch.zeros(1, 0, dtype=torch.long, device=device)
-
-        avail = int(ref_btq.size(1))
-        fps = float(self.cfg.mimi_fps)
-
-        if prefix_sec_fixed is not None and prefix_sec_fixed > 0:
-            P = min(avail, int(round(prefix_sec_fixed * fps)))
-        else:
-            P = min(avail, max(1, int(round(self.cfg.preprompt_sec_max * fps))))
-
-        if P <= 0:
-            return torch.zeros(1, 0, dtype=torch.long, device=device)
-        return ref_btq[:, :P, 0].contiguous()
-
     @torch.no_grad()
     def ar_stream(
         self,
@@ -505,200 +227,178 @@ class SoproTTSModel(nn.Module):
         loop_streak: int = 8,
         recovery_top_p: float = 0.85,
         recovery_temp: float = 1.2,
-        use_prefix: bool = True,
-        prefix_sec_fixed: Optional[float] = 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[Tuple[int, int, Optional[float]]]:
-        device = prep["cond_all"].device
-        cond_all = prep["cond_all"]
+    ) -> Iterator[Tuple[int, int, bool]]:
+        device = prep["cond_ar"].device
+        cond_ar = prep["cond_ar"]
         txt_seq = prep["txt_seq"]
         text_mask = prep["text_mask"]
-        ref_btq = prep["ref_btq"]
-
-        R_AR = self.rf_ar()
-
-        stop_head = self.stop_head
-        if use_stop_head is not None:
-            if not bool(use_stop_head):
-                stop_head = None
 
-        eff_stop_patience = int(
-            stop_patience if stop_patience is not None else self.cfg.stop_patience
-        )
-        eff_stop_threshold = float(
-            stop_threshold if stop_threshold is not None else self.cfg.stop_threshold
-        )
-        eff_min_gen_frames = int(
+        eos_id = int(self.eos_id)
+        eff_min_gen = int(
             min_gen_frames if min_gen_frames is not None else self.cfg.min_gen_frames
         )
 
-        A_prefix = self.build_ar_prefix(
-            ref_btq, device, prefix_sec_fixed, use_prefix=use_prefix
-        )
-        P = int(A_prefix.size(1))
+        max_steps = int(max_frames) + 1
+        ctx_ids = torch.zeros((1, max_steps), dtype=torch.long, device=device)
 
-        ctx_ids = torch.zeros(
-            (1, P + int(max_frames) + 1), dtype=torch.long, device=device
+        ar_state = self.ar.init_stream_state(
+            batch_size=1,
+            device=device,
+            dtype=cond_ar.dtype,
+            text_emb=txt_seq,
+            text_mask=text_mask,
         )
-        if P > 0:
-            ctx_ids[:, :P] = A_prefix
 
-        hist_A: List[int] = []
+        hist: List[int] = []
         loop_streak_count = 0
-        stop_streak_count = 0
         last_a: Optional[int] = None
 
-        gen_len = 0
-
-        for t in range(int(max_frames)):
-            if prep["cond_all"].size(1) < (t + 1):
-                self.ensure_cond_upto(prep, t, chunk_size=int(cond_chunk_size or 64))
-            cond_all = prep["cond_all"]
-
-            L_ar = min(t + 1, R_AR)
-            s_ar = t + 1 - L_ar
-            cond_win_ar = cond_all[:, s_ar : t + 1, :]
+        if self.cb_embed.bos_id is None:
+            raise RuntimeError(
+                "CodebookEmbedding.use_bos must be True for streaming AR cache"
+            )
+        bos_idx = torch.full(
+            (1, 1), int(self.cb_embed.bos_id), device=device, dtype=torch.long
+        )
 
-            total_len = P + gen_len + 1
-            A_ctx_full = ctx_ids[:, :total_len]
+        for t in range(max_steps):
+            if t == 0:
+                prev_emb = self.cb_embed.emb(bos_idx)
+            else:
+                prev_tok = ctx_ids[:, t - 1 : t]
+                prev_emb = self.cb_embed.embed_tokens(prev_tok, cb_index=0)
 
-            prev_ctx_full = self._ar_prev_from_seq(A_ctx_full)
-            prev_ctx_win = prev_ctx_full[:, -L_ar:, :]
+            x_t = cond_ar[:, t : t + 1, :] + prev_emb
 
             cur_top_p, cur_temp = top_p, temperature
             if anti_loop:
-                if repeated_tail(hist_A, max_n=16):
+                if repeated_tail(hist, max_n=16):
                     cur_top_p, cur_temp = recovery_top_p, recovery_temp
                 elif last_a is not None and loop_streak_count >= loop_streak:
                     cur_top_p, cur_temp = recovery_top_p, recovery_temp
 
-            ar_logits_win = self.ar(
-                cond_win_ar + prev_ctx_win, text_emb=txt_seq, text_mask=text_mask
+            logits_t, ar_state = self.ar.step(
+                x_t, ar_state, text_emb=txt_seq, text_mask=text_mask
             )
-            ar_logits_t = ar_logits_win[:, -1:, :]
-
-            rvq1_id = sample_token(
-                ar_logits_t,
-                history=hist_A,
+            tok = sample_token(
+                logits_t,
+                history=hist,
                 top_p=cur_top_p,
                 temperature=cur_temp,
                 top_k=50,
                 repetition_penalty=1.1,
             )
 
-            ctx_ids[0, P + gen_len] = int(rvq1_id)
-            gen_len += 1
+            ctx_ids[0, t] = int(tok)
+            hist.append(int(tok))
 
-            hist_A.append(int(rvq1_id))
             loop_streak_count = (
-                (loop_streak_count + 1)
-                if (last_a is not None and rvq1_id == last_a)
-                else 0
+                (loop_streak_count + 1) if (last_a is not None and tok == last_a) else 0
             )
-            last_a = int(rvq1_id)
-
-            p_stop: Optional[float] = None
-            if stop_head is not None:
-                A_now = torch.tensor([[rvq1_id]], device=device, dtype=torch.long)
-                stop_inp = (
-                    cond_all[:, t : t + 1, :]
-                    + self.cb_embed.embed_tokens(A_now, cb_index=0).detach()
-                )
-                stop_logits = stop_head(stop_inp)
-                p_stop = float(torch.sigmoid(stop_logits).item())
-
-                if t + 1 >= eff_min_gen_frames and p_stop > eff_stop_threshold:
-                    stop_streak_count += 1
-                else:
-                    stop_streak_count = 0
-
-            yield t, int(rvq1_id), p_stop
-
-            if stop_head is not None and stop_streak_count >= eff_stop_patience:
+            last_a = int(tok)
+
+            is_eos = int(tok) == eos_id
+            yield t, int(tok), bool(is_eos)
+
+            if is_eos and (t + 1) >= eff_min_gen:
                 break
 
     @torch.no_grad()
     def nar_refine(
-        self, cond_seq: torch.Tensor, tokens_A_1xT: torch.Tensor
+        self, cond_seq: torch.Tensor, rvq1_1xT: torch.Tensor
     ) -> torch.Tensor:
-        preds_all: List[torch.Tensor] = [tokens_A_1xT.unsqueeze(-1)]
-        prev_tokens_list: List[torch.Tensor] = [tokens_A_1xT.unsqueeze(-1)]
+        B, T, D = cond_seq.shape
+        Q = int(self.cfg.num_codebooks)
+
+        out_btq = torch.zeros((B, T, Q), device=cond_seq.device, dtype=torch.long)
+        out_btq[:, :, 0] = rvq1_1xT
+
+        prev_tokens_list: List[torch.Tensor] = [rvq1_1xT.unsqueeze(-1)]
         prev_cb_list: List[List[int]] = [[0]]
 
-        for stage_name in ["B", "C", "D", "E"]:
-            idxs = self.stage_indices[stage_name]
+        for stage in self.stage_order:
+            idxs = self.stage_indices[stage]
+            if len(idxs) == 0:
+                continue
+
             prev_tokens_cat = torch.cat(prev_tokens_list, dim=-1)
             prev_cbs_cat = sum(prev_cb_list, [])
-            prev_emb_sum = self.cb_embed.sum_embed_subset(prev_tokens_cat, prev_cbs_cat)
 
-            h = self.stages[stage_name].forward_hidden(cond_seq, prev_emb_sum)
-            logits_list = self.stages[stage_name].forward_heads(h)
-            preds = torch.stack([x.argmax(dim=-1) for x in logits_list], dim=-1)
+            prev_emb_sum = self.cb_embed.sum_embed_subset(
+                prev_tokens_cat,
+                prev_cbs_cat,
+                keep_mask=None,
+                cb_weights=self.nar_prev_cb_weights,
+            )
+
+            logits_list = self.nar.forward_stage(stage, cond_seq, prev_emb_sum)
+            if len(logits_list) == 0:
+                continue
+
+            preds = torch.stack([lg.argmax(dim=-1) for lg in logits_list], dim=-1)
+
+            for k, cb in enumerate(idxs):
+                out_btq[:, :, cb] = preds[:, :, k]
 
-            preds_all.append(preds)
-            prev_tokens_list.append(preds)
+            prev_tokens_list.append(preds.detach())
             prev_cb_list.append(idxs)
 
-        tokens_btq = torch.cat(preds_all, dim=-1)
-        return tokens_btq
+        return out_btq
 
     @torch.no_grad()
     def generate_tokens(
         self,
         text_ids_1d: torch.Tensor,
-        ref_tokens_tq: torch.Tensor,
+        ref: PreparedReference,
         *,
         max_frames: int,
         device: torch.device,
         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: float = 1.0,
-        use_stop_head: Optional[bool] = None,
-        stop_patience: Optional[int] = None,
-        stop_threshold: Optional[float] = None,
+        style_strength: float = 1.2,
         min_gen_frames: Optional[int] = None,
     ) -> torch.Tensor:
         prep = self.prepare_conditioning(
             text_ids_1d,
-            ref_tokens_tq,
+            ref,
             max_frames=max_frames,
             device=device,
             style_strength=style_strength,
         )
 
-        hist_A: List[int] = []
-        for _t, rvq1, _p_stop in self.ar_stream(
+        eos_id = int(self.eos_id)
+        hist: List[int] = []
+        for _t, tok, is_eos in self.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,
-            use_stop_head=use_stop_head,
-            stop_patience=stop_patience,
-            stop_threshold=stop_threshold,
             min_gen_frames=min_gen_frames,
         ):
-            hist_A.append(rvq1)
+            hist.append(tok)
+            if is_eos:
+                break
+
+        Tfull = len(hist)
+        cut = Tfull
+        for i, v in enumerate(hist):
+            if int(v) == eos_id:
+                cut = i
+                break
 
-        T = len(hist_A)
-        if T == 0:
+        T = int(cut)
+        if T <= 0:
             return torch.zeros(
-                0, self.cfg.num_codebooks, dtype=torch.long, device=device
+                (0, int(self.cfg.num_codebooks)), dtype=torch.long, device=device
             )
 
-        tokens_A = torch.tensor(hist_A, device=device, dtype=torch.long).unsqueeze(0)
-        cond_seq = prep["cond_all"][:, :T, :]
-        tokens_btq_1xTQ = self.nar_refine(cond_seq, tokens_A)
-        return tokens_btq_1xTQ.squeeze(0)
+        rvq1 = torch.tensor(hist[:T], device=device, dtype=torch.long).unsqueeze(0)
+        cond_seq = prep["cond_ar"][:, :T, :]
+        tokens_1xTQ = self.nar_refine(cond_seq, rvq1)
+        return tokens_1xTQ.squeeze(0)
 
 
 class SoproTTS:
@@ -738,16 +438,14 @@ class SoproTTS:
             raise FileNotFoundError(f"Expected {model_path} in repo snapshot.")
 
         cfg = load_cfg_from_safetensors(model_path)
-
         tokenizer = TextTokenizer(model_name=local_dir)
 
         model = SoproTTSModel(cfg, tokenizer).to(dev).eval()
         state = load_state_dict_from_safetensors(model_path)
 
-        model.load_state_dict(state)
+        model.load_state_dict(state, strict=False)
 
         codec = MimiCodec(num_quantizers=cfg.num_codebooks, device=device)
-
         return cls(
             model=model, cfg=cfg, tokenizer=tokenizer, codec=codec, device=device
         )
@@ -756,6 +454,26 @@ class SoproTTS:
         ids = self.tokenizer.encode(text)
         return torch.tensor(ids, dtype=torch.long, device=self.device)
 
+    @torch.inference_mode()
+    def encode_speaker(
+        self,
+        *,
+        ref_audio_path: Optional[str] = None,
+        ref_tokens_tq: Optional[torch.Tensor] = None,
+        ref_seconds: Optional[float] = None,
+    ) -> torch.Tensor:
+        ref = self.encode_reference(
+            ref_audio_path=ref_audio_path,
+            ref_tokens_tq=ref_tokens_tq,
+            ref_seconds=ref_seconds,
+        )
+        ref_btq = ref.unsqueeze(0)
+        lengths = torch.tensor(
+            [int(ref_btq.size(1))], device=self.device, dtype=torch.long
+        )
+        sv = self.model.token2sv(ref_btq, lengths=lengths)
+        return sv.squeeze(0).detach()
+
     def encode_reference(
         self,
         *,
@@ -763,6 +481,8 @@ class SoproTTS:
         ref_tokens_tq: Optional[torch.Tensor] = None,
         ref_seconds: Optional[float] = None,
     ) -> torch.Tensor:
+        from .sampling import center_crop_tokens
+
         if (ref_tokens_tq is None) and (ref_audio_path is None):
             raise RuntimeError(
                 "SoproTTS requires a reference. Provide ref_audio_path=... or ref_tokens_tq=..."
@@ -773,11 +493,11 @@ class SoproTTS:
             )
 
         if ref_seconds is None:
-            ref_seconds = float(self.cfg.ref_seconds_max)
+            ref_seconds = 12.0
 
         if ref_tokens_tq is not None:
             ref = ref_tokens_tq.to(self.device).long()
-            if ref_seconds > 0:
+            if ref_seconds and ref_seconds > 0:
                 fps = float(self.cfg.mimi_fps)
                 win = max(1, int(round(ref_seconds * fps)))
                 ref = center_crop_tokens(ref, win)
@@ -793,51 +513,61 @@ class SoproTTS:
         )
         return ref
 
+    @torch.inference_mode()
+    def prepare_reference(
+        self,
+        *,
+        ref_audio_path: Optional[str] = None,
+        ref_tokens_tq: Optional[torch.Tensor] = None,
+        ref_seconds: Optional[float] = None,
+    ) -> PreparedReference:
+        tokens_tq = self.encode_reference(
+            ref_audio_path=ref_audio_path,
+            ref_tokens_tq=ref_tokens_tq,
+            ref_seconds=ref_seconds,
+        )
+        return self.model.prepare_reference(tokens_tq, device=self.device)
+
     @torch.inference_mode()
     def synthesize(
         self,
         text: str,
         *,
+        ref: Optional[PreparedReference] = None,
         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,
-        use_stop_head: Optional[bool] = None,
-        stop_patience: Optional[int] = None,
-        stop_threshold: Optional[float] = None,
         min_gen_frames: Optional[int] = None,
     ) -> torch.Tensor:
         text_ids = self.encode_text(text)
-        ref = self.encode_reference(
-            ref_audio_path=ref_audio_path,
-            ref_tokens_tq=ref_tokens_tq,
-            ref_seconds=ref_seconds,
-        )
+
+        if ref is None:
+            ref = self.prepare_reference(
+                ref_audio_path=ref_audio_path,
+                ref_tokens_tq=ref_tokens_tq,
+                ref_seconds=ref_seconds,
+            )
+
+        text_ids = self.encode_text(text)
 
         tokens_tq = self.model.generate_tokens(
             text_ids,
-            ref,
+            ref=ref,
             max_frames=max_frames,
             device=self.device,
             top_p=top_p,
             temperature=temperature,
             anti_loop=anti_loop,
-            use_prefix=use_prefix,
-            prefix_sec_fixed=prefix_sec_fixed,
             style_strength=float(
                 style_strength
                 if style_strength is not None
                 else self.cfg.style_strength
             ),
-            use_stop_head=use_stop_head,
-            stop_patience=stop_patience,
-            stop_threshold=stop_threshold,
             min_gen_frames=min_gen_frames,
         )