Update app.py

app.py

@@ -1,46 +1,47 @@
 # app.py  ─────────────────────────────────────────────────────────────
-import os, json, asyncio, torch
+import os, json, torch, asyncio
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
-from transformers import (AutoTokenizer, AutoModelForCausalLM, LogitsProcessor)
+from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
 from transformers.generation.utils import Cache
 from snac import SNAC
 
-# ── 0. HF‑Login & Device ─────────────────────────────────────────────
+# ── 0. HF‑Auth & Device ──────────────────────────────────────────────
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(HF_TOKEN)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Flash‑Attention‑Bug in PyTorch 2.2.x umgehen
+# 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                      # pro mini‑generate
+CHUNK_TOKENS      = 50                    # „Mini‑Generate“‑Länge
 START_TOKEN       = 128259
 NEW_BLOCK_TOKEN   = 128257
 EOS_TOKEN         = 128258
-AUDIO_BASE        = 128266                 # erster Audio‑Code
+AUDIO_BASE        = 128266
 VALID_AUDIO_IDS   = torch.arange(AUDIO_BASE, AUDIO_BASE + 4096)
 
 # ── 2. Dynamischer Logit‑Masker ──────────────────────────────────────
 class DynamicAudioMask(LogitsProcessor):
-    """
-    blockt EOS, bis mindestens `min_audio_blocks` gesendet wurden
-    """
+    """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):
         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
+        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:              # jetzt darf EOS dazu
-            allowed = torch.cat([allowed, torch.tensor([EOS_TOKEN], device=scores.device)])
+        if self.blocks_done >= self.min_blocks:
+            allowed = torch.cat([allowed,
+                                 torch.tensor([EOS_TOKEN],
+                                 device=scores.device)])
         mask = torch.full_like(scores, float("-inf"))
         mask[:, allowed] = 0
         return scores + mask
@@ -50,7 +51,7 @@ app = FastAPI()
 
 @app.get("/")
 async def ping():
-    return {"msg": "Orpheus‑TTS up & running"}
+    return {"msg": "Orpheus‑TTS OK"}
 
 @app.on_event("startup")
 async def load_models():
@@ -79,30 +80,23 @@ def build_inputs(text: str, voice: str):
     ids = torch.cat(
         [ torch.tensor([[START_TOKEN]], device=device),
           ids,
-          torch.tensor([[128009, 128260]], device=device) ],
-        dim=1,
-    )
-    attn = torch.ones_like(ids)
-    return ids, attn
+          torch.tensor([[128009, 128260]], device=device) ], dim=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])
+    l3.extend([b[2] -  8192, b[3] - 12288])
     l2.append(b[4] - 16384)
     l3.extend([b[5] - 20480, b[6] - 24576])
 
-    codes = [
-        torch.tensor(l1, device=device).unsqueeze(0),
-        torch.tensor(l2, device=device).unsqueeze(0),
-        torch.tensor(l3, device=device).unsqueeze(0),
-    ]
+    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‑Endpoint ───────────────────────────────────────
+# ── 5. WebSocket‑Endpoint ────────────────────────────────────────────
 @app.websocket("/ws/tts")
 async def tts(ws: WebSocket):
     await ws.accept()
@@ -111,9 +105,9 @@ async def tts(ws: WebSocket):
         text  = req.get("text", "")
         voice = req.get("voice", "Jakob")
 
-        ids, attn = build_inputs(text, voice)   # vollständiger Prompt
+        ids, attn = build_inputs(text, voice)
         past      = None
-        last_tok  = None                       # <- NEU
+        last_tok  = None
         buf       = []
 
         while True:
@@ -126,24 +120,22 @@ async def tts(ws: WebSocket):
                 do_sample=True, temperature=0.7, top_p=0.95,
                 return_dict_in_generate=True,
                 use_cache=True,
-                return_legacy_cache=True,       # <- Warnung unterdrücken
+                return_legacy_cache=True       # verhindert Cache‑Warnung
             )
 
-            # ----- Cache & neue Token --------------------------------------
             pkv = out.past_key_values
-            if isinstance(pkv, Cache):          # HF >= 4.47
+            if isinstance(pkv, Cache):         # HF ≥ 4.47
                 pkv = pkv.to_legacy()
             past = pkv
 
-            new = out.sequences[0, -out.num_generated_tokens :].tolist()
-            print("new tokens:", new[:20])      # Debug‑Print
+            new_toks = out.sequences[0, -out.num_generated_tokens:].tolist()
+            print("new tokens:", new_toks[:32])  # Debug‑Ausgabe
 
-            if not new:                         # Safety – nichts erzeugt
+            if not new_toks:
                 raise StopIteration
 
-            # ----- Token‑Handling ------------------------------------------
-            for t in new:
-                last_tok = t                   # speichern für nächste Runde
+            for t in new_toks:
+                last_tok = t
 
                 if t == EOS_TOKEN:
                     raise StopIteration
@@ -156,9 +148,9 @@ async def tts(ws: WebSocket):
                 if len(buf) == 7:
                     await ws.send_bytes(decode_block(buf))
                     buf.clear()
-                    masker.blocks_done += 1     # nach 1. Block darf EOS
+                    masker.blocks_done += 1
 
-            # ab nächster Runde nur 1 Token + Cache
+            # ab jetzt nur noch 1 Token + Cache
             ids, attn = None, None
 
     except (StopIteration, WebSocketDisconnect):
@@ -174,7 +166,7 @@ async def tts(ws: WebSocket):
             except RuntimeError:
                 pass
 
-# ── 6. Lokaler Start (uvicorn) ───────────────────────────────────────
+# ── 6. Lokaler Test‑Start ───────────────────────────────────────────
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("app:app", host="0.0.0.0", port=7860)