Update app.py

app.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, UploadFile, File, WebSocket
 from pywhispercpp.model import Model
 import uvicorn, tempfile, os
 from time import time
-
+import wave
 app = FastAPI(title="pyWhisperCPP API")
 model = Model("base.en")
 
@@ -27,33 +27,67 @@ async def transcribe(file: UploadFile = File(...)):
 # ================================
 # 🔥 Real-time streaming endpoint
 # ================================
+
+def save_pcm16_as_wav(data: bytes, sample_rate=16000, channels=1):
+    """Convert PCM16 stream into a valid WAV file"""
+    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+    with wave.open(temp.name, "wb") as wf:
+        wf.setnchannels(channels)
+        wf.setsampwidth(2)  # PCM16 = 2 bytes per sample
+        wf.setframerate(sample_rate)
+        wf.writeframes(data)
+    return temp.name
+
+
 @app.websocket("/ws/live")
 async def websocket_live(websocket: WebSocket):
     await websocket.accept()
+    print("Client connected!")
+
     buffer = b""
+    SAMPLE_RATE = 16000
+    MIN_PCM_SIZE = SAMPLE_RATE * 2 * 2  # 2 seconds PCM16
+
+    try:
+        while True:
+            data = await websocket.receive()
+
+            # Text (control)
+            if data.get("text"):
+                if data["text"] == "__END__":
+                    break
+                continue
+
+            # Binary PCM audio
+            chunk = data.get("bytes", b"")
+            if not chunk:
+                continue
+
+            buffer += chunk
+
+            if len(buffer) >= MIN_PCM_SIZE:
+                print(f"Transcribing {len(buffer)} bytes")
+
+                # Convert PCM16 → WAV with header
+                audio_path = save_pcm16_as_wav(buffer, SAMPLE_RATE)
 
-    while True:
-        data = await websocket.receive_bytes()
-        if data == b"__END__":
-            break
+                segments = model.transcribe(audio_path)
+                text = " ".join(seg.text for seg in segments).strip()
 
-        buffer += data
-        if len(buffer) > 32000:  # ~1 second PCM16 16kHz
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp:
-                temp.write(buffer)
-                temp.flush()
-                audio_path = temp.name
+                os.remove(audio_path)
+                buffer = b""
 
-            segments = model.transcribe(audio_path)
-            text = " ".join(seg.text for seg in segments)
+                if text:
+                    await websocket.send_text(text)
 
-            await websocket.send_text(text)
-            buffer = b""
-            os.remove(audio_path)
+        await websocket.send_text("[END]")
 
-    await websocket.send_text("[END]")
-    await websocket.close()
+    except Exception as e:
+        print("WebSocket Error:", e)
 
+    finally:
+        await websocket.close()
+        print("Client disconnected")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)