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
127 lines
8.2 KiB
C++
127 lines
8.2 KiB
C++
#include "app/cli/args.h"
|
|
#include "engine/models/pocket_tts_small/assets.h"
|
|
#include "engine/models/pocket_tts_small/mimi_decoder.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/audio/chunking.h"
|
|
#include "engine/framework/debug/profiler.h"
|
|
#include "engine/framework/io/safetensors.h"
|
|
#include "engine/framework/io/json.h"
|
|
#include "engine/models/pocket_tts_small/conditioning.h"
|
|
#include <cstring>
|
|
#include <iomanip>
|
|
#include <sstream>
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <iostream>
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
const auto model = minitts::cli::find_arg(argc, argv, "--model").value();
|
|
const auto reference = minitts::cli::find_arg(argc, argv, "--reference").value();
|
|
const auto backend = minitts::cli::find_arg(argc, argv, "--backend").value_or("cuda");
|
|
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 auto log = minitts::cli::find_arg(argc, argv, "--log-file").value();
|
|
if (!minitts::cli::has_arg(argc, argv, "--log")) throw std::runtime_error("--log is required");
|
|
const auto chunk_frames = minitts::cli::parse_int_arg(argc, argv, "--chunk-frames", 4);
|
|
const bool full_decoder = minitts::cli::has_arg(argc, argv, "--full-decoder");
|
|
minitts::cli::require_known_args(argc, argv);
|
|
engine::debug::configure_logging({true, log});
|
|
const auto index = engine::io::load_safetensors_index(reference);
|
|
const auto source = engine::assets::open_tensor_source(reference);
|
|
const auto latents = source->require_f32("latents", index.tensors.at("latents").shape);
|
|
const auto expected = source->require_f32("waveform");
|
|
auto assets = engine::models::pocket_tts_small::load_assets(model);
|
|
engine::core::ExecutionContext execution({minitts::cli::parse_backend(backend), 0, 8});
|
|
const auto weights = engine::models::pocket_tts_small::load_pocket_tts_small_backend_weights(
|
|
*assets, execution.backend(), execution.backend_type(), engine::assets::TensorStorageType::Native,
|
|
engine::assets::TensorStorageType::Native, 8 << 20, 64 << 20, 16 << 20);
|
|
const auto texts = engine::io::json::parse(index.metadata.at("prepared_text")).as_array();
|
|
engine::models::pocket_tts_small::TextConditioner conditioner({assets->model_config.flow_dim});
|
|
const auto chunks = conditioner.split_into_sentence_chunks(*assets, index.metadata.at("text"),
|
|
std::stoll(index.metadata.at("max_tokens")));
|
|
if (chunks.size() != texts.size()) throw std::runtime_error("Official text chunk count differs");
|
|
for (size_t i = 0; i < texts.size(); ++i) {
|
|
std::ostringstream name;
|
|
name << "text." << std::setw(4) << std::setfill('0') << i << ".token_ids";
|
|
const auto tokens = source->require_tensor_data(name.str());
|
|
const auto prepared = conditioner.prepare(*assets, weights->host, chunks[i], execution.backend(), 8);
|
|
if (prepared.prepared_text != texts[i].as_string())
|
|
throw std::runtime_error("Official prepared text differs");
|
|
if (tokens.metadata.dtype != "I64" || tokens.bytes.size() != prepared.tokens.size() * sizeof(int64_t))
|
|
throw std::runtime_error("Official text token count or type differs");
|
|
for (size_t j = 0; j < prepared.tokens.size(); ++j) {
|
|
int64_t token;
|
|
std::memcpy(&token, tokens.bytes.data() + j * sizeof(token), sizeof(token));
|
|
if (token != prepared.tokens[j]) throw std::runtime_error("Official text tokens differ");
|
|
}
|
|
}
|
|
const auto voice = engine::assets::open_tensor_source(assets->default_voice_path.value());
|
|
for (const auto & [name, info] : index.tensors) {
|
|
(void) info;
|
|
if (name.rfind("voice.", 0) != 0) continue;
|
|
const auto expected_voice = source->require_tensor_data(name);
|
|
const auto actual_voice = voice->require_tensor_data(name.substr(6));
|
|
if (expected_voice.metadata.shape != actual_voice.metadata.shape ||
|
|
expected_voice.metadata.dtype != actual_voice.metadata.dtype || expected_voice.bytes != actual_voice.bytes)
|
|
throw std::runtime_error("Official voice input differs: " + name);
|
|
}
|
|
const auto config = assets->model_config.decoder;
|
|
const auto frames = static_cast<int64_t>(latents.size()) / config.latent_size;
|
|
engine::models::pocket_tts_small::MimiDecoder decoder(config);
|
|
const auto started = std::chrono::steady_clock::now();
|
|
std::vector<float> actual;
|
|
const auto fade = engine::audio::make_linear_fade_window(240, 120);
|
|
if (mode == "offline") {
|
|
if (texts.size() != 1) throw std::runtime_error("Offline replay expects one text chunk");
|
|
actual = decoder.decode(execution.backend(), 8, *assets, *weights, latents,
|
|
frames, 32 << 20, 96 << 20, 64 << 20, chunk_frames, 900, full_decoder);
|
|
for (size_t i = 0; i < std::min<size_t>(120, actual.size()); ++i) actual[i] *= fade[i];
|
|
} else {
|
|
const auto calls = engine::io::json::parse(index.metadata.at("decoder_calls")).as_array();
|
|
for (const auto & call : calls) {
|
|
if (call.require("reset").as_bool()) decoder.reset_streaming_state();
|
|
std::ostringstream name;
|
|
name << "decoder." << std::setw(4) << std::setfill('0') << call.require("index").as_i64() << ".latent";
|
|
const auto input = source->require_f32(name.str());
|
|
std::vector<float> samples;
|
|
for (size_t offset = 0; offset < input.size(); offset += config.latent_size) {
|
|
const std::vector<float> frame(input.begin() + offset, input.begin() + offset + config.latent_size);
|
|
auto output = decoder.decode_streaming(execution.backend(), 8, *assets, *weights,
|
|
frame, 32 << 20, 96 << 20, 64 << 20);
|
|
samples.insert(samples.end(), output.begin(), output.end());
|
|
}
|
|
if (samples.size() != static_cast<size_t>(call.require("samples").as_i64()))
|
|
throw std::runtime_error("Streaming decoder call length differs");
|
|
if (call.require("reset").as_bool())
|
|
for (size_t i = 0; i < std::min<size_t>(120, samples.size()); ++i) samples[i] *= fade[i];
|
|
actual.insert(actual.end(), samples.begin(), samples.end());
|
|
}
|
|
}
|
|
if (actual.size() != expected.size()) throw std::runtime_error("Teacher-forced waveform length differs");
|
|
double aa = 0, bb = 0, ab = 0, max_abs = 0;
|
|
for (size_t i = 0; i < actual.size(); ++i) {
|
|
if (!std::isfinite(actual[i])) throw std::runtime_error("Nonfinite teacher-forced waveform");
|
|
aa += double(actual[i]) * actual[i]; bb += double(expected[i]) * expected[i];
|
|
ab += double(actual[i]) * expected[i];
|
|
max_abs = std::max(max_abs, std::abs(double(actual[i]) - expected[i]));
|
|
}
|
|
double first_aa = 0, first_bb = 0, first_ab = 0;
|
|
for (size_t i = 0; i < std::min<size_t>(1920, actual.size()); ++i) {
|
|
first_aa += double(actual[i]) * actual[i];
|
|
first_bb += double(expected[i]) * expected[i];
|
|
first_ab += double(actual[i]) * expected[i];
|
|
}
|
|
engine::debug::trace_log_scalar("pocket_tts_small.reference.first_frame_waveform_cosine",
|
|
first_ab / std::sqrt(first_aa * first_bb));
|
|
const double cosine = ab / std::sqrt(aa * bb);
|
|
engine::debug::trace_log_scalar("pocket_tts_small.reference.waveform_cosine", cosine);
|
|
engine::debug::trace_log_scalar("pocket_tts_small.reference.waveform_max_abs", max_abs);
|
|
engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(started));
|
|
std::cout << "waveform_cosine=" << cosine << " samples=" << actual.size() << '\n';
|
|
if (!(cosine > .99999)) throw std::runtime_error("Teacher-forced waveform parity failed");
|
|
} catch (const std::exception & error) {
|
|
std::cerr << error.what() << '\n'; return 1;
|
|
}
|
|
}
|