#ifndef parler_model_h #define parler_model_h #include "dac_model.h" #include "ttst5_encoder_model.h" #include "ttssampler.h" enum parler_tensor { PARLER_EMBD, PARLER_EMBD_PROMPTS, PARLER_TEXT_ENCODING, PARLER_POSITIONAL_EMBD, PARLER_HEAD, PARLER_NORM, PARLER_NORM_BIAS, PARLER_LAYER_SELF_ATTN_Q, PARLER_LAYER_SELF_ATTN_K, PARLER_LAYER_SELF_ATTN_V, PARLER_LAYER_SELF_ATTN_O, PARLER_LAYER_SELF_ATTN_NORM, PARLER_LAYER_SELF_ATTN_NORM_BIAS, PARLER_LAYER_ATTN_Q, PARLER_LAYER_ATTN_K, PARLER_LAYER_ATTN_V, PARLER_LAYER_ATTN_O, PARLER_LAYER_ATTN_NORM, PARLER_LAYER_ATTN_NORM_BIAS, PARLER_LAYER_FC1, PARLER_LAYER_FC2, PARLER_LAYER_OUT_NORM, PARLER_LAYER_OUT_NORM_BIAS, }; struct parler_layer { struct ggml_tensor * self_attn_k_proj; struct ggml_tensor * self_attn_q_proj; struct ggml_tensor * self_attn_v_proj; struct ggml_tensor * self_attn_o_proj; struct ggml_tensor * self_attn_norm; struct ggml_tensor * self_attn_norm_bias; struct ggml_tensor * attn_k_proj; struct ggml_tensor * attn_q_proj; struct ggml_tensor * attn_v_proj; struct ggml_tensor * attn_o_proj; struct ggml_tensor * attn_norm; struct ggml_tensor * attn_norm_bias; struct ggml_tensor * cross_k; struct ggml_tensor * cross_v; struct ggml_tensor * fc1; struct ggml_tensor * fc2; struct ggml_tensor * final_norm; struct ggml_tensor * final_norm_bias; }; struct parler_tts_model : tts_model { // These default configurations are based on the configuration of Parler TTS Mini (version 1.0) uint32_t n_output_heads = 9; uint32_t n_encode_length; uint32_t max_encode_length = 512; // This corresponds with the max token length of the conditional prompt uint32_t hidden_size = 1024; uint32_t max_ctx_length = 4096; uint32_t n_attn_heads = 16; uint32_t head_size = 64; uint32_t output_vocab_size = 1088; uint32_t eos_token_id = 1024; uint32_t audio_vocab_size = 1024; uint32_t max_generation_size = 2580; uint32_t n_layers = 24; uint32_t bos_token_id = 1025; uint32_t max_cross_nodes = 32; uint32_t prompt_vocab_size; bool use_cross_attn = true; std::vector embds; std::vector layers; std::vector heads; struct ggml_tensor * precomputed_input_emb; struct ggml_tensor * precomputed_positional_embds; struct ggml_tensor * layer_norm; struct ggml_tensor * layer_norm_bias; struct ggml_tensor * prompt_embd; void assign_weight(std::string name, ggml_tensor * tensor); void prep_constants(gguf_context * meta); void prep_layers(gguf_context * meta); void prep_cross_key_values(int n_threads, struct tts_response * conditional_prompt = nullptr); void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only) { prep_constants(meta_ctx); prep_layers(meta_ctx); tts_model::setup_from_file(meta_ctx, load_context, cpu_only, "decoder", 1.30, max_encode_length*hidden_size*sizeof(float)*n_layers*2); } }; // For assigning weights to the parler model from a gguf file. void assign_parler_layer(parler_tts_model * model, parler_layer & layer, std::string name, ggml_tensor * tensor); void assign_to_decoder(parler_tts_model * model, const std::string name, ggml_tensor * tensor); struct parler_context : runner_context { parler_context(parler_tts_model * model, int n_threads): runner_context(n_threads), model(model) {}; struct parler_tts_model * model; std::vector eos_seen; bool use_cache = true; size_t output_size = 0; // capacity (of tokens positions) for the output buffers int32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch uint32_t current_position = 0; // current position in the active sequence uint32_t prompt_end_position = 0; // the position of the text prompt termination (used for adjusting the cache when incrementally generating) int32_t seq_id; // a unique identifier associated with the active sequence. std::vector output_tokens; struct ggml_tensor * inp_tokens; struct ggml_tensor * audio_inp_tokens; struct ggml_tensor * positions; struct ggml_tensor * attn_mask; struct ggml_tensor * attn_mask_cross; void build_schedule() { runner_context::build_schedule(model->max_nodes()); } void reset(int32_t n_output_heads); }; struct parler_kv_cache { int32_t seq_id; ggml_type type_k = GGML_TYPE_F32; ggml_type type_v = GGML_TYPE_F32; std::vector k_l; std::vector v_l; struct ggml_context * ctx; ggml_backend_buffer_type_t buft; ggml_backend_buffer_t buf; void free() { ggml_free(ctx); ggml_backend_buffer_free(buf); } ~parler_kv_cache() { free(); } }; struct parler_ubatch { parler_ubatch(bool audio_generation, size_t n_tokens, size_t n_audio_tokens, size_t sequence_length, uint32_t * tokens, uint32_t * audio_tokens, uint32_t * positions, uint32_t * true_order, int current_step): audio_generation(audio_generation), n_tokens(n_tokens), n_audio_tokens(n_audio_tokens), sequence_length(sequence_length), tokens(tokens), audio_tokens(audio_tokens), positions(positions), true_order(true_order), current_step(current_step) {}; parler_ubatch() {}; bool audio_generation; // whether we are receiving codebook decoded tokens or text tokens size_t n_tokens; // total sentence tokens size_t n_audio_tokens; // total audio tokens size_t sequence_length; // for just audio tokens the sequence length should be the total_tokens / num_heads; in general this should be n_tokens + n_audio_tokens / num_heads uint32_t * tokens; // [n_tokens] uint32_t * audio_tokens; // [n_audio_tokens] uint32_t * positions; // [sequence_length] uint32_t * true_order; int current_step = 0; // total_generations }; struct parler_context * build_new_parler_context(struct parler_tts_model * model, int n_threads, bool use_cpu = true); static bool parler_kv_cache_init(struct parler_kv_cache * cache, parler_tts_model * model, parler_context * pctx, int32_t seq_id); struct ggml_tensor * parler_build_inp_embd(struct ggml_context * ctx, struct parler_context * pctx, parler_tts_model * model, const parler_ubatch & batch); struct ggml_tensor * parler_build_layer_norm(struct ggml_context * ctx, struct ggml_tensor * inputs, struct ggml_tensor * weight, struct ggml_tensor * bias); void parler_build_kv_store(struct ggml_context * ctx, const parler_kv_cache * kv, struct ggml_cgraph * graph, struct ggml_tensor * k_cur, struct ggml_tensor * v_cur, int32_t n_tokens, int32_t kv_head, int32_t index, int32_t n_embd_gqa); struct ggml_tensor * parler_build_head_outputs(struct ggml_context * ctx, parler_tts_model * model, struct ggml_tensor * cur); struct ggml_tensor * build_attn_mask(ggml_context * ctx, parler_context * pctx, parler_ubatch & batch); struct ggml_tensor * build_attn_mask_cross(ggml_context * ctx, parler_context * pctx, parler_tts_model * model, parler_ubatch & batch); static struct parler_ubatch batch_from_sentence(std::string sentence, parler_tts_model * model, unigram_tokenizer * tokenizer); // This struct is intended to support end-to-end TTS generation. As such, it manages the parler tts model compilation, compute and generation process, // the tokenization and sampling process, and uses the dac_runner struct to encode audio outputs. struct parler_tts_runner : tts_runner { parler_tts_runner(parler_tts_model * model, dac_runner * audio_decoder, parler_context * pctx, unigram_tokenizer * ut, sampler * samp, parler_kv_cache * cache): model(model), dac_runner(audio_decoder), pctx(pctx), tokenizer(ut), sampler(samp), kv_self(cache) {}; ~parler_tts_runner() { if (ctx) { ggml_free(ctx); } model->free(); delete model; delete kv_self; delete dac_runner; delete pctx; delete sampler; } struct parler_tts_model * model; struct dac_runner * dac_runner; struct parler_context * pctx; struct unigram_tokenizer * tokenizer; struct parler_kv_cache * kv_self = nullptr; struct sampler * sampler; void init_build() { tts_runner::init_build(&pctx->buf_compute_meta); } void configure_generation(generation_configuration * config); void assign_weight(std::string name, ggml_tensor * tensor); parler_ubatch build_worst_case_batch(); struct ggml_cgraph * build_parler_graph(parler_ubatch & batch); void set_inputs(parler_ubatch & batch); int decode(parler_ubatch & batch); void prepare_post_load(); bool adjust_for_sequence_continuation(struct parler_ubatch & batch); int generate(std::string sentence, struct tts_response * response, int32_t seq_id = -1); bool check_stopping(); void adjust_output_tokens(std::vector & output_tokens, std::vector & filtered); int generate_from_batch(parler_ubatch & batch, struct tts_response * output); void parler_graph_compute(ggml_cgraph * gf); void just_audio_token_decode(uint32_t * tokens, int32_t sq_len, struct tts_response * output); int generate_audio_tokens(std::string sentence); void update_conditional_prompt(const std::string file_path, const std::string prompt, int n_threads, bool cpu_only = true); }; #endif