Update main.py

main.py

@@ -7,11 +7,8 @@ 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')
 
@@ -22,8 +19,8 @@ async def get_index():
 @app.websocket("/ws/chat")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
-    while True:
-        try:
+    try:
+        while True:
             user_msg = await websocket.receive_text()
             # Generate response
             response = llm(f"User: {user_msg}\nBella:", max_new_tokens=30)[0]['generated_text']
@@ -31,9 +28,7 @@ async def websocket_endpoint(websocket: WebSocket):
             
             # 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
+    except Exception as e:
+        print(f"Connection closed: {e}")