Update app.py

app.py

@@ -12,18 +12,18 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(HF_TOKEN)
 
-# — Device wählen —
+# — Device auswählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # — FastAPI instanziieren —
 app = FastAPI()
 
-# — Hello‑Route, damit kein 404 bei GET / mehr kommt —
+# — Hello‑Route, damit GET / nicht 404 gibt —
 @app.get("/")
 async def read_root():
     return {"message": "Hello, world!"}
 
-# — Modelle bei Startup laden —
+# — Modelle beim Startup laden —
 @app.on_event("startup")
 async def load_models():
     global tokenizer, model, snac
@@ -34,99 +34,122 @@ async def load_models():
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        device_map={"": 0} if device == "cuda" else None,
-        torch_dtype=torch.bfloat16 if device == "cuda" else None,
+        device_map="auto" if device=="cuda" else None,
+        torch_dtype=torch.bfloat16 if device=="cuda" else None,
         low_cpu_mem_usage=True
-    )
-    # Pad‑ID auf EOS einstellen
+    ).to(device)
     model.config.pad_token_id = model.config.eos_token_id
 
-# — Hilfsfunktionen —
+# — Input‑Vorbereitung —
 def prepare_inputs(text: str, voice: str):
     prompt = f"{voice}: {text}"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    # Start‑/End‑Marker
     start = torch.tensor([[128259]], dtype=torch.int64, device=device)
     end   = torch.tensor([[128009, 128260]], dtype=torch.int64, device=device)
     ids   = torch.cat([start, input_ids, end], dim=1)
-    mask  = torch.ones_like(ids)
+    mask  = torch.ones_like(ids, device=device)
     return ids, mask
 
-def decode_block(block_tokens: list[int]):
-    # aus 7 Tokens einen SNAC‑Decode‑Block bauen
-    layer1, layer2, layer3 = [], [], []
-    b = block_tokens
-    layer1.append(b[0])
-    layer2.append(b[1] -   4096)
-    layer3.append(b[2] - 2*4096)
-    layer3.append(b[3] - 3*4096)
-    layer2.append(b[4] - 4*4096)
-    layer3.append(b[5] - 5*4096)
-    layer3.append(b[6] - 6*4096)
+# — SNAC‑Dekodierung eines 7‑Token‑Blocks →
+def decode_block(tokens: list[int]) -> bytes:
+    l1, l2, l3 = [], [], []
+    b = tokens
+    l1.append(b[0])
+    l2.append(b[1]-4096)
+    l3.append(b[2]-2*4096)
+    l3.append(b[3]-3*4096)
+    l2.append(b[4]-4*4096)
+    l3.append(b[5]-5*4096)
+    l3.append(b[6]-6*4096)
     codes = [
-        torch.tensor(layer1, device=device).unsqueeze(0),
-        torch.tensor(layer2, device=device).unsqueeze(0),
-        torch.tensor(layer3, device=device).unsqueeze(0),
+        torch.tensor(l1, device=device).unsqueeze(0),
+        torch.tensor(l2, device=device).unsqueeze(0),
+        torch.tensor(l3, device=device).unsqueeze(0),
     ]
-    # ergibt FloatTensor shape (1, N), @24 kHz
     audio = snac.decode(codes).squeeze().cpu().numpy()
-    # in PCM16 umwandeln
     return (audio * 32767).astype("int16").tobytes()
 
-# — WebSocket Endpoint für TTS Streaming —
+# — WebSocket‑Endpoint mit Chunked‑Generate (max_new_tokens=50) —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        # erst die Anfrage als JSON empfangen
+        # 1) Anfrage einlesen
         msg = await ws.receive_text()
         req = json.loads(msg)
         text  = req.get("text", "")
         voice = req.get("voice", "Jakob")
 
-        # Inputs bauen
+        # 2) Inputs bauen
         input_ids, attention_mask = prepare_inputs(text, voice)
         past_kvs = None
-        collected = []
+        buffer_codes: list[int] = []
+
+        # 3) Chunk‑Generate‑Loop
+        chunk_size = 50
+        eos_id     = model.config.eos_token_id
+
+        # Wir tracken bisher erzeugte Länge, um abzugrenzen, was neu ist
+        prev_len = 0
 
-        # Token‑für‑Token loop
         while True:
-            out = model(
-                input_ids=input_ids if past_kvs is None else None,
+            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,
+                max_new_tokens=chunk_size,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.1,
+                eos_token_id=eos_id,
                 use_cache=True,
+                return_dict_in_generate=True,
+                output_scores=False,
+                past_key_values=past_kvs
             )
-            logits = out.logits[:, -1, :]
+            # Update past_kvs und sequences
             past_kvs = out.past_key_values
+            seqs     = out.sequences  # (1, total_length)
+            total_len = seqs.shape[1]
 
-            # Sampling
-            probs = torch.softmax(logits, dim=-1)
-            nxt   = torch.multinomial(probs, num_samples=1).item()
+            # 4) Neue Tokens extrahieren
+            new_tokens = seqs[0, prev_len:total_len].tolist()
+            prev_len = total_len
 
-            # Ende, wenn EOS
-            if nxt == model.config.eos_token_id:
+            # 5) Jeden neuen Token aufbereiten
+            for tok in new_tokens:
+                if tok == eos_id:
+                    # Ende
+                    new_tokens = []  # clean up
+                    break
+                if tok == 128257:
+                    buffer_codes.clear()
+                    continue
+                # offset und puffern
+                buffer_codes.append(tok - 128266)
+                # sobald 7 Codes gesammelt, dekodieren & senden
+                if len(buffer_codes) >= 7:
+                    block = buffer_codes[:7]
+                    buffer_codes = buffer_codes[7:]
+                    pcm = decode_block(block)
+                    await ws.send_bytes(pcm)
+
+            # 6) Abbruch, wenn EOS im Chunk war
+            if eos_id in new_tokens:
                 break
-            # Reset bei neuem Start‑Marker
-            if nxt == 128257:
-                collected = []
-                continue
-
-            # Audio‑Code offsetten und sammeln
-            collected.append(nxt - 128266)
-            # sobald 7 Stück, direkt dekodieren und senden
-            if len(collected) == 7:
-                pcm = decode_block(collected)
-                collected = []
-                await ws.send_bytes(pcm)
-
-        # nach Ende sauber schließen
-        await ws.close()
 
+            # Inputs für nächsten Durchgang nur beim ersten Mal
+            input_ids = attention_mask = None
+
+        # 7) Zum Schluss sauber schließen
+        await ws.close()
     except WebSocketDisconnect:
-        # Client hat disconnectet
-        pass
+        return
     except Exception as e:
-        # bei Fehlern 1011 senden
         print("Error in /ws/tts:", e)
         await ws.close(code=1011)
+
+# — Main für lokalen Test —
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)