Update app.py

app.py

@@ -1,186 +1,139 @@
 import os
 import json
+import asyncio
 import torch
-import numpy as np
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
-from transformers import AutoModelForCausalLM, AutoTokenizer
 from snac import SNAC
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# — HF‑Token & Login (falls gesetzt) —
+# — HF‑Token & Login —
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(HF_TOKEN)
 
-# — Device auswählen —
+# — Device wählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# — FastAPI instanziieren —
 app = FastAPI()
 
+# — Hello‑Route, damit GET / nicht 404 wirft —
 @app.get("/")
 async def read_root():
     return {"message": "Hello, world!"}
 
-# — Globale Modelle —
-model = None
-tokenizer = None
-snac_model = None
-
+# — Modelle bei Startup laden —
 @app.on_event("startup")
 async def load_models():
-    global model, tokenizer, snac_model
-
-    # 1) SNAC laden
-    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-
-    # 2) Orpheus‑TTS (public “natural”-Variante)
+    global tokenizer, model, snac
+    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
     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,
+        device_map="auto",
         torch_dtype=torch.bfloat16 if device == "cuda" else None,
         low_cpu_mem_usage=True
-    ).to(device)
+    )
+    # Für pad-token fallback auf eos
     model.config.pad_token_id = model.config.eos_token_id
 
-# — Marker und Offsets —
-START_TOKEN  = 128259
-END_TOKENS   = [128009, 128260]
-AUDIO_OFFSET = 128266
-
-def process_single_prompt(prompt: str, voice: str) -> list[int]:
-    # Prompt zusammenstellen
-    text = f"{voice}: {prompt}" if voice and voice != "in_prompt" else prompt
-
-    # Tokenize + Marker
-    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)
+# — Hilfsfunktionen — 
+START_TOKEN       = 128259
+END_TOKENS        = [128009, 128260]
+RESET_TOKEN       = 128257
+AUDIO_OFFSET      = 128266
+EOS_TOKEN         = model.config.eos_token_id if 'model' in globals() else 128258
+
+def prepare_inputs(text: str, voice: str):
+    prompt = f"{voice}: {text}"
+    ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    start = torch.tensor([[START_TOKEN]], device=device)
+    end   = torch.tensor([END_TOKENS], device=device)
     input_ids = torch.cat([start, ids, end], dim=1)
     attention_mask = torch.ones_like(input_ids)
-
-    # Generieren
-    gen = model.generate(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        max_new_tokens=4000,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.95,
-        repetition_penalty=1.1,
-        eos_token_id=128258,
-        use_cache=True,
-    )
-
-    # 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:
-        cropped = gen[:, idxs[-1] + 1 :]
-    else:
-        cropped = gen
-
-    # Padding entfernen & Länge auf Vielfaches von 7 bringen
-    row = cropped[0][cropped[0] != token_to_remove]
-    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:
-    # 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]
-        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)
-
+    return input_ids, attention_mask
+
+def decode_block(block: list[int]):
+    # aus genau 7 Audio‑Codes ein PCM‑Byte‑Block bauen
+    l1, l2, l3 = [], [], []
+    b = block
+    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)
     codes = [
-        torch.tensor(layer1, device=device).unsqueeze(0),
-        torch.tensor(layer2, device=device).unsqueeze(0),
-        torch.tensor(layer3, device=device).unsqueeze(0),
+        torch.tensor(l1, device=device).unsqueeze(0),
+        torch.tensor(l2, device=device).unsqueeze(0),
+        torch.tensor(l3, device=device).unsqueeze(0),
     ]
-    audio = snac_model.decode(codes).squeeze().cpu().numpy()
-    return audio  # float32 @24 kHz
+    audio = snac.decode(codes).squeeze().cpu().numpy()
+    return (audio * 32767).astype("int16").tobytes()
 
+# — WebSocket‑Endpoint für TTS Streaming —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        msg   = await ws.receive_text()
+        msg = await ws.receive_text()
         req   = json.loads(msg)
         text  = req.get("text", "")
-        voice = req.get("voice", "")
+        voice = req.get("voice", "Jakob")
 
-        # 1) Prompt vorbereiten
         input_ids, attention_mask = prepare_inputs(text, voice)
         past_kvs = None
-        buffer   = []
+        collected = []
 
-        # 2) Token‑für‑Token (oder in kleinen Blöcken)
+        # Token‑für‑Token mit eigener Sampling‑Schleife
         while True:
-            # Nur max_new_tokens=50 pro Aufruf
-            out = model.generate(
-                input_ids=input_ids   if past_kvs is None else None,
+            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,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.1,
-                max_new_tokens=50,
-                eos_token_id=128258,
-                return_dict_in_generate=True,
-                output_past_key_values=True,
-                return_legacy_cache=True,  # falls Ihr noch das alte past_key_values-Format braucht
             )
-
-            # Extrahiere neue Token (ohne die already generated ones)
-            new_ids = out.sequences[0, input_ids.shape[-1]:].tolist()
+            logits = out.logits[:, -1, :]
             past_kvs = out.past_key_values
 
-            for tok in new_ids:
-                if tok == model.config.eos_token_id:
-                    # Stream zu Ende
-                    break
-                if tok == 128257:    # Reset-Start‑Marker
-                    buffer = []
-                    continue
-                buffer.append(tok - AUDIO_OFFSET)
-
-                # Sobald wir 7 Audio‑Codes gesammelt haben → dekodieren & schicken
-                if len(buffer) == 7:
-                    pcm = decode_block(buffer)
-                    buffer = []
-                    await ws.send_bytes(pcm)
-
-            # Wenn EOS im Chunk war, abbrechen
-            if model.config.eos_token_id in new_ids:
-                break
+            # Sampling
+            probs = torch.softmax(logits, dim=-1)
+            nxt   = torch.multinomial(probs, num_samples=1).item()
 
-            # Danach weiter mit nächsten 50 Tokens,
-            # input_ids & attention_mask nur beim ersten Aufruf nötig
+            # EOS → fertig
+            if nxt == EOS_TOKEN:
+                break
+            # RESET → alte Sammlung verwerfen
+            if nxt == RESET_TOKEN:
+                collected = []
+                # und input_ids für nächsten Durchlauf auf None setzen
+                input_ids = None
+                attention_mask = None
+                continue
+
+            # Audio‑Code abziehen & sammeln
+            collected.append(nxt - AUDIO_OFFSET)
+            # jede 7 Codes → dekodieren & streamen
+            if len(collected) == 7:
+                pcm = decode_block(collected)
+                collected = []
+                await ws.send_bytes(pcm)
+
+            # nur beim allerersten Schritt mit IDs arbeiten
             input_ids = None
             attention_mask = None
 
-        # 3) Am Ende WebSocket sauber schließen
+        # Stream sauber beenden
         await ws.close()
 
     except WebSocketDisconnect:
+        # Client hat Disconnect gemacht → nichts tun
         pass
+
     except Exception as e:
+        # auf Fehler 1011 senden
         print("Error 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)