WIP: feat: support for per-tensor quant types

Allows changing weight quants by regex, with a similar syntax to
the llama.cpp tensor overrides. This is useful for instance to reduce
unet memory requirements without changing the more sensitive clip
weights.

I only did the conversion part for now, but I intend to adapt the inference
code afterwards.

BTW, if we consider a plain quant name as matching any tensor, this could
reuse the --type command line argument instead of demanding another one.
Same thing for the type parameters for conversion and inference, which
in principle could be simply changed to a const char*.

Author: wbruna

examples/cli/main.cpp

@@ -84,6 +84,7 @@ struct SDParams {
     std::string stacked_id_embeddings_path;
     std::string input_id_images_path;
     sd_type_t wtype = SD_TYPE_COUNT;
+    std::string tensor_wtype;
     std::string lora_model_dir;
     std::string output_path = "output.png";
     std::string input_path;
@@ -204,6 +205,7 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --upscale-repeats                  Run the ESRGAN upscaler this many times (default 1)\n");
     printf("  --type [TYPE]                      weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, q4_K)\n");
     printf("                                     If not specified, the default is the type of the weight file\n");
+    printf("  --tensor_type [EXPRESSION]         weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")\n");
     printf("  --lora-model-dir [DIR]             lora model directory\n");
     printf("  -i, --init-img [IMAGE]             path to the input image, required by img2img\n");
     printf("  --mask [MASK]                      path to the mask image, required by img2img with mask\n");
@@ -381,6 +383,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                         valid_types.c_str());
                 exit(1);
             }
+        } else if (arg == "--tensor_type") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.tensor_wtype = argv[i];
         } else if (arg == "--lora-model-dir") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -800,7 +808,7 @@ int main(int argc, const char* argv[]) {
     }
 
     if (params.mode == CONVERT) {
-        bool success = convert(params.model_path.c_str(), params.vae_path.c_str(), params.output_path.c_str(), params.wtype);
+        bool success = convert(params.model_path.c_str(), params.vae_path.c_str(), params.output_path.c_str(), params.wtype, params.tensor_wtype.c_str());
         if (!success) {
             fprintf(stderr,
                     "convert '%s'/'%s' to '%s' failed\n",

model.cpp

@@ -1946,6 +1946,44 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
     return true;
 }
 
+std::vector<std::pair<std::string,ggml_type> > parse_quant_overrides (const std::string & overrides)
+{
+    std::vector<std::pair<std::string, ggml_type> > result;
+    for (const auto & item : splitString(overrides, ',')) {
+        std::string::size_type pos = item.find('=');
+        if (pos == std::string::npos) {
+            LOG_WARN("ignoring invalid quant override \"%s\"", item.c_str());
+            continue;
+        }
+        std::string tensor_pattern = item.substr(0, pos);
+        std::string quant_name = item.substr(pos + 1);
+
+        ggml_type over_type = GGML_TYPE_COUNT;
+
+        if (quant_name == "f32") {
+            over_type = GGML_TYPE_F32;
+        }
+        else {
+            for (size_t i = 0; i < SD_TYPE_COUNT; i++) {
+                auto trait = ggml_get_type_traits((ggml_type)i);
+                if (trait->to_float && trait->type_size && quant_name == trait->type_name) {
+                    over_type = (ggml_type)i;
+                }
+            }
+        }
+
+        if (over_type != GGML_TYPE_COUNT) {
+            //LOG_INFO("adding quant override \"%s\" ==> %s", tensor_pattern.c_str(), ggml_type_name(over_type));
+            result.emplace_back(tensor_pattern, over_type);
+        }
+        else {
+            LOG_WARN("ignoring invalid quant override \"%s\"", item.c_str());
+        }
+
+    }
+    return result;
+}
+
 bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) {
     const std::string& name = tensor_storage.name;
     if (type != GGML_TYPE_COUNT) {
@@ -1977,7 +2015,7 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
     return false;
 }
 
-bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type) {
+bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type, const char * overrides) {
     auto backend    = ggml_backend_cpu_init();
     size_t mem_size = 1 * 1024 * 1024;  // for padding
     mem_size += tensor_storages.size() * ggml_tensor_overhead();
@@ -1987,12 +2025,27 @@ bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type
 
     gguf_context* gguf_ctx = gguf_init_empty();
 
+    if (overrides == nullptr)
+        overrides = "";
+    auto quant_overrides = parse_quant_overrides(overrides);
+
     auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
         const std::string& name = tensor_storage.name;
-
         ggml_type tensor_type = tensor_storage.type;
-        if (tensor_should_be_converted(tensor_storage, type)) {
-            tensor_type = type;
+        ggml_type change_type = type;
+
+        for (const auto & quant_override : quant_overrides) {
+            std::regex pattern(quant_override.first);
+            if (std::regex_search(name, pattern)) {
+                change_type = quant_override.second;
+                //LOG_DEBUG("%s quant overriden to %s by \"%s\"",
+                //    name.c_str(), ggml_type_name(quant_override.second), quant_override.first.c_str());
+                break;
+            }
+        }
+
+        if (tensor_should_be_converted(tensor_storage, change_type)) {
+            tensor_type = change_type;
         }
 
         ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
@@ -2051,7 +2104,8 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type)
     return mem_size;
 }
 
-bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type) {
+bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type,
+    const char * output_tensor_type) {
     ModelLoader model_loader;
 
     if (!model_loader.init_from_file(input_path)) {
@@ -2065,6 +2119,6 @@ bool convert(const char* input_path, const char* vae_path, const char* output_pa
             return false;
         }
     }
-    bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type);
+    bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type, output_tensor_type);
     return success;
 }

model.h

@@ -221,7 +221,7 @@ class ModelLoader {
                       ggml_backend_t backend,
                       std::set<std::string> ignore_tensors = {});
 
-    bool save_to_gguf_file(const std::string& file_path, ggml_type type);
+    bool save_to_gguf_file(const std::string& file_path, ggml_type type, const char * overrides);
     bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
     int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
     ~ModelLoader() = default;

stable-diffusion.h

@@ -228,7 +228,7 @@ 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 bool convert(const char* input_path, const char* vae_path, const char* output_path, enum sd_type_t output_type);
+SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, enum sd_type_t output_type, const char * output_tensor_type);
 
 SD_API uint8_t* preprocess_canny(uint8_t* img,
                                  int width,