Update app/main.py

app/main.py

@@ -1,5 +1,3 @@
-# app/main.py
-
 import asyncio
 import json
 import torch
@@ -29,26 +27,25 @@ print("[BOOT] Model loaded")
 app = FastAPI(title="LFM2.5 Speech-to-Speech", version="3.0")
 
 
-def wav_header(sample_rate: int, channels: int = 1, bits: int = 16) -> bytes:
-    byte_rate = sample_rate * channels * bits // 8
-    block_align = channels * bits // 8
+def wav_header(sr=24000, ch=1, bits=16):
+    byte_rate = sr * ch * bits // 8
+    block_align = ch * bits // 8
     return (
         b"RIFF"
-        + (b"\xff\xff\xff\xff")
+        + b"\xff\xff\xff\xff"
         + b"WAVEfmt "
         + (16).to_bytes(4, "little")
         + (1).to_bytes(2, "little")
-        + channels.to_bytes(2, "little")
-        + sample_rate.to_bytes(4, "little")
+        + ch.to_bytes(2, "little")
+        + sr.to_bytes(4, "little")
         + byte_rate.to_bytes(4, "little")
         + block_align.to_bytes(2, "little")
         + bits.to_bytes(2, "little")
         + b"data"
-        + (b"\xff\xff\xff\xff")
+        + b"\xff\xff\xff\xff"
     )
 
 
-
 async def generate_response(websocket: WebSocket, audio_np: np.ndarray):
 
     chat = ChatState(processor)
@@ -58,7 +55,7 @@ async def generate_response(websocket: WebSocket, audio_np: np.ndarray):
     chat.end_turn()
 
     chat.new_turn("user")
-    chat.add_audio(audio_np)
+    chat.add_audio(audio_np, sampling_rate=SAMPLE_RATE)  
     chat.end_turn()
 
     chat.new_turn("assistant")
@@ -68,14 +65,12 @@ async def generate_response(websocket: WebSocket, audio_np: np.ndarray):
     audio_buffer = []
 
     with torch.inference_mode():
-
         for token in model.generate_interleaved(
             **chat,
             max_new_tokens=4096,
             audio_temperature=0.8,
             audio_top_k=4,
         ):
-
             if token.numel() == 1:
                 continue
 
@@ -85,7 +80,6 @@ async def generate_response(websocket: WebSocket, audio_np: np.ndarray):
                 audio_buffer.append(token)
 
             if len(audio_buffer) >= CHUNK_SIZE:
-
                 audio_codes = (
                     torch.stack(audio_buffer, dim=1)
                     .unsqueeze(0)
@@ -100,27 +94,38 @@ async def generate_response(websocket: WebSocket, audio_np: np.ndarray):
 
                 waveform = waveform.squeeze().cpu().numpy()
                 waveform = np.clip(waveform, -1.0, 1.0)
-
                 audio_int16 = (waveform * 32767).astype(np.int16)
 
                 await websocket.send_bytes(audio_int16.tobytes())
-
                 audio_buffer.clear()
 
+    # flush remaining
+    if len(audio_buffer) > 1:
+        audio_codes = (
+            torch.stack(audio_buffer, dim=1)
+            .unsqueeze(0)
+            .to(DEVICE)
+        )
+        try:
+            waveform = processor.decode(audio_codes)
+            waveform = waveform.squeeze().cpu().numpy()
+            waveform = np.clip(waveform, -1.0, 1.0)
+            audio_int16 = (waveform * 32767).astype(np.int16)
+            await websocket.send_bytes(audio_int16.tobytes())
+        except Exception:
+            pass
+
     await websocket.send_text(json.dumps({"type": "done"}))
 
 
 @app.websocket("/ws/s2s")
 async def websocket_s2s(websocket: WebSocket):
-
     await websocket.accept()
 
     try:
-
         audio_bytes = bytearray()
 
         while True:
-
             message = await websocket.receive()
 
             if "text" in message:
@@ -129,11 +134,9 @@ async def websocket_s2s(websocket: WebSocket):
                 if payload["type"] == "start":
                     audio_bytes.clear()
 
-                if payload["type"] == "end":
-
+                elif payload["type"] == "end":
                     audio_np = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32)
                     audio_np /= 32767.0
-
                     await generate_response(websocket, audio_np)
 
             elif "bytes" in message: