Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
161 lines
10 KiB
C++
161 lines
10 KiB
C++
#include "engine/models/firered_vad/runtime.h"
|
|
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/io/json.h"
|
|
#include "engine/framework/model_spec/package.h"
|
|
#include "engine/framework/debug/trace.h"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <chrono>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <limits>
|
|
|
|
using namespace engine;
|
|
|
|
void write_floats(const std::string & path, const std::vector<float> & values) {
|
|
std::ofstream file(path, std::ios::binary);
|
|
file.write(reinterpret_cast<const char *>(values.data()), values.size() * sizeof(float));
|
|
if (!file) throw std::runtime_error("Cannot write " + path);
|
|
}
|
|
|
|
std::vector<float> read_floats(const std::string & path) {
|
|
std::ifstream file(path, std::ios::binary | std::ios::ate);
|
|
if (!file || file.tellg() % sizeof(float)) throw std::runtime_error("Cannot read " + path);
|
|
std::vector<float> values(static_cast<size_t>(file.tellg()) / sizeof(float));
|
|
file.seekg(0);
|
|
file.read(reinterpret_cast<char *>(values.data()), values.size() * sizeof(float));
|
|
if (!file) throw std::runtime_error("Cannot read " + path);
|
|
return values;
|
|
}
|
|
|
|
template<class F> double median_ms(F fn) {
|
|
for (int i = 0; i < 5; ++i) fn();
|
|
std::vector<double> times;
|
|
for (int i = 0; i < 21; ++i) {
|
|
auto started = std::chrono::steady_clock::now();
|
|
fn();
|
|
times.push_back(std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - started).count());
|
|
}
|
|
std::sort(times.begin(), times.end());
|
|
return times[times.size() / 2];
|
|
}
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
const bool log = argc > 1 && std::string(argv[argc - 1]) == "--log";
|
|
if (log) --argc;
|
|
debug::configure_logging({log});
|
|
if (argc != 7 && argc != 8) throw std::runtime_error("Usage: probe MODEL AUDIO cpu|metal|cuda|vulkan THREADS OUTPUT_PREFIX REFERENCE_PREFIX [MAX_SPEECH_FRAMES] [--log]");
|
|
auto bundle = model_spec::load_resource_bundle_for_family(argv[1], "firered_vad");
|
|
core::BackendConfig backend;
|
|
const std::string name = argv[3];
|
|
if (name == "cpu") backend.type = core::BackendType::Cpu;
|
|
else if (name == "metal") backend.type = core::BackendType::Metal;
|
|
else if (name == "cuda") backend.type = core::BackendType::Cuda;
|
|
else if (name == "vulkan") backend.type = core::BackendType::Vulkan;
|
|
else throw std::runtime_error("Unsupported backend: " + name);
|
|
backend.threads = std::stoi(argv[4]);
|
|
core::ExecutionContext execution(backend);
|
|
models::firered_vad::FireRedDfsmnRuntime model(bundle.open_tensor_source("weights"), execution);
|
|
auto wav = audio::read_wav_f32(std::filesystem::path(argv[2]));
|
|
runtime::AudioBuffer input{wav.sample_rate, wav.channels, std::move(wav.samples)};
|
|
const std::string output = argv[5], reference = argv[6];
|
|
models::firered_vad::DetectionOptions options;
|
|
if (model.causal()) { options.threshold = 0.5; options.min_speech_frames = 8; }
|
|
if (argc == 8) options.max_speech_frames = std::stoi(argv[7]);
|
|
const auto cold_start = std::chrono::steady_clock::now();
|
|
auto result = model.detect(input, options);
|
|
const double cold_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - cold_start).count();
|
|
auto features = model.extract_features(input);
|
|
write_floats(output + ".features.f32", features);
|
|
write_floats(output + ".probs.f32", model.infer_features(features));
|
|
auto reference_features = read_floats(reference + ".features.f32");
|
|
if (reference_features.size() != features.size()) throw std::runtime_error("Frontend frame count mismatch");
|
|
write_floats(output + ".reference_features_probs.f32", model.infer_features(reference_features));
|
|
io::json::Value::Array segments;
|
|
for (const auto & segment : result.speech_segments)
|
|
segments.push_back(io::json::Value::make_array({io::json::Value::make_number(segment.span.start_sample / 16000.0), io::json::Value::make_number(segment.span.end_sample / 16000.0)}));
|
|
std::ofstream(output + ".segments.json") << io::json::stringify(io::json::Value::make_array(std::move(segments))) << '\n';
|
|
if (model.causal()) {
|
|
model.reset(options);
|
|
std::vector<float> chunked;
|
|
for (size_t start = 0; start < reference_features.size(); start += 16 * 80) {
|
|
const size_t end = std::min(reference_features.size(), start + 16 * 80);
|
|
auto values = model.infer_features({reference_features.begin() + start, reference_features.begin() + end}, true);
|
|
chunked.insert(chunked.end(), values.begin(), values.end());
|
|
}
|
|
write_floats(output + ".chunked_probs.f32", chunked);
|
|
const auto full = model.infer_features(reference_features);
|
|
for (size_t chunk_frames : {3, 4, 15, 16, 17, 31, 32, 33, 63, 64, 65, 255, 256, 257}) {
|
|
model.reset(options);
|
|
for (size_t start = 0; start < reference_features.size(); start += chunk_frames * 80) {
|
|
const size_t end = std::min(reference_features.size(), start + chunk_frames * 80);
|
|
const auto values = model.infer_features({reference_features.begin() + start, reference_features.begin() + end}, true);
|
|
for (size_t frame = 0; frame < values.size(); ++frame)
|
|
if (std::abs(values[frame] - full[start / 80 + frame]) > 0.002f)
|
|
throw std::runtime_error("Streaming probability mismatch at convolution block boundary");
|
|
}
|
|
}
|
|
std::cout << "block_boundary_parity=passed\n";
|
|
model.reset(options);
|
|
const auto stream_start = std::chrono::steady_clock::now();
|
|
double first_event_ms = -1, worst_chunk_ms = 0;
|
|
size_t first_event_samples = 0;
|
|
for (size_t start = 0; start < input.samples.size(); start += 2560) {
|
|
const size_t end = std::min(input.samples.size(), start + 2560);
|
|
const auto chunk_start = std::chrono::steady_clock::now();
|
|
auto event = model.process({16000, 1, static_cast<int64_t>(start), {input.samples.begin() + start, input.samples.begin() + end}});
|
|
worst_chunk_ms = std::max(worst_chunk_ms, std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - chunk_start).count());
|
|
if (first_event_ms < 0 && !event.voice_activity.empty()) {
|
|
first_event_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - stream_start).count();
|
|
first_event_samples = end;
|
|
}
|
|
}
|
|
auto streamed = model.finalize();
|
|
if (streamed.speech_segments.size() != result.speech_segments.size()) throw std::runtime_error("Streaming segment count mismatch");
|
|
for (size_t i = 0; i < streamed.speech_segments.size(); ++i)
|
|
if (streamed.speech_segments[i].span.start_sample != result.speech_segments[i].span.start_sample ||
|
|
streamed.speech_segments[i].span.end_sample != result.speech_segments[i].span.end_sample)
|
|
throw std::runtime_error("Streaming timestamp mismatch");
|
|
model.reset(options);
|
|
const std::array<size_t, 6> packets{1, 399, 160, 640, 2560, 7777};
|
|
for (size_t start = 0, packet = 0; start < input.samples.size(); ++packet) {
|
|
const size_t end = std::min(input.samples.size(), start + packets[packet % packets.size()]);
|
|
model.process({16000, 1, static_cast<int64_t>(start), {input.samples.begin() + start, input.samples.begin() + end}});
|
|
start = end;
|
|
}
|
|
streamed = model.finalize();
|
|
if (streamed.speech_segments.size() != result.speech_segments.size()) throw std::runtime_error("Irregular streaming segment count mismatch");
|
|
for (size_t i = 0; i < streamed.speech_segments.size(); ++i)
|
|
if (streamed.speech_segments[i].span.start_sample != result.speech_segments[i].span.start_sample ||
|
|
streamed.speech_segments[i].span.end_sample != result.speech_segments[i].span.end_sample)
|
|
throw std::runtime_error("Irregular streaming timestamp mismatch");
|
|
std::cout << "streaming_parity=passed first_event_compute_ms=" << first_event_ms
|
|
<< " first_event_input_sec=" << first_event_samples / 16000.0
|
|
<< " worst_chunk_ms=" << worst_chunk_ms << '\n';
|
|
}
|
|
if (!model.detect({16000, 1, {}}, options).speech_segments.empty() ||
|
|
!model.detect({16000, 1, std::vector<float>(399, 0.1f)}, options).speech_segments.empty() ||
|
|
!model.detect({16000, 1, std::vector<float>(16000, 0.0f)}, options).speech_segments.empty())
|
|
throw std::runtime_error("Empty, short or silent audio produced speech");
|
|
bool rejected = false;
|
|
try { model.extract_features({8000, 1, {0.1f}}); } catch (const std::runtime_error &) { rejected = true; }
|
|
if (!rejected) throw std::runtime_error("Wrong sample rate was accepted");
|
|
rejected = false;
|
|
try { model.extract_features({16000, 1, {std::numeric_limits<float>::quiet_NaN()}}); } catch (const std::runtime_error &) { rejected = true; }
|
|
if (!rejected) throw std::runtime_error("Non-finite audio was accepted");
|
|
std::cout << "input_validation=passed\n";
|
|
const double frontend = median_ms([&] { model.extract_features(input); });
|
|
const double network = median_ms([&] { model.infer_features(features); });
|
|
const double total = median_ms([&] { model.detect(input, options); });
|
|
debug::timing_log_scalar("session.wall_ms", total);
|
|
std::cout << "frames=" << features.size() / 80 << " backend=" << argv[3] << " threads=" << backend.threads
|
|
<< " cold_ms=" << cold_ms << " frontend_ms=" << frontend << " network_ms=" << network << " total_ms=" << total
|
|
<< " rtf=" << total * 16.0 / input.samples.size() << '\n';
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << error.what() << '\n';
|
|
return 1;
|
|
}
|
|
}
|