configurable resolution limit (#1586)

* refactor image gen configuration screen

* make image size limit configurable

* fix resolution limits and keep dimensions closer to the original ratio

* use 0.0 for the configured default image size limit

This prevents the current default value from being saved into the
config files, in case we later decide to adopt a different value.

* export image model version when loading

* restore model-specific default image size limit

* change the image area restriction to be specified by a square side

* move image resolution limits down to the C++ level

* Revert "export image model version when loading"

This reverts commit fa65b23.

* Linting Fixes:
PY:
- Inconsistent var name sd_restrict_square -> sd_restrict_square_var
- GUI swap back to using absolute row numbers for now.
- fstring fix
- size_limit -> side_limit inconsistency
C++:
- roundup_64 standalone function
- refactor sd_fix_resolution variable names for clarity
- move "anti crashing" hard total megapixel limit always to be applied after soft total megapixel limit instead of conditionally only when sd_restrict_square is unset

* allow unsafe resolutions if debugmode is on

---------

Co-authored-by: Concedo <[email protected]>

Author: wbruna

expose.h

@@ -169,6 +169,8 @@ struct sd_load_model_inputs
     const char * vae_filename = nullptr;
     const char * lora_filename = nullptr;
     const float lora_multiplier = 1.0f;
+    const int side_limit = 0;
+    const int square_limit = 0;
     const bool quiet = false;
     const int debugmode = 0;
 };

koboldcpp.py

@@ -274,6 +274,8 @@ class sd_load_model_inputs(ctypes.Structure):
                 ("vae_filename", ctypes.c_char_p),
                 ("lora_filename", ctypes.c_char_p),
                 ("lora_multiplier", ctypes.c_float),
+                ("side_limit", ctypes.c_int),
+                ("square_limit", ctypes.c_int),
                 ("quiet", ctypes.c_bool),
                 ("debugmode", ctypes.c_int)]
 
@@ -1539,6 +1541,8 @@ def sd_load_model(model_filename,vae_filename,lora_filename,t5xxl_filename,clipl
     inputs.t5xxl_filename = t5xxl_filename.encode("UTF-8")
     inputs.clipl_filename = clipl_filename.encode("UTF-8")
     inputs.clipg_filename = clipg_filename.encode("UTF-8")
+    inputs.side_limit = args.sdclamped
+    inputs.square_limit = args.sdrestrictsquare
     inputs = set_backend_props(inputs)
     ret = handle.sd_load_model(inputs)
     return ret
@@ -1618,27 +1622,11 @@ def sd_generate(genparams):
     clip_skip = tryparseint(genparams.get("clip_skip", -1),-1)
 
     #clean vars
-    width = width - (width%64)
-    height = height - (height%64)
     cfg_scale = (1 if cfg_scale < 1 else (25 if cfg_scale > 25 else cfg_scale))
     sample_steps = (1 if sample_steps < 1 else (forced_steplimit if sample_steps > forced_steplimit else sample_steps))
-    reslimit = 1024
-    width = (64 if width < 64 else width)
-    height = (64 if height < 64 else height)
 
     if args.sdclamped:
         sample_steps = (40 if sample_steps > 40 else sample_steps)
-        reslimit = int(args.sdclamped)
-        reslimit = (512 if reslimit<512 else reslimit)
-        print(f"\nImgGen: Clamped Mode (For Shared Use). Step counts and resolution are clamped to {reslimit}x{reslimit}.")
-
-    biggest = max(width,height)
-    if biggest > reslimit:
-        scaler = biggest / reslimit
-        width = int(width / scaler)
-        height = int(height / scaler)
-        width = width - (width%64)
-        height = height - (height%64)
 
     inputs = sd_generation_inputs()
     inputs.prompt = prompt.encode("UTF-8")
@@ -4268,6 +4256,7 @@ def hide_tooltip(event):
     sd_vaeauto_var = ctk.IntVar(value=0)
     sd_notile_var = ctk.IntVar(value=0)
     sd_clamped_var = ctk.StringVar(value="0")
+    sd_restrict_square_var = ctk.StringVar(value="0")
     sd_threads_var = ctk.StringVar(value=str(default_threads))
     sd_quant_var = ctk.IntVar(value=0)
 
@@ -4965,20 +4954,22 @@ def togglehorde(a,b,c):
     images_tab = tabcontent["Image Gen"]
     makefileentry(images_tab, "Image Gen. Model (safetensors/gguf):", "Select Image Gen Model File", sd_model_var, 1, width=280, singlecol=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")], tooltiptxt="Select a .safetensors or .gguf Image Generation model file on disk to be loaded.")
     makelabelentry(images_tab, "Clamped Mode (Limit Resolution):", sd_clamped_var, 4, 50, padx=290,singleline=True,tooltip="Limit generation steps and resolution settings for shared use.\nSet to 0 to disable, otherwise value is the size limit (min 512px).")
-    makelabelentry(images_tab, "Image Threads:" , sd_threads_var, 6, 50,padx=290,singleline=True,tooltip="How many threads to use during image generation.\nIf left blank, uses same value as threads.")
+    makelabelentry(images_tab, "Restrict Square Size:", sd_restrict_square_var, 6, 50, padx=290,singleline=True,tooltip="Square image size restriction, to protect the server against memory crashes.\nAllows width-height tradeoffs, eg. 640 allows 640x640 and 512x768\nLeave at 0 for the default value: 832 for SD1.5/SD2, 1024 otherwise.")
+    makelabelentry(images_tab, "Image Threads:" , sd_threads_var, 8, 50,padx=290,singleline=True,tooltip="How many threads to use during image generation.\nIf left blank, uses same value as threads.")
     sd_model_var.trace("w", gui_changed_modelfile)
+    makecheckbox(images_tab, "Compress Weights (Saves Memory)", sd_quant_var, 10,tooltiptxt="Quantizes the SD model weights to save memory. May degrade quality.")
+    sd_quant_var.trace("w", changed_gpulayers_estimate)
+
+    makefileentry(images_tab, "Image LoRA (safetensors/gguf):", "Select SD lora file",sd_lora_var, 20, width=280, singlecol=True, filetypes=[("*.safetensors *.gguf", "*.safetensors *.gguf")],tooltiptxt="Select a .safetensors or .gguf SD LoRA model file to be loaded. Should be unquantized!")
+    makelabelentry(images_tab, "Image LoRA Multiplier:" , sd_loramult_var, 22, 50,padx=290,singleline=True,tooltip="What mutiplier value to apply the SD LoRA with.")
 
-    makefileentry(images_tab, "Image LoRA (safetensors/gguf):", "Select SD lora file",sd_lora_var, 10, width=280, singlecol=True, filetypes=[("*.safetensors *.gguf", "*.safetensors *.gguf")],tooltiptxt="Select a .safetensors or .gguf SD LoRA model file to be loaded. Should be unquantized!")
-    makelabelentry(images_tab, "Image LoRA Multiplier:" , sd_loramult_var, 12, 50,padx=290,singleline=True,tooltip="What mutiplier value to apply the SD LoRA with.")
 
-    makecheckbox(images_tab, "Compress Weights (Saves Memory)", sd_quant_var, 8,tooltiptxt="Quantizes the SD model weights to save memory. May degrade quality.")
-    sd_quant_var.trace("w", changed_gpulayers_estimate)
 
-    makefileentry(images_tab, "T5-XXL File:", "Select Optional T5-XXL model file (SD3 or flux)",sd_t5xxl_var, 14, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")],tooltiptxt="Select a .safetensors t5xxl file to be loaded.")
-    makefileentry(images_tab, "Clip-L File:", "Select Optional Clip-L model file (SD3 or flux)",sd_clipl_var, 16, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")],tooltiptxt="Select a .safetensors t5xxl file to be loaded.")
-    makefileentry(images_tab, "Clip-G File:", "Select Optional Clip-G model file (SD3)",sd_clipg_var, 18, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")],tooltiptxt="Select a .safetensors t5xxl file to be loaded.")
+    makefileentry(images_tab, "T5-XXL File:", "Select Optional T5-XXL model file (SD3 or flux)",sd_t5xxl_var, 24, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")],tooltiptxt="Select a .safetensors t5xxl file to be loaded.")
+    makefileentry(images_tab, "Clip-L File:", "Select Optional Clip-L model file (SD3 or flux)",sd_clipl_var, 26, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")],tooltiptxt="Select a .safetensors t5xxl file to be loaded.")
+    makefileentry(images_tab, "Clip-G File:", "Select Optional Clip-G model file (SD3)",sd_clipg_var, 28, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf","*.safetensors *.gguf")],tooltiptxt="Select a .safetensors t5xxl file to be loaded.")
 
-    sdvaeitem1,sdvaeitem2,sdvaeitem3 = makefileentry(images_tab, "Image VAE:", "Select Optional SD VAE file",sd_vae_var, 20, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf", "*.safetensors *.gguf")],tooltiptxt="Select a .safetensors or .gguf SD VAE file to be loaded.")
+    sdvaeitem1,sdvaeitem2,sdvaeitem3 = makefileentry(images_tab, "Image VAE:", "Select Optional SD VAE file",sd_vae_var, 30, width=280, singlerow=True, filetypes=[("*.safetensors *.gguf", "*.safetensors *.gguf")],tooltiptxt="Select a .safetensors or .gguf SD VAE file to be loaded.")
     def toggletaesd(a,b,c):
         if sd_vaeauto_var.get()==1:
             sdvaeitem1.grid_remove()
@@ -4989,7 +4980,7 @@ def toggletaesd(a,b,c):
                 sdvaeitem1.grid()
                 sdvaeitem2.grid()
                 sdvaeitem3.grid()
-    makecheckbox(images_tab, "Use TAE SD (AutoFix Broken VAE)", sd_vaeauto_var, 22,command=toggletaesd,tooltiptxt="Replace VAE with TAESD. May fix bad VAE.")
+    makecheckbox(images_tab, "Use TAE SD (AutoFix Broken VAE)", sd_vaeauto_var, 32,command=toggletaesd,tooltiptxt="Replace VAE with TAESD. May fix bad VAE.")
     makecheckbox(images_tab, "No VAE Tiling", sd_notile_var, 24,tooltiptxt="Disables VAE tiling, may not work for large images.")
 
     # audio tab
@@ -5222,6 +5213,7 @@ def export_vars():
 
         args.sdthreads = (0 if sd_threads_var.get()=="" else int(sd_threads_var.get()))
         args.sdclamped = (0 if int(sd_clamped_var.get())<=0 else int(sd_clamped_var.get()))
+        args.sdrestrictsquare = (0 if int(sd_restrict_square_var.get())<=0 else int(sd_restrict_square_var.get()))
         args.sdnotile = (True if sd_notile_var.get()==1 else False)
         if sd_vaeauto_var.get()==1:
             args.sdvaeauto = True
@@ -5432,6 +5424,7 @@ def import_vars(dict):
 
         sd_model_var.set(dict["sdmodel"] if ("sdmodel" in dict and dict["sdmodel"]) else "")
         sd_clamped_var.set(int(dict["sdclamped"]) if ("sdclamped" in dict and dict["sdclamped"]) else 0)
+        sd_restrict_square_var.set(int(dict["sdrestrictsquare"]) if ("sdrestrictsquare" in dict and dict["sdrestrictsquare"]) else 0)
         sd_threads_var.set(str(dict["sdthreads"]) if ("sdthreads" in dict and dict["sdthreads"]) else str(default_threads))
         sd_quant_var.set(1 if ("sdquant" in dict and dict["sdquant"]) else 0)
         sd_vae_var.set(dict["sdvae"] if ("sdvae" in dict and dict["sdvae"]) else "")
@@ -7168,6 +7161,7 @@ def range_checker(arg: str):
     sdparsergroup.add_argument("--sdmodel", metavar=('[filename]'), help="Specify an image generation safetensors or gguf model to enable image generation.", default="")
     sdparsergroup.add_argument("--sdthreads", metavar=('[threads]'), help="Use a different number of threads for image generation if specified. Otherwise, has the same value as --threads.", type=int, default=0)
     sdparsergroup.add_argument("--sdclamped", metavar=('[maxres]'), help="If specified, limit generation steps and resolution settings for shared use. Accepts an extra optional parameter that indicates maximum resolution (eg. 768 clamps to 768x768, min 512px, disabled if 0).", nargs='?', const=512, type=int, default=0)
+    sdparsergroup.add_argument("--sdrestrictsquare", metavar=('[maxres]'), help="If specified, restrict square image sides to this value, in pixels, to avoid server crashes related to excessive memory usage. Similar to --sdclamped, but allows trade-offs between width and height (e.g. 640 would allow 640x640, 512x768 and 768x512 images). If 0 or unspecified, use a model-specific safe value: 832 for SD1.5/SD2, 1024 otherwise. Total resolution cannot exceed 1MP.", type=int, default=0)
     sdparsergroup.add_argument("--sdt5xxl", metavar=('[filename]'), help="Specify a T5-XXL safetensors model for use in SD3 or Flux. Leave blank if prebaked or unused.", default="")
     sdparsergroup.add_argument("--sdclipl", metavar=('[filename]'), help="Specify a Clip-L safetensors model for use in SD3 or Flux. Leave blank if prebaked or unused.", default="")
     sdparsergroup.add_argument("--sdclipg", metavar=('[filename]'), help="Specify a Clip-G safetensors model for use in SD3. Leave blank if prebaked or unused.", default="")

otherarch/sdcpp/sdtype_adapter.cpp

@@ -119,6 +119,8 @@ static uint8_t * input_mask_buffer = NULL;
 
 static std::string sdplatformenv, sddeviceenv, sdvulkandeviceenv;
 static bool notiling = false;
+static int cfg_square_limit = 0;
+static int cfg_side_limit = 0;
 static bool sd_is_quiet = false;
 static std::string sdmodelfilename = "";
 
@@ -133,6 +135,8 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
     std::string clipl_filename = inputs.clipl_filename;
     std::string clipg_filename = inputs.clipg_filename;
     notiling = inputs.notile;
+    cfg_side_limit = inputs.side_limit;
+    cfg_square_limit = inputs.square_limit;
     printf("\nImageGen Init - Load Model: %s\n",inputs.model_filename);
 
     if(lorafilename!="")
@@ -307,6 +311,99 @@ static std::string get_image_params(const SDParams& params) {
     return parameter_string;
 }
 
+static inline int rounddown_64(int n) {
+    return n - n % 64;
+}
+
+static inline int roundup_64(int n) {
+    return ((n + 63) / 64) * 64;
+}
+
+//scale dimensions to ensure width and height stay within limits
+//side_limit = sdclamped, hard size limit per side, no side can exceed this
+//square limit = total NxN resolution based limit to also apply
+static void sd_fix_resolution(int &width, int &height, int side_limit, int square_limit) {
+
+    // sanitize the original values
+    width = std::max(std::min(width, 8192), 64);
+    height = std::max(std::min(height, 8192), 64);
+
+    bool is_landscape = (width > height);
+    int long_side = is_landscape ? width : height;
+    int short_side = is_landscape ? height : width;
+    float original_ratio = static_cast<float>(long_side) / short_side;
+
+    // for the initial rounding, don't bother comparing to the original
+    // requested ratio, since the user can choose those values directly
+    long_side = rounddown_64(long_side);
+    short_side = rounddown_64(short_side);
+    side_limit = rounddown_64(side_limit);
+
+    //enforce sdclamp side limit
+    if (long_side > side_limit) {
+        short_side = static_cast<int>(short_side * side_limit / static_cast<float>(long_side));
+        long_side = side_limit;
+        if (short_side <= 64) {
+            short_side = 64;
+        } else {
+            int down = rounddown_64(short_side);
+            int up = roundup_64(short_side);
+            float longf = static_cast<float>(long_side);
+            // Choose better ratio match between rounding up or down
+            short_side = (longf / down - original_ratio < original_ratio - longf / up) ? down : up;
+        }
+    }
+
+    //enforce sd_restrict_square area limit
+    int area_limit = square_limit * square_limit;
+    if (long_side * short_side > area_limit) {
+        float scale = std::sqrt(static_cast<float>(area_limit) / (long_side * short_side));
+        int new_short = static_cast<int>(short_side * scale);
+        int new_long = static_cast<int>(long_side * scale);
+
+        if (new_short <= 64) {
+            short_side = 64;
+            long_side = rounddown_64(area_limit / short_side);
+        } else {
+            int new_long_down = rounddown_64(new_long);
+            int new_short_down = rounddown_64(new_short);
+            int new_short_up = roundup_64(new_short);
+            int new_long_up = roundup_64(new_long);
+            long_side = new_long_down;
+            short_side = new_short_down;
+
+            // we may get a ratio closer to the original if we still stay below the
+            // limit when rounding up one of the dimensions, so check both cases
+            float rdiff = std::fabs(static_cast<float>(new_long_down) / new_short_down - original_ratio);
+
+            if (new_long_down * new_short_up < area_limit) {
+                float newrdiff = std::fabs(static_cast<float>(new_long_down) / new_short_up - original_ratio);
+                if (newrdiff < rdiff) {
+                    long_side = new_long_down;
+                    short_side = new_short_up;
+                    rdiff = newrdiff;
+                }
+            }
+
+            if (new_long_up * new_short_down < area_limit) {
+                float newrdiff = std::fabs(static_cast<float>(new_long_up) / new_short_down - original_ratio);
+                if (newrdiff < rdiff) {
+                    long_side = new_long_up;
+                    short_side = new_short_down;
+                }
+            }
+        }
+    }
+
+    if (is_landscape) {
+        width = long_side;
+        height = short_side;
+    } else {
+        width = short_side;
+        height = long_side;
+    }
+}
+
 sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
 {
     sd_generation_outputs output;
@@ -339,8 +436,6 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
     sd_params->clip_skip = inputs.clip_skip;
     sd_params->mode = (img2img_data==""?SDMode::TXT2IMG:SDMode::IMG2IMG);
 
-    //ensure unsupported dimensions are fixed
-    int biggestdim = (sd_params->width>sd_params->height?sd_params->width:sd_params->height);
     auto loadedsdver = get_loaded_sd_version(sd_ctx);
     if (loadedsdver == SDVersion::VERSION_FLUX)
     {
@@ -351,21 +446,34 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
             sampler = "euler";  //euler a broken on flux
         }
     }
-    int reslimit = (loadedsdver==SDVersion::VERSION_SD1 || loadedsdver==SDVersion::VERSION_SD2)?832:1024;
-    if(biggestdim > reslimit)
-    {
-        float scaler = (float)biggestdim / (float)reslimit;
-        int newwidth = (int)((float)sd_params->width / scaler);
-        int newheight = (int)((float)sd_params->height / scaler);
-        newwidth = newwidth - (newwidth%64);
-        newheight = newheight - (newheight%64);
-        sd_params->width = newwidth;
-        sd_params->height = newheight;
-        if(!sd_is_quiet && sddebugmode==1)
-        {
-            printf("\nDownscale to %dx%d as %d > %d\n",newwidth,newheight,biggestdim,reslimit);
-        }
+
+    const int default_res_limit = 8192; // arbitrary, just to simplify the code
+    // avoid crashes due to bugs/limitations on certain models
+    // although it can be possible for a single side to exceed 1024, the total resolution of the image
+    // cannot exceed (832x832) for sd1/sd2 or (1024x1024) for sdxl/sd3/flux, to prevent crashing the server
+    const int hard_megapixel_res_limit = (loadedsdver==SDVersion::VERSION_SD1 || loadedsdver==SDVersion::VERSION_SD2)?832:1024;
+
+    int side_limit = default_res_limit;
+    if (cfg_side_limit > 0) {
+        side_limit = std::max(std::min(cfg_side_limit, default_res_limit), 64);
+    }
+
+    int square_limit = default_res_limit;
+    if (cfg_square_limit > 0) {
+        square_limit = std::max(std::min(cfg_square_limit, default_res_limit), 64);
     }
+
+    if (cfg_square_limit > 0 && sddebugmode == 1) {
+        square_limit = std::min(hard_megapixel_res_limit * 2, square_limit);  //double the limit for debugmode if cfg_square_limit is set
+    } else {
+        square_limit = std::min(hard_megapixel_res_limit, square_limit);
+    }
+
+    sd_fix_resolution(sd_params->width, sd_params->height, side_limit, square_limit);
+    if (inputs.width != sd_params->width || inputs.height != sd_params->height) {
+        printf("\nKCPP SD: Requested dimensions %dx%d changed to %dx%d\n", inputs.width, inputs.height, sd_params->width, sd_params->height);
+    }
+
     bool dotile = (sd_params->width>768 || sd_params->height>768) && !notiling;
     set_sd_vae_tiling(sd_ctx,dotile); //changes vae tiling, prevents memory related crash/oom