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
173 lines
11 KiB
C++
173 lines
11 KiB
C++
#include "engine/framework/runtime/registry.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/audio/wav_writer.h"
|
|
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/io/json.h"
|
|
#include "engine/framework/debug/trace.h"
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <filesystem>
|
|
#include <iostream>
|
|
#include <stdexcept>
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
if (argc < 3) throw std::runtime_error("usage: kitten_tts2_session_probe MODEL OUTPUT_DIR [THREADS] [WEIGHT_TYPE] [REFERENCE_MODEL_DIR|-] [cpu|cuda|metal] [EXEC_PLAN_JSON]");
|
|
if (std::filesystem::path(argv[1]).extension() == ".gguf") {
|
|
// Check stored precision, not just the runtime's upcast tensor types.
|
|
// A misspelled namespace in converter overrides silently quantizes
|
|
// these tensors and can otherwise escape an end-to-end smoke test.
|
|
namespace assets = engine::assets;
|
|
auto source = assets::open_tensor_source(argv[1]);
|
|
const auto body_type = assets::tensor_storage_type_for_dtype(
|
|
source->require_metadata("language_model/model.layers.0.self_attn.q_proj.weight").dtype);
|
|
if (body_type != assets::TensorStorageType::BF16 && body_type != assets::TensorStorageType::Q8_0 &&
|
|
body_type != assets::TensorStorageType::Q4_0)
|
|
throw std::runtime_error("incorrect packaged LM body precision");
|
|
const auto embedding_type = body_type == assets::TensorStorageType::Q4_0
|
|
? assets::TensorStorageType::Q8_0 : assets::TensorStorageType::BF16;
|
|
const std::pair<const char *, assets::TensorStorageType> expected[] = {
|
|
{"language_model/model.embed_tokens.weight", embedding_type},
|
|
{"language_model/spk_proj.0.weight", assets::TensorStorageType::BF16},
|
|
{"s3gen/flow.encoder_proj.weight", assets::TensorStorageType::F32},
|
|
{"speaker/embedding.weight", assets::TensorStorageType::F32},
|
|
};
|
|
for (const auto & [name, type] : expected)
|
|
if (assets::tensor_storage_type_for_dtype(source->require_metadata(name).dtype) != type)
|
|
throw std::runtime_error(std::string("incorrect packaged precision: ") + name);
|
|
std::cout << "PASS: packaged tensor precision" << std::endl;
|
|
}
|
|
engine::debug::configure_logging({true,std::nullopt});
|
|
engine::runtime::ModelLoadRequest load;
|
|
load.model_path = argv[1];
|
|
load.family_hint = "kitten_tts2";
|
|
auto registry = engine::runtime::make_default_registry();
|
|
auto model = registry.load(load);
|
|
engine::runtime::SessionOptions options;
|
|
options.backend.type = engine::core::BackendType::Cpu;
|
|
if (argc > 6) {
|
|
const std::string backend = argv[6];
|
|
if (backend == "cuda") options.backend.type = engine::core::BackendType::Cuda;
|
|
else if (backend == "metal") options.backend.type = engine::core::BackendType::Metal;
|
|
else if (backend != "cpu") throw std::runtime_error("probe backend must be cpu, cuda or metal");
|
|
}
|
|
options.backend.threads = argc > 3 ? std::stoi(argv[3]) : 8;
|
|
if (argc > 4 && std::string(argv[4]) != "default") options.options["kitten_tts2.weight_type"] = argv[4];
|
|
if (argc > 7) options.options["kitten_tts2.exec_plan"] = argv[7];
|
|
const engine::runtime::TaskSpec task{engine::runtime::VoiceTaskKind::Tts, engine::runtime::RunMode::Offline};
|
|
for (const auto * plan : {
|
|
R"({"unknown":{"impl":"ggml"}})",
|
|
R"({"vocoder":{"impl":"unknown"}})",
|
|
R"({"vocoder":{"params":{"unknown":true}}})",
|
|
R"({"vocoder":{"impl":"neon","backend":"metal"}})",
|
|
R"({"vocoder":{"impl":"neon","params":{"threads":1}}})"}) {
|
|
auto invalid = options;
|
|
invalid.options["kitten_tts2.exec_plan"] = plan;
|
|
bool rejected = false;
|
|
try { model->create_task_session(task, invalid); }
|
|
catch (const std::runtime_error &) { rejected = true; }
|
|
if (!rejected) throw std::runtime_error("invalid Kitten execution plan was accepted");
|
|
}
|
|
std::cout << "PASS: execution-plan validation" << std::endl;
|
|
auto base = model->create_task_session(task, options);
|
|
auto * session = dynamic_cast<engine::runtime::IOfflineVoiceTaskSession *>(base.get());
|
|
if (!session) throw std::runtime_error("expected offline session");
|
|
base->prepare({});
|
|
std::filesystem::create_directories(argv[2]);
|
|
engine::runtime::TaskRequest request;
|
|
request.text_input = engine::runtime::Transcript{"Hello there. This is a test of Kitten speech.", "en"};
|
|
request.options = {{"seed", "1234"}, {"voice_id", "Bruno"}};
|
|
std::vector<float> first;
|
|
for (int i = 0; i < 4; ++i) {
|
|
if (i == 2) request.options["voice_id"] = "Bella";
|
|
if (i == 3) {
|
|
request.options["voice_id"] = "Bruno";
|
|
request.options["text_chunk_size"] = "100";
|
|
request.text_input->text = "The morning sun fills the room with light. A small bird sings outside the window. "
|
|
"Today we are testing long form speech, with each sentence joining the next. "
|
|
"The same model and decoder stay loaded for all of these requests.";
|
|
}
|
|
const auto started = std::chrono::steady_clock::now();
|
|
auto result = session->run(request);
|
|
const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now() - started).count();
|
|
if (!result.audio_output || result.audio_output->sample_rate != 24000 || result.audio_output->channels != 1)
|
|
throw std::runtime_error("incorrect audio format");
|
|
const auto & audio = *result.audio_output;
|
|
if (audio.samples.size() < 2400) throw std::runtime_error("audio is too short");
|
|
double energy = 0;
|
|
for (float value : audio.samples) {
|
|
if (!std::isfinite(value)) throw std::runtime_error("non-finite audio");
|
|
energy += value * value;
|
|
}
|
|
if (energy / audio.samples.size() < 1e-8) throw std::runtime_error("silent audio");
|
|
const auto file = std::filesystem::path(argv[2]) / ("request_" + std::to_string(i) + ".wav");
|
|
engine::audio::WavWriteOptions wav_options;
|
|
wav_options.format = engine::audio::WavSampleFormat::Float32;
|
|
engine::audio::write_wav(file, audio.sample_rate, audio.channels, audio.samples, wav_options);
|
|
if (i == 0) first = audio.samples;
|
|
if (i == 1 && first != audio.samples) throw std::runtime_error("same-session seeded repeat differs");
|
|
if (i == 2 && first == audio.samples) throw std::runtime_error("voice switch did not change output");
|
|
std::cout << "request=" << i << " seconds=" << seconds << " audio_sec=" << audio.samples.size()/24000.0
|
|
<< " rtf=" << seconds/(audio.samples.size()/24000.0) << std::endl;
|
|
}
|
|
auto invalid = request;
|
|
invalid.options["voice_id"] = "NoSuchVoice";
|
|
bool rejected = false;
|
|
try { session->run(invalid); } catch (const std::runtime_error &) { rejected = true; }
|
|
if (!rejected) throw std::runtime_error("unknown voice was not rejected");
|
|
std::cout << "PASS: seeded repeat, voice switch, long-form chunks, and unknown-voice error" << std::endl;
|
|
if (argc > 5 && std::string(argv[5]) != "-") {
|
|
const auto root = std::filesystem::path(argv[5]);
|
|
const auto voices = engine::io::json::parse_file(root/"voices/voices.json");
|
|
auto reference = [&](const char * name) {
|
|
const auto & item = voices.require(name);
|
|
const auto wav = engine::audio::read_wav_f32(root/"voices"/item.require("reference").as_string());
|
|
request.audio_input = engine::runtime::AudioBuffer{wav.sample_rate,wav.channels,wav.samples};
|
|
request.options["reference_text"] = item.require("transcript").as_string();
|
|
};
|
|
request.text_input->text = "Hello there. This is a test of native voice cloning.";
|
|
request.options = {{"seed","1234"}};
|
|
reference("Bruno");
|
|
auto invalid = request;
|
|
invalid.options.erase("reference_text");
|
|
rejected = false;
|
|
try { session->run(invalid); } catch (const std::runtime_error &) { rejected = true; }
|
|
if (!rejected) throw std::runtime_error("missing reference transcript was not rejected");
|
|
invalid = request;
|
|
invalid.audio_input->samples.resize(100);
|
|
rejected = false;
|
|
try { session->run(invalid); } catch (const std::runtime_error &) { rejected = true; }
|
|
if (!rejected) throw std::runtime_error("short reference was not rejected");
|
|
std::vector<float> cloned;
|
|
for (int i = 0; i < 3; ++i) {
|
|
if (i == 2) reference("Bella");
|
|
const auto started = std::chrono::steady_clock::now();
|
|
auto result = session->run(request);
|
|
const double seconds = std::chrono::duration<double>(std::chrono::steady_clock::now()-started).count();
|
|
if (!result.audio_output || result.audio_output->samples.size() < 2400)
|
|
throw std::runtime_error("cloning returned no audio");
|
|
const auto & audio = *result.audio_output;
|
|
double energy = 0;
|
|
for (float value : audio.samples) {
|
|
if (!std::isfinite(value)) throw std::runtime_error("cloning produced non-finite audio");
|
|
energy += value*value;
|
|
}
|
|
if (energy/audio.samples.size() < 1e-8) throw std::runtime_error("cloning produced silence");
|
|
if (i == 0) cloned = audio.samples;
|
|
if (i == 1 && cloned != audio.samples) throw std::runtime_error("cloned voice seeded repeat differs");
|
|
if (i == 2 && cloned == audio.samples) throw std::runtime_error("changed reference reused old voice");
|
|
engine::audio::WavWriteOptions wav_options;
|
|
wav_options.format = engine::audio::WavSampleFormat::Float32;
|
|
engine::audio::write_wav(std::filesystem::path(argv[2])/("clone_"+std::to_string(i)+".wav"),
|
|
audio.sample_rate, audio.channels, audio.samples, wav_options);
|
|
std::cout << "clone=" << i << " seconds=" << seconds << " audio_sec=" << audio.samples.size()/24000.0
|
|
<< " rtf=" << seconds/(audio.samples.size()/24000.0) << std::endl;
|
|
}
|
|
std::cout << "PASS: native cloning, cached repeat, changed reference, missing transcript and short clip errors" << std::endl;
|
|
}
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << error.what() << std::endl;
|
|
return 1;
|
|
}
|
|
}
|