Refactor imatrix api, fix build shared libs

Author: stduhpf

examples/cli/main.cpp

@@ -22,10 +22,6 @@
 #define STB_IMAGE_RESIZE_STATIC
 #include "stb_image_resize.h"
 
-#define IMATRIX_IMPL
-#include "imatrix.hpp"
-static IMatrixCollector g_collector;
-
 const char* rng_type_to_str[] = {
     "std_default",
     "cuda",
@@ -663,7 +659,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         }
     }
 
-    if (params.imatrix_out.size() > 0 && file_exists(params.imatrix_out)) {
+    if (params.imatrix_out.size() > 0 && std::ifstream(params.imatrix_out).good()) {
         // imatrix file already exists
         if (std::find(params.imatrix_in.begin(), params.imatrix_in.end(), params.imatrix_out) == params.imatrix_in.end()) {
             printf("\n IMPORTANT: imatrix file %s already exists, but wasn't found in the imatrix inputs.\n", params.imatrix_out.c_str());
@@ -823,10 +819,6 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
     fflush(out_stream);
 }
 
-static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
-    return g_collector.collect_imatrix(t, ask, user_data);
-}
-
 int main(int argc, const char* argv[]) {
     SDParams params;
 
@@ -840,13 +832,12 @@ int main(int argc, const char* argv[]) {
     }
 
     if (params.imatrix_out != "") {
-        sd_set_backend_eval_callback((sd_graph_eval_callback_t)collect_imatrix, &params);
+        enableImatrixCollection();
     }
     if (params.imatrix_out != "" || params.mode == CONVERT || params.wtype != SD_TYPE_COUNT) {
-        setConvertImatrixCollector((void*)&g_collector);
         for (const auto& in_file : params.imatrix_in) {
             printf("loading imatrix from '%s'\n", in_file.c_str());
-            if (!g_collector.load_imatrix(in_file.c_str())) {
+            if (!loadImatrix(in_file.c_str())) {
                 printf("Failed to load %s\n", in_file.c_str());
             }
         }
@@ -1165,7 +1156,7 @@ int main(int argc, const char* argv[]) {
         results[i].data = NULL;
     }
     if (params.imatrix_out != "") {
-        g_collector.save_imatrix(params.imatrix_out);
+        saveImatrix(params.imatrix_out.c_str());
     }
     free(results);
     free_sd_ctx(sd_ctx);

model.cpp

@@ -29,7 +29,7 @@
 
 #define ST_HEADER_SIZE_LEN 8
 
-static IMatrixCollector* imatrix_collector = NULL;
+static IMatrixCollector imatrix_collector;
 
 uint64_t read_u64(uint8_t* buffer) {
     // little endian
@@ -1842,7 +1842,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
 
                     auto processed_name = convert_tensor_name(tensor_storage.name);
                     // LOG_DEBUG("%s",processed_name.c_str());
-                    std::vector<float> imatrix = imatrix_collector ? imatrix_collector->get_values(processed_name) : std::vector<float>{};
+                    std::vector<float> imatrix = imatrix_collector.get_values(processed_name);
 
                     convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
                                    dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], (int)tensor_storage.ne[0], imatrix);
@@ -1869,7 +1869,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, ggml_backend
                     // convert first, then copy to device memory
                     auto processed_name = convert_tensor_name(tensor_storage.name);
                     // LOG_DEBUG("%s",processed_name.c_str());
-                    std::vector<float> imatrix = imatrix_collector ? imatrix_collector->get_values(processed_name) : std::vector<float>{};
+                    std::vector<float> imatrix = imatrix_collector.get_values(processed_name);
 
                     convert_buffer.resize(ggml_nbytes(dst_tensor));
                     convert_tensor((void*)read_buffer.data(), tensor_storage.type,
@@ -2069,10 +2069,6 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type)
     return mem_size;
 }
 
-void setConvertImatrixCollector(void* collector) {
-    imatrix_collector = ((IMatrixCollector*)collector);
-}
-
 bool convert(const char* model_path, const char* clip_l_path, const char* clip_g_path, const char* t5xxl_path, const char* diffusion_model_path, const char* vae_path, const char* output_path, sd_type_t output_type) {
     ModelLoader model_loader;
 
@@ -2120,3 +2116,19 @@ bool convert(const char* model_path, const char* clip_l_path, const char* clip_g
     bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type);
     return success;
 }
+
+bool loadImatrix(const char* imatrix_path) {
+    return imatrix_collector.load_imatrix(imatrix_path);
+}
+void saveImatrix(const char* imatrix_path) {
+    imatrix_collector.save_imatrix(imatrix_path);
+}
+static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) {
+    return imatrix_collector.collect_imatrix(t, ask, user_data);
+}
+void enableImatrixCollection() {
+    sd_set_backend_eval_callback((sd_graph_eval_callback_t)collect_imatrix, NULL);
+}
+void disableImatrixCollection() {
+    sd_set_backend_eval_callback(NULL, NULL);
+}

stable-diffusion.h

@@ -230,7 +230,6 @@ SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
 SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor);
 
-SD_API void setConvertImatrixCollector(void * collector);
 SD_API bool convert(const char* model_path, const char* clip_l_path, const char* clip_g_path, const char* t5xxl_path, const char* diffusion_model_path, const char* vae_path, const char* output_path, enum sd_type_t output_type);
 
 SD_API uint8_t* preprocess_canny(uint8_t* img,
@@ -242,6 +241,11 @@ SD_API uint8_t* preprocess_canny(uint8_t* img,
                                  float strong,
                                  bool inverse);
 
+SD_API bool loadImatrix(const char * imatrix_path);
+SD_API void saveImatrix(const char * imatrix_path);
+SD_API void enableImatrixCollection();
+SD_API void disableImatrixCollection();
+
 #ifdef __cplusplus
 }
 #endif