/ examples / llava / llava.cpp
llava.cpp
  1  #include "clip.h"
  2  #include "common.h"
  3  #include "llama.h"
  4  #include "llava.h"
  5  #include "base64.hpp"
  6  
  7  #include <cstdio>
  8  #include <cstdlib>
  9  #include <vector>
 10  #include <numeric>
 11  
 12  // RGB uint8 image
 13  struct clip_image_u8 {
 14      int nx;
 15      int ny;
 16  
 17      std::vector<uint8_t> buf;
 18  };
 19  
 20  // RGB float32 image (NHWC)
 21  // Memory layout: RGBRGBRGB...
 22  struct clip_image_f32 {
 23      int nx;
 24      int ny;
 25  
 26      std::vector<float> buf;
 27  };
 28  
 29  struct clip_image_grid_shape {
 30      int first;
 31      int second;
 32  };
 33  
 34  /**
 35   * Selects the best resolution from a list of possible resolutions based on the original size.
 36   *
 37   * @param original_size The original size of the image in the format (width, height).
 38   * @param possible_resolutions A list of possible resolutions in the format [(width1, height1), (width2, height2), ...].
 39   * @return The best fit resolution in the format (width, height).
 40   */
 41  static std::pair<int, int> select_best_resolution(const std::pair<int, int>& original_size, const std::vector<std::pair<int, int>>& possible_resolutions) {
 42      int original_width  = original_size.first;
 43      int original_height = original_size.second;
 44  
 45      std::pair<int, int> best_fit;
 46      int max_effective_resolution = 0;
 47      int min_wasted_resolution = std::numeric_limits<int>::max();
 48  
 49      for (const auto& resolution : possible_resolutions) {
 50          int width = resolution.first;
 51          int height = resolution.second;
 52          float scale = std::min(static_cast<float>(width) / original_width, static_cast<float>(height) / original_height);
 53          int downscaled_width  = static_cast<int>(original_width * scale);
 54          int downscaled_height = static_cast<int>(original_height * scale);
 55          int effective_resolution = std::min(downscaled_width * downscaled_height, original_width * original_height);
 56          int wasted_resolution = (width * height) - effective_resolution;
 57          // LOG_TEE("resolution: %d %d, scale: %f, downscaled: %d %d, effective: %d, wasted: %d\n", width, height, scale, downscaled_width, downscaled_height, effective_resolution, wasted_resolution);
 58          if (effective_resolution > max_effective_resolution || (effective_resolution == max_effective_resolution && wasted_resolution < min_wasted_resolution)) {
 59              max_effective_resolution = effective_resolution;
 60              min_wasted_resolution = wasted_resolution;
 61              best_fit = resolution;
 62          }
 63      }
 64  
 65      return best_fit;
 66  }
 67  
 68  /**
 69   * @brief Get the anyres image grid shape object
 70   *
 71   * @param image_size
 72   * @param grid_pinpoints
 73   * @param image_patch_size
 74   * @return <int, int>
 75   */
 76  static struct clip_image_grid_shape get_anyres_image_grid_shape(const std::pair<int, int> & image_size, const std::vector<std::pair<int, int>> & grid_pinpoints, int image_patch_size) {
 77      /**
 78          Conversion from gguf flat array to vector:
 79          std::vector<std::pair<int, int>> possible_resolutions;
 80          for (int i = 0; i < 32 && params.image_grid_pinpoints[i] != 0; i+=2) {
 81              possible_resolutions.push_back({params.image_grid_pinpoints[i], params.image_grid_pinpoints[i+1]});
 82          }
 83       */
 84      auto best_resolution = select_best_resolution(image_size, grid_pinpoints);
 85      return {best_resolution.first / image_patch_size, best_resolution.second / image_patch_size};
 86  }
 87  
 88  // Take the image segments in a grid configuration and return the embeddings and the number of embeddings into preallocated memory (image_embd_out)
 89  static bool clip_llava_handle_patches(clip_ctx * ctx_clip, std::vector<float *> & image_embd_v, struct clip_image_grid_shape grid_shape, float * image_embd_out, int * n_img_pos_out) {
 90      struct {
 91          struct ggml_context * ctx;
 92      } model;
 93  
 94      const int32_t image_size = clip_image_size(ctx_clip);
 95      const int32_t patch_size = clip_patch_size(ctx_clip);
 96  
 97      int32_t num_patches_per_side = image_size / patch_size; // 336 / 14 = 24 - used for embedding-patching boxes (24*24 = 576 patches)
 98  
 99      int num_patches_width  = grid_shape.first;  // grid 1-4
100      int num_patches_height = grid_shape.second; // grid 1-4
101  
102      const size_t num_images = num_patches_width * num_patches_height + 1;
103  
104      // TODO: size calculation is not calculated - it's only tens of MB
105      size_t ctx_size = 0;
106  
107      {
108          ctx_size += clip_embd_nbytes(ctx_clip) * num_images * 8; // image_features
109          ctx_size += 1024*1024 * ggml_type_size(GGML_TYPE_F32);
110      }
111  
112      struct ggml_init_params params {
113          /*.mem_size   =*/ ctx_size,
114          /*.mem_buffer =*/ NULL,
115          /*.no_alloc   =*/ false, // NOTE: this should be false when using the legacy API
116      };
117  
118      // Python reference code for full unpad:
119      /*
120          base_image_feature = image_feature[0]
121          image_feature = image_feature[1:]
122          image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous()
123          image_feature = image_feature.flatten(1, 2).flatten(2, 3)
124          image_feature = unpad_image(image_feature, image_sizes[image_idx])
125          image_feature = torch.cat((
126              image_feature,
127              self.model.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1)
128          ), dim=-1)
129          image_feature = image_feature.flatten(1, 2).transpose(0, 1)
130          image_feature = torch.cat((base_image_feature, image_feature), dim=0)
131      */
132      // We now have two options: unpad or no unpad. Unpad removes tokens for faster llm eval.
133      // In terms of result quality it appears to make no difference, so we'll start with the easier approach given 5D tensors are not supported in ggml yet.
134      // Without unpad we have to split the sub-image embeddings into patches of 24 features each and permute them.
135      // Once all images are processed to prepended the base_image_features without any changes.
136  
137      // Pytorch reference simplified, modified for ggml compatibility - confirmed identical output in python (for a 2x2 grid image (676x676 scaling))
138      /*
139          image_feature = image_feature.view(2, 2, 24, 24, 4096)
140          image_feature = image_feature.permute(0, 2, 1, 3, 4).contiguous()
141          image_feature = image_feature.view(2, 24, 2, 24, 4096)
142          image_feature = image_feature.flatten(0, 3)
143  
144          // Reshape to 4D tensor by merging the last two dimensions
145          image_feature = image_feature.view(2, 2, 24, 24*4096)
146          image_feature = image_feature.permute(0, 2, 1, 3).contiguous()
147          image_feature = image_feature.view(-1, 4096)
148      */
149  
150      model.ctx = ggml_init(params);
151  
152      struct ggml_tensor * image_features = ggml_new_tensor_3d(model.ctx, GGML_TYPE_F32, clip_n_mmproj_embd(ctx_clip), clip_n_patches(ctx_clip), num_images - 1); // example: 4096 x 576 x 4
153      // ggml_tensor_printf(image_features,"image_features",__LINE__,false,false);
154      // fill it with the image embeddings, ignoring the base
155      for (size_t i = 1; i < num_images; i++) {
156          size_t offset = (i-1) * clip_embd_nbytes(ctx_clip);
157          memcpy((uint8_t *)(image_features->data) + offset, image_embd_v[i], clip_embd_nbytes(ctx_clip));
158      }
159  
160      struct ggml_cgraph  * gf = ggml_new_graph(model.ctx);
161      size_t size_ele = ggml_type_size(GGML_TYPE_F32);
162  
163      struct ggml_tensor *image_features_patchview = ggml_view_4d(model.ctx, image_features,
164                                                                  num_patches_per_side * clip_n_mmproj_embd(ctx_clip),
165                                                                  num_patches_per_side,
166                                                                  num_patches_width,
167                                                                  num_patches_height,
168                                                                  size_ele * num_patches_per_side * clip_n_mmproj_embd(ctx_clip),
169                                                                  size_ele * num_patches_per_side * clip_n_mmproj_embd(ctx_clip) * num_patches_per_side,
170                                                                  size_ele * num_patches_per_side * clip_n_mmproj_embd(ctx_clip) * num_patches_per_side * num_patches_width, 0);
171      // ggml_tensor_printf(image_features_patchview,"image_features_patchview",__LINE__,false,false);
172      struct ggml_tensor *permuted_cont = ggml_cont(model.ctx, ggml_permute(model.ctx, image_features_patchview, 0, 2, 1, 3));
173      /**
174       At the end of each row we have to add the row_end embeddings, which are the same as the newline embeddings
175           image_feature = torch.cat((
176          image_feature,
177          self.model.image_newline[:, None, None].expand(*image_feature.shape[:-1], 1).to(image_feature.device)
178      ), dim=-1)
179       *
180       */
181  
182      // ggml_tensor_printf(permuted_cont,"permuted_cont",__LINE__,false,false);
183      struct ggml_tensor *flatten = ggml_view_2d(model.ctx, permuted_cont, clip_n_mmproj_embd(ctx_clip), num_patches_height * num_patches_width * num_patches_per_side * num_patches_per_side,  size_ele * clip_n_mmproj_embd(ctx_clip), 0);
184      // ggml_tensor_printf(flatten,"flatten",__LINE__,false,false);
185      ggml_build_forward_expand(gf, flatten);
186      ggml_graph_compute_with_ctx(model.ctx, gf, 1);
187      struct ggml_tensor* result = gf->nodes[gf->n_nodes - 1];
188  
189      memcpy(image_embd_out, image_embd_v[0], clip_embd_nbytes(ctx_clip)); // main image as global context
190      // append without newline tokens (default behavior in llava_arch when not using unpad ):
191      memcpy(image_embd_out + clip_n_patches(ctx_clip) * clip_n_mmproj_embd(ctx_clip), (float*)result->data, clip_embd_nbytes(ctx_clip) * (num_images-1)); // grid patches
192      *n_img_pos_out = static_cast<int>(result->ne[1]+clip_n_patches(ctx_clip));
193  
194      // Debug: Test single segments
195      // Current findings: sending base image, sending a segment embedding all works similar to python
196      // However, permuted embeddings do not work yet (stride issue?)
197      // memcpy(image_embd_out, image_embd_v[0], clip_embd_nbytes(ctx_clip)); // main image as context
198      // memcpy(image_embd_out, (float*)prepared_cont->data, clip_embd_nbytes(ctx_clip)); // main image as context
199      // *n_img_pos_out=576;
200  
201      ggml_free(model.ctx);
202      return true;
203  }
204  
205  
206  static bool encode_image_with_clip(clip_ctx * ctx_clip, int n_threads, const clip_image_u8 * img, float * image_embd, int * n_img_pos) {
207      // std::vector<clip_image_f32*> img_res_v; // format VectN x H x W x RGB (N x 336 x 336 x 3), so interleaved RGB - different to the python implementation which is N x 3 x 336 x 336
208      clip_image_f32_batch img_res_v;
209      img_res_v.size = 0;
210      img_res_v.data = nullptr;
211      if (!clip_image_preprocess(ctx_clip, img, &img_res_v)) {
212          LOG_TEE("%s: unable to preprocess image\n", __func__);
213          delete[] img_res_v.data;
214          return false;
215      }
216  
217      const int64_t t_img_enc_start_us = ggml_time_us();
218  
219      const char * mm_patch_merge_type = clip_patch_merge_type(ctx_clip);
220  
221      if (strcmp(mm_patch_merge_type, "spatial_unpad") != 0) {
222          // flat / default llava-1.5 type embedding
223          *n_img_pos = clip_n_patches(ctx_clip);
224          bool encoded = clip_image_encode(ctx_clip, n_threads, &img_res_v.data[0], image_embd); // image_embd shape is 576 x 4096
225          delete[] img_res_v.data;
226          if (!encoded) {
227              LOG_TEE("Unable to encode image\n");
228  
229              return false;
230          }
231      } else {
232          // spatial_unpad llava-1.6 type embedding
233          // TODO: CLIP needs batching support - in HF the llm projection is separate after encoding, which might be a solution to quickly get batching working
234          std::vector<float *> image_embd_v;
235          image_embd_v.resize(img_res_v.size);
236          for (size_t i = 0; i < img_res_v.size; i++) {
237              image_embd_v[i] = (float *)malloc(clip_embd_nbytes(ctx_clip)); // 576 patches * 4096 embeddings * 4 bytes = 9437184
238              const bool encoded = clip_image_encode(ctx_clip, n_threads, &img_res_v.data[i], image_embd_v[i]); // image data is in 3x336x336 format and will be converted to 336x336x3 inside
239              if (!encoded) {
240                  LOG_TEE("Unable to encode image - spatial_unpad - subimage %d of %d\n", (int) i+1, (int) img_res_v.size);
241                  return false;
242              }
243          }
244          const int64_t t_img_enc_batch_us = ggml_time_us();
245          LOG_TEE("%s: %d segments encoded in %8.2f ms\n", __func__, (int)img_res_v.size, (t_img_enc_batch_us - t_img_enc_start_us) / 1000.0);
246  
247          const int32_t * image_grid = clip_image_grid(ctx_clip);
248  
249          std::vector<std::pair<int, int>> grid_pinpoints;
250          for (int i = 0; i < 32 && image_grid[i] != 0; i += 2) {
251              grid_pinpoints.push_back({image_grid[i], image_grid[i+1]});
252          }
253  
254          // free all img_res_v - not needed anymore
255          delete[] img_res_v.data;
256          img_res_v.size = 0;
257          img_res_v.data = nullptr;
258  
259          const int32_t image_size = clip_image_size(ctx_clip);
260  
261          struct clip_image_grid_shape grid_shape = get_anyres_image_grid_shape({img->nx,img->ny}, grid_pinpoints, image_size);
262  
263          int n_img_pos_out;
264          clip_llava_handle_patches(ctx_clip, image_embd_v, grid_shape, image_embd, &n_img_pos_out);
265          *n_img_pos = n_img_pos_out;
266  
267          for (size_t i = 0; i < image_embd_v.size(); i++) {
268              free(image_embd_v[i]);
269          }
270          image_embd_v.clear();
271  
272          // debug image/segment/normalization content:
273          // clip_image_u8 * tmp = clip_image_u8_init();
274          // clip_image_convert_f32_to_u8(*image_feature, *tmp);
275          // clip_image_save_to_bmp(*tmp, "image_feature.bmp");
276      }
277  
278      LOG_TEE("%s: image embedding created: %d tokens\n", __func__, *n_img_pos);
279  
280      const int64_t t_img_enc_end_us = ggml_time_us();
281      float t_img_enc_ms = (t_img_enc_end_us - t_img_enc_start_us) / 1000.0;
282  
283      LOG_TEE("\n%s: image encoded in %8.2f ms by CLIP (%8.2f ms per image patch)\n", __func__, t_img_enc_ms, t_img_enc_ms / *n_img_pos);
284  
285      return true;
286  }
287  
288  bool llava_validate_embed_size(const llama_context * ctx_llama, const clip_ctx * ctx_clip) {
289          // make sure that the correct mmproj was used, i.e., compare apples to apples
290      int n_llama_embd = llama_n_embd(llama_get_model(ctx_llama));
291      auto n_image_embd = clip_n_mmproj_embd(ctx_clip);
292      if (n_image_embd != n_llama_embd) {
293          LOG_TEE("%s: embedding dim of the multimodal projector (%d) is not equal to that of LLaMA (%d). Make sure that you use the correct mmproj file.\n", __func__, n_image_embd, n_llama_embd);
294          return false;
295      }
296      return true;
297  }
298  
299  bool llava_image_embed_make_with_clip_img(clip_ctx * ctx_clip, int n_threads, const clip_image_u8 * img, float ** image_embd_out, int * n_img_pos_out) {
300      float * image_embd = (float *)malloc(clip_embd_nbytes(ctx_clip)*6); // TODO: base on gridsize/llava model
301      if (!image_embd) {
302          LOG_TEE("Unable to allocate memory for image embeddings\n");
303          return false;
304      }
305  
306      int n_img_pos;
307      if (!encode_image_with_clip(ctx_clip, n_threads, img, image_embd, &n_img_pos)) {
308          LOG_TEE("%s: cannot encode image, aborting\n", __func__);
309          free(image_embd);
310          return false;
311      }
312      *image_embd_out = image_embd;
313      *n_img_pos_out = n_img_pos;
314  
315      return true;
316  }
317  
318  bool llava_eval_image_embed(llama_context * ctx_llama, const struct llava_image_embed * image_embed, int n_batch, int * n_past) {
319      int n_embd  = llama_n_embd(llama_get_model(ctx_llama));
320  
321      for (int i = 0; i < image_embed->n_image_pos; i += n_batch) {
322          int n_eval = image_embed->n_image_pos - i;
323          if (n_eval > n_batch) {
324              n_eval = n_batch;
325          }
326          llama_batch batch = {int32_t(n_eval), nullptr, (image_embed->embed+i*n_embd), nullptr, nullptr, nullptr, nullptr, *n_past, 1, 0, };
327          if (llama_decode(ctx_llama, batch)) {
328              LOG_TEE("%s : failed to eval\n", __func__);
329              return false;
330          }
331          *n_past += n_eval;
332      }
333      return true;
334  }
335  
336  struct llava_image_embed * llava_image_embed_make_with_bytes(struct clip_ctx * ctx_clip, int n_threads, const unsigned char * image_bytes, int image_bytes_length) {
337      clip_image_u8 * img = clip_image_u8_init();
338      if (!clip_image_load_from_bytes(image_bytes, image_bytes_length, img)) {
339          clip_image_u8_free(img);
340          LOG_TEE("%s: can't load image from bytes, is it a valid image?", __func__);
341          return NULL;
342      }
343  
344      float* image_embed = NULL;
345      int n_image_pos = 0;
346      bool image_embed_result = llava_image_embed_make_with_clip_img(ctx_clip, n_threads, img, &image_embed, &n_image_pos);
347      if (!image_embed_result) {
348          clip_image_u8_free(img);
349          LOG_TEE("%s: coulnd't embed the image\n", __func__);
350          return NULL;
351      }
352  
353      clip_image_u8_free(img);
354      auto result = (llava_image_embed*)malloc(sizeof(llava_image_embed));
355      result->embed = image_embed;
356      result->n_image_pos = n_image_pos;
357      return result;
358  }
359  
360  static bool load_file_to_bytes(const char* path, unsigned char** bytesOut, long *sizeOut) {
361      auto file = fopen(path, "rb");
362      if (file == NULL) {
363          LOG_TEE("%s: can't read file %s\n", __func__, path);
364          return false;
365      }
366  
367      fseek(file, 0, SEEK_END);
368      auto fileSize = ftell(file);
369      fseek(file, 0, SEEK_SET);
370  
371      auto buffer = (unsigned char *)malloc(fileSize); // Allocate memory to hold the file data
372      if (buffer == NULL) {
373          LOG_TEE("%s: failed to alloc %ld bytes for file %s\n", __func__, fileSize, path);
374          perror("Memory allocation error");
375          fclose(file);
376          return false;
377      }
378      errno = 0;
379      size_t ret = fread(buffer, 1, fileSize, file); // Read the file into the buffer
380      if (ferror(file)) {
381          die_fmt("read error: %s", strerror(errno));
382      }
383      if (ret != (size_t) fileSize) {
384          die("unexpectedly reached end of file");
385      }
386      fclose(file); // Close the file
387  
388      *bytesOut = buffer;
389      *sizeOut = fileSize;
390      return true;
391  }
392  
393  struct llava_image_embed * llava_image_embed_make_with_filename(struct clip_ctx * ctx_clip, int n_threads, const char * image_path) {
394      unsigned char* image_bytes;
395      long image_bytes_length;
396      auto loaded = load_file_to_bytes(image_path, &image_bytes, &image_bytes_length);
397      if (!loaded) {
398          LOG_TEE("%s: failed to load %s\n", __func__, image_path);
399          return NULL;
400      }
401  
402      llava_image_embed *embed = llava_image_embed_make_with_bytes(ctx_clip, n_threads, image_bytes, image_bytes_length);
403      free(image_bytes);
404  
405      return embed;
406  }
407  
408  void llava_image_embed_free(struct llava_image_embed * embed) {
409      free(embed->embed);
410      free(embed);
411  }