Support Custom ESRGAN tile size

Author: pedroCabrera

esrgan.hpp

@@ -156,9 +156,11 @@ 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 {

examples/cli/main.cpp

@@ -118,6 +118,7 @@ struct SDParams {
     bool canny_preprocess      = false;
     bool color                 = false;
     int upscale_repeats        = 1;
+    int upscale_tile_size      = 128;
 
     // Photo Maker
     std::string photo_maker_path;
@@ -204,6 +205,7 @@ void print_params(SDParams params) {
     printf("    vae_tiling:                        %s\n", params.vae_tiling_params.enabled ? "true" : "false");
     printf("    force_sdxl_vae_conv_scale:         %s\n", params.force_sdxl_vae_conv_scale ? "true" : "false");
     printf("    upscale_repeats:                   %d\n", params.upscale_repeats);
+    printf("    upscale_tile_size:                 %d\n", params.upscale_tile_size);
     printf("    chroma_use_dit_mask:               %s\n", params.chroma_use_dit_mask ? "true" : "false");
     printf("    chroma_use_t5_mask:                %s\n", params.chroma_use_t5_mask ? "true" : "false");
     printf("    chroma_t5_mask_pad:                %d\n", params.chroma_t5_mask_pad);
@@ -605,6 +607,10 @@ void parse_args(int argc, const char** argv, SDParams& params) {
          "--upscale-repeats",
          "Run the ESRGAN upscaler this many times (default: 1)",
          &params.upscale_repeats},
+        {"",
+         "--upscale-tile-size",
+         "tile size for ESRGAN upscaling (default: 128)",
+         &params.upscale_tile_size},
         {"-H",
          "--height",
          "image height, in pixel space (default: 512)",
@@ -1187,6 +1193,11 @@ void parse_args(int argc, const char** argv, SDParams& params) {
         exit(1);
     }
 
+    if (params.upscale_tile_size < 1) {
+        fprintf(stderr, "error: upscale tile size must be at least 1\n");
+        exit(1);
+    }
+
     if (params.mode == UPSCALE) {
         if (params.esrgan_path.length() == 0) {
             fprintf(stderr, "error: upscale mode needs an upscaler model (--upscale-model)\n");
@@ -1757,7 +1768,8 @@ int main(int argc, const char* argv[]) {
         upscaler_ctx_t* upscaler_ctx = new_upscaler_ctx(params.esrgan_path.c_str(),
                                                         params.offload_params_to_cpu,
                                                         params.diffusion_conv_direct,
-                                                        params.n_threads);
+                                                        params.n_threads,
+                                                        params.upscale_tile_size);
 
         if (upscaler_ctx == nullptr) {
             printf("new_upscaler_ctx failed\n");

stable-diffusion.h

@@ -293,7 +293,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,

upscaler.cpp

@@ -10,11 +10,14 @@ struct UpscalerGGML {
     std::string esrgan_path;
     int n_threads;
     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<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) {
             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;
     }