yangrongzhao

Inflect-Nano-v2 for AX650 NPU3

a44ca9d 3 months ago

3.99 kB

#include "inflect_tts.hpp"
#include "ax_engine.h"
#include "ax_sys.h"
#include <cstring>
#include <stdexcept>
#include <cmath>
#include <algorithm>

// ---------------------------------------------------------------------------
// TTSEngine — NPU lifecycle
// ---------------------------------------------------------------------------

TTSEngine::TTSEngine(const std::string& encoder_path,
                     const std::string& decoder_path) {
    auto* eng = ax_engine_init();
    if (!eng)
        throw std::runtime_error("ax_engine_init failed");
    enc_engine_ = ax_engine_load_model(eng, encoder_path.c_str());
    dec_engine_ = ax_engine_load_model(eng, decoder_path.c_str());
    if (!enc_engine_ || !dec_engine_)
        throw std::runtime_error("Failed to load AX models");
}

TTSEngine::~TTSEngine() {
    if (enc_context_) ax_engine_destroy_context(enc_context_);
    if (dec_context_) ax_engine_destroy_context(dec_context_);
}

// ---------------------------------------------------------------------------
// encode — run encoder NPU inference
//   tokens [token_len]  →  4 output tensors {m_p, logs_p, x, x_mask}
// ---------------------------------------------------------------------------
std::vector<float> TTSEngine::encode(const int64_t* tokens, int token_len) {
    // 1) Embedding placeholder — in production, port PyTorch emb via libtorch.
    //    For now, zeros (caller must supply pre-computed x_emb).
    (void)tokens; (void)token_len;

    std::vector<float> empty;
    return empty;  // caller fills via direct x_emb
}

// ---------------------------------------------------------------------------
// decode — run decoder NPU inference
//   z_p [kInterChannels * mel_len]  →  waveform [mel_len * kHopLength]
// ---------------------------------------------------------------------------
std::vector<float> TTSEngine::decode(const float* z_p, int mel_len) {
    auto* engine = static_cast<ax_engine_t*>(dec_engine_);
    ax_engine_io_t io;
    ax_engine_get_io(engine, &io);

    // Zero-pad inputs to fixed shapes
    std::vector<float> z_p_padded(kInterChannels * kMaxMelFrames, 0.0f);
    std::vector<float> y_mask_padded(kMaxMelFrames, 0.0f);
    for (int i = 0; i < std::min(mel_len, kMaxMelFrames); ++i) {
        for (int c = 0; c < kInterChannels; ++c)
            z_p_padded[c * kMaxMelFrames + i] = z_p[c * mel_len + i];
        y_mask_padded[i] = 1.0f;
    }

    std::memcpy(io.inputs[0].data, z_p_padded.data(),
                kInterChannels * kMaxMelFrames * sizeof(float));
    std::memcpy(io.inputs[1].data, y_mask_padded.data(),
                kMaxMelFrames * sizeof(float));

    ax_engine_run(engine, &io);

    int out_len = std::min(mel_len, kMaxMelFrames) * kHopLength;
    std::vector<float> waveform(out_len);
    std::memcpy(waveform.data(), io.outputs[0].data, out_len * sizeof(float));
    return waveform;
}

// ---------------------------------------------------------------------------
// synthesize — full pipeline
// ---------------------------------------------------------------------------
std::vector<float> TTSEngine::synthesize(const std::string& text,
                                          float speed, float variation) {
    // Full pipeline requires:
    //   1. Text → phonemes → token IDs          (CPU, Python frontend port)
    //   2. Token IDs → embedding                (CPU, libtorch)
    //   3. Embedding → NPU encoder → stats      (NPU)
    //   4. Duration prediction + alignment      (CPU, libtorch)
    //   5. z_p → NPU decoder → waveform         (NPU)
    //
    // For a complete C++ pipeline, build with libtorch and the
    // Python frontend logic ported to C++. See model_convert/ for
    // ONNX sources and compile configs.
    //
    // This skeleton demonstrates the NPU API — plug in your text
    // processor and return raw waveform.

    (void)text; (void)speed; (void)variation;
    return {};  // skeleton — fill with your full pipeline
}

	#include "inflect_tts.hpp"
	#include "ax_engine.h"
	#include "ax_sys.h"
	#include <cstring>
	#include <stdexcept>
	#include <cmath>
	#include <algorithm>

	// ---------------------------------------------------------------------------
	// TTSEngine — NPU lifecycle
	// ---------------------------------------------------------------------------

	TTSEngine::TTSEngine(const std::string& encoder_path,
	const std::string& decoder_path) {
	auto* eng = ax_engine_init();
	if (!eng)
	throw std::runtime_error("ax_engine_init failed");
	enc_engine_ = ax_engine_load_model(eng, encoder_path.c_str());
	dec_engine_ = ax_engine_load_model(eng, decoder_path.c_str());
	if (!enc_engine_ \|\| !dec_engine_)
	throw std::runtime_error("Failed to load AX models");
	}

	TTSEngine::~TTSEngine() {
	if (enc_context_) ax_engine_destroy_context(enc_context_);
	if (dec_context_) ax_engine_destroy_context(dec_context_);
	}

	// ---------------------------------------------------------------------------
	// encode — run encoder NPU inference
	// tokens [token_len] → 4 output tensors {m_p, logs_p, x, x_mask}
	// ---------------------------------------------------------------------------
	std::vector<float> TTSEngine::encode(const int64_t* tokens, int token_len) {
	// 1) Embedding placeholder — in production, port PyTorch emb via libtorch.
	// For now, zeros (caller must supply pre-computed x_emb).
	(void)tokens; (void)token_len;

	std::vector<float> empty;
	return empty; // caller fills via direct x_emb
	}

	// ---------------------------------------------------------------------------
	// decode — run decoder NPU inference
	// z_p [kInterChannels * mel_len] → waveform [mel_len * kHopLength]
	// ---------------------------------------------------------------------------
	std::vector<float> TTSEngine::decode(const float* z_p, int mel_len) {
	auto* engine = static_cast<ax_engine_t*>(dec_engine_);
	ax_engine_io_t io;
	ax_engine_get_io(engine, &io);

	// Zero-pad inputs to fixed shapes
	std::vector<float> z_p_padded(kInterChannels * kMaxMelFrames, 0.0f);
	std::vector<float> y_mask_padded(kMaxMelFrames, 0.0f);
	for (int i = 0; i < std::min(mel_len, kMaxMelFrames); ++i) {
	for (int c = 0; c < kInterChannels; ++c)
	z_p_padded[c * kMaxMelFrames + i] = z_p[c * mel_len + i];
	y_mask_padded[i] = 1.0f;
	}

	std::memcpy(io.inputs[0].data, z_p_padded.data(),
	kInterChannels * kMaxMelFrames * sizeof(float));
	std::memcpy(io.inputs[1].data, y_mask_padded.data(),
	kMaxMelFrames * sizeof(float));

	ax_engine_run(engine, &io);

	int out_len = std::min(mel_len, kMaxMelFrames) * kHopLength;
	std::vector<float> waveform(out_len);
	std::memcpy(waveform.data(), io.outputs[0].data, out_len * sizeof(float));
	return waveform;
	}

	// ---------------------------------------------------------------------------
	// synthesize — full pipeline
	// ---------------------------------------------------------------------------
	std::vector<float> TTSEngine::synthesize(const std::string& text,
	float speed, float variation) {
	// Full pipeline requires:
	// 1. Text → phonemes → token IDs (CPU, Python frontend port)
	// 2. Token IDs → embedding (CPU, libtorch)
	// 3. Embedding → NPU encoder → stats (NPU)
	// 4. Duration prediction + alignment (CPU, libtorch)
	// 5. z_p → NPU decoder → waveform (NPU)
	//
	// For a complete C++ pipeline, build with libtorch and the
	// Python frontend logic ported to C++. See model_convert/ for
	// ONNX sources and compile configs.
	//
	// This skeleton demonstrates the NPU API — plug in your text
	// processor and return raw waveform.

	(void)text; (void)speed; (void)variation;
	return {}; // skeleton — fill with your full pipeline
	}