echodict/llama.cpp

1// fix problem with std::min and std::max2#if defined(_WIN32)3#define WIN32_LEAN_AND_MEAN4#ifndef NOMINMAX5#   define NOMINMAX6#endif7#include <windows.h>8#endif9 10#include "mtmd.h"11#include "mtmd-helper.h"12#include "llama.h"13 14#include <algorithm>15#include <cinttypes>16#include <vector>17 18//#define MTMD_AUDIO_DEBUG19 20#define MINIAUDIO_IMPLEMENTATION21#ifndef MTMD_AUDIO_DEBUG22#   define MA_NO_ENCODING23#endif24#define MA_NO_DEVICE_IO25#define MA_NO_RESOURCE_MANAGER26#define MA_NO_NODE_GRAPH27#define MA_NO_ENGINE28#define MA_NO_GENERATION29#define MA_API static30#include "miniaudio/miniaudio.h"31 32#define STB_IMAGE_IMPLEMENTATION33#include "stb/stb_image.h"34 35#ifdef MTMD_INTERNAL_HEADER36#error "mtmd-helper is a public library outside of mtmd. it must not include internal headers"37#endif38 39//40// internal logging functions41//42 43struct mtmd_helper_logger {44    ggml_log_callback default_callback = [](ggml_log_level level, const char * text, void * user_data) {45        (void) level;46        (void) user_data;47        fputs(text, stderr);48        fflush(stderr);49    };50 51    ggml_log_callback log_callback = default_callback;52    void * log_callback_user_data;53 54    void log_v(enum ggml_log_level level, const char * format, va_list args) {55        if (format == NULL) {56            return;57        }58        va_list args_copy;59        va_copy(args_copy, args);60        char buffer[128];61        int len = vsnprintf(buffer, 128, format, args);62        if (len < 128) {63            log_callback(level, buffer, log_callback_user_data);64        } else {65            char * buffer2 = (char *) calloc(len + 1, sizeof(char));66            vsnprintf(buffer2, len + 1, format, args_copy);67            buffer2[len] = 0;68            log_callback(level, buffer2, log_callback_user_data);69            free(buffer2);70        }71        va_end(args_copy);72    }73 74    void log(enum ggml_log_level level, const char * format, ...) {75        va_list args;76        va_start(args, format);77        log_v(level, format, args);78        va_end(args);79    }80} g_logger;81 82#define LOG_INF(...) g_logger.log(GGML_LOG_LEVEL_INFO,  __VA_ARGS__)83#define LOG_WRN(...) g_logger.log(GGML_LOG_LEVEL_WARN,  __VA_ARGS__)84#define LOG_ERR(...) g_logger.log(GGML_LOG_LEVEL_ERROR, __VA_ARGS__)85 86void mtmd_helper_log_set(ggml_log_callback log_callback, void * user_data) {87    if (log_callback == nullptr) {88        log_callback = g_logger.default_callback;89    }90    g_logger.log_callback = log_callback;91    g_logger.log_callback_user_data = user_data;92    mtmd_log_set(log_callback, user_data);93}94 95//96// helper functions97//98 99size_t mtmd_helper_get_n_tokens(const mtmd_input_chunks * chunks) {100    size_t n_tokens = 0;101    for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {102        auto chunk = mtmd_input_chunks_get(chunks, i);103        n_tokens += mtmd_input_chunk_get_n_tokens(chunk);104    }105    return n_tokens;106}107 108llama_pos mtmd_helper_get_n_pos(const mtmd_input_chunks * chunks) {109    llama_pos n_pos = 0;110    for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {111        auto chunk = mtmd_input_chunks_get(chunks, i);112        n_pos += mtmd_input_chunk_get_n_pos(chunk);113    }114    return n_pos;115}116 117void mtmd_helper_image_get_decoder_pos(const mtmd_image_tokens * chunks, llama_pos pos_0, mtmd_decoder_pos * out_pos) {118    size_t n_tokens = mtmd_image_tokens_get_n_tokens(chunks);119    for (size_t i = 0; i < n_tokens; i++) {120        out_pos[i] = mtmd_image_tokens_get_decoder_pos(chunks, pos_0, i);121    }122}123 124// helper struct to make working with embd batch easier125// note: this will be removed after llama_batch_ext refactoring126struct decode_embd_batch {127    int n_pos_per_embd;128    int n_mmproj_embd;129    std::vector<llama_pos>      pos;130    std::vector<llama_pos>      pos_view; // used by mrope131    std::vector<int32_t>        n_seq_id;132    std::vector<llama_seq_id>   seq_id_0;133    std::vector<llama_seq_id *> seq_ids;134    std::vector<int8_t>         logits;135    llama_batch batch;136    decode_embd_batch(float * embd, int32_t n_tokens, int n_pos_per_embd, int n_mmproj_embd) : n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) {137        GGML_ASSERT(n_tokens > 0 && n_pos_per_embd > 0 && n_mmproj_embd > 0);138        pos     .resize(n_tokens * n_pos_per_embd);139        n_seq_id.resize(n_tokens);140        seq_ids .resize(n_tokens + 1);141        logits  .resize(n_tokens);142        seq_id_0.resize(1);143        seq_ids [n_tokens] = nullptr;144        batch = {145            /*n_tokens       =*/ n_tokens,146            /*tokens         =*/ nullptr,147            /*embd           =*/ embd,148            /*pos            =*/ pos.data(),149            /*n_seq_id       =*/ n_seq_id.data(),150            /*seq_id         =*/ seq_ids.data(),151            /*logits         =*/ logits.data(),152        };153    }154 155    void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) {156        seq_id_0[0] = seq_id;157        for (int i = 0; i < batch.n_tokens; i++) {158            batch.pos     [i] = pos_0 + i;159            batch.n_seq_id[i] = 1;160            batch.seq_id  [i] = seq_id_0.data();161            batch.logits  [i] = false;162        }163    }164 165    // M-RoPE for image166    void set_position_mrope_2d(const std::vector<mtmd_decoder_pos> & rel_pos, llama_seq_id seq_id) {167        GGML_ASSERT(n_pos_per_embd == 4);168        GGML_ASSERT(!rel_pos.empty() && (int32_t)rel_pos.size() == batch.n_tokens);169        seq_id_0[0] = seq_id;170        for (int32_t i = 0; i < batch.n_tokens; i++) {171            pos[i                     ] = rel_pos[i].t;172            pos[i + batch.n_tokens    ] = rel_pos[i].y;173            pos[i + batch.n_tokens * 2] = rel_pos[i].x;174            pos[i + batch.n_tokens * 3] = rel_pos[i].z;175        }176        for (int i = 0; i < batch.n_tokens; i++) {177            batch.n_seq_id[i] = 1;178            batch.seq_id  [i] = seq_id_0.data();179            batch.logits  [i] = false;180        }181    }182 183    // M-RoPE for audio184    void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) {185        GGML_ASSERT(n_pos_per_embd == 4);186        seq_id_0[0] = seq_id;187        for (int i = 0; i < batch.n_tokens; i++) {188            pos[i                     ] = pos_0 + i;189            pos[i + batch.n_tokens    ] = pos_0 + i;190            pos[i + batch.n_tokens * 2] = pos_0 + i;191            pos[i + batch.n_tokens * 3] = pos_0 + i;192        }193        for (int i = 0; i < batch.n_tokens; i++) {194            batch.n_seq_id[i] = 1;195            batch.seq_id  [i] = seq_id_0.data();196            batch.logits  [i] = false;197        }198    }199 200    llama_batch get_view(int offset, int n_tokens) {201        GGML_ASSERT(offset >= 0 && n_tokens > 0 && offset + n_tokens <= batch.n_tokens);202        llama_pos * pos_ptr;203        pos_view.clear();204        pos_view.reserve(n_tokens * n_pos_per_embd);205        if (n_pos_per_embd > 1) {206            // mrope207            // for example, with layout of src: 1234...1234...1234...1234...208            //       offset 2 will give us dst: 34...34...34...34...209            for (int i = 0; i < n_pos_per_embd; i++) {210                // assume n_tokens is less than or equal to batch.n_tokens211                // batch.n_tokens is number of **total** tokens212                // n_tokens is number of viewed token213                size_t src_idx = i * batch.n_tokens + offset;214                pos_view.insert(pos_view.end(),215                    pos.data() + src_idx,216                    pos.data() + src_idx + n_tokens);217            }218            pos_ptr = pos_view.data();219        } else {220            // normal221            pos_ptr = pos.data() + offset;222        }223        return {224            /*n_tokens       =*/ n_tokens,225            /*tokens         =*/ nullptr,226            /*embd           =*/ batch.embd     + offset * n_mmproj_embd,227            /*pos            =*/ pos_ptr,228            /*n_seq_id       =*/ batch.n_seq_id + offset,229            /*seq_id         =*/ batch.seq_id   + offset,230            /*logits         =*/ batch.logits   + offset,231        };232    }233};234 235// Helper function for decoding an image whose embeddings have already been calculated236int32_t mtmd_helper_decode_image_chunk(237        mtmd_context * ctx,238        struct llama_context * lctx,239        const mtmd_input_chunk * chunk,240        float * encoded_embd,241        llama_pos n_past,242        llama_seq_id seq_id,243        int32_t n_batch,244        llama_pos * new_n_past) {245    GGML_ASSERT(n_batch > 0);246    auto chunk_type = mtmd_input_chunk_get_type(chunk);247    const char * name = chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "image" : "audio";248    if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {249        LOG_ERR("failed to decode chunk: input chunk not of image/audio type\n");250        return -1;251    }252 253    const llama_model * model = llama_get_model(lctx);254    int n_mmproj_embd = llama_model_n_embd_inp(model);255    int n_pos_per_embd = mtmd_decode_use_mrope(ctx) ? 4 : 1;256 257    int32_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk);258    int32_t i_batch = 0;259    int32_t n_img_batches = (n_tokens + n_batch - 1) / n_batch;260    decode_embd_batch batch_embd(encoded_embd, n_tokens, n_pos_per_embd, n_mmproj_embd);261 262    if (mtmd_decode_use_mrope(ctx)) {263        if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE) {264            const auto image_tokens = mtmd_input_chunk_get_tokens_image(chunk);265            if (!image_tokens) {266                LOG_ERR("failed to decode chunk: image tokens are null\n");267                return -1;268            }269            const auto n_tokens = mtmd_image_tokens_get_n_tokens(image_tokens);270            std::vector<mtmd_decoder_pos> rel_pos(n_tokens);271            mtmd_helper_image_get_decoder_pos(image_tokens, n_past, rel_pos.data());272            batch_embd.set_position_mrope_2d(rel_pos, seq_id);273        } else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {274            batch_embd.set_position_mrope_1d(n_past, seq_id);275        } else {276            GGML_ABORT("invalid chunk type for M-RoPE");277        }278    } else {279        batch_embd.set_position_normal(n_past, seq_id);280    }281 282    const bool use_non_causal = mtmd_decode_use_non_causal(ctx, chunk);283    if (use_non_causal) {284        llama_set_causal_attn(lctx, false);285        // TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image286    }287 288    while (i_batch < n_img_batches) { // split into batches289        int pos_offset = i_batch*n_batch;290        int n_tokens_batch = std::min(n_batch, n_tokens - pos_offset);291        llama_batch batch_embd_view = batch_embd.get_view(pos_offset, n_tokens_batch);292 293        LOG_INF("decoding %s batch %d/%d, n_tokens_batch = %d\n", name, i_batch+1, n_img_batches, n_tokens_batch);294 295        int64_t t1 = ggml_time_ms();296        int32_t ret = llama_decode(lctx, batch_embd_view);297        if (ret != 0) {298            LOG_ERR("failed to decode %s\n", name);299            llama_set_causal_attn(lctx, true); // restore causal attn300            return ret;301        }302 303        LOG_INF("%s decoded (batch %d/%d) in %" PRId64 " ms\n", name, i_batch+1, n_img_batches, ggml_time_ms() - t1);304 305        i_batch++;306    }307 308    n_past += mtmd_input_chunk_get_n_pos(chunk);309    *new_n_past = n_past;310 311    if (use_non_causal) {312        llama_set_causal_attn(lctx, true);313    }314    return 0;315}316 317int32_t mtmd_helper_eval_chunk_single(mtmd_context * ctx,318        struct llama_context * lctx,319        const mtmd_input_chunk * chunk,320        llama_pos n_past,321        llama_seq_id seq_id,322        int32_t n_batch,323        bool logits_last,324        llama_pos * new_n_past) {325    GGML_ASSERT(n_batch > 0);326    int32_t ret;327    llama_batch text_batch = llama_batch_init(n_batch, 0, 1);328    auto chunk_type = mtmd_input_chunk_get_type(chunk);329 330    if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {331        size_t n_tokens;332        const auto tokens = mtmd_input_chunk_get_tokens_text(chunk, &n_tokens);333        // LOG_INF("decoding text chunk, n_tokens = %zu\n", n_tokens);334        size_t i = 0;335        while (i < n_tokens) { // split into batches336            text_batch.n_tokens = 0; // clear the batch337            for (; i < n_tokens && text_batch.n_tokens < n_batch; i++) {338                int32_t j = text_batch.n_tokens;339                text_batch.token   [j]    = tokens[i];340                text_batch.pos     [j]    = n_past++;341                text_batch.n_seq_id[j]    = 1;342                text_batch.seq_id  [j][0] = seq_id;343                text_batch.logits  [j]    = false;344 345                text_batch.n_tokens++;346            }347            bool is_last_token = (i == n_tokens);348            if (logits_last && is_last_token) {349                text_batch.logits[text_batch.n_tokens - 1] = true;350            }351            ret = llama_decode(lctx, text_batch);352            if (ret != 0) {353                LOG_ERR("failed to decode text\n");354                llama_batch_free(text_batch);355                return ret;356            }357            *new_n_past += text_batch.n_tokens;358        }359 360    } else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE || chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {361        const char * name = chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "image" : "audio";362        int64_t t0 = ggml_time_ms();363 364        LOG_INF("encoding %s slice...\n", name);365 366        ret = mtmd_encode_chunk(ctx, chunk);367        if (ret != 0) {368            LOG_ERR("failed to encode %s slice\n", name);369            llama_batch_free(text_batch);370            return ret;371        }372 373        LOG_INF("%s slice encoded in %" PRId64 " ms\n", name, ggml_time_ms() - t0);374 375        float * embd = mtmd_get_output_embd(ctx);376        ret = mtmd_helper_decode_image_chunk(ctx, lctx, chunk, embd, n_past, seq_id, n_batch, new_n_past);377        if (ret != 0) {378            LOG_ERR("failed to decode %s\n", name);379            llama_batch_free(text_batch);380            return ret;381        }382    } else {383        GGML_ABORT("chunk type not supported");384    }385 386    llama_batch_free(text_batch);387    return 0;388}389 390int32_t mtmd_helper_eval_chunks(mtmd_context * ctx,391                                struct llama_context * lctx,392                                const mtmd_input_chunks * chunks,393                                llama_pos n_past,394                                llama_seq_id seq_id,395                                int32_t n_batch,396                                bool logits_last,397                                llama_pos * new_n_past) {398    size_t n_chunks = mtmd_input_chunks_size(chunks);399    if (n_chunks == 0) {400        LOG_WRN("no chunks to eval\n");401        return 0;402    }403 404    for (size_t i = 0; i < n_chunks; i++) {405        bool chunk_logits_last = (i == n_chunks - 1) && logits_last;406        auto chunk = mtmd_input_chunks_get(chunks, i);407 408        int32_t res = mtmd_helper_eval_chunk_single(ctx, lctx, chunk, n_past, seq_id, n_batch, chunk_logits_last, &n_past);409        if (res != 0) {410            LOG_ERR("failed to eval chunk %zu\n", i);411            return res;412        }413        *new_n_past = n_past;414    }415 416    return 0;417}418 419namespace audio_helpers {420 421static bool is_audio_file(const char * buf, size_t len) {422    if (len < 12) {423        return false;424    }425 426    // RIFF ref: https://en.wikipedia.org/wiki/Resource_Interchange_File_Format427    // WAV ref: https://www.mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html428    bool is_wav = memcmp(buf, "RIFF", 4) == 0 && memcmp(buf + 8, "WAVE", 4) == 0;429    bool is_mp3 = len >= 3 && (430        memcmp(buf, "ID3", 3) == 0 ||431        // Check for MPEG sync word (simplified check)432        ((unsigned char)buf[0] == 0xFF && ((unsigned char)buf[1] & 0xE0) == 0xE0)433    );434    bool is_flac = memcmp(buf, "fLaC", 4) == 0;435 436    return is_wav || is_mp3 || is_flac;437}438 439// returns true if the buffer is a valid audio file440static bool decode_audio_from_buf(const unsigned char * buf_in, size_t len, int target_sampler_rate, std::vector<float> & pcmf32_mono) {441    ma_result result;442    const int channels = 1;443    ma_decoder_config decoder_config = ma_decoder_config_init(ma_format_f32, channels, target_sampler_rate);444    ma_decoder decoder;445 446    result = ma_decoder_init_memory(buf_in, len, &decoder_config, &decoder);447    if (result != MA_SUCCESS) {448        return false;449    }450 451    ma_uint64 frame_count;452    ma_uint64 frames_read;453    result = ma_decoder_get_length_in_pcm_frames(&decoder, &frame_count);454    if (result != MA_SUCCESS) {455        ma_decoder_uninit(&decoder);456        return false;457    }458 459    pcmf32_mono.resize(frame_count);460    result = ma_decoder_read_pcm_frames(&decoder, pcmf32_mono.data(), frame_count, &frames_read);461    if (result != MA_SUCCESS) {462        ma_decoder_uninit(&decoder);463        return false;464    }465 466#ifdef MTMD_AUDIO_DEBUG467    // save audio to wav file468    ma_encoder_config config = ma_encoder_config_init(ma_encoding_format_wav, ma_format_f32, 1, target_sampler_rate);469    ma_encoder encoder;470    ma_encoder_init_file("output.wav", &config, &encoder);471    ma_encoder_write_pcm_frames(&encoder, pcmf32_mono.data(), pcmf32_mono.size(), &frames_read);472    ma_encoder_uninit(&encoder);473#endif474 475    ma_decoder_uninit(&decoder);476    return true;477}478 479} // namespace audio_helpers480 481mtmd_bitmap * mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len) {482    if (audio_helpers::is_audio_file((const char *)buf, len)) {483        std::vector<float> pcmf32;484        const int sample_rate = mtmd_get_audio_sample_rate(ctx);485        if (sample_rate < 0) {486            LOG_ERR("This model does not support audio input\n");487            return nullptr;488        }489        if (!audio_helpers::decode_audio_from_buf(buf, len, sample_rate, pcmf32)) {490            LOG_ERR("Unable to read WAV audio file from buffer\n");491            return nullptr;492        }493        return mtmd_bitmap_init_from_audio(pcmf32.size(), pcmf32.data());494    }495 496    // otherwise, we assume it's an image497    mtmd_bitmap * result = nullptr;498    {499        int nx, ny, nc;500        auto * data = stbi_load_from_memory(buf, len, &nx, &ny, &nc, 3);501        if (!data) {502            LOG_ERR("%s: failed to decode image bytes\n", __func__);503            return nullptr;504        }505        result = mtmd_bitmap_init(nx, ny, data);506        stbi_image_free(data);507    }508    return result;509}510 511mtmd_bitmap * mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname) {512    std::vector<unsigned char> buf;513    FILE * f = fopen(fname, "rb");514    if (!f) {515        LOG_ERR("Unable to open file %s: %s\n", fname, strerror(errno));516        return nullptr;517    }518 519    fseek(f, 0, SEEK_END);520    long file_size = ftell(f);521    fseek(f, 0, SEEK_SET);522    if (file_size < 0) {523        LOG_ERR("Failed to get file size of %s\n", fname);524        fclose(f);525        return nullptr;526    }527    buf.resize(file_size);528 529    size_t n_read = fread(buf.data(), 1, file_size, f);530    fclose(f);531    if (n_read != (size_t)file_size) {532        LOG_ERR("Failed to read entire file %s", fname);533        return nullptr;534    }535 536    return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size());537}538