diff --git a/esrgan.hpp b/esrgan.hpp index 4cac956..961e84f 100644 --- a/esrgan.hpp +++ b/esrgan.hpp @@ -156,9 +156,10 @@ struct ESRGAN : public GGMLRunner { ESRGAN(ggml_backend_t backend, bool offload_params_to_cpu, + int tile_size = 128, const String2TensorStorage& tensor_storage_map = {}) : 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 { diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index c55da35..49b202f 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -1079,7 +1079,8 @@ struct SDGenerationParams { std::string pm_id_embed_path; float pm_style_strength = 20.f; - int upscale_repeats = 1; + int upscale_repeats = 1; + int upscale_tile_size = 128; std::map lora_map; std::map high_noise_lora_map; @@ -1176,6 +1177,10 @@ struct SDGenerationParams { "--upscale-repeats", "Run the ESRGAN upscaler this many times (default: 1)", &upscale_repeats}, + {"", + "--upscale-tile-size", + "tile size for ESRGAN upscaling (default: 128)", + &upscale_tile_size}, }; options.float_options = { @@ -1635,6 +1640,10 @@ struct SDGenerationParams { return false; } + if (upscale_tile_size < 1) { + return false; + } + if (mode == UPSCALE) { if (init_image_path.length() == 0) { fprintf(stderr, "error: upscale mode needs an init image (--init-img)\n"); @@ -1720,6 +1729,7 @@ struct SDGenerationParams { << " control_strength: " << control_strength << ",\n" << " seed: " << seed << ",\n" << " upscale_repeats: " << upscale_repeats << ",\n" + << " upscale_tile_size: " << upscale_tile_size << ",\n" << "}"; free(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(), ctx_params.offload_params_to_cpu, ctx_params.diffusion_conv_direct, - ctx_params.n_threads); + ctx_params.n_threads, + gen_params.upscale_tile_size); if (upscaler_ctx == nullptr) { printf("new_upscaler_ctx failed\n"); diff --git a/stable-diffusion.h b/stable-diffusion.h index 601b79f..2da70bd 100644 --- a/stable-diffusion.h +++ b/stable-diffusion.h @@ -347,7 +347,8 @@ typedef struct upscaler_ctx_t upscaler_ctx_t; SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path, bool offload_params_to_cpu, bool direct, - int n_threads); + int n_threads, + int tile_size); SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx); SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, diff --git a/upscaler.cpp b/upscaler.cpp index 62c0d29..29ac981 100644 --- a/upscaler.cpp +++ b/upscaler.cpp @@ -9,12 +9,15 @@ struct UpscalerGGML { std::shared_ptr esrgan_upscaler; std::string esrgan_path; int n_threads; - bool direct = false; + bool direct = false; + int tile_size = 128; UpscalerGGML(int n_threads, - bool direct = false) + bool direct = false, + int tile_size = 128) : n_threads(n_threads), - direct(direct) { + direct(direct), + tile_size(tile_size) { } bool load_from_file(const std::string& esrgan_path, @@ -51,7 +54,7 @@ struct UpscalerGGML { backend = ggml_backend_cpu_init(); } LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); - esrgan_upscaler = std::make_shared(backend, offload_params_to_cpu, model_loader.get_tensor_storage_map()); + esrgan_upscaler = std::make_shared(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map()); if (direct) { 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, bool offload_params_to_cpu, bool direct, - int n_threads) { + int n_threads, + int tile_size) { upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t)); if (upscaler_ctx == nullptr) { return nullptr; } 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) { return nullptr; }