Update app.py

app.py

@@ -10,10 +10,10 @@ from fastapi.responses import StreamingResponse, Response, HTMLResponse
 from fastapi.middleware import Middleware
 from fastapi.middleware.gzip import GZipMiddleware
 
-from kokoro import KPipeline
+from kokoro import StreamKPipeline, KModel, KPipeline  # Import StreamKPipeline and KModel
 
 app = FastAPI(
-    title="Kokoro TTS FastAPI",
+    title="Kokoro Streaming TTS FastAPI",
     middleware=[
         Middleware(GZipMiddleware, compresslevel=9)  # Add GZip compression
     ]
@@ -23,7 +23,8 @@ app = FastAPI(
 # Global Pipeline Instance
 # ------------------------------------------------------------------------------
 # Create one pipeline instance for the entire app.
-pipeline = KPipeline(lang_code="a")
+model = KModel() # Initialize KModel
+pipeline = StreamKPipeline(lang_code="a", model=model) # Initialize StreamKPipeline, passing the model
 
 
 # ------------------------------------------------------------------------------
@@ -47,12 +48,12 @@ def generate_wav_header(sample_rate: int, num_channels: int, sample_width: int,
     return header + fmt_chunk + data_chunk_header
 
 
-def audio_tensor_to_pcm_bytes(audio_tensor: torch.Tensor) -> bytes:
+def audio_chunk_to_pcm_bytes(audio_chunk: torch.Tensor) -> bytes:
     """
-    Convert a torch.FloatTensor (with values in [-1, 1]) to raw 16-bit PCM bytes.
+    Convert a torch.FloatTensor audio chunk (values in [-1, 1]) to raw 16-bit PCM bytes.
     """
     # Ensure tensor is on CPU and flatten if necessary.
-    audio_np = audio_tensor.cpu().numpy()
+    audio_np = audio_chunk.cpu().numpy()
     if audio_np.ndim > 1:
         audio_np = audio_np.flatten()
     # Scale to int16 range.
@@ -60,9 +61,9 @@ def audio_tensor_to_pcm_bytes(audio_tensor: torch.Tensor) -> bytes:
     return audio_int16.tobytes()
 
 
-def audio_tensor_to_opus_bytes(audio_tensor: torch.Tensor, sample_rate: int = 24000, bitrate: int = 32000) -> bytes:
+def audio_chunk_to_opus_bytes(audio_chunk: torch.Tensor, sample_rate: int = 24000, bitrate: int = 32000) -> bytes:
     """
-    Convert a torch.FloatTensor to Opus encoded bytes.
+    Convert a torch.FloatTensor audio chunk to Opus encoded bytes.
     Requires the 'opuslib' package: pip install opuslib
     """
     try:
@@ -70,7 +71,7 @@ def audio_tensor_to_opus_bytes(audio_tensor: torch.Tensor, sample_rate: int = 24
     except ImportError:
         raise ImportError("opuslib is not installed. Please install it with: pip install opuslib")
 
-    audio_np = audio_tensor.cpu().numpy()
+    audio_np = audio_chunk.cpu().numpy()
     if audio_np.ndim > 1:
         audio_np = audio_np.flatten()
     # Scale to int16 range.  Important for opus.
@@ -94,55 +95,51 @@ def audio_tensor_to_opus_bytes(audio_tensor: torch.Tensor, sample_rate: int = 24
 
 
 # ------------------------------------------------------------------------------
-# Endpoints
+# Streaming TTS Endpoint
 # ------------------------------------------------------------------------------
 
-@app.get("/tts/streaming", summary="True Streaming TTS")
+@app.get("/tts/streaming", summary="Streaming TTS")
 def tts_streaming(text: str, voice: str = "af_heart", speed: float = 1.0, format: str = "opus"):
     """
-    True Streaming TTS endpoint that returns a continuous audio stream.
-    It processes text and generates audio token by token (or small chunks as KPipeline yields),
-    providing a more responsive streaming experience.
+    Streaming TTS endpoint that returns a continuous audio stream.
     Supports WAV (PCM) and Opus formats. Opus offers significantly better compression.
 
     The endpoint first yields a WAV header (with a dummy length) for WAV,
-    then yields encoded audio data for each token's audio as soon as it is generated.
+    then yields encoded audio data chunks as they are generated.
     """
     sample_rate = 24000
     num_channels = 1
     sample_width = 2  # 16-bit PCM
 
-    def audio_generator():
+    def audio_chunk_generator():
         if format.lower() == "wav":
-            # Yield the WAV header first.
+            # Yield the WAV header first for PCM WAV format.
             header = generate_wav_header(sample_rate, num_channels, sample_width)
             yield header
 
-        try:
-            results = pipeline(text, voice=voice, speed=speed, split_pattern=None) # split_pattern=None to avoid splitting here, let KPipeline handle
-            for result in results:
-                if result.audio is not None:
-                    if format.lower() == "wav":
-                        yield audio_tensor_to_pcm_bytes(result.audio)
-                    elif format.lower() == "opus":
-                        yield audio_tensor_to_opus_bytes(result.audio, sample_rate=sample_rate)
-                    else:
-                        raise ValueError(f"Unsupported audio format: {format}")
+        # Stream audio chunks from the pipeline.
+        for audio_chunk in pipeline(text=text, voice=voice, speed=speed):
+            if audio_chunk is not None and audio_chunk.numel() > 0:
+                if format.lower() == "wav":
+                    yield audio_chunk_to_pcm_bytes(audio_chunk)
+                elif format.lower() == "opus":
+                    yield audio_chunk_to_opus_bytes(audio_chunk, sample_rate=sample_rate)
                 else:
-                    print("No audio generated for a token/chunk") # Debugging, remove in production if not needed
-        except Exception as e:
-            print(f"Error during TTS processing: {e}")
-            yield b''  # Important: yield empty bytes to keep stream alive, or handle error sound
+                    raise ValueError(f"Unsupported audio format: {format}")
 
     media_type = "audio/wav" if format.lower() == "wav" else "audio/opus"
 
     return StreamingResponse(
-        audio_generator(),
+        audio_chunk_generator(),
         media_type=media_type,
         headers={"Cache-Control": "no-cache"},
     )
 
 
+# ------------------------------------------------------------------------------
+# Full TTS Endpoint (unchanged from your original code)
+# ------------------------------------------------------------------------------
+
 @app.get("/tts/full", summary="Full TTS")
 def tts_full(text: str, voice: str = "af_heart", speed: float = 1.0, format: str = "wav"):
     """
@@ -186,6 +183,9 @@ def tts_full(text: str, voice: str = "af_heart", speed: float = 1.0, format: str
         raise HTTPException(status_code=400, detail=f"Unsupported audio format: {format}")
 
 
+# ------------------------------------------------------------------------------
+# HTML Demo Page Endpoint (unchanged from your original code, but updated to call new streaming endpoint)
+# ------------------------------------------------------------------------------
 
 @app.get("/", response_class=HTMLResponse)
 def index():
@@ -199,10 +199,10 @@ def index():
     <!DOCTYPE html>
     <html>
     <head>
-        <title>Kokoro TTS Demo</title>
+        <title>Kokoro Streaming TTS Demo</title>
     </head>
     <body>
-        <h1>Kokoro TTS Demo</h1>
+        <h1>Kokoro Streaming TTS Demo</h1>
         <textarea id="text" rows="4" cols="50" placeholder="Enter text here"></textarea><br>
         <label for="voice">Voice:</label>
         <input type="text" id="voice" value="af_heart"><br>