apg: add experimental threshold smoothing parameter

Author: stduhpf

examples/cli/main.cpp

@@ -130,6 +130,7 @@ struct SDParams {
     float apg_eta            = 1.0f;
     float apg_momentum       = 0.0f;
     float apg_norm_threshold = 0.0f;
+    float apg_norm_smoothing = 0.0f;
 };
 
 void print_params(SDParams params) {
@@ -214,6 +215,8 @@ void print_usage(int argc, const char* argv[]) {
     printf("  --apg-eta VALUE                    parallel projected guidance scale for APG (default: 1.0, recommended: between 0 and 1)\n");
     printf("  --apg-momentum VALUE               CFG update direction momentum for APG (default: 0, recommended: around -0.5)\n");
     printf("  --apg-nt, --apg-rescale VALUE      CFG update direction norm threshold for APG (default: 0 = disabled, recommended: 4-15)\n");
+    printf("  --apg-nt-smoothing VALUE           EXPERIMENTAL! Norm threshold smoothing for APG (default: 0 = disabled)\n");
+    printf("                                     (replaces saturation with a smooth approximation)\n");
     printf("  --slg-scale SCALE                  skip layer guidance (SLG) scale, only for DiT models: (default: 0)\n");
     printf("                                     0 means disabled, a value of 2.5 is nice for sd3.5 medium\n");
     printf("  --skip-layers LAYERS               Layers to skip for SLG steps: (default: [7,8,9])\n");
@@ -641,6 +644,12 @@ void parse_args(int argc, const char** argv, SDParams& params) {
                 break;
             }
             params.apg_norm_threshold = std::stof(argv[i]);
+        } else if (arg == "--apg-nt-smoothing") {
+            if (++i >= argc) {
+                invalid_arg = true;
+                break;
+            }
+            params.apg_norm_smoothing = std::stof(argv[i]);
         } else {
             fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
             print_usage(argc, argv);
@@ -739,6 +748,9 @@ std::string get_image_params(SDParams params, int64_t seed) {
     }
     if (params.apg_norm_threshold != 0) {
         parameter_string += "CFG normalization threshold: " + std::to_string(params.apg_norm_threshold) + ", ";
+        if (params.apg_norm_smoothing != 0) {
+            parameter_string += "CFG normalization threshold: " + std::to_string(params.apg_norm_smoothing) + ", ";
+        }
     }
     if (params.slg_scale != 0 && params.skip_layers.size() != 0) {
         parameter_string += "SLG scale: " + std::to_string(params.cfg_scale) + ", ";
@@ -989,7 +1001,8 @@ int main(int argc, const char* argv[]) {
                                           params.skip_layer_end},
                           sd_apg_params_t{params.apg_eta,
                                           params.apg_momentum,
-                                          params.apg_norm_threshold});
+                                          params.apg_norm_threshold,
+                                          params.apg_norm_smoothing});
     } else {
         sd_image_t input_image = {(uint32_t)params.width,
                                   (uint32_t)params.height,
@@ -1060,7 +1073,8 @@ int main(int argc, const char* argv[]) {
                                               params.skip_layer_end},
                               sd_apg_params_t{params.apg_eta,
                                               params.apg_momentum,
-                                              params.apg_norm_threshold});
+                                              params.apg_norm_threshold,
+                                              params.apg_norm_smoothing});
         }
     }
 

stable-diffusion.cpp

@@ -797,7 +797,7 @@ class StableDiffusionGGML {
                         int start_merge_step,
                         SDCondition id_cond,
                         sd_slg_params_t slg_params = {NULL, 0, 0, 0, 0},
-                        sd_apg_params_t apg_params = {1, 0, 0},
+                        sd_apg_params_t apg_params = {1, 0, 0, 0},
                         ggml_tensor* noise_mask    = nullptr) {
         std::vector<int> skip_layers(slg_params.skip_layers, slg_params.skip_layers + slg_params.skip_layers_count);
 
@@ -976,8 +976,14 @@ class StableDiffusionGGML {
                     deltas[i] = delta;
                 }
                 if (apg_params.norm_treshold > 0) {
-                    diff_norm        = sqrtf(diff_norm);
-                    apg_scale_factor = std::min(1.0f, apg_params.norm_treshold / diff_norm);
+                    diff_norm = sqrtf(diff_norm);
+                    if (apg_params.norm_treshold_smoothing <= 0) {
+                        apg_scale_factor = std::min(1.0f, apg_params.norm_treshold / diff_norm);
+                    } else {
+                        // Experimental: smooth saturate
+                        float x          = apg_params.norm_treshold / diff_norm;
+                        apg_scale_factor = x / std::pow(1 + std::pow(x, 1.0 / apg_params.norm_treshold_smoothing), apg_params.norm_treshold_smoothing);
+                    }
                 }
                 if (apg_params.eta != 1.0f) {
                     dot *= apg_scale_factor;

stable-diffusion.h

@@ -126,6 +126,7 @@ typedef struct {
     float eta;
     float momentum;
     float norm_treshold;
+    float norm_treshold_smoothing;
 } sd_apg_params_t;
 
 typedef struct {