Update app.py

app.py

@@ -1,154 +1,162 @@
 import os
 import json
 import asyncio
-import numpy as np
 import torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.responses import PlainTextResponse
 from dotenv import load_dotenv
 from snac import SNAC
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import login, snapshot_download
+from peft import PeftModel
 
 # — ENV & HF‑AUTH —
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
-    login(token=HF_TOKEN)
+    # automatisch über huggingface-cli eingeloggt
+    os.environ["HUGGINGFACE_HUB_TOKEN"] = HF_TOKEN
 
-# — Device —
+# — FastAPI →
+app = FastAPI()
+
+@app.get("/")
+async def hello():
+    return PlainTextResponse("Hallo Welt!")
+
+# — Device konfigurieren —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# — Modelle laden —
-print("Loading SNAC model...")
+# — SNAC laden —
+print("Loading SNAC model…")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
+# — Orpheus/Kartoffel‑3B über PEFT laden —
 model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-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.json", "tokenizer_config.json",
-        "special_tokens_map.json", "vocab.json", "merges.txt", "tokenizer.*"
-    ]
+print(f"Loading base LM + PEFT from {model_name}…")
+base = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
 )
-
-print("Loading Orpheus model...")
-model = AutoModelForCausalLM.from_pretrained(
+model = PeftModel.from_pretrained(
+    base,
     model_name,
-    torch_dtype=torch.bfloat16
-).to(device)
-model.config.pad_token_id = model.config.eos_token_id
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+    device_map="auto",
+)
+model.eval()
 
-# — Konstanten —
-AUDIO_TOKEN_OFFSET = 128266  # globaler Offset der Audio‑Tokens
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+# sicherstellen, dass pad_token_id gesetzt ist
+model.config.pad_token_id = model.config.eos_token_id
 
 # — Hilfsfunktionen —
-
-def process_prompt(text: str, voice: str):
-    """Bereitet input_ids und attention_mask für das Modell vor."""
-    prompt = f"{voice}: {text}"
-    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):
-    """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]
+def prepare_prompt(text: str, voice: str):
+    """Setzt Start‑ und End‑Marker um den eigentlichen Prompt."""
+    if voice:
+        full = f"{voice}: {text}"
+    else:
+        full = text
+    start = torch.tensor([[128259]], dtype=torch.int64)  # BOS für Audio
+    end   = torch.tensor([[128009, 128260]], dtype=torch.int64)  # ggf. Speaker‑ID + Marker
+    enc = tokenizer(full, return_tensors="pt").input_ids
+    seq = torch.cat([start, enc, end], dim=1).to(device)
+    mask = torch.ones_like(seq).to(device)
+    return seq, mask
+
+def extract_audio_tokens(generated: torch.LongTensor):
+    """Croppe alles bis zum echten Audio-Start, entferne EOS und mache 7er-Batches."""
+    bos_tok = 128257
+    eos_tok = 128258
+
+    # letzten Start‑Token finden und ab da weiter
+    idxs = (generated == bos_tok).nonzero(as_tuple=True)[1]
     if idxs.numel() > 0:
         cut = idxs[-1].item() + 1
-        cropped = generated_ids[:, cut:]
+        cropped = generated[:, cut:]
     else:
-        cropped = generated_ids
-
-    row = cropped[0]
-    # entferne das EOS‑Marker‑Token
-    return row[row != token_to_remove].tolist()
-
-def redistribute_codes(raw_codes: list[int], snac_model: SNAC):
-    """
-    Subtrahiere erst den globalen Offset, dann packe in 7er-Blöcke und dekodiere.
-    Unvollständige Reste (<7 Tokens) werden verworfen.
-    """
-    # 1) Offset abziehen
-    codes = [c - AUDIO_TOKEN_OFFSET for c in raw_codes]
-
-    # 2) Nur ganze 7er‑Blöcke
-    n_blocks = len(codes) // 7
-    if n_blocks == 0:
-        return np.zeros(0, dtype=np.float32)
-
-    layer1, layer2, layer3 = [], [], []
-    for i in range(n_blocks):
-        b = codes[7*i : 7*i+7]
-        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)
-
-    # 3) SNAC‑Layer‑Tensors bauen und dekodieren
-    dev = next(snac_model.parameters()).device
-    t1 = torch.tensor(layer1, device=dev).unsqueeze(0)
-    t2 = torch.tensor(layer2, device=dev).unsqueeze(0)
-    t3 = torch.tensor(layer3, device=dev).unsqueeze(0)
-    audio = snac_model.decode([t1, t2, t3])
-
-    return audio.detach().squeeze().cpu().numpy()
-
-# — FastAPI Setup —
-
-app = FastAPI()
-
-@app.get("/")
-async def hello():
-    return {"message": "Hello World"}
-
+        cropped = generated
+
+    # EOS‑Marker entfernen
+    flat = cropped[0][cropped[0] != eos_tok]
+
+    # nur ein Vielfaches von 7 behalten
+    length = (flat.size(0) // 7) * 7
+    flat = flat[:length]
+
+    # Die Audio‑Token beginnen ab Offset 128266
+    return [(t.item() - 128266) for t in flat]
+
+def decode_and_stream(tokens: list[int], ws: WebSocket):
+    """Wandelt 7er‑Gruppen in Wave‑Samples um und streamt in 0.1 s Chunks."""
+    # gruppiere nach 7 und dekodiere jeweils
+    pcm16 = bytearray()
+    offset = 0
+    while offset + 7 <= len(tokens):
+        block = tokens[offset:offset+7]
+        offset += 7
+
+        # SNAC‑Input vorbereiten
+        # Layer‑1: direkt, Layer‑2/3 mit Offsets
+        l1, l2, l3 = [], [], []
+        l1.append(block[0])
+        l2.append(block[1] -   4096)
+        l3.append(block[2] - 2*4096)
+        l3.append(block[3] - 3*4096)
+        l2.append(block[4] - 4*4096)
+        l3.append(block[5] - 5*4096)
+        l3.append(block[6] - 6*4096)
+
+        t1 = torch.tensor(l1, device=device).unsqueeze(0)
+        t2 = torch.tensor(l2, device=device).unsqueeze(0)
+        t3 = torch.tensor(l3, device=device).unsqueeze(0)
+        audio = snac.decode([t1, t2, t3]).squeeze().cpu().numpy()
+
+        # in PCM16 @24 kHz
+        pcm = (audio * 32767).astype("int16").tobytes()
+        pcm16.extend(pcm)
+
+    # in 0.1 s‑Chunks (2400 Samples ×2 Bytes)
+    chunk_size = 2400 * 2
+    for i in range(0, len(pcm16), chunk_size):
+        ws.send_bytes(pcm16[i : i+chunk_size])
+        # ohne Pause kann das WebSocket überlastet werden
+        asyncio.sleep(0.1)
+
+# — WebSocket TTS Endpoint —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
         while True:
-            # Empfang: {"text":"...", "voice":"Jakob"}
-            data = json.loads(await ws.receive_text())
-            text  = data.get("text", "")
-            voice = data.get("voice", "Jakob")
+            raw = await ws.receive_text()
+            req = json.loads(raw)
+            text  = req.get("text", "")
+            voice = req.get("voice", "")
 
-            # 1) Eingabe → Tokens
-            ids, mask = process_prompt(text, voice)
+            # Prompt vorbereiten
+            ids, mask = prepare_prompt(text, voice)
 
-            # 2) Generierung
-            gen_ids = model.generate(
+            # Audio‑Token generieren
+            gen = model.generate(
                 input_ids=ids,
                 attention_mask=mask,
-                max_new_tokens=2000,    # nach Bedarf hochsetzen
+                max_new_tokens=4000,
                 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,
+                forced_bos_token_id=128259,
+                use_cache=True,
             )
 
-            # 3) Tokens → Audio‑Codes → PCM
-            raw_codes = parse_output(gen_ids)
-            audio_np  = redistribute_codes(raw_codes, snac)
-            pcm16     = (audio_np * 32767).astype("int16").tobytes()
+            codes = extract_audio_tokens(gen)
+            # stream synchron
+            await decode_and_stream(codes, ws)
 
-            # 4) Stream in 0.1 s‑Chunks
-            chunk = 2400 * 2  # 2400 Samples @24 kHz = 0.1 s * 2 Bytes
-            for i in range(0, len(pcm16), chunk):
-                await ws.send_bytes(pcm16[i : i+chunk])
-                await asyncio.sleep(0.1)
+            # sauber schließen
+            await ws.close(code=1000)
+            break
 
     except WebSocketDisconnect:
         print("Client disconnected")
@@ -156,6 +164,7 @@ async def tts_ws(ws: WebSocket):
         print("Error in /ws/tts:", e)
         await ws.close(code=1011)
 
+# — Lokal starten —
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("app:app", host="0.0.0.0", port=7860)