Update app.py

app.py

@@ -1,53 +1,52 @@
-# app.py  ─────────────────────────────────────────────────────────────
-import os, json, asyncio, torch
+# app.py  -------------------------------------------------------------
+import os, json, torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
-from transformers import AutoModelForCausalLM, AutoTokenizer, LogitsProcessor
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
 from transformers.generation.utils import Cache
 from snac import SNAC
 
-# ── 0. HF‑Auth & Device ──────────────────────────────────────────────
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN:
-    login(HF_TOKEN)
+# ── 0. Auth & Device ────────────────────────────────────────────────
+if (tok := os.getenv("HF_TOKEN")):
+    login(tok)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Flash‑Attention‑Bug in PyTorch 2.2.x
-torch.backends.cuda.enable_flash_sdp(False)
-
-# ── 1. Konstanten ────────────────────────────────────────────────────
-REPO              = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-CHUNK_TOKENS      = 50
-START_TOKEN       = 128259
-NEW_BLOCK_TOKEN   = 128257
-EOS_TOKEN         = 128258
-AUDIO_BASE        = 128266
-VALID_AUDIO_IDS   = torch.arange(AUDIO_BASE, AUDIO_BASE + 4096)
-
-# ── 2. Logit‑Processor zum Maskieren ────────────────────────────────
-class AudioLogitMask(LogitsProcessor):
-    def __init__(self, allowed_ids: torch.Tensor):
-        super().__init__()
-        self.allowed = allowed_ids
-
-    def __call__(self, input_ids, scores):
-        # scores shape: [batch, vocab]
+torch.backends.cuda.enable_flash_sdp(False)      # PyTorch‑2.2 fix
+
+# ── 1. Konstanten ───────────────────────────────────────────────────
+REPO             = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
+CHUNK_TOKENS     = 50           #  ≤ 50 → < 1 s Latenz
+START_TOKEN      = 128259
+NEW_BLOCK_TOKEN  = 128257
+EOS_TOKEN        = 128258
+AUDIO_BASE       = 128266
+VALID_AUDIO_IDS  = torch.arange(AUDIO_BASE, AUDIO_BASE + 4096)
+
+# ── 2. Logit‑Maske  (nur Audio‑ und Steuer‑Token)  ──────────────────
+class AudioMask(LogitsProcessor):
+    def __init__(self, allowed: torch.Tensor):        # allowed @device!
+        self.allowed = allowed
+
+    def __call__(self, _ids, scores):
         mask = torch.full_like(scores, float("-inf"))
-        mask[:, self.allowed] = 0
+        mask[:, self.allowed] = 0.0
         return scores + mask
 
 ALLOWED_IDS = torch.cat(
-    [VALID_AUDIO_IDS, torch.tensor([NEW_BLOCK_TOKEN, EOS_TOKEN])]
+    [VALID_AUDIO_IDS,
+     torch.tensor([NEW_BLOCK_TOKEN, EOS_TOKEN])]
 ).to(device)
-MASKER = AudioLogitMask(ALLOWED_IDS)
+MASKER = AudioMask(ALLOWED_IDS)
 
-# ── 3. FastAPI ‑ Grundgerüst ─────────────────────────────────────────
+# ── 3. FastAPI Grundgerüst ──────────────────────────────────────────
 app = FastAPI()
 
 @app.get("/")
-async def ping():
-    return {"msg": "Orpheus‑TTS OK"}
+async def root():
+    return {"msg": "Orpheus‑TTS ready"}
+
+# global handles
+tok = model = snac = None
 
 @app.on_event("startup")
 async def load_models():
@@ -63,11 +62,11 @@ async def load_models():
     model.config.pad_token_id = model.config.eos_token_id
     model.config.use_cache    = True
 
-# ── 4. Hilfs‑Funktionen ─────────────────────────────────────────────
-def build_prompt(text:str, voice:str):
-    base = f"{voice}: {text}"
-    ids  = tok(base, return_tensors="pt").input_ids.to(device)
-    ids  = torch.cat(
+# ── 4. Helper ───────────────────────────────────────────────────────
+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,
@@ -77,29 +76,32 @@ def build_prompt(text:str, voice:str):
     )
     return ids, torch.ones_like(ids)
 
-def decode_snac(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(b7: list[int]) -> bytes:
+    l1, l2, l3 = [], [], []
+    l1.append(b7[0])
+    l2.append(b7[1] - 4096)
+    l3.extend([b7[2] - 8192, b7[3] - 12288])
+    l2.append(b7[4] - 16384)
+    l3.extend([b7[5] - 20480, b7[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()
+
+def new_tokens_only(full_seq, prev_len):
+    """liefert Liste der Tokens, die *neu* hinzukamen"""
+    return full_seq[prev_len:].tolist()
+
+# ── 5. WebSocket‑Endpoint ───────────────────────────────────────────
 @app.websocket("/ws/tts")
 async def tts(ws: WebSocket):
     await ws.accept()
     try:
         req = json.loads(await ws.receive_text())
         ids, attn = build_inputs(req.get("text", ""), req.get("voice", "Jakob"))
+        prompt_len = ids.size(1)               # Länge des Prompts
 
-        past  = None        # Cache
+        past  = None
         buf   = []
 
         while True:
@@ -112,14 +114,18 @@ async def tts(ws: WebSocket):
                 do_sample=True, top_p=0.95, temperature=0.7,
                 return_dict_in_generate=True,
                 use_cache=True,
-                return_legacy_cache=True,       # ⚡ wichtig
+                return_legacy_cache=True,      # wichtig <4.49
             )
 
-            # ⚡ Legacy‑Cache weitergeben
-            past = gen.past_key_values
+            # Cache für den nächsten Loop
+            past = gen.past_key_values if not isinstance(gen.past_key_values, Cache) else gen.past_key_values.to_legacy()
+
+            seq = gen.sequences[0].tolist()
+            new_tok = new_tokens_only(seq, prompt_len)
+            prompt_len = len(seq)              # nächstes Delta
 
-            # die tatsächlich erzeugten neuen Tokens
-            new_tok = gen.sequences[0, -gen.num_generated_tokens :].tolist()
+            if not new_tok:                    # (selten) nichts erzeugt ⇒ weiter
+                continue
 
             for t in new_tok:
                 if t == EOS_TOKEN:
@@ -132,24 +138,20 @@ async def tts(ws: WebSocket):
                     await ws.send_bytes(decode_block(buf))
                     buf.clear()
 
-            # ab jetzt nur Cache, keine neuen IDs mehr nötig
-            ids, attn = None, None
+            ids = None; attn = None            # ab jetzt nur noch Cache
 
-    except (StopIteration, WebSocketDisconnect):
-        pass                           # normales Ende
+    except (StopAsyncIteration, WebSocketDisconnect):
+        pass
     except Exception as e:
         print("WS‑Error:", e)
-        if ws.client_state.name != "DISCONNECTED":
-            await ws.close(code=1011)  # Fehlercode nur, falls noch offen
+        if ws.client_state.name == "CONNECTED":
+            await ws.close(code=1011)
     finally:
-        try:
-            if ws.client_state.name != "DISCONNECTED":
-                await ws.close()       # sauberes Close
-        except RuntimeError:
-            # Starlette hat bereits ein Close‑Frame verschickt
-            pass
-
-# ── 6. Lokaler Test ─────────────────────────────────────────────────
+        if ws.client_state.name == "CONNECTED":
+            await ws.close()
+
+# ── 6. Local run ────────────────────────────────────────────────────
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)
+    import uvicorn, sys
+    port = int(sys.argv[1]) if len(sys.argv) > 1 else 7860
+    uvicorn.run("app:app", host="0.0.0.0", port=port)