Skip to content

Commit b99ac72

Browse files
stduhpfstduhpf
committed
Handle graph compute failures more gracefully
1 parent d586038 commit b99ac72

File tree

16 files changed

+100
-45
lines changed

16 files changed

+100
-45
lines changed

‎clip.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -936,7 +936,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
936936
return gf;
937937
}
938938

939-
void compute(const int n_threads,
939+
bool compute(const int n_threads,
940940
struct ggml_tensor* input_ids,
941941
int num_custom_embeddings,
942942
void* custom_embeddings_data,
@@ -947,7 +947,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
947947
auto get_graph = [&]() -> struct ggml_cgraph* {
948948
return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled);
949949
};
950-
GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
950+
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
951951
}
952952
};
953953

‎conditioner.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -642,14 +642,14 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
642642
return gf;
643643
}
644644

645-
void compute(const int n_threads,
645+
bool compute(const int n_threads,
646646
ggml_tensor* pixel_values,
647647
ggml_tensor** output,
648648
ggml_context* output_ctx) {
649649
auto get_graph = [&]() -> struct ggml_cgraph* {
650650
return build_graph(pixel_values);
651651
};
652-
GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
652+
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
653653
}
654654
};
655655

‎control.hpp‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -410,7 +410,7 @@ struct ControlNet : public GGMLRunner {
410410
return gf;
411411
}
412412

413-
void compute(int n_threads,
413+
bool compute(int n_threads,
414414
struct ggml_tensor* x,
415415
struct ggml_tensor* hint,
416416
struct ggml_tensor* timesteps,
@@ -426,8 +426,11 @@ struct ControlNet : public GGMLRunner {
426426
return build_graph(x, hint, timesteps, context, y);
427427
};
428428

429-
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
429+
if (!GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx)) {
430+
return false;
431+
}
430432
guided_hint_cached = true;
433+
return true;
431434
}
432435

433436
bool load_from_file(const std::string& file_path) {

‎denoiser.hpp‎

Lines changed: 49 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -352,7 +352,7 @@ struct EDMVDenoiser : public CompVisVDenoiser {
352352
}
353353

354354
float t_to_sigma(float t) {
355-
return std::exp(t * 4/(float)TIMESTEPS);
355+
return std::exp(t * 4 / (float)TIMESTEPS);
356356
}
357357

358358
float sigma_to_t(float s) {
@@ -491,7 +491,7 @@ struct FluxFlowDenoiser : public Denoiser {
491491
typedef std::function<ggml_tensor*(ggml_tensor*, float, int)> denoise_cb_t;
492492

493493
// k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
494-
static void sample_k_diffusion(sample_method_t method,
494+
static bool sample_k_diffusion(sample_method_t method,
495495
denoise_cb_t model,
496496
ggml_context* work_ctx,
497497
ggml_tensor* x,
@@ -510,6 +510,9 @@ static void sample_k_diffusion(sample_method_t method,
510510

511511
// denoise
512512
ggml_tensor* denoised = model(x, sigma, i + 1);
513+
if (denoised == NULL) {
514+
return false;
515+
}
513516

514517
// d = (x - denoised) / sigma
515518
{
@@ -563,6 +566,9 @@ static void sample_k_diffusion(sample_method_t method,
563566

564567
// denoise
565568
ggml_tensor* denoised = model(x, sigma, i + 1);
569+
if (denoised == NULL) {
570+
return false;
571+
}
566572

567573
// d = (x - denoised) / sigma
568574
{
@@ -594,6 +600,9 @@ static void sample_k_diffusion(sample_method_t method,
594600
for (int i = 0; i < steps; i++) {
595601
// denoise
596602
ggml_tensor* denoised = model(x, sigmas[i], -(i + 1));
603+
if (denoised == NULL) {
604+
return false;
605+
}
597606

598607
// d = (x - denoised) / sigma
599608
{
@@ -628,7 +637,10 @@ static void sample_k_diffusion(sample_method_t method,
628637
}
629638

630639
ggml_tensor* denoised = model(x2, sigmas[i + 1], i + 1);
631-
float* vec_denoised = (float*)denoised->data;
640+
if (denoised == NULL) {
641+
return false;
642+
}
643+
float* vec_denoised = (float*)denoised->data;
632644
for (int j = 0; j < ggml_nelements(x); j++) {
633645
float d2 = (vec_x2[j] - vec_denoised[j]) / sigmas[i + 1];
634646
vec_d[j] = (vec_d[j] + d2) / 2;
@@ -644,6 +656,9 @@ static void sample_k_diffusion(sample_method_t method,
644656
for (int i = 0; i < steps; i++) {
645657
// denoise
646658
ggml_tensor* denoised = model(x, sigmas[i], i + 1);
659+
if (denoised == NULL) {
660+
return false;
661+
}
647662

648663
// d = (x - denoised) / sigma
649664
{
@@ -680,7 +695,10 @@ static void sample_k_diffusion(sample_method_t method,
680695
}
681696

682697
ggml_tensor* denoised = model(x2, sigma_mid, i + 1);
683-
float* vec_denoised = (float*)denoised->data;
698+
if (denoised == NULL) {
699+
return false;
700+
}
701+
float* vec_denoised = (float*)denoised->data;
684702
for (int j = 0; j < ggml_nelements(x); j++) {
685703
float d2 = (vec_x2[j] - vec_denoised[j]) / sigma_mid;
686704
vec_x[j] = vec_x[j] + d2 * dt_2;
@@ -697,6 +715,9 @@ static void sample_k_diffusion(sample_method_t method,
697715
for (int i = 0; i < steps; i++) {
698716
// denoise
699717
ggml_tensor* denoised = model(x, sigmas[i], i + 1);
718+
if (denoised == NULL) {
719+
return false;
720+
}
700721

701722
// get_ancestral_step
702723
float sigma_up = std::min(sigmas[i + 1],
@@ -741,6 +762,9 @@ static void sample_k_diffusion(sample_method_t method,
741762
}
742763

743764
ggml_tensor* denoised = model(x2, sigmas[i + 1], i + 1);
765+
if (denoised == NULL) {
766+
return false;
767+
}
744768

745769
// Second half-step
746770
for (int j = 0; j < ggml_nelements(x); j++) {
@@ -771,6 +795,9 @@ static void sample_k_diffusion(sample_method_t method,
771795
for (int i = 0; i < steps; i++) {
772796
// denoise
773797
ggml_tensor* denoised = model(x, sigmas[i], i + 1);
798+
if (denoised == NULL) {
799+
return false;
800+
}
774801

775802
float t = t_fn(sigmas[i]);
776803
float t_next = t_fn(sigmas[i + 1]);
@@ -810,6 +837,9 @@ static void sample_k_diffusion(sample_method_t method,
810837
for (int i = 0; i < steps; i++) {
811838
// denoise
812839
ggml_tensor* denoised = model(x, sigmas[i], i + 1);
840+
if (denoised == NULL) {
841+
return false;
842+
}
813843

814844
float t = t_fn(sigmas[i]);
815845
float t_next = t_fn(sigmas[i + 1]);
@@ -860,7 +890,10 @@ static void sample_k_diffusion(sample_method_t method,
860890

861891
// Denoising step
862892
ggml_tensor* denoised = model(x_cur, sigma, i + 1);
863-
float* vec_denoised = (float*)denoised->data;
893+
if (denoised == NULL) {
894+
return false;
895+
}
896+
float* vec_denoised = (float*)denoised->data;
864897
// d_cur = (x_cur - denoised) / sigma
865898
struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x_cur);
866899
float* vec_d_cur = (float*)d_cur->data;
@@ -1003,6 +1036,9 @@ static void sample_k_diffusion(sample_method_t method,
10031036

10041037
// denoise
10051038
ggml_tensor* denoised = model(x, sigma, i + 1);
1039+
if (denoised == NULL) {
1040+
return false;
1041+
}
10061042

10071043
// x = denoised
10081044
{
@@ -1129,6 +1165,9 @@ static void sample_k_diffusion(sample_method_t method,
11291165
// p. 8 (7), compare also p. 38 (226) therein.
11301166
struct ggml_tensor* model_output =
11311167
model(x, sigma, i + 1);
1168+
if (model_output == NULL) {
1169+
return false;
1170+
}
11321171
// Here model_output is still the k-diffusion denoiser
11331172
// output, not the U-net output F_theta(c_in(sigma) x;
11341173
// ...) in Karras et al. (2022), whereas Diffusers'
@@ -1288,6 +1327,9 @@ static void sample_k_diffusion(sample_method_t method,
12881327
}
12891328
struct ggml_tensor* model_output =
12901329
model(x, sigma, i + 1);
1330+
if (model_output == NULL) {
1331+
return false;
1332+
}
12911333
{
12921334
float* vec_x = (float*)x->data;
12931335
float* vec_model_output =
@@ -1395,8 +1437,9 @@ static void sample_k_diffusion(sample_method_t method,
13951437

13961438
default:
13971439
LOG_ERROR("Attempting to sample with nonexisting sample method %i", method);
1398-
abort();
1440+
return false;
13991441
}
1442+
return true;
14001443
}
14011444

14021445
#endif // __DENOISER_HPP__

‎diffusion_model.hpp‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
#include "unet.hpp"
77

88
struct DiffusionModel {
9-
virtual void compute(int n_threads,
9+
virtual bool compute(int n_threads,
1010
struct ggml_tensor* x,
1111
struct ggml_tensor* timesteps,
1212
struct ggml_tensor* context,
@@ -61,7 +61,7 @@ struct UNetModel : public DiffusionModel {
6161
return unet.unet.adm_in_channels;
6262
}
6363

64-
void compute(int n_threads,
64+
bool compute(int n_threads,
6565
struct ggml_tensor* x,
6666
struct ggml_tensor* timesteps,
6767
struct ggml_tensor* context,
@@ -111,7 +111,7 @@ struct MMDiTModel : public DiffusionModel {
111111
return 768 + 1280;
112112
}
113113

114-
void compute(int n_threads,
114+
bool compute(int n_threads,
115115
struct ggml_tensor* x,
116116
struct ggml_tensor* timesteps,
117117
struct ggml_tensor* context,
@@ -162,7 +162,7 @@ struct FluxModel : public DiffusionModel {
162162
return 768;
163163
}
164164

165-
void compute(int n_threads,
165+
bool compute(int n_threads,
166166
struct ggml_tensor* x,
167167
struct ggml_tensor* timesteps,
168168
struct ggml_tensor* context,

‎esrgan.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,14 +183,14 @@ struct ESRGAN : public GGMLRunner {
183183
return gf;
184184
}
185185

186-
void compute(const int n_threads,
186+
bool compute(const int n_threads,
187187
struct ggml_tensor* x,
188188
ggml_tensor** output,
189189
ggml_context* output_ctx = NULL) {
190190
auto get_graph = [&]() -> struct ggml_cgraph* {
191191
return build_graph(x);
192192
};
193-
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
193+
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
194194
}
195195
};
196196

‎flux.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1163,7 +1163,7 @@ namespace Flux {
11631163
return gf;
11641164
}
11651165

1166-
void compute(int n_threads,
1166+
bool compute(int n_threads,
11671167
struct ggml_tensor* x,
11681168
struct ggml_tensor* timesteps,
11691169
struct ggml_tensor* context,
@@ -1182,7 +1182,7 @@ namespace Flux {
11821182
return build_graph(x, timesteps, context, c_concat, y, guidance, skip_layers);
11831183
};
11841184

1185-
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
1185+
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
11861186
}
11871187

11881188
void test() {

‎ggml_extend.hpp‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -383,19 +383,19 @@ __STATIC_INLINE__ void sd_apply_mask(struct ggml_tensor* image_data,
383383
ggml_tensor_set_f32(mask, m, ix, iy);
384384
}
385385
}
386-
float rescale_mx = mask->ne[0]/output->ne[0];
387-
float rescale_my = mask->ne[1]/output->ne[1];
386+
float rescale_mx = mask->ne[0] / output->ne[0];
387+
float rescale_my = mask->ne[1] / output->ne[1];
388388
GGML_ASSERT(output->type == GGML_TYPE_F32);
389389
for (int ix = 0; ix < width; ix++) {
390390
for (int iy = 0; iy < height; iy++) {
391-
int mx = (int)(ix * rescale_mx);
392-
int my = (int)(iy * rescale_my);
391+
int mx = (int)(ix * rescale_mx);
392+
int my = (int)(iy * rescale_my);
393393
float m = ggml_tensor_get_f32(mask, mx, my);
394394
m = round(m); // inpaint models need binary masks
395395
ggml_tensor_set_f32(mask, m, mx, my);
396396
for (int k = 0; k < channels; k++) {
397397
float value = ggml_tensor_get_f32(image_data, ix, iy, k);
398-
value = (1 - m) * (value - masked_value) + masked_value;
398+
value = (1 - m) * (value - masked_value) + masked_value;
399399
ggml_tensor_set_f32(output, value, ix, iy, k);
400400
}
401401
}
@@ -1319,12 +1319,14 @@ struct GGMLRunner {
13191319
}
13201320
}
13211321

1322-
void compute(get_graph_cb_t get_graph,
1322+
bool compute(get_graph_cb_t get_graph,
13231323
int n_threads,
13241324
bool free_compute_buffer_immediately = true,
13251325
struct ggml_tensor** output = NULL,
13261326
struct ggml_context* output_ctx = NULL) {
1327-
alloc_compute_buffer(get_graph);
1327+
if (!alloc_compute_buffer(get_graph)) {
1328+
return false;
1329+
}
13281330
reset_compute_ctx();
13291331
struct ggml_cgraph* gf = get_graph();
13301332
GGML_ASSERT(ggml_gallocr_alloc_graph(compute_allocr, gf));
@@ -1382,6 +1384,7 @@ struct GGMLRunner {
13821384
if (free_compute_buffer_immediately) {
13831385
free_compute_buffer();
13841386
}
1387+
return true;
13851388
}
13861389
};
13871390

‎lora.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -835,11 +835,11 @@ struct LoraModel : public GGMLRunner {
835835
return gf;
836836
}
837837

838-
void apply(std::map<std::string, struct ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
838+
bool apply(std::map<std::string, struct ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
839839
auto get_graph = [&]() -> struct ggml_cgraph* {
840840
return build_lora_graph(model_tensors, version);
841841
};
842-
GGMLRunner::compute(get_graph, n_threads, true);
842+
return GGMLRunner::compute(get_graph, n_threads, true);
843843
}
844844
};
845845

‎mmdit.hpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -910,7 +910,7 @@ struct MMDiTRunner : public GGMLRunner {
910910
return gf;
911911
}
912912

913-
void compute(int n_threads,
913+
bool compute(int n_threads,
914914
struct ggml_tensor* x,
915915
struct ggml_tensor* timesteps,
916916
struct ggml_tensor* context,
@@ -926,7 +926,7 @@ struct MMDiTRunner : public GGMLRunner {
926926
return build_graph(x, timesteps, context, y, skip_layers);
927927
};
928928

929-
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
929+
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
930930
}
931931

932932
void test() {

0 commit comments

Comments
 (0)