feat(tx): control freeing compute buffer

thxCode · thxCode · commit 1b3faeb43a90 · 2024-12-03T19:25:46.000+08:00
Signed-off-by: thxCode &lt;thxcode0824@gmail.com&gt;
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -840,6 +840,7 @@ int main(int argc, const char* argv[]) {
                                   vae_decode_only,
                                   params.vae_tiling,
                                   true,
+                                  true,
                                   params.n_threads,
                                   params.wtype,
                                   params.rng_type,
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -139,8 +139,9 @@ class StableDiffusionGGML {
     ggml_backend_t vae_backend         = NULL;
 
     SDVersion version;
-    bool vae_decode_only         = false;
-    bool free_params_immediately = false;
+    bool vae_decode_only          = false;
+    bool free_params_immediately  = false;
+    bool free_compute_immediately = true;
 
     rng_type_t rng_type = STD_DEFAULT_RNG;
     int n_threads       = -1;
@@ -174,11 +175,13 @@ class StableDiffusionGGML {
     StableDiffusionGGML(int n_threads,
                         bool vae_decode_only,
                         bool free_params_immediately,
+                        bool free_compute_immediately,
                         std::string lora_model_dir,
                         rng_type_t rng_type)
         : n_threads(n_threads),
           vae_decode_only(vae_decode_only),
           free_params_immediately(free_params_immediately),
+          free_compute_immediately(free_compute_immediately),
           lora_model_dir(lora_model_dir),
           rng_type(rng_type) {
     }
@@ -1094,7 +1097,9 @@ class StableDiffusionGGML {
             control_net->free_control_ctx();
             control_net->free_compute_buffer();
         }
-        diffusion_model->free_compute_buffer();
+        if (free_compute_immediately) {
+            diffusion_model->free_compute_buffer();
+        }
         return x;
     }
 
@@ -1164,7 +1169,9 @@ class StableDiffusionGGML {
             } else {
                 first_stage_model->compute(n_threads, x, decode, &result);
             }
-            first_stage_model->free_compute_buffer();
+            if (free_compute_immediately) {
+                first_stage_model->free_compute_buffer();
+            }
             if (decode) {
                 ggml_tensor_scale_output(result);
             }
@@ -1178,7 +1185,9 @@ class StableDiffusionGGML {
             } else {
                 tae_first_stage->compute(n_threads, x, decode, &result);
             }
-            tae_first_stage->free_compute_buffer();
+            if (free_compute_immediately) {
+                tae_first_stage->free_compute_buffer();
+            }
         }
 
         int64_t t1 = ggml_time_ms();
@@ -1218,6 +1227,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                      bool vae_decode_only,
                      bool vae_tiling,
                      bool free_params_immediately,
+                     bool free_compute_immediately,
                      int n_threads,
                      enum sd_type_t wtype,
                      enum rng_type_t rng_type,
@@ -1246,6 +1256,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
     sd_ctx->sd = new StableDiffusionGGML(n_threads,
                                          vae_decode_only,
                                          free_params_immediately,
+                                         free_compute_immediately,
                                          lora_model_dir,
                                          rng_type);
     if (sd_ctx->sd == NULL) {
@@ -2541,6 +2552,9 @@ bool sd_sampling_stream_sample(sd_ctx_t* sd_ctx, sd_sampling_stream_t* stream) {
         return true;
     }
     stream->x = sd_ctx->sd->denoiser->inverse_noise_scaling(stream->sigmas[stream->sigmas.size() - 1], stream->x);
+    if (sd_ctx->sd->free_compute_immediately) {
+        sd_ctx->sd->diffusion_model->free_compute_buffer();
+    }
 
     size_t sampling_end = ggml_time_ms();
     LOG_INFO("sampling completed, taking %.2fs", (sampling_end - stream->sampling_start) * 1.0f / 1000);
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -149,6 +149,7 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
                             bool vae_decode_only,
                             bool vae_tiling,
                             bool free_params_immediately,
+                            bool free_compute_immediately,
                             int n_threads,
                             enum sd_type_t wtype,
                             enum rng_type_t rng_type,
