Update app.py

app.py

@@ -115,21 +115,64 @@ def redistribute_codes(code_list: list[int]) -> np.ndarray:
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        msg = await ws.receive_text()
-        req = json.loads(msg)
+        msg   = await ws.receive_text()
+        req   = json.loads(msg)
         text  = req.get("text", "")
         voice = req.get("voice", "")
 
-        # 1) Prompt → Codes → Audio
-        with torch.no_grad():
-            codes = process_single_prompt(text, voice)
-            audio_np = redistribute_codes(codes)
-
-        # 2) In PCM16 wandeln & senden
-        pcm16 = (audio_np * 32767).astype(np.int16).tobytes()
-        await ws.send_bytes(pcm16)
-
-        # 3) sauber schließen
+        # 1) Prompt vorbereiten
+        input_ids, attention_mask = prepare_inputs(text, voice)
+        past_kvs = None
+        buffer   = []
+
+        # 2) Token‑für‑Token (oder in kleinen Blöcken)
+        while True:
+            # Nur max_new_tokens=50 pro Aufruf
+            out = model.generate(
+                input_ids=input_ids   if past_kvs is None else None,
+                attention_mask=attention_mask if past_kvs is None else None,
+                past_key_values=past_kvs,
+                use_cache=True,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.1,
+                max_new_tokens=50,
+                eos_token_id=128258,
+                return_dict_in_generate=True,
+                output_past_key_values=True,
+                return_legacy_cache=True,  # falls Ihr noch das alte past_key_values-Format braucht
+            )
+
+            # Extrahiere neue Token (ohne die already generated ones)
+            new_ids = out.sequences[0, input_ids.shape[-1]:].tolist()
+            past_kvs = out.past_key_values
+
+            for tok in new_ids:
+                if tok == model.config.eos_token_id:
+                    # Stream zu Ende
+                    break
+                if tok == 128257:    # Reset-Start‑Marker
+                    buffer = []
+                    continue
+                buffer.append(tok - AUDIO_OFFSET)
+
+                # Sobald wir 7 Audio‑Codes gesammelt haben → dekodieren & schicken
+                if len(buffer) == 7:
+                    pcm = decode_block(buffer)
+                    buffer = []
+                    await ws.send_bytes(pcm)
+
+            # Wenn EOS im Chunk war, abbrechen
+            if model.config.eos_token_id in new_ids:
+                break
+
+            # Danach weiter mit nächsten 50 Tokens,
+            # input_ids & attention_mask nur beim ersten Aufruf nötig
+            input_ids = None
+            attention_mask = None
+
+        # 3) Am Ende WebSocket sauber schließen
         await ws.close()
 
     except WebSocketDisconnect: