From 8053aefdf7c5940290cdf50182d44113aaf050f0 Mon Sep 17 00:00:00 2001
From: bssrdf <bssrdf@gmail.com>
Date: Wed, 21 Feb 2024 12:44:04 -0500
Subject: [PATCH] added an option to offload vae decoder to CPU for mem-limited
 gpus

---
 examples/cli/main.cpp |  7 ++++++-
 stable-diffusion.cpp  | 19 +++++++++++++++----
 stable-diffusion.h    |  3 ++-
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
index f68c2959..56e62722 100644
--- a/examples/cli/main.cpp
+++ b/examples/cli/main.cpp
@@ -90,6 +90,7 @@ struct SDParams {
     bool verbose                  = false;
     bool vae_tiling               = false;
     bool control_net_cpu          = false;
+    bool vae_on_cpu               = false;
     bool canny_preprocess         = false;
 };
 
@@ -111,6 +112,7 @@ void print_params(SDParams params) {
     printf("    init_img:          %s\n", params.input_path.c_str());
     printf("    control_image:     %s\n", params.control_image_path.c_str());
     printf("    controlnet cpu:    %s\n", params.control_net_cpu ? "true" : "false");
+    printf("    vae decoder cpu:   %s\n", params.vae_on_cpu ? "true" : "false");
     printf("    strength(control): %.2f\n", params.control_strength);
     printf("    prompt:            %s\n", params.prompt.c_str());
     printf("    negative_prompt:   %s\n", params.negative_prompt.c_str());
@@ -365,6 +367,8 @@ void parse_args(int argc, const char** argv, SDParams& params) {
             params.vae_tiling = true;
         } else if (arg == "--control-net-cpu") {
             params.control_net_cpu = true;
+        } else if (arg == "--vae-on-cpu") {
+            params.vae_on_cpu = true; // will slow down latent decoding but necessary for low MEM GPUs
         } else if (arg == "--canny") {
             params.canny_preprocess = true;
         } else if (arg == "-b" || arg == "--batch-count") {
@@ -608,7 +612,8 @@ int main(int argc, const char* argv[]) {
                                   params.wtype,
                                   params.rng_type,
                                   params.schedule,
-                                  params.control_net_cpu);
+                                  params.control_net_cpu,
+                                  params.vae_on_cpu);
 
     if (sd_ctx == NULL) {
         printf("new_sd_ctx_t failed\n");
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
index 0eef52d7..efad98c6 100644
--- a/stable-diffusion.cpp
+++ b/stable-diffusion.cpp
@@ -126,7 +126,8 @@ public:
                         bool vae_tiling_,
                         ggml_type wtype,
                         schedule_t schedule,
-                        bool control_net_cpu) {
+                        bool control_net_cpu,
+                        bool vae_on_cpu) {
         use_tiny_autoencoder = taesd_path.size() > 0;
 #ifdef SD_USE_CUBLAS
         LOG_DEBUG("Using CUDA backend");
@@ -238,10 +239,18 @@ public:
             vae_type = GGML_TYPE_F32;  // avoid nan, not work...
         }
 
-        if (!use_tiny_autoencoder && !first_stage_model.alloc_params_buffer(backend, vae_type)) {
+        ggml_backend_t vae_backend = NULL;
+        if (vae_on_cpu && !ggml_backend_is_cpu(backend)) {
+            LOG_INFO("VAE Autoencoder: Using CPU backend");
+            vae_backend = ggml_backend_cpu_init();
+        }else{
+            vae_backend = backend;
+        }
+        if (!use_tiny_autoencoder && !first_stage_model.alloc_params_buffer(vae_backend, vae_type)) {
             return false;
         }
 
+
         LOG_DEBUG("preparing memory for the weights");
         // prepare memory for the weights
         {
@@ -1384,7 +1393,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                      enum sd_type_t wtype,
                      enum rng_type_t rng_type,
                      enum schedule_t s,
-                     bool keep_control_net_cpu) {
+                     bool keep_control_net_cpu,
+                     bool keep_vae_on_cpu) {
     sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
     if (sd_ctx == NULL) {
         return NULL;
@@ -1415,7 +1425,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                                     vae_tiling,
                                     (ggml_type)wtype,
                                     s,
-                                    keep_control_net_cpu)) {
+                                    keep_control_net_cpu,
+                                    keep_vae_on_cpu)) {
         delete sd_ctx->sd;
         sd_ctx->sd = NULL;
         free(sd_ctx);
diff --git a/stable-diffusion.h b/stable-diffusion.h
index 4a7c518a..8a69a6a0 100644
--- a/stable-diffusion.h
+++ b/stable-diffusion.h
@@ -116,7 +116,8 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
                             enum sd_type_t wtype,
                             enum rng_type_t rng_type,
                             enum schedule_t s,
-                            bool keep_control_net_cpu);
+                            bool keep_control_net_cpu,
+                            bool keep_vae_on_cpu);
 
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);