Update app.py

app.py

@@ -7,12 +7,12 @@ from huggingface_hub import login
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from snac import SNAC
 
-# — HF‑Token & Login (wenn gesetzt) —
+# — HF‑Token & Login (falls gesetzt) —
 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"
 
 app = FastAPI()
@@ -26,39 +26,38 @@ model = None
 tokenizer = None
 snac_model = None
 
-# — Startup: SNAC & Orpheus laden —
 @app.on_event("startup")
 async def load_models():
     global model, tokenizer, snac_model
-    # 1) SNAC
+
+    # 1) SNAC laden
     snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-    # 2) Orpheus‑TTS
-    REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_synthetic-v0.1"
+
+    # 2) Orpheus‑TTS (public “natural”-Variante)
+    REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
     tokenizer = AutoTokenizer.from_pretrained(REPO)
     model = AutoModelForCausalLM.from_pretrained(
         REPO,
-        device_map="auto" 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
     ).to(device)
     model.config.pad_token_id = model.config.eos_token_id
 
-# — Marker und Offsets aus der Vorlage —
+# — Marker und Offsets —
 START_TOKEN  = 128259
 END_TOKENS   = [128009, 128260]
 AUDIO_OFFSET = 128266
 
 def process_single_prompt(prompt: str, voice: str) -> list[int]:
-    # Prompt zusammenbauen
-    if voice and voice != "in_prompt":
-        text = f"{voice}: {prompt}"
-    else:
-        text = prompt
+    # Prompt zusammenstellen
+    text = f"{voice}: {prompt}" if voice and voice != "in_prompt" else prompt
+
     # Tokenize + Marker
-    ids = tokenizer(text, return_tensors="pt").input_ids
-    start = torch.tensor([[START_TOKEN]], dtype=torch.int64)
-    end   = torch.tensor([END_TOKENS], dtype=torch.int64)
-    input_ids = torch.cat([start, ids, end], dim=1).to(device)
+    ids = tokenizer(text, return_tensors="pt").input_ids.to(device)
+    start = torch.tensor([[START_TOKEN]], dtype=torch.int64, device=device)
+    end   = torch.tensor([END_TOKENS],   dtype=torch.int64, device=device)
+    input_ids = torch.cat([start, ids, end], dim=1)
     attention_mask = torch.ones_like(input_ids)
 
     # Generieren
@@ -74,8 +73,8 @@ def process_single_prompt(prompt: str, voice: str) -> list[int]:
         use_cache=True,
     )
 
-    # letzten START_TOKEN finden & croppen
-    token_to_find   = 128257
+    # Nach letztem START_TOKEN croppen
+    token_to_find = 128257
     token_to_remove = 128258
     idxs = (gen == token_to_find).nonzero(as_tuple=True)[1]
     if idxs.numel() > 0:
@@ -83,16 +82,16 @@ def process_single_prompt(prompt: str, voice: str) -> list[int]:
     else:
         cropped = gen
 
-    # Padding entfernen
+    # Padding entfernen & Länge auf Vielfaches von 7 bringen
     row = cropped[0][cropped[0] != token_to_remove]
-    # Aus Länge ein Vielfaches von 7 machen
     new_len = (row.size(0) // 7) * 7
     trimmed = row[:new_len].tolist()
+
     # Offset abziehen
     return [t - AUDIO_OFFSET for t in trimmed]
 
 def redistribute_codes(code_list: list[int]) -> np.ndarray:
-    # Die 7er‑Blöcke auf 3 Layer verteilen und dekodieren
+    # 7er‑Blöcke auf 3 Layer verteilen
     layer1, layer2, layer3 = [], [], []
     for i in range(len(code_list) // 7):
         b = code_list[7*i : 7*i+7]
@@ -112,27 +111,25 @@ def redistribute_codes(code_list: list[int]) -> np.ndarray:
     audio = snac_model.decode(codes).squeeze().cpu().numpy()
     return audio  # float32 @24 kHz
 
-# — WebSocket‑Endpoint für TTS —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        # 1) Text + Voice empfangen
         msg = await ws.receive_text()
         req = json.loads(msg)
         text  = req.get("text", "")
         voice = req.get("voice", "")
 
-        # 2) Prompt → Code‑Liste
+        # 1) Prompt → Codes → Audio
         with torch.no_grad():
             codes = process_single_prompt(text, voice)
             audio_np = redistribute_codes(codes)
 
-        # 3) In PCM16 konvertieren und senden
+        # 2) In PCM16 wandeln & senden
         pcm16 = (audio_np * 32767).astype(np.int16).tobytes()
         await ws.send_bytes(pcm16)
 
-        # 4) sauber schließen
+        # 3) sauber schließen
         await ws.close()
 
     except WebSocketDisconnect: