diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 3198907c..6c9218fe 100644 --- a/examples/cli/main.cpp +++ b/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 input_id_images; + if (params.stacked_id_embeddings_path.size() > 0 && params.input_id_images_path.size() > 0) { + std::vector 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(), diff --git a/util.cpp b/util.cpp index c5f3f861..50ff4bf5 100644 --- a/util.cpp +++ b/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 get_files_from_dir(const std::string &dir){ + + std::vector 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 diff --git a/util.h b/util.h index c1b035f1..14345085 100644 --- a/util.h +++ b/util.h @@ -3,6 +3,7 @@ #include #include +#include #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 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);