Kazeia-engine/dist/jni/kazeia_mel.cpp

237 lines
9.3 KiB
C++

// Implementation kazeia_mel — voir kazeia_mel.h pour l'API publique.
#include "kazeia_mel.h"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <fstream>
namespace kazeia_mel {
// ============================================================================
// Configurations pré-faites
// ============================================================================
MelConfig config_qwen3_tts_speaker() {
MelConfig c;
c.sample_rate = 24000;
c.n_fft = 1024;
c.hop_size = 256;
c.win_size = 1024;
c.n_mels = 128;
c.pad_amount = (1024 - 256) / 2; // 384, reflect (TTS convention)
c.fixed_n_frames = 0;
c.clip_eps = 1e-5f;
c.window = WindowKind::HANN_SYMMETRIC;
c.comp = Compression::LOG_NATURAL;
c.post = PostNorm::NONE;
return c;
}
MelConfig config_whisper() {
MelConfig c;
c.sample_rate = 16000;
c.n_fft = 400;
c.hop_size = 160;
c.win_size = 400;
c.n_mels = 80;
c.pad_amount = 200; // n_fft/2 (HF whisper)
c.fixed_n_frames = 3000; // 30s padding fixe
c.clip_eps = 1e-10f;
c.window = WindowKind::HANN_PERIODIC;
c.comp = Compression::LOG10;
c.post = PostNorm::WHISPER;
return c;
}
// ============================================================================
// Helpers internes
// ============================================================================
// FFT radix-2 in-place (n doit être puissance de 2). Si n_fft pas puissance de 2
// (cas Whisper N_FFT=400), on zero-pad jusqu'à la puissance de 2 supérieure.
static void fft_radix2(float* real, float* imag, int n) {
int j = 0;
for (int i = 1; i < n; i++) {
int bit = n >> 1;
while (j & bit) { j ^= bit; bit >>= 1; }
j ^= bit;
if (i < j) { std::swap(real[i], real[j]); std::swap(imag[i], imag[j]); }
}
for (int len = 2; len <= n; len <<= 1) {
int half = len / 2;
double angle = -2.0 * M_PI / len;
float wR = (float)std::cos(angle), wI = (float)std::sin(angle);
for (int i = 0; i < n; i += len) {
float cR = 1.f, cI = 0.f;
for (int k = 0; k < half; k++) {
float tR = cR * real[i+k+half] - cI * imag[i+k+half];
float tI = cR * imag[i+k+half] + cI * real[i+k+half];
real[i+k+half] = real[i+k] - tR;
imag[i+k+half] = imag[i+k] - tI;
real[i+k] += tR;
imag[i+k] += tI;
float ncR = cR * wR - cI * wI;
cI = cR * wI + cI * wR;
cR = ncR;
}
}
}
}
static int next_pow2(int n) {
int p = 1; while (p < n) p <<= 1; return p;
}
static std::vector<float> make_window(int n, WindowKind kind) {
std::vector<float> w(n);
if (kind == WindowKind::HANN_SYMMETRIC) {
// torch.hann_window(N) default = symmetric : 0.5 * (1 - cos(2*pi*n/(N-1)))
for (int i = 0; i < n; i++)
w[i] = 0.5f * (1.0f - std::cos(2.0f * (float)M_PI * (float)i / (float)(n - 1)));
} else {
// Periodic (HF whisper, numpy default) : 0.5 * (1 - cos(2*pi*n/N))
for (int i = 0; i < n; i++)
w[i] = 0.5f * (1.0f - std::cos(2.0f * (float)M_PI * (float)i / (float)n));
}
return w;
}
// Pad reflect : miroir interne (n'inclut PAS le sample bord). Équivalent torch reflect.
static std::vector<float> pad_reflect(const std::vector<float>& x, int padding) {
const int n = (int)x.size();
std::vector<float> out((size_t)n + 2 * padding);
for (int i = 0; i < padding; i++) out[i] = x[padding - i];
for (int i = 0; i < n; i++) out[padding + i] = x[i];
for (int i = 0; i < padding; i++) out[padding + n + i] = x[n - 2 - i];
return out;
}
// ============================================================================
// load_mel_basis : fichier binaire f32 [n_mels, n_fft/2+1]
// ============================================================================
bool load_mel_basis(const char * path, const MelConfig & cfg, std::vector<float> & out) {
const int fft_bins = cfg.n_fft / 2 + 1;
const size_t expected = (size_t)cfg.n_mels * fft_bins;
std::ifstream f(path, std::ios::binary | std::ios::ate);
if (!f) { fprintf(stderr, "kazeia_mel::load_mel_basis: open %s FAIL\n", path); return false; }
size_t n = (size_t)f.tellg() / sizeof(float);
if (n != expected) {
fprintf(stderr, "kazeia_mel::load_mel_basis: %s contient %zu f32, attendu %zu (=%d mels x %d bins)\n",
path, n, expected, cfg.n_mels, fft_bins);
return false;
}
f.seekg(0);
out.assign(expected, 0.0f);
f.read((char*)out.data(), expected * sizeof(float));
return (bool)f;
}
// ============================================================================
// compute : STFT (radix-2 zero-pad) + projection mel + log + post-norm
// ============================================================================
std::vector<float> compute(const std::vector<float> & wav_in,
const std::vector<float> & mel_basis,
const MelConfig & cfg,
int & out_T) {
const int fft_bins = cfg.n_fft / 2 + 1;
const int pad_amount = (cfg.pad_amount >= 0) ? cfg.pad_amount : (cfg.n_fft - cfg.hop_size) / 2;
// Optionnel : padding audio à T_fixe (Whisper 30s)
std::vector<float> wav = wav_in;
if (cfg.fixed_n_frames > 0) {
int target_samples = cfg.fixed_n_frames * cfg.hop_size;
if ((int)wav.size() > target_samples) wav.resize(target_samples);
else if ((int)wav.size() < target_samples) wav.resize(target_samples, 0.0f);
}
auto pad = pad_reflect(wav, pad_amount);
auto hann = make_window(cfg.win_size, cfg.window);
const int len_pad = (int)pad.size();
if (len_pad < cfg.n_fft) {
fprintf(stderr, "kazeia_mel::compute: wav trop court (%d < %d)\n", len_pad, cfg.n_fft);
out_T = 0; return {};
}
// T_compute : sans fixed_n_frames -> calcule autant que possible (center=False).
// Avec fixed_n_frames -> exactement N_FRAMES (le padding ci-dessus garantit la place).
int T;
if (cfg.fixed_n_frames > 0) T = cfg.fixed_n_frames;
else T = (len_pad - cfg.n_fft) / cfg.hop_size + 1;
out_T = T;
// FFT en puissance de 2 >= n_fft (Whisper N_FFT=400 -> FFT 512).
const int fftN = next_pow2(cfg.n_fft);
std::vector<float> spec((size_t)fft_bins * T, 0.0f);
std::vector<float> re(fftN), im(fftN);
for (int t = 0; t < T; t++) {
const int off = t * cfg.hop_size;
// Window sur win_size samples, zero-pad jusqu'à fftN
for (int i = 0; i < cfg.win_size; i++) {
re[i] = pad[off + i] * hann[i];
im[i] = 0.0f;
}
for (int i = cfg.win_size; i < fftN; i++) { re[i] = 0.0f; im[i] = 0.0f; }
fft_radix2(re.data(), im.data(), fftN);
// Power spectrum [0..n_fft/2]
for (int k = 0; k < fft_bins; k++) {
spec[(size_t)k * T + t] = re[k]*re[k] + im[k]*im[k];
}
}
// mel = basis [n_mels, fft_bins] @ spec [fft_bins, T] -> [n_mels, T]
std::vector<float> mel((size_t)cfg.n_mels * T, 0.0f);
for (int m = 0; m < cfg.n_mels; m++) {
for (int t = 0; t < T; t++) {
float s = 0.0f;
for (int k = 0; k < fft_bins; k++) {
s += mel_basis[m * fft_bins + k] * spec[(size_t)k * T + t];
}
mel[(size_t)m * T + t] = std::max(s, cfg.clip_eps);
}
}
// Compression (log naturel ou log10) + sqrt si TTS (variante: TTS prend sqrt magnitude
// AVANT projection mel ; Whisper prend power direct projetée puis log10).
// Distinction empirique du speaker encoder Qwen3-TTS : il calcule magnitude = sqrt(re²+im²+1e-9)
// puis projette puis log. Pour rester strictement bit-correct avec le code spk d'origine,
// on adapte si compression == LOG_NATURAL : sqrt avant log.
if (cfg.comp == Compression::LOG_NATURAL) {
// TTS : sqrt(power+1e-9) puis mel_proj puis log(max(., eps))
// Refait correctement : revert la projection power, refais en magnitude.
// Implémenter ainsi pour bit-match :
std::vector<float> mag_spec((size_t)fft_bins * T);
for (int k = 0; k < fft_bins; k++)
for (int t = 0; t < T; t++)
mag_spec[(size_t)k * T + t] = std::sqrt(spec[(size_t)k * T + t] + 1e-9f);
for (int m = 0; m < cfg.n_mels; m++) {
for (int t = 0; t < T; t++) {
float s = 0.0f;
for (int k = 0; k < fft_bins; k++)
s += mel_basis[m * fft_bins + k] * mag_spec[(size_t)k * T + t];
mel[(size_t)m * T + t] = std::log(std::max(s, cfg.clip_eps));
}
}
} else { // LOG10 (Whisper)
for (auto & v : mel) v = std::log10(v);
}
// Post-norm (Whisper : clamp à max-8, puis (x+4)/4)
if (cfg.post == PostNorm::WHISPER) {
float maxv = *std::max_element(mel.begin(), mel.end());
float floor_v = maxv - 8.0f;
for (auto & v : mel) {
if (v < floor_v) v = floor_v;
v = (v + 4.0f) * 0.25f;
}
}
return mel;
}
} // namespace kazeia_mel