Update app.py

app.py

@@ -14,17 +14,15 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-# — Debug: CPU‑Modus zum Entwickeln, später wieder "cuda" —
+# — Device —
 device = "cuda" if torch.cuda.is_available() else "cpu"
-#device = "cpu"
 
 # — Modelle laden —
-print("Loading SNAC model...")
+print("Loading SNAC model…")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
 model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-# optional: explizites snapshot_download (entfernt große Dateien)
-
+# Nur die Konfig + Safetensors, alles andere wird ignoriert
 snapshot_download(
     repo_id=model_name,
     allow_patterns=["config.json", "*.safetensors", "model.safetensors.index.json"],
@@ -35,19 +33,24 @@ snapshot_download(
     ]
 )
 
-print("Loading Orpheus model...")
+print("Loading Orpheus model…")
 model = AutoModelForCausalLM.from_pretrained(
     model_name, torch_dtype=torch.bfloat16
 ).to(device)
 model.config.pad_token_id = model.config.eos_token_id
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# — Konstanten für Audio‑Token →
+# (muss übereinstimmen mit Deinem Training; hier 128266)
+AUDIO_TOKEN_OFFSET = 128266
+
 # — Hilfsfunktionen —
 
 def process_prompt(text: str, voice: str):
     prompt = f"{voice}: {text}"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    # Laut Spezifikation:
+    #  start_token=128259, end_tokens=(128009,128260)
     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)
@@ -55,30 +58,32 @@ def process_prompt(text: str, voice: str):
     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
-
-    # 1) Finde letztes Start-Token, croppe
-    idxs = (generated_ids == token_to_find).nonzero(as_tuple=True)[1]
+    """
+    Croppt nach dem letzten 128257-Start-Token, entfernt Padding (128258)
+    und zieht dann den Audio‑Offset ab, um echte Code‑IDs zu bekommen.
+    """
+    # finde letztes Audio‑Start‑Token
+    token_to_start  = 128257
+    token_to_remove = model.config.eos_token_id  # 128258
+
+    idxs = (generated_ids == token_to_start).nonzero(as_tuple=True)[1]
     if idxs.numel() > 0:
         cut = idxs[-1].item() + 1
         cropped = generated_ids[:, cut:]
     else:
         cropped = generated_ids
 
-    # 2) Entferne Padding-Markierungen
-    rows = []
-    for row in cropped:
-        rows.append(row[row != token_to_remove])
-
-    # 3) Flache Liste zurückgeben
-    return rows[0].tolist()
+    # flatten & remove PAD, dann Offset abziehen
+    flat = cropped[0][cropped[0] != token_to_remove]
+    codes = [(int(t) - AUDIO_TOKEN_OFFSET) for t in flat]
+    return codes
 
 def redistribute_codes(code_list: list[int], snac_model: SNAC):
-    """Verteile die Codes auf drei Layer, dekodiere in Audio."""
+    """
+    Verteilt die flache Code‑Liste in 3 Layers und dekodiert mit SNAC.
+    """
     layer1, layer2, layer3 = [], [], []
-    for i in range((len(code_list) + 1) // 7):
+    for i in range(len(code_list) // 7):
         base = code_list[7*i : 7*i+7]
         layer1.append(base[0])
         layer2.append(base[1] -   4096)
@@ -89,13 +94,11 @@ def redistribute_codes(code_list: list[int], snac_model: SNAC):
         layer3.append(base[6] - 6*4096)
 
     dev = next(snac_model.parameters()).device
-    codes = [
-        torch.tensor(layer1, device=dev).unsqueeze(0),
-        torch.tensor(layer2, device=dev).unsqueeze(0),
-        torch.tensor(layer3, device=dev).unsqueeze(0),
-    ]
-    audio = snac_model.decode(codes)
-    return audio.detach().squeeze().cpu().numpy()  # float32 @24 kHz
+    c1 = torch.tensor(layer1, device=dev).unsqueeze(0)
+    c2 = torch.tensor(layer2, device=dev).unsqueeze(0)
+    c3 = torch.tensor(layer3, device=dev).unsqueeze(0)
+    audio = snac_model.decode([c1, c2, c3])
+    return audio.detach().squeeze().cpu().numpy()
 
 # — FastAPI + WebSocket-Endpoint —
 app = FastAPI()
@@ -110,28 +113,28 @@ async def tts_ws(ws: WebSocket):
             text  = data.get("text", "")
             voice = data.get("voice", "Jakob")
 
-            # 1) Prompt → Tokens
+            # 1) Prompt → Token‑Tensoren
             ids, mask = process_prompt(text, voice)
 
-            # 2) Token-Generation (erst klein testen!)
+            # 2) Generation
             gen_ids = model.generate(
                 input_ids=ids,
                 attention_mask=mask,
-                max_new_tokens=200,      # zum Debuggen klein halten
+                max_new_tokens=200,  # zum Debug
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.95,
                 repetition_penalty=1.1,
-                eos_token_id=128258,
+                eos_token_id=model.config.eos_token_id,
             )
 
-            # 3) Tokens → Code-Liste → Audio
+            # 3) Token → Code‑Liste → Audio (Float32 @24 kHz)
             code_list = parse_output(gen_ids)
             audio_np  = redistribute_codes(code_list, snac)
 
-            # 4) PCM16-Bytes und Stream in 0.1s-Chunks
+            # 4) In 0.1 s‑Chunks (2400 Samples) als PCM16 streamen
             pcm16 = (audio_np * 32767).astype("int16").tobytes()
-            chunk = 2400 * 2  # 2400 samples @24kHz → 0.1s * 2 bytes
+            chunk = 2400 * 2
             for i in range(0, len(pcm16), chunk):
                 await ws.send_bytes(pcm16[i : i+chunk])
                 await asyncio.sleep(0.1)