File size: 4,685 Bytes
bf6d986 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | import json
import os.path
import tempfile
import sys
import re
import uuid
import requests
import librosa
import numpy as np
import torch
import uvicorn
import torchaudio
import base64
import io
from argparse import ArgumentParser
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import WhisperFeatureExtractor, AutoTokenizer
from speech_tokenizer.modeling_whisper import WhisperVQEncoder
sys.path.insert(0, "./cosyvoice")
sys.path.insert(0, "./third_party/Matcha-TTS")
from speech_tokenizer.utils import extract_speech_token
from flow_inference import AudioDecoder
# Initialize FastAPI
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
class AudioRequest(BaseModel):
audio_data: str # Base64 encoded audio
sample_rate: int
class AudioResponse(BaseModel):
audio_data: str # Base64 encoded response audio
text_transcript: str
# Global initialization
DEVICE = "cuda"
audio_decoder = None
whisper_model = None
feature_extractor = None
glm_tokenizer = None
def initialize_models():
global audio_decoder, feature_extractor, whisper_model, glm_tokenizer
# Initialize Whisper and GLM components
glm_tokenizer = AutoTokenizer.from_pretrained("THUDM/glm-4-voice-9b", trust_remote_code=True)
# Initialize Whisper model
whisper_model = WhisperVQEncoder.from_pretrained("THUDM/glm-4-voice-tokenizer").eval().to(DEVICE)
feature_extractor = WhisperFeatureExtractor.from_pretrained("THUDM/glm-4-voice-tokenizer")
# Initialize AudioDecoder
audio_decoder = AudioDecoder(
config_path="./glm-4-voice-decoder/config.yaml",
flow_ckpt_path="./glm-4-voice-decoder/flow.pt",
hift_ckpt_path="./glm-4-voice-decoder/hift.pt",
device=DEVICE
)
@app.on_event("startup")
async def startup_event():
try:
initialize_models()
except Exception as e:
raise HTTPException(status_code=500, detail=f"Model initialization failed: {str(e)}")
def process_audio(audio_bytes: bytes, target_sr: int = 16000):
# Convert bytes to numpy array
audio_np = np.frombuffer(audio_bytes, dtype=np.int16)
# Resample if necessary
if target_sr != 16000: # Whisper's default sample rate
audio_np = librosa.resample(audio_np, orig_sr=target_sr, target_sr=16000)
return audio_np
@app.post("/api/voice_chat")
async def voice_chat(request: AudioRequest):
try:
# Decode and process audio
audio_bytes = base64.b64decode(request.audio_data)
audio_np = process_audio(audio_bytes, request.sample_rate)
# Extract tokens
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = os.path.join(tmp_dir, "audio.wav")
torchaudio.save(tmp_path, torch.from_numpy(audio_np).unsqueeze(0), request.sample_rate)
audio_tokens = extract_speech_token(whisper_model, feature_extractor, [tmp_path])[0]
if not audio_tokens:
raise HTTPException(400, "No speech detected")
# Generate response
response = requests.post(
"http://localhost:10000/generate_stream",
json={
"prompt": f"<|system|>Respond<|user|>{' '.join(f'<|audio_{x}|>' for x in audio_tokens)}<|assistant|>",
"temperature": 0.7,
"top_p": 0.9,
"max_new_tokens": 256
},
stream=True
)
# Process response
text_tokens = []
audio_tokens = []
audio_offset = glm_tokenizer.convert_tokens_to_ids('<|audio_0|>')
for chunk in response.iter_lines():
token_id = json.loads(chunk)["token_id"]
if token_id >= audio_offset:
audio_tokens.append(token_id - audio_offset)
else:
text_tokens.append(token_id)
# Generate audio
tts_token = torch.tensor(audio_tokens, device=DEVICE).unsqueeze(0)
tts_speech, _ = audio_decoder.token2wav(tts_token)
# Prepare response
buffer = io.BytesIO()
torchaudio.save(buffer, tts_speech.cpu(), 22050, format="wav")
return AudioResponse(
audio_data=base64.b64encode(buffer.getvalue()).decode(),
text_transcript=glm_tokenizer.decode(text_tokens, skip_special_tokens=True)
)
except Exception as e:
raise HTTPException(500, str(e))
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000) |