up1

app.py

@@ -14,15 +14,14 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-# — Device —
+# — Gerät wählen —
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # — Modelle laden —
-print("Loading SNAC model…")
+print("Loading SNAC model...")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
 model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"
-print("Downloading Orpheus weights (konfig + safetensors)…")
 snapshot_download(
     repo_id=model_name,
     allow_patterns=["config.json", "*.safetensors", "model.safetensors.index.json"],
@@ -33,112 +32,120 @@ snapshot_download(
     ]
 )
 
-print("Loading Orpheus model…")
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
+print("Loading Orpheus model...")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, torch_dtype=torch.bfloat16
+).to(device)
 model.config.pad_token_id = model.config.eos_token_id
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# — Hilfsfunktionen — 
+# — Konstanten für Token‑Mapping —
+AUDIO_TOKEN_OFFSET = 128266
+START_TOKEN  = 128259
+SOS_TOKEN    = 128257
+EOS_TOKEN    = 128258
+
+# — Hilfsfunktionen —
 def process_prompt(text: str, voice: str):
-    """Erzeuge input_ids und attention_mask für einen Prompt."""
     prompt = f"{voice}: {text}"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
-    start = torch.tensor([[128259]], dtype=torch.int64, device=device)
+    start = torch.tensor([[START_TOKEN]], 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)
+    ids  = torch.cat([start, input_ids, end], dim=1)
+    mask = torch.ones_like(ids, dtype=torch.int64, device=device)
     return ids, mask
 
-def parse_output(generated_ids: torch.LongTensor) -> list[int]:
-    """Extrahiere rohe Tokenliste nach dem letzten 128257-Start-Token."""
-    token_to_find   = 128257
-    token_to_remove = 128258
-
-    idxs = (generated_ids == token_to_find).nonzero(as_tuple=True)[1]
-    if idxs.numel() > 0:
-        cropped = generated_ids[:, idxs[-1].item() + 1 :]
-    else:
-        cropped = generated_ids
-
-    row = cropped[0]
-    row = row[row != token_to_remove]
-    return row.tolist()
-
-def redistribute_codes(code_list: list[int]) -> bytes:
-    """Verteile die Codes auf die drei SNAC-Layer und dekodiere zu PCM16-Bytes."""
+def redistribute_codes(block: list[int], snac_model: SNAC):
+    # exakt wie vorher: 7 Codes → 3 Layer → SNAC.decode → NumPy float32 @24 kHz
     l1, l2, l3 = [], [], []
-    for i in range((len(code_list) + 1) // 7):
-        base = code_list[7*i : 7*i+7]
-        l1.append(base[0])
-        l2.append(base[1] -   4096)
-        l3.append(base[2] - 2*4096)
-        l3.append(base[3] - 3*4096)
-        l2.append(base[4] - 4*4096)
-        l3.append(base[5] - 5*4096)
-        l3.append(base[6] - 6*4096)
-
-    dev = next(snac.parameters()).device
+    for i in range(len(block)//7):
+        b = block[7*i:7*i+7]
+        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)
+    dev = next(snac_model.parameters()).device
     codes = [
         torch.tensor(l1, device=dev).unsqueeze(0),
         torch.tensor(l2, device=dev).unsqueeze(0),
         torch.tensor(l3, device=dev).unsqueeze(0),
     ]
-    audio = snac.decode(codes).squeeze().cpu().numpy()  # float32 @24 kHz
-    pcm16 = (audio * 32767).astype("int16").tobytes()
-    return pcm16
+    audio = snac_model.decode(codes)  # → Tensor[1, T]
+    return audio.squeeze().cpu().numpy()
 
-# — FastAPI + WebSocket-Endpoint —
+# — FastAPI Setup —
 app = FastAPI()
 
+# 1) Hello‑World Endpoint
+@app.get("/")
+async def root():
+    return {"message": "Hallo Welt"}
+
+# 2) WebSocket Token‑für‑Token TTS
 @app.websocket("/ws/tts")
 async def tts_ws(ws: WebSocket):
     await ws.accept()
     try:
         while True:
-            # 1) Nachricht empfangen
-            msg = await ws.receive_text()
-            data  = json.loads(msg)
-            text  = data.get("text", "")
-            voice = data.get("voice", "Jakob")
-
-            # 2) Prompt → IDs/Mask
+            # JSON mit Text & Voice empfangen
+            raw = await ws.receive_text()
+            req = json.loads(raw)
+            text, voice = req.get("text", ""), req.get("voice", "Jakob")
             ids, mask = process_prompt(text, voice)
 
-            # 3) Token-Generation
-            gen_ids = model.generate(
-                input_ids=ids,
-                attention_mask=mask,
-                max_new_tokens=2000,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.1,
-                eos_token_id=128258,
-            )
-
-            # 4) Parse + SNAC → PCM16‑Bytes
-            codes   = parse_output(gen_ids)
-            pcm16   = redistribute_codes(codes)
-            chunk_sz = 2400 * 2  # 0.1 s @24 kHz
-
-            # 5) Stream audio‑Chunks
-            for i in range(0, len(pcm16), chunk_sz):
-                await ws.send_bytes(pcm16[i : i + chunk_sz])
-                await asyncio.sleep(0.1)
-
-            # 6) Ende‑Signal
-            await ws.send_json({"event": "eos"})
-
-            # (Verbindung bleibt offen für nächste Anfrage)
+            past_kv = None
+            collected = []
+
+            # im Sampling‑Loop Token für Token generieren
+            with torch.no_grad():
+                for _ in range(2000):  # max 200 Tokens
+                    out = model(
+                        input_ids=ids if past_kv is None else None,
+                        attention_mask=mask if past_kv is None else None,
+                        past_key_values=past_kv,
+                        use_cache=True,
+                    )
+                    logits = out.logits[:, -1, :]
+                    next_id = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
+                    past_kv = out.past_key_values
+
+                    token = next_id.item()
+                    # Ende
+                    if token == EOS_TOKEN:
+                        break
+                    # Reset bei SOS
+                    if token == SOS_TOKEN:
+                        collected = []
+                        continue
+
+                    # in Audio‑Code konvertieren
+                    collected.append(token - AUDIO_TOKEN_OFFSET)
+
+                    # sobald 7 Codes → direkt dekodieren & streamen
+                    if len(collected) >= 7:
+                        block = collected[:7]
+                        collected = collected[7:]
+                        audio_np = redistribute_codes(block, snac)
+                        pcm16 = (audio_np * 32767).astype("int16").tobytes()
+                        await ws.send_bytes(pcm16)
+
+                    # ab jetzt nur noch past_kv verwenden
+                    ids  = None
+                    mask = None
+
+            # zum Schluss End‑Of‑Stream signalisieren
+            await ws.send_text(json.dumps({"event": "eos"}))
 
     except WebSocketDisconnect:
         print("Client disconnected")
     except Exception as e:
         print("Error in /ws/tts:", e)
-        # Schließe erst, nachdem Fehler gemeldet
         await ws.close(code=1011)
 
+# zum lokalen Test
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("app:app", host="0.0.0.0", port=7860)