builds but crashes

otherarch/ttscpp/src/args.cpp

@@ -0,0 +1,164 @@
+#include "ttsargs.h"
+
+std::string arg::help_text() {
+    std::string htxt = full_name;
+    if (abbreviation != "") {
+        htxt += " (" + abbreviation + ")";
+    }
+    htxt += ":\n    ";
+    if (description != "") {
+        htxt += description + "\n";
+    } else {
+        htxt += "is a " + (std::string)(required ? "required " : "optional ") + "parameter.\n";
+    }
+    return htxt;
+}
+
+int string_arg::parse(int argc, const char ** argv) {
+    required = false;
+    value.assign(argv[0]);
+    return 1;
+}
+
+int int_arg::parse(int argc, const char ** argv) {
+    if (required) {
+        required = false;
+    }
+    int val = atoi(argv[0]);
+    *value = val;
+    return 1;
+}
+
+int float_arg::parse(int argc, const char ** argv) {
+    if (required) {
+        required = false;
+    }
+    float val = strtof(argv[0], nullptr);
+    *value = val;
+    return 1;
+}
+
+void arg_list::help() {
+    std::string help_text = "";
+    for (auto arg : fargs) {
+        help_text += arg.help_text();
+    }
+    for (auto arg : iargs) {
+        help_text += arg.help_text();
+
+    }
+    for (auto arg : bargs) {
+        help_text += arg.help_text();
+
+    }
+    for (auto arg : sargs) {
+        help_text += arg.help_text();
+
+    }
+    fprintf(stdout, "%s", help_text.c_str());
+}
+
+void arg_list::validate() {
+    for (auto arg : fargs) {
+        if (arg.required) {
+            fprintf(stderr, "argument '%s' is required.\n", arg.full_name.c_str());
+            exit(1);
+        }
+    }
+    for (auto arg : iargs) {
+        if (arg.required) {
+            fprintf(stderr, "argument '%s' is required.\n", arg.full_name.c_str());
+            exit(1);
+        }
+    }
+    for (auto arg : bargs) {
+        if (arg.required) {
+            fprintf(stderr, "argument '%s' is required.\n", arg.full_name.c_str());
+            exit(1);
+        }
+    }
+    for (auto arg : sargs) {
+        if (arg.required) {
+            fprintf(stderr, "argument '%s' is required.\n", arg.full_name.c_str());
+            exit(1);
+        }
+    }
+}
+
+void arg_list::parse(int argc, const char ** argv) {
+    int current_arg = 1;
+    while (current_arg < argc) {
+        std::string name(argv[current_arg]);
+        if (name == "--help") {
+            for_help = true;
+            return;
+        }
+        current_arg += 1;
+        current_arg += find_and_parse(name, argc - current_arg, argv + current_arg);
+    }
+}
+
+int arg_list::find_and_parse(std::string name, int argc, const char ** argv) {
+    for (int i = 0; i < fargs.size(); i++) {
+        if (fargs[i].full_name == name || fargs[i].abbreviation == name) {
+            return fargs[i].parse(argc, argv);
+        }
+    }
+    for (int i = 0; i < iargs.size(); i++) {
+        if (iargs[i].full_name == name || iargs[i].abbreviation == name) {
+            return iargs[i].parse(argc, argv);
+        }
+    }
+    for (int i = 0; i < bargs.size(); i++) {
+        if (bargs[i].full_name == name || bargs[i].abbreviation == name) {
+            bargs[i].value = !bargs[i].value;
+            bargs[i].required = false;
+            return 0;
+        }
+
+    }
+    for (int i = 0; i < sargs.size(); i++) {
+        if (sargs[i].full_name == name || sargs[i].abbreviation == name) {
+            return sargs[i].parse(argc, argv);
+        }
+    }
+    fprintf(stderr, "argument '%s' is not a valid argument. Call '--help' for information on all valid arguments.\n", name.c_str());
+    exit(1);
+}
+
+std::string arg_list::get_string_param(std::string full_name) {
+    for (auto arg : sargs) {
+        if (arg.full_name == full_name) {
+            return arg.value;
+        }
+    }
+    return "";
+}
+
+int * arg_list::get_int_param(std::string full_name) {
+    for (auto arg : iargs) {
+        if (arg.full_name == full_name) {
+            return arg.value;
+        }
+    }
+    return nullptr;
+}
+
+float * arg_list::get_float_param(std::string full_name) {
+    for (auto arg : fargs) {
+        if (arg.full_name == full_name) {
+            return arg.value;
+        }
+    }
+    return nullptr;
+}
+
+bool arg_list::get_bool_param(std::string full_name) {
+    for (auto arg : bargs) {
+        if (arg.full_name == full_name) {
+            return arg.value;
+        }
+    }
+    return false;
+}
+

otherarch/ttscpp/src/dac_model.cpp

@@ -0,0 +1,212 @@
+#include "dac_model.h"
+#include <algorithm>
+#include <stdexcept>
+
+// For loading DAC model from gguf file.
+static const std::map<std::string, dac_tensor> DAC_TENSOR_GGUF_LOOKUP = {
+    {"initial.bias", DAC_ENCODER_IN_BIAS},
+    {"initial.weight", DAC_ENCODER_IN_KERNEL},
+    {"final.bias", DAC_ENCODER_OUT_BIAS},
+    {"final.weight", DAC_ENCODER_OUT_KERNEL},
+    {"final.alpha", DAC_ENCODER_SNAKE_ALPHA},
+};
+
+void dac_model::prep_constants(gguf_context * meta) {
+    int output_heads_key = search_for_gguf_keys(meta, {"parler-tts.decoder.output_heads", "output_heads", "dia.decoder.output_heads"});
+    if (output_heads_key != -1) {
+        n_heads = gguf_get_val_u32(meta, output_heads_key);
+    }
+
+    int sampling_factor_key = search_for_gguf_keys(meta, {"dac.up_sampling_factor", "up_sampling_factor"});
+    if (sampling_factor_key != -1) {
+        up_sampling_factor = gguf_get_val_u32(meta, sampling_factor_key);
+    }
+
+    int max_gen_key = search_for_gguf_keys(meta, {"parler-tts.decoder.max_generation", "max_generation", "dia.decoder.max_generation"});
+    if (max_gen_key != -1) {
+        max_generation_size = gguf_get_val_u32(meta, max_gen_key);
+    }
+}
+
+void dac_model::prep_layers(gguf_context * meta) {
+    for (int i = 0; i < n_heads; i++) {
+        quantizer_layers.push_back(general_neural_audio_codec::residual_vector_quantize_layer{});
+    }
+
+    for (int i = 0; i < n_layers; i++) {
+        std::string stride_key = "dac_layer_stride_" + std::to_string(i);
+        std::string padding_key = "dac_layer_padding_" + std::to_string(i);
+        int layer_stride_key = search_for_gguf_keys(meta, {"dac." + stride_key, stride_key});
+        if (layer_stride_key == -1) {
+            TTS_ABORT("key %s must be specified in gguf file inorder to initialize the DAC audio decoder.", stride_key.c_str());
+        }
+        int layer_padding_key = search_for_gguf_keys(meta, {"dac." + padding_key, padding_key});
+        if (layer_padding_key == -1) {
+            TTS_ABORT("key %s must be specified in gguf file inorder to initialize the DAC audio decoder.", padding_key.c_str());
+        }
+        layers.push_back(
+            general_neural_audio_codec::layer{
+                gguf_get_val_u32(meta, layer_padding_key),
+                gguf_get_val_u32(meta, layer_stride_key),
+            }
+        );
+    }
+}
+
+void dac_model::assign_weight(std::string name, ggml_tensor * tensor) {
+    assign_to_audio_encoder(this, name, tensor);
+}
+
+void assign_to_audio_encoder(dac_model * model, std::string name, ggml_tensor * tensor) {
+    if (DAC_TENSOR_GGUF_LOOKUP.find(name) != DAC_TENSOR_GGUF_LOOKUP.end()) {
+        switch(DAC_TENSOR_GGUF_LOOKUP.at(name)) {
+            case DAC_ENCODER_IN_BIAS:
+                model->in_conv_bias = ggml_dup_tensor(model->ctx, ggml_transpose(model->ctx, tensor));
+                model->set_tensor(model->in_conv_bias, tensor);
+                break;
+            case DAC_ENCODER_IN_KERNEL:
+                model->in_conv_kernel = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->in_conv_kernel, tensor);
+                break;
+            case DAC_ENCODER_OUT_BIAS:
+                model->out_conv_bias = ggml_dup_tensor(model->ctx, ggml_transpose(model->ctx, tensor));
+                model->set_tensor(model->out_conv_bias, tensor);
+                break;
+            case DAC_ENCODER_OUT_KERNEL:
+                model->out_conv_kernel = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->out_conv_kernel, tensor);
+                break;
+            case DAC_ENCODER_SNAKE_ALPHA:
+                model->snake_alpha = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->snake_alpha, tensor);
+                break;
+            default:
+                fprintf(stdout, "unassigned tensor %s\n", name.c_str());
+                break;
+        }
+    } else if (std::find_if(name.begin(), name.end(), ::isdigit) != name.end())  {
+        auto pair = parse_layer_count(name);
+        int l = pair.first;
+        std::string lt_name = pair.second;
+        if (name.find("quantizers") != std::string::npos) {
+            general_neural_audio_codec::assign_to_quantize_layer((tts_model *) model, model->quantizer_layers[l], lt_name, tensor);
+        } else {
+            general_neural_audio_codec::assign_to_layer((tts_model *) model, model->layers[l - 1], lt_name, tensor);
+        }
+    }
+}
+
+static struct ggml_tensor * dac_build_audio_inputs(struct ggml_context * ctx, struct dac_context * dctx, const dac_ubatch & batch, std::vector<general_neural_audio_codec::residual_vector_quantize_layer> layers) {
+    struct ggml_tensor * embd;
+
+    dctx->inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.sequence_length*dctx->model->n_heads);
+    ggml_set_input(dctx->inp_tokens);
+
+    if (dctx->backend) {
+        ggml_backend_sched_set_tensor_backend(dctx->sched, dctx->inp_tokens, dctx->backend);
+    }
+
+    for(int i = 0; i < dctx->model->n_heads; i++) {
+        auto quantize_layer = dctx->model->quantizer_layers[i];
+        struct ggml_tensor * code = ggml_cont(ctx, ggml_view_2d(ctx, dctx->inp_tokens, 1, batch.sequence_length, dctx->model->n_heads*ggml_type_size(GGML_TYPE_I32), i*ggml_type_size(GGML_TYPE_I32)));
+        code = ggml_reshape_1d(ctx, code, batch.sequence_length);
+        code = general_neural_audio_codec::build_quantize_layer(ctx, code, quantize_layer);
+
+        if (i == 0) {
+            embd = code;
+        } else {
+            embd = ggml_add(ctx, embd, code);
+        }
+    }
+    return embd;
+}
+
+struct dac_context * build_new_dac_context(struct dac_model * model, int n_threads, bool use_cpu) {
+    dac_context * dctx = new dac_context(model, n_threads);
+    if (!use_cpu) {
+#ifdef GGML_USE_METAL
+        dctx->backend = ggml_backend_metal_init();
+#endif
+    }
+    dctx->backend_cpu = ggml_backend_cpu_init();
+    dctx->set_threads();
+    dctx->build_schedule();
+    dctx->buf_compute_meta.resize(ggml_tensor_overhead()*model->max_nodes() + ggml_graph_overhead_custom(model->max_nodes(), false));
+    return dctx;
+}
+
+void dac_runner::prepare_post_load() {
+    dac_ubatch batch;
+    batch.sequence_length = model->max_generation_size;
+    ggml_cgraph * gf = build_dac_graph(batch);
+    dctx->prep_schedule(gf);
+}
+
+struct ggml_cgraph * dac_runner::build_dac_graph(dac_ubatch & batch) {
+    init_build();
+    // splitting this out from the primary graph so that we can better manage streaming (i.e. sentence chunks are better performed this way)
+    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
+
+    struct ggml_tensor * cur;
+    struct ggml_tensor * inputs;
+
+    inputs = dac_build_audio_inputs(ctx, dctx, batch, model->quantizer_layers);
+    ggml_set_name(inputs, "quanitzed_inputs");
+
+    // everything besides the inputs is just a forward pass
+    cur = ggml_conv_1d_tts(ctx, model->in_conv_kernel, inputs, 1, 3, 1);
+    cur = ggml_add(ctx, cur, model->in_conv_bias);
+    for (auto l : model->layers) {
+        cur = general_neural_audio_codec::build_layer(ctx, cur, l);
+    }
+    cur = snake_1d(ctx, model->snake_alpha, cur);
+    cur = ggml_conv_1d_tts(ctx, model->out_conv_kernel, cur, 1, 3, 1);
+    cur = ggml_add(ctx, cur, model->out_conv_bias);
+    cur = ggml_tanh(ctx, cur);
+    ggml_build_forward_expand(gf, cur);
+    free_build();
+    return gf;
+}
+
+void dac_runner::run(uint32_t * input_tokens, uint32_t sequence_length, struct tts_response * outputs) {
+    dac_ubatch batch;
+    batch.input_tokens = input_tokens;
+    batch.sequence_length = sequence_length;
+    ggml_backend_sched_reset(dctx->sched);
+
+    const size_t prev_size = dctx->buf_output ? ggml_backend_buffer_get_size(dctx->buf_output) : 0;
+    const size_t new_size = model->max_generation_size * model->up_sampling_factor * sizeof(float);
+
+    if (!dctx->buf_output || prev_size < new_size) {
+        if (dctx->buf_output) {
+            ggml_backend_buffer_free(dctx->buf_output);
+            dctx->buf_output = nullptr;
+            dctx->logits = nullptr;
+        }
+
+        dctx->buf_output = ggml_backend_buft_alloc_buffer(dctx->backend_cpu_buffer, new_size);
+    }
+
+    outputs->data = (float *) ggml_backend_buffer_get_base(dctx->buf_output);
+    ggml_backend_buffer_clear(dctx->buf_output, 0);
+
+    struct ggml_cgraph * gf = NULL;
+    gf = build_dac_graph(batch);
+
+    // the output is always the last tensor in the graph
+    struct ggml_tensor * result = gf->nodes[gf->n_nodes - 1];
+    ggml_backend_sched_alloc_graph(dctx->sched, gf);
+
+    ggml_backend_tensor_set(dctx->inp_tokens, batch.input_tokens, 0, batch.sequence_length*model->n_heads*ggml_element_size(dctx->inp_tokens));
+
+    ggml_backend_sched_graph_compute_async(dctx->sched, gf);
+
+    dctx->get_ggml_node_data(result, outputs->data, batch.sequence_length*sizeof(float)*model->up_sampling_factor);
+
+    // Reset state for the next token before backend sync, to allow the CPU activities in the reset to
+    // overlap with device computation.
+    ggml_backend_sched_reset(dctx->sched);
+    outputs->n_outputs = sequence_length * model->up_sampling_factor;
+    return;
+}
+

otherarch/ttscpp/src/dac_model.h

@@ -0,0 +1,98 @@
+#ifndef dac_model_h
+#define dac_model_h
+
+#include "general_neural_audio_codec.h"
+#include <map>
+
+enum dac_tensor {
+    DAC_ENCODER_IN_KERNEL,
+    DAC_ENCODER_IN_BIAS,
+    DAC_ENCODER_OUT_KERNEL,
+    DAC_ENCODER_OUT_BIAS,
+    DAC_ENCODER_SNAKE_ALPHA,
+};
+
+struct dac_quantize_layer {
+    struct ggml_tensor * out_proj_kernel;
+    struct ggml_tensor * out_proj_bias;
+    struct ggml_tensor * codebook;
+};
+
+// DAC, Descript Audio Codec, is a channel token to audio autoencoder model (though we only use its decoder functionality).
+// this struct maintains the static tensors for the dac audio decoder graph.
+// As such, this is designed to contain basic configuration and ggml tensor support for DAC.
+// The dac_runner describes how the graph is built and run.
+struct dac_model : tts_model {    
+    // These configs  are essentially built for the 44khZ 8kbps standard DAC model audio encoder and decoder
+    uint32_t n_layers = 4;
+    uint32_t n_heads = 9;
+    uint32_t up_sampling_factor = 512;
+    uint32_t max_generation_size = 2580;
+    
+    struct ggml_tensor * in_conv_kernel;
+    struct ggml_tensor * in_conv_bias;
+    struct ggml_tensor * out_conv_kernel;
+    struct ggml_tensor * out_conv_bias;
+    struct ggml_tensor * snake_alpha;
+    std::vector<general_neural_audio_codec::layer> layers;
+    std::vector<general_neural_audio_codec::residual_vector_quantize_layer> quantizer_layers;
+
+    void assign_weight(std::string name, ggml_tensor * weight);
+    void prep_constants(gguf_context * meta);
+    void prep_layers(gguf_context * meta);
+    void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only) {
+        prep_layers(meta_ctx);
+        prep_constants(meta_ctx);
+        tts_model::setup_from_file(meta_ctx, load_context, cpu_only, "audio_encoder");
+    }
+};
+
+// for loading DAC model from gguf file
+void assign_to_audio_encoder(dac_model * model, std::string name, ggml_tensor * tensor);
+
+// the context used for running the dac model
+struct dac_context : runner_context {
+    dac_context(dac_model * model, int n_threads): runner_context(n_threads), model(model) {};
+    
+    struct dac_model * model;
+        
+    struct ggml_tensor * inp_tokens;
+    
+    void build_schedule() {
+        runner_context::build_schedule(model->max_nodes());
+    }
+};
+
+struct dac_context * build_new_dac_context(struct dac_model * model, int n_threads, bool use_cpu = true);
+
+struct dac_ubatch {
+    uint32_t * input_tokens;
+    uint32_t sequence_length;
+};
+
+static struct ggml_tensor * dac_build_audio_inputs(struct ggml_context * ctx, struct dac_context * dctx, const dac_ubatch & batch, std::vector<general_neural_audio_codec::residual_vector_quantize_layer> layers);
+
+// This struct is intended to manage the dac model's graph compilation and compute function.
+struct dac_runner : tts_runner {
+    dac_runner(dac_model * model, dac_context * context): model(model), dctx(context) {};
+    ~dac_runner() {
+        if (ctx) {
+            ggml_free(ctx);
+        }
+        model->free();
+        delete model;
+        delete dctx;
+    }
+    dac_model * model;
+    dac_context * dctx;
+    
+    void init_build() {
+        tts_runner::init_build(&dctx->buf_compute_meta);
+    }
+    
+    void prepare_post_load();
+    struct ggml_cgraph * build_dac_graph(dac_ubatch & batch);
+    void run(uint32_t * input_tokens, uint32_t sequence_length, struct tts_response * outputs);
+};
+
+#endif

otherarch/ttscpp/src/dia_model.cpp

@@ -0,0 +1,911 @@
+#include "dia_model.h"
+
+void dia_model::assign_weight(std::string name, struct ggml_tensor * tensor) {
+    std::vector<std::string> parts = split(name, ".");
+    TTS_ASSERT(parts.size() >= 3);
+
+    if (parts[1] == "encoder") {
+        assign_to_encoder(parts, tensor, name);
+    } else if (parts[1] == "decoder"){
+        assign_to_decoder(parts, tensor, name);
+    } else {
+        TTS_ABORT("Unrecognized tensor '%s' when loading Dia from GGUF file.", name.c_str());
+    }
+}
+
+void dia_model::assign_to_encoder(std::vector<std::string> parts, struct ggml_tensor * tensor, std::string name) {
+    if (parts[2] == "embedding") {
+        encoder->embedding = ggml_dup_tensor(ctx, tensor);
+        set_tensor(encoder->embedding, tensor);
+    } else if (parts[2] == "norm") {
+        encoder->norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(encoder->norm, tensor);
+    } else if (parts[2] == "layers") {
+        TTS_ASSERT(parts.size() >= 4);
+        int index = std::stoi(parts[3]);
+        TTS_ASSERT(index < decoder->layers.size());
+        assign_to_encoder_layer(parts[4], encoder->layers[index], tensor);
+    } else {
+        TTS_ABORT("Unrecognized tensor '%s' when loading Dia from GGUF file.", name.c_str());
+    }
+}
+
+void dia_model::assign_to_decoder(std::vector<std::string> parts, struct ggml_tensor * tensor, std::string name) {
+    if (parts[2] == "embeddings") {
+        TTS_ASSERT(parts.size() > 2);
+        int index = std::stoi(parts[3]);
+        TTS_ASSERT(index < decoder->embds.size());
+        decoder->embds[index] = ggml_dup_tensor(ctx, tensor);
+        set_tensor(decoder->embds[index], tensor);
+    } else if (parts[2] == "norm") {
+        decoder->norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(decoder->norm, tensor);
+    } else if (parts[2] == "heads") {
+        TTS_ASSERT(parts.size() > 2);
+        int index = std::stoi(parts[3]);
+        TTS_ASSERT(index < decoder->heads.size());
+        decoder->heads[index] = ggml_dup_tensor(ctx, tensor);
+        set_tensor(decoder->heads[index], tensor);
+    } else if (parts[2] == "layers") {
+        TTS_ASSERT(parts.size() >= 4);
+        int index = std::stoi(parts[3]);
+        TTS_ASSERT(index < decoder->layers.size());
+        assign_to_decoder_layer(parts[4], decoder->layers[index], tensor);
+    } else {
+        TTS_ABORT("Unrecognized tensor '%s' when loading Dia from GGUF file.", name.c_str());
+    }
+}
+
+void dia_model::assign_to_encoder_layer(std::string part, dia_encoder_layer * layer, struct ggml_tensor * tensor) {
+    if (part == "q_proj") {
+        layer->q = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->q, tensor);
+    } else if (part == "k_proj") {
+        layer->k = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->k, tensor);
+    } else if (part == "v_proj") {
+        layer->v = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->v, tensor);
+    } else if (part == "o_proj") {
+        layer->o = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->o, tensor);
+    } else if (part == "pre_sa_norm") {
+        layer->self_attn_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->self_attn_norm, tensor);
+    } else if (part == "post_sa_norm") {
+        layer->mlp_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->mlp_norm, tensor);
+    } else if (part == "gate") {
+        layer->gate = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->gate, tensor);
+    } else if (part == "up") {
+        layer->up = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->up, tensor);
+    } else if (part == "wo") {
+        layer->out = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->out, tensor);
+    } else {
+        TTS_ABORT("Unrecognized tensor '%s' for encoder layer when loading Dia from GGUF file.", part.c_str());
+    }
+}
+
+void dia_model::assign_to_decoder_layer(std::string part, dia_decoder_layer * layer, struct ggml_tensor * tensor) {
+    if (part == "self_q_proj") {
+        layer->self_attn_q = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->self_attn_q, tensor);
+    } else if (part == "self_k_proj") {
+        layer->self_attn_k = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->self_attn_k, tensor);
+    } else if (part == "self_v_proj") {
+        layer->self_attn_v = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->self_attn_v, tensor);
+    } else if (part == "self_o_proj") {
+        layer->self_attn_o = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->self_attn_o, tensor);
+    } else if (part == "cross_q_proj") {
+        layer->cross_attn_q = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->cross_attn_q, tensor);
+    } else if (part == "cross_k_proj") {
+        layer->cross_attn_k = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->cross_attn_k, tensor);
+    } else if (part == "cross_v_proj") {
+        layer->cross_attn_v = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->cross_attn_v, tensor);
+    } else if (part == "cross_o_proj") {
+        layer->cross_attn_o = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->cross_attn_o, tensor);
+    } else if (part == "pre_sa_norm") {
+        layer->self_attn_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->self_attn_norm, tensor);
+    } else if (part == "pre_mlp_norm") {
+        layer->mlp_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->mlp_norm, tensor);    
+    } else if (part == "pre_ca_norm") {
+        layer->cross_attn_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->cross_attn_norm, tensor);
+    } else if (part == "gate") {
+        layer->gate = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->gate, tensor);
+    } else if (part == "up") {
+        layer->up = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->up, tensor);
+    } else if (part == "wo") {
+        layer->out = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer->out, tensor);
+    } else {
+        TTS_ABORT("Unrecognized tensor '%s' for encoder layer when loading Dia from GGUF file.", part.c_str());
+    }
+}
+
+void dia_model::prep_layers() {
+    encoder = new dia_encoder;
+    decoder = new dia_decoder;
+    encoder->layers.reserve((size_t) n_encoder_layers);
+    for (int i = 0; i < (int) n_encoder_layers; i++) {
+        dia_encoder_layer * l = new dia_encoder_layer;
+        encoder->layers.push_back(l);
+    }
+
+    decoder->layers.reserve((size_t) n_decoder_layers);
+    for (int i = 0; i < (int) n_decoder_layers; i++) {
+        dia_decoder_layer * l = new dia_decoder_layer;
+        decoder->layers.push_back(l);
+    }
+    
+    decoder->embds.reserve((size_t) n_output_heads);
+    decoder->heads.reserve((size_t) n_output_heads);
+    for (int i = 0; i < n_output_heads; i++) {
+        struct ggml_tensor * h = nullptr;
+        struct ggml_tensor * embd = nullptr;
+        decoder->embds.push_back(embd);
+        decoder->heads.push_back(h);
+    }
+}
+
+void dia_model::prep_constants(gguf_context * meta) {
+    int output_heads_key = gguf_find_key(meta, "dia.decoder.output_heads");
+    if (output_heads_key != -1) {
+        n_output_heads = gguf_get_val_u32(meta, output_heads_key);
+    }
+
+    int decoder_layers_key = gguf_find_key(meta, "dia.decoder.layers");
+    if (decoder_layers_key != -1) {
+        n_decoder_layers = gguf_get_val_u32(meta, decoder_layers_key);
+    }
+
+    int encoder_layers_key = gguf_find_key(meta, "dia.encoder.layers");
+    if (encoder_layers_key != -1) {
+        n_encoder_layers = gguf_get_val_u32(meta, encoder_layers_key);
+    }
+
+    int decoder_hidden_size_key = gguf_find_key(meta, "dia.decoder.hidden_size");
+    if (decoder_hidden_size_key != -1) {
+        decoder_hidden_size = gguf_get_val_u32(meta, decoder_hidden_size_key);
+    }
+
+    int decoder_attn_heads_key = gguf_find_key(meta, "dia.decoder.attn_heads");
+    if (decoder_attn_heads_key != -1) {
+        decoder_attn_heads = gguf_get_val_u32(meta, decoder_attn_heads_key);
+    }
+
+    int decoder_query_heads_key = gguf_find_key(meta, "dia.decoder.query_heads");
+    if (decoder_query_heads_key != -1) {
+        decoder_query_heads = gguf_get_val_u32(meta, decoder_query_heads_key);
+    }
+
+    int encoder_attn_heads_key = gguf_find_key(meta, "dia.encoder.attn_heads");
+    if (encoder_attn_heads_key != -1) {
+        encoder_attn_heads = gguf_get_val_u32(meta, encoder_attn_heads_key);
+    }    
+
+    int head_size_key = gguf_find_key(meta, "dia.attn_head_size");
+    if (head_size_key != -1) {
+        head_size = gguf_get_val_u32(meta, head_size_key);
+    }
+
+    int eos_token_id_key = gguf_find_key(meta, "dia.eos_token_id");
+    if (eos_token_id_key != -1) {
+        eos_token_id = gguf_get_val_u32(meta, eos_token_id_key);
+    }
+
+    int bos_token_id_key = gguf_find_key(meta, "dia.bos_token_id");
+    if (bos_token_id_key != -1) {
+        bos_token_id = gguf_get_val_u32(meta, bos_token_id_key);
+    }
+
+    int pad_token_id_key = gguf_find_key(meta, "dia.pad_token_id");
+    if (pad_token_id_key != -1) {
+        pad_token_id = gguf_get_val_u32(meta, pad_token_id_key);
+    }
+
+    int max_context_key = gguf_find_key(meta, "dia.encoder.max_context_length");
+    if (max_context_key != -1) {
+        max_encoder_context_length = gguf_get_val_u32(meta, max_context_key);
+    }
+
+    int output_vocab_size_key = gguf_find_key(meta, "dia.decoder.output_vocab_size");
+    if (output_vocab_size_key != -1) {
+        output_vocab_size = gguf_get_val_u32(meta, output_vocab_size_key);
+    }
+
+    int audio_vocab_size_key = gguf_find_key(meta, "dia.decoder.audio_vocab_size");
+    if (audio_vocab_size_key != -1) {
+        audio_vocab_size = gguf_get_val_u32(meta, audio_vocab_size_key);
+    }
+
+    int max_generation_size_key = gguf_find_key(meta, "dia.decoder.max_generation_size");
+    if (max_generation_size_key != -1) {
+        max_generation_size = gguf_get_val_u32(meta, max_generation_size_key);
+    }
+    int max_delay_key = gguf_find_key(meta, "dia.max_delay");
+    if (max_delay_key != -1) {
+        max_delay = gguf_get_val_u32(meta, max_delay_key);
+    }
+
+    // please note that this value is not currently set in the gguf encoder as it effectively only exists as a default
+    // python parameter (rather than an attribute in the model config) for the python Dia model.
+    int cfg_scale_key = gguf_find_key(meta, "dia.cfg_scale");
+    if (cfg_scale_key != -1) {
+        cfg_scale_data[0] = gguf_get_val_f32(meta, cfg_scale_key);
+    }
+}
+
+void dia_context::reset() {
+    current_position = 0;
+    prompt_size = 0;
+    output_tokens.clear();
+    delay_steps = -1;
+}
+
+struct dia_context * build_new_dia_context(struct dia_model * model, int n_threads, bool use_cpu) {
+    dia_context * dctx = new dia_context(model, n_threads);
+    if (!use_cpu) {
+#ifdef GGML_USE_METAL
+        dctx->backend = ggml_backend_metal_init();
+#endif
+    }
+    dctx->backend_cpu = ggml_backend_cpu_init();
+    dctx->set_threads();
+    dctx->build_schedule();
+    dctx->buf_compute_meta.resize(ggml_tensor_overhead()*model->max_nodes() + ggml_graph_overhead_custom(model->max_nodes(), false));
+    return dctx;
+}
+
+static bool dia_kv_cache_init(struct dia_kv_cache * cache, dia_model * model, dia_context * dctx) {    
+    ggml_backend_buffer_type_t buft = nullptr;
+    // this will only really support cpu or metal for the time being;
+    if (dctx->backend != nullptr) {
+#ifdef GGML_USE_METAL
+        buft = ggml_backend_metal_buffer_type();
+#endif
+    } else {
+        buft = ggml_backend_cpu_buffer_type();
+    }
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ (4u * model->n_decoder_layers + 1) * ggml_tensor_overhead(),
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true,
+    };
+    ggml_context * ctx = ggml_init(params);
+    if (!ctx) {
+        return false;
+    }
+    cache->ctx = ctx;
+
+    cache->k_l.reserve(model->n_decoder_layers);
+    cache->v_l.reserve(model->n_decoder_layers);
+    cache->cross_k_l.reserve(model->n_decoder_layers);
+    cache->cross_v_l.reserve(model->n_decoder_layers);
+
+    for (int i = 0; i < (int) model->n_decoder_layers; i++) {
+        struct ggml_tensor * k = ggml_new_tensor_1d(cache->ctx, cache->tensor_type, model->head_size * model->decoder_attn_heads * model->max_generation_size * 2);
+        struct ggml_tensor * v = ggml_new_tensor_1d(cache->ctx, cache->tensor_type, model->head_size * model->decoder_attn_heads * model->max_generation_size * 2);
+        struct ggml_tensor * cross_k = ggml_new_tensor_1d(cache->ctx, cache->tensor_type, model->head_size * model->decoder_attn_heads * model->max_encoder_context_length * 2);
+        struct ggml_tensor * cross_v = ggml_new_tensor_1d(cache->ctx, cache->tensor_type, model->head_size * model->decoder_attn_heads * model->max_encoder_context_length * 2);
+        ggml_format_name(k, "cache_k_l%d", i);
+        ggml_format_name(v, "cache_v_l%d", i);
+        ggml_format_name(cross_k, "cache_cross_k_l%d", i);
+        ggml_format_name(cross_v, "cache_cross_v_l%d", i);
+        cache->k_l.push_back(k);
+        cache->v_l.push_back(v);
+        cache->cross_k_l.push_back(cross_k);
+        cache->cross_v_l.push_back(cross_v);
+    }
+
+    // allocate tensors and initialize the buffers to avoid NaNs in the padding
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(cache->ctx, buft);
+    if (!buf) {
+        return false;
+    }
+    ggml_backend_buffer_clear(buf, 0);
+    cache->buf = buf;
+
+    return true;
+}
+
+static struct ggml_tensor * build_dia_decoder_inp_embd(struct ggml_context * ctx, dia_context *dctx, dia_decoder * decoder, dia_ubatch & batch, uint32_t n_output_heads) {
+    struct ggml_tensor * input_embs;
+
+    dctx->audio_inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_output_heads * 2);
+    ggml_set_input(dctx->audio_inp_tokens);
+    for (int i = 0; i < n_output_heads; i++) {
+        struct ggml_tensor * view = ggml_view_1d(ctx, dctx->audio_inp_tokens, 2, i * ggml_element_size(dctx->audio_inp_tokens));
+        view->nb[0] = n_output_heads * ggml_element_size(dctx->audio_inp_tokens);
+        if (i == 0) {
+            input_embs = ggml_get_rows(ctx, decoder->embds[i], view);
+        } else {
+            input_embs = ggml_add(ctx, ggml_get_rows(ctx, decoder->embds[i], view), input_embs);
+        }
+    }
+    return input_embs;
+}
+
+static struct ggml_tensor * dia_layer_norm(struct ggml_context * ctx, struct ggml_tensor * inputs, struct ggml_tensor * weight) {
+    // dia always uses 1e-5 as the default eps
+    float eps = 0.00001;
+    inputs = ggml_rms_norm(ctx, inputs, eps);
+    return ggml_mul(ctx, inputs, weight);
+}
+
+static struct ggml_tensor * build_dia_encoder_attn_mask(ggml_context * ctx, struct dia_context * dctx, dia_model * model) {
+    dctx->encode_attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) model->max_encoder_context_length, (int64_t) model->max_encoder_context_length);
+    ggml_set_input(dctx->encode_attn_mask);
+
+    return dctx->encode_attn_mask;
+}
+
+static struct ggml_tensor * build_dia_head_outputs(struct ggml_context * ctx, dia_model * model, struct ggml_tensor * cur) {
+    // going to cat the heads together and then reshape them
+    struct ggml_tensor * out;
+    for (int i = 0; i < model->n_output_heads; i++) {
+        if (i == 0) {
+            out = ggml_mul_mat(ctx, model->decoder->heads[i], cur);
+        } else {
+            out = ggml_concat(ctx, out, ggml_mul_mat(ctx, model->decoder->heads[i], cur), 2);
+        }
+    }
+    struct ggml_tensor * cond = ggml_cont(ctx, ggml_view_2d(ctx, out, out->ne[0], out->ne[2], out->nb[2], 0));
+    struct ggml_tensor * uncond = ggml_cont(ctx, ggml_view_2d(ctx, out, out->ne[0], out->ne[2], out->nb[2], out->nb[1]));
+    return ggml_map_custom2(ctx, cond, uncond, &cfg_scale, out->ne[0], &model->cfg_scale_data);
+}
+
+static struct ggml_tensor * build_dia_encoder(ggml_context * ctx, dia_model * model, dia_context * dctx, dia_ubatch & batch) {
+    dctx->inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, model->max_encoder_context_length*2);
+    ggml_set_input(dctx->inp_tokens);
+
+    dctx->encode_positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, model->max_encoder_context_length);
+    ggml_set_input(dctx->encode_positions);
+
+    struct ggml_tensor * attn_mask = build_dia_encoder_attn_mask(ctx, dctx, model);
+
+    struct ggml_tensor * cur = ggml_reshape_3d(ctx, ggml_get_rows(ctx, model->encoder->embedding, dctx->inp_tokens), model->encoder_hidden_size, model->max_encoder_context_length, 2);
+    for (auto layer : model->encoder->layers) {
+        struct ggml_tensor * residual = cur;
+        
+        cur = dia_layer_norm(ctx, cur, layer->self_attn_norm);
+        // self-attention
+        {
+            struct ggml_tensor * Qcur = ggml_mul_mat(ctx, layer->q, cur);
+            struct ggml_tensor * Kcur = ggml_mul_mat(ctx, layer->k, cur);
+            struct ggml_tensor * Vcur = ggml_mul_mat(ctx, layer->v, cur);
+
+            // Strangely Dia follows the neoX Rotary Positional Embeddings Protocol
+            Qcur = ggml_rope(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, Qcur, model->head_size, model->encoder_attn_heads, model->max_encoder_context_length, 2)), dctx->encode_positions, model->head_size, 2);
+            Kcur = ggml_rope(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, Kcur, model->head_size, model->encoder_attn_heads, model->max_encoder_context_length, 2)), dctx->encode_positions, model->head_size, 2);
+            struct ggml_tensor * q = ggml_cont(ctx, ggml_permute(ctx, Qcur, 0, 2, 1, 3));
+            struct ggml_tensor * k = ggml_cont(ctx, ggml_permute(ctx, Kcur, 0, 2, 1, 3));
+            struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q);
+            kq = ggml_soft_max_ext(ctx, kq, attn_mask, 1.0f, 0.0f);
+            struct ggml_tensor * v = ggml_cont_4d(ctx, ggml_transpose(ctx, Vcur), model->max_encoder_context_length, model->head_size, model->encoder_attn_heads, 2);
+            struct ggml_tensor * kqv = ggml_mul_mat(ctx, kq, v);
+            struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 2, 0, 1, 3);
+
+            // It is unclear why the attention ops in Dia's encoder don't project to the embedding dimension size as is standard. Instead they up project to the decoder's embedding dimension
+            // then down project back the the encoder embedding dimension. 
+            cur = ggml_cont_3d(ctx, kqv_merged, model->decoder_hidden_size, model->max_encoder_context_length, 2);
+            cur = ggml_mul_mat(ctx, layer->o, cur);
+        }
+
+        cur = ggml_add(ctx, cur, residual);
+        struct ggml_tensor * residual_mlp = cur;
+
+        cur = dia_layer_norm(ctx, cur, layer->mlp_norm);
+        // mlp
+        {
+            cur = ggml_mul(ctx, ggml_silu(ctx, ggml_mul_mat(ctx, layer->gate, cur)), ggml_mul_mat(ctx, layer->up, cur));
+            cur = ggml_mul_mat(ctx, layer->out, cur);
+        }
+
+        cur = ggml_add(ctx, cur, residual_mlp);
+    }
+
+    cur = dia_layer_norm(ctx, cur, model->encoder->norm);
+    return cur;
+}
+
+static struct ggml_tensor * repeat_interleave_dim1(ggml_context * ctx, struct ggml_tensor * a, int repeat) {
+    //return ggml_repeat(ctx, a, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, a->ne[0], 4*a->ne[1], a->ne[2], a->ne[3]));
+    struct ggml_tensor * running;
+    for (int i = 0; i < a->ne[1]; i++) {
+        int offset = i * a->nb[1];
+        struct ggml_tensor * t = ggml_cont(ctx, ggml_view_4d(ctx, a, a->ne[0], 1, a->ne[2], a->ne[3], a->nb[1], a->nb[2], a->nb[3], offset));
+        t = ggml_repeat(ctx, t, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, a->ne[0], repeat, a->ne[2], a->ne[3]));
+        if (i == 0) {
+            running = t;
+        } else {
+            running = ggml_concat(ctx, running, t, 1);
+        }
+    }
+    return running;
+}
+
+static void build_dia_self_kv_store(ggml_context * ctx, dia_context * dctx, dia_model * model, dia_kv_cache * kv, ggml_cgraph * gf, struct ggml_tensor * k, struct ggml_tensor * v, dia_ubatch & batch, int layer_index) {
+    int64_t attn_size = model->head_size * model->decoder_attn_heads;
+
+    struct ggml_tensor * k_cache_view = 
+        ggml_view_2d(
+                ctx, kv->k_l[layer_index], attn_size, 2, 
+                attn_size * model->max_generation_size * ggml_element_size(kv->k_l[layer_index]), 
+                attn_size*dctx->current_position*ggml_element_size(kv->k_l[layer_index]));
+
+    k = ggml_rope(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, k, model->head_size, model->decoder_attn_heads / model->decoder_query_heads, batch.sequence_length, 2)), dctx->positions, model->head_size, 2);
+    // Since the sequence length should always be 1 here this is the most pertinent time to repeat the heads for grouped query attention.
+    // If GGML supported a repeat_interleave op then it would be more optimal to store just the groups in the cache and interleave the attention heads after recalling
+    // from the cache
+    k = repeat_interleave_dim1(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, k, model->head_size, model->decoder_attn_heads / model->decoder_query_heads, batch.sequence_length, 2)), model->decoder_query_heads);
+    k = ggml_cont(ctx, ggml_reshape_2d(ctx, k, attn_size, 2));
+
+    ggml_build_forward_expand(gf, ggml_cpy(ctx, k, k_cache_view));
+
+    struct ggml_tensor * v_cache_view = nullptr;
+
+    v_cache_view = ggml_view_2d(
+            ctx, kv->v_l[layer_index], attn_size, 2, 
+            attn_size * model->max_generation_size * ggml_element_size(kv->v_l[layer_index]), 
+            attn_size*dctx->current_position*ggml_element_size(kv->v_l[layer_index]));
+
+    // Since the sequence length should always be 1 here this is the most pertinent time to repeat the heads for grouped query attention.
+    // If GGML supported a repeat_interleave op then it would be more optimal to store just the groups in the cache and interleave the attention heads after recalling
+    // from the cache
+    v = repeat_interleave_dim1(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, v, model->head_size, model->decoder_attn_heads / model->decoder_query_heads, batch.sequence_length, 2)), model->decoder_query_heads);
+
+    ggml_build_forward_expand(gf, ggml_cpy(ctx, v, v_cache_view));
+}
+
+static void build_dia_cross_kv_store(ggml_context * ctx, dia_context * dctx, dia_model * model, dia_kv_cache * kv, ggml_cgraph * gf, struct ggml_tensor * encoder_hidden_states, int layer_index) {
+    dia_decoder_layer * layer = model->decoder->layers[layer_index];
+    struct ggml_tensor * encoder_states_key_view = ggml_cont(ctx, ggml_view_3d(
+        ctx, 
+        encoder_hidden_states, 
+        model->encoder_hidden_size, 
+        dctx->prompt_size, 
+        2, 
+        model->encoder_hidden_size * ggml_element_size(encoder_hidden_states), model->encoder_hidden_size * model->max_encoder_context_length * ggml_element_size(encoder_hidden_states), 0));
+
+    struct ggml_tensor * k = ggml_mul_mat(ctx, layer->cross_attn_k, encoder_states_key_view);
+    struct ggml_tensor * positions_view = ggml_view_1d(ctx, dctx->encode_positions, dctx->prompt_size, 0);
+
+    k = ggml_rope(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, k, model->head_size, model->decoder_attn_heads, dctx->prompt_size, 2)), positions_view, model->head_size, 2);
+    k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 1, 3, 2));
+
+    struct ggml_tensor * k_cache_view =
+        ggml_view_4d(
+                ctx, kv->cross_k_l[layer_index], model->head_size, model->decoder_attn_heads, 2, dctx->prompt_size, 
+                model->head_size*ggml_element_size(kv->cross_k_l[layer_index]), 
+                model->head_size*model->decoder_attn_heads*ggml_element_size(kv->cross_k_l[layer_index]),
+                model->head_size*model->decoder_attn_heads*2*ggml_element_size(kv->cross_k_l[layer_index]),
+                0);
+
+    ggml_build_forward_expand(gf, ggml_cpy(ctx, k, k_cache_view));
+
+    struct ggml_tensor * v = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, layer->cross_attn_v, encoder_hidden_states)));
+    v = ggml_cont_4d(ctx, v, model->max_encoder_context_length, model->head_size, model->decoder_attn_heads, 2);
+
+    struct ggml_tensor * v_cache_view =
+        ggml_view_4d(
+                ctx, kv->cross_v_l[layer_index], model->max_encoder_context_length, model->head_size, model->decoder_attn_heads, 2, 
+                model->max_encoder_context_length*ggml_element_size(kv->cross_v_l[layer_index]), 
+                model->head_size*model->max_encoder_context_length*ggml_element_size(kv->cross_v_l[layer_index]), 
+                model->head_size*model->max_encoder_context_length*model->decoder_attn_heads*ggml_element_size(kv->cross_v_l[layer_index]), 
+                0);
+
+    ggml_build_forward_expand(gf, ggml_cpy(ctx, v, v_cache_view));
+}
+
+static struct ggml_tensor * build_dia_decoder(
+        ggml_cgraph * gf,
+        ggml_context * ctx, 
+        dia_model * model, 
+        dia_context * dctx, 
+        dia_kv_cache * cache, 
+        dia_ubatch & batch, 
+        struct ggml_tensor * encoder_hidden_states) {
+    dctx->positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.sequence_length);
+    ggml_set_input(dctx->positions);
+    struct ggml_tensor * cur = build_dia_decoder_inp_embd(ctx, dctx, model->decoder, batch, model->n_output_heads);
+
+    for (int l = 0; l < model->decoder->layers.size(); l++){
+        dia_decoder_layer * layer = model->decoder->layers[l];
+        struct ggml_tensor * residual = cur;
+        
+        cur = dia_layer_norm(ctx, cur, layer->self_attn_norm);
+        // self-attention
+        {
+            struct ggml_tensor * Qcur = ggml_mul_mat(ctx, layer->self_attn_q, cur);
+            struct ggml_tensor * Kcur = ggml_mul_mat(ctx, layer->self_attn_k, cur);
+            struct ggml_tensor * Vcur = ggml_mul_mat(ctx, layer->self_attn_v, cur);
+
+            build_dia_self_kv_store(ctx, dctx, model, cache, gf, Kcur, Vcur, batch, l);
+            struct ggml_tensor * k =
+                ggml_view_4d(ctx, cache->k_l[l],
+                        model->head_size, model->decoder_attn_heads, dctx->current_position + 1, 2,
+                        ggml_element_size(cache->k_l[l]) * model->head_size,
+                        ggml_element_size(cache->k_l[l]) * model->decoder_attn_heads * model->head_size,
+                        ggml_element_size(cache->k_l[l]) * model->decoder_attn_heads * model->head_size * model->max_generation_size,
+                        0);
+            k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3));
+
+            struct ggml_tensor * v = 
+                ggml_view_3d(ctx, cache->v_l[l],
+                        model->head_size * model->decoder_attn_heads, dctx->current_position + 1, 2,
+                        ggml_element_size(cache->v_l[l]) * model->decoder_attn_heads * model->head_size,
+                        ggml_element_size(cache->v_l[l]) * model->decoder_attn_heads * model->head_size * model->max_generation_size,
+                        0);
+            v = ggml_cont_4d(ctx, ggml_transpose(ctx, v), dctx->current_position + 1, model->head_size, model->decoder_attn_heads, 2); 
+
+            // As noted in the encoder Dia uses the Neo-X protocol for RoPE.
+            Qcur = ggml_rope(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, Qcur, model->head_size, model->decoder_attn_heads, batch.sequence_length, 2)), dctx->positions, model->head_size, 2);
+            struct ggml_tensor * q = ggml_cont(ctx, ggml_permute(ctx, Qcur, 0, 2, 1, 3));
+            struct ggml_tensor * kq = ggml_mul_mat(ctx, ggml_cont(ctx, k), q);
+
+            // given that attention bias, scaling and masking are not used for decoding, it might be faster to prefer the #ggml_soft_max op here,
+            kq = ggml_soft_max_ext(ctx, kq, nullptr, 1.0f, 0.0f);
+            struct ggml_tensor * kqv = ggml_mul_mat(ctx, kq, v);
+            struct ggml_tensor * kqv_merged = ggml_cont(ctx, ggml_permute(ctx, kqv, 2, 0, 1, 3));
+            cur = ggml_cont_3d(ctx, kqv_merged, model->decoder_hidden_size, batch.sequence_length, 2);
+            cur = ggml_mul_mat(ctx, layer->self_attn_o, cur);
+        }
+
+
+        // if we ever need to support multiple step decoder runs then this reshape will need to be replaced with permutation.
+        cur = ggml_cont_2d(ctx, cur, cur->ne[0], 2);
+        cur = ggml_add(ctx, cur, residual);
+        struct ggml_tensor * residual_cross = cur;
+
+        cur = dia_layer_norm(ctx, cur, layer->cross_attn_norm);
+        // cross-attention
+        {
+            struct ggml_tensor * cross_Qcur = ggml_mul_mat(ctx, layer->cross_attn_q, cur);
+
+            // only load the cross attention kv store when performing the encoding step
+            if (batch.encoder_step) {
+                build_dia_cross_kv_store(ctx, dctx, model, cache, gf, encoder_hidden_states, l);
+            }
+
+            struct ggml_tensor * cross_k = 
+                ggml_view_4d(
+                        ctx, cache->cross_k_l[l], model->head_size, model->decoder_attn_heads, 2,
+                        model->max_encoder_context_length, model->head_size*ggml_element_size(cache->cross_k_l[l]), 
+                        model->head_size*model->decoder_attn_heads*ggml_element_size(cache->cross_k_l[l]), 
+                        model->head_size*model->decoder_attn_heads*2*ggml_element_size(cache->cross_k_l[l]),                 
+                        0);
+            // the double permute operation shouldn't be necessary here, but it seems that currently ggml permute only currently alows for a single
+            // axis pair to be transposed.
+            cross_k = ggml_cont(ctx, ggml_permute(ctx, ggml_permute(ctx, cross_k, 0, 1, 3, 2), 0, 2, 1, 3));
+
+            struct ggml_tensor * cross_v = 
+                ggml_cont(ctx, ggml_view_4d(
+                        ctx, cache->cross_v_l[l], model->max_encoder_context_length, model->head_size, model->decoder_attn_heads, 2,
+                        model->max_encoder_context_length*ggml_element_size(cache->cross_v_l[l]), 
+                        model->head_size*model->max_encoder_context_length*ggml_element_size(cache->cross_v_l[l]), 
+                        model->head_size*model->max_encoder_context_length*model->decoder_attn_heads*ggml_element_size(cache->cross_v_l[l]),
+                        0));
+
+            // As noted in the encoder Dia uses the Neo-X protocol for RoPE.
+            cross_Qcur = ggml_rope(ctx, ggml_cont(ctx, ggml_reshape_4d(ctx, cross_Qcur, model->head_size, model->decoder_attn_heads, batch.sequence_length, 2)), dctx->positions, model->head_size, 2);
+            struct ggml_tensor * cross_q = ggml_cont(ctx, ggml_permute(ctx, cross_Qcur, 0, 2, 1, 3));
+            struct ggml_tensor * cross_kq = ggml_mul_mat(ctx, cross_k, cross_q);
+
+            // given that attention bias, scaling and masking are not used for decoding, it might be faster to prefer the #ggml_soft_max op here,
+            cross_kq = ggml_soft_max_ext(ctx, cross_kq, nullptr, 1.0f, 0.0f);
+            struct ggml_tensor * cross_kqv = ggml_mul_mat(ctx, cross_kq, cross_v);
+            struct ggml_tensor * cross_kqv_merged = ggml_cont(ctx, ggml_permute(ctx, cross_kqv, 2, 0, 1, 3));
+            cur = ggml_cont_3d(ctx, cross_kqv_merged, model->decoder_hidden_size, batch.sequence_length, 2);
+            cur = ggml_mul_mat(ctx, layer->cross_attn_o, cur);
+        }
+
+
+        // if we ever need to support multiple step decoder runs then this reshape will need to be replaced with permutation.
+        cur = ggml_cont_2d(ctx, cur, cur->ne[0], 2);
+        cur = ggml_add(ctx, cur, residual_cross);
+        struct ggml_tensor * residual_mlp = cur;
+
+        cur = dia_layer_norm(ctx, cur, layer->mlp_norm);
+        // mlp
+        {
+            cur = ggml_mul(ctx, ggml_silu(ctx, ggml_mul_mat(ctx, layer->gate, cur)), ggml_mul_mat(ctx, layer->up, cur));
+            cur = ggml_mul_mat(ctx, layer->out, cur);
+        }
+
+        cur = ggml_add(ctx, cur, residual_mlp);
+    }
+
+    cur = dia_layer_norm(ctx, cur, model->decoder->norm);
+    cur = build_dia_head_outputs(ctx, model, cur);
+    return cur;
+}
+
+void dia_runner::tokenize_sentence(std::string sentence, dia_ubatch & batch) {
+    // Dia's tokenization process is unusual. Essentially Dia takes the byte value for each character and uses that as 
+    // a token array. Additionally, because Dia performs a cfg-scale adjustment before sampling tokens, it is necessary to 
+    // generate with a conditioned context (i.e. with the text) and an unconditioned context (i.e. without any text) so that
+    // proper adjustments can be perfored at each generation step. This means that we need to pad the end of our tokens to the 
+    // max context size for both the conditional and unconditional sequence.
+
+    // if the sentence isn't prepended by dialogue start tokens, [S1] or [S2], then append one.
+    sentence = strip(sentence);
+    std::string start = sentence.substr(0, 4);
+    if (start != "[S1]" && start != "[S2]") {
+        sentence = "[S1] " + sentence;
+    }
+    if (sentence[sentence.size() - 1] != '.') {
+        sentence += ".";
+    }
+
+    // [S1] and [S2] are special character sequences that are replaced with the special tokens 0x01 and 0x02 respectively.
+    std::string r1(1, 1);
+    std::string r2(1, 2);
+    while (sentence.find("[S1]") != std::string::npos) {
+        size_t pos = sentence.find("[S1]");
+        sentence.replace(pos, 4, r1);
+    }
+    while (sentence.find("[S2]") != std::string::npos) {
+        size_t pos = sentence.find("[S2]");
+        sentence.replace(pos, 4, r2);
+    }
+
+    if (sentence.size() > model->max_encoder_context_length) {
+        TTS_ABORT("Dia currently only supports a max of %d characters and received an input of %d characters.", model->max_encoder_context_length, sentence.size());
+    }
+    batch.tokens.reserve(model->max_encoder_context_length * 2);
+    for (auto character : sentence) {
+        batch.tokens.push_back((uint32_t) character);
+    }
+    batch.sentence_length = batch.tokens.size();
+    // this 100 token warning is arbitrarily chosen based on spot checking small prompt performance
+    if (batch.sentence_length <= 100) {
+        fprintf(stdout, "Your prompt has fewer than 100 tokens. Please note that Dia's generation with prompts that are fewer than 100 tokens is highly inconsistent.\n");
+    }
+
+    for (int i = (int) batch.tokens.size(); i < model->max_encoder_context_length * 2; i++) {
+        batch.tokens.push_back(0u);
+    }
+ }
+
+dia_ubatch dia_runner::batch_from_sentence(std::string sentence) {
+    // if we are generating a new batch from tokens then we need to run the encoder step;
+    struct dia_ubatch batch{ 1, true};
+    tokenize_sentence(sentence, batch);
+    batch.audio_tokens.reserve(model->n_output_heads);
+    for (int i = 0; i < model->n_output_heads; i++) {
+        batch.audio_tokens.push_back(model->bos_token_id);
+    }
+    return batch;
+}
+
+/*
+ * There are two unique features of Dia's model architecture:
+ * 1.  Dia cleans its output generation by adding the difference between its text based output (its conditional output) and its unconditional output
+ *     to the conditional ouput before sampling. This is why the batch is set to two throughout the graph.
+ *
+ * 2.  Dia's decoder attends across the entire encoded space including the pad buffer which receives a unique attention mask. This is why the 
+ *     encoder sequence is always max length.
+ */
+struct ggml_cgraph * dia_runner::build_dia_graph(dia_ubatch & batch) {
+    init_build();
+    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
+    struct ggml_tensor * encoded_states = nullptr;
+
+    if (batch.encoder_step) {
+        encoded_states = build_dia_encoder(ctx, model, dctx, batch);
+        ggml_build_forward_expand(gf, encoded_states);
+    }
+
+    struct ggml_tensor * cur = build_dia_decoder(gf, ctx, model, dctx, kv_cross_self, batch, encoded_states);
+    ggml_set_name(cur, "decoder_output");
+    ggml_build_forward_expand(gf, cur);
+    free_build();
+    
+    return gf;
+}
+
+void dia_runner::configure_generation(generation_configuration * config) {
+    GGML_ASSERT(config->max_tokens == 0 || config->max_tokens > model->max_delay);
+    decode_sampler->temperature = config->temperature;
+    decode_sampler->repetition_penalty = config->repetition_penalty;
+    decode_sampler->do_sample = config->sample;
+    decode_sampler->top_k = config->top_k;
+    decode_sampler->top_p = config->top_p;
+    dctx->max_generation_size = config->max_tokens > model->max_delay ? config->max_tokens : model->max_generation_size;
+}
+
+void dia_runner::set_inputs(dia_ubatch & batch) {
+    if (batch.encoder_step) {
+        ggml_backend_tensor_set(dctx->inp_tokens, batch.tokens.data(), 0, batch.tokens.size()*ggml_element_size(dctx->inp_tokens));
+        int32_t * ep = (int32_t*) dctx->encode_positions->data;
+        float * mask = (float*) dctx->encode_attn_mask->data;
+        for (int i = 0; i < model->max_encoder_context_length; i++) {
+            ep[i] = (int32_t) i;
+            for (int ii = 0; ii < model->max_encoder_context_length; ii++) {
+                if (i < batch.sentence_length) {
+                    mask[i*model->max_encoder_context_length + ii] = ii < batch.sentence_length ? 0.0 : -INFINITY;
+                } else {
+                    mask[i*model->max_encoder_context_length + ii] = ii >= batch.sentence_length ? 0.0 : -INFINITY;
+                }
+            }
+        }
+    }
+    // The audio tokens need to be repeated in the input in order to support cfg-scaling. I.E we need duplicate inputs for conditional and unconditional logits.
+    ggml_backend_tensor_set(dctx->audio_inp_tokens, batch.audio_tokens.data(), 0, batch.audio_tokens.size()*ggml_element_size(dctx->audio_inp_tokens));
+    ggml_backend_tensor_set(dctx->audio_inp_tokens, batch.audio_tokens.data(), batch.audio_tokens.size()*ggml_element_size(dctx->audio_inp_tokens), batch.audio_tokens.size()*ggml_element_size(dctx->audio_inp_tokens));
+    ((int32_t*) dctx->positions->data)[0] = dctx->current_position;
+}
+
+int dia_runner::decode(dia_ubatch & batch) {
+    if (batch.encoder_step) {
+        dctx->prompt_size = batch.sentence_length;
+        dctx->output_tokens.reserve(dctx->max_generation_size * model->n_output_heads);
+    }
+    ggml_backend_sched_reset(dctx->sched);
+        
+    const size_t logits_size = model->output_vocab_size * dctx->max_generation_size * model->n_output_heads;
+    const size_t prev_size = dctx->buf_output ? ggml_backend_buffer_get_size(dctx->buf_output) : 0;
+    const size_t new_size  = logits_size * sizeof(float);
+    
+    if (!dctx->buf_output || prev_size < new_size) {
+        if (dctx->buf_output) {
+            ggml_backend_buffer_free(dctx->buf_output);
+            dctx->buf_output = nullptr;
+            dctx->logits = nullptr;
+        }
+
+        dctx->buf_output = ggml_backend_buft_alloc_buffer(dctx->backend_cpu_buffer, new_size);
+    }
+    
+    dctx->logits = (float *) ggml_backend_buffer_get_base(dctx->buf_output);
+
+    ggml_cgraph * gf = build_dia_graph(batch);
+
+    // the output is always the last tensor in the graph
+    struct ggml_tensor * res = gf->nodes[gf->n_nodes - 1];
+    std::string resname = ggml_get_name(res);
+    ggml_backend_sched_alloc_graph(dctx->sched, gf);
+
+    set_inputs(batch);
+
+    ggml_backend_sched_graph_compute_async(dctx->sched, gf);
+
+    float * logits_out = dctx->logits + dctx->current_position * model->output_vocab_size * model->n_output_heads;
+    dctx->get_ggml_node_data(res, logits_out, model->output_vocab_size * model->n_output_heads * sizeof(float));
+
+    // Reset state for the next token before backend sync, to allow the CPU activities in the reset to
+    // overlap with device computation.
+    ggml_backend_sched_reset(dctx->sched);
+
+    return 0;
+}
+
+dia_ubatch dia_runner::build_worst_case_batch()  {
+    struct dia_ubatch batch{ 1, true };
+    batch.tokens.resize(model->max_encoder_context_length * 2);
+    batch.audio_tokens.resize(model->n_output_heads);
+    return batch;
+}
+
+void dia_runner::prepare_post_load() {
+    dac_runner->prepare_post_load();
+    dia_kv_cache_init(kv_cross_self, model, dctx);
+    auto batch = build_worst_case_batch();
+    batch.sentence_length = model->max_encoder_context_length;
+    dctx->prompt_size = model->max_encoder_context_length;
+    auto gf = build_dia_graph(batch);
+    dctx->prep_schedule(gf);
+}
+
+bool dia_runner::check_stopping(dia_ubatch & batch) {
+    if (dctx->delay_steps == -1 && (batch.audio_tokens[0] == model->eos_token_id || dctx->current_position >= dctx->max_generation_size - model->max_delay)) {
+        dctx->delay_steps = model->max_delay;
+    }
+    
+    if (dctx->delay_steps > 0) {
+        int step_after_eos = model->max_delay - dctx->delay_steps;
+        for (int i = 0; i < model->delay_pattern.size(); i++) {
+            if (step_after_eos == model->delay_pattern[i]) {
+                batch.audio_tokens[i] = model->eos_token_id;
+            } else if (step_after_eos > model->delay_pattern[i]) {
+                batch.audio_tokens[i] = model->pad_token_id;
+            }
+        }
+        dctx->delay_steps -= 1;
+    }
+    return dctx->delay_steps == 0;
+}
+
+void dia_runner::adjust_output_tokens(std::vector<uint32_t> & output_tokens, std::vector<uint32_t> & filtered) {
+    // currently this is applying sliding window over the heads and filtering out bad tokens.
+    // If we convert the DAC model's quantizer layers to support by row + column embeddings then we will need to transpose
+    // the heads and the sequence here, but right now simplying using a strided view is more peformant.
+    size_t size = output_tokens.size();
+    filtered.reserve(size);
+    for (int i = 0; i < (size / model->n_output_heads) - model->max_delay; i++) {
+        bool skip_step = false;
+        for (int ii = 0; ii < model->n_output_heads; ii++) {
+            int next_index = i*model->n_output_heads+model->delay_pattern[ii]*model->n_output_heads+ii;
+            if (next_index > size || output_tokens[next_index] >= model->audio_vocab_size) {
+                skip_step = true;
+                break;
+            }
+        }
+        if (!skip_step) {
+            for (int ii = 0; ii < model->n_output_heads; ii++) {
+                int next_index = i*model->n_output_heads+model->delay_pattern[ii]*model->n_output_heads+ii;
+                filtered.push_back(output_tokens[next_index]);
+            }
+        }
+    }
+}
+
+int dia_runner::generate_from_batch(dia_ubatch & batch, struct tts_response * output) {
+    while (!check_stopping(batch)) {
+        int state = decode(batch);
+        if (state != 0) {
+            return state;
+        }
+        decode_sampler->sample(dctx->logits + dctx->current_position * model->n_output_heads * model->output_vocab_size, dctx->output_tokens);
+        dctx->current_position += batch.sequence_length;
+        batch = dia_ubatch{ 1 };
+        uint32_t * last_outputs = (dctx->output_tokens.data() + (int) dctx->output_tokens.size() - model->n_output_heads);
+        batch.audio_tokens.reserve(model->n_output_heads);
+        for (int i = 0; i < model->n_output_heads; i++) {
+            batch.audio_tokens.push_back(dctx->current_position > i ? last_outputs[i] : model->bos_token_id);
+        }
+    }
+
+    std::vector<uint32_t> filtered_output_tokens;
+    adjust_output_tokens(dctx->output_tokens, filtered_output_tokens);
+
+    dac_runner->run(filtered_output_tokens.data(), (int32_t) filtered_output_tokens.size() / model->n_output_heads, output);
+    return 0;
+}
+
+int dia_runner::generate(std::string sentence, struct tts_response * output) {
+    dia_ubatch batch = batch_from_sentence(sentence);
+    dctx->reset();
+    decode_sampler->reset();
+    dctx->current_position = 0;
+    if (!kv_cross_self) {
+        kv_cross_self = new dia_kv_cache;
+        if (!dia_kv_cache_init(kv_cross_self, model, dctx)) {
+            return 1;
+        }
+    }
+    return generate_from_batch(batch, output);
+}
+
+void dia_runner::assign_weight(std::string name, ggml_tensor * tensor) {
+    if (tensor->data == NULL) {
+        return;
+    }
+
+    if (name.size() == 0) {
+        // handles the top level meta tensor
+        return;
+    }
+
+    if (name.size() > 14 && name.substr(0, 14) == "audio_encoder.") {
+        dac_runner->model->assign_weight(name.substr(14), tensor);
+    } else {
+        model->assign_weight(name, tensor);
+    }   
+}

otherarch/ttscpp/src/dia_model.h

@@ -0,0 +1,206 @@
+#pragma once
+
+#include "dac_model.h"
+#include "sampler.h"
+
+struct dia_encoder_layer {
+    struct ggml_tensor * k;
+    struct ggml_tensor * q;
+    struct ggml_tensor * v;
+    struct ggml_tensor * o;
+    struct ggml_tensor * self_attn_norm;
+
+    struct ggml_tensor * gate;
+    struct ggml_tensor * up;
+    struct ggml_tensor * out;
+    struct ggml_tensor * mlp_norm;
+};
+
+struct dia_decoder_layer {
+    struct ggml_tensor * self_attn_k;
+    struct ggml_tensor * self_attn_q;
+    struct ggml_tensor * self_attn_v;
+    struct ggml_tensor * self_attn_o;
+    struct ggml_tensor * self_attn_norm;
+    
+    struct ggml_tensor * cross_attn_k;
+    struct ggml_tensor * cross_attn_q;
+    struct ggml_tensor * cross_attn_v;
+    struct ggml_tensor * cross_attn_o;
+    struct ggml_tensor * cross_attn_norm;
+
+    struct ggml_tensor * gate;
+    struct ggml_tensor * up;
+    struct ggml_tensor * out;
+    struct ggml_tensor * mlp_norm;
+
+    struct ggml_tensor * pad_attn_values;
+};
+
+struct dia_encoder {
+    struct ggml_tensor * norm;
+    struct ggml_tensor * embedding;
+    std::vector<dia_encoder_layer*> layers;
+};
+
+struct dia_decoder {
+    struct ggml_tensor * norm;
+    std::vector<struct ggml_tensor*> embds;
+    std::vector<struct ggml_tensor*> heads;
+    std::vector<dia_decoder_layer*> layers;
+};
+
+struct dia_model : tts_model {
+    // These default configurations are based on the default configuration for the Dia 1.68b param model.
+    uint32_t n_output_heads = 9;
+    uint32_t n_encoder_layers = 12;
+    uint32_t n_decoder_layers = 18;
+    uint32_t encoder_hidden_size = 1024;
+    uint32_t decoder_hidden_size = 2048;
+    uint32_t encoder_attn_heads = 16;
+    uint32_t decoder_attn_heads = 16;
+    uint32_t decoder_query_heads = 4;
+    uint32_t head_size = 128;
+    uint32_t eos_token_id = 1024;
+    uint32_t pad_token_id = 1025;
+    uint32_t bos_token_id = 1026;
+    uint32_t output_vocab_size = 1028;
+    uint32_t audio_vocab_size = 1024;
+    uint32_t max_generation_size = 3072;
+    uint32_t max_encoder_context_length = 1024;
+
+
+    float cfg_scale_data[2] = {3.0, 1024.0};
+    uint32_t max_delay = 15;
+    std::vector<uint32_t> delay_pattern = {0, 8, 9, 10, 11, 12, 13, 14, 15};
+
+    dia_encoder * encoder;
+    dia_decoder * decoder;
+    
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    void assign_to_encoder(std::vector<std::string> parts, struct ggml_tensor * tensor, std::string name);
+    void assign_to_decoder(std::vector<std::string> parts, struct ggml_tensor * tensor, std::string name);
+    void assign_to_encoder_layer(std::string part, dia_encoder_layer * layer, struct ggml_tensor * tensor);
+    void assign_to_decoder_layer(std::string part, dia_decoder_layer * layer, struct ggml_tensor * tensor);
+    void prep_constants(gguf_context * meta);
+    void prep_layers();
+    void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only) {
+        prep_constants(meta_ctx);
+        prep_layers();
+        tts_model::setup_from_file(meta_ctx, load_context, cpu_only, "dia", 1.30);
+    }
+};
+
+struct dia_context : runner_context {
+    dia_context(dia_model * model, int n_threads): runner_context(n_threads), model(model) {
+        max_generation_size = model->max_generation_size;
+    };
+
+    uint32_t current_position = 0;  // current position in the active sequence
+    int delay_steps           = -1; // the max remaining steps to take before terminating; is set after an eos token is seen on the first output channel
+    size_t prompt_size        = 0;
+
+    uint32_t max_generation_size; // this is set by the generation context or defaults to the config set on dia model.
+
+    std::vector<uint32_t> output_tokens;
+    struct dia_model * model;    
+    
+    struct ggml_tensor * inp_tokens;
+    struct ggml_tensor * audio_inp_tokens;
+    struct ggml_tensor * positions;
+    struct ggml_tensor * encode_positions;
+    struct ggml_tensor * encode_attn_mask;
+    struct ggml_tensor * cross_attn_mask;
+    
+    void build_schedule() {
+        runner_context::build_schedule(model->max_nodes());
+    }
+    void reset();
+};
+
+struct dia_kv_cache {
+    ggml_type tensor_type = GGML_TYPE_F32;
+
+    std::vector<struct ggml_tensor *> cross_k_l;
+    std::vector<struct ggml_tensor *> cross_v_l;
+
+    std::vector<struct ggml_tensor *> k_l;
+    std::vector<struct ggml_tensor *> 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);
+    }
+
+    ~dia_kv_cache() {
+        free();
+    }
+};
+
+struct dia_ubatch {
+    dia_ubatch(size_t sequence_length, bool encoder_step = false): sequence_length(sequence_length), encoder_step(encoder_step) {};
+    bool encoder_step; // whether we are performing the prompt encoding in this step.
+    size_t sequence_length; // for just audio tokens the sequence length should be the total_tokens / num_heads; for normal generation this should always be 1.
+    size_t sentence_length; // the number of non padded tokens in the conditional context
+    std::vector<uint32_t> tokens; // character tokens for the encoder
+    std::vector<uint32_t> audio_tokens; // audio tokens from the last generation
+};
+
+struct dia_context * build_new_dia_context(struct dia_model * model, int n_threads, bool use_cpu = true);
+static bool dia_kv_cache_init(struct dia_kv_cache * cache, dia_model * model, dia_context * dctx) ;
+static struct ggml_tensor * build_dia_decoder_inp_embd(struct ggml_context * ctx, dia_context *dctx, dia_decoder * decoder, dia_ubatch & batch, uint32_t n_output_heads);
+static struct ggml_tensor * dia_layer_norm(struct ggml_context * ctx, struct ggml_tensor * inputs, struct ggml_tensor * weight);
+static struct ggml_tensor * build_dia_encoder_attn_mask(ggml_context * ctx, struct dia_context * dctx, dia_model * model);
+static struct ggml_tensor * build_dia_decoder_attn_mask(ggml_context * ctx, struct dia_context * dctx, dia_ubatch & batch);
+static struct ggml_tensor * build_dia_decoder_cross_attn_mask(ggml_context * ctx, struct dia_context * dctx, dia_ubatch & batch);
+static struct ggml_tensor * build_dia_head_outputs(struct ggml_context * ctx, dia_model * model, struct ggml_tensor * cur);
+static struct ggml_tensor * build_dia_encoder(ggml_context * ctx, dia_model * model, dia_context * dctx, dia_ubatch & batch);
+static void build_dia_self_kv_store(ggml_context * ctx, dia_context * dctx, dia_model * model, dia_kv_cache * kv, ggml_cgraph * gf, struct ggml_tensor * k, struct ggml_tensor * v, dia_ubatch & batch, int layer_index);
+static void build_dia_cross_kv_store(ggml_context * ctx, dia_context * dctx, dia_model * model, dia_kv_cache * kv, ggml_cgraph * gf, struct ggml_tensor * encoder_hidden_states, int layer_index);
+static struct ggml_tensor * build_dia_decoder( ggml_cgraph * gf, ggml_context * ctx, dia_model * model,  dia_context * dctx,  dia_kv_cache * cache, dia_ubatch & batch, struct ggml_tensor * encoder_hidden_states);
+
+// This struct is intended to support end-to-end TTS generation for the Dia model. As such, it manages Dia's model compilation, compute, generation,
+// tokenizationm and sampling process, and uses the dac_runner struct to encode audio outputs.
+struct dia_runner : tts_runner {
+    dia_runner(dia_model * model, dac_runner * audio_decoder, dia_context * dctx, sampler * samp, dia_kv_cache * cache): model(model), dac_runner(audio_decoder), dctx(dctx), decode_sampler(samp), kv_cross_self(cache) {
+        decode_sampler->vocab_size = model->output_vocab_size;
+    };
+    ~dia_runner() {
+        if (ctx) {
+            ggml_free(ctx);
+        }
+        model->free();
+        delete model;
+        delete kv_cross_self;
+        delete dac_runner;
+        delete dctx;
+        delete decode_sampler;
+    }
+    struct dia_model * model;
+    struct dac_runner * dac_runner;
+    struct dia_context * dctx;
+    struct dia_kv_cache * kv_cross_self = nullptr;
+    struct sampler * decode_sampler;
+
+    void init_build() {
+        tts_runner::init_build(&dctx->buf_compute_meta);
+    }
+
+    void tokenize_sentence(std::string sentence, dia_ubatch & tokens);
+    dia_ubatch batch_from_sentence(std::string sentence);
+    void configure_generation(generation_configuration * config);
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    dia_ubatch build_worst_case_batch();
+    struct ggml_cgraph * build_dia_graph(dia_ubatch & batch);
+    void set_inputs(dia_ubatch & batch);
+    int decode(dia_ubatch & batch);
+    void prepare_post_load();
+    int generate(std::string sentence, struct tts_response * response);
+    bool check_stopping(dia_ubatch & batch);
+    void adjust_output_tokens(std::vector<uint32_t> & output_tokens, std::vector<uint32_t> & filtered);
+    int generate_from_batch(dia_ubatch & batch, struct tts_response * output);
+};

otherarch/ttscpp/src/general_neural_audio_codec.cpp

@@ -0,0 +1,172 @@
+#include "general_neural_audio_codec.h"
+#include <algorithm>
+#include <stdexcept>
+#include <map>
+
+namespace general_neural_audio_codec {
+    // This contains a mapping between string names and gguf_tensor enum values for the purposes of assigning the weights from a gguf file
+    // to the general_neural_audio_codec::layer.
+    // Please note that some gguf_tensor values have multiple keys; this is to support backwards compatibility with original DAC settings.
+    static const std::map<std::string, gguf_tensor> GENERAL_NEURAL_AUDIO_CODEC_TENSOR_LOOKUP = {
+        {".final.alpha", LAYER_ALPHA},
+        {".final.bias", LAYER_INPUT_BIAS},
+        {".final.weight", LAYER_INPUT_KERNEL},
+        {".alpha", LAYER_ALPHA},
+        {".bias", LAYER_INPUT_BIAS},
+        {".weight", LAYER_INPUT_KERNEL},
+        {".noise_weight", LAYER_NOISE_KERNEL},
+        {".res.initial.alpha", RESIDUAL_UNIT_INPUT_ALPHA},
+        {".res.initial.bias", RESIDUAL_UNIT_INPUT_BIAS},
+        {".res.initial.weight", RESIDUAL_UNIT_INPUT_KERNEL},
+        {".res.final.alpha", RESIDUAL_UNIT_OUTPUT_ALPHA},
+        {".res.final.bias", RESIDUAL_UNIT_OUTPUT_BIAS},
+        {".res.final.weight", RESIDUAL_UNIT_OUTPUT_KERNEL},
+        {".in_alpha", RESIDUAL_UNIT_INPUT_ALPHA},
+        {".in_bias", RESIDUAL_UNIT_INPUT_BIAS},
+        {".in_weight", RESIDUAL_UNIT_INPUT_KERNEL},
+        {".out_alpha", RESIDUAL_UNIT_OUTPUT_ALPHA},
+        {".out_bias", RESIDUAL_UNIT_OUTPUT_BIAS},
+        {".out_weight", RESIDUAL_UNIT_OUTPUT_KERNEL},
+        {".out_proj.bias", QUANTIZER_LAYER_OUT_BIAS},
+        {".out_proj.weight", QUANTIZER_LAYER_OUT_KERNEL},
+        {".codebook.weight", QUANTIZER_LAYER_CODEBOOK},
+    };
+
+    void assign_to_residual_unit(tts_model * model, residual_unit & unit, std::string name, struct ggml_tensor * tensor) {
+        try {
+            gguf_tensor tensor_type = GENERAL_NEURAL_AUDIO_CODEC_TENSOR_LOOKUP.at(name);
+            switch (tensor_type) {
+                case RESIDUAL_UNIT_INPUT_ALPHA:
+                    unit.in_alpha = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(unit.in_alpha, tensor);
+                    break;
+                case RESIDUAL_UNIT_OUTPUT_ALPHA:
+                    unit.out_alpha = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(unit.out_alpha, tensor);
+                    break;
+                case RESIDUAL_UNIT_INPUT_KERNEL:
+                    unit.in_conv_kernel = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(unit.in_conv_kernel, tensor);
+                    break;
+                case RESIDUAL_UNIT_OUTPUT_KERNEL:
+                    unit.out_conv_kernel = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(unit.out_conv_kernel, tensor);
+                    break;
+                case RESIDUAL_UNIT_INPUT_BIAS:
+                    unit.in_conv_bias = ggml_dup_tensor(model->ctx, ggml_transpose(model->ctx, tensor));
+                    model->set_tensor(unit.in_conv_bias, tensor);
+                    break;
+                case RESIDUAL_UNIT_OUTPUT_BIAS:
+                    unit.out_conv_bias = ggml_dup_tensor(model->ctx, ggml_transpose(model->ctx, tensor));
+                    model->set_tensor(unit.out_conv_bias, tensor);
+                    break;
+                default:
+                    fprintf(stdout, "residual unit unassigned tensor %s\n", name.c_str());
+                    break;
+            }
+        } catch (const std::out_of_range& e) {
+            TTS_ABORT("Tensor, '%s', is not a valid tensor general_neural_audio_codec::residual_unit tensor.", name.c_str());
+        }
+    }
+
+    void assign_to_layer(tts_model * model, layer & l, std::string name, struct ggml_tensor * tensor) {
+        if (GENERAL_NEURAL_AUDIO_CODEC_TENSOR_LOOKUP.find(name) != GENERAL_NEURAL_AUDIO_CODEC_TENSOR_LOOKUP.end()) {
+            switch(GENERAL_NEURAL_AUDIO_CODEC_TENSOR_LOOKUP.at(name)) {
+                case LAYER_ALPHA:
+                    l.in_alpha = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(l.in_alpha, tensor);
+                    break;
+                case LAYER_INPUT_KERNEL:
+                    l.in_conv_kernel = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(l.in_conv_kernel, tensor);
+                    break;
+                case LAYER_INPUT_BIAS:
+                    l.in_conv_bias = ggml_dup_tensor(model->ctx, ggml_transpose(model->ctx, tensor));
+                    model->set_tensor(l.in_conv_bias, tensor);
+                    break;
+                case LAYER_NOISE_KERNEL:
+                    l.noise_conv_kernel = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(l.noise_conv_kernel, tensor);
+                    break;
+                default:
+                    fprintf(stdout, "layer unassigned tensor %s\n", name.c_str());
+                    break;
+            }
+        } else if (std::find_if(name.begin(), name.end(), ::isdigit) != name.end())  {
+            auto pair = parse_layer_count(name);
+            int i = pair.first;
+            std::string lt_name = pair.second;
+            assign_to_residual_unit(model, l.residual_blocks[i], lt_name, tensor);
+        } else {
+            TTS_ABORT("Tensor, '%s', is not a valid tensor general_neural_audio_codec::layer tensor.", name.c_str());
+        }
+    }
+
+    void assign_to_quantize_layer(tts_model * model, residual_vector_quantize_layer & l, std::string name, struct ggml_tensor * tensor) {
+        try {
+            switch(GENERAL_NEURAL_AUDIO_CODEC_TENSOR_LOOKUP.at(name)) {
+                case QUANTIZER_LAYER_OUT_KERNEL:
+                    l.out_proj_kernel = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(l.out_proj_kernel, tensor);
+                    break;
+                case QUANTIZER_LAYER_OUT_BIAS:
+                    l.out_proj_bias = ggml_dup_tensor(model->ctx, ggml_transpose(model->ctx, tensor));
+                    model->set_tensor(l.out_proj_bias, tensor);
+                    break;
+                case QUANTIZER_LAYER_CODEBOOK:
+                    l.codebook = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(l.codebook, tensor);
+                    break;
+                default:
+                    fprintf(stdout, "quantized layer unassigned tensor %s\n", name.c_str());
+                    break;
+            }
+        } catch (const std::out_of_range& e) {
+            // older GGUF files still have the unused in_proj convolutional layer, so ignore it if we find it.
+            if (!has_prefix(name, ".in_proj")) {
+                TTS_ABORT("Error: %s\nTensor, '%s', is not a valid tensor.", e.what(), name.c_str());
+            }
+        }
+    }
+
+    struct ggml_tensor * build_residual_unit(ggml_context * ctx, struct ggml_tensor * cur, residual_unit & unit) {
+        struct ggml_tensor * residual = cur;
+        cur = snake_1d(ctx, unit.in_alpha, cur);
+        if (unit.groups > 1) {
+            // depthwise 1d convolution is equivalent to convolution in which grouping is equal to filter size.
+            // If there is a divergence between filter size and grouping then the kernel's output filters will not be zero.
+            TTS_ASSERT(unit.in_conv_kernel->ne[1] == 1);
+            cur = ggml_conv_1d_dw_tts(ctx, unit.in_conv_kernel, cur, 1, unit.padding, unit.dilation);
+        } else {
+            cur = ggml_conv_1d_tts(ctx, unit.in_conv_kernel, cur, 1, unit.padding, unit.dilation);
+        }
+        cur = ggml_add(ctx, cur, unit.in_conv_bias);
+        cur = snake_1d(ctx, unit.out_alpha, cur);
+        cur = ggml_conv_1d_tts(ctx, unit.out_conv_kernel, cur, 1, 0, 1);
+        cur = ggml_add(ctx, cur, unit.out_conv_bias);
+        return ggml_add(ctx, cur, residual);
+    }
+
+    struct ggml_tensor * build_layer(ggml_context * ctx, struct ggml_tensor * cur, layer & l, struct ggml_tensor * noise) {
+        cur = snake_1d(ctx, l.in_alpha, cur);
+        cur = ggml_conv_transpose_1d_tts(ctx, l.in_conv_kernel, cur, l.stride, l.padding, 1, 0, 1);
+        cur = ggml_add(ctx, cur, l.in_conv_bias);
+        if (l.noise_conv_kernel && noise) {
+            struct ggml_tensor * x = ggml_conv_1d_tts(ctx, l.noise_conv_kernel, cur, 1, 0, 1);
+            x = ggml_mul(ctx, x, noise);
+            cur = ggml_add(ctx, cur, x);
+        }
+        for (int i = 0; i < l.residual_blocks.size(); i++) {
+           cur = build_residual_unit(ctx, cur, l.residual_blocks[i]);
+        }
+        return cur;
+    }
+
+    struct ggml_tensor * build_quantize_layer(ggml_context * ctx, struct ggml_tensor * cur, residual_vector_quantize_layer & l) {
+        cur = ggml_get_rows(ctx, l.codebook, cur);
+        cur = ggml_cont(ctx, ggml_transpose(ctx, cur));
+        cur = ggml_conv_1d_tts(ctx, l.out_proj_kernel, cur, 1, 0, 1);
+        cur = ggml_add(ctx, cur, l.out_proj_bias);
+        return cur;
+    }
+}

otherarch/ttscpp/src/general_neural_audio_codec.h

@@ -0,0 +1,67 @@
+#pragma once
+
+#include "tts_model.h"
+
+// This namespace implements a general abstraction of the core functionality used in common neural audio codecs like DAC and SNAC.
+namespace general_neural_audio_codec {
+    enum gguf_tensor {
+        LAYER_ALPHA,
+        LAYER_INPUT_KERNEL,
+        LAYER_INPUT_BIAS,
+        LAYER_NOISE_KERNEL,
+        RESIDUAL_UNIT_INPUT_ALPHA,
+        RESIDUAL_UNIT_OUTPUT_ALPHA,
+        RESIDUAL_UNIT_INPUT_KERNEL,
+        RESIDUAL_UNIT_OUTPUT_KERNEL,
+        RESIDUAL_UNIT_INPUT_BIAS,
+        RESIDUAL_UNIT_OUTPUT_BIAS,
+        QUANTIZER_LAYER_OUT_KERNEL,
+        QUANTIZER_LAYER_OUT_BIAS,
+        QUANTIZER_LAYER_CODEBOOK
+    };
+
+    struct residual_vector_quantize_layer {
+        struct ggml_tensor * out_proj_kernel;
+        struct ggml_tensor * out_proj_bias;
+        struct ggml_tensor * codebook;
+    };
+
+    struct residual_unit {
+        residual_unit(uint32_t padding, uint32_t dilation, uint32_t groups = 1): padding(padding), dilation(dilation), groups(groups) {}
+        struct ggml_tensor * in_alpha;
+        struct ggml_tensor * in_conv_kernel;
+        struct ggml_tensor * in_conv_bias;
+        struct ggml_tensor * out_alpha;
+        struct ggml_tensor * out_conv_kernel;
+        struct ggml_tensor * out_conv_bias;
+
+        uint32_t padding;
+        uint32_t dilation;
+        uint32_t groups;
+    };
+
+    struct layer {
+        layer(uint32_t padding, uint32_t stride, uint32_t groups = 1): padding(padding), stride(stride) {
+            for (int i = 0; i < 3; i++) {
+                residual_blocks.push_back(residual_unit{(uint32_t) pow(3, (i + 1)), (uint32_t) pow(3, i), groups});
+            }
+        }
+        struct ggml_tensor * in_alpha;
+        struct ggml_tensor * in_conv_kernel;
+        struct ggml_tensor * in_conv_bias;
+        struct ggml_tensor * noise_conv_kernel = nullptr;
+
+        uint32_t padding;
+        uint32_t stride;
+        
+        std::vector<residual_unit> residual_blocks;
+    };
+
+    void assign_to_residual_unit(tts_model * model, residual_unit & unit, std::string name, struct ggml_tensor * tensor);
+    void assign_to_layer(tts_model * model, layer & l, std::string name, struct ggml_tensor * tensor);
+    void assign_to_quantize_layer(tts_model * model, residual_vector_quantize_layer & l, std::string name, struct ggml_tensor * tensor);
+
+    struct ggml_tensor * build_residual_unit(ggml_context * ctx, struct ggml_tensor * cur, residual_unit & unit);
+    struct ggml_tensor * build_layer(ggml_context * ctx, struct ggml_tensor * cur, layer & l, struct ggml_tensor * noise = nullptr);
+    struct ggml_tensor * build_quantize_layer(ggml_context * ctx, struct ggml_tensor * cur, residual_vector_quantize_layer & l);
+}

otherarch/ttscpp/src/kokoro_model.h

@@ -0,0 +1,462 @@
+#ifndef kokoro_model_h
+#define kokoro_model_h
+
+#include <stdlib.h>
+#include "tts_model.h"
+#include "tokenizer.h"
+#include "phonemizer.h"
+
+// Rather than using ISO 639-2 language codes, Kokoro voice pack specify their corresponding language via their first letter.
+// Below is a map that describes the relationship between those designations and espeak-ng's voice identifiers so that the 
+// appropriate phonemization protocol can inferred from the Kokoro voice.
+static std::map<char, std::string> KOKORO_LANG_TO_ESPEAK_ID = {
+	{'a', "gmw/en-US"},
+	{'b', "gmw/en"},
+	{'e', "roa/es"},
+	{'f', "roa/fr"},
+	{'h', "inc/hi"},
+	{'i', "roa/it"},
+	{'j', "jpx/ja"},
+	{'p', "roa/pt-BR"},
+	{'z', "sit/cmn"}
+};
+
+struct lstm_cell {
+	std::vector<ggml_tensor*> weights; 
+	std::vector<ggml_tensor*> biases;
+	std::vector<ggml_tensor*> reverse_weights;
+	std::vector<ggml_tensor*> reverse_biases;
+};
+
+struct lstm {
+	std::vector<ggml_tensor*> hidden;
+	std::vector<ggml_tensor*> states;
+
+	bool bidirectional = false;
+	std::vector<lstm_cell*> cells;
+};
+
+struct duration_predictor_layer {
+	lstm * rnn;
+	struct ggml_tensor * ada_norm_gamma_weight;
+	struct ggml_tensor * ada_norm_gamma_bias;
+	struct ggml_tensor * ada_norm_beta_weight;
+	struct ggml_tensor * ada_norm_beta_bias;
+};
+
+struct ada_residual_conv_block {
+	struct ggml_tensor * conv1;
+	struct ggml_tensor * conv1_bias;
+	struct ggml_tensor * conv2;
+	struct ggml_tensor * conv2_bias;
+	struct ggml_tensor * norm1_gamma;
+	struct ggml_tensor * norm1_gamma_bias;
+	struct ggml_tensor * norm1_beta;
+	struct ggml_tensor * norm1_beta_bias;
+	struct ggml_tensor * norm2_gamma;
+	struct ggml_tensor * norm2_gamma_bias;
+	struct ggml_tensor * norm2_beta;
+	struct ggml_tensor * norm2_beta_bias;
+	struct ggml_tensor * pool = nullptr;
+	struct ggml_tensor * pool_bias = nullptr;
+	struct ggml_tensor * upsample = nullptr;
+	struct ggml_tensor * upsample_bias = nullptr;
+};
+
+struct duration_predictor {
+	struct ggml_tensor * albert_encode;
+	struct ggml_tensor * albert_encode_bias;
+	std::vector<duration_predictor_layer*> layers;
+	lstm * duration_proj_lstm;
+	struct ggml_tensor * duration_proj;
+	struct ggml_tensor * duration_proj_bias;
+	struct ggml_tensor * n_proj_kernel;
+	struct ggml_tensor * n_proj_bias;
+	struct ggml_tensor * f0_proj_kernel;
+	struct ggml_tensor * f0_proj_bias;
+	lstm * shared_lstm;
+	std::vector<ada_residual_conv_block*> f0_blocks;
+	std::vector<ada_residual_conv_block*> n_blocks;
+};
+
+struct kokoro_text_encoder_conv_layer {
+	struct ggml_tensor * norm_gamma;
+	struct ggml_tensor * norm_beta;
+	struct ggml_tensor * conv_weight;
+	struct ggml_tensor * conv_bias;
+};
+
+struct kokoro_text_encoder {
+	struct ggml_tensor * embd;
+	std::vector<kokoro_text_encoder_conv_layer*> conv_layers;
+	lstm * out_lstm;
+};
+
+struct kokoro_generator_residual_block {
+	std::vector<uint32_t> conv1_dilations;
+	std::vector<uint32_t> conv1_paddings;
+
+	std::vector<ggml_tensor*> adain1d_1_gamma_weights;
+	std::vector<ggml_tensor*> adain1d_2_gamma_weights;
+	std::vector<ggml_tensor*> adain1d_1_gamma_biases;
+	std::vector<ggml_tensor*> adain1d_2_gamma_biases;
+	std::vector<ggml_tensor*> adain1d_1_beta_weights;
+	std::vector<ggml_tensor*> adain1d_2_beta_weights;
+	std::vector<ggml_tensor*> adain1d_1_beta_biases;
+	std::vector<ggml_tensor*> adain1d_2_beta_biases;
+	std::vector<ggml_tensor*> input_alphas;
+	std::vector<ggml_tensor*> output_alphas;
+	std::vector<ggml_tensor*> convs1_weights;
+	std::vector<ggml_tensor*> convs1_biases;
+	std::vector<ggml_tensor*> convs2_weights;
+	std::vector<ggml_tensor*> convs2_biases;
+};
+
+struct kokoro_noise_residual_block {
+	uint32_t input_conv_stride;
+	uint32_t input_conv_padding;
+
+	struct ggml_tensor * input_conv;
+	struct ggml_tensor * input_conv_bias;
+	struct kokoro_generator_residual_block * res_block;
+};
+
+struct kokoro_generator_upsample_block {
+	uint32_t padding;
+	uint32_t stride;
+
+	// these are just conv transpose layers
+	struct ggml_tensor * upsample_weight;
+	struct ggml_tensor * upsample_bias;
+};
+
+struct kokoro_generator {
+	// unfortunately the squared sum of the windows needs to be computed dynamically per run because it is dependent
+	// on the sequence size of the generation and the hop is typically less than half the size of our window.
+	struct ggml_tensor * window;
+
+	struct ggml_tensor * m_source_weight;
+	struct ggml_tensor * m_source_bias;
+	struct ggml_tensor * out_conv_weight;
+	struct ggml_tensor * out_conv_bias;
+	std::vector<kokoro_noise_residual_block*> noise_blocks;
+	std::vector<kokoro_generator_residual_block*> res_blocks;
+	std::vector<kokoro_generator_upsample_block*> ups;
+};
+
+struct kokoro_decoder {
+	struct ggml_tensor * f0_conv;
+	struct ggml_tensor * f0_conv_bias;
+	struct ggml_tensor * n_conv;
+	struct ggml_tensor * n_conv_bias;
+	struct ggml_tensor * asr_conv;
+	struct ggml_tensor * asr_conv_bias;
+	std::vector<ada_residual_conv_block*> decoder_blocks;
+	ada_residual_conv_block* encoder_block;
+	kokoro_generator * generator;
+};
+
+struct albert_layer {
+	struct ggml_tensor * ffn;
+	struct ggml_tensor * ffn_out;
+	struct ggml_tensor * ffn_bias;
+	struct ggml_tensor * ffn_out_bias;
+	struct ggml_tensor * layer_output_norm_weight;
+	struct ggml_tensor * layer_output_norm_bias;
+	struct ggml_tensor * q;
+	struct ggml_tensor * k;
+	struct ggml_tensor * v;
+	struct ggml_tensor * o;
+	struct ggml_tensor * q_bias;
+	struct ggml_tensor * k_bias;
+	struct ggml_tensor * v_bias;
+	struct ggml_tensor * o_bias;
+	struct ggml_tensor * attn_norm_weight;
+	struct ggml_tensor * attn_norm_bias;
+};
+
+struct kokoro_model : tts_model {
+	// standard configruation for Kokoro's Albert model
+	// tokenization
+	uint32_t bos_token_id = 0;
+	uint32_t eos_token_id = 0;
+	uint32_t space_token_id = 16;
+	// duration prediction
+	uint32_t max_context_length = 512;
+	uint32_t vocab_size = 178;
+	uint32_t hidden_size = 768;
+	uint32_t n_attn_heads = 12;
+	uint32_t n_layers = 1;
+	uint32_t n_recurrence = 12;
+	uint32_t head_size = 64;
+	uint32_t duration_hidden_size = 512;
+	uint32_t up_sampling_factor;
+	float upsample_scale = 300.0f;
+	float scale = 0.125f;
+
+	// standard configuration for duration prediction
+	uint32_t f0_n_blocks = 3;
+	uint32_t n_duration_prediction_layers = 3;
+	// while it is technically possible for the duration predictor to assign 50 values per token there is no practical need to 
+	// allocate that many items to the sequence as it is impossible for all tokens to require such long durations and each 
+	// allocation increases node allocation size by O(N)
+	uint32_t max_duration_per_token = 20;
+	uint32_t style_half_size = 128;
+
+	// standard text encoding configuration
+	uint32_t n_conv_layers = 3;
+
+	// standard decoder configuration
+	uint32_t n_kernels = 3;
+	uint32_t n_upsamples = 2;
+	uint32_t n_decoder_blocks = 4;
+	uint32_t n_res_blocks = 6;
+	uint32_t n_noise_blocks = 2;
+	uint32_t out_conv_padding = 3;
+	uint32_t post_n_fft = 11;
+	uint32_t true_n_fft = 20;
+	uint32_t stft_hop = 5;
+	uint32_t harmonic_num = 8;
+	float sin_amp = 0.1f;
+	float noise_std = 0.003f;
+	float voice_threshold = 10.0f;
+	float sample_rate = 24000.0f;
+	std::string window = "hann"; 
+
+	// It is really annoying that ggml doesn't allow using non ggml tensors as the operator for simple math ops.
+	// This is just the constant defined above as a tensor.
+	struct ggml_tensor * n_kernels_tensor;
+
+	// Kokoro loads albert with use_pooling = true but doesn't use the pooling outputs.
+	bool uses_pooling = false;
+	bool static_token_types = true;
+
+	std::map<std::string, struct ggml_tensor *> voices;
+
+	// Albert portion of the model
+	struct ggml_tensor * embd_hidden;
+	struct ggml_tensor * embd_hidden_bias;
+	struct ggml_tensor * token_type_embd = nullptr;
+	struct ggml_tensor * token_embd;
+	struct ggml_tensor * position_embd;
+	struct ggml_tensor * input_norm_weight;
+	struct ggml_tensor * input_norm_bias;
+	struct ggml_tensor * static_token_type_values = nullptr;
+	struct ggml_tensor * pool = nullptr;
+	struct ggml_tensor * pool_bias = nullptr;
+	std::vector<albert_layer*> layers;
+
+	struct ggml_tensor * harmonic_sampling_norm = nullptr; // a static 1x9 harmonic multiplier
+	struct ggml_tensor * sampling_factor_scalar = nullptr; // a static scalar
+	struct ggml_tensor * sqrt_tensor = nullptr; // static tensor for constant division
+
+	// Prosody Predictor portion of the model
+	struct duration_predictor * prosody_pred;
+
+	// Text encoding portion of the model
+	struct kokoro_text_encoder * text_encoder;
+
+	// Decoding and Generation portion of the model
+	struct kokoro_decoder * decoder;
+
+	// the default hidden states need to be initialized 
+	std::vector<lstm*> lstms;
+
+	size_t duration_node_counter = 0;
+	size_t generation_node_counter = 0;
+	// setting this is likely unnecessary as it is precomputed by the post load function.
+	uint32_t post_load_tensor_bytes = 13000;
+
+	size_t max_gen_nodes();
+	size_t max_duration_nodes();
+
+	lstm * prep_lstm();
+	// helper functions for assigning tensors to substructs
+	void assign_lstm(lstm * rnn, std::string name, ggml_tensor * tensor);
+	void assign_generator_weight(kokoro_generator * generator, std::string name, ggml_tensor * tensor);
+	void assign_gen_resblock(kokoro_generator_residual_block * block, std::string name, ggml_tensor * tensor);
+	void assign_ada_res_block(ada_residual_conv_block * block, std::string name, ggml_tensor * tensor);
+	void assign_decoder_weight(std::string name, ggml_tensor * tensor);
+	void assign_duration_weight(std::string name, ggml_tensor * tensor);
+	void assign_text_encoder_weight(std::string name, ggml_tensor * tensor);
+	void assign_albert_weight(std::string name, ggml_tensor * tensor);
+
+
+	void post_load_assign();
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    void prep_layers(gguf_context * meta);
+    void prep_constants(gguf_context * meta);
+    void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only = true) {
+    	std::function<void (ggml_tensor *)> fn = ([&](ggml_tensor* cur) {
+    		std::string name = ggml_get_name(cur);
+    		size_t increment = 1;
+    		if (name.find("lstm") != std::string::npos) {
+    			increment = max_context_length;
+    		}
+    		if (name.find("duration_predictor") != std::string::npos) {
+    			duration_node_counter += increment;
+    		} else {
+    			generation_node_counter += increment;
+    		}
+    	});
+    	compute_tensor_meta_cb = &fn;
+        prep_constants(meta_ctx);
+        prep_layers(meta_ctx);
+        tts_model::setup_from_file(meta_ctx, load_context, cpu_only, "kokoro", 1.6, post_load_tensor_bytes);
+    }
+};
+
+struct kokoro_ubatch {
+    size_t n_tokens; // the number of tokens in our encoded sequence
+    uint32_t * input_tokens;    // [n_tokens]
+    struct kokoro_duration_response * resp = nullptr;
+};
+
+struct kokoro_duration_context : runner_context {
+    kokoro_duration_context(kokoro_model * model, int n_threads): runner_context(n_threads), model(model) {};
+    ~kokoro_duration_context() {
+        ggml_backend_buffer_free(buf_len_output);
+    }
+    
+    std::string voice = "af_alloy";
+    struct kokoro_model * model;
+    ggml_backend_buffer_t buf_len_output = nullptr;
+
+    
+    size_t  logits_size = 0; // capacity (of floats) for logits
+    float * lens 		= nullptr;
+    
+    struct ggml_tensor * inp_tokens;
+    struct ggml_tensor * positions;
+    struct ggml_tensor * attn_mask;
+    struct ggml_tensor * token_types = nullptr;
+
+    void build_schedule() {
+        runner_context::build_schedule(model->max_duration_nodes()*5);
+    }
+};
+
+static struct ggml_tensor * build_albert_attn_mask(ggml_context * ctx, struct kokoro_duration_context *kctx, const kokoro_ubatch & batch);
+static struct ggml_tensor * build_albert_inputs(ggml_context * ctx, kokoro_model * model, ggml_tensor * input_tokens, ggml_tensor * positions, ggml_tensor * token_types);
+static struct ggml_tensor * build_albert_norm(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * weight, ggml_tensor * bias);
+static struct ggml_tensor * build_ada_residual_conv(ggml_context * ctx, struct ggml_tensor * x, ada_residual_conv_block * block, struct ggml_tensor * style, struct ggml_tensor * sqrt_tensor);
+static struct ggml_tensor * build_kokoro_generator_res_block(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * style, kokoro_generator_residual_block * block);
+static struct ggml_tensor * build_noise_block(ggml_context * ctx, kokoro_noise_residual_block * block, struct ggml_tensor * x, struct ggml_tensor * style);
+static kokoro_generator_residual_block * build_res_block_from_file(gguf_context * meta, std::string base_config_key);
+static kokoro_noise_residual_block * build_noise_block_from_file(gguf_context * meta, int index);
+static kokoro_generator_upsample_block* kokoro_generator_upsample_block(gguf_context * meta, int index);
+
+std::string get_espeak_id_from_kokoro_voice(std::string voice);
+struct kokoro_duration_context * build_new_duration_kokoro_context(struct kokoro_model * model, int n_threads, bool use_cpu = true);
+
+struct kokoro_duration_response {
+	size_t n_outputs;
+	float * lengths;
+	float * hidden_states;
+};
+
+// This struct is intended to manage graph and compute for the duration prediction portion of the kokoro model.
+// Duration computation and speech generation are separated into distinct graphs because the precomputed graph structure of ggml doesn't 
+// support the tensor dependent views that would otherwise be necessary.
+struct kokoro_duration_runner : tts_runner {
+    kokoro_duration_runner(kokoro_model * model, kokoro_duration_context * context, single_pass_tokenizer * tokenizer): model(model), kctx(context), tokenizer(tokenizer) {};
+    ~kokoro_duration_runner() {
+        if (ctx) {
+            ggml_free(ctx);
+        }
+        model->free();
+        delete model;
+        delete kctx;
+    }
+    struct single_pass_tokenizer * tokenizer;
+    kokoro_model * model;
+    kokoro_duration_context * kctx;
+
+    void init_build() {
+        tts_runner::init_build(&kctx->buf_compute_meta);
+    }
+    
+    void prepare_post_load();
+    struct kokoro_ubatch build_worst_case_batch();
+    void set_inputs(kokoro_ubatch & batch);
+    struct ggml_cgraph * build_kokoro_duration_graph(kokoro_ubatch & batch);
+    void run(kokoro_ubatch & ubatch);
+};
+
+struct kokoro_context : runner_context {
+    kokoro_context(kokoro_model * model, int n_threads): runner_context(n_threads), model(model) {};
+    ~kokoro_context() {
+        ggml_backend_sched_free(sched);
+        ggml_backend_free(backend_cpu);
+        if (backend) {
+            ggml_backend_free(backend);
+        }
+        if (buf_output) {
+            ggml_backend_buffer_free(buf_output);
+        }
+    }
+
+    std::string voice = "af_alloy";
+    
+    struct kokoro_model * model;
+
+    uint32_t total_duration;
+    uint32_t sequence_length;
+
+    struct ggml_tensor * inp_tokens;
+    struct ggml_tensor * duration_pred;
+    struct ggml_tensor * duration_mask;
+    struct ggml_tensor * window_sq_sum; // needs to be calculatd from the generator window.
+    struct ggml_tensor * uv_noise_data;
+    
+    void build_schedule() {
+        runner_context::build_schedule(model->max_gen_nodes()*30);
+    }
+};
+
+// TODO: now that we are passing the context down to these methods we should clean up their parameters
+static struct ggml_tensor * build_generator(ggml_context * ctx, kokoro_model * model, kokoro_context * kctx, struct ggml_tensor * x, struct ggml_tensor * style, struct ggml_tensor * f0_curve, kokoro_generator* generator, int sequence_length, struct ggml_tensor * window_sq_sum, ggml_cgraph * gf);
+static struct ggml_tensor * build_sin_gen(ggml_context * ctx, kokoro_model * model, kokoro_context * kctx, struct ggml_tensor * x, int harmonic_num, int sequence_length, float voice_threshold, float sin_amp, float noise_std);
+
+struct kokoro_context * build_new_kokoro_context(struct kokoro_model * model, int n_threads, bool use_cpu = true);
+
+// This manages the graph compilation of computation for the Kokoro model.
+struct kokoro_runner : tts_runner {
+    kokoro_runner(kokoro_model * model, kokoro_context * context, single_pass_tokenizer * tokenizer, kokoro_duration_runner * drunner, phonemizer * phmzr): model(model), kctx(context), tokenizer(tokenizer), drunner(drunner), phmzr(phmzr) {
+    	tts_runner::sampling_rate = 24000.0f;
+    	tts_runner::supports_voices = true;
+    };
+    ~kokoro_runner() {
+        if (ctx) {
+            ggml_free(ctx);
+        }
+        delete drunner;
+        model->free();
+        delete model;
+        delete kctx;
+        delete phmzr;
+    }
+    struct single_pass_tokenizer * tokenizer;
+    kokoro_model * model;
+    kokoro_context * kctx;
+    kokoro_duration_runner * drunner;
+    phonemizer * phmzr;
+
+    std::string default_voice = "af_alloy";
+
+    void init_build() {
+        tts_runner::init_build(&kctx->buf_compute_meta);
+    }
+
+    std::vector<std::string> list_voices();
+    std::vector<std::vector<uint32_t>> tokenize_chunks(std::vector<std::string> clauses);
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    void prepare_post_load();
+    kokoro_ubatch build_worst_case_batch();
+    void set_inputs(kokoro_ubatch & batch, uint32_t total_size);
+    struct ggml_cgraph * build_kokoro_graph(kokoro_ubatch & batch);
+    void run(kokoro_ubatch & batch, struct tts_response * outputs);
+    int generate(std::string prompt, struct tts_response * response, std::string voice, std::string voice_code = "");
+};
+
+#endif

otherarch/ttscpp/src/orpheus_model.cpp

@@ -0,0 +1,475 @@
+#include "orpheus_model.h"
+
+#include <array>
+
+// These tokens and variables aren't defined in the Orpheus' model configuration but instead are defined inline in various python functions.
+// As such, they are not discoverable so defining them as unconfigurable constants should be fine.
+static constexpr std::array<const char *, 7> orpheus_voices{"zoe", "zac","jess", "leo", "mia", "julia", "leah"};
+static constexpr std::array<uint32_t, 2> orpheus_prepended_tokens = { 128259, 128000 };
+static constexpr std::array<uint32_t, 4> orpheus_appended_tokens = { 128009, 128260, 128261, 128257 };
+
+void orpheus_model::assign_weight(std::string name, struct ggml_tensor * tensor) {
+    if (name == "norm") {
+        output_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(output_norm, tensor);
+    } else if (name == "lm_head") {
+        head = ggml_dup_tensor(ctx, tensor);
+        set_tensor(head, tensor);
+    } else if (name == "embed_tokens") {
+        embd = ggml_dup_tensor(ctx, tensor);
+        set_tensor(embd, tensor);
+    } else if (name == "rope_frequencies") {
+        rope_frequencies = ggml_dup_tensor(ctx, tensor);
+        set_tensor(rope_frequencies, tensor);
+    } else if (has_prefix(name, "layers")) {
+        auto lpair = parse_layer_count(name);
+        int l = lpair.first;
+        std::string lt_name = lpair.second;
+        assign_to_layer(lt_name, layers[l], tensor);
+    }
+}
+
+void orpheus_model::assign_to_layer(std::string part, orpheus_layer & layer, struct ggml_tensor * tensor) {
+    if (part == ".self_attn.k_proj") {
+        layer.k = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.k, tensor);
+    } else if (part == ".self_attn.q_proj") {
+        layer.q = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.q, tensor);
+    } else if (part == ".self_attn.v_proj") {
+        layer.v = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.v, tensor);
+    } else if (part == ".self_attn.o_proj") {
+        layer.o = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.o, tensor);
+    } else if (part == ".mlp.gate_proj") {
+        layer.gate = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.gate, tensor);
+    } else if (part == ".mlp.up_proj") {
+        layer.up = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.up, tensor);
+    } else if (part == ".mlp.down_proj") {
+        layer.down = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.down, tensor);
+    } else if (part == ".input_layernorm") {
+        layer.input_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.input_norm, tensor);
+    } else if (part == ".post_attention_layernorm") {
+        layer.post_attention_norm = ggml_dup_tensor(ctx, tensor);
+        set_tensor(layer.post_attention_norm, tensor);
+    }
+}
+
+void orpheus_model::prep_constants(gguf_context * meta) {
+    // get constants for orpheus
+    int vocab_size_key = gguf_find_key(meta, "orpheus.vocab_size");
+    if (vocab_size_key != -1) {
+        vocab_size = gguf_get_val_u32(meta, vocab_size_key);
+    }
+
+    int attn_heads_key = gguf_find_key(meta, "orpheus.attn_heads");
+    if (attn_heads_key != -1) {
+        n_attn_heads = gguf_get_val_u32(meta, attn_heads_key);
+    }
+
+    int kv_attn_heads_key = gguf_find_key(meta, "orpheus.kv_attn_heads");
+    if (kv_attn_heads_key != -1) {
+        n_kv_attn_heads = gguf_get_val_u32(meta, kv_attn_heads_key);
+    }
+
+    int head_size_key = gguf_find_key(meta, "orpheus.head_dim");
+    if (head_size_key != -1) {
+        head_size = gguf_get_val_u32(meta, head_size_key);
+    }
+
+    int stopping_token_key = gguf_find_key(meta, "orpheus.stopping_token_id");
+    if (stopping_token_key != -1) {
+        stopping_token_id = gguf_get_val_u32(meta, stopping_token_key);;
+    }
+
+    int eos_token_id_key = gguf_find_key(meta, "tokenizer.ggml.eos_token_id");
+    if (eos_token_id_key != -1) {
+        eos_token_id = gguf_get_val_u32(meta, eos_token_id_key);
+    }
+
+    int bos_token_id_key = gguf_find_key(meta, "tokenizer.ggml.bos_token_id");
+    if (bos_token_id_key != -1) {
+        bos_token_id = gguf_get_val_u32(meta, bos_token_id_key);
+    }
+
+    int hidden_size_key = gguf_find_key(meta, "orpheus.hidden_size");
+    if (hidden_size_key != -1) {
+        hidden_size = gguf_get_val_u32(meta, hidden_size_key);
+    }
+
+    int kv_hidden_size_key = gguf_find_key(meta, "orpheus.kv_hidden_size");
+    if (kv_hidden_size_key != -1) {
+        kv_hidden_size = gguf_get_val_u32(meta, kv_hidden_size_key);
+    }
+}
+
+void orpheus_model::prep_layers(gguf_context * meta) {
+    int n_layers_key = gguf_find_key(meta, "orpheus.layers");
+    if (n_layers_key == -1) {
+        TTS_ABORT("the 'orpheus.layers' must be specified in the GGUF file.");
+    }
+    n_layers = (int) gguf_get_val_u32(meta, n_layers_key);
+    for (int i = 0; i < n_layers; i++) {
+        layers.push_back(orpheus_layer{});
+    }
+}
+
+struct ggml_tensor * orpheus_build_layer_norm(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * weight) {
+    float eps = 0.00001;
+    return ggml_mul(ctx, ggml_rms_norm(ctx, x, eps), weight);
+}
+
+struct ggml_tensor * build_attn_mask(ggml_context * ctx, orpheus_context * octx, orpheus_ubatch & batch) {
+    octx->attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) octx->current_position + batch.n_tokens, (int64_t) octx->current_position + batch.n_tokens);
+    ggml_set_input(octx->attn_mask);
+    return octx->attn_mask;
+}
+
+ void orpheus_context::reset() {
+    output_tokens.clear();
+    current_position = 0;
+    n_outputs = 0;
+ }
+
+orpheus_context * build_new_orpheus_context(orpheus_model * model, int n_threads, bool use_cpu) {
+    orpheus_context * octx = new orpheus_context(model, n_threads);
+    if (!use_cpu) {
+#ifdef GGML_USE_METAL
+        octx->backend = ggml_backend_metal_init();
+#endif
+    }
+    octx->backend_cpu = ggml_backend_cpu_init();
+    octx->set_threads();
+    octx->build_schedule();
+    octx->buf_compute_meta.resize(ggml_tensor_overhead()*model->max_nodes() + ggml_graph_overhead_custom(model->max_nodes(), false));
+    return octx;
+}
+
+void orpheus_runner::orpheus_kv_cache_init() {    
+    ggml_backend_buffer_type_t buft = nullptr;
+    if (octx->backend != nullptr) {
+#ifdef GGML_USE_METAL
+        buft = ggml_backend_metal_buffer_type();
+#endif
+    } else {
+        buft = ggml_backend_cpu_buffer_type();
+    }
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ (2u * model->layers.size() + 1)*ggml_tensor_overhead(),
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true,
+    };
+    ggml_context * ctx = ggml_init(params);
+    if (!ctx) {
+        TTS_ABORT("%s: failed to initialze ggml context for key value cache.\n", __func__);
+    }
+    if (!kv_self) {
+        kv_self = new orpheus_kv_cache;
+    }
+    kv_self->ctx = ctx;
+    kv_self->k_l.reserve(model->layers.size());
+    kv_self->v_l.reserve(model->layers.size());
+
+    for (int i = 0; i < (int) model->layers.size(); i++) {
+        ggml_tensor * k = ggml_new_tensor_1d(kv_self->ctx, kv_self->cache_type, model->hidden_size * (model->max_context_length + model->max_generation_size));
+        ggml_tensor * v = ggml_new_tensor_1d(kv_self->ctx, kv_self->cache_type, model->hidden_size * (model->max_context_length + model->max_generation_size));
+        ggml_format_name(k, "cache_k_l%d", i);
+        ggml_format_name(v, "cache_v_l%d", i);
+        kv_self->k_l.push_back(k);
+        kv_self->v_l.push_back(v);
+    }
+
+    // allocate tensors and initialize the buffers to avoid NaNs in the padding
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(kv_self->ctx, buft);
+    ggml_backend_buffer_clear(buf, 0);
+    kv_self->buf = buf;
+ }
+
+ void orpheus_runner::orpheus_build_kv_store(struct ggml_context * ctx, struct ggml_cgraph * graph, struct ggml_tensor * k_cur, struct ggml_tensor * v_cur, int index, uint32_t n_tokens, int repeat) {
+    k_cur = ggml_rope_ext(ctx, ggml_cont(ctx, ggml_reshape_3d(ctx, k_cur, model->head_size, model->n_kv_attn_heads, n_tokens)), octx->positions, model->rope_frequencies, 
+                model->head_size, 2,0, 500000.0f,
+                1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+
+    // A performance comparison between this method, i.e. performing 3 incremental copy operations in order to achieve repeat_interleave,
+    // and performing the repeat operation upfront before performign a single copy needs to be performed in order to better optimize this function.
+    // Additionally, it might be more performant for the values transposition to be performed prior to appending it to the cache, as it would save us 
+    // from incrementally larger transpositions with generation.
+    for (int i = 0; i < repeat; i++) {
+        struct ggml_tensor * k_cache_view = ggml_view_3d(
+            ctx, 
+            kv_self->k_l[index], 
+            model->head_size,
+            model->n_kv_attn_heads,
+            n_tokens, 
+            ggml_element_size(kv_self->k_l[index]) * model->head_size * repeat,
+            ggml_element_size(kv_self->k_l[index]) * model->n_attn_heads * model->head_size,
+            ggml_element_size(kv_self->k_l[index]) * model->n_attn_heads * model->head_size * octx->current_position + i * ggml_element_size(kv_self->k_l[index]) * model->head_size
+        );
+        ggml_build_forward_expand(graph, ggml_cpy(ctx, k_cur, k_cache_view));
+
+        struct ggml_tensor * v_cache_view = ggml_view_3d(
+            ctx,
+            kv_self->v_l[index],
+            model->head_size,
+            model->n_kv_attn_heads,
+            n_tokens,
+            ggml_element_size(kv_self->k_l[index]) * model->head_size * repeat,
+            ggml_element_size(kv_self->k_l[index]) * model->n_attn_heads * model->head_size,
+            ggml_element_size(kv_self->k_l[index]) * model->n_attn_heads * model->head_size * octx->current_position + i * ggml_element_size(kv_self->k_l[index]) * model->head_size
+        );
+        ggml_build_forward_expand(graph, ggml_cpy(ctx, v_cur, v_cache_view));
+    }
+}
+
+struct ggml_cgraph * orpheus_runner::build_orpheus_graph(orpheus_ubatch & batch) {
+    init_build();
+    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
+    
+    struct ggml_tensor * cur;
+    struct ggml_tensor * inpL;
+    
+    const int32_t full_sequence_length = octx->current_position + (uint32_t) batch.n_tokens;
+    octx->positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.n_tokens);
+    ggml_set_input(octx->positions);
+    octx->inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.n_tokens);
+    ggml_set_input(octx->inp_tokens);
+    inpL = ggml_get_rows(ctx, model->embd, octx->inp_tokens);
+    
+    struct ggml_tensor * KQ_mask_dec = build_attn_mask(ctx, octx, batch);
+    
+    for (int l = 0; l < model->n_layers; l++) {
+        struct ggml_tensor * residual = inpL;
+        cur = orpheus_build_layer_norm(ctx, inpL, model->layers[l].input_norm);
+
+        struct ggml_tensor * attn_out;
+
+        // self-attention
+        {
+            struct ggml_tensor * Qcur = ggml_mul_mat(ctx, model->layers[l].q, cur);
+            struct ggml_tensor * Kcur = ggml_mul_mat(ctx, model->layers[l].k, cur);
+            struct ggml_tensor * Vcur = ggml_mul_mat(ctx, model->layers[l].v, cur);
+
+            orpheus_build_kv_store(ctx, gf, Kcur, Vcur, l, batch.n_tokens, 3);
+            struct ggml_tensor * k =
+                ggml_cont(ctx, ggml_view_3d(ctx, kv_self->k_l[l],
+                        model->head_size, full_sequence_length, model->n_attn_heads,
+                        ggml_element_size(kv_self->k_l[l]) * model->n_attn_heads * model->head_size,
+                        ggml_element_size(kv_self->k_l[l]) * model->head_size,
+                        0));            
+            
+            struct ggml_tensor * v =
+                ggml_view_2d(ctx, kv_self->v_l[l],
+                        model->hidden_size, full_sequence_length,
+                        ggml_element_size(kv_self->k_l[l]) * model->hidden_size,
+                        0);
+
+            v = ggml_cont_3d(ctx, ggml_transpose(ctx, v), full_sequence_length, model->head_size, model->n_attn_heads);
+
+            Qcur = ggml_rope_ext(
+                ctx, ggml_cont(ctx, ggml_reshape_3d(ctx, Qcur, model->head_size, model->n_attn_heads, batch.n_tokens)), 
+                octx->positions, model->rope_frequencies, model->head_size, 2, 0, 500000.0f, // rope theta
+                1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+
+            struct ggml_tensor * q = ggml_cont(ctx, ggml_permute(ctx, Qcur, 0, 2, 1, 3));
+            struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q);
+            kq = ggml_soft_max_ext(ctx, kq, KQ_mask_dec, 1.0f/sqrtf(model->head_size), 0.0f);
+            struct ggml_tensor * kqv = ggml_mul_mat(ctx, kq, v);
+            struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 2, 0, 1, 3);
+            attn_out = ggml_cont_2d(ctx, kqv_merged, model->hidden_size, batch.n_tokens);
+            attn_out = ggml_mul_mat(ctx, model->layers[l].o, attn_out);
+        }
+
+        cur = ggml_add(ctx, attn_out, residual);
+        
+        struct ggml_tensor * residualffn = cur;
+
+        // mlp
+        {
+            cur = orpheus_build_layer_norm(ctx, cur, model->layers[l].post_attention_norm);
+            cur = ggml_mul(ctx, ggml_silu(ctx, ggml_mul_mat(ctx, model->layers[l].gate, cur)), ggml_mul_mat(ctx, model->layers[l].up, cur));
+            cur = ggml_mul_mat(ctx, model->layers[l].down, cur);
+        }
+        cur = ggml_add(ctx, cur, residualffn);
+        inpL = cur;
+    }
+    
+    cur = orpheus_build_layer_norm(ctx, cur, model->output_norm);
+    // only about 40k of the output head is actually uses for generation purposes. Ideally the head tensor should be shrunk and sampled tokens should be incremented.
+    cur = ggml_mul_mat(ctx, model->head, cur);
+    if (batch.n_tokens > 1) {
+        cur = ggml_cont(ctx, ggml_view_1d(ctx, cur, model->vocab_size, ggml_element_size(cur) * (cur->ne[1] - 1) * model->vocab_size));
+    }
+    ggml_build_forward_expand(gf, cur);
+    free_build();
+    
+    return gf;
+}
+
+void orpheus_runner::decode(orpheus_ubatch & batch) {
+    ggml_backend_sched_reset(octx->sched);
+    
+    octx->output_tokens.reserve(model->max_generation_size);
+    
+    const size_t new_size  = model->vocab_size * model->max_generation_size * sizeof(float);
+    octx->prep_output_buffer(new_size);
+
+    ggml_cgraph * gf = build_orpheus_graph(batch);
+
+    // the output is always the last tensor in the graph
+    struct ggml_tensor * res = gf->nodes[gf->n_nodes - 1];
+    ggml_backend_sched_alloc_graph(octx->sched, gf);
+    
+    set_inputs(batch);
+    ggml_backend_sched_graph_compute_async(octx->sched, gf);
+ 
+    float * logits_out = octx->logits + octx->n_outputs * model->vocab_size;
+    octx->get_ggml_node_data(res, logits_out, model->vocab_size * sizeof(float));
+
+    // update the total number of outputs retrieved and the current position
+    octx->current_position += batch.n_tokens;
+
+    // Reset state for the next token before backend sync, to allow the CPU activities in the reset to
+    // overlap with device computation.
+    ggml_backend_sched_reset(octx->sched);
+}
+
+void orpheus_runner::set_inputs(orpheus_ubatch & batch) {
+    ggml_backend_tensor_set(octx->inp_tokens, batch.tokens.data(), 0, batch.tokens.size()*ggml_element_size(octx->inp_tokens));
+    int32_t * pos = (int32_t*) octx->positions->data;
+    float * mask = (float*) octx->attn_mask->data;
+    uint32_t max_pos = octx->current_position + batch.n_tokens;
+    for (int i = 0; i < batch.n_tokens; i++) {
+        pos[i] = (int32_t) octx->current_position + i;
+        for (int ii = 0; ii < max_pos; ii++) {
+            mask[i*max_pos + ii] = ii > pos[i] ? -INFINITY : 0.0f;
+        }
+    }
+}
+
+orpheus_ubatch orpheus_runner::batch_from_sentence(std::string sentence) {
+    struct orpheus_ubatch batch;
+    for (auto t : orpheus_prepended_tokens) {
+        batch.tokens.push_back(t);
+    }
+    if (!octx->voice.empty()) {
+        sentence = octx->voice  + ": " + sentence;
+    }
+    tokenizer->tokenize(sentence, batch.tokens);
+    for (auto t : orpheus_appended_tokens) {
+        batch.tokens.push_back(t);
+    }
+    batch.n_tokens = batch.tokens.size();
+    return batch;
+}
+
+std::vector<std::vector<uint32_t>> orpheus_runner::prepare_output_tokens() {
+    size_t chunks = octx->output_tokens.size() / 7;
+    std::vector<std::vector<uint32_t>> output_tokens;
+    for (int i = 0; i < model->audio_heads; i++) {
+        output_tokens.push_back(std::vector<uint32_t>{});
+    }
+    for (int i = 0; i < chunks; i++) {
+        for (int ii = 0; ii < 7; ii++) {
+            uint32_t thead = model->heads[ii];
+            // the manipulations below are not configured because they are performed inline via undocumented constants in the Orpheus codebase.
+            // Essentially this is how Orpheus converts discrete samples from the output shape to the audio input shape.
+            uint32_t t = octx->output_tokens[i*7 + ii] - 128266 - ((ii % 7) * 4096);
+            output_tokens[thead].push_back(t);
+        }
+    }
+    return output_tokens;
+}
+
+void orpheus_runner::generate_from_batch(orpheus_ubatch & batch, struct tts_response * output) {
+    while ((octx->output_tokens.size() == 0 || octx->output_tokens.back() != model->stopping_token_id) && octx->output_tokens.size() < model->max_generation_size) {
+        decode(batch);
+        generation_sampler->sample(octx->logits + octx->n_outputs * model->vocab_size, octx->output_tokens);
+        // only increment the output count after sampling
+        octx->n_outputs++;
+        batch = orpheus_ubatch{
+            1, {octx->output_tokens.back()}
+        };
+    }
+    // this case could be better addressed by adding spliting to the generation process.
+    if (octx->output_tokens.size() >= model->max_generation_size) {
+        fprintf(stdout, "Warning: generation hit its max default length. The generated audio may not contain the entire prompt.\n");
+    }
+    std::vector<std::vector<uint32_t>> processed_output_tokens = prepare_output_tokens();
+    srunner->run(processed_output_tokens, output);
+}
+
+int orpheus_runner::generate(std::string sentence, struct tts_response * response) {
+    orpheus_ubatch batch = batch_from_sentence(sentence);
+    // it should be possible to update the max context window size, but currently it is extremely unlikely that a single prompt will
+    // surpass the default size.
+    if (batch.tokens.size() > model->max_context_length) {
+        TTS_ABORT("The prompt was too large for the default context window. Try splitting up or shortenning the prompt.");
+    }
+    octx->reset();
+    generation_sampler->reset();
+    if  (!kv_self) {
+        orpheus_kv_cache_init();
+    }
+    generate_from_batch(batch, response);
+    return 0;
+}
+
+void orpheus_runner::configure_generation(generation_configuration * config) {
+    generation_sampler->temperature = config->temperature;
+    generation_sampler->repetition_penalty = config->repetition_penalty;
+    generation_sampler->do_sample = config->sample;
+    generation_sampler->top_k = config->top_k;
+    generation_sampler->top_p = config->top_p;
+    if (std::find(orpheus_voices.begin(), orpheus_voices.end(), config->voice) == orpheus_voices.end() && !config->voice.empty()) {
+        TTS_ABORT("Voice '%s' is not a valid voice for Orpheus.", config->voice.c_str());
+    }
+    octx->voice = config->voice;
+}
+
+orpheus_ubatch orpheus_runner::build_worst_case_batch() {
+    orpheus_ubatch batch;
+    batch.n_tokens = model->max_context_length;
+    return batch;
+}
+
+void orpheus_runner::assign_weight(std::string name, ggml_tensor * tensor) {
+    if (tensor->data == NULL) {
+        return;
+    }
+
+    if (name.size() == 0) {
+        // handles the top level meta tensor
+        return;
+    }
+
+    if (name.size() > 5 && name.substr(0, 5) == "snac.") {
+        srunner->model->assign_weight(name.substr(5), tensor);
+    } else if (name.size() > 8 && name.substr(0, 8) == "orpheus.") {
+        model->assign_weight(name.substr(8), tensor);
+    } else {
+        fprintf(stdout, "Warning: function %s encountered an unhandled tensor named '%s'.\n", __func__, name.c_str());
+    }
+}
+
+void orpheus_runner::prepare_post_load() {
+    srunner->prepare_post_load();
+    orpheus_kv_cache_init();
+    auto batch = build_worst_case_batch();
+    auto gf = build_orpheus_graph(batch);
+    octx->prep_schedule(gf);
+}
+
+std::vector<std::string> list_voices() {
+	std::vector<std::string> voices;
+	voices.reserve(orpheus_voices.size());
+	for (auto voice : orpheus_voices) {
+		voices.push_back(voice);
+	}
+	return voices;
+}

otherarch/ttscpp/src/orpheus_model.h

@@ -0,0 +1,146 @@
+#pragma once
+
+#include "sampler.h"
+#include "tokenizer.h"
+#include "snac_model.h"
+
+// Orpheus uses vLLM with a llama-3 architecture. The only critical difference from the normal llama architecture is the use of kv heads.
+
+struct orpheus_layer {
+    struct ggml_tensor * input_norm;
+    struct ggml_tensor * post_attention_norm;
+    struct ggml_tensor * q;
+    struct ggml_tensor * k;
+    struct ggml_tensor * v;
+    struct ggml_tensor * o;
+    struct ggml_tensor * gate;
+    struct ggml_tensor * up;
+    struct ggml_tensor * down;
+};
+
+struct orpheus_model : tts_model {
+    uint32_t vocab_size = 156940;
+    uint32_t n_attn_heads = 24;
+    uint32_t n_kv_attn_heads = 8;
+    uint32_t head_size = 128;
+    uint32_t max_context_length = 1024;
+    // the generation size is technically arbitrary as the model can handle a large context. This size comes out to being 25.6 seconds.
+    uint32_t max_generation_size = 2100;
+    uint32_t stopping_token_id = 128258;
+    uint32_t eos_token_id = 128001;
+    uint32_t bos_token_id = 128000;
+    uint32_t hidden_size = 3072;
+    uint32_t kv_hidden_size = 1024;
+    uint32_t audio_heads = 3;
+    uint32_t heads[7] = {0, 1, 2, 2, 1, 2, 2};
+
+    int n_layers = 28;
+
+    struct std::vector<orpheus_layer> layers;
+    struct ggml_tensor * head;
+    struct ggml_tensor * embd;
+    struct ggml_tensor * output_norm;
+    struct ggml_tensor * rope_frequencies;
+
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    void assign_to_layer(std::string part, orpheus_layer & layer, struct ggml_tensor * tensor);
+    void prep_constants(gguf_context * meta);
+    void prep_layers(gguf_context * meta);
+    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, "orpheus", 1.30);
+    }
+};
+
+struct orpheus_context : runner_context {
+    orpheus_context(orpheus_model * model, int n_threads): runner_context(n_threads), model(model) {};
+    struct orpheus_model * model;
+
+    uint32_t current_position = 0; // current position in the active sequence
+    uint32_t n_outputs = 0; // the position of the text prompt termination (used for adjusting the cache when incrementally generating)
+    std::string voice;
+
+    std::vector<uint32_t> output_tokens;
+
+    void reset();
+    void build_schedule() {
+        runner_context::build_schedule(model->max_nodes());
+    }
+
+    struct ggml_tensor * inp_tokens;
+    struct ggml_tensor * attn_mask;
+    struct ggml_tensor * positions;
+};
+
+struct orpheus_kv_cache {    
+    ggml_type cache_type = GGML_TYPE_F32;
+
+    std::vector<struct ggml_tensor *> k_l;
+    std::vector<struct ggml_tensor *> 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);
+    }
+
+    ~orpheus_kv_cache() {
+        free();
+    }
+};
+
+struct orpheus_context * build_new_orpheus_context(struct orpheus_model * model, int n_threads, bool use_cpu = true);
+
+struct orpheus_ubatch {
+    orpheus_ubatch() = default;
+    orpheus_ubatch(size_t n_tokens, std::vector<uint32_t> tokens): n_tokens(n_tokens), tokens(tokens) {};
+    size_t n_tokens; // total sentence tokens
+    std::vector<uint32_t> tokens;    // [n_tokens]
+};
+
+struct orpheus_runner : tts_runner {
+    orpheus_runner(
+            orpheus_model * model, 
+            snac_runner * audio_decoder, 
+            orpheus_context * octx, 
+            bpe_tokenizer * bt, 
+            sampler * samp, 
+            orpheus_kv_cache * cache): model(model), srunner(audio_decoder), octx(octx), tokenizer(bt), generation_sampler(samp), kv_self(cache) {
+        tts_runner::sampling_rate = 24000.0f;
+        generation_sampler->n_output_heads = 1;
+        generation_sampler->vocab_size = model->vocab_size;
+        generation_sampler->eos_token_id = model->eos_token_id;
+    }
+    orpheus_model * model;
+    snac_runner * srunner;
+    orpheus_context * octx;
+    bpe_tokenizer * tokenizer;
+    orpheus_kv_cache * kv_self;
+    sampler * generation_sampler;
+
+    void init_build() {
+        tts_runner::init_build(&octx->buf_compute_meta);
+    }
+
+    std::vector<std::string> list_voices();
+    struct ggml_cgraph * build_orpheus_graph(orpheus_ubatch & batch);
+    void orpheus_kv_cache_init();
+    void orpheus_build_kv_store(struct ggml_context * ctx, struct ggml_cgraph * graph, struct ggml_tensor * k_cur, struct ggml_tensor * v_cur, int index, uint32_t n_tokens, int repeat);
+    void configure_generation(generation_configuration * config);
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    std::vector<std::vector<uint32_t>> prepare_output_tokens();
+    orpheus_ubatch build_worst_case_batch();
+    orpheus_ubatch batch_from_sentence(std::string sentence);
+    void set_inputs(orpheus_ubatch & batch);
+    void decode(orpheus_ubatch & batch);
+    void prepare_post_load();
+    int generate(std::string sentence, struct tts_response * response);
+    void generate_from_batch(orpheus_ubatch & batch, struct tts_response * output);
+};
+
+static struct ggml_tensor * orpheus_build_layer_norm(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * weight);
+static struct ggml_tensor * build_attn_mask(ggml_context * ctx, orpheus_context * octx, orpheus_ubatch & batch);

otherarch/ttscpp/src/parler_model.cpp

@@ -0,0 +1,874 @@
+#include "parler_model.h"
+
+// For loading parler model from gguf file.
+static const std::map<std::string, parler_tensor> PARLER_TENSOR_GGUF_LOOKUP = {
+    {"layer_norm.weight", PARLER_NORM},
+    {"layer_norm.bias", PARLER_NORM_BIAS},
+    {"embed_prompts", PARLER_EMBD_PROMPTS},
+    {"text_encoding", PARLER_TEXT_ENCODING},
+    {"positional_embed", PARLER_POSITIONAL_EMBD},
+    {".self_attn.q_proj.weight", PARLER_LAYER_SELF_ATTN_Q},
+    {".self_attn.k_proj.weight", PARLER_LAYER_SELF_ATTN_K},
+    {".self_attn.v_proj.weight", PARLER_LAYER_SELF_ATTN_V},
+    {".self_attn.out_proj.weight", PARLER_LAYER_SELF_ATTN_O},
+    {".self_attn_layer_norm.weight", PARLER_LAYER_SELF_ATTN_NORM},
+    {".self_attn_layer_norm.bias", PARLER_LAYER_SELF_ATTN_NORM_BIAS},
+    {".encoder_attn.q_proj.weight", PARLER_LAYER_ATTN_Q},
+    {".encoder_attn.k_proj.weight", PARLER_LAYER_ATTN_K},
+    {".encoder_attn.v_proj.weight", PARLER_LAYER_ATTN_V},
+    {".encoder_attn.out_proj.weight", PARLER_LAYER_ATTN_O},
+    {".encoder_attn_layer_norm.weight", PARLER_LAYER_ATTN_NORM},
+    {".encoder_attn_layer_norm.bias", PARLER_LAYER_ATTN_NORM_BIAS},
+    {".fc1.weight", PARLER_LAYER_FC1},
+    {".fc2.weight", PARLER_LAYER_FC2},
+    {".final_layer_norm.weight", PARLER_LAYER_OUT_NORM},
+    {".final_layer_norm.bias", PARLER_LAYER_OUT_NORM_BIAS},
+    {".weight", PARLER_EMBD},
+    {".weight.head", PARLER_HEAD}
+};
+
+void parler_tts_model::assign_weight(std::string name, ggml_tensor * tensor) {
+    assign_to_decoder(this, name, tensor);
+}
+
+void parler_tts_model::prep_layers(gguf_context * meta_ctx) {
+    layers.reserve((size_t) n_layers);
+    for (int i = 0; i < (int) n_layers; i++) {
+        parler_layer * l = new parler_layer{};
+        layers.push_back(l);
+    }
+
+    embds.reserve((size_t) n_output_heads);
+    heads.reserve((size_t) n_output_heads);
+    for (int i = 0; i < n_output_heads; i++) {
+        struct ggml_tensor * h = nullptr;
+        struct ggml_tensor * embd = nullptr;
+        embds.push_back(embd);
+        heads.push_back(h);
+    }
+}
+
+void parler_tts_model::prep_constants(gguf_context * meta) {
+    int encode_length_key = search_for_gguf_keys(meta, {"parler-tts.decoder.encode_length", "encode_length"});
+    if (encode_length_key == -1) {
+        TTS_ABORT("key 'parler-tts.decoder.encode_length' must be specified in gguf file.");
+    }
+    n_encode_length = gguf_get_val_u32(meta, encode_length_key);
+
+    int hidden_size_key = search_for_gguf_keys(meta, {"parler-tts.decoder.hidden_size", "hidden_size"});
+    if (hidden_size_key != -1) {
+        hidden_size = gguf_get_val_u32(meta, hidden_size_key);
+    }
+
+    int output_heads_key = search_for_gguf_keys(meta, {"parler-tts.decoder.output_heads", "output_heads"});
+    if (output_heads_key != -1) {
+        n_output_heads = gguf_get_val_u32(meta, output_heads_key);
+    }
+    int ctx_length_key = search_for_gguf_keys(meta, {"parler-tts.decoder.context_length", "ctx_length"});
+    if (ctx_length_key != -1) {
+        max_ctx_length = gguf_get_val_u32(meta, ctx_length_key);
+    }
+
+    int attn_heads_key = search_for_gguf_keys(meta, {"parler-tts.decoder.attention.head_count", "attn_heads"});
+    if (attn_heads_key != -1) {
+        n_attn_heads = gguf_get_val_u32(meta, attn_heads_key);
+    }
+    head_size = hidden_size / n_attn_heads;
+    max_cross_nodes = n_attn_heads * 2;
+
+    int output_vocab_size_key = search_for_gguf_keys(meta, {"parler-tts.decoder.out_vocab_size", "out_vocab_size"});
+    if (output_vocab_size_key != -1) {
+        output_vocab_size = gguf_get_val_u32(meta, output_vocab_size_key);
+    }
+
+    int audio_vocab_size_key = search_for_gguf_keys(meta, {"parler-tts.decoder.audio_vocab_size", "audio_vocab_size"});
+    if (audio_vocab_size_key != -1) {
+        audio_vocab_size = gguf_get_val_u32(meta, audio_vocab_size_key);
+    }
+
+    int max_gen_key = search_for_gguf_keys(meta, {"parler-tts.decoder.max_generation", "max_generation"});
+    if (max_gen_key != -1) {
+        max_generation_size = gguf_get_val_u32(meta, max_gen_key);
+    }
+
+    int n_layers_key = search_for_gguf_keys(meta, {"parler-tts.decoder.num_hidden_layers", "num_hidden_layers"});
+    if (n_layers_key != -1) {
+        n_layers = gguf_get_val_u32(meta, n_layers_key);
+    }
+
+    int bos_token_id_key = search_for_gguf_keys(meta, {"audio.bos_token_id", "bos_token_id"});
+    if (bos_token_id_key != -1) {
+        bos_token_id = gguf_get_val_u32(meta, bos_token_id_key);
+    }
+
+    int eos_token_id_key = search_for_gguf_keys(meta, {"audio.eos_token_id", "eos_token_id"});
+    if (eos_token_id_key != -1) {
+        eos_token_id = gguf_get_val_u32(meta, eos_token_id_key);
+    }
+}
+
+void parler_tts_model::prep_cross_key_values(int n_threads, struct tts_response * conditional_prompt) {
+    ggml_backend_t backend_cpu = ggml_backend_cpu_init();
+    ggml_backend_buffer_type_t backend_cpu_buffer = ggml_backend_cpu_buffer_type();
+    // Let it create a disposable threadpool just this once
+    ggml_backend_cpu_set_n_threads(backend_cpu, n_threads);
+    std::vector<ggml_backend_buffer_type_t> bufs = {backend_cpu_buffer};
+    std::vector<ggml_backend_t> backs = {backend_cpu};
+    ggml_backend_sched_t sched = ggml_backend_sched_new(backs.data(), bufs.data(), 1, max_cross_nodes*n_layers, false, false);
+
+    std::vector<uint8_t> buf_compute_meta;
+    buf_compute_meta.resize(max_cross_nodes*n_layers*ggml_tensor_overhead() + ggml_graph_overhead_custom(max_cross_nodes*n_layers, false));
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ buf_compute_meta.size(),
+        /*.mem_buffer =*/ buf_compute_meta.data(),
+        /*.no_alloc   =*/ true,
+    };
+    struct ggml_context * cctx = ggml_init(params);
+    struct ggml_cgraph * gf = ggml_new_graph_custom(cctx, 4096, false);
+    if (conditional_prompt) {
+        // If we are updating the conditional prompt then we have to reset the tensor offsets into the ggml_context otherwise we could overflow the assigned buffer and lose our prompt.
+        // These offsets are assigned by #set_tensor below.
+        offset -= n_encode_length*hidden_size*sizeof(float)*n_layers*2;
+        precomputed_input_emb = ggml_new_tensor_2d(cctx, GGML_TYPE_F32, conditional_prompt->hidden_size, conditional_prompt->n_outputs);
+        ggml_set_input(precomputed_input_emb);
+        n_encode_length = conditional_prompt->n_outputs;
+    }
+
+    for (int i = 0; i < layers.size(); i++) {
+        struct ggml_tensor * Kcur = ggml_mul_mat(cctx, layers[i]->attn_k_proj, precomputed_input_emb);
+        struct ggml_tensor * Vcur = ggml_mul_mat(cctx, layers[i]->attn_v_proj, precomputed_input_emb);
+
+        Kcur = ggml_reshape_3d(cctx, Kcur, head_size, n_attn_heads, n_encode_length);
+        Vcur = ggml_transpose(cctx, Vcur);
+
+        struct ggml_tensor * k = ggml_cont(cctx, ggml_permute(cctx, Kcur, 0, 2, 1, 3));
+        ggml_set_name(k, ("cross_key_" + std::to_string(i)).c_str());
+        ggml_build_forward_expand(gf, k);
+
+        struct ggml_tensor * v = ggml_cont_3d(cctx, Vcur, n_encode_length, head_size, n_attn_heads);
+        ggml_set_name(v, ("cross_value_" + std::to_string(i)).c_str());
+        ggml_build_forward_expand(gf, v);
+    }
+
+    ggml_free(cctx);
+    ggml_backend_sched_reserve(sched, gf);
+    ggml_backend_sched_alloc_graph(sched, gf);
+    if (conditional_prompt) {
+        ggml_backend_tensor_set(precomputed_input_emb, conditional_prompt->data, 0, conditional_prompt->n_outputs*conditional_prompt->hidden_size*ggml_element_size(precomputed_input_emb));
+    }
+
+    ggml_backend_sched_graph_compute_async(sched, gf);
+
+    for (int i = 0; i < layers.size(); i++) {
+        struct ggml_tensor * k = ggml_graph_get_tensor(gf, ("cross_key_" + std::to_string(i)).c_str());
+        layers[i]->cross_k = ggml_dup_tensor(ctx, k);
+        set_tensor(layers[i]->cross_k, k);
+        struct ggml_tensor * v = ggml_graph_get_tensor(gf, ("cross_value_" + std::to_string(i)).c_str());
+        layers[i]->cross_v = ggml_dup_tensor(ctx, v);
+        set_tensor(layers[i]->cross_v, v);
+    }
+    ggml_backend_sched_free(sched);
+    ggml_backend_free(backend_cpu);
+}
+
+void assign_parler_layer(parler_tts_model * model, parler_layer * layer, std::string name, ggml_tensor * tensor) {
+    try {
+        switch(PARLER_TENSOR_GGUF_LOOKUP.at(name)) {
+            case PARLER_LAYER_SELF_ATTN_Q:
+                layer->self_attn_q_proj = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->self_attn_q_proj, tensor);
+                break;
+            case PARLER_LAYER_SELF_ATTN_K:
+                layer->self_attn_k_proj = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->self_attn_k_proj, tensor);
+                break;
+            case PARLER_LAYER_SELF_ATTN_V:
+                layer->self_attn_v_proj = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->self_attn_v_proj, tensor);
+                break;
+            case PARLER_LAYER_SELF_ATTN_O:
+                layer->self_attn_o_proj = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->self_attn_o_proj, tensor);
+                break;
+            case PARLER_LAYER_SELF_ATTN_NORM:
+                layer->self_attn_norm = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->self_attn_norm, tensor);
+                break;
+            case PARLER_LAYER_SELF_ATTN_NORM_BIAS:
+                layer->self_attn_norm_bias = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->self_attn_norm_bias, tensor);
+                break;
+            case PARLER_LAYER_ATTN_Q:
+                if (model->use_cross_attn) {
+                    layer->attn_q_proj = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(layer->attn_q_proj, tensor);
+                }
+                break;
+            case PARLER_LAYER_ATTN_K:
+                if (model->use_cross_attn) {
+                    layer->attn_k_proj = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(layer->attn_k_proj, tensor);
+                }
+                break;
+            case PARLER_LAYER_ATTN_V:
+                if (model->use_cross_attn) {
+                    layer->attn_v_proj = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(layer->attn_v_proj, tensor);
+                }
+                break;
+            case PARLER_LAYER_ATTN_O:
+                if (model->use_cross_attn) {
+                    layer->attn_o_proj = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(layer->attn_o_proj, tensor);
+                }
+                break;
+            case PARLER_LAYER_ATTN_NORM:
+                if (model->use_cross_attn) {
+                    layer->attn_norm = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(layer->attn_norm, tensor);
+                }
+                break;
+            case PARLER_LAYER_ATTN_NORM_BIAS:
+                if (model->use_cross_attn) {
+                    layer->attn_norm_bias = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(layer->attn_norm_bias, tensor);
+                }
+                break;
+            case PARLER_LAYER_FC1:
+                layer->fc1 = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->fc1, tensor);
+                break;
+            case PARLER_LAYER_FC2:
+                layer->fc2 = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->fc2, tensor);
+                break;
+            case PARLER_LAYER_OUT_NORM:
+                layer->final_norm = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->final_norm, tensor);
+                break;
+            case PARLER_LAYER_OUT_NORM_BIAS:
+                layer->final_norm_bias = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer->final_norm_bias, tensor);
+                break;
+            default:
+                fprintf(stdout, "unassigned tensor %s\n", name.c_str());
+                break;
+        }
+    } catch (const std::out_of_range& e) {
+        TTS_ABORT("Error: %s\nTensor, '%s', is not a valid tensor.", e.what(), name.c_str());
+    }
+}
+
+void assign_to_decoder(parler_tts_model * model, const std::string name, ggml_tensor * tensor) {
+    if (PARLER_TENSOR_GGUF_LOOKUP.find(name) != PARLER_TENSOR_GGUF_LOOKUP.end()) {
+        try {
+            switch (PARLER_TENSOR_GGUF_LOOKUP.at(name)) {
+                case PARLER_NORM:
+                    model->layer_norm = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(model->layer_norm, tensor);
+                    break;
+                case PARLER_NORM_BIAS:
+                    model->layer_norm_bias = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(model->layer_norm_bias, tensor);
+                    break;
+                case PARLER_EMBD_PROMPTS:
+                    model->prompt_embd = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(model->prompt_embd, tensor);
+                    break;
+                case PARLER_TEXT_ENCODING:
+                    if (model->use_cross_attn) {
+                        model->precomputed_input_emb = ggml_dup_tensor(model->ctx, tensor);
+                        model->set_tensor(model->precomputed_input_emb, tensor);
+                    }
+                    break;
+                case PARLER_POSITIONAL_EMBD:
+                    model->precomputed_positional_embds = ggml_dup_tensor(model->ctx, tensor);
+                    model->set_tensor(model->precomputed_positional_embds, tensor);
+                    break;
+                default:
+                    fprintf(stdout, "unassigned tensor %s\n", name.c_str());
+                    break;
+            }
+        } catch (const std::out_of_range& e) {
+            TTS_ABORT("Error: %s\nTensor, '%s', is not a valid tensor.", e.what(), name.c_str());
+        }
+    } else if (std::find_if(name.begin(), name.end(), ::isdigit) != name.end())  {
+        auto pair = parse_layer_count(name);
+        int layer = pair.first;
+        std::string lt_name = pair.second;
+        if (name.find("embed_tokens") != std::string::npos) {
+            model->embds[layer] = ggml_dup_tensor(model->ctx, tensor);
+            model->set_tensor(model->embds[layer], tensor);
+        } else if (name.find("lm_heads") != std::string::npos) {
+            model->heads[layer] = ggml_dup_tensor(model->ctx, tensor);
+            model->set_tensor(model->heads[layer], tensor);
+        } else {
+            assign_parler_layer(model, model->layers[layer], lt_name, tensor);
+        }
+    }
+}
+
+void parler_context::reset(int32_t n_output_heads) {
+    n_outputs = 0;
+    prompt_end_position = 0;
+    current_position = 0;
+    output_size = 0;
+    output_tokens.clear();
+    eos_seen.clear();
+    for (int i = 0; i < (int) n_output_heads; i++) {
+        eos_seen.push_back(false);
+    }
+}
+
+struct parler_context * build_new_parler_context(struct parler_tts_model * model, int n_threads, bool use_cpu) {
+    parler_context * pctx = new parler_context(model, n_threads);
+    if (!use_cpu) {
+#ifdef GGML_USE_METAL
+        pctx->backend = ggml_backend_metal_init();
+#endif
+    }
+    pctx->eos_seen.reserve(model->n_output_heads);
+    pctx->backend_cpu = ggml_backend_cpu_init();
+    pctx->set_threads();
+    pctx->build_schedule();
+    pctx->buf_compute_meta.resize(ggml_tensor_overhead()*model->max_nodes() + ggml_graph_overhead_custom(model->max_nodes(), false));
+    return pctx;
+}
+
+static bool parler_kv_cache_init(struct parler_kv_cache * cache, parler_tts_model * model, parler_context * pctx, int32_t seq_id) {
+    const int64_t n_layer = (int64_t) model->layers.size();
+    cache->seq_id = seq_id;
+
+    ggml_backend_buffer_type_t buft = nullptr;
+    // this will only really support cpu or metal for the time being;
+    if (pctx->backend != nullptr) {
+#ifdef GGML_USE_METAL
+        buft = ggml_backend_metal_buffer_type();
+#endif
+    } else {
+        buft = ggml_backend_cpu_buffer_type();
+    }
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ (2u*model->n_layers+1)*ggml_tensor_overhead(),
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true,
+    };
+    ggml_context * ctx = ggml_init(params);
+    if (!ctx) {
+        return false;
+    }
+    cache->ctx = ctx;
+
+
+    cache->k_l.reserve(n_layer);
+    cache->v_l.reserve(n_layer);
+
+    for (int i = 0; i < (int) n_layer; i++) {
+        ggml_tensor * k = ggml_new_tensor_1d(cache->ctx, cache->type_k, model->hidden_size*model->max_ctx_length);
+        ggml_tensor * v = ggml_new_tensor_1d(cache->ctx, cache->type_v, model->hidden_size*model->max_ctx_length);
+        ggml_format_name(k, "cache_k_l%d", i);
+        ggml_format_name(v, "cache_v_l%d", i);
+        cache->k_l.push_back(k);
+        cache->v_l.push_back(v);
+    }
+
+    // allocate tensors and initialize the buffers to avoid NaNs in the padding
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(cache->ctx, buft);
+    if (!buf) {
+        return false;
+    }
+    ggml_backend_buffer_clear(buf, 0);
+    cache->buf = buf;
+
+    return true;
+}
+
+struct ggml_tensor * parler_build_inp_embd(struct ggml_context * ctx, struct parler_context * pctx, parler_tts_model * model, parler_ubatch & batch) {
+    // Parler has two embedding schemas one for the text input and one for generative audio tokens. These two schemas have effectively distinct shapes (i.e. [batch_size, sequence_length] and [batch_size, sequence_lenghth, num_codebooks] respectively).
+    // This means that depending on where we are in generation we need to follow a distinct pattern
+    struct ggml_tensor * input_embs;
+    pctx->positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.sequence_length);
+    ggml_set_input(pctx->positions);
+    if (batch.audio_generation) {
+        pctx->audio_inp_tokens = ggml_reshape_2d(ctx, ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.n_audio_tokens), batch.n_audio_tokens / model->n_output_heads, model->n_output_heads);
+        ggml_set_input(pctx->audio_inp_tokens);
+        struct ggml_tensor * audio_tokens = ggml_reshape_2d(ctx, pctx->audio_inp_tokens, batch.n_audio_tokens / model->n_output_heads, model->n_output_heads);
+        for (int i = 0; i < model->n_output_heads; i++) {
+            if (i == 0) {
+                input_embs = ggml_get_rows(ctx, model->embds[i], ggml_view_2d(ctx, audio_tokens, 1, batch.n_audio_tokens / model->n_output_heads, audio_tokens->nb[1], i*sizeof(int32_t)));
+            } else {
+                input_embs = ggml_add(ctx, ggml_get_rows(ctx, model->embds[i], ggml_view_2d(ctx, audio_tokens, 1, batch.n_audio_tokens / model->n_output_heads, audio_tokens->nb[1], i*sizeof(int32_t))), input_embs);
+            }
+        }
+    } else {
+        pctx->inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.n_tokens);
+        ggml_set_input(pctx->inp_tokens);
+        input_embs = ggml_get_rows(ctx, model->prompt_embd, pctx->inp_tokens);
+    }
+    return ggml_add(ctx, input_embs, ggml_get_rows(ctx, model->precomputed_positional_embds, pctx->positions));
+}
+
+struct ggml_tensor * parler_build_layer_norm(struct ggml_context * ctx, struct ggml_tensor * inputs, struct ggml_tensor * weight, struct ggml_tensor * bias) {
+    // parler always uses default eps
+    float eps = 0.00001;
+    inputs = ggml_norm(ctx, inputs, eps);
+    inputs = ggml_mul(ctx, inputs, weight);
+    return ggml_add(ctx, inputs, bias);
+}
+
+void parler_build_kv_store(struct ggml_context * ctx, 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) {
+    // this is the max context size;
+    const int64_t n_ctx = 4096;
+
+    struct ggml_tensor * k_cache_view = ggml_view_1d(ctx, kv->k_l[index], n_tokens*n_embd_gqa, ggml_row_size(kv->k_l[index]->type, n_embd_gqa)*kv_head);
+
+    ggml_build_forward_expand(graph, ggml_cpy(ctx, k_cur, k_cache_view));
+
+    assert(v_cur->ne[0] == n_embd_gqa && v_cur->ne[1] == n_tokens);
+
+    struct ggml_tensor * v_cache_view = nullptr;
+
+    v_cache_view = ggml_view_2d(ctx, kv->v_l[index], n_tokens, n_embd_gqa,
+            (  n_ctx)*ggml_element_size(kv->v_l[index]),
+            (kv_head)*ggml_element_size(kv->v_l[index]));
+
+    v_cur = ggml_cont(ctx, ggml_transpose(ctx, v_cur));
+
+    ggml_build_forward_expand(graph, ggml_cpy(ctx, v_cur, v_cache_view));
+}
+
+struct ggml_tensor * parler_build_head_outputs(struct ggml_context * ctx, parler_tts_model * model, struct ggml_tensor * cur) {
+    // going to cat the heads together and then reshape them;
+    // honestly ggml doesn't provide good support for stacking and discrete tensor access
+    struct ggml_tensor * out;
+    for (int i = 0; i < model->n_output_heads; i++) {
+        if (i == 0) {
+            out = ggml_mul_mat(ctx, model->heads[i], cur);
+        } else {
+            out = ggml_concat(ctx, out, ggml_mul_mat(ctx, model->heads[i], cur), 1);
+        }
+    }
+    ggml_set_name(out, "final_out");
+    //out = ggml_cont(ctx, ggml_transpose(ctx, out));
+
+    int32_t sql_len = (int32_t) (ggml_nelements(out) / (model->output_vocab_size * model->n_output_heads));
+    return ggml_cont_3d(ctx, out, model->output_vocab_size, sql_len, model->n_output_heads);
+}
+
+struct ggml_tensor * build_attn_mask(ggml_context * ctx, parler_context * pctx, parler_ubatch & batch) {
+    pctx->attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) pctx->current_position + batch.sequence_length, (int64_t) pctx->current_position + batch.sequence_length);
+    ggml_set_input(pctx->attn_mask);
+
+    return pctx->attn_mask;
+}
+
+struct ggml_tensor * build_attn_mask_cross(ggml_context * ctx, parler_context * pctx, parler_tts_model * model, parler_ubatch & batch) {
+    pctx->attn_mask_cross = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) model->n_encode_length, (int64_t) batch.sequence_length);
+    ggml_set_input(pctx->attn_mask_cross);
+
+    return pctx->attn_mask_cross;
+}
+
+static struct parler_ubatch batch_from_sentence(std::string sentence, parler_tts_model * model, unigram_tokenizer * tokenizer) {
+    struct parler_ubatch batch;
+    batch.audio_generation = false;
+    std::vector<uint32_t>* token_ids = new std::vector<uint32_t>;
+    tokenizer->tokenize(sentence, *token_ids);
+    token_ids->push_back(tokenizer->eos_token);
+    batch.current_step = 0;
+    batch.n_tokens = token_ids->size();
+    batch.n_audio_tokens = 0;
+    batch.sequence_length = batch.n_tokens; // sequence_length is equal to the number of tokens for non-audio generation
+    std::vector<uint32_t>* position = new std::vector<uint32_t>;
+    for (uint32_t i = 0; i < batch.sequence_length; i++) {
+        position->push_back(i);
+    }
+    std::vector<uint32_t>* order = new std::vector<uint32_t>;
+    for (int i = 0; i < batch.sequence_length; i++) {
+        if (i >= batch.sequence_length - 1) {
+            order->push_back(0);
+        } else {
+            order->push_back(i+1);
+        }
+    }
+    batch.positions = position->data();
+    batch.tokens = token_ids->data();
+    return batch;
+}
+
+void parler_tts_runner::assign_weight(std::string name, ggml_tensor * tensor) {
+    std::string::size_type pos = name.find(".", 0);
+    std::string top_level(name.substr(0, pos));
+    std::string value(name.substr(pos + 1));
+    if (tensor->data == NULL) {
+        return;
+    }
+    if (top_level == "audio_encoder") {
+        dac_runner->model->assign_weight(value, tensor);
+    } else if (top_level == "decoder") {
+        model->assign_weight(value, tensor);
+    } else {
+        return;
+    }
+}
+
+void parler_tts_runner::update_conditional_prompt(const std::string file_path, const std::string prompt, int n_threads, bool cpu_only) {
+    t5_runner * text_encoder = text_encoder_from_file(file_path, n_threads, tokenizer, cpu_only);
+    tts_response* response;
+    text_encoder->generate(prompt, response);
+    model->prep_cross_key_values(n_threads, response);
+    delete text_encoder;
+    return;
+}
+
+
+struct ggml_cgraph * parler_tts_runner::build_parler_graph(parler_ubatch & batch) {
+    init_build();
+    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
+
+    struct ggml_tensor * cur;
+    struct ggml_tensor * inpL;
+
+    const int32_t full_sequence_length = pctx->current_position + (uint32_t) batch.sequence_length;
+
+    inpL = parler_build_inp_embd(ctx, pctx, model, batch);
+
+    struct ggml_tensor * KQ_mask_dec = build_attn_mask(ctx, pctx, batch);
+    struct ggml_tensor * KQ_mask_cross = build_attn_mask_cross(ctx, pctx, model, batch);
+
+    for (int l = 0; l < model->n_layers; l++) {
+        struct ggml_tensor * residual = inpL;
+        ggml_set_name(inpL, ("layer_" + std::to_string(l) + "_input").c_str());
+
+        cur = parler_build_layer_norm(ctx, inpL, model->layers[l]->self_attn_norm, model->layers[l]->self_attn_norm_bias);
+
+        struct ggml_tensor * attn_out;
+
+        // self-attention
+        {
+            struct ggml_tensor * Qcur = ggml_mul_mat(ctx, model->layers[l]->self_attn_q_proj, cur);
+            struct ggml_tensor * Kcur = ggml_mul_mat(ctx, model->layers[l]->self_attn_k_proj, cur);
+            struct ggml_tensor * Vcur = ggml_mul_mat(ctx, model->layers[l]->self_attn_v_proj, cur);
+
+            parler_build_kv_store(ctx, kv_self, gf, Kcur, Vcur, (int32_t) batch.sequence_length, pctx->current_position, l, model->hidden_size);
+            struct ggml_tensor * k =
+                ggml_view_3d(ctx, kv_self->k_l[l],
+                        model->head_size, full_sequence_length, model->n_attn_heads,
+                        ggml_row_size(kv_self->k_l[l]->type, model->hidden_size),
+                        ggml_row_size(kv_self->k_l[l]->type, model->head_size),
+                        0);
+
+
+            struct ggml_tensor * v =
+                ggml_view_3d(ctx, kv_self->v_l[l],
+                        full_sequence_length, model->head_size, model->n_attn_heads,
+                        ggml_element_size(kv_self->v_l[l])*model->max_ctx_length,
+                        ggml_element_size(kv_self->v_l[l])*model->max_ctx_length*model->head_size,
+                        0);
+
+            Qcur = ggml_reshape_3d(ctx, Qcur, model->head_size, model->n_attn_heads, batch.sequence_length);
+            struct ggml_tensor * q = ggml_cont(ctx, ggml_permute(ctx, Qcur, 0, 2, 1, 3));
+            struct ggml_tensor * kq = ggml_mul_mat(ctx, ggml_cont(ctx, k), q);
+            kq = ggml_soft_max_ext(ctx, kq, KQ_mask_dec, 1.0f/sqrtf(model->head_size), 0.0f);
+            struct ggml_tensor * kqv = ggml_mul_mat(ctx, kq, v);
+            struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 2, 0, 1, 3);
+            attn_out = ggml_cont_2d(ctx, kqv_merged, model->hidden_size, batch.sequence_length);
+            attn_out = ggml_mul_mat(ctx, model->layers[l]->self_attn_o_proj, attn_out);
+        }
+
+        cur = ggml_add(ctx, attn_out, residual);
+
+        if (model->use_cross_attn) {
+            struct ggml_tensor * residuala = cur;
+
+            // norm
+            cur = parler_build_layer_norm(ctx, cur, model->layers[l]->attn_norm, model->layers[l]->attn_norm_bias);
+
+            //cross-attention
+            struct ggml_tensor * Qcur = ggml_mul_mat(ctx, model->layers[l]->attn_q_proj, cur);
+            Qcur = ggml_reshape_3d(ctx, Qcur, model->head_size, model->n_attn_heads, batch.sequence_length);
+
+            struct ggml_tensor * q = ggml_cont(ctx, ggml_permute(ctx, Qcur, 0, 2, 1, 3));
+
+            struct ggml_tensor * kq = ggml_mul_mat(ctx, model->layers[l]->cross_k, q);
+            kq = ggml_soft_max_ext(ctx, kq, KQ_mask_cross, 1.0f/sqrtf(model->head_size), 0.0f);
+
+            struct ggml_tensor * kqv  = ggml_mul_mat(ctx, kq, model->layers[l]->cross_v);
+            struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 2, 0, 1, 3);
+            cur = ggml_cont_2d(ctx, kqv_merged, model->hidden_size, batch.sequence_length);
+            cur = ggml_mul_mat(ctx, model->layers[l]->attn_o_proj, cur);
+            cur = ggml_add(ctx, cur, residuala);
+        }
+
+        struct ggml_tensor * residualffn = cur;
+
+        cur = parler_build_layer_norm(ctx, cur, model->layers[l]->final_norm, model->layers[l]->final_norm_bias);
+        cur = ggml_mul_mat(ctx, model->layers[l]->fc1, cur);
+        cur = ggml_gelu(ctx, cur);
+        cur = ggml_mul_mat(ctx, model->layers[l]->fc2, cur);
+        cur = ggml_add(ctx, cur, residualffn);
+        inpL = cur;
+    }
+
+    cur = parler_build_layer_norm(ctx, cur, model->layer_norm, model->layer_norm_bias);
+    cur = parler_build_head_outputs(ctx, model, cur);
+    ggml_build_forward_expand(gf, cur);
+    free_build();
+
+    return gf;
+}
+
+void parler_tts_runner::configure_generation(generation_configuration * config) {
+    sampler->temperature = config->temperature;
+    sampler->repetition_penalty = config->repetition_penalty;
+    sampler->do_sample = config->sample;
+    sampler->top_k = config->top_k;
+    sampler->top_p = config->top_p;
+    model->use_cross_attn = config->use_cross_attn;
+}
+
+void parler_tts_runner::set_inputs(parler_ubatch & batch) {
+    if (batch.audio_generation) {
+        ggml_backend_tensor_set(pctx->audio_inp_tokens, batch.audio_tokens, 0, batch.n_audio_tokens*ggml_element_size(pctx->audio_inp_tokens));
+    } else {
+        ggml_backend_tensor_set(pctx->inp_tokens, batch.tokens, 0, batch.n_tokens*ggml_element_size(pctx->inp_tokens));
+    }
+    ggml_backend_tensor_set(pctx->positions, batch.positions, 0, batch.sequence_length*ggml_element_size(pctx->positions));
+    float * d = nullptr;
+    d = (float *) pctx->attn_mask->data;
+    uint32_t max_pos = pctx->current_position + batch.sequence_length;
+    for (int i = 0; i < batch.sequence_length; i++) {
+        uint32_t pos = batch.positions[i];
+        for (int ii = 0; ii < max_pos; ii++) {
+            d[i*max_pos + ii] = ii > pos ? -INFINITY : 0.0f;
+        }
+    }
+
+    if (model->use_cross_attn) {
+        float * d2 = nullptr;
+        d2 = (float *) pctx->attn_mask_cross->data;
+        for (int i = 0; i < model->n_encode_length; i++) {
+            for (int ii = 0; ii < batch.sequence_length; ii++) {
+                d2[i*batch.sequence_length + ii] = 0.0f;
+            }
+        }
+    }
+
+}
+void parler_tts_runner::parler_graph_compute(ggml_cgraph * gf) {
+    ggml_backend_sched_graph_compute_async(pctx->sched, gf);
+}
+
+int parler_tts_runner::decode(parler_ubatch & batch) {
+    ggml_backend_sched_reset(pctx->sched);
+
+    pctx->output_tokens.reserve(model->max_generation_size);
+
+    const size_t logits_size = model->output_vocab_size*model->max_generation_size*model->n_output_heads;
+    const size_t prev_size = pctx->buf_output ? ggml_backend_buffer_get_size(pctx->buf_output) : 0;
+    const size_t new_size  = logits_size * sizeof(float);
+
+    if (!pctx->buf_output || prev_size < new_size) {
+        if (pctx->buf_output) {
+            ggml_backend_buffer_free(pctx->buf_output);
+            pctx->buf_output = nullptr;
+            pctx->logits = nullptr;
+        }
+
+        pctx->buf_output = ggml_backend_buft_alloc_buffer(pctx->backend_cpu_buffer, new_size);
+    }
+
+    pctx->logits = (float *) ggml_backend_buffer_get_base(pctx->buf_output);
+    //ggml_backend_buffer_clear(pctx->buf_output, 0);
+
+    ggml_cgraph * gf = build_parler_graph(batch);
+
+    // the output is always the last tensor in the graph
+    struct ggml_tensor * res = gf->nodes[gf->n_nodes - 1];
+    ggml_backend_sched_alloc_graph(pctx->sched, gf);
+
+    // use the sequence_length variable here so that audio input tokens are handled correctly.
+    size_t n_outputs_new = batch.sequence_length;
+
+    set_inputs(batch);
+    parler_graph_compute(gf);
+
+    float * logits_out = pctx->logits + pctx->n_outputs * model->output_vocab_size * model->n_output_heads;
+    pctx->get_ggml_node_data(res, logits_out, n_outputs_new*model->output_vocab_size*model->n_output_heads*sizeof(float));
+
+    // set to total number of outputs in the batch*/
+    pctx->n_outputs += n_outputs_new;
+
+    // Reset state for the next token before backend sync, to allow the CPU activities in the reset to
+    // overlap with device computation.
+    ggml_backend_sched_reset(pctx->sched);
+
+    return 0;
+}
+
+parler_ubatch parler_tts_runner::build_worst_case_batch()  {
+    struct parler_ubatch batch;
+    batch.audio_generation = false;
+    batch.n_tokens = model->max_ctx_length;
+    batch.n_audio_tokens = 0;
+    batch.sequence_length = model->max_ctx_length;
+    return batch;
+}
+
+void parler_tts_runner::prepare_post_load() {
+    dac_runner->prepare_post_load();
+    parler_kv_cache_init(kv_self, model, pctx, std::mt19937(std::random_device{}())());
+    auto batch = build_worst_case_batch();
+    auto gf = build_parler_graph(batch);
+    pctx->prep_schedule(gf);
+}
+
+bool parler_tts_runner::adjust_for_sequence_continuation(struct parler_ubatch & batch) {
+    return false; // not implemneted
+}
+
+bool parler_tts_runner::check_stopping() {
+    int32_t token_position = (int32_t) pctx->output_tokens.size() - (int32_t) model->n_output_heads;
+    if (token_position < 0) {
+        return false;
+    }
+    if (pctx->current_position >= model->max_generation_size) {
+        return true;
+    }
+
+    bool channels_complete = true;
+    for (int i = 0; i < model->n_output_heads; i++) {
+        pctx->eos_seen[i] = pctx->eos_seen[i] || pctx->output_tokens[token_position+i] == model->eos_token_id;
+        if (channels_complete) {
+            channels_complete = pctx->eos_seen[i];
+        }
+    }
+    return channels_complete;
+}
+
+void parler_tts_runner::adjust_output_tokens(std::vector<uint32_t> & output_tokens, std::vector<uint32_t> & filtered) {
+    // currently this is applying sliding window over the heads and filtering out bad tokens.
+    // If we convert the DAC model's quantizer layers to support by row + column embeddings then we will need to transpose
+    // the heads and the sequence here, but right now simplying using a strided view is more peformant.
+    size_t size = output_tokens.size();
+    filtered.reserve(size);
+    for (int i = 0; i < size / model->n_output_heads; i++) {
+        bool remove = false;
+        for (int ii = 0; ii < model->n_output_heads; ii++) {
+            int next_index = i*model->n_output_heads+ii*model->n_output_heads+ii;
+            if (next_index > size || output_tokens[next_index] >= model->audio_vocab_size) {
+                remove = true;
+                break;
+            }
+        }
+        if (!remove) {
+            for (int ii = 0; ii < model->n_output_heads; ii++) {
+                int next_index = i*model->n_output_heads+ii*model->n_output_heads+ii;
+                if (next_index > size) {
+                    filtered.push_back(model->eos_token_id);
+                } else {
+                    filtered.push_back(output_tokens[next_index]);
+                }
+            }
+        }
+    }
+}
+
+int parler_tts_runner::generate_from_batch(parler_ubatch & batch, struct tts_response * output) {
+    std::vector<uint32_t> next_decoder_token_ids;
+    next_decoder_token_ids.reserve(model->n_output_heads);
+
+    while (!check_stopping()) {
+        int state = decode(batch);
+        if (state != 0) {
+            return state;
+        }
+        if (!batch.audio_generation) {
+            pctx->prompt_end_position += batch.sequence_length;
+        }
+        if (batch.audio_generation) {
+            sampler->sample(pctx->logits + pctx->current_position * model->n_output_heads * model->output_vocab_size, pctx->output_tokens);
+        }
+        pctx->current_position += batch.sequence_length;
+        next_decoder_token_ids.clear();
+        uint32_t * last_outputs = (pctx->output_tokens.data() + (int) pctx->output_tokens.size() - model->n_output_heads);
+        for (int i = 0; i < model->n_output_heads; i++) {
+            next_decoder_token_ids.push_back(batch.current_step > i ? pctx->eos_seen[i] ? model->eos_token_id : last_outputs[i] : model->bos_token_id);
+        }
+        batch = parler_ubatch{
+            true, 0, 9, 1, nullptr, next_decoder_token_ids.data(), &pctx->current_position, nullptr, batch.current_step+1
+        };
+    }
+
+    std::vector<uint32_t> filtered_output_tokens;
+    adjust_output_tokens(pctx->output_tokens, filtered_output_tokens);
+    dac_runner->run(filtered_output_tokens.data(), (int32_t) filtered_output_tokens.size() / model->n_output_heads, output);
+    return 0;
+}
+
+int parler_tts_runner::generate_audio_tokens(std::string sentence) {
+    parler_ubatch batch = batch_from_sentence(sentence, model, tokenizer);
+    pctx->reset(model->n_output_heads);
+    sampler->reset();
+    int32_t seq_id = std::mt19937(std::random_device{}())();
+    if (!kv_self) {
+        kv_self = new parler_kv_cache;
+        if (!parler_kv_cache_init(kv_self, model, pctx, seq_id)) {
+            return 1;
+        }
+    }
+
+    std::vector<uint32_t> next_decoder_token_ids;
+    next_decoder_token_ids.reserve(model->n_output_heads);
+
+    while (!check_stopping()) {
+        int state = decode(batch);
+        if (state != 0) {
+            return state;
+        }
+        if (!batch.audio_generation) {
+            pctx->prompt_end_position += batch.sequence_length;
+        }
+        if (batch.audio_generation) {
+            sampler->sample(pctx->logits + pctx->current_position * model->n_output_heads * model->output_vocab_size, pctx->output_tokens);
+        }
+        pctx->current_position += batch.sequence_length;
+        next_decoder_token_ids.clear();
+        uint32_t * last_outputs = (pctx->output_tokens.data() + (int) pctx->output_tokens.size() - model->n_output_heads);
+        for (int i = 0; i < model->n_output_heads; i++) {
+            next_decoder_token_ids.push_back(batch.current_step > i ? pctx->eos_seen[i] ? model->eos_token_id : last_outputs[i] : model->bos_token_id);
+        }
+        batch = parler_ubatch{
+            true, 0, 9, 1, nullptr, next_decoder_token_ids.data(), &pctx->current_position, nullptr, batch.current_step+1
+        };
+    }
+
+    return 0;
+}
+
+void parler_tts_runner::just_audio_token_decode(uint32_t * tokens, int32_t sq_len, struct tts_response * outputs) {
+    dac_runner->run(tokens, sq_len, outputs);
+}
+
+int parler_tts_runner::generate(std::string sentence, struct tts_response * output, int32_t seq_id) {
+    parler_ubatch batch = batch_from_sentence(sentence, model, tokenizer);
+    pctx->reset(model->n_output_heads);
+    sampler->reset();
+    if (pctx->seq_id != seq_id || seq_id == -1) {
+        seq_id = std::mt19937(std::random_device{}())();
+        pctx->current_position = 0;
+        if (!kv_self) {
+            kv_self = new parler_kv_cache;
+            if (!parler_kv_cache_init(kv_self, model, pctx, seq_id)) {
+                return 1;
+            }
+        }
+    } else {
+        if (!adjust_for_sequence_continuation(batch)) {
+            return 2;
+        }
+    }
+    return generate_from_batch(batch, output);
+}

otherarch/ttscpp/src/parler_model.h

@@ -0,0 +1,225 @@
+#ifndef parler_model_h
+#define parler_model_h
+
+#include "dac_model.h"
+#include "t5_encoder_model.h"
+#include "sampler.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<struct ggml_tensor*> embds;
+    std::vector<parler_layer*> layers;
+    std::vector<struct ggml_tensor*> 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<bool> 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<uint32_t> 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<struct ggml_tensor *> k_l;
+    std::vector<struct ggml_tensor *> 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<uint32_t> & output_tokens, std::vector<uint32_t> & 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

otherarch/ttscpp/src/sampler.cpp

@@ -0,0 +1,204 @@
+#include "sampler.h"
+
+void sampler::sample(float * logits, std::vector<uint32_t> & output_tokens) {
+    // assume that we are pointing to the start of the first token output;
+    if (!do_sample) {
+        return max(logits, output_tokens);
+    }
+    std::vector<uint32_t> max_vals;
+    // the max_head_probs variable is used when top-p is applied but exists to address the case in which top-k and top-p cause the cumulative probability of the nucleus to beless than or 
+    // equal to top_p;
+    std::vector<float> max_head_probs;
+
+    // This allows us to perform an effective softmax without logarithms or big number calculations.
+    // Additionally by avoiding large number division we drastically improve the stability of
+    // our softmax implementation;
+    max(logits, max_vals);
+
+    std::vector<std::vector<size_t>> picks;
+    bool use_nucleus_sampling = false;
+    bool performed_softmax = false;
+
+    if (top_p < 1.0) {
+        // if we are nucleus sampling via top-p then we need to perform softmax over the samples before getting top_k samples, so that we don't trim beyond top_p.
+        // Otherwise, if we are not performing top-p sampling then it is more efficient to perform softmax after getting the top_k nucleus.
+        softmax(logits, picks, max_vals);
+        performed_softmax = true;
+    }
+    if (top_k > 0 && top_k < vocab_size) {
+        picks = topk(logits, performed_softmax);
+        use_nucleus_sampling = true;
+    }
+
+    if (top_p >= 1.0) {
+        softmax(logits, picks, max_vals);
+        performed_softmax = true;
+    }
+
+    if (top_p < 1.0) {
+        topp(logits, picks, max_head_probs);
+        use_nucleus_sampling = true;
+    }
+
+    bool has_repetition_penalty = repetition_penalty != 1.0;
+    if (has_repetition_penalty && (last_token_ids.size() == 0 || repetition_counts.size() == 0)) {
+        reset();
+    }
+    std::minstd_rand gen(std::random_device{}());
+    std::uniform_real_distribution<float> dist(0.0f, 1.0f);
+    for (int i = 0; i < n_output_heads; i++) {
+        float assignment =  top_p < 1.0 ? dist(gen) * max_head_probs[i] : dist(gen);
+        float cumulative = 0.0f;
+        for (uint32_t j = 0; j < (use_nucleus_sampling ? picks[i].size() : vocab_size); j++) {
+            int ii = use_nucleus_sampling ? (int) picks[i][j] : j;
+            cumulative += *(logits+(i*vocab_size+ii));
+            // with top_k and top_p it is possible for the assignment to be greater than the cumulative value
+            if (assignment <= cumulative || ii >= vocab_size + 1 || j >= picks[i].size() - 1) {
+                if (has_repetition_penalty) {
+                    if (last_token_ids[i] != ii) {
+                        repetition_counts[i] = 0;
+                    }
+                    last_token_ids[i] = ii;
+                    repetition_counts[i] += 1;
+                }
+                output_tokens.push_back(ii);
+                break;
+            }
+        }
+    }
+}
+
+void sampler::reset() {
+    if (repetition_penalty != 1.0) {
+        last_token_ids.clear();
+        repetition_counts.clear();
+        for (int i = 0; i < n_output_heads; i++) {
+            last_token_ids.push_back(-1);
+            repetition_counts.push_back(0);
+        }
+    }
+}
+
+void sampler::softmax(float * logits, std::vector<std::vector<size_t>> picks, std::vector<uint32_t> max_indices) {
+    bool use_nucleus_sampling = picks.size() > 0;
+    bool has_repetition_penalty = repetition_penalty != 1.0f;
+    bool has_temperature = temperature != 1.0f;
+    for (int i = 0; i < n_output_heads; i++) {
+        float cumsum = 0.0;
+        float max_val = logits[i*vocab_size + max_indices[i]];
+        if (has_repetition_penalty && last_token_ids[i] == max_indices[i]) {
+            max_val /= (pow(repetition_penalty, repetition_counts[i]));
+        }
+        if (has_temperature) {
+            max_val /= temperature;
+        }
+        for (int j = 0; j < (use_nucleus_sampling ? picks[i].size() : vocab_size); j++) {
+            int ii = use_nucleus_sampling ? (int) picks[i][j] : j;
+            int index = i * vocab_size + ii;
+            float v = *(logits + index);
+            if (has_repetition_penalty && last_token_ids[i] == ii) {
+                v /= (pow(repetition_penalty, repetition_counts[i]));
+            }
+            if (has_temperature) {
+                v /= temperature;
+            }
+            v = expf(v - max_val);
+            cumsum += v;
+            logits[index] = v;
+        }
+        for (int j = 0; j < (use_nucleus_sampling ? picks[i].size() : vocab_size); j++) {
+            int ii = use_nucleus_sampling ? picks[i][j] : j;
+            int index = i * vocab_size + ii;
+            float v = *(logits + index);
+            logits[index] = v / cumsum;
+        }
+    }
+}
+
+void sampler::topp(float * logits, std::vector<std::vector<size_t>> & picks, std::vector<float> & max_head_probs) {
+    if (picks.empty()) {
+        // we need to get the softmaxed logits ordered
+        for (int i = 0; i < n_output_heads; i++) {
+            std::vector<size_t> head_picks(vocab_size);
+            iota(head_picks.begin(), head_picks.end(), 0);
+            // have to sort with repetition penalty applied so not to inavertently trim our nucleus size.
+            std::sort(head_picks.begin(), head_picks.end(), [&logits, &i, this](size_t s1, size_t s2) {
+                float v1 = logits[i*vocab_size+s1];
+                float v2 = logits[i*vocab_size+s2];
+                return v1 > v2;
+            });
+
+            picks.push_back(head_picks);
+        }
+    }
+    // if we didn't already perform topk or if the probable sum of topk logits is greater than top_p then we need to trim.
+    for (int i = 0; i < n_output_heads; i++) {
+        float prob_sum = 0.0f;
+        int trim_to = -1;
+        for (int ii = 0; ii < picks[i].size(); ii++) {
+            prob_sum += logits[i*vocab_size+picks[i][ii]];
+            if (prob_sum >= top_p) {
+                trim_to = ii+1;
+                break;
+            }
+        }
+        max_head_probs.push_back(std::min(prob_sum, top_p));
+        if (trim_to > 0) {
+            picks[i] = std::vector<size_t>(picks[i].begin(), picks[i].begin()+trim_to);
+        }
+    }
+}
+
+std::vector<std::vector<size_t>> sampler::topk(float * logits, bool performed_softmax) {
+    bool has_repetition_penalty = repetition_penalty != 1.0f;
+    std::vector<std::vector<size_t>> head_picks;
+    if (vocab_size < top_k) {
+        // technically we should never get here, but lets be protective.
+        for (int i = 0; i < n_output_heads; i++) {
+            std::vector<size_t> picks(vocab_size);
+            iota(picks.begin(), picks.end(), 0);
+            head_picks.push_back(picks);
+        }
+        return head_picks;
+    }
+    for (int i = 0; i < n_output_heads; i++) {
+        std::vector<size_t> picks(vocab_size);
+        iota(picks.begin(), picks.end(), 0);
+        // have to sort with repetition penalty applied so not to inavertently trim our nucleus size.
+        std::sort(picks.begin(), picks.end(), [&logits, &i, &has_repetition_penalty, &performed_softmax, this](size_t s1, size_t s2) {
+            float v1 = logits[i*vocab_size+s1];
+            float v2 = logits[i*vocab_size+s2];
+            if (!performed_softmax) {
+                if (has_repetition_penalty && last_token_ids[i] == s1) {
+                    v1 /= (pow(repetition_penalty, repetition_counts[i]));
+                } else if (has_repetition_penalty && last_token_ids[i] == s2) {
+                    v2 /= (pow(repetition_penalty, repetition_counts[i]));
+                }
+            }
+            return v1 > v2;
+        });
+        head_picks.push_back(std::vector<size_t>(picks.begin(), picks.begin() + top_k));
+    }
+    return head_picks;
+}
+
+void sampler::max(float * logits, std::vector<uint32_t> & output_tokens) {
+    bool has_repetition_penalty = repetition_penalty != 1.0f;
+    for (int i = 0; i < n_output_heads; i++) {
+        float max = -INFINITY;
+        uint32_t token_id = 0;
+        for (uint32_t ii = 0; ii < vocab_size; ii++) {
+            float v = *(logits+i*vocab_size+ii);
+            // while repetition penalty will never be used for maximum token selection, it is used for the logarithmic stabilization of 
+            // the softmax function in which case it is possible for repetition counts to be set.
+            if (has_repetition_penalty && last_token_ids[i] == ii) {
+                v /= (pow(repetition_penalty, repetition_counts[i]));
+            }
+            if (v > max) {
+                max = v;
+                token_id = ii;
+            }
+        }
+        output_tokens.push_back(token_id);
+    }
+}

otherarch/ttscpp/src/sampler.h

@@ -0,0 +1,33 @@
+#ifndef sampler_h
+#define sampler_h
+
+#include <stdint.h>
+#include <vector>
+#include <random>
+#include <numeric>
+#include <algorithm>
+
+// currently this is only built to support single sequence output sampling without beam search.
+struct sampler {
+    // These default configurations are based on the generation configuration for Parler TTS Mini (version 1.0)
+    uint32_t n_output_heads = 9;
+    uint32_t eos_token_id = 1024;
+    uint32_t vocab_size = 1088;
+    float temperature = 1.0f;
+    uint32_t top_k = 0;
+    float top_p = 1.0f;
+    float repetition_penalty = 1.0f;
+    std::vector<int32_t> last_token_ids;
+    std::vector<uint32_t> repetition_counts;
+    bool do_sample = true;
+    bool apply_softmax = true;
+    
+    void sample(float * logits, std::vector<uint32_t> & output_tokens);
+    void softmax(float * logits, std::vector<std::vector<size_t>> picks, std::vector<uint32_t> max_indices);
+    void max(float * logits, std::vector<uint32_t> & output_tokens);
+    std::vector<std::vector<size_t>> topk(float * logits, bool performed_softmax);
+    void topp(float * logits, std::vector<std::vector<size_t>> & picks, std::vector<float> & max_head_probs);
+    void reset();
+};
+
+#endif

otherarch/ttscpp/src/snac_model.cpp

@@ -0,0 +1,209 @@
+#include "snac_model.h"
+
+void snac_model::prep_constants(gguf_context * meta) {
+    int heads_key = gguf_find_key(meta, "snac.audio_token_channels");
+    if (heads_key != -1) {
+        n_heads = gguf_get_val_u32(meta, heads_key);
+    }
+
+    int sampling_factor_key = gguf_find_key(meta, "snac.up_sampling_factor");
+    if (sampling_factor_key != -1) {
+        up_sampling_factor = gguf_get_val_u32(meta, sampling_factor_key);
+    }
+
+    int max_gen_key = gguf_find_key(meta, "snac.max_generation_size");
+    if (max_gen_key != -1) {
+        max_generation_size = gguf_get_val_u32(meta, max_gen_key);
+    }
+}
+
+void snac_model::prep_layers(gguf_context * meta) {
+    for (int i = 0; i < n_heads; i++) {
+        quantizer_layers.push_back(general_neural_audio_codec::residual_vector_quantize_layer{});
+    }
+
+    for (int i = 0; i < n_layers; i++) {
+        std::string stride_key = "snac.snac_layer_stride_" + std::to_string(i);
+        std::string padding_key = "snac.snac_layer_padding_" + std::to_string(i);
+        std::string grouping_key = "snac.snac_layer_grouping_" + std::to_string(i);
+        int layer_stride_key = gguf_find_key(meta, stride_key.c_str());
+        if (layer_stride_key == -1) {
+            TTS_ABORT("key %s must be specified in gguf file inorder to initialize the SNAC audio decoder.", stride_key.c_str());
+        }
+        int layer_padding_key = gguf_find_key(meta, padding_key.c_str());
+        if (layer_padding_key == -1) {
+            TTS_ABORT("key %s must be specified in gguf file inorder to initialize the SNAC audio decoder.", padding_key.c_str());
+        }
+        int layer_grouping_key = gguf_find_key(meta, grouping_key.c_str());
+        if (layer_grouping_key == -1) {
+            TTS_ABORT("key %s must be specified in gguf file inorder to initialize the SNAC audio decoder.", grouping_key.c_str());
+        }
+        layers.push_back(
+            general_neural_audio_codec::layer{
+                gguf_get_val_u32(meta, layer_padding_key),
+                gguf_get_val_u32(meta, layer_stride_key),
+                gguf_get_val_u32(meta, layer_grouping_key)
+            }
+        );
+    }
+}
+
+void snac_model::assign_weight(std::string name, ggml_tensor * tensor) {
+    if (name == "alpha_out") {
+        snake_alpha = ggml_dup_tensor(ctx, tensor);
+        set_tensor(snake_alpha, tensor);
+    } else if (name == "in.weight") {
+        in_conv_kernel = ggml_dup_tensor(ctx, tensor);
+        set_tensor(in_conv_kernel, tensor);
+    } else if (name == "in.bias") {
+        in_conv_bias = ggml_dup_tensor(ctx, ggml_transpose(ctx, tensor));
+        set_tensor(in_conv_bias, tensor);
+    } else if (name == "up.weight") {
+        up_conv_kernel = ggml_dup_tensor(ctx, tensor);
+        set_tensor(up_conv_kernel, tensor);
+    } else if (name == "up.bias") {
+        up_conv_bias = ggml_dup_tensor(ctx, ggml_transpose(ctx, tensor));
+        set_tensor(up_conv_bias, tensor);
+    } else if (name == "final.weight") {
+        out_conv_kernel = ggml_dup_tensor(ctx, tensor);
+        set_tensor(out_conv_kernel, tensor);
+    } else if (name == "final.bias") {
+        out_conv_bias = ggml_dup_tensor(ctx, ggml_transpose(ctx, tensor));
+        set_tensor(out_conv_bias, tensor);
+    } else if (has_prefix(name, "layers")) {
+        auto pair = parse_layer_count(name);
+        int l = pair.first;
+        std::string lt_name = pair.second;
+        general_neural_audio_codec::assign_to_layer((tts_model *) this, layers[l], lt_name, tensor);
+    } else if (has_prefix(name, "quantizers")) {
+        auto pair = parse_layer_count(name);
+        int l = pair.first;
+        std::string lt_name = pair.second;
+        general_neural_audio_codec::assign_to_quantize_layer((tts_model *) this, quantizer_layers[l], lt_name, tensor);
+    }
+}
+
+static struct ggml_tensor * snac_build_audio_inputs(struct ggml_context * ctx, struct snac_context * sctx, size_t sequence_length, std::vector<general_neural_audio_codec::residual_vector_quantize_layer> layers) {
+    struct ggml_tensor * embd;
+    // these devisors represent the discreate repeats performed against each of the three input heads.
+    sctx->inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sequence_length / 4 + sequence_length / 2 + sequence_length);
+    ggml_set_input(sctx->inp_tokens);
+    size_t last_stride = 0;
+    for(int i = 0; i < sctx->model->n_heads; i++) {
+        auto quantize_layer = sctx->model->quantizer_layers[i];
+        struct ggml_tensor * inp_head = ggml_cont(ctx, ggml_view_1d(ctx, sctx->inp_tokens, sequence_length / sctx->model->repeats[i], last_stride));
+        last_stride += (sequence_length / sctx->model->repeats[i]) * ggml_element_size(sctx->inp_tokens);
+        struct ggml_tensor * code = general_neural_audio_codec::build_quantize_layer(ctx, inp_head, quantize_layer);
+        if (sctx->model->repeats[i] > 1) {
+            // this manipulation is equivalent to repeat_interleave against the first dimension of the tensor
+            code = ggml_repeat(ctx, ggml_cont_3d(ctx, code, 1, code->ne[0], code->ne[1]), ggml_new_tensor_3d(ctx, GGML_TYPE_F32, sctx->model->repeats[i], code->ne[0], sctx->model->embd));
+            code = ggml_cont_2d(ctx, code, sequence_length, code->ne[2]);
+        }
+        if (i == 0) {
+            embd = code;
+        } else {
+            embd = ggml_add(ctx, embd, code);
+        }
+    }
+    return embd;
+}
+
+snac_context * build_new_snac_context(struct snac_model * model, int n_threads, bool use_cpu) {
+    snac_context * sctx = new snac_context(model, n_threads);
+    if (!use_cpu) {
+#ifdef GGML_USE_METAL
+        sctx->backend = ggml_backend_metal_init();
+#endif
+    }
+    sctx->backend_cpu = ggml_backend_cpu_init();
+    sctx->set_threads();
+    sctx->build_schedule();
+    sctx->buf_compute_meta.resize(ggml_tensor_overhead()*model->max_nodes() + ggml_graph_overhead_custom(model->max_nodes(), false));
+    return sctx;
+}
+
+void snac_runner::prepare_post_load() {
+    ggml_cgraph * gf = build_snac_graph(model->max_generation_size);
+    sctx->prep_schedule(gf);
+}
+
+struct ggml_cgraph * snac_runner::build_snac_graph(size_t sequence_length) {
+    init_build();
+    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
+
+    struct ggml_tensor * cur;
+    struct ggml_tensor * inputs;
+
+    sctx->noise = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, model->noise_steps_sum * sequence_length);
+    ggml_set_input(sctx->noise);
+
+    inputs = snac_build_audio_inputs(ctx, sctx, sequence_length, model->quantizer_layers);
+    cur = ggml_conv_1d_dw_tts(ctx, model->in_conv_kernel, inputs, 1, 3, 1);
+    cur = ggml_add(ctx, cur, model->in_conv_bias);
+    cur = ggml_conv_1d_tts(ctx, model->up_conv_kernel, cur, 1, 0, 1);
+    cur = ggml_add(ctx, cur, model->up_conv_bias);
+    size_t noise_offset = 0;
+    for (int l = 0; l < model->layers.size(); l++) {
+        auto layer = model->layers[l];
+        struct ggml_tensor * noise = ggml_cont(ctx, ggml_view_1d(ctx, sctx->noise, model->noise_steps[l] * sequence_length, noise_offset));
+        noise_offset += model->noise_steps[l] * sequence_length * sizeof(float);
+        cur = general_neural_audio_codec::build_layer(ctx, cur, layer, noise);
+    }
+    cur = snake_1d(ctx, model->snake_alpha, cur);
+    cur = ggml_conv_1d_tts(ctx, model->out_conv_kernel, cur, 1, 3, 1);
+    cur = ggml_add(ctx, cur, model->out_conv_bias);
+    cur = ggml_tanh(ctx, cur);
+    ggml_build_forward_expand(gf, cur);
+    free_build();
+    return gf;
+}
+
+void snac_runner::set_inputs(std::vector<std::vector<uint32_t>> & tokens) {
+    ggml_backend_tensor_set(
+        sctx->inp_tokens, tokens[0].data(), 0,
+        tokens[0].size()*ggml_element_size(sctx->inp_tokens)
+    );
+
+    ggml_backend_tensor_set(
+        sctx->inp_tokens, tokens[1].data(), tokens[0].size() * ggml_element_size(sctx->inp_tokens),
+        tokens[1].size() * ggml_element_size(sctx->inp_tokens)
+    );
+
+    ggml_backend_tensor_set(
+        sctx->inp_tokens, tokens[2].data(),
+        tokens[1].size()*ggml_element_size(sctx->inp_tokens)+tokens[0].size()*ggml_element_size(sctx->inp_tokens),
+        tokens[2].size()*ggml_element_size(sctx->inp_tokens)
+    );
+    size_t sequence_length = tokens[2].size();
+    random_normal_gen(model->noise_steps_sum * sequence_length, (float*) sctx->noise->data);
+}
+
+void snac_runner::run(std::vector<std::vector<uint32_t>> & tokens, struct tts_response * outputs) {
+    size_t sequence_length = tokens[2].size();
+    ggml_backend_sched_reset(sctx->sched);
+
+    sctx->prep_output_buffer(model->max_generation_size * model->up_sampling_factor * sizeof(float));
+
+    outputs->data = sctx->logits;
+    ggml_backend_buffer_clear(sctx->buf_output, 0);
+
+    struct ggml_cgraph * gf = NULL;
+    gf = build_snac_graph(sequence_length);
+
+    // the output is always the last tensor in the graph
+    struct ggml_tensor * result = gf->nodes[gf->n_nodes - 1];
+    ggml_backend_sched_alloc_graph(sctx->sched, gf);
+
+    set_inputs(tokens);
+
+    ggml_backend_sched_graph_compute_async(sctx->sched, gf);
+
+    sctx->get_ggml_node_data(result, outputs->data, sequence_length*sizeof(float)*model->up_sampling_factor);
+
+    // Reset state for the next token before backend sync, to allow the CPU activities in the reset to
+    // overlap with device computation.
+    ggml_backend_sched_reset(sctx->sched);
+    outputs->n_outputs = sequence_length * model->up_sampling_factor;
+    return;
+}
+

otherarch/ttscpp/src/snac_model.h

@@ -0,0 +1,86 @@
+#pragma once
+
+#include "general_neural_audio_codec.h"
+
+// SNAC, Scale Neural Audio Codec, is another neural audio codec much like DAC.
+// The key differences are that it uses grouping in the residual units of its layers,
+// performs a repeat_interleave over the second and third input channels, applies 
+// a noise convolutional layer after input encoding for each layer, and applies
+// an extra convolutional layer before residual layers are applied.
+struct snac_model : tts_model {
+    // general configuration from SNAC as used by Orpheus
+    uint32_t n_layers = 4;
+    uint32_t n_heads = 3;
+    uint32_t up_sampling_factor = 512;
+    uint32_t embd = 768;
+    size_t max_generation_size = 2580;
+    uint32_t repeats[3] = {4, 2, 1};
+    // configuration for adding noise
+    uint32_t noise_steps[4] = {8, 64, 256, 512};
+    uint32_t noise_steps_sum = 840;
+    bool use_noise = true;
+    
+    struct ggml_tensor * repeat_interleave_buffer;
+
+    struct ggml_tensor * in_conv_kernel;
+    struct ggml_tensor * in_conv_bias;
+    struct ggml_tensor * up_conv_kernel;
+    struct ggml_tensor * up_conv_bias;
+    struct ggml_tensor * out_conv_kernel;
+    struct ggml_tensor * out_conv_bias;
+    struct ggml_tensor * snake_alpha;
+    std::vector<general_neural_audio_codec::layer> layers;
+    std::vector<general_neural_audio_codec::residual_vector_quantize_layer> quantizer_layers;
+
+    void assign_weight(std::string name, ggml_tensor * weight);
+    void prep_constants(gguf_context * meta);
+    void prep_layers(gguf_context * meta);
+    void post_load_assign();
+    void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only) {
+        prep_layers(meta_ctx);
+        prep_constants(meta_ctx);
+        tts_model::setup_from_file(meta_ctx, load_context, cpu_only, "snac");
+    }
+};
+
+// the context used for running the snac model
+struct snac_context : runner_context {
+    snac_context(snac_model * model, int n_threads): runner_context(n_threads), model(model) {};
+    
+    struct snac_model * model;
+        
+    struct ggml_tensor * inp_tokens;
+    struct ggml_tensor * noise;
+    
+    void build_schedule() {
+        runner_context::build_schedule(model->max_nodes());
+    }
+};
+
+snac_context * build_new_snac_context(struct snac_model * model, int n_threads, bool use_cpu = true);
+
+static struct ggml_tensor * snac_build_audio_inputs(struct ggml_context * ctx, struct snac_context * sctx, size_t sequence_length, std::vector<general_neural_audio_codec::residual_vector_quantize_layer> layers);
+
+// This struct is intended to manage the snac model's graph compilation and compute function.
+struct snac_runner : tts_runner {
+    snac_runner(snac_model * model, snac_context * context): model(model), sctx(context) {};
+    ~snac_runner() {
+        if (ctx) {
+            ggml_free(ctx);
+        }
+        model->free();
+        delete model;
+        delete sctx;
+    }
+    snac_model * model;
+    snac_context * sctx;
+    
+    void init_build() {
+        tts_runner::init_build(&sctx->buf_compute_meta);
+    }
+    
+    void set_inputs(std::vector<std::vector<uint32_t>> & tokens);
+    void prepare_post_load();
+    struct ggml_cgraph * build_snac_graph(size_t sequence_length);
+    void run(std::vector<std::vector<uint32_t>> & tokens, struct tts_response * outputs);
+};

otherarch/ttscpp/src/t5_encoder_model.cpp

@@ -0,0 +1,402 @@
+#include "t5_encoder_model.h"
+
+static const std::map<std::string, t5_tensor> T5_TENSOR_GGUF_LOOKUP = {
+    {"t5encoder.token_embd", T5_EMBD},
+    {"t5encoder.enc.final_layer_norm", T5_NORM},
+    {"t5encoder.down_proj", T5_DOWN_PROJ},
+    {"t5encoder.down_proj_bias", T5_DOWN_PROJ_BIAS},
+    {".attn_norm", T5_LAYER_ATTN_NORM},
+    {".attn_q", T5_LAYER_ATTN_Q},
+    {".attn_k", T5_LAYER_ATTN_K},
+    {".attn_v", T5_LAYER_ATTN_V},
+    {".attn_o", T5_LAYER_ATTN_O},
+    {".attn_rel_b", T5_RELATIVE_BIAS},
+    {".ffn_norm", T5_LAYER_OUT_NORM},
+    {".ffn_gate", T5_LAYER_WI_1},
+    {".ffn_down", T5_LAYER_WO},
+    {".ffn_up", T5_LAYER_WI_0},
+};
+
+void assign_to_t5_layer(t5_encoder * model, t5_layer & layer, std::string name, ggml_tensor * tensor) {
+    try {
+        switch(T5_TENSOR_GGUF_LOOKUP.at(name)) {
+            case T5_LAYER_ATTN_NORM:
+                layer.attn_norm = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.attn_norm, tensor);
+                break;
+            case T5_LAYER_ATTN_Q:
+                layer.q = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.q, tensor);
+                break;
+            case T5_LAYER_ATTN_K:
+                layer.k = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.k, tensor);
+                break;
+            case T5_LAYER_ATTN_V:
+                layer.v = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.v, tensor);
+                break;
+            case T5_LAYER_ATTN_O:
+                layer.o = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.o, tensor);
+                break;
+            case T5_LAYER_OUT_NORM:
+                layer.mlp_norm = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.mlp_norm, tensor);
+                break;
+            case T5_LAYER_WI_1:
+                layer.wi_1 = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.wi_1, tensor);
+                break;
+            case T5_LAYER_WI_0:
+                layer.wi_0 = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.wi_0, tensor);
+                break;
+            case T5_LAYER_WO:
+                layer.wo = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(layer.wo, tensor);
+                break;
+            case T5_RELATIVE_BIAS:
+                model->relative_attn_bias = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->relative_attn_bias, tensor);
+                break;
+            default:
+                fprintf(stdout, "unassigned tensor %s\n", name.c_str());
+                break;
+        }
+    } catch (const std::out_of_range& e) {
+        TTS_ABORT("Error: %s\nTensor, '%s', is not a valid tensor.", e.what(), name.c_str());
+    }
+}
+
+void assign_to_t5_encoder(t5_encoder * model, const std::string name, ggml_tensor * tensor) {
+    if (tensor->data == NULL) {
+        return;
+    }
+    std::string::size_type pos = name.find(".", 0);
+    std::string top_level(name.substr(0, pos));
+    if (T5_TENSOR_GGUF_LOOKUP.find(name) != T5_TENSOR_GGUF_LOOKUP.end()) {
+        switch (T5_TENSOR_GGUF_LOOKUP.at(name)) {
+            case T5_EMBD:
+                model->embd = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->embd, tensor);
+                break;
+            case T5_NORM:
+                model->out_norm = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->out_norm, tensor);
+                break;
+            case T5_DOWN_PROJ:
+                model->down_proj = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->down_proj, tensor);
+                break;
+            case T5_DOWN_PROJ_BIAS:
+                model->down_proj_bias = ggml_dup_tensor(model->ctx, tensor);
+                model->set_tensor(model->down_proj_bias, tensor);
+                break;
+            default:
+                fprintf(stdout, "unassigned tensor %s\n", name.c_str());
+                break;
+        }
+    } else if (top_level == "t5encoder") {
+        auto pair = parse_layer_count(name, 2);
+        int l = pair.first;
+        std::string lt_name = pair.second;
+
+        assign_to_t5_layer(model, model->layers[l], lt_name, tensor);
+    } else {
+        return;
+    }
+}
+
+void t5_encoder::prep_layers(gguf_context * meta) {
+	for (uint32_t i = 0; i < n_layers; i++) {
+		t5_layer l;
+		layers.push_back(l);
+	}
+}
+
+void t5_encoder::prep_constants(gguf_context * meta) {
+	int n_layers_key = gguf_find_key(meta, "t5encoder.block_count");
+    if (n_layers_key != -1) {
+        n_layers = gguf_get_val_u32(meta, n_layers_key);
+    }
+
+	int hidden_size_key = gguf_find_key(meta, "t5encoder.embedding_length");
+    if (hidden_size_key != -1) {
+        hidden_size = gguf_get_val_u32(meta, hidden_size_key);
+    }
+
+	int attn_heads_key = gguf_find_key(meta, "t5encoder.attention.head_count");
+    if (attn_heads_key != -1) {
+        n_attn_heads = gguf_get_val_u32(meta, attn_heads_key);
+    }
+
+    int context_size_key = gguf_find_key(meta, "t5encoder.context_length");
+    if (context_size_key != -1) {
+        max_context_length = gguf_get_val_u32(meta, context_size_key);
+    }
+
+    int bos_token_id_key = gguf_find_key(meta, "tokenizer.ggml.bos_token_id");
+    if (bos_token_id_key != -1) {
+        bos_token_id = gguf_get_val_u32(meta, bos_token_id_key);
+    }    
+
+    int eos_token_id_key = gguf_find_key(meta, "tokenizer.ggml.eos_token_id");
+    if (eos_token_id_key != -1) {
+        eos_token_id = gguf_get_val_u32(meta, eos_token_id_key);
+    }
+
+    int vocab_size_key = gguf_find_key(meta, "t5encoder.vocab_size");
+    if (vocab_size_key == -1) {
+        TTS_ABORT("key 't5encoder.vocab_size' must be specified in gguf file.");
+    }
+    vocab_size = gguf_get_val_u32(meta, vocab_size_key);
+
+    int output_size_key = gguf_find_key(meta, "t5encoder.output_size");
+    if (output_size_key != -1) {
+        output_size = gguf_get_val_u32(meta, output_size_key);
+    }
+}
+
+void t5_encoder::assign_weight(std::string name, ggml_tensor * tensor) {
+    assign_to_t5_encoder(this, name, tensor);
+}
+
+struct t5_context * build_new_t5_context(struct t5_encoder * model, int n_threads, bool use_cpu) {
+	t5_context * t5ctx = new t5_context(model, n_threads);
+    if (!use_cpu) {
+#ifdef GGML_USE_METAL
+        t5ctx->backend = ggml_backend_metal_init();
+#endif
+    }
+    t5ctx->backend_cpu = ggml_backend_cpu_init();
+    t5ctx->set_threads();
+    t5ctx->build_schedule();
+    t5ctx->buf_compute_meta.resize(ggml_tensor_overhead()*model->max_nodes() + ggml_graph_overhead_custom(model->max_nodes(), false));
+    return t5ctx;
+}
+
+static struct ggml_tensor * build_t5_norm(struct ggml_context * ctx, struct ggml_tensor * cur, struct ggml_tensor * weight) {
+	// this is static for all versions of t5 flan
+    float eps = 0.000001;
+    cur = ggml_rms_norm(ctx, cur, eps);
+    cur = ggml_mul(ctx, cur, weight);
+    return cur;
+}
+
+static struct ggml_tensor * build_t5_attn_mask(ggml_context * ctx, struct t5_context *t5ctx, const t5_ubatch & batch) {
+    t5ctx->attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) batch.n_tokens, (int64_t) batch.n_tokens);
+    ggml_set_input(t5ctx->attn_mask);
+
+    return t5ctx->attn_mask;
+}
+
+static struct ggml_tensor * build_t5_pos_bias(ggml_context * ctx, struct ggml_tensor * pos_bucket, struct ggml_tensor * relative_attn_bias) {
+    struct ggml_tensor * pos_bucket_1d = ggml_view_1d(ctx, pos_bucket, pos_bucket->ne[0] * pos_bucket->ne[1], 0);
+    struct ggml_tensor * pos_bias = ggml_get_rows(ctx, relative_attn_bias, pos_bucket_1d);
+
+    pos_bias = ggml_view_3d(ctx, pos_bias, pos_bias->ne[0], pos_bucket->ne[0], pos_bucket->ne[1], ggml_element_size(pos_bias) * pos_bias->ne[0], ggml_element_size(pos_bias) * pos_bias->ne[0] * pos_bucket->ne[0],  0);
+    pos_bias = ggml_permute(ctx, pos_bias, 2, 1, 0, 3);
+    pos_bias = ggml_cont(ctx, pos_bias);
+    return pos_bias;
+}
+
+t5_ubatch t5_runner::build_worst_case_batch()  {
+    struct t5_ubatch batch;
+    batch.n_tokens = model->max_context_length;
+    return batch;
+}
+
+void t5_runner::prepare_post_load() {
+    auto batch = build_worst_case_batch();
+    auto gf = build_t5_graph(batch);
+    t5ctx->prep_schedule(gf);
+}
+
+struct ggml_cgraph * t5_runner::build_t5_graph(t5_ubatch & batch) {
+    init_build();
+    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
+
+    struct ggml_tensor * cur;
+    struct ggml_tensor * inpL;
+    
+    //t5ctx->positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.n_tokens);
+    //ggml_set_input(t5ctx->positions);
+
+    t5ctx->inp_pos_bucket = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, batch.n_tokens, batch.n_tokens);
+    ggml_set_input(t5ctx->inp_pos_bucket);
+
+    t5ctx->inp_tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, batch.n_tokens);
+    ggml_set_input(t5ctx->inp_tokens);
+
+    inpL = ggml_get_rows(ctx, model->embd, t5ctx->inp_tokens);
+
+    struct ggml_tensor * KQ_mask_dec = build_t5_attn_mask(ctx, t5ctx, batch);
+    struct ggml_tensor * pos_bias = build_t5_pos_bias(ctx, t5ctx->inp_pos_bucket, model->relative_attn_bias);
+    
+    for (int l = 0; l < model->n_layers; l++) {
+        struct ggml_tensor * residual = inpL;
+
+        cur = build_t5_norm(ctx, inpL, model->layers[l].attn_norm);
+
+        struct ggml_tensor * attn_out;
+
+        // self-attention
+        {
+            struct ggml_tensor * Qcur = ggml_mul_mat(ctx, model->layers[l].q, cur);
+            struct ggml_tensor * Kcur = ggml_mul_mat(ctx, model->layers[l].k, cur);
+            struct ggml_tensor * Vcur = ggml_mul_mat(ctx, model->layers[l].v, cur);
+
+			Qcur = ggml_reshape_3d(ctx, Qcur, model->head_size, model->n_attn_heads, batch.n_tokens);
+            Kcur = ggml_reshape_3d(ctx, Kcur, model->head_size, model->n_attn_heads, batch.n_tokens);
+
+            struct ggml_tensor * q = ggml_permute(ctx, Qcur, 0, 2, 1, 3);
+            struct ggml_tensor * k = ggml_cont(ctx, ggml_permute(ctx, Kcur, 0, 2, 1, 3));
+
+            struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q);
+            kq = ggml_add(ctx, kq, pos_bias);
+
+            kq = ggml_soft_max_ext(ctx, kq, KQ_mask_dec, 1.0f, 0.0f);
+
+            struct ggml_tensor * v = ggml_cont_3d(ctx, ggml_transpose(ctx, Vcur), batch.n_tokens, model->head_size, model->n_attn_heads);
+            struct ggml_tensor * kqv = ggml_mul_mat(ctx, kq, v);
+            struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 2, 0, 1, 3);
+            attn_out = ggml_cont_2d(ctx, kqv_merged, model->hidden_size, batch.n_tokens);
+            attn_out = ggml_mul_mat(ctx, model->layers[l].o, attn_out);
+        }
+
+        cur = ggml_add(ctx, attn_out, residual);
+        struct ggml_tensor * residualmlp = cur;
+
+        // mlp
+        {
+        	cur = build_t5_norm(ctx, cur, model->layers[l].mlp_norm);
+        	struct ggml_tensor * gate_proj = ggml_mul_mat(ctx, model->layers[l].wi_1, cur);
+        	cur = ggml_mul(ctx, ggml_gelu(ctx, ggml_mul_mat(ctx, model->layers[l].wi_0, cur)), gate_proj);
+        	cur = ggml_mul_mat(ctx, model->layers[l].wo, cur);
+        }
+
+		cur = ggml_add(ctx, cur, residualmlp);
+        inpL = cur;
+    }
+
+    cur = build_t5_norm(ctx, cur, model->out_norm);
+
+    if (model->down_proj) {
+        cur = ggml_mul_mat(ctx, model->down_proj, cur);
+    }
+
+    if (model->down_proj_bias) {
+        cur = ggml_add(ctx, cur, model->down_proj_bias);
+    }
+
+    ggml_build_forward_expand(gf, cur);
+
+    free_build();
+    
+    return gf;
+}
+
+void t5_runner::set_inputs(t5_ubatch & batch) {
+    ggml_backend_tensor_set(t5ctx->inp_tokens, batch.input_tokens, 0, batch.n_tokens*ggml_element_size(t5ctx->inp_tokens));
+    float * attn_mask = nullptr;
+    uint32_t * positions = nullptr;
+    uint32_t * pos_bucket = nullptr;
+    attn_mask = (float *) t5ctx->attn_mask->data;
+    positions = (uint32_t *) t5ctx->positions->data;
+    pos_bucket = (uint32_t *) t5ctx->inp_pos_bucket->data;
+    int n_buckets = (int) model->relative_attn_buckets / 2;
+    int max_exact = (int) n_buckets / 2;
+	float logarithmic_denominator = log(128.0 / max_exact);
+    for (int i = 0; i < batch.n_tokens; i++) {
+        for (int ii = 0; ii < batch.n_tokens; ii++) {
+        	int ab_rpos = abs(i - ii);
+        	int rpos = i - ii;
+            attn_mask[i*batch.n_tokens + ii] = 0.0f; //ii > i ? -INFINITY : 0.0f; 
+            pos_bucket[i*batch.n_tokens + ii] = (uint32_t) (rpos > 0 ? n_buckets : 0) + (ab_rpos < max_exact ? ab_rpos : std::min((n_buckets - 1), (max_exact + (int)((log((ab_rpos / max_exact)) / logarithmic_denominator) * max_exact))));
+        }
+    }
+
+}
+
+void t5_runner::run(uint32_t * input_tokens, uint32_t sequence_length, struct tts_response * outputs) {
+	t5_ubatch batch;
+    batch.input_tokens = input_tokens;
+    batch.n_tokens = sequence_length;
+    ggml_backend_sched_reset(t5ctx->sched);
+    
+    const size_t prev_size = t5ctx->buf_output ? ggml_backend_buffer_get_size(t5ctx->buf_output) : 0;
+    const size_t new_size = model->max_context_length * model->output_size * sizeof(float);
+    
+    if (!t5ctx->buf_output || prev_size < new_size) {
+        if (t5ctx->buf_output) {
+            ggml_backend_buffer_free(t5ctx->buf_output);
+            t5ctx->buf_output = nullptr;
+            t5ctx->logits = nullptr;
+        }
+
+        t5ctx->buf_output = ggml_backend_buft_alloc_buffer(t5ctx->backend_cpu_buffer, new_size);
+    }
+    
+    outputs->data = (float *) ggml_backend_buffer_get_base(t5ctx->buf_output);
+    ggml_backend_buffer_clear(t5ctx->buf_output, 0);
+    struct ggml_cgraph * gf = NULL;
+    gf = build_t5_graph(batch);
+    // the output is always the last tensor in the graph
+    struct ggml_tensor * result = gf->nodes[gf->n_nodes - 1];
+    ggml_backend_sched_alloc_graph(t5ctx->sched, gf);
+    set_inputs(batch);
+
+    ggml_backend_sched_graph_compute_async(t5ctx->sched, gf);
+
+    t5ctx->get_ggml_node_data(result, outputs->data, batch.n_tokens*sizeof(float)*model->output_size);
+
+    // Reset state for the next token before backend sync, to allow the CPU activities in the reset to
+    // overlap with device computation.
+    ggml_backend_sched_reset(t5ctx->sched);
+    outputs->n_outputs = sequence_length;
+    outputs->hidden_size = model->output_size;
+    return;
+}
+
+int t5_runner::generate(std::string prompt, tts_response *response) {
+	std::vector<uint32_t> tokens;
+	tokenizer->tokenize(prompt, tokens);
+    tokens.push_back(model->eos_token_id);
+	run(tokens.data(), (uint32_t) tokens.size(), response);
+	return 0;
+}
+
+struct t5_runner * text_encoder_from_file(std::string file_path, int n_threads, unigram_tokenizer * tokenizer, bool cpu_only) {
+    t5_encoder * model = new t5_encoder;
+    ggml_context * weight_ctx = NULL;
+
+    struct gguf_init_params params = {
+        /*.no_alloc   =*/ false,
+        /*.ctx        =*/ &weight_ctx,
+    };
+    gguf_context * meta_ctx = gguf_init_from_file(file_path.c_str(), params);
+    if (!meta_ctx) {
+        TTS_ABORT("%s failed for file %s\n", __func__, file_path.c_str());
+    }
+    if (!tokenizer) {
+        tokenizer = unigram_tokenizer_from_gguf(meta_ctx);
+    }
+    if (!tokenizer->init) {
+        tokenizer->initialize_tokenizer();
+    }
+    model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+
+    // TODO: change this weight assignment pattern to mirror llama.cpp
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        model->assign_weight(cur->name, cur);
+    }
+
+    struct t5_context * t5ctx = build_new_t5_context(model, n_threads, cpu_only);
+    struct t5_runner * runner = new t5_runner(model, t5ctx, tokenizer);
+    runner->prepare_post_load();
+    gguf_free(meta_ctx);
+    ggml_free(weight_ctx);
+
+    return runner;
+}

otherarch/ttscpp/src/t5_encoder_model.h

@@ -0,0 +1,130 @@
+#ifndef t5_encoder_model_h
+#define t5_encoder_model_h
+
+#include "tts_model.h"
+#include "tokenizer.h"
+
+
+enum t5_tensor {
+    T5_EMBD,
+    T5_NORM,
+    T5_DOWN_PROJ,
+    T5_DOWN_PROJ_BIAS,
+    T5_RELATIVE_BIAS,
+    T5_LAYER_ATTN_Q,
+    T5_LAYER_ATTN_K,
+    T5_LAYER_ATTN_V,
+    T5_LAYER_ATTN_O,
+    T5_LAYER_ATTN_NORM,
+    T5_LAYER_WI_0,
+    T5_LAYER_WI_1,
+    T5_LAYER_WO,
+    T5_LAYER_OUT_NORM,
+};
+
+struct t5_layer {
+    struct ggml_tensor * q;
+    struct ggml_tensor * k;
+    struct ggml_tensor * v;
+    struct ggml_tensor * o;
+    struct ggml_tensor * attn_norm;
+    struct ggml_tensor * wi_0;
+    struct ggml_tensor * wi_1;
+    struct ggml_tensor * wo;
+    struct ggml_tensor * mlp_norm;
+};
+
+// this struct maintains the static tensors for a t5_encoder model
+// the defautl configuration is form copied from standard configuration for
+// flan-t5-xl. Note this model is slightly different from a standard t5 encoder.
+// Specifically this model has a down projection which converts the text encoder's
+// hidden size to the hidden size of the parler decoder.
+struct t5_encoder : tts_model {
+    // These configs  are essentially built for the 44khZ 8kbps standard DAC model audio encoder and decoder
+    uint32_t n_layers = 24;
+    uint32_t n_attn_heads = 32;
+    uint32_t head_size = 64;
+    uint32_t hidden_size = 2048;
+    uint32_t relative_attn_buckets = 32;
+    uint32_t eos_token_id = 1;
+    uint32_t bos_token_id = 0;
+    uint32_t max_context_length = 512;
+    uint32_t output_size = 1536;
+    uint32_t vocab_size;
+
+    struct ggml_tensor * embd;
+    struct ggml_tensor * relative_attn_bias;
+    struct ggml_tensor * out_norm;
+    struct ggml_tensor * down_proj = nullptr;
+    struct ggml_tensor * down_proj_bias = nullptr;
+    std::vector<t5_layer> layers;
+
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    void prep_layers(gguf_context * meta);
+    void prep_constants(gguf_context * meta);
+    void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only = true) {
+        prep_constants(meta_ctx);
+        prep_layers(meta_ctx);
+        tts_model::setup_from_file(meta_ctx, load_context, cpu_only, "t5encoder", 1.25);
+    }
+};
+
+// For assigning weights from gguf file to local model.
+void assign_to_t5_encoder(t5_encoder * model, const std::string name, ggml_tensor * tensor);
+void assign_to_t5_layer(t5_encoder * model, t5_layer & layer, std::string name, ggml_tensor * tensor);
+
+struct t5_context : runner_context {
+    t5_context(t5_encoder * model, int n_threads): runner_context(n_threads), model(model) {};
+    
+    struct t5_encoder * model;
+    
+    struct ggml_tensor * inp_tokens;
+    struct ggml_tensor * positions;
+    struct ggml_tensor * attn_mask;
+    struct ggml_tensor * inp_pos_bucket;
+    
+    void build_schedule() {
+        runner_context::build_schedule(model->max_nodes());
+    }
+};
+
+struct t5_context * build_new_t5_context(struct t5_encoder * model, int n_threads, bool use_cpu = true);
+
+struct t5_ubatch {
+    size_t n_tokens; // the number of tokens in our encoded sequence
+    uint32_t * input_tokens;    // [n_tokens]
+};
+
+static struct ggml_tensor * build_t5_norm(struct ggml_context * ctx, struct ggml_tensor * cur, struct ggml_tensor * weight);
+static struct ggml_tensor * build_t5_attn_mask(ggml_context * ctx, struct t5_context *t5ctx, const t5_ubatch & batch);
+
+// This struct is intended to manage the t5 encoder model's graph compilation and compute function.
+struct t5_runner : tts_runner {
+    t5_runner(t5_encoder * model, t5_context * context, unigram_tokenizer * tokenizer): model(model), t5ctx(context), tokenizer(tokenizer) {};
+    ~t5_runner() {
+        if (ctx) {
+            ggml_free(ctx);
+        }
+        model->free();
+        delete model;
+        delete t5ctx;
+    }
+    struct unigram_tokenizer * tokenizer;
+    t5_encoder * model;
+    t5_context * t5ctx;
+
+    void init_build() {
+        tts_runner::init_build(&t5ctx->buf_compute_meta);
+    }
+    
+    void prepare_post_load();
+    struct t5_ubatch build_worst_case_batch();
+    void set_inputs(t5_ubatch & batch);
+    struct ggml_cgraph * build_t5_graph(t5_ubatch & batch);
+    void run(uint32_t * input_tokens, uint32_t sequence_length, struct tts_response * outputs);
+    int generate(std::string prompt, struct tts_response * response);
+};
+
+struct t5_runner * text_encoder_from_file(std::string file_path, int n_threads, unigram_tokenizer * tokenizer, bool cpu_only = true);
+
+#endif

otherarch/ttscpp/src/tokenizer.cpp

@@ -0,0 +1,331 @@
+#include "tokenizer.h"
+
+void token_trie::add(const std::string & gram, uint32_t token) {
+    _add(gram, token, 0);
+}
+
+void token_trie::_add(const std::string & gram, uint32_t new_token, size_t index) {
+    if (index >= gram.size()) {
+        has_value = true;
+        token = new_token;
+        return;
+    }
+    const char c = gram[index];
+    auto res = children.find(c);
+    if (res != children.end()) {
+        res->second._add(gram, new_token, index + 1);
+    } else {
+        struct token_trie nt{};
+        nt._add(gram, new_token, index + 1);
+        children[c] = nt;
+    }
+}
+
+const struct token_trie * token_trie::traverse(const char c) const {
+    auto res = children.find(c);
+    if (res != children.end()) {
+        return &res->second;
+    }
+
+    return NULL;
+}
+
+size_t unicode_len_utf8_tts(char src) {
+    const size_t lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 };
+    uint8_t highbits = static_cast<uint8_t>(src) >> 4;
+    return lookup[highbits];
+}
+
+void unigram_tokenizer::initialize_tokenizer() {
+    for (const auto it : vocab) {
+        root_trie.add(it.first, it.second);
+    }
+    init = true;
+}
+
+// the general approach here is to find the character grams that sum to the max possible value over the entire text sequence.
+// The particular algorithm used here effectively works by walking the text and at each index storing the max value of all possible gram combinations
+// we can then reverse that sequence to pick the best possible tokens.
+void unigram_tokenizer::tokenize(const std::string & text, std::vector<uint32_t> & tokens) {
+    if (!init) {
+        TTS_ABORT("Error: %s\nTokenizer must be initialized before #tokenize is called.");
+    }
+    // the parler tokenizer's normalizer (i.e. the bert normalizer implemented by huggingface tokenizers libs) only deduplicates and strips extra spaces and
+    // optionally handles chinese characters and accents (neither of which are currently supported here).
+    std::string normalized = text;
+    if (dedupe_spaces) {
+        normalized = " " + std::regex_replace(text, duped_spaces, " ");
+    }
+
+    size_t text_length = normalized.size();
+
+    // initialize score_sum to neg infinity so it will be always lower than sums of token scores
+    std::vector<struct result> results(text_length + 1, {unk_token, 0, -INFINITY});
+    results[0] = { unk_token, 0, 0 };
+
+    size_t offset = 0;
+
+    while (offset < text_length) {
+        size_t current_offset = offset;
+        // pulled this directly from llama.cpp; I suspect that this is for handling of non-utf8 steps (to be marked as unknown tokens)
+        size_t n_utf8_code_units = std::min<size_t>(unicode_len_utf8_tts(normalized[offset]), text_length - offset);
+
+        bool found_unknown = true;
+        const struct result & current_best = results[offset];
+
+        // find the current branch in the trie
+        const struct token_trie * node = root_trie.traverse(normalized[current_offset++]);
+        // search for the next token
+        while (current_offset <= text_length && node != NULL) {
+            // check if this is a complete token (it could just be an unkown step between two tokens).
+            if (node->has_value) {
+                // check if it corresponds to the whole utf8 step
+                if (current_offset - offset == n_utf8_code_units) {
+                    found_unknown = false;
+                }
+                float score = current_best.score + scores[node->token];
+                struct result & current_champ = results[current_offset];
+                if (score > current_champ.score) {
+                    struct result challenger = { node->token, offset, score };
+                    current_champ = challenger;
+                }
+            }
+            node = node->traverse(normalized[current_offset++]);
+        }
+
+        // if we found an unknown token, process it
+        if (found_unknown) {
+            current_offset = offset + n_utf8_code_units;
+            struct result & current_champ = results[current_offset];
+            float score = current_best.score + unk_token_score;
+            if (score > current_champ.score) {
+                struct result challenger = { unk_token, offset, score };
+                current_champ = challenger;
+            }
+        }
+
+        // move one utf8 step
+        offset += n_utf8_code_units;
+    }
+
+    // if we have more than on unknown token in a row, we can join them.
+    bool is_prev_unknown = false;
+    // iterate from the last result backwards and get the best performing tokens
+    for (struct result & result = results[text_length]; ; result = results[result.offset]) {
+        bool is_unknown = result.token == unk_token;
+        if (!(is_prev_unknown && is_unknown)) {
+            tokens.push_back(result.token);
+        }
+        if (result.offset == 0) {
+            break;
+        }
+        is_prev_unknown = is_unknown;
+    }
+
+    // reverse the tokens since we added tokens starting from the end of the input
+    std::reverse(tokens.begin(), tokens.end());
+}
+
+// loading the vocab to the tokenizer from gguf file.
+unigram_tokenizer * unigram_tokenizer_from_gguf(gguf_context * meta) {
+    std::unordered_map<std::string, uint32_t> vocab;
+    std::vector<float> scores;
+    int vocab_key = gguf_find_key(meta, "tokenizer.ggml.tokens");
+    int vocab_size = gguf_get_arr_n(meta, vocab_key);
+    scores.reserve(vocab_size);
+    for (int i = 0; i < vocab_size; i++) {
+        std::string val = gguf_get_arr_str(meta, vocab_key, i);
+        vocab[val] = (uint32_t) i;
+    }
+    int scores_key = gguf_find_key(meta, "tokenizer.ggml.scores");
+    int scores_size = gguf_get_arr_n(meta, scores_key);
+    assert(scores_size == vocab_size);
+    float * data = (float*) gguf_get_arr_data(meta, scores_key);
+    for (int i = 0; i < scores_size; i++) {
+        scores.push_back(data[i]);
+    }
+    int unkown_token_key = gguf_find_key(meta, "tokenizer.ggml.unknown_token_id");
+    uint32_t token = gguf_get_val_u32(meta, unkown_token_key);
+
+    auto tokenizer =  new unigram_tokenizer(vocab, token, scores[token], scores);
+
+    uint32_t eos_token_key = gguf_find_key(meta, "tokenizer.ggml.eos_token_id");
+    if (eos_token_key != -1) {
+        tokenizer->eos_token = gguf_get_val_u32(meta, eos_token_key);
+    }
+    return tokenizer;
+}
+
+void single_pass_tokenizer::tokenize(const std::string & text, std::vector<uint32_t> & token_ids) {
+    std::string remaining = text;
+    while (remaining.size() > 0) {
+        uint32_t token_id = unknown_id;
+        for (int i = 1; i < std::min(remaining.size()+1, max_size+1); i++) {
+            std::string part = remaining.substr(0, i);
+            ptrdiff_t pos = std::distance(tokens.begin(), std::find(tokens.begin(), tokens.end(), part));
+            if (pos < tokens.size()) {
+                token_id = (uint32_t) pos;
+                remaining = remaining.substr(part.size(), remaining.size() - part.size());
+                break;
+            }
+        }
+        if (token_id == unknown_id) {
+            remaining = remaining.substr(1, remaining.size() - 1);
+        }
+        token_ids.push_back(token_id);
+    }
+}
+
+void single_pass_tokenizer::token_split(const std::string & text, std::vector<std::string> & tokens) {
+    std::string remaining = text;
+    while (remaining.size() > 0) {
+        // String copying is much slower than using a std::string_view, but the former is simpler to implement for now.
+        std::string token = remaining.substr(0, 1);
+        for (int i = 1; i < remaining.size(); i++) {
+            std::string part = remaining.substr(0, i+1);
+            if (token_vocab.find(part) == token_vocab.end()) {
+                break;
+            }
+            token = part;
+        }
+        tokens.push_back(token);
+        remaining = remaining.substr(token.size(), remaining.size() - token.size());
+    }
+}
+
+struct single_pass_tokenizer * single_pass_tokenizer_from_gguf(gguf_context * meta, std::string key_name) {
+    int tokens_key = gguf_find_key(meta, key_name.c_str());
+    if (tokens_key == -1) {
+        TTS_ABORT("The '%s' key must be set in order to support single pass tokenization.", key_name.c_str());
+    }
+    std::vector<std::string> tokens;
+    int token_count = gguf_get_arr_n(meta, tokens_key);
+    for (int i = 0; i < token_count; i++) {
+        tokens.push_back(gguf_get_arr_str(meta, tokens_key, i));
+    }
+    return new single_pass_tokenizer(tokens);
+}
+
+void bpe_symbol::add_merges(std::priority_queue<bpe_merge, std::vector<bpe_merge>, bpe_merge_comp> & merges, std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> & rank_map, bool only_forward) {
+    if (!only_forward && last) {
+        auto rid = std::make_pair<std::string, std::string>(last->as_str(), as_str());
+        if (rank_map.find(rid) != rank_map.end()) {
+            bpe_merge m{last, this, rank_map[rid], last->size + size};
+            merges.push(m);
+        }
+    }
+
+    if (next) {
+        auto rid = std::make_pair<std::string, std::string>(as_str(), next->as_str());
+        if (rank_map.find(rid) != rank_map.end()) {
+            bpe_merge m{this, next, rank_map[rid], size + next->size};
+            merges.push(m);
+        }
+    }
+}
+
+std::string bpe_symbol::as_str() {
+    return std::string(token, size);
+}
+
+bool bpe_merge_comp::operator() (const bpe_merge & a, const bpe_merge & b) {
+    return a.rank > b.rank || (a.rank == b.rank && a.a && b.a && a.a->pos > b.a->pos);
+}
+
+size_t pair_hash::operator() (const std::pair<std::string, std::string> & p) const {
+    return std::hash<std::string>{}(p.first) ^ (std::hash<std::string>{}(p.second) << 1);
+}
+
+bpe_symbol * bpe_merge::merge() {
+    a->size += b->size;
+    b->size = -1;
+    a->next = b->next;
+    if (a->next) {
+        a->next->last = a;
+    }
+    return a;
+}
+
+void pair_builder::join_pairs(std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> & rank_map) {
+    std::priority_queue<bpe_merge, std::vector<bpe_merge>, bpe_merge_comp> merges;
+    for (auto part : parts) {
+        part->add_merges(merges, rank_map, true);
+    }
+    while (!merges.empty()) {
+        auto m = merges.top();
+        merges.pop();
+        if (m.a->size > 0 && m.b->size > 0 && m.new_size == m.a->size + m.b->size) {
+            m.merge();
+            m.a->add_merges(merges, rank_map);
+        }
+
+    }
+}
+
+void bpe_tokenizer::tokenize(const std::string & text, std::vector<uint32_t> & token_ids) {
+    std::vector<std::string> chunks = split(text, " ", true);
+    bool space_prior = false;
+    for (auto chunk : chunks) {
+        if (chunk != " ") {
+            bpe_tokenize(space_prior ? "Ġ" + chunk : chunk, token_ids);
+        } else {
+            space_prior = true;
+        }
+    }
+}
+
+void bpe_tokenizer::bpe_tokenize(std::string chunk, std::vector<uint32_t> & token_ids) {
+    if (tokens_to_ids.find(chunk) != tokens_to_ids.end()) {
+        token_ids.push_back(tokens_to_ids[chunk]);
+        return;
+    }
+    auto pb = pair_builder{chunk};
+    pb.join_pairs(ranks);
+    bpe_symbol * next = pb.parts[0];
+    while (next) {
+        token_ids.push_back(tokens_to_ids[next->as_str()]);
+        next = next->next;
+    }
+}
+
+bpe_tokenizer * bpe_tokenizer_from_gguf(gguf_context * meta, std::string base_name) {
+    int vocab_key = gguf_find_key(meta, (base_name + ".tokens").c_str());
+    if (vocab_key == -1) {
+        TTS_ABORT("The '%s' key must be set in order to support BPE tokenization.", (base_name + ".tokens").c_str());
+    }
+    int merges_key = gguf_find_key(meta, (base_name + ".merges").c_str());
+    if (merges_key == -1) {
+        TTS_ABORT("The '%s' key must be set in order to support BPE tokenization.", (base_name + ".merges").c_str());
+    }
+    int eos_token_id_key = gguf_find_key(meta, (base_name + ".eos_token_id").c_str());
+    if (eos_token_id_key == -1) {
+        TTS_ABORT("The '%s' key must be set in order to support BPE tokenization.", (base_name + ".eos_token_id").c_str());
+    }
+    int bos_token_id_key = gguf_find_key(meta, (base_name + ".bos_token_id").c_str());
+    if (bos_token_id_key == -1) {
+        TTS_ABORT("The '%s' key must be set in order to support BPE tokenization.", (base_name + ".bos_token_id").c_str());
+    }
+
+    uint32_t bos_token_id = gguf_get_val_u32(meta, bos_token_id_key);
+    uint32_t eos_token_id = gguf_get_val_u32(meta, eos_token_id_key);
+
+    std::unordered_map<std::string, uint32_t> vocab;
+    int token_count = gguf_get_arr_n(meta, vocab_key);
+    for (int i = 0; i < token_count; i++) {
+        vocab[gguf_get_arr_str(meta, vocab_key, i)] = (uint32_t) i;
+    }
+
+    std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> ranks;
+    int merge_count = gguf_get_arr_n(meta, merges_key);
+
+    for (int i = 0; i < merge_count; i++) {
+        auto raw_merge = gguf_get_arr_str(meta, merges_key, i);
+        std::vector<std::string> pair = split(raw_merge, " ");
+        if (pair.size() != 2) {
+            TTS_ABORT("Invalid pair, '%s', found in BPE merges, '%s', at index %d.", raw_merge, (base_name + ".merges").c_str(), i);
+        }
+        ranks[std::make_pair<>(pair[0], pair[1])] = i;
+    }
+
+    return new bpe_tokenizer(vocab, ranks, bos_token_id, eos_token_id);
+}

otherarch/ttscpp/src/tokenizer.h

@@ -0,0 +1,154 @@
+#ifndef tokenizer_h
+#define tokenizer_h
+
+#include <unordered_map>
+#include <stdint.h>
+#include <map>
+#include <unordered_set>
+#include <regex>
+#include <queue>
+#include "ttsutil.h"
+
+struct token_trie {
+    bool has_value = false;
+    uint32_t token;
+    std::map<char, struct token_trie> children;
+
+    void add(const std::string & gram, uint32_t token);
+    void _add(const std::string & gram, uint32_t new_token, size_t index);
+    const struct token_trie * traverse(const char c) const;
+};
+
+static std::regex duped_spaces("\\s{2,}");
+static std::regex spaces("\\s");
+
+struct result {
+    uint32_t token;
+    size_t offset;
+    float score;
+};
+
+// much of this is implemented in llama.cpp, but in order to simplify this for my use case, I reimplementing here.
+// There are several important simplifications here:
+// 1. I only implement unigram tokenization
+// 2. I don't need to support detokenization
+struct unigram_tokenizer {
+    unigram_tokenizer(std::unordered_map<std::string, uint32_t> vocab, uint32_t unk_token, float unk_token_score, std::vector<float> scores): vocab(vocab), unk_token(unk_token), unk_token_score(unk_token_score), scores(scores) {};
+    ~unigram_tokenizer() = default;
+
+    std::unordered_map<std::string, uint32_t> vocab;
+    std::vector<float> scores;
+    struct token_trie root_trie;
+    uint32_t unk_token;
+    float unk_token_score;
+    uint32_t eos_token = 1;
+    bool dedupe_spaces = true;
+    bool init = false;
+
+    void initialize_tokenizer();
+    void tokenize(const std::string & text, std::vector<uint32_t> & tokens);
+};
+
+// For intializing a new tokenizer from a gguf file meta
+unigram_tokenizer * unigram_tokenizer_from_gguf(gguf_context * meta);
+
+// While this functions like a tokenizer, no token ids are assigned as the token ids never need to be used in the context in which this is
+// currently being used. This tokenizer pattern is currently being used by the phonemizer to break up a word into its relevant graphemes.
+// As such, only the graphemes need to be returned.
+struct single_pass_tokenizer {
+    single_pass_tokenizer(std::vector<std::string> tkns): tokens(tkns) {
+        max_size = 0;
+        for (auto token : tkns) {
+            token_vocab.insert(token);
+            if (token.size() > max_size) {
+                max_size = token.size();
+            }
+        }
+    }
+    size_t max_size;
+    uint32_t unknown_id = 0;
+    std::vector<std::string> tokens;
+    std::unordered_set<std::string> token_vocab;
+    void tokenize(const std::string & text, std::vector<uint32_t> & token_ids);
+    void token_split(const std::string & text, std::vector<std::string> & tokens);
+};
+
+single_pass_tokenizer * single_pass_tokenizer_from_gguf(gguf_context * meta, std::string key_name = "phonemizer.graphemes");
+
+struct bpe_symbol;
+
+struct bpe_merge {
+    bpe_symbol * a;
+    bpe_symbol * b;
+    int rank;
+    int new_size;
+
+    bpe_symbol * merge();
+};
+
+struct bpe_merge_comp{
+    bool operator() (const bpe_merge & a, const bpe_merge & b);
+};
+
+struct pair_hash {
+    size_t operator() (const std::pair<std::string, std::string> & p) const;
+};
+
+struct bpe_symbol {
+    bpe_symbol(const char * token): token(token) {};
+    const char* token;
+    int size = 1;
+    int pos;
+    bpe_symbol * next = nullptr;
+    bpe_symbol * last = nullptr;
+
+    void add_merges(std::priority_queue<bpe_merge, std::vector<bpe_merge>, bpe_merge_comp> & merges, std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> & rank_map, bool only_forward = false);
+    std::string as_str();
+};
+
+struct pair_builder {
+    pair_builder(std::string word) {
+        bpe_symbol * last = nullptr;
+        for (int i = 0; i < word.size(); i++) {
+            int increment = 0;
+            // make sure we process each utf-8 character.
+            while(i + increment + 1 < word.size() && (word[i+increment+1] & 0b11000000) == 0b10000000) {
+                ++increment;
+            }
+            bpe_symbol * part = new bpe_symbol(word.data()+i);
+            part->pos = i;
+            part->size += increment;
+            i += increment;
+            if (last) {
+                last->next = part;
+                part->last = last;
+            }
+            last = part;
+            parts.push_back(part);
+        }
+    }
+
+    ~pair_builder() {
+        for (auto p : parts) {
+            delete p;
+        }
+    }
+
+    void join_pairs(std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> & rank_map);
+    std::vector<bpe_symbol*> parts;
+};
+
+struct bpe_tokenizer {
+    bpe_tokenizer(std::unordered_map<std::string, uint32_t> & tokens_to_ids, std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> & ranks, uint32_t bos, uint32_t eos): tokens_to_ids(tokens_to_ids), ranks(ranks), eos_token_id(eos), bos_token_id(bos) {};
+    std::unordered_map<std::string, uint32_t> tokens_to_ids;
+    std::unordered_map<std::pair<std::string, std::string>, int, pair_hash> ranks;
+    uint32_t eos_token_id;
+    uint32_t bos_token_id;
+
+    void tokenize(const std::string & text, std::vector<uint32_t> & token_ids);
+    void bpe_tokenize(std::string chunk, std::vector<uint32_t> & token_ids);
+};
+
+bpe_tokenizer * bpe_tokenizer_from_gguf(gguf_context * meta, std::string base_name = "tokenizer.ggml");
+
+#endif

otherarch/ttscpp/src/tts.cpp

@@ -0,0 +1,445 @@
+#include "tts.h"
+#include <mutex>
+
+// A list of all of the top level GGUF names under kokoro.duration_predictor that have quantization compatible tensors.
+static constexpr std::array<const char *, 5> DURATION_PREDICTOR_QUANTIZATION_COMPATIBLE_PARTS = {
+    "duration_proj",
+    "encode",
+    "shared_lstm",
+    "duration_lstm",
+    "layers"
+};
+
+struct tts_runner * orpheus_from_file(gguf_context * meta_ctx, ggml_context * weight_ctx, int n_threads, generation_configuration * config, tts_arch arch, bool cpu_only) {
+    orpheus_model * model = new orpheus_model;
+    snac_model * audio_model = new snac_model;
+    bpe_tokenizer * bt = bpe_tokenizer_from_gguf(meta_ctx);
+    model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    audio_model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    sampler * samp = new sampler;
+    snac_context * sctx = build_new_snac_context(audio_model, n_threads, cpu_only);
+    snac_runner * audio_decoder = new snac_runner(audio_model, sctx);
+    orpheus_context * octx = build_new_orpheus_context(model, n_threads, cpu_only);
+    orpheus_kv_cache * cache = new orpheus_kv_cache;
+    orpheus_runner * runner = new orpheus_runner(model, audio_decoder, octx, bt, samp, cache);
+
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        runner->assign_weight(cur->name, cur);
+    }
+
+    runner->prepare_post_load();
+
+    gguf_free(meta_ctx);
+    ggml_free(weight_ctx);
+    runner->arch = arch;
+
+    return (tts_runner*)runner;
+}
+
+struct tts_runner * parler_tts_from_file(gguf_context * meta_ctx, ggml_context * weight_ctx, int n_threads, generation_configuration * config, tts_arch arch, bool cpu_only) {
+    parler_tts_model * model = new parler_tts_model;
+    dac_model * audio_model = new dac_model;
+    unigram_tokenizer * ut = unigram_tokenizer_from_gguf(meta_ctx);
+    ut->initialize_tokenizer();
+    model->use_cross_attn = config->use_cross_attn;
+    model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    audio_model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    struct sampler * samp = new sampler;
+    struct dac_context * dctx = build_new_dac_context(audio_model, n_threads, cpu_only);
+    struct dac_runner * audio_decoder = new dac_runner(audio_model, dctx);
+    struct parler_context * pctx = build_new_parler_context(model, n_threads, cpu_only);
+    struct parler_kv_cache * cache = new parler_kv_cache;
+    struct parler_tts_runner * runner = new parler_tts_runner(model, audio_decoder, pctx, ut, samp, cache);
+
+    // TODO: change this weight assignment pattern to mirror llama.cpp
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        runner->assign_weight(cur->name, cur);
+    }
+
+    if (config->use_cross_attn) {
+        runner->model->prep_cross_key_values(n_threads);
+    }
+
+    runner->prepare_post_load();
+
+    gguf_free(meta_ctx);
+    ggml_free(weight_ctx);
+    runner->arch = arch;
+
+    return (tts_runner*)runner;
+}
+
+struct tts_runner * kokoro_from_file(gguf_context * meta_ctx, ggml_context * weight_ctx, int n_threads, generation_configuration * config, tts_arch arch, bool cpu_only) {
+    kokoro_model * model = new kokoro_model;
+    single_pass_tokenizer * spt = single_pass_tokenizer_from_gguf(meta_ctx, "tokenizer.ggml.tokens");
+    model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    struct kokoro_duration_context * kdctx = build_new_duration_kokoro_context(model, n_threads, cpu_only);
+    struct kokoro_duration_runner * duration_runner = new kokoro_duration_runner(model, kdctx, spt);
+    struct kokoro_context * kctx = build_new_kokoro_context(model, n_threads, cpu_only);
+    // if an espeak voice id wasn't specifically set infer it from the kokoro voice, if it was override it, otherwise fallback to American English.
+    std::string espeak_voice_id = config->espeak_voice_id;
+    if (espeak_voice_id.empty()) {
+        espeak_voice_id = !config->voice.empty() && KOKORO_LANG_TO_ESPEAK_ID.find(config->voice.at(0)) != KOKORO_LANG_TO_ESPEAK_ID.end() ? KOKORO_LANG_TO_ESPEAK_ID[config->voice.at(0)] : "gmw/en-US";
+    }
+    struct phonemizer * phmzr = phonemizer_from_gguf(meta_ctx, espeak_voice_id);
+    struct kokoro_runner * runner = new kokoro_runner(model, kctx, spt, duration_runner, phmzr);
+
+    // TODO: change this weight assignment pattern to mirror llama.cpp
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        runner->assign_weight(cur->name, cur);
+    }
+
+    runner->prepare_post_load();
+
+    gguf_free(meta_ctx);
+    ggml_free(weight_ctx);
+    runner->arch = arch;
+
+    return (tts_runner*)runner;
+}
+
+struct tts_runner * dia_from_file(gguf_context * meta_ctx, ggml_context * weight_ctx, int n_threads, generation_configuration * config, tts_arch arch, bool cpu_only) {
+    dia_model * model = new dia_model;
+    dac_model * audio_model = new dac_model;
+    model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    audio_model->setup_from_file(meta_ctx, weight_ctx, cpu_only);
+    struct sampler * samp = new sampler;
+    struct dac_context * dctx = build_new_dac_context(audio_model, n_threads, cpu_only);
+    struct dac_runner * audio_decoder = new dac_runner(audio_model, dctx);
+    struct dia_context * diactx = build_new_dia_context(model, n_threads, cpu_only);
+    struct dia_kv_cache * cache = new dia_kv_cache;
+    struct dia_runner * runner = new dia_runner(model, audio_decoder, diactx, samp, cache);
+
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        runner->assign_weight(cur->name, cur);
+    }
+
+    runner->prepare_post_load();
+
+    gguf_free(meta_ctx);
+    ggml_free(weight_ctx);
+    runner->arch = arch;
+
+    return (tts_runner*)runner;
+}
+
+// currently only metal and cpu devices are supported, so cpu_only only describes whether or not to try to load and run on metal.
+struct tts_runner * runner_from_file(const std::string & fname, int n_threads, generation_configuration * config, bool cpu_only) {
+    ggml_context * weight_ctx = NULL;
+
+    struct gguf_init_params params = {
+        /*.no_alloc   =*/ false,
+        /*.ctx        =*/ &weight_ctx,
+    };
+    gguf_context * meta_ctx = gguf_init_from_file(fname.c_str(), params);
+    if (!meta_ctx) {
+        TTS_ABORT("%s failed for file %s\n", __func__, fname.c_str());
+    }
+    int arch_key = gguf_find_key(meta_ctx, "general.architecture");
+    if (arch_key == -1) {
+        TTS_ABORT("%s failed for file %s. No architecture is set.\n", __func__, fname.c_str());
+    }
+    std::string arch = std::string(gguf_get_val_str(meta_ctx, arch_key));
+    if (SUPPORTED_ARCHITECTURES.find(arch) == SUPPORTED_ARCHITECTURES.end()) {
+        TTS_ABORT("%s failed for file %s. The architecture '%s' is not supported.", __func__, fname.c_str(), arch.c_str());
+    }
+    tts_arch arch_type = SUPPORTED_ARCHITECTURES.at(arch);
+    switch(arch_type) {
+        case PARLER_TTS_ARCH:
+            return parler_tts_from_file(meta_ctx, weight_ctx, n_threads, config, arch_type, cpu_only);
+        case KOKORO_ARCH:
+            return kokoro_from_file(meta_ctx, weight_ctx, n_threads, config, arch_type, cpu_only);
+        case DIA_ARCH:
+            return dia_from_file(meta_ctx, weight_ctx, n_threads, config, arch_type, cpu_only);
+        case ORPHEUS_ARCH:
+            return orpheus_from_file(meta_ctx, weight_ctx, n_threads, config, arch_type, cpu_only);
+        default:
+            TTS_ABORT("%s failed for file %s. The architecture '%s' is not supported.", __func__, fname.c_str(), arch.c_str());
+    }
+}
+
+int generate(tts_runner * runner, std::string sentence, struct tts_response * response, generation_configuration * config) {
+    switch(runner->arch) {
+        case PARLER_TTS_ARCH:
+            ((parler_tts_runner*)runner)->configure_generation(config);
+            return ((parler_tts_runner*)runner)->generate(sentence, response);
+        case KOKORO_ARCH:
+            return ((kokoro_runner*)runner)->generate(sentence, response, config->voice, config->espeak_voice_id);
+        case DIA_ARCH:
+            ((dia_runner*)runner)->configure_generation(config);
+            return ((dia_runner*)runner)->generate(sentence, response);
+        case ORPHEUS_ARCH:
+            ((orpheus_runner*)runner)->configure_generation(config);
+            return ((orpheus_runner*)runner)->generate(sentence, response);
+        default:
+            TTS_ABORT("%s failed. The architecture '%d' is not supported.", __func__, runner->arch);
+    }
+}
+
+std::vector<std::string> list_voices(tts_runner * runner) {
+    switch(runner->arch) {
+        case KOKORO_ARCH:
+            return ((kokoro_runner*)runner)->list_voices();
+        default:
+            TTS_ABORT("%s failed. The architecture '%d' does not support #list_voices supported.", __func__, runner->arch);
+    }
+}
+
+void update_conditional_prompt(tts_runner * runner, const std::string file_path, const std::string prompt, bool cpu_only) {
+    int n_threads = ((parler_tts_runner*)runner)->pctx->n_threads;
+    ((parler_tts_runner*)runner)->update_conditional_prompt(file_path, prompt, n_threads, cpu_only);
+}
+
+bool kokoro_is_f16_compatible(std::string name) {
+    return name.find("voice_tensors") == std::string::npos &&
+           name.find("bias") == std::string::npos &&
+           name.find("gamma") == std::string::npos &&
+           name.find("beta") == std::string::npos &&
+           name.find("alpha") == std::string::npos &&
+           !has_suffix(name, "embd") &&
+           !has_suffix(name, "norm");
+}
+
+bool kokoro_is_quantizable(std::string name, struct quantization_params * params) {
+    if (kokoro_is_f16_compatible(name)) {
+        if (has_prefix(name, "kokoro.albert") || has_prefix(name, "kokoro.text_encoder.lstm")) {
+            return true;
+        } else if (has_prefix(name, "kokoro.duration_predictor.")) {
+            std::vector<std::string> parts = split(name, ".");
+            for (std::string part : DURATION_PREDICTOR_QUANTIZATION_COMPATIBLE_PARTS) {
+                if (part == parts[2]) {
+                    return true;
+                }
+            }
+        }
+    }
+    return false;
+}
+
+bool dia_is_quantizable(std::string name, struct quantization_params * params) {
+    // The DAC audio encoder / decoder is not compatible with quantization and normalization tensors should not be quantized.
+    bool quantizable = !has_prefix(name, "audio_encoder") && !has_suffix(name, "norm");
+    if (!params->quantize_output_heads) {
+        quantizable = quantizable && !has_prefix(name, "dia.decoder.heads");
+    }
+    return quantizable;
+}
+
+bool parler_is_quanitizable(std::string name, struct quantization_params * params) {
+    // the DAC audio encoder / decoder is not compatible with quantization, normalization weight shouldn't be quantized, and the text encoding shouldn't be normalized.
+    bool quantizable = !has_prefix(name, "audio_encoder") && !has_suffix(name, "norm.weight") && !has_suffix(name, "text_encoding") && !has_suffix(name, "positional_embed") && !has_suffix(name, "norm.bias");
+    if (!params->quantize_output_heads) {
+        quantizable = quantizable && !has_suffix(name, "weight.head");
+    }
+    if (!params->quantize_text_embeddings) {
+        quantizable = quantizable && !has_suffix(name, "embed_prompts");
+    }
+    if (!params->quantize_cross_attn_kv) {
+        quantizable = quantizable && !has_suffix(name, "encoder_attn.k_proj.weight") && !has_suffix(name, "encoder_attn.v_proj.weight");
+    }
+    return quantizable;
+}
+
+bool is_quantizable(tts_arch arch, std::string name, struct quantization_params * params) {
+    switch(arch) {
+        case PARLER_TTS_ARCH:
+            return parler_is_quanitizable(name, params);
+        case DIA_ARCH:
+            return dia_is_quantizable(name, params);
+        case KOKORO_ARCH:
+            return kokoro_is_quantizable(name, params);
+        default:
+            TTS_ABORT("%s failed. The architecture '%d' is not supported.", __func__, arch);
+    }
+}
+
+size_t quantize_tensor(void * new_data, struct ggml_tensor * tensor, const float * imatrix, enum ggml_type qtype, uint32_t n_threads) {
+    // much of this is form copied from llama.cpp
+    int chunk_size_multiplier = 1;
+    if (qtype == GGML_TYPE_Q4_0_4_4 || qtype == GGML_TYPE_Q4_0_4_8 || qtype == GGML_TYPE_Q4_0_8_8) {
+        if ((qtype == GGML_TYPE_Q4_0_8_8) && (tensor->ne[1] % 8 != 0)) qtype = GGML_TYPE_Q4_0;
+        else if (tensor->ne[1] % 4 != 0) qtype = GGML_TYPE_Q4_0;
+        if (qtype == GGML_TYPE_Q4_0_8_8) chunk_size_multiplier = 8;
+        else if (qtype == GGML_TYPE_Q4_0_4_4 || qtype == GGML_TYPE_Q4_0_4_8) chunk_size_multiplier = 4;
+    }
+    size_t out_size = 0;
+    const int32_t d3_step = tensor->ne[0] * tensor->ne[1];
+    const int32_t n_per_row = tensor->ne[0];
+    const int32_t nrows = tensor->ne[1];
+    static const int32_t min_chunk_size = 32 * 512;
+    const int32_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row)) * chunk_size_multiplier;
+    uint32_t thread_count = std::max(1, std::min((int)n_threads, (int)(d3_step + chunk_size - 1) / chunk_size));
+    std::mutex mutex;
+
+    for (int32_t d3_index = 0; d3_index < tensor->ne[2]; d3_index++) {
+        const float * f32_data_d3 = ((float *) tensor->data) + d3_index * d3_step;
+        void * new_data_d3 = (char *)new_data + ggml_row_size(qtype, tensor->ne[0]) * d3_index * nrows;
+        const float * imatrix_03 = imatrix ? imatrix + d3_index * tensor->ne[0] : nullptr;
+        if (thread_count <= 1) {
+            // not threaded
+            out_size += ggml_quantize_chunk(qtype, f32_data_d3, new_data_d3, 0, nrows, n_per_row, imatrix);
+        } else {
+            std::vector <std::thread> threads;
+            int64_t counter = 0;
+            size_t new_size = 0;
+            bool valid = true;
+            for (uint32_t t = 0; t < thread_count; t++) {
+                auto func = [&mutex, &counter, &new_size, &valid, qtype, f32_data_d3, new_data_d3, chunk_size, nrows, n_per_row, imatrix]() {
+                    const int64_t nrows_per_chunk = chunk_size / n_per_row;
+                    size_t local_size = 0;
+                    while (true) {
+                        std::unique_lock<std::mutex> lock(mutex);
+                        int64_t first_row = counter;
+                        counter += nrows_per_chunk;
+                        if (first_row >= nrows) {
+                            if (local_size > 0) {
+                                new_size += local_size;
+                            }
+                            break;
+                        }
+                        lock.unlock();
+                        const int64_t this_nrow = std::min(nrows - first_row, nrows_per_chunk);
+                        size_t this_size = ggml_quantize_chunk(qtype, f32_data_d3, new_data_d3, first_row * n_per_row, this_nrow, n_per_row, imatrix);
+                        local_size += this_size;
+
+                        // validate the quantized data; I am not sure how this would occur, but there is always the safe fallback on doing this single threaded.
+                        const size_t row_size  = ggml_row_size(qtype, n_per_row);
+                        void * this_data = (char *) new_data_d3 + first_row * row_size;
+                        if (!ggml_validate_row_data(qtype, this_data, this_size)) {
+                            std::unique_lock<std::mutex> lock(mutex);
+                            valid = false;
+                            break;
+                        }
+                    }
+                };
+                threads.push_back(std::thread(func));
+            }
+            for (auto & t : threads) t.join();
+
+            if (!valid) {
+                TTS_ABORT("Validation of quantized data failed. Please try again and/or switch to single thread quantization.\n");
+            }
+            out_size += new_size;
+        }
+    }
+    return out_size;
+}
+
+static void zeros(std::ofstream & file, size_t n) {
+    char zero = 0;
+    for (size_t i = 0; i < n; ++i) {
+        file.write(&zero, 1);
+    }
+}
+
+template <typename T>
+struct no_init {
+    T value;
+    no_init() { /* do nothing */ }
+};
+
+void quantize_gguf(const std::string & ifile, const std::string & ofile, struct quantization_params * params) {
+    ggml_context * weight_ctx = NULL;
+    struct gguf_init_params gguf_params = {
+        /*.no_alloc   =*/ false,
+        /*.ctx        =*/ &weight_ctx,
+    };
+    gguf_context * meta_ctx = gguf_init_from_file(ifile.c_str(), gguf_params);
+    std::string arch = "parler-tts"; // only parler-tts gguf files should lack an explicit architecture.
+
+    int arch_key = gguf_find_key(meta_ctx, "general.architecture");
+    if (arch_key != -1) {
+        arch = std::string(gguf_get_val_str(meta_ctx, arch_key));
+    }
+    tts_arch arch_type = SUPPORTED_ARCHITECTURES.at(arch);
+
+    if (params->quantize_type != GGML_TYPE_Q5_0 && params->quantize_type != GGML_TYPE_Q8_0 && params->quantize_type != GGML_TYPE_F16 && params->quantize_type != GGML_TYPE_Q4_0) {
+        fprintf(stdout, "Warning, %s is untested for quantization type '%d'. Use at your own risk.\n", arch.c_str(), params->quantize_type);
+    }
+
+    const size_t align = GGUF_DEFAULT_ALIGNMENT;
+    gguf_context_ptr ctx_out { gguf_init_empty() };
+
+    // copy the KV pairs from the input file
+    gguf_set_kv(ctx_out.get(), meta_ctx);
+    gguf_set_val_u32(ctx_out.get(), "general.quantization_version", GGML_QNT_VERSION);
+    gguf_set_val_u32(ctx_out.get(), "general.quantization_type", params->quantize_type);
+    for (ggml_tensor * tensor = ggml_get_first_tensor(weight_ctx); tensor; tensor = ggml_get_next_tensor(weight_ctx, tensor)) {
+        std::string name = ggml_get_name(tensor);
+        if (name.size() != 0) {
+            gguf_add_tensor(ctx_out.get(), tensor);
+        }
+    }
+
+    std::vector<no_init<uint8_t>> work;
+
+    std::ofstream fout;
+    auto close_ofstream = [&]() {
+        // Write metadata and close file handler
+        if (fout.is_open()) {
+            fout.seekp(0);
+            std::vector<uint8_t> data(gguf_get_meta_size(ctx_out.get()));
+            gguf_get_meta_data(ctx_out.get(), data.data());
+            fout.write((const char *) data.data(), data.size());
+            fout.close();
+        }
+    };
+    auto new_ofstream = [&]() {
+        std::string fname = ofile;
+        fout = std::ofstream(fname, std::ios::binary);
+        fout.exceptions(std::ofstream::failbit); // fail fast on write errors
+        const size_t meta_size = gguf_get_meta_size(ctx_out.get());
+        // placeholder for the meta data
+        ::zeros(fout, meta_size);
+    };
+    new_ofstream();
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        enum ggml_type new_type;
+        void * new_data;
+        size_t new_size;
+        std::string name = ggml_get_name(cur);
+
+        if (name.size() == 0) {
+            continue;
+        }
+
+        if (is_quantizable(arch_type, name, params)) {
+            if ((cur->type) != GGML_TYPE_F32) {
+                TTS_ABORT("ERROR: All quantized tensors must be transformed from 32bit floats. Tensor, '%s', has improper type, '%d'\n", cur->name, cur->type);
+            }
+            new_type = params->quantize_type;
+            if ((new_type >= GGML_TYPE_IQ2_XXS && new_type <= GGML_TYPE_IQ4_XS)) {
+                TTS_ABORT("ERROR: Quantization type '%d' requires an importance matrix.\n", new_type);
+            }
+            const int64_t nelement_size = ggml_nelements(cur) * 4;
+            if (work.size() < (size_t)nelement_size) {
+                work.resize(nelement_size); // upper bound on size
+            }
+            new_data = work.data();
+            new_size = quantize_tensor(new_data, cur, nullptr, new_type, params->n_threads);
+        } else if ((params->convert_non_quantizable_to_f16 && kokoro_is_f16_compatible(name)) || (params->convert_dac_to_f16 && has_prefix(name, "audio_encoder") && !has_suffix(name, "alpha"))) {
+            if ((cur->type) != GGML_TYPE_F32) {
+                TTS_ABORT("ERROR: All converted tensors must be transformed from 32bit floats. Tensor, '%s', has improper type, '%d'\n", cur->name, cur->type);
+            }
+            new_type = GGML_TYPE_F16;
+            const int64_t nelement_size = ggml_nelements(cur) * 4;
+            if (work.size() < (size_t)nelement_size) {
+                work.resize(nelement_size); // upper bound on size
+            }
+            new_data = work.data();
+            new_size = quantize_tensor(new_data, cur, nullptr, new_type, params->n_threads);
+        } else {
+            new_type = cur->type;
+            new_data = cur->data;
+            new_size = ggml_nbytes(cur);
+        }
+
+        gguf_set_tensor_type(ctx_out.get(), name.c_str(), new_type);
+        gguf_set_tensor_data(ctx_out.get(), name.c_str(), new_data);
+        fprintf(stdout, "At tensor: '%s' with new size: %zu bytes\n", name.c_str(), new_size);
+        // write tensor data + padding
+        fout.write((const char *) new_data, new_size);
+        zeros(fout, GGML_PAD(new_size, align) - new_size);
+    }
+    close_ofstream();
+}

otherarch/ttscpp/src/tts_model.cpp

@@ -0,0 +1,157 @@
+#include "tts_model.h"
+#include "ggml-backend.h"
+#include "ggml-cpu.h"
+
+void append_to_response(struct tts_response * response, struct tts_response * to_append) {
+    float * new_data = (float *) malloc((response->n_outputs + to_append->n_outputs) * sizeof(float));
+    if (response->n_outputs > 0) {
+        std::memcpy(new_data, response->data, response->n_outputs*sizeof(float));
+    }
+    if (to_append->n_outputs > 0) {
+        float * next_loc = new_data + response->n_outputs;
+        std::memcpy(next_loc, to_append->data, to_append->n_outputs*sizeof(float));
+    }
+    response->data = new_data;
+    response->n_outputs += to_append->n_outputs;
+}
+
+/*
+ * Pulls output_size to prepped buffer 'output' from 'output_node' tensor. If no buffer is passed will default to the existing output buffer present
+ * on runner_context.
+ */
+void runner_context::get_ggml_node_data(struct ggml_tensor * output_node, float * output, size_t output_size, ggml_backend_buffer_t buffer) {
+    if (buffer == nullptr) {
+        buffer = buf_output;
+    }
+    if (ggml_backend_buffer_get_size(buffer) < output_size) {
+        TTS_ABORT("Output buffer overflow of %d / %d for output node '%s'\n", output_size, ggml_backend_buffer_get_size(buffer), ggml_get_name(output_node));
+    } else if (ggml_nbytes(output_node) < output_size) {
+        TTS_ABORT("Output node, '%s', with %d bytes is too small for #ggml_backend_tensor_get_async with size of %d.\n", ggml_get_name(output_node), ggml_nbytes(output_node), output_size);
+    }
+    ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched, output_node);
+    ggml_backend_tensor_get_async(backend_res, output_node, output, 0, output_size);
+}
+
+void runner_context::set_threads() {
+    if (backend != nullptr) {
+#ifdef GGML_USE_METAL
+        // this is form copied from llama.cpp, but has since been removed. I don't know if this should be tuned.
+        ggml_backend_metal_set_n_cb(backend, 1);
+#endif
+    }
+    if (backend_cpu != nullptr) {
+        ggml_backend_cpu_set_n_threads(backend_cpu, n_threads);
+        struct ggml_threadpool_params ttp = ggml_threadpool_params_default(n_threads);
+        threadpool = ggml_threadpool_new(&ttp);
+        ggml_backend_cpu_set_threadpool(backend_cpu, threadpool);
+    }
+}
+
+void runner_context::build_schedule(size_t max_nodes) {
+    backend_cpu_buffer = ggml_backend_cpu_buffer_type();
+    if (backend != nullptr) {
+#ifdef GGML_USE_METAL
+        backend_buffer = ggml_backend_metal_buffer_type();
+#endif
+        std::vector<ggml_backend_buffer_type_t> bufs = {backend_buffer, backend_cpu_buffer};
+        std::vector<ggml_backend_t> backs = {backend, backend_cpu};
+        sched = ggml_backend_sched_new(backs.data(), bufs.data(), 2, max_nodes, false, false);
+    } else {
+        std::vector<ggml_backend_buffer_type_t> bufs = {backend_cpu_buffer};
+        std::vector<ggml_backend_t> backs = {backend_cpu};
+        sched = ggml_backend_sched_new(backs.data(), bufs.data(), 1, max_nodes, false, false);
+    }
+}
+
+bool runner_context::prep_schedule(struct ggml_cgraph * gf) {
+    return ggml_backend_sched_reserve(sched, gf);
+}
+
+void runner_context::prep_output_buffer(size_t new_size) {
+    const size_t prev_size = buf_output ? ggml_backend_buffer_get_size(buf_output) : 0;
+    if (!buf_output || prev_size < new_size) {
+        if (buf_output) {
+            ggml_backend_buffer_free(buf_output);
+            buf_output = nullptr;
+            logits = nullptr;
+        }
+        buf_output = ggml_backend_buft_alloc_buffer(backend_cpu_buffer, new_size);
+    }
+    logits = (float *) ggml_backend_buffer_get_base(buf_output);
+}
+
+void tts_runner::init_build(std::vector<uint8_t>* buf_compute_meta) {
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ buf_compute_meta->size(),
+        /*.mem_buffer =*/ buf_compute_meta->data(),
+        /*.no_alloc   =*/ true,
+    };
+
+    ctx = ggml_init(params);
+}
+
+void tts_runner::free_build() {
+    if (ctx) {
+        ggml_free(ctx);
+        ctx = nullptr;
+    }
+}
+
+void tts_model::prep_buffers_and_context(bool cpu_only, float size_offset, uint32_t dedicated_add_on_size) {
+    // currently DAC is only supported on cpu because the ops are not implemented on other devices;
+    if (cpu_only) {
+        backend = ggml_backend_cpu_init();
+        buffer = ggml_backend_cpu_buffer_type();
+    } else {
+#ifdef GGML_USE_METAL
+        backend = ggml_backend_metal_init();
+        buffer = ggml_backend_metal_buffer_type();
+#endif
+        // if use metal is not installed then we need to warn here
+        if (!backend || !buffer) {
+            TTS_ABORT("'GGML_USE_METAL' is not defined either set the model to use CPU only or install ggml with metal support.");
+        }
+    }
+    size_t ctx_size = ggml_tensor_overhead() * (tensor_meta.n_tensors * size_offset);
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ ctx_size,
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true,
+    };
+    ctx = ggml_init(params);
+    buf = ggml_backend_buft_alloc_buffer(buffer, tensor_meta.n_bytes + dedicated_add_on_size);
+}
+
+void tts_model::assign_weight(std::string name, ggml_tensor * tensor) {
+	TTS_ABORT("%s received name, %s, tensor without being defined. %s must be defined for all implementations of tts_model. \n", __func__, name.c_str(), __func__);
+}
+
+void tts_model::set_tensor(struct ggml_tensor * tensor, struct ggml_tensor * target) {
+    tensor->buffer = buf;
+    tensor->data = (void *)((uint8_t *) ggml_backend_buffer_get_base(buf) + offset);
+    size_t size = ggml_nbytes(target);
+    ggml_backend_tensor_set(tensor, target->data, 0, size);
+    ggml_set_name(tensor, target->name);
+    offset += size;
+}
+
+void tts_model::setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only, std::string model_prefix, float size_offset, uint32_t dedicated_add_on_size) {
+    tensor_meta = compute_tensor_meta(model_prefix, load_context, compute_tensor_meta_cb);
+    prep_buffers_and_context(cpu_only, size_offset, dedicated_add_on_size);
+}
+
+size_t tts_model::max_nodes() {
+    return std::max<size_t>(8192, tensor_meta.n_tensors*5);
+}
+
+void tts_model::free() {
+    if (ctx) {
+        ggml_free(ctx);
+    }
+    if (buf) {
+        ggml_backend_buffer_free(buf);
+    }
+    if (backend) {
+        ggml_backend_free(backend);
+    }
+}

otherarch/ttscpp/src/tts_model.h

@@ -0,0 +1,69 @@
+#ifndef tts_model_h
+#define tts_model_h
+
+#include <cstring>
+#include <functional>
+#include "ttsutil.h"
+#include "ttscommon.h"
+
+void append_to_response(struct tts_response * response, struct tts_response * to_append);
+
+using tensor_meta_callback = std::function<void(ggml_tensor*)>*;
+
+struct runner_context {
+    runner_context(int n_threads): n_threads(n_threads) {};
+    virtual ~runner_context() {
+        ggml_backend_sched_free(sched);
+        ggml_threadpool_free(threadpool);
+        ggml_backend_free(backend_cpu);
+        ggml_backend_free(backend);
+        ggml_backend_buffer_free(buf_output);
+    }
+    // TODO: extend the backend and buffer support out to all devices
+    ggml_backend_t backend = nullptr;
+    ggml_backend_buffer_type_t backend_buffer = nullptr;
+
+    ggml_backend_t backend_cpu = nullptr;
+    ggml_backend_buffer_type_t backend_cpu_buffer = nullptr;
+
+    std::vector<uint8_t> buf_compute_meta;
+    ggml_backend_buffer_t buf_output = nullptr;
+    ggml_backend_sched_t sched = nullptr;
+    ggml_threadpool_t threadpool = nullptr;
+    float * logits = nullptr;
+    int n_threads;
+
+    void get_ggml_node_data(struct ggml_tensor * output_tensor, float * output, size_t output_size, ggml_backend_buffer_t buffer = nullptr);
+    void set_threads();
+    void build_schedule(size_t max_nodes);
+    bool prep_schedule(ggml_cgraph * gf);
+    void prep_output_buffer(size_t new_size);
+};
+
+struct tts_model {
+    struct model_tensor_meta tensor_meta;
+
+    // this is the current byte offset into the model's buffer.
+    size_t offset = 0;
+
+    bool use_cross_attn = true;
+
+    ggml_backend_buffer_type_t buffer = nullptr;
+    ggml_backend_t backend = nullptr;
+    ggml_backend_buffer_t buf = nullptr;
+
+    // it is quite common for implementations of tts_model to need to update attributes or perform distinct operations
+    // when computing the tensor meta of the loaded model. This callback allows this as it will receive each processed tensor.
+    tensor_meta_callback compute_tensor_meta_cb = nullptr;
+
+    struct ggml_context * ctx;
+
+    void prep_buffers_and_context(bool cpu_only, float size_offset, uint32_t dedicated_add_on_size);
+    void setup_from_file(gguf_context * meta_ctx, ggml_context * load_context, bool cpu_only, std::string model_prefix, float size_offset = 1.4, uint32_t dedicated_add_on_size = 0);
+    void set_tensor(struct ggml_tensor * tensor, struct ggml_tensor * target);
+    size_t max_nodes();
+    void assign_weight(std::string name, ggml_tensor * tensor);
+    void free();
+};
+
+#endif

otherarch/ttscpp/src/ttsutil.cpp

@@ -0,0 +1,308 @@
+#include "ttsutil.h"
+
+#include <algorithm>
+#include <cstdio>
+#include <stdarg.h>
+#ifdef __APPLE__
+#include <sys/sysctl.h>
+#elif __linux__
+#include <unistd.h>
+#else
+// windows stuff
+#endif
+
+void tts_abort(const char * file, int line, const char * fmt, ...) {
+    fflush(stdout);
+    fprintf(stderr, "%s:%d: ", file, line);
+    va_list args;
+    va_start(args, fmt);
+    vfprintf(stderr, fmt, args);
+    va_end(args);
+    abort();
+}
+
+// Simple helper function for getting layer count from tensor name
+std::pair<int, std::string> parse_layer_count(std::string name, int skip) {
+    bool found = false;
+    bool after_layer = false;
+    std::string digit_chars = "";
+    std::string after_layer_name = "";
+    int count = 0;
+    for (char& c : name) {
+        if (count < skip) {
+            count += 1;
+            continue;
+        }
+        count += 1;
+        if (after_layer) {
+            after_layer_name += c;
+        } else if (std::isdigit(c)) {
+            found = true;
+            digit_chars += c;
+        } else if (!found) {
+
+        } else {
+            after_layer = true;
+            after_layer_name += c;
+        }
+    }
+    if (digit_chars.size() == 0) {
+        return std::make_pair(-1, name);
+    }
+    return std::make_pair(std::stoi(digit_chars), after_layer_name);
+}
+
+int search_for_gguf_keys(gguf_context * meta, std::vector<std::string> possible_keys) {
+    int gguf_key = -1;
+    for (auto key : possible_keys) {
+        gguf_key = gguf_find_key(meta, key.c_str());
+        if (gguf_key != -1) {
+            return gguf_key;
+        }
+    }
+    return gguf_key;
+}
+
+void random_uniform_gen(int count, float * tgt, float min, float max) {
+    static std::default_random_engine e;
+    static std::uniform_real_distribution<float> dis(min, max);
+    for (int i = 0; i < count; i++) {
+        tgt[i] = dis(e);
+    }
+}
+
+void random_normal_gen(int count, float * tgt, float mean, float std) {
+    static std::default_random_engine e;
+    static std::normal_distribution<float> dis(mean, std);
+    for (int i = 0; i < count; i++) {
+        tgt[i] = dis(e);
+    }
+}
+
+float round_to_float(double v) {
+    return roundf(v * powl(10, 6)) / powl(10, 6);
+}
+
+struct ggml_tensor * reciprocal(ggml_context * ctx, struct ggml_tensor * x) {
+    TTS_ASSERT(x->ne[0] == 1);
+    static constexpr float one = 1.0f;
+    ggml_tensor * numerator = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, x->ne[1]);
+    // stride trick so that the scalar numerator can be divided by x.
+    numerator->nb[1] = 0;
+    numerator->data = const_cast<float *>(&one);
+    return ggml_div(ctx, numerator, x);
+}
+
+// Described in https://arxiv.org/abs/2006.08195
+// Snake1d is a common tunable activation function used in the DAC model.
+struct ggml_tensor * snake_1d(ggml_context * ctx, struct ggml_tensor * alpha, struct ggml_tensor * a) {
+    assert(a->ne[2] == 1 && a->ne[3] == 1);
+    return ggml_add(ctx, a, ggml_mul(ctx, ggml_sqr(ctx, ggml_sin(ctx, ggml_mul(ctx, a, alpha))), reciprocal(ctx, alpha)));
+}
+
+bool has_suffix(std::string value, std::string suffix) {
+    return value.size() >= suffix.size() && value.compare(value.size()-suffix.size(), suffix.size(), suffix) == 0;
+}
+
+bool has_prefix(std::string value, std::string prefix) {
+    return value.size() >= prefix.size() && value.compare(0, prefix.size(), prefix) == 0;
+}
+
+struct ggml_tensor * stft(ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * window, size_t n_fft, size_t hop, bool abs_and_angle, bool one_sided) {
+    if (window->ne[0] != n_fft) {
+        TTS_ABORT("For #stft the window_size, %d, must be either equal to n_fft, %d, or, when one sided, n_fft / 2 + 1, %d.\n", a->ne[0], n_fft, n_fft/2+1);
+    }
+    struct ggml_tensor * cur = ggml_stft(ctx, a, window, n_fft, hop, abs_and_angle);
+    if (one_sided) {
+        cur = ggml_cont(ctx, ggml_view_4d(ctx, cur, ((int64_t) n_fft / 2) + 1, cur->ne[1], cur->ne[2], cur->ne[3], cur->nb[1], cur->nb[2], cur->nb[3], 0));
+    }
+
+    return cur;
+}
+
+struct ggml_tensor * istft(ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * window_squared_sum, struct ggml_tensor * window, size_t n_fft, size_t hop, bool abs_and_angle, bool one_sided) {
+    if ((!one_sided && a->ne[0] != n_fft) || (one_sided && a->ne[0] != n_fft / 2 + 1)) {
+        TTS_ABORT("For #istft the window_size, %d, must be either equal to n_fft, %d, or, when one sided, n_fft / 2 + 1, %d.\n", a->ne[0], n_fft, n_fft/2+1);
+    }
+    struct ggml_tensor * cur = ggml_istft(ctx, a, window, n_fft, hop, abs_and_angle);
+    cur = ggml_div(ctx, cur, window_squared_sum);
+    return cur;
+}
+
+void hann_window(size_t n_fft, std::vector<float> & tgt) {
+    for (int i = 0; i < n_fft; i++) {
+        float v = pow(sin(M_PI * (double)i / (double) n_fft), 2.0);
+        tgt.push_back(v);
+    }
+}
+
+// This is a custom map op for computing noise and relevant voiced sections.
+void uv_noise_compute(struct ggml_tensor * dst, const struct ggml_tensor * a, const struct ggml_tensor * b, const struct ggml_tensor * c, int ith, int nth, void * userdata) {
+    float voice_threshold = ((float *) c->data)[0];
+    float noise_std = ((float *) c->data)[1];
+    float sin_amp = ((float *) c->data)[2];
+    float sin_amp_div = ((float *) c->data)[3];
+    float * rand_init = ((float *) c->data) + 4;
+
+    const int rpt = (b->ne[0] + nth - 1)/nth;
+    const int start = ith * rpt;
+    const int end = MIN((ith + 1) * rpt, b->ne[0]);
+
+    float * uv_dst = (float *) dst->data;
+    float * noise_dst = (float *)((char*)dst->data + dst->nb[2]);
+    float * tgt = (float *) b->data;
+
+    for(int bt = 0; bt < b->ne[2]; bt++) {
+        for(int r = start; r < end; r++) {
+            if (tgt[r] > voice_threshold) {
+                for (int h = 0; h < a->ne[1]; h++) {
+                    int index = h*dst->ne[0]+r;
+                    uv_dst[index] = sin_amp;
+                    noise_dst[index] = noise_std * rand_init[index];
+                }
+            } else {
+                for (int h = 0; h < a->ne[1]; h++) {
+                    int index = h*dst->ne[0]+r;
+                    uv_dst[index] = 0.0f;
+                    noise_dst[index] = sin_amp_div * rand_init[index];
+                }
+            }
+        }
+    }
+}
+
+// This is a custom map op for applying cfg scale. It is used at the terminus of logit generation in Dia.
+void cfg_scale(struct ggml_tensor * dst, const struct ggml_tensor * a, const struct ggml_tensor * b, int ith, int nth, void * userdata) {
+    const float scale = ((float *) userdata)[0];
+    const float max_output = ((float*) userdata)[1];
+    const int rpt = (b->ne[0] + nth - 1)/nth;
+    const int start = ith * rpt;
+    const int end = MIN((ith + 1) * rpt, b->ne[0]);
+
+    float * output = (float *) dst->data;
+    float * cond = (float *) a->data;
+    float * uncond = (float *) b->data;
+
+    for(int bt = 0; bt < b->ne[2]; bt++) {
+        for (int h = 0; h < b->ne[1]; h++) {
+            int i = (h * b->ne[0]) + (bt * b->ne[0] * b->ne[1]);
+            for(int r = start; r < end; r++) {
+                // only let the output heads yield tokens up to EOS
+                if (r > max_output) {
+                    output[i+r] = -INFINITY;
+                }
+                const float cr = cond[i+r];
+                const float ur = uncond[i+r];
+                output[i+r] = cr + scale * (cr - ur);
+            }
+        }
+    }
+}
+
+// currently this assumes a center view in which the output vector is reflectively padded by n_fft / 2 on each side.
+void compute_window_squared_sum(size_t n_fft, size_t hop, size_t n_frames, float * tgt, float * window) {
+    size_t cutoff = n_frames * hop;
+    size_t half = n_fft / 2;
+    std::memset(tgt, 0, cutoff*sizeof(float));
+    // istft applies half / hop steps before the beginning of the sequence. We need to account for these accumulated windows.
+    for (int i = 0; i < n_frames + (half / hop); i++) {
+        for (int ii = 0; ii < n_fft; ii++) {
+            int index = ii + i*hop - half;
+            if (index < 0 || index >= cutoff) {
+                continue;
+            }
+            tgt[index] += powf(window[ii], 2);
+        }
+    }
+}
+
+std::vector<std::string> split(std::string target, std::string split_on, bool include_split_characters) {
+    std::vector<std::string> output;
+    size_t last = 0;
+
+    for (int i = 0; i < target.size(); i++) {
+        if (i > last && split_on.find(target[i]) != std::string::npos) {
+            std::string part(target.substr(last, i - last));
+            output.push_back(part);
+            if (include_split_characters) {
+                output.push_back(target.substr(i, 1));
+            }
+            last = i+1;
+        } else if (i == last && split_on.find(target[i]) != std::string::npos) {
+            if (include_split_characters) {
+                output.push_back(target.substr(i, 1));
+            }
+            last = i+1;
+        }
+    }
+    if (last < target.size()) {
+        std::string part(target.substr(last));
+        output.push_back(part);
+    }
+
+    return output;
+}
+
+std::vector<std::string> split(std::string target, const char split_on, bool include_split_characters) {
+    std::vector<std::string> output;
+    size_t last = 0;
+
+    for (int i = 0; i < target.size(); i++) {
+        if (i > last && split_on == target[i]) {
+            std::string part(target.substr(last, i - last));
+            output.push_back(part);
+            if (include_split_characters) {
+                output.push_back(target.substr(i, 1));
+            }
+            last = i+1;
+        } else if (i == last && split_on == target[i]) {
+            if (include_split_characters) {
+                output.push_back(target.substr(i, 1));
+            }
+            last = i+1;
+        }
+    }
+    if (last < target.size()) {
+        std::string part(target.substr(last));
+        output.push_back(part);
+    }
+
+    return output;
+}
+
+std::string strip(std::string target, std::string vals) {
+    target.erase(target.begin(), std::find_if(target.begin(), target.end(), [&vals](unsigned char ch) {
+        return vals.find(ch) == std::string::npos;
+    }));
+    target.erase(std::find_if(target.rbegin(), target.rend(), [&vals](unsigned char ch) {
+        return vals.find(ch) == std::string::npos;
+    }).base(), target.end());
+    return target;
+}
+
+std::string replace_any(std::string target, std::string to_replace, std::string replacement) {
+    for (int i = 0; i < to_replace.size(); i++) {
+        size_t position = target.find(to_replace[i]);
+        while (position != std::string::npos) {
+            target.replace(position, 1, replacement);
+            position = target.find(to_replace[i]);
+        }
+    }
+    return target;
+}
+
+struct model_tensor_meta compute_tensor_meta(std::string name_prefix, ggml_context * weight_ctx, std::function<void(ggml_tensor*)>* callback) {
+    model_tensor_meta meta;
+    for (ggml_tensor * cur = ggml_get_first_tensor(weight_ctx); cur; cur = ggml_get_next_tensor(weight_ctx, cur)) {
+        if (callback) {
+            (*callback)(cur);
+        }
+        std::string::size_type pos = std::string(cur->name).find(".", 0);
+        std::string top_level(std::string(cur->name).substr(0, pos));
+        if (top_level == name_prefix) {
+            meta.n_tensors += 1;
+            meta.n_bytes += ggml_nbytes_pad(cur);
+        }
+    }
+    return meta;
+}

otherarch/ttscpp/src/ttsutil.h

@@ -0,0 +1,71 @@
+#ifndef util_h
+#define util_h
+
+#define _USE_MATH_DEFINES
+#include <cmath>
+#include <functional>
+#include <random>
+#include <stdio.h>
+#include <string>
+#include <cstring>
+#include <vector>
+#include <stdint.h>
+#include <sys/types.h>
+#include "ggml-metal.h"
+#include "ggml-backend.h"
+#include "ggml-alloc.h"
+#include "ggml-cpu.h"
+#include "ggml.h"
+#include "ggml-impl.h"
+#include "ggml-cpp.h"
+
+#define TTS_ABORT(...) tts_abort(__FILE__, __LINE__, __VA_ARGS__)
+#define TTS_ASSERT(x) if (!(x)) TTS_ABORT("TTS_ASSERT(%s) failed", #x)
+
+struct model_tensor_meta {
+	uint32_t n_tensors = 0;
+	size_t n_bytes = 0;
+};
+
+/**
+ * Both of these random fill the tgt array with count random floating point values.
+ * the default parameter values are consistent with pytorch random function defaults.
+ */
+void random_uniform_gen(int count, float * tgt, float min = 0.0f, float max = 1.0f);
+void random_normal_gen(int count, float * tgt, float mean = 0.0f, float std = 1.0f);
+
+std::pair<int, std::string> parse_layer_count(std::string name, int skip = 0);
+
+struct model_tensor_meta compute_tensor_meta(std::string name_prefix, ggml_context * weight_ctx, std::function<void(ggml_tensor*)>* callback = nullptr);
+struct ggml_tensor * snake_1d(ggml_context * ctx, struct ggml_tensor * alpha, struct ggml_tensor * a);
+int search_for_gguf_keys(gguf_context * meta, std::vector<std::string> possible_keys);
+
+// a simple window function for stft
+void hann_window(size_t n_fft, std::vector<float>& tgt);
+
+// currently this assumes a center view in which the output vector is reflectively padded by n_fft / 2 on each side.
+void compute_window_squared_sum(size_t n_fft, size_t hop, size_t n_frames, float * tgt, float * window);
+
+// these functions wrap the stft and istft ggml ops and compute the necessary view and division ops for their indepentent settings.
+struct ggml_tensor * stft(ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * window, size_t n_fft, size_t hop, bool abs_and_angle, bool one_sided);
+struct ggml_tensor * istft(ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * window_squared_sum, struct ggml_tensor * window, size_t n_fft, size_t hop, bool abs_and_angle, bool one_sided);
+
+// This is a custom op for sine_generation in the Kokoro model.
+void uv_noise_compute(struct ggml_tensor * dst, const struct ggml_tensor * a, const struct ggml_tensor * b, const struct ggml_tensor * c, int ith, int nth, void * userdata);
+
+// This is a custom op for logit correction in the Dia model.
+void cfg_scale(struct ggml_tensor * dst, const struct ggml_tensor * a, const struct ggml_tensor * b, int ith, int nth, void * userdata);
+
+struct ggml_tensor * reciprocal(ggml_context * ctx, struct ggml_tensor * x);
+
+bool has_suffix(std::string value, std::string suffix);
+bool has_prefix(std::string value, std::string prefix);
+
+std::vector<std::string> split(std::string target, std::string split_on, bool include_split_characters = false);
+std::vector<std::string> split(std::string target, const char split_on, bool include_split_characters = false);
+std::string strip(std::string target, std::string vals = " ");
+std::string replace_any(std::string target, std::string to_replace, std::string replacement);
+
+[[noreturn]] void tts_abort(const char * file, int line, const char * fmt, ...);
+
+#endif