app.py

import os
import json
import asyncio
import torch
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, snapshot_download

# ─── ENV & HF TOKEN ────────────────────────────────────────────────────────────
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    login(token=HF_TOKEN)

# ─── DEVICE ─────────────────────────────────────────────────────────────────────
device = "cuda" if torch.cuda.is_available() else "cpu"

# ─── SNAC ───────────────────────────────────────────────────────────────────────
print("Loading SNAC model…")
snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)

# ─── ORPHEUS ────────────────────────────────────────────────────────────────────
model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1"

# pre‑download only the config + safetensors, damit das Image schlank bleibt
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.*", "vocab.json", "merges.txt"
    ]
)

print("Loading Orpheus model…")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16  # optional: beschleunigt das FP16‑ähnliche Rechnen
)
model = model.to(device)
model.config.pad_token_id = model.config.eos_token_id

tokenizer = AutoTokenizer.from_pretrained(model_name)

# ─── HILFSFUNKTIONEN ──────────────────────────────────────────────────────────

def process_prompt(text: str, voice: str):
    """
    Baut aus Text+Voice ein batch‑Tensor input_ids für `model.generate`.
    """
    prompt = f"{voice}: {text}"
    tok = tokenizer(prompt, return_tensors="pt").to(device)
    start = torch.tensor([[128259]], device=device)
    end   = torch.tensor([[128009, 128260]], device=device)
    return torch.cat([start, tok.input_ids, end], dim=1)

def parse_output(generated_ids: torch.LongTensor):
    """
    Schneidet bis zum letzten 128257 und entfernt 128258, gibt reine Token‑Liste zurück.
    """
    START, PAD = 128257, 128258
    idxs = (generated_ids == START).nonzero(as_tuple=True)[1]
    if idxs.numel() > 0:
        cropped = generated_ids[:, idxs[-1].item()+1:]
    else:
        cropped = generated_ids
    row = cropped[0][cropped[0] != PAD]
    return row.tolist()

def redistribute_codes(code_list: list[int], snac_model: SNAC):
    """
    Verteilt 7er‑Blöcke auf die drei SNAC‑Layer und dekodiert zu Audio (numpy float32).
    """
    layer1, layer2, layer3 = [], [], []
    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.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=dev).unsqueeze(0),
        torch.tensor(layer2, device=dev).unsqueeze(0),
        torch.tensor(layer3, device=dev).unsqueeze(0),
    ]
    audio = snac_model.decode(codes)
    return audio.detach().squeeze().cpu().numpy()

# ─── FASTAPI ────────────────────────────────────────────────────────────────────

app = FastAPI()

@app.get("/")
async def healthcheck():
    return {"status": "ok", "msg": "Hello, Orpheus TTS up!"}

@app.websocket("/ws/tts")
async def tts_ws(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            # 1) Eintreffende JSON‐Nachricht parsen
            data = json.loads(await ws.receive_text())
            text  = data.get("text", "")
            voice = data.get("voice", "Jakob")

            # 2) Prompt → input_ids
            ids = process_prompt(text, voice)

            # 3) Token‐Erzeugung
            gen_ids = model.generate(
                input_ids=ids,
                max_new_tokens=2000,      # hier z.B. 20k geht auch, wird aber speicherintensiv
                do_sample=True,
                temperature=0.7,
                top_p=0.95,
                repetition_penalty=1.1,
                eos_token_id=model.config.eos_token_id,
            )

            # 4) Tokens → Code‐Liste → Audio
            codes    = parse_output(gen_ids)
            audio_np = redistribute_codes(codes, snac)

            # 5) PCM16‐Stream in 0.1 s‐Blöcken
            pcm16 = (audio_np * 32767).astype("int16").tobytes()
            chunk = 2400 * 2
            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:
        print("Error in /ws/tts:", e)
        await ws.close(code=1011)

# ─── START ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")