Skip to content

Commit 6aa3b9d

Browse files
author
Stéphane du Hamel
committed
improve handling of VAE decode failures
1 parent a48b4a3 commit 6aa3b9d

4 files changed

Lines changed: 48 additions & 24 deletions

File tree

‎examples/cli/main.cpp‎

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -409,7 +409,7 @@ bool save_results(const SDCliParams& cli_params,
409409
auto write_image = [&](const fs::path& path, int idx) {
410410
const sd_image_t& img = results[idx];
411411
if (!img.data)
412-
return;
412+
return false;
413413

414414
std::string params = get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + idx);
415415
int ok = 0;
@@ -419,8 +419,11 @@ bool save_results(const SDCliParams& cli_params,
419419
ok = stbi_write_png(path.string().c_str(), img.width, img.height, img.channel, img.data, 0, params.c_str());
420420
}
421421
LOG_INFO("save result image %d to '%s' (%s)", idx, path.string().c_str(), ok ? "success" : "failure");
422+
return ok != 0;
422423
};
423424

425+
int sucessful_reults = 0;
426+
424427
if (std::regex_search(cli_params.output_path, format_specifier_regex)) {
425428
if (!is_jpg && ext_lower != ".png")
426429
ext = ".png";
@@ -429,19 +432,26 @@ bool save_results(const SDCliParams& cli_params,
429432

430433
for (int i = 0; i < num_results; ++i) {
431434
fs::path img_path = format_frame_idx(pattern.string(), output_begin_idx + i);
432-
write_image(img_path, i);
435+
if (write_image(img_path, i)) {
436+
sucessful_reults++;
437+
}
433438
}
434-
return true;
439+
LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
440+
return sucessful_reults != 0;
435441
}
436442

437443
if (cli_params.mode == VID_GEN && num_results > 1) {
438444
if (ext_lower != ".avi")
439445
ext = ".avi";
440446
fs::path video_path = base_path;
441447
video_path += ext;
442-
create_mjpg_avi_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps);
443-
LOG_INFO("save result MJPG AVI video to '%s'", video_path.string().c_str());
444-
return true;
448+
if (create_mjpg_avi_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
449+
LOG_INFO("save result MJPG AVI video to '%s'", video_path.string().c_str());
450+
return true;
451+
} else {
452+
LOG_ERROR("Failed to save result MPG AVI video to '%s'", video_path.string().c_str());
453+
return false;
454+
}
445455
}
446456

447457
if (!is_jpg && ext_lower != ".png")
@@ -453,10 +463,12 @@ bool save_results(const SDCliParams& cli_params,
453463
img_path += "_" + std::to_string(output_begin_idx + i);
454464
}
455465
img_path += ext;
456-
write_image(img_path, i);
466+
if (write_image(img_path, i)) {
467+
sucessful_reults++;
468+
}
457469
}
458-
459-
return true;
470+
LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
471+
return sucessful_reults != 0;
460472
}
461473

462474
int main(int argc, const char* argv[]) {

‎ggml_extend.hpp‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -778,7 +778,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor*
778778
return x;
779779
}
780780

781-
typedef std::function<void(ggml_tensor*, ggml_tensor*, bool)> on_tile_process;
781+
typedef std::function<bool(ggml_tensor*, ggml_tensor*, bool)> on_tile_process;
782782

783783
__STATIC_INLINE__ void sd_tiling_calc_tiles(int& num_tiles_dim,
784784
float& tile_overlap_factor_dim,
@@ -929,12 +929,15 @@ __STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input,
929929

930930
int64_t t1 = ggml_time_ms();
931931
ggml_ext_tensor_split_2d(input, input_tile, x_in, y_in);
932-
on_processing(input_tile, output_tile, false);
933-
ggml_ext_tensor_merge_2d(output_tile, output, x_out, y_out, overlap_x_out, overlap_y_out, dx, dy);
932+
if (on_processing(input_tile, output_tile, false)) {
933+
ggml_ext_tensor_merge_2d(output_tile, output, x_out, y_out, overlap_x_out, overlap_y_out, dx, dy);
934934

935-
int64_t t2 = ggml_time_ms();
936-
last_time = (t2 - t1) / 1000.0f;
937-
pretty_progress(tile_count, num_tiles, last_time);
935+
int64_t t2 = ggml_time_ms();
936+
last_time = (t2 - t1) / 1000.0f;
937+
pretty_progress(tile_count, num_tiles, last_time);
938+
} else {
939+
LOG_ERROR("Failed to process patch %d at (%d, %d)", tile_count, x, y);
940+
}
938941
tile_count++;
939942
}
940943
last_x = false;

‎stable-diffusion.cpp‎

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1542,7 +1542,7 @@ class StableDiffusionGGML {
15421542
if (vae_tiling_params.enabled) {
15431543
// split latent in 32x32 tiles and compute in several steps
15441544
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
1545-
first_stage_model->compute(n_threads, in, true, &out, nullptr);
1545+
return first_stage_model->compute(n_threads, in, true, &out, nullptr);
15461546
};
15471547
silent_tiling(latents, result, get_vae_scale_factor(), 32, 0.5f, on_tiling);
15481548

@@ -1561,7 +1561,7 @@ class StableDiffusionGGML {
15611561
if (vae_tiling_params.enabled) {
15621562
// split latent in 64x64 tiles and compute in several steps
15631563
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
1564-
tae_first_stage->compute(n_threads, in, true, &out, nullptr);
1564+
return tae_first_stage->compute(n_threads, in, true, &out, nullptr);
15651565
};
15661566
silent_tiling(latents, result, get_vae_scale_factor(), 64, 0.5f, on_tiling);
15671567
} else {
@@ -2530,7 +2530,7 @@ class StableDiffusionGGML {
25302530
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
25312531

25322532
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
2533-
first_stage_model->compute(n_threads, in, false, &out, work_ctx);
2533+
return first_stage_model->compute(n_threads, in, false, &out, work_ctx);
25342534
};
25352535
sd_tiling_non_square(x, result, vae_scale_factor, tile_size_x, tile_size_y, tile_overlap, on_tiling);
25362536
} else {
@@ -2541,7 +2541,7 @@ class StableDiffusionGGML {
25412541
if (vae_tiling_params.enabled && !encode_video) {
25422542
// split latent in 32x32 tiles and compute in several steps
25432543
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
2544-
tae_first_stage->compute(n_threads, in, false, &out, nullptr);
2544+
return tae_first_stage->compute(n_threads, in, false, &out, nullptr);
25452545
};
25462546
sd_tiling(x, result, vae_scale_factor, 64, 0.5f, on_tiling);
25472547
} else {
@@ -2659,23 +2659,31 @@ class StableDiffusionGGML {
26592659

26602660
// split latent in 32x32 tiles and compute in several steps
26612661
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
2662-
first_stage_model->compute(n_threads, in, true, &out, nullptr);
2662+
return first_stage_model->compute(n_threads, in, true, &out, nullptr);
26632663
};
26642664
sd_tiling_non_square(x, result, vae_scale_factor, tile_size_x, tile_size_y, tile_overlap, on_tiling);
26652665
} else {
2666-
first_stage_model->compute(n_threads, x, true, &result, work_ctx);
2666+
if(!first_stage_model->compute(n_threads, x, true, &result, work_ctx)){
2667+
LOG_ERROR("Failed to decode latetnts");
2668+
first_stage_model->free_compute_buffer();
2669+
return nullptr;
2670+
}
26672671
}
26682672
first_stage_model->free_compute_buffer();
26692673
process_vae_output_tensor(result);
26702674
} else {
26712675
if (vae_tiling_params.enabled) {
26722676
// split latent in 64x64 tiles and compute in several steps
26732677
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
2674-
tae_first_stage->compute(n_threads, in, true, &out);
2678+
return tae_first_stage->compute(n_threads, in, true, &out);
26752679
};
26762680
sd_tiling(x, result, vae_scale_factor, 64, 0.5f, on_tiling);
26772681
} else {
2678-
tae_first_stage->compute(n_threads, x, true, &result);
2682+
if(!tae_first_stage->compute(n_threads, x, true, &result)){
2683+
LOG_ERROR("Failed to decode latetnts");
2684+
tae_first_stage->free_compute_buffer();
2685+
return nullptr;
2686+
}
26792687
}
26802688
tae_first_stage->free_compute_buffer();
26812689
}
@@ -3440,6 +3448,7 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx,
34403448
ggml_free(work_ctx);
34413449
return nullptr;
34423450
}
3451+
memset(result_images, 0, batch_count * sizeof(sd_image_t));
34433452

34443453
for (size_t i = 0; i < decoded_images.size(); i++) {
34453454
result_images[i].width = width;

‎upscaler.cpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ struct UpscalerGGML {
8989

9090
ggml_tensor* upscaled = ggml_new_tensor_4d(upscale_ctx, GGML_TYPE_F32, output_width, output_height, 3, 1);
9191
auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
92-
esrgan_upscaler->compute(n_threads, in, &out);
92+
return esrgan_upscaler->compute(n_threads, in, &out);
9393
};
9494
int64_t t0 = ggml_time_ms();
9595
sd_tiling(input_image_tensor, upscaled, esrgan_upscaler->scale, esrgan_upscaler->tile_size, 0.25f, on_tiling);

0 commit comments

Comments
 (0)