Latent preview support for img2img and img2vid

stduhpf · thxCode · commit 4f62a9cb703b · 2024-12-13T10:09:02.000+08:00
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -968,7 +968,8 @@ int main(int argc, const char* argv[]) {
                               params.sample_method,
                               params.sample_steps,
                               params.strength,
-                              params.seed);
+                              params.seed,
+                              (step_callback_t)step_callback);
             if (results == NULL) {
                 printf("generate failed\n");
                 free_sd_ctx(sd_ctx);
@@ -1014,7 +1015,8 @@ int main(int argc, const char* argv[]) {
                               params.skip_layers.size(),
                               params.slg_scale,
                               params.skip_layer_start,
-                              params.skip_layer_end);
+                              params.skip_layer_end,
+                              (step_callback_t)step_callback);
         }
     }
 
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -1320,10 +1320,10 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx,
                            float style_ratio,
                            bool normalize_input,
                            std::string input_id_images_path,
-                           std::vector<int> skip_layers = {},
-                           float slg_scale              = 0,
-                           float skip_layer_start       = 0.01,
-                           float skip_layer_end         = 0.2,
+                           std::vector<int> skip_layers                                    = {},
+                           float slg_scale                                                 = 0,
+                           float skip_layer_start                                          = 0.01,
+                           float skip_layer_end                                            = 0.2,
                            std::function<void(int, ggml_tensor*, SDVersion)> step_callback = nullptr) {
     if (seed < 0) {
         // Generally, when using the provided command line, the seed is always >0.
@@ -1624,7 +1624,7 @@ sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                     float slg_scale          = 0,
                     float skip_layer_start   = 0.01,
                     float skip_layer_end     = 0.2,
-                    step_callback_t step_callback) {
+                    step_callback_t step_callback = NULL) {
     std::vector<int> skip_layers_vec(skip_layers, skip_layers + skip_layers_count);
     LOG_DEBUG("txt2img %dx%d", width, height);
     if (sd_ctx == NULL) {
@@ -1730,7 +1730,8 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
                     size_t skip_layers_count = 0,
                     float slg_scale          = 0,
                     float skip_layer_start   = 0.01,
-                    float skip_layer_end     = 0.2) {
+                    float skip_layer_end     = 0.2,
+                    step_callback_t step_callback = NULL) {
     std::vector<int> skip_layers_vec(skip_layers, skip_layers + skip_layers_count);
     LOG_DEBUG("img2img %dx%d", width, height);
     if (sd_ctx == NULL) {
@@ -1805,7 +1806,8 @@ sd_image_t* img2img(sd_ctx_t* sd_ctx,
                                                skip_layers_vec,
                                                slg_scale,
                                                skip_layer_start,
-                                               skip_layer_end);
+                                               skip_layer_end,
+                                               step_callback);
 
     size_t t2 = ggml_time_ms();
 
@@ -1827,7 +1829,8 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
                            enum sample_method_t sample_method,
                            int sample_steps,
                            float strength,
-                           int64_t seed) {
+                           int64_t seed,
+                           step_callback_t step_callback = NULL) {
     if (sd_ctx == NULL) {
         return NULL;
     }
@@ -1903,7 +1906,10 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
                                                  sigmas,
                                                  -1,
                                                  SDCondition(NULL, NULL, NULL),
-                                                 rng);
+                                                 rng,
+                                                 {},
+                                                 0, 0, 0,
+                                                 step_callback);
 
     int64_t t2 = ggml_time_ms();
     LOG_INFO("sampling completed, taking %.2fs", (t2 - t1) * 1.0f / 1000);
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -210,7 +210,8 @@ SD_API sd_image_t* img2img(sd_ctx_t* sd_ctx,
                            size_t skip_layers_count,
                            float slg_scale,
                            float skip_layer_start,
-                           float skip_layer_end);
+                           float skip_layer_end,
+                           step_callback_t step_callback);
 
 SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
                            sd_image_t init_image,
@@ -225,7 +226,8 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
                            enum sample_method_t sample_method,
                            int sample_steps,
                            float strength,
-                           int64_t seed);
+                           int64_t seed,
+                           step_callback_t step_callback);
 
 typedef struct sd_lora_adapter_container_t {
     const char* path;
