Update app.py

app.py

@@ -1,144 +1,138 @@
 import os
 import json
-import asyncio
-
 import torch
+import numpy as np
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from huggingface_hub import login
-from snac import SNAC
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from snac import SNAC
 
-# — ENV & AUTH —
-HF_TOKEN = os.getenv("HF_TOKEN", "")
+# — HF‑Token & Login (wenn gesetzt) —
+HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(HF_TOKEN)
 
-# — DEVICE SETUP —
+# — Device wählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# — FASTAPI INSTANCE —
 app = FastAPI()
 
-# — HEALTHCHECK / ROOT —
 @app.get("/")
 async def read_root():
-    return {"message": "TTS WebSocket up and running!"}
+    return {"message": "Hello, world!"}
 
-# — LOAD MODELS ON STARTUP —
+# — Globale Modelle —
+model = None
+tokenizer = None
+snac_model = None
+
+# — Startup: SNAC & Orpheus laden —
 @app.on_event("startup")
-async def startup_event():
-    global tokenizer, model, snac
-    # 1) SNAC vocoder
-    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
-    # 2) TTS model & tokenizer
-    model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+async def load_models():
+    global model, tokenizer, snac_model
+    # 1) SNAC
+    snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
+    # 2) Orpheus‑TTS
+    REPO = "SebastianBodza/Kartoffel_Orpheus-3B_german_synthetic-v0.1"
+    tokenizer = AutoTokenizer.from_pretrained(REPO)
     model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",
-        torch_dtype=torch.bfloat16 if device == "cuda" else None,
+        REPO,
+        device_map="auto" if device=="cuda" else None,
+        torch_dtype=torch.bfloat16 if device=="cuda" else None,
         low_cpu_mem_usage=True
-    )
-    # make pad == eos
+    ).to(device)
     model.config.pad_token_id = model.config.eos_token_id
 
-# — 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}"
-    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
-
-def decode_seven(tokens: list[int]) -> bytes:
-    """Take exactly 7 audio‑codes, build SNAC input and decode to PCM16 bytes."""
-    b = tokens
-    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 ]
+# — Marker und Offsets aus der Vorlage —
+START_TOKEN  = 128259
+END_TOKENS   = [128009, 128260]
+AUDIO_OFFSET = 128266
+
+def process_single_prompt(prompt: str, voice: str) -> list[int]:
+    # Prompt zusammenbauen
+    if voice and voice != "in_prompt":
+        text = f"{voice}: {prompt}"
+    else:
+        text = prompt
+    # Tokenize + Marker
+    ids = tokenizer(text, return_tensors="pt").input_ids
+    start = torch.tensor([[START_TOKEN]], dtype=torch.int64)
+    end   = torch.tensor([END_TOKENS], dtype=torch.int64)
+    input_ids = torch.cat([start, ids, end], dim=1).to(device)
+    attention_mask = torch.ones_like(input_ids)
+
+    # Generieren
+    gen = model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        max_new_tokens=4000,
+        do_sample=True,
+        temperature=0.6,
+        top_p=0.95,
+        repetition_penalty=1.1,
+        eos_token_id=128258,
+        use_cache=True,
+    )
+
+    # letzten START_TOKEN finden & croppen
+    token_to_find   = 128257
+    token_to_remove = 128258
+    idxs = (gen == token_to_find).nonzero(as_tuple=True)[1]
+    if idxs.numel() > 0:
+        cropped = gen[:, idxs[-1] + 1 :]
+    else:
+        cropped = gen
+
+    # Padding entfernen
+    row = cropped[0][cropped[0] != token_to_remove]
+    # Aus Länge ein Vielfaches von 7 machen
+    new_len = (row.size(0) // 7) * 7
+    trimmed = row[:new_len].tolist()
+    # Offset abziehen
+    return [t - AUDIO_OFFSET for t in trimmed]
+
+def redistribute_codes(code_list: list[int]) -> np.ndarray:
+    # Die 7er‑Blöcke auf 3 Layer verteilen und dekodieren
+    layer1, layer2, layer3 = [], [], []
+    for i in range(len(code_list) // 7):
+        b = code_list[7*i : 7*i+7]
+        layer1.append(b[0])
+        layer2.append(b[1] -   4096)
+        layer3.append(b[2] - 2*4096)
+        layer3.append(b[3] - 3*4096)
+        layer2.append(b[4] - 4*4096)
+        layer3.append(b[5] - 5*4096)
+        layer3.append(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),
+        torch.tensor(layer1, device=device).unsqueeze(0),
+        torch.tensor(layer2, device=device).unsqueeze(0),
+        torch.tensor(layer3, device=device).unsqueeze(0),
     ]
-    audio = snac.decode(codes).squeeze().cpu().numpy()
-    pcm16 = (audio * 32767).astype("int16").tobytes()
-    return pcm16
+    audio = snac_model.decode(codes).squeeze().cpu().numpy()
+    return audio  # float32 @24 kHz
 
-# — WEBSOCKET ENDPOINT —
+# — WebSocket‑Endpoint für TTS —
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
-        # 1) receive JSON request
+        # 1) Text + Voice empfangen
         msg = await ws.receive_text()
         req = json.loads(msg)
         text  = req.get("text", "")
-        voice = req.get("voice", "Jakob")
-
-        # 2) prepare prompt
-        input_ids, attention_mask = prepare_inputs(text, voice)
-        prompt_len = input_ids.size(1)
-
-        # 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,
-                attention_mask=attention_mask if past_kvs is None else None,
-                max_new_tokens=50,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.1,
-                eos_token_id=EOS_TOKEN,
-                pad_token_id=EOS_TOKEN,
-                use_cache=True,
-                return_dict_in_generate=False,
-                return_legacy_cache=True,
-                past_key_values=past_kvs,
-            )
-            # 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 t == RESET_MARKER:
-                    buffer.clear()
-                    continue
-                # 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
-
-        # 7) clean close
+        voice = req.get("voice", "")
+
+        # 2) Prompt → Code‑Liste
+        with torch.no_grad():
+            codes = process_single_prompt(text, voice)
+            audio_np = redistribute_codes(codes)
+
+        # 3) In PCM16 konvertieren und senden
+        pcm16 = (audio_np * 32767).astype(np.int16).tobytes()
+        await ws.send_bytes(pcm16)
+
+        # 4) sauber schließen
         await ws.close()
 
     except WebSocketDisconnect:
@@ -146,3 +140,7 @@ async def tts_ws(ws: WebSocket):
     except Exception as e:
         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)