Fabio
c36510c21e
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
168 lines
10 KiB
C++
168 lines
10 KiB
C++
#include "app/cli/args.h"
|
|
#include "app/streaming/streaming.h"
|
|
#include "engine/models/pocket_tts_small/session.h"
|
|
#include "engine/framework/runtime/registry.h"
|
|
#include "engine/framework/audio/wav_writer.h"
|
|
#include "engine/framework/debug/profiler.h"
|
|
#include "engine/framework/io/json.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <filesystem>
|
|
#include <iostream>
|
|
|
|
int main(int argc, char ** argv) {
|
|
using namespace engine;
|
|
try {
|
|
const auto model = minitts::cli::find_arg(argc, argv, "--model").value();
|
|
const auto requests_file = minitts::cli::find_arg(argc, argv, "--requests").value();
|
|
const auto out = std::filesystem::path(minitts::cli::find_arg(argc, argv, "--out-dir").value());
|
|
const auto log = minitts::cli::find_arg(argc, argv, "--log-file").value();
|
|
const auto backend = minitts::cli::find_arg(argc, argv, "--backend").value_or("cuda");
|
|
const int batch_size = minitts::cli::parse_int_arg(argc, argv, "--batch-size", 4);
|
|
const int threads = minitts::cli::parse_int_arg(argc, argv, "--threads", 8);
|
|
const int seed = minitts::cli::parse_int_arg(argc, argv, "--seed", 42);
|
|
const auto mode = minitts::cli::find_arg(argc, argv, "--mode").value_or("offline");
|
|
if (mode != "offline" && mode != "streaming") throw std::runtime_error("Expected offline or streaming mode");
|
|
const bool streaming = mode == "streaming";
|
|
const bool sequential = minitts::cli::has_arg(argc, argv, "--sequential");
|
|
const auto session_options = minitts::cli::find_arg(argc, argv, "--session-options");
|
|
if (!minitts::cli::has_arg(argc, argv, "--log") || batch_size < 1 || threads < 1)
|
|
throw std::runtime_error("Require --log and positive batch size/threads");
|
|
minitts::cli::require_known_args(argc, argv);
|
|
debug::configure_logging({true, log});
|
|
std::filesystem::create_directories(out);
|
|
const auto fixture = io::json::parse_file(requests_file);
|
|
std::vector<runtime::TaskRequest> requests;
|
|
for (const auto & item : fixture.as_array()) {
|
|
runtime::TaskRequest request;
|
|
request.text_input = runtime::Transcript{item.require("text").as_string(), "en"};
|
|
runtime::VoiceReference speaker;
|
|
speaker.cached_voice_id = item.require("voice_id").as_string();
|
|
request.voice = runtime::VoiceCondition{speaker, std::nullopt};
|
|
if (const auto options = item.find("options")) {
|
|
for (const auto & [key, value] : options->as_object())
|
|
request.options[key] = value.is_string() ? value.as_string() : io::json::stringify(value);
|
|
}
|
|
requests.push_back(std::move(request));
|
|
}
|
|
if (requests.empty()) throw std::runtime_error("Require at least one request");
|
|
runtime::ModelRegistry registry;
|
|
registry.register_loader(models::pocket_tts_small::make_pocket_tts_small_loader());
|
|
runtime::ModelLoadRequest load;
|
|
load.model_path = model;
|
|
load.family_hint = "pocket_tts_small";
|
|
auto loaded = registry.load(load);
|
|
runtime::SessionOptions options;
|
|
options.backend = {minitts::cli::parse_backend(backend), 0, threads};
|
|
if (session_options) {
|
|
const auto parsed = io::json::parse(*session_options);
|
|
for (const auto & [key, value] : parsed.as_object())
|
|
options.options[key] = value.is_string() ? value.as_string() : io::json::stringify(value);
|
|
}
|
|
auto owned = loaded->create_task_session({runtime::VoiceTaskKind::Tts,
|
|
streaming ? runtime::RunMode::Streaming : runtime::RunMode::Offline}, options);
|
|
auto * single = dynamic_cast<runtime::IOfflineVoiceTaskSession *>(owned.get());
|
|
auto * batched = dynamic_cast<runtime::IBatchedOfflineVoiceTaskSession *>(owned.get());
|
|
auto * stream = dynamic_cast<runtime::IStreamingVoiceTaskSession *>(owned.get());
|
|
if (streaming ? stream == nullptr : (!single || !batched)) throw std::runtime_error("Missing requested session interfaces");
|
|
runtime::SessionPreparationRequest preparation;
|
|
preparation.voice = requests.front().voice;
|
|
preparation.options["seed"] = std::to_string(seed);
|
|
owned->prepare(preparation);
|
|
double wall_ms = 0, seconds = 0;
|
|
size_t completed = 0;
|
|
audio::WavWriteOptions wav_options;
|
|
wav_options.format = audio::WavSampleFormat::Float32;
|
|
std::string prepared_voice = *requests.front().voice->speaker->cached_voice_id;
|
|
auto prepared_options = preparation.options;
|
|
for (int pass = 0; pass < 2; ++pass) {
|
|
debug::trace_log_scalar("pocket_tts_small.test.measured_pass", pass);
|
|
for (size_t start = 0; start < requests.size(); start += (sequential || streaming) ? 1 : batch_size) {
|
|
const auto end = std::min(requests.size(), start + ((sequential || streaming) ? 1 : batch_size));
|
|
std::vector<runtime::TaskRequest> group(requests.begin() + start, requests.begin() + end);
|
|
const auto started = std::chrono::steady_clock::now();
|
|
std::vector<runtime::TaskResult> results;
|
|
std::optional<double> ttfa_ms;
|
|
size_t audio_events = 0;
|
|
if (sequential || streaming) {
|
|
auto row_options = group.front().options;
|
|
row_options.try_emplace("seed", std::to_string(seed));
|
|
const auto & voice = *group.front().voice->speaker->cached_voice_id;
|
|
if (voice != prepared_voice || row_options != prepared_options) {
|
|
preparation.voice = group.front().voice;
|
|
preparation.options = row_options;
|
|
owned->prepare(preparation);
|
|
prepared_voice = voice;
|
|
prepared_options = std::move(row_options);
|
|
}
|
|
if (streaming) {
|
|
runtime::AudioBuffer emitted{24000, 1, {}};
|
|
auto result = minitts::app::run_streaming_task(*stream, group.front(),
|
|
[&](const runtime::StreamEvent & event) {
|
|
const bool has_audio = (event.audio_output && !event.audio_output->samples.empty()) ||
|
|
std::any_of(event.named_audio_outputs.begin(), event.named_audio_outputs.end(),
|
|
[](const auto & output) { return !output.audio.samples.empty(); });
|
|
if (!has_audio) return;
|
|
if (!ttfa_ms) ttfa_ms = debug::elapsed_ms(started);
|
|
++audio_events;
|
|
if (event.audio_output && !event.audio_output->samples.empty())
|
|
runtime::append_audio_buffer(emitted, *event.audio_output);
|
|
for (const auto & output : event.named_audio_outputs)
|
|
if (!output.audio.samples.empty()) runtime::append_audio_buffer(emitted, output.audio);
|
|
});
|
|
if (!ttfa_ms || !result.audio_output || emitted.samples != result.audio_output->samples)
|
|
throw std::runtime_error("Streaming events differ from the final waveform or contain no audio");
|
|
if (result.audio_output->samples.size() > 1920 && audio_events < 2)
|
|
throw std::runtime_error("Streaming did not emit incremental audio");
|
|
results.push_back(std::move(result));
|
|
} else {
|
|
results.push_back(single->run(group.front()));
|
|
}
|
|
} else {
|
|
results = batched->run_batch(group);
|
|
}
|
|
const auto elapsed = debug::elapsed_ms(started);
|
|
if (results.size() != group.size()) throw std::runtime_error("Batch result count differs");
|
|
debug::trace_log_scalar("pocket_tts_small.test.request_index", static_cast<int64_t>(start));
|
|
if (ttfa_ms) {
|
|
debug::timing_log_scalar("session.ttfa_ms", *ttfa_ms);
|
|
debug::timing_log_scalar("pocket_tts_small.test.request_wall_ms", elapsed);
|
|
debug::trace_log_scalar("pocket_tts_small.test.audio_events", static_cast<int64_t>(audio_events));
|
|
}
|
|
if (pass == 0) continue;
|
|
wall_ms += elapsed;
|
|
for (size_t row = 0; row < results.size(); ++row) {
|
|
if (!results[row].audio_output) throw std::runtime_error("Missing batch waveform");
|
|
const auto & wav = *results[row].audio_output;
|
|
if (wav.samples.empty() || wav.sample_rate != 24000) throw std::runtime_error("Invalid batch waveform shape");
|
|
for (const float value : wav.samples) if (!std::isfinite(value)) throw std::runtime_error("Nonfinite batch waveform");
|
|
seconds += double(wav.samples.size()) / wav.sample_rate;
|
|
audio::write_wav(out / (std::to_string(start + row) + ".wav"), wav.sample_rate, wav.channels,
|
|
wav.samples, wav_options);
|
|
++completed;
|
|
}
|
|
}
|
|
}
|
|
if (!sequential && !streaming) {
|
|
const auto result = single->run(requests.front());
|
|
if (!result.audio_output || result.audio_output->samples.empty())
|
|
throw std::runtime_error("Single request after batching produced no audio");
|
|
const auto & wav = *result.audio_output;
|
|
for (float value : wav.samples) if (!std::isfinite(value))
|
|
throw std::runtime_error("Single request after batching produced nonfinite audio");
|
|
std::filesystem::create_directories(out / "single-after-batch");
|
|
audio::write_wav(out / "single-after-batch" / "0.wav", wav.sample_rate, wav.channels,
|
|
wav.samples, wav_options);
|
|
debug::trace_log_scalar("pocket_tts_small.test.single_after_batch", 1);
|
|
}
|
|
debug::timing_log_scalar("session.wall_ms", wall_ms);
|
|
debug::trace_log_scalar("pocket_tts_small.test.audio_seconds", seconds);
|
|
debug::trace_log_scalar("pocket_tts_small.test.rtf", wall_ms / (1000 * seconds));
|
|
std::cout << "requests=" << completed << " seconds=" << seconds << " wall_ms=" << wall_ms
|
|
<< " rtf=" << wall_ms / (1000 * seconds) << '\n';
|
|
} catch (const std::exception & error) {
|
|
std::cerr << error.what() << '\n';
|
|
return 1;
|
|
}
|
|
}
|