fix: tensor loading thread count (#854)

* fix: avoid segfault for pix2pix models with no reference images

* fix: default to empty reference on pix2pix models to avoid segfault

* use resize instead of reserve

* format code

---------

Co-authored-by: leejet <leejet714@gmail.com>

.github/workflows/build.yml

@@ -254,7 +254,7 @@ jobs:
 
       - name: Copy and pack Cuda runtime
         id: pack_cuda_runtime
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' && matrix.build == 'cuda12' ) || github.event.inputs.create_release == 'true' }}
+        if: ${{ matrix.build == 'cuda12' && (github.event_name == 'push' && github.ref == 'refs/heads/master' || github.event.inputs.create_release == 'true') }}
         run: |
           echo "Cuda install location: ${{steps.cuda-toolkit.outputs.CUDA_PATH}}"
           $dst='.\build\bin\cudart\'
@@ -262,7 +262,7 @@ jobs:
           7z a cudart-sd-bin-win-cu12-x64.zip $dst\*
 
       - name: Upload Cuda runtime
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' && matrix.build == 'cuda12' ) || github.event.inputs.create_release == 'true' }}
+        if: ${{ matrix.build == 'cuda12' && (github.event_name == 'push' && github.ref == 'refs/heads/master' || github.event.inputs.create_release == 'true') }}
         uses: actions/upload-artifact@v4
         with:
           name: sd-cudart-sd-bin-win-cu12-x64.zip
@@ -288,6 +288,11 @@ jobs:
       - windows-latest-cmake
 
     steps:
+      - name: Clone
+        uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
       - name: Download artifacts
         id: download-artifact
         uses: actions/download-artifact@v4
@@ -298,7 +303,8 @@ jobs:
 
       - name: Get commit count
         id: commit_count
-        run: echo "count=$(git rev-list --count HEAD)" >> $GITHUB_OUTPUT
+        run: |
+          echo "count=$(git rev-list --count HEAD)" >> $GITHUB_OUTPUT
 
       - name: Get commit hash
         id: commit
@@ -306,17 +312,16 @@ jobs:
 
       - name: Create release
         id: create_release
+        if: ${{ github.event_name == 'workflow_dispatch' || github.ref_name == 'master' }}
         uses: anzz1/action-create-release@v1
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         with:
-          tag_name: >
-            ${{ github.ref_name == 'master' &&
-                format('release_{0}_{1}', steps.commit_count.outputs.count, steps.commit.outputs.short) ||
-                format('{0}-{1}', env.BRANCH_NAME, steps.commit.outputs.short) }}
+          tag_name: ${{ format('{0}-{1}-{2}', env.BRANCH_NAME, steps.commit_count.outputs.count, steps.commit.outputs.short) }}
 
       - name: Upload release
         id: upload_release
+        if: ${{ github.event_name == 'workflow_dispatch' || github.ref_name == 'master' }}
         uses: actions/github-script@v3
         with:
           github-token: ${{secrets.GITHUB_TOKEN}}

README.md

@@ -125,13 +125,14 @@ cmake --build . --config Release
 
 ##### Using HipBLAS
 This provides BLAS acceleration using the ROCm cores of your AMD GPU. Make sure to have the ROCm toolkit installed.
+To build for another GPU architecture than installed in your system, set `$GFX_NAME` manually to the desired architecture (replace first command). This is also necessary if your GPU is not officially supported by ROCm, for example you have to set `$GFX_NAME` manually to `gfx1030` for consumer RDNA2 cards.
 
 Windows User Refer to [docs/hipBLAS_on_Windows.md](docs%2FhipBLAS_on_Windows.md) for a comprehensive guide.
 
 ```
-export GFX_NAME=$(rocminfo | grep -m 1 -E "gfx[^0]{1}" | sed -e 's/ *Name: *//' | awk '{$1=$1; print}' || echo "rocminfo missing")
-echo $GFX_NAME
-cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DGPU_TARGETS=$GFX_NAME -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON
+if command -v rocminfo; then export GFX_NAME=$(rocminfo | awk '/ *Name: +gfx[1-9]/ {print $2; exit}'); else echo "rocminfo missing!"; fi
+if [ -z "${GFX_NAME}" ]; then echo "Error: Couldn't detect GPU!"; else echo "Building for GPU: ${GFX_NAME}"; fi
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DGPU_TARGETS=$GFX_NAME -DAMDGPU_TARGETS=$GFX_NAME -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_POSITION_INDEPENDENT_CODE=ON
 cmake --build . --config Release
 ```
 
@@ -384,7 +385,6 @@ arguments:
   --pm-id-images-dir [DIR]           path to PHOTOMAKER input id images dir
   --pm-id-embed-path [PATH]          path to PHOTOMAKER v2 id embed
   --pm-style-strength                strength for keeping PHOTOMAKER input identity (default: 20)
-  --normalize-input                  normalize PHOTOMAKER input id images
   -v, --verbose                      print extra info
 ```
 

conditioner.hpp

@@ -141,7 +141,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
             }
             return true;
         };
-        model_loader.load_tensors(on_load);
+        model_loader.load_tensors(on_load, 1);
         readed_embeddings.push_back(embd_name);
         if (embd) {
             int64_t hidden_size = text_model->model.hidden_size;

control.hpp

@@ -445,7 +445,7 @@ struct ControlNet : public GGMLRunner {
         guided_hint_cached = true;
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading control net from '%s'", file_path.c_str());
         alloc_params_buffer();
         std::map<std::string, ggml_tensor*> tensors;
@@ -458,7 +458,7 @@ struct ControlNet : public GGMLRunner {
             return false;
         }
 
-        bool success = model_loader.load_tensors(tensors, ignore_tensors);
+        bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);
 
         if (!success) {
             LOG_ERROR("load control net tensors from model loader failed");

docs/lora.md

@@ -20,20 +20,30 @@ Here's a simple example:
 
 NOTE: The other backends may have different support.
 
-| Quant / Type | CUDA |
-|--------------|------|
-| F32          | ✔️   |
-| F16          | ✔️   |
-| BF16         | ✔️   |
-| I32          | ✔️   |
-| Q4_0         | ✔️   |
-| Q4_1         | ✔️   |
-| Q5_0         | ✔️   |
-| Q5_1         | ✔️   |
-| Q8_0         | ✔️   |
-| Q2_K         | ❌   |
-| Q3_K         | ❌   |
-| Q4_K         | ❌   |
-| Q5_K         | ❌   |
-| Q6_K         | ❌   |
-| Q8_K         | ❌   |
+| Quant / Type | CUDA | Vulkan |
+|--------------|------|--------|
+| F32          | ✔️   | ✔️   |
+| F16          | ✔️   | ✔️   |
+| BF16         | ✔️   | ✔️   |
+| I32          | ✔️   | ❌   |
+| Q4_0         | ✔️   | ✔️   |
+| Q4_1         | ✔️   | ✔️   |
+| Q5_0         | ✔️   | ✔️   |
+| Q5_1         | ✔️   | ✔️   |
+| Q8_0         | ✔️   | ✔️   |
+| Q2_K         | ❌   | ❌   |
+| Q3_K         | ❌   | ❌   |
+| Q4_K         | ❌   | ❌   |
+| Q5_K         | ❌   | ❌   |
+| Q6_K         | ❌   | ❌   |
+| Q8_K         | ❌   | ❌   |
+| IQ1_S        | ❌   | ✔️   |
+| IQ1_M        | ❌   | ✔️   |
+| IQ2_XXS      | ❌   | ✔️   |
+| IQ2_XS       | ❌   | ✔️   |
+| IQ2_S        | ❌   | ✔️   |
+| IQ3_XXS      | ❌   | ✔️   |
+| IQ3_S        | ❌   | ✔️   |
+| IQ4_XS       | ❌   | ✔️   |
+| IQ4_NL       | ❌   | ✔️   |
+| MXFP4        | ❌   | ✔️   |

esrgan.hpp

@@ -164,7 +164,7 @@ struct ESRGAN : public GGMLRunner {
         return "esrgan";
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading esrgan from '%s'", file_path.c_str());
 
         alloc_params_buffer();
@@ -177,7 +177,7 @@ struct ESRGAN : public GGMLRunner {
             return false;
         }
 
-        bool success = model_loader.load_tensors(esrgan_tensors);
+        bool success = model_loader.load_tensors(esrgan_tensors, {}, n_threads);
 
         if (!success) {
             LOG_ERROR("load esrgan tensors from model loader failed");

examples/cli/main.cpp

@@ -103,7 +103,6 @@ struct SDParams {
     bool verbose               = false;
     bool offload_params_to_cpu = false;
     bool control_net_cpu       = false;
-    bool normalize_input       = false;
     bool clip_on_cpu           = false;
     bool vae_on_cpu            = false;
     bool diffusion_flash_attn  = false;
@@ -156,7 +155,6 @@ void print_params(SDParams params) {
     printf("    pm_id_images_dir:                  %s\n", params.pm_id_images_dir.c_str());
     printf("    pm_id_embed_path:                  %s\n", params.pm_id_embed_path.c_str());
     printf("    pm_style_strength:                 %.2f\n", params.pm_style_strength);
-    printf("    normalize input image:             %s\n", params.normalize_input ? "true" : "false");
     printf("    output_path:                       %s\n", params.output_path.c_str());
     printf("    init_image_path:                   %s\n", params.init_image_path.c_str());
     printf("    end_image_path:                    %s\n", params.end_image_path.c_str());
@@ -306,7 +304,6 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --pm-id-images-dir [DIR]           path to PHOTOMAKER input id images dir\n");
     printf("  --pm-id-embed-path [PATH]          path to PHOTOMAKER v2 id embed\n");
     printf("  --pm-style-strength                strength for keeping PHOTOMAKER input identity (default: 20)\n");
-    printf("  --normalize-input                  normalize PHOTOMAKER input id images\n");
     printf("  -v, --verbose                      print extra info\n");
 }
 
@@ -552,7 +549,6 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         {"", "--vae-tiling", "", true, &params.vae_tiling_params.enabled},
         {"", "--offload-to-cpu", "", true, &params.offload_params_to_cpu},
         {"", "--control-net-cpu", "", true, &params.control_net_cpu},
-        {"", "--normalize-input", "", true, &params.normalize_input},
         {"", "--clip-on-cpu", "", true, &params.clip_on_cpu},
         {"", "--vae-on-cpu", "", true, &params.vae_on_cpu},
         {"", "--diffusion-fa", "", true, &params.diffusion_flash_attn},
@@ -1379,7 +1375,6 @@ int main(int argc, const char* argv[]) {
             params.batch_count,
             control_image,
             params.control_strength,
-            params.normalize_input,
             {
                 pmid_images.data(),
                 (int)pmid_images.size(),

lora.hpp

@@ -116,7 +116,7 @@ struct LoraModel : public GGMLRunner {
         return "lora";
     }
 
-    bool load_from_file(bool filter_tensor = false, int n_threads = 0) {
+    bool load_from_file(bool filter_tensor, int n_threads) {
         LOG_INFO("loading LoRA from '%s'", file_path.c_str());
 
         if (load_failed) {

model.cpp

@@ -1,4 +1,5 @@
 #include <stdarg.h>
+#include <algorithm>
 #include <atomic>
 #include <chrono>
 #include <fstream>
@@ -1956,7 +1957,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     std::atomic<int64_t> copy_to_backend_time_ms(0);
     std::atomic<int64_t> convert_time_ms(0);
 
-    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : (int)std::thread::hardware_concurrency();
+    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : get_num_physical_cores();
+    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
 
     int64_t start_time = ggml_time_ms();
     std::vector<TensorStorage> processed_tensor_storages;
@@ -2006,13 +2008,25 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
             w.join();
         }
 
-        std::unordered_map<std::string, IndexedStorage> latest_map;
+        std::vector<IndexedStorage> deduplicated;
+        deduplicated.reserve(all_results.size());
+        std::unordered_map<std::string, size_t> name_to_pos;
         for (auto& entry : all_results) {
-            latest_map[entry.ts.name] = entry;
+            auto it = name_to_pos.find(entry.ts.name);
+            if (it == name_to_pos.end()) {
+                name_to_pos.emplace(entry.ts.name, deduplicated.size());
+                deduplicated.push_back(entry);
+            } else if (deduplicated[it->second].index < entry.index) {
+                deduplicated[it->second] = entry;
+            }
         }
 
-        processed_tensor_storages.reserve(latest_map.size());
-        for (auto& [name, entry] : latest_map) {
+        std::sort(deduplicated.begin(), deduplicated.end(), [](const IndexedStorage& a, const IndexedStorage& b) {
+            return a.index < b.index;
+        });
+
+        processed_tensor_storages.reserve(deduplicated.size());
+        for (auto& entry : deduplicated) {
             processed_tensor_storages.push_back(entry.ts);
         }
     }
@@ -2427,6 +2441,7 @@ bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type
 
     auto tensor_type_rules = parse_tensor_type_rules(tensor_type_rules_str);
 
+    std::mutex tensor_mutex;
     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;
@@ -2444,6 +2459,7 @@ bool ModelLoader::save_to_gguf_file(const std::string& file_path, ggml_type type
             tensor_type = dst_type;
         }
 
+        std::lock_guard<std::mutex> lock(tensor_mutex);
         ggml_tensor* tensor = ggml_new_tensor(ggml_ctx, tensor_type, tensor_storage.n_dims, tensor_storage.ne);
         if (tensor == NULL) {
             LOG_ERROR("ggml_new_tensor failed");

pmid.hpp

@@ -591,7 +591,7 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
         return "id_embeds";
     }
 
-    bool load_from_file(bool filter_tensor = false) {
+    bool load_from_file(bool filter_tensor, int n_threads) {
         LOG_INFO("loading PhotoMaker ID Embeds from '%s'", file_path.c_str());
 
         if (load_failed) {
@@ -599,7 +599,8 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             return false;
         }
 
-        bool dry_run          = true;
+        bool dry_run = true;
+        std::mutex tensor_mutex;
         auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
             const std::string& name = tensor_storage.name;
 
@@ -608,6 +609,7 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
                 return true;
             }
             if (dry_run) {
+                std::lock_guard<std::mutex> lock(tensor_mutex);
                 struct ggml_tensor* real = ggml_new_tensor(params_ctx,
                                                            tensor_storage.type,
                                                            tensor_storage.n_dims,
@@ -621,11 +623,11 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             return true;
         };
 
-        model_loader->load_tensors(on_new_tensor_cb);
+        model_loader->load_tensors(on_new_tensor_cb, n_threads);
         alloc_params_buffer();
 
         dry_run = false;
-        model_loader->load_tensors(on_new_tensor_cb);
+        model_loader->load_tensors(on_new_tensor_cb, n_threads);
 
         LOG_DEBUG("finished loading PhotoMaker ID Embeds ");
         return true;

stable-diffusion.cpp

@@ -443,6 +443,10 @@ public:
             diffusion_model->alloc_params_buffer();
             diffusion_model->get_param_tensors(tensors);
 
+            if (sd_version_is_unet_edit(version)) {
+                vae_decode_only = false;
+            }
+
             if (high_noise_diffusion_model) {
                 high_noise_diffusion_model->alloc_params_buffer();
                 high_noise_diffusion_model->get_param_tensors(tensors);
@@ -527,7 +531,7 @@ public:
             }
             if (strlen(SAFE_STR(sd_ctx_params->photo_maker_path)) > 0) {
                 pmid_lora = std::make_shared<LoraModel>(backend, sd_ctx_params->photo_maker_path, "");
-                if (!pmid_lora->load_from_file(true)) {
+                if (!pmid_lora->load_from_file(true, n_threads)) {
                     LOG_WARN("load photomaker lora tensors from %s failed", sd_ctx_params->photo_maker_path);
                     return false;
                 }
@@ -595,14 +599,14 @@ public:
             if (!use_tiny_autoencoder) {
                 vae_params_mem_size = first_stage_model->get_params_buffer_size();
             } else {
-                if (!tae_first_stage->load_from_file(taesd_path)) {
+                if (!tae_first_stage->load_from_file(taesd_path, n_threads)) {
                     return false;
                 }
                 vae_params_mem_size = tae_first_stage->get_params_buffer_size();
             }
             size_t control_net_params_mem_size = 0;
             if (control_net) {
-                if (!control_net->load_from_file(SAFE_STR(sd_ctx_params->control_net_path))) {
+                if (!control_net->load_from_file(SAFE_STR(sd_ctx_params->control_net_path), n_threads)) {
                     return false;
                 }
                 control_net_params_mem_size = control_net->get_params_buffer_size();
@@ -748,15 +752,15 @@ public:
                 denoiser->scheduler->version = version;
                 break;
             case SGM_UNIFORM:
-                    LOG_INFO("Running with SGM Uniform schedule");
-                    denoiser->scheduler          = std::make_shared<SGMUniformSchedule>();
-                    denoiser->scheduler->version = version;
-                    break;
+                LOG_INFO("Running with SGM Uniform schedule");
+                denoiser->scheduler          = std::make_shared<SGMUniformSchedule>();
+                denoiser->scheduler->version = version;
+                break;
             case SIMPLE:
-                    LOG_INFO("Running with Simple schedule");
-                    denoiser->scheduler          = std::make_shared<SimpleSchedule>();
-                    denoiser->scheduler->version = version;
-                    break;
+                LOG_INFO("Running with Simple schedule");
+                denoiser->scheduler          = std::make_shared<SimpleSchedule>();
+                denoiser->scheduler->version = version;
+                break;
             case SMOOTHSTEP:
                 LOG_INFO("Running with SmoothStep scheduler");
                 denoiser->scheduler = std::make_shared<SmoothStepSchedule>();
@@ -832,7 +836,7 @@ public:
             return;
         }
         LoraModel lora(backend, file_path, is_high_noise ? "model.high_noise_" : "");
-        if (!lora.load_from_file()) {
+        if (!lora.load_from_file(false, n_threads)) {
             LOG_WARN("load lora tensors from %s failed", file_path.c_str());
             return;
         }
@@ -1053,7 +1057,7 @@ public:
                         ggml_tensor* denoise_mask             = NULL,
                         ggml_tensor* vace_context             = NULL,
                         float vace_strength                   = 1.f) {
-         if (shifted_timestep > 0 && !sd_version_is_sdxl(version)) {
+        if (shifted_timestep > 0 && !sd_version_is_sdxl(version)) {
             LOG_WARN("timestep shifting is only supported for SDXL models!");
             shifted_timestep = 0;
         }
@@ -1127,7 +1131,7 @@ public:
             } else {
                 timesteps_vec.assign(1, t);
             }
-            
+
             timesteps_vec  = process_timesteps(timesteps_vec, init_latent, denoise_mask);
             auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec);
             std::vector<float> guidance_vec(1, guidance.distilled_guidance);
@@ -1790,7 +1794,6 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_img_gen_params->seed              = -1;
     sd_img_gen_params->batch_count       = 1;
     sd_img_gen_params->control_strength  = 0.9f;
-    sd_img_gen_params->normalize_input   = false;
     sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
     sd_img_gen_params->vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
 }
@@ -1816,7 +1819,6 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "ref_images_count: %d\n"
              "increase_ref_index: %s\n"
              "control_strength: %.2f\n"
-             "normalize_input: %s\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
              "VAE tiling: %s\n",
              SAFE_STR(sd_img_gen_params->prompt),
@@ -1831,7 +1833,6 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->ref_images_count,
              BOOL_STR(sd_img_gen_params->increase_ref_index),
              sd_img_gen_params->control_strength,
-             BOOL_STR(sd_img_gen_params->normalize_input),
              sd_img_gen_params->pm_params.style_strength,
              sd_img_gen_params->pm_params.id_images_count,
              SAFE_STR(sd_img_gen_params->pm_params.id_embed_path),
@@ -1915,7 +1916,6 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
                                     int batch_count,
                                     sd_image_t control_image,
                                     float control_strength,
-                                    bool normalize_input,
                                     sd_pm_params_t pm_params,
                                     std::vector<ggml_tensor*> ref_latents,
                                     bool increase_ref_index,
@@ -2387,19 +2387,35 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
         init_latent = generate_init_latent(sd_ctx, work_ctx, width, height);
     }
 
-    if (sd_img_gen_params->ref_images_count > 0) {
+    sd_guidance_params_t guidance = sd_img_gen_params->sample_params.guidance;
+    std::vector<sd_image_t*> ref_images;
+    for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) {
+        ref_images.push_back(&sd_img_gen_params->ref_images[i]);
+    }
+
+    std::vector<uint8_t> empty_image_data;
+    sd_image_t empty_image = {(uint32_t)width, (uint32_t)height, 3, nullptr};
+    if (ref_images.empty() && sd_version_is_unet_edit(sd_ctx->sd->version)) {
+        LOG_WARN("This model needs at least one reference image; using an empty reference");
+        empty_image_data.resize(width * height * 3);
+        ref_images.push_back(&empty_image);
+        empty_image.data = empty_image_data.data();
+        guidance.img_cfg = 0.f;
+    }
+
+    if (ref_images.size() > 0) {
         LOG_INFO("EDIT mode");
     }
 
     std::vector<ggml_tensor*> ref_latents;
-    for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) {
+    for (int i = 0; i < ref_images.size(); i++) {
         ggml_tensor* img = ggml_new_tensor_4d(work_ctx,
                                               GGML_TYPE_F32,
-                                              sd_img_gen_params->ref_images[i].width,
-                                              sd_img_gen_params->ref_images[i].height,
+                                              ref_images[i]->width,
+                                              ref_images[i]->height,
                                               3,
                                               1);
-        sd_image_to_tensor(sd_img_gen_params->ref_images[i], img);
+        sd_image_to_tensor(*ref_images[i], img);
 
         ggml_tensor* latent = NULL;
         if (sd_ctx->sd->use_tiny_autoencoder) {
@@ -2437,7 +2453,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
                                                         SAFE_STR(sd_img_gen_params->prompt),
                                                         SAFE_STR(sd_img_gen_params->negative_prompt),
                                                         sd_img_gen_params->clip_skip,
-                                                        sd_img_gen_params->sample_params.guidance,
+                                                        guidance,
                                                         sd_img_gen_params->sample_params.eta,
                                                         sd_img_gen_params->sample_params.shifted_timestep,
                                                         width,
@@ -2448,7 +2464,6 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
                                                         sd_img_gen_params->batch_count,
                                                         sd_img_gen_params->control_image,
                                                         sd_img_gen_params->control_strength,
-                                                        sd_img_gen_params->normalize_input,
                                                         sd_img_gen_params->pm_params,
                                                         ref_latents,
                                                         sd_img_gen_params->increase_ref_index,

stable-diffusion.h

@@ -212,7 +212,6 @@ typedef struct {
     int batch_count;
     sd_image_t control_image;
     float control_strength;
-    bool normalize_input;
     sd_pm_params_t pm_params;
     sd_tiling_params_t vae_tiling_params;
 } sd_img_gen_params_t;

tae.hpp

@@ -222,7 +222,7 @@ struct TinyAutoEncoder : public GGMLRunner {
         return "taesd";
     }
 
-    bool load_from_file(const std::string& file_path) {
+    bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading taesd from '%s', decode_only = %s", file_path.c_str(), decode_only ? "true" : "false");
         alloc_params_buffer();
         std::map<std::string, ggml_tensor*> taesd_tensors;
@@ -238,7 +238,7 @@ struct TinyAutoEncoder : public GGMLRunner {
             return false;
         }
 
-        bool success = model_loader.load_tensors(taesd_tensors, ignore_tensors);
+        bool success = model_loader.load_tensors(taesd_tensors, ignore_tensors, n_threads);
 
         if (!success) {
             LOG_ERROR("load tae tensors from model loader failed");

upscaler.cpp

@@ -18,7 +18,8 @@ struct UpscalerGGML {
     }
 
     bool load_from_file(const std::string& esrgan_path,
-                        bool offload_params_to_cpu) {
+                        bool offload_params_to_cpu,
+                        int n_threads) {
         ggml_log_set(ggml_log_callback_default, nullptr);
 #ifdef SD_USE_CUDA
         LOG_DEBUG("Using CUDA backend");
@@ -54,7 +55,7 @@ struct UpscalerGGML {
         if (direct) {
             esrgan_upscaler->enable_conv2d_direct();
         }
-        if (!esrgan_upscaler->load_from_file(esrgan_path)) {
+        if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) {
             return false;
         }
         return true;
@@ -124,7 +125,7 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
         return NULL;
     }
 
-    if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, offload_params_to_cpu)) {
+    if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, offload_params_to_cpu, n_threads)) {
         delete upscaler_ctx->upscaler;
         upscaler_ctx->upscaler = NULL;
         free(upscaler_ctx);