Refactor preview to match the other callbacks

stduhpf · stduhpf · commit 1bef24dc5c42 · 2025-07-15T19:33:42.000+02:00
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -117,10 +117,10 @@ struct SDParams {
     bool chroma_use_t5_mask  = false;
     int chroma_t5_mask_pad   = 1;
 
-    sd_preview_policy_t preview_method = SD_PREVIEW_NONE;
-    int preview_interval               = 1;
-    std::string preview_path           = "preview.png";
-    bool taesd_preview                 = false;
+    sd_preview_t preview_method = SD_PREVIEW_NONE;
+    int preview_interval        = 1;
+    std::string preview_path    = "preview.png";
+    bool taesd_preview          = false;
 };
 
 void print_params(SDParams params) {
@@ -595,7 +595,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 preview);
             return -1;
         }
-        params.preview_method = (sd_preview_policy_t)preview_method;
+        params.preview_method = (sd_preview_t)preview_method;
         return 1;
     };
 
@@ -796,6 +796,7 @@ int main(int argc, const char* argv[]) {
                                             }};
 
     sd_set_log_callback(sd_log_cb, (void*)&params);
+    sd_set_preview_callback((sd_preview_cb_t)step_callback, params.preview_method, params.preview_interval);
 
     if (params.verbose) {
         print_params(params);
@@ -1018,7 +1019,7 @@ int main(int argc, const char* argv[]) {
             params.input_id_images_path.c_str(),
         };
 
-        results              = generate_image(sd_ctx, &img_gen_params, params.preview_method, params.preview_interval,(step_callback_t)step_callback);
+        results              = generate_image(sd_ctx, &img_gen_params);
         expected_num_results = params.batch_count;
     } else if (params.mode == VID_GEN) {
         sd_vid_gen_params_t vid_gen_params = {
@@ -1036,7 +1037,7 @@ int main(int argc, const char* argv[]) {
             params.augmentation_level,
         };
 
-        results              = generate_video(sd_ctx, &vid_gen_params, (step_callback_t)step_callback);
+        results              = generate_video(sd_ctx, &vid_gen_params);
         expected_num_results = params.video_frames;
     }
 
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -858,7 +858,7 @@ class StableDiffusionGGML {
                        int step,
                        struct ggml_tensor* latents,
                        enum SDVersion version,
-                       sd_preview_policy_t preview_mode,
+                       sd_preview_t preview_mode,
                        ggml_tensor* result,
                        std::function<void(int, sd_image_t)> step_callback) {
         const uint32_t channel = 3;
@@ -969,10 +969,7 @@ class StableDiffusionGGML {
                         int start_merge_step,
                         SDCondition id_cond,
                         std::vector<ggml_tensor*> ref_latents = {},
-                        ggml_tensor* denoise_mask             = nullptr,
-                        sd_preview_policy_t preview_mode                   = SD_PREVIEW_NONE,
-                        int preview_interval                               = 1,
-                        std::function<void(int, sd_image_t)> step_callback = nullptr) {
+                        ggml_tensor* denoise_mask             = nullptr) {
         std::vector<int> skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count);
 
         float cfg_scale     = guidance.txt_cfg;
@@ -1034,7 +1031,8 @@ class StableDiffusionGGML {
         struct ggml_tensor* denoised = ggml_dup_tensor(work_ctx, x);
 
         struct ggml_tensor* preview_tensor = NULL;
-        if (preview_mode != SD_PREVIEW_NONE && preview_mode != SD_PREVIEW_PROJ) {
+        auto sd_preview_mode = sd_get_preview_mode();
+        if (sd_preview_mode != SD_PREVIEW_NONE && sd_preview_mode != SD_PREVIEW_PROJ) {
             preview_tensor = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32,
                                                 (denoised->ne[0] * 8),
                                                 (denoised->ne[1] * 8),
@@ -1216,10 +1214,11 @@ class StableDiffusionGGML {
                 pretty_progress(step, (int)steps, (t1 - t0) / 1000000.f);
                 // LOG_INFO("step %d sampling completed taking %.2fs", step, (t1 - t0) * 1.0f / 1000000);
             }
-
-            if (step_callback != nullptr) {
-                if (step % preview_interval == 0) {
-                    preview_image(work_ctx, step, denoised, version, preview_mode, preview_tensor, step_callback);
+            auto sd_preview_cb = sd_get_preview_callback();
+            auto sd_preview_mode = sd_get_preview_mode();
+            if (sd_preview_cb != NULL) {
+                if (step % sd_get_preview_interval() == 0) {
+                    preview_image(work_ctx, step, denoised, version, sd_preview_mode, preview_tensor, sd_preview_cb);
                 }
             }
             return denoised;
@@ -1671,10 +1670,7 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
                                     std::string input_id_images_path,
                                     std::vector<ggml_tensor*> ref_latents,
                                     ggml_tensor* concat_latent                                      = NULL,
-                                    ggml_tensor* denoise_mask                                       = NULL,
-                                    sd_preview_policy_t preview_mode                   = SD_PREVIEW_NONE,
-                                    int preview_interval                               = 1,
-                                    std::function<void(int, sd_image_t)> step_callback = nullptr) {
+                                    ggml_tensor* denoise_mask                                       = NULL) {
     if (seed < 0) {
         // Generally, when using the provided command line, the seed is always >0.
         // However, to prevent potential issues if 'stable-diffusion.cpp' is invoked as a library
@@ -1943,10 +1939,7 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
                                                      start_merge_step,
                                                      id_cond,
                                                      ref_latents,
-                                                     denoise_mask,
-                                                     preview_mode,
-                                                     preview_interval,
-                                                     step_callback);
+                                                     denoise_mask);
 
         // struct ggml_tensor* x_0 = load_tensor_from_file(ctx, "samples_ddim.bin");
         // print_ggml_tensor(x_0);
@@ -2020,7 +2013,7 @@ ggml_tensor* generate_init_latent(sd_ctx_t* sd_ctx,
     return init_latent;
 }
 
-sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, sd_preview_policy_t preview_mode, int preview_interval, step_callback_t step_callback) {
+sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
     int width  = sd_img_gen_params->width;
     int height = sd_img_gen_params->height;
     LOG_DEBUG("generate_image %dx%d", width, height);
@@ -2039,7 +2032,8 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
     if (sd_ctx->sd->stacked_id) {
         params.mem_size += static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
     }
-    if (preview_mode != SD_PREVIEW_NONE && preview_mode != SD_PREVIEW_PROJ) {
+    auto sd_preview_mode = sd_get_preview_mode();
+    if (sd_preview_mode != SD_PREVIEW_NONE && sd_preview_mode != SD_PREVIEW_PROJ) {
         params.mem_size *= 2;
     }
     params.mem_size += width * height * 3 * sizeof(float) * 3;
@@ -2223,10 +2217,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
                                                         sd_img_gen_params->input_id_images_path,
                                                         ref_latents,
                                                         concat_latent,
-                                                        denoise_mask,
-                                                        preview_mode,
-                                                        preview_interval,
-                                                        step_callback);
+                                                        denoise_mask);
 
     size_t t2 = ggml_time_ms();
 
@@ -2235,7 +2226,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g
     return result_images;
 }
 
-SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, step_callback_t step_callback) {
+SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params) {
     if (sd_ctx == NULL || sd_vid_gen_params == NULL) {
         return NULL;
     }
@@ -2319,8 +2310,6 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                  -1,
                                                  SDCondition(NULL, NULL, NULL),
                                                  {},
-                                                 NULL,
-                                                 (sd_preview_policy_t)0, 1,
                                                  NULL);
 
     int64_t t2 = ggml_time_ms();
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -111,7 +111,7 @@ enum sd_log_level_t {
     SD_LOG_ERROR
 };
 
-enum sd_preview_policy_t {
+enum sd_preview_t {
     SD_PREVIEW_NONE,
     SD_PREVIEW_PROJ,
     SD_PREVIEW_TAE,
@@ -214,11 +214,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, sd_image_t);
 
 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 sd_progress_cb_t sd_get_progress_callback();
-SD_API void* sd_get_progress_callback_data();
+SD_API void sd_set_preview_callback(sd_preview_cb_t cb, sd_preview_t mode, int interval);
 SD_API int32_t get_num_physical_cores();
 SD_API const char* sd_get_system_info();
 
@@ -237,14 +237,12 @@ SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
 SD_API sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params);
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
 
-typedef void (*step_callback_t)(int, sd_image_t);
-
 SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params);
 SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params);
-SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, sd_preview_policy_t preview_mode, int preview_interval, step_callback_t step_callback);
+SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
 
 SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params);
-SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, step_callback_t step_callback);  // broken
+SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params);  // broken
 
 typedef struct upscaler_ctx_t upscaler_ctx_t;
 
diff --git a/util.cpp b/util.cpp
@@ -247,6 +247,10 @@ int32_t get_num_physical_cores() {
 static sd_progress_cb_t sd_progress_cb = NULL;
 void* sd_progress_cb_data              = NULL;
 
+static sd_preview_cb_t sd_preview_cb = NULL;
+sd_preview_t sd_preview_mode         = SD_PREVIEW_NONE;
+int sd_preview_interval              = 1;
+
 std::u32string utf8_to_utf32(const std::string& utf8_str) {
     std::wstring_convert<std::codecvt_utf8<char32_t>, char32_t> converter;
     return converter.from_bytes(utf8_str);
@@ -420,10 +424,27 @@ void sd_set_progress_callback(sd_progress_cb_t cb, void* data) {
     sd_progress_cb      = cb;
     sd_progress_cb_data = data;
 }
-sd_progress_cb_t sd_get_progress_callback(){
+void sd_set_preview_callback(sd_preview_cb_t cb, sd_preview_t mode = SD_PREVIEW_PROJ, int interval = 1) {
+    sd_preview_cb       = cb;
+    sd_preview_mode     = mode;
+    sd_preview_interval = interval;
+}
+
+sd_preview_cb_t sd_get_preview_callback() {
+    return sd_preview_cb;
+}
+
+sd_preview_t sd_get_preview_mode() {
+    return sd_preview_mode;
+}
+int sd_get_preview_interval() {
+    return sd_preview_interval;
+}
+
+sd_progress_cb_t sd_get_progress_callback() {
     return sd_progress_cb;
 }
-void* sd_get_progress_callback_data(){
+void* sd_get_progress_callback_data() {
     return sd_progress_cb_data;
 }
 const char* sd_get_system_info() {
diff --git a/util.h b/util.h
@@ -57,6 +57,13 @@ std::string trim(const std::string& s);
 
 std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::string& text);
 
+sd_progress_cb_t sd_get_progress_callback();
+void* sd_get_progress_callback_data();
+
+sd_preview_cb_t sd_get_preview_callback();
+sd_preview_t sd_get_preview_mode();
+int sd_get_preview_interval();
+
 #define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
