sopro 1.0.1__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.1"
+__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")
@@ -97,74 +97,59 @@ def main() -> None:
         arr = np.load(args.ref_tokens)
         ref_tokens_tq = torch.from_numpy(arr).long()
 
-    text_ids = tts.encode_text(args.text)
-    ref = tts.encode_reference(
-        ref_audio_path=args.ref_audio,
-        ref_tokens_tq=ref_tokens_tq,
-        ref_seconds=args.ref_seconds,
-    )
+    with torch.inference_mode():
+        text_ids = tts.encode_text(args.text)
+        ref = tts.prepare_reference(
+            ref_audio_path=args.ref_audio,
+            ref_tokens_tq=ref_tokens_tq,
+            ref_seconds=args.ref_seconds,
+        )
 
-    prep = tts.model.prepare_conditioning(
-        text_ids,
-        ref,
-        max_frames=args.max_frames,
-        device=tts.device,
-        style_strength=float(
-            args.style_strength
-            if args.style_strength is not None
-            else cfg.style_strength
-        ),
-    )
+        prep = tts.model.prepare_conditioning(
+            text_ids,
+            ref,
+            max_frames=args.max_frames,
+            device=tts.device,
+            style_strength=float(
+                args.style_strength
+                if args.style_strength is not None
+                else cfg.style_strength
+            ),
+        )
 
-    t_start = time.perf_counter()
+        t_start = time.perf_counter()
 
     hist_A: list[int] = []
     pbar = tqdm(
-        total=args.max_frames,
-        desc="AR sampling",
-        unit="frame",
-        disable=args.quiet,
+        total=args.max_frames + 1, desc="AR sampling", unit="step", disable=args.quiet
     )
 
-    for _t, rvq1, p_stop in tts.model.ar_stream(
+    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),
-        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))
+        if is_eos:
+            pbar.set_postfix(eos="yes")
+            pbar.update(1)
+            break
+        hist_A.append(int(tok))
         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}")
+
+    t_after_sampling = time.perf_counter()
 
     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, :]
+    cond_seq = prep["cond_ar"][:, :T, :]
     tokens_1xTQ = tts.model.nar_refine(cond_seq, tokens_A)
     tokens_tq = tokens_1xTQ.squeeze(0)
 

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]: