Spaces:
Paused
Paused
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‑AUTH — | |
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" | |
# — Modelle laden — | |
print("Loading SNAC model…") | |
snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device) | |
model_name = "SebastianBodza/Kartoffel_Orpheus-3B_german_natural-v0.1" | |
# Nur die Konfig + Safetensors, alles andere wird ignoriert | |
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 | |
).to(device) | |
model.config.pad_token_id = model.config.eos_token_id | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
# — Konstanten für Audio‑Token → | |
# (muss übereinstimmen mit Deinem Training; hier 128266) | |
AUDIO_TOKEN_OFFSET = 128266 | |
# — Hilfsfunktionen — | |
def process_prompt(text: str, voice: str): | |
prompt = f"{voice}: {text}" | |
input_ids = tokenizer(prompt, return_tensors="pt").input_ids | |
# Laut Spezifikation: | |
# start_token=128259, end_tokens=(128009,128260) | |
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) | |
mask = torch.ones_like(ids).to(device) | |
return ids, mask | |
def parse_output(generated_ids: torch.LongTensor): | |
""" | |
Croppt nach dem letzten 128257-Start-Token, entfernt Padding (128258) | |
und zieht dann den Audio‑Offset ab, um echte Code‑IDs zu bekommen. | |
""" | |
# finde letztes Audio‑Start‑Token | |
token_to_start = 128257 | |
token_to_remove = model.config.eos_token_id # 128258 | |
idxs = (generated_ids == token_to_start).nonzero(as_tuple=True)[1] | |
if idxs.numel() > 0: | |
cut = idxs[-1].item() + 1 | |
cropped = generated_ids[:, cut:] | |
else: | |
cropped = generated_ids | |
# flatten & remove PAD, dann Offset abziehen | |
flat = cropped[0][cropped[0] != token_to_remove] | |
codes = [(int(t) - AUDIO_TOKEN_OFFSET) for t in flat] | |
return codes | |
def redistribute_codes(code_list: list[int], snac_model: SNAC): | |
""" | |
Verteilt die flache Code‑Liste in 3 Layers und dekodiert mit SNAC. | |
""" | |
layer1, layer2, layer3 = [], [], [] | |
for i in range(len(code_list) // 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 | |
c1 = torch.tensor(layer1, device=dev).unsqueeze(0) | |
c2 = torch.tensor(layer2, device=dev).unsqueeze(0) | |
c3 = torch.tensor(layer3, device=dev).unsqueeze(0) | |
audio = snac_model.decode([c1, c2, c3]) | |
return audio.detach().squeeze().cpu().numpy() | |
# — FastAPI + WebSocket-Endpoint — | |
app = FastAPI() | |
async def tts_ws(ws: WebSocket): | |
await ws.accept() | |
try: | |
while True: | |
msg = await ws.receive_text() | |
data = json.loads(msg) | |
text = data.get("text", "") | |
voice = data.get("voice", "Jakob") | |
# 1) Prompt → Token‑Tensoren | |
ids, mask = process_prompt(text, voice) | |
# 2) Generation | |
gen_ids = model.generate( | |
input_ids=ids, | |
attention_mask=mask, | |
max_new_tokens=200, # zum Debug | |
do_sample=True, | |
temperature=0.7, | |
top_p=0.95, | |
repetition_penalty=1.1, | |
eos_token_id=model.config.eos_token_id, | |
) | |
# 3) Token → Code‑Liste → Audio (Float32 @24 kHz) | |
code_list = parse_output(gen_ids) | |
audio_np = redistribute_codes(code_list, snac) | |
# 4) In 0.1 s‑Chunks (2400 Samples) als PCM16 streamen | |
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) | |
if __name__ == "__main__": | |
import uvicorn | |
uvicorn.run("app:app", host="0.0.0.0", port=7860) | |