up5

app.py

@@ -1,6 +1,7 @@
 import os
 import json
 import asyncio
+import numpy as np
 import torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from dotenv import load_dotenv
@@ -8,74 +9,78 @@ from snac import SNAC
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login, snapshot_download
 
-# ─── ENV & HF TOKEN ────────────────────────────────────────────────────────────
+# — ENV & HF‑AUTH —
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-# ─── DEVICE ─────────────────────────────────────────────────────────────────────
+# — Device —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ─── SNAC ───────────────────────────────────────────────────────────────────────
-print("Loading SNAC model…")
+# — Modelle laden —
+print("Loading SNAC model...")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
-# ─── ORPHEUS ────────────────────────────────────────────────────────────────────
 model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-
-# pre‑download only the config + safetensors, damit das Image schlank bleibt
+print("Downloading model weights (config + safetensors)...")
 snapshot_download(
     repo_id=model_name,
     allow_patterns=["config.json", "*.safetensors", "model.safetensors.index.json"],
     ignore_patterns=[
         "optimizer.pt", "pytorch_model.bin", "training_args.bin",
-        "scheduler.pt", "tokenizer.*", "vocab.json", "merges.txt"
+        "scheduler.pt", "tokenizer.json", "tokenizer_config.json",
+        "special_tokens_map.json", "vocab.json", "merges.txt", "tokenizer.*"
     ]
 )
 
-print("Loading Orpheus model…")
+print("Loading Orpheus model...")
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.bfloat16  # optional: beschleunigt das FP16‑ähnliche Rechnen
-)
-model = model.to(device)
+    torch_dtype=torch.bfloat16
+).to(device)
 model.config.pad_token_id = model.config.eos_token_id
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# ─── HILFSFUNKTIONEN ──────────────────────────────────────────────────────────
+# — Hilfsfunktionen —
 
 def process_prompt(text: str, voice: str):
-    """
-    Baut aus Text+Voice ein batch‑Tensor input_ids für `model.generate`.
-    """
+    """Bereitet input_ids und attention_mask für das Modell vor."""
     prompt = f"{voice}: {text}"
-    tok = tokenizer(prompt, return_tensors="pt").to(device)
-    start = torch.tensor([[128259]], device=device)
-    end   = torch.tensor([[128009, 128260]], device=device)
-    return torch.cat([start, tok.input_ids, end], dim=1)
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    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)
+    return ids, mask
 
 def parse_output(generated_ids: torch.LongTensor):
-    """
-    Schneidet bis zum letzten 128257 und entfernt 128258, gibt reine Token‑Liste zurück.
-    """
-    START, PAD = 128257, 128258
-    idxs = (generated_ids == START).nonzero(as_tuple=True)[1]
+    """Extrahiere rohe Tokenliste nach dem letzten 128257-Start-Token."""
+    token_to_find   = 128257
+    token_to_remove = 128258
+
+    idxs = (generated_ids == token_to_find).nonzero(as_tuple=True)[1]
     if idxs.numel() > 0:
-        cropped = generated_ids[:, idxs[-1].item()+1:]
+        cut = idxs[-1].item() + 1
+        cropped = generated_ids[:, cut:]
     else:
         cropped = generated_ids
-    row = cropped[0][cropped[0] != PAD]
-    return row.tolist()
+
+    # Entferne EOS‑Token
+    row = cropped[0]
+    return row[row != token_to_remove].tolist()
 
 def redistribute_codes(code_list: list[int], snac_model: SNAC):
     """
-    Verteilt 7er‑Blöcke auf die drei SNAC‑Layer und dekodiert zu Audio (numpy float32).
+    Verteilt die Token nur in kompletten 7er‑Blöcken auf die drei SNAC‑Layer
+    und dekodiert in Audio. Unvollständige Reste (<7 Tokens) werden verworfen.
     """
+    n_blocks = len(code_list) // 7
     layer1, layer2, layer3 = [], [], []
-    for i in range((len(code_list) + 1) // 7):
-        base = code_list[7*i : 7*i+7]
+
+    for i in range(n_blocks):
+        base = code_list[7*i : 7*i + 7]
         layer1.append(base[0])
         layer2.append(base[1] -   4096)
         layer3.append(base[2] - 2*4096)
@@ -83,6 +88,11 @@ def redistribute_codes(code_list: list[int], snac_model: SNAC):
         layer2.append(base[4] - 4*4096)
         layer3.append(base[5] - 5*4096)
         layer3.append(base[6] - 6*4096)
+
+    if not layer1:
+        # kein kompletter Block → leeres Audio
+        return np.zeros(0, dtype=np.float32)
+
     dev = next(snac_model.parameters()).device
     codes = [
         torch.tensor(layer1, device=dev).unsqueeze(0),
@@ -92,31 +102,32 @@ def redistribute_codes(code_list: list[int], snac_model: SNAC):
     audio = snac_model.decode(codes)
     return audio.detach().squeeze().cpu().numpy()
 
-# ─── FASTAPI ────────────────────────────────────────────────────────────────────
+# — FastAPI Setup —
 
 app = FastAPI()
 
 @app.get("/")
-async def healthcheck():
-    return {"status": "ok", "msg": "Hello, Orpheus TTS up!"}
+def greet_json():
+    return {"Hello": "World!"}
 
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
         while True:
-            # 1) Eintreffende JSON‐Nachricht parsen
+            # Erwartet JSON: {"text": "...", "voice": "Jakob"}
             data = json.loads(await ws.receive_text())
             text  = data.get("text", "")
             voice = data.get("voice", "Jakob")
 
-            # 2) Prompt → input_ids
-            ids = process_prompt(text, voice)
+            # 1) Tokens vorbereiten
+            ids, mask = process_prompt(text, voice)
 
-            # 3) Token‐Erzeugung
+            # 2) Generierung
             gen_ids = model.generate(
                 input_ids=ids,
-                max_new_tokens=2000,      # hier z.B. 20k geht auch, wird aber speicherintensiv
+                attention_mask=mask,
+                max_new_tokens=2000,    # hier nach Bedarf anpassen
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.95,
@@ -124,25 +135,24 @@ async def tts_ws(ws: WebSocket):
                 eos_token_id=model.config.eos_token_id,
             )
 
-            # 4) Tokens → Code‐Liste → Audio
-            codes    = parse_output(gen_ids)
+            # 3) Tokens → Code-Liste → Audio
+            codes   = parse_output(gen_ids)
             audio_np = redistribute_codes(codes, snac)
 
-            # 5) PCM16‐Stream in 0.1 s‐Blöcken
+            # 4) in 0.1s‑Stücken PCM16 streamen
             pcm16 = (audio_np * 32767).astype("int16").tobytes()
-            chunk = 2400 * 2
+            chunk = 2400 * 2  # 2400 samples @24kHz = 0.1s * 2 bytes
             for i in range(0, len(pcm16), chunk):
                 await ws.send_bytes(pcm16[i : i+chunk])
                 await asyncio.sleep(0.1)
 
+        # Ende der while‐Schleife
     except WebSocketDisconnect:
         print("Client disconnected")
     except Exception as e:
         print("Error in /ws/tts:", e)
         await ws.close(code=1011)
 
-# ─── START ──────────────────────────────────────────────────────────────────────
-
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)