From 16c9da0ccadcb0be94ee1dad7e48dc5abf5fba81 Mon Sep 17 00:00:00 2001 From: bssrdf Date: Wed, 28 Feb 2024 16:29:32 -0500 Subject: [PATCH] added a WIN32 get_files_from_dir; turn off Photomakder if receiving no input images --- examples/cli/main.cpp | 6 ++-- stable-diffusion.cpp | 77 +++++++++++++++++++++++-------------------- util.cpp | 39 ++++++++++++++++++++++ 3 files changed, 83 insertions(+), 39 deletions(-) diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index f42c6677..45b1859e 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -706,9 +706,9 @@ int main(int argc, const char* argv[]) { } sd_image_t* input_image = NULL; input_image = new sd_image_t{(uint32_t)width, - (uint32_t)height, - 3, - input_image_buffer}; + (uint32_t)height, + 3, + input_image_buffer}; input_image = preprocess_id_image(input_image); if(input_image == NULL){ LOG_ERROR("preprocess input id image from '%s' failed", img_file.c_str()); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index 5708af44..5dda7e71 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -1679,47 +1679,52 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx, // ggml_tensor* class_tokens_mask = NULL; std::vector class_tokens_mask; if(sd_ctx->sd->stacked_id){ - sd_ctx->sd->pmid_model->style_strength = style_ratio; - int32_t w = input_id_images[0]->width; - int32_t h = input_id_images[0]->height; - int32_t channels = input_id_images[0]->channel; - int32_t num_input_images = input_id_images.size(); - init_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, w, h, channels, num_input_images); - // TODO: move these to somewhere else and be user settable - float mean[] = {0.48145466, 0.4578275, 0.40821073}; - float std[] = {0.26862954, 0.26130258, 0.27577711}; - for(int i = 0; i < num_input_images; i++) { - sd_image_t* init_image = input_id_images[i]; - if(normalize_input) - sd_mul_images_to_tensor(init_image->data, init_img, i, mean, std); - else - sd_mul_images_to_tensor(init_image->data, init_img, i, NULL, NULL); - } - t0 = ggml_time_ms(); - auto cond_tup = sd_ctx->sd->get_learned_condition_with_trigger(work_ctx, prompt, - clip_skip, width, height, num_input_images ); - prompts_embeds = std::get<0>(cond_tup); - pooled_prompts_embeds = std::get<1>(cond_tup); // [adm_in_channels, ] - class_tokens_mask = std::get<2>(cond_tup); // + if(input_id_images.size() > 0){ + sd_ctx->sd->pmid_model->style_strength = style_ratio; + int32_t w = input_id_images[0]->width; + int32_t h = input_id_images[0]->height; + int32_t channels = input_id_images[0]->channel; + int32_t num_input_images = input_id_images.size(); + init_img = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, w, h, channels, num_input_images); + // TODO: move these to somewhere else and be user settable + float mean[] = {0.48145466, 0.4578275, 0.40821073}; + float std[] = {0.26862954, 0.26130258, 0.27577711}; + for(int i = 0; i < num_input_images; i++) { + sd_image_t* init_image = input_id_images[i]; + if(normalize_input) + sd_mul_images_to_tensor(init_image->data, init_img, i, mean, std); + else + sd_mul_images_to_tensor(init_image->data, init_img, i, NULL, NULL); + } + t0 = ggml_time_ms(); + auto cond_tup = sd_ctx->sd->get_learned_condition_with_trigger(work_ctx, prompt, + clip_skip, width, height, num_input_images ); + prompts_embeds = std::get<0>(cond_tup); + pooled_prompts_embeds = std::get<1>(cond_tup); // [adm_in_channels, ] + class_tokens_mask = std::get<2>(cond_tup); // - prompts_embeds = sd_ctx->sd->id_encoder(work_ctx, init_img, prompts_embeds, class_tokens_mask); - t1 = ggml_time_ms(); - LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->pmid_model->free_params_buffer(); - } - // Encode input prompt without the trigger word for delayed conditioning - prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt); - // printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str()); - prompt = prompt_text_only; // - if(sample_steps < 50){ - LOG_INFO("sampling steps increases from %d to 50 for PHOTOMAKER", sample_steps); - sample_steps = 50; + prompts_embeds = sd_ctx->sd->id_encoder(work_ctx, init_img, prompts_embeds, class_tokens_mask); + t1 = ggml_time_ms(); + LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0); + if (sd_ctx->sd->free_params_immediately) { + sd_ctx->sd->pmid_model->free_params_buffer(); + } + // Encode input prompt without the trigger word for delayed conditioning + prompt_text_only = sd_ctx->sd->remove_trigger_from_prompt(work_ctx, prompt); + // printf("%s || %s \n", prompt.c_str(), prompt_text_only.c_str()); + prompt = prompt_text_only; // + if(sample_steps < 50){ + LOG_INFO("sampling steps increases from %d to 50 for PHOTOMAKER", sample_steps); + sample_steps = 50; + } + }else{ + LOG_WARN("Provided PhotoMaker model file, but NO input ID images"); + LOG_WARN("Turn off PhotoMaker"); + sd_ctx->sd->stacked_id = false; } } - t0 = ggml_time_ms(); auto cond_pair = sd_ctx->sd->get_learned_condition(work_ctx, prompt, clip_skip, width, height); ggml_tensor* c = cond_pair.first; diff --git a/util.cpp b/util.cpp index dd96375e..f95e4dd6 100644 --- a/util.cpp +++ b/util.cpp @@ -91,6 +91,45 @@ std::string get_full_path(const std::string& dir, const std::string& filename) { } } +std::vector get_files_from_dir(const std::string& dir) { + + std::vector files; + + WIN32_FIND_DATA findFileData; + HANDLE hFind; + + char currentDirectory[MAX_PATH]; + GetCurrentDirectory(MAX_PATH, currentDirectory); + + char directoryPath[MAX_PATH]; // this is absolute path + sprintf(directoryPath, "%s\\%s\\*", currentDirectory, dir.c_str()); + + // Find the first file in the directory + hFind = FindFirstFile(directoryPath, &findFileData); + + // Check if the directory was found + if (hFind == INVALID_HANDLE_VALUE) { + printf("Unable to find directory.\n"); + return files; + } + + // Loop through all files in the directory + do { + // Check if the found file is a regular file (not a directory) + if (!(findFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)) { + files.push_back(std::string(currentDirectory) + "\\" + dir + "\\" + std::string(findFileData.cFileName)); + } + } while (FindNextFile(hFind, &findFileData) != 0); + + // Close the handle + FindClose(hFind); + + + sort(files.begin(), files.end()); + + return files; +} + #else // Unix #include #include