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
99 lines
6.3 KiB
C++
99 lines
6.3 KiB
C++
#include "engine/framework/codecs/s3gen_runtime.h"
|
|
#include "engine/community_models/kitten_tts2/speaker.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/audio/conversion.h"
|
|
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/audio/wav_writer.h"
|
|
#include "engine/framework/io/json.h"
|
|
#include "engine/framework/modules/vocoders/hift_vocoder.h"
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
if (argc < 5 || argc > 7) throw std::runtime_error("usage: kitten_tts2_components_probe S3GEN VOICES_JSON CODES_JSON OUTPUT_DIR [cpu|cuda [REFERENCE_MODEL_DIR]]");
|
|
namespace s3 = engine::codecs::s3gen;
|
|
namespace json = engine::io::json;
|
|
engine::core::BackendConfig config;
|
|
config.type = engine::core::BackendType::Cpu;
|
|
if (argc >= 6) {
|
|
const std::string backend = argv[5];
|
|
if (backend == "cuda") config.type = engine::core::BackendType::Cuda;
|
|
else if (backend != "cpu") throw std::runtime_error("probe backend must be cpu or cuda");
|
|
}
|
|
config.threads = 8;
|
|
engine::core::ExecutionContext execution(config);
|
|
std::filesystem::create_directories(argv[4]);
|
|
if (argc == 7) {
|
|
// Test-only speaker comparison through its public component API.
|
|
const auto root = std::filesystem::path(argv[6]);
|
|
auto speaker_source = engine::assets::open_tensor_source(root / "speaker/model.safetensors");
|
|
auto lm_source = engine::assets::open_tensor_source(root / "lm/model.safetensors");
|
|
engine::community_models::kitten_tts2::SpeakerEncoder speaker(*speaker_source, *lm_source, execution);
|
|
const auto voices = json::parse_file(root / "voices/voices.json");
|
|
const auto wav = engine::audio::read_wav_f32(root / "voices" / voices.require("Bruno").require("reference").as_string());
|
|
const auto mono = engine::audio::mixdown_interleaved_to_mono_average(wav.samples, wav.channels);
|
|
const auto input = engine::audio::resample_mono_torchaudio_sinc_hann(mono, wav.sample_rate, 16000,
|
|
engine::audio::torchaudio_sinc_hann_float32_options());
|
|
const auto embedding = speaker.embed(input);
|
|
const auto projected = speaker.project(embedding);
|
|
const auto array = [](const std::vector<float> & values) {
|
|
json::Value::Array result;
|
|
for (float value : values) result.push_back(json::Value::make_number(value));
|
|
return json::Value::make_array(std::move(result));
|
|
};
|
|
json::Value::Object result;
|
|
result.emplace("embedding", array(embedding));
|
|
result.emplace("projection", array(projected));
|
|
std::ofstream(std::filesystem::path(argv[4]) / "speaker.json") << json::stringify(json::Value::make_object(std::move(result)));
|
|
engine::audio::WavWriteOptions options;
|
|
options.format = engine::audio::WavSampleFormat::Float32;
|
|
engine::audio::write_wav(std::filesystem::path(argv[4]) / "speaker_input.wav", 16000, 1, input, options);
|
|
}
|
|
auto source = engine::assets::open_tensor_source(argv[1]);
|
|
s3::S3GenConfig s3gen_config;
|
|
s3gen_config.decoder.variant = s3::S3FlowVariant::MeanFlow;
|
|
s3::S3GenRuntime runtime(source, execution, s3gen_config, {engine::assets::TensorStorageType::F32});
|
|
const auto all = json::parse_file(argv[2]);
|
|
const auto & voice = all.require("Bruno");
|
|
s3::S3GenConditioning ref;
|
|
for (const auto & x : voice.require("prompt_token").as_array()[0].as_array()) ref.prompt_tokens.push_back(x.as_number());
|
|
for (const auto & row : voice.require("prompt_feat").as_array()[0].as_array())
|
|
for (const auto & x : row.as_array()) ref.prompt_feat.push_back(x.as_number());
|
|
for (const auto & x : voice.require("embedding").as_array()[0].as_array()) ref.embedding.push_back(x.as_number());
|
|
ref.prompt_token_count = ref.prompt_tokens.size();
|
|
ref.prompt_feat_frames = ref.prompt_feat.size()/80;
|
|
ref.prompt_feat_dims = 80;
|
|
ref.embedding_size = 192;
|
|
std::vector<int32_t> tokens;
|
|
const auto codes = json::parse_file(argv[3]);
|
|
for (const auto & x : codes.as_array()) tokens.push_back(x.as_number());
|
|
tokens.insert(tokens.end(), 3, 4299);
|
|
const int64_t frames = 2 * (ref.prompt_token_count + tokens.size());
|
|
// Isolate the encoder through the runtime's stage API.
|
|
auto table = source->require_f32("flow.input_embedding.weight", {6561,512});
|
|
auto all_tokens = ref.prompt_tokens;
|
|
all_tokens.insert(all_tokens.end(), tokens.begin(), tokens.end());
|
|
std::vector<float> embedded;
|
|
for (auto token : all_tokens) for (int i = 0; i < 512; ++i)
|
|
embedded.push_back(table[token*512+i]);
|
|
auto encoded = runtime.encode(embedded, all_tokens.size(), all_tokens.size(), 512);
|
|
std::ofstream hidden(std::filesystem::path(argv[4])/"encoder.f32", std::ios::binary);
|
|
hidden.write(reinterpret_cast<const char *>(encoded.hidden.data()), encoded.hidden.size()*sizeof(float));
|
|
auto mel = runtime.token_to_mel(ref, tokens, tokens.size(),
|
|
2, 0, false, std::vector<float>(80*frames, 0), 0);
|
|
std::filesystem::create_directories(argv[4]);
|
|
std::ofstream out(std::filesystem::path(argv[4])/"zero_noise_mel.f32", std::ios::binary);
|
|
out.write(reinterpret_cast<const char *>(mel.mel.data()), mel.mel.size()*sizeof(float));
|
|
const auto & vocoder = runtime.vocoder();
|
|
auto f0 = vocoder.predict_f0(mel.mel, mel.frames);
|
|
std::vector<float> random(9 + 9*480*mel.frames, 0);
|
|
auto audio = vocoder.synthesize(mel.mel, mel.frames, 0, 0, &random);
|
|
std::ofstream pitch(std::filesystem::path(argv[4])/"f0.f32", std::ios::binary);
|
|
pitch.write(reinterpret_cast<const char *>(f0.data()), f0.size()*sizeof(float));
|
|
std::ofstream wave(std::filesystem::path(argv[4])/"vocoder.f32", std::ios::binary);
|
|
wave.write(reinterpret_cast<const char *>(audio.waveform.data()), audio.waveform.size()*sizeof(float));
|
|
std::cout << "mel_frames=" << mel.frames << std::endl;
|
|
} catch (const std::exception & e) { std::cerr << e.what() << std::endl; return 1; }
|
|
}
|