Update app.py

app.py

@@ -1,6 +1,8 @@
 import os
 import json
 import asyncio
+import logging
+
 import torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
@@ -10,142 +12,145 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # — HF‑Token & Login —
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
-    login(HF_TOKEN)
+    login(token=HF_TOKEN)
 
-# — Gerät wählen —
+# — Device auswählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# — Modell‑Parameter —
-MODEL_NAME         = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-START_MARKER       = 128259        # <|startoftranscript|>
-RESTART_MARKER     = 128257        # <|startoftranscript_again|>
-EOS_TOKEN          = 128258        # <|endoftranscript|>
-AUDIO_TOKEN_OFFSET = 128266        # Offset zum Zurückrechnen
-BLOCK_TOKENS       = 7             # SNAC erwartet 7 Audio‑Tokens pro Block
-CHUNK_TOKENS       = 50            # Anzahl neuer Tokens pro Generate‑Runde
-
 # — FastAPI instanziieren —
 app = FastAPI()
 
-# — Damit GET / nicht 404 wirft —
+# — Einfacher GET‑Endpunkt, damit / keine 404 liefert —
 @app.get("/")
-async def read_root():
-    return {"message": "Orpheus TTS Server ist live 🎙️"}
+async def root():
+    return {"message": "Hello, world!"}
 
 # — Modelle bei Startup laden —
 @app.on_event("startup")
 async def load_models():
     global tokenizer, model, snac
-    # SNAC laden
+    logging.info("Lade SNAC...")
     snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-    # Tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    # TTS‑LM
+
+    REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
+    logging.info("Lade TTS‑Modell...")
+    tokenizer = AutoTokenizer.from_pretrained(REPO)
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map="auto",
+        REPO,
+        device_map="auto" if device=="cuda" else None,
         torch_dtype=torch.bfloat16 if device=="cuda" else None,
         low_cpu_mem_usage=True
-    )
-    model.config.pad_token_id = EOS_TOKEN
+    ).to(device)
+    model.config.pad_token_id = model.config.eos_token_id
+    logging.info("Modelle geladen ✔️")
 
-# — Eingabe aufbereiten —
+# — Konstanten für Audio‑Token und SNAC‑Blockgröße —
+AUDIO_TOKEN_OFFSET = 128266
+AUDIO_CODE_SIZE    = 4096
+BLOCK_SIZE         = 7
+
+# — 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)
-    start      = torch.tensor([[START_MARKER]], device=device)
-    end        = torch.tensor([[128009, EOS_TOKEN]], device=device)
-    ids        = torch.cat([start, input_ids, end], dim=1)
-    attn_mask  = torch.ones_like(ids)
-    return ids, attn_mask
-
-# — Aus 7 Audio‑Tokens ein PCM‑Block erzeugen —
-def decode_block(block: list[int]) -> bytes:
+    prompt = f"{voice}: {text}"
+    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
+
+# — 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 = block
+    b = clean
     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)
+    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])
     codes = [
-        torch.tensor(l1, device=device).unsqueeze(0),
-        torch.tensor(l2, device=device).unsqueeze(0),
-        torch.tensor(l3, device=device).unsqueeze(0),
+        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),
     ]
     audio = snac.decode(codes).squeeze().cpu().numpy()
-    pcm16 = (audio * 32767).astype("int16").tobytes()
-    return pcm16
-
-# — Generator: kleine Chunks token‑weise erzeugen und block‑weise dekodieren —
-async def generate_and_stream(ws: WebSocket, ids, attn_mask):
-    buffer: list[int] = []
-    past_kvs = None
-
-    while True:
-        # wir rufen model.generate in Häppchen auf
-        outputs = model.generate(
-            input_ids      = ids      if past_kvs is None else None,
-            attention_mask = attn_mask if past_kvs is None else None,
-            past_key_values= past_kvs,
-            use_cache      = True,
-            max_new_tokens = CHUNK_TOKENS,
-            do_sample      = True,
-            temperature    = 0.7,
-            top_p          = 0.95,
-            repetition_penalty = 1.1,
-            eos_token_id   = EOS_TOKEN,
-            pad_token_id   = EOS_TOKEN,
-            return_dict_in_generate = True,
-            output_scores           = False,
-        )
-
-        # update past_kvs
-        past_kvs = outputs.past_key_values
-
-        # erhalte nur die gerade neu generierten Token
-        seq       = outputs.sequences[0]
-        new_tokens = seq[-CHUNK_TOKENS:].tolist() if past_kvs is not None else seq[ids.shape[-1]:].tolist()
-
-        for tok in new_tokens:
-            # Neustart bei erneutem START‑Marker
-            if tok == RESTART_MARKER:
-                buffer = []
-                continue
-            # Ende
-            if tok == EOS_TOKEN:
-                return
-            # Audio‑Code berechnen
-            buffer.append(tok - AUDIO_TOKEN_OFFSET)
-            # sobald 7 Audio‑Tokens, dekodieren und streamen
-            if len(buffer) >= BLOCK_TOKENS:
-                block = buffer[:BLOCK_TOKENS]
-                buffer = buffer[BLOCK_TOKENS:]
-                pcm   = decode_block(block)
-                await ws.send_bytes(pcm)
+    return (audio * 32767).astype("int16").tobytes()
 
-# — 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:
-        data = await ws.receive_text()
-        req  = json.loads(data)
-        text  = req.get("text", "")
-        voice = req.get("voice", "Jakob")
+        # 1) Input empfangen
+        msg = await ws.receive_text()
+        data = json.loads(msg)
+        text  = data.get("text", "")
+        voice = data.get("voice", "Jakob")
+
+        # 2) Prompt → Input‑Tensors
+        input_ids, attention_mask = prepare_inputs(text, voice)
+        past_kvs = None
+        buffer   = []
+
+        # 3) Token‑Loop (du kannst hier auch max_new_tokens=50 fahren,
+        #    indem Du in jedem Durchgang bis zu 50 Token samplet und aufsummierst)
+        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,
+            )
+            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()
+
+            # Ende‑Bedingungen
+            if next_token == model.config.eos_token_id:
+                break
+            if next_token == 128257:
+                # neuer Start → Buffer resetten
+                buffer = []
+                continue
 
-        ids, attn_mask = prepare_inputs(text, voice)
-        await generate_and_stream(ws, ids, attn_mask)
+            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
+                await ws.send_bytes(pcm)
+
+            # Input nur beim ersten Schritt mitgeben
+            input_ids = None
+            attention_mask = None
 
+        # 4) nach Ende sauber schließen
         await ws.close()
     except WebSocketDisconnect:
-        pass
+        logging.info("Client hat WS geschlossen")
     except Exception as e:
-        print("Error in /ws/tts:", e)
+        logging.error(f"Unbehandelter Fehler in /ws/tts: {e}")
         await ws.close(code=1011)
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)