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
210 lines
7.9 KiB
C++
210 lines
7.9 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/core/backend_weight_store.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/modules/attention/types.h"
|
|
#include "engine/framework/modules/conv_modules.h"
|
|
#include "engine/framework/modules/linear_module.h"
|
|
#include "engine/framework/modules/norm_modules.h"
|
|
#include "engine/framework/runtime/kv_cache.h"
|
|
#include "engine/framework/modules/streaming_conv_modules.h"
|
|
#include "engine/models/index_tts2/assets.h"
|
|
|
|
#include "ggml-backend.h"
|
|
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace engine::models::index_tts2 {
|
|
|
|
struct IndexTTS2GptConditionSubsamplingWeights {
|
|
engine::modules::Conv2dWeights conv;
|
|
engine::modules::LinearWeights out;
|
|
engine::core::TensorValue pos_enc;
|
|
};
|
|
|
|
struct IndexTTS2GptConditionLayerWeights {
|
|
engine::modules::NormWeights norm_ff;
|
|
engine::modules::NormWeights norm_mha;
|
|
engine::modules::NormWeights norm_conv;
|
|
engine::modules::NormWeights norm_final;
|
|
engine::modules::LinearWeights feed_forward_in;
|
|
engine::modules::LinearWeights feed_forward_out;
|
|
engine::modules::RelativeAttentionWeights self_attn;
|
|
engine::modules::Conv1dWeights conv_pointwise_in;
|
|
engine::modules::DepthwiseConv1dWeights conv_depthwise;
|
|
engine::modules::NormWeights conv_norm;
|
|
engine::modules::Conv1dWeights conv_pointwise_out;
|
|
};
|
|
|
|
struct IndexTTS2GptConditionEncoderWeights {
|
|
IndexTTS2GptConditionSubsamplingWeights subsampling;
|
|
std::vector<IndexTTS2GptConditionLayerWeights> layers;
|
|
engine::modules::NormWeights after_norm;
|
|
};
|
|
|
|
struct IndexTTS2PerceiverAttentionWeights {
|
|
engine::modules::LinearWeights q;
|
|
engine::modules::LinearWeights kv;
|
|
engine::modules::LinearWeights out;
|
|
};
|
|
|
|
struct IndexTTS2PerceiverFeedForwardWeights {
|
|
engine::modules::LinearWeights in;
|
|
engine::modules::LinearWeights out;
|
|
};
|
|
|
|
struct IndexTTS2PerceiverLayerWeights {
|
|
IndexTTS2PerceiverAttentionWeights attention;
|
|
IndexTTS2PerceiverFeedForwardWeights feed_forward;
|
|
};
|
|
|
|
struct IndexTTS2PerceiverWeights {
|
|
engine::core::TensorValue latents;
|
|
engine::modules::LinearWeights project_context;
|
|
std::vector<IndexTTS2PerceiverLayerWeights> layers;
|
|
engine::core::TensorValue norm_gamma;
|
|
};
|
|
|
|
struct IndexTTS2Gpt2LayerWeights {
|
|
engine::modules::NormWeights attn_norm;
|
|
engine::modules::LinearWeights qkv;
|
|
engine::modules::LinearWeights attn_out;
|
|
engine::modules::NormWeights mlp_norm;
|
|
engine::modules::LinearWeights mlp_in;
|
|
engine::modules::LinearWeights mlp_out;
|
|
};
|
|
|
|
struct IndexTTS2GptWeights {
|
|
std::shared_ptr<engine::core::BackendWeightStore> store;
|
|
IndexTTS2GptConditionEncoderWeights speaker_conditioner;
|
|
IndexTTS2GptConditionEncoderWeights emotion_conditioner;
|
|
IndexTTS2PerceiverWeights speaker_perceiver;
|
|
IndexTTS2PerceiverWeights emotion_perceiver;
|
|
engine::core::TensorValue text_embedding;
|
|
engine::core::TensorValue mel_embedding;
|
|
engine::core::TensorValue text_pos_embedding;
|
|
engine::core::TensorValue mel_pos_embedding;
|
|
engine::modules::LinearWeights emotion_vec_projection;
|
|
engine::modules::LinearWeights emotion_layer;
|
|
std::vector<float> speed_embedding_values;
|
|
// v2.5 campplus speaker conditioning (spk_cond_mode="campplus" in the
|
|
// official model_v2.py): projects the 192-dim CAMPPlus speaker/style
|
|
// embedding into a GPT speaker token.
|
|
engine::modules::LinearWeights spk_emb_proj;
|
|
// v2.5: row of this table is added to every text embedding during prefill.
|
|
engine::core::TensorValue lang_embedding;
|
|
std::vector<IndexTTS2Gpt2LayerWeights> gpt_layers;
|
|
engine::modules::NormWeights gpt_final_norm;
|
|
engine::modules::NormWeights final_norm;
|
|
engine::modules::LinearWeights mel_head;
|
|
engine::modules::LinearWeights text_head;
|
|
};
|
|
|
|
struct IndexTTS2GptLatent {
|
|
std::vector<float> values;
|
|
int64_t frames = 0;
|
|
int64_t dims = 0;
|
|
};
|
|
|
|
struct IndexTTS2GptGeneration {
|
|
std::vector<int32_t> codes;
|
|
IndexTTS2GptLatent speech_conditioning_latent;
|
|
uint64_t rng_offset_blocks = 0;
|
|
};
|
|
|
|
struct IndexTTS2GptGenerationRequest {
|
|
std::vector<int32_t> text_tokens;
|
|
// v2 speaker conditioning: wav2vec2bert semantic features of the reference.
|
|
std::vector<float> speaker_semantic;
|
|
int64_t speaker_frames = 0;
|
|
// v2.5 speaker conditioning: 192-dim CAMPPlus speaker embedding, projected
|
|
// by spk_emb_proj inside the prefill graph.
|
|
std::vector<float> speaker_style;
|
|
// v2.5: row of the GPT lang_embedding table added to every text embedding.
|
|
int32_t lang_id = 0;
|
|
std::vector<float> emotion_semantic;
|
|
int64_t emotion_frames = 0;
|
|
std::vector<float> emotion_vector;
|
|
float top_p = 0.8F;
|
|
int top_k = 30;
|
|
float temperature = 0.8F;
|
|
float repetition_penalty = 10.0F;
|
|
bool do_sample = true;
|
|
float length_penalty = 0.0F;
|
|
int num_beams = 3;
|
|
int max_mel_tokens = 1500;
|
|
uint32_t seed = 0;
|
|
};
|
|
|
|
// Mirrors the valid_mask filtering in the official v2.5 prepare_gpt_inputs:
|
|
// any start/stop text tokens in the segment (including the trailing pad
|
|
// appended by the tokenizer) are dropped before the start/stop pair is
|
|
// re-added around it.
|
|
std::vector<int32_t> align_index_tts2_gpt_text_tokens(const std::vector<int32_t> & text_tokens);
|
|
|
|
std::shared_ptr<const IndexTTS2GptWeights> load_index_tts2_gpt_weights(
|
|
const IndexTTS2Assets & assets,
|
|
ggml_backend_t backend,
|
|
engine::core::BackendType backend_type,
|
|
engine::assets::TensorStorageType matmul_storage_type,
|
|
engine::assets::TensorStorageType conv_storage_type,
|
|
size_t weight_context_bytes);
|
|
|
|
class IndexTTS2GptRuntime {
|
|
public:
|
|
IndexTTS2GptRuntime(
|
|
std::shared_ptr<const IndexTTS2Assets> assets,
|
|
engine::core::ExecutionContext & execution,
|
|
size_t graph_arena_bytes,
|
|
size_t weight_context_bytes,
|
|
engine::assets::TensorStorageType matmul_storage_type,
|
|
engine::assets::TensorStorageType conv_storage_type);
|
|
~IndexTTS2GptRuntime();
|
|
|
|
IndexTTS2GptRuntime(const IndexTTS2GptRuntime &) = delete;
|
|
IndexTTS2GptRuntime & operator=(const IndexTTS2GptRuntime &) = delete;
|
|
|
|
void prepare_speaker_conditioning(int64_t frames);
|
|
void prepare_emotion_conditioning(int64_t frames);
|
|
void prepare_generation(int64_t text_tokens, int64_t max_mel_tokens, int64_t num_beams);
|
|
IndexTTS2GptLatent speaker_conditioning(const std::vector<float> & semantic_btc, int64_t frames);
|
|
IndexTTS2GptLatent emotion_conditioning(const std::vector<float> & semantic_btc, int64_t frames);
|
|
std::vector<float> project_emotion_vector(const IndexTTS2GptLatent & emotion_conditioning);
|
|
std::vector<float> merge_emotion_vector(
|
|
const std::vector<float> & speaker_semantic,
|
|
int64_t speaker_frames,
|
|
const std::vector<float> & emotion_semantic,
|
|
int64_t emotion_frames,
|
|
float alpha);
|
|
IndexTTS2GptGeneration generate_speech(const IndexTTS2GptGenerationRequest & request);
|
|
IndexTTS2GptLatent forward_latent(
|
|
const IndexTTS2GptLatent & speech_conditioning_latent,
|
|
const std::vector<int32_t> & text_tokens,
|
|
const std::vector<int32_t> & codes,
|
|
const std::vector<float> & emotion_semantic,
|
|
int64_t emotion_frames,
|
|
const std::vector<float> & emotion_vector);
|
|
void release_conditioning_graphs();
|
|
void release_generation_graphs();
|
|
|
|
private:
|
|
class ConditioningGraph;
|
|
class EmotionVectorGraph;
|
|
class PrefillGraph;
|
|
class DecodeGraph;
|
|
class ForwardGraph;
|
|
|
|
std::shared_ptr<const IndexTTS2Assets> assets_;
|
|
engine::core::ExecutionContext * execution_ = nullptr;
|
|
size_t graph_arena_bytes_ = 0;
|
|
std::shared_ptr<const IndexTTS2GptWeights> weights_;
|
|
std::unique_ptr<ConditioningGraph> speaker_conditioning_graph_;
|
|
std::unique_ptr<ConditioningGraph> emotion_conditioning_graph_;
|
|
std::unique_ptr<EmotionVectorGraph> emotion_vector_graph_;
|
|
std::unique_ptr<PrefillGraph> prefill_graph_;
|
|
std::unique_ptr<DecodeGraph> decode_graph_;
|
|
std::unique_ptr<ForwardGraph> forward_graph_;
|
|
};
|
|
|
|
} // namespace engine::models::index_tts2
|