refactor: support main gpu

thxCode · thxCode · commit c7988622ddc9 · 2024-11-11T10:52:40.000+08:00
Signed-off-by: thxCode &lt;thxcode0824@gmail.com&gt;
diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp
@@ -829,7 +829,7 @@ int main(int argc, const char* argv[]) {
 
     if (results == NULL) {
         printf("generate failed\n");
-        free_sd_ctx(sd_ctx);
+        sd_ctx_free(sd_ctx);
         return 1;
     }
 
@@ -875,7 +875,7 @@ int main(int argc, const char* argv[]) {
         results[i].data = NULL;
     }
     free(results);
-    free_sd_ctx(sd_ctx);
+    sd_ctx_free(sd_ctx);
     free(control_image_buffer);
     free(input_image_buffer);
 
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -193,23 +193,24 @@ class StableDiffusionGGML {
                         schedule_t schedule,
                         bool clip_on_cpu,
                         bool control_net_cpu,
-                        bool vae_on_cpu) {
+                        bool vae_on_cpu,
+                        int main_gpu) {
         use_tiny_autoencoder = taesd_path.size() > 0;
 #ifdef SD_USE_CUBLAS
         LOG_DEBUG("Using CUDA backend");
-        backend = ggml_backend_cuda_init(0);
+        backend = ggml_backend_cuda_init(main_gpu);
 #endif
 #ifdef SD_USE_METAL
         LOG_DEBUG("Using Metal backend");
         backend = ggml_backend_metal_init();
 #endif
 #ifdef SD_USE_CANN
         LOG_DEBUG("Using CANN backend");
-        backend = ggml_backend_cann_init(0);
+        backend = ggml_backend_cann_init(main_gpu);
 #endif
 #ifdef SD_USE_SYCL
         LOG_DEBUG("Using SYCL backend");
-        backend = ggml_backend_sycl_init(0);
+        backend = ggml_backend_sycl_init(main_gpu);
 #endif
 
         if (!backend) {
@@ -1058,7 +1059,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                      enum schedule_t s,
                      bool keep_clip_on_cpu,
                      bool keep_control_net_cpu,
-                     bool keep_vae_on_cpu) {
+                     bool keep_vae_on_cpu,
+                     int main_gpu) {
     sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
     if (sd_ctx == NULL) {
         return NULL;
@@ -1080,9 +1082,6 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                                          free_params_immediately,
                                          lora_model_dir,
                                          rng_type);
-    if (sd_ctx->sd == NULL) {
-        return NULL;
-    }
 
     if (!sd_ctx->sd->load_from_file(model_path,
                                     clip_l_path,
@@ -1099,7 +1098,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
                                     s,
                                     keep_clip_on_cpu,
                                     keep_control_net_cpu,
-                                    keep_vae_on_cpu)) {
+                                    keep_vae_on_cpu,
+                                    main_gpu)) {
         delete sd_ctx->sd;
         sd_ctx->sd = NULL;
         free(sd_ctx);
@@ -1108,7 +1108,7 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
     return sd_ctx;
 }
 
-void free_sd_ctx(sd_ctx_t* sd_ctx) {
+void sd_ctx_free(sd_ctx_t* sd_ctx) {
     if (sd_ctx->sd != NULL) {
         delete sd_ctx->sd;
         sd_ctx->sd = NULL;
diff --git a/stable-diffusion.h b/stable-diffusion.h
@@ -139,9 +139,10 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
                             enum schedule_t s,
                             bool keep_clip_on_cpu,
                             bool keep_control_net_cpu,
-                            bool keep_vae_on_cpu);
+                            bool keep_vae_on_cpu,
+                            int main_gpu = 0);
 
-SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
+SD_API void sd_ctx_free(sd_ctx_t* sd_ctx);
 
 SD_API sd_image_t* txt2img(sd_ctx_t* sd_ctx,
                            const char* prompt,

Original file line number	Diff line number	Diff line change
`@@ -829,7 +829,7 @@ int main(int argc, const char* argv[]) {`
`829`	`829`
`830`	`830`	`if (results == NULL) {`
`831`	`831`	`printf("generate failed\n");`
`832`		`- free_sd_ctx(sd_ctx);`
	`832`	`+ sd_ctx_free(sd_ctx);`
`833`	`833`	`return 1;`
`834`	`834`	`}`
`835`	`835`
`@@ -875,7 +875,7 @@ int main(int argc, const char* argv[]) {`
`875`	`875`	`results[i].data = NULL;`
`876`	`876`	`}`
`877`	`877`	`free(results);`
`878`		`- free_sd_ctx(sd_ctx);`
	`878`	`+ sd_ctx_free(sd_ctx);`
`879`	`879`	`free(control_image_buffer);`
`880`	`880`	`free(input_image_buffer);`
`881`	`881`