added input id image loading

examples/cli/main.cpp

@@ -64,6 +64,7 @@ struct SDParams {
     std::string controlnet_path;
     std::string embeddings_path;
     std::string stacked_id_embeddings_path;
+    std::string input_id_images_path;
     sd_type_t wtype = SD_TYPE_COUNT;
     std::string lora_model_dir;
     std::string output_path = "output.png";
@@ -103,6 +104,7 @@ void print_params(SDParams params) {
     printf("    controlnet_path:   %s\n", params.controlnet_path.c_str());
     printf("    embeddings_path:   %s\n", params.embeddings_path.c_str());
     printf("    stacked_id_embeddings_path:   %s\n", params.stacked_id_embeddings_path.c_str());
+    printf("    input_id_images_path:   %s\n", params.input_id_images_path.c_str());
     printf("    output_path:       %s\n", params.output_path.c_str());
     printf("    init_img:          %s\n", params.input_path.c_str());
     printf("    control_image:     %s\n", params.control_image_path.c_str());
@@ -137,7 +139,8 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --taesd [TAESD_PATH]               path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)\n");
     printf("  --control-net [CONTROL_PATH]       path to control net model\n");
     printf("  --embd-dir [EMBEDDING_PATH]        path to embeddings.\n");
-    printf("  --stacked-id-embd-dir [ID_EMBEDDING_PATH]  path to photomakerstacked id embeddings.\n");
+    printf("  --stacked-id-embd-dir [DIR]        path to PHOTOMAKER stacked id embeddings.\n");
+    printf("  --input-id-images-dir [DIR]        path to PHOTOMAKER input id images dir.\n");
     printf("  --upscale-model [ESRGAN_PATH]      path to esrgan model. Upscale images after generate, just RealESRGAN_x4plus_anime_6B supported by now.\n");
     printf("  --type [TYPE]                      weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)\n");
     printf("                                     If not specified, the default is the type of the weight file.\n");
@@ -240,6 +243,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.stacked_id_embeddings_path = argv[i];
+        } else if (arg == "--input-id-images-dir") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.input_id_images_path = argv[i];
         } else if (arg == "--type") {
             if (++i >= argc) {
                 invalid_arg = true;
@@ -616,6 +625,27 @@ int main(int argc, const char* argv[]) {
                 control_image->data = preprocess_canny(control_image->data, control_image->width, control_image->height);
             }
         }
+        std::vector<sd_image_t*> input_id_images;
+        if (params.stacked_id_embeddings_path.size() > 0 && params.input_id_images_path.size() > 0) {
+            std::vector<std::string> img_files = get_files_from_dir(params.input_id_images_path);
+            for(std::string img_file : img_files){
+                int c = 0;
+                int width, height;
+                input_image_buffer = stbi_load(img_file.c_str(), &width, &height, &c, 3);
+                if (input_image_buffer == NULL) {
+                    fprintf(stderr, "load image from '%s' failed\n", img_file.c_str());
+                    return 1;
+                }else{
+                    fprintf(stderr, "successfully loaded image from '%s' \n", img_file.c_str());
+                }
+                sd_image_t* input_image = NULL;
+                input_image = new sd_image_t{(uint32_t)width,
+                                             (uint32_t)height,
+                                             3,
+                                             input_image_buffer};
+                input_id_images.push_back(input_image);
+            }
+        }
         results = txt2img(sd_ctx,
                           params.prompt.c_str(),
                           params.negative_prompt.c_str(),

util.cpp

@@ -100,6 +100,7 @@ bool is_directory(const std::string& path) {
     return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
+// TODO: add windows version 
 std::string get_full_path(const std::string& dir, const std::string& filename) {
     DIR* dp = opendir(dir.c_str());
 
@@ -119,6 +120,29 @@ std::string get_full_path(const std::string& dir, const std::string& filename) {
     return "";
 }
 
+std::vector<std::string> get_files_from_dir(const std::string &dir){
+
+    std::vector<std::string> files;
+
+    DIR* dp = opendir(dir.c_str());
+
+    if (dp != nullptr) {
+        struct dirent* entry;
+
+        while ((entry = readdir(dp)) != nullptr) {
+            std::string fname = dir + "/" + entry->d_name;
+            if (!is_directory(fname))
+                files.push_back(fname);
+        }
+        closedir(dp);
+    }
+
+    sort(files.begin(), files.end());
+
+    return files;
+
+}
+
 #endif
 
 // get_num_physical_cores is copy from

util.h

@@ -3,6 +3,7 @@
 
 #include <cstdint>
 #include <string>
+#include <vector>
 
 #include "stable-diffusion.h"
 
@@ -17,6 +18,8 @@ bool file_exists(const std::string& filename);
 bool is_directory(const std::string& path);
 std::string get_full_path(const std::string& dir, const std::string& filename);
 
+std::vector<std::string> get_files_from_dir(const std::string &dir);
+
 std::u32string utf8_to_utf32(const std::string& utf8_str);
 std::string utf32_to_utf8(const std::u32string& utf32_str);
 std::u32string unicode_value_to_utf32(int unicode_value);