initial commit

.gitignore

@@ -0,0 +1,22 @@
+pretrained_models/Kokoro/*.pth
+pretrained_models/Kokoro/voices/*.pt
+pretrained_models/Kokoro/config.json
+output.wav
+.env
+# Ommit the DS_Store folder automatically created by macOS
+.DS_Store/
+
+# python virtual environment folder
+.venv/               
+.vscode/
+
+# process log file
+process_log.txt*
+process_log.txt
+
+# Python cache
+__pycache__/
+/*/__pycache__/
+/*/*/__pycache__/
+
+

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -8,7 +8,164 @@ sdk_version: 5.9.1
 app_file: app.py
 pinned: false
 license: apache-2.0
-short_description: Simple Space for the Kokoro Model
+short_description: Simple Space for the comparing the espeak-ng and openphonemizer with the Kokoro Model
+base_model:
+- yl4579/StyleTTS2-LJSpeech
+- hexgrad/Kokoro-82M 
+- openphonemizer/ckpt
+pipeline_tag: text-to-speech
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+❤️ Kokoro Discord Server: <https://discord.gg/QuGxSWBfQy>
+
+<audio controls><source src="https://huggingface.co/hexgrad/Kokoro-82M/resolve/main/demo/HEARME.wav" type="audio/wav"></audio>
+
+**Kokoro** is a frontier TTS model for its size of **82 million parameters** (text in/audio out).
+
+On 25 Dec 2024, Kokoro v0.19 weights were permissively released in full fp32 precision along with 2 voicepacks (Bella and Sarah), all under an Apache 2.0 license.
+
+As of 28 Dec 2024, **8 unique Voicepacks have been released**: 2F 2M each for American and British English.
+
+At the time of release, Kokoro v0.19 was the #1🥇 ranked model in [TTS Spaces Arena](https://huggingface.co/spaces/Pendrokar/TTS-Spaces-Arena). Kokoro had achieved higher Elo in this single-voice Arena setting over other models, using fewer parameters and less data:
+
+1. **Kokoro v0.19: 82M params, Apache, trained on <100 hours of audio, for <20 epochs**
+2. XTTS v2: 467M, CPML, >10k hours
+3. Edge TTS: Microsoft, proprietary
+4. MetaVoice: 1.2B, Apache, 100k hours
+5. Parler Mini: 880M, Apache, 45k hours
+6. Fish Speech: ~500M, CC-BY-NC-SA, 1M hours
+
+Kokoro's ability to top this Elo ladder suggests that the scaling law (Elo vs compute/data/params) for traditional TTS models might have a steeper slope than previously expected.
+
+You can find a hosted demo at [hf.co/spaces/hexgrad/Kokoro-TTS](https://huggingface.co/spaces/hexgrad/Kokoro-TTS).
+
+### Usage
+
+The following can be run in a single cell on [Google Colab](https://colab.research.google.com/).
+
+```py
+# 1️⃣ Install dependencies silently
+!git clone https://huggingface.co/hexgrad/Kokoro-82M
+%cd Kokoro-82M
+!apt-get -qq -y install espeak-ng > /dev/null 2>&1
+!pip install -q phonemizer torch transformers scipy munch
+
+# 2️⃣ Build the model and load the default voicepack
+from models import build_model
+import torch
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+MODEL = build_model('kokoro-v0_19.pth', device)
+VOICE_NAME = [
+    'af', # Default voice is a 50-50 mix of af_bella & af_sarah
+    'af_bella', 'af_sarah', 'am_adam', 'am_michael',
+    'bf_emma', 'bf_isabella', 'bm_george', 'bm_lewis',
+][0]
+VOICEPACK = torch.load(f'voices/{VOICE_NAME}.pt', weights_only=True).to(device)
+print(f'Loaded voice: {VOICE_NAME}')
+
+# 3️⃣ Call generate, which returns a 24khz audio waveform and a string of output phonemes
+from kokoro import generate
+text = "How could I know? It's an unanswerable question. Like asking an unborn child if they'll lead a good life. They haven't even been born."
+audio, out_ps = generate(MODEL, text, VOICEPACK, lang=VOICE_NAME[0])
+# Language is determined by the first letter of the VOICE_NAME:
+# 🇺🇸 'a' => American English => en-us
+# 🇬🇧 'b' => British English => en-gb
+
+# 4️⃣ Display the 24khz audio and print the output phonemes
+from IPython.display import display, Audio
+display(Audio(data=audio, rate=24000, autoplay=True))
+print(out_ps)
+```
+
+The inference code was quickly hacked together on Christmas Day. It is not clean code and leaves a lot of room for improvement. If you'd like to contribute, feel free to open a PR.
+
+### Model Facts
+
+No affiliation can be assumed between parties on different lines.
+
+**Architecture:**
+
+- StyleTTS 2: <https://arxiv.org/abs/2306.07691>
+- ISTFTNet: <https://arxiv.org/abs/2203.02395>
+- Decoder only: no diffusion, no encoder release
+
+**Architected by:** Li et al @ <https://github.com/yl4579/StyleTTS2>
+
+**Trained by**: `@rzvzn` on Discord
+
+**Supported Languages:** American English, British English
+
+**Model SHA256 Hash:** `3b0c392f87508da38fad3a2f9d94c359f1b657ebd2ef79f9d56d69503e470b0a`
+
+### Releases
+
+- 25 Dec 2024: Model v0.19, `af_bella`, `af_sarah`
+- 26 Dec 2024: `am_adam`, `am_michael`
+- 28 Dec 2024: `bf_emma`, `bf_isabella`, `bm_george`, `bm_lewis`
+
+### Licenses
+
+- Apache 2.0 weights in this repository
+- MIT inference code in [spaces/hexgrad/Kokoro-TTS](https://huggingface.co/spaces/hexgrad/Kokoro-TTS) adapted from [yl4579/StyleTTS2](https://github.com/yl4579/StyleTTS2)
+- GPLv3 dependency in [espeak-ng](https://github.com/espeak-ng/espeak-ng)
+
+The inference code was originally MIT licensed by the paper author. Note that this card applies only to this model, Kokoro. Original models published by the paper author can be found at [hf.co/yl4579](https://huggingface.co/yl4579).
+
+### Evaluation
+
+**Metric:** Elo rating
+
+**Leaderboard:** [hf.co/spaces/Pendrokar/TTS-Spaces-Arena](https://huggingface.co/spaces/Pendrokar/TTS-Spaces-Arena)
+
+![TTS-Spaces-Arena-25-Dec-2024](demo/TTS-Spaces-Arena-25-Dec-2024.png)
+
+The voice ranked in the Arena is a 50-50 mix of Bella and Sarah. For your convenience, this mix is included in this repository as `af.pt`, but you can trivially reproduce it like this:
+
+```py
+import torch
+bella = torch.load('voices/af_bella.pt', weights_only=True)
+sarah = torch.load('voices/af_sarah.pt', weights_only=True)
+af = torch.mean(torch.stack([bella, sarah]), dim=0)
+assert torch.equal(af, torch.load('voices/af.pt', weights_only=True))
+```
+
+### Training Details
+
+**Compute:** Kokoro was trained on A100 80GB vRAM instances rented from [Vast.ai](https://cloud.vast.ai/?ref_id=79907) (referral link). Vast was chosen over other compute providers due to its competitive on-demand hourly rates. The average hourly cost for the A100 80GB vRAM instances used for training was below $1/hr per GPU, which was around half the quoted rates from other providers at the time.
+
+**Data:** Kokoro was trained exclusively on **permissive/non-copyrighted audio data** and IPA phoneme labels. Examples of permissive/non-copyrighted audio include:
+
+- Public domain audio
+- Audio licensed under Apache, MIT, etc
+- Synthetic audio<sup>[1]</sup> generated by closed<sup>[2]</sup> TTS models from large providers<br/>
+[1] <https://copyright.gov/ai/ai_policy_guidance.pdf><br/>
+[2] No synthetic audio from open TTS models or "custom voice clones"
+
+**Epochs:** Less than **20 epochs**
+
+**Total Dataset Size:** Less than **100 hours** of audio
+
+### Limitations
+
+Kokoro v0.19 is limited in some specific ways, due to its training set and/or architecture:
+
+- [Data] Lacks voice cloning capability, likely due to small <100h training set
+- [Arch] Relies on external g2p (espeak-ng), which introduces a class of g2p failure modes
+- [Data] Training dataset is mostly long-form reading and narration, not conversation
+- [Arch] At 82M params, Kokoro almost certainly falls to a well-trained 1B+ param diffusion transformer, or a many-billion-param MLLM like GPT-4o / Gemini 2.0 Flash
+- [Data] Multilingual capability is architecturally feasible, but training data is mostly English
+
+Refer to the [Philosophy discussion](https://huggingface.co/hexgrad/Kokoro-82M/discussions/5) to better understand these limitations.
+
+**Will the other voicepacks be released?** There is currently no release date scheduled for the other voicepacks, but in the meantime you can try them in the hosted demo at [hf.co/spaces/hexgrad/Kokoro-TTS](https://huggingface.co/spaces/hexgrad/Kokoro-TTS).
+
+### Acknowledgements
+
+- [@yl4579](https://huggingface.co/yl4579) for architecting StyleTTS 2
+- [@Pendrokar](https://huggingface.co/Pendrokar) for adding Kokoro as a contender in the TTS Spaces Arena
+
+### Model Card Contact
+
+`@rzvzn` on Discord. Server invite: <https://discord.gg/QuGxSWBfQy>
+
+<img src="https://static0.gamerantimages.com/wordpress/wp-content/uploads/2024/08/terminator-zero-41-1.jpg" width="400" alt="kokoro" />

app.py

@@ -0,0 +1,166 @@
+import gradio as gr
+import os
+import torch
+
+# Import eSpeak TTS pipeline
+from tts_cli import (
+    build_model as build_model_espeak,
+    generate_long_form_tts as generate_long_form_tts_espeak,
+)
+
+# Import OpenPhonemizer TTS pipeline
+from tts_cli_op import (
+    build_model as build_model_open,
+    generate_long_form_tts as generate_long_form_tts_open,
+)
+from pretrained_models import Kokoro
+
+# ---------------------------------------------------------------------
+# Path to models and voicepacks
+# ---------------------------------------------------------------------
+MODELS_DIR = "pretrained_models/Kokoro"
+VOICES_DIR = "pretrained_models/Kokoro/voices"
+
+
+# ---------------------------------------------------------------------
+# List the models (.pth) and voices (.pt)
+# ---------------------------------------------------------------------
+def get_models():
+    return sorted([f for f in os.listdir(MODELS_DIR) if f.endswith(".pth")])
+
+
+def get_voices():
+    return sorted([f for f in os.listdir(VOICES_DIR) if f.endswith(".pt")])
+
+
+# ---------------------------------------------------------------------
+# We'll map engine selection -> (build_model_func, generate_func)
+# ---------------------------------------------------------------------
+ENGINES = {
+    "espeak": (build_model_espeak, generate_long_form_tts_espeak),
+    "openphonemizer": (build_model_open, generate_long_form_tts_open),
+}
+
+
+# ---------------------------------------------------------------------
+# The main inference function called by Gradio
+# ---------------------------------------------------------------------
+def tts_inference(text, engine, model_file, voice_file, speed=1.0):
+    """
+    text:        Input string
+    engine:      "espeak" or "openphonemizer"
+    model_file:  Selected .pth from the models folder
+    voice_file:  Selected .pt from the voices folder
+    speed:       Speech speed
+    """
+    # 1) Map engine to the correct build_model + generate_long_form_tts
+    build_fn, gen_fn = ENGINES[engine]
+
+    # 2) Prepare paths
+    model_path = os.path.join(MODELS_DIR, model_file)
+    voice_path = os.path.join(VOICES_DIR, voice_file)
+
+    # 3) Decide device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # 4) Load model
+    model = build_fn(model_path, device=device)
+    # Set submodules eval
+    for k, subm in model.items():
+        if hasattr(subm, "eval"):
+            subm.eval()
+
+    # 5) Load voicepack
+    voicepack = torch.load(voice_path, map_location=device)
+    if hasattr(voicepack, "eval"):
+        voicepack.eval()
+
+    # 6) Generate TTS
+    audio, phonemes = gen_fn(model, text, voicepack, speed=speed)
+    sr = 22050  # or your actual sample rate
+
+    return (sr, audio)  # Gradio expects (sample_rate, np_array)
+
+
+# ---------------------------------------------------------------------
+# Build Gradio App
+# ---------------------------------------------------------------------
+def create_gradio_app():
+    model_list = get_models()
+    voice_list = get_voices()
+
+    css = """
+    h4 {
+        text-align: center;
+        display:block;
+    }
+    h2 {
+        text-align: center;
+        display:block;
+    }
+    """
+    with gr.Blocks(theme=gr.themes.Ocean(), css=css) as demo:
+        gr.Markdown("## Kokoro TTS Demo: Choose engine, model, and voice")
+
+        # Row 1: Text input
+        text_input = gr.Textbox(
+            label="Input Text",
+            value="Hello, world! Testing both eSpeak and OpenPhonemizer. Can you believe that we live in 2024 and have access to advanced AI?",
+            lines=3,
+        )
+
+        # Row 2: Engine selection
+        engine_dropdown = gr.Dropdown(
+            choices=["espeak", "openphonemizer"],
+            value="openphonemizer",
+            label="Phonemizer",
+        )
+
+        # Row 3: Model dropdown
+        model_dropdown = gr.Dropdown(
+            choices=model_list,
+            value=model_list[0] if model_list else None,
+            label="Model (.pth)",
+        )
+
+        # Row 4: Voice dropdown
+        voice_dropdown = gr.Dropdown(
+            choices=voice_list,
+            value=voice_list[0] if voice_list else None,
+            label="Voice (.pt)",
+        )
+
+        # Row 5: Speed slider
+        speed_slider = gr.Slider(
+            minimum=0.5, maximum=2.0, value=1.0, step=0.1, label="Speech Speed"
+        )
+
+        # Generate button + audio output
+        generate_btn = gr.Button("Generate")
+        tts_output = gr.Audio(label="TTS Output")
+
+        # Connect the button to our inference function
+        generate_btn.click(
+            fn=tts_inference,
+            inputs=[
+                text_input,
+                engine_dropdown,
+                model_dropdown,
+                voice_dropdown,
+                speed_slider,
+            ],
+            outputs=tts_output,
+        )
+
+        gr.Markdown(
+            "#### Kokoro TTS Demo based on [Kokoro-82M](https://huggingface.co/hexgrad/Kokoro-82M)"
+        )
+    return demo
+
+
+# ---------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------
+if __name__ == "__main__":
+    app = create_gradio_app()
+    app.launch()

espeak_util.py

@@ -0,0 +1,206 @@
+import platform
+import subprocess
+import shutil
+from pathlib import Path
+import os
+from typing import Optional, Tuple
+from phonemizer.backend.espeak.wrapper import EspeakWrapper
+
+
+class EspeakConfig:
+    """Utility class for configuring espeak-ng library and binary."""
+
+    @staticmethod
+    def find_espeak_binary() -> tuple[bool, Optional[str]]:
+        """
+        Find espeak-ng binary using multiple methods.
+
+        Returns:
+            tuple: (bool indicating if espeak is available, path to espeak binary if found)
+        """
+        # Common binary names
+        binary_names = ["espeak-ng", "espeak"]
+        if platform.system() == "Windows":
+            binary_names = ["espeak-ng.exe", "espeak.exe"]
+
+        # Common installation directories for Linux
+        linux_paths = [
+            "/usr/bin",
+            "/usr/local/bin",
+            "/usr/lib/espeak-ng",
+            "/usr/local/lib/espeak-ng",
+            "/opt/espeak-ng/bin",
+        ]
+
+        # First check if it's in PATH
+        for name in binary_names:
+            espeak_path = shutil.which(name)
+            if espeak_path:
+                return True, espeak_path
+
+        # For Linux, check common installation directories
+        if platform.system() == "Linux":
+            for directory in linux_paths:
+                for name in binary_names:
+                    path = Path(directory) / name
+                    if path.exists():
+                        return True, str(path)
+
+        # Try running the command directly as a last resort
+        try:
+            subprocess.run(
+                ["espeak-ng", "--version"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                check=True,
+            )
+            return True, "espeak-ng"
+        except (subprocess.SubprocessError, FileNotFoundError):
+            pass
+
+        return False, None
+
+    @staticmethod
+    def find_library_path() -> Optional[str]:
+        """
+        Find the espeak-ng library using multiple search methods.
+
+        Returns:
+            Optional[str]: Path to the library if found, None otherwise
+        """
+        system = platform.system()
+
+        if system == "Linux":
+            lib_names = ["libespeak-ng.so", "libespeak-ng.so.1"]
+            common_paths = [
+                # Debian/Ubuntu paths
+                "/usr/lib/x86_64-linux-gnu",
+                "/usr/lib/aarch64-linux-gnu",  # For ARM64
+                "/usr/lib/arm-linux-gnueabihf",  # For ARM32
+                "/usr/lib",
+                "/usr/local/lib",
+                # Fedora/RHEL paths
+                "/usr/lib64",
+                "/usr/lib32",
+                # Common additional paths
+                "/usr/lib/espeak-ng",
+                "/usr/local/lib/espeak-ng",
+                "/opt/espeak-ng/lib",
+            ]
+
+            # Check common locations first
+            for path in common_paths:
+                for lib_name in lib_names:
+                    lib_path = Path(path) / lib_name
+                    if lib_path.exists():
+                        return str(lib_path)
+
+            # Search system library paths
+            try:
+                # Use ldconfig to find the library
+                result = subprocess.run(
+                    ["ldconfig", "-p"], capture_output=True, text=True, check=True
+                )
+                for line in result.stdout.splitlines():
+                    if "libespeak-ng.so" in line:
+                        # Extract path from ldconfig output
+                        return line.split("=>")[-1].strip()
+            except (subprocess.SubprocessError, FileNotFoundError):
+                pass
+
+        elif system == "Darwin":  # macOS
+            common_paths = [
+                Path("/opt/homebrew/lib/libespeak-ng.dylib"),
+                Path("/usr/local/lib/libespeak-ng.dylib"),
+                *list(
+                    Path("/opt/homebrew/Cellar/espeak-ng").glob(
+                        "*/lib/libespeak-ng.dylib"
+                    )
+                ),
+                *list(
+                    Path("/usr/local/Cellar/espeak-ng").glob("*/lib/libespeak-ng.dylib")
+                ),
+            ]
+
+            for path in common_paths:
+                if path.exists():
+                    return str(path)
+
+        elif system == "Windows":
+            common_paths = [
+                Path(os.environ.get("PROGRAMFILES", "C:\\Program Files"))
+                / "eSpeak NG"
+                / "libespeak-ng.dll",
+                Path(os.environ.get("PROGRAMFILES(X86)", "C:\\Program Files (x86)"))
+                / "eSpeak NG"
+                / "libespeak-ng.dll",
+                *[
+                    Path(p) / "libespeak-ng.dll"
+                    for p in os.environ.get("PATH", "").split(os.pathsep)
+                ],
+            ]
+
+            for path in common_paths:
+                if path.exists():
+                    return str(path)
+
+        return None
+
+    @classmethod
+    def configure_espeak(cls) -> Tuple[bool, str]:
+        """
+        Configure espeak-ng for use with the phonemizer.
+
+        Returns:
+            Tuple[bool, str]: (Success status, Status message)
+        """
+        # First check if espeak binary is available
+        espeak_available, espeak_path = cls.find_espeak_binary()
+        if not espeak_available:
+            raise FileNotFoundError(
+                "Could not find espeak-ng binary. Please install espeak-ng:\n"
+                "Ubuntu/Debian: sudo apt-get install espeak-ng espeak-ng-data\n"
+                "Fedora: sudo dnf install espeak-ng\n"
+                "Arch: sudo pacman -S espeak-ng\n"
+                "MacOS: brew install espeak-ng\n"
+                "Windows: Download from https://github.com/espeak-ng/espeak-ng/releases"
+            )
+
+        # Find the library
+        library_path = cls.find_library_path()
+        if not library_path:
+            # On Linux, we might not need to explicitly set the library path
+            if platform.system() == "Linux":
+                return True, f"Using system espeak-ng installation at: {espeak_path}"
+            else:
+                raise FileNotFoundError(
+                    "Could not find espeak-ng library. Please ensure espeak-ng is properly installed."
+                )
+
+        # Try to set the library path
+        try:
+            EspeakWrapper.set_library(library_path)
+            return True, f"Successfully configured espeak-ng library at: {library_path}"
+        except Exception as e:
+            if platform.system() == "Linux":
+                # On Linux, try to continue without explicit library path
+                return True, f"Using system espeak-ng installation at: {espeak_path}"
+            else:
+                raise RuntimeError(f"Failed to configure espeak-ng library: {str(e)}")
+
+
+def setup_espeak():
+    """
+    Set up espeak-ng for use with the phonemizer.
+    Raises appropriate exceptions if setup fails.
+    """
+    try:
+        success, message = EspeakConfig.configure_espeak()
+        print(message)
+    except Exception as e:
+        print(f"Error configuring espeak-ng: {str(e)}")
+        raise
+
+
+# Replace the original set_espeak_library function with this
+set_espeak_library = setup_espeak

istftnet.py

@@ -0,0 +1,523 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/Modules/istftnet.py
+from scipy.signal import get_window
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Modules/utils.py
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+LRELU_SLOPE = 0.1
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm1d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class AdaINResBlock1(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), style_dim=64):
+        super(AdaINResBlock1, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+        
+        self.adain1 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        
+        self.adain2 = nn.ModuleList([
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+            AdaIN1d(style_dim, channels),
+        ])
+        
+        self.alpha1 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs1))])
+        self.alpha2 = nn.ParameterList([nn.Parameter(torch.ones(1, channels, 1)) for i in range(len(self.convs2))])
+
+
+    def forward(self, x, s):
+        for c1, c2, n1, n2, a1, a2 in zip(self.convs1, self.convs2, self.adain1, self.adain2, self.alpha1, self.alpha2):
+            xt = n1(x, s)
+            xt = xt + (1 / a1) * (torch.sin(a1 * xt) ** 2)  # Snake1D
+            xt = c1(xt)
+            xt = n2(xt, s)
+            xt = xt + (1 / a2) * (torch.sin(a2 * xt) ** 2)  # Snake1D
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+            
+class TorchSTFT(torch.nn.Module):
+    def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = torch.from_numpy(get_window(window, win_length, fftbins=True).astype(np.float32))
+
+    def transform(self, input_data):
+        forward_transform = torch.stft(
+            input_data,
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(input_data.device),
+            return_complex=True)
+
+        return torch.abs(forward_transform), torch.angle(forward_transform)
+
+    def inverse(self, magnitude, phase):
+        inverse_transform = torch.istft(
+            magnitude * torch.exp(phase * 1j),
+            self.filter_length, self.hop_length, self.win_length, window=self.window.to(magnitude.device))
+
+        return inverse_transform.unsqueeze(-2)  # unsqueeze to stay consistent with conv_transpose1d implementation
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+    
+class SineGen(torch.nn.Module):
+    """ Definition of sine generator
+    SineGen(samp_rate, harmonic_num = 0,
+            sine_amp = 0.1, noise_std = 0.003,
+            voiced_threshold = 0,
+            flag_for_pulse=False)
+    samp_rate: sampling rate in Hz
+    harmonic_num: number of harmonic overtones (default 0)
+    sine_amp: amplitude of sine-wavefrom (default 0.1)
+    noise_std: std of Gaussian noise (default 0.003)
+    voiced_thoreshold: F0 threshold for U/V classification (default 0)
+    flag_for_pulse: this SinGen is used inside PulseGen (default False)
+    Note: when flag_for_pulse is True, the first time step of a voiced
+        segment is always sin(np.pi) or cos(0)
+    """
+
+    def __init__(self, samp_rate, upsample_scale, harmonic_num=0,
+                 sine_amp=0.1, noise_std=0.003,
+                 voiced_threshold=0,
+                 flag_for_pulse=False):
+        super(SineGen, self).__init__()
+        self.sine_amp = sine_amp
+        self.noise_std = noise_std
+        self.harmonic_num = harmonic_num
+        self.dim = self.harmonic_num + 1
+        self.sampling_rate = samp_rate
+        self.voiced_threshold = voiced_threshold
+        self.flag_for_pulse = flag_for_pulse
+        self.upsample_scale = upsample_scale
+
+    def _f02uv(self, f0):
+        # generate uv signal
+        uv = (f0 > self.voiced_threshold).type(torch.float32)
+        return uv
+
+    def _f02sine(self, f0_values):
+        """ f0_values: (batchsize, length, dim)
+            where dim indicates fundamental tone and overtones
+        """
+        # convert to F0 in rad. The interger part n can be ignored
+        # because 2 * np.pi * n doesn't affect phase
+        rad_values = (f0_values / self.sampling_rate) % 1
+
+        # initial phase noise (no noise for fundamental component)
+        rand_ini = torch.rand(f0_values.shape[0], f0_values.shape[2], \
+                              device=f0_values.device)
+        rand_ini[:, 0] = 0
+        rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
+
+        # instantanouse phase sine[t] = sin(2*pi \sum_i=1 ^{t} rad)
+        if not self.flag_for_pulse:
+#             # for normal case
+
+#             # To prevent torch.cumsum numerical overflow,
+#             # it is necessary to add -1 whenever \sum_k=1^n rad_value_k > 1.
+#             # Buffer tmp_over_one_idx indicates the time step to add -1.
+#             # This will not change F0 of sine because (x-1) * 2*pi = x * 2*pi
+#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
+#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
+#             cumsum_shift = torch.zeros_like(rad_values)
+#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+
+#             phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
+            rad_values = torch.nn.functional.interpolate(rad_values.transpose(1, 2), 
+                                                         scale_factor=1/self.upsample_scale, 
+                                                         mode="linear").transpose(1, 2)
+    
+#             tmp_over_one = torch.cumsum(rad_values, 1) % 1
+#             tmp_over_one_idx = (padDiff(tmp_over_one)) < 0
+#             cumsum_shift = torch.zeros_like(rad_values)
+#             cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
+    
+            phase = torch.cumsum(rad_values, dim=1) * 2 * np.pi
+            phase = torch.nn.functional.interpolate(phase.transpose(1, 2) * self.upsample_scale, 
+                                                    scale_factor=self.upsample_scale, mode="linear").transpose(1, 2)
+            sines = torch.sin(phase)
+            
+        else:
+            # If necessary, make sure that the first time step of every
+            # voiced segments is sin(pi) or cos(0)
+            # This is used for pulse-train generation
+
+            # identify the last time step in unvoiced segments
+            uv = self._f02uv(f0_values)
+            uv_1 = torch.roll(uv, shifts=-1, dims=1)
+            uv_1[:, -1, :] = 1
+            u_loc = (uv < 1) * (uv_1 > 0)
+
+            # get the instantanouse phase
+            tmp_cumsum = torch.cumsum(rad_values, dim=1)
+            # different batch needs to be processed differently
+            for idx in range(f0_values.shape[0]):
+                temp_sum = tmp_cumsum[idx, u_loc[idx, :, 0], :]
+                temp_sum[1:, :] = temp_sum[1:, :] - temp_sum[0:-1, :]
+                # stores the accumulation of i.phase within
+                # each voiced segments
+                tmp_cumsum[idx, :, :] = 0
+                tmp_cumsum[idx, u_loc[idx, :, 0], :] = temp_sum
+
+            # rad_values - tmp_cumsum: remove the accumulation of i.phase
+            # within the previous voiced segment.
+            i_phase = torch.cumsum(rad_values - tmp_cumsum, dim=1)
+
+            # get the sines
+            sines = torch.cos(i_phase * 2 * np.pi)
+        return sines
+
+    def forward(self, f0):
+        """ sine_tensor, uv = forward(f0)
+        input F0: tensor(batchsize=1, length, dim=1)
+                  f0 for unvoiced steps should be 0
+        output sine_tensor: tensor(batchsize=1, length, dim)
+        output uv: tensor(batchsize=1, length, 1)
+        """
+        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
+                             device=f0.device)
+        # fundamental component
+        fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
+
+        # generate sine waveforms
+        sine_waves = self._f02sine(fn) * self.sine_amp
+
+        # generate uv signal
+        # uv = torch.ones(f0.shape)
+        # uv = uv * (f0 > self.voiced_threshold)
+        uv = self._f02uv(f0)
+
+        # noise: for unvoiced should be similar to sine_amp
+        #        std = self.sine_amp/3 -> max value ~ self.sine_amp
+        # .       for voiced regions is self.noise_std
+        noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
+        noise = noise_amp * torch.randn_like(sine_waves)
+
+        # first: set the unvoiced part to 0 by uv
+        # then: additive noise
+        sine_waves = sine_waves * uv + noise
+        return sine_waves, uv, noise
+
+
+class SourceModuleHnNSF(torch.nn.Module):
+    """ SourceModule for hn-nsf
+    SourceModule(sampling_rate, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0)
+    sampling_rate: sampling_rate in Hz
+    harmonic_num: number of harmonic above F0 (default: 0)
+    sine_amp: amplitude of sine source signal (default: 0.1)
+    add_noise_std: std of additive Gaussian noise (default: 0.003)
+        note that amplitude of noise in unvoiced is decided
+        by sine_amp
+    voiced_threshold: threhold to set U/V given F0 (default: 0)
+    Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+    F0_sampled (batchsize, length, 1)
+    Sine_source (batchsize, length, 1)
+    noise_source (batchsize, length 1)
+    uv (batchsize, length, 1)
+    """
+
+    def __init__(self, sampling_rate, upsample_scale, harmonic_num=0, sine_amp=0.1,
+                 add_noise_std=0.003, voiced_threshod=0):
+        super(SourceModuleHnNSF, self).__init__()
+
+        self.sine_amp = sine_amp
+        self.noise_std = add_noise_std
+
+        # to produce sine waveforms
+        self.l_sin_gen = SineGen(sampling_rate, upsample_scale, harmonic_num,
+                                 sine_amp, add_noise_std, voiced_threshod)
+
+        # to merge source harmonics into a single excitation
+        self.l_linear = torch.nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = torch.nn.Tanh()
+
+    def forward(self, x):
+        """
+        Sine_source, noise_source = SourceModuleHnNSF(F0_sampled)
+        F0_sampled (batchsize, length, 1)
+        Sine_source (batchsize, length, 1)
+        noise_source (batchsize, length 1)
+        """
+        # source for harmonic branch
+        with torch.no_grad():
+            sine_wavs, uv, _ = self.l_sin_gen(x)
+        sine_merge = self.l_tanh(self.l_linear(sine_wavs))
+
+        # source for noise branch, in the same shape as uv
+        noise = torch.randn_like(uv) * self.sine_amp / 3
+        return sine_merge, noise, uv
+def padDiff(x):
+    return F.pad(F.pad(x, (0,0,-1,1), 'constant', 0) - x, (0,0,0,-1), 'constant', 0)
+
+    
+class Generator(torch.nn.Module):
+    def __init__(self, style_dim, resblock_kernel_sizes, upsample_rates, upsample_initial_channel, resblock_dilation_sizes, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size):
+        super(Generator, self).__init__()
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        resblock = AdaINResBlock1
+
+        self.m_source = SourceModuleHnNSF(
+                    sampling_rate=24000,
+                    upsample_scale=np.prod(upsample_rates) * gen_istft_hop_size,
+                    harmonic_num=8, voiced_threshod=10)
+        self.f0_upsamp = torch.nn.Upsample(scale_factor=np.prod(upsample_rates) * gen_istft_hop_size)
+        self.noise_convs = nn.ModuleList()
+        self.noise_res = nn.ModuleList()
+        
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(weight_norm(
+                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
+                                k, u, padding=(k-u)//2)))
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel//(2**(i+1))
+            for j, (k, d) in enumerate(zip(resblock_kernel_sizes,resblock_dilation_sizes)):
+                self.resblocks.append(resblock(ch, k, d, style_dim))
+                
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            
+            if i + 1 < len(upsample_rates):  #
+                stride_f0 = np.prod(upsample_rates[i + 1:])
+                self.noise_convs.append(Conv1d(
+                    gen_istft_n_fft + 2, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2))
+                self.noise_res.append(resblock(c_cur, 7, [1,3,5], style_dim))
+            else:
+                self.noise_convs.append(Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
+                self.noise_res.append(resblock(c_cur, 11, [1,3,5], style_dim))
+                
+                
+        self.post_n_fft = gen_istft_n_fft
+        self.conv_post = weight_norm(Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
+        self.stft = TorchSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+        
+        
+    def forward(self, x, s, f0):
+        with torch.no_grad():
+            f0 = self.f0_upsamp(f0[:, None]).transpose(1, 2)  # bs,n,t
+
+            har_source, noi_source, uv = self.m_source(f0)
+            har_source = har_source.transpose(1, 2).squeeze(1)
+            har_spec, har_phase = self.stft.transform(har_source)
+            har = torch.cat([har_spec, har_phase], dim=1)
+        
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x_source = self.noise_convs[i](har)
+            x_source = self.noise_res[i](x_source, s)
+
+            x = self.ups[i](x)
+            if i == self.num_upsamples - 1:
+                x = self.reflection_pad(x)
+
+            x = x + x_source
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
+        return self.stft.inverse(spec, phase)
+    
+    def fw_phase(self, x, s):
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x, s)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x, s)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.reflection_pad(x)
+        x = self.conv_post(x)
+        spec = torch.exp(x[:,:self.post_n_fft // 2 + 1, :])
+        phase = torch.sin(x[:, self.post_n_fft // 2 + 1:, :])
+        return spec, phase
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+        
+class AdainResBlk1d(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=64, actv=nn.LeakyReLU(0.2),
+                 upsample='none', dropout_p=0.0):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        self.dropout = nn.Dropout(dropout_p)
+        
+        if upsample == 'none':
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(nn.ConvTranspose1d(dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1, output_padding=1))
+        
+        
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(self.dropout(x))
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(self.dropout(x))
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / np.sqrt(2)
+        return out
+    
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == 'none':
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode='nearest')
+
+class Decoder(nn.Module):
+    def __init__(self, dim_in=512, F0_channel=512, style_dim=64, dim_out=80, 
+                resblock_kernel_sizes = [3,7,11],
+                upsample_rates = [10, 6],
+                upsample_initial_channel=512,
+                resblock_dilation_sizes=[[1,3,5], [1,3,5], [1,3,5]],
+                upsample_kernel_sizes=[20, 12], 
+                gen_istft_n_fft=20, gen_istft_hop_size=5):
+        super().__init__()
+        
+        self.decode = nn.ModuleList()
+        
+        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+        
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 1024, style_dim))
+        self.decode.append(AdainResBlk1d(1024 + 2 + 64, 512, style_dim, upsample=True))
+
+        self.F0_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        
+        self.N_conv = weight_norm(nn.Conv1d(1, 1, kernel_size=3, stride=2, groups=1, padding=1))
+        
+        self.asr_res = nn.Sequential(
+            weight_norm(nn.Conv1d(512, 64, kernel_size=1)),
+        )
+        
+        
+        self.generator = Generator(style_dim, resblock_kernel_sizes, upsample_rates, 
+                                   upsample_initial_channel, resblock_dilation_sizes, 
+                                   upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size)
+        
+    def forward(self, asr, F0_curve, N, s):
+        F0 = self.F0_conv(F0_curve.unsqueeze(1))
+        N = self.N_conv(N.unsqueeze(1))
+        
+        x = torch.cat([asr, F0, N], axis=1)
+        x = self.encode(x, s)
+        
+        asr_res = self.asr_res(asr)
+        
+        res = True
+        for block in self.decode:
+            if res:
+                x = torch.cat([x, asr_res, F0, N], axis=1)
+            x = block(x, s)
+            if block.upsample_type != "none":
+                res = False
+                
+        x = self.generator(x, s, F0_curve)
+        return x

kokoro.py

@@ -0,0 +1,187 @@
+import phonemizer
+import re
+import torch
+from espeak_util import set_espeak_library
+
+set_espeak_library()
+
+
+def split_num(num):
+    num = num.group()
+    if "." in num:
+        return num
+    elif ":" in num:
+        h, m = [int(n) for n in num.split(":")]
+        if m == 0:
+            return f"{h} o'clock"
+        elif m < 10:
+            return f"{h} oh {m}"
+        return f"{h} {m}"
+    year = int(num[:4])
+    if year < 1100 or year % 1000 < 10:
+        return num
+    left, right = num[:2], int(num[2:4])
+    s = "s" if num.endswith("s") else ""
+    if 100 <= year % 1000 <= 999:
+        if right == 0:
+            return f"{left} hundred{s}"
+        elif right < 10:
+            return f"{left} oh {right}{s}"
+    return f"{left} {right}{s}"
+
+
+def flip_money(m):
+    m = m.group()
+    bill = "dollar" if m[0] == "$" else "pound"
+    if m[-1].isalpha():
+        return f"{m[1:]} {bill}s"
+    elif "." not in m:
+        s = "" if m[1:] == "1" else "s"
+        return f"{m[1:]} {bill}{s}"
+    b, c = m[1:].split(".")
+    s = "" if b == "1" else "s"
+    c = int(c.ljust(2, "0"))
+    coins = (
+        f"cent{'' if c == 1 else 's'}"
+        if m[0] == "$"
+        else ("penny" if c == 1 else "pence")
+    )
+    return f"{b} {bill}{s} and {c} {coins}"
+
+
+def point_num(num):
+    a, b = num.group().split(".")
+    return " point ".join([a, " ".join(b)])
+
+
+def normalize_text(text):
+    text = text.replace(chr(8216), "'").replace(chr(8217), "'")
+    text = text.replace("«", chr(8220)).replace("»", chr(8221))
+    text = text.replace(chr(8220), '"').replace(chr(8221), '"')
+    text = text.replace("(", "«").replace(")", "»")
+    for a, b in zip("、。！，：；？", ",.!,:;?"):
+        text = text.replace(a, b + " ")
+    text = re.sub(r"[^\S \n]", " ", text)
+    text = re.sub(r"  +", " ", text)
+    text = re.sub(r"(?<=\n) +(?=\n)", "", text)
+    text = re.sub(r"\bD[Rr]\.(?= [A-Z])", "Doctor", text)
+    text = re.sub(r"\b(?:Mr\.|MR\.(?= [A-Z]))", "Mister", text)
+    text = re.sub(r"\b(?:Ms\.|MS\.(?= [A-Z]))", "Miss", text)
+    text = re.sub(r"\b(?:Mrs\.|MRS\.(?= [A-Z]))", "Mrs", text)
+    text = re.sub(r"\betc\.(?! [A-Z])", "etc", text)
+    text = re.sub(r"(?i)\b(y)eah?\b", r"\1e'a", text)
+    text = re.sub(
+        r"\d*\.\d+|\b\d{4}s?\b|(?<!:)\b(?:[1-9]|1[0-2]):[0-5]\d\b(?!:)", split_num, text
+    )
+    text = re.sub(r"(?<=\d),(?=\d)", "", text)
+    text = re.sub(
+        r"(?i)[$£]\d+(?:\.\d+)?(?: hundred| thousand| (?:[bm]|tr)illion)*\b|[$£]\d+\.\d\d?\b",
+        flip_money,
+        text,
+    )
+    text = re.sub(r"\d*\.\d+", point_num, text)
+    text = re.sub(r"(?<=\d)-(?=\d)", " to ", text)
+    text = re.sub(r"(?<=\d)S", " S", text)
+    text = re.sub(r"(?<=[BCDFGHJ-NP-TV-Z])'?s\b", "'S", text)
+    text = re.sub(r"(?<=X')S\b", "s", text)
+    text = re.sub(
+        r"(?:[A-Za-z]\.){2,} [a-z]", lambda m: m.group().replace(".", "-"), text
+    )
+    text = re.sub(r"(?i)(?<=[A-Z])\.(?=[A-Z])", "-", text)
+    return text.strip()
+
+
+def get_vocab():
+    _pad = "$"
+    _punctuation = ';:,.!?¡¿—…"«»“” '
+    _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+    symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+    dicts = {}
+    for i in range(len((symbols))):
+        dicts[symbols[i]] = i
+    return dicts
+
+
+VOCAB = get_vocab()
+
+
+def tokenize(ps):
+    return [i for i in map(VOCAB.get, ps) if i is not None]
+
+
+phonemizers = dict(
+    a=phonemizer.backend.EspeakBackend(
+        language="en-us", preserve_punctuation=True, with_stress=True
+    ),
+    b=phonemizer.backend.EspeakBackend(
+        language="en-gb", preserve_punctuation=True, with_stress=True
+    ),
+)
+
+
+def phonemize(text, lang, norm=True):
+    if norm:
+        text = normalize_text(text)
+    ps = phonemizers[lang].phonemize([text])
+    ps = ps[0] if ps else ""
+    # https://en.wiktionary.org/wiki/kokoro#English
+    ps = ps.replace("kəkˈoːɹoʊ", "kˈoʊkəɹoʊ").replace("kəkˈɔːɹəʊ", "kˈəʊkəɹəʊ")
+    ps = ps.replace("ʲ", "j").replace("r", "ɹ").replace("x", "k").replace("ɬ", "l")
+    ps = re.sub(r"(?<=[a-zɹː])(?=hˈʌndɹɪd)", " ", ps)
+    ps = re.sub(r' z(?=[;:,.!?¡¿—…"«»“” ]|$)', "z", ps)
+    if lang == "a":
+        ps = re.sub(r"(?<=nˈaɪn)ti(?!ː)", "di", ps)
+    ps = "".join(filter(lambda p: p in VOCAB, ps))
+    return ps.strip()
+
+
+def length_to_mask(lengths):
+    mask = (
+        torch.arange(lengths.max())
+        .unsqueeze(0)
+        .expand(lengths.shape[0], -1)
+        .type_as(lengths)
+    )
+    mask = torch.gt(mask + 1, lengths.unsqueeze(1))
+    return mask
+
+
+@torch.no_grad()
+def forward(model, tokens, ref_s, speed):
+    device = ref_s.device
+    tokens = torch.LongTensor([[0, *tokens, 0]]).to(device)
+    input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
+    text_mask = length_to_mask(input_lengths).to(device)
+    bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
+    d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
+    s = ref_s[:, 128:]
+    d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+    x, _ = model.predictor.lstm(d)
+    duration = model.predictor.duration_proj(x)
+    duration = torch.sigmoid(duration).sum(axis=-1) / speed
+    pred_dur = torch.round(duration).clamp(min=1).long()
+    pred_aln_trg = torch.zeros(input_lengths, pred_dur.sum().item())
+    c_frame = 0
+    for i in range(pred_aln_trg.size(0)):
+        pred_aln_trg[i, c_frame : c_frame + pred_dur[0, i].item()] = 1
+        c_frame += pred_dur[0, i].item()
+    en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device)
+    F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+    t_en = model.text_encoder(tokens, input_lengths, text_mask)
+    asr = t_en @ pred_aln_trg.unsqueeze(0).to(device)
+    return model.decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze().cpu().numpy()
+
+
+def generate(model, text, voicepack, lang="a", speed=1):
+    ps = phonemize(text, lang)
+    tokens = tokenize(ps)
+    if not tokens:
+        return None
+    elif len(tokens) > 510:
+        tokens = tokens[:510]
+        print("Truncated to 510 tokens")
+    ref_s = voicepack[len(tokens)]
+    out = forward(model, tokens, ref_s, speed)
+    ps = "".join(next(k for k, v in VOCAB.items() if i == v) for i in tokens)
+    return out, ps

models.py

@@ -0,0 +1,738 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/models.py
+from istftnet import Decoder
+from munch import Munch
+from pathlib import Path
+from plbert import load_plbert
+from torch.nn.utils import weight_norm, spectral_norm
+import json
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LearnedDownSample(nn.Module):
+    def __init__(self, layer_type, dim_in):
+        super().__init__()
+        self.layer_type = layer_type
+
+        if self.layer_type == "none":
+            self.conv = nn.Identity()
+        elif self.layer_type == "timepreserve":
+            self.conv = spectral_norm(
+                nn.Conv2d(
+                    dim_in,
+                    dim_in,
+                    kernel_size=(3, 1),
+                    stride=(2, 1),
+                    groups=dim_in,
+                    padding=(1, 0),
+                )
+            )
+        elif self.layer_type == "half":
+            self.conv = spectral_norm(
+                nn.Conv2d(
+                    dim_in,
+                    dim_in,
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    groups=dim_in,
+                    padding=1,
+                )
+            )
+        else:
+            raise RuntimeError(
+                "Got unexpected donwsampletype %s, expected is [none, timepreserve, half]"
+                % self.layer_type
+            )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class LearnedUpSample(nn.Module):
+    def __init__(self, layer_type, dim_in):
+        super().__init__()
+        self.layer_type = layer_type
+
+        if self.layer_type == "none":
+            self.conv = nn.Identity()
+        elif self.layer_type == "timepreserve":
+            self.conv = nn.ConvTranspose2d(
+                dim_in,
+                dim_in,
+                kernel_size=(3, 1),
+                stride=(2, 1),
+                groups=dim_in,
+                output_padding=(1, 0),
+                padding=(1, 0),
+            )
+        elif self.layer_type == "half":
+            self.conv = nn.ConvTranspose2d(
+                dim_in,
+                dim_in,
+                kernel_size=(3, 3),
+                stride=(2, 2),
+                groups=dim_in,
+                output_padding=1,
+                padding=1,
+            )
+        else:
+            raise RuntimeError(
+                "Got unexpected upsampletype %s, expected is [none, timepreserve, half]"
+                % self.layer_type
+            )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class DownSample(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == "none":
+            return x
+        elif self.layer_type == "timepreserve":
+            return F.avg_pool2d(x, (2, 1))
+        elif self.layer_type == "half":
+            if x.shape[-1] % 2 != 0:
+                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
+            return F.avg_pool2d(x, 2)
+        else:
+            raise RuntimeError(
+                "Got unexpected donwsampletype %s, expected is [none, timepreserve, half]"
+                % self.layer_type
+            )
+
+
+class UpSample(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == "none":
+            return x
+        elif self.layer_type == "timepreserve":
+            return F.interpolate(x, scale_factor=(2, 1), mode="nearest")
+        elif self.layer_type == "half":
+            return F.interpolate(x, scale_factor=2, mode="nearest")
+        else:
+            raise RuntimeError(
+                "Got unexpected upsampletype %s, expected is [none, timepreserve, half]"
+                % self.layer_type
+            )
+
+
+class ResBlk(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        actv=nn.LeakyReLU(0.2),
+        normalize=False,
+        downsample="none",
+    ):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = DownSample(downsample)
+        self.downsample_res = LearnedDownSample(downsample, dim_in)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = spectral_norm(nn.Conv2d(dim_in, dim_in, 3, 1, 1))
+        self.conv2 = spectral_norm(nn.Conv2d(dim_in, dim_out, 3, 1, 1))
+        if self.normalize:
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        if self.learned_sc:
+            self.conv1x1 = spectral_norm(
+                nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+            )
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = self.downsample(x)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        x = self.downsample_res(x)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / np.sqrt(2)  # unit variance
+
+
+class LinearNorm(torch.nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, w_init_gain="linear"):
+        super(LinearNorm, self).__init__()
+        self.linear_layer = torch.nn.Linear(in_dim, out_dim, bias=bias)
+
+        torch.nn.init.xavier_uniform_(
+            self.linear_layer.weight, gain=torch.nn.init.calculate_gain(w_init_gain)
+        )
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+class Discriminator2d(nn.Module):
+    def __init__(self, dim_in=48, num_domains=1, max_conv_dim=384, repeat_num=4):
+        super().__init__()
+        blocks = []
+        blocks += [spectral_norm(nn.Conv2d(1, dim_in, 3, 1, 1))]
+
+        for lid in range(repeat_num):
+            dim_out = min(dim_in * 2, max_conv_dim)
+            blocks += [ResBlk(dim_in, dim_out, downsample="half")]
+            dim_in = dim_out
+
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [spectral_norm(nn.Conv2d(dim_out, dim_out, 5, 1, 0))]
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [nn.AdaptiveAvgPool2d(1)]
+        blocks += [spectral_norm(nn.Conv2d(dim_out, num_domains, 1, 1, 0))]
+        self.main = nn.Sequential(*blocks)
+
+    def get_feature(self, x):
+        features = []
+        for l in self.main:
+            x = l(x)
+            features.append(x)
+        out = features[-1]
+        out = out.view(out.size(0), -1)  # (batch, num_domains)
+        return out, features
+
+    def forward(self, x):
+        out, features = self.get_feature(x)
+        out = out.squeeze()  # (batch)
+        return out, features
+
+
+class ResBlk1d(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        actv=nn.LeakyReLU(0.2),
+        normalize=False,
+        downsample="none",
+        dropout_p=0.2,
+    ):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample_type = downsample
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out)
+        self.dropout_p = dropout_p
+
+        if self.downsample_type == "none":
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(
+                nn.Conv1d(
+                    dim_in, dim_in, kernel_size=3, stride=2, groups=dim_in, padding=1
+                )
+            )
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_in, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        if self.normalize:
+            self.norm1 = nn.InstanceNorm1d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm1d(dim_in, affine=True)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def downsample(self, x):
+        if self.downsample_type == "none":
+            return x
+        else:
+            if x.shape[-1] % 2 != 0:
+                x = torch.cat([x, x[..., -1].unsqueeze(-1)], dim=-1)
+            return F.avg_pool1d(x, 2)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        x = self.downsample(x)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = F.dropout(x, p=self.dropout_p, training=self.training)
+
+        x = self.conv1(x)
+        x = self.pool(x)
+        if self.normalize:
+            x = self.norm2(x)
+
+        x = self.actv(x)
+        x = F.dropout(x, p=self.dropout_p, training=self.training)
+
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / np.sqrt(2)  # unit variance
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)):
+        super().__init__()
+        self.embedding = nn.Embedding(n_symbols, channels)
+
+        padding = (kernel_size - 1) // 2
+        self.cnn = nn.ModuleList()
+        for _ in range(depth):
+            self.cnn.append(
+                nn.Sequential(
+                    weight_norm(
+                        nn.Conv1d(
+                            channels, channels, kernel_size=kernel_size, padding=padding
+                        )
+                    ),
+                    LayerNorm(channels),
+                    actv,
+                    nn.Dropout(0.2),
+                )
+            )
+        # self.cnn = nn.Sequential(*self.cnn)
+
+        self.lstm = nn.LSTM(
+            channels, channels // 2, 1, batch_first=True, bidirectional=True
+        )
+
+    def forward(self, x, input_lengths, m):
+        x = self.embedding(x)  # [B, T, emb]
+        x = x.transpose(1, 2)  # [B, emb, T]
+        m = m.to(input_lengths.device).unsqueeze(1)
+        x.masked_fill_(m, 0.0)
+
+        for c in self.cnn:
+            x = c(x)
+            x.masked_fill_(m, 0.0)
+
+        x = x.transpose(1, 2)  # [B, T, chn]
+
+        input_lengths = input_lengths.cpu().numpy()
+        x = nn.utils.rnn.pack_padded_sequence(
+            x, input_lengths, batch_first=True, enforce_sorted=False
+        )
+
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+
+        x = x.transpose(-1, -2)
+        x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]])
+
+        x_pad[:, :, : x.shape[-1]] = x
+        x = x_pad.to(x.device)
+
+        x.masked_fill_(m, 0.0)
+
+        return x
+
+    def inference(self, x):
+        x = self.embedding(x)
+        x = x.transpose(1, 2)
+        x = self.cnn(x)
+        x = x.transpose(1, 2)
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        return x
+
+    def length_to_mask(self, lengths):
+        mask = (
+            torch.arange(lengths.max())
+            .unsqueeze(0)
+            .expand(lengths.shape[0], -1)
+            .type_as(lengths)
+        )
+        mask = torch.gt(mask + 1, lengths.unsqueeze(1))
+        return mask
+
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm1d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features * 2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, layer_type):
+        super().__init__()
+        self.layer_type = layer_type
+
+    def forward(self, x):
+        if self.layer_type == "none":
+            return x
+        else:
+            return F.interpolate(x, scale_factor=2, mode="nearest")
+
+
+class AdainResBlk1d(nn.Module):
+    def __init__(
+        self,
+        dim_in,
+        dim_out,
+        style_dim=64,
+        actv=nn.LeakyReLU(0.2),
+        upsample="none",
+        dropout_p=0.0,
+    ):
+        super().__init__()
+        self.actv = actv
+        self.upsample_type = upsample
+        self.upsample = UpSample1d(upsample)
+        self.learned_sc = dim_in != dim_out
+        self._build_weights(dim_in, dim_out, style_dim)
+        self.dropout = nn.Dropout(dropout_p)
+
+        if upsample == "none":
+            self.pool = nn.Identity()
+        else:
+            self.pool = weight_norm(
+                nn.ConvTranspose1d(
+                    dim_in,
+                    dim_in,
+                    kernel_size=3,
+                    stride=2,
+                    groups=dim_in,
+                    padding=1,
+                    output_padding=1,
+                )
+            )
+
+    def _build_weights(self, dim_in, dim_out, style_dim):
+        self.conv1 = weight_norm(nn.Conv1d(dim_in, dim_out, 3, 1, 1))
+        self.conv2 = weight_norm(nn.Conv1d(dim_out, dim_out, 3, 1, 1))
+        self.norm1 = AdaIN1d(style_dim, dim_in)
+        self.norm2 = AdaIN1d(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = weight_norm(nn.Conv1d(dim_in, dim_out, 1, 1, 0, bias=False))
+
+    def _shortcut(self, x):
+        x = self.upsample(x)
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        x = self.pool(x)
+        x = self.conv1(self.dropout(x))
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(self.dropout(x))
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / np.sqrt(2)
+        return out
+
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, style_dim, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.fc = nn.Linear(style_dim, channels * 2)
+
+    def forward(self, x, s):
+        x = x.transpose(-1, -2)
+        x = x.transpose(1, -1)
+
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        gamma, beta = gamma.transpose(1, -1), beta.transpose(1, -1)
+
+        x = F.layer_norm(x, (self.channels,), eps=self.eps)
+        x = (1 + gamma) * x + beta
+        return x.transpose(1, -1).transpose(-1, -2)
+
+
+class ProsodyPredictor(nn.Module):
+    def __init__(self, style_dim, d_hid, nlayers, max_dur=50, dropout=0.1):
+        super().__init__()
+
+        self.text_encoder = DurationEncoder(
+            sty_dim=style_dim, d_model=d_hid, nlayers=nlayers, dropout=dropout
+        )
+
+        self.lstm = nn.LSTM(
+            d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True
+        )
+        self.duration_proj = LinearNorm(d_hid, max_dur)
+
+        self.shared = nn.LSTM(
+            d_hid + style_dim, d_hid // 2, 1, batch_first=True, bidirectional=True
+        )
+        self.F0 = nn.ModuleList()
+        self.F0.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.F0.append(
+            AdainResBlk1d(
+                d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout
+            )
+        )
+        self.F0.append(
+            AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout)
+        )
+
+        self.N = nn.ModuleList()
+        self.N.append(AdainResBlk1d(d_hid, d_hid, style_dim, dropout_p=dropout))
+        self.N.append(
+            AdainResBlk1d(
+                d_hid, d_hid // 2, style_dim, upsample=True, dropout_p=dropout
+            )
+        )
+        self.N.append(
+            AdainResBlk1d(d_hid // 2, d_hid // 2, style_dim, dropout_p=dropout)
+        )
+
+        self.F0_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+        self.N_proj = nn.Conv1d(d_hid // 2, 1, 1, 1, 0)
+
+    def forward(self, texts, style, text_lengths, alignment, m):
+        d = self.text_encoder(texts, style, text_lengths, m)
+
+        batch_size = d.shape[0]
+        text_size = d.shape[1]
+
+        # predict duration
+        input_lengths = text_lengths.cpu().numpy()
+        x = nn.utils.rnn.pack_padded_sequence(
+            d, input_lengths, batch_first=True, enforce_sorted=False
+        )
+
+        m = m.to(text_lengths.device).unsqueeze(1)
+
+        self.lstm.flatten_parameters()
+        x, _ = self.lstm(x)
+        x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+
+        x_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]])
+
+        x_pad[:, : x.shape[1], :] = x
+        x = x_pad.to(x.device)
+
+        duration = self.duration_proj(
+            nn.functional.dropout(x, 0.5, training=self.training)
+        )
+
+        en = d.transpose(-1, -2) @ alignment
+
+        return duration.squeeze(-1), en
+
+    def F0Ntrain(self, x, s):
+        x, _ = self.shared(x.transpose(-1, -2))
+
+        F0 = x.transpose(-1, -2)
+        for block in self.F0:
+            F0 = block(F0, s)
+        F0 = self.F0_proj(F0)
+
+        N = x.transpose(-1, -2)
+        for block in self.N:
+            N = block(N, s)
+        N = self.N_proj(N)
+
+        return F0.squeeze(1), N.squeeze(1)
+
+    def length_to_mask(self, lengths):
+        mask = (
+            torch.arange(lengths.max())
+            .unsqueeze(0)
+            .expand(lengths.shape[0], -1)
+            .type_as(lengths)
+        )
+        mask = torch.gt(mask + 1, lengths.unsqueeze(1))
+        return mask
+
+
+class DurationEncoder(nn.Module):
+    def __init__(self, sty_dim, d_model, nlayers, dropout=0.1):
+        super().__init__()
+        self.lstms = nn.ModuleList()
+        for _ in range(nlayers):
+            self.lstms.append(
+                nn.LSTM(
+                    d_model + sty_dim,
+                    d_model // 2,
+                    num_layers=1,
+                    batch_first=True,
+                    bidirectional=True,
+                    dropout=dropout,
+                )
+            )
+            self.lstms.append(AdaLayerNorm(sty_dim, d_model))
+
+        self.dropout = dropout
+        self.d_model = d_model
+        self.sty_dim = sty_dim
+
+    def forward(self, x, style, text_lengths, m):
+        masks = m.to(text_lengths.device)
+
+        x = x.permute(2, 0, 1)
+        s = style.expand(x.shape[0], x.shape[1], -1)
+        x = torch.cat([x, s], axis=-1)
+        x.masked_fill_(masks.unsqueeze(-1).transpose(0, 1), 0.0)
+
+        x = x.transpose(0, 1)
+        input_lengths = text_lengths.cpu().numpy()
+        x = x.transpose(-1, -2)
+
+        for block in self.lstms:
+            if isinstance(block, AdaLayerNorm):
+                x = block(x.transpose(-1, -2), style).transpose(-1, -2)
+                x = torch.cat([x, s.permute(1, -1, 0)], axis=1)
+                x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0)
+            else:
+                x = x.transpose(-1, -2)
+                x = nn.utils.rnn.pack_padded_sequence(
+                    x, input_lengths, batch_first=True, enforce_sorted=False
+                )
+                block.flatten_parameters()
+                x, _ = block(x)
+                x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
+                x = F.dropout(x, p=self.dropout, training=self.training)
+                x = x.transpose(-1, -2)
+
+                x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]])
+
+                x_pad[:, :, : x.shape[-1]] = x
+                x = x_pad.to(x.device)
+
+        return x.transpose(-1, -2)
+
+    def inference(self, x, style):
+        x = self.embedding(x.transpose(-1, -2)) * np.sqrt(self.d_model)
+        style = style.expand(x.shape[0], x.shape[1], -1)
+        x = torch.cat([x, style], axis=-1)
+        src = self.pos_encoder(x)
+        output = self.transformer_encoder(src).transpose(0, 1)
+        return output
+
+    def length_to_mask(self, lengths):
+        mask = (
+            torch.arange(lengths.max())
+            .unsqueeze(0)
+            .expand(lengths.shape[0], -1)
+            .type_as(lengths)
+        )
+        mask = torch.gt(mask + 1, lengths.unsqueeze(1))
+        return mask
+
+
+# https://github.com/yl4579/StyleTTS2/blob/main/utils.py
+def recursive_munch(d):
+    if isinstance(d, dict):
+        return Munch((k, recursive_munch(v)) for k, v in d.items())
+    elif isinstance(d, list):
+        return [recursive_munch(v) for v in d]
+    else:
+        return d
+
+
+def build_model(path, device):
+    config = Path(path).parent / "config.json"
+    assert config.exists(), f"Config path incorrect: config.json not found at {config}"
+    with open(config, "r") as r:
+        args = recursive_munch(json.load(r))
+    assert args.decoder.type == "istftnet", f"Unknown decoder type: {args.decoder.type}"
+    decoder = Decoder(
+        dim_in=args.hidden_dim,
+        style_dim=args.style_dim,
+        dim_out=args.n_mels,
+        resblock_kernel_sizes=args.decoder.resblock_kernel_sizes,
+        upsample_rates=args.decoder.upsample_rates,
+        upsample_initial_channel=args.decoder.upsample_initial_channel,
+        resblock_dilation_sizes=args.decoder.resblock_dilation_sizes,
+        upsample_kernel_sizes=args.decoder.upsample_kernel_sizes,
+        gen_istft_n_fft=args.decoder.gen_istft_n_fft,
+        gen_istft_hop_size=args.decoder.gen_istft_hop_size,
+    )
+    text_encoder = TextEncoder(
+        channels=args.hidden_dim,
+        kernel_size=5,
+        depth=args.n_layer,
+        n_symbols=args.n_token,
+    )
+    predictor = ProsodyPredictor(
+        style_dim=args.style_dim,
+        d_hid=args.hidden_dim,
+        nlayers=args.n_layer,
+        max_dur=args.max_dur,
+        dropout=args.dropout,
+    )
+    bert = load_plbert()
+    bert_encoder = nn.Linear(bert.config.hidden_size, args.hidden_dim)
+    for parent in [bert, bert_encoder, predictor, decoder, text_encoder]:
+        for child in parent.children():
+            if isinstance(child, nn.RNNBase):
+                child.flatten_parameters()
+    model = Munch(
+        bert=bert.to(device).eval(),
+        bert_encoder=bert_encoder.to(device).eval(),
+        predictor=predictor.to(device).eval(),
+        decoder=decoder.to(device).eval(),
+        text_encoder=text_encoder.to(device).eval(),
+    )
+    for key, state_dict in torch.load(path, map_location="cpu", weights_only=True)[
+        "net"
+    ].items():
+        assert key in model, key
+        try:
+            model[key].load_state_dict(state_dict)
+        except:
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+            model[key].load_state_dict(state_dict, strict=False)
+    return model

packages.txt

@@ -0,0 +1 @@
+espeak-ng

plbert.py

@@ -0,0 +1,15 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py
+from transformers import AlbertConfig, AlbertModel
+
+class CustomAlbert(AlbertModel):
+    def forward(self, *args, **kwargs):
+        # Call the original forward method
+        outputs = super().forward(*args, **kwargs)
+        # Only return the last_hidden_state
+        return outputs.last_hidden_state
+
+def load_plbert():
+    plbert_config = {'vocab_size': 178, 'hidden_size': 768, 'num_attention_heads': 12, 'intermediate_size': 2048, 'max_position_embeddings': 512, 'num_hidden_layers': 12, 'dropout': 0.1}
+    albert_base_configuration = AlbertConfig(**plbert_config)
+    bert = CustomAlbert(albert_base_configuration)
+    return bert

pretrained_models/Kokoro/__init__.py

@@ -0,0 +1,3 @@
+from .downloader import initialize_files
+
+initialize_files()

pretrained_models/Kokoro/downloader.py

@@ -0,0 +1,61 @@
+import os
+import shutil
+from huggingface_hub import snapshot_download
+
+
+def download_files(repo_id, target_base_dir, patterns):
+    """
+    Downloads files matching patterns from a Hugging Face repository and organizes them in a directory structure.
+
+    :param repo_id: Hugging Face repository ID.
+    :param target_base_dir: Base directory where files should be stored.
+    :param patterns: A dictionary mapping subdirectories to file patterns.
+                     Example: {"root": ["config.json", "*.pth"], "voices": ["*.pt"]}
+    """
+    # Ensure target base directory exists
+    if not os.path.exists(target_base_dir):
+        os.makedirs(target_base_dir)
+
+    # Download the snapshot containing all matching files
+    snapshot_dir = snapshot_download(repo_id=repo_id, allow_patterns="*")
+
+    # Loop through patterns and subdirectories
+    for subdir, file_patterns in patterns.items():
+        # Set target directory for root-level files
+        target_dir = (
+            target_base_dir
+            if subdir == "root"
+            else os.path.join(target_base_dir, subdir)
+        )
+        os.makedirs(target_dir, exist_ok=True)
+
+        for file_pattern in file_patterns:
+            # Walk through the snapshot directory to find matching files
+            for root, _, files in os.walk(snapshot_dir):
+                for file in files:
+                    if file.endswith(file_pattern.lstrip("*")):  # Match pattern
+                        source_path = os.path.join(root, file)
+                        target_file_path = os.path.join(target_dir, file)
+
+                        # Check if file already exists
+                        if not os.path.exists(target_file_path):
+                            # Copy the file to the target directory
+                            shutil.copy(source_path, target_file_path)
+                            print(f"Downloaded and saved: {file} to {target_dir}")
+                        else:
+                            print(f"File already exists, skipping: {target_file_path}")
+
+
+def initialize_files():
+    repo_id = "hexgrad/Kokoro-82M"
+    target_base_dir = "pretrained_models/Kokoro"  # Base directory for files
+    file_patterns = {
+        "root": ["config.json", "*.pth"],  # Files for the root directory
+        "voices": ["*.pt"],  # Wildcard for voice pack files
+    }
+
+    download_files(repo_id, target_base_dir, file_patterns)
+
+
+if __name__ == "__main__":
+    initialize_files()

pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "kokoro-studio"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+  "gradio>=5.9.1",
+  "munch>=4.0.0",
+  "openphonemizer>=0.1.2",
+  "phonemizer>=3.3.0",
+  "pyyaml>=6.0.2",
+  "scipy>=1.14.1",
+  "soundfile>=0.12.1",
+  "torch>=2.5.1",
+  "transformers>=4.47.1",
+]

requirements.txt

@@ -0,0 +1,9 @@
+gradio>=5.9.1
+munch>=4.0.0
+openphonemizer>=0.1.2
+phonemizer>=3.3.0
+pyyaml>=6.0.2
+scipy>=1.14.1
+soundfile>=0.12.1
+torch>=2.5.1
+transformers>=4.47.1

tts_cli.py

@@ -0,0 +1,510 @@
+#!/usr/bin/env python3
+# tts_cli.py
+"""
+Example CLI for generating audio with Kokoro-StyleTTS2.
+
+Usage:
+  python tts_cli.py \
+    --model /path/to/kokoro-v0_19.pth \
+    --config /path/to/config.json \
+    --text "Hello, my stinking friends from 1906! You stink." \
+    --voicepack /path/to/af.pt \
+    --output output.wav
+
+Make sure:
+1. `models.py` is in the same folder (with `build_model`, `Decoder`, etc.).
+2. You have installed the needed libraries:
+   pip install torch phonemizer munch soundfile pyyaml
+3. The model is a checkpoint that your `build_model` can load.
+
+Adapt as needed!
+"""
+
+import argparse
+import os
+import re
+import torch
+import soundfile as sf
+import numpy as np
+from phonemizer import backend as phonemizer_backend
+
+# If you use eSpeak library:
+try:
+    from espeak_util import set_espeak_library
+
+    set_espeak_library()
+except ImportError:
+    pass
+
+# --------------------------------------------------------------------
+# Import from your local `models.py` (requires that file to be present).
+#   This example assumes `build_model` loads the entire TTS submodules
+#   (bert, bert_encoder, predictor, decoder, text_encoder).
+# --------------------------------------------------------------------
+from models import build_model
+
+
+def resplit_strings(arr):
+    """
+    Given a list of string tokens (e.g. words, phrases), tries to
+    split them into two sub-lists whose total lengths are as balanced
+    as possible. The goal is to chunk a large string in half without
+    splitting in the middle of a word.
+    """
+    if not arr:
+        return "", ""
+    if len(arr) == 1:
+        return arr[0], ""
+
+    min_diff = float("inf")
+    best_split = 0
+    lengths = [len(s) for s in arr]
+    spaces = len(arr) - 1
+    left_len = 0
+    right_len = sum(lengths) + spaces
+
+    for i in range(1, len(arr)):
+        # Add current word + space to left side
+        left_len += lengths[i - 1] + (1 if i > 1 else 0)
+        # Remove from right side
+        right_len -= lengths[i - 1] + 1
+        diff = abs(left_len - right_len)
+        if diff < min_diff:
+            min_diff = diff
+            best_split = i
+
+    return " ".join(arr[:best_split]), " ".join(arr[best_split:])
+
+
+def recursive_split(text, lang="a"):
+    """
+    Splits a piece of text into smaller segments so that
+    each segment's phoneme length < some ~limit (~500 tokens).
+    """
+    # We'll reuse your existing `phonemize_text` + `tokenize` from script 1
+    # to see if it is < 512 tokens. If it is, return it as a single chunk.
+    # Otherwise, split on punctuation or whitespace and recurse.
+
+    # 1. Phonemize first, check length
+    ps = phonemize_text(text, lang=lang, do_normalize=True)
+    tokens = tokenize(ps)
+    if len(tokens) < 512:
+        return [(text, ps)]
+
+    # If too large, we split on certain punctuation or fallback to whitespace
+    # We'll look for punctuation that often indicates sentence boundaries
+    # If none found, fallback to space-split
+    for punctuation in [r"[.?!…]", r"[:,;—]"]:
+        pattern = f"(?:(?<={punctuation})|(?<={punctuation}[\"'»])) "
+        # Attempt to split on that punctuation
+        splits = re.split(pattern, text)
+        if len(splits) > 1:
+            break
+    else:
+        # If we didn't break out, just do whitespace split
+        splits = text.split(" ")
+
+    # Use resplit_strings to chunk it about halfway
+    left, right = resplit_strings(splits)
+    # Recurse
+    return recursive_split(left, lang=lang) + recursive_split(right, lang=lang)
+
+
+def segment_and_tokenize(long_text, lang="a"):
+    """
+    Takes a large text, optionally normalizes or cleans it,
+    then breaks it into a list of (segment_text, segment_phonemes).
+    """
+    # Additional cleaning if you want:
+    # long_text = normalize_text(long_text) # your existing function
+    # We chunk it up using recursive_split
+    segments = recursive_split(long_text, lang=lang)
+    return segments
+
+
+# -------------- Normalization & Phonemization Routines -------------- #
+def parens_to_angles(s):
+    return s.replace("(", "«").replace(")", "»")
+
+
+def split_num(num):
+    num = num.group()
+    if "." in num:
+        return num
+    elif ":" in num:
+        h, m = [int(n) for n in num.split(":")]
+        if m == 0:
+            return f"{h} o'clock"
+        elif m < 10:
+            return f"{h} oh {m}"
+        return f"{h} {m}"
+    year = int(num[:4])
+    if year < 1100 or year % 1000 < 10:
+        return num
+    left, right = num[:2], int(num[2:4])
+    s = "s" if num.endswith("s") else ""
+    if 100 <= year % 1000 <= 999:
+        if right == 0:
+            return f"{left} hundred{s}"
+        elif right < 10:
+            return f"{left} oh {right}{s}"
+    return f"{left} {right}{s}"
+
+
+def flip_money(m):
+    m = m.group()
+    bill = "dollar" if m[0] == "$" else "pound"
+    if m[-1].isalpha():
+        return f"{m[1:]} {bill}s"
+    elif "." not in m:
+        s = "" if m[1:] == "1" else "s"
+        return f"{m[1:]} {bill}{s}"
+    b, c = m[1:].split(".")
+    s = "" if b == "1" else "s"
+    c = int(c.ljust(2, "0"))
+    coins = (
+        f"cent{'' if c == 1 else 's'}"
+        if m[0] == "$"
+        else ("penny" if c == 1 else "pence")
+    )
+    return f"{b} {bill}{s} and {c} {coins}"
+
+
+def point_num(num):
+    a, b = num.group().split(".")
+    return " point ".join([a, " ".join(b)])
+
+
+def normalize_text(text):
+    text = text.replace(chr(8216), "'").replace(chr(8217), "'")
+    text = text.replace("«", chr(8220)).replace("»", chr(8221))
+    text = text.replace(chr(8220), '"').replace(chr(8221), '"')
+    text = parens_to_angles(text)
+
+    # Replace some common full-width punctuation in CJK:
+    for a, b in zip("、。！，：；？", ",.!,:;?"):
+        text = text.replace(a, b + " ")
+
+    text = re.sub(r"[^\S \n]", " ", text)
+    text = re.sub(r"  +", " ", text)
+    text = re.sub(r"(?<=\n) +(?=\n)", "", text)
+    text = re.sub(r"\bD[Rr]\.(?= [A-Z])", "Doctor", text)
+    text = re.sub(r"\b(?:Mr\.|MR\.(?= [A-Z]))", "Mister", text)
+    text = re.sub(r"\b(?:Ms\.|MS\.(?= [A-Z]))", "Miss", text)
+    text = re.sub(r"\b(?:Mrs\.|MRS\.(?= [A-Z]))", "Mrs", text)
+    text = re.sub(r"\betc\.(?! [A-Z])", "etc", text)
+    text = re.sub(r"(?i)\b(y)eah?\b", r"\1e'a", text)
+    text = re.sub(
+        r"\d*\.\d+|\b\d{4}s?\b|(?<!:)\b(?:[1-9]|1[0-2]):[0-5]\d\b(?!:)",
+        split_num,
+        text,
+    )
+    text = re.sub(r"(?<=\d),(?=\d)", "", text)
+    text = re.sub(
+        r"(?i)[$£]\d+(?:\.\d+)?(?: hundred| thousand| (?:[bm]|tr)illion)*\b|[$£]\d+\.\d\d?\b",
+        flip_money,
+        text,
+    )
+    text = re.sub(r"\d*\.\d+", point_num, text)
+    text = re.sub(r"(?<=\d)-(?=\d)", " to ", text)  # Could be minus; adjust if needed
+    text = re.sub(r"(?<=\d)S", " S", text)
+    text = re.sub(r"(?<=[BCDFGHJ-NP-TV-Z])'?s\b", "'S", text)
+    text = re.sub(r"(?<=X')S\b", "s", text)
+    text = re.sub(
+        r"(?:[A-Za-z]\.){2,} [a-z]", lambda m: m.group().replace(".", "-"), text
+    )
+    text = re.sub(r"(?i)(?<=[A-Z])\.(?=[A-Z])", "-", text)
+    return text.strip()
+
+
+# -------------------------------------------------------------------
+# Vocab and Symbol Mapping
+# -------------------------------------------------------------------
+def get_vocab():
+    _pad = "$"
+    _punctuation = ';:,.!?¡¿—…"«»“” '
+    _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+    symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+    dicts = {}
+    for i, s in enumerate(symbols):
+        dicts[s] = i
+    return dicts
+
+
+VOCAB = get_vocab()
+
+
+def tokenize(ps: str):
+    """Convert the phoneme string into integer tokens based on VOCAB."""
+    return [VOCAB.get(p) for p in ps if p in VOCAB]
+
+
+# -------------------------------------------------------------------
+# Initialize a simple phonemizer
+#   For English:
+#       'a' ~ en-us
+#       'b' ~ en-gb
+# -------------------------------------------------------------------
+phonemizers = dict(
+    a=phonemizer_backend.EspeakBackend(
+        language="en-us", preserve_punctuation=True, with_stress=True
+    ),
+    b=phonemizer_backend.EspeakBackend(
+        language="en-gb", preserve_punctuation=True, with_stress=True
+    ),
+    # You can add more, e.g. 'j': some Japanese phonemizer, etc.
+)
+
+
+def phonemize_text(text, lang="a", do_normalize=True):
+    if do_normalize:
+        text = normalize_text(text)
+    ps_list = phonemizers[lang].phonemize([text])
+    ps = ps_list[0] if ps_list else ""
+
+    # Some custom replacements (from your code)
+    ps = ps.replace("kəkˈoːɹoʊ", "kˈoʊkəɹoʊ").replace("kəkˈɔːɹəʊ", "kˈəʊkəɹəʊ")
+    ps = ps.replace("ʲ", "j").replace("r", "ɹ").replace("x", "k").replace("ɬ", "l")
+    # Example: insert space before "hˈʌndɹɪd" if there's a letter, e.g. "nˈaɪn" => "nˈaɪn hˈʌndɹɪd"
+    ps = re.sub(r"(?<=[a-zɹː])(?=hˈʌndɹɪd)", " ", ps)
+    # "z" at the end of a word -> remove space (just your snippet)
+    ps = re.sub(r' z(?=[;:,.!?¡¿—…"«»“” ]|$)', "z", ps)
+    # If lang is 'a', handle "ninety" => "ninedi"? Just from your snippet:
+    if lang == "a":
+        ps = re.sub(r"(?<=nˈaɪn)ti(?!ː)", "di", ps)
+
+    # Only keep valid symbols
+    ps = "".join(p for p in ps if p in VOCAB)
+    return ps.strip()
+
+
+# -------------------------------------------------------------------
+# Utility for generating text masks
+# -------------------------------------------------------------------
+def length_to_mask(lengths):
+    # lengths is a Tensor of shape [B], containing the text length for each batch
+    max_len = lengths.max()
+    row_ids = torch.arange(max_len, device=lengths.device).unsqueeze(0)
+    mask = row_ids.expand(lengths.shape[0], -1)
+    return (mask + 1) > lengths.unsqueeze(1)
+
+
+# -------------------------------------------------------------------
+# The forward pass for inference (from your snippet).
+# This version references `model.predictor`, `model.decoder`, etc.
+# -------------------------------------------------------------------
+@torch.no_grad()
+def forward_tts(model, tokens, ref_s, speed=1.0):
+    """
+    model:   Munch with submodels: bert, bert_encoder, predictor, decoder, text_encoder
+    tokens:  list[int], the tokenized input (without [0, ... , 0] yet)
+    ref_s:   reference embedding (torch.Tensor)
+    speed:   float, speed factor
+    """
+    device = ref_s.device
+    tokens_t = torch.LongTensor([[0, *tokens, 0]]).to(device)  # add boundary tokens
+    input_lengths = torch.LongTensor([tokens_t.shape[-1]]).to(device)
+    text_mask = length_to_mask(input_lengths).to(device)
+
+    # 1. Encode with BERT
+    bert_dur = model.bert(tokens_t, attention_mask=(~text_mask).int())
+    d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
+
+    # 2. Prosody predictor
+    s = ref_s[
+        :, 128:
+    ]  # from your snippet: the last 128 is ???, or the first 128 is ???
+
+    d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+    x, _ = model.predictor.lstm(d)
+    duration = model.predictor.duration_proj(x)
+    duration = torch.sigmoid(duration).sum(axis=-1) / speed
+    pred_dur = torch.round(duration).clamp(min=1).long()
+
+    # 3. Expand alignment
+    total_len = pred_dur.sum().item()
+    pred_aln_trg = torch.zeros(input_lengths, total_len, device=device)
+    c_frame = 0
+    for i in range(pred_aln_trg.size(0)):
+        n = pred_dur[0, i].item()
+        pred_aln_trg[i, c_frame : c_frame + n] = 1
+        c_frame += n
+
+    # 4. Run F0 + Noise predictor
+    en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0)
+    F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+
+    # 5. Text encoder -> asr
+    t_en = model.text_encoder(tokens_t, input_lengths, text_mask)
+    asr = t_en @ pred_aln_trg.unsqueeze(0)
+
+    # 6. Decode audio
+    audio = model.decoder(asr, F0_pred, N_pred, ref_s[:, :128])  # B x audio_len
+    return audio.squeeze().cpu().numpy()
+
+
+def generate_tts(model, text, voicepack, lang="a", speed=1.0):
+    """
+    model:    the Munch returned by build_model(...)
+    text:     the input text (string)
+    voicepack: the torch Tensor reference embedding, or a dict of them
+    lang:     'a' or 'b' or etc. from your phonemizers
+    speed:    speech speed factor
+    sample_rate: sampling rate for the output
+    """
+    # 1. Phonemize
+    ps = phonemize_text(text, lang=lang, do_normalize=True)
+    tokens = tokenize(ps)
+    if not tokens:
+        return None, ps
+
+    # 2. Retrieve reference style
+    #    If your voicepack is a single embedding for all lengths, adapt as needed.
+    #    If your voicepack is something like `voicepack[len(tokens)]`, do that.
+    #    If you have multiple voices, you might do something else.
+    try:
+        ref_s = voicepack[len(tokens)]
+    except:
+        # fallback if len(tokens) is out of range
+        ref_s = voicepack[-1]
+    ref_s = ref_s.to("cpu" if not next(model.bert.parameters()).is_cuda else "cuda")
+
+    # 3. Generate
+    audio = forward_tts(model, tokens, ref_s, speed=speed)
+    return audio, ps
+
+
+def generate_long_form_tts(model, full_text, voicepack, lang="a", speed=1.0):
+    """
+    Generate TTS for a large `full_text`, splitting it into smaller segments
+    and concatenating the resulting audio.
+
+    Returns: (np.float32 array) final_audio, list_of_segment_phonemes
+    """
+    # 1. Segment the text
+    segments = segment_and_tokenize(full_text, lang=lang)
+    # segments is a list of (seg_text, seg_phonemes)
+
+    # 2. For each segment, call `generate_tts(...)`
+    audio_chunks = []
+    all_phonemes = []
+    for i, (seg_text, seg_ps) in enumerate(segments, 1):
+        print(f"[LongForm] Generating chunk {i}/{len(segments)}: {seg_text[:40]}...")
+        audio, used_phonemes = generate_tts(
+            model, seg_text, voicepack, lang=lang, speed=speed
+        )
+        if audio is not None:
+            audio_chunks.append(audio)
+            all_phonemes.append(used_phonemes)
+        else:
+            print(f"[LongForm] Skipped empty segment {i}...")
+
+    if not audio_chunks:
+        return None, []
+
+    # 3. Concatenate the audio
+    final_audio = np.concatenate(audio_chunks, axis=0)
+    return final_audio, all_phonemes
+
+
+# -------------------------------------------------------------------
+# Main CLI
+# -------------------------------------------------------------------
+def main():
+    parser = argparse.ArgumentParser(description="Kokoro-StyleTTS2 CLI Example")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="pretrained_models/Kokoro/kokoro-v0_19.pth",
+        help="Path to your model checkpoint (e.g. kokoro-v0_19.pth).",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="pretrained_models/Kokoro/config.json",
+        help="Path to config.json (used by build_model).",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Hello world! This is Kokoro, a new text-to-speech model based on StyleTTS2 from 2024!",
+        help="Text to be converted into speech.",
+    )
+    parser.add_argument(
+        "--voicepack",
+        type=str,
+        default="pretrained_models/Kokoro/voices/af.pt",
+        help="Path to a .pt file for your reference embedding(s).",
+    )
+    parser.add_argument(
+        "--output", type=str, default="output.wav", help="Output WAV filename."
+    )
+    parser.add_argument(
+        "--speed",
+        type=float,
+        default=1.0,
+        help="Speech speed factor, e.g. 0.8 slower, 1.2 faster, etc.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        choices=["cpu", "cuda"],
+        help="Device to run inference on.",
+    )
+    args = parser.parse_args()
+
+    # 1. Build model using your local build_model function
+    #    (which loads TextEncoder, Decoder, etc. and returns a Munch).
+    if not os.path.isfile(args.config):
+        raise FileNotFoundError(f"config.json not found: {args.config}")
+
+    # Optionally load config as Munch (depends on your build_model usage)
+    # But your snippet does something like:
+    #   with open(config, 'r') as r: ...
+    #   ...
+    #   model = build_model(path, device)
+    # We'll do the same but in a simpler form:
+    device = (
+        args.device if (args.device == "cuda" and torch.cuda.is_available()) else "cpu"
+    )
+    print(f"Loading model from: {args.model}")
+    model = build_model(
+        args.model, device
+    )  # This requires that `args.model` is the checkpoint path
+
+    # Because `build_model` returns a Munch (dict of submodules),
+    # we can't just do `model.eval()`, we must set each submodule to eval:
+    for k, subm in model.items():
+        if isinstance(subm, torch.nn.Module):
+            subm.eval()
+
+    # 2. Load voicepack
+    if not os.path.isfile(args.voicepack):
+        raise FileNotFoundError(f"Voicepack file not found: {args.voicepack}")
+    print(f"Loading voicepack from: {args.voicepack}")
+    vp = torch.load(args.voicepack, map_location=device)
+    # If your voicepack is an nn.Module, set it to eval as well
+    if isinstance(vp, torch.nn.Module):
+        vp.eval()
+
+    # 3. Generate audio
+    print(f"Generating speech for text: {args.text}")
+    audio, phonemes = generate_long_form_tts(
+        model, args.text, vp, lang="a", speed=args.speed
+    )
+    if audio is None:
+        print("No tokens were generated (maybe empty text?). Exiting.")
+        return
+
+    # 4. Write WAV
+    print(f"Writing output to: {args.output}")
+    sf.write(args.output, audio, 22050)
+
+    print("Finished!")
+    print(f"Phonemes used: {phonemes}")
+
+
+if __name__ == "__main__":
+    main()

tts_cli_op.py

@@ -0,0 +1,569 @@
+#!/usr/bin/env python3
+# tts_cli_op.py
+"""
+Example CLI for generating audio with Kokoro-StyleTTS2.
+
+Usage:
+  python tts_cli.py \
+    --model /path/to/kokoro-v0_19.pth \
+    --config /path/to/config.json \
+    --text "Hello, my stinking friends from 1906! You stink." \
+    --voicepack /path/to/af.pt \
+    --output output.wav
+
+Make sure:
+1. `models.py` is in the same folder (with `build_model`, `Decoder`, etc.).
+2. You have installed the needed libraries:
+   pip install torch phonemizer munch soundfile pyyaml
+3. The model is a checkpoint that your `build_model` can load.
+
+Adapt as needed!
+"""
+
+import argparse
+import os
+import re
+import torch
+import soundfile as sf
+import numpy as np
+from openphonemizer import OpenPhonemizer
+from typing import List
+import joblib
+
+# --------------------------------------------------------------------
+# Import from your local `models.py` (requires that file to be present).
+#   This example assumes `build_model` loads the entire TTS submodules
+#   (bert, bert_encoder, predictor, decoder, text_encoder).
+# --------------------------------------------------------------------
+from models import build_model
+
+
+def resplit_strings(arr):
+    """
+    Given a list of string tokens (e.g. words, phrases), tries to
+    split them into two sub-lists whose total lengths are as balanced
+    as possible. The goal is to chunk a large string in half without
+    splitting in the middle of a word.
+    """
+    if not arr:
+        return "", ""
+    if len(arr) == 1:
+        return arr[0], ""
+
+    min_diff = float("inf")
+    best_split = 0
+    lengths = [len(s) for s in arr]
+    spaces = len(arr) - 1
+    left_len = 0
+    right_len = sum(lengths) + spaces
+
+    for i in range(1, len(arr)):
+        # Add current word + space to left side
+        left_len += lengths[i - 1] + (1 if i > 1 else 0)
+        # Remove from right side
+        right_len -= lengths[i - 1] + 1
+        diff = abs(left_len - right_len)
+        if diff < min_diff:
+            min_diff = diff
+            best_split = i
+
+    return " ".join(arr[:best_split]), " ".join(arr[best_split:])
+
+
+def recursive_split(text, lang="a"):
+    """
+    Splits a piece of text into smaller segments so that
+    each segment's phoneme length < some ~limit (~500 tokens).
+    """
+    # We'll reuse your existing `phonemize_text` + `tokenize` from script 1
+    # to see if it is < 512 tokens. If it is, return it as a single chunk.
+    # Otherwise, split on punctuation or whitespace and recurse.
+
+    # 1. Phonemize first, check length
+    ps = phonemize_text(text, do_normalize=True)
+    tokens = tokenize(ps)
+    if len(tokens) < 512:
+        return [(text, ps)]
+
+    # If too large, we split on certain punctuation or fallback to whitespace
+    # We'll look for punctuation that often indicates sentence boundaries
+    # If none found, fallback to space-split
+    for punctuation in [r"[.?!…]", r"[:,;—]"]:
+        pattern = f"(?:(?<={punctuation})|(?<={punctuation}[\"'»])) "
+        # Attempt to split on that punctuation
+        splits = re.split(pattern, text)
+        if len(splits) > 1:
+            break
+    else:
+        # If we didn't break out, just do whitespace split
+        splits = text.split(" ")
+
+    # Use resplit_strings to chunk it about halfway
+    left, right = resplit_strings(splits)
+    # Recurse
+    return recursive_split(left) + recursive_split(right)
+
+
+def segment_and_tokenize(long_text, lang="a"):
+    """
+    Takes a large text, optionally normalizes or cleans it,
+    then breaks it into a list of (segment_text, segment_phonemes).
+    """
+    # Additional cleaning if you want:
+    # long_text = normalize_text(long_text) # your existing function
+    # We chunk it up using recursive_split
+    segments = recursive_split(long_text)
+    return segments
+
+
+# -------------- Normalization & Phonemization Routines -------------- #
+def parens_to_angles(s):
+    return s.replace("(", "«").replace(")", "»")
+
+
+def split_num(num):
+    num = num.group()
+    if "." in num:
+        return num
+    elif ":" in num:
+        h, m = [int(n) for n in num.split(":")]
+        if m == 0:
+            return f"{h} o'clock"
+        elif m < 10:
+            return f"{h} oh {m}"
+        return f"{h} {m}"
+    year = int(num[:4])
+    if year < 1100 or year % 1000 < 10:
+        return num
+    left, right = num[:2], int(num[2:4])
+    s = "s" if num.endswith("s") else ""
+    if 100 <= year % 1000 <= 999:
+        if right == 0:
+            return f"{left} hundred{s}"
+        elif right < 10:
+            return f"{left} oh {right}{s}"
+    return f"{left} {right}{s}"
+
+
+def flip_money(m):
+    m = m.group()
+    bill = "dollar" if m[0] == "$" else "pound"
+    if m[-1].isalpha():
+        return f"{m[1:]} {bill}s"
+    elif "." not in m:
+        s = "" if m[1:] == "1" else "s"
+        return f"{m[1:]} {bill}{s}"
+    b, c = m[1:].split(".")
+    s = "" if b == "1" else "s"
+    c = int(c.ljust(2, "0"))
+    coins = (
+        f"cent{'' if c == 1 else 's'}"
+        if m[0] == "$"
+        else ("penny" if c == 1 else "pence")
+    )
+    return f"{b} {bill}{s} and {c} {coins}"
+
+
+def point_num(num):
+    a, b = num.group().split(".")
+    return " point ".join([a, " ".join(b)])
+
+
+def normalize_text(text):
+    text = text.replace(chr(8216), "'").replace(chr(8217), "'")
+    text = text.replace("«", chr(8220)).replace("»", chr(8221))
+    text = text.replace(chr(8220), '"').replace(chr(8221), '"')
+    text = parens_to_angles(text)
+
+    # Replace some common full-width punctuation in CJK:
+    for a, b in zip("、。！，：；？", ",.!,:;?"):
+        text = text.replace(a, b + " ")
+
+    text = re.sub(r"[^\S \n]", " ", text)
+    text = re.sub(r"  +", " ", text)
+    text = re.sub(r"(?<=\n) +(?=\n)", "", text)
+    text = re.sub(r"\bD[Rr]\.(?= [A-Z])", "Doctor", text)
+    text = re.sub(r"\b(?:Mr\.|MR\.(?= [A-Z]))", "Mister", text)
+    text = re.sub(r"\b(?:Ms\.|MS\.(?= [A-Z]))", "Miss", text)
+    text = re.sub(r"\b(?:Mrs\.|MRS\.(?= [A-Z]))", "Mrs", text)
+    text = re.sub(r"\betc\.(?! [A-Z])", "etc", text)
+    text = re.sub(r"(?i)\b(y)eah?\b", r"\1e'a", text)
+    text = re.sub(
+        r"\d*\.\d+|\b\d{4}s?\b|(?<!:)\b(?:[1-9]|1[0-2]):[0-5]\d\b(?!:)",
+        split_num,
+        text,
+    )
+    text = re.sub(r"(?<=\d),(?=\d)", "", text)
+    text = re.sub(
+        r"(?i)[$£]\d+(?:\.\d+)?(?: hundred| thousand| (?:[bm]|tr)illion)*\b|[$£]\d+\.\d\d?\b",
+        flip_money,
+        text,
+    )
+    text = re.sub(r"\d*\.\d+", point_num, text)
+    text = re.sub(r"(?<=\d)-(?=\d)", " to ", text)  # Could be minus; adjust if needed
+    text = re.sub(r"(?<=\d)S", " S", text)
+    text = re.sub(r"(?<=[BCDFGHJ-NP-TV-Z])'?s\b", "'S", text)
+    text = re.sub(r"(?<=X')S\b", "s", text)
+    text = re.sub(
+        r"(?:[A-Za-z]\.){2,} [a-z]", lambda m: m.group().replace(".", "-"), text
+    )
+    text = re.sub(r"(?i)(?<=[A-Z])\.(?=[A-Z])", "-", text)
+    return text.strip()
+
+
+# -------------------------------------------------------------------
+# Vocab and Symbol Mapping
+# -------------------------------------------------------------------
+def get_vocab():
+    _pad = "$"
+    _punctuation = ';:,.!?¡¿—…"«»“” '
+    _letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+    _letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+    symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+    dicts = {}
+    for i, s in enumerate(symbols):
+        dicts[s] = i
+    return dicts
+
+
+VOCAB = get_vocab()
+
+
+def tokenize(ps: str):
+    """Convert the phoneme string into integer tokens based on VOCAB."""
+    return [VOCAB.get(p) for p in ps if p in VOCAB]
+
+
+# -------------------------------------------------------------------
+# Initialize a simple phonemizer
+#   For English:
+#       'a' ~ en-us
+#       'b' ~ en-gb
+# -------------------------------------------------------------------
+
+# Wrapper around OpenPhonemizer to enable batch phonemization
+open_phonemizer = OpenPhonemizer()
+
+
+class Phonemizer:
+    def __init__(self, open_phonemizer):
+        """
+        open_phonemizer: an instance of OpenPhonemizer()
+        """
+        self.open_phonemizer = open_phonemizer
+
+    def phonemize(
+        self,
+        text: List[str],
+        njobs: int = 4,
+        strip: bool = False,
+    ) -> List[str]:
+        """
+        Phonemizes a list of input strings using OpenPhonemizer (which
+        itself only supports single-string input).
+
+        Parameters
+        ----------
+        text : list of str
+            Each element is an utterance (or line) to be phonemized.
+        njobs : int
+            The number of parallel jobs for phonemization.
+        strip : bool
+            Not used by OpenPhonemizer directly, but you can implement
+            an additional final “strip” step if needed.
+
+        Returns
+        -------
+        list of str
+            The phonemized text strings, in the same order as input.
+        """
+
+        if not isinstance(text, list) or any(not isinstance(x, str) for x in text):
+            raise ValueError("`text` must be a list of strings.")
+
+        # Optionally do any pre-processing you want here ...
+        # e.g. text = [self._some_preprocess(line) for line in text]
+
+        # If we only have one job, do it in a single loop
+        if njobs == 1:
+            return [self._phonemize_single(utterance, strip) for utterance in text]
+
+        # Otherwise, we can split `text` into chunks and process in parallel
+        # The easiest approach is to chunk the entire list into smaller sublists
+        # and phonemize each chunk. Then flatten them.
+        # For large corpora, you might want a more sophisticated approach.
+        chunked_results = joblib.Parallel(n_jobs=njobs)(
+            joblib.delayed(self._phonemize_single)(t, strip) for t in text
+        )
+        return chunked_results
+
+    def _phonemize_single(self, line: str, strip: bool) -> str:
+        """
+        Phonemize a single line using openphonemizer.
+        """
+        # OpenPhonemizer usage:
+        #   out_str = self.open_phonemizer(line)
+        #   (That returns the phonemes as a string.)
+        phonemes = self.open_phonemizer(line)
+
+        # Implement a post-strip if you want to mimic removing trailing delimiters
+        # (this is just a placeholder for demonstration).
+        if strip:
+            phonemes = phonemes.rstrip()
+
+        return phonemes
+
+
+# initiate the wrapper:
+phonemizer_wrapper = Phonemizer(open_phonemizer)
+
+
+def phonemize_text(text, lang="a", do_normalize=True):
+    if do_normalize:
+        text = normalize_text(text)
+    ps_list = phonemizer_wrapper.phonemize([text])
+    ps = ps_list[0] if ps_list else ""
+
+    # Some custom replacements
+    ps = ps.replace("kəkˈoːɹoʊ", "kˈoʊkəɹoʊ").replace("kəkˈɔːɹəʊ", "kˈəʊkəɹəʊ")
+    ps = ps.replace("ʲ", "j").replace("r", "ɹ").replace("x", "k").replace("ɬ", "l")
+    # Example: insert space before "hˈʌndɹɪd" if there's a letter, e.g. "nˈaɪn" => "nˈaɪn hˈʌndɹɪd"
+    ps = re.sub(r"(?<=[a-zɹː])(?=hˈʌndɹɪd)", " ", ps)
+    # "z" at the end of a word -> remove space (just your snippet)
+    ps = re.sub(r' z(?=[;:,.!?¡¿—…"«»“” ]|$)', "z", ps)
+    # Handle "ninety" => "ninedi"? Just from your snippet:
+    # If lang is 'a', handle "ninety" => "ninedi"? Just from your snippet:
+    if lang == "a":
+        ps = re.sub(r"(?<=nˈaɪn)ti(?!ː)", "di", ps)
+
+    # Only keep valid symbols
+    ps = "".join(p for p in ps if p in VOCAB)
+    return ps.strip()
+
+
+# -------------------------------------------------------------------
+# Utility for generating text masks
+# -------------------------------------------------------------------
+def length_to_mask(lengths):
+    # lengths is a Tensor of shape [B], containing the text length for each batch
+    max_len = lengths.max()
+    row_ids = torch.arange(max_len, device=lengths.device).unsqueeze(0)
+    mask = row_ids.expand(lengths.shape[0], -1)
+    return (mask + 1) > lengths.unsqueeze(1)
+
+
+# -------------------------------------------------------------------
+# The forward pass for inference (from your snippet).
+# This version references `model.predictor`, `model.decoder`, etc.
+# -------------------------------------------------------------------
+@torch.no_grad()
+def forward_tts(model, tokens, ref_s, speed=1.0):
+    """
+    model:   Munch with submodels: bert, bert_encoder, predictor, decoder, text_encoder
+    tokens:  list[int], the tokenized input (without [0, ... , 0] yet)
+    ref_s:   reference embedding (torch.Tensor)
+    speed:   float, speed factor
+    """
+    device = ref_s.device
+    tokens_t = torch.LongTensor([[0, *tokens, 0]]).to(device)  # add boundary tokens
+    input_lengths = torch.LongTensor([tokens_t.shape[-1]]).to(device)
+    text_mask = length_to_mask(input_lengths).to(device)
+
+    # 1. Encode with BERT
+    bert_dur = model.bert(tokens_t, attention_mask=(~text_mask).int())
+    d_en = model.bert_encoder(bert_dur).transpose(-1, -2)
+
+    # 2. Prosody predictor
+    s = ref_s[
+        :, 128:
+    ]  # from your snippet: the last 128 is ???, or the first 128 is ???
+
+    d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+    x, _ = model.predictor.lstm(d)
+    duration = model.predictor.duration_proj(x)
+    duration = torch.sigmoid(duration).sum(axis=-1) / speed
+    pred_dur = torch.round(duration).clamp(min=1).long()
+
+    # 3. Expand alignment
+    total_len = pred_dur.sum().item()
+    pred_aln_trg = torch.zeros(input_lengths, total_len, device=device)
+    c_frame = 0
+    for i in range(pred_aln_trg.size(0)):
+        n = pred_dur[0, i].item()
+        pred_aln_trg[i, c_frame : c_frame + n] = 1
+        c_frame += n
+
+    # 4. Run F0 + Noise predictor
+    en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0)
+    F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
+
+    # 5. Text encoder -> asr
+    t_en = model.text_encoder(tokens_t, input_lengths, text_mask)
+    asr = t_en @ pred_aln_trg.unsqueeze(0)
+
+    # 6. Decode audio
+    audio = model.decoder(asr, F0_pred, N_pred, ref_s[:, :128])  # B x audio_len
+    return audio.squeeze().cpu().numpy()
+
+
+def generate_tts(model, text, voicepack, lang="a", speed=1.0):
+    """
+    model:    the Munch returned by build_model(...)
+    text:     the input text (string)
+    voicepack: the torch Tensor reference embedding, or a dict of them
+    speed:    speech speed factor
+    sample_rate: sampling rate for the output
+    """
+    # 1. Phonemize
+    ps = phonemize_text(text, do_normalize=True)
+    tokens = tokenize(ps)
+    if not tokens:
+        return None, ps
+
+    # 2. Retrieve reference style
+    #    If your voicepack is a single embedding for all lengths, adapt as needed.
+    #    If your voicepack is something like `voicepack[len(tokens)]`, do that.
+    #    If you have multiple voices, you might do something else.
+    try:
+        ref_s = voicepack[len(tokens)]
+    except:
+        # fallback if len(tokens) is out of range
+        ref_s = voicepack[-1]
+    ref_s = ref_s.to("cpu" if not next(model.bert.parameters()).is_cuda else "cuda")
+
+    # 3. Generate
+    audio = forward_tts(model, tokens, ref_s, speed=speed)
+    return audio, ps
+
+
+def generate_long_form_tts(model, full_text, voicepack, lang="a", speed=1.0):
+    """
+    Generate TTS for a large `full_text`, splitting it into smaller segments
+    and concatenating the resulting audio.
+
+    Returns: (np.float32 array) final_audio, list_of_segment_phonemes
+    """
+    # 1. Segment the text
+    segments = segment_and_tokenize(full_text)
+    # segments is a list of (seg_text, seg_phonemes)
+
+    # 2. For each segment, call `generate_tts(...)`
+    audio_chunks = []
+    all_phonemes = []
+    for i, (seg_text, seg_ps) in enumerate(segments, 1):
+        print(f"[LongForm] Generating chunk {i}/{len(segments)}: {seg_text[:40]}...")
+        audio, used_phonemes = generate_tts(model, seg_text, voicepack, speed=speed)
+        if audio is not None:
+            audio_chunks.append(audio)
+            all_phonemes.append(used_phonemes)
+        else:
+            print(f"[LongForm] Skipped empty segment {i}...")
+
+    if not audio_chunks:
+        return None, []
+
+    # 3. Concatenate the audio
+    final_audio = np.concatenate(audio_chunks, axis=0)
+    return final_audio, all_phonemes
+
+
+# -------------------------------------------------------------------
+# Main CLI
+# -------------------------------------------------------------------
+def main():
+    parser = argparse.ArgumentParser(description="Kokoro-StyleTTS2 CLI Example")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="pretrained_models/Kokoro/kokoro-v0_19.pth",
+        help="Path to your model checkpoint (e.g. kokoro-v0_19.pth).",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="pretrained_models/Kokoro/config.json",
+        help="Path to config.json (used by build_model).",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Hello world! This is Kokoro, a new text-to-speech model based on StyleTTS2 from 2024!",
+        help="Text to be converted into speech.",
+    )
+    parser.add_argument(
+        "--voicepack",
+        type=str,
+        default="pretrained_models/Kokoro/voices/af.pt",
+        help="Path to a .pt file for your reference embedding(s).",
+    )
+    parser.add_argument(
+        "--output", type=str, default="output.wav", help="Output WAV filename."
+    )
+    parser.add_argument(
+        "--speed",
+        type=float,
+        default=1.2,
+        help="Speech speed factor, e.g. 0.8 slower, 1.2 faster, etc.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        choices=["cpu", "cuda"],
+        help="Device to run inference on.",
+    )
+    args = parser.parse_args()
+
+    # 1. Build model using your local build_model function
+    #    (which loads TextEncoder, Decoder, etc. and returns a Munch).
+    if not os.path.isfile(args.config):
+        raise FileNotFoundError(f"config.json not found: {args.config}")
+
+    # Optionally load config as Munch (depends on your build_model usage)
+    # But your snippet does something like:
+    #   with open(config, 'r') as r: ...
+    #   ...
+    #   model = build_model(path, device)
+    # We'll do the same but in a simpler form:
+    device = (
+        args.device if (args.device == "cuda" and torch.cuda.is_available()) else "cpu"
+    )
+    print(f"Loading model from: {args.model}")
+    model = build_model(
+        args.model, device
+    )  # This requires that `args.model` is the checkpoint path
+
+    # Because `build_model` returns a Munch (dict of submodules),
+    # we can't just do `model.eval()`, we must set each submodule to eval:
+    for k, subm in model.items():
+        if isinstance(subm, torch.nn.Module):
+            subm.eval()
+
+    # 2. Load voicepack
+    if not os.path.isfile(args.voicepack):
+        raise FileNotFoundError(f"Voicepack file not found: {args.voicepack}")
+    print(f"Loading voicepack from: {args.voicepack}")
+    vp = torch.load(args.voicepack, map_location=device)
+    # If your voicepack is an nn.Module, set it to eval as well
+    if isinstance(vp, torch.nn.Module):
+        vp.eval()
+
+    # 3. Generate audio
+    print(f"Generating speech for text: {args.text}")
+    audio, phonemes = generate_long_form_tts(
+        model, args.text, vp, lang="a", speed=args.speed
+    )
+    if audio is None:
+        print("No tokens were generated (maybe empty text?). Exiting.")
+        return
+
+    # 4. Write WAV
+    print(f"Writing output to: {args.output}")
+    sf.write(args.output, audio, 22050)
+
+    print("Finished!")
+    print(f"Phonemes used: {phonemes}")
+
+
+if __name__ == "__main__":
+    main()