Update main.py

main.py

@@ -1,22 +1,39 @@
 from fastapi import FastAPI, WebSocket
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
 from transformers import pipeline
 from kokoro import KPipeline
 import asyncio
+import numpy as np
 
 app = FastAPI()
+
+# Mount the static folder
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# Initialize models
 llm = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-360M-Instruct")
 tts = KPipeline(lang_code='a', model='shahid202/Kokoro-82M-TTS')
 
+@app.get("/")
+async def get_index():
+    return FileResponse('static/index.html')
+
 @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()
-        
-        # 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
+        try:
+            user_msg = await websocket.receive_text()
+            # Generate response
+            response = llm(f"User: {user_msg}\nBella:", max_new_tokens=30)[0]['generated_text']
+            text = response.split("Bella:")[-1].strip()
+            
+            # Stream audio
+            for _, _, audio in tts(text, voice="af_heart", speed=1.0):
+                # Ensure float32 format
+                await websocket.send_bytes(audio.astype(np.float32).tobytes())
+                await asyncio.sleep(0.01)
+        except Exception as e:
+            print(f"Error: {e}")
+            break