diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 28548b4..ed3f849 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -151,6 +151,7 @@ struct SDParams { preview_t preview_method = PREVIEW_NONE; int preview_interval = 1; std::string preview_path = "preview.png"; + float preview_fps = 16; bool taesd_preview = false; bool preview_noisy = false; @@ -1638,25 +1639,22 @@ bool load_images_from_dir(const std::string dir, return true; } -std::string preview_path; -float preview_fps; - -void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy) { +void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, void* data) { (void)step; (void)is_noisy; + SDParams* params = (SDParams*)data; // is_noisy is set to true if the preview corresponds to noisy latents, false if it's denoised latents // unused in this app, it will either be always noisy or always denoised here if (frame_count == 1) { - stbi_write_png(preview_path.c_str(), image->width, image->height, image->channel, image->data, 0); + stbi_write_png(params->preview_path.c_str(), image->width, image->height, image->channel, image->data, 0); } else { - create_mjpg_avi_from_sd_images(preview_path.c_str(), image, frame_count, preview_fps); + create_mjpg_avi_from_sd_images(params->preview_path.c_str(), image, frame_count, params->preview_fps); } } int main(int argc, const char* argv[]) { SDParams params; parse_args(argc, argv, params); - preview_path = params.preview_path; if (params.video_frames > 4) { size_t last_dot_pos = params.preview_path.find_last_of("."); std::string base_path = params.preview_path; @@ -1667,12 +1665,12 @@ int main(int argc, const char* argv[]) { std::transform(file_ext.begin(), file_ext.end(), file_ext.begin(), ::tolower); } if (file_ext == ".png") { - preview_path = base_path + ".avi"; + params.preview_path = base_path + ".avi"; } } - preview_fps = params.fps; + params.preview_fps = params.fps; if (params.preview_method == PREVIEW_PROJ) - preview_fps /= 4.0f; + params.preview_fps /= 4.0f; params.sample_params.guidance.slg.layers = params.skip_layers.data(); params.sample_params.guidance.slg.layer_count = params.skip_layers.size(); @@ -1680,7 +1678,7 @@ int main(int argc, const char* argv[]) { params.high_noise_sample_params.guidance.slg.layer_count = params.high_noise_skip_layers.size(); sd_set_log_callback(sd_log_cb, (void*)¶ms); - sd_set_preview_callback((sd_preview_cb_t)step_callback, params.preview_method, params.preview_interval, !params.preview_noisy, params.preview_noisy); + sd_set_preview_callback(step_callback, params.preview_method, params.preview_interval, !params.preview_noisy, params.preview_noisy, (void*)¶ms); if (params.verbose) { print_params(params); diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index a7ba034..ef798c2 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -1307,7 +1307,8 @@ public: enum SDVersion version, preview_t preview_mode, ggml_tensor* result, - std::function step_callback, + std::function step_callback, + void* step_callback_data, bool is_noisy) { const uint32_t channel = 3; uint32_t width = latents->ne[0]; @@ -1378,7 +1379,7 @@ public: for (int i = 0; i < frames; i++) { images[i] = {width, height, channel, data + i * width * height * channel}; } - step_callback(step, frames, images, is_noisy); + step_callback(step, frames, images, is_noisy, step_callback_data); free(data); free(images); } else { @@ -1432,7 +1433,7 @@ public: images[i].data = ggml_tensor_to_sd_image(result, i, ggml_n_dims(latents) == 4); } - step_callback(step, frames, images, is_noisy); + step_callback(step, frames, images, is_noisy, step_callback_data); ggml_ext_tensor_scale_inplace(result, 0); for (int i = 0; i < frames; i++) { @@ -1581,8 +1582,9 @@ public: } auto denoise = [&](ggml_tensor* input, float sigma, int step) -> ggml_tensor* { - auto sd_preview_cb = sd_get_preview_callback(); - auto sd_preview_mode = sd_get_preview_mode(); + auto sd_preview_cb = sd_get_preview_callback(); + auto sd_preview_cb_data = sd_get_preview_callback_data(); + auto sd_preview_mode = sd_get_preview_mode(); if (step == 1 || step == -1) { pretty_progress(0, (int)steps, 0); } @@ -1651,7 +1653,7 @@ public: } if (sd_preview_cb != nullptr && sd_should_preview_noisy()) { if (step % sd_get_preview_interval() == 0) { - preview_image(work_ctx, step, noised_input, version, sd_preview_mode, preview_tensor, sd_preview_cb, true); + preview_image(work_ctx, step, noised_input, version, sd_preview_mode, preview_tensor, sd_preview_cb, sd_preview_cb_data, true); } } @@ -1799,7 +1801,7 @@ public: if (sd_preview_cb != nullptr && sd_should_preview_denoised()) { if (step % sd_get_preview_interval() == 0) { - preview_image(work_ctx, step, denoised, version, sd_preview_mode, preview_tensor, sd_preview_cb, false); + preview_image(work_ctx, step, denoised, version, sd_preview_mode, preview_tensor, sd_preview_cb, sd_preview_cb_data, false); } } diff --git a/stable-diffusion.h b/stable-diffusion.h index 4e3f8ea..0cbefd9 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -283,11 +283,11 @@ typedef struct sd_ctx_t sd_ctx_t; typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data); typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data); -typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy); +typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data); SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data); SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data); -SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy); +SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data); SD_API int32_t get_num_physical_cores(); SD_API const char* sd_get_system_info(); diff --git a/util.cpp b/util.cpp index c462166..26cfde4 100644 --- a/util.cpp +++ b/util.cpp @@ -187,6 +187,7 @@ static sd_progress_cb_t sd_progress_cb = nullptr; void* sd_progress_cb_data = nullptr; static sd_preview_cb_t sd_preview_cb = nullptr; +static void* sd_preview_cb_data = nullptr; preview_t sd_preview_mode = PREVIEW_NONE; int sd_preview_interval = 1; bool sd_preview_denoised = true; @@ -335,8 +336,9 @@ void sd_set_progress_callback(sd_progress_cb_t cb, void* data) { sd_progress_cb = cb; sd_progress_cb_data = data; } -void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode = PREVIEW_PROJ, int interval = 1, bool denoised = true, bool noisy = false) { +void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, bool denoised, bool noisy, void* data) { sd_preview_cb = cb; + sd_preview_cb_data = data; sd_preview_mode = mode; sd_preview_interval = interval; sd_preview_denoised = denoised; @@ -346,6 +348,9 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode = PREVIEW_PROJ, sd_preview_cb_t sd_get_preview_callback() { return sd_preview_cb; } +void* sd_get_preview_callback_data() { + return sd_preview_cb_data; +} preview_t sd_get_preview_mode() { return sd_preview_mode; diff --git a/util.h b/util.h index 5bd69a6..2721f29 100644 --- a/util.h +++ b/util.h @@ -58,6 +58,7 @@ sd_progress_cb_t sd_get_progress_callback(); void* sd_get_progress_callback_data(); sd_preview_cb_t sd_get_preview_callback(); +void* sd_get_preview_callback_data(); preview_t sd_get_preview_mode(); int sd_get_preview_interval(); bool sd_should_preview_denoised();