Update app.py

app.py

@@ -1,9 +1,9 @@
-from kokoro import KModel, KPipeline
 import gradio as gr
 import os
 import torch
 import logging
 import soundfile as sf
+from kokoro import KModel, KPipeline
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -23,30 +23,22 @@ CUDA_AVAILABLE = torch.cuda.is_available()
 device = "cuda" if CUDA_AVAILABLE else "cpu"
 logger.info(f"Using hardware: {device}")
 
-# Load models for CPU and GPU (if available)
-models = {gpu: KModel("hexgrad/Kokoro-82M").to("cuda" if gpu else "cpu").eval() for gpu in [False] + ([True] if CUDA_AVAILABLE else [])}
+# Load a single model instance
+model = KModel("hexgrad/Kokoro-82M").to(device).eval()
 
 # Define pipelines for American ('a') and British ('b') English
 pipelines = {
-    'a': KPipeline(model=models[False], lang_code='a', device='cpu'),  # American English
-    'b': KPipeline(model=models[False], lang_code='b', device='cpu')   # British English
+    'a': KPipeline(model=model, lang_code='a', device=device),  # American English
+    'b': KPipeline(model=model, lang_code='b', device=device)   # British English
 }
 
-# Set custom pronunciations for "kokoro" in both American and British modes
+# Set custom pronunciations for "kokoro"
 try:
     pipelines["a"].g2p.lexicon.golds["kokoro"] = "kˈOkəɹO"
     pipelines["b"].g2p.lexicon.golds["kokoro"] = "kˈQkəɹQ"
 except AttributeError as e:
     logger.warning(f"Could not set custom pronunciations: {e}")
 
-def forward_gpu(text, voice_path, speed):
-    pipeline = pipelines[voice_path[0]]
-    pipeline.model = models[True]  # Switch to GPU model
-    generator = pipeline(text, voice=voice_path, speed=speed)
-    for _, _, audio in generator:
-        return audio
-    return None
-
 def generate_first(text, voice="af_bella.pt", speed=1, use_gpu=CUDA_AVAILABLE):
     voice_path = os.path.join(VOICE_DIR, voice)
     if not os.path.exists(voice_path):
@@ -55,18 +47,14 @@ def generate_first(text, voice="af_bella.pt", speed=1, use_gpu=CUDA_AVAILABLE):
     pipeline = pipelines[voice[0]]
     use_gpu = use_gpu and CUDA_AVAILABLE
     try:
-        if use_gpu:
-            audio = forward_gpu(text, voice_path, speed)
-        else:
-            pipeline.model = models[False]  # Ensure CPU model is used
-            generator = pipeline(text, voice=voice_path, speed=speed)
-            for _, ps, audio in generator:
-                return (24000, audio.numpy()), ps
+        generator = pipeline(text, voice=voice_path, speed=speed)
+        for _, ps, audio in generator:
+            return (24000, audio.numpy()), ps
     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.")
-            pipeline.model = models[False]  # Switch to CPU model
+            model.to("cpu")
             generator = pipeline(text, voice=voice_path, speed=speed)
             for _, ps, audio in generator:
                 return (24000, audio.numpy()), ps
@@ -74,9 +62,6 @@ def generate_first(text, voice="af_bella.pt", speed=1, use_gpu=CUDA_AVAILABLE):
             raise gr.Error(e)
     return None, ""
 
-def predict(text, voice="af_bella.pt", speed=1):
-    return generate_first(text, voice, speed, use_gpu=False)[0]
-
 def tokenize_first(text, voice="af_bella.pt"):
     voice_path = os.path.join(VOICE_DIR, voice)
     if not os.path.exists(voice_path):
@@ -96,10 +81,8 @@ def generate_all(text, voice="af_bella.pt", speed=1, use_gpu=CUDA_AVAILABLE):
     pipeline = pipelines[voice[0]]
     use_gpu = use_gpu and CUDA_AVAILABLE
     first = True
-    if use_gpu:
-        pipeline.model = models[True]  # Switch to GPU model
-    else:
-        pipeline.model = models[False]  # Switch to CPU model
+    if not use_gpu:
+        model.to("cpu")
     generator = pipeline(text, voice=voice_path, speed=speed)
     for _, _, audio in generator:
         yield 24000, audio.numpy()
@@ -107,7 +90,7 @@ def generate_all(text, voice="af_bella.pt", speed=1, use_gpu=CUDA_AVAILABLE):
             first = False
             yield 24000, torch.zeros(1).numpy()
 
-# Dynamically load all .pt voice files from VOICE_DIR
+# Dynamically load .pt voice files from VOICE_DIR
 def load_voice_choices():
     voice_files = [f for f in os.listdir(VOICE_DIR) if f.endswith('.pt')]
     choices = {}
@@ -159,7 +142,6 @@ with gr.Blocks() as generate_tab:
                             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)
 
 with gr.Blocks() as stream_tab:
     out_stream = gr.Audio(label="Output Audio Stream", interactive=False, streaming=True, autoplay=True)
@@ -167,8 +149,7 @@ with gr.Blocks() as stream_tab:
         stream_btn = gr.Button("Stream", variant="primary")
         stop_btn = gr.Button("Stop", variant="stop")
     with gr.Accordion("Note", open=True):
-        gr.Markdown("⚠️ There is an unknown Gradio bug that might yield no audio the first time you click Stream.")
-        gr.DuplicateButton()
+        gr.Markdown("⚠️ There may be delays in streaming audio due to processing limitations.")
 
 with gr.Blocks() as app:
     with gr.Row():
@@ -178,7 +159,7 @@ with gr.Blocks() as app:
                 voice = gr.Dropdown(list(CHOICES.items()), value="af_bella.pt" if "af_bella.pt" in CHOICES.values() else list(CHOICES.values())[0], label="Voice",
                                     info="Quality and availability vary by language")
                 use_gpu = gr.Dropdown(
-                    [("GPU 🚀", True), ("CPU 🐌", False)],
+                    [("GPU �-held", True), ("CPU 🐌", False)],
                     value=CUDA_AVAILABLE,
                     label="Hardware",
                     info="GPU is usually faster, but may require CUDA support",
@@ -192,7 +173,6 @@ with gr.Blocks() as app:
     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().launch()