Voice_Assistant.AXERA / zipvoice_cpp_server /zipvoice_resident_server.cpp
HY-2012's picture
Upload ASR and TTS models
a37072b verified
Raw
History Blame Contribute Delete
24.4 kB
#include "zipvoice_resident_server.hpp"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <sstream>
#include "src/EngineWrapper.hpp"
#include "src/fbank.hpp"
#include "src/tokenizer.hpp"
#include "src/vocoder.hpp"
#include "src/zipvoice_engine.hpp"
#include "src/wav_writer.hpp"
// ---------------------------------------------------------------------------
// Simple WAV writer (inlined; avoids external soundfile dependency)
// ---------------------------------------------------------------------------
static bool WriteWav(const std::string& path, const std::vector<float>& samples,
int sample_rate, int bits_per_sample = 16) {
std::ofstream file(path, std::ios::binary);
if (!file.is_open()) return false;
int num_samples = (int)samples.size();
int byte_rate = sample_rate * (bits_per_sample / 8);
int block_align = bits_per_sample / 8;
int data_size = num_samples * block_align;
// RIFF header
file.write("RIFF", 4);
uint32_t chunk_size = 36 + data_size;
file.write(reinterpret_cast<const char*>(&chunk_size), 4);
file.write("WAVE", 4);
// fmt subchunk
file.write("fmt ", 4);
uint32_t subchunk1_size = 16;
uint16_t audio_format = 1; // PCM
uint16_t num_channels = 1; // mono
uint16_t bps = (uint16_t)bits_per_sample;
file.write(reinterpret_cast<const char*>(&subchunk1_size), 4);
file.write(reinterpret_cast<const char*>(&audio_format), 2);
file.write(reinterpret_cast<const char*>(&num_channels), 2);
uint32_t sr = (uint32_t)sample_rate;
file.write(reinterpret_cast<const char*>(&sr), 4);
uint32_t br = (uint32_t)byte_rate;
file.write(reinterpret_cast<const char*>(&br), 4);
file.write(reinterpret_cast<const char*>(&block_align), 2);
file.write(reinterpret_cast<const char*>(&bps), 2);
// data subchunk
file.write("data", 4);
file.write(reinterpret_cast<const char*>(&data_size), 4);
// Write samples as int16 PCM
for (float s : samples) {
float clamped = std::max(-1.0f, std::min(1.0f, s));
int16_t v = (int16_t)(clamped * 32767.0f);
file.write(reinterpret_cast<const char*>(&v), sizeof(v));
}
file.close();
return true;
}
// ---------------------------------------------------------------------------
// WAV file helpers (ported from zipvoice.cpp)
// ---------------------------------------------------------------------------
static int ReadWavFile(const std::string& path, std::vector<float>& samples, int& sample_rate) {
std::ifstream file(path, std::ios::binary);
if (!file.is_open()) {
fprintf(stderr, "Failed to open: %s\n", path.c_str());
return -1;
}
char riff[5] = {};
file.read(riff, 4);
if (std::strncmp(riff, "RIFF", 4) != 0) {
fprintf(stderr, "Not a valid WAV: %s\n", path.c_str());
return -1;
}
uint32_t file_size;
file.read(reinterpret_cast<char*>(&file_size), 4);
char wave[5] = {};
file.read(wave, 4);
if (std::strncmp(wave, "WAVE", 4) != 0) {
fprintf(stderr, "Not a valid WAV: %s\n", path.c_str());
return -1;
}
int num_channels = 1, bits_per_sample = 16;
sample_rate = 24000;
uint32_t data_size = 0;
while (file.good()) {
char chunk_id[5] = {};
file.read(chunk_id, 4);
uint32_t chunk_size;
file.read(reinterpret_cast<char*>(&chunk_size), 4);
if (std::strncmp(chunk_id, "fmt ", 4) == 0) {
uint16_t fmt, ch, bps; uint32_t sr, br; uint16_t ba;
file.read(reinterpret_cast<char*>(&fmt), 2);
file.read(reinterpret_cast<char*>(&ch), 2);
file.read(reinterpret_cast<char*>(&sr), 4);
file.read(reinterpret_cast<char*>(&br), 4);
file.read(reinterpret_cast<char*>(&ba), 2);
file.read(reinterpret_cast<char*>(&bps), 2);
num_channels = ch; sample_rate = sr; bits_per_sample = bps;
if (chunk_size > 16) file.seekg(chunk_size - 16, std::ios::cur);
} else if (std::strncmp(chunk_id, "data", 4) == 0) {
data_size = chunk_size;
break;
} else {
file.seekg(chunk_size, std::ios::cur);
}
}
if (data_size == 0) { fprintf(stderr, "No data chunk\n"); return -1; }
int num_samples = data_size / (bits_per_sample / 8) / num_channels;
if (bits_per_sample == 16) {
std::vector<int16_t> raw(num_samples * num_channels);
file.read(reinterpret_cast<char*>(raw.data()), data_size);
samples.resize(num_samples);
for (int i = 0; i < num_samples; ++i)
samples[i] = raw[i * num_channels] / 32768.0f;
} else if (bits_per_sample == 32) {
samples.resize(num_samples * num_channels);
file.read(reinterpret_cast<char*>(samples.data()), data_size);
std::vector<float> mono(num_samples);
for (int i = 0; i < num_samples; ++i) mono[i] = samples[i * num_channels];
samples = std::move(mono);
} else {
fprintf(stderr, "Unsupported bit depth: %d\n", bits_per_sample);
return -1;
}
return 0;
}
static std::vector<float> ResampleLinear(const std::vector<float>& samples,
int orig_sr, int target_sr) {
if (orig_sr == target_sr) return samples;
int old_len = (int)samples.size();
int new_len = std::max(1, (int)std::round((double)old_len * target_sr / orig_sr));
std::vector<float> result(new_len);
for (int i = 0; i < new_len; ++i) {
double pos = (double)i * (old_len - 1) / (new_len - 1);
int idx = (int)pos;
double frac = pos - idx;
if (idx + 1 < old_len)
result[i] = (float)(samples[idx] * (1.0 - frac) + samples[idx + 1] * frac);
else
result[i] = samples[old_len - 1];
}
return result;
}
static float ComputeRms(const std::vector<float>& samples) {
if (samples.empty()) return 0.0f;
float sum_sq = 0.0f;
for (float s : samples) sum_sq += s * s;
return std::sqrt(sum_sq / (float)samples.size());
}
static void RmsNormalize(std::vector<float>& samples, float target_rms) {
float rms = ComputeRms(samples);
if (rms < target_rms && rms > 1e-10f) {
float gain = target_rms / rms;
for (float& s : samples) s *= gain;
}
}
// ---------------------------------------------------------------------------
// UTF-8 / Chinese helpers (ported from zipvoice.cpp)
// ---------------------------------------------------------------------------
static bool IsUtf8Lead(unsigned char c) { return (c & 0xC0) != 0x80; }
static std::vector<std::string> SplitUtf8Chars(const std::string& text) {
std::vector<std::string> chars;
for (size_t i = 0; i < text.size();) {
unsigned char c = (unsigned char)text[i];
size_t len = 1;
if ((c & 0x80) == 0) len = 1;
else if ((c & 0xE0) == 0xC0) len = 2;
else if ((c & 0xF0) == 0xE0) len = 3;
else if ((c & 0xF8) == 0xF0) len = 4;
chars.push_back(text.substr(i, len));
i += len;
}
return chars;
}
static bool IsChineseUtf8Char(const std::string& ch) {
if (ch.size() != 3) return false;
unsigned char b0 = (unsigned char)ch[0], b1 = (unsigned char)ch[1], b2 = (unsigned char)ch[2];
uint32_t cp = ((b0 & 0x0F) << 12) | ((b1 & 0x3F) << 6) | (b2 & 0x3F);
return cp >= 0x4E00 && cp <= 0x9FFF;
}
static bool IsSplitPunct(const std::string& ch) {
static const char* puncts[] = {
".", "!", "?", ";", ",", ":",
"。", "!", "?", ";", ",", "、", ":"
};
for (auto* p : puncts) if (ch == p) return true;
return false;
}
static std::string TrimAsciiSpaces(const std::string& s) {
size_t start = 0, end = s.size();
while (start < end && (s[start] == ' ' || s[start] == '\t' || s[start] == '\n' || s[start] == '\r')) start++;
while (end > start && (s[end-1] == ' ' || s[end-1] == '\t' || s[end-1] == '\n' || s[end-1] == '\r')) end--;
return s.substr(start, end - start);
}
static std::string JoinUnits(const std::string& left, const std::string& right) {
if (left.empty()) return TrimAsciiSpaces(right);
std::string r = TrimAsciiSpaces(right);
if (r.empty()) return TrimAsciiSpaces(left);
auto lchars = SplitUtf8Chars(left), rchars = SplitUtf8Chars(r);
bool zh = (!lchars.empty() && IsChineseUtf8Char(lchars.back())) ||
(!rchars.empty() && IsChineseUtf8Char(rchars.front()));
return zh ? (TrimAsciiSpaces(left) + r) : (TrimAsciiSpaces(left) + " " + r);
}
static std::vector<std::string> SplitUnits(const std::string& text) {
std::vector<std::string> units;
auto chars = SplitUtf8Chars(TrimAsciiSpaces(text));
std::string current;
for (const auto& ch : chars) {
current += ch;
if (IsSplitPunct(ch)) {
std::string t = TrimAsciiSpaces(current);
if (!t.empty()) units.push_back(t);
current.clear();
}
}
current = TrimAsciiSpaces(current);
if (!current.empty()) units.push_back(current);
if (units.empty() && !text.empty()) units.push_back(TrimAsciiSpaces(text));
return units;
}
// ---------------------------------------------------------------------------
// Segment estimation / building (ported from zipvoice.cpp BuildSegmentsCpp)
// ---------------------------------------------------------------------------
struct SegmentInfo {
std::string text;
int text_tokens = 0;
int raw_features_len = 0;
int features_len = 0;
int generated_frames = 0;
};
static int TokenCount(Tokenizer& tok, const std::string& text) {
return (int)tok.TextToTokenIds(text).size();
}
static std::vector<std::string> SplitLongUnit(Tokenizer& tok, const std::string& unit,
int max_text_tokens) {
if (TokenCount(tok, unit) <= max_text_tokens) return {unit};
std::vector<std::string> chunks;
if (unit.find(' ') != std::string::npos) {
std::stringstream ss(unit);
std::string piece, current;
while (ss >> piece) {
std::string candidate = JoinUnits(current, piece);
if (!current.empty() && TokenCount(tok, candidate) > max_text_tokens) {
chunks.push_back(current); current = piece;
} else { current = candidate; }
}
if (!current.empty()) chunks.push_back(current);
return chunks;
}
auto chars = SplitUtf8Chars(unit);
std::string current;
for (const auto& ch : chars) {
std::string candidate = current + ch;
if (!current.empty() && TokenCount(tok, candidate) > max_text_tokens) {
chunks.push_back(current); current = ch;
} else { current = candidate; }
}
if (!current.empty()) chunks.push_back(current);
return chunks;
}
static SegmentInfo EstimateSegment(Tokenizer& tok, const std::string& text,
int prompt_frames, int prompt_tokens_len,
float speed, int max_feat_len) {
SegmentInfo s;
s.text = text;
s.text_tokens = TokenCount(tok, text);
s.raw_features_len = (int)std::ceil((double)prompt_frames / prompt_tokens_len *
(prompt_tokens_len + s.text_tokens) / speed);
s.features_len = std::min(s.raw_features_len, max_feat_len);
s.generated_frames = s.features_len - prompt_frames;
if (s.generated_frames <= 0) s.generated_frames = s.features_len;
return s;
}
static std::vector<SegmentInfo> BuildSegments(Tokenizer& tok, const std::string& text,
int prompt_frames, int prompt_tokens_len,
float speed, int max_feat_len,
int max_text_tokens, int min_gen_frames,
int max_gen_frames, double max_raw_ratio) {
auto raw_units = SplitUnits(text);
std::vector<std::string> units;
for (const auto& u : raw_units) {
auto split = SplitLongUnit(tok, u, max_text_tokens);
units.insert(units.end(), split.begin(), split.end());
}
std::vector<SegmentInfo> segments;
std::string current;
for (const auto& unit : units) {
std::string candidate = JoinUnits(current, unit);
auto cand = EstimateSegment(tok, candidate, prompt_frames, prompt_tokens_len,
speed, max_feat_len);
bool raw_too_long = cand.raw_features_len > (int)(max_feat_len * max_raw_ratio);
bool too_long = cand.text_tokens > max_text_tokens ||
cand.generated_frames > max_gen_frames || raw_too_long;
if (!current.empty() && too_long) {
segments.push_back(EstimateSegment(tok, current, prompt_frames,
prompt_tokens_len, speed, max_feat_len));
current = unit;
} else { current = candidate; }
}
if (!current.empty())
segments.push_back(EstimateSegment(tok, current, prompt_frames,
prompt_tokens_len, speed, max_feat_len));
if (segments.size() >= 2 && segments.back().generated_frames < min_gen_frames) {
std::string merged = JoinUnits(segments[segments.size()-2].text, segments.back().text);
auto m = EstimateSegment(tok, merged, prompt_frames, prompt_tokens_len, speed, max_feat_len);
bool raw_ok = m.raw_features_len <= (int)(max_feat_len * max_raw_ratio);
if (m.text_tokens <= max_text_tokens && m.generated_frames <= max_gen_frames && raw_ok) {
segments[segments.size()-2] = m;
segments.pop_back();
}
}
return segments;
}
// ---------------------------------------------------------------------------
// ZipVoiceResidentServer implementation
// ---------------------------------------------------------------------------
static double GetCurrentTimeMs() {
struct timeval tv; gettimeofday(&tv, nullptr);
return tv.tv_sec * 1000.0 + tv.tv_usec / 1000.0;
}
ZipVoiceResidentServer::ZipVoiceResidentServer() = default;
ZipVoiceResidentServer::~ZipVoiceResidentServer() = default;
bool ZipVoiceResidentServer::Init(const std::string& root_dir,
const std::string& model_dir,
const std::string& token_file,
const std::string& vocoder_model,
int max_tokens, int max_feat_len,
int num_step, float default_speed, int seed) {
root_dir_ = root_dir;
model_dir_ = model_dir;
token_file_ = token_file;
vocoder_model_ = vocoder_model;
max_tokens_ = max_tokens;
max_feat_len_ = max_feat_len;
num_step_ = num_step;
default_speed_ = default_speed;
seed_ = seed;
// 1. Init AX_SYS
fprintf(stderr, "[zipvoice_daemon] Init AX_SYS...\n");
if (!InitAxSys()) { fprintf(stderr, "[zipvoice_daemon] AX_SYS_Init FAILED\n"); return false; }
fprintf(stderr, "[zipvoice_daemon] AX_SYS OK\n");
// 2. Load tokenizer
fprintf(stderr, "[zipvoice_daemon] Loading tokenizer...\n");
if (!LoadTokenizer()) { fprintf(stderr, "[zipvoice_daemon] LoadTokenizer FAILED\n"); return false; }
fprintf(stderr, "[zipvoice_daemon] Tokenizer OK\n");
// 3. Load ZipVoice engine (encoder + decoder4)
fprintf(stderr, "[zipvoice_daemon] Loading engine...\n");
if (!LoadRuntime()) { fprintf(stderr, "[zipvoice_daemon] LoadRuntime FAILED\n"); return false; }
fprintf(stderr, "[zipvoice_daemon] Engine OK\n");
// 4. Load vocoder
fprintf(stderr, "[zipvoice_daemon] Loading vocoder...\n");
if (!LoadVocoder()) { fprintf(stderr, "[zipvoice_daemon] LoadVocoder FAILED\n"); return false; }
fprintf(stderr, "[zipvoice_daemon] Vocoder OK\n");
// 5. Cache both prompt feature sets
fprintf(stderr, "[zipvoice_daemon] Caching prompts...\n");
if (!CachePrompts()) { fprintf(stderr, "[zipvoice_daemon] CachePrompts FAILED\n"); return false; }
fprintf(stderr, "[zipvoice_daemon] Prompts cached OK\n");
return true;
}
bool ZipVoiceResidentServer::InitAxSys() {
int ret = AX_SYS_Init();
if (ret != 0) {
fprintf(stderr, "AX_SYS_Init failed: 0x%x\n", ret);
return false;
}
AX_ENGINE_NPU_ATTR_T npu_attr;
memset(&npu_attr, 0, sizeof(npu_attr));
npu_attr.eHardMode = static_cast<AX_ENGINE_NPU_MODE_T>(0);
ret = AX_ENGINE_Init(&npu_attr);
if (ret != 0) {
fprintf(stderr, "AX_ENGINE_Init failed: 0x%x\n", ret);
return false;
}
ax_inited_ = true;
return true;
}
bool ZipVoiceResidentServer::LoadTokenizer() {
auto tok = std::make_unique<Tokenizer>();
if (tok->Load(token_file_) != 0) {
fprintf(stderr, "Tokenizer load failed: %s\n", token_file_.c_str());
return false;
}
tokenizer_ = std::move(tok);
return true;
}
bool ZipVoiceResidentServer::LoadRuntime() {
auto eng = std::make_unique<ZipVoiceEngine>();
if (eng->Init(model_dir_) != 0) {
fprintf(stderr, "ZipVoiceEngine init failed: %s\n", model_dir_.c_str());
return false;
}
engine_ = std::move(eng);
return true;
}
bool ZipVoiceResidentServer::LoadVocoder() {
auto v = std::make_unique<Vocoder>();
Vocoder::Config vcfg;
vcfg.model_path = vocoder_model_;
if (v->Init(vcfg) != 0) {
fprintf(stderr, "Vocoder init failed: %s\n", vocoder_model_.c_str());
return false;
}
vocoder_ = std::move(v);
return true;
}
bool ZipVoiceResidentServer::CachePrompts() {
// Prompt definitions: zh + en
struct PromptDef {
std::string key, wav, text;
};
std::vector<PromptDef> defs = {
{"zh", root_dir_ + "/assets/moss_prompts/zh_1_4p5s.wav",
"不管怎么样我和汤姆还是要感谢贝尔卡金的援手"},
{"en", root_dir_ + "/assets/moss_prompts/en_4_4p5s.wav",
"This is almost twice the current industry production level per train."},
};
MelFilterBank fbank;
fbank.Init();
for (const auto& d : defs) {
ZipVoicePromptCache cache;
cache.key = d.key;
cache.prompt_text = d.text;
cache.prompt_wav = d.wav;
// Tokenize
cache.prompt_tokens = tokenizer_->TextToTokenIds(d.text);
// Read prompt wav
std::vector<float> wav;
int wav_sr = 0;
if (ReadWavFile(d.wav, wav, wav_sr) != 0) return false;
// Resample to 24kHz if needed
std::vector<float> resampled = ResampleLinear(wav, wav_sr, 24000);
// RMS normalize
float target_rms = 0.1f;
RmsNormalize(resampled, target_rms);
float prompt_rms = ComputeRms(resampled);
// Extract mel features
auto features = fbank.Extract(resampled, 24000);
int T = (int)features.size() / 100;
std::vector<float> feat_flat(features.size());
std::copy(features.begin(), features.end(), feat_flat.begin());
for (auto& f : feat_flat) f *= 0.1f; // feat_scale
cache.prompt_features = feat_flat;
cache.prompt_frames = T;
cache.prompt_rms = prompt_rms;
prompts_[d.key] = cache;
}
return true;
}
ZipVoiceResponse ZipVoiceResidentServer::Infer(const ZipVoiceRequest& req) {
ZipVoiceResponse resp;
resp.id = req.id;
auto it = prompts_.find(req.prompt);
if (it == prompts_.end()) {
resp.ok = false;
resp.error = "Unknown prompt: " + req.prompt;
return resp;
}
auto& prompt = it->second;
float speed = req.speed > 0.01f ? req.speed : default_speed_;
float feat_scale = 0.1f, target_rms = 0.1f;
int min_gen = 360, max_gen = 620;
double max_raw_ratio = 1.2;
double t0 = GetCurrentTimeMs();
// Build segments
int max_text_tokens = max_tokens_ - (int)prompt.prompt_tokens.size() - 1;
if (max_text_tokens <= 0) {
resp.ok = false;
resp.error = "prompt tokens leave no room";
return resp;
}
auto segments = BuildSegments(*tokenizer_, req.text, prompt.prompt_frames,
(int)prompt.prompt_tokens.size(), speed,
max_feat_len_, max_text_tokens, min_gen, max_gen, max_raw_ratio);
// Infer each segment
std::vector<float> all_audio;
double model_ms = 0.0;
for (size_t si = 0; si < segments.size(); ++si) {
auto& seg = segments[si];
// Tokenize
auto text_tokens = tokenizer_->TextToTokenIds(seg.text);
// Build cat_tokens
int pad_id = tokenizer_->GetPadId();
std::vector<int> cat = prompt.prompt_tokens;
cat.insert(cat.end(), text_tokens.begin(), text_tokens.end());
cat.push_back(pad_id);
std::vector<int32_t> cat_padded(max_tokens_, pad_id);
for (size_t i = 0; i < cat.size() && i < (size_t)max_tokens_; ++i)
cat_padded[i] = (int32_t)cat[i];
// Run sample
ZipVoiceEngine::Timing timing;
std::vector<float> output_features;
if (engine_->Sample(cat_padded, (int)prompt.prompt_tokens.size(),
(int)text_tokens.size(), prompt.prompt_features,
prompt.prompt_frames, speed, 1.0f, seed_ + (int)si,
output_features, timing) != 0) {
resp.ok = false;
resp.error = "Inference failed on segment " + std::to_string(si);
return resp;
}
model_ms += timing.total_time_sec * 1000.0;
// Vocoder decode
std::vector<float> audio;
if (vocoder_->Decode(output_features, timing.generated_frames, feat_scale, audio) != 0) {
resp.ok = false;
resp.error = "Vocoder decode failed on segment " + std::to_string(si);
return resp;
}
// RMS normalize
RmsNormalize(audio, target_rms);
if (prompt.prompt_rms < target_rms) {
float scale = prompt.prompt_rms / target_rms;
for (float& s : audio) s *= scale;
}
// Silence between segments
if (!all_audio.empty()) {
std::vector<float> silence(24000 * 140 / 1000, 0.0f); // 140ms
all_audio.insert(all_audio.end(), silence.begin(), silence.end());
}
all_audio.insert(all_audio.end(), audio.begin(), audio.end());
}
// Write WAV
if (!WriteWav(req.output_wav, all_audio, 24000, 16)) {
resp.ok = false;
resp.error = "WriteWav failed: " + req.output_wav;
return resp;
}
double wall_ms = GetCurrentTimeMs() - t0;
double audio_sec = all_audio.size() / 24000.0;
resp.ok = true;
resp.path = req.output_wav;
resp.wall_sec = wall_ms / 1000.0;
resp.model_sec = model_ms / 1000.0;
resp.rtf = audio_sec > 0 ? (model_ms / 1000.0 / audio_sec) : 0.0;
return resp;
}
bool ZipVoiceResidentServer::CachePrompt(const std::string& key,
const std::string& wav_path,
const std::string& text) {
MelFilterBank fbank;
fbank.Init();
ZipVoicePromptCache cache;
cache.key = key;
cache.prompt_text = text;
cache.prompt_wav = wav_path;
cache.prompt_tokens = tokenizer_->TextToTokenIds(text);
std::vector<float> wav;
int wav_sr = 0;
if (ReadWavFile(wav_path, wav, wav_sr) != 0) return false;
std::vector<float> resampled = ResampleLinear(wav, wav_sr, 24000);
float target_rms = 0.1f;
RmsNormalize(resampled, target_rms);
float prompt_rms = ComputeRms(resampled);
auto features = fbank.Extract(resampled, 24000);
int T = (int)features.size() / 100;
std::vector<float> feat_flat(features.size());
std::copy(features.begin(), features.end(), feat_flat.begin());
for (auto& f : feat_flat) f *= 0.1f;
cache.prompt_features = feat_flat;
cache.prompt_frames = T;
cache.prompt_rms = prompt_rms;
prompts_[key] = cache;
fprintf(stderr, "[zipvoice_daemon] Custom prompt cached: key=%s tokens=%d frames=%d\n",
key.c_str(), (int)cache.prompt_tokens.size(), T);
return true;
}