Skip to content

Commit 23991ff

Browse files
thxCodethxCode
committed
feat: support main gpu
Signed-off-by: thxCode <thxcode0824@gmail.com>
1 parent 2a1a90a commit 23991ff

File tree

3 files changed

+22
-16
lines changed

3 files changed

+22
-16
lines changed

‎stable-diffusion.cpp

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,8 @@ class StableDiffusionGGML {
197197
bool clip_on_cpu,
198198
bool control_net_cpu,
199199
bool vae_on_cpu,
200-
bool diffusion_flash_attn) {
200+
bool diffusion_flash_attn,
201+
int main_gpu) {
201202
use_tiny_autoencoder = taesd_path.size() > 0;
202203
#ifdef SD_USE_CUDA
203204
#ifdef SD_USE_HIP
@@ -209,7 +210,7 @@ class StableDiffusionGGML {
209210
LOG_DEBUG("Using CUDA backend");
210211
#endif
211212
#endif
212-
backend = ggml_backend_cuda_init(0);
213+
backend = ggml_backend_cuda_init(main_gpu);
213214
if (!backend) {
214215
LOG_ERROR("CUDA backend init failed");
215216
}
@@ -223,21 +224,21 @@ class StableDiffusionGGML {
223224
#endif
224225
#ifdef SD_USE_VULKAN
225226
LOG_DEBUG("Using Vulkan backend");
226-
backend = ggml_backend_vk_init(0);
227+
backend = ggml_backend_vk_init(main_gpu);
227228
if (!backend) {
228229
LOG_ERROR("Vulkan backend init failed");
229230
}
230231
#endif
231232
#ifdef SD_USE_SYCL
232233
LOG_DEBUG("Using SYCL backend");
233-
backend = ggml_backend_sycl_init(0);
234+
backend = ggml_backend_sycl_init(main_gpu);
234235
if (!backend) {
235236
LOG_ERROR("SYCL backend init failed");
236237
}
237238
#endif
238239
#ifdef SD_USE_CANN
239240
LOG_DEBUG("Using CANN backend");
240-
backend = ggml_backend_cann_init(0);
241+
backend = ggml_backend_cann_init(main_gpu);
241242
if (!backend) {
242243
LOG_ERROR("CANN backend init failed");
243244
}
@@ -1144,7 +1145,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
11441145
bool keep_clip_on_cpu,
11451146
bool keep_control_net_cpu,
11461147
bool keep_vae_on_cpu,
1147-
bool diffusion_flash_attn) {
1148+
bool diffusion_flash_attn,
1149+
int main_gpu) {
11481150
sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
11491151
if (sd_ctx == NULL) {
11501152
return NULL;
@@ -1186,7 +1188,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
11861188
keep_clip_on_cpu,
11871189
keep_control_net_cpu,
11881190
keep_vae_on_cpu,
1189-
diffusion_flash_attn)) {
1191+
diffusion_flash_attn,
1192+
main_gpu)) {
11901193
delete sd_ctx->sd;
11911194
sd_ctx->sd = NULL;
11921195
free(sd_ctx);

‎stable-diffusion.h

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -182,7 +182,8 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
182182
bool keep_clip_on_cpu,
183183
bool keep_control_net_cpu,
184184
bool keep_vae_on_cpu,
185-
bool diffusion_flash_attn);
185+
bool diffusion_flash_attn,
186+
int main_gpu = 0);
186187

187188
SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
188189

@@ -253,7 +254,8 @@ typedef struct upscaler_ctx_t upscaler_ctx_t;
253254

254255
SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
255256
int n_threads,
256-
enum sd_type_t wtype);
257+
enum sd_type_t wtype,
258+
int main_gpu = 0);
257259
SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
258260

259261
SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor);

‎upscaler.cpp

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ struct UpscalerGGML {
1414
: n_threads(n_threads) {
1515
}
1616

17-
bool load_from_file(const std::string& esrgan_path) {
17+
bool load_from_file(const std::string& esrgan_path, int main_gpu) {
1818
#ifdef SD_USE_CUDA
1919
#ifdef SD_USE_HIP
2020
LOG_DEBUG("Using HIP backend");
@@ -25,7 +25,7 @@ struct UpscalerGGML {
2525
LOG_DEBUG("Using CUDA backend");
2626
#endif
2727
#endif
28-
backend = ggml_backend_cuda_init(0);
28+
backend = ggml_backend_cuda_init(main_gpu);
2929
if (!backend) {
3030
LOG_ERROR("CUDA backend init failed");
3131
}
@@ -39,21 +39,21 @@ struct UpscalerGGML {
3939
#endif
4040
#ifdef SD_USE_VULKAN
4141
LOG_DEBUG("Using Vulkan backend");
42-
backend = ggml_backend_vk_init(0);
42+
backend = ggml_backend_vk_init(main_gpu);
4343
if (!backend) {
4444
LOG_ERROR("Vulkan backend init failed");
4545
}
4646
#endif
4747
#ifdef SD_USE_SYCL
4848
LOG_DEBUG("Using SYCL backend");
49-
backend = ggml_backend_sycl_init(0);
49+
backend = ggml_backend_sycl_init(main_gpu);
5050
if (!backend) {
5151
LOG_ERROR("SYCL backend init failed");
5252
}
5353
#endif
5454
#ifdef SD_USE_CANN
5555
LOG_DEBUG("Using CANN backend");
56-
backend = ggml_backend_cann_init(0);
56+
backend = ggml_backend_cann_init(main_gpu);
5757
if (!backend) {
5858
LOG_ERROR("CANN backend init failed");
5959
}
@@ -124,7 +124,8 @@ struct upscaler_ctx_t {
124124

125125
upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
126126
int n_threads,
127-
enum sd_type_t wtype) {
127+
enum sd_type_t wtype,
128+
int main_gpu) {
128129
upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t));
129130
if (upscaler_ctx == NULL) {
130131
return NULL;
@@ -136,7 +137,7 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
136137
return NULL;
137138
}
138139

139-
if (!upscaler_ctx->upscaler->load_from_file(esrgan_path)) {
140+
if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, main_gpu)) {
140141
delete upscaler_ctx->upscaler;
141142
upscaler_ctx->upscaler = NULL;
142143
free(upscaler_ctx);

0 commit comments

Comments
 (0)