Allow setting the VAE tile size using --vae-tile-size wxh

Author: cmdr2

include/image_generator.h

@@ -113,6 +113,7 @@ class ImageGenerator {
     // SD context parameters from CLI
     bool vae_on_cpu_;
     bool vae_tiling_;
+    std::string vae_tile_size_;
     bool offload_to_cpu_;
     bool diffusion_fa_;
     bool control_net_cpu_;

include/server.h

@@ -16,6 +16,7 @@ struct ServerParams {
     std::shared_ptr<ModelManager> model_manager;
     bool vae_on_cpu = false;
     bool vae_tiling = false;
+    std::string vae_tile_size;
     bool offload_to_cpu = false;
     bool diffusion_fa = false;
     bool control_net_cpu = false;

src/image_generator.cpp

@@ -2,6 +2,7 @@
 
 #include <cstring>
 #include <stdexcept>
+#include <string>
 
 #define STB_IMAGE_WRITE_IMPLEMENTATION
 #define STB_IMAGE_WRITE_STATIC
@@ -84,6 +85,7 @@ ImageGenerator::ImageGenerator(std::shared_ptr<TaskStateManager> task_state_mana
       interrupted_(false),
       vae_on_cpu_(server_params.vae_on_cpu),
       vae_tiling_(server_params.vae_tiling),
+      vae_tile_size_(server_params.vae_tile_size),
       offload_to_cpu_(server_params.offload_to_cpu),
       diffusion_fa_(server_params.diffusion_fa),
       control_net_cpu_(server_params.control_net_cpu),
@@ -202,12 +204,33 @@ std::vector<std::string> ImageGenerator::generateInternal(const ImageGenerationP
     // VAE tiling (if enabled via CLI)
     if (vae_tiling_) {
         gen_params.vae_tiling_params.enabled = true;
-        gen_params.vae_tiling_params.tile_size_x = 512;  // Default tile size
-        gen_params.vae_tiling_params.tile_size_y = 512;
+
+        // Parse tile size
+        int tile_size_x = 256;
+        int tile_size_y = 256;
+        if (!vae_tile_size_.empty()) {
+            size_t x_pos = vae_tile_size_.find('x');
+            try {
+                if (x_pos != std::string::npos) {
+                    std::string tile_x_str = vae_tile_size_.substr(0, x_pos);
+                    std::string tile_y_str = vae_tile_size_.substr(x_pos + 1);
+                    tile_size_x = std::stoi(tile_x_str);
+                    tile_size_y = std::stoi(tile_y_str);
+                } else {
+                    tile_size_x = tile_size_y = std::stoi(vae_tile_size_);
+                }
+            } catch (const std::exception& e) {
+                LOG_WARNING("Invalid VAE tile size '%s', using default 256x256", vae_tile_size_.c_str());
+                tile_size_x = tile_size_y = 256;
+            }
+        }
+
+        gen_params.vae_tiling_params.tile_size_x = tile_size_x;
+        gen_params.vae_tiling_params.tile_size_y = tile_size_y;
         gen_params.vae_tiling_params.target_overlap = 0.5f;
         gen_params.vae_tiling_params.rel_size_x = 1.0f;
         gen_params.vae_tiling_params.rel_size_y = 1.0f;
-        LOG_INFO("VAE tiling enabled with tile size 512x512");
+        LOG_INFO("VAE tiling enabled with tile size %dx%d", tile_size_x, tile_size_y);
     }
 
     // Generate images

src/main.cpp

@@ -91,6 +91,7 @@ struct CommandLineArgs {
     std::string text_encoder_dir;
     bool vae_on_cpu = false;
     bool vae_tiling = false;
+    std::string vae_tile_size;
     bool offload_to_cpu = false;
     bool diffusion_fa = false;
     bool control_net_cpu = false;
@@ -117,6 +118,8 @@ void print_usage(const char* program_name) {
     std::cerr << "  --text-encoder-dir <path>          Text encoder models directory" << std::endl;
     std::cerr << "  --vae-on-cpu                       Keep VAE on CPU (default: false)" << std::endl;
     std::cerr << "  --vae-tiling                       Enable VAE tiling (default: false)" << std::endl;
+    std::cerr << "  --vae-tile-size <size>             VAE tile size (in pixels), format [X]x[Y] (default: 256x256)"
+              << std::endl;
     std::cerr << "  --offload-to-cpu                   Offload parameters to CPU (default: false)" << std::endl;
     std::cerr << "  --diffusion-fa                     Enable diffusion flash attention (default: false)" << std::endl;
     std::cerr << "  --control-net-cpu                  Keep ControlNet on CPU (default: false)" << std::endl;
@@ -171,6 +174,8 @@ CommandLineArgs parse_args(int argc, char* argv[]) {
             args.vae_on_cpu = true;
         } else if (arg == "--vae-tiling") {
             args.vae_tiling = true;
+        } else if (arg == "--vae-tile-size" && i + 1 < argc) {
+            args.vae_tile_size = argv[++i];
         } else if (arg == "--offload-to-cpu") {
             args.offload_to_cpu = true;
         } else if (arg == "--diffusion-fa") {
@@ -255,6 +260,7 @@ int main(int argc, char* argv[]) {
         server_params.model_manager = model_manager;
         server_params.vae_on_cpu = args.vae_on_cpu;
         server_params.vae_tiling = args.vae_tiling;
+        server_params.vae_tile_size = args.vae_tile_size;
         server_params.offload_to_cpu = args.offload_to_cpu;
         server_params.diffusion_fa = args.diffusion_fa;
         server_params.control_net_cpu = args.control_net_cpu;