53 lines
2.5 KiB
Python
53 lines
2.5 KiB
Python
"""Headless check of the server's generation path: TTS -> ASR -> voice chat, with streaming."""
|
|
|
|
import asyncio
|
|
import sys
|
|
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import librosa
|
|
import soundfile as sf
|
|
import torch
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
import server # noqa: E402
|
|
|
|
OUT = Path(__file__).resolve().parent.parent / "smoke_out"
|
|
OUT.mkdir(exist_ok=True)
|
|
|
|
|
|
def run(engine, conv, clips, prefill, mode):
|
|
loop = asyncio.new_event_loop()
|
|
queue = asyncio.Queue()
|
|
stats, cancel = {}, threading.Event()
|
|
emit = server.Emitter(loop, queue, time.perf_counter())
|
|
emit.send = lambda item: None
|
|
_, inputs = engine.build_inputs(conv, clips, prefill)
|
|
text = server.TextStream(engine, emit, cancel, stats)
|
|
audio = server.AudioStream(engine, emit, cancel, stats) if mode != "asr" else None
|
|
t = time.perf_counter()
|
|
with torch.inference_mode():
|
|
engine.model.generate(**inputs, max_new_tokens=512, streamer=text, audio_streamer=audio,
|
|
return_dict_in_generate=True, **({"generation_mode": "text"} if mode == "asr" else {}))
|
|
gen = time.perf_counter() - t
|
|
pcm = torch.cat([torch.from_numpy(c) for c in audio.collected]).numpy() if audio and audio.collected else None
|
|
dur = pcm.size / server.OUT_RATE if pcm is not None else 0
|
|
print(f"[{mode}] prompt={inputs['input_ids'].shape[1]} gen={gen:.2f}s audio={dur:.2f}s stats={stats}\n text={text.text!r}")
|
|
return text.text, pcm
|
|
|
|
|
|
engine = server.Engine()
|
|
print(engine.info())
|
|
_, pcm = run(engine, [{"role": "system", "content": "Perform TTS in English using a female voice."},
|
|
{"role": "user", "content": "What is the tallest mountain in Europe?"}], [], "<|audio_start|>", "tts")
|
|
sf.write(OUT / "tts.wav", pcm, server.OUT_RATE)
|
|
q16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
|
|
run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [q16], None, "asr")
|
|
reply, pcm = run(engine, [{"role": "system", "content": "Respond with interleaved text and audio."},
|
|
{"role": "user", "content": [{"type": "audio"}]}], [q16], "<|mixed_start|>", "chat")
|
|
sf.write(OUT / "chat.wav", pcm, server.OUT_RATE)
|
|
# sanity-check the chat reply's speech by transcribing it
|
|
r16 = librosa.resample(pcm, orig_sr=server.OUT_RATE, target_sr=server.IN_RATE)
|
|
run(engine, [{"role": "system", "content": "Perform ASR."}, {"role": "user", "content": [{"type": "audio"}]}], [r16], None, "asr")
|