feat: additional --preview-interval values (#1915)

Co-authored-by: leejet <leejet714@gmail.com>

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",

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();

src/runtime/preview_interval.h

@@ -0,0 +1,45 @@
+#ifndef __SD_RUNTIME_PREVIEW_INTERVAL_H__
+#define __SD_RUNTIME_PREVIEW_INTERVAL_H__
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+namespace sd::preview {
+
+    constexpr std::uint64_t logical_sample_step(int step) {
+        return step < 0 ? static_cast<std::uint64_t>(-static_cast<std::int64_t>(step))
+                        : static_cast<std::uint64_t>(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<std::uint64_t>(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<std::uint64_t>(-static_cast<std::int64_t>(interval));
+            return logical_step == requested_step;
+        }
+        return preview_final_step && logical_step == static_cast<std::uint64_t>(total_steps);
+    }
+}  // namespace sd::preview
+
+#endif  // __SD_RUNTIME_PREVIEW_INTERVAL_H__

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 @@ public:
                                     sd_get_preview_mode()};
     }
 
-    void report_sample_progress(int step, size_t total_steps, int64_t* last_progress_us) {
+    void report_sample_progress(int step,
-        if (step > 0 || step == -(int)total_steps) {
+                                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 @@ public:
                              int audio_length,
                              float frame_rate,
                              const sd_cache_params_t* cache_params,
+                             bool preview_final_step,
                              const sd::Tensor<float>& video_positions = {}) {
         struct RunnerDoneOnExit {
             GGMLRunner* runner = nullptr;
@@ -2589,8 +2594,9 @@ public:
             }
         }
 
-        size_t steps       = sigmas.size() - 1;
+        size_t steps                = sigmas.size() - 1;
-        bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
+        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 @@ public:
             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<float> base_timesteps_vec = prepare_sample_timesteps(sigma, shifted_timestep);
             std::vector<float> timesteps_vec      = base_timesteps_vec;
             sd::Tensor<float> audio_timesteps_tensor;
@@ -2669,21 +2682,17 @@ public:
                 if (!denoise_mask.empty()) {
                     denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
                 }
-                if (sd_should_preview_denoised() && preview.callback != nullptr) {
+                if (preview_needed && sd_should_preview_denoised()) {
-                    if (step % sd_get_preview_interval() == 0) {
+                    preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
-                        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 (preview_needed && sd_should_preview_noisy()) {
-                if (step % sd_get_preview_interval() == 0) {
+                preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
-                    preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
-                }
             }
 
             sd::Tensor<float> cond_out;
@@ -2896,12 +2905,10 @@ public:
             if (!denoise_mask.empty()) {
                 denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
             }
-            if (sd_should_preview_denoised() && preview.callback != nullptr) {
+            if (preview_needed && sd_should_preview_denoised()) {
-                if (step % sd_get_preview_interval() == 0) {
+                preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false);
-                    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<float>(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<float>(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<float>(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<float>(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<float>(hires_request.fps),
                                             hires_request.cache_params,
+                                            false,
                                             hires_video_positions);
         sampling_end   = ggml_time_ms();
         if (final_latent.empty()) {