Update main.py

main.py

@@ -4,34 +4,19 @@ from kokoro import KPipeline
 import asyncio
 
 app = FastAPI()
-
-# Initialize models
 llm = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-360M-Instruct")
 tts = KPipeline(lang_code='a', model='shahid202/Kokoro-82M-TTS')
 
-def get_llm_response(text):
-    # Very short, witty generation
-    prompt = f"User: {text}\nBella:"
-    return llm(prompt, max_new_tokens=30)[0]['generated_text'].split("Bella:")[-1].strip()
-
 @app.websocket("/ws/chat")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     while True:
         user_msg = await websocket.receive_text()
+        # Get LLM text
+        response_text = llm(f"User: {user_msg}\nBella:", max_new_tokens=30)[0]['generated_text']
+        text = response_text.split("Bella:")[-1].strip()
         
-        # 1. Get response from LLM
-        full_text = get_llm_response(user_msg)
-        
-        # 2. Simple split by punctuation to trigger TTS immediately
-        sentences = [s.strip() + "." for s in full_text.split('.') if s]
-        
-        for sentence in sentences:
-            # 3. Generate audio for the sentence
-            generator = tts(sentence, voice="af_heart", speed=1.0)
-            for _, _, audio in generator:
-                # 4. Send audio bytes over WebSocket
-                await websocket.send_bytes(audio.tobytes())
-                # Small pause to mimic natural speech flow
-                await asyncio.sleep(0.1) 
-                
+        # Stream audio chunk by chunk
+        for _, _, audio in tts(text, voice="af_heart", speed=1.0):
+            await websocket.send_bytes(audio.tobytes())
+            await asyncio.sleep(0.05) # Keeps flow smooth