sopro 1.0.0__py3-none-any.whl

sopro/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from .model import SoproTTS
+
+__all__ = ["SoproTTS"]
+__version__ = "1.0.0"

sopro/audio.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import os
+from typing import Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from .constants import TARGET_SR
+
+try:
+    import torchaudio
+    import torchaudio.functional as AF
+except Exception:
+    torchaudio = None
+    AF = None
+
+try:
+    import soundfile as sf
+except Exception:
+    sf = None
+
+
+def device_str() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+def trim_silence_energy(
+    wav: torch.Tensor,
+    sr: int,
+    frame_ms: float = 25.0,
+    hop_ms: float = 10.0,
+    thresh_db_floor: float = -40.0,
+    prepad_ms: float = 30.0,
+    postpad_ms: float = 30.0,
+    min_keep_sec: float = 0.5,
+) -> torch.Tensor:
+    orig_1d = wav.ndim == 1
+    if orig_1d:
+        wav = wav.unsqueeze(0)
+
+    C, T = wav.shape
+    if T == 0:
+        return wav.squeeze(0) if orig_1d else wav
+    if T < int(sr * 0.1):
+        return wav.squeeze(0) if orig_1d else wav
+
+    frame_len = max(1, int(sr * frame_ms / 1000.0))
+    hop = max(1, int(sr * hop_ms / 1000.0))
+    if T < frame_len:
+        return wav.squeeze(0) if orig_1d else wav
+
+    mono = wav.mean(dim=0, keepdim=True)
+    frames = mono.unfold(-1, frame_len, hop)
+    energy = frames.pow(2).mean(dim=-1).squeeze(0)
+
+    eps = 1e-10
+    energy_db = 10.0 * torch.log10(energy + eps)
+    max_db = float(energy_db.max().item())
+
+    rel_thresh = max_db + thresh_db_floor
+    thresh_db = max(rel_thresh, thresh_db_floor)
+
+    voiced = energy_db > thresh_db
+    idx = torch.nonzero(voiced, as_tuple=False)
+    if idx.numel() == 0:
+        return wav.squeeze(0) if orig_1d else wav
+
+    first_frame = int(idx[0, 0].item())
+    last_frame = int(idx[-1, 0].item())
+
+    prepad_samples = int(sr * prepad_ms / 1000.0)
+    postpad_samples = int(sr * postpad_ms / 1000.0)
+
+    start = max(0, first_frame * hop - prepad_samples)
+
+    end = min(T, last_frame * hop + frame_len + postpad_samples)
+
+    min_keep = int(min_keep_sec * sr)
+    if end <= start or (end - start) < min_keep:
+        return wav.squeeze(0) if orig_1d else wav
+
+    out = wav[:, start:end]
+    return out.squeeze(0) if orig_1d else out
+
+
+def load_audio_file(path: str) -> Tuple[torch.Tensor, int]:
+    if sf is not None:
+        wav_np, sr = sf.read(path, dtype="float32", always_2d=True)
+        wav = torch.from_numpy(wav_np).transpose(0, 1)
+    elif torchaudio is not None:
+        wav, sr = torchaudio.load(path)
+        if wav.dtype != torch.float32:
+            if wav.dtype == torch.int16:
+                wav = wav.float() / (2**15)
+            else:
+                wav = wav.float()
+    else:
+        raise RuntimeError("Install 'soundfile' or 'torchaudio' to read audio.")
+
+    if wav.size(0) > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+    return wav, sr
+
+
+def resample(
+    wav: torch.Tensor, sr_in: int, sr_out: int, device: Optional[str] = None
+) -> torch.Tensor:
+    device = device or device_str()
+    wav = wav.to(device)
+    if sr_in == sr_out:
+        return wav
+    if sr_in != sr_out and AF is None:
+        raise RuntimeError("Resampling requires torchaudio. pip install torchaudio")
+    return AF.resample(wav, sr_in, sr_out)
+
+
+def save_audio(path: str, wav: torch.Tensor, sr: int = TARGET_SR) -> None:
+    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
+
+    wav = wav.detach().cpu()
+
+    if wav.ndim == 1:
+        wav = wav.unsqueeze(0)
+    elif wav.ndim == 2:
+        pass
+    elif wav.ndim == 3:
+        wav = wav[0]
+    else:
+        raise ValueError(f"Expected wav with 1-3 dims, got shape {tuple(wav.shape)}")
+
+    if wav.size(0) > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+
+    if sf is not None:
+        sf.write(path, wav[0].numpy(), sr)
+        return
+
+    if torchaudio is not None:
+        torchaudio.save(path, wav, sample_rate=sr)
+        return
+
+    raise RuntimeError("Install 'soundfile' or 'torchaudio' to write audio.")
+
+
+def center_crop_audio(wav: torch.Tensor, win_samples: int) -> torch.Tensor:
+    if win_samples <= 0:
+        return wav
+    T = int(wav.shape[-1])
+    if T <= win_samples:
+        return wav
+    s = (T - win_samples) // 2
+    return wav[..., s : s + win_samples]

sopro/cli.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+from tqdm.auto import tqdm
+
+from .audio import save_audio
+from .constants import TARGET_SR
+from .model import SoproTTS
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="SoproTTS cli inference")
+
+    ap.add_argument("--repo_id", type=str, default="samuel-vitorino/sopro")
+    ap.add_argument(
+        "--revision", type=str, default=None, help="Optional git revision/branch/tag"
+    )
+    ap.add_argument("--cache_dir", type=str, default=None, help="Optional HF cache dir")
+    ap.add_argument("--hf_token", type=str, default=None, help="HF token")
+
+    ap.add_argument("--text", type=str, required=True)
+    ap.add_argument("--ref_audio", type=str, default=None)
+    ap.add_argument("--ref_tokens", type=str, default=None, help="Path to .npy [T,Q]")
+    ap.add_argument("--out", type=str, required=True)
+
+    ap.add_argument("--max_frames", type=int, default=400)
+
+    ap.add_argument("--top_p", type=float, default=0.9)
+    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)
+
+    ap.add_argument("--seed", type=int, default=None, help="Random seed for sampling")
+
+    ap.add_argument(
+        "--no_stop_head", action="store_true", help="Disable stop head early stopping"
+    )
+    ap.add_argument(
+        "--stop_patience", type=int, default=None, help="Override cfg.stop_patience"
+    )
+    ap.add_argument(
+        "--stop_threshold", type=float, default=None, help="Override cfg.stop_threshold"
+    )
+
+    ap.add_argument(
+        "--device",
+        type=str,
+        default=None,
+        choices=["cpu", "cuda", "mps"],
+        help="Device to run on (default: cuda if available else cpu)",
+    )
+
+    ap.add_argument("--quiet", action="store_true")
+
+    args = ap.parse_args()
+
+    default_device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = args.device or default_device
+
+    if device == "cuda" and not torch.cuda.is_available():
+        raise SystemExit("Error: --device cuda requested but CUDA is not available.")
+    if device == "mps" and not (
+        hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+    ):
+        raise SystemExit("Error: --device mps requested but MPS is not available.")
+
+    if args.seed is not None:
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(args.seed)
+
+    t0 = time.perf_counter()
+    tts = SoproTTS.from_pretrained(
+        args.repo_id,
+        revision=args.revision,
+        cache_dir=args.cache_dir,
+        token=args.hf_token,
+        device=device,
+    )
+    t1 = time.perf_counter()
+    if not args.quiet:
+        print(f"[Load] {t1 - t0:.2f}s")
+
+    cfg = tts.cfg
+
+    ref_tokens_tq = None
+    if args.ref_tokens is not None:
+        import numpy as np
+
+        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,
+    )
+
+    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()
+
+    hist_A: list[int] = []
+    pbar = tqdm(
+        total=args.max_frames,
+        desc="AR sampling",
+        unit="frame",
+        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))
+        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)
+
+    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}")
+
+
+if __name__ == "__main__":
+    main()

sopro/codec/__init__.py

@@ -0,0 +1,3 @@
+from .mimi import MimiCodec, MimiDecodeState, MimiStreamDecoder
+
+__all__ = ["MimiCodec", "MimiStreamDecoder", "MimiDecodeState"]

sopro/codec/mimi.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+
+import torch
+
+from ..audio import center_crop_audio, load_audio_file, resample, trim_silence_energy
+from ..constants import DEFAULT_MIMI_ID, TARGET_SR
+
+try:
+    from transformers import MimiConfig, MimiModel
+except Exception:
+    MimiModel = None
+    MimiConfig = None
+
+
+class MimiCodec:
+    def __init__(
+        self, num_quantizers: int, device: str = "cuda", model_id: str = DEFAULT_MIMI_ID
+    ):
+        if MimiModel is None or MimiConfig is None:
+            raise RuntimeError(
+                "MimiModel missing. Install a recent transformers version."
+            )
+
+        self.device = torch.device(device)
+        cfg = MimiConfig.from_pretrained(model_id, num_quantizers=int(num_quantizers))
+        self.model = (
+            MimiModel.from_pretrained(model_id, config=cfg).to(self.device).eval()
+        )
+
+    @property
+    def codebook_size(self) -> int:
+        return int(getattr(self.model.config, "codebook_size", 2048))
+
+    @property
+    def num_quantizers(self) -> int:
+        return int(getattr(self.model.config, "num_quantizers", 32))
+
+    @torch.no_grad()
+    def encode_file(
+        self, wav_path: str, *, crop_seconds: Optional[float] = None
+    ) -> torch.Tensor:
+        wav, sr = load_audio_file(wav_path)
+
+        wav = trim_silence_energy(wav, sr)
+
+        cfg = self.model.config
+        sr_target = int(getattr(cfg, "sampling_rate", TARGET_SR))
+        wav = resample(wav, sr, sr_target, device=str(self.device))
+
+        if crop_seconds is not None and crop_seconds > 0:
+            fps = float(getattr(cfg, "frame_rate", 12.5))
+            hop = int(round(sr_target / fps))
+            win_frames = max(1, int(round(crop_seconds * fps)))
+            win_samples = win_frames * hop
+            wav = center_crop_audio(wav, win_samples)
+
+        wav = wav.unsqueeze(0)
+        out = self.model.encode(wav, return_dict=True)
+        codes_bqt = out.audio_codes
+        return codes_bqt[0].permute(1, 0).contiguous()
+
+    @torch.no_grad()
+    def decode_full(self, codes_tq: torch.Tensor) -> torch.Tensor:
+        audio_codes = codes_tq.to(self.device).permute(1, 0).unsqueeze(0).contiguous()
+        out = self.model.decode(audio_codes=audio_codes, return_dict=True)
+        wav = out.audio_values
+        if wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+        return wav
+
+
+@dataclass
+class MimiDecodeState:
+    decoder_past_key_values: Optional[object] = None
+    frames_seen: int = 0
+    samples_emitted: int = 0
+    tail_codes_tq: Optional[torch.Tensor] = None
+
+
+class MimiStreamDecoder:
+    def __init__(self, codec: MimiCodec):
+        self.codec = codec
+        self.codec.model.config.use_cache = True
+
+    def drop_cache_tail(self, pkv: Any, n: int):
+        if pkv is None or n <= 0:
+            return pkv
+
+        if hasattr(pkv, "to_legacy_cache") and hasattr(type(pkv), "from_legacy_cache"):
+            legacy = pkv.to_legacy_cache()
+
+            trimmed = tuple(
+                (
+                    k[..., : max(0, k.shape[-2] - n), :].contiguous(),
+                    v[..., : max(0, v.shape[-2] - n), :].contiguous(),
+                )
+                for (k, v) in legacy
+            )
+            return type(pkv).from_legacy_cache(trimmed)
+
+        if isinstance(pkv, tuple):
+            return tuple(
+                (
+                    k[..., : max(0, k.shape[-2] - n), :].contiguous(),
+                    v[..., : max(0, v.shape[-2] - n), :].contiguous(),
+                )
+                for (k, v) in pkv
+            )
+
+        return pkv
+
+    @torch.inference_mode()
+    def decode_step(
+        self,
+        codes_chunk_tq: torch.Tensor,
+        state: Optional[MimiDecodeState] = None,
+        *,
+        overlap_frames: int = 2,
+    ) -> Tuple[torch.Tensor, MimiDecodeState]:
+        if state is None:
+            state = MimiDecodeState()
+
+        cfg = self.codec.model.config
+        sr = int(getattr(cfg, "sampling_rate", 24000))
+        fps = float(getattr(cfg, "frame_rate", 12.5))
+        hop = int(round(sr / fps))
+
+        n_new = int(codes_chunk_tq.size(0))
+        if n_new == 0:
+            return torch.zeros(1, 0, device=self.codec.device), state
+
+        tail = state.tail_codes_tq
+        ov = 0
+        if overlap_frames > 0 and tail is not None and tail.numel() > 0:
+            ov = min(int(overlap_frames), int(tail.size(0)))
+            tail = tail[-ov:]
+            codes_in_tq = torch.cat([tail, codes_chunk_tq], dim=0)
+        else:
+            codes_in_tq = codes_chunk_tq
+
+        pkv = state.decoder_past_key_values
+        if ov > 0 and pkv is not None:
+            pkv = self.drop_cache_tail(pkv, ov)
+
+        audio_codes = (
+            codes_in_tq.to(self.codec.device).permute(1, 0).unsqueeze(0).contiguous()
+        )
+
+        out = self.codec.model.decode(
+            audio_codes=audio_codes,
+            decoder_past_key_values=pkv,
+            return_dict=True,
+        )
+
+        wav = out.audio_values
+        if wav.ndim == 1:
+            wav = wav.unsqueeze(0)
+        else:
+            wav = wav.reshape(1, -1)
+
+        expected_total = int((ov + n_new) * hop)
+        if wav.size(1) >= expected_total:
+            wav = wav[:, :expected_total]
+
+        ov_samp = min(int(ov * hop), int(wav.size(1)))
+        wav_new = wav[:, ov_samp:]
+
+        state.decoder_past_key_values = getattr(out, "decoder_past_key_values", None)
+        state.frames_seen += n_new
+        state.samples_emitted += int(wav_new.size(1))
+
+        if overlap_frames > 0:
+            keep = min(int(overlap_frames), int(codes_in_tq.size(0)))
+            state.tail_codes_tq = codes_in_tq[-keep:].detach()
+        else:
+            state.tail_codes_tq = None
+
+        return wav_new, state

sopro/config.py

@@ -0,0 +1,48 @@
+from dataclasses import dataclass
+from typing import Tuple
+
+from sopro.constants import TARGET_SR
+
+
+@dataclass
+class SoproTTSConfig:
+    num_codebooks: int = 32
+    codebook_size: int = 2048
+    mimi_fps: float = 12.5
+    max_frames: int = 400
+    audio_sr: int = TARGET_SR
+
+    d_model: int = 384
+    n_layers_text: int = 4
+    n_layers_ar: int = 6
+    n_layers_nar: int = 6
+    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
+    min_gen_frames: int = 12
+
+    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

sopro/constants.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+TARGET_SR: int = 24000
+
+DEFAULT_MIMI_ID: str = "kyutai/mimi"

sopro/hub.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import json
+import struct
+from typing import Any, Dict, Optional
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+from sopro.config import SoproTTSConfig
+
+
+def download_repo(
+    repo_id: str,
+    *,
+    revision: Optional[str] = None,
+    cache_dir: Optional[str] = None,
+    token: Optional[str] = None,
+) -> str:
+    return snapshot_download(
+        repo_id=repo_id,
+        revision=revision,
+        cache_dir=cache_dir,
+        token=token,
+    )
+
+
+def _read_safetensors_metadata(path: str) -> Dict[str, str]:
+    with open(path, "rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_len).decode("utf-8"))
+    meta = header.get("__metadata__", {}) or {}
+    return {str(k): str(v) for k, v in meta.items()}
+
+
+def load_cfg_from_safetensors(path: str) -> SoproTTSConfig:
+    meta = _read_safetensors_metadata(path)
+    if "cfg" not in meta:
+        raise RuntimeError(f"No 'cfg' metadata found in {path}.")
+
+    cfg_dict = json.loads(meta["cfg"])
+    init: Dict[str, Any] = {}
+    for k in SoproTTSConfig.__annotations__.keys():
+        if k in cfg_dict:
+            init[k] = cfg_dict[k]
+
+    cfg = SoproTTSConfig(**init)
+    return cfg
+
+
+def load_state_dict_from_safetensors(path: str) -> Dict[str, torch.Tensor]:
+    return load_file(path)