HY-2012

Upload ASR and TTS models

a37072b verified 3 months ago

5.34 kB

/**************************************************************************************************
 * ZipVoice AXERA C++ Port
 *
 * ZipVoiceEngine: the core TTS inference engine.
 *
 * Manages encoder + split-decoder4 axmodel inference, duration expansion,
 * and flow-matching sampling. Mirrors the Python Decoder4ZipVoiceBoardRuntime.
 *
 * Model pipeline:
 *   1. encoder.axmodel:  cat_tokens[n] → encoded[n, feat_dim]
 *   2. Duration expand:  repeat tokens → text_condition[1, seq_len, feat_dim]
 *   3. decoder_part0..3: flow-matching ODE solver (Euler, num_step steps)
 *      part0: t, x, text_cond, speech_cond, guidance_scale, padding_mask
 *             → decoder_hidden_p0, time_hidden, padding_mask2, cfg_scale
 *      part1: decoder_hidden_p0, time_hidden, padding_mask2 → decoder_hidden_p1
 *      part2: decoder_hidden_p1, time_hidden, padding_mask2 → decoder_hidden_p2
 *      part3: decoder_hidden_p2, time_hidden, padding_mask2, cfg_scale → v
 *      x = x + v * dt
 **************************************************************************************************/

#pragma once

#include <string>
#include <vector>
#include <map>
#include <memory>
#include <cstdint>

#include "EngineWrapper.hpp"

class ZipVoiceEngine {
public:
    struct Config {
        std::string model_dir;
        int max_feat_len = 1024;
        int max_tokens = 384;
        int feat_dim = 100;
        int sampling_rate = 24000;
        int hop_length = 256;
        int num_step = 10;
        float t_shift = 0.5f;
    };

    struct ModelInfo {
        std::string name;
        std::string file;
        std::vector<std::string> inputs;
        std::vector<std::string> outputs;
    };

    struct Timing {
        float encoder_time_sec = 0.0f;
        float duration_expand_time_sec = 0.0f;
        float decoder_time_sec = 0.0f;
        float total_time_sec = 0.0f;
        int generated_frames = 0;
        int features_len = 0;
    };

    ZipVoiceEngine();
    ~ZipVoiceEngine();

    /**
     * Initialize from a model directory containing:
     *   - runtime_config.json
     *   - decoder4_split_manifest.json
     *   - encoder.axmodel
     *   - decoder_part0.axmodel .. decoder_part3.axmodel
     */
    int Init(const std::string& model_dir);

    /**
     * Run encoder: cat_tokens[1, max_tokens] → encoded[1, max_tokens, feat_dim].
     * out_encoded should be pre-allocated with size: 1 * max_tokens * feat_dim floats.
     */
    int RunEncoder(const std::vector<int32_t>& cat_tokens, std::vector<float>& out_encoded);

    /**
     * Duration expand: token-level features → frame-level features.
     */
    int DurationExpand(const std::vector<float>& encoded,
                       int prompt_tokens_len,
                       int text_tokens_len,
                       int prompt_features_len,
                       float speed,
                       std::vector<float>& out_text_condition,
                       int& out_features_len);

    /**
     * Full sampling: encoder → duration_expand → decoder flow-matching.
     *
     * @param cat_tokens          [max_tokens] int32 token IDs
     * @param prompt_tokens_len   number of prompt tokens
     * @param text_tokens_len     number of text tokens
     * @param prompt_features     [prompt_frames, feat_dim] mel features from prompt wav
     * @param prompt_features_len number of prompt feature frames
     * @param speed               speed factor (1.0 = normal)
     * @param guidance_scale      classifier-free guidance scale
     * @param seed                random seed
     * @param out_features        output mel features [generated_frames, feat_dim]
     * @param out_timing          timing breakdown
     */
    int Sample(const std::vector<int32_t>& cat_tokens,
               int prompt_tokens_len,
               int text_tokens_len,
               const std::vector<float>& prompt_features,
               int prompt_features_len,
               float speed,
               float guidance_scale,
               int seed,
               std::vector<float>& out_features,
               Timing& out_timing);

    const Config& GetConfig() const { return m_config; }
    bool HasInit() const { return m_has_init; }

    /** Override num_step for subsequent Sample() calls. -1 to reset to config default. */
    void SetNumStep(int n) { m_num_step_override = n; }

private:
    Config m_config;
    bool m_has_init;

    ModelInfo m_encoder_info;
    std::vector<ModelInfo> m_decoder_parts;

    // Model sessions (one per axmodel)
    std::map<std::string, std::unique_ptr<EngineWrapper>> m_sessions;

    // Decoder metadata
    int m_decoder_seq_len;
    bool m_decoder_has_padding_mask;
    int m_num_step_override = -1;  // -1 = use config default

    int LoadConfig(const std::string& model_dir);
    int LoadManifest(const std::string& model_dir);
    int LoadModels(const std::string& model_dir);
    int LoadDecoderMetadata();

    std::vector<float> GetTimesteps(int num_step, float t_shift) const;

    int RunDecoderPart(const ModelInfo& part,
                       std::map<std::string, std::vector<float>>& values,
                       const std::vector<uint8_t>* padding_mask_data = nullptr,
                       std::vector<uint8_t>* padding_mask2_out = nullptr);

    static double GetCurrentTimeMs();
};

	/**************************************************************************************************
	* ZipVoice AXERA C++ Port
	*
	* ZipVoiceEngine: the core TTS inference engine.
	*
	* Manages encoder + split-decoder4 axmodel inference, duration expansion,
	* and flow-matching sampling. Mirrors the Python Decoder4ZipVoiceBoardRuntime.
	*
	* Model pipeline:
	* 1. encoder.axmodel: cat_tokens[n] → encoded[n, feat_dim]
	* 2. Duration expand: repeat tokens → text_condition[1, seq_len, feat_dim]
	* 3. decoder_part0..3: flow-matching ODE solver (Euler, num_step steps)
	* part0: t, x, text_cond, speech_cond, guidance_scale, padding_mask
	* → decoder_hidden_p0, time_hidden, padding_mask2, cfg_scale
	* part1: decoder_hidden_p0, time_hidden, padding_mask2 → decoder_hidden_p1
	* part2: decoder_hidden_p1, time_hidden, padding_mask2 → decoder_hidden_p2
	* part3: decoder_hidden_p2, time_hidden, padding_mask2, cfg_scale → v
	* x = x + v * dt
	**************************************************************************************************/

	#pragma once

	#include <string>
	#include <vector>
	#include <map>
	#include <memory>
	#include <cstdint>

	#include "EngineWrapper.hpp"

	class ZipVoiceEngine {
	public:
	struct Config {
	std::string model_dir;
	int max_feat_len = 1024;
	int max_tokens = 384;
	int feat_dim = 100;
	int sampling_rate = 24000;
	int hop_length = 256;
	int num_step = 10;
	float t_shift = 0.5f;
	};

	struct ModelInfo {
	std::string name;
	std::string file;
	std::vector<std::string> inputs;
	std::vector<std::string> outputs;
	};

	struct Timing {
	float encoder_time_sec = 0.0f;
	float duration_expand_time_sec = 0.0f;
	float decoder_time_sec = 0.0f;
	float total_time_sec = 0.0f;
	int generated_frames = 0;
	int features_len = 0;
	};

	ZipVoiceEngine();
	~ZipVoiceEngine();

	/**
	* Initialize from a model directory containing:
	* - runtime_config.json
	* - decoder4_split_manifest.json
	* - encoder.axmodel
	* - decoder_part0.axmodel .. decoder_part3.axmodel
	*/
	int Init(const std::string& model_dir);

	/**
	* Run encoder: cat_tokens[1, max_tokens] → encoded[1, max_tokens, feat_dim].
	* out_encoded should be pre-allocated with size: 1 * max_tokens * feat_dim floats.
	*/
	int RunEncoder(const std::vector<int32_t>& cat_tokens, std::vector<float>& out_encoded);

	/**
	* Duration expand: token-level features → frame-level features.
	*/
	int DurationExpand(const std::vector<float>& encoded,
	int prompt_tokens_len,
	int text_tokens_len,
	int prompt_features_len,
	float speed,
	std::vector<float>& out_text_condition,
	int& out_features_len);

	/**
	* Full sampling: encoder → duration_expand → decoder flow-matching.
	*
	* @param cat_tokens [max_tokens] int32 token IDs
	* @param prompt_tokens_len number of prompt tokens
	* @param text_tokens_len number of text tokens
	* @param prompt_features [prompt_frames, feat_dim] mel features from prompt wav
	* @param prompt_features_len number of prompt feature frames
	* @param speed speed factor (1.0 = normal)
	* @param guidance_scale classifier-free guidance scale
	* @param seed random seed
	* @param out_features output mel features [generated_frames, feat_dim]
	* @param out_timing timing breakdown
	*/
	int Sample(const std::vector<int32_t>& cat_tokens,
	int prompt_tokens_len,
	int text_tokens_len,
	const std::vector<float>& prompt_features,
	int prompt_features_len,
	float speed,
	float guidance_scale,
	int seed,
	std::vector<float>& out_features,
	Timing& out_timing);

	const Config& GetConfig() const { return m_config; }
	bool HasInit() const { return m_has_init; }

	/** Override num_step for subsequent Sample() calls. -1 to reset to config default. */
	void SetNumStep(int n) { m_num_step_override = n; }

	private:
	Config m_config;
	bool m_has_init;

	ModelInfo m_encoder_info;
	std::vector<ModelInfo> m_decoder_parts;

	// Model sessions (one per axmodel)
	std::map<std::string, std::unique_ptr<EngineWrapper>> m_sessions;

	// Decoder metadata
	int m_decoder_seq_len;
	bool m_decoder_has_padding_mask;
	int m_num_step_override = -1; // -1 = use config default

	int LoadConfig(const std::string& model_dir);
	int LoadManifest(const std::string& model_dir);
	int LoadModels(const std::string& model_dir);
	int LoadDecoderMetadata();

	std::vector<float> GetTimesteps(int num_step, float t_shift) const;

	int RunDecoderPart(const ModelInfo& part,
	std::map<std::string, std::vector<float>>& values,
	const std::vector<uint8_t>* padding_mask_data = nullptr,
	std::vector<uint8_t>* padding_mask2_out = nullptr);

	static double GetCurrentTimeMs();
	};