Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py  ─────────────────────────────────────────────────────────────
+# app.py ──────────────────────────────────────────────────────────────
 import os, json, torch, asyncio
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
@@ -6,109 +6,97 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
 from transformers.generation.utils import Cache
 from snac import SNAC
 
-# ── 0. HF‑Auth & Device ──────────────────────────────────────────────
+# ── 0 · Login & Device ───────────────────────────────────────────────
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(HF_TOKEN)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Flash‑SDP‑Bug (PyTorch 2.2) deaktivieren
-torch.backends.cuda.enable_flash_sdp(False)
-
-# ── 1. Konstanten ────────────────────────────────────────────────────
-REPO              = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-CHUNK_TOKENS      = 50                    # „Mini‑Generate“‑Länge
-START_TOKEN       = 128259
-NEW_BLOCK_TOKEN   = 128257
-EOS_TOKEN         = 128258
-AUDIO_BASE        = 128266
-VALID_AUDIO_IDS   = torch.arange(AUDIO_BASE, AUDIO_BASE + 4096)
-
-# ── 2. Dynamischer Logit‑Masker ──────────────────────────────────────
+torch.backends.cuda.enable_flash_sdp(False)        # CUDA‑Assert‑Fix
+
+# ── 1 · Konstanten ───────────────────────────────────────────────────
+REPO           = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
+CHUNK_TOKENS   = 50
+START_TOKEN    = 128259
+NEW_BLOCK      = 128257
+EOS_TOKEN      = 128258
+AUDIO_BASE     = 128266
+VALID_AUDIO    = torch.arange(AUDIO_BASE, AUDIO_BASE+4096)
+
+# ── 2 · Logit‑Masker ─────────────────────────────────────────────────
 class DynamicAudioMask(LogitsProcessor):
-    """Lässt zu Beginn nur Audio‑ und NEW_BLOCK‑Tokens zu;
-       EOS erst, wenn min_audio_blocks fertig sind."""
-    def __init__(self, audio_ids: torch.Tensor, min_audio_blocks: int = 1):
+    def __init__(self, audio_ids: torch.Tensor, min_blocks:int=1):
         super().__init__()
-        self.audio_ids   = audio_ids
-        self.ctrl_ids    = torch.tensor([NEW_BLOCK_TOKEN], device=audio_ids.device)
-        self.min_blocks  = min_audio_blocks
-        self.blocks_done = 0
-
-    def __call__(self, input_ids, scores):
-        allowed = torch.cat([self.audio_ids, self.ctrl_ids])
-        if self.blocks_done >= self.min_blocks:
-            allowed = torch.cat([allowed,
-                                 torch.tensor([EOS_TOKEN],
-                                 device=scores.device)])
+        self.audio_ids  = audio_ids
+        self.ctrl_ids   = torch.tensor([NEW_BLOCK], device=audio_ids.device)
+        self.min_blocks = min_blocks
+        self.blocks     = 0
+    def __call__(self, inp, scores):
+        allow = torch.cat([self.audio_ids, self.ctrl_ids])
+        if self.blocks >= self.min_blocks:
+            allow = torch.cat([allow,
+                               torch.tensor([EOS_TOKEN], device=scores.device)])
         mask = torch.full_like(scores, float("-inf"))
-        mask[:, allowed] = 0
+        mask[:, allow] = 0
         return scores + mask
 
-# ── 3. FastAPI Grundgerüst ───────────────────────────────────────────
+# ── 3 · FastAPI‑App ──────────────────────────────────────────────────
 app = FastAPI()
 
 @app.get("/")
-async def ping():
-    return {"msg": "Orpheus‑TTS OK"}
+async def root():
+    return {"msg": "Orpheus‑TTS alive"}
 
 @app.on_event("startup")
-async def load_models():
+async def load():
     global tok, model, snac, masker
     print("⏳ Lade Modelle …")
-
-    tok  = AutoTokenizer.from_pretrained(REPO)
-    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-
+    tok   = AutoTokenizer.from_pretrained(REPO)
+    snac  = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
     model = AutoModelForCausalLM.from_pretrained(
         REPO,
         low_cpu_mem_usage=True,
-        device_map={"": 0} if device == "cuda" else None,
-        torch_dtype=torch.bfloat16 if device == "cuda" else None,
-    )
+        device_map={"":0} if device=="cuda" else None,
+        torch_dtype=torch.bfloat16 if device=="cuda" else None)
     model.config.pad_token_id = model.config.eos_token_id
     model.config.use_cache    = True
-
-    masker = DynamicAudioMask(VALID_AUDIO_IDS.to(device))
+    masker = DynamicAudioMask(VALID_AUDIO.to(device))
     print("✅ Modelle geladen")
 
-# ── 4. Hilfs‑Funktionen ──────────────────────────────────────────────
-def build_inputs(text: str, voice: str):
+# ── 4 · Hilfsfunktionen ──────────────────────────────────────────────
+def build_inputs(text:str, voice:str):
     prompt = f"{voice}: {text}"
     ids = tok(prompt, return_tensors="pt").input_ids.to(device)
-    ids = torch.cat(
-        [ torch.tensor([[START_TOKEN]], device=device),
-          ids,
-          torch.tensor([[128009, 128260]], device=device) ], dim=1)
+    ids = torch.cat([torch.tensor([[START_TOKEN]], device=device),
+                     ids,
+                     torch.tensor([[128009,128260]], device=device)],1)
     return ids, torch.ones_like(ids)
 
-def decode_block(block7: list[int]) -> bytes:
-    l1, l2, l3 = [], [], []
-    b = block7
-    l1.append(b[0])
-    l2.append(b[1] -   4096)
-    l3.extend([b[2] -  8192, b[3] - 12288])
-    l2.append(b[4] - 16384)
-    l3.extend([b[5] - 20480, b[6] - 24576])
-
-    codes = [torch.tensor(x, device=device).unsqueeze(0) for x in (l1,l2,l3)]
-    audio = snac.decode(codes).squeeze().cpu().numpy()
-    return (audio * 32767).astype("int16").tobytes()
-
-# ── 5. WebSocket‑Endpoint ────────────────────────────────────────────
+def decode_block(block):
+    l1,l2,l3=[],[],[]
+    l1.append(block[0])
+    l2.append(block[1]-4096)
+    l3.extend([block[2]-8192, block[3]-12288])
+    l2.append(block[4]-16384)
+    l3.extend([block[5]-20480, block[6]-24576])
+    codes=[torch.tensor(x,device=device).unsqueeze(0) for x in (l1,l2,l3)]
+    audio=snac.decode(codes).squeeze().cpu().numpy()
+    return (audio*32767).astype("int16").tobytes()
+
+# ── 5 · WebSocket‑TTS ────────────────────────────────────────────────
 @app.websocket("/ws/tts")
-async def tts(ws: WebSocket):
+async def tts(ws:WebSocket):
     await ws.accept()
     try:
-        req   = json.loads(await ws.receive_text())
-        text  = req.get("text", "")
-        voice = req.get("voice", "Jakob")
+        req = json.loads(await ws.receive_text())
+        text  = req.get("text","")
+        voice = req.get("voice","Jakob")
 
-        ids, attn = build_inputs(text, voice)
-        past      = None
-        last_tok  = None
-        buf       = []
+        ids, attn   = build_inputs(text, voice)
+        total_len   = ids.shape[1]            # Länge des Prompts
+        past        = None
+        last_tok    = None
+        buf         = []
 
         while True:
             out = model.generate(
@@ -118,40 +106,32 @@ async def tts(ws: WebSocket):
                 max_new_tokens  = CHUNK_TOKENS,
                 logits_processor= [masker],
                 do_sample=True, temperature=0.7, top_p=0.95,
-                return_dict_in_generate=True,
-                use_cache=True,
-                return_legacy_cache=True       # verhindert Cache‑Warnung
-            )
+                use_cache=True, return_dict_in_generate=True,
+                return_legacy_cache=True)
 
             pkv = out.past_key_values
-            if isinstance(pkv, Cache):         # HF ≥ 4.47
-                pkv = pkv.to_legacy()
+            if isinstance(pkv, Cache): pkv = pkv.to_legacy()
             past = pkv
 
-            new_toks = out.sequences[0, -out.num_generated_tokens:].tolist()
-            print("new tokens:", new_toks[:32])  # Debug‑Ausgabe
+            seq   = out.sequences[0].tolist()
+            new   = seq[total_len:]            # alles *nach* Prompt
+            total_len = len(seq)               # fürs nächste Mal
+            print("new tokens:", new[:32])
 
-            if not new_toks:
+            if not new:                        # nichts generiert
                 raise StopIteration
 
-            for t in new_toks:
+            for t in new:
                 last_tok = t
-
-                if t == EOS_TOKEN:
-                    raise StopIteration
-
-                if t == NEW_BLOCK_TOKEN:
-                    buf.clear()
-                    continue
-
-                buf.append(t - AUDIO_BASE)
-                if len(buf) == 7:
+                if t == EOS_TOKEN: raise StopIteration
+                if t == NEW_BLOCK:
+                    buf.clear();  continue
+                buf.append(t-AUDIO_BASE)
+                if len(buf)==7:
                     await ws.send_bytes(decode_block(buf))
                     buf.clear()
-                    masker.blocks_done += 1
-
-            # ab jetzt nur noch 1 Token + Cache
-            ids, attn = None, None
+                    masker.blocks += 1
+            ids, attn = None, None             # ab jetzt 1‑Token‑Step
 
     except (StopIteration, WebSocketDisconnect):
         pass
@@ -161,12 +141,10 @@ async def tts(ws: WebSocket):
             await ws.close(code=1011)
     finally:
         if ws.client_state.name != "DISCONNECTED":
-            try:
-                await ws.close()
-            except RuntimeError:
-                pass
+            try: await ws.close()
+            except RuntimeError: pass
 
-# ── 6. Lokaler Test‑Start ───────────────────────────────────────────
+# ── 6 · local run ────────────────────────────────────────────────────
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("app:app", host="0.0.0.0", port=7860)