added an option to offload vae decoder to CPU for mem-limited gpus

This commit is contained in:
bssrdf
2024-02-21 12:44:04 -05:00
parent 78ec6f7e51
commit 8053aefdf7
3 changed files with 23 additions and 6 deletions

View File

@@ -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");

View File

@@ -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);

View File

@@ -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);