diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 953a44e9d..43719abbc 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -94,7 +94,7 @@ struct SDCliParams { options.int_options = { {"", "--preview-interval", - "interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at every step)", + "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1", &preview_interval}, {"", "--output-begin-idx", diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 6bea445ab..d5cda6d9a 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -446,6 +446,9 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* 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); +// In each sampling pass, a positive interval previews every Nth denoiser step, while a +// negative interval previews only completed logical step -interval. Zero previews the final +// completed step of the first sampling pass (base-resolution or high-noise). 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 void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data); SD_API int32_t sd_get_num_physical_cores(); diff --git a/src/runtime/preview_interval.h b/src/runtime/preview_interval.h new file mode 100644 index 000000000..aab997682 --- /dev/null +++ b/src/runtime/preview_interval.h @@ -0,0 +1,45 @@ +#ifndef __SD_RUNTIME_PREVIEW_INTERVAL_H__ +#define __SD_RUNTIME_PREVIEW_INTERVAL_H__ + +#include +#include +#include + +namespace sd::preview { + + constexpr std::uint64_t logical_sample_step(int step) { + return step < 0 ? static_cast(-static_cast(step)) + : static_cast(step); + } + + constexpr bool sample_step_is_complete(int step, + std::size_t total_steps, + bool terminal_sigma_is_zero) { + return step > 0 || + (terminal_sigma_is_zero && + step < 0 && + logical_sample_step(step) == static_cast(total_steps)); + } + + constexpr bool should_preview_sample_step(int step, + std::size_t total_steps, + bool terminal_sigma_is_zero, + int interval, + bool preview_final_step) { + if (interval > 0) { + return step % interval == 0; + } + if (!sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) { + return false; + } + + std::uint64_t logical_step = logical_sample_step(step); + if (interval < 0) { + std::uint64_t requested_step = static_cast(-static_cast(interval)); + return logical_step == requested_step; + } + return preview_final_step && logical_step == static_cast(total_steps); + } +} // namespace sd::preview + +#endif // __SD_RUNTIME_PREVIEW_INTERVAL_H__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 3ff0df980..2c1a18b12 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -61,6 +61,7 @@ #include "model/vae/wan_vae.hpp" #include "runtime/denoiser.hpp" #include "runtime/guidance.h" +#include "runtime/preview_interval.h" #include "runtime/sample-cache.h" #include "upscaler.h" @@ -2467,8 +2468,11 @@ class StableDiffusionGGML { sd_get_preview_mode()}; } - void report_sample_progress(int step, size_t total_steps, int64_t* last_progress_us) { - if (step > 0 || step == -(int)total_steps) { + void report_sample_progress(int step, + size_t total_steps, + bool terminal_sigma_is_zero, + int64_t* last_progress_us) { + if (sd::preview::sample_step_is_complete(step, total_steps, terminal_sigma_is_zero)) { int64_t now = ggml_time_us(); int showstep = std::abs(step); float step_seconds = last_progress_us != nullptr && *last_progress_us > 0 @@ -2530,6 +2534,7 @@ class StableDiffusionGGML { int audio_length, float frame_rate, const sd_cache_params_t* cache_params, + bool preview_final_step, const sd::Tensor& video_positions = {}) { struct RunnerDoneOnExit { GGMLRunner* runner = nullptr; @@ -2589,8 +2594,9 @@ class StableDiffusionGGML { } } - size_t steps = sigmas.size() - 1; - bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty(); + size_t steps = sigmas.size() - 1; + bool terminal_sigma_is_zero = sigmas.back() == 0.f; + bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty(); if (has_skiplayer && !sd_version_is_dit(version)) { has_skiplayer = false; LOG_WARN("SLG is incompatible with this model type"); @@ -2639,6 +2645,13 @@ class StableDiffusionGGML { float c_out = scaling[1]; float c_in = scaling[2]; + bool preview_needed = preview.callback != nullptr && + sd::preview::should_preview_sample_step(step, + steps, + terminal_sigma_is_zero, + sd_get_preview_interval(), + preview_final_step); + std::vector base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep); std::vector timesteps_vec = base_timesteps_vec; sd::Tensor audio_timesteps_tensor; @@ -2669,21 +2682,17 @@ class StableDiffusionGGML { if (!denoise_mask.empty()) { denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask); } - if (sd_should_preview_denoised() && preview.callback != nullptr) { - if (step % sd_get_preview_interval() == 0) { - preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false); - } + if (preview_needed && sd_should_preview_denoised()) { + preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false); } - report_sample_progress(step, steps, &last_progress_us); + report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us); sd::guidance::GuiderOutput output; output.pred = denoised; return output; } - if (sd_should_preview_noisy() && preview.callback != nullptr) { - if (step % sd_get_preview_interval() == 0) { - preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true); - } + if (preview_needed && sd_should_preview_noisy()) { + preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true); } sd::Tensor cond_out; @@ -2896,12 +2905,10 @@ class StableDiffusionGGML { if (!denoise_mask.empty()) { denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask); } - if (sd_should_preview_denoised() && preview.callback != nullptr) { - if (step % sd_get_preview_interval() == 0) { - preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false); - } + if (preview_needed && sd_should_preview_denoised()) { + preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false); } - report_sample_progress(step, steps, &last_progress_us); + report_sample_progress(step, steps, terminal_sigma_is_zero, &last_progress_us); output.pred = denoised; return output; }; @@ -5719,7 +5726,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, 1.f, 0, static_cast(request.fps), - request.cache_params); + request.cache_params, + true); int64_t sampling_end = ggml_time_ms(); if (!x_0.empty()) { LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); @@ -5840,7 +5848,8 @@ SD_API bool generate_image(sd_ctx_t* sd_ctx, 1.f, 0, static_cast(request.fps), - request.cache_params); + request.cache_params, + false); int64_t hires_sample_end = ggml_time_ms(); if (!x_0.empty()) { LOG_INFO("hires sampling %d/%d completed, taking %.2fs", @@ -6974,6 +6983,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, latents.audio_length, static_cast(request.fps), request.cache_params, + true, latents.video_positions); int64_t sampling_end = ggml_time_ms(); if (x_t_sampled.empty()) { @@ -7016,6 +7026,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, latents.audio_length, static_cast(request.fps), request.cache_params, + plan.high_noise_sample_steps <= 0, latents.video_positions); int64_t sampling_end = ggml_time_ms(); @@ -7154,6 +7165,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, latents.audio_length, static_cast(hires_request.fps), hires_request.cache_params, + false, hires_video_positions); sampling_end = ggml_time_ms(); if (final_latent.empty()) {