mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-11 03:24:30 +00:00
feat: add VAE encoding tiling support and adaptive overlap (#484)
* implement tiling vae encode support * Tiling (vae/upscale): adaptative overlap * Tiling: fix edge case * Tiling: fix crash when less than 2 tiles per dim * remove extra dot * Tiling: fix edge cases for adaptative overlap * tiling: fix edge case * set vae tile size via env var * vae tiling: refactor again, base on smaller buffer for alignment * Use bigger tiles for encode (to match compute buffer size) * Fix edge case when tile is bigger than latent * non-square VAE tiling (#3) * refactor tile number calculation * support non-square tiles * add env var to change tile overlap * add safeguards and better error messages for SD_TILE_OVERLAP * add safeguards and include overlapping factor for SD_TILE_SIZE * avoid rounding issues when specifying SD_TILE_SIZE as a factor * lower SD_TILE_OVERLAP limit * zero-init empty output buffer * Fix decode latent size * fix encode * tile size params instead of env * Tiled vae parameter validation (#6) * avoid crash with invalid tile sizes, use 0 for default * refactor default tile size, limit overlap factor * remove explicit parameter for relative tile size * limit encoding tile to latent size * unify code style and format code * update docs * fix get_tile_sizes in decode_first_stage --------- Co-authored-by: Wagner Bruna <wbruna@users.noreply.github.com> Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
5 files changed
+303
-58
No files matched your search
+57
-5
@@ -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);
|
||||
|
||||
Reference in new issue
Block a user