init

Dockerfile

@@ -1,7 +1,7 @@
 # Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
 # you will also find guides on how best to write your Dockerfile
 
-FROM python:3.9
+FROM python:3.11
 
 RUN useradd -m -u 1000 user
 USER user

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README copy.md

@@ -0,0 +1,130 @@
+# kokoro
+
+An inference library for [Kokoro-82M](https://huggingface.co/hexgrad/Kokoro-82M). You can [`pip install kokoro`](https://pypi.org/project/kokoro/).
+
+> **Kokoro** is an open-weight TTS model with 82 million parameters. Despite its lightweight architecture, it delivers comparable quality to larger models while being significantly faster and more cost-efficient. With Apache-licensed weights, Kokoro can be deployed anywhere from production environments to personal projects.
+
+### Usage
+You can run this basic cell on [Google Colab](https://colab.research.google.com/). [Listen to samples](https://huggingface.co/hexgrad/Kokoro-82M/blob/main/SAMPLES.md).
+```py
+!pip install -q kokoro>=0.9.4 soundfile
+!apt-get -qq -y install espeak-ng > /dev/null 2>&1
+from kokoro import KPipeline
+from IPython.display import display, Audio
+import soundfile as sf
+import torch
+pipeline = KPipeline(lang_code='a')
+text = '''
+[Kokoro](/kˈOkəɹO/) is an open-weight TTS model with 82 million parameters. Despite its lightweight architecture, it delivers comparable quality to larger models while being significantly faster and more cost-efficient. With Apache-licensed weights, [Kokoro](/kˈOkəɹO/) can be deployed anywhere from production environments to personal projects.
+'''
+generator = pipeline(text, voice='af_heart')
+for i, (gs, ps, audio) in enumerate(generator):
+    print(i, gs, ps)
+    display(Audio(data=audio, rate=24000, autoplay=i==0))
+    sf.write(f'{i}.wav', audio, 24000)
+```
+Under the hood, `kokoro` uses [`misaki`](https://pypi.org/project/misaki/), a G2P library at https://github.com/hexgrad/misaki
+
+### Advanced Usage
+You can run this advanced cell on [Google Colab](https://colab.research.google.com/).
+```py
+# 1️⃣ Install kokoro
+!pip install -q kokoro>=0.9.4 soundfile
+# 2️⃣ Install espeak, used for English OOD fallback and some non-English languages
+!apt-get -qq -y install espeak-ng > /dev/null 2>&1
+
+# 3️⃣ Initalize a pipeline
+from kokoro import KPipeline
+from IPython.display import display, Audio
+import soundfile as sf
+import torch
+# 🇺🇸 'a' => American English, 🇬🇧 'b' => British English
+# 🇪🇸 'e' => Spanish es
+# 🇫🇷 'f' => French fr-fr
+# 🇮🇳 'h' => Hindi hi
+# 🇮🇹 'i' => Italian it
+# 🇯🇵 'j' => Japanese: pip install misaki[ja]
+# 🇧🇷 'p' => Brazilian Portuguese pt-br
+# 🇨🇳 'z' => Mandarin Chinese: pip install misaki[zh]
+pipeline = KPipeline(lang_code='a') # <= make sure lang_code matches voice, reference above.
+
+# This text is for demonstration purposes only, unseen during training
+text = '''
+The sky above the port was the color of television, tuned to a dead channel.
+"It's not like I'm using," Case heard someone say, as he shouldered his way through the crowd around the door of the Chat. "It's like my body's developed this massive drug deficiency."
+It was a Sprawl voice and a Sprawl joke. The Chatsubo was a bar for professional expatriates; you could drink there for a week and never hear two words in Japanese.
+
+These were to have an enormous impact, not only because they were associated with Constantine, but also because, as in so many other areas, the decisions taken by Constantine (or in his name) were to have great significance for centuries to come. One of the main issues was the shape that Christian churches were to take, since there was not, apparently, a tradition of monumental church buildings when Constantine decided to help the Christian church build a series of truly spectacular structures. The main form that these churches took was that of the basilica, a multipurpose rectangular structure, based ultimately on the earlier Greek stoa, which could be found in most of the great cities of the empire. Christianity, unlike classical polytheism, needed a large interior space for the celebration of its religious services, and the basilica aptly filled that need. We naturally do not know the degree to which the emperor was involved in the design of new churches, but it is tempting to connect this with the secular basilica that Constantine completed in the Roman forum (the so-called Basilica of Maxentius) and the one he probably built in Trier, in connection with his residence in the city at a time when he was still caesar.
+
+[Kokoro](/kˈOkəɹO/) is an open-weight TTS model with 82 million parameters. Despite its lightweight architecture, it delivers comparable quality to larger models while being significantly faster and more cost-efficient. With Apache-licensed weights, [Kokoro](/kˈOkəɹO/) can be deployed anywhere from production environments to personal projects.
+'''
+# text = '「もしおれがただ偶然、そしてこうしようというつもりでなくここに立っているのなら、ちょっとばかり絶望するところだな」と、そんなことが彼の頭に思い浮かんだ。'
+# text = '中國人民不信邪也不怕邪，不惹事也不怕事，任何外國不要指望我們會拿自己的核心利益做交易，不要指望我們會吞下損害我國主權、安��、發展利益的苦果！'
+# text = 'Los partidos políticos tradicionales compiten con los populismos y los movimientos asamblearios.'
+# text = 'Le dromadaire resplendissant déambulait tranquillement dans les méandres en mastiquant de petites feuilles vernissées.'
+# text = 'ट्रांसपोर्टरों की हड़ताल लगातार पांचवें दिन जारी, दिसंबर से इलेक्ट्रॉनिक टोल कलेक्शनल सिस्टम'
+# text = "Allora cominciava l'insonnia, o un dormiveglia peggiore dell'insonnia, che talvolta assumeva i caratteri dell'incubo."
+# text = 'Elabora relatórios de acompanhamento cronológico para as diferentes unidades do Departamento que propõem contratos.'
+
+# 4️⃣ Generate, display, and save audio files in a loop.
+generator = pipeline(
+    text, voice='af_heart', # <= change voice here
+    speed=1, split_pattern=r'\n+'
+)
+# Alternatively, load voice tensor directly:
+# voice_tensor = torch.load('path/to/voice.pt', weights_only=True)
+# generator = pipeline(
+#     text, voice=voice_tensor,
+#     speed=1, split_pattern=r'\n+'
+# )
+
+for i, (gs, ps, audio) in enumerate(generator):
+    print(i)  # i => index
+    print(gs) # gs => graphemes/text
+    print(ps) # ps => phonemes
+    display(Audio(data=audio, rate=24000, autoplay=i==0))
+    sf.write(f'{i}.wav', audio, 24000) # save each audio file
+```
+
+### Windows Installation
+To install espeak-ng on Windows:
+1. Go to [espeak-ng releases](https://github.com/espeak-ng/espeak-ng/releases)
+2. Click on **Latest release** 
+3. Download the appropriate `*.msi` file (e.g. **espeak-ng-20191129-b702b03-x64.msi**)
+4. Run the downloaded installer
+
+For advanced configuration and usage on Windows, see the [official espeak-ng Windows guide](https://github.com/espeak-ng/espeak-ng/blob/master/docs/guide.md)
+
+### MacOS Apple Silicon GPU Acceleration
+
+On Mac M1/M2/M3/M4 devices, you can explicitly specify the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to enable GPU acceleration.
+
+```bash
+PYTORCH_ENABLE_MPS_FALLBACK=1 python run-your-kokoro-script.py
+```
+
+### Conda Environment
+Use the following conda `environment.yml` if you're facing any dependency issues.
+```yaml
+name: kokoro
+channels:
+  - defaults
+dependencies:
+  - python==3.9       
+  - libstdcxx~=12.4.0 # Needed to load espeak correctly. Try removing this if you're facing issues with Espeak fallback. 
+  - pip:
+      - kokoro>=0.3.1
+      - soundfile
+      - misaki[en]
+```
+
+### 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.
+- 📊 Thank you to everyone who contributed synthetic training data.
+- ❤️ Special thanks to all compute sponsors.
+- 👾 Discord server: https://discord.gg/QuGxSWBfQy
+- 🪽 Kokoro is a Japanese word that translates to "heart" or "spirit". Kokoro is also a [character in the Terminator franchise](https://terminator.fandom.com/wiki/Kokoro) along with [Misaki](https://github.com/hexgrad/misaki?tab=readme-ov-file#acknowledgements).
+
+<img src="https://static0.gamerantimages.com/wordpress/wp-content/uploads/2024/08/terminator-zero-41-1.jpg" width="400" alt="kokoro" />
+# kokoro

demo/README.md

@@ -0,0 +1,15 @@
+---
+title: Kokoro TTS
+emoji: ❤️
+colorFrom: indigo
+colorTo: pink
+sdk: gradio
+sdk_version: 5.12.0
+app_file: app.py
+pinned: true
+license: apache-2.0
+short_description: Upgraded to v1.0!
+disable_embedding: true
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

demo/app.py

@@ -0,0 +1,182 @@
+import spaces
+from kokoro import KModel, KPipeline
+import gradio as gr
+import os
+import random
+import torch
+
+CUDA_AVAILABLE = torch.cuda.is_available()
+models = {gpu: KModel().to('cuda' if gpu else 'cpu').eval() for gpu in [False] + ([True] if CUDA_AVAILABLE else [])}
+pipelines = {lang_code: KPipeline(lang_code=lang_code, model=False) for lang_code in 'ab'}
+pipelines['a'].g2p.lexicon.golds['kokoro'] = 'kˈOkəɹO'
+pipelines['b'].g2p.lexicon.golds['kokoro'] = 'kˈQkəɹQ'
+
+@spaces.GPU(duration=30)
+def forward_gpu(ps, ref_s, speed):
+    return models[True](ps, ref_s, speed)
+
+def generate_first(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
+    pipeline = pipelines[voice[0]]
+    pack = pipeline.load_voice(voice)
+    use_gpu = use_gpu and CUDA_AVAILABLE
+    for _, ps, _ in pipeline(text, voice, speed):
+        ref_s = pack[len(ps)-1]
+        try:
+            if use_gpu:
+                audio = forward_gpu(ps, ref_s, speed)
+            else:
+                audio = models[False](ps, ref_s, speed)
+        except gr.exceptions.Error as e:
+            if use_gpu:
+                gr.Warning(str(e))
+                gr.Info('Retrying with CPU. To avoid this error, change Hardware to CPU.')
+                audio = models[False](ps, ref_s, speed)
+            else:
+                raise gr.Error(e)
+        return (24000, audio.numpy()), ps
+    return None, ''
+
+# Arena API
+def predict(text, voice='af_heart', speed=1):
+    return generate_first(text, voice, speed, use_gpu=False)[0]
+
+def tokenize_first(text, voice='af_heart'):
+    pipeline = pipelines[voice[0]]
+    for _, ps, _ in pipeline(text, voice):
+        return ps
+    return ''
+
+def generate_all(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
+    pipeline = pipelines[voice[0]]
+    pack = pipeline.load_voice(voice)
+    use_gpu = use_gpu and CUDA_AVAILABLE
+    first = True
+    for _, ps, _ in pipeline(text, voice, speed):
+        ref_s = pack[len(ps)-1]
+        try:
+            if use_gpu:
+                audio = forward_gpu(ps, ref_s, speed)
+            else:
+                audio = models[False](ps, ref_s, speed)
+        except gr.exceptions.Error as e:
+            if use_gpu:
+                gr.Warning(str(e))
+                gr.Info('Switching to CPU')
+                audio = models[False](ps, ref_s, speed)
+            else:
+                raise gr.Error(e)
+        yield 24000, audio.numpy()
+        if first:
+            first = False
+            yield 24000, torch.zeros(1).numpy()
+
+with open('en.txt', 'r') as r:
+    random_quotes = [line.strip() for line in r]
+
+def get_random_quote():
+    return random.choice(random_quotes)
+
+def get_gatsby():
+    with open('gatsby5k.md', 'r') as r:
+        return r.read().strip()
+
+def get_frankenstein():
+    with open('frankenstein5k.md', 'r') as r:
+        return r.read().strip()
+
+CHOICES = {
+'🇺🇸 🚺 Heart ❤️': 'af_heart',
+'🇺🇸 🚺 Bella 🔥': 'af_bella',
+'🇺🇸 🚺 Nicole 🎧': 'af_nicole',
+'🇺🇸 🚺 Aoede': 'af_aoede',
+'🇺🇸 🚺 Kore': 'af_kore',
+'🇺🇸 🚺 Sarah': 'af_sarah',
+'🇺🇸 🚺 Nova': 'af_nova',
+'🇺🇸 🚺 Sky': 'af_sky',
+'🇺🇸 🚺 Alloy': 'af_alloy',
+'🇺🇸 🚺 Jessica': 'af_jessica',
+'🇺🇸 🚺 River': 'af_river',
+'🇺🇸 🚹 Michael': 'am_michael',
+'🇺🇸 🚹 Fenrir': 'am_fenrir',
+'🇺🇸 🚹 Puck': 'am_puck',
+'🇺🇸 🚹 Echo': 'am_echo',
+'🇺🇸 🚹 Eric': 'am_eric',
+'🇺🇸 🚹 Liam': 'am_liam',
+'🇺🇸 🚹 Onyx': 'am_onyx',
+'🇺🇸 🚹 Santa': 'am_santa',
+'🇺🇸 🚹 Adam': 'am_adam',
+'🇬🇧 🚺 Emma': 'bf_emma',
+'🇬🇧 🚺 Isabella': 'bf_isabella',
+'🇬🇧 🚺 Alice': 'bf_alice',
+'🇬🇧 🚺 Lily': 'bf_lily',
+'🇬🇧 🚹 George': 'bm_george',
+'🇬🇧 🚹 Fable': 'bm_fable',
+'🇬🇧 🚹 Lewis': 'bm_lewis',
+'🇬🇧 🚹 Daniel': 'bm_daniel',
+}
+for v in CHOICES.values():
+    pipelines[v[0]].load_voice(v)
+
+TOKEN_NOTE = '''
+💡 Customize pronunciation with Markdown link syntax and /slashes/ like `[Kokoro](/kˈOkəɹO/)`
+
+💬 To adjust intonation, try punctuation `;:,.!?—…"()“”` or stress `ˈ` and `ˌ`
+
+⬇️ Lower stress `[1 level](-1)` or `[2 levels](-2)`
+
+⬆️ Raise stress 1 level `[or](+2)` 2 levels (only works on less stressed, usually short words)
+'''
+
+with gr.Blocks() as generate_tab:
+    out_audio = gr.Audio(label='Output Audio', interactive=False, streaming=False, autoplay=True)
+    generate_btn = gr.Button('Generate', variant='primary')
+    with gr.Accordion('Output Tokens', open=True):
+        out_ps = gr.Textbox(interactive=False, show_label=False, info='Tokens used to generate the audio, up to 510 context length.')
+        tokenize_btn = gr.Button('Tokenize', variant='secondary')
+        gr.Markdown(TOKEN_NOTE)
+        predict_btn = gr.Button('Predict', variant='secondary', visible=False)
+
+STREAM_NOTE = ['⚠️ There is an unknown Gradio bug that might yield no audio the first time you click `Stream`.']
+STREAM_NOTE = '\n\n'.join(STREAM_NOTE)
+
+with gr.Blocks() as stream_tab:
+    out_stream = gr.Audio(label='Output Audio Stream', interactive=False, streaming=True, autoplay=True)
+    with gr.Row():
+        stream_btn = gr.Button('Stream', variant='primary')
+        stop_btn = gr.Button('Stop', variant='stop')
+    with gr.Accordion('Note', open=True):
+        gr.Markdown(STREAM_NOTE)
+        gr.DuplicateButton()
+
+API_OPEN = True
+with gr.Blocks() as app:
+    with gr.Row():
+        with gr.Column():
+            text = gr.Textbox(label='Input Text', info=f"Arbitrarily many characters supported")
+            with gr.Row():
+                voice = gr.Dropdown(list(CHOICES.items()), value='af_heart', label='Voice', info='Quality and availability vary by language')
+                use_gpu = gr.Dropdown(
+                    [('ZeroGPU 🚀', True), ('CPU 🐌', False)],
+                    value=CUDA_AVAILABLE,
+                    label='Hardware',
+                    info='GPU is usually faster, but has a usage quota',
+                    interactive=CUDA_AVAILABLE
+                )
+            speed = gr.Slider(minimum=0.5, maximum=2, value=1, step=0.1, label='Speed')
+            random_btn = gr.Button('🎲 Random Quote 💬', variant='secondary')
+            with gr.Row():
+                gatsby_btn = gr.Button('🥂 Gatsby 📕', variant='secondary')
+                frankenstein_btn = gr.Button('💀 Frankenstein 📗', variant='secondary')
+        with gr.Column():
+            gr.TabbedInterface([generate_tab, stream_tab], ['Generate', 'Stream'])
+    random_btn.click(fn=get_random_quote, inputs=[], outputs=[text])
+    gatsby_btn.click(fn=get_gatsby, inputs=[], outputs=[text])
+    frankenstein_btn.click(fn=get_frankenstein, inputs=[], outputs=[text])
+    generate_btn.click(fn=generate_first, inputs=[text, voice, speed, use_gpu], outputs=[out_audio, out_ps])
+    tokenize_btn.click(fn=tokenize_first, inputs=[text, voice], outputs=[out_ps])
+    stream_event = stream_btn.click(fn=generate_all, inputs=[text, voice, speed, use_gpu], outputs=[out_stream])
+    stop_btn.click(fn=None, cancels=stream_event)
+    predict_btn.click(fn=predict, inputs=[text, voice, speed], outputs=[out_audio])
+
+if __name__ == '__main__':
+    app.queue(api_open=API_OPEN).launch(server_name="0.0.0.0", server_port=40001, show_api=API_OPEN)

demo/frankenstein5k.md

@@ -0,0 +1,11 @@
+You will rejoice to hear that no disaster has accompanied the commencement of an enterprise which you have regarded with such evil forebodings. I arrived here yesterday, and my first task is to assure my dear sister of my welfare and increasing confidence in the success of my undertaking.
+
+I am already far north of London, and as I walk in the streets of Petersburgh, I feel a cold northern breeze play upon my cheeks, which braces my nerves and fills me with delight. Do you understand this feeling? This breeze, which has travelled from the regions towards which I am advancing, gives me a foretaste of those icy climes. Inspirited by this wind of promise, my daydreams become more fervent and vivid. I try in vain to be persuaded that the pole is the seat of frost and desolation; it ever presents itself to my imagination as the region of beauty and delight. There, Margaret, the sun is for ever visible, its broad disk just skirting the horizon and diffusing a perpetual splendour. There—for with your leave, my sister, I will put some trust in preceding navigators—there snow and frost are banished; and, sailing over a calm sea, we may be wafted to a land surpassing in wonders and in beauty every region hitherto discovered on the habitable globe. Its productions and features may be without example, as the phenomena of the heavenly bodies undoubtedly are in those undiscovered solitudes. What may not be expected in a country of eternal light? I may there discover the wondrous power which attracts the needle and may regulate a thousand celestial observations that require only this voyage to render their seeming eccentricities consistent for ever. I shall satiate my ardent curiosity with the sight of a part of the world never before visited, and may tread a land never before imprinted by the foot of man. These are my enticements, and they are sufficient to conquer all fear of danger or death and to induce me to commence this laborious voyage with the joy a child feels when he embarks in a little boat, with his holiday mates, on an expedition of discovery up his native river. But supposing all these conjectures to be false, you cannot contest the inestimable benefit which I shall confer on all mankind, to the last generation, by discovering a passage near the pole to those countries, to reach which at present so many months are requisite; or by ascertaining the secret of the magnet, which, if at all possible, can only be effected by an undertaking such as mine.
+
+These reflections have dispelled the agitation with which I began my letter, and I feel my heart glow with an enthusiasm which elevates me to heaven, for nothing contributes so much to tranquillise the mind as a steady purpose—a point on which the soul may fix its intellectual eye. This expedition has been the favourite dream of my early years. I have read with ardour the accounts of the various voyages which have been made in the prospect of arriving at the North Pacific Ocean through the seas which surround the pole. You may remember that a history of all the voyages made for purposes of discovery composed the whole of our good Uncle Thomas’s library. My education was neglected, yet I was passionately fond of reading. These volumes were my study day and night, and my familiarity with them increased that regret which I had felt, as a child, on learning that my father’s dying injunction had forbidden my uncle to allow me to embark in a seafaring life.
+
+These visions faded when I perused, for the first time, those poets whose effusions entranced my soul and lifted it to heaven. I also became a poet and for one year lived in a paradise of my own creation; I imagined that I also might obtain a niche in the temple where the names of Homer and Shakespeare are consecrated. You are well acquainted with my failure and how heavily I bore the disappointment. But just at that time I inherited the fortune of my cousin, and my thoughts were turned into the channel of their earlier bent.
+
+Six years have passed since I resolved on my present undertaking. I can, even now, remember the hour from which I dedicated myself to this great enterprise. I commenced by inuring my body to hardship. I accompanied the whale-fishers on several expeditions to the North Sea; I voluntarily endured cold, famine, thirst, and want of sleep; I often worked harder than the common sailors during the day and devoted my nights to the study of mathematics, the theory of medicine, and those branches of physical science from which a naval adventurer might derive the greatest practical advantage. Twice I actually hired myself as an under-mate in a Greenland whaler, and acquitted myself to admiration. I must own I felt a little proud when my captain offered me the second dignity in the vessel and entreated me to remain with the greatest earnestness, so valuable did he consider my services.
+
+And now, dear Margaret, do I not deserve to accomplish some great purpose?

demo/gatsby5k.md

@@ -0,0 +1,17 @@
+In my younger and more vulnerable years my father gave me some advice that I’ve been turning over in my mind ever since.
+
+“Whenever you feel like criticizing anyone,” he told me, “just remember that all the people in this world haven’t had the advantages that you’ve had.”
+
+He didn’t say any more, but we’ve always been unusually communicative in a reserved way, and I understood that he meant a great deal more than that. In consequence, I’m inclined to reserve all judgements, a habit that has opened up many curious natures to me and also made me the victim of not a few veteran bores. The abnormal mind is quick to detect and attach itself to this quality when it appears in a normal person, and so it came about that in college I was unjustly accused of being a politician, because I was privy to the secret griefs of wild, unknown men. Most of the confidences were unsought—frequently I have feigned sleep, preoccupation, or a hostile levity when I realized by some unmistakable sign that an intimate revelation was quivering on the horizon; for the intimate revelations of young men, or at least the terms in which they express them, are usually plagiaristic and marred by obvious suppressions. Reserving judgements is a matter of infinite hope. I am still a little afraid of missing something if I forget that, as my father snobbishly suggested, and I snobbishly repeat, a sense of the fundamental decencies is parcelled out unequally at birth.
+
+And, after boasting this way of my tolerance, I come to the admission that it has a limit. Conduct may be founded on the hard rock or the wet marshes, but after a certain point I don’t care what it’s founded on. When I came back from the East last autumn I felt that I wanted the world to be in uniform and at a sort of moral attention forever; I wanted no more riotous excursions with privileged glimpses into the human heart. Only Gatsby, the man who gives his name to this book, was exempt from my reaction—Gatsby, who represented everything for which I have an unaffected scorn. If personality is an unbroken series of successful gestures, then there was something gorgeous about him, some heightened sensitivity to the promises of life, as if he were related to one of those intricate machines that register earthquakes ten thousand miles away. This responsiveness had nothing to do with that flabby impressionability which is dignified under the name of the “creative temperament”—it was an extraordinary gift for hope, a romantic readiness such as I have never found in any other person and which it is not likely I shall ever find again. No—Gatsby turned out all right at the end; it is what preyed on Gatsby, what foul dust floated in the wake of his dreams that temporarily closed out my interest in the abortive sorrows and short-winded elations of men.
+
+My family have been prominent, well-to-do people in this Middle Western city for three generations. The Carraways are something of a clan, and we have a tradition that we’re descended from the Dukes of Buccleuch, but the actual founder of my line was my grandfather’s brother, who came here in fifty-one, sent a substitute to the Civil War, and started the wholesale hardware business that my father carries on today.
+
+I never saw this great-uncle, but I’m supposed to look like him—with special reference to the rather hard-boiled painting that hangs in father’s office. I graduated from New Haven in 1915, just a quarter of a century after my father, and a little later I participated in that delayed Teutonic migration known as the Great War. I enjoyed the counter-raid so thoroughly that I came back restless. Instead of being the warm centre of the world, the Middle West now seemed like the ragged edge of the universe—so I decided to go East and learn the bond business. Everybody I knew was in the bond business, so I supposed it could support one more single man. All my aunts and uncles talked it over as if they were choosing a prep school for me, and finally said, “Why—[ye-es](/jˈɛ ɛs/),” with very grave, hesitant faces. Father agreed to finance me for a year, and after various delays I came East, permanently, I thought, in the spring of twenty-two.
+
+The practical thing was to find rooms in the city, but it was a warm season, and I had just left a country of wide lawns and friendly trees, so when a young man at the office suggested that we take a house together in a commuting town, it sounded like a great idea. He found the house, a weather-beaten cardboard bungalow at eighty a month, but at the last minute the firm ordered him to Washington, and I went out to the country alone. I had a dog—at least I had him for a few days until he ran away—and an old Dodge and a Finnish woman, who made my bed and cooked breakfast and muttered Finnish wisdom to herself over the electric stove.
+
+It was lonely for a day or so until one morning some man, more recently arrived than I, stopped me on the road.
+
+“How do you get to West Egg village?” he asked helplessly.

demo/packages.txt

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

demo/requirements.txt

@@ -0,0 +1,3 @@
+kokoro>=0.7.13
+gradio
+pip

examples/device_examples.py

@@ -0,0 +1,45 @@
+"""
+Quick example to show how device selection can be controlled, and was checked
+"""
+import time
+from kokoro import KPipeline
+from loguru import logger
+
+def generate_audio(pipeline, text):
+    for _, _, audio in pipeline(text, voice='af_bella'):
+        samples = audio.shape[0] if audio is not None else 0
+        assert samples > 0, "No audio generated"
+        return samples
+
+def time_synthesis(device=None):
+    try:
+        start = time.perf_counter()
+        pipeline = KPipeline(lang_code='a', device=device)
+        samples = generate_audio(pipeline, "The quick brown fox jumps over the lazy dog.")
+        ms = (time.perf_counter() - start) * 1000
+        logger.info(f"✓ {device or 'auto':<6} | {ms:>5.1f}ms total | {samples:>6,d} samples")
+    except RuntimeError as e:
+        logger.error(f"✗ {'cuda' if 'CUDA' in str(e) else device or 'auto':<6} | {'not available' if 'CUDA' in str(e) else str(e)}")
+
+def compare_shared_model():
+    try:
+        start = time.perf_counter()
+        en_us = KPipeline(lang_code='a')
+        en_uk = KPipeline(lang_code='a', model=en_us.model)
+        
+        for pipeline in [en_us, en_uk]:
+            generate_audio(pipeline, "Testing model reuse.")
+                
+        ms = (time.perf_counter() - start) * 1000
+        logger.info(f"✓ reuse  | {ms:>5.1f}ms for both models")
+    except Exception as e:
+        logger.error(f"✗ reuse  | {str(e)}")
+
+if __name__ == '__main__':
+    logger.info("Device Selection & Performance")
+    logger.info("-" * 40)
+    time_synthesis()
+    time_synthesis('cuda')
+    time_synthesis('cpu') 
+    logger.info("-" * 40)
+    compare_shared_model()

examples/export.py

@@ -0,0 +1,149 @@
+import argparse
+import os
+import torch
+import onnx
+import onnxruntime as ort
+import sounddevice as sd
+
+from kokoro import KModel, KPipeline
+from kokoro.model import KModelForONNX
+
+def export_onnx(model, output):
+    onnx_file = output + "/" + "kokoro.onnx"
+
+    input_ids = torch.randint(1, 100, (48,)).numpy()
+    input_ids = torch.LongTensor([[0, *input_ids, 0]])
+    style = torch.randn(1, 256)
+    speed = torch.randint(1, 10, (1,)).int()
+
+    torch.onnx.export(
+        model, 
+        args = (input_ids, style, speed), 
+        f = onnx_file, 
+        export_params = True, 
+        verbose = True, 
+        input_names = [ 'input_ids', 'style', 'speed' ], 
+        output_names = [ 'waveform', 'duration' ],
+        opset_version = 17, 
+        dynamic_axes = {
+            'input_ids': {0: "batch_size", 1: 'input_ids_len' }, 
+            'style': {0: "batch_size"}, 
+            "speed": {0: "batch_size"}
+        }, 
+        do_constant_folding = True, 
+    )
+
+    print('export kokoro.onnx ok!')
+
+    onnx_model = onnx.load(onnx_file)
+    onnx.checker.check_model(onnx_model)
+    print('onnx check ok!')
+
+def load_input_ids(pipeline, text):
+    if pipeline.lang_code in 'ab':
+        _, tokens = pipeline.g2p(text)
+        for gs, ps, tks in pipeline.en_tokenize(tokens):
+            if not ps:
+                continue
+    else:
+        ps, _ = pipeline.g2p(text)
+
+    if len(ps) > 510:
+        ps = ps[:510]
+
+    input_ids = list(filter(lambda i: i is not None, map(lambda p: pipeline.model.vocab.get(p), ps)))
+    print(f"text: {text} -> phonemes: {ps} -> input_ids: {input_ids}")
+    input_ids = torch.LongTensor([[0, *input_ids, 0]]).to(pipeline.model.device)
+    return ps, input_ids
+
+def load_voice(pipeline, voice, phonemes):
+    pack = pipeline.load_voice(voice).to('cpu')
+    return pack[len(phonemes) - 1]
+
+def load_sample(model):
+    pipeline = KPipeline(lang_code='a', model=model.kmodel, device='cpu')
+    text = '''
+    In today's fast-paced tech world, building software applications has never been easier — thanks to AI-powered coding assistants.'
+    '''
+    text = '''
+    The sky above the port was the color of television, tuned to a dead channel.
+    '''
+    voice = 'checkpoints/voices/af_heart.pt'
+
+    pipeline = KPipeline(lang_code='z', model=model.kmodel, device='cpu')
+    text = '''
+    2月15日晚，猫眼专业版数据显示，截至发稿，《哪吒之魔童闹海》（或称《哪吒2》）今日票房已达7.8亿元，累计票房（含预售）超过114亿元。
+    '''
+    voice = 'checkpoints/voices/zf_xiaoxiao.pt'
+
+    phonemes, input_ids = load_input_ids(pipeline, text)
+    style = load_voice(pipeline, voice, phonemes)
+    speed = torch.IntTensor([1])
+
+    return input_ids, style, speed
+
+def inference_onnx(model, output):
+    onnx_file = output + "/" + "kokoro.onnx"
+    session = ort.InferenceSession(onnx_file)
+
+    input_ids, style, speed = load_sample(model)
+
+    outputs = session.run(None, {
+        'input_ids': input_ids.numpy(), 
+        'style': style.numpy(), 
+        'speed': speed.numpy(), 
+    })
+
+    output = torch.from_numpy(outputs[0])
+    print(f'output: {output.shape}')
+    print(output)
+
+    audio = output.numpy()
+    sd.play(audio, 24000)
+    sd.wait()
+
+def check_model(model):
+    input_ids, style, speed = load_sample(model)
+    output, duration = model(input_ids, style, speed)
+
+    print(f'output: {output.shape}')
+    print(f'duration: {duration.shape}')
+    print(output)
+
+    audio = output.numpy()
+    sd.play(audio, 24000)
+    sd.wait()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("Export kokoro Model to ONNX", add_help=True)
+    parser.add_argument("--inference", "-t", help="test kokoro.onnx model", action="store_true")
+    parser.add_argument("--check", "-m", help="check kokoro model", action="store_true")
+    parser.add_argument(
+        "--config_file", "-c", type=str, default="checkpoints/config.json", help="path to config file"
+    )
+    parser.add_argument(
+        "--checkpoint_path", "-p", type=str, default="checkpoints/kokoro-v1_0.pth", help="path to checkpoint file"
+    )
+    parser.add_argument(
+        "--output_dir", "-o", type=str, default="onnx", help="output directory"
+    )
+
+    args = parser.parse_args()
+
+    # cfg
+    config_file = args.config_file  # change the path of the model config file
+    checkpoint_path = args.checkpoint_path  # change the path of the model
+    output_dir = args.output_dir
+    
+    # make dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    kmodel = KModel(config=config_file, model=checkpoint_path, disable_complex=True)
+    model = KModelForONNX(kmodel).eval()
+
+    if args.inference:
+        inference_onnx(model, output_dir)
+    elif args.check:
+        check_model(model)
+    else:
+        export_onnx(model, output_dir)

examples/make_triton_compatible.py

@@ -0,0 +1,92 @@
+"""
+This script makes the ONNX model compatible with Triton inference server.
+"""
+
+import sys
+import numpy as np
+import onnx
+import onnxruntime as ort
+import onnx_graphsurgeon as gs
+
+
+def add_squeeze(graph, speed_input, speed_unsqueezed):
+    """
+    Add squeeze operation to the speed input to change shape from [batch_size, 1] to [batch_size]
+    """
+    # Create a squeeze node
+    squeeze_node = gs.Node(
+        op="Squeeze",
+        name="speed_squeeze",
+        inputs=[speed_unsqueezed],
+        outputs=[gs.Variable(name="speed_squeezed", dtype=speed_unsqueezed.dtype)]
+    )
+    
+    ## Find first node that has speed_unsqueezed as input
+    insert_idx = 0
+    for idx, node in enumerate(graph.nodes):
+        for i, input_name in enumerate(node.inputs):
+            if input_name.name == speed_unsqueezed.name:
+                insert_idx = idx
+                break
+        if insert_idx != 0:
+            break
+    
+    ## Add squeeze node to the graph
+    insert_idx = min(0, insert_idx - 1)
+    graph.nodes.insert(insert_idx, squeeze_node)
+    
+    # Update the speed input to point to the squeezed output
+    for node in graph.nodes:
+        for i, input_name in enumerate(node.inputs):
+            if input_name.name == speed_input.name and not node.name == "speed_squeeze":
+                node.inputs[i] = squeeze_node.outputs[0]
+                
+    return graph
+
+
+def main():
+    if len(sys.argv) != 2:
+        print("Usage: python make_triton_compatible.py <onnx_model_path>")
+        sys.exit(1)
+
+    onnx_model_path = sys.argv[1]
+    onnx_model = onnx.load(onnx_model_path)
+    onnx.checker.check_model(onnx_model)
+    print("Model is valid")
+    
+    graph = gs.import_onnx(onnx_model)
+
+    ## get input_id for speed
+    speed_idx, speed = None, None
+    for idx, input_ in enumerate(graph.inputs):
+        if input_.name=="speed":
+            speed_idx = idx
+            speed = input_
+
+    # Update the speed input to have shape [batch_size, 1]
+    speed_unsqueezed = gs.Variable(name="speed", dtype=speed.dtype, shape=[speed.shape[0], 1])
+    graph.inputs[speed_idx] = speed_unsqueezed
+        
+    ## Add squeeze to change speed shape from [batch_size, 1] to [batch_size]
+    if speed is not None:
+        print(f"Found speed input: {speed.name}")
+        print(f"Found speed input shape: {speed.shape}")
+        print(f"Found speed input dtype: {speed.dtype}")
+        print(f"Found speed input: {speed}")
+        print(f"Found speed input: {type(speed)}")
+        graph = add_squeeze(graph, speed, speed_unsqueezed)
+        
+        # Export the modified graph back to ONNX
+        modified_model = gs.export_onnx(graph)
+        onnx.checker.check_model(modified_model)
+        
+        # Save the modified model
+        output_path = onnx_model_path.replace('.onnx', '_triton.onnx')
+        onnx.save(modified_model, output_path)
+        print(f"Modified model saved to: {output_path}")
+    else:
+        print("Speed input not found in the model")
+
+
+if __name__ == "__main__":
+    main()

examples/phoneme_example.py

@@ -0,0 +1,62 @@
+from kokoro import KPipeline, KModel
+import torch
+from scipy.io import wavfile
+
+def save_audio(audio: torch.Tensor, filename: str):
+    """Helper function to save audio tensor as WAV file"""
+    if audio is not None:
+        # Ensure audio is on CPU and in the right format
+        audio_cpu = audio.cpu().numpy()
+        
+        # Save using scipy.io.wavfile
+        wavfile.write(
+            filename,
+            24000,  # Kokoro uses 24kHz sample rate
+            audio_cpu
+        )
+        print(f"Audio saved as '{filename}'")
+    else:
+        print("No audio was generated")
+
+def main():
+    # Initialize pipeline with American English
+    pipeline = KPipeline(lang_code='a')
+    
+    # The phoneme string for:
+    # "How are you today? I am doing reasonably well, thank you for asking"
+    phonemes = "hˌW ɑɹ ju tədˈA? ˌI ɐm dˈuɪŋ ɹˈizənəbli wˈɛl, θˈæŋk ju fɔɹ ˈæskɪŋ"
+    
+    try:
+        print("\nExample 1: Using generate_from_tokens with raw phonemes")
+        results = list(pipeline.generate_from_tokens(
+            tokens=phonemes,
+            voice="af_bella",
+            speed=1.0
+        ))
+        if results:
+            save_audio(results[0].audio, 'phoneme_output_new.wav')
+        
+        # Example 2: Using generate_from_tokens with pre-processed tokens
+        print("\nExample 2: Using generate_from_tokens with pre-processed tokens")
+        #  get the tokens through G2P or any other method
+        text = "How are you today? I am doing reasonably well, thank you for asking"
+        _, tokens = pipeline.g2p(text)
+        
+        # Then generate from tokens
+        for result in pipeline.generate_from_tokens(
+            tokens=tokens,
+            voice="af_bella",
+            speed=1.0
+        ):
+            # Each result may contain timestamps if available
+            if result.tokens:
+                for token in result.tokens:
+                    if hasattr(token, 'start_ts') and hasattr(token, 'end_ts'):
+                        print(f"Token: {token.text} ({token.start_ts:.2f}s - {token.end_ts:.2f}s)")
+            save_audio(result.audio, f'token_output_{hash(result.phonemes)}.wav')
+            
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+
+if __name__ == "__main__":
+    main()

kokoro/__init__.py

@@ -0,0 +1,23 @@
+__version__ = '0.9.4'
+
+from loguru import logger
+import sys
+
+# Remove default handler
+logger.remove()
+
+# Add custom handler with clean format including module and line number
+logger.add(
+    sys.stderr,
+    format="<green>{time:HH:mm:ss}</green> | <cyan>{module:>16}:{line}</cyan> | <level>{level: >8}</level> | <level>{message}</level>",
+    colorize=True,
+    level="INFO" # "DEBUG" to enable logger.debug("message") and up prints 
+                 # "ERROR" to enable only logger.error("message") prints
+                 # etc
+)
+
+# Disable before release or as needed
+logger.disable("kokoro")
+
+from .model import KModel
+from .pipeline import KPipeline

kokoro/__main__.py

@@ -0,0 +1,148 @@
+"""Kokoro TTS CLI
+Example usage:
+python3 -m kokoro --text "The sky above the port was the color of television, tuned to a dead channel." -o file.wav --debug
+
+echo "Bom dia mundo, como vão vocês" > text.txt
+python3 -m kokoro -i text.txt -l p --voice pm_alex > audio.wav
+
+Common issues:
+pip not installed: `uv pip install pip`
+(Temporary workaround while https://github.com/explosion/spaCy/issues/13747 is not fixed)
+
+espeak not installed: `apt-get install espeak-ng`
+"""
+
+import argparse
+import wave
+from pathlib import Path
+from typing import Generator, TYPE_CHECKING
+
+import numpy as np
+from loguru import logger
+
+languages = [
+    "a",  # American English
+    "b",  # British English
+    "h",  # Hindi
+    "e",  # Spanish
+    "f",  # French
+    "i",  # Italian
+    "p",  # Brazilian Portuguese
+    "j",  # Japanese
+    "z",  # Mandarin Chinese
+]
+
+if TYPE_CHECKING:
+    from kokoro import KPipeline
+
+
+def generate_audio(
+    text: str, kokoro_language: str, voice: str, speed=1
+) -> Generator["KPipeline.Result", None, None]:
+    from kokoro import KPipeline
+
+    if not voice.startswith(kokoro_language):
+        logger.warning(f"Voice {voice} is not made for language {kokoro_language}")
+    pipeline = KPipeline(lang_code=kokoro_language)
+    yield from pipeline(text, voice=voice, speed=speed, split_pattern=r"\n+")
+
+
+def generate_and_save_audio(
+    output_file: Path, text: str, kokoro_language: str, voice: str, speed=1
+) -> None:
+    with wave.open(str(output_file.resolve()), "wb") as wav_file:
+        wav_file.setnchannels(1)  # Mono audio
+        wav_file.setsampwidth(2)  # 2 bytes per sample (16-bit audio)
+        wav_file.setframerate(24000)  # Sample rate
+
+        for result in generate_audio(
+            text, kokoro_language=kokoro_language, voice=voice, speed=speed
+        ):
+            logger.debug(result.phonemes)
+            if result.audio is None:
+                continue
+            audio_bytes = (result.audio.numpy() * 32767).astype(np.int16).tobytes()
+            wav_file.writeframes(audio_bytes)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-m",
+        "--voice",
+        default="af_heart",
+        help="Voice to use",
+    )
+    parser.add_argument(
+        "-l",
+        "--language",
+        help="Language to use (defaults to the one corresponding to the voice)",
+        choices=languages,
+    )
+    parser.add_argument(
+        "-o",
+        "--output-file",
+        "--output_file",
+        type=Path,
+        help="Path to output WAV file",
+        required=True,
+    )
+    parser.add_argument(
+        "-i",
+        "--input-file",
+        "--input_file",
+        type=Path,
+        help="Path to input text file (default: stdin)",
+    )
+    parser.add_argument(
+        "-t",
+        "--text",
+        help="Text to use instead of reading from stdin",
+    )
+    parser.add_argument(
+        "-s",
+        "--speed",
+        type=float,
+        default=1.0,
+        help="Speech speed",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Print DEBUG messages to console",
+    )
+    args = parser.parse_args()
+    if args.debug:
+        logger.level("DEBUG")
+    logger.debug(args)
+
+    lang = args.language or args.voice[0]
+
+    if args.text is not None and args.input_file is not None:
+        raise Exception("You cannot specify both 'text' and 'input_file'")
+    elif args.text:
+        text = args.text
+    elif args.input_file:
+        file: Path = args.input_file
+        text = file.read_text()
+    else:
+        import sys
+        print("Press Ctrl+D to stop reading input and start generating", flush=True)
+        text = '\n'.join(sys.stdin)
+
+    logger.debug(f"Input text: {text!r}")
+
+    out_file: Path = args.output_file
+    if not out_file.suffix == ".wav":
+        logger.warning("The output file name should end with .wav")
+    generate_and_save_audio(
+        output_file=out_file,
+        text=text,
+        kokoro_language=lang,
+        voice=args.voice,
+        speed=args.speed,
+    )
+
+
+if __name__ == "__main__":
+    main()

kokoro/custom_stft.py

@@ -0,0 +1,197 @@
+from attr import attr
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CustomSTFT(nn.Module):
+    """
+    STFT/iSTFT without unfold/complex ops, using conv1d and conv_transpose1d.
+
+    - forward STFT => Real-part conv1d + Imag-part conv1d
+    - inverse STFT => Real-part conv_transpose1d + Imag-part conv_transpose1d + sum
+    - avoids F.unfold, so easier to export to ONNX
+    - uses replicate or constant padding for 'center=True' to approximate 'reflect' 
+      (reflect is not supported for dynamic shapes in ONNX)
+    """
+
+    def __init__(
+        self,
+        filter_length=800,
+        hop_length=200,
+        win_length=800,
+        window="hann",
+        center=True,
+        pad_mode="replicate",  # or 'constant'
+    ):
+        super().__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_fft = filter_length
+        self.center = center
+        self.pad_mode = pad_mode
+
+        # Number of frequency bins for real-valued STFT with onesided=True
+        self.freq_bins = self.n_fft // 2 + 1
+
+        # Build window
+        assert window == 'hann', window
+        window_tensor = torch.hann_window(win_length, periodic=True, dtype=torch.float32)
+        if self.win_length < self.n_fft:
+            # Zero-pad up to n_fft
+            extra = self.n_fft - self.win_length
+            window_tensor = F.pad(window_tensor, (0, extra))
+        elif self.win_length > self.n_fft:
+            window_tensor = window_tensor[: self.n_fft]
+        self.register_buffer("window", window_tensor)
+
+        # Precompute forward DFT (real, imag)
+        # PyTorch stft uses e^{-j 2 pi k n / N} => real=cos(...), imag=-sin(...)
+        n = np.arange(self.n_fft)
+        k = np.arange(self.freq_bins)
+        angle = 2 * np.pi * np.outer(k, n) / self.n_fft  # shape (freq_bins, n_fft)
+        dft_real = np.cos(angle)
+        dft_imag = -np.sin(angle)  # note negative sign
+
+        # Combine window and dft => shape (freq_bins, filter_length)
+        # We'll make 2 conv weight tensors of shape (freq_bins, 1, filter_length).
+        forward_window = window_tensor.numpy()  # shape (n_fft,)
+        forward_real = dft_real * forward_window  # (freq_bins, n_fft)
+        forward_imag = dft_imag * forward_window
+
+        # Convert to PyTorch
+        forward_real_torch = torch.from_numpy(forward_real).float()
+        forward_imag_torch = torch.from_numpy(forward_imag).float()
+
+        # Register as Conv1d weight => (out_channels, in_channels, kernel_size)
+        # out_channels = freq_bins, in_channels=1, kernel_size=n_fft
+        self.register_buffer(
+            "weight_forward_real", forward_real_torch.unsqueeze(1)
+        )
+        self.register_buffer(
+            "weight_forward_imag", forward_imag_torch.unsqueeze(1)
+        )
+
+        # Precompute inverse DFT
+        # Real iFFT formula => scale = 1/n_fft, doubling for bins 1..freq_bins-2 if n_fft even, etc.
+        # For simplicity, we won't do the "DC/nyquist not doubled" approach here. 
+        # If you want perfect real iSTFT, you can add that logic. 
+        # This version just yields good approximate reconstruction with Hann + typical overlap.
+        inv_scale = 1.0 / self.n_fft
+        n = np.arange(self.n_fft)
+        angle_t = 2 * np.pi * np.outer(n, k) / self.n_fft  # shape (n_fft, freq_bins)
+        idft_cos = np.cos(angle_t).T  # => (freq_bins, n_fft)
+        idft_sin = np.sin(angle_t).T  # => (freq_bins, n_fft)
+
+        # Multiply by window again for typical overlap-add
+        # We also incorporate the scale factor 1/n_fft
+        inv_window = window_tensor.numpy() * inv_scale
+        backward_real = idft_cos * inv_window  # (freq_bins, n_fft)
+        backward_imag = idft_sin * inv_window
+
+        # We'll implement iSTFT as real+imag conv_transpose with stride=hop.
+        self.register_buffer(
+            "weight_backward_real", torch.from_numpy(backward_real).float().unsqueeze(1)
+        )
+        self.register_buffer(
+            "weight_backward_imag", torch.from_numpy(backward_imag).float().unsqueeze(1)
+        )
+        
+
+
+    def transform(self, waveform: torch.Tensor):
+        """
+        Forward STFT => returns magnitude, phase
+        Output shape => (batch, freq_bins, frames)
+        """
+        # waveform shape => (B, T).  conv1d expects (B, 1, T).
+        # Optional center pad
+        if self.center:
+            pad_len = self.n_fft // 2
+            waveform = F.pad(waveform, (pad_len, pad_len), mode=self.pad_mode)
+
+        x = waveform.unsqueeze(1)  # => (B, 1, T)
+        # Convolution to get real part => shape (B, freq_bins, frames)
+        real_out = F.conv1d(
+            x,
+            self.weight_forward_real,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        # Imag part
+        imag_out = F.conv1d(
+            x,
+            self.weight_forward_imag,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        # magnitude, phase
+        magnitude = torch.sqrt(real_out**2 + imag_out**2 + 1e-14)
+        phase = torch.atan2(imag_out, real_out)
+        # Handle the case where imag_out is 0 and real_out is negative to correct ONNX atan2 to match PyTorch
+        # In this case, PyTorch returns pi, ONNX returns -pi
+        correction_mask = (imag_out == 0) & (real_out < 0)
+        phase[correction_mask] = torch.pi
+        return magnitude, phase
+
+
+    def inverse(self, magnitude: torch.Tensor, phase: torch.Tensor, length=None):
+        """
+        Inverse STFT => returns waveform shape (B, T).
+        """
+        # magnitude, phase => (B, freq_bins, frames)
+        # Re-create real/imag => shape (B, freq_bins, frames)
+        real_part = magnitude * torch.cos(phase)
+        imag_part = magnitude * torch.sin(phase)
+
+        # conv_transpose wants shape (B, freq_bins, frames). We'll treat "frames" as time dimension
+        # so we do (B, freq_bins, frames) => (B, freq_bins, frames)
+        # But PyTorch conv_transpose1d expects (B, in_channels, input_length)
+        real_part = real_part  # (B, freq_bins, frames)
+        imag_part = imag_part
+
+        # real iSTFT => convolve with "backward_real", "backward_imag", and sum
+        # We'll do 2 conv_transpose calls, each giving (B, 1, time),
+        # then add them => (B, 1, time).
+        real_rec = F.conv_transpose1d(
+            real_part,
+            self.weight_backward_real,  # shape (freq_bins, 1, filter_length)
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        imag_rec = F.conv_transpose1d(
+            imag_part,
+            self.weight_backward_imag,
+            bias=None,
+            stride=self.hop_length,
+            padding=0,
+        )
+        # sum => (B, 1, time)
+        waveform = real_rec - imag_rec  # typical real iFFT has minus for imaginary part
+
+        # If we used "center=True" in forward, we should remove pad
+        if self.center:
+            pad_len = self.n_fft // 2
+            # Because of transposed convolution, total length might have extra samples
+            # We remove `pad_len` from start & end if possible
+            waveform = waveform[..., pad_len:-pad_len]
+
+        # If a specific length is desired, clamp
+        if length is not None:
+            waveform = waveform[..., :length]
+
+        # shape => (B, T)
+        return waveform
+
+    def forward(self, x: torch.Tensor):
+        """
+        Full STFT -> iSTFT pass: returns time-domain reconstruction.
+        Same interface as your original code.
+        """
+        mag, phase = self.transform(x)
+        return self.inverse(mag, phase, length=x.shape[-1])

kokoro/istftnet.py

@@ -0,0 +1,421 @@
+# ADAPTED from https://github.com/yl4579/StyleTTS2/blob/main/Modules/istftnet.py
+from kokoro.custom_stft import CustomSTFT
+from torch.nn.utils.parametrizations import weight_norm
+import math
+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)
+
+
+class AdaIN1d(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        # affine should be False, however there's a bug in the old torch.onnx.export (not newer dynamo) that causes the channel dimension to be lost if affine=False. When affine is true, there's additional learnably parameters. This shouldn't really matter setting it to True, since we're in inference mode
+        self.norm = nn.InstanceNorm1d(num_features, affine=True)
+        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(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(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                                  padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                                  padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(nn.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(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1))),
+            weight_norm(nn.Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                                  padding=get_padding(kernel_size, 1))),
+            weight_norm(nn.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
+
+
+class TorchSTFT(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
+        assert window == 'hann', window
+        self.window = torch.hann_window(win_length, periodic=True, dtype=torch.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(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(torch.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 * torch.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:
+            rad_values = F.interpolate(rad_values.transpose(1, 2), scale_factor=1/self.upsample_scale, mode="linear").transpose(1, 2)
+            phase = torch.cumsum(rad_values, dim=1) * 2 * torch.pi
+            phase = F.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 * torch.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(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 = nn.Linear(harmonic_num + 1, 1)
+        self.l_tanh = 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
+
+
+class Generator(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, disable_complex=False):
+        super(Generator, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.m_source = SourceModuleHnNSF(
+                    sampling_rate=24000,
+                    upsample_scale=math.prod(upsample_rates) * gen_istft_hop_size,
+                    harmonic_num=8, voiced_threshod=10)
+        self.f0_upsamp = nn.Upsample(scale_factor=math.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(
+                nn.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(AdaINResBlock1(ch, k, d, style_dim))
+            c_cur = upsample_initial_channel // (2 ** (i + 1))
+            if i + 1 < len(upsample_rates):
+                stride_f0 = math.prod(upsample_rates[i + 1:])
+                self.noise_convs.append(nn.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(AdaINResBlock1(c_cur, 7, [1,3,5], style_dim))
+            else:
+                self.noise_convs.append(nn.Conv1d(gen_istft_n_fft + 2, c_cur, kernel_size=1))
+                self.noise_res.append(AdaINResBlock1(c_cur, 11, [1,3,5], style_dim))
+        self.post_n_fft = gen_istft_n_fft
+        self.conv_post = weight_norm(nn.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 = nn.ReflectionPad1d((1, 0))
+        self.stft = (
+            CustomSTFT(filter_length=gen_istft_n_fft, hop_length=gen_istft_hop_size, win_length=gen_istft_n_fft)
+            if disable_complex
+            else 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, negative_slope=0.1) 
+            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)
+
+
+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)) * torch.rsqrt(torch.tensor(2))
+        return out
+
+
+class Decoder(nn.Module):
+    def __init__(self, dim_in, style_dim, dim_out, 
+                 resblock_kernel_sizes,
+                 upsample_rates,
+                 upsample_initial_channel,
+                 resblock_dilation_sizes,
+                 upsample_kernel_sizes,
+                 gen_istft_n_fft, gen_istft_hop_size,
+                 disable_complex=False):
+        super().__init__()
+        self.encode = AdainResBlk1d(dim_in + 2, 1024, style_dim)
+        self.decode = nn.ModuleList()
+        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, disable_complex=disable_complex)
+
+    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/model.py

@@ -0,0 +1,151 @@
+from .istftnet import Decoder
+from .modules import CustomAlbert, ProsodyPredictor, TextEncoder
+from dataclasses import dataclass
+from huggingface_hub import hf_hub_download
+from loguru import logger
+from transformers import AlbertConfig
+from typing import Dict, Optional, Union
+import json
+import torch
+
+class KModel(torch.nn.Module):
+    '''
+    KModel is a torch.nn.Module with 2 main responsibilities:
+    1. Init weights, downloading config.json + model.pth from HF if needed
+    2. forward(phonemes: str, ref_s: FloatTensor) -> (audio: FloatTensor)
+
+    You likely only need one KModel instance, and it can be reused across
+    multiple KPipelines to avoid redundant memory allocation.
+
+    Unlike KPipeline, KModel is language-blind.
+
+    KModel stores self.vocab and thus knows how to map phonemes -> input_ids,
+    so there is no need to repeatedly download config.json outside of KModel.
+    '''
+
+    MODEL_NAMES = {
+        'hexgrad/Kokoro-82M': 'kokoro-v1_0.pth',
+        'hexgrad/Kokoro-82M-v1.1-zh': 'kokoro-v1_1-zh.pth',
+    }
+
+    def __init__(
+        self,
+        repo_id: Optional[str] = None,
+        config: Union[Dict, str, None] = None,
+        model: Optional[str] = None,
+        disable_complex: bool = False
+    ):
+        super().__init__()
+        if repo_id is None:
+            repo_id = 'hexgrad/Kokoro-82M'
+            print(f"WARNING: Defaulting repo_id to {repo_id}. Pass repo_id='{repo_id}' to suppress this warning.")
+        self.repo_id = repo_id
+        if not isinstance(config, dict):
+            if not config:
+                logger.debug("No config provided, downloading from HF")
+                config = hf_hub_download(repo_id=repo_id, filename='config.json')
+            with open(config, 'r', encoding='utf-8') as r:
+                config = json.load(r)
+                logger.debug(f"Loaded config: {config}")
+        self.vocab = config['vocab']
+        self.bert = CustomAlbert(AlbertConfig(vocab_size=config['n_token'], **config['plbert']))
+        self.bert_encoder = torch.nn.Linear(self.bert.config.hidden_size, config['hidden_dim'])
+        self.context_length = self.bert.config.max_position_embeddings
+        self.predictor = ProsodyPredictor(
+            style_dim=config['style_dim'], d_hid=config['hidden_dim'],
+            nlayers=config['n_layer'], max_dur=config['max_dur'], dropout=config['dropout']
+        )
+        self.text_encoder = TextEncoder(
+            channels=config['hidden_dim'], kernel_size=config['text_encoder_kernel_size'],
+            depth=config['n_layer'], n_symbols=config['n_token']
+        )
+        self.decoder = Decoder(
+            dim_in=config['hidden_dim'], style_dim=config['style_dim'],
+            dim_out=config['n_mels'], disable_complex=disable_complex, **config['istftnet']
+        )
+        if not model:
+            model = hf_hub_download(repo_id=repo_id, filename=KModel.MODEL_NAMES[repo_id])
+        for key, state_dict in torch.load(model, map_location='cpu', weights_only=True).items():
+            assert hasattr(self, key), key
+            try:
+                getattr(self, key).load_state_dict(state_dict)
+            except:
+                logger.debug(f"Did not load {key} from state_dict")
+                state_dict = {k[7:]: v for k, v in state_dict.items()}
+                getattr(self, key).load_state_dict(state_dict, strict=False)
+
+    @property
+    def device(self):
+        return self.bert.device
+
+    @dataclass
+    class Output:
+        audio: torch.FloatTensor
+        pred_dur: Optional[torch.LongTensor] = None
+
+    @torch.no_grad()
+    def forward_with_tokens(
+        self,
+        input_ids: torch.LongTensor,
+        ref_s: torch.FloatTensor,
+        speed: float = 1
+    ) -> tuple[torch.FloatTensor, torch.LongTensor]:
+        input_lengths = torch.full(
+            (input_ids.shape[0],), 
+            input_ids.shape[-1], 
+            device=input_ids.device,
+            dtype=torch.long
+        )
+
+        text_mask = torch.arange(input_lengths.max()).unsqueeze(0).expand(input_lengths.shape[0], -1).type_as(input_lengths)
+        text_mask = torch.gt(text_mask+1, input_lengths.unsqueeze(1)).to(self.device)
+        bert_dur = self.bert(input_ids, attention_mask=(~text_mask).int())
+        d_en = self.bert_encoder(bert_dur).transpose(-1, -2)
+        s = ref_s[:, 128:]
+        d = self.predictor.text_encoder(d_en, s, input_lengths, text_mask)
+        x, _ = self.predictor.lstm(d)
+        duration = self.predictor.duration_proj(x)
+        duration = torch.sigmoid(duration).sum(axis=-1) / speed
+        pred_dur = torch.round(duration).clamp(min=1).long().squeeze()
+        indices = torch.repeat_interleave(torch.arange(input_ids.shape[1], device=self.device), pred_dur)
+        pred_aln_trg = torch.zeros((input_ids.shape[1], indices.shape[0]), device=self.device)
+        pred_aln_trg[indices, torch.arange(indices.shape[0])] = 1
+        pred_aln_trg = pred_aln_trg.unsqueeze(0).to(self.device)
+        en = d.transpose(-1, -2) @ pred_aln_trg
+        F0_pred, N_pred = self.predictor.F0Ntrain(en, s)
+        t_en = self.text_encoder(input_ids, input_lengths, text_mask)
+        asr = t_en @ pred_aln_trg
+        audio = self.decoder(asr, F0_pred, N_pred, ref_s[:, :128]).squeeze()
+        return audio, pred_dur
+
+    def forward(
+        self,
+        phonemes: str,
+        ref_s: torch.FloatTensor,
+        speed: float = 1,
+        return_output: bool = False
+    ) -> Union['KModel.Output', torch.FloatTensor]:
+        input_ids = list(filter(lambda i: i is not None, map(lambda p: self.vocab.get(p), phonemes)))
+        logger.debug(f"phonemes: {phonemes} -> input_ids: {input_ids}")
+        assert len(input_ids)+2 <= self.context_length, (len(input_ids)+2, self.context_length)
+        input_ids = torch.LongTensor([[0, *input_ids, 0]]).to(self.device)
+        ref_s = ref_s.to(self.device)
+        audio, pred_dur = self.forward_with_tokens(input_ids, ref_s, speed)
+        audio = audio.squeeze().cpu()
+        pred_dur = pred_dur.cpu() if pred_dur is not None else None
+        logger.debug(f"pred_dur: {pred_dur}")
+        return self.Output(audio=audio, pred_dur=pred_dur) if return_output else audio
+
+class KModelForONNX(torch.nn.Module):
+    def __init__(self, kmodel: KModel):
+        super().__init__()
+        self.kmodel = kmodel
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        ref_s: torch.FloatTensor,
+        speed: float = 1
+    ) -> tuple[torch.FloatTensor, torch.LongTensor]:
+        waveform, duration = self.kmodel.forward_with_tokens(input_ids, ref_s, speed)
+        return waveform, duration

kokoro/modules.py

@@ -0,0 +1,183 @@
+# https://github.com/yl4579/StyleTTS2/blob/main/models.py
+from .istftnet import AdainResBlk1d
+from torch.nn.utils.parametrizations import weight_norm
+from transformers import AlbertModel
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LinearNorm(nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, w_init_gain='linear'):
+        super(LinearNorm, self).__init__()
+        self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)
+        nn.init.xavier_uniform_(self.linear_layer.weight, gain=nn.init.calculate_gain(w_init_gain))
+
+    def forward(self, x):
+        return self.linear_layer(x)
+
+
+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.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.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]
+        lengths = input_lengths if input_lengths.device == torch.device('cpu') else input_lengths.to('cpu')
+        x = nn.utils.rnn.pack_padded_sequence(x, 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]], device=x.device)
+        x_pad[:, :, :x.shape[-1]] = x
+        x = x_pad
+        x.masked_fill_(m, 0.0)
+        return x
+
+
+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)
+        m = m.unsqueeze(1)
+        lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu')
+        x = nn.utils.rnn.pack_padded_sequence(d, 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_pad = torch.zeros([x.shape[0], m.shape[-1], x.shape[-1]], device=x.device)
+        x_pad[:, :x.shape[1], :] = x
+        x = x_pad
+        duration = self.duration_proj(nn.functional.dropout(x, 0.5, training=False))
+        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)
+
+
+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))
+            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
+        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)
+        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, 2, 0)], axis=1)
+                x.masked_fill_(masks.unsqueeze(-1).transpose(-1, -2), 0.0)
+            else:
+                lengths = text_lengths if text_lengths.device == torch.device('cpu') else text_lengths.to('cpu')
+                x = x.transpose(-1, -2)
+                x = nn.utils.rnn.pack_padded_sequence(
+                    x, 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=False)
+                x = x.transpose(-1, -2)
+                x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device)
+                x_pad[:, :, :x.shape[-1]] = x
+                x = x_pad
+
+        return x.transpose(-1, -2)
+
+
+# https://github.com/yl4579/StyleTTS2/blob/main/Utils/PLBERT/util.py
+class CustomAlbert(AlbertModel):
+    def forward(self, *args, **kwargs):
+        outputs = super().forward(*args, **kwargs)
+        return outputs.last_hidden_state

kokoro/pipeline.py

@@ -0,0 +1,442 @@
+from .model import KModel
+from dataclasses import dataclass
+from huggingface_hub import hf_hub_download
+from loguru import logger
+from misaki import en, espeak
+from typing import Callable, Generator, List, Optional, Tuple, Union
+import re
+import torch
+import os
+
+ALIASES = {
+    'en-us': 'a',
+    'en-gb': 'b',
+    'es': 'e',
+    'fr-fr': 'f',
+    'hi': 'h',
+    'it': 'i',
+    'pt-br': 'p',
+    'ja': 'j',
+    'zh': 'z',
+}
+
+LANG_CODES = dict(
+    # pip install misaki[en]
+    a='American English',
+    b='British English',
+
+    # espeak-ng
+    e='es',
+    f='fr-fr',
+    h='hi',
+    i='it',
+    p='pt-br',
+
+    # pip install misaki[ja]
+    j='Japanese',
+
+    # pip install misaki[zh]
+    z='Mandarin Chinese',
+)
+
+class KPipeline:
+    '''
+    KPipeline is a language-aware support class with 2 main responsibilities:
+    1. Perform language-specific G2P, mapping (and chunking) text -> phonemes
+    2. Manage and store voices, lazily downloaded from HF if needed
+
+    You are expected to have one KPipeline per language. If you have multiple
+    KPipelines, you should reuse one KModel instance across all of them.
+
+    KPipeline is designed to work with a KModel, but this is not required.
+    There are 2 ways to pass an existing model into a pipeline:
+    1. On init: us_pipeline = KPipeline(lang_code='a', model=model)
+    2. On call: us_pipeline(text, voice, model=model)
+
+    By default, KPipeline will automatically initialize its own KModel. To
+    suppress this, construct a "quiet" KPipeline with model=False.
+
+    A "quiet" KPipeline yields (graphemes, phonemes, None) without generating
+    any audio. You can use this to phonemize and chunk your text in advance.
+
+    A "loud" KPipeline _with_ a model yields (graphemes, phonemes, audio).
+    '''
+    def __init__(
+        self,
+        lang_code: str,
+        repo_id: Optional[str] = None,
+        model: Union[KModel, bool] = True,
+        trf: bool = False,
+        en_callable: Optional[Callable[[str], str]] = None,
+        device: Optional[str] = None
+    ):
+        """Initialize a KPipeline.
+        
+        Args:
+            lang_code: Language code for G2P processing
+            model: KModel instance, True to create new model, False for no model
+            trf: Whether to use transformer-based G2P
+            device: Override default device selection ('cuda' or 'cpu', or None for auto)
+                   If None, will auto-select cuda if available
+                   If 'cuda' and not available, will explicitly raise an error
+        """
+        if repo_id is None:
+            repo_id = 'hexgrad/Kokoro-82M'
+            print(f"WARNING: Defaulting repo_id to {repo_id}. Pass repo_id='{repo_id}' to suppress this warning.")
+        self.repo_id = repo_id
+        lang_code = lang_code.lower()
+        lang_code = ALIASES.get(lang_code, lang_code)
+        assert lang_code in LANG_CODES, (lang_code, LANG_CODES)
+        self.lang_code = lang_code
+        self.model = None
+        if isinstance(model, KModel):
+            self.model = model
+        elif model:
+            if device == 'cuda' and not torch.cuda.is_available():
+                raise RuntimeError("CUDA requested but not available")
+            if device == 'mps' and not torch.backends.mps.is_available():
+                raise RuntimeError("MPS requested but not available")
+            if device == 'mps' and os.environ.get('PYTORCH_ENABLE_MPS_FALLBACK') != '1':
+                raise RuntimeError("MPS requested but fallback not enabled")
+            if device is None:
+                if torch.cuda.is_available():
+                    device = 'cuda'
+                elif os.environ.get('PYTORCH_ENABLE_MPS_FALLBACK') == '1' and torch.backends.mps.is_available():
+                    device = 'mps'
+                else:
+                    device = 'cpu'
+            try:
+                self.model = KModel(repo_id=repo_id).to(device).eval()
+            except RuntimeError as e:
+                if device == 'cuda':
+                    raise RuntimeError(f"""Failed to initialize model on CUDA: {e}. 
+                                       Try setting device='cpu' or check CUDA installation.""")
+                raise
+        self.voices = {}
+        if lang_code in 'ab':
+            try:
+                fallback = espeak.EspeakFallback(british=lang_code=='b')
+            except Exception as e:
+                logger.warning("EspeakFallback not Enabled: OOD words will be skipped")
+                logger.warning({str(e)})
+                fallback = None
+            self.g2p = en.G2P(trf=trf, british=lang_code=='b', fallback=fallback, unk='')
+        elif lang_code == 'j':
+            try:
+                from misaki import ja
+                self.g2p = ja.JAG2P()
+            except ImportError:
+                logger.error("You need to `pip install misaki[ja]` to use lang_code='j'")
+                raise
+        elif lang_code == 'z':
+            try:
+                from misaki import zh
+                self.g2p = zh.ZHG2P(
+                    version=None if repo_id.endswith('/Kokoro-82M') else '1.1',
+                    en_callable=en_callable
+                )
+            except ImportError:
+                logger.error("You need to `pip install misaki[zh]` to use lang_code='z'")
+                raise
+        else:
+            language = LANG_CODES[lang_code]
+            logger.warning(f"Using EspeakG2P(language='{language}'). Chunking logic not yet implemented, so long texts may be truncated unless you split them with '\\n'.")
+            self.g2p = espeak.EspeakG2P(language=language)
+
+    def load_single_voice(self, voice: str):
+        if voice in self.voices:
+            return self.voices[voice]
+        if voice.endswith('.pt'):
+            f = voice
+        else:
+            f = hf_hub_download(repo_id=self.repo_id, filename=f'voices/{voice}.pt')
+            if not voice.startswith(self.lang_code):
+                v = LANG_CODES.get(voice, voice)
+                p = LANG_CODES.get(self.lang_code, self.lang_code)
+                logger.warning(f'Language mismatch, loading {v} voice into {p} pipeline.')
+        pack = torch.load(f, weights_only=True)
+        self.voices[voice] = pack
+        return pack
+
+    """
+    load_voice is a helper function that lazily downloads and loads a voice:
+    Single voice can be requested (e.g. 'af_bella') or multiple voices (e.g. 'af_bella,af_jessica').
+    If multiple voices are requested, they are averaged.
+    Delimiter is optional and defaults to ','.
+    """
+    def load_voice(self, voice: Union[str, torch.FloatTensor], delimiter: str = ",") -> torch.FloatTensor:
+        if isinstance(voice, torch.FloatTensor):
+            return voice
+        if voice in self.voices:
+            return self.voices[voice]
+        logger.debug(f"Loading voice: {voice}")
+        packs = [self.load_single_voice(v) for v in voice.split(delimiter)]
+        if len(packs) == 1:
+            return packs[0]
+        self.voices[voice] = torch.mean(torch.stack(packs), dim=0)
+        return self.voices[voice]
+
+    @staticmethod
+    def tokens_to_ps(tokens: List[en.MToken]) -> str:
+        return ''.join(t.phonemes + (' ' if t.whitespace else '') for t in tokens).strip()
+
+    @staticmethod
+    def waterfall_last(
+        tokens: List[en.MToken],
+        next_count: int,
+        waterfall: List[str] = ['!.?…', ':;', ',—'],
+        bumps: List[str] = [')', '”']
+    ) -> int:
+        for w in waterfall:
+            z = next((i for i, t in reversed(list(enumerate(tokens))) if t.phonemes in set(w)), None)
+            if z is None:
+                continue
+            z += 1
+            if z < len(tokens) and tokens[z].phonemes in bumps:
+                z += 1
+            if next_count - len(KPipeline.tokens_to_ps(tokens[:z])) <= 510:
+                return z
+        return len(tokens)
+
+    @staticmethod
+    def tokens_to_text(tokens: List[en.MToken]) -> str:
+        return ''.join(t.text + t.whitespace for t in tokens).strip()
+
+    def en_tokenize(
+        self,
+        tokens: List[en.MToken]
+    ) -> Generator[Tuple[str, str, List[en.MToken]], None, None]:
+        tks = []
+        pcount = 0
+        for t in tokens:
+            # American English: ɾ => T
+            t.phonemes = '' if t.phonemes is None else t.phonemes#.replace('ɾ', 'T')
+            next_ps = t.phonemes + (' ' if t.whitespace else '')
+            next_pcount = pcount + len(next_ps.rstrip())
+            if next_pcount > 510:
+                z = KPipeline.waterfall_last(tks, next_pcount)
+                text = KPipeline.tokens_to_text(tks[:z])
+                logger.debug(f"Chunking text at {z}: '{text[:30]}{'...' if len(text) > 30 else ''}'")
+                ps = KPipeline.tokens_to_ps(tks[:z])
+                yield text, ps, tks[:z]
+                tks = tks[z:]
+                pcount = len(KPipeline.tokens_to_ps(tks))
+                if not tks:
+                    next_ps = next_ps.lstrip()
+            tks.append(t)
+            pcount += len(next_ps)
+        if tks:
+            text = KPipeline.tokens_to_text(tks)
+            ps = KPipeline.tokens_to_ps(tks)
+            yield ''.join(text).strip(), ''.join(ps).strip(), tks
+
+    @staticmethod
+    def infer(
+        model: KModel,
+        ps: str,
+        pack: torch.FloatTensor,
+        speed: Union[float, Callable[[int], float]] = 1
+    ) -> KModel.Output:
+        if callable(speed):
+            speed = speed(len(ps))
+        return model(ps, pack[len(ps)-1], speed, return_output=True)
+
+    def generate_from_tokens(
+        self,
+        tokens: Union[str, List[en.MToken]],
+        voice: str,
+        speed: float = 1,
+        model: Optional[KModel] = None
+    ) -> Generator['KPipeline.Result', None, None]:
+        """Generate audio from either raw phonemes or pre-processed tokens.
+        
+        Args:
+            tokens: Either a phoneme string or list of pre-processed MTokens
+            voice: The voice to use for synthesis
+            speed: Speech speed modifier (default: 1)
+            model: Optional KModel instance (uses pipeline's model if not provided)
+        
+        Yields:
+            KPipeline.Result containing the input tokens and generated audio
+            
+        Raises:
+            ValueError: If no voice is provided or token sequence exceeds model limits
+        """
+        model = model or self.model
+        if model and voice is None:
+            raise ValueError('Specify a voice: pipeline.generate_from_tokens(..., voice="af_heart")')
+        
+        pack = self.load_voice(voice).to(model.device) if model else None
+
+        # Handle raw phoneme string
+        if isinstance(tokens, str):
+            logger.debug("Processing phonemes from raw string")
+            if len(tokens) > 510:
+                raise ValueError(f'Phoneme string too long: {len(tokens)} > 510')
+            output = KPipeline.infer(model, tokens, pack, speed) if model else None
+            yield self.Result(graphemes='', phonemes=tokens, output=output)
+            return
+        
+        logger.debug("Processing MTokens")
+        # Handle pre-processed tokens
+        for gs, ps, tks in self.en_tokenize(tokens):
+            if not ps:
+                continue
+            elif len(ps) > 510:
+                logger.warning(f"Unexpected len(ps) == {len(ps)} > 510 and ps == '{ps}'")
+                logger.warning("Truncating to 510 characters")
+                ps = ps[:510]
+            output = KPipeline.infer(model, ps, pack, speed) if model else None
+            if output is not None and output.pred_dur is not None:
+                KPipeline.join_timestamps(tks, output.pred_dur)
+            yield self.Result(graphemes=gs, phonemes=ps, tokens=tks, output=output)
+
+    @staticmethod
+    def join_timestamps(tokens: List[en.MToken], pred_dur: torch.LongTensor):
+        # Multiply by 600 to go from pred_dur frames to sample_rate 24000
+        # Equivalent to dividing pred_dur frames by 40 to get timestamp in seconds
+        # We will count nice round half-frames, so the divisor is 80
+        MAGIC_DIVISOR = 80
+        if not tokens or len(pred_dur) < 3:
+            # We expect at least 3: <bos>, token, <eos>
+            return
+        # We track 2 counts, measured in half-frames: (left, right)
+        # This way we can cut space characters in half
+        # TODO: Is -3 an appropriate offset?
+        left = right = 2 * max(0, pred_dur[0].item() - 3)
+        # Updates:
+        # left = right + (2 * token_dur) + space_dur
+        # right = left + space_dur
+        i = 1
+        for t in tokens:
+            if i >= len(pred_dur)-1:
+                break
+            if not t.phonemes:
+                if t.whitespace:
+                    i += 1
+                    left = right + pred_dur[i].item()
+                    right = left + pred_dur[i].item()
+                    i += 1
+                continue
+            j = i + len(t.phonemes)
+            if j >= len(pred_dur):
+                break
+            t.start_ts = left / MAGIC_DIVISOR
+            token_dur = pred_dur[i: j].sum().item()
+            space_dur = pred_dur[j].item() if t.whitespace else 0
+            left = right + (2 * token_dur) + space_dur
+            t.end_ts = left / MAGIC_DIVISOR
+            right = left + space_dur
+            i = j + (1 if t.whitespace else 0)
+
+    @dataclass
+    class Result:
+        graphemes: str
+        phonemes: str
+        tokens: Optional[List[en.MToken]] = None
+        output: Optional[KModel.Output] = None
+        text_index: Optional[int] = None
+
+        @property
+        def audio(self) -> Optional[torch.FloatTensor]:
+            return None if self.output is None else self.output.audio
+
+        @property
+        def pred_dur(self) -> Optional[torch.LongTensor]:
+            return None if self.output is None else self.output.pred_dur
+
+        ### MARK: BEGIN BACKWARD COMPAT ###
+        def __iter__(self):
+            yield self.graphemes
+            yield self.phonemes
+            yield self.audio
+
+        def __getitem__(self, index):
+            return [self.graphemes, self.phonemes, self.audio][index]
+
+        def __len__(self):
+            return 3
+        #### MARK: END BACKWARD COMPAT ####
+
+    def __call__(
+        self,
+        text: Union[str, List[str]],
+        voice: Optional[str] = None,
+        speed: Union[float, Callable[[int], float]] = 1,
+        split_pattern: Optional[str] = r'\n+',
+        model: Optional[KModel] = None
+    ) -> Generator['KPipeline.Result', None, None]:
+        model = model or self.model
+        if model and voice is None:
+            raise ValueError('Specify a voice: en_us_pipeline(text="Hello world!", voice="af_heart")')
+        pack = self.load_voice(voice).to(model.device) if model else None
+        
+        # Convert input to list of segments
+        if isinstance(text, str):
+            text = re.split(split_pattern, text.strip()) if split_pattern else [text]
+            
+        # Process each segment
+        for graphemes_index, graphemes in enumerate(text):
+            if not graphemes.strip():  # Skip empty segments
+                continue
+                
+            # English processing (unchanged)
+            if self.lang_code in 'ab':
+                logger.debug(f"Processing English text: {graphemes[:50]}{'...' if len(graphemes) > 50 else ''}")
+                _, tokens = self.g2p(graphemes)
+                for gs, ps, tks in self.en_tokenize(tokens):
+                    if not ps:
+                        continue
+                    elif len(ps) > 510:
+                        logger.warning(f"Unexpected len(ps) == {len(ps)} > 510 and ps == '{ps}'")
+                        ps = ps[:510]
+                    output = KPipeline.infer(model, ps, pack, speed) if model else None
+                    if output is not None and output.pred_dur is not None:
+                        KPipeline.join_timestamps(tks, output.pred_dur)
+                    yield self.Result(graphemes=gs, phonemes=ps, tokens=tks, output=output, text_index=graphemes_index)
+            
+            # Non-English processing with chunking
+            else:
+                # Split long text into smaller chunks (roughly 400 characters each)
+                # Using sentence boundaries when possible
+                chunk_size = 400
+                chunks = []
+                
+                # Try to split on sentence boundaries first
+                sentences = re.split(r'([.!?]+)', graphemes)
+                current_chunk = ""
+                
+                for i in range(0, len(sentences), 2):
+                    sentence = sentences[i]
+                    # Add the punctuation back if it exists
+                    if i + 1 < len(sentences):
+                        sentence += sentences[i + 1]
+                        
+                    if len(current_chunk) + len(sentence) <= chunk_size:
+                        current_chunk += sentence
+                    else:
+                        if current_chunk:
+                            chunks.append(current_chunk.strip())
+                        current_chunk = sentence
+                
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                
+                # If no chunks were created (no sentence boundaries), fall back to character-based chunking
+                if not chunks:
+                    chunks = [graphemes[i:i+chunk_size] for i in range(0, len(graphemes), chunk_size)]
+                
+                # Process each chunk
+                for chunk in chunks:
+                    if not chunk.strip():
+                        continue
+                        
+                    ps, _ = self.g2p(chunk)
+                    if not ps:
+                        continue
+                    elif len(ps) > 510:
+                        logger.warning(f'Truncating len(ps) == {len(ps)} > 510')
+                        ps = ps[:510]
+                        
+                    output = KPipeline.infer(model, ps, pack, speed) if model else None
+                    yield self.Result(graphemes=chunk, phonemes=ps, output=output, text_index=graphemes_index)

pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "kokoro"
+version = "0.9.4"
+description = "TTS"
+readme = "README.md"
+authors = [
+    { name="hexgrad", email="[email protected]" }
+]
+license = { file = "LICENSE" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: OS Independent"
+]
+requires-python = ">=3.10, <3.14"
+dependencies = [
+    "fastapi>=0.121.3",
+    "huggingface_hub",
+    "loguru",
+    "misaki[en]>=0.9.4",
+    "numpy",
+    "torch",
+    "transformers",
+    "uvicorn>=0.38.0",
+]
+
+[project.scripts]
+kokoro = "kokoro.__main__:main"
+
+[tool.hatch.build.targets.wheel]
+only-include = ["kokoro"]
+only-packages = true
+
+[project.urls]
+Homepage = "https://github.com/hexgrad/kokoro"
+Repository = "https://github.com/hexgrad/kokoro"

run.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# Get the directory where the script is located
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
+
+# Enable MPS fallback for macOS Apple Silicon
+export PYTORCH_ENABLE_MPS_FALLBACK=1
+
+# Run the python script with uv using existing .venv
+cd "$SCRIPT_DIR" && VIRTUAL_ENV="$SCRIPT_DIR/.venv" uv run --no-sync talk.py "$@"

talk.py

@@ -0,0 +1,65 @@
+import sys
+import os
+import soundfile as sf
+from kokoro import KPipeline
+import subprocess
+import argparse
+
+def play_audio(file_path):
+    """Plays audio using the macOS afplay command."""
+    try:
+        subprocess.run(["afplay", file_path], check=True)
+    except FileNotFoundError:
+        print("Error: 'afplay' command not found. Are you on macOS?")
+    except Exception as e:
+        print(f"Error playing audio: {e}")
+
+def main():
+    parser = argparse.ArgumentParser(description='Kokoro TTS Generator')
+    parser.add_argument('text', nargs='*', help='Text to speak')
+    parser.add_argument('-l', '--lang', default='a', help='Language code (a=Am. English, b=Br. English, e=Spanish, f=French, h=Hindi, i=Italian, j=Japanese, p=Portuguese, z=Chinese)')
+    parser.add_argument('-v', '--voice', default='af_heart', help='Voice ID (default: af_heart)')
+    
+    args = parser.parse_args()
+
+    # Combine text arguments or read from stdin
+    if args.text:
+        text = " ".join(args.text)
+    else:
+        print("Enter text to speak (Ctrl+D to finish):")
+        try:
+            text = sys.stdin.read()
+        except KeyboardInterrupt:
+            print("\nExiting.")
+            sys.exit(0)
+
+    if not text.strip():
+        print("No text provided.")
+        return
+
+    # Initialize pipeline
+    try:
+        print(f"Initializing pipeline for language '{args.lang}'...")
+        pipeline = KPipeline(lang_code=args.lang)
+    except Exception as e:
+        print(f"Failed to initialize KPipeline: {e}")
+        print("Note: Japanese ('j') and Chinese ('z') require extra dependencies: pip install 'misaki[ja]' or 'misaki[zh]'")
+        sys.exit(1)
+    
+    print(f"Generating audio with voice '{args.voice}'...")
+    
+    # Generate audio
+    try:
+        generator = pipeline(text, voice=args.voice, speed=1, split_pattern=r'\n+')
+        
+        for i, (gs, ps, audio) in enumerate(generator):
+            filename = f'output_{i}.wav'
+            sf.write(filename, audio, 24000)
+            print(f"Playing segment {i}...")
+            play_audio(filename)
+            
+    except Exception as e:
+        print(f"Error during generation: {e}")
+
+if __name__ == "__main__":
+    main()