support non-square tiles

Author: wbruna

ggml_extend.hpp

@@ -635,7 +635,9 @@ sd_tiling_calc_tiles(int &num_tiles_dim, float& tile_overlap_factor_dim, int sma
 }
 
 // Tiling
-__STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const int scale, const int tile_size, const float tile_overlap_factor, on_tile_process on_processing) {
+__STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, ggml_tensor* output, const int scale,
+                                            const int p_tile_size_x, const int p_tile_size_y,
+                                            const float tile_overlap_factor, on_tile_process on_processing) {
     int input_width   = (int)input->ne[0];
     int input_height  = (int)input->ne[1];
     int output_width  = (int)output->ne[0];
@@ -656,25 +658,25 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const
 
     int num_tiles_x;
     float tile_overlap_factor_x;
-    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, tile_size, tile_overlap_factor);
+    sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor);
 
     int num_tiles_y;
     float tile_overlap_factor_y;
-    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, tile_size, tile_overlap_factor);
+    sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor);
 
     LOG_DEBUG("num tiles : %d, %d ", num_tiles_x, num_tiles_y);
     LOG_DEBUG("optimal overlap : %f, %f (targeting %f)", tile_overlap_factor_x, tile_overlap_factor_y, tile_overlap_factor);
 
     GGML_ASSERT(input_width % 2 == 0 && input_height % 2 == 0 && output_width % 2 == 0 && output_height % 2 == 0);  // should be multiple of 2
 
-    int tile_overlap_x     = (int32_t)(tile_size * tile_overlap_factor_x);
-    int non_tile_overlap_x = tile_size - tile_overlap_x;
+    int tile_overlap_x     = (int32_t)(p_tile_size_x * tile_overlap_factor_x);
+    int non_tile_overlap_x = p_tile_size_x - tile_overlap_x;
 
-    int tile_overlap_y     = (int32_t)(tile_size * tile_overlap_factor_y);
-    int non_tile_overlap_y = tile_size - tile_overlap_y;
+    int tile_overlap_y     = (int32_t)(p_tile_size_y * tile_overlap_factor_y);
+    int non_tile_overlap_y = p_tile_size_y - tile_overlap_y;
 
-    int tile_size_x = tile_size < small_width ? tile_size : small_width;
-    int tile_size_y = tile_size < small_height ? tile_size : small_height;
+    int tile_size_x = p_tile_size_x < small_width ? p_tile_size_x : small_width;
+    int tile_size_y = p_tile_size_y < small_height ? p_tile_size_y : small_height;
 
     int input_tile_size_x  = tile_size_x;
     int input_tile_size_y  = tile_size_y;
@@ -763,6 +765,11 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const
     ggml_free(tiles_ctx);
 }
 
+__STATIC_INLINE__ void sd_tiling(ggml_tensor* input, ggml_tensor* output, const int scale,
+    const int tile_size, const float tile_overlap_factor, on_tile_process on_processing) {
+    sd_tiling_non_square(input, output, scale, tile_size, tile_size, tile_overlap_factor, on_processing);
+}
+
 __STATIC_INLINE__ struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx,
                                                          struct ggml_tensor* a) {
     const float eps = 1e-6f;  // default eps parameter

stable-diffusion.cpp

@@ -1427,13 +1427,30 @@ class StableDiffusionGGML {
                                                  x->ne[3]);  // channels
         int64_t t0          = ggml_time_ms();
 
-        int tile_size = 32;
+        int tile_size_x = 32;
+        int tile_size_y = 32;
         // TODO: arg instead of env?
         const char* SD_TILE_SIZE = getenv("SD_TILE_SIZE");
         if (SD_TILE_SIZE != nullptr) {
             std::string sd_tile_size_str = SD_TILE_SIZE;
+            size_t dot_pos = sd_tile_size_str.find('.');
+            size_t x_pos = sd_tile_size_str.find('x');
             try {
-                tile_size = std::stoi(sd_tile_size_str);
+                if (dot_pos != std::string::npos) {
+                    float factor = std::stof(sd_tile_size_str);
+                    tile_size_x = W * (decode ? 8 : 1) / factor;
+                    tile_size_y = H * (decode ? 8 : 1) / factor;
+                }
+                else {
+                    int tmp = std::stoi(sd_tile_size_str.substr(0, x_pos));
+                    if (x_pos != std::string::npos) {
+                        tile_size_y = std::stoi(sd_tile_size_str.substr(x_pos + 1));
+                    }
+                    else {
+                        tile_size_y = tmp;
+                    }
+                    tile_size_x = tmp;
+                }
             } catch (const std::invalid_argument&) {
                 LOG_WARN("Invalid");
             } catch (const std::out_of_range&) {
@@ -1443,7 +1460,8 @@ class StableDiffusionGGML {
         if(!decode){
             // TODO: also use and arg for this one?
             // to keep the compute buffer size consistent
-            tile_size*=1.30539;
+            tile_size_x*=1.30539;
+            tile_size_y*=1.30539;
         }
         if (!use_tiny_autoencoder) {
             if (decode) {
@@ -1452,11 +1470,14 @@ class StableDiffusionGGML {
                 ggml_tensor_scale_input(x);
             }
             if (vae_tiling) {
+                if (SD_TILE_SIZE != nullptr) {
+                    LOG_INFO("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
+                }
                 // split latent in 32x32 tiles and compute in several steps
                 auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) {
                     first_stage_model->compute(n_threads, in, decode, &out);
                 };
-                sd_tiling(x, result, 8, tile_size, 0.5f, on_tiling);
+                sd_tiling_non_square(x, result, 8, tile_size_x, tile_size_y, 0.5f, on_tiling);
             } else {
                 first_stage_model->compute(n_threads, x, decode, &result);
             }