Update app.py

app.py

@@ -1,12 +1,105 @@
-import zipfile
+from kokoro import KModel, KPipeline
+import gradio as gr
 import os
+import random
+import torch
+import logging
 
-zip_path = "model/kokoro-v1_0.zip"
-extract_dir = "model/"
+# Configuration
+VOICE_DIR = "model/voices"
+OUTPUT_DIR = "output_audio"
+TEXT = "Hello, this is a test of the Kokoro TTS system."
 
-if zipfile.is_zipfile(zip_path):
-    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-        zip_ref.extractall(extract_dir)
-        print("Extraction completed.")
-else:
-    print("Not a valid ZIP file.")
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Device setup
+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 [])}
+
+# 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
+}
+
+# Set custom pronunciations for "kokoro" in both American and British modes
+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}")
+
+# Core functions for voice generation
+def forward_gpu(text, voice_path, speed):
+    pipeline = pipelines[voice_path[0]]
+    pipeline.model = models[True]  # Use 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):
+        raise FileNotFoundError(f"Voice file not found: {voice_path}")
+
+    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]
+            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))
+            pipeline.model = models[False]
+            generator = pipeline(text, voice=voice_path, speed=speed)
+            for _, ps, audio in generator:
+                return (24000, audio.numpy()), ps
+        else:
+            raise gr.Error(e)
+    return None, ""
+
+# Load available voices
+def load_voice_choices():
+    if not os.path.exists(VOICE_DIR):
+        os.makedirs(VOICE_DIR)
+    voice_files = [f for f in os.listdir(VOICE_DIR) if f.endswith('.pt')]
+    choices = {voice_file: voice_file for voice_file in voice_files}
+    return choices
+
+CHOICES = load_voice_choices()
+
+if not CHOICES:
+    logger.warning("No voice files found in VOICE_DIR. Adding a placeholder.")
+    CHOICES = {"Bella": "af_bella.pt"}
+
+TOKEN_NOTE = '''
+💡 Customize pronunciation with Markdown link syntax and /slashes/ like [Kokoro](/kˈOkəɹO/)
+⬆️ Adjust stress levels using special notations.
+'''
+
+# Gradio Interface
+with gr.Blocks() as app:
+    with gr.Row():
+        text = gr.Textbox(label="Input Text", value=TEXT)
+        voice = gr.Dropdown(list(CHOICES.values()), label="Voice", value=list(CHOICES.values())[0])
+        speed = gr.Slider(0.5, 2, value=1, label="Speed")
+    output_audio = gr.Audio(label="Output Audio", interactive=False)
+    generate_btn = gr.Button("Generate")
+
+    generate_btn.click(fn=generate_first, inputs=[text, voice, speed], outputs=[output_audio])
+
+# Run the app
+if __name__ == "__main__":
+    app.launch(server_name="0.0.0.0", server_port=7860)