up6

app.py

@@ -40,9 +40,11 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16
 ).to(device)
 model.config.pad_token_id = model.config.eos_token_id
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# — Konstanten —
+AUDIO_TOKEN_OFFSET = 128266  # globaler Offset der Audio‑Tokens
+
 # — Hilfsfunktionen —
 
 def process_prompt(text: str, voice: str):
@@ -67,39 +69,41 @@ def parse_output(generated_ids: torch.LongTensor):
     else:
         cropped = generated_ids
 
-    # Entferne EOS‑Token
     row = cropped[0]
+    # entferne das EOS‑Marker‑Token
     return row[row != token_to_remove].tolist()
 
-def redistribute_codes(code_list: list[int], snac_model: SNAC):
+def redistribute_codes(raw_codes: list[int], snac_model: SNAC):
     """
-    Verteilt die Token nur in kompletten 7er‑Blöcken auf die drei SNAC‑Layer
-    und dekodiert in Audio. Unvollständige Reste (<7 Tokens) werden verworfen.
+    Subtrahiere erst den globalen Offset, dann packe in 7er-Blöcke und dekodiere.
+    Unvollständige Reste (<7 Tokens) werden verworfen.
     """
-    n_blocks = len(code_list) // 7
-    layer1, layer2, layer3 = [], [], []
+    # 1) Offset abziehen
+    codes = [c - AUDIO_TOKEN_OFFSET for c in raw_codes]
 
-    for i in range(n_blocks):
-        base = code_list[7*i : 7*i + 7]
-        layer1.append(base[0])
-        layer2.append(base[1] -   4096)
-        layer3.append(base[2] - 2*4096)
-        layer3.append(base[3] - 3*4096)
-        layer2.append(base[4] - 4*4096)
-        layer3.append(base[5] - 5*4096)
-        layer3.append(base[6] - 6*4096)
-
-    if not layer1:
-        # kein kompletter Block → leeres Audio
+    # 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
-    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)
+    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 —
@@ -107,27 +111,27 @@ def redistribute_codes(code_list: list[int], snac_model: SNAC):
 app = FastAPI()
 
 @app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+async def hello():
+    return {"message": "Hello World"}
 
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
         while True:
-            # Erwartet JSON: {"text": "...", "voice": "Jakob"}
+            # Empfang: {"text":"...", "voice":"Jakob"}
             data = json.loads(await ws.receive_text())
             text  = data.get("text", "")
             voice = data.get("voice", "Jakob")
 
-            # 1) Tokens vorbereiten
+            # 1) Eingabe → Tokens
             ids, mask = process_prompt(text, voice)
 
             # 2) Generierung
             gen_ids = model.generate(
                 input_ids=ids,
                 attention_mask=mask,
-                max_new_tokens=2000,    # hier nach Bedarf anpassen
+                max_new_tokens=2000,    # nach Bedarf hochsetzen
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.95,
@@ -135,18 +139,17 @@ async def tts_ws(ws: WebSocket):
                 eos_token_id=model.config.eos_token_id,
             )
 
-            # 3) Tokens → Code-Liste → Audio
-            codes   = parse_output(gen_ids)
-            audio_np = redistribute_codes(codes, snac)
+            # 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()
 
-            # 4) in 0.1s‑Stücken PCM16 streamen
-            pcm16 = (audio_np * 32767).astype("int16").tobytes()
-            chunk = 2400 * 2  # 2400 samples @24kHz = 0.1s * 2 bytes
+            # 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)
 
-        # Ende der while‐Schleife
     except WebSocketDisconnect:
         print("Client disconnected")
     except Exception as e: