diff --git a/README.md b/README.md index ed6d15acc..a5c86b0f6 100644 --- a/README.md +++ b/README.md @@ -356,6 +356,9 @@ arguments: --clip-skip N ignore last_dot_pos layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1) <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x --vae-tiling process vae in tiles to reduce memory usage + --vae-tile-size [X]x[Y] tile size for vae tiling (default: 32x32) + --vae-relative-tile-size [X]x[Y] relative tile size for vae tiling, in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size) + --vae-tile-overlap OVERLAP tile overlap for vae tiling, in fraction of tile size (default: 0.5) --vae-on-cpu keep vae in cpu (for low vram) --clip-on-cpu keep clip in cpu (for low vram) --diffusion-fa use flash attention in the diffusion model (for low vram) diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 10cb22b87..1c3c15ba5 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -101,7 +101,6 @@ struct SDParams { rng_type_t rng_type = CUDA_RNG; int64_t seed = 42; bool verbose = false; - bool vae_tiling = false; bool offload_params_to_cpu = false; bool control_net_cpu = false; bool normalize_input = false; @@ -119,6 +118,8 @@ struct SDParams { int chroma_t5_mask_pad = 1; float flow_shift = INFINITY; + sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f}; + SDParams() { sd_sample_params_init(&sample_params); sd_sample_params_init(&high_noise_sample_params); @@ -180,7 +181,7 @@ void print_params(SDParams params) { printf(" rng: %s\n", sd_rng_type_name(params.rng_type)); printf(" seed: %ld\n", params.seed); printf(" batch_count: %d\n", params.batch_count); - printf(" vae_tiling: %s\n", params.vae_tiling ? "true" : "false"); + printf(" vae_tiling: %s\n", params.vae_tiling_params.enabled ? "true" : "false"); printf(" upscale_repeats: %d\n", params.upscale_repeats); printf(" chroma_use_dit_mask: %s\n", params.chroma_use_dit_mask ? "true" : "false"); printf(" chroma_use_t5_mask: %s\n", params.chroma_use_t5_mask ? "true" : "false"); @@ -268,6 +269,9 @@ void print_usage(int argc, const char* argv[]) { printf(" --clip-skip N ignore last_dot_pos layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1)\n"); printf(" <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x\n"); printf(" --vae-tiling process vae in tiles to reduce memory usage\n"); + printf(" --vae-tile-size [X]x[Y] tile size for vae tiling (default: 32x32)\n"); + printf(" --vae-relative-tile-size [X]x[Y] relative tile size for vae tiling, in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)\n"); + printf(" --vae-tile-overlap OVERLAP tile overlap for vae tiling, in fraction of tile size (default: 0.5)\n"); printf(" --vae-on-cpu keep vae in cpu (for low vram)\n"); printf(" --clip-on-cpu keep clip in cpu (for low vram)\n"); printf(" --diffusion-fa use flash attention in the diffusion model (for low vram)\n"); @@ -485,7 +489,6 @@ void parse_args(int argc, const char** argv, SDParams& params) { {"-o", "--output", "", ¶ms.output_path}, {"-p", "--prompt", "", ¶ms.prompt}, {"-n", "--negative-prompt", "", ¶ms.negative_prompt}, - {"", "--upscale-model", "", ¶ms.esrgan_path}, }; @@ -523,10 +526,11 @@ void parse_args(int argc, const char** argv, SDParams& params) { {"", "--control-strength", "", ¶ms.control_strength}, {"", "--moe-boundary", "", ¶ms.moe_boundary}, {"", "--flow-shift", "", ¶ms.flow_shift}, + {"", "--vae-tile-overlap", "", ¶ms.vae_tiling_params.target_overlap}, }; options.bool_options = { - {"", "--vae-tiling", "", true, ¶ms.vae_tiling}, + {"", "--vae-tiling", "", true, ¶ms.vae_tiling_params.enabled}, {"", "--offload-to-cpu", "", true, ¶ms.offload_params_to_cpu}, {"", "--control-net-cpu", "", true, ¶ms.control_net_cpu}, {"", "--normalize-input", "", true, ¶ms.normalize_input}, @@ -726,6 +730,52 @@ void parse_args(int argc, const char** argv, SDParams& params) { return 1; }; + auto on_tile_size_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + std::string tile_size_str = argv[index]; + size_t x_pos = tile_size_str.find('x'); + try { + if (x_pos != std::string::npos) { + std::string tile_x_str = tile_size_str.substr(0, x_pos); + std::string tile_y_str = tile_size_str.substr(x_pos + 1); + params.vae_tiling_params.tile_size_x = std::stoi(tile_x_str); + params.vae_tiling_params.tile_size_y = std::stoi(tile_y_str); + } else { + params.vae_tiling_params.tile_size_x = params.vae_tiling_params.tile_size_y = std::stoi(tile_size_str); + } + } catch (const std::invalid_argument& e) { + return -1; + } catch (const std::out_of_range& e) { + return -1; + } + return 1; + }; + + auto on_relative_tile_size_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + std::string rel_size_str = argv[index]; + size_t x_pos = rel_size_str.find('x'); + try { + if (x_pos != std::string::npos) { + std::string rel_x_str = rel_size_str.substr(0, x_pos); + std::string rel_y_str = rel_size_str.substr(x_pos + 1); + params.vae_tiling_params.rel_size_x = std::stof(rel_x_str); + params.vae_tiling_params.rel_size_y = std::stof(rel_y_str); + } else { + params.vae_tiling_params.rel_size_x = params.vae_tiling_params.rel_size_y = std::stof(rel_size_str); + } + } catch (const std::invalid_argument& e) { + return -1; + } catch (const std::out_of_range& e) { + return -1; + } + return 1; + }; + options.manual_options = { {"-M", "--mode", "", on_mode_arg}, {"", "--type", "", on_type_arg}, @@ -739,6 +789,8 @@ void parse_args(int argc, const char** argv, SDParams& params) { {"", "--high-noise-skip-layers", "", on_high_noise_skip_layers_arg}, {"-r", "--ref-image", "", on_ref_image_arg}, {"-h", "--help", "", on_help_arg}, + {"", "--vae-tile-size", "", on_tile_size_arg}, + {"", "--vae-relative-tile-size", "", on_relative_tile_size_arg}, }; if (!parse_options(argc, argv, options)) { @@ -1176,7 +1228,6 @@ int main(int argc, const char* argv[]) { params.embedding_dir.c_str(), params.stacked_id_embed_dir.c_str(), vae_decode_only, - params.vae_tiling, true, params.n_threads, params.wtype, @@ -1229,6 +1280,7 @@ int main(int argc, const char* argv[]) { params.style_ratio, params.normalize_input, params.input_id_images_path.c_str(), + params.vae_tiling_params, }; results = generate_image(sd_ctx, &img_gen_params); diff --git a/ggml_extend.hpp b/ggml_extend.hpp index b88344e18..ca7efd4d7 100644 --- a/ggml_extend.hpp +++ b/ggml_extend.hpp @@ -494,7 +494,10 @@ __STATIC_INLINE__ void ggml_merge_tensor_2d(struct ggml_tensor* input, struct ggml_tensor* output, int x, int y, - int overlap) { + int overlap_x, + int overlap_y, + int x_skip = 0, + int y_skip = 0) { int64_t width = input->ne[0]; int64_t height = input->ne[1]; int64_t channels = input->ne[2]; @@ -503,17 +506,17 @@ __STATIC_INLINE__ void ggml_merge_tensor_2d(struct ggml_tensor* input, int64_t img_height = output->ne[1]; GGML_ASSERT(input->type == GGML_TYPE_F32 && output->type == GGML_TYPE_F32); - for (int iy = 0; iy < height; iy++) { - for (int ix = 0; ix < width; ix++) { + for (int iy = y_skip; iy < height; iy++) { + for (int ix = x_skip; ix < width; ix++) { for (int k = 0; k < channels; k++) { float new_value = ggml_tensor_get_f32(input, ix, iy, k); - if (overlap > 0) { // blend colors in overlapped area + if (overlap_x > 0 || overlap_y > 0) { // blend colors in overlapped area float old_value = ggml_tensor_get_f32(output, x + ix, y + iy, k); - const float x_f_0 = (x > 0) ? ix / float(overlap) : 1; - const float x_f_1 = (x < (img_width - width)) ? (width - ix) / float(overlap) : 1; - const float y_f_0 = (y > 0) ? iy / float(overlap) : 1; - const float y_f_1 = (y < (img_height - height)) ? (height - iy) / float(overlap) : 1; + const float x_f_0 = (overlap_x > 0 && x > 0) ? (ix - x_skip) / float(overlap_x) : 1; + const float x_f_1 = (overlap_x > 0 && x < (img_width - width)) ? (width - ix) / float(overlap_x) : 1; + const float y_f_0 = (overlap_y > 0 && y > 0) ? (iy - y_skip) / float(overlap_y) : 1; + const float y_f_1 = (overlap_y > 0 && y < (img_height - height)) ? (height - iy) / float(overlap_y) : 1; const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f); const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f); @@ -745,22 +748,102 @@ __STATIC_INLINE__ std::vector ggml_chunk(struct ggml_contex typedef std::function on_tile_process; +__STATIC_INLINE__ void sd_tiling_calc_tiles(int& num_tiles_dim, + float& tile_overlap_factor_dim, + int small_dim, + int tile_size, + const float tile_overlap_factor) { + int tile_overlap = (tile_size * tile_overlap_factor); + int non_tile_overlap = tile_size - tile_overlap; + + num_tiles_dim = (small_dim - tile_overlap) / non_tile_overlap; + int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim; + + if ((overshoot_dim != non_tile_overlap) && (overshoot_dim <= num_tiles_dim * (tile_size / 2 - tile_overlap))) { + // if tiles don't fit perfectly using the desired overlap + // and there is enough room to squeeze an extra tile without overlap becoming >0.5 + num_tiles_dim++; + } + + tile_overlap_factor_dim = (float)(tile_size * num_tiles_dim - small_dim) / (float)(tile_size * (num_tiles_dim - 1)); + if (num_tiles_dim <= 2) { + if (small_dim <= tile_size) { + num_tiles_dim = 1; + tile_overlap_factor_dim = 0; + } else { + num_tiles_dim = 2; + tile_overlap_factor_dim = (2 * tile_size - small_dim) / (float)tile_size; + } + } +} + // Tiling -__STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const int scale, const int tile_size, const float tile_overlap_factor, on_tile_process on_processing) { +__STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, + ggml_tensor* output, + const int scale, + const int p_tile_size_x, + const int p_tile_size_y, + const float tile_overlap_factor, + on_tile_process on_processing) { output = ggml_set_f32(output, 0); int input_width = (int)input->ne[0]; int input_height = (int)input->ne[1]; int output_width = (int)output->ne[0]; int output_height = (int)output->ne[1]; + + GGML_ASSERT(((input_width / output_width) == (input_height / output_height)) && + ((output_width / input_width) == (output_height / input_height))); + GGML_ASSERT(((input_width / output_width) == scale) || + ((output_width / input_width) == scale)); + + int small_width = output_width; + int small_height = output_height; + + bool decode = output_width > input_width; + if (decode) { + small_width = input_width; + small_height = input_height; + } + + int num_tiles_x; + float tile_overlap_factor_x; + sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor); + + int num_tiles_y; + float tile_overlap_factor_y; + sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor); + + LOG_DEBUG("num tiles : %d, %d ", num_tiles_x, num_tiles_y); + LOG_DEBUG("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor); + GGML_ASSERT(input_width % 2 == 0 && input_height % 2 == 0 && output_width % 2 == 0 && output_height % 2 == 0); // should be multiple of 2 - int tile_overlap = (int32_t)(tile_size * tile_overlap_factor); - int non_tile_overlap = tile_size - tile_overlap; + int tile_overlap_x = (int32_t)(p_tile_size_x * tile_overlap_factor_x); + int non_tile_overlap_x = p_tile_size_x - tile_overlap_x; + + int tile_overlap_y = (int32_t)(p_tile_size_y * tile_overlap_factor_y); + int non_tile_overlap_y = p_tile_size_y - tile_overlap_y; + + int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width; + int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height; + + int input_tile_size_x = tile_size_x; + int input_tile_size_y = tile_size_y; + int output_tile_size_x = tile_size_x; + int output_tile_size_y = tile_size_y; + + if (decode) { + output_tile_size_x *= scale; + output_tile_size_y *= scale; + } else { + input_tile_size_x *= scale; + input_tile_size_y *= scale; + } struct ggml_init_params params = {}; - params.mem_size += tile_size * tile_size * input->ne[2] * sizeof(float); // input chunk - params.mem_size += (tile_size * scale) * (tile_size * scale) * output->ne[2] * sizeof(float); // output chunk + params.mem_size += input_tile_size_x * input_tile_size_y * input->ne[2] * sizeof(float); // input chunk + params.mem_size += output_tile_size_x * output_tile_size_y * output->ne[2] * sizeof(float); // output chunk params.mem_size += 3 * ggml_tensor_overhead(); params.mem_buffer = NULL; params.no_alloc = false; @@ -775,29 +858,50 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const } // tiling - ggml_tensor* input_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size, tile_size, input->ne[2], 1); - ggml_tensor* output_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, tile_size * scale, tile_size * scale, output->ne[2], 1); - on_processing(input_tile, NULL, true); - int num_tiles = ceil((float)input_width / non_tile_overlap) * ceil((float)input_height / non_tile_overlap); + ggml_tensor* input_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, input_tile_size_x, input_tile_size_y, input->ne[2], 1); + ggml_tensor* output_tile = ggml_new_tensor_4d(tiles_ctx, GGML_TYPE_F32, output_tile_size_x, output_tile_size_y, output->ne[2], 1); + int num_tiles = num_tiles_x * num_tiles_y; LOG_INFO("processing %i tiles", num_tiles); - pretty_progress(1, num_tiles, 0.0f); + pretty_progress(0, num_tiles, 0.0f); int tile_count = 1; bool last_y = false, last_x = false; float last_time = 0.0f; - for (int y = 0; y < input_height && !last_y; y += non_tile_overlap) { - if (y + tile_size >= input_height) { - y = input_height - tile_size; + for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { + int dy = 0; + if (y + tile_size_y >= small_height) { + int _y = y; + y = small_height - tile_size_y; + dy = _y - y; + if (decode) { + dy *= scale; + } last_y = true; } - for (int x = 0; x < input_width && !last_x; x += non_tile_overlap) { - if (x + tile_size >= input_width) { - x = input_width - tile_size; + for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) { + int dx = 0; + if (x + tile_size_x >= small_width) { + int _x = x; + x = small_width - tile_size_x; + dx = _x - x; + if (decode) { + dx *= scale; + } last_x = true; } + + int x_in = decode ? x : scale * x; + int y_in = decode ? y : scale * y; + int x_out = decode ? x * scale : x; + int y_out = decode ? y * scale : y; + + int overlap_x_out = decode ? tile_overlap_x * scale : tile_overlap_x; + int overlap_y_out = decode ? tile_overlap_y * scale : tile_overlap_y; + int64_t t1 = ggml_time_ms(); - ggml_split_tensor_2d(input, input_tile, x, y); + ggml_split_tensor_2d(input, input_tile, x_in, y_in); on_processing(input_tile, output_tile, false); - ggml_merge_tensor_2d(output_tile, output, x * scale, y * scale, tile_overlap * scale); + ggml_merge_tensor_2d(output_tile, output, x_out, y_out, overlap_x_out, overlap_y_out, dx, dy); + int64_t t2 = ggml_time_ms(); last_time = (t2 - t1) / 1000.0f; pretty_progress(tile_count, num_tiles, last_time); @@ -811,6 +915,15 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const ggml_free(tiles_ctx); } +__STATIC_INLINE__ void sd_tiling(ggml_tensor* input, + ggml_tensor* output, + const int scale, + const int tile_size, + const float tile_overlap_factor, + on_tile_process on_processing) { + sd_tiling_non_square(input, output, scale, tile_size, tile_size, tile_overlap_factor, on_processing); +} + __STATIC_INLINE__ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx, struct ggml_tensor* a) { const float eps = 1e-6f; // default eps parameter diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp index e104d3f0a..3749b1007 100644 --- a/stable-diffusion.cpp +++ b/stable-diffusion.cpp @@ -108,10 +108,10 @@ class StableDiffusionGGML { std::shared_ptr pmid_id_embeds; std::string taesd_path; - bool use_tiny_autoencoder = false; - bool vae_tiling = false; - bool offload_params_to_cpu = false; - bool stacked_id = false; + bool use_tiny_autoencoder = false; + sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0, 0}; + bool offload_params_to_cpu = false; + bool stacked_id = false; bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; @@ -183,7 +183,6 @@ class StableDiffusionGGML { lora_model_dir = SAFE_STR(sd_ctx_params->lora_model_dir); taesd_path = SAFE_STR(sd_ctx_params->taesd_path); use_tiny_autoencoder = taesd_path.size() > 0; - vae_tiling = sd_ctx_params->vae_tiling; offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu; if (sd_ctx_params->rng_type == STD_DEFAULT_RNG) { @@ -1297,15 +1296,77 @@ class StableDiffusionGGML { return latent; } - ggml_tensor* encode_first_stage(ggml_context* work_ctx, ggml_tensor* x, bool decode_video = false) { + void get_tile_sizes(int& tile_size_x, + int& tile_size_y, + float& tile_overlap, + const sd_tiling_params_t& params, + int latent_x, + int latent_y, + float encoding_factor = 1.0f) { + tile_overlap = std::max(std::min(params.target_overlap, 0.5f), 0.0f); + auto get_tile_size = [&](int requested_size, float factor, int latent_size) { + const int default_tile_size = 32; + const int min_tile_dimension = 4; + int tile_size = default_tile_size; + // factor <= 1 means simple fraction of the latent dimension + // factor > 1 means number of tiles across that dimension + if (factor > 0.f) { + if (factor > 1.0) + factor = 1 / (factor - factor * tile_overlap + tile_overlap); + tile_size = std::round(latent_size * factor); + } else if (requested_size >= min_tile_dimension) { + tile_size = requested_size; + } + tile_size *= encoding_factor; + return std::max(std::min(tile_size, latent_size), min_tile_dimension); + }; + + tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x); + tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y); + } + + ggml_tensor* encode_first_stage(ggml_context* work_ctx, ggml_tensor* x, bool encode_video = false) { int64_t t0 = ggml_time_ms(); ggml_tensor* result = NULL; + int W = x->ne[0] / 8; + int H = x->ne[1] / 8; + if (vae_tiling_params.enabled && !encode_video) { + // TODO wan2.2 vae support? + int C = sd_version_is_dit(version) ? 16 : 4; + if (!use_tiny_autoencoder) { + C *= 2; + } + result = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, W, H, C, x->ne[3]); + } + if (!use_tiny_autoencoder) { + float tile_overlap; + int tile_size_x, tile_size_y; + // multiply tile size for encode to keep the compute buffer size consistent + get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, vae_tiling_params, W, H, 1.30539f); + + LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + process_vae_input_tensor(x); - first_stage_model->compute(n_threads, x, false, &result, work_ctx); + if (vae_tiling_params.enabled && !encode_video) { + auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { + first_stage_model->compute(n_threads, in, false, &out, work_ctx); + }; + sd_tiling_non_square(x, result, 8, tile_size_x, tile_size_y, tile_overlap, on_tiling); + } else { + first_stage_model->compute(n_threads, x, false, &result, work_ctx); + } first_stage_model->free_compute_buffer(); } else { - tae_first_stage->compute(n_threads, x, false, &result, work_ctx); + if (vae_tiling_params.enabled && !encode_video) { + // split latent in 32x32 tiles and compute in several steps + auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { + tae_first_stage->compute(n_threads, in, false, &out, NULL); + }; + sd_tiling(x, result, 8, 64, 0.5f, on_tiling); + } else { + tae_first_stage->compute(n_threads, x, false, &result, work_ctx); + } tae_first_stage->free_compute_buffer(); } @@ -1422,24 +1483,29 @@ class StableDiffusionGGML { C, x->ne[3]); } - int64_t t0 = ggml_time_ms(); if (!use_tiny_autoencoder) { + float tile_overlap; + int tile_size_x, tile_size_y; + get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, vae_tiling_params, x->ne[0], x->ne[1]); + + LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y); + process_latent_out(x); // x = load_tensor_from_file(work_ctx, "wan_vae_z.bin"); - if (vae_tiling && !decode_video) { + if (vae_tiling_params.enabled && !decode_video) { // split latent in 32x32 tiles and compute in several steps auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { first_stage_model->compute(n_threads, in, true, &out, NULL); }; - sd_tiling(x, result, 8, 32, 0.5f, on_tiling); + sd_tiling_non_square(x, result, 8, tile_size_x, tile_size_y, tile_overlap, on_tiling); } else { first_stage_model->compute(n_threads, x, true, &result, work_ctx); } first_stage_model->free_compute_buffer(); process_vae_output_tensor(result); } else { - if (vae_tiling && !decode_video) { + if (vae_tiling_params.enabled && !decode_video) { // split latent in 64x64 tiles and compute in several steps auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { tae_first_stage->compute(n_threads, in, true, &out); @@ -1561,7 +1627,6 @@ enum scheduler_t str_to_schedule(const char* str) { void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { *sd_ctx_params = {}; sd_ctx_params->vae_decode_only = true; - sd_ctx_params->vae_tiling = false; sd_ctx_params->free_params_immediately = true; sd_ctx_params->n_threads = get_num_physical_cores(); sd_ctx_params->wtype = SD_TYPE_COUNT; @@ -1625,7 +1690,6 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { SAFE_STR(sd_ctx_params->embedding_dir), SAFE_STR(sd_ctx_params->stacked_id_embed_dir), BOOL_STR(sd_ctx_params->vae_decode_only), - BOOL_STR(sd_ctx_params->vae_tiling), BOOL_STR(sd_ctx_params->free_params_immediately), sd_ctx_params->n_threads, sd_type_name(sd_ctx_params->wtype), @@ -1692,16 +1756,17 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { *sd_img_gen_params = {}; sd_sample_params_init(&sd_img_gen_params->sample_params); - sd_img_gen_params->clip_skip = -1; - sd_img_gen_params->ref_images_count = 0; - sd_img_gen_params->width = 512; - sd_img_gen_params->height = 512; - sd_img_gen_params->strength = 0.75f; - sd_img_gen_params->seed = -1; - sd_img_gen_params->batch_count = 1; - sd_img_gen_params->control_strength = 0.9f; - sd_img_gen_params->style_strength = 20.f; - sd_img_gen_params->normalize_input = false; + sd_img_gen_params->clip_skip = -1; + sd_img_gen_params->ref_images_count = 0; + sd_img_gen_params->width = 512; + sd_img_gen_params->height = 512; + sd_img_gen_params->strength = 0.75f; + sd_img_gen_params->seed = -1; + sd_img_gen_params->batch_count = 1; + sd_img_gen_params->control_strength = 0.9f; + sd_img_gen_params->style_strength = 20.f; + sd_img_gen_params->normalize_input = false; + sd_img_gen_params->vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f}; } char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { @@ -1721,6 +1786,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { "sample_params: %s\n" "strength: %.2f\n" "seed: %" PRId64 + "VAE tiling:" "\n" "batch_count: %d\n" "ref_images_count: %d\n" @@ -1737,6 +1803,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { SAFE_STR(sample_params_str), sd_img_gen_params->strength, sd_img_gen_params->seed, + BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled), sd_img_gen_params->batch_count, sd_img_gen_params->ref_images_count, BOOL_STR(sd_img_gen_params->increase_ref_index), @@ -2173,8 +2240,9 @@ ggml_tensor* generate_init_latent(sd_ctx_t* sd_ctx, } sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) { - int width = sd_img_gen_params->width; - int height = sd_img_gen_params->height; + sd_ctx->sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; + int width = sd_img_gen_params->width; + int height = sd_img_gen_params->height; if (sd_version_is_dit(sd_ctx->sd->version)) { if (width % 16 || height % 16) { LOG_ERROR("Image dimensions must be must be a multiple of 16 on each axis for %s models. (Got %dx%d)", diff --git a/stable-diffusion.h b/stable-diffusion.h index dbd53840a..0a6493f95 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -114,6 +114,15 @@ enum sd_log_level_t { SD_LOG_ERROR }; +typedef struct { + bool enabled; + int tile_size_x; + int tile_size_y; + float target_overlap; + float rel_size_x; + float rel_size_y; +} sd_tiling_params_t; + typedef struct { const char* model_path; const char* clip_l_path; @@ -129,7 +138,6 @@ typedef struct { const char* embedding_dir; const char* stacked_id_embed_dir; bool vae_decode_only; - bool vae_tiling; bool free_params_immediately; int n_threads; enum sd_type_t wtype; @@ -197,6 +205,7 @@ typedef struct { float style_strength; bool normalize_input; const char* input_id_images_path; + sd_tiling_params_t vae_tiling_params; } sd_img_gen_params_t; typedef struct {