Update app.py

app.py

@@ -10,247 +10,285 @@ from fastapi.responses import StreamingResponse, Response, HTMLResponse
 from fastapi.middleware import Middleware
 from fastapi.middleware.gzip import GZipMiddleware
 
-from kokoro import StreamKPipeline, KModel, KPipeline  # Import StreamKPipeline and KModel
+from kokoro import KPipeline
 
 app = FastAPI(
-    title="Kokoro Streaming TTS FastAPI",
-    middleware=[
-        Middleware(GZipMiddleware, compresslevel=9)  # Add GZip compression
-    ]
+title="Kokoro TTS FastAPI",
+middleware=[
+Middleware(GZipMiddleware, compresslevel=9)  # Add GZip compression
+]
 )
 
-# ------------------------------------------------------------------------------
-# Global Pipeline Instance
-# ------------------------------------------------------------------------------
-# Create one pipeline instance for the entire app.
-model = KModel() # Initialize KModel
-stream_pipeline = StreamKPipeline(lang_code="a", model=model, chunk_size=2) # Initialize StreamKPipeline, passing the model
-pipeline = KPipeline(lang_code="a", model=model) # Initialize KPipeline, passing the model
+------------------------------------------------------------------------------
+Global Pipeline Instance
+------------------------------------------------------------------------------
+Create one pipeline instance for the entire app.
 
+pipeline = KPipeline(lang_code="a")
 
-
-# ------------------------------------------------------------------------------
-# Helper Functions
-# ------------------------------------------------------------------------------
+------------------------------------------------------------------------------
+Helper Functions
+------------------------------------------------------------------------------
 
 def generate_wav_header(sample_rate: int, num_channels: int, sample_width: int, data_size: int = 0x7FFFFFFF) -> bytes:
-    """
-    Generate a WAV header for streaming.
-    Since we don't know the final audio size, we set the data chunk size to a large dummy value.
-    This header is sent only once at the start of the stream.
-    """
-    bits_per_sample = sample_width * 8
-    byte_rate = sample_rate * num_channels * sample_width
-    block_align = num_channels * sample_width
-    # total file size = 36 + data_size (header is 44 bytes total)
-    total_size = 36 + data_size
-    header = struct.pack('<4sI4s', b'RIFF', total_size, b'WAVE')
-    fmt_chunk = struct.pack('<4sIHHIIHH', b'fmt ', 16, 1, num_channels, sample_rate, byte_rate, block_align, bits_per_sample)
-    data_chunk_header = struct.pack('<4sI', b'data', data_size)
-    return header + fmt_chunk + data_chunk_header
-
-
-def audio_chunk_to_pcm_bytes(audio_chunk: torch.Tensor) -> 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_chunk.cpu().numpy()
-    if audio_np.ndim > 1:
-        audio_np = audio_np.flatten()
-    # Scale to int16 range.
-    audio_int16 = np.int16(audio_np * 32767)
-    return audio_int16.tobytes()
-
-
-def audio_chunk_to_opus_bytes(audio_chunk: torch.Tensor, sample_rate: int = 24000, bitrate: int = 32000) -> bytes:
-    """
-    Convert a torch.FloatTensor audio chunk to Opus encoded bytes.
-    Requires the 'opuslib' package: pip install opuslib
-    """
-    try:
-        import opuslib
-    except ImportError:
-        raise ImportError("opuslib is not installed. Please install it with: pip install opuslib")
-
-    audio_np = audio_chunk.cpu().numpy()
-    if audio_np.ndim > 1:
-        audio_np = audio_np.flatten()
-    # Scale to int16 range.  Important for opus.
-    audio_int16 = np.int16(audio_np * 32767)
-
-    encoder = opuslib.Encoder(sample_rate, 1, opuslib.APPLICATION_VOIP) # 1 channel for mono.
-
-    # Calculate the number of frames to encode. Opus frames are 2.5, 5, 10, or 20 ms long.
-    frame_size = int(sample_rate * 0.020)  # 20ms frame size
-
-    encoded_data = b''
-    for i in range(0, len(audio_int16), frame_size):
-        frame = audio_int16[i:i + frame_size]
-        if len(frame) < frame_size:
-            # Pad the last frame with zeros if needed.
-            frame = np.pad(frame, (0, frame_size - len(frame)), 'constant')
-        encoded_frame = encoder.encode(frame.tobytes(), frame_size) # Encode the frame.
-        encoded_data += encoded_frame
-
-    return encoded_data
-
-
-# ------------------------------------------------------------------------------
-# Streaming TTS Endpoint
-# ------------------------------------------------------------------------------
+"""
+Generate a WAV header for streaming.
+Since we don't know the final audio size, we set the data chunk size to a large dummy value.
+This header is sent only once at the start of the stream.
+"""
+bits_per_sample = sample_width * 8
+byte_rate = sample_rate * num_channels * sample_width
+block_align = num_channels * sample_width
+# total file size = 36 + data_size (header is 44 bytes total)
+total_size = 36 + data_size
+header = struct.pack('<4sI4s', b'RIFF', total_size, b'WAVE')
+fmt_chunk = struct.pack('<4sIHHIIHH', b'fmt ', 16, 1, num_channels, sample_rate, byte_rate, block_align, bits_per_sample)
+data_chunk_header = struct.pack('<4sI', b'data', data_size)
+return header + fmt_chunk + data_chunk_header
+
+def custom_split_text(text: str) -> list:
+"""
+Custom splitting:
+- Start with a chunk size of 2 words.
+- For each chunk, if a period (".") is found in any word (except if it’s the very last word),
+then split the chunk at that word (include words up to that word).
+- Otherwise, use the current chunk size.
+- For subsequent chunks, increase the chunk size by 2.
+- If there are fewer than the desired number of words for a full chunk, add all remaining words.
+"""
+words = text.split()
+chunks = []
+chunk_size = 2
+start = 0
+while start < len(words):
+candidate_end = start + chunk_size
+if candidate_end > len(words):
+candidate_end = len(words)
+chunk_words = words[start:candidate_end]
+# Look for a period in any word except the last one.
+split_index = None
+for i in range(len(chunk_words) - 1):
+if '.' in chunk_words[i]:
+split_index = i
+break
+if split_index is not None:
+candidate_end = start + split_index + 1
+chunk_words = words[start:candidate_end]
+chunks.append(" ".join(chunk_words))
+start = candidate_end
+chunk_size += 2  # Increase the chunk size by 2 for the next iteration.
+return chunks
+
+def audio_tensor_to_pcm_bytes(audio_tensor: torch.Tensor) -> bytes:
+"""
+Convert a torch.FloatTensor (with 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()
+if audio_np.ndim > 1:
+audio_np = audio_np.flatten()
+# Scale to int16 range.
+audio_int16 = np.int16(audio_np * 32767)
+return audio_int16.tobytes()
+
+def audio_tensor_to_opus_bytes(audio_tensor: torch.Tensor, sample_rate: int = 24000, bitrate: int = 32000) -> bytes:
+"""
+Convert a torch.FloatTensor to Opus encoded bytes.
+Requires the 'opuslib' package: pip install opuslib
+"""
+try:
+import opuslib
+except ImportError:
+raise ImportError("opuslib is not installed. Please install it with: pip install opuslib")
+
+audio_np = audio_tensor.cpu().numpy()
+if audio_np.ndim > 1:
+    audio_np = audio_np.flatten()
+# Scale to int16 range.  Important for opus.
+audio_int16 = np.int16(audio_np * 32767)
+
+encoder = opuslib.Encoder(sample_rate, 1, opuslib.APPLICATION_VOIP) # 1 channel for mono.
+
+# Calculate the number of frames to encode. Opus frames are 2.5, 5, 10, or 20 ms long.
+frame_size = int(sample_rate * 0.020)  # 20ms frame size
+
+encoded_data = b''
+for i in range(0, len(audio_int16), frame_size):
+    frame = audio_int16[i:i + frame_size]
+    if len(frame) < frame_size:
+        # Pad the last frame with zeros if needed.
+        frame = np.pad(frame, (0, frame_size - len(frame)), 'constant')
+    encoded_frame = encoder.encode(frame.tobytes(), frame_size) # Encode the frame.
+    encoded_data += encoded_frame
+
+return encoded_data
+content_copy
+download
+ Use code with caution.
+------------------------------------------------------------------------------
+Endpoints
+------------------------------------------------------------------------------
 
 @app.get("/tts/streaming", summary="Streaming TTS")
 def tts_streaming(text: str, voice: str = "af_heart", speed: float = 1.0, format: str = "opus"):
-    """
-    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 chunks as they are generated.
-    """
-    sample_rate = 24000
-    num_channels = 1
-    sample_width = 2  # 16-bit PCM
-
-    def audio_chunk_generator():
-        if format.lower() == "wav":
-            # Yield the WAV header first for PCM WAV format.
-            header = generate_wav_header(sample_rate, num_channels, sample_width)
-            yield header
-
-        # Stream audio chunks from the pipeline.
-        for audio_chunk in stream_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)
+"""
+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 text chunk as soon as it is generated.
+"""
+# Split the input text using the custom doubling strategy.
+chunks = custom_split_text(text)
+sample_rate = 24000
+num_channels = 1
+sample_width = 2  # 16-bit PCM
+
+def audio_generator():
+    if format.lower() == "wav":
+        # Yield the WAV header first.
+        header = generate_wav_header(sample_rate, num_channels, sample_width)
+        yield header
+    # Process and yield each chunk's audio data.
+    for i, chunk in enumerate(chunks):
+        print(f"Processing chunk {i}: {chunk}")  # Debugging
+        try:
+            results = list(pipeline(chunk, voice=voice, speed=speed, split_pattern=None))
+            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}")
                 else:
-                    raise ValueError(f"Unsupported audio format: {format}")
-
+                    print(f"Chunk {i}: No audio generated")
+        except Exception as e:
+            print(f"Error processing chunk {i}: {e}")
+            yield b'' # important so that streaming continues.  Consider returning an error sound.
 
-    return StreamingResponse(
-        audio_chunk_generator(),
-        media_type="audio/wav",
-        headers={"Cache-Control": "no-cache"},
-    )
+media_type = "audio/wav" if format.lower() == "wav" else "audio/opus"
 
-
-# ------------------------------------------------------------------------------
-# Full TTS Endpoint (unchanged from your original code)
-# ------------------------------------------------------------------------------
+return StreamingResponse(
+    audio_generator(),
+    media_type=media_type,
+    headers={"Cache-Control": "no-cache"},
+)
+content_copy
+download
+ Use code with caution.
 
 @app.get("/tts/full", summary="Full TTS")
 def tts_full(text: str, voice: str = "af_heart", speed: float = 1.0, format: str = "wav"):
-    """
-    Full TTS endpoint that synthesizes the entire text, concatenates the audio,
-    and returns a complete WAV or Opus file.
-    """
-    # Use newline-based splitting via the pipeline's split_pattern.
-    results = list(pipeline(text, voice=voice, speed=speed, split_pattern=r"\n+"))
-    audio_segments = []
-    for result in results:
-        if result.audio is not None:
-            audio_np = result.audio.cpu().numpy()
-            if audio_np.ndim > 1:
-                audio_np = audio_np.flatten()
-            audio_segments.append(audio_np)
-
-    if not audio_segments:
-        raise HTTPException(status_code=500, detail="No audio generated.")
-
-    # Concatenate all audio segments.
-    full_audio = np.concatenate(audio_segments)
-
-    # Write the concatenated audio to an in-memory WAV or Opus file.
-    sample_rate = 24000
-    num_channels = 1
-    sample_width = 2  # 16-bit PCM -> 2 bytes per sample
-    if format.lower() == "wav":
-        wav_io = io.BytesIO()
-        with wave.open(wav_io, "wb") as wav_file:
-            wav_file.setnchannels(num_channels)
-            wav_file.setsampwidth(sample_width)
-            wav_file.setframerate(sample_rate)
-            full_audio_int16 = np.int16(full_audio * 32767)
-            wav_file.writeframes(full_audio_int16.tobytes())
-        wav_io.seek(0)
-        return Response(content=wav_io.read(), media_type="audio/wav")
-    elif format.lower() == "opus":
-        opus_data = audio_tensor_to_opus_bytes(torch.from_numpy(full_audio), sample_rate=sample_rate)
-        return Response(content=opus_data, media_type="audio/opus")
-    else:
-        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)
-# ------------------------------------------------------------------------------
+"""
+Full TTS endpoint that synthesizes the entire text, concatenates the audio,
+and returns a complete WAV or Opus file.
+"""
+# Use newline-based splitting via the pipeline's split_pattern.
+results = list(pipeline(text, voice=voice, speed=speed, split_pattern=r"\n+"))
+audio_segments = []
+for result in results:
+if result.audio is not None:
+audio_np = result.audio.cpu().numpy()
+if audio_np.ndim > 1:
+audio_np = audio_np.flatten()
+audio_segments.append(audio_np)
+
+if not audio_segments:
+    raise HTTPException(status_code=500, detail="No audio generated.")
+
+# Concatenate all audio segments.
+full_audio = np.concatenate(audio_segments)
+
+# Write the concatenated audio to an in-memory WAV or Opus file.
+sample_rate = 24000
+num_channels = 1
+sample_width = 2  # 16-bit PCM -> 2 bytes per sample
+if format.lower() == "wav":
+    wav_io = io.BytesIO()
+    with wave.open(wav_io, "wb") as wav_file:
+        wav_file.setnchannels(num_channels)
+        wav_file.setsampwidth(sample_width)
+        wav_file.setframerate(sample_rate)
+        full_audio_int16 = np.int16(full_audio * 32767)
+        wav_file.writeframes(full_audio_int16.tobytes())
+    wav_io.seek(0)
+    return Response(content=wav_io.read(), media_type="audio/wav")
+elif format.lower() == "opus":
+    opus_data = audio_tensor_to_opus_bytes(torch.from_numpy(full_audio), sample_rate=sample_rate)
+    return Response(content=opus_data, media_type="audio/opus")
+else:
+    raise HTTPException(status_code=400, detail=f"Unsupported audio format: {format}")
+content_copy
+download
+ Use code with caution.
 
 @app.get("/", response_class=HTMLResponse)
 def index():
-    """
-    HTML demo page for Kokoro TTS.
-
-    This page provides a simple UI to enter text, choose a voice and speed,
-    and play synthesized audio from both the streaming and full endpoints.
-    """
-    return """
-    <!DOCTYPE html>
-    <html>
-    <head>
-        <title>Kokoro Streaming TTS Demo</title>
-    </head>
-    <body>
-        <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>
-        <label for="speed">Speed:</label>
-        <input type="number" step="0.1" id="speed" value="1.0"><br>
-        <label for="format">Format:</label>
-        <select id="format">
-            <option value="wav">WAV</option>
-            <option value="opus" selected>Opus</option>
-        </select><br><br>
-        <button onclick="playStreaming()">Play Streaming TTS</button>
-        <button onclick="playFull()">Play Full TTS</button>
-        <br><br>
-        <audio id="audio" controls autoplay></audio>
-        <script>
-            function playStreaming() {
-                const text = document.getElementById('text').value;
-                const voice = document.getElementById('voice').value;
-                const speed = document.getElementById('speed').value;
-                const format = document.getElementById('format').value;
-                const audio = document.getElementById('audio');
-                // Set the audio element's source to the streaming endpoint.
-                audio.src = `/tts/streaming?text=${encodeURIComponent(text)}&voice=${encodeURIComponent(voice)}&speed=${speed}&format=${format}`;
-                audio.type = format === 'wav' ? 'audio/wav' : 'audio/opus';
-                audio.play();
-            }
-            function playFull() {
-                const text = document.getElementById('text').value;
-                const voice = document.getElementById('voice').value;
-                const speed = document.getElementById('speed').value;
-                const format = document.getElementById('format').value;
-                const audio = document.getElementById('audio');
-                // Set the audio element's source to the full TTS endpoint.
-                audio.src = `/tts/full?text=${encodeURIComponent(text)}&voice=${encodeURIComponent(voice)}&speed=${speed}&format=${format}`;
-                audio.type = format === 'wav' ? 'audio/wav' : 'audio/opus';
-                audio.play();
-            }
-        </script>
-    </body>
-    </html>
-    """
-
-
-# ------------------------------------------------------------------------------
-# Run with: uvicorn app:app --reload
-# ------------------------------------------------------------------------------
-if __name__ == "__main__":
-    import uvicorn
-
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
+"""
+HTML demo page for Kokoro TTS.
+
+This page provides a simple UI to enter text, choose a voice and speed,
+and play synthesized audio from both the streaming and full endpoints.
+"""
+return """
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Kokoro TTS Demo</title>
+</head>
+<body>
+    <h1>Kokoro 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>
+    <label for="speed">Speed:</label>
+    <input type="number" step="0.1" id="speed" value="1.0"><br>
+    <label for="format">Format:</label>
+    <select id="format">
+        <option value="wav">WAV</option>
+        <option value="opus" selected>Opus</option>
+    </select><br><br>
+    <button onclick="playStreaming()">Play Streaming TTS</button>
+    <button onclick="playFull()">Play Full TTS</button>
+    <br><br>
+    <audio id="audio" controls autoplay></audio>
+    <script>
+        function playStreaming() {
+            const text = document.getElementById('text').value;
+            const voice = document.getElementById('voice').value;
+            const speed = document.getElementById('speed').value;
+            const format = document.getElementById('format').value;
+            const audio = document.getElementById('audio');
+            // Set the audio element's source to the streaming endpoint.
+            audio.src = `/tts/streaming?text=${encodeURIComponent(text)}&voice=${encodeURIComponent(voice)}&speed=${speed}&format=${format}`;
+            audio.type = format === 'wav' ? 'audio/wav' : 'audio/opus';
+            audio.play();
+        }
+        function playFull() {
+            const text = document.getElementById('text').value;
+            const voice = document.getElementById('voice').value;
+            const speed = document.getElementById('speed').value;
+            const format = document.getElementById('format').value;
+            const audio = document.getElementById('audio');
+            // Set the audio element's source to the full TTS endpoint.
+            audio.src = `/tts/full?text=${encodeURIComponent(text)}&voice=${encodeURIComponent(voice)}&speed=${speed}&format=${format}`;
+            audio.type = format === 'wav' ? 'audio/wav' : 'audio/opus';
+            audio.play();
+        }
+    </script>
+</body>
+</html>
+"""
+content_copy
+download
+ Use code with caution.
+------------------------------------------------------------------------------
+Run with: uvicorn app:app --reload
+------------------------------------------------------------------------------
+
+if name == "main":
+import uvicorn
+
+uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)