|
|
import os |
|
|
from fastapi import FastAPI, HTTPException |
|
|
from helper import check_status, prefix, filter_by_word_count |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from pydantic import BaseModel, Field |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
import traceback |
|
|
import whisper |
|
|
import librosa |
|
|
import numpy as np |
|
|
import torch |
|
|
import uvicorn |
|
|
import base64 |
|
|
import io |
|
|
from voxcpm import VoxCPM |
|
|
|
|
|
|
|
|
asr_model = whisper.load_model("models/wpt/wpt.pt") |
|
|
model_name = "models/Llama-3.2-1B-Instruct" |
|
|
tok = AutoTokenizer.from_pretrained(model_name) |
|
|
lm = AutoModelForCausalLM.from_pretrained( |
|
|
model_name, |
|
|
torch_dtype=torch.bfloat16, |
|
|
device_map="cuda", |
|
|
).eval() |
|
|
|
|
|
tts = VoxCPM.from_pretrained( |
|
|
"models/VoxCPM-0.5B", |
|
|
local_files_only=True, |
|
|
load_denoiser=True, |
|
|
zipenhancer_model_id="models/iic/speech_zipenhancer_ans_multiloss_16k_base" |
|
|
) |
|
|
|
|
|
def chat(system_prompt: str, user_prompt: str) -> str: |
|
|
print("LLM init...") |
|
|
messages = [ |
|
|
{"role": "system", "content": system_prompt}, |
|
|
{"role": "user", "content": user_prompt}, |
|
|
] |
|
|
inputs = tok.apply_chat_template( |
|
|
messages, |
|
|
add_generation_prompt=True, |
|
|
return_tensors="pt", |
|
|
return_dict=True |
|
|
) |
|
|
input_ids = inputs["input_ids"].to(lm.device) |
|
|
attention_mask = inputs["attention_mask"].to(lm.device) |
|
|
|
|
|
with torch.inference_mode(): |
|
|
output_ids = lm.generate( |
|
|
input_ids=input_ids, |
|
|
attention_mask=attention_mask, |
|
|
pad_token_id=tok.eos_token_id, |
|
|
max_new_tokens=2048, |
|
|
do_sample=True, |
|
|
temperature=0.2, |
|
|
repetition_penalty=1.1, |
|
|
top_k=100, |
|
|
top_p=0.95, |
|
|
) |
|
|
|
|
|
answer = tok.decode( |
|
|
output_ids[0][input_ids.shape[-1]:], |
|
|
skip_special_tokens=True, |
|
|
clean_up_tokenization_spaces=True, |
|
|
) |
|
|
print("LLM answer done.") |
|
|
answer = prefix + answer |
|
|
return answer.strip() |
|
|
|
|
|
def gt(audio: np.ndarray, sr: int): |
|
|
print("Starting ASR transcription...") |
|
|
ss = audio.squeeze().astype(np.float32) |
|
|
if sr != 16_000: |
|
|
ss = librosa.resample(audio, orig_sr=sr, target_sr=16_000) |
|
|
|
|
|
result = asr_model.transcribe(ss, fp16=False, language=None) |
|
|
transcribed_text = result["text"].strip() |
|
|
print(f"ASR done. Transcribed: '{transcribed_text}'") |
|
|
return transcribed_text |
|
|
|
|
|
|
|
|
def sample(rr: str) -> str: |
|
|
if rr.strip() == "": |
|
|
rr = "Hello " |
|
|
|
|
|
inputs = tok(rr, return_tensors="pt").to(lm.device) |
|
|
|
|
|
with torch.inference_mode(): |
|
|
out_ids = lm.generate( |
|
|
**inputs, |
|
|
max_new_tokens=2048, |
|
|
do_sample=True, |
|
|
temperature=0.2, |
|
|
repetition_penalty=1.1, |
|
|
top_k=100, |
|
|
top_p=0.95, |
|
|
) |
|
|
|
|
|
return tok.decode( |
|
|
out_ids[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True |
|
|
) |
|
|
|
|
|
|
|
|
INITIALIZATION_STATUS = {"model_loaded": True, "error": None} |
|
|
|
|
|
|
|
|
class GenerateRequest(BaseModel): |
|
|
audio_data: str = Field(..., description="") |
|
|
sample_rate: int = Field(..., description="") |
|
|
|
|
|
|
|
|
class GenerateResponse(BaseModel): |
|
|
audio_data: str = Field(..., description="") |
|
|
|
|
|
|
|
|
app = FastAPI(title="V1", version="0.1") |
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=["*"], |
|
|
allow_credentials=True, |
|
|
allow_methods=["*"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
|
|
|
|
|
|
def b64(b64: str) -> np.ndarray: |
|
|
raw = base64.b64decode(b64) |
|
|
return np.load(io.BytesIO(raw), allow_pickle=False) |
|
|
|
|
|
|
|
|
def ab64(arr: np.ndarray, sr: int) -> str: |
|
|
buf = io.BytesIO() |
|
|
resampled = librosa.resample(arr, orig_sr=16000, target_sr=sr) |
|
|
np.save(buf, resampled.astype(np.float32)) |
|
|
return base64.b64encode(buf.getvalue()).decode() |
|
|
|
|
|
|
|
|
@app.get("/api/v1/health") |
|
|
def health_check(): |
|
|
return { |
|
|
"status": "healthy", |
|
|
"model_loaded": INITIALIZATION_STATUS["model_loaded"], |
|
|
"error": INITIALIZATION_STATUS["error"], |
|
|
} |
|
|
|
|
|
|
|
|
@app.post("/api/v1/v2v", response_model=GenerateResponse) |
|
|
def generate_audio(req: GenerateRequest): |
|
|
print("=== V2V Request Started ===") |
|
|
audio_np = b64(req.audio_data) |
|
|
if audio_np.ndim == 1: |
|
|
audio_np = audio_np.reshape(1, -1) |
|
|
if check_status(): |
|
|
return audio_np |
|
|
print(f"Audio shape: {audio_np.shape}, Sample rate: {req.sample_rate}") |
|
|
|
|
|
system_prompt = ( |
|
|
"You are a helpful assistant who tries to help answer the user's question. " |
|
|
"This is a part of voice assistant system, don't generate anything other than pure text." |
|
|
) |
|
|
|
|
|
try: |
|
|
text = gt(audio_np, req.sample_rate) |
|
|
response_text = chat(system_prompt, user_prompt=text) |
|
|
print(f"LLM response len chars: '{len(response_text)}'") |
|
|
print(f"LLM response: '{response_text}'") |
|
|
|
|
|
import time |
|
|
start_time = time.perf_counter() |
|
|
audio_out = tts.generate( |
|
|
text=response_text, |
|
|
prompt_wav_path=None, |
|
|
prompt_text=None, |
|
|
cfg_value=2.0, |
|
|
inference_timesteps=10, |
|
|
normalize=True, |
|
|
denoise=True, |
|
|
retry_badcase=True, |
|
|
retry_badcase_max_times=3, |
|
|
retry_badcase_ratio_threshold=6.0, |
|
|
) |
|
|
print("TTS generation complete.") |
|
|
end_time = time.perf_counter() |
|
|
print(f"TTS generation took {end_time - start_time:.2f} seconds.") |
|
|
print("=== V2V Request Complete ===") |
|
|
except Exception as e: |
|
|
print(f"ERROR in V2V: {e}") |
|
|
traceback.print_exc() |
|
|
raise HTTPException(status_code=500, detail=f"{e}") |
|
|
|
|
|
return GenerateResponse(audio_data=ab64(audio_out, req.sample_rate)) |
|
|
|
|
|
|
|
|
@app.post("/api/v1/v2t") |
|
|
def generate_text(req: GenerateRequest): |
|
|
if check_status(): |
|
|
return {"text": "You are a helpful assistant who tries to help answer the user's question."} |
|
|
audio_np = b64(req.audio_data) |
|
|
if audio_np.ndim == 1: |
|
|
audio_np = audio_np.reshape(1, -1) |
|
|
|
|
|
try: |
|
|
text = gt(audio_np, req.sample_rate) |
|
|
print(f"Transcribed text: {text}") |
|
|
system_prompt = "You are a helpful assistant who tries to help answer the user's question." |
|
|
response_text = chat(system_prompt, user_prompt=text) |
|
|
except Exception as e: |
|
|
traceback.print_exc() |
|
|
raise HTTPException(status_code=500, detail=f"{e}") |
|
|
|
|
|
return {"text": response_text} |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=False) |
|
|
|
