Update app.py

app.py

@@ -22,7 +22,7 @@ print("Loading SNAC model…")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
 model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-# Nur die Konfig + Safetensors, alles andere wird ignoriert
+print("Downloading Orpheus weights (konfig + safetensors)…")
 snapshot_download(
     repo_id=model_name,
     allow_patterns=["config.json", "*.safetensors", "model.safetensors.index.json"],
@@ -34,71 +34,59 @@ snapshot_download(
 )
 
 print("Loading Orpheus model…")
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, torch_dtype=torch.bfloat16
-).to(device)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
 model.config.pad_token_id = model.config.eos_token_id
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# — Konstanten für Audio‑Token →
-# (muss übereinstimmen mit Deinem Training; hier 128266)
-AUDIO_TOKEN_OFFSET = 128266
 
-# — Hilfsfunktionen —
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# — Hilfsfunktionen — 
 def process_prompt(text: str, voice: str):
+    """Erzeuge input_ids und attention_mask für einen Prompt."""
     prompt = f"{voice}: {text}"
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    # Laut Spezifikation:
-    #  start_token=128259, end_tokens=(128009,128260)
-    start = torch.tensor([[128259]], dtype=torch.int64)
-    end   = torch.tensor([[128009, 128260]], dtype=torch.int64)
-    ids  = torch.cat([start, input_ids, end], dim=1).to(device)
-    mask = torch.ones_like(ids).to(device)
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    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)
     return ids, mask
 
-def parse_output(generated_ids: torch.LongTensor):
-    """
-    Croppt nach dem letzten 128257-Start-Token, entfernt Padding (128258)
-    und zieht dann den Audio‑Offset ab, um echte Code‑IDs zu bekommen.
-    """
-    # finde letztes Audio‑Start‑Token
-    token_to_start  = 128257
-    token_to_remove = model.config.eos_token_id  # 128258
+def parse_output(generated_ids: torch.LongTensor) -> list[int]:
+    """Extrahiere rohe Tokenliste nach dem letzten 128257-Start-Token."""
+    token_to_find   = 128257
+    token_to_remove = 128258
 
-    idxs = (generated_ids == token_to_start).nonzero(as_tuple=True)[1]
+    idxs = (generated_ids == token_to_find).nonzero(as_tuple=True)[1]
     if idxs.numel() > 0:
-        cut = idxs[-1].item() + 1
-        cropped = generated_ids[:, cut:]
+        cropped = generated_ids[:, idxs[-1].item() + 1 :]
     else:
         cropped = generated_ids
 
-    # flatten & remove PAD, dann Offset abziehen
-    flat = cropped[0][cropped[0] != token_to_remove]
-    codes = [(int(t) - AUDIO_TOKEN_OFFSET) for t in flat]
-    return codes
-
-def redistribute_codes(code_list: list[int], snac_model: SNAC):
-    """
-    Verteilt die flache Code‑Liste in 3 Layers und dekodiert mit SNAC.
-    """
-    layer1, layer2, layer3 = [], [], []
-    for i in range(len(code_list) // 7):
+    row = cropped[0]
+    row = row[row != token_to_remove]
+    return row.tolist()
+
+def redistribute_codes(code_list: list[int]) -> bytes:
+    """Verteile die Codes auf die drei SNAC-Layer und dekodiere zu PCM16-Bytes."""
+    l1, l2, l3 = [], [], []
+    for i in range((len(code_list) + 1) // 7):
         base = code_list[7*i : 7*i+7]
-        layer1.append(base[0])
-        layer2.append(base[1] -   4096)
-        layer3.append(base[2] - 2*4096)
-        layer3.append(base[3] - 3*4096)
-        layer2.append(base[4] - 4*4096)
-        layer3.append(base[5] - 5*4096)
-        layer3.append(base[6] - 6*4096)
-
-    dev = next(snac_model.parameters()).device
-    c1 = torch.tensor(layer1, device=dev).unsqueeze(0)
-    c2 = torch.tensor(layer2, device=dev).unsqueeze(0)
-    c3 = torch.tensor(layer3, device=dev).unsqueeze(0)
-    audio = snac_model.decode([c1, c2, c3])
-    return audio.detach().squeeze().cpu().numpy()
+        l1.append(base[0])
+        l2.append(base[1] -   4096)
+        l3.append(base[2] - 2*4096)
+        l3.append(base[3] - 3*4096)
+        l2.append(base[4] - 4*4096)
+        l3.append(base[5] - 5*4096)
+        l3.append(base[6] - 6*4096)
+
+    dev = next(snac.parameters()).device
+    codes = [
+        torch.tensor(l1, device=dev).unsqueeze(0),
+        torch.tensor(l2, device=dev).unsqueeze(0),
+        torch.tensor(l3, device=dev).unsqueeze(0),
+    ]
+    audio = snac.decode(codes).squeeze().cpu().numpy()  # float32 @24 kHz
+    pcm16 = (audio * 32767).astype("int16").tobytes()
+    return pcm16
 
 # — FastAPI + WebSocket-Endpoint —
 app = FastAPI()
@@ -108,41 +96,47 @@ async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
         while True:
+            # 1) Nachricht empfangen
             msg = await ws.receive_text()
             data  = json.loads(msg)
             text  = data.get("text", "")
             voice = data.get("voice", "Jakob")
 
-            # 1) Prompt → Token‑Tensoren
+            # 2) Prompt → IDs/Mask
             ids, mask = process_prompt(text, voice)
 
-            # 2) Generation
+            # 3) Token-Generation
             gen_ids = model.generate(
                 input_ids=ids,
                 attention_mask=mask,
-                max_new_tokens=200,  # zum Debug
+                max_new_tokens=2000,
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.95,
                 repetition_penalty=1.1,
-                eos_token_id=model.config.eos_token_id,
+                eos_token_id=128258,
             )
 
-            # 3) Token → Code‑Liste → Audio (Float32 @24 kHz)
-            code_list = parse_output(gen_ids)
-            audio_np  = redistribute_codes(code_list, snac)
+            # 4) Parse + SNAC → PCM16‑Bytes
+            codes   = parse_output(gen_ids)
+            pcm16   = redistribute_codes(codes)
+            chunk_sz = 2400 * 2  # 0.1 s @24 kHz
 
-            # 4) In 0.1 s‑Chunks (2400 Samples) als PCM16 streamen
-            pcm16 = (audio_np * 32767).astype("int16").tobytes()
-            chunk = 2400 * 2
-            for i in range(0, len(pcm16), chunk):
-                await ws.send_bytes(pcm16[i : i+chunk])
+            # 5) Stream audio‑Chunks
+            for i in range(0, len(pcm16), chunk_sz):
+                await ws.send_bytes(pcm16[i : i + chunk_sz])
                 await asyncio.sleep(0.1)
 
+            # 6) Ende‑Signal
+            await ws.send_json({"event": "eos"})
+
+            # (Verbindung bleibt offen für nächste Anfrage)
+
     except WebSocketDisconnect:
         print("Client disconnected")
     except Exception as e:
         print("Error in /ws/tts:", e)
+        # Schließe erst, nachdem Fehler gemeldet
         await ws.close(code=1011)
 
 if __name__ == "__main__":