Skip to content

Commit e362c90

Browse files
thxCodethxCode
committed
feat: support main gpu
Signed-off-by: thxCode <thxcode0824@gmail.com>
1 parent 222d74d commit e362c90

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
@@ -219,7 +219,8 @@ class StableDiffusionGGML {
219219
bool clip_on_cpu,
220220
bool control_net_cpu,
221221
bool vae_on_cpu,
222-
bool diffusion_flash_attn) {
222+
bool diffusion_flash_attn,
223+
int main_gpu) {
223224
use_tiny_autoencoder = taesd_path.size() > 0;
224225
#ifdef SD_USE_CUDA
225226
#ifdef SD_USE_HIP
@@ -229,7 +230,7 @@ class StableDiffusionGGML {
229230
#else
230231
LOG_DEBUG("Using CUDA backend");
231232
#endif
232-
backend = ggml_backend_cuda_init(0);
233+
backend = ggml_backend_cuda_init(main_gpu);
233234
if (!backend) {
234235
LOG_ERROR("CUDA backend init failed");
235236
}
@@ -243,21 +244,21 @@ class StableDiffusionGGML {
243244
#endif
244245
#ifdef SD_USE_VULKAN
245246
LOG_DEBUG("Using Vulkan backend");
246-
backend = ggml_backend_vk_init(0);
247+
backend = ggml_backend_vk_init(main_gpu);
247248
if (!backend) {
248249
LOG_ERROR("Vulkan backend init failed");
249250
}
250251
#endif
251252
#ifdef SD_USE_SYCL
252253
LOG_DEBUG("Using SYCL backend");
253-
backend = ggml_backend_sycl_init(0);
254+
backend = ggml_backend_sycl_init(main_gpu);
254255
if (!backend) {
255256
LOG_ERROR("SYCL backend init failed");
256257
}
257258
#endif
258259
#ifdef SD_USE_CANN
259260
LOG_DEBUG("Using CANN backend");
260-
backend = ggml_backend_cann_init(0);
261+
backend = ggml_backend_cann_init(main_gpu);
261262
if (!backend) {
262263
LOG_ERROR("CANN backend init failed");
263264
}
@@ -1170,7 +1171,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
11701171
bool keep_clip_on_cpu,
11711172
bool keep_control_net_cpu,
11721173
bool keep_vae_on_cpu,
1173-
bool diffusion_flash_attn) {
1174+
bool diffusion_flash_attn,
1175+
int main_gpu) {
11741176
sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
11751177
if (sd_ctx == NULL) {
11761178
return NULL;
@@ -1212,7 +1214,8 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
12121214
keep_clip_on_cpu,
12131215
keep_control_net_cpu,
12141216
keep_vae_on_cpu,
1215-
diffusion_flash_attn)) {
1217+
diffusion_flash_attn,
1218+
main_gpu)) {
12161219
delete sd_ctx->sd;
12171220
sd_ctx->sd = NULL;
12181221
free(sd_ctx);

‎stable-diffusion.h

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,8 @@ SD_API sd_ctx_t* new_sd_ctx(const char* model_path,
155155
bool keep_clip_on_cpu,
156156
bool keep_control_net_cpu,
157157
bool keep_vae_on_cpu,
158-
bool diffusion_flash_attn);
158+
bool diffusion_flash_attn,
159+
int main_gpu = 0);
159160

160161
SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
161162

@@ -225,7 +226,8 @@ SD_API sd_image_t* img2vid(sd_ctx_t* sd_ctx,
225226
typedef struct upscaler_ctx_t upscaler_ctx_t;
226227

227228
SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
228-
int n_threads);
229+
int n_threads,
230+
int main_gpu = 0);
229231
SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
230232

231233
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");
@@ -23,7 +23,7 @@ struct UpscalerGGML {
2323
#else
2424
LOG_DEBUG("Using CUDA backend");
2525
#endif
26-
backend = ggml_backend_cuda_init(0);
26+
backend = ggml_backend_cuda_init(main_gpu);
2727
if (!backend) {
2828
LOG_ERROR("CUDA backend init failed");
2929
}
@@ -37,21 +37,21 @@ struct UpscalerGGML {
3737
#endif
3838
#ifdef SD_USE_VULKAN
3939
LOG_DEBUG("Using Vulkan backend");
40-
backend = ggml_backend_vk_init(0);
40+
backend = ggml_backend_vk_init(main_gpu);
4141
if (!backend) {
4242
LOG_ERROR("Vulkan backend init failed");
4343
}
4444
#endif
4545
#ifdef SD_USE_SYCL
4646
LOG_DEBUG("Using SYCL backend");
47-
backend = ggml_backend_sycl_init(0);
47+
backend = ggml_backend_sycl_init(main_gpu);
4848
if (!backend) {
4949
LOG_ERROR("SYCL backend init failed");
5050
}
5151
#endif
5252
#ifdef SD_USE_CANN
5353
LOG_DEBUG("Using CANN backend");
54-
backend = ggml_backend_cann_init(0);
54+
backend = ggml_backend_cann_init(main_gpu);
5555
if (!backend) {
5656
LOG_ERROR("CANN backend init failed");
5757
}
@@ -125,7 +125,8 @@ struct upscaler_ctx_t {
125125
};
126126

127127
upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
128-
int n_threads) {
128+
int n_threads,
129+
int main_gpu) {
129130
upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t));
130131
if (upscaler_ctx == NULL) {
131132
return NULL;
@@ -137,7 +138,7 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
137138
return NULL;
138139
}
139140

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

0 commit comments

Comments
 (0)