feat: support custom upscale tile size (#896)

This commit is contained in:
Pedrito 2025-12-10 15:25:19 +01:00 committed by GitHub
parent d939f6e86a
commit 1ac5a616de
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 27 additions and 10 deletions

View File

@ -156,9 +156,10 @@ struct ESRGAN : public GGMLRunner {
ESRGAN(ggml_backend_t backend, ESRGAN(ggml_backend_t backend,
bool offload_params_to_cpu, bool offload_params_to_cpu,
int tile_size = 128,
const String2TensorStorage& tensor_storage_map = {}) const String2TensorStorage& tensor_storage_map = {})
: GGMLRunner(backend, offload_params_to_cpu) { : GGMLRunner(backend, offload_params_to_cpu) {
// rrdb_net will be created in load_from_file this->tile_size = tile_size;
} }
std::string get_desc() override { std::string get_desc() override {

View File

@ -1079,7 +1079,8 @@ struct SDGenerationParams {
std::string pm_id_embed_path; std::string pm_id_embed_path;
float pm_style_strength = 20.f; float pm_style_strength = 20.f;
int upscale_repeats = 1; int upscale_repeats = 1;
int upscale_tile_size = 128;
std::map<std::string, float> lora_map; std::map<std::string, float> lora_map;
std::map<std::string, float> high_noise_lora_map; std::map<std::string, float> high_noise_lora_map;
@ -1176,6 +1177,10 @@ struct SDGenerationParams {
"--upscale-repeats", "--upscale-repeats",
"Run the ESRGAN upscaler this many times (default: 1)", "Run the ESRGAN upscaler this many times (default: 1)",
&upscale_repeats}, &upscale_repeats},
{"",
"--upscale-tile-size",
"tile size for ESRGAN upscaling (default: 128)",
&upscale_tile_size},
}; };
options.float_options = { options.float_options = {
@ -1635,6 +1640,10 @@ struct SDGenerationParams {
return false; return false;
} }
if (upscale_tile_size < 1) {
return false;
}
if (mode == UPSCALE) { if (mode == UPSCALE) {
if (init_image_path.length() == 0) { if (init_image_path.length() == 0) {
fprintf(stderr, "error: upscale mode needs an init image (--init-img)\n"); fprintf(stderr, "error: upscale mode needs an init image (--init-img)\n");
@ -1720,6 +1729,7 @@ struct SDGenerationParams {
<< " control_strength: " << control_strength << ",\n" << " control_strength: " << control_strength << ",\n"
<< " seed: " << seed << ",\n" << " seed: " << seed << ",\n"
<< " upscale_repeats: " << upscale_repeats << ",\n" << " upscale_repeats: " << upscale_repeats << ",\n"
<< " upscale_tile_size: " << upscale_tile_size << ",\n"
<< "}"; << "}";
free(sample_params_str); free(sample_params_str);
free(high_noise_sample_params_str); free(high_noise_sample_params_str);
@ -2336,7 +2346,8 @@ int main(int argc, const char* argv[]) {
upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(ctx_params.esrgan_path.c_str(), upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(ctx_params.esrgan_path.c_str(),
ctx_params.offload_params_to_cpu, ctx_params.offload_params_to_cpu,
ctx_params.diffusion_conv_direct, ctx_params.diffusion_conv_direct,
ctx_params.n_threads); ctx_params.n_threads,
gen_params.upscale_tile_size);
if (upscaler_ctx == nullptr) { if (upscaler_ctx == nullptr) {
printf("new_upscaler_ctx failed\n"); printf("new_upscaler_ctx failed\n");

View File

@ -347,7 +347,8 @@ typedef struct upscaler_ctx_t upscaler_ctx_t;
SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path, SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
bool offload_params_to_cpu, bool offload_params_to_cpu,
bool direct, bool direct,
int n_threads); int n_threads,
int tile_size);
SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx); SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,

View File

@ -9,12 +9,15 @@ struct UpscalerGGML {
std::shared_ptr<ESRGAN> esrgan_upscaler; std::shared_ptr<ESRGAN> esrgan_upscaler;
std::string esrgan_path; std::string esrgan_path;
int n_threads; int n_threads;
bool direct = false; bool direct = false;
int tile_size = 128;
UpscalerGGML(int n_threads, UpscalerGGML(int n_threads,
bool direct = false) bool direct = false,
int tile_size = 128)
: n_threads(n_threads), : n_threads(n_threads),
direct(direct) { direct(direct),
tile_size(tile_size) {
} }
bool load_from_file(const std::string& esrgan_path, bool load_from_file(const std::string& esrgan_path,
@ -51,7 +54,7 @@ struct UpscalerGGML {
backend = ggml_backend_cpu_init(); backend = ggml_backend_cpu_init();
} }
LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, model_loader.get_tensor_storage_map()); esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
if (direct) { if (direct) {
esrgan_upscaler->set_conv2d_direct_enabled(true); esrgan_upscaler->set_conv2d_direct_enabled(true);
} }
@ -113,14 +116,15 @@ struct upscaler_ctx_t {
upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str, upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
bool offload_params_to_cpu, bool offload_params_to_cpu,
bool direct, bool direct,
int n_threads) { int n_threads,
int tile_size) {
upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t)); upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t));
if (upscaler_ctx == nullptr) { if (upscaler_ctx == nullptr) {
return nullptr; return nullptr;
} }
std::string esrgan_path(esrgan_path_c_str); std::string esrgan_path(esrgan_path_c_str);
upscaler_ctx->upscaler = new UpscalerGGML(n_threads, direct); upscaler_ctx->upscaler = new UpscalerGGML(n_threads, direct, tile_size);
if (upscaler_ctx->upscaler == nullptr) { if (upscaler_ctx->upscaler == nullptr) {
return nullptr; return nullptr;
} }