Add WebSocket endpoint for true real-time bidirectional streaming

api.py

@@ -1,5 +1,6 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
 from fastapi.responses import StreamingResponse
+import json
 import torch
 import numpy as np
 import re
@@ -179,10 +180,13 @@ async def root():
         "speakers": list(SPEAKERS.keys()),
         "sample_rate": SAMPLE_RATE,
         "endpoints": {
-            "/tts": "Standard TTS (wait for full audio)",
-            "/tts/stream": "Streaming TTS (audio chunks in real-time - recommended for Pipecat)",
-            "/speakers": "List available speakers"
-        }
+            "POST /tts": "Standard TTS (wait for full audio)",
+            "POST /tts/stream": "HTTP Streaming TTS (audio chunks in real-time)",
+            "WS /ws/tts": "WebSocket TTS (BEST FOR PIPECAT - true real-time bidirectional)",
+            "GET /speakers": "List available speakers"
+        },
+        "device": DEVICE,
+        "optimization": "fp16 (half precision)"
     }
 
 
@@ -375,6 +379,78 @@ async def get_speakers():
     return {"speakers": list(SPEAKERS.keys())}
 
 
+@app.websocket("/ws/tts")
+async def websocket_tts(websocket: WebSocket):
+    """WebSocket endpoint for real-time audio streaming.
+
+    Usage:
+    1. Connect to ws://localhost:7860/ws/tts
+    2. Send JSON: {"text": "سلام دنیا", "speaker": "Divya"}
+    3. Receive audio chunks in real-time
+    4. Connection closes when generation completes
+    """
+    await websocket.accept()
+    try:
+        data = await websocket.receive_text()
+        request_data = json.loads(data)
+
+        text = request_data.get("text", "").strip()
+        speaker = request_data.get("speaker", "Divya")
+        pitch = request_data.get("pitch", "Moderate")
+        rate = request_data.get("rate", "Moderate")
+        temperature = request_data.get("temperature", 0.8)
+        do_sample = request_data.get("do_sample", True)
+
+        # Validate inputs
+        if not text:
+            await websocket.send_json({"error": "Text cannot be empty"})
+            await websocket.close()
+            return
+
+        if speaker not in SPEAKERS:
+            await websocket.send_json({
+                "error": f"Invalid speaker. Choose from: {list(SPEAKERS.keys())}"
+            })
+            await websocket.close()
+            return
+
+        # Send status message
+        await websocket.send_json({
+            "status": "generating",
+            "message": f"Generating audio for speaker {speaker}..."
+        })
+
+        # Generate and stream audio chunks
+        chunk_count = 0
+        try:
+            for audio_chunk in generate_audio_chunks_streaming(
+                text, speaker, pitch, rate, temperature, do_sample
+            ):
+                # Send audio chunk as binary data
+                await websocket.send_bytes(audio_chunk.tobytes())
+                chunk_count += 1
+
+            # Send completion message
+            await websocket.send_json({
+                "status": "complete",
+                "chunks_sent": chunk_count
+            })
+
+        except Exception as e:
+            await websocket.send_json({
+                "error": f"Generation failed: {str(e)}"
+            })
+
+    except WebSocketDisconnect:
+        print("WebSocket client disconnected")
+    except json.JSONDecodeError:
+        await websocket.send_json({"error": "Invalid JSON format"})
+        await websocket.close()
+    except Exception as e:
+        print(f"WebSocket error: {e}")
+        await websocket.close()
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)