voice_audio_multi/tests/sopro_tts/sopro_probe.cpp
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
Initial import: audio.cpp with Orange Pi documentation
2026-11-04 19:05:22 +01:00

383 lines
17 KiB
C++

// Stage probe for the sopro_tts port.
//
// The offline pipeline has five stages and a broken one is hard to localise
// from the final waveform. This binary exercises them in isolation:
//
// mel analysis mel -> vocoder -> waveform. Vocos is trained to invert
// its own analysis mel, so a clean round trip proves the mel front
// end and the whole vocoder graph at once.
// semantic reference waveform -> FSQ token ids (histogram + first ids).
// speaker reference waveform -> id/style/style-ctrl embedding statistics.
//
// Usage: sopro_probe <model-dir> <reference.wav> [out-dir]
#include "engine/community_models/sopro_tts/acoustic.h"
#include "engine/community_models/sopro_tts/assets.h"
#include "engine/community_models/sopro_tts/reference.h"
#include "engine/community_models/sopro_tts/semantic_encoder.h"
#include "engine/community_models/sopro_tts/speaker_encoder.h"
#include "engine/community_models/sopro_tts/vocoder.h"
#include "engine/framework/assets/tensor_source.h"
#include "engine/framework/audio/conversion.h"
#include "engine/framework/audio/resampling.h"
#include "engine/framework/audio/wav_reader.h"
#include "engine/framework/audio/wav_writer.h"
#include "engine/framework/core/backend.h"
#include "engine/framework/core/execution_context.h"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <filesystem>
#include <map>
#include <random>
#include <string>
#include <vector>
namespace sopro = engine::community_models::sopro_tts;
namespace {
struct Stats {
float min = 0.0F;
float max = 0.0F;
double mean = 0.0;
double rms = 0.0;
size_t nonfinite = 0;
};
Stats describe(const std::vector<float> & values) {
Stats out;
if (values.empty()) {
return out;
}
out.min = out.max = values.front();
double sum = 0.0;
double square_sum = 0.0;
for (const float value : values) {
if (!std::isfinite(value)) {
++out.nonfinite;
continue;
}
out.min = std::min(out.min, value);
out.max = std::max(out.max, value);
sum += value;
square_sum += static_cast<double>(value) * value;
}
const auto count = static_cast<double>(values.size());
out.mean = sum / count;
out.rms = std::sqrt(square_sum / count);
return out;
}
void print_stats(const char * label, const std::vector<float> & values) {
const auto stats = describe(values);
std::printf(
" %-22s n=%-8zu min=%+.4f max=%+.4f mean=%+.5f rms=%.5f nonfinite=%zu\n",
label, values.size(), stats.min, stats.max, stats.mean, stats.rms, stats.nonfinite);
}
// Segmental SNR between two aligned signals, in dB.
double snr_db(const std::vector<float> & reference, const std::vector<float> & test) {
const size_t n = std::min(reference.size(), test.size());
double signal = 0.0;
double noise = 0.0;
for (size_t i = 0; i < n; ++i) {
const double r = reference[i];
const double d = r - test[i];
signal += r * r;
noise += d * d;
}
if (noise <= 0.0) {
return 999.0;
}
return 10.0 * std::log10(std::max(signal, 1e-20) / noise);
}
} // namespace
int main(int argc, char ** argv) {
if (argc < 3) {
std::fprintf(stderr, "usage: %s <model-dir> <reference.wav> [out-dir]\n", argv[0]);
return 2;
}
const std::filesystem::path model_path = argv[1];
const std::filesystem::path reference_path = argv[2];
const std::filesystem::path out_dir = argc > 3 ? std::filesystem::path(argv[3])
: std::filesystem::path(".");
std::filesystem::create_directories(out_dir);
auto assets = sopro::load_sopro_tts_assets(model_path);
const auto & config = assets->config;
const auto sample_rate = static_cast<int>(config.sample_rate);
engine::core::BackendConfig backend_config;
backend_config.type = engine::core::BackendType::Cpu;
backend_config.threads = 8;
engine::core::ExecutionContext execution(backend_config);
constexpr size_t kWeightBytes = 512ull * 1024ull * 1024ull;
constexpr size_t kGraphBytes = 1024ull * 1024ull * 1024ull;
const auto storage = engine::assets::TensorStorageType::F32;
// ---- input ----
const auto wav = engine::audio::read_wav_f32(reference_path);
std::vector<float> mono = wav.channels == 1
? wav.samples
: engine::audio::mixdown_interleaved_to_mono_average(wav.samples, wav.channels);
if (wav.sample_rate != sample_rate) {
mono = engine::audio::resample_mono_torchaudio_sinc_hann(mono, wav.sample_rate, sample_rate);
}
for (auto & value : mono) {
value = std::min(1.0F, std::max(-1.0F, value));
}
std::printf("reference: %s\n", reference_path.string().c_str());
std::printf(" source_rate=%d channels=%d samples=%zu -> %.3f s at %d Hz\n",
wav.sample_rate, wav.channels, mono.size(),
static_cast<double>(mono.size()) / sample_rate, sample_rate);
// ---- crop / level normalisation, as prepare_reference does ----
std::mt19937_64 rng(1234);
auto cropped = sopro::audio_ops::crop_on_pause(
mono, config.generation.ref_seconds, sample_rate, rng);
const auto level = sopro::audio_ops::speech_level_db(mono, sample_rate);
std::printf(" speech_level=%.2f dB active=%.3f s\n", level.level_db, level.active_seconds);
std::printf(" crop_on_pause(%.1f s): %zu -> %zu samples (%.3f s)\n",
config.generation.ref_seconds, mono.size(), cropped.size(),
static_cast<double>(cropped.size()) / sample_rate);
const auto crop_level = sopro::audio_ops::speech_level_db(cropped, sample_rate);
float crop_peak = 0.0F;
for (const float value : cropped) {
crop_peak = std::max(crop_peak, std::fabs(value));
}
auto normalisation = sopro::audio_ops::normalize_reference(cropped, sample_rate);
const auto normalised = std::move(normalisation.wav);
std::printf(" normalize_reference: level %.2f -> %.2f dB (gain %+.2f dB, crop peak %.3f)\n",
crop_level.level_db, normalisation.level_db,
normalisation.level_db - crop_level.level_db, crop_peak);
// ---- stage: vocoder round trip ----
{
std::printf("\n[mel] analysis mel -> vocoder round trip\n");
sopro::SoproVocosRuntime vocoder(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
const auto mel = vocoder.log_mel(normalised);
const int64_t n_mels = vocoder.n_mels();
const int64_t frames = static_cast<int64_t>(mel.size()) / n_mels;
std::printf(" frames=%lld (expected %lld)\n",
static_cast<long long>(frames),
static_cast<long long>(vocoder.mel_frames(static_cast<int64_t>(normalised.size()))));
print_stats("log_mel", mel);
// Dump for an independent numpy check of the front end.
{
std::FILE * fh = std::fopen((out_dir / "probe_input_logmel.f32").string().c_str(), "wb");
if (fh != nullptr) {
std::fwrite(mel.data(), sizeof(float), mel.size(), fh);
std::fclose(fh);
}
std::FILE * wh = std::fopen((out_dir / "probe_input_audio.f32").string().c_str(), "wb");
if (wh != nullptr) {
std::fwrite(normalised.data(), sizeof(float), normalised.size(), wh);
std::fclose(wh);
}
}
const auto audio = vocoder.decode(mel, frames);
print_stats("round_trip_audio", audio);
engine::audio::write_pcm16_wav(out_dir / "probe_vocoder_roundtrip.wav", sample_rate, 1, audio);
{
std::FILE * fh = std::fopen((out_dir / "probe_roundtrip_audio.f32").string().c_str(), "wb");
if (fh != nullptr) {
std::fwrite(audio.data(), sizeof(float), audio.size(), fh);
std::fclose(fh);
}
}
// Waveform SNR is phase-sensitive and Vocos predicts phase, so the
// meaningful check is whether re-analysing the output reproduces the
// mel the vocoder was asked to render.
std::printf(" waveform SNR vs input: %.2f dB (phase-sensitive, informative only)\n",
snr_db(normalised, audio));
const auto remel = vocoder.log_mel(audio);
const int64_t reframes = static_cast<int64_t>(remel.size()) / n_mels;
const int64_t common = std::min(frames, reframes);
double abs_sum = 0.0;
double abs_max = 0.0;
size_t counted = 0;
for (int64_t c = 0; c < n_mels; ++c) {
for (int64_t t = 0; t < common; ++t) {
const double a = mel[static_cast<size_t>(c * frames + t)];
const double b = remel[static_cast<size_t>(c * reframes + t)];
// Ignore bins pinned at the log floor; they carry no signal.
if (a <= -16.0) {
continue;
}
const double d = std::fabs(a - b);
abs_sum += d;
abs_max = std::max(abs_max, d);
++counted;
}
}
const double mae = counted > 0 ? abs_sum / static_cast<double>(counted) : 0.0;
std::printf(" mel round-trip MAE=%.4f max=%.4f over %zu bins (log units)\n",
mae, abs_max, counted);
std::printf(" -> %s\n",
mae < 0.35 ? "vocoder + mel front end look CORRECT"
: "vocoder or mel front end is WRONG");
}
// ---- stage: semantic encoder ----
{
std::printf("\n[semantic] FSQ tokeniser\n");
sopro::SoproSemanticEncoderRuntime encoder(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
const auto tokens = encoder.encode(normalised);
const int64_t expected = (static_cast<int64_t>(normalised.size()) +
config.semantic_encoder.token_samples_24k - 1) /
config.semantic_encoder.token_samples_24k;
std::printf(" tokens=%zu (expected %lld)\n", tokens.size(),
static_cast<long long>(expected));
std::map<int32_t, int> histogram;
for (const int32_t token : tokens) {
++histogram[token];
}
int32_t lo = tokens.empty() ? 0 : *std::min_element(tokens.begin(), tokens.end());
int32_t hi = tokens.empty() ? 0 : *std::max_element(tokens.begin(), tokens.end());
std::printf(" distinct=%zu range=[%d, %d] of [0, %lld)\n",
histogram.size(), lo, hi,
static_cast<long long>(config.model.semantic_vocab_size));
std::printf(" first ids:");
for (size_t i = 0; i < std::min<size_t>(16, tokens.size()); ++i) {
std::printf(" %d", tokens[i]);
}
std::printf("\n");
// A collapsed tokeniser (one id repeated) means the encoder is broken.
int most = 0;
for (const auto & [id, count] : histogram) {
(void) id;
most = std::max(most, count);
}
const double share = tokens.empty() ? 0.0 : static_cast<double>(most) / tokens.size();
std::printf(" most common id share: %.1f%% %s\n", share * 100.0,
share > 0.5 ? "(COLLAPSED - encoder likely wrong)" : "(looks healthy)");
}
// ---- stage: speaker encoder ----
{
std::printf("\n[speaker] identity / style embeddings\n");
sopro::SoproSpeakerEncoderRuntime encoder(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
const auto wav16 = engine::audio::resample_mono_torchaudio_sinc_hann(
normalised, sample_rate, static_cast<int>(encoder.sample_rate()));
const auto embeddings = encoder.encode(wav16);
print_stats("id_emb", embeddings.id_emb);
print_stats("style_emb", embeddings.style_emb);
print_stats("style_ctrl", embeddings.style_ctrl);
double norm = 0.0;
for (const float value : embeddings.id_emb) {
norm += static_cast<double>(value) * value;
}
std::printf(" |id_emb| = %.6f (must be 1.0)\n", std::sqrt(norm));
}
// ---- stage: acoustic head self-reconstruction ----
// Re-render the second half of the reference from its own semantic tokens,
// conditioned on the first half as the prompt. The acoustic head is lossy
// (it only sees FSQ ids) but a correct one tracks the real mel closely;
// a broken one produces something uncorrelated with it.
{
std::printf("\n[acoustic] self-reconstruction from the reference's own tokens\n");
sopro::SoproVocosRuntime vocoder(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
sopro::SoproSpeakerEncoderRuntime speaker(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
sopro::SoproSemanticEncoderRuntime semantic(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
sopro::SoproReferenceBuilder builder(*assets, speaker, semantic, vocoder);
std::mt19937_64 build_rng(1234);
const auto voice = builder.build(mono, config.generation.ref_seconds, build_rng);
const int64_t n_mels = config.model.acoustic_mel_n_mels;
const int64_t hop_ratio = config.hop_ratio();
const auto total_tokens = static_cast<int64_t>(voice.semantic_tokens.size());
const int64_t split = total_tokens / 2;
const int64_t prompt_frames = std::min(voice.mel_frames, split * hop_ratio);
const int64_t gen_tokens = total_tokens - split;
const int64_t total_frames = prompt_frames + gen_tokens * hop_ratio;
std::printf(" reference: %lld tokens, %lld mel frames\n",
static_cast<long long>(total_tokens), static_cast<long long>(voice.mel_frames));
std::printf(" prompt: %lld frames, regenerating %lld tokens -> %lld frames\n",
static_cast<long long>(prompt_frames), static_cast<long long>(gen_tokens),
static_cast<long long>(total_frames));
sopro::SoproAcousticRequest request;
request.semantic_tokens = voice.semantic_tokens;
request.cond_vec = voice.cond_vec;
request.prompt_mel.assign(static_cast<size_t>(n_mels * prompt_frames), 0.0F);
for (int64_t c = 0; c < n_mels; ++c) {
for (int64_t t = 0; t < prompt_frames; ++t) {
request.prompt_mel[static_cast<size_t>(c * prompt_frames + t)] =
voice.mel[static_cast<size_t>(c * voice.mel_frames + t)];
}
}
request.prompt_frames = prompt_frames;
request.total_frames = total_frames;
request.steps = 32; // many steps: isolate the field from the schedule
request.seed = 1234;
sopro::SoproAcousticRuntime acoustic(
*assets, execution, kWeightBytes, kGraphBytes, storage, storage);
const auto solved = acoustic.solve(request);
print_stats("solved_mel(normalised)", solved);
// Compare only the regenerated span against the true reference mel.
const int64_t compare_end = std::min(total_frames, voice.mel_frames);
double abs_sum = 0.0;
double ref_sq = 0.0;
double err_sq = 0.0;
size_t counted = 0;
for (int64_t c = 0; c < n_mels; ++c) {
for (int64_t t = prompt_frames; t < compare_end; ++t) {
const double truth = voice.mel[static_cast<size_t>(c * voice.mel_frames + t)];
const double got = solved[static_cast<size_t>(c * total_frames + t)];
abs_sum += std::fabs(truth - got);
ref_sq += truth * truth;
err_sq += (truth - got) * (truth - got);
++counted;
}
}
if (counted > 0) {
const double mae = abs_sum / static_cast<double>(counted);
const double nmse = err_sq / std::max(ref_sq, 1e-12);
std::printf(" regenerated span vs true mel: MAE=%.4f NMSE=%.4f (%zu bins)\n",
mae, nmse, counted);
std::printf(" -> %s\n",
nmse < 0.6 ? "acoustic head tracks the reference (looks CORRECT)"
: "acoustic head output is uncorrelated with the reference (WRONG)");
}
std::FILE * fh = std::fopen((out_dir / "probe_acoustic_solved.f32").string().c_str(), "wb");
if (fh != nullptr) {
std::fwrite(solved.data(), sizeof(float), solved.size(), fh);
std::fclose(fh);
}
std::FILE * th = std::fopen((out_dir / "probe_ref_tokens.i32").string().c_str(), "wb");
if (th != nullptr) {
std::fwrite(voice.semantic_tokens.data(), sizeof(int32_t), voice.semantic_tokens.size(), th);
std::fclose(th);
}
std::FILE * ch = std::fopen((out_dir / "probe_cond_vec.f32").string().c_str(), "wb");
if (ch != nullptr) {
std::fwrite(voice.cond_vec.data(), sizeof(float), voice.cond_vec.size(), ch);
std::fclose(ch);
}
std::FILE * mh = std::fopen((out_dir / "probe_ref_mel.f32").string().c_str(), "wb");
if (mh != nullptr) {
std::fwrite(voice.mel.data(), sizeof(float), voice.mel.size(), mh);
std::fclose(mh);
}
}
std::printf("\nwrote %s\n", (out_dir / "probe_vocoder_roundtrip.wav").string().c_str());
return 0;
}