up4

app.py

@@ -1,4 +1,6 @@
-import os, json, asyncio
+import os
+import json
+import asyncio
 import torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from dotenv import load_dotenv
@@ -6,87 +8,141 @@ from snac import SNAC
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import login, snapshot_download
 
+# ─── ENV & HF TOKEN ────────────────────────────────────────────────────────────
 load_dotenv()
-if (tok := os.getenv("HF_TOKEN")):
-    login(token=tok)
+HF_TOKEN = os.getenv("HF_TOKEN")
+if HF_TOKEN:
+    login(token=HF_TOKEN)
 
+# ─── DEVICE ─────────────────────────────────────────────────────────────────────
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print("Loading SNAC…")
+# ─── SNAC ───────────────────────────────────────────────────────────────────────
+print("Loading SNAC model…")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
+# ─── ORPHEUS ────────────────────────────────────────────────────────────────────
 model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
+
+# pre‑download only the config + safetensors, damit das Image schlank bleibt
 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.*", "vocab.json", "merges.txt" ]
+    ignore_patterns=[
+        "optimizer.pt", "pytorch_model.bin", "training_args.bin",
+        "scheduler.pt", "tokenizer.*", "vocab.json", "merges.txt"
+    ]
 )
 
-print("Loading Orpheus…")
+print("Loading Orpheus model…")
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.bfloat16
+    torch_dtype=torch.bfloat16  # optional: beschleunigt das FP16‑ähnliche Rechnen
 )
 model = model.to(device)
 model.config.pad_token_id = model.config.eos_token_id
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# — Helper Functions (wie gehabt) —
+# ─── HILFSFUNKTIONEN ──────────────────────────────────────────────────────────
 
 def process_prompt(text: str, voice: str):
+    """
+    Baut aus Text+Voice ein batch‑Tensor input_ids für `model.generate`.
+    """
     prompt = f"{voice}: {text}"
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    tok = tokenizer(prompt, return_tensors="pt").to(device)
     start = torch.tensor([[128259]], device=device)
     end   = torch.tensor([[128009, 128260]], device=device)
-    return torch.cat([start, inputs.input_ids, end], dim=1)
+    return torch.cat([start, tok.input_ids, end], dim=1)
 
-def parse_output(ids: torch.LongTensor):
-    st, rm = 128257, 128258
-    idxs = (ids==st).nonzero(as_tuple=True)[1]
-    cropped = ids[:, idxs[-1].item()+1:] if idxs.numel()>0 else ids
-    row = cropped[0][cropped[0]!=rm]
+def parse_output(generated_ids: torch.LongTensor):
+    """
+    Schneidet bis zum letzten 128257 und entfernt 128258, gibt reine Token‑Liste zurück.
+    """
+    START, PAD = 128257, 128258
+    idxs = (generated_ids == START).nonzero(as_tuple=True)[1]
+    if idxs.numel() > 0:
+        cropped = generated_ids[:, idxs[-1].item()+1:]
+    else:
+        cropped = generated_ids
+    row = cropped[0][cropped[0] != PAD]
     return row.tolist()
 
-def redistribute_codes(codes: list[int], snac_model: SNAC):
-    # … genau wie vorher …
-    # return numpy array
+def redistribute_codes(code_list: list[int], snac_model: SNAC):
+    """
+    Verteilt 7er‑Blöcke auf die drei SNAC‑Layer und dekodiert zu Audio (numpy float32).
+    """
+    layer1, layer2, layer3 = [], [], []
+    for i in range((len(code_list) + 1) // 7):
+        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)
+    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()
+
+# ─── FASTAPI ───────���────────────────────────────────────────────────────────────
 
 app = FastAPI()
 
 @app.get("/")
-async def root():
-    return {"status":"ok","msg":"Hello, Orpheus TTS up!"}
+async def healthcheck():
+    return {"status": "ok", "msg": "Hello, Orpheus TTS up!"}
 
 @app.websocket("/ws/tts")
-async def ws_tts(ws: WebSocket):
+async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        msg = json.loads(await ws.receive_text())
-        text, voice = msg.get("text",""), msg.get("voice","Jakob")
-        ids = process_prompt(text, voice)
-        gen = model.generate(
-            input_ids=ids,
-            max_new_tokens=2000,
-            do_sample=True, temperature=0.7, top_p=0.95,
-            repetition_penalty=1.1,
-            eos_token_id=model.config.eos_token_id,
-        )
-        codes    = parse_output(gen)
-        audio_np = redistribute_codes(codes, snac)
-        pcm16    = (audio_np*32767).astype("int16").tobytes()
-        chunk    = 2400*2
-        for i in range(0,len(pcm16),chunk):
-            await ws.send_bytes(pcm16[i:i+chunk])
-            await asyncio.sleep(0.1)
-        await ws.close()
+        while True:
+            # 1) Eintreffende JSON‐Nachricht parsen
+            data = json.loads(await ws.receive_text())
+            text  = data.get("text", "")
+            voice = data.get("voice", "Jakob")
+
+            # 2) Prompt → input_ids
+            ids = process_prompt(text, voice)
+
+            # 3) Token‐Erzeugung
+            gen_ids = model.generate(
+                input_ids=ids,
+                max_new_tokens=2000,      # hier z.B. 20k geht auch, wird aber speicherintensiv
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.1,
+                eos_token_id=model.config.eos_token_id,
+            )
+
+            # 4) Tokens → Code‐Liste → Audio
+            codes    = parse_output(gen_ids)
+            audio_np = redistribute_codes(codes, snac)
+
+            # 5) PCM16‐Stream in 0.1 s‐Blöcken
+            pcm16 = (audio_np * 32767).astype("int16").tobytes()
+            chunk = 2400 * 2
+            for i in range(0, len(pcm16), chunk):
+                await ws.send_bytes(pcm16[i : i+chunk])
+                await asyncio.sleep(0.1)
+
     except WebSocketDisconnect:
-        print("Client left")
+        print("Client disconnected")
     except Exception as e:
-        print("Error in /ws/tts:",e)
+        print("Error in /ws/tts:", e)
         await ws.close(code=1011)
 
-if __name__=="__main__":
+# ─── START ──────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app:app",host="0.0.0.0",port=7860)
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")