whisperspeech2 0.9.0__py3-none-any.whl

whisperspeech2/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.9.0"

whisperspeech2/a2wav.py

@@ -0,0 +1,72 @@
+__all__ = ['Vocoder']
+
+from vocos import Vocos
+from whisperspeech2 import inference
+import torch
+
+class Vocoder:
+    def __init__(self, repo_id="charactr/vocos-encodec-24khz", device=None, cache_dir=None):
+        if device is None: device = inference.get_compute_device()
+        if device == 'mps': device = 'cpu'
+        self.device = device
+        self.vocos = Vocos.from_pretrained(repo_id).to(device)
+
+    def is_notebook(self):
+        try:
+            return get_ipython().__class__.__name__ == "ZMQInteractiveShell"
+        except:
+            return False
+
+    @torch.no_grad()
+    def decode(self, atoks):
+        if len(atoks.shape) == 3:
+            b,q,t = atoks.shape
+            atoks = atoks.permute(1,0,2)
+        else:
+            q,t = atoks.shape
+        atoks = atoks.to(self.device)
+        features = self.vocos.codes_to_features(atoks)
+        bandwidth_id = torch.tensor({2: 0, 4: 1, 8: 2}[q]).to(self.device)
+        return self.vocos.decode(features, bandwidth_id=bandwidth_id)
+
+    def _save_audio(self, fname, audio_tensor, sample_rate=24000):
+        try:
+            import torchaudio
+            torchaudio.save(fname, audio_tensor, sample_rate, backend="soundfile")
+            return
+        except (ImportError, RuntimeError, TypeError):
+            pass
+        
+        try:
+            import torchaudio
+            torchaudio.save(fname, audio_tensor, sample_rate)
+            return
+        except (ImportError, RuntimeError):
+            pass
+        
+        try:
+            import soundfile as sf
+            audio_np = audio_tensor.numpy().T
+            sf.write(fname, audio_np, sample_rate)
+            return
+        except ImportError:
+            pass
+        
+        raise ImportError(
+            "No audio backend available. Please install either torchaudio or soundfile:\n"
+            "  pip install torchaudio\n"
+            "or\n"
+            "  pip install soundfile"
+        )
+
+    def decode_to_file(self, fname, atoks):
+        audio = self.decode(atoks)
+        self._save_audio(fname, audio.cpu(), 24000)
+        if self.is_notebook():
+            from IPython.display import display, HTML, Audio
+            display(HTML(f'<a href="{fname}" target="_blank">Listen to {fname}</a>'))
+
+    def decode_to_notebook(self, atoks):
+        from IPython.display import display, HTML, Audio
+        audio = self.decode(atoks)
+        display(Audio(audio.cpu().numpy(), rate=24000))

whisperspeech2/inference.py

@@ -0,0 +1,56 @@
+__all__ = ['get_compute_device']
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from contextlib import nullcontext
+
+def get_default_compute_device():
+    if torch.cuda.is_available() and (torch.version.cuda or torch.version.hip):
+        return 'cuda'
+    elif torch.backends.mps.is_available():
+        return 'mps'
+    else:
+        return 'cpu'
+
+preferred_device = None
+
+def get_compute_device():
+    global preferred_device
+    if preferred_device is None: preferred_device = get_default_compute_device()
+    return preferred_device
+
+def load_model(ref=None, spec=None, device='cpu', cache_dir=None):
+    if spec is not None: return spec
+    if ":" in ref:
+        repo_id, filename = ref.split(":", 1)
+        local_filename = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
+    else:
+        local_filename = ref
+    return torch.load(local_filename, map_location=device)
+
+def inference_context():
+    if torch.cuda.is_available():
+        return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
+    else:
+        return nullcontext()
+
+def multinomial_sample_one_no_sync(probs_sort):
+    q = torch.empty_like(probs_sort).exponential_(1)
+    return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+def logits_to_probs(logits, T=1.0, top_k=None):
+    logits = logits / max(T, 1e-5)
+
+    if top_k is not None:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        pivot = v.select(-1, -1).unsqueeze(-1)
+        logits = torch.where(logits < pivot, -float("Inf"), logits)
+
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    return probs
+
+def sample(logits, T=1.0, top_k=None):
+    probs = logits_to_probs(logits, T, top_k)
+    idx_next = multinomial_sample_one_no_sync(probs)
+    return idx_next

whisperspeech2/languages.py

@@ -0,0 +1,131 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B. Languages.ipynb.
+
+# %% auto 0
+__all__ = ['to_id']
+
+# %% ../nbs/B. Languages.ipynb 3
+LANGUAGES = {
+    "en": "english",
+    "zh": "chinese",
+    "de": "german",
+    "es": "spanish",
+    "ru": "russian",
+    "ko": "korean",
+    "fr": "french",
+    "ja": "japanese",
+    "pt": "portuguese",
+    "tr": "turkish",
+    "pl": "polish",
+    "ca": "catalan",
+    "nl": "dutch",
+    "ar": "arabic",
+    "sv": "swedish",
+    "it": "italian",
+    "id": "indonesian",
+    "hi": "hindi",
+    "fi": "finnish",
+    "vi": "vietnamese",
+    "he": "hebrew",
+    "uk": "ukrainian",
+    "el": "greek",
+    "ms": "malay",
+    "cs": "czech",
+    "ro": "romanian",
+    "da": "danish",
+    "hu": "hungarian",
+    "ta": "tamil",
+    "no": "norwegian",
+    "th": "thai",
+    "ur": "urdu",
+    "hr": "croatian",
+    "bg": "bulgarian",
+    "lt": "lithuanian",
+    "la": "latin",
+    "mi": "maori",
+    "ml": "malayalam",
+    "cy": "welsh",
+    "sk": "slovak",
+    "te": "telugu",
+    "fa": "persian",
+    "lv": "latvian",
+    "bn": "bengali",
+    "sr": "serbian",
+    "az": "azerbaijani",
+    "sl": "slovenian",
+    "kn": "kannada",
+    "et": "estonian",
+    "mk": "macedonian",
+    "br": "breton",
+    "eu": "basque",
+    "is": "icelandic",
+    "hy": "armenian",
+    "ne": "nepali",
+    "mn": "mongolian",
+    "bs": "bosnian",
+    "kk": "kazakh",
+    "sq": "albanian",
+    "sw": "swahili",
+    "gl": "galician",
+    "mr": "marathi",
+    "pa": "punjabi",
+    "si": "sinhala",
+    "km": "khmer",
+    "sn": "shona",
+    "yo": "yoruba",
+    "so": "somali",
+    "af": "afrikaans",
+    "oc": "occitan",
+    "ka": "georgian",
+    "be": "belarusian",
+    "tg": "tajik",
+    "sd": "sindhi",
+    "gu": "gujarati",
+    "am": "amharic",
+    "yi": "yiddish",
+    "lo": "lao",
+    "uz": "uzbek",
+    "fo": "faroese",
+    "ht": "haitian creole",
+    "ps": "pashto",
+    "tk": "turkmen",
+    "nn": "nynorsk",
+    "mt": "maltese",
+    "sa": "sanskrit",
+    "lb": "luxembourgish",
+    "my": "myanmar",
+    "bo": "tibetan",
+    "tl": "tagalog",
+    "mg": "malagasy",
+    "as": "assamese",
+    "tt": "tatar",
+    "haw": "hawaiian",
+    "ln": "lingala",
+    "ha": "hausa",
+    "ba": "bashkir",
+    "jw": "javanese",
+    "su": "sundanese",
+}
+
+# %% ../nbs/B. Languages.ipynb 4
+# language code lookup by name, with a few language aliases
+TO_LANGUAGE_CODE = {
+    **{language: code for code, language in LANGUAGES.items()},
+    "burmese": "my",
+    "valencian": "ca",
+    "flemish": "nl",
+    "haitian": "ht",
+    "letzeburgesch": "lb",
+    "pushto": "ps",
+    "panjabi": "pa",
+    "moldavian": "ro",
+    "moldovan": "ro",
+    "sinhalese": "si",
+    "castilian": "es",
+}
+
+# %% ../nbs/B. Languages.ipynb 5
+languages = tuple(LANGUAGES.keys())
+
+# %% ../nbs/B. Languages.ipynb 6
+def to_id(lang):
+    return languages.index(TO_LANGUAGE_CODE.get(lang, lang))

whisperspeech2/modules.py

@@ -0,0 +1,303 @@
+__all__ = ['LayerNorm', 'LinearHead', 'QueryHead', 'init_transformer', 'sinusoids', 'MultiHeadAttention',
+           'ResidualAttentionBlock', 'BaseDecoder', 'EmbeddingProjector', 'FlexEmbeddings']
+
+import torch
+import numpy as np
+import math
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+from typing import Dict, Iterable, Optional
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+class LinearHead(nn.Linear):
+    pass
+
+class QueryHead(nn.Linear):
+    pass
+
+def init_transformer(m):
+    if isinstance(m, (nn.Linear, nn.Embedding)):
+        torch.nn.init.trunc_normal_(m.weight, std=.02)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            torch.nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.LayerNorm):
+        torch.nn.init.constant_(m.bias, 0)
+        torch.nn.init.constant_(m.weight, 1.0)
+
+def sinusoids(length, channels, max_timescale=10000):
+    assert channels % 2 == 0
+    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, n_state: int, n_head: int, qk_scale: float = 1, rope: bool = False, cross=False):
+        super().__init__()
+        self.n_state = n_state
+        self.n_head = n_head
+        self.sqrt_qk_scale = math.sqrt(qk_scale)
+        self.query = QueryHead(n_state, n_state)
+        self.key = nn.Linear(n_state, n_state, bias=False)
+        self.value = nn.Linear(n_state, n_state)
+        self.out = nn.Linear(n_state, n_state)
+        self.cross = cross
+        self.query_subsampling = 1
+        self.key_subsampling = 1
+
+        self.cached_kvx = None
+        self.register_buffer('k_cache', None)
+        self.register_buffer('v_cache', None)
+        
+        self.rotary = None
+        if rope:
+            self.rotary = Rotary(n_state // n_head)
+        self.qkv = None
+        self.kv = None
+
+    def setup_kv_cache(self, max_batch_size, max_seq_len, dtype=torch.float32):
+        cache_shape = (max_batch_size, self.n_head, max_seq_len, self.n_state//self.n_head)
+        self.k_cache = torch.zeros(cache_shape, dtype=dtype, device=self.key.weight.device)
+        self.v_cache = torch.zeros(cache_shape, dtype=dtype, device=self.value.weight.device)
+
+    def merge_linears(self, layers, mults):
+        bias = [x.bias for x in layers if x.bias is not None][0]
+        din, dout = layers[0].weight.shape
+        new = nn.Linear(din, len(layers) * dout).to(layers[0].weight.device)
+        with torch.no_grad():
+            new.weight[:] = torch.cat([x.weight * m for x,m in zip(layers, mults)])
+            new.bias[:] = torch.cat([torch.zeros_like(bias) if x.bias is None else x.bias * m for x, m in zip(layers, mults)])
+        return new
+
+    def convert_for_eval(self):
+        if self.qkv or self.kv: raise AttributeError("already converted")
+        
+        self.odim = self.key.weight.shape[1]
+        if self.cross:
+            self.q = self.merge_linears([self.query], [self.sqrt_qk_scale])
+            self.kv = self.merge_linears([self.key, self.value],
+                                         [self.sqrt_qk_scale, 1])
+        else:
+            self.qkv = self.merge_linears([self.query, self.key, self.value],
+                                          [self.sqrt_qk_scale, self.sqrt_qk_scale, 1])
+        
+    def split_heads(self, x, x_positions, rope=False, subsampling=1):
+        x = x.view(*x.shape[:2], self.n_head, -1)
+        if rope:
+            x = rope_rotate(x, x_positions * subsampling, *self.rotary(x))
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        qx,
+        q_positions,
+        kvx,
+        kv_positions,
+        causal = False,
+        mask=None,
+    ):
+        if self.k_cache is not None:
+            assert qx.shape[0] <= self.k_cache.shape[0], "please pass in a larger max_batch_size to setup_kv_cache"
+        if self.qkv:
+            q,k,v = self.qkv(qx).split(self.odim, dim=-1)
+        elif self.kv:
+            q = self.q(qx)
+            k,v = self.kv(kvx).split(self.odim, dim=-1)
+        else:
+            q,k,v = None,None,None
+        
+        if q is None: q = self.query(qx) * self.sqrt_qk_scale
+        q = self.split_heads(q, q_positions, rope = self.rotary, subsampling = self.query_subsampling)
+
+        if kvx is not self.cached_kvx:
+            if k is None: k = self.key(kvx) * self.sqrt_qk_scale
+            k = self.split_heads(k, kv_positions, rope = self.rotary, subsampling = self.key_subsampling)
+            if v is None: v = self.value(kvx)
+            v = self.split_heads(v, kv_positions)
+            if self.k_cache is not None:
+                self.k_cache[:k.shape[0],:,kv_positions] = k
+                self.v_cache[:v.shape[0],:,kv_positions] = v
+
+        if self.k_cache is not None:
+            k, v = self.k_cache[:k.shape[0]], self.v_cache[:v.shape[0]]
+
+        if mask is not None:
+            mask = mask[q_positions,:k.shape[-2]]
+            
+        wv = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0, is_causal=causal)
+        
+        return self.out(wv.permute(0, 2, 1, 3).flatten(start_dim=2))
+
+class Rotary(torch.nn.Module):
+    def __init__(self, dim, base=10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.seq_len_cached = None
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def forward(self, x, seq_dim=1):
+        seq_len = x.shape[seq_dim]
+        if not self.seq_len_cached or seq_len > self.seq_len_cached:
+            self.seq_len_cached = 2500
+            
+            t = torch.arange(self.seq_len_cached, device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            self.cos_cached = emb.cos()[None, :, None, :]
+            self.sin_cached = emb.sin()[None, :, None, :]
+        return self.cos_cached, self.sin_cached
+
+def rotate_half(x):
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat(
+        (-x2, x1), dim=len(x.shape)-1
+    )
+
+def rope_rotate(x, positions, cos, sin):
+    return x * cos[:,positions] + rotate_half(x) * sin[:,positions]
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, n_state: int, n_head: int, cross_attention: bool = False, rope: bool = False,
+                 qk_scale: float = 1, ffn_mult: int = 4):
+        super().__init__()
+        self.attn = MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope)
+        self.attn_ln = LayerNorm(n_state)
+
+        self.cross_attn = (
+            MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope, cross=True) if cross_attention else None
+        )
+        self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
+
+        n_mlp = n_state * ffn_mult
+        self.mlp = nn.Sequential(
+            nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state)
+        )
+        self.mlp_ln = LayerNorm(n_state)
+    
+    def setup_kv_cache(self, max_batch_size, max_seq_len, max_cross_seq_len=None):
+        self.attn.setup_kv_cache(max_batch_size, max_seq_len)
+        if self.cross_attn:
+            self.cross_attn.setup_kv_cache(max_batch_size, max_cross_seq_len)
+    
+    def forward(
+        self,
+        x: Tensor,
+        x_positions: Tensor = None,
+        xa: Optional[Tensor] = None,
+        xa_positions: Optional[Tensor] = None,
+        causal = False,
+        mask=None,
+    ):
+        lnx = self.attn_ln(x)
+        x = x + self.attn(lnx, x_positions, lnx, x_positions, causal=causal, mask=mask)
+        if self.cross_attn:
+            lnx = self.cross_attn_ln(x)
+            x = x + self.cross_attn(lnx, x_positions, xa, xa_positions)
+        x = x + self.mlp(self.mlp_ln(x))
+        return x
+
+class BaseDecoder(nn.Module):
+    def __init__(self, depth=6, n_head=6, width=384, qk_scale=1, ffn_mult=4, length=2250, rope=False):
+        super().__init__()
+        self.length = length
+        self.width = width
+        self.layers = nn.ModuleList([
+            ResidualAttentionBlock(
+                self.width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, cross_attention=True, rope=rope
+            ) for _ in range(math.floor(depth))
+        ])
+
+        self.ln_post = LayerNorm(width)
+        
+        mask = torch.empty(length, length).fill_(-torch.inf).triu_(1)
+        self.register_buffer("mask", mask, persistent=False)
+
+    def forward(self, x, x_positions, xenc, xenc_positions):
+        for i,l in enumerate(self.layers):
+            x = l(x, x_positions, xenc, xenc_positions, causal=self.training, mask=self.mask if not self.training else None)
+
+        x = self.ln_post(x)
+
+        return x
+
+class EmbeddingProjector(nn.Linear):
+    pass
+
+class FlexEmbeddings(nn.Module):
+    def __init__(self, codes, width, special_codes=None, frozen_width=None, special_embedding=None, unembed=True):
+        super().__init__()
+        self.codes = codes
+        self.special_codes = special_codes
+        if frozen_width is None: frozen_width = width
+        
+        self.main = nn.Embedding(codes, frozen_width or width)
+        self.emb_to_hidden = EmbeddingProjector(frozen_width, width) if frozen_width != width else None
+        self.hidden_to_emb = EmbeddingProjector(width, frozen_width) if unembed and frozen_width != width else None
+        if special_codes:
+            self.special = special_embedding or nn.Embedding(special_codes, width)
+            
+        self.register_buffer('merged_in', None)
+        self.register_buffer('merged_out', None)
+        self.register_buffer('bias_out', None)
+    
+    def set_frozen_embeddings(self, values):
+        with torch.no_grad():
+            self.main.weight[:] = values
+            self.main.lr_scale = 0
+    
+    @torch.no_grad()
+    def convert_for_eval(self):
+        if not self.special_codes: return
+        main_w = self.main.weight
+        if self.emb_to_hidden is not None: main_w = self.emb_to_hidden(main_w)
+        weight = torch.cat([main_w, self.special.weight], dim=0)
+        self.merged_in = nn.Embedding(*weight.shape, _weight=weight)
+        
+        weight = self.main.weight
+        if self.hidden_to_emb: weight = weight @ self.hidden_to_emb.weight
+        self.merged_out = torch.cat([weight.T, self.special.weight.T], dim=1).T.contiguous()
+        if self.hidden_to_emb:
+            self.bias_out = torch.cat([
+                self.hidden_to_emb.bias @ self.main.weight.T,
+                torch.zeros(self.special.weight.shape[0], device=weight.device, dtype=weight.dtype)
+            ], dim=0)
+        else:
+            self.bias_out = None
+
+    def forward(self, toks):
+        if not self.training and self.merged_in is not None:
+            return self.merged_in(toks)
+        
+        if self.special_codes:
+            special_mask = toks >= self.codes
+            embs = self.main(torch.where(special_mask, 0, toks))
+        else:
+            embs = self.main(toks)
+        
+        if self.emb_to_hidden: embs = self.emb_to_hidden(embs)
+        
+        if self.special_codes:
+            embs[special_mask] = self.special(toks[special_mask] - self.codes).to(embs.dtype)
+        
+        return embs
+    
+    def unembed(self, embs):
+        if not self.training and self.merged_out is not None:
+            return F.linear(embs, self.merged_out, self.bias_out)
+
+        orig_embs = embs
+        if self.hidden_to_emb: embs = self.hidden_to_emb(embs)
+        
+        main_logits = (embs @ self.main.weight.to(embs.dtype).T).float()
+        
+        if not self.special_codes:
+            return main_logits
+        
+        special_logits = (orig_embs @ self.special.weight.to(orig_embs.dtype).T).float()
+        return torch.cat([main_logits, special_logits], dim=-1)

whisperspeech2/pipeline.py

@@ -0,0 +1,107 @@
+from os.path import expanduser
+import torch
+from whisperspeech2.t2s_up_wds_mlang_enclm import TSARTransformer
+from whisperspeech2.s2a_delar_mup_wds_mlang import SADelARTransformer
+from whisperspeech2.a2wav import Vocoder
+from whisperspeech2 import inference, s2a_delar_mup_wds_mlang_cond
+import traceback
+from pathlib import Path
+
+__all__ = ['Pipeline']
+
+class Pipeline:
+    default_speaker = torch.tensor(
+       [-0.2929, -0.4503,  0.4155, -0.1417,  0.0473, -0.1624, -0.2322,  0.7071,
+         0.4800,  0.5496,  0.0410,  0.6236,  0.4729,  0.0587,  0.2194, -0.0466,
+        -0.3036,  0.0497,  0.5028, -0.1703,  0.5039, -0.6464,  0.3857, -0.7350,
+        -0.1605,  0.4808,  0.5397, -0.4851,  0.1774, -0.8712,  0.5789,  0.1785,
+        -0.1417,  0.3039,  0.4232, -0.0186,  0.2685,  0.6153, -0.3103, -0.5706,
+        -0.4494,  0.3394, -0.6184, -0.3617,  1.1041, -0.1178, -0.1885,  0.1997,
+         0.5571, -0.2906, -0.0477, -0.4048, -0.1062,  1.4779,  0.1639, -0.3712,
+        -0.1776, -0.0568, -0.6162,  0.0110, -0.0207, -0.1319, -0.3854,  0.7248,
+         0.0343,  0.5724,  0.0670,  0.0486, -0.3813,  0.1738,  0.3017,  1.0502,
+         0.1550,  0.5708,  0.0366,  0.5093,  0.0294, -0.7091, -0.8220, -0.1583,
+        -0.2343,  0.1366,  0.7372, -0.0631,  0.1505,  0.4600, -0.1252, -0.5245,
+         0.7523, -0.0386, -0.2587,  1.0066, -0.2037,  0.1617, -0.3800,  0.2790,
+         0.0184, -0.5111, -0.7291,  0.1627,  0.2367, -0.0192,  0.4822, -0.4458,
+         0.1457, -0.5884,  0.1909,  0.2563, -0.2035, -0.0377,  0.7771,  0.2139,
+         0.3801,  0.6047, -0.6043, -0.2563, -0.0726,  0.3856,  0.3217,  0.0823,
+        -0.1302,  0.3287,  0.5693,  0.2453,  0.8231,  0.0072,  1.0327,  0.6065,
+        -0.0620, -0.5572,  0.5220,  0.2485,  0.1520,  0.0222, -0.2179, -0.7392,
+        -0.3855,  0.1822,  0.1042,  0.7133,  0.3583,  0.0606, -0.0424, -0.9189,
+        -0.4882, -0.5480, -0.5719, -0.1660, -0.3439, -0.5814, -0.2542,  0.0197,
+         0.4942,  0.0915, -0.0420, -0.0035,  0.5578,  0.1051, -0.0891,  0.2348,
+         0.6876, -0.6685,  0.8215, -0.3692, -0.3150, -0.0462, -0.6806, -0.2661,
+        -0.0308, -0.0050,  0.6756, -0.1647,  1.0734,  0.0049,  0.4969,  0.0259,
+        -0.8949,  0.0731,  0.0886,  0.3442, -0.1433, -0.6804,  0.2204,  0.1859,
+         0.2702,  0.1699, -0.1443, -0.9614,  0.3261,  0.1718,  0.3545, -0.0686]
+    )
+    
+    def __init__(self, t2s_ref=None, s2a_ref=None, optimize=True, torch_compile=False, device=None, cache_dir=None):
+        if device is None: device = inference.get_compute_device()
+        self.device = device
+        args = dict(device = device, cache_dir=cache_dir)
+        try:
+            if t2s_ref:
+                args["ref"] = t2s_ref
+            self.t2s = TSARTransformer.load_model(**args)
+            if optimize: self.t2s.optimize(torch_compile=torch_compile)
+        except:
+            print("Failed to load the T2S model:")
+            print(traceback.format_exc())
+        args = dict(device = device, cache_dir=cache_dir)
+        try:
+            if s2a_ref:
+                spec = inference.load_model(ref=s2a_ref, device=device, cache_dir=cache_dir)
+                if [x for x in spec['state_dict'].keys() if x.startswith('cond_embeddings.')]:
+                    cls = s2a_delar_mup_wds_mlang_cond.SADelARTransformer
+                    args['spec'] = spec
+                else:
+                    cls = SADelARTransformer
+                    args['spec'] = spec
+            else:
+                cls = SADelARTransformer
+            self.s2a = cls.load_model(**args)
+            if optimize: self.s2a.optimize(torch_compile=torch_compile)
+        except:
+            print("Failed to load the S2A model:")
+            print(traceback.format_exc())
+
+        self.vocoder = Vocoder(device=device)
+        self.encoder = None
+
+    def extract_spk_emb(self, fname):
+        import torchaudio
+        if self.encoder is None:
+            device = self.device
+            if device == 'mps': device = 'cpu'
+            from speechbrain.pretrained import EncoderClassifier
+            self.encoder = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
+                                                          savedir=expanduser("~/.cache/speechbrain/"),
+                                                          run_opts={"device": device})
+        audio_info = torchaudio.info(fname)
+        actual_sample_rate = audio_info.sample_rate
+        num_frames = actual_sample_rate * 30
+        samples, sr = torchaudio.load(fname, num_frames=num_frames)
+        samples = samples[:, :num_frames]
+        samples = self.encoder.audio_normalizer(samples[0], sr)
+        spk_emb = self.encoder.encode_batch(samples.unsqueeze(0))
+        
+        return spk_emb[0,0].to(self.device)
+        
+    def generate_atoks(self, text, speaker=None, lang='en', cps=15, step_callback=None):
+        if speaker is None: speaker = self.default_speaker
+        elif isinstance(speaker, (str, Path)): speaker = self.extract_spk_emb(speaker)
+        text = text.replace("\n", " ")
+        stoks = self.t2s.generate(text, cps=cps, lang=lang, step=step_callback)[0]
+        atoks = self.s2a.generate(stoks, speaker.unsqueeze(0), step=step_callback)
+        return atoks
+        
+    def generate(self, text, speaker=None, lang='en', cps=15, step_callback=None):
+        return self.vocoder.decode(self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=step_callback))
+    
+    def generate_to_file(self, fname, text, speaker=None, lang='en', cps=15, step_callback=None):
+        self.vocoder.decode_to_file(fname, self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=None))
+        
+    def generate_to_notebook(self, text, speaker=None, lang='en', cps=15, step_callback=None):
+        self.vocoder.decode_to_notebook(self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=None))

whisperspeech2/quick_test.py

@@ -0,0 +1,8 @@
+from whisperspeech2.pipeline import Pipeline
+
+tts_pipe = Pipeline(s2a_ref='collabora/whisperspeech:s2a-q4-tiny-en+pl.model') # uncomment the line for the model you want to use
+# tts_pipe = Pipeline(s2a_ref='collabora/whisperspeech:s2a-q4-base-en+pl.model') # uncomment the line for the model you want to use
+# tts_pipe = Pipeline(s2a_ref='collabora/whisperspeech:s2a-q4-small-en+pl.model') # uncomment the line for the model you want to use
+
+save_path = 'output.wav' # change the file extension to .mp3, .flac, .ogg etc. to save to a different file format
+tts_pipe.generate_to_file(save_path, "This is a test")