Update app.py

app.py

@@ -14,113 +14,24 @@ if not IS_DUPLICATE:
 
 CHAR_LIMIT = None if IS_DUPLICATE else 5000
 
-# Initialize models
+# --- Models: CPU & GPU ---
 models = {gpu: KModel().to('cuda' if gpu else 'cpu').eval() for gpu in [False] + ([True] if CUDA_AVAILABLE else [])}
 
-# Supported language codes from docs
-LANG_CODES = {
-    "🇺🇸 American English": "a",
-    "🇬🇧 British English": "b",
-    "🇯🇵 Japanese": "j",
-    "🇨🇳 Mandarin Chinese": "z",
-    "🇪🇸 Spanish": "e",
-    "🇫🇷 French": "f",
-    "🇮🇳 Hindi": "h",
-    "🇮🇹 Italian": "i",
-    "🇧🇷 Brazilian Portuguese": "p"
-}
-
-# Initialize all pipelines
-pipelines = {code: KPipeline(lang_code=code, model=False) for code in LANG_CODES.values()}
-
-# Voice dictionary grouped by language
-VOICES = {
-    "a": {  # American English
-        '🇺🇸 🚺 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',
-    },
-    "b": {  # British English
-        '🇬🇧 🚺 Emma': 'bf_emma',
-        '🇬🇧 🚺 Isabella': 'bf_isabella',
-        '🇬🇧 🚺 Alice': 'bf_alice',
-        '🇬🇧 🚺 Lily': 'bf_lily',
-        '🇬🇧 🚹 George': 'bm_george',
-        '🇬🇧 🚹 Fable': 'bm_fable',
-        '🇬🇧 🚹 Lewis': 'bm_lewis',
-        '🇬🇧 🚹 Daniel': 'bm_daniel',
-    },
-    "j": {  # Japanese
-        '🇯🇵 🚺 Alpha': 'jf_alpha',
-        '🇯🇵 🚺 Gongitsune': 'jf_gongitsune',
-        '🇯🇵 🚺 Nezumi': 'jf_nezumi',
-        '🇯🇵 🚺 Tebukuro': 'jf_tebukuro',
-        '🇯🇵 🚹 Kumo': 'jm_kumo',
-    },
-    "z": {  # Mandarin Chinese
-        '🇨🇳 🚺 Xiaobei': 'zf_xiaobei',
-        '🇨🇳 🚺 Xiaoni': 'zf_xiaoni',
-        '🇨🇳 🚺 Xiaoxiao': 'zf_xiaoxiao',
-        '🇨🇳 🚺 Xiaoyi': 'zf_xiaoyi',
-        '🇨🇳 🚹 Yunjian': 'zm_yunjian',
-        '🇨🇳 🚹 Yunxi': 'zm_yunxi',
-        '🇨🇳 🚹 Yunxia': 'zm_yunxia',
-        '🇨🇳 🚹 Yunyang': 'zm_yunyang',
-    },
-    "e": {  # Spanish
-        '🇪🇸 🚺 Dora': 'ef_dora',
-        '🇪🇸 🚹 Alex': 'em_alex',
-        '🇪🇸 🚹 Santa': 'em_santa',
-    },
-    "f": {  # French
-        '🇫🇷 🚺 Siwis': 'ff_siwis',
-    },
-    "h": {  # Hindi
-        '🇮🇳 🚺 Alpha': 'hf_alpha',
-        '🇮🇳 🚺 Beta': 'hf_beta',
-        '🇮🇳 🚹 Omega': 'hm_omega',
-        '🇮🇳 🚹 Psi': 'hm_psi',
-    },
-    "i": {  # Italian
-        '🇮🇹 🚺 Sara': 'if_sara',
-        '🇮🇹 🚹 Nicola': 'im_nicola',
-    },
-    "p": {  # Brazilian Portuguese
-        '🇧🇷 🚺 Dora': 'pf_dora',
-        '🇧🇷 🚹 Alex': 'pm_alex',
-        '🇧🇷 🚹 Santa': 'pm_santa',
-    }
-}
+# --- Pipelines: American English (a), British English (b), Hindi (h) ---
+pipelines = {lang_code: KPipeline(lang_code=lang_code, model=False) for lang_code in 'abh'}
 
-# Preload voices
-for lang_code, voices in VOICES.items():
-    for v in voices.values():
-        pipelines[lang_code].load_voice(v)
+# Example custom phoneme golds for English voices
+pipelines['a'].g2p.lexicon.golds['kokoro'] = 'kˈOkəɹO'
+pipelines['b'].g2p.lexicon.golds['kokoro'] = 'kˈQkəɹQ'
+# Hindi pipeline can have custom entries if needed
 
 @spaces.GPU(duration=30)
 def forward_gpu(ps, ref_s, speed):
     return models[True](ps, ref_s, speed)
 
-def generate_first(text, lang, voice, speed=1, use_gpu=CUDA_AVAILABLE):
+def generate_first(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
     text = text if CHAR_LIMIT is None else text.strip()[:CHAR_LIMIT]
-    pipeline = pipelines[lang]
+    pipeline = pipelines[voice[0]]
     pack = pipeline.load_voice(voice)
     use_gpu = use_gpu and CUDA_AVAILABLE
     for _, ps, _ in pipeline(text, voice, speed):
@@ -133,22 +44,26 @@ def generate_first(text, lang, voice, speed=1, use_gpu=CUDA_AVAILABLE):
         except gr.exceptions.Error as e:
             if use_gpu:
                 gr.Warning(str(e))
-                gr.Info('Retrying with CPU.')
+                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, ''
 
-def tokenize_first(text, lang, voice):
-    pipeline = pipelines[lang]
+# 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, lang, voice, speed=1, use_gpu=CUDA_AVAILABLE):
+def generate_all(text, voice='af_heart', speed=1, use_gpu=CUDA_AVAILABLE):
     text = text if CHAR_LIMIT is None else text.strip()[:CHAR_LIMIT]
-    pipeline = pipelines[lang]
+    pipeline = pipelines[voice[0]]
     pack = pipeline.load_voice(voice)
     use_gpu = use_gpu and CUDA_AVAILABLE
     first = True
@@ -185,7 +100,51 @@ def get_frankenstein():
     with open('frankenstein5k.md', 'r') as r:
         return r.read().strip()
 
-# UI
+# --- Voice Choices ---
+CHOICES = {
+    # American English
+    '🇺🇸 🚺 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',
+
+    # British English
+    '🇬🇧 🚺 Emma': 'bf_emma',
+    '🇬🇧 🚺 Isabella': 'bf_isabella',
+    '🇬🇧 🚺 Alice': 'bf_alice',
+    '🇬🇧 🚺 Lily': 'bf_lily',
+    '🇬🇧 🚹 George': 'bm_george',
+    '🇬🇧 🚹 Fable': 'bm_fable',
+    '🇬🇧 🚹 Lewis': 'bm_lewis',
+    '🇬🇧 🚹 Daniel': 'bm_daniel',
+
+    # Hindi
+    '🇮🇳 🚺 Alpha': 'hf_alpha',
+    '🇮🇳 🚺 Beta': 'hf_beta',
+    '🇮🇳 🚹 Omega': 'hm_omega',
+    '🇮🇳 🚹 Psi': 'hm_psi',
+}
+
+# Preload all voices
+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/)`
 
@@ -200,22 +159,31 @@ 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)
+        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 = '⚠️ Streaming might fail the first time. Try again if no audio plays.'
+STREAM_NOTE = ['⚠️ There is an unknown Gradio bug that might yield no audio the first time you click `Stream`.']
+if CHAR_LIMIT is not None:
+    STREAM_NOTE.append(f'✂️ Each stream is capped at {CHAR_LIMIT} characters.')
+    STREAM_NOTE.append('🚀 Want more characters? You can [use Kokoro directly](https://huggingface.co/hexgrad/Kokoro-82M#usage) or duplicate this space:')
+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()
 
 BANNER_TEXT = '''
-[***Kokoro*** **TTS model - multilingual**](https://huggingface.co/hexgrad/Kokoro-82M)
-'''
+[***Kokoro*** **is an open-weight TTS model with 82 million parameters.**](https://huggingface.co/hexgrad/Kokoro-82M)
 
+This demo only showcases English and Hindi, but you can directly use the model to access other languages.
+'''
 API_OPEN = os.getenv('SPACE_ID') != 'hexgrad/Kokoro-TTS'
 API_NAME = None if API_OPEN else False
 with gr.Blocks() as app:
@@ -223,42 +191,31 @@ with gr.Blocks() as app:
         gr.Markdown(BANNER_TEXT, container=True)
     with gr.Row():
         with gr.Column():
-            text = gr.Textbox(label='Input Text')
-            lang = gr.Dropdown(list(LANG_CODES.keys()), value="🇺🇸 American English", label="Language")
-            voice = gr.Dropdown([], label='Voice')  # Populated dynamically
-
-            def update_voices(selected_lang):
-                code = LANG_CODES[selected_lang]
-                return gr.update(choices=list(VOICES[code].keys()), value=list(VOICES[code].keys())[0])
-
-            lang.change(fn=update_voices, inputs=[lang], outputs=[voice])
-            # Initialize with American English voices
-            voice.choices = list(VOICES["a"].keys())
-            voice.value = list(VOICES["a"].keys())[0]
-
-            use_gpu = gr.Dropdown(
-                [('ZeroGPU 🚀', True), ('CPU 🐌', False)],
-                value=CUDA_AVAILABLE,
-                label='Hardware',
-                interactive=CUDA_AVAILABLE
-            )
-            speed = gr.Slider(0.5, 2, 1, 0.1, label='Speed')
-            random_btn = gr.Button('🎲 Random Quote')
-            gatsby_btn = gr.Button('🥂 Gatsby')
-            frankenstein_btn = gr.Button('💀 Frankenstein')
+            text = gr.Textbox(label='Input Text', info=f"Up to ~500 characters per Generate, or {'∞' if CHAR_LIMIT is None else CHAR_LIMIT} characters per Stream")
+            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(get_random_quote, outputs=[text])
-    gatsby_btn.click(get_gatsby, outputs=[text])
-    frankenstein_btn.click(get_frankenstein, outputs=[text])
-    generate_btn.click(lambda t, l, v, s, g: generate_first(t, LANG_CODES[l], VOICES[LANG_CODES[l]][v], s, g),
-                       inputs=[text, lang, voice, speed, use_gpu], outputs=[out_audio, out_ps])
-    tokenize_btn.click(lambda t, l, v: tokenize_first(t, LANG_CODES[l], VOICES[LANG_CODES[l]][v]),
-                       inputs=[text, lang, voice], outputs=[out_ps])
-    stream_event = stream_btn.click(lambda t, l, v, s, g: generate_all(t, LANG_CODES[l], VOICES[LANG_CODES[l]][v], s, g),
-                                    inputs=[text, lang, voice, speed, use_gpu], outputs=[out_stream])
+    random_btn.click(fn=get_random_quote, inputs=[], outputs=[text], api_name=API_NAME)
+    gatsby_btn.click(fn=get_gatsby, inputs=[], outputs=[text], api_name=API_NAME)
+    frankenstein_btn.click(fn=get_frankenstein, inputs=[], outputs=[text], api_name=API_NAME)
+    generate_btn.click(fn=generate_first, inputs=[text, voice, speed, use_gpu], outputs=[out_audio, out_ps], api_name=API_NAME)
+    tokenize_btn.click(fn=tokenize_first, inputs=[text, voice], outputs=[out_ps], api_name=API_NAME)
+    stream_event = stream_btn.click(fn=generate_all, inputs=[text, voice, speed, use_gpu], outputs=[out_stream], api_name=API_NAME)
     stop_btn.click(fn=None, cancels=stream_event)
+    predict_btn.click(fn=predict, inputs=[text, voice, speed], outputs=[out_audio], api_name=API_NAME)
 
 if __name__ == '__main__':
     app.queue(api_open=API_OPEN).launch(show_api=API_OPEN, ssr_mode=True)