Update app.py

app.py

@@ -6,21 +6,34 @@ from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 from dotenv import load_dotenv
 from snac import SNAC
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import login
+from huggingface_hub import login, snapshot_download
 
-# — Environment & HF‑Auth —
+# — ENV & HF‑AUTH —
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-# — Device & Modelle laden —
+# — Debug: CPU‑Modus zum Entwickeln, später wieder "cuda" —
 device = "cuda" if torch.cuda.is_available() else "cpu"
+# device = "cpu"
 
+# — Modelle laden —
 print("Loading SNAC model...")
 snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
 
 model_name = "canopylabs/3b-de-ft-research_release"
+# optional: explizites snapshot_download (entfernt große Dateien)
+snapshot_download(
+    repo_id=model_name,
+    allow_patterns=["config.json", "*.safetensors", "model.safetensors.index.json"],
+    ignore_patterns=[
+        "optimizer.pt", "pytorch_model.bin", "training_args.bin",
+        "scheduler.pt", "tokenizer.json", "tokenizer_config.json",
+        "special_tokens_map.json", "vocab.json", "merges.txt", "tokenizer.*"
+    ]
+)
+
 print("Loading Orpheus model...")
 model = AutoModelForCausalLM.from_pretrained(
     model_name, torch_dtype=torch.bfloat16
@@ -29,49 +42,61 @@ model.config.pad_token_id = model.config.eos_token_id
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# — Hilfsfunktionen — 
+# — Hilfsfunktionen —
+
 def process_prompt(text: str, voice: str):
     prompt = f"{voice}: {text}"
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids
     start = torch.tensor([[128259]], dtype=torch.int64)
-    end = torch.tensor([[128009, 128260]], dtype=torch.int64)
-    ids = torch.cat([start, input_ids, end], dim=1).to(device)
+    end   = torch.tensor([[128009, 128260]], dtype=torch.int64)
+    ids  = torch.cat([start, input_ids, end], dim=1).to(device)
     mask = torch.ones_like(ids).to(device)
     return ids, mask
 
 def parse_output(generated_ids: torch.LongTensor):
-    token_to_find = 128257
+    """Extrahiere rohe Tokenliste nach dem letzten 128257-Start-Token."""
+    token_to_find   = 128257
     token_to_remove = 128258
+
+    # 1) Finde letztes Start-Token, croppe
     idxs = (generated_ids == token_to_find).nonzero(as_tuple=True)[1]
     if idxs.numel() > 0:
-        last = idxs[-1].item()
-        cropped = generated_ids[:, last+1:]
+        cut = idxs[-1].item() + 1
+        cropped = generated_ids[:, cut:]
     else:
         cropped = generated_ids
-    # remove padding token markers
+
+    # 2) Entferne Padding-Markierungen
     rows = []
     for row in cropped:
-        row = row[row != token_to_remove]
-        rows.append(row)
-    flat = rows[0].tolist()
-    # adjust and regroup
+        rows.append(row[row != token_to_remove])
+
+    # 3) Flache Liste zurückgeben
+    return rows[0].tolist()
+
+def redistribute_codes(code_list: list[int], snac_model: SNAC):
+    """Verteile die Codes auf drei Layer, dekodiere in Audio."""
     layer1, layer2, layer3 = [], [], []
-    for i in range(len(flat)//7):
-        base = flat[7*i:7*i+7]
+    for i in range((len(code_list) + 1) // 7):
+        base = code_list[7*i : 7*i+7]
         layer1.append(base[0])
-        layer2.append(base[1]-4096)
-        layer3.extend([base[2]-(2*4096), base[3]-(3*4096)])
-        layer2.append(base[4]-4*4096)
-        layer3.extend([base[5]-(5*4096), base[6]-(6*4096)])
+        layer2.append(base[1] -   4096)
+        layer3.append(base[2] - 2*4096)
+        layer3.append(base[3] - 3*4096)
+        layer2.append(base[4] - 4*4096)
+        layer3.append(base[5] - 5*4096)
+        layer3.append(base[6] - 6*4096)
+
+    dev = next(snac_model.parameters()).device
     codes = [
-        torch.tensor(layer1, device=device).unsqueeze(0),
-        torch.tensor(layer2, device=device).unsqueeze(0),
-        torch.tensor(layer3, device=device).unsqueeze(0),
+        torch.tensor(layer1, device=dev).unsqueeze(0),
+        torch.tensor(layer2, device=dev).unsqueeze(0),
+        torch.tensor(layer3, device=dev).unsqueeze(0),
     ]
-    audio = snac.decode(codes).detach().squeeze().cpu().numpy()
-    return audio  # float32 numpy at 24000 Hz
+    audio = snac_model.decode(codes)
+    return audio.detach().squeeze().cpu().numpy()  # float32 @24 kHz
 
-# — FastAPI + WebSocket-Endpoint — 
+# — FastAPI + WebSocket-Endpoint —
 app = FastAPI()
 
 @app.websocket("/ws/tts")
@@ -80,31 +105,36 @@ async def tts_ws(ws: WebSocket):
     try:
         while True:
             msg = await ws.receive_text()
-            data = json.loads(msg)
-            text = data.get("text", "")
+            data  = json.loads(msg)
+            text  = data.get("text", "")
             voice = data.get("voice", "jana")
-            # Generate tokens
+
+            # 1) Prompt → Tokens
             ids, mask = process_prompt(text, voice)
-            with torch.no_grad():
-                gen_ids = model.generate(
-                    input_ids=ids,
-                    attention_mask=mask,
-                    max_new_tokens=1200,
-                    do_sample=True,
-                    temperature=0.7,
-                    top_p=0.95,
-                    repetition_penalty=1.1,
-                    eos_token_id=128258,
-                )
-            # Convert to waveform
-            audio = parse_output(gen_ids)
-            # PCM16 conversion & chunking
-            pcm16 = (audio * 32767).astype('int16').tobytes()
-            # 0.1 s @24 kHz = 2400 samples = 4800 bytes
-            chunk_size = 2400 * 2
-            for i in range(0, len(pcm16), chunk_size):
-                await ws.send_bytes(pcm16[i:i+chunk_size])
-                await asyncio.sleep(0.1)  # pacing
+
+            # 2) Token-Generation (erst klein testen!)
+            gen_ids = model.generate(
+                input_ids=ids,
+                attention_mask=mask,
+                max_new_tokens=200,      # zum Debuggen klein halten
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.1,
+                eos_token_id=128258,
+            )
+
+            # 3) Tokens → Code-Liste → Audio
+            code_list = parse_output(gen_ids)
+            audio_np  = redistribute_codes(code_list, snac)
+
+            # 4) PCM16-Bytes und Stream in 0.1s-Chunks
+            pcm16 = (audio_np * 32767).astype("int16").tobytes()
+            chunk = 2400 * 2  # 2400 samples @24kHz → 0.1s * 2 bytes
+            for i in range(0, len(pcm16), chunk):
+                await ws.send_bytes(pcm16[i : i+chunk])
+                await asyncio.sleep(0.1)
+
     except WebSocketDisconnect:
         print("Client disconnected")
     except Exception as e: