app.py

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")
if HF_TOKEN:
    login(HF_TOKEN)

# — Device auswählen —
device = "cuda" if torch.cuda.is_available() else "cpu"

# — FastAPI instanziieren —
app = FastAPI()

# — Hello‑Route, damit GET / nicht 404 gibt —
@app.get("/")
async def read_root():
    return {"message": "Hello, world!"}

# — Modelle beim Startup laden —
@app.on_event("startup")
async def load_models():
    global tokenizer, model, snac
    # SNAC laden
    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
    # TTS‑Modell laden
    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,
        low_cpu_mem_usage=True
    ).to(device)
    model.config.pad_token_id = model.config.eos_token_id

# — Input‑Vorbereitung —
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)
    return ids, mask

# — SNAC‑Dekodierung eines 7‑Token‑Blocks →
def decode_block(tokens: list[int]) -> bytes:
    l1, l2, l3 = [], [], []
    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)
    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()

# — WebSocket‑Endpoint mit Chunked‑Generate (max_new_tokens=50) —
@app.websocket("/ws/tts")
async def tts_ws(ws: WebSocket):
    await ws.accept()
    try:
        # 1) Anfrage einlesen
        msg = await ws.receive_text()
        req = json.loads(msg)
        text  = req.get("text", "")
        voice = req.get("voice", "Jakob")

        # 2) Inputs bauen
        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

        # Wir tracken bisher erzeugte Länge, um abzugrenzen, was neu ist
        prev_len = 0

        while True:
            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=chunk_size,
                do_sample=True,
                temperature=0.7,
                top_p=0.95,
                repetition_penalty=1.1,
                eos_token_id=eos_id,
                use_cache=True,
                return_dict_in_generate=True,
                output_scores=False,
                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
                    break
                if tok == 128257:
                    buffer_codes.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:
                break

            # Inputs für nächsten Durchgang nur beim ersten Mal
            input_ids = attention_mask = None

        # 7) Zum Schluss sauber schließen
        await ws.close()
    except WebSocketDisconnect:
        return
    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)