Update app.py

app.py

@@ -1,155 +1,148 @@
 import os
 import json
 import asyncio
+
 import torch
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
 from snac import SNAC
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# — HF‑Token & Login —
-HF_TOKEN = os.getenv("HF_TOKEN")
+# — ENV & AUTH —
+HF_TOKEN = os.getenv("HF_TOKEN", "")
 if HF_TOKEN:
     login(HF_TOKEN)
 
-# — Device auswählen —
+# — DEVICE SETUP —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# — FastAPI instanziieren —
+# — FASTAPI INSTANCE —
 app = FastAPI()
 
-# — Hello‑Route, damit GET / nicht 404 gibt —
+# — HEALTHCHECK / ROOT —
 @app.get("/")
 async def read_root():
-    return {"message": "Hello, world!"}
+    return {"message": "TTS WebSocket up and running!"}
 
-# — Modelle beim Startup laden —
+# — LOAD MODELS ON STARTUP —
 @app.on_event("startup")
-async def load_models():
+async def startup_event():
     global tokenizer, model, snac
-    # SNAC laden
+    # 1) SNAC vocoder
     snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-    # TTS‑Modell laden
+    # 2) TTS model & tokenizer
     model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        device_map="auto" if device=="cuda" else None,
-        torch_dtype=torch.bfloat16 if device=="cuda" else None,
+        device_map="auto",
+        torch_dtype=torch.bfloat16 if device == "cuda" else None,
         low_cpu_mem_usage=True
-    ).to(device)
+    )
+    # make pad == eos
     model.config.pad_token_id = model.config.eos_token_id
 
-# — Input‑Vorbereitung —
+# — HELPERS —
+START_TOKEN      = 128259
+END_TOKENS       = [128009, 128260]
+RESET_MARKER     = 128257
+EOS_TOKEN        = 128258
+AUDIO_TOKEN_OFFSET = 128266  # to subtract from token→audio code
+
 def prepare_inputs(text: str, voice: str):
     prompt = f"{voice}: {text}"
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    start = torch.tensor([[128259]], dtype=torch.int64, device=device)
-    end   = torch.tensor([[128009, 128260]], dtype=torch.int64, device=device)
-    ids   = torch.cat([start, input_ids, end], dim=1)
-    mask  = torch.ones_like(ids, device=device)
+    in_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    start = torch.tensor([[START_TOKEN]], dtype=torch.int64, device=device)
+    end   = torch.tensor([END_TOKENS], dtype=torch.int64, device=device)
+    ids   = torch.cat([start, in_ids, end], dim=1)
+    mask  = torch.ones_like(ids)
     return ids, mask
 
-# — SNAC‑Dekodierung eines 7‑Token‑Blocks →
-def decode_block(tokens: list[int]) -> bytes:
-    l1, l2, l3 = [], [], []
+def decode_seven(tokens: list[int]) -> bytes:
+    """Take exactly 7 audio‑codes, build SNAC input and decode to PCM16 bytes."""
     b = tokens
-    l1.append(b[0])
-    l2.append(b[1]-4096)
-    l3.append(b[2]-2*4096)
-    l3.append(b[3]-3*4096)
-    l2.append(b[4]-4*4096)
-    l3.append(b[5]-5*4096)
-    l3.append(b[6]-6*4096)
+    l1 = [ b[0] ]
+    l2 = [ b[1] - 1*4096, b[4] - 4*4096 ]
+    l3 = [ b[2] - 2*4096, b[3] - 3*4096, b[5] - 5*4096, b[6] - 6*4096 ]
     codes = [
         torch.tensor(l1, device=device).unsqueeze(0),
         torch.tensor(l2, device=device).unsqueeze(0),
         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
 
-# — WebSocket‑Endpoint mit Chunked‑Generate (max_new_tokens=50) —
+# — WEBSOCKET ENDPOINT —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        # 1) Anfrage einlesen
+        # 1) receive JSON request
         msg = await ws.receive_text()
         req = json.loads(msg)
         text  = req.get("text", "")
         voice = req.get("voice", "Jakob")
 
-        # 2) Inputs bauen
+        # 2) prepare prompt
         input_ids, attention_mask = prepare_inputs(text, voice)
-        past_kvs = None
-        buffer_codes: list[int] = []
-
-        # 3) Chunk‑Generate‑Loop
-        chunk_size = 50
-        eos_id     = model.config.eos_token_id
+        prompt_len = input_ids.size(1)
 
-        # Wir tracken bisher erzeugte Länge, um abzugrenzen, was neu ist
-        prev_len = 0
+        # 3) chunked generation setup
+        past_kvs = None
+        buffer: list[int] = []
+        generated_offset = 0
 
         while True:
+            # 4) generate up to 50 new tokens at a time
             out = model.generate(
-                input_ids = input_ids     if past_kvs is None else None,
+                input_ids= input_ids if past_kvs is None else None,
                 attention_mask=attention_mask if past_kvs is None else None,
-                max_new_tokens=chunk_size,
+                max_new_tokens=50,
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.95,
                 repetition_penalty=1.1,
-                eos_token_id=eos_id,
+                eos_token_id=EOS_TOKEN,
+                pad_token_id=EOS_TOKEN,
                 use_cache=True,
-                return_dict_in_generate=True,
-                output_scores=False,
-                past_key_values=past_kvs
+                return_dict_in_generate=False,
+                return_legacy_cache=True,
+                past_key_values=past_kvs,
             )
-            # Update past_kvs und sequences
-            past_kvs = out.past_key_values
-            seqs     = out.sequences  # (1, total_length)
-            total_len = seqs.shape[1]
-
-            # 4) Neue Tokens extrahieren
-            new_tokens = seqs[0, prev_len:total_len].tolist()
-            prev_len = total_len
-
-            # 5) Jeden neuen Token aufbereiten
-            for tok in new_tokens:
-                if tok == eos_id:
-                    # Ende
-                    new_tokens = []  # clean up
+            # out is a tuple: (generated_ids, new_past_kvs)
+            gen_ids, past_kvs = out
+
+            # 5) extract only newly generated tokens
+            seq = gen_ids[0]
+            new_seq = seq[prompt_len + generated_offset :]
+            generated_offset += new_seq.size(0)
+
+            # 6) process each new token
+            stop = False
+            for t in new_seq.tolist():
+                if t == EOS_TOKEN:
+                    stop = True
                     break
-                if tok == 128257:
-                    buffer_codes.clear()
+                if t == RESET_MARKER:
+                    buffer.clear()
                     continue
-                # offset und puffern
-                buffer_codes.append(tok - 128266)
-                # sobald 7 Codes gesammelt, dekodieren & senden
-                if len(buffer_codes) >= 7:
-                    block = buffer_codes[:7]
-                    buffer_codes = buffer_codes[7:]
-                    pcm = decode_block(block)
-                    await ws.send_bytes(pcm)
-
-            # 6) Abbruch, wenn EOS im Chunk war
-            if eos_id in new_tokens:
+                # convert to audio-code
+                buffer.append(t - AUDIO_TOKEN_OFFSET)
+                # once we have 7 codes, decode & stream
+                if len(buffer) >= 7:
+                    block = buffer[:7]
+                    buffer = buffer[7:]
+                    pcm_bytes = decode_seven(block)
+                    await ws.send_bytes(pcm_bytes)
+            if stop:
                 break
 
-            # Inputs für nächsten Durchgang nur beim ersten Mal
-            input_ids = attention_mask = None
-
-        # 7) Zum Schluss sauber schließen
+        # 7) clean close
         await ws.close()
+
     except WebSocketDisconnect:
-        return
+        pass
     except Exception as e:
         print("Error in /ws/tts:", e)
         await ws.close(code=1011)
-
-# — Main für lokalen Test —
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860)