Update app/main.py

app/main.py

@@ -1,18 +1,14 @@
-
-
 import asyncio
 import json
 import torch
 import numpy as np
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
-from pydantic import BaseModel
 from liquid_audio import (
     LFM2AudioModel,
     LFM2AudioProcessor,
     ChatState,
 )
 
-
 HF_REPO = "LiquidAI/LFM2.5-Audio-1.5B"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 SAMPLE_RATE = 24_000
@@ -25,26 +21,18 @@ else:
 
 torch.backends.cuda.matmul.allow_tf32 = True
 
-
-
-
+print(f"[BOOT] Loading model on {DEVICE} with dtype {DTYPE}...")
 
 processor = LFM2AudioProcessor.from_pretrained(HF_REPO)
-model = LFM2AudioModel.from_pretrained(
-    HF_REPO,
-    torch_dtype=DTYPE,
-).to(DEVICE).eval()
-
-print(f"[BOOT]  LFM2.5 Loaded on {DEVICE}")
-
+model = LFM2AudioModel.from_pretrained(HF_REPO).to(dtype=DTYPE, device=DEVICE).eval()
 
+print(f"[BOOT] LFM2.5 Loaded on {DEVICE}")
 
 
 
 app = FastAPI(title="LFM2.5 WebSocket TTS", version="2.0.0")
 
 
-# WAV HEADER
 
 
 def wav_header(sample_rate: int, channels: int = 1, bits: int = 16) -> bytes:
@@ -66,8 +54,7 @@ def wav_header(sample_rate: int, channels: int = 1, bits: int = 16) -> bytes:
     )
 
 
-# STREAM CORE
-
+#  Stream core 
 
 async def stream_lfm_tts(websocket: WebSocket, text: str):
     chat = ChatState(processor)
@@ -96,7 +83,7 @@ async def stream_lfm_tts(websocket: WebSocket, text: str):
                 if data.get("type") == "stop":
                     stop_flag = True
                     break
-        except:
+        except Exception:
             stop_flag = True
 
     listener_task = asyncio.create_task(listen_for_stop())
@@ -123,7 +110,6 @@ async def stream_lfm_tts(websocket: WebSocket, text: str):
                         .unsqueeze(0)
                         .to(DEVICE)
                     )
-
                     waveform = processor.decode(audio_codes)
                     waveform = waveform.squeeze().cpu().numpy()
                     waveform = np.clip(waveform, -1.0, 1.0)
@@ -132,7 +118,7 @@ async def stream_lfm_tts(websocket: WebSocket, text: str):
                     await websocket.send_bytes(audio_int16.tobytes())
                     audio_buffer.clear()
 
-        # flush
+        # flush remaining
         if not stop_flag and len(audio_buffer) > 1:
             audio_codes = (
                 torch.stack(audio_buffer[:-1], dim=1)
@@ -143,7 +129,6 @@ async def stream_lfm_tts(websocket: WebSocket, text: str):
             waveform = waveform.squeeze().cpu().numpy()
             waveform = np.clip(waveform, -1.0, 1.0)
             audio_int16 = (waveform * 32767.0).astype(np.int16)
-
             await websocket.send_bytes(audio_int16.tobytes())
 
         await websocket.send_text(json.dumps({"type": "done"}))
@@ -152,13 +137,11 @@ async def stream_lfm_tts(websocket: WebSocket, text: str):
         listener_task.cancel()
 
 
-# WEBSOCKET ENDPOINT
-
+# WebSocket endpoint 
 
 @app.websocket("/ws/tts")
 async def websocket_tts(websocket: WebSocket):
     await websocket.accept()
-
     try:
         while True:
             message = await websocket.receive_text()
@@ -176,4 +159,11 @@ async def websocket_tts(websocket: WebSocket):
                 await stream_lfm_tts(websocket, text)
 
     except WebSocketDisconnect:
-        print("Client disconnected")
+        print("[WS] Client disconnected")
+
+
+
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "device": DEVICE}