feat: add configurable image input preprocessing (#2028)

This commit is contained in:
leejet
2026-09-23 02:12:49 +08:00
committed by GitHub
parent 2bb72947cb
commit 28b454bda1
19 changed files with 898 additions and 227 deletions
+11 -51
View File
@@ -41,7 +41,6 @@ struct SDCliParams {
std::string metadata_format = "text";
sd_log_level_t log_level = SD_LOG_INFO;
bool canny_preprocess = false;
bool convert_name = false;
preview_t preview_method = PREVIEW_NONE;
@@ -107,10 +106,6 @@ struct SDCliParams {
};
options.bool_options = {
{"",
"--canny",
"apply canny preprocessor (edge detection)",
true, &canny_preprocess},
{"",
"--convert-name",
"convert tensor name (for convert mode)",
@@ -268,7 +263,6 @@ struct SDCliParams {
<< " metadata_format: \"" << metadata_format << "\",\n"
<< " log_level: " << log_level_name(log_level) << ",\n"
<< " color: " << (color ? "true" : "false") << ",\n"
<< " canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
<< " convert_name: " << (convert_name ? "true" : "false") << ",\n"
<< " preview_method: " << previews_str[preview_method] << ",\n"
<< " preview_interval: " << preview_interval << ",\n"
@@ -328,9 +322,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
bool load_images_from_dir(const std::string dir,
std::vector<SDImageOwner>& images,
int expected_width = 0,
int expected_height = 0,
int max_image_num = 0) {
int max_image_num = 0) {
if (!fs::exists(dir) || !fs::is_directory(dir)) {
LOG_ERROR("'%s' is not a valid directory\n", dir.c_str());
return false;
@@ -358,7 +350,7 @@ bool load_images_from_dir(const std::string dir,
int width = 0;
int height = 0;
int loaded_channel = 0;
uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, expected_width, expected_height);
uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, 0, 0);
if (image_buffer == nullptr) {
LOG_ERROR("load image from '%s' failed", path.c_str());
return false;
@@ -654,8 +646,8 @@ int main(int argc, const char* argv[]) {
SDContextParams ctx_params;
SDGenerationParams gen_params;
parse_args(argc, argv, cli_params, ctx_params, gen_params);
sd_set_log_callback(sd_log_cb, (void*)&cli_params);
parse_args(argc, argv, cli_params, ctx_params, gen_params);
if (cli_params.mode == METADATA) {
MetadataReadOptions options;
@@ -751,16 +743,8 @@ int main(int argc, const char* argv[]) {
auto load_image_and_update_size = [&](const std::string& path,
SDImageOwner& image,
bool resize_image = true,
int expected_channel = 3) -> bool {
int expected_width = 0;
int expected_height = 0;
if (resize_image && gen_params.width_and_height_are_set()) {
expected_width = gen_params.width;
expected_height = gen_params.height;
}
if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) {
if (!load_sd_image_from_file(image.put(), path.c_str(), 0, 0, expected_channel)) {
LOG_ERROR("load image from '%s' failed", path.c_str());
return false;
}
@@ -783,7 +767,7 @@ int main(int argc, const char* argv[]) {
if (gen_params.init_image_path.size() > 0) {
const bool native_init = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, true, native_init ? 0 : 3)) {
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, native_init ? 0 : 3)) {
return 1;
}
}
@@ -798,7 +782,7 @@ int main(int argc, const char* argv[]) {
gen_params.ref_images.clear();
for (auto& path : gen_params.ref_image_paths) {
SDImageOwner ref_image({0, 0, 0, nullptr});
if (!load_image_and_update_size(path, ref_image, false, 0)) {
if (!load_image_and_update_size(path, ref_image, 0)) {
return 1;
}
gen_params.ref_images.push_back(std::move(ref_image));
@@ -839,41 +823,22 @@ int main(int argc, const char* argv[]) {
if (gen_params.mask_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.mask_image.put(),
gen_params.mask_image_path.c_str(),
gen_params.get_resolved_width(),
gen_params.get_resolved_height(),
0,
0,
1)) {
LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
return 1;
}
} else {
sd_image_t generated_mask = {0, 0, 1, nullptr};
generated_mask.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
if (generated_mask.data == nullptr) {
LOG_ERROR("malloc mask image failed");
return 1;
}
generated_mask.width = gen_params.get_resolved_width();
generated_mask.height = gen_params.get_resolved_height();
memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
gen_params.mask_image.reset(generated_mask);
}
if (gen_params.control_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.control_image.put(),
gen_params.control_image_path.c_str(),
gen_params.get_resolved_width(),
gen_params.get_resolved_height())) {
0,
0)) {
LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
return 1;
}
if (cli_params.canny_preprocess) { // apply preprocessor
preprocess_canny(gen_params.control_image.get(),
0.08f,
0.08f,
0.8f,
1.0f,
false);
}
}
if (gen_params.ip_adapter_image_path.size() > 0) {
@@ -890,8 +855,6 @@ int main(int argc, const char* argv[]) {
gen_params.control_frames.clear();
if (!load_images_from_dir(gen_params.control_video_path,
gen_params.control_frames,
gen_params.get_resolved_width(),
gen_params.get_resolved_height(),
gen_params.video_frames)) {
return 1;
}
@@ -900,10 +863,7 @@ int main(int argc, const char* argv[]) {
if (!gen_params.pm_id_images_dir.empty()) {
gen_params.pm_id_images.clear();
if (!load_images_from_dir(gen_params.pm_id_images_dir,
gen_params.pm_id_images,
0,
0,
0)) {
gen_params.pm_id_images)) {
return 1;
}
}