Update app.py

app.py

@@ -1,11 +1,11 @@
 import gradio as gr
+import io
 import os
 import tempfile
 import time
 import wave
 import struct
 import numpy as np
-import io
 from openai import OpenAI
 from elevenlabs.client import ElevenLabs
 from elevenlabs import stream, play
@@ -113,7 +113,7 @@ def get_elevenlabs_voices(api_key):
         "Sam": "yoZ06aMxZJJ28mfd3POQ"
     }
 
-# Kokoro TTS (CPU/GPU)
+# Kokoro TTS (CPU-only)
 _KOKORO_STATE = { "initialized": False, "device": "cpu", "model": None, "pipelines": {} }
 
 def _init_kokoro() -> None:
@@ -121,14 +121,8 @@ def _init_kokoro() -> None:
         return
     if KModel is None or KPipeline is None:
         raise gr.Error("Kokoro is not installed. Please add 'kokoro>=0.9.4' and 'torch' to requirements and install.")
-    # Prefer CUDA if available, otherwise CPU
+    
     device = "cpu"
-    try:
-        if torch is not None and hasattr(torch, "cuda") and torch.cuda.is_available():
-            device = "cuda"
-    except Exception:
-        device = "cpu"
-
     model = KModel().to(device).eval()
     pipelines = {"a": KPipeline(lang_code="a", model=False)}
     try:
@@ -138,16 +132,6 @@ def _init_kokoro() -> None:
 
     _KOKORO_STATE.update({"initialized": True, "device": device, "model": model, "pipelines": pipelines})
 
-def _int16_wav_bytes(audio_int16: np.ndarray, sr: int = 24_000) -> bytes:
-    """Encode a mono int16 numpy array to WAV bytes (single chunk)."""
-    with io.BytesIO() as buf:
-        with wave.open(buf, "wb") as wf:
-            wf.setnchannels(1)
-            wf.setsampwidth(2)
-            wf.setframerate(sr)
-            wf.writeframes(audio_int16.tobytes())
-        return buf.getvalue()
-
 def get_kokoro_voices():
     """Get list of available Kokoro voice IDs."""
     try:
@@ -171,11 +155,33 @@ def get_kokoro_voices():
             "zf_xiaobei", "zf_xiaoni", "zf_xiaoxiao", "zf_xiaoyi", "zm_yunjian", "zm_yunxi", "zm_yunxia", "zm_yunyang"
         ]
 
-def kokoro_tts_stream(text: str, speed: float, voice: str):
-    """Stream speech with Kokoro-82M by yielding WAV bytes per segment.
+def _audio_np_to_int16(audio_np: np.ndarray) -> np.ndarray:
+    audio_clipped = np.clip(audio_np, -1.0, 1.0)
+    return (audio_clipped * 32767.0).astype(np.int16)
+
 
-    Yields: bytes (WAV) chunks suitable for gr.Audio(streaming=True).
-    """
+def _write_wav_file(audio_int16: np.ndarray, sample_rate: int = 24_000) -> str:
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        path = tmp.name
+    with wave.open(path, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(audio_int16.tobytes())
+    return path
+
+
+def _wav_bytes_from_int16(audio_int16: np.ndarray, sample_rate: int = 24_000) -> bytes:
+    buffer = io.BytesIO()
+    with wave.open(buffer, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(audio_int16.tobytes())
+    return buffer.getvalue()
+
+
+def _kokoro_segment_generator(text: str, speed: float, voice: str):
     if not text or not text.strip():
         raise gr.Error("Please enter text to synthesize.")
 
@@ -186,51 +192,80 @@ def kokoro_tts_stream(text: str, speed: float, voice: str):
     if pipeline is None:
         raise gr.Error("Kokoro English pipeline not initialized.")
 
-    sr = 24_000
     pack = pipeline.load_voice(voice)
 
     try:
-        # Iterate lazily; do not materialize all segments
         for idx, (_, ps, _) in enumerate(pipeline(text, voice, speed)):
             ref_s = pack[len(ps) - 1]
             try:
-                audio = model(ps, ref_s, float(speed))  # torch tensor possibly on GPU/CPU
+                audio = model(ps, ref_s, float(speed))
                 audio_np = audio.detach().cpu().numpy()
-                audio_clipped = np.clip(audio_np, -1.0, 1.0)
-                audio_int16 = (audio_clipped * 32767.0).astype(np.int16)
-                yield _int16_wav_bytes(audio_int16, sr)
+                yield audio_np
             except Exception as e:
                 raise gr.Error(f"Error generating audio for segment {idx + 1}: {str(e)[:200]}...")
     except gr.Error:
         raise
     except Exception as e:
-        raise gr.Error(f"Error during streaming generation: {str(e)[:200]}...")
+        raise gr.Error(f"Error during speech generation: {str(e)[:200]}...")
+
+
+def kokoro_tts(text: str, speed: float, voice: str) -> str:
+    sr = 24_000
+    segments = list(_kokoro_segment_generator(text, speed, voice))
+    if not segments:
+        raise gr.Error("No audio was generated.")
+
+    audio_np = segments[0] if len(segments) == 1 else np.concatenate(segments, axis=0)
+    audio_int16 = _audio_np_to_int16(audio_np)
+    return _write_wav_file(audio_int16, sr)
+
+
+def kokoro_tts_stream(text: str, speed: float, voice: str):
+    sr = 24_000
+    produced_any = False
+
+    for audio_np in _kokoro_segment_generator(text, speed, voice):
+        produced_any = True
+        audio_int16 = _audio_np_to_int16(audio_np)
+        chunk_bytes = _wav_bytes_from_int16(audio_int16, sr)
+        yield chunk_bytes
+
+    if not produced_any:
+        raise gr.Error("No audio was generated.")
 
 # Main dispatcher function to handle all services
+def _read_file_bytes(path: str) -> bytes:
+    with open(path, "rb") as file:
+        data = file.read()
+    return data
+
+
 def generate_tts(text, service, openai_api_key, openai_model, openai_voice, 
                  elevenlabs_api_key, elevenlabs_voice, voice_dict,
                  kokoro_speed, kokoro_voice):
-    """Route to appropriate TTS service.
-
-    Acts as a generator to support streaming audio output.
-    """
+    """Route to appropriate TTS service based on selection"""
     if service == "Kokoro":
-        for chunk in kokoro_tts_stream(text, kokoro_speed, kokoro_voice):
-            yield chunk
-        return
-    elif service == "OpenAI":
-        final_path = openai_tts(text, openai_model, openai_voice, openai_api_key)
-        yield final_path
+        yield from kokoro_tts_stream(text, kokoro_speed, kokoro_voice)
         return
+
+    if service == "OpenAI":
+        file_path = openai_tts(text, openai_model, openai_voice, openai_api_key)
     elif service == "ElevenLabs":
         voice_id = voice_dict.get(elevenlabs_voice, elevenlabs_voice)
-        final_path = elevenlabs_tts(text, voice_id, elevenlabs_api_key)
-        yield final_path
-        return
+        file_path = elevenlabs_tts(text, voice_id, elevenlabs_api_key)
     else:
-        # Fallback in case of an unknown service
         raise gr.Error(f"Unknown service selected: {service}")
 
+    try:
+        audio_bytes = _read_file_bytes(file_path)
+    finally:
+        try:
+            os.remove(file_path)
+        except OSError:
+            pass
+
+    yield audio_bytes
+
 # Function to update ElevenLabs voices when API key changes
 def update_elevenlabs_voices(api_key):
     """Update voice dropdown when API key is entered"""
@@ -337,6 +372,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme') as demo:
         label="Generated Speech",
         streaming=True,
         autoplay=True,
+        show_download_button=True,
     )
     
     # ==========================