1// fix problem with std::min and std::max
2#if defined(_WIN32)
3#define WIN32_LEAN_AND_MEAN
4#ifndef NOMINMAX
5# define NOMINMAX
6#endif
7#include <windows.h>
8#endif
9
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_DEBUG
19
20#define MINIAUDIO_IMPLEMENTATION
21#ifndef MTMD_AUDIO_DEBUG
22# define MA_NO_ENCODING
23#endif
24#define MA_NO_DEVICE_IO
25#define MA_NO_RESOURCE_MANAGER
26#define MA_NO_NODE_GRAPH
27#define MA_NO_ENGINE
28#define MA_NO_GENERATION
29#define MA_API static
30#include "miniaudio/miniaudio.h"
31
32#define STB_IMAGE_IMPLEMENTATION
33#include "stb/stb_image.h"
34
35#ifdef MTMD_INTERNAL_HEADER
36#error "mtmd-helper is a public library outside of mtmd. it must not include internal headers"
37#endif
38
39//
40// internal logging functions
41//
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 functions
97//
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
117// helper struct to make working with embd batch easier
118// note: this will be removed after llama_batch_ext refactoring
119struct decode_embd_batch {
120 int n_pos_per_embd;
121 int n_mmproj_embd;
122 std::vector<llama_pos> pos;
123 std::vector<llama_pos> pos_view; // used by mrope
124 std::vector<int32_t> n_seq_id;
125 std::vector<llama_seq_id> seq_id_0;
126 std::vector<llama_seq_id *> seq_ids;
127 std::vector<int8_t> logits;
128 llama_batch batch;
129 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) {
130 pos .resize(n_tokens * n_pos_per_embd);
131 n_seq_id.resize(n_tokens);
132 seq_ids .resize(n_tokens + 1);
133 logits .resize(n_tokens);
134 seq_id_0.resize(1);
135 seq_ids [n_tokens] = nullptr;
136 batch = {
137 /*n_tokens =*/ n_tokens,
138 /*tokens =*/ nullptr,
139 /*embd =*/ embd,
140 /*pos =*/ pos.data(),
141 /*n_seq_id =*/ n_seq_id.data(),
142 /*seq_id =*/ seq_ids.data(),
143 /*logits =*/ logits.data(),
144 };
145 }
146
147 void set_position_normal(llama_pos pos_0, llama_seq_id seq_id) {
148 seq_id_0[0] = seq_id;
149 for (int i = 0; i < batch.n_tokens; i++) {
150 batch.pos [i] = pos_0 + i;
151 batch.n_seq_id[i] = 1;
152 batch.seq_id [i] = seq_id_0.data();
153 batch.logits [i] = false;
154 }
155 }
156
157 // M-RoPE for image
158 void set_position_mrope_2d(llama_pos pos_0, int nx, int ny, llama_seq_id seq_id) {
159 GGML_ASSERT(n_pos_per_embd == 4);
160 seq_id_0[0] = seq_id;
161 for (int y = 0; y < ny; y++) {
162 for (int x = 0; x < nx; x++) {
163 int i = y * nx + x;
164 pos[i ] = pos_0;
165 pos[i + batch.n_tokens ] = pos_0 + y;
166 pos[i + batch.n_tokens * 2] = pos_0 + x;
167 pos[i + batch.n_tokens * 3] = 0; // last pos dim is unused
168 }
169 }
170 for (int i = 0; i < batch.n_tokens; i++) {
171 batch.n_seq_id[i] = 1;
172 batch.seq_id [i] = seq_id_0.data();
173 batch.logits [i] = false;
174 }
175 }
176
177 // M-RoPE for audio
178 void set_position_mrope_1d(llama_pos pos_0, llama_seq_id seq_id) {
179 GGML_ASSERT(n_pos_per_embd == 4);
180 seq_id_0[0] = seq_id;
181 for (int i = 0; i < batch.n_tokens; i++) {
182 pos[i ] = pos_0 + i;
183 pos[i + batch.n_tokens ] = pos_0 + i;
184 pos[i + batch.n_tokens * 2] = pos_0 + i;
185 pos[i + batch.n_tokens * 3] = 0; // last pos dim is unused
186 }
187 for (int i = 0; i < batch.n_tokens; i++) {
188 batch.n_seq_id[i] = 1;
189 batch.seq_id [i] = seq_id_0.data();
190 batch.logits [i] = false;
191 }
192 }
193
194 llama_batch get_view(int offset, int n_tokens) {
195 llama_pos * pos_ptr;
196 pos_view.clear();
197 pos_view.reserve(n_tokens * n_pos_per_embd);
198 if (n_pos_per_embd > 1) {
199 // mrope
200 // for example, with layout of src: 1234...1234...1234...1234...
201 // offset 2 will give us dst: 34...34...34...34...
202 for (int i = 0; i < n_pos_per_embd; i++) {
203 // assume n_tokens is less than or equal to batch.n_tokens
204 // batch.n_tokens is number of **total** tokens
205 // n_tokens is number of viewed token
206 size_t src_idx = i * batch.n_tokens + offset;
207 pos_view.insert(pos_view.end(),
208 pos.data() + src_idx,
209 pos.data() + src_idx + n_tokens);
210 }
211 pos_ptr = pos_view.data();
212 } else {
213 // normal
214 pos_ptr = pos.data() + offset;
215 }
216 return {
217 /*n_tokens =*/ n_tokens,
218 /*tokens =*/ nullptr,
219 /*embd =*/ batch.embd + offset * n_mmproj_embd,
220 /*pos =*/ pos_ptr,
221 /*n_seq_id =*/ batch.n_seq_id + offset,
222 /*seq_id =*/ batch.seq_id + offset,
223 /*logits =*/ batch.logits + offset,
224 };
225 }
226};
227
228// Helper function for decoding an image whose embeddings have already been calculated
229int32_t mtmd_helper_decode_image_chunk(
230 mtmd_context * ctx,
231 struct llama_context * lctx,
232 const mtmd_input_chunk * chunk,
233 float * encoded_embd,
234 llama_pos n_past,
235 llama_seq_id seq_id,
236 int32_t n_batch,
237 llama_pos * new_n_past) {
238 auto chunk_type = mtmd_input_chunk_get_type(chunk);
239 const char * name = chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "image" : "audio";
240 if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
241 LOG_ERR("failed to decode chunk: input chunk not of image/audio type\n");
242 return -1;
243 }
244
245 const llama_model * model = llama_get_model(lctx);
246 int n_mmproj_embd = llama_model_n_embd_inp(model);
247 int n_pos_per_embd = mtmd_decode_use_mrope(ctx) ? 4 : 1;
248
249 int32_t n_tokens = mtmd_input_chunk_get_n_tokens(chunk);
250 int32_t i_batch = 0;
251 int32_t n_img_batches = GGML_PAD(n_tokens, n_batch) / n_batch;
252 decode_embd_batch batch_embd(encoded_embd, n_tokens, n_pos_per_embd, n_mmproj_embd);
253
254 if (mtmd_decode_use_mrope(ctx)) {
255 if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE) {
256 const auto image_tokens = mtmd_input_chunk_get_tokens_image(chunk);
257 if (!image_tokens) {
258 LOG_ERR("failed to decode chunk: image tokens are null\n");
259 return -1;
260 }
261 const int nx = mtmd_image_tokens_get_nx(image_tokens);
262 const int ny = mtmd_image_tokens_get_ny(image_tokens);
263 batch_embd.set_position_mrope_2d(n_past, nx, ny, seq_id);
264 } else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
265 batch_embd.set_position_mrope_1d(n_past, seq_id);
266 } else {
267 GGML_ABORT("invalid chunk type for M-RoPE");
268 }
269 } else {
270 batch_embd.set_position_normal(n_past, seq_id);
271 }
272
273 if (mtmd_decode_use_non_causal(ctx)) {
274 llama_set_causal_attn(lctx, false);
275 // TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image
276 }
277
278 while (i_batch < n_img_batches) { // split into batches
279 int pos_offset = i_batch*n_batch;
280 int n_tokens_batch = std::min(n_batch, n_tokens - pos_offset);
281 llama_batch batch_embd_view = batch_embd.get_view(pos_offset, n_tokens_batch);
282
283 LOG_INF("decoding %s batch %d/%d, n_tokens_batch = %d\n", name, i_batch+1, n_img_batches, n_tokens_batch);
284
285 int64_t t1 = ggml_time_ms();
286 int32_t ret = llama_decode(lctx, batch_embd_view);
287 if (ret != 0) {
288 LOG_ERR("failed to decode %s\n", name);
289 llama_set_causal_attn(lctx, true); // restore causal attn
290 return ret;
291 }
292
293 LOG_INF("%s decoded (batch %d/%d) in %" PRId64 " ms\n", name, i_batch+1, n_img_batches, ggml_time_ms() - t1);
294
295 i_batch++;
296 }
297
298 n_past += mtmd_input_chunk_get_n_pos(chunk);
299 *new_n_past = n_past;
300
301 if (mtmd_decode_use_non_causal(ctx)) {
302 llama_set_causal_attn(lctx, true);
303 }
304 return 0;
305}
306
307int32_t mtmd_helper_eval_chunk_single(mtmd_context * ctx,
308 struct llama_context * lctx,
309 const mtmd_input_chunk * chunk,
310 llama_pos n_past,
311 llama_seq_id seq_id,
312 int32_t n_batch,
313 bool logits_last,
314 llama_pos * new_n_past) {
315 int32_t ret;
316 llama_batch text_batch = llama_batch_init(n_batch, 0, 1);
317 auto chunk_type = mtmd_input_chunk_get_type(chunk);
318
319 if (chunk_type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
320 size_t n_tokens;
321 const auto tokens = mtmd_input_chunk_get_tokens_text(chunk, &n_tokens);
322 // LOG_INF("decoding text chunk, n_tokens = %zu\n", n_tokens);
323 size_t i = 0;
324 while (i < n_tokens) { // split into batches
325 text_batch.n_tokens = 0; // clear the batch
326 for (; i < n_tokens && text_batch.n_tokens < n_batch; i++) {
327 int32_t j = text_batch.n_tokens;
328 text_batch.token [j] = tokens[i];
329 text_batch.pos [j] = n_past++;
330 text_batch.n_seq_id[j] = 1;
331 text_batch.seq_id [j][0] = seq_id;
332 text_batch.logits [j] = false;
333
334 text_batch.n_tokens++;
335 }
336 bool is_last_token = (i == n_tokens);
337 if (logits_last && is_last_token) {
338 text_batch.logits[text_batch.n_tokens - 1] = true;
339 }
340 ret = llama_decode(lctx, text_batch);
341 if (ret != 0) {
342 LOG_ERR("failed to decode text\n");
343 llama_batch_free(text_batch);
344 return ret;
345 }
346 *new_n_past += text_batch.n_tokens;
347 }
348
349 } else if (chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE || chunk_type == MTMD_INPUT_CHUNK_TYPE_AUDIO) {
350 const char * name = chunk_type == MTMD_INPUT_CHUNK_TYPE_IMAGE ? "image" : "audio";
351 int64_t t0 = ggml_time_ms();
352
353 LOG_INF("encoding %s slice...\n", name);
354
355 ret = mtmd_encode_chunk(ctx, chunk);
356 if (ret != 0) {
357 LOG_ERR("failed to encode %s slice\n", name);
358 llama_batch_free(text_batch);
359 return ret;
360 }
361
362 LOG_INF("%s slice encoded in %" PRId64 " ms\n", name, ggml_time_ms() - t0);
363
364 float * embd = mtmd_get_output_embd(ctx);
365 ret = mtmd_helper_decode_image_chunk(ctx, lctx, chunk, embd, n_past, seq_id, n_batch, new_n_past);
366 if (ret != 0) {
367 LOG_ERR("failed to decode %s\n", name);
368 llama_batch_free(text_batch);
369 return ret;
370 }
371 } else {
372 GGML_ABORT("chunk type not supported");
373 }
374
375 llama_batch_free(text_batch);
376 return 0;
377}
378
379int32_t mtmd_helper_eval_chunks(mtmd_context * ctx,
380 struct llama_context * lctx,
381 const mtmd_input_chunks * chunks,
382 llama_pos n_past,
383 llama_seq_id seq_id,
384 int32_t n_batch,
385 bool logits_last,
386 llama_pos * new_n_past) {
387 size_t n_chunks = mtmd_input_chunks_size(chunks);
388 if (n_chunks == 0) {
389 LOG_WRN("no chunks to eval\n");
390 return 0;
391 }
392
393 for (size_t i = 0; i < n_chunks; i++) {
394 bool chunk_logits_last = (i == n_chunks - 1) && logits_last;
395 auto chunk = mtmd_input_chunks_get(chunks, i);
396
397 int32_t res = mtmd_helper_eval_chunk_single(ctx, lctx, chunk, n_past, seq_id, n_batch, chunk_logits_last, &n_past);
398 if (res != 0) {
399 LOG_ERR("failed to eval chunk %zu\n", i);
400 return res;
401 }
402 *new_n_past = n_past;
403 }
404
405 return 0;
406}
407
408namespace audio_helpers {
409
410static bool is_audio_file(const char * buf, size_t len) {
411 if (len < 12) {
412 return false;
413 }
414
415 // RIFF ref: https://en.wikipedia.org/wiki/Resource_Interchange_File_Format
416 // WAV ref: https://www.mmsp.ece.mcgill.ca/Documents/AudioFormats/WAVE/WAVE.html
417 bool is_wav = memcmp(buf, "RIFF", 4) == 0 && memcmp(buf + 8, "WAVE", 4) == 0;
418 bool is_mp3 = len >= 3 && (
419 memcmp(buf, "ID3", 3) == 0 ||
420 // Check for MPEG sync word (simplified check)
421 ((unsigned char)buf[0] == 0xFF && ((unsigned char)buf[1] & 0xE0) == 0xE0)
422 );
423 bool is_flac = memcmp(buf, "fLaC", 4) == 0;
424
425 return is_wav || is_mp3 || is_flac;
426}
427
428// returns true if the buffer is a valid audio file
429static bool decode_audio_from_buf(const unsigned char * buf_in, size_t len, int target_sampler_rate, std::vector<float> & pcmf32_mono) {
430 ma_result result;
431 const int channels = 1;
432 ma_decoder_config decoder_config = ma_decoder_config_init(ma_format_f32, channels, target_sampler_rate);
433 ma_decoder decoder;
434
435 result = ma_decoder_init_memory(buf_in, len, &decoder_config, &decoder);
436 if (result != MA_SUCCESS) {
437 return false;
438 }
439
440 ma_uint64 frame_count;
441 ma_uint64 frames_read;
442 result = ma_decoder_get_length_in_pcm_frames(&decoder, &frame_count);
443 if (result != MA_SUCCESS) {
444 ma_decoder_uninit(&decoder);
445 return false;
446 }
447
448 pcmf32_mono.resize(frame_count);
449 result = ma_decoder_read_pcm_frames(&decoder, pcmf32_mono.data(), frame_count, &frames_read);
450 if (result != MA_SUCCESS) {
451 ma_decoder_uninit(&decoder);
452 return false;
453 }
454
455#ifdef MTMD_AUDIO_DEBUG
456 // save audio to wav file
457 ma_encoder_config config = ma_encoder_config_init(ma_encoding_format_wav, ma_format_f32, 1, target_sampler_rate);
458 ma_encoder encoder;
459 ma_encoder_init_file("output.wav", &config, &encoder);
460 ma_encoder_write_pcm_frames(&encoder, pcmf32_mono.data(), pcmf32_mono.size(), &frames_read);
461 ma_encoder_uninit(&encoder);
462#endif
463
464 ma_decoder_uninit(&decoder);
465 return true;
466}
467
468} // namespace audio_helpers
469
470mtmd_bitmap * mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len) {
471 if (audio_helpers::is_audio_file((const char *)buf, len)) {
472 std::vector<float> pcmf32;
473 int bitrate = mtmd_get_audio_bitrate(ctx);
474 if (bitrate < 0) {
475 LOG_ERR("This model does not support audio input\n");
476 return nullptr;
477 }
478 if (!audio_helpers::decode_audio_from_buf(buf, len, bitrate, pcmf32)) {
479 LOG_ERR("Unable to read WAV audio file from buffer\n");
480 return nullptr;
481 }
482 return mtmd_bitmap_init_from_audio(pcmf32.size(), pcmf32.data());
483 }
484
485 // otherwise, we assume it's an image
486 mtmd_bitmap * result = nullptr;
487 {
488 int nx, ny, nc;
489 auto * data = stbi_load_from_memory(buf, len, &nx, &ny, &nc, 3);
490 if (!data) {
491 LOG_ERR("%s: failed to decode image bytes\n", __func__);
492 return nullptr;
493 }
494 result = mtmd_bitmap_init(nx, ny, data);
495 stbi_image_free(data);
496 }
497 return result;
498}
499
500mtmd_bitmap * mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname) {
501 std::vector<unsigned char> buf;
502 FILE * f = fopen(fname, "rb");
503 if (!f) {
504 LOG_ERR("Unable to open file %s: %s\n", fname, strerror(errno));
505 return nullptr;
506 }
507
508 fseek(f, 0, SEEK_END);
509 long file_size = ftell(f);
510 fseek(f, 0, SEEK_SET);
511 buf.resize(file_size);
512
513 size_t n_read = fread(buf.data(), 1, file_size, f);
514 fclose(f);
515 if (n_read != (size_t)file_size) {
516 LOG_ERR("Failed to read entire file %s", fname);
517 return nullptr;
518 }
519
520 return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size());
521}