Skip to content

Commit 2bb0389

Browse files
leejetleejet
authored
refactor: return bool from image and upscale APIs (#1728)
1 parent ccda89e commit 2bb0389

7 files changed

Lines changed: 106 additions & 40 deletions

File tree

‎examples/cli/main.cpp‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -766,8 +766,12 @@ int main(int argc, const char* argv[]) {
766766
if (cli_params.mode == IMG_GEN) {
767767
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
768768

769-
num_results = gen_params.batch_count;
770-
results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
769+
sd_image_t* generated_images = nullptr;
770+
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
771+
generated_images = nullptr;
772+
num_results = 0;
773+
}
774+
results.adopt(generated_images, num_results);
771775
} else if (cli_params.mode == VID_GEN) {
772776
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
773777
sd_image_t* generated_video = nullptr;
@@ -802,15 +806,21 @@ int main(int argc, const char* argv[]) {
802806
SDImageOwner current_image(results[i]);
803807
results[i] = {0, 0, 0, nullptr};
804808
for (int u = 0; u < gen_params.upscale_repeats; ++u) {
805-
sd_image_t* upscaled_images = upscale(upscaler_ctx.get(), current_image.get(), upscale_factor);
806-
if (upscaled_images == nullptr || upscaled_images[0].data == nullptr) {
807-
free_sd_images(upscaled_images, 1);
809+
sd_image_t* upscaled_images = nullptr;
810+
int upscaled_count = 0;
811+
bool upscale_ok = upscale(upscaler_ctx.get(),
812+
current_image.get(),
813+
upscale_factor,
814+
&upscaled_images,
815+
&upscaled_count);
816+
if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) {
817+
free_sd_images(upscaled_images, upscaled_count);
808818
LOG_ERROR("upscale failed");
809819
break;
810820
}
811821
sd_image_t upscaled_image = upscaled_images[0];
812822
upscaled_images[0] = {0, 0, 0, nullptr};
813-
free_sd_images(upscaled_images, 1);
823+
free_sd_images(upscaled_images, upscaled_count);
814824
current_image.reset(upscaled_image);
815825
}
816826
results[i] = current_image.release(); // Set the final upscaled image as the result

‎examples/server/async_jobs.cpp‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,8 +173,13 @@ bool execute_img_gen_job(ServerRuntime& runtime,
173173

174174
{
175175
std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
176-
sd_image_t* raw_results = generate_image(runtime.sd_ctx, &params);
177-
results.adopt(raw_results, params.batch_count);
176+
sd_image_t* raw_results = nullptr;
177+
int num_results = 0;
178+
if (!generate_image(runtime.sd_ctx, &params, &raw_results, &num_results)) {
179+
raw_results = nullptr;
180+
num_results = 0;
181+
}
182+
results.adopt(raw_results, num_results);
178183
}
179184

180185
const int num_results = results.count();

‎examples/server/routes_openai.cpp‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -229,8 +229,11 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime,
229229

230230
{
231231
std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
232-
sd_image_t* raw_results = generate_image(runtime.sd_ctx, &img_gen_params);
233-
num_results = request.gen_params.batch_count;
232+
sd_image_t* raw_results = nullptr;
233+
if (!generate_image(runtime.sd_ctx, &img_gen_params, &raw_results, &num_results)) {
234+
raw_results = nullptr;
235+
num_results = 0;
236+
}
234237
results.adopt(raw_results, num_results);
235238
}
236239

‎examples/server/routes_sdapi.cpp‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -292,8 +292,11 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
292292

293293
{
294294
std::lock_guard<std::mutex> lock(*runtime->sd_ctx_mutex);
295-
sd_image_t* raw_results = generate_image(runtime->sd_ctx, &img_gen_params);
296-
num_results = request.gen_params.batch_count;
295+
sd_image_t* raw_results = nullptr;
296+
if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) {
297+
raw_results = nullptr;
298+
num_results = 0;
299+
}
297300
results.adopt(raw_results, num_results);
298301
}
299302

‎include/stable-diffusion.h‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -454,7 +454,10 @@ SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sa
454454

455455
SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params);
456456
SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params);
457-
SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
457+
SD_API bool generate_image(sd_ctx_t* sd_ctx,
458+
const sd_img_gen_params_t* sd_img_gen_params,
459+
sd_image_t** images_out,
460+
int* num_images_out);
458461

459462
enum sd_cancel_mode_t {
460463
// Stop the current generation as soon as possible.
@@ -484,9 +487,11 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
484487
const char* params_backend);
485488
SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
486489

487-
SD_API sd_image_t* upscale(upscaler_ctx_t* upscaler_ctx,
488-
sd_image_t input_image,
489-
uint32_t upscale_factor);
490+
SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
491+
sd_image_t input_image,
492+
uint32_t upscale_factor,
493+
sd_image_t** images_out,
494+
int* num_images_out);
490495

491496
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
492497

‎src/stable-diffusion.cpp‎

Lines changed: 41 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -4278,7 +4278,8 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
42784278

42794279
static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
42804280
const GenerationRequest& request,
4281-
const std::vector<sd::Tensor<float>>& final_latents) {
4281+
const std::vector<sd::Tensor<float>>& final_latents,
4282+
int* num_images_out) {
42824283
if (final_latents.empty()) {
42834284
LOG_ERROR("no latent images to decode");
42844285
return nullptr;
@@ -4320,11 +4321,14 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
43204321
return nullptr;
43214322
}
43224323

4323-
sd_image_t* result_images = (sd_image_t*)calloc(request.batch_count, sizeof(sd_image_t));
4324+
int image_count = static_cast<int>(decoded_images.size());
4325+
sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t));
43244326
if (result_images == nullptr) {
43254327
return nullptr;
43264328
}
4327-
memset(result_images, 0, request.batch_count * sizeof(sd_image_t));
4329+
if (num_images_out != nullptr) {
4330+
*num_images_out = image_count;
4331+
}
43284332

43294333
for (size_t i = 0; i < decoded_images.size(); i++) {
43304334
result_images[i] = tensor_to_sd_image(decoded_images[i]);
@@ -4517,9 +4521,18 @@ static std::vector<float> make_hires_sigma_schedule(sd_ctx_t* sd_ctx,
45174521
sigmas.end());
45184522
}
45194523

4520-
SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
4524+
SD_API bool generate_image(sd_ctx_t* sd_ctx,
4525+
const sd_img_gen_params_t* sd_img_gen_params,
4526+
sd_image_t** images_out,
4527+
int* num_images_out) {
4528+
if (images_out != nullptr) {
4529+
*images_out = nullptr;
4530+
}
4531+
if (num_images_out != nullptr) {
4532+
*num_images_out = 0;
4533+
}
45214534
if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
4522-
return nullptr;
4535+
return false;
45234536
}
45244537

45254538
sd_ctx->sd->reset_cancel_flag();
@@ -4542,7 +4555,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
45424555
&request,
45434556
&plan);
45444557
if (!latents_opt.has_value()) {
4545-
return nullptr;
4558+
return false;
45464559
}
45474560
ImageGenerationLatents latents = std::move(*latents_opt);
45484561

@@ -4552,7 +4565,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
45524565
&plan,
45534566
&latents);
45544567
if (!embeds_opt.has_value()) {
4555-
return nullptr;
4568+
return false;
45564569
}
45574570
ImageGenerationEmbeds embeds = std::move(*embeds_opt);
45584571

@@ -4562,7 +4575,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
45624575
sd_cancel_mode_t cancel = sd_ctx->sd->get_cancel_flag();
45634576
if (cancel == SD_CANCEL_ALL) {
45644577
LOG_ERROR("cancelling generation");
4565-
return nullptr;
4578+
return false;
45664579
}
45674580
if (cancel == SD_CANCEL_NEW_LATENTS) {
45684581
LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents",
@@ -4614,29 +4627,29 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
46144627
b + 1,
46154628
request.batch_count,
46164629
(sampling_end - sampling_start) * 1.0f / 1000);
4617-
return nullptr;
4630+
return false;
46184631
}
46194632
int64_t denoise_end = ggml_time_ms();
46204633
LOG_INFO("generating %zu latent images completed, taking %.2fs",
46214634
final_latents.size(),
46224635
(denoise_end - denoise_start) * 1.0f / 1000);
46234636
if (final_latents.empty()) {
46244637
LOG_ERROR("no latent images generated");
4625-
return nullptr;
4638+
return false;
46264639
}
46274640

46284641
if (request.hires.enabled && request.hires.target_width > 0) {
46294642
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
46304643
LOG_ERROR("cancelling generation before hires fix");
4631-
return nullptr;
4644+
return false;
46324645
}
46334646
LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
46344647

46354648
std::unique_ptr<UpscalerGGML> hires_upscaler;
46364649
if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
46374650
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
46384651
LOG_ERROR("cancelling generation before hires model load");
4639-
return nullptr;
4652+
return false;
46404653
}
46414654
LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
46424655
hires_upscaler = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
@@ -4649,7 +4662,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
46494662
if (!hires_upscaler->load_from_file(request.hires.model_path,
46504663
sd_ctx->sd->n_threads)) {
46514664
LOG_ERROR("load hires model upscaler failed");
4652-
return nullptr;
4665+
return false;
46534666
}
46544667
}
46554668

@@ -4673,7 +4686,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
46734686
for (int b = 0; b < (int)final_latents.size(); b++) {
46744687
if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
46754688
LOG_ERROR("cancelling generation during hires fix");
4676-
return nullptr;
4689+
return false;
46774690
}
46784691
int64_t cur_seed = request.seed + b;
46794692
sd_ctx->sd->rng->manual_seed(cur_seed);
@@ -4684,7 +4697,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
46844697
request,
46854698
hires_upscaler.get());
46864699
if (upscaled.empty()) {
4687-
return nullptr;
4700+
return false;
46884701
}
46894702

46904703
sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
@@ -4738,24 +4751,33 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
47384751
b + 1,
47394752
(int)final_latents.size(),
47404753
(hires_sample_end - hires_sample_start) * 1.0f / 1000);
4741-
return nullptr;
4754+
return false;
47424755
}
47434756
int64_t hires_denoise_end = ggml_time_ms();
47444757
LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
47454758

47464759
final_latents = std::move(hires_final_latents);
47474760
}
47484761

4749-
auto result = decode_image_outputs(sd_ctx, request, final_latents);
4762+
int num_images = 0;
4763+
auto result = decode_image_outputs(sd_ctx, request, final_latents, &num_images);
47504764
if (result == nullptr) {
4751-
return nullptr;
4765+
return false;
47524766
}
47534767

47544768
sd_ctx->sd->lora_stat();
47554769

47564770
int64_t t1 = ggml_time_ms();
47574771
LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000);
4758-
return result;
4772+
if (num_images_out != nullptr) {
4773+
*num_images_out = num_images;
4774+
}
4775+
if (images_out != nullptr) {
4776+
*images_out = result;
4777+
} else {
4778+
free_sd_images(result, num_images);
4779+
}
4780+
return true;
47594781
}
47604782

47614783
static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd_ctx_t* sd_ctx,

‎src/upscaler.cpp‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -199,23 +199,41 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
199199
return upscaler_ctx;
200200
}
201201

202-
sd_image_t* upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor) {
202+
bool upscale(upscaler_ctx_t* upscaler_ctx,
203+
sd_image_t input_image,
204+
uint32_t upscale_factor,
205+
sd_image_t** images_out,
206+
int* num_images_out) {
207+
if (images_out != nullptr) {
208+
*images_out = nullptr;
209+
}
210+
if (num_images_out != nullptr) {
211+
*num_images_out = 0;
212+
}
203213
if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) {
204-
return nullptr;
214+
return false;
205215
}
206216

207217
sd_image_t* result_images = (sd_image_t*)calloc(1, sizeof(sd_image_t));
208218
if (result_images == nullptr) {
209-
return nullptr;
219+
return false;
210220
}
211221

212222
result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
213223
if (result_images[0].data == nullptr) {
214224
free(result_images);
215-
return nullptr;
225+
return false;
216226
}
217227

218-
return result_images;
228+
if (num_images_out != nullptr) {
229+
*num_images_out = 1;
230+
}
231+
if (images_out != nullptr) {
232+
*images_out = result_images;
233+
} else {
234+
free_sd_images(result_images, 1);
235+
}
236+
return true;
219237
}
220238

221239
int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {

0 commit comments

Comments
 (0)