mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-07-29 06:10:52 -05:00
added an option to offload vae decoder to CPU for mem-limited gpus
This commit is contained in:
@@ -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");
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user