tiny-tts 0.2.0__py3-none-any.whl

tiny_tts/__init__.py

@@ -0,0 +1,102 @@
+import os
+import torch
+import soundfile as sf
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text.advanced_normalization import advanced_normalize
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.models.synthesizer import VoiceSynthesizer
+from tiny_tts.text.symbols import symbols
+from tiny_tts.utils.config import (
+    SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
+    N_SPEAKERS, SPK2ID, MODEL_PARAMS,
+)
+from tiny_tts.infer import load_engine
+
+class TinyTTS:
+    def __init__(self, checkpoint_path=None, device=None):
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+            
+        if checkpoint_path is None:
+            # Look for default checkpoint in pacakage
+            pkg_dir = os.path.dirname(os.path.abspath(__file__))
+            default_ckpt = os.path.join(os.path.dirname(pkg_dir), "checkpoints", "G.pth")
+            # 2. Check HuggingFace Cache / Download
+            if not os.path.exists(default_ckpt):
+                try:
+                    from huggingface_hub import hf_hub_download
+                    print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
+                    default_ckpt = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
+                except ImportError:
+                    raise ImportError("huggingface_hub is required to auto-download the model. Run: pip install huggingface_hub")
+                except Exception as e:
+                    raise ValueError(f"Failed to download checkpoint from Hugging Face: {e}")
+
+            checkpoint_path = default_ckpt
+                
+        self.model = load_engine(checkpoint_path, self.device)
+
+    def speak(self, text, output_path="output.wav", speaker="MALE", speed=1.0, use_advanced_normalization=True):
+        """Synthesize text to speech and save to output_path.
+        
+        Args:
+            text: Input text
+            output_path: Output audio file path
+            speaker: Speaker ID (default: "MALE")
+            speed: Speech speed (1.0=normal, 1.5=faster, 0.7=slower)
+            use_advanced_normalization: Enable advanced text normalization (URLs, emails, etc.)
+        """
+        print(f"Synthesizing: {text}")
+
+        # Normalize text
+        if use_advanced_normalization:
+            normalized = advanced_normalize(text)
+        else:
+            normalized = normalize_text(text)
+
+        # Phonemize
+        phones, tones, word2ph = grapheme_to_phoneme(normalized)
+
+        # Convert to sequence
+        phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+        # Add blanks
+        if ADD_BLANK:
+            phone_ids = commons.insert_blanks(phone_ids, 0)
+            tone_ids = commons.insert_blanks(tone_ids, 0)
+            lang_ids = commons.insert_blanks(lang_ids, 0)
+
+        x = torch.LongTensor(phone_ids).unsqueeze(0).to(self.device)
+        x_lengths = torch.LongTensor([len(phone_ids)]).to(self.device)
+        tone = torch.LongTensor(tone_ids).unsqueeze(0).to(self.device)
+        language = torch.LongTensor(lang_ids).unsqueeze(0).to(self.device)
+
+        # Speaker ID
+        if speaker not in SPK2ID:
+            print(f"Warning: Speaker '{speaker}' not found, using ID 0. Available: {list(SPK2ID.keys())}")
+            sid = torch.LongTensor([0]).to(self.device)
+        else:
+            sid = torch.LongTensor([SPK2ID[speaker]]).to(self.device)
+
+        # BERT features (disabled - using zero tensors)
+        bert = torch.zeros(1024, len(phone_ids)).to(self.device).unsqueeze(0)
+        ja_bert = torch.zeros(768, len(phone_ids)).to(self.device).unsqueeze(0)
+
+        # speed > 1.0 = faster speech, < 1.0 = slower speech
+        length_scale = 1.0 / speed
+
+        with torch.no_grad():
+            audio, *_ = self.model.infer(
+                x, x_lengths, sid, tone, language, bert, ja_bert,
+                noise_scale=0.667,
+                noise_scale_w=0.8,
+                length_scale=length_scale
+            )
+
+        audio_np = audio[0, 0].cpu().numpy()
+        sf.write(output_path, audio_np, SAMPLING_RATE)
+        print(f"Saved audio to {output_path}")
+        return audio_np

tiny_tts/alignment/__init__.py

@@ -0,0 +1,16 @@
+from numpy import zeros, int32, float32
+from torch import from_numpy
+
+from .core import viterbi_decode_kernel
+
+
+def viterbi_decode(neg_cent, mask):
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+    path = zeros(neg_cent.shape, dtype=int32)
+
+    t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+    t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+    viterbi_decode_kernel(path, neg_cent, t_t_max, t_s_max)
+    return from_numpy(path).to(device=device, dtype=dtype)

tiny_tts/alignment/core.py

@@ -0,0 +1,46 @@
+import numba
+
+
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython=True,
+    nogil=True,
+)
+def viterbi_decode_kernel(paths, values, t_ys, t_xs):
+    b = paths.shape[0]
+    max_neg_val = -1e9
+    for i in range(int(b)):
+        path = paths[i]
+        value = values[i]
+        t_y = t_ys[i]
+        t_x = t_xs[i]
+
+        v_prev = v_cur = 0.0
+        index = t_x - 1
+
+        for y in range(t_y):
+            for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+                if x == y:
+                    v_cur = max_neg_val
+                else:
+                    v_cur = value[y - 1, x]
+                if x == 0:
+                    if y == 0:
+                        v_prev = 0.0
+                    else:
+                        v_prev = max_neg_val
+                else:
+                    v_prev = value[y - 1, x - 1]
+                value[y, x] += max(v_prev, v_cur)
+
+        for y in range(t_y - 1, -1, -1):
+            path[y, index] = 1
+            if index != 0 and (
+                index == y or value[y - 1, index] < value[y - 1, index - 1]
+            ):
+                index = index - 1

tiny_tts/infer.py

@@ -0,0 +1,191 @@
+import os
+import sys
+import re
+import torch
+import soundfile as sf
+import argparse
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.models import VoiceSynthesizer
+from tiny_tts.text.symbols import symbols
+from tiny_tts.utils import (
+    SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
+    N_SPEAKERS, SPK2ID, MODEL_PARAMS,
+)
+
+
+def load_engine(checkpoint_path, device='cuda'):
+    print(f"Loading model from {checkpoint_path}")
+    net_g = VoiceSynthesizer(
+        len(symbols),
+        SPEC_CHANNELS,
+        SEGMENT_FRAMES,
+        n_speakers=N_SPEAKERS,
+        **MODEL_PARAMS
+    ).to(device)
+
+    # Count model parameters
+    total_params = sum(p.numel() for p in net_g.parameters())
+    trainable_params = sum(p.numel() for p in net_g.parameters() if p.requires_grad)
+    print(f"Model parameters: {total_params/1e6:.2f}M total, {trainable_params/1e6:.2f}M trainable")
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    state_dict = checkpoint['model']
+
+    # Remove module. prefix and filter shape mismatches
+    model_state = net_g.state_dict()
+    new_state_dict = {}
+    skipped = []
+    for k, v in state_dict.items():
+        key = k[7:] if k.startswith('module.') else k
+        if key in model_state:
+            if v.shape == model_state[key].shape:
+                new_state_dict[key] = v
+            else:
+                skipped.append(f"{key}: ckpt{v.shape} vs model{model_state[key].shape}")
+        else:
+            new_state_dict[key] = v
+
+    if skipped:
+        print(f"Skipped {len(skipped)} mismatched keys:")
+        for s in skipped[:5]:
+            print(f"  {s}")
+        if len(skipped) > 5:
+            print(f"  ... and {len(skipped)-5} more")
+
+    net_g.load_state_dict(new_state_dict, strict=False)
+    net_g.eval()
+
+    # Fold weight_norm into weight tensors for faster inference (~18% speedup)
+    net_g.dec.remove_weight_norm()
+
+    return net_g
+
+
+def synthesize(text, output_path, model, speaker="MALE", device='cuda', speed=1.0):
+    print(f"Synthesizing: {text}")
+
+    # Normalize text
+    normalized = normalize_text(text)
+
+    # Phonemize
+    phones, tones, word2ph = grapheme_to_phoneme(normalized)
+
+    # Convert to sequence
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+    # Add blanks
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+
+    x = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
+    x_lengths = torch.LongTensor([len(phone_ids)]).to(device)
+    tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
+    language = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
+
+    # Speaker ID
+    if speaker not in SPK2ID:
+        print(f"Warning: Speaker {speaker} not found, using ID 0")
+        sid = torch.LongTensor([0]).to(device)
+    else:
+        sid = torch.LongTensor([SPK2ID[speaker]]).to(device)
+
+    # BERT features (disabled - using zero tensors)
+    bert = torch.zeros(1024, len(phone_ids)).to(device).unsqueeze(0)
+    ja_bert = torch.zeros(768, len(phone_ids)).to(device).unsqueeze(0)
+
+    # speed > 1.0 = faster speech, < 1.0 = slower speech
+    length_scale = 1.0 / speed
+
+    with torch.no_grad():
+        audio, *_ = model.infer(
+            x, x_lengths, sid, tone, language, bert, ja_bert,
+            noise_scale=0.667,
+            noise_scale_w=0.8,
+            length_scale=length_scale
+        )
+
+    audio = audio[0, 0].cpu().numpy()
+    sf.write(output_path, audio, SAMPLING_RATE)
+    print(f"Saved audio to {output_path}")
+
+
+def get_latest_checkpoint(checkpoint_dir):
+    """Finds the latest G_*.pth checkpoint in the given directory."""
+    checkpoints = [f for f in os.listdir(checkpoint_dir) if f.startswith('G_') and f.endswith('.pth')]
+    if not checkpoints:
+        return None
+
+    def get_step(filename):
+        match = re.search(r'_(\d+)\.pth', filename)
+        return int(match.group(1)) if match else -1
+
+    latest_ckpt = max(checkpoints, key=get_step)
+    return os.path.join(checkpoint_dir, latest_ckpt)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="TinyTTS — English Text-to-Speech Inference")
+    parser.add_argument("--text", "-t", type=str, default="The weather is nice today, and I feel very relaxed.", help="Text to synthesize")
+    parser.add_argument("--checkpoint", "-c", type=str, default=None, help="Path to checkpoint. Auto-downloads if not provided.")
+    parser.add_argument("--output", "-o", type=str, default="output.wav", help="Output audio file path")
+    parser.add_argument("--speaker", "-s", type=str, default="MALE", help="Speaker ID")
+    parser.add_argument("--speed", type=float, default=1.0, help="Speech speed (1.0=normal, 1.5=faster, 0.7=slower)")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to use (cuda or cpu)")
+
+    args = parser.parse_args()
+
+    if args.checkpoint is None:
+        try:
+            from huggingface_hub import hf_hub_download
+            print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
+            args.checkpoint = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
+        except ImportError:
+            print("Error: huggingface_hub is required for auto-download. Run: pip install huggingface_hub")
+            sys.exit(1)
+        except Exception as e:
+            print(f"Error downloading checkpoint: {e}")
+            sys.exit(1)
+
+    if not os.path.exists(args.checkpoint):
+        print(f"Error: Checkpoint or directory not found at {args.checkpoint}")
+        sys.exit(1)
+
+    if os.path.isdir(args.checkpoint):
+        latest_ckpt = get_latest_checkpoint(args.checkpoint)
+        if not latest_ckpt:
+            print(f"Error: No G_*.pth checkpoints found in directory {args.checkpoint}")
+            sys.exit(1)
+        args.checkpoint = latest_ckpt
+        print(f"Auto-detected latest checkpoint: {args.checkpoint}")
+
+    # Extract step from checkpoint filename
+    ckpt_basename = os.path.basename(args.checkpoint)
+    match = re.search(r'_(\d+)\.pth', ckpt_basename)
+    step_str = match.group(1) if match else "unknown"
+
+    # Save to output folder
+    out_dir = "infer_outputs"
+    os.makedirs(out_dir, exist_ok=True)
+
+    out_name = os.path.basename(args.output)
+    name, ext = os.path.splitext(out_name)
+    model = load_engine(args.checkpoint, args.device)
+
+    if args.speaker.lower() == "all":
+        if not SPK2ID:
+            print("Error: No speakers found")
+            sys.exit(1)
+        print(f"Synthesizing for all {len(SPK2ID)} speakers...")
+        for spk in SPK2ID.keys():
+            final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{spk}{ext}")
+            synthesize(args.text, final_output, model, speaker=spk, device=args.device, speed=args.speed)
+    else:
+        final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{args.speaker}{ext}")
+        synthesize(args.text, final_output, model, speaker=args.speaker, device=args.device, speed=args.speed)
+
+if __name__ == "__main__":
+    main()

tiny_tts/infer_onnx.py

@@ -0,0 +1,199 @@
+"""
+ONNX Runtime inference engine for TinyTTS.
+
+Replaces the PyTorch VoiceSynthesizer.infer() with equivalent
+ONNX Runtime sessions + NumPy ops for the non-exported parts
+(alignment path computation).
+"""
+import os
+import numpy as np
+import soundfile as sf
+
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.utils.config import (
+    SAMPLING_RATE, ADD_BLANK, SPK2ID,
+)
+
+try:
+    import onnxruntime as ort
+except ImportError:
+    raise ImportError("onnxruntime is required. Run: pip install onnxruntime")
+
+
+def _build_session(path: str, use_gpu: bool = False):
+    """Create an ORT InferenceSession with optional GPU support."""
+    providers = (
+        ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        if use_gpu else
+        ["CPUExecutionProvider"]
+    )
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    opts.intra_op_num_threads = os.cpu_count() or 4
+    return ort.InferenceSession(path, sess_options=opts, providers=providers)
+
+
+def _create_length_mask_np(lengths, max_len=None):
+    """NumPy equivalent of commons.create_length_mask."""
+    if max_len is None:
+        max_len = int(lengths.max())
+    ids = np.arange(max_len, dtype=np.float32)     # [T]
+    mask = (ids[None, :] < lengths[:, None]).astype(np.float32)  # [B, T]
+    return mask
+
+
+def _compute_alignment_path_np(w_ceil, attn_mask):
+    """
+    Monotonic alignment path - vectorized via cumsum (much faster than Python loops).
+    w_ceil:    [B, 1, T_x] — integer duration per phone
+    attn_mask: [B, 1, T_y, T_x] — joint mask
+    Returns attn: [B, 1, T_y, T_x]
+    """
+    B, _, T_x = w_ceil.shape
+    T_y = attn_mask.shape[2]
+
+    # Build duration matrix: for each phone column expand the duration
+    # cumulative sum of durations gives us the end frame index for each phone
+    dur = w_ceil[:, 0, :]                            # [B, T_x]
+    cum_dur = np.cumsum(dur, axis=1)                 # [B, T_x]  — end frame (1-indexed)
+    cum_dur_prev = np.pad(cum_dur[:, :-1], ((0,0),(1,0)))  # [B, T_x] — start frame
+
+    # Frame indices: [1, T_y, 1]
+    frame_idx = np.arange(T_y, dtype=np.float32)[None, :, None]   # [1, T_y, 1]
+    # For each phone, mark frames [start, end)
+    # cum_dur_prev: [B,1,T_x], cum_dur: [B,1,T_x]
+    start = cum_dur_prev[:, None, :]                 # [B, 1, T_x]
+    end   = cum_dur[:, None, :]                      # [B, 1, T_x]
+    attn  = ((frame_idx >= start) & (frame_idx < end)).astype(np.float32)  # [B, T_y, T_x]
+    attn  = attn[:, None, :, :]                      # [B, 1, T_y, T_x]
+    return attn * attn_mask
+
+
+class OnnxTinyTTS:
+    """
+    Inference using ONNX Runtime.
+
+    Args:
+        onnx_dir: directory containing the 4 .onnx files
+        use_gpu:  if True, try CUDAExecutionProvider
+    """
+
+    def __init__(self, onnx_dir: str = "onnx", use_gpu: bool = False):
+        onnx_dir = os.path.abspath(onnx_dir)
+        print(f"Loading ONNX sessions from: {onnx_dir}")
+
+        self._enc  = _build_session(os.path.join(onnx_dir, "text_encoder.onnx"), use_gpu)
+        self._dp   = _build_session(os.path.join(onnx_dir, "duration_predictor.onnx"), use_gpu)
+        self._flow = _build_session(os.path.join(onnx_dir, "flow.onnx"), use_gpu)
+        self._dec  = _build_session(os.path.join(onnx_dir, "decoder.onnx"), use_gpu)
+
+        print("ONNX sessions ready ✅")
+
+    def _text_to_ids(self, text: str):
+        normalized = normalize_text(text)
+        phones, tones, _ = grapheme_to_phoneme(normalized)
+        phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+        if ADD_BLANK:
+            phone_ids = commons.insert_blanks(phone_ids, 0)
+            tone_ids  = commons.insert_blanks(tone_ids, 0)
+            lang_ids  = commons.insert_blanks(lang_ids, 0)
+
+        return phone_ids, tone_ids, lang_ids
+
+    def speak(
+        self,
+        text: str,
+        output_path: str = "onnx_output.wav",
+        speaker: str = "female",
+        noise_scale: float = 0.667,
+        noise_scale_w: float = 0.8,
+        length_scale: float = 1.0,
+        output_sr: int = None,
+    ) -> np.ndarray:
+        """Synthesize speech and save to output_path.
+        
+        Args:
+            output_sr: If set (e.g. 22050), resample the output from 44100 Hz.
+                       Useful to reduce file size while keeping quality.
+        """
+        print(f"[ONNX] Synthesizing: {text}")
+
+        phone_ids, tone_ids, lang_ids = self._text_to_ids(text)
+        T = len(phone_ids)
+
+        # Prepare inputs as float32 / int64 arrays
+        x       = np.array(phone_ids, dtype=np.int64)[None, :]        # [1, T]
+        x_len   = np.array([T], dtype=np.int64)                        # [1]
+        tone    = np.array(tone_ids, dtype=np.int64)[None, :]          # [1, T]
+        lang    = np.array(lang_ids, dtype=np.int64)[None, :]          # [1, T]
+        bert    = np.zeros((1, 1024, T), dtype=np.float32)
+        ja_bert = np.zeros((1, 768,  T), dtype=np.float32)
+        sid_val = SPK2ID.get(speaker, 0)
+        sid     = np.array([sid_val], dtype=np.int64)                  # [1]
+
+        # ── 1. Text Encoder ──────────────────────────────────────────────
+        x_enc, m_p, logs_p, x_mask, g = self._enc.run(
+            None,
+            {
+                "phone_ids":    x,
+                "phone_lengths":x_len,
+                "tone_ids":     tone,
+                "language_ids": lang,
+                "bert":         bert,
+                "ja_bert":      ja_bert,
+                "speaker_id":   sid,
+            },
+        )
+
+        # ── 2. Duration Predictor ─────────────────────────────────────────
+        logw = self._dp.run(None, {"x": x_enc, "x_mask": x_mask, "g": g})[0]
+
+        # ── 3. Alignment Path (NumPy) ─────────────────────────────────────
+        w        = np.exp(logw) * x_mask * length_scale          # [1, 1, T]
+        w_ceil   = np.ceil(w)                                    # [1, 1, T]
+        y_len    = max(1, int(w_ceil.sum()))
+        y_lens   = np.array([y_len], dtype=np.int64)
+
+        y_mask   = _create_length_mask_np(y_lens, y_len)         # [1, T_y]
+        y_mask   = y_mask[:, None, :]                            # [1, 1, T_y]
+        # attn_mask: [1, 1, T_y, T_x]  (outer product of frame mask and phone mask)
+        attn_mask = y_mask[:, :, :, None] * x_mask[:, :, None, :]   # [1,1,T_y,T_x]
+        attn     = _compute_alignment_path_np(w_ceil, attn_mask) # [1, 1, T_y, T_x]
+
+        # Expand prior stats via alignment
+        m_p_exp    = np.matmul(attn[:, 0], m_p.transpose(0, 2, 1)).transpose(0, 2, 1)
+        logs_p_exp = np.matmul(attn[:, 0], logs_p.transpose(0, 2, 1)).transpose(0, 2, 1)
+
+        # ── 4. Sample z_p ─────────────────────────────────────────────────
+        z_p = m_p_exp + np.random.randn(*m_p_exp.shape).astype(np.float32) * \
+              np.exp(logs_p_exp) * noise_scale
+
+        # ── 5. Flow (reverse) ─────────────────────────────────────────────
+        z = self._flow.run(
+            None,
+            {"z_p": z_p, "y_mask": y_mask.astype(np.float32), "g": g},
+        )[0]
+
+        # ── 6. Decoder ────────────────────────────────────────────────────
+        z_masked = (z * y_mask).astype(np.float32)
+        audio = self._dec.run(None, {"z": z_masked, "g": g})[0]  # [1, 1, samples]
+
+        audio_np = audio[0, 0]
+        save_sr = SAMPLING_RATE
+        if output_sr is not None and output_sr != SAMPLING_RATE:
+            try:
+                import torchaudio
+                import torch
+                wav_t = torch.from_numpy(audio_np).unsqueeze(0)
+                resampler = torchaudio.transforms.Resample(SAMPLING_RATE, output_sr)
+                audio_np = resampler(wav_t).squeeze(0).numpy()
+                save_sr = output_sr
+            except Exception as e:
+                print(f"[ONNX] Resampling failed ({e}), saving at {SAMPLING_RATE}Hz")
+
+        sf.write(output_path, audio_np, save_sr)
+        print(f"[ONNX] Saved: {output_path} ({save_sr}Hz)")
+        return audio_np

tiny_tts/integrations/__init__.py

@@ -0,0 +1,6 @@
+"""
+TinyTTS Integrations Package
+"""
+from .langchain import TinyTTSTool, TinyTTSVoiceTool, create_tts_chain
+
+__all__ = ["TinyTTSTool", "TinyTTSVoiceTool", "create_tts_chain"]