Update app.py

app.py

@@ -12,51 +12,55 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(HF_TOKEN)
 
-# — Device wählen —
+# — Gerät wählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# — Modell‑Parameter —
+MODEL_NAME         = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
+START_MARKER       = 128259        # <|startoftranscript|>
+RESTART_MARKER     = 128257        # <|startoftranscript_again|>
+EOS_TOKEN          = 128258        # <|endoftranscript|>
+AUDIO_TOKEN_OFFSET = 128266        # Offset zum Zurückrechnen
+BLOCK_TOKENS       = 7             # SNAC erwartet 7 Audio‑Tokens pro Block
+CHUNK_TOKENS       = 50            # Anzahl neuer Tokens pro Generate‑Runde
+
 # — FastAPI instanziieren —
 app = FastAPI()
 
-# — Hello‑Route, damit GET / nicht 404 wirft —
+# — Damit GET / nicht 404 wirft —
 @app.get("/")
 async def read_root():
-    return {"message": "Hello, world!"}
+    return {"message": "Orpheus TTS Server ist live 🎙️"}
 
 # — Modelle bei Startup laden —
 @app.on_event("startup")
 async def load_models():
     global tokenizer, model, snac
+    # SNAC laden
     snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-    REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-    tokenizer = AutoTokenizer.from_pretrained(REPO)
+    # Tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    # TTS‑LM
     model = AutoModelForCausalLM.from_pretrained(
-        REPO,
+        MODEL_NAME,
         device_map="auto",
-        torch_dtype=torch.bfloat16 if device == "cuda" else None,
+        torch_dtype=torch.bfloat16 if device=="cuda" else None,
         low_cpu_mem_usage=True
     )
-    # Für pad-token fallback auf eos
-    model.config.pad_token_id = model.config.eos_token_id
-
-# — Hilfsfunktionen — 
-START_TOKEN       = 128259
-END_TOKENS        = [128009, 128260]
-RESET_TOKEN       = 128257
-AUDIO_OFFSET      = 128266
-EOS_TOKEN         = model.config.eos_token_id if 'model' in globals() else 128258
+    model.config.pad_token_id = EOS_TOKEN
 
+# — Eingabe aufbereiten —
 def prepare_inputs(text: str, voice: str):
-    prompt = f"{voice}: {text}"
-    ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    start = torch.tensor([[START_TOKEN]], device=device)
-    end   = torch.tensor([END_TOKENS], device=device)
-    input_ids = torch.cat([start, ids, end], dim=1)
-    attention_mask = torch.ones_like(input_ids)
-    return input_ids, attention_mask
-
-def decode_block(block: list[int]):
-    # aus genau 7 Audio‑Codes ein PCM‑Byte‑Block bauen
+    prompt     = f"{voice}: {text}"
+    input_ids  = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    start      = torch.tensor([[START_MARKER]], device=device)
+    end        = torch.tensor([[128009, EOS_TOKEN]], device=device)
+    ids        = torch.cat([start, input_ids, end], dim=1)
+    attn_mask  = torch.ones_like(ids)
+    return ids, attn_mask
+
+# — Aus 7 Audio‑Tokens ein PCM‑Block erzeugen —
+def decode_block(block: list[int]) -> bytes:
     l1, l2, l3 = [], [], []
     b = block
     l1.append(b[0])
@@ -72,68 +76,76 @@ def decode_block(block: list[int]):
         torch.tensor(l3, device=device).unsqueeze(0),
     ]
     audio = snac.decode(codes).squeeze().cpu().numpy()
-    return (audio * 32767).astype("int16").tobytes()
+    pcm16 = (audio * 32767).astype("int16").tobytes()
+    return pcm16
+
+# — Generator: kleine Chunks token‑weise erzeugen und block‑weise dekodieren —
+async def generate_and_stream(ws: WebSocket, ids, attn_mask):
+    buffer: list[int] = []
+    past_kvs = None
+
+    while True:
+        # wir rufen model.generate in Häppchen auf
+        outputs = model.generate(
+            input_ids      = ids      if past_kvs is None else None,
+            attention_mask = attn_mask if past_kvs is None else None,
+            past_key_values= past_kvs,
+            use_cache      = True,
+            max_new_tokens = CHUNK_TOKENS,
+            do_sample      = True,
+            temperature    = 0.7,
+            top_p          = 0.95,
+            repetition_penalty = 1.1,
+            eos_token_id   = EOS_TOKEN,
+            pad_token_id   = EOS_TOKEN,
+            return_dict_in_generate = True,
+            output_scores           = False,
+        )
+
+        # update past_kvs
+        past_kvs = outputs.past_key_values
+
+        # erhalte nur die gerade neu generierten Token
+        seq       = outputs.sequences[0]
+        new_tokens = seq[-CHUNK_TOKENS:].tolist() if past_kvs is not None else seq[ids.shape[-1]:].tolist()
+
+        for tok in new_tokens:
+            # Neustart bei erneutem START‑Marker
+            if tok == RESTART_MARKER:
+                buffer = []
+                continue
+            # Ende
+            if tok == EOS_TOKEN:
+                return
+            # Audio‑Code berechnen
+            buffer.append(tok - AUDIO_TOKEN_OFFSET)
+            # sobald 7 Audio‑Tokens, dekodieren und streamen
+            if len(buffer) >= BLOCK_TOKENS:
+                block = buffer[:BLOCK_TOKENS]
+                buffer = buffer[BLOCK_TOKENS:]
+                pcm   = decode_block(block)
+                await ws.send_bytes(pcm)
 
 # — WebSocket‑Endpoint für TTS Streaming —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        msg = await ws.receive_text()
-        req   = json.loads(msg)
+        data = await ws.receive_text()
+        req  = json.loads(data)
         text  = req.get("text", "")
         voice = req.get("voice", "Jakob")
 
-        input_ids, attention_mask = prepare_inputs(text, voice)
-        past_kvs = None
-        collected = []
-
-        # Token‑für‑Token mit eigener Sampling‑Schleife
-        while True:
-            out = model(
-                input_ids=input_ids if past_kvs is None else None,
-                attention_mask=attention_mask if past_kvs is None else None,
-                past_key_values=past_kvs,
-                use_cache=True,
-            )
-            logits = out.logits[:, -1, :]
-            past_kvs = out.past_key_values
-
-            # Sampling
-            probs = torch.softmax(logits, dim=-1)
-            nxt   = torch.multinomial(probs, num_samples=1).item()
-
-            # EOS → fertig
-            if nxt == EOS_TOKEN:
-                break
-            # RESET → alte Sammlung verwerfen
-            if nxt == RESET_TOKEN:
-                collected = []
-                # und input_ids für nächsten Durchlauf auf None setzen
-                input_ids = None
-                attention_mask = None
-                continue
+        ids, attn_mask = prepare_inputs(text, voice)
+        await generate_and_stream(ws, ids, attn_mask)
 
-            # Audio‑Code abziehen & sammeln
-            collected.append(nxt - AUDIO_OFFSET)
-            # jede 7 Codes → dekodieren & streamen
-            if len(collected) == 7:
-                pcm = decode_block(collected)
-                collected = []
-                await ws.send_bytes(pcm)
-
-            # nur beim allerersten Schritt mit IDs arbeiten
-            input_ids = None
-            attention_mask = None
-
-        # Stream sauber beenden
         await ws.close()
-
     except WebSocketDisconnect:
-        # Client hat Disconnect gemacht → nichts tun
         pass
-
     except Exception as e:
-        # auf Fehler 1011 senden
         print("Error in /ws/tts:", e)
         await ws.close(code=1011)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860)