Send a callback every step, but the preview image only at the defined interval. -1 preview interval to disable image previews

cmdr2 · cmdr2 · commit ae4586804541 · 2024-11-19T16:20:02.000+05:30
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -1026,32 +1026,36 @@ class StableDiffusionGGML {
         result_cb(number, image_data, result_cb_data);
     }
 
-    sd_result_step_cb_t result_step_cb = nullptr;
-    int result_step_cb_interval        = 1;
-    void* result_step_cb_data          = nullptr;
+    sd_result_step_cb_t result_step_cb   = nullptr;
+    int result_step_cb_preview_interval  = -1;
+    void* result_step_cb_data            = nullptr;
 
     void send_result_step_callback(ggml_context* task_work_ctx, ggml_tensor* x, size_t number, size_t step) {
-        if (result_step_cb == nullptr || step % result_step_cb_interval != 0) {
+        if (result_step_cb == nullptr) {
             return;
         }
 
-        struct ggml_init_params params {};
-        params.mem_size   = ggml_get_mem_size(task_work_ctx);
-        params.mem_buffer = nullptr;
-        params.no_alloc   = false;
+        if (result_step_cb_preview_interval != -1 && step % result_step_cb_preview_interval == 0) {
+            struct ggml_init_params params {};
+            params.mem_size   = ggml_get_mem_size(task_work_ctx);
+            params.mem_buffer = nullptr;
+            params.no_alloc   = false;
 
-        struct ggml_context* work_ctx = ggml_init(params);
-        if (!work_ctx) {
-            return;
-        }
+            struct ggml_context* work_ctx = ggml_init(params);
+            if (!work_ctx) {
+                return;
+            }
 
-        struct ggml_tensor* result = ggml_dup_tensor(work_ctx, x);
-        copy_ggml_tensor(result, x);
+            struct ggml_tensor* result = ggml_dup_tensor(work_ctx, x);
+            copy_ggml_tensor(result, x);
 
-        struct ggml_tensor* img = decode_first_stage(work_ctx, result);
-        result_step_cb(number, step, sd_tensor_to_image(img), result_step_cb_data);
+            struct ggml_tensor* img = decode_first_stage(work_ctx, result);
+            result_step_cb(number, step, sd_tensor_to_image(img), result_step_cb_data);
 
-        ggml_free(work_ctx);
+            ggml_free(work_ctx);
+        } else {
+            result_step_cb(number, step, nullptr, result_step_cb_data);
+        }
     }
 };
 
@@ -1140,13 +1144,13 @@ void free_sd_ctx(sd_ctx_t* sd_ctx) {
 }
 
 void sd_ctx_set_result_callback(sd_ctx_t* sd_ctx, sd_result_cb_t cb, void* data) {
-    sd_ctx->sd->result_cb      = cb;
+    sd_ctx->sd->result_cb = cb;
     sd_ctx->sd->result_cb_data = data;
 }
 
-void sd_ctx_set_result_step_callback(sd_ctx_t* sd_ctx, sd_result_step_cb_t cb, int cb_interval, void* data) {
-    sd_ctx->sd->result_step_cb      = cb;
-    sd_ctx->sd->result_step_cb_interval = cb_interval;
+void sd_ctx_set_result_step_callback(sd_ctx_t* sd_ctx, sd_result_step_cb_t cb, int cb_preview_interval, void* data) {
+    sd_ctx->sd->result_step_cb = cb;
+    sd_ctx->sd->result_step_cb_preview_interval = cb_preview_interval;
     sd_ctx->sd->result_step_cb_data = data;
 }
 
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -148,7 +148,7 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
 
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
 SD_API void sd_ctx_set_result_callback(sd_ctx_t* sd_ctx, sd_result_cb_t cb, void* data);
-SD_API void sd_ctx_set_result_step_callback(sd_ctx_t* sd_ctx, sd_result_step_cb_t cb, int cb_interval, void* data);
+SD_API void sd_ctx_set_result_step_callback(sd_ctx_t* sd_ctx, sd_result_step_cb_t cb, int cb_preview_interval, void* data);
 
 SD_API sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                            const char* prompt,
