Update app.py

app.py

@@ -1,9 +1,10 @@
 import os
 import json
 import asyncio
-import logging
-
 import torch
+# Bugfix für PyTorch 2.2.x Flash‑SDP‑Assertion
+torch.backends.cuda.enable_flash_sdp(False)
+
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
 from snac import SNAC
@@ -12,145 +13,153 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # — HF‑Token & Login —
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
-    login(token=HF_TOKEN)
+    login(HF_TOKEN)
 
-# — Device auswählen —
+# — Device wählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# — FastAPI instanziieren —
+# — FastAPI instanzieren —
 app = FastAPI()
 
-# — Einfacher GET‑Endpunkt, damit / keine 404 liefert —
+# — Hello‑Route, damit GET / kein 404 mehr gibt —
 @app.get("/")
-async def root():
-    return {"message": "Hello, world!"}
+async def read_root():
+    return {"message": "Orpheus TTS WebSocket Server läuft"}
 
-# — Modelle bei Startup laden —
+# — Modelle beim Startup laden —
 @app.on_event("startup")
 async def load_models():
     global tokenizer, model, snac
-    logging.info("Lade SNAC...")
+
+    # SNAC für Audio‑Decoding
     snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
+    # Orpheus‑TTS Base
     REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-    logging.info("Lade TTS‑Modell...")
     tokenizer = AutoTokenizer.from_pretrained(REPO)
     model = AutoModelForCausalLM.from_pretrained(
         REPO,
-        device_map="auto" if device=="cuda" else None,
+        device_map={"": 0} if device=="cuda" else None,
         torch_dtype=torch.bfloat16 if device=="cuda" else None,
-        low_cpu_mem_usage=True
+        low_cpu_mem_usage=True,
+        return_legacy_cache=True  # für compatibility mit past_key_values als Tuple
     ).to(device)
     model.config.pad_token_id = model.config.eos_token_id
-    logging.info("Modelle geladen ✔️")
 
-# — Konstanten für Audio‑Token und SNAC‑Blockgröße —
-AUDIO_TOKEN_OFFSET = 128266
-AUDIO_CODE_SIZE    = 4096
-BLOCK_SIZE         = 7
+    # --- Logit‑Masking vorbereiten ---
+    # reine Audio‑Tokens laufen von 128266 bis 128266+4096-1
+    AUDIO_OFFSET = 128266
+    AUDIO_COUNT  = 4096
+    valid_audio = torch.arange(AUDIO_OFFSET, AUDIO_OFFSET + AUDIO_COUNT, device=device)
+    ctrl_tokens = torch.tensor([128257, model.config.eos_token_id], device=device)
+    global ALLOWED_IDS
+    ALLOWED_IDS = torch.cat([valid_audio, ctrl_tokens])
+
+def sample_from_logits(logits: torch.Tensor) -> int:
+    """
+    Maskt alle IDs außer ALLOWED_IDS und sampelt dann einen Token.
+    """
+    # logits: [1, vocab_size]
+    mask = torch.full_like(logits, float("-inf"))
+    mask[0, ALLOWED_IDS] = 0.0
+    probs = torch.softmax(logits + mask, dim=-1)
+    return torch.multinomial(probs, num_samples=1).item()
 
-# — Hilfsfunktion: Prompt in Token/Mask umwandeln —
 def prepare_inputs(text: str, voice: str):
     prompt = f"{voice}: {text}"
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    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)
-    return ids, mask
-
-# — Hilfsfunktion: Dekodiere genau 7 Audio‑Codes →
-def decode_block(block_tokens: list[int]):
-    # Filter invalid
-    clean = []
-    for t in block_tokens:
-        code = t - AUDIO_TOKEN_OFFSET
-        if 0 <= code < AUDIO_CODE_SIZE:
-            clean.append(code)
-        else:
-            logging.warning(f"Ungültiger Audio‑Token {t}, skippe ihn")
-    if len(clean) != BLOCK_SIZE:
-        # Hier werfen wir raus, um nicht per CUDA‑Assertion zu crashen
-        logging.error(f"Block nicht gültig (saubere Codes={clean}), werfe Exception")
-        raise ValueError(f"Audio‑Block muss {BLOCK_SIZE} sauber haben, habe {len(clean)}")
-    # Baue SNAC‑Eingabe
-    l1, l2, l3 = [], [], []
-    b = clean
-    l1.append(b[0])
-    l2.append(b[1])
-    # das Original verschachtelte Layer‑Mapping
-    l3.append(b[2])
-    l3.append(b[3])
-    l2.append(b[4])
-    l3.append(b[5])
-    l3.append(b[6])
+    input_ids = torch.cat([start, ids, end], dim=1)
+    attention_mask = torch.ones_like(input_ids, device=device)
+    return input_ids, attention_mask
+
+def decode_block(block: list[int]) -> bytes:
+    """
+    Aus 7 gesampelten Audio‑Codes einen PCM‑16‑Byte‐Block dekodieren.
+    Hier erwarten wir block[i] = raw_token - 128266.
+    """
+    layer1, layer2, layer3 = [], [], []
+    b = block
+    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)
+
+    dev = next(snac.parameters()).device
     codes = [
-        torch.tensor(l1, dtype=torch.int64, device=device).unsqueeze(0),
-        torch.tensor(l2, dtype=torch.int64, device=device).unsqueeze(0),
-        torch.tensor(l3, dtype=torch.int64, device=device).unsqueeze(0),
+        torch.tensor(layer1, device=dev).unsqueeze(0),
+        torch.tensor(layer2, device=dev).unsqueeze(0),
+        torch.tensor(layer3, device=dev).unsqueeze(0),
     ]
     audio = snac.decode(codes).squeeze().cpu().numpy()
-    return (audio * 32767).astype("int16").tobytes()
+    # in PCM16 umwandeln
+    pcm16 = (audio * 32767).astype("int16").tobytes()
+    return pcm16
 
-# — WebSocket‑Endpoint für TTS‑Streaming —
+# — WebSocket Endpoint für TTS Streaming —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        # 1) Input empfangen
         msg = await ws.receive_text()
-        data = json.loads(msg)
-        text  = data.get("text", "")
-        voice = data.get("voice", "Jakob")
+        req = json.loads(msg)
+        text  = req.get("text", "")
+        voice = req.get("voice", "Jakob")
 
-        # 2) Prompt → Input‑Tensors
+        # Inputs vorbereiten
         input_ids, attention_mask = prepare_inputs(text, voice)
         past_kvs = None
-        buffer   = []
+        buffer   = []  # sammelt die 7 Audio‑Codes
 
-        # 3) Token‑Loop (du kannst hier auch max_new_tokens=50 fahren,
-        #    indem Du in jedem Durchgang bis zu 50 Token samplet und aufsummierst)
+        # Token‑für‑Token Loop
         while True:
             out = model(
                 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,
+                return_dict=True
             )
-            logits     = out.logits[:, -1, :]
-            past_kvs   = out.past_key_values
-            probs      = torch.softmax(logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1).item()
+            past_kvs = out.past_key_values
+            next_tok = sample_from_logits(out.logits[:, -1, :])
 
-            # Ende‑Bedingungen
-            if next_token == model.config.eos_token_id:
+            # Ende?
+            if next_tok == model.config.eos_token_id:
                 break
-            if next_token == 128257:
-                # neuer Start → Buffer resetten
-                buffer = []
+
+            # Reset bei neuem Audio‑Block‑Start
+            if next_tok == 128257:
+                buffer.clear()
+                input_ids      = torch.tensor([[next_tok]], device=device)
+                attention_mask = torch.ones_like(input_ids)
                 continue
 
-            buffer.append(next_token)
-            # immer, wenn wir ≥7 Codes sammeln, → dekodieren + senden
-            while len(buffer) >= BLOCK_SIZE:
-                block = buffer[:BLOCK_SIZE]
-                buffer = buffer[BLOCK_SIZE:]
-                try:
-                    pcm = decode_block(block)
-                except Exception as e:
-                    logging.error(f"Fehler beim Dekodieren: {e}")
-                    await ws.close(code=1011)
-                    return
+            # Audio‑Code sammeln (Offset abziehen)
+            buffer.append(next_tok - 128266)
+            # sobald wir 7 Codes haben → dekodieren & senden
+            if len(buffer) == 7:
+                pcm = decode_block(buffer)
+                buffer.clear()
                 await ws.send_bytes(pcm)
 
-            # Input nur beim ersten Schritt mitgeben
-            input_ids = None
-            attention_mask = None
+            # nächster Schritt: genau diesen Token wieder einspeisen
+            input_ids      = torch.tensor([[next_tok]], device=device)
+            attention_mask = torch.ones_like(input_ids)
 
-        # 4) nach Ende sauber schließen
+        # sauber beenden
         await ws.close()
     except WebSocketDisconnect:
-        logging.info("Client hat WS geschlossen")
+        pass
     except Exception as e:
-        logging.error(f"Unbehandelter Fehler in /ws/tts: {e}")
+        print("Error in /ws/tts:", e)
         await ws.close(code=1011)
+
+# — CLI zum lokalen Testen —
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)