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

sopro/__init__.py

@@ -3,4 +3,4 @@ from __future__ import annotations
 from .model import SoproTTS
 
 __all__ = ["SoproTTS"]
-__version__ = "1.0.2"
+__version__ = "1.5.0"

sopro/cli.py

@@ -32,8 +32,6 @@ def main() -> None:
     ap.add_argument("--temperature", type=float, default=1.05)
     ap.add_argument("--no_anti_loop", action="store_true")
 
-    ap.add_argument("--no_prefix", action="store_true")
-    ap.add_argument("--prefix_sec", type=float, default=None)
     ap.add_argument("--style_strength", type=float, default=None)
     ap.add_argument("--ref_seconds", type=float, default=None)
 
@@ -77,6 +75,7 @@ def main() -> None:
             torch.cuda.manual_seed_all(args.seed)
 
     t0 = time.perf_counter()
+
     tts = SoproTTS.from_pretrained(
         args.repo_id,
         revision=args.revision,
@@ -84,6 +83,7 @@ def main() -> None:
         token=args.hf_token,
         device=device,
     )
+
     t1 = time.perf_counter()
     if not args.quiet:
         print(f"[Load] {t1 - t0:.2f}s")
@@ -99,7 +99,7 @@ def main() -> None:
 
     with torch.inference_mode():
         text_ids = tts.encode_text(args.text)
-        ref = tts.encode_reference(
+        ref = tts.prepare_reference(
             ref_audio_path=args.ref_audio,
             ref_tokens_tq=ref_tokens_tq,
             ref_seconds=args.ref_seconds,
@@ -119,58 +119,42 @@ def main() -> None:
 
         t_start = time.perf_counter()
 
-        hist_A: list[int] = []
-        pbar = tqdm(
-            total=args.max_frames,
-            desc="AR sampling",
-            unit="frame",
-            disable=args.quiet,
-        )
+    hist_A: list[int] = []
+    pbar = tqdm(
+        total=args.max_frames + 1, desc="AR sampling", unit="step", disable=args.quiet
+    )
 
-        for _t, rvq1, p_stop in tts.model.ar_stream(
-            prep,
-            max_frames=args.max_frames,
-            top_p=args.top_p,
-            temperature=args.temperature,
-            anti_loop=(not args.no_anti_loop),
-            use_prefix=(not args.no_prefix),
-            prefix_sec_fixed=args.prefix_sec,
-            use_stop_head=(False if args.no_stop_head else None),
-            stop_patience=args.stop_patience,
-            stop_threshold=args.stop_threshold,
-        ):
-            hist_A.append(int(rvq1))
+    for _t, tok, is_eos in tts.model.ar_stream(
+        prep,
+        max_frames=args.max_frames,
+        top_p=args.top_p,
+        temperature=args.temperature,
+        anti_loop=(not args.no_anti_loop),
+    ):
+        if is_eos:
+            pbar.set_postfix(eos="yes")
             pbar.update(1)
-            if p_stop is None:
-                pbar.set_postfix(p_stop="off")
-            else:
-                pbar.set_postfix(p_stop=f"{float(p_stop):.2f}")
-
-        pbar.n = len(hist_A)
-        pbar.close()
-
-        t_after_sampling = time.perf_counter()
-
-        T = len(hist_A)
-        if T == 0:
-            save_audio(args.out, torch.zeros(1, 0), sr=TARGET_SR)
-            t_end = time.perf_counter()
-            if not args.quiet:
-                print(
-                    f"[Timing] sampling={t_after_sampling - t_start:.2f}s, "
-                    f"postproc+decode+save={t_end - t_after_sampling:.2f}s, "
-                    f"total={t_end - t_start:.2f}s"
-                )
-                print(f"[Done] Wrote {args.out}")
-            return
-
-        tokens_A = torch.tensor(hist_A, device=tts.device, dtype=torch.long).unsqueeze(0)
-        cond_seq = prep["cond_all"][:, :T, :]
-        tokens_1xTQ = tts.model.nar_refine(cond_seq, tokens_A)
-        tokens_tq = tokens_1xTQ.squeeze(0)
-
-        wav = tts.codec.decode_full(tokens_tq)
-        save_audio(args.out, wav, sr=TARGET_SR)
+            break
+        hist_A.append(int(tok))
+        pbar.update(1)
+
+    t_after_sampling = time.perf_counter()
+
+    pbar.n = len(hist_A)
+    pbar.close()
+
+    T = len(hist_A)
+    if T == 0:
+        save_audio(args.out, torch.zeros(1, 0), sr=TARGET_SR)
+        return
+
+    tokens_A = torch.tensor(hist_A, device=tts.device, dtype=torch.long).unsqueeze(0)
+    cond_seq = prep["cond_ar"][:, :T, :]
+    tokens_1xTQ = tts.model.nar_refine(cond_seq, tokens_A)
+    tokens_tq = tokens_1xTQ.squeeze(0)
+
+    wav = tts.codec.decode_full(tokens_tq)
+    save_audio(args.out, wav, sr=TARGET_SR)
 
     t_end = time.perf_counter()
     if not args.quiet:

sopro/config.py

@@ -13,36 +13,31 @@ class SoproTTSConfig:
     audio_sr: int = TARGET_SR
 
     d_model: int = 384
-    n_layers_text: int = 4
-    n_layers_ar: int = 6
-    n_layers_nar: int = 6
+    n_layers_text: int = 2
     dropout: float = 0.05
-
     pos_emb_max: int = 4096
     max_text_len: int = 2048
 
-    nar_head_dim: int = 256
-
-    use_stop_head: bool = True
-    stop_threshold: float = 0.8
-    stop_patience: int = 5
+    n_layers_ar: int = 6
+    ar_kernel: int = 13
+    ar_dilation_cycle: Tuple[int, ...] = (1, 2, 4, 1)
+    ar_text_attn_freq: int = 2
     min_gen_frames: int = 12
 
+    n_layers_nar: int = 6
+    nar_head_dim: int = 256
+    nar_kernel_size: int = 11
+    nar_dilation_cycle: Tuple[int, ...] = (1, 2, 4, 8)
+
     stage_B: Tuple[int, int] = (2, 4)
     stage_C: Tuple[int, int] = (5, 8)
     stage_D: Tuple[int, int] = (9, 16)
     stage_E: Tuple[int, int] = (17, 32)
 
-    ar_lookback: int = 4
-    ar_kernel: int = 13
-    ar_dilation_cycle: Tuple[int, ...] = (1, 2, 4, 1)
-
-    ar_text_attn_freq: int = 2
-
-    ref_attn_heads: int = 2
-    ref_seconds_max: float = 12.0
-
-    preprompt_sec_max: float = 4.0
-
     sv_student_dim: int = 192
     style_strength: float = 1.0
+
+    ref_enc_layers: int = 2
+    ref_xattn_heads: int = 2
+    ref_xattn_layers: int = 3
+    ref_xattn_gmax: float = 0.35

sopro/hub.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import json
+import os
 import struct
 from typing import Any, Dict, Optional
 
@@ -44,9 +45,7 @@ def load_cfg_from_safetensors(path: str) -> SoproTTSConfig:
     for k in SoproTTSConfig.__annotations__.keys():
         if k in cfg_dict:
             init[k] = cfg_dict[k]
-
-    cfg = SoproTTSConfig(**init)
-    return cfg
+    return SoproTTSConfig(**init)
 
 
 def load_state_dict_from_safetensors(path: str) -> Dict[str, torch.Tensor]: