feat: add beta scheduler (#811)

Co-authored-by: phil2sat <phil2sat@users.noreply.github.com>

examples/cli/main.cpp

@@ -766,8 +766,12 @@ int main(int argc, const char* argv[]) {
         if (cli_params.mode == IMG_GEN) {
             sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
 
-            num_results = gen_params.batch_count;
+            sd_image_t* generated_images = nullptr;
-            results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
+            if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
+                generated_images = nullptr;
+                num_results      = 0;
+            }
+            results.adopt(generated_images, num_results);
         } else if (cli_params.mode == VID_GEN) {
             sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
             sd_image_t* generated_video        = nullptr;
@@ -802,12 +806,22 @@ int main(int argc, const char* argv[]) {
                 SDImageOwner current_image(results[i]);
                 results[i] = {0, 0, 0, nullptr};
                 for (int u = 0; u < gen_params.upscale_repeats; ++u) {
-                    SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor));
+                    sd_image_t* upscaled_images = nullptr;
-                    if (upscaled_image.get().data == nullptr) {
+                    int upscaled_count          = 0;
+                    bool upscale_ok             = upscale(upscaler_ctx.get(),
+                                                          current_image.get(),
+                                                          upscale_factor,
+                                                          &upscaled_images,
+                                                          &upscaled_count);
+                    if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) {
+                        free_sd_images(upscaled_images, upscaled_count);
                         LOG_ERROR("upscale failed");
                         break;
                     }
-                    current_image = std::move(upscaled_image);
+                    sd_image_t upscaled_image = upscaled_images[0];
+                    upscaled_images[0]        = {0, 0, 0, nullptr};
+                    free_sd_images(upscaled_images, upscaled_count);
+                    current_image.reset(upscaled_image);
                 }
                 results[i] = current_image.release();  // Set the final upscaled image as the result
             }

examples/common/common.cpp

@@ -1475,7 +1475,7 @@ ArgOptions SDGenerationParams::get_options() {
          on_high_noise_sample_method_arg},
         {"",
          "--scheduler",
-         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux], alias: normal=discrete, default: model-specific",
+         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific",
          on_scheduler_arg},
         {"",
          "--sigmas",

examples/server/async_jobs.cpp

@@ -173,8 +173,13 @@ bool execute_img_gen_job(ServerRuntime& runtime,
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
-        sd_image_t* raw_results = generate_image(runtime.sd_ctx, &params);
+        sd_image_t* raw_results = nullptr;
-        results.adopt(raw_results, params.batch_count);
+        int num_results         = 0;
+        if (!generate_image(runtime.sd_ctx, &params, &raw_results, &num_results)) {
+            raw_results = nullptr;
+            num_results = 0;
+        }
+        results.adopt(raw_results, num_results);
     }
 
     const int num_results = results.count();

examples/server/routes_openai.cpp

@@ -229,8 +229,11 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime,
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
-        sd_image_t* raw_results = generate_image(runtime.sd_ctx, &img_gen_params);
+        sd_image_t* raw_results = nullptr;
-        num_results             = request.gen_params.batch_count;
+        if (!generate_image(runtime.sd_ctx, &img_gen_params, &raw_results, &num_results)) {
+            raw_results = nullptr;
+            num_results = 0;
+        }
         results.adopt(raw_results, num_results);
     }
 

examples/server/routes_sdapi.cpp

@@ -292,8 +292,11 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
 
             {
                 std::lock_guard<std::mutex> lock(*runtime->sd_ctx_mutex);
-                sd_image_t* raw_results = generate_image(runtime->sd_ctx, &img_gen_params);
+                sd_image_t* raw_results = nullptr;
-                num_results             = request.gen_params.batch_count;
+                if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) {
+                    raw_results = nullptr;
+                    num_results = 0;
+                }
                 results.adopt(raw_results, num_results);
             }
 

include/stable-diffusion.h

@@ -73,6 +73,7 @@ enum scheduler_t {
     LOGIT_NORMAL_SCHEDULER,
     FLUX2_SCHEDULER,
     FLUX_SCHEDULER,
+    BETA_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -454,7 +455,10 @@ SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sa
 
 SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params);
 SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params);
-SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
+SD_API bool generate_image(sd_ctx_t* sd_ctx,
+                           const sd_img_gen_params_t* sd_img_gen_params,
+                           sd_image_t** images_out,
+                           int* num_images_out);
 
 enum sd_cancel_mode_t {
     // Stop the current generation as soon as possible.
@@ -484,9 +488,11 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
                                         const char* params_backend);
 SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
-SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,
+SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
-                          sd_image_t input_image,
+                    sd_image_t input_image,
-                          uint32_t upscale_factor);
+                    uint32_t upscale_factor,
+                    sd_image_t** images_out,
+                    int* num_images_out);
 
 SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
 

script/remove_utf8_bom.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""Remove UTF-8 BOMs from files under a directory.
+
+By default this scans the current working directory recursively and skips
+repository areas that should not be touched by ordinary maintenance scripts.
+Only files whose first three bytes are the UTF-8 BOM are rewritten.
+"""
+
+import argparse
+import os
+import shutil
+import sys
+import tempfile
+from pathlib import Path
+
+
+UTF8_BOM = b"\xef\xbb\xbf"
+
+DEFAULT_EXCLUDED_DIR_NAMES = {
+    ".git",
+    ".hg",
+    ".svn",
+    ".mypy_cache",
+    ".pytest_cache",
+    "__pycache__",
+    "test",
+}
+
+DEFAULT_EXCLUDED_DIR_PREFIXES = {
+    "build",
+}
+
+DEFAULT_EXCLUDED_REL_DIRS = {
+    "examples/server/frontend",
+    "ggml",
+    "models",
+    "src/vocab",
+    "thirdparty",
+}
+
+
+def rel_posix(path: Path, root: Path) -> str:
+    try:
+        return path.relative_to(root).as_posix()
+    except ValueError:
+        return path.as_posix()
+
+
+def should_skip_dir(
+    path: Path,
+    root: Path,
+    excluded_rel_dirs: set[str],
+    excluded_names: set[str],
+    excluded_prefixes: set[str],
+) -> bool:
+    rel = rel_posix(path, root)
+    return (
+        path.name in excluded_names
+        or rel in excluded_rel_dirs
+        or any(path.name.startswith(prefix) for prefix in excluded_prefixes)
+    )
+
+
+def iter_files(
+    root: Path,
+    recursive: bool,
+    excluded_rel_dirs: set[str],
+    excluded_names: set[str],
+    excluded_prefixes: set[str],
+    follow_symlinks: bool,
+):
+    if recursive:
+        for dirpath, dirnames, filenames in os.walk(root, followlinks=follow_symlinks):
+            current_dir = Path(dirpath)
+            dirnames[:] = [
+                name
+                for name in dirnames
+                if not should_skip_dir(
+                    current_dir / name,
+                    root,
+                    excluded_rel_dirs,
+                    excluded_names,
+                    excluded_prefixes,
+                )
+            ]
+            for filename in filenames:
+                path = current_dir / filename
+                if path.is_symlink() and not follow_symlinks:
+                    continue
+                yield path
+    else:
+        for path in root.iterdir():
+            if path.is_file() and (follow_symlinks or not path.is_symlink()):
+                yield path
+
+
+def has_utf8_bom(path: Path) -> bool:
+    with path.open("rb") as f:
+        return f.read(len(UTF8_BOM)) == UTF8_BOM
+
+
+def strip_utf8_bom(path: Path) -> None:
+    tmp_path = None
+    try:
+        with path.open("rb") as src:
+            if src.read(len(UTF8_BOM)) != UTF8_BOM:
+                return
+
+            fd, tmp_name = tempfile.mkstemp(
+                prefix=f".{path.name}.",
+                suffix=".tmp",
+                dir=str(path.parent),
+            )
+            tmp_path = Path(tmp_name)
+            with os.fdopen(fd, "wb") as dst:
+                shutil.copyfileobj(src, dst, length=1024 * 1024)
+
+        shutil.copystat(path, tmp_path, follow_symlinks=False)
+        os.replace(tmp_path, path)
+        tmp_path = None
+    finally:
+        if tmp_path is not None:
+            try:
+                tmp_path.unlink()
+            except FileNotFoundError:
+                pass
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Scan files and convert UTF-8 BOM files to UTF-8 without BOM.",
+    )
+    parser.add_argument(
+        "root",
+        nargs="?",
+        default=".",
+        help="Directory to scan. Defaults to the current directory.",
+    )
+    parser.add_argument(
+        "-n",
+        "--dry-run",
+        action="store_true",
+        help="Only list files that would be converted.",
+    )
+    parser.add_argument(
+        "--no-recursive",
+        action="store_true",
+        help="Only scan files directly under root.",
+    )
+    parser.add_argument(
+        "--include-repo-excluded",
+        action="store_true",
+        help="Do not skip default repository excluded directories.",
+    )
+    parser.add_argument(
+        "--exclude-dir",
+        action="append",
+        default=[],
+        metavar="DIR",
+        help="Additional directory name or root-relative path to skip. Can be used multiple times.",
+    )
+    parser.add_argument(
+        "--follow-symlinks",
+        action="store_true",
+        help="Follow symlinked directories and files.",
+    )
+    parser.add_argument(
+        "-q",
+        "--quiet",
+        action="store_true",
+        help="Only print the final summary.",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    root = Path(args.root).resolve()
+
+    if not root.is_dir():
+        print(f"error: not a directory: {root}", file=sys.stderr)
+        return 2
+
+    excluded_names = set()
+    excluded_rel_dirs = set()
+    excluded_prefixes = set()
+    if not args.include_repo_excluded:
+        excluded_names.update(DEFAULT_EXCLUDED_DIR_NAMES)
+        excluded_rel_dirs.update(DEFAULT_EXCLUDED_REL_DIRS)
+        excluded_prefixes.update(DEFAULT_EXCLUDED_DIR_PREFIXES)
+
+    for item in args.exclude_dir:
+        normalized = Path(item).as_posix().strip("/")
+        if "/" in normalized:
+            excluded_rel_dirs.add(normalized)
+        else:
+            excluded_names.add(normalized)
+
+    scanned = 0
+    converted = 0
+    errors = 0
+
+    for path in iter_files(
+        root=root,
+        recursive=not args.no_recursive,
+        excluded_rel_dirs=excluded_rel_dirs,
+        excluded_names=excluded_names,
+        excluded_prefixes=excluded_prefixes,
+        follow_symlinks=args.follow_symlinks,
+    ):
+        scanned += 1
+        try:
+            if not has_utf8_bom(path):
+                continue
+            converted += 1
+            rel = rel_posix(path, root)
+            if args.dry_run:
+                if not args.quiet:
+                    print(f"would convert: {rel}")
+            else:
+                strip_utf8_bom(path)
+                if not args.quiet:
+                    print(f"converted: {rel}")
+        except OSError as exc:
+            errors += 1
+            print(f"error: {rel_posix(path, root)}: {exc}", file=sys.stderr)
+
+    action = "would convert" if args.dry_run else "converted"
+    print(f"scanned {scanned} file(s), {action} {converted}, errors {errors}")
+    return 1 if errors else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

src/conditioning/conditioner.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
+#ifndef __SD_CONDITIONING_CONDITIONER_HPP__
 #define __SD_CONDITIONING_CONDITIONER_HPP__
 
 #include <cmath>

src/model/diffusion/control.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
 #define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
 #include "model/common/block.hpp"

src/model/diffusion/hidream_o1.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
+#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 #define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 
 #include <algorithm>

src/model/diffusion/model.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
 #define __SD_MODEL_DIFFUSION_MODEL_HPP__
 
 #include <string>

src/model/te/clip.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_TE_CLIP_HPP__
+#ifndef __SD_MODEL_TE_CLIP_HPP__
 #define __SD_MODEL_TE_CLIP_HPP__
 
 #include "core/ggml_extend.hpp"

src/model/te/llm.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_TE_LLM_HPP__
+#ifndef __SD_MODEL_TE_LLM_HPP__
 #define __SD_MODEL_TE_LLM_HPP__
 
 #include <algorithm>

src/model/te/t5.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_TE_T5_HPP__
+#ifndef __SD_MODEL_TE_T5_HPP__
 #define __SD_MODEL_TE_T5_HPP__
 
 #include <cfloat>

src/model/upscaler/esrgan.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
+#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
 #define __SD_MODEL_UPSCALER_ESRGAN_HPP__
 
 #include <algorithm>

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
+#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
 #define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
 
 #include <algorithm>

src/model/vae/ltx_audio_vae.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
+#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 #define __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 
 #include <cmath>

src/model/vae/vae.hpp

@@ -1,4 +1,4 @@
-#ifndef __SD_MODEL_VAE_VAE_HPP__
+#ifndef __SD_MODEL_VAE_VAE_HPP__
 #define __SD_MODEL_VAE_VAE_HPP__
 
 #include "core/tensor_ggml.hpp"

src/runtime/denoiser.hpp

@@ -302,6 +302,137 @@ struct KarrasScheduler : SigmaScheduler {
     }
 };
 
+struct BetaScheduler : SigmaScheduler {
+    static constexpr double alpha = 0.6;
+    static constexpr double beta  = 0.6;
+
+    static double log_beta(double a, double b) {
+        return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
+    }
+
+    static double incbeta(double x, double a, double b) {
+        if (x <= 0.0) {
+            return 0.0;
+        }
+        if (x >= 1.0) {
+            return 1.0;
+        }
+
+        // Continued fraction approximation using Lentz's method.
+        const int max_iter   = 200;
+        const double epsilon = 3.0e-7;
+        const double tiny    = 1e-30;
+
+        const double qab = a + b;
+        const double qap = a + 1.0;
+        const double qam = a - 1.0;
+
+        double c = 1.0;
+        double d = 1.0 - qab * x / qap;
+        if (std::abs(d) < tiny) {
+            d = tiny;
+        }
+        d        = 1.0 / d;
+        double h = d;
+
+        for (int m = 1; m <= max_iter; m++) {
+            const int m2 = 2 * m;
+
+            double aa = m * (b - m) * x / ((qam + m2) * (a + m2));
+            d         = 1.0 + aa * d;
+            if (std::abs(d) < tiny) {
+                d = tiny;
+            }
+            c = 1.0 + aa / c;
+            if (std::abs(c) < tiny) {
+                c = tiny;
+            }
+            d = 1.0 / d;
+            h *= d * c;
+
+            aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2));
+            d  = 1.0 + aa * d;
+            if (std::abs(d) < tiny) {
+                d = tiny;
+            }
+            c = 1.0 + aa / c;
+            if (std::abs(c) < tiny) {
+                c = tiny;
+            }
+            d                = 1.0 / d;
+            const double del = d * c;
+            h *= del;
+
+            if (std::abs(del - 1.0) < epsilon) {
+                break;
+            }
+        }
+
+        return std::exp(a * std::log(x) + b * std::log(1.0 - x) - log_beta(a, b)) / a * h;
+    }
+
+    static double beta_cdf(double x, double a, double b) {
+        if (x == 0.0) {
+            return 0.0;
+        }
+        if (x == 1.0) {
+            return 1.0;
+        }
+        if (x < (a + 1.0) / (a + b + 2.0)) {
+            return incbeta(x, a, b);
+        }
+        return 1.0 - incbeta(1.0 - x, b, a);
+    }
+
+    static double beta_ppf(double u, double a, double b, int max_iter = 30) {
+        double x = 0.5;
+        for (int i = 0; i < max_iter; i++) {
+            const double f = beta_cdf(x, a, b) - u;
+            if (std::abs(f) < 1e-10) {
+                break;
+            }
+            const double df = std::exp((a - 1.0) * std::log(x) + (b - 1.0) * std::log(1.0 - x) - log_beta(a, b));
+            x -= f / df;
+            if (x <= 0.0) {
+                x = 1e-10;
+            }
+            if (x >= 1.0) {
+                x = 1.0 - 1e-10;
+            }
+        }
+        return x;
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t t_to_sigma) override {
+        std::vector<float> result;
+        result.reserve(n + 1);
+
+        const int t_max = TIMESTEPS - 1;
+        if (n == 0) {
+            return result;
+        } else if (n == 1) {
+            result.push_back(t_to_sigma(static_cast<float>(t_max)));
+            result.push_back(0.f);
+            return result;
+        }
+
+        int last_t = -1;
+        for (uint32_t i = 0; i < n; i++) {
+            const double u      = 1.0 - static_cast<double>(i) / static_cast<double>(n);
+            const double t_cont = beta_ppf(u, alpha, beta) * t_max;
+            const int t         = static_cast<int>(std::lround(t_cont));
+
+            if (t != last_t) {
+                result.push_back(t_to_sigma(static_cast<float>(t)));
+                last_t = t;
+            }
+        }
+
+        result.push_back(0.f);
+        return result;
+    }
+};
+
 struct SimpleScheduler : SigmaScheduler {
     std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
         std::vector<float> result_sigmas;
@@ -895,6 +1026,10 @@ struct Denoiser {
                 LOG_INFO("get_sigmas with Karras scheduler");
                 scheduler = std::make_shared<KarrasScheduler>();
                 break;
+            case BETA_SCHEDULER:
+                LOG_INFO("get_sigmas with Beta scheduler");
+                scheduler = std::make_shared<BetaScheduler>();
+                break;
             case EXPONENTIAL_SCHEDULER:
                 LOG_INFO("get_sigmas exponential scheduler");
                 scheduler = std::make_shared<ExponentialScheduler>();

src/stable-diffusion.cpp

@@ -2562,6 +2562,7 @@ const char* scheduler_to_str[] = {
     "logit_normal",
     "flux2",
     "flux",
+    "beta",
 };
 
 const char* sd_scheduler_name(enum scheduler_t scheduler) {
@@ -4278,7 +4279,8 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
 
 static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
                                         const GenerationRequest& request,
-                                        const std::vector<sd::Tensor<float>>& final_latents) {
+                                        const std::vector<sd::Tensor<float>>& final_latents,
+                                        int* num_images_out) {
     if (final_latents.empty()) {
         LOG_ERROR("no latent images to decode");
         return nullptr;
@@ -4320,11 +4322,14 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
         return nullptr;
     }
 
-    sd_image_t* result_images = (sd_image_t*)calloc(request.batch_count, sizeof(sd_image_t));
+    int image_count           = static_cast<int>(decoded_images.size());
+    sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t));
     if (result_images == nullptr) {
         return nullptr;
     }
-    memset(result_images, 0, request.batch_count * sizeof(sd_image_t));
+    if (num_images_out != nullptr) {
+        *num_images_out = image_count;
+    }
 
     for (size_t i = 0; i < decoded_images.size(); i++) {
         result_images[i] = tensor_to_sd_image(decoded_images[i]);
@@ -4517,9 +4522,18 @@ static std::vector<float> make_hires_sigma_schedule(sd_ctx_t* sd_ctx,
                               sigmas.end());
 }
 
-SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) {
+SD_API bool generate_image(sd_ctx_t* sd_ctx,
+                           const sd_img_gen_params_t* sd_img_gen_params,
+                           sd_image_t** images_out,
+                           int* num_images_out) {
+    if (images_out != nullptr) {
+        *images_out = nullptr;
+    }
+    if (num_images_out != nullptr) {
+        *num_images_out = 0;
+    }
     if (sd_ctx == nullptr || sd_img_gen_params == nullptr) {
-        return nullptr;
+        return false;
     }
 
     sd_ctx->sd->reset_cancel_flag();
@@ -4542,7 +4556,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                         &request,
                                                         &plan);
     if (!latents_opt.has_value()) {
-        return nullptr;
+        return false;
     }
     ImageGenerationLatents latents = std::move(*latents_opt);
 
@@ -4552,7 +4566,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                       &plan,
                                                       &latents);
     if (!embeds_opt.has_value()) {
-        return nullptr;
+        return false;
     }
     ImageGenerationEmbeds embeds = std::move(*embeds_opt);
 
@@ -4562,7 +4576,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         sd_cancel_mode_t cancel = sd_ctx->sd->get_cancel_flag();
         if (cancel == SD_CANCEL_ALL) {
             LOG_ERROR("cancelling generation");
-            return nullptr;
+            return false;
         }
         if (cancel == SD_CANCEL_NEW_LATENTS) {
             LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents",
@@ -4614,7 +4628,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                   b + 1,
                   request.batch_count,
                   (sampling_end - sampling_start) * 1.0f / 1000);
-        return nullptr;
+        return false;
     }
     int64_t denoise_end = ggml_time_ms();
     LOG_INFO("generating %zu latent images completed, taking %.2fs",
@@ -4622,13 +4636,13 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
              (denoise_end - denoise_start) * 1.0f / 1000);
     if (final_latents.empty()) {
         LOG_ERROR("no latent images generated");
-        return nullptr;
+        return false;
     }
 
     if (request.hires.enabled && request.hires.target_width > 0) {
         if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
             LOG_ERROR("cancelling generation before hires fix");
-            return nullptr;
+            return false;
         }
         LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height);
 
@@ -4636,7 +4650,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) {
             if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
                 LOG_ERROR("cancelling generation before hires model load");
-                return nullptr;
+                return false;
             }
             LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path);
             hires_upscaler                    = std::make_unique<UpscalerGGML>(sd_ctx->sd->n_threads,
@@ -4649,7 +4663,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
             if (!hires_upscaler->load_from_file(request.hires.model_path,
                                                 sd_ctx->sd->n_threads)) {
                 LOG_ERROR("load hires model upscaler failed");
-                return nullptr;
+                return false;
             }
         }
 
@@ -4673,7 +4687,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         for (int b = 0; b < (int)final_latents.size(); b++) {
             if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
                 LOG_ERROR("cancelling generation during hires fix");
-                return nullptr;
+                return false;
             }
             int64_t cur_seed = request.seed + b;
             sd_ctx->sd->rng->manual_seed(cur_seed);
@@ -4684,7 +4698,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                               request,
                                                               hires_upscaler.get());
             if (upscaled.empty()) {
-                return nullptr;
+                return false;
             }
 
             sd::Tensor<float> noise = sd::randn_like<float>(upscaled, sd_ctx->sd->rng);
@@ -4738,7 +4752,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                       b + 1,
                       (int)final_latents.size(),
                       (hires_sample_end - hires_sample_start) * 1.0f / 1000);
-            return nullptr;
+            return false;
         }
         int64_t hires_denoise_end = ggml_time_ms();
         LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000);
@@ -4746,16 +4760,25 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
         final_latents = std::move(hires_final_latents);
     }
 
-    auto result = decode_image_outputs(sd_ctx, request, final_latents);
+    int num_images = 0;
+    auto result    = decode_image_outputs(sd_ctx, request, final_latents, &num_images);
     if (result == nullptr) {
-        return nullptr;
+        return false;
     }
 
     sd_ctx->sd->lora_stat();
 
     int64_t t1 = ggml_time_ms();
     LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000);
-    return result;
+    if (num_images_out != nullptr) {
+        *num_images_out = num_images;
+    }
+    if (images_out != nullptr) {
+        *images_out = result;
+    } else {
+        free_sd_images(result, num_images);
+    }
+    return true;
 }
 
 static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd_ctx_t* sd_ctx,

src/upscaler.cpp

@@ -4,6 +4,7 @@
 #include "model_loader.h"
 #include "stable-diffusion.h"
 
+#include <cstdlib>
 #include <utility>
 
 UpscalerGGML::UpscalerGGML(int n_threads,
@@ -198,8 +199,41 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
     return upscaler_ctx;
 }
 
-sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor) {
+bool upscale(upscaler_ctx_t* upscaler_ctx,
-    return upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
+             sd_image_t input_image,
+             uint32_t upscale_factor,
+             sd_image_t** images_out,
+             int* num_images_out) {
+    if (images_out != nullptr) {
+        *images_out = nullptr;
+    }
+    if (num_images_out != nullptr) {
+        *num_images_out = 0;
+    }
+    if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) {
+        return false;
+    }
+
+    sd_image_t* result_images = (sd_image_t*)calloc(1, sizeof(sd_image_t));
+    if (result_images == nullptr) {
+        return false;
+    }
+
+    result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor);
+    if (result_images[0].data == nullptr) {
+        free(result_images);
+        return false;
+    }
+
+    if (num_images_out != nullptr) {
+        *num_images_out = 1;
+    }
+    if (images_out != nullptr) {
+        *images_out = result_images;
+    } else {
+        free_sd_images(result_images, 1);
+    }
+    return true;
 }
 
 int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {