Function bark_eval_fine_encoder

bark.cpp:2010–2108

}

static bool bark_eval_fine_encoder(struct bark_context* bctx, int n_threads) {
    // input shape: [N, n_codes]
    bark_codes input = bctx->coarse_tokens;

    std::vector<float> logits;
    logits.resize(1024 * 1056);

    auto& model = bctx->text_model.fine_model;
    auto& hparams = model.hparams;
    auto& params = bctx->params;

    float temp = params.fine_temp;

    int32_t n_coarse_codebooks = params.n_coarse_codebooks;
    int32_t n_fine_codebooks = params.n_fine_codebooks;
    int32_t codebook_size = params.codebook_size;

    int n_coarse = input[0].size();
    int original_seq_len = input.size();
    int n_remove_from_end = 0;

    // channel padding
    for (int i = 0; i < (int)input.size(); i++) {
        for (int j = n_coarse_codebooks; j < n_fine_codebooks; j++) {
            input[i].push_back(codebook_size);
        }
    }

    // spatial padding if sequence is too short
    if (original_seq_len < 1024) {
        n_remove_from_end = 1024 - original_seq_len;
        for (int i = original_seq_len; i < 1024; i++) {
            bark_sequence _tmp(n_fine_codebooks, codebook_size);
            input.push_back(_tmp);
        }
    }

    int n_loops = std::max(0, (int)ceilf((input.size() - 1024) / 512.f)) + 1;

    bark_codes in_arr = input;  // [seq_length, n_codes]

    for (int n = 0; n < n_loops; n++) {
        int start_idx = std::min(n * 512, (int)in_arr.size() - 1024);
        int start_fill_idx = std::min(n * 512, (int)in_arr.size() - 512);
        int rel_start_fill_idx = start_fill_idx - start_idx;

        // in_buffer: [n_codes*seq_length] (sequences are contiguous)
        bark_sequence in_buffer;
        for (int i = 0; i < n_fine_codebooks; i++) {
            for (int j = start_idx; j < start_idx + 1024; j++) {
                in_buffer.push_back(in_arr[j][i]);
            }
        }

        for (int nn = n_coarse; nn < n_fine_codebooks; nn++) {
            if (params.progress_callback) {
                const int progress_cur = 100*(n*(n_fine_codebooks-n_coarse)+(nn-n_coarse+1))/(n_loops*(n_fine_codebooks-n_coarse));