diff --git a/expose.h b/expose.h index 9b7347219..84da17d32 100644 --- a/expose.h +++ b/expose.h @@ -226,6 +226,8 @@ struct sd_generation_inputs const bool remove_limits = false; const bool circular_x = false; const bool circular_y = false; + const char * cache_mode = nullptr; + const char * cache_options = nullptr; const bool upscale = false; const int lora_len = 0; const float * lora_multipliers = nullptr; diff --git a/koboldcpp.py b/koboldcpp.py index 858f42df9..a19e5a02e 100755 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -358,6 +358,8 @@ class sd_generation_inputs(ctypes.Structure): ("remove_limits", ctypes.c_bool), ("circular_x", ctypes.c_bool), ("circular_y", ctypes.c_bool), + ("cache_mode", ctypes.c_char_p), + ("cache_options", ctypes.c_char_p), ("upscale", ctypes.c_bool), ("lora_len", ctypes.c_int), ("lora_multipliers", ctypes.POINTER(ctypes.c_float))] @@ -2112,7 +2114,26 @@ def sd_comfyui_tranform_params(genparams): return genparams # json with top-level dict -def gendefaults_parse_meta_field(input_str): +def parse_json_object(value, field): + broken = False + if isinstance(value, str): + try: # Try parsing as-is + value = json.loads(value) + except json.JSONDecodeError: + # Try wrapping in braces for loose key/value strings + try: + value = json.loads(f"{{{value}}}") + except json.JSONDecodeError: + broken = True + if isinstance(value, dict): + return value + elif broken: + print(f"Warning: couldn't parse {field} field.") + else: + print(f"Warning: {field} field - not a JSON object.") + return None + +def gendefaults_parse_meta_field(value): alias_map = { 'cfg-scale': 'cfg_scale', 'guidance': 'distilled_guidance', @@ -2120,22 +2141,13 @@ def gendefaults_parse_meta_field(input_str): 'sampling-method': 'sampler_name', 'timestep-shift': 'shifted_timestep', 'flow-shift': 'flow_shift', + 'cache-mode': 'cache_mode', + 'cache-options': 'cache_options', + # match sd.cpp flag + 'cache-option': 'cache_options', + 'cache_option': 'cache_options', } - if not isinstance(input_str, str) or not input_str.strip(): - return {} - parsed = None - try: # Try parsing as-is - parsed = json.loads(input_str) - except json.JSONDecodeError: - # Try wrapping in braces for loose key/value strings - try: - parsed = json.loads(f"{{{input_str}}}") - except json.JSONDecodeError: - print("Warning: couldn't parse gendefaults_parse_meta_field.") - return {} - if not isinstance(parsed, dict): - print("Warning: gendefaults_parse_meta_field - not a JSON object.") - return {} + parsed = parse_json_object(value, 'gendefaults') or {} result = {} # First pass: apply aliases only if canonical key is not explicitly present for key, value in parsed.items(): @@ -2264,6 +2276,8 @@ def sd_generate(genparams): vid_req_frames = tryparseint(genparams.get("frames", 1),1) vid_req_frames = 1 if (not vid_req_frames or vid_req_frames < 1) else vid_req_frames video_output_type = genparams.get("video_output_type", 0) + cache_mode = str(genparams.get("cache_mode", "")) + cache_options = str(genparams.get("cache_options", "")) extra_images_arr = genparams.get("extra_images", []) extra_images_arr = ([] if not extra_images_arr else extra_images_arr) extra_images_arr = [img for img in extra_images_arr if img not in (None, "")] @@ -2316,6 +2330,8 @@ def sd_generate(genparams): inputs.remove_limits = allow_remove_limits inputs.circular_x = tryparseint(adapter_obj.get("circular_x", genparams.get("circular_x",0)),0) inputs.circular_y = tryparseint(adapter_obj.get("circular_y", genparams.get("circular_y",0)),0) + inputs.cache_mode = cache_mode.encode("UTF-8") + inputs.cache_options = cache_options.encode("UTF-8") inputs.upscale = (True if tryparseint(genparams.get("enable_hr", 0),0) else False) lora_multipliers = prepare_lora_multipliers(genparams.get("lora", [])) @@ -5746,6 +5762,11 @@ def save_config_dict(filename, savdict, template): filenamestr += ".kcppt" do_not_save = {'analyze', 'config', 'exportconfig', 'exporttemplate', 'testmemory', 'unpack', 'version'} filtered = {k: v for k, v in savdict.items() if k not in do_not_save} + if 'gendefaults' in filtered: + gendefaults = parse_json_object(filtered['gendefaults'], 'gendefaults') + if isinstance(gendefaults, dict): + filtered['gendefaults'] = gendefaults + # keep it as-is if it's a broken string with open(filenamestr, 'w') as file: file.write(json.dumps(filtered,indent=2)) return filenamestr @@ -7468,7 +7489,10 @@ def show_gui(): else: sd_lora_var.set("") sd_loramult_var.set(" ".join(f"{n:.3f}".rstrip('0').rstrip('.') for n in dict.get("sdloramult", []))) - gen_defaults_var.set(dict["gendefaults"] if ("gendefaults" in dict and dict["gendefaults"]) else "") + gendefaults = (dict["gendefaults"] if ("gendefaults" in dict and dict["gendefaults"]) else "") + if isinstance(gendefaults, type({})): + gendefaults = json.dumps(gendefaults) + gen_defaults_var.set(gendefaults) gen_defaults_overwrite_var.set(1 if "gendefaultsoverwrite" in dict and dict["gendefaultsoverwrite"] else 0) whisper_model_var.set(dict["whispermodel"] if ("whispermodel" in dict and dict["whispermodel"]) else "") diff --git a/otherarch/sdcpp/common/common.hpp b/otherarch/sdcpp/common/common.hpp index 4150bf9d7..76d1342dd 100644 --- a/otherarch/sdcpp/common/common.hpp +++ b/otherarch/sdcpp/common/common.hpp @@ -1427,8 +1427,8 @@ struct SDGenerationParams { } cache_mode = argv_to_utf8(index, argv); if (cache_mode != "easycache" && cache_mode != "ucache" && - cache_mode != "dbcache" && cache_mode != "taylorseer" && cache_mode != "cache-dit") { - fprintf(stderr, "error: invalid cache mode '%s', must be 'easycache', 'ucache', 'dbcache', 'taylorseer', or 'cache-dit'\n", cache_mode.c_str()); + cache_mode != "dbcache" && cache_mode != "taylorseer" && cache_mode != "cache-dit" && cache_mode != "spectrum") { + fprintf(stderr, "error: invalid cache mode '%s', must be 'easycache', 'ucache', 'dbcache', 'taylorseer', 'cache-dit', or 'spectrum'\n", cache_mode.c_str()); return -1; } return 1; @@ -1784,7 +1784,23 @@ struct SDGenerationParams { } else if (key == "Bn" || key == "bn") { cache_params.Bn_compute_blocks = std::stoi(val); } else if (key == "warmup") { - cache_params.max_warmup_steps = std::stoi(val); + if (cache_mode == "spectrum") { + cache_params.spectrum_warmup_steps = std::stoi(val); + } else { + cache_params.max_warmup_steps = std::stoi(val); + } + } else if (key == "w") { + cache_params.spectrum_w = std::stof(val); + } else if (key == "m") { + cache_params.spectrum_m = std::stoi(val); + } else if (key == "lam") { + cache_params.spectrum_lam = std::stof(val); + } else if (key == "window") { + cache_params.spectrum_window_size = std::stoi(val); + } else if (key == "flex") { + cache_params.spectrum_flex_window = std::stof(val); + } else if (key == "stop") { + cache_params.spectrum_stop_percent = std::stof(val); } else { LOG_ERROR("error: unknown cache parameter '%s'", key.c_str()); return false; @@ -1832,6 +1848,15 @@ struct SDGenerationParams { cache_params.Bn_compute_blocks = 0; cache_params.residual_diff_threshold = 0.08f; cache_params.max_warmup_steps = 8; + } else if (cache_mode == "spectrum") { + cache_params.mode = SD_CACHE_SPECTRUM; + cache_params.spectrum_w = 0.40f; + cache_params.spectrum_m = 3; + cache_params.spectrum_lam = 1.0f; + cache_params.spectrum_window_size = 2; + cache_params.spectrum_flex_window = 0.50f; + cache_params.spectrum_warmup_steps = 4; + cache_params.spectrum_stop_percent = 0.9f; } if (!cache_option.empty()) { diff --git a/otherarch/sdcpp/ggml_extend.hpp b/otherarch/sdcpp/ggml_extend.hpp index 5265aeed4..02d8d44a6 100644 --- a/otherarch/sdcpp/ggml_extend.hpp +++ b/otherarch/sdcpp/ggml_extend.hpp @@ -491,12 +491,16 @@ __STATIC_INLINE__ void ggml_ext_tensor_split_2d(struct ggml_tensor* input, int64_t height = output->ne[1]; int64_t channels = output->ne[2]; int64_t ne3 = output->ne[3]; + + int64_t input_width = input->ne[0]; + int64_t input_height = input->ne[1]; + GGML_ASSERT(input->type == GGML_TYPE_F32 && output->type == GGML_TYPE_F32); for (int iy = 0; iy < height; iy++) { for (int ix = 0; ix < width; ix++) { for (int k = 0; k < channels; k++) { for (int l = 0; l < ne3; l++) { - float value = ggml_ext_tensor_get_f32(input, ix + x, iy + y, k, l); + float value = ggml_ext_tensor_get_f32(input, (ix + x) % input_width, (iy + y) % input_height, k, l); ggml_ext_tensor_set_f32(output, value, ix, iy, k, l); } } @@ -516,6 +520,8 @@ __STATIC_INLINE__ void ggml_ext_tensor_merge_2d(struct ggml_tensor* input, int y, int overlap_x, int overlap_y, + bool circular_x, + bool circular_y, int x_skip = 0, int y_skip = 0) { int64_t width = input->ne[0]; @@ -533,12 +539,12 @@ __STATIC_INLINE__ void ggml_ext_tensor_merge_2d(struct ggml_tensor* input, for (int l = 0; l < ne3; l++) { float new_value = ggml_ext_tensor_get_f32(input, ix, iy, k, l); if (overlap_x > 0 || overlap_y > 0) { // blend colors in overlapped area - float old_value = ggml_ext_tensor_get_f32(output, x + ix, y + iy, k, l); + float old_value = ggml_ext_tensor_get_f32(output, (x + ix) % img_width, (y + iy) % img_height, k, l); - const float x_f_0 = (overlap_x > 0 && x > 0) ? (ix - x_skip) / float(overlap_x) : 1; - const float x_f_1 = (overlap_x > 0 && x < (img_width - width)) ? (width - ix) / float(overlap_x) : 1; - const float y_f_0 = (overlap_y > 0 && y > 0) ? (iy - y_skip) / float(overlap_y) : 1; - const float y_f_1 = (overlap_y > 0 && y < (img_height - height)) ? (height - iy) / float(overlap_y) : 1; + const float x_f_0 = (circular_x || (overlap_x > 0 && x > 0)) ? (ix - x_skip) / float(overlap_x) : 1; + const float x_f_1 = (circular_x || (overlap_x > 0 && x < (img_width - width))) ? (width - ix) / float(overlap_x) : 1; + const float y_f_0 = (circular_y || (overlap_y > 0 && y > 0)) ? (iy - y_skip) / float(overlap_y) : 1; + const float y_f_1 = (circular_y || (overlap_y > 0 && y < (img_height - height))) ? (height - iy) / float(overlap_y) : 1; const float x_f = std::min(std::min(x_f_0, x_f_1), 1.f); const float y_f = std::min(std::min(y_f_0, y_f_1), 1.f); @@ -546,9 +552,9 @@ __STATIC_INLINE__ void ggml_ext_tensor_merge_2d(struct ggml_tensor* input, ggml_ext_tensor_set_f32( output, old_value + new_value * smootherstep_f32(y_f) * smootherstep_f32(x_f), - x + ix, y + iy, k, l); + (x + ix) % img_width, (y + iy) % img_height, k, l); } else { - ggml_ext_tensor_set_f32(output, new_value, x + ix, y + iy, k, l); + ggml_ext_tensor_set_f32(output, new_value, (x + ix) % img_width, (y + iy) % img_height, k, l); } } } @@ -773,10 +779,31 @@ __STATIC_INLINE__ void sd_tiling_calc_tiles(int& num_tiles_dim, float& tile_overlap_factor_dim, int small_dim, int tile_size, - const float tile_overlap_factor) { + const float tile_overlap_factor, + bool circular) { int tile_overlap = static_cast(tile_size * tile_overlap_factor); int non_tile_overlap = tile_size - tile_overlap; + if (circular) { + // circular means the last and first tile are overlapping (wraping around) + num_tiles_dim = small_dim / non_tile_overlap; + + if (num_tiles_dim < 1) { + num_tiles_dim = 1; + } + + tile_overlap_factor_dim = (tile_size - small_dim / num_tiles_dim) / (float)tile_size; + + // if single tile and tile_overlap_factor is not 0, add one to ensure we have at least two overlapping tiles + if (num_tiles_dim == 1 && tile_overlap_factor_dim > 0) { + num_tiles_dim++; + tile_overlap_factor_dim = 0.5; + } + + return; + } + // else, non-circular means the last and first tile are not overlapping + num_tiles_dim = (small_dim - tile_overlap) / non_tile_overlap; int overshoot_dim = ((num_tiles_dim + 1) * non_tile_overlap + tile_overlap) % small_dim; @@ -805,6 +832,8 @@ __STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, const int p_tile_size_x, const int p_tile_size_y, const float tile_overlap_factor, + const bool circular_x, + const bool circular_y, on_tile_process on_processing) { output = ggml_set_f32(output, 0); @@ -829,11 +858,11 @@ __STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, int num_tiles_x; float tile_overlap_factor_x; - sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor); + sd_tiling_calc_tiles(num_tiles_x, tile_overlap_factor_x, small_width, p_tile_size_x, tile_overlap_factor, circular_x); int num_tiles_y; float tile_overlap_factor_y; - sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor); + sd_tiling_calc_tiles(num_tiles_y, tile_overlap_factor_y, small_height, p_tile_size_y, tile_overlap_factor, circular_y); 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); @@ -887,7 +916,7 @@ __STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, float last_time = 0.0f; for (int y = 0; y < small_height && !last_y; y += non_tile_overlap_y) { int dy = 0; - if (y + tile_size_y >= small_height) { + if (!circular_y && y + tile_size_y >= small_height) { int _y = y; y = small_height - tile_size_y; dy = _y - y; @@ -898,7 +927,7 @@ __STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, } for (int x = 0; x < small_width && !last_x; x += non_tile_overlap_x) { int dx = 0; - if (x + tile_size_x >= small_width) { + if (!circular_x && x + tile_size_x >= small_width) { int _x = x; x = small_width - tile_size_x; dx = _x - x; @@ -919,7 +948,7 @@ __STATIC_INLINE__ void sd_tiling_non_square(ggml_tensor* input, int64_t t1 = ggml_time_ms(); ggml_ext_tensor_split_2d(input, input_tile, x_in, y_in); if (on_processing(input_tile, output_tile, false)) { - ggml_ext_tensor_merge_2d(output_tile, output, x_out, y_out, overlap_x_out, overlap_y_out, dx, dy); + ggml_ext_tensor_merge_2d(output_tile, output, x_out, y_out, overlap_x_out, overlap_y_out, circular_x, circular_y, dx, dy); int64_t t2 = ggml_time_ms(); last_time = (t2 - t1) / 1000.0f; @@ -942,8 +971,10 @@ __STATIC_INLINE__ void sd_tiling(ggml_tensor* input, const int scale, const int tile_size, const float tile_overlap_factor, + const bool circular_x, + const bool circular_y, on_tile_process on_processing) { - sd_tiling_non_square(input, output, scale, tile_size, tile_size, tile_overlap_factor, on_processing); + sd_tiling_non_square(input, output, scale, tile_size, tile_size, tile_overlap_factor, circular_x, circular_y, on_processing); } __STATIC_INLINE__ struct ggml_tensor* ggml_ext_group_norm_32(struct ggml_context* ctx, diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index fe3850b57..36b1d2d13 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -82,6 +82,9 @@ struct SDParams { std::vector lora_paths; std::vector lora_multipliers; bool lora_dynamic = false; + + std::string cache_mode; + std::string cache_options; }; //shared @@ -765,6 +768,121 @@ static enum scheduler_t scheduler_from_name(const char * scheduler) return scheduler_t::SCHEDULER_COUNT; } +static void parse_cache_options(sd_cache_params_t & params, const std::string& cache_mode, + const std::string& cache_options) { + + sd_cache_params_init(¶ms); + if (cache_mode == "easycache") { + params.mode = SD_CACHE_EASYCACHE; + } else if (cache_mode == "ucache") { + params.mode = SD_CACHE_UCACHE; + // this is the only difference from the defaults right now + params.reuse_threshold = 1.0f; + } else if (cache_mode == "dbcache") { + params.mode = SD_CACHE_DBCACHE; + } else if (cache_mode == "taylorseer") { + params.mode = SD_CACHE_TAYLORSEER; + } else if (cache_mode == "cache-dit") { + params.mode = SD_CACHE_CACHE_DIT; + } else if (cache_mode == "spectrum") { + params.mode = SD_CACHE_SPECTRUM; + } else if (cache_mode != "" && cache_mode != "disabled") { + printf("warning: unknown cache mode '%s'", cache_mode.c_str()); + } + + if (params.mode == SD_CACHE_DISABLED) + return; + + if (cache_options == "") + return; + + sd_cache_params_t cache_params = params; + + // from examples/common/common.hpp + auto parse_named_params = [&](const std::string& opt_str) -> bool { + std::stringstream ss(opt_str); + std::string token; + while (std::getline(ss, token, ',')) { + size_t eq_pos = token.find('='); + if (eq_pos == std::string::npos) { + printf("error: cache option '%s' missing '=' separator", token.c_str()); + return false; + } + std::string key = token.substr(0, eq_pos); + std::string val = token.substr(eq_pos + 1); + try { + if (key == "threshold") { + if (cache_mode == "easycache" || cache_mode == "ucache") { + cache_params.reuse_threshold = std::stof(val); + } else { + cache_params.residual_diff_threshold = std::stof(val); + } + } else if (key == "start") { + cache_params.start_percent = std::stof(val); + } else if (key == "end") { + cache_params.end_percent = std::stof(val); + } else if (key == "decay") { + cache_params.error_decay_rate = std::stof(val); + } else if (key == "relative") { + cache_params.use_relative_threshold = (std::stof(val) != 0.0f); + } else if (key == "reset") { + cache_params.reset_error_on_compute = (std::stof(val) != 0.0f); + } else if (key == "Fn" || key == "fn") { + cache_params.Fn_compute_blocks = std::stoi(val); + } else if (key == "Bn" || key == "bn") { + cache_params.Bn_compute_blocks = std::stoi(val); + } else if (key == "warmup") { + if (cache_mode == "spectrum") { + cache_params.spectrum_warmup_steps = std::stoi(val); + } else { + cache_params.max_warmup_steps = std::stoi(val); + } + } else if (key == "w") { + cache_params.spectrum_w = std::stof(val); + } else if (key == "m") { + cache_params.spectrum_m = std::stoi(val); + } else if (key == "lam") { + cache_params.spectrum_lam = std::stof(val); + } else if (key == "window") { + cache_params.spectrum_window_size = std::stoi(val); + } else if (key == "flex") { + cache_params.spectrum_flex_window = std::stof(val); + } else if (key == "stop") { + cache_params.spectrum_stop_percent = std::stof(val); + } else { + printf("error: unknown cache parameter '%s'", key.c_str()); + return false; + } + } catch (const std::exception&) { + printf("error: invalid value '%s' for parameter '%s'", val.c_str(), key.c_str()); + return false; + } + } + + switch (cache_params.mode) { + case SD_CACHE_EASYCACHE: + case SD_CACHE_UCACHE: + if (cache_params.reuse_threshold < 0.0f) { + printf("error: cache threshold must be non-negative"); + return false; + } + if (cache_params.start_percent < 0.0f || cache_params.start_percent >= 1.0f || + cache_params.end_percent <= 0.0f || cache_params.end_percent > 1.0f || + cache_params.start_percent >= cache_params.end_percent) { + printf("error: cache start/end percents must satisfy 0.0 <= start < end <= 1.0"); + return false; + } + break; + default: ; + } + return true; + }; + + if (parse_named_params(cache_options)) { + params = cache_params; + } +} + sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) { sd_generation_outputs output; @@ -812,6 +930,9 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) SetCircularAxesAll(sd_ctx, inputs.circular_x, inputs.circular_y); + sd_params->cache_mode = inputs.cache_mode ? inputs.cache_mode : ""; + sd_params->cache_options = inputs.cache_options ? inputs.cache_options : ""; + auto loadedsdver = get_loaded_sd_version(sd_ctx); bool is_img2img = img2img_data != ""; bool is_wan = (loadedsdver == SDVersion::VERSION_WAN2 || loadedsdver == SDVersion::VERSION_WAN2_2_I2V || loadedsdver == SDVersion::VERSION_WAN2_2_TI2V); @@ -1042,6 +1163,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) params.seed = sd_params->seed; params.strength = sd_params->strength; params.vae_tiling_params.enabled = dotile; + parse_cache_options(params.cache, sd_params->cache_mode, sd_params->cache_options); params.batch_count = 1; std::vector lora_specs; diff --git a/otherarch/sdcpp/spectrum.hpp b/otherarch/sdcpp/spectrum.hpp new file mode 100644 index 000000000..0b206c180 --- /dev/null +++ b/otherarch/sdcpp/spectrum.hpp @@ -0,0 +1,195 @@ +#ifndef __SPECTRUM_HPP__ +#define __SPECTRUM_HPP__ + +#include +#include +#include + +#include "ggml_extend.hpp" + +struct SpectrumConfig { + float w = 0.40f; + int m = 3; + float lam = 1.0f; + int window_size = 2; + float flex_window = 0.50f; + int warmup_steps = 4; + float stop_percent = 0.9f; +}; + +struct SpectrumState { + SpectrumConfig config; + int cnt = 0; + int num_cached = 0; + float curr_ws = 2.0f; + int K = 6; + int stop_step = 0; + int total_steps_skipped = 0; + + std::vector> H_buf; + std::vector T_buf; + + void init(const SpectrumConfig& cfg, size_t total_steps) { + config = cfg; + cnt = 0; + num_cached = 0; + curr_ws = (float)cfg.window_size; + K = std::max(cfg.m + 1, 6); + stop_step = (int)(cfg.stop_percent * (float)total_steps); + total_steps_skipped = 0; + H_buf.clear(); + T_buf.clear(); + } + + float taus(int step_cnt) const { + return (step_cnt / 50.0f) * 2.0f - 1.0f; + } + + bool should_predict() { + if (cnt < config.warmup_steps) + return false; + if (stop_step > 0 && cnt >= stop_step) + return false; + if ((int)H_buf.size() < 2) + return false; + + int ws = std::max(1, (int)std::floor(curr_ws)); + return (num_cached + 1) % ws != 0; + } + + void update(const struct ggml_tensor* denoised) { + int64_t ne = ggml_nelements(denoised); + const float* data = (const float*)denoised->data; + + H_buf.emplace_back(data, data + ne); + T_buf.push_back(taus(cnt)); + + while ((int)H_buf.size() > K) { + H_buf.erase(H_buf.begin()); + T_buf.erase(T_buf.begin()); + } + + if (cnt >= config.warmup_steps) + curr_ws += config.flex_window; + + num_cached = 0; + cnt++; + } + + void predict(struct ggml_tensor* denoised) { + int64_t F = (int64_t)H_buf[0].size(); + int K_curr = (int)H_buf.size(); + int M1 = config.m + 1; + float tau_at = taus(cnt); + + // Design matrix X: K_curr x M1 (Chebyshev basis) + std::vector X(K_curr * M1); + for (int i = 0; i < K_curr; i++) { + X[i * M1] = 1.0f; + if (M1 > 1) + X[i * M1 + 1] = T_buf[i]; + for (int j = 2; j < M1; j++) + X[i * M1 + j] = 2.0f * T_buf[i] * X[i * M1 + j - 1] - X[i * M1 + j - 2]; + } + + // x_star: Chebyshev basis at current tau + std::vector x_star(M1); + x_star[0] = 1.0f; + if (M1 > 1) + x_star[1] = tau_at; + for (int j = 2; j < M1; j++) + x_star[j] = 2.0f * tau_at * x_star[j - 1] - x_star[j - 2]; + + // XtX = X^T X + lambda I + std::vector XtX(M1 * M1, 0.0f); + for (int i = 0; i < M1; i++) { + for (int j = 0; j < M1; j++) { + float sum = 0.0f; + for (int k = 0; k < K_curr; k++) + sum += X[k * M1 + i] * X[k * M1 + j]; + XtX[i * M1 + j] = sum + (i == j ? config.lam : 0.0f); + } + } + + // Cholesky decomposition + std::vector L(M1 * M1, 0.0f); + if (!cholesky_decompose(XtX.data(), L.data(), M1)) { + float trace = 0.0f; + for (int i = 0; i < M1; i++) + trace += XtX[i * M1 + i]; + for (int i = 0; i < M1; i++) + XtX[i * M1 + i] += 1e-4f * trace / M1; + cholesky_decompose(XtX.data(), L.data(), M1); + } + + // Solve XtX v = x_star + std::vector v(M1); + cholesky_solve(L.data(), x_star.data(), v.data(), M1); + + // Prediction weights per history entry + std::vector weights(K_curr, 0.0f); + for (int k = 0; k < K_curr; k++) + for (int j = 0; j < M1; j++) + weights[k] += X[k * M1 + j] * v[j]; + + // Blend Chebyshev and Taylor predictions + float* out = (float*)denoised->data; + float w_cheb = config.w; + float w_taylor = 1.0f - w_cheb; + const float* h_last = H_buf.back().data(); + const float* h_prev = H_buf[H_buf.size() - 2].data(); + + for (int64_t f = 0; f < F; f++) { + float pred_cheb = 0.0f; + for (int k = 0; k < K_curr; k++) + pred_cheb += weights[k] * H_buf[k][f]; + + float pred_taylor = h_last[f] + 0.5f * (h_last[f] - h_prev[f]); + + out[f] = w_taylor * pred_taylor + w_cheb * pred_cheb; + } + + num_cached++; + total_steps_skipped++; + cnt++; + } + +private: + static bool cholesky_decompose(const float* A, float* L, int n) { + std::memset(L, 0, n * n * sizeof(float)); + for (int i = 0; i < n; i++) { + for (int j = 0; j <= i; j++) { + float sum = 0.0f; + for (int k = 0; k < j; k++) + sum += L[i * n + k] * L[j * n + k]; + if (i == j) { + float diag = A[i * n + i] - sum; + if (diag <= 0.0f) + return false; + L[i * n + j] = std::sqrt(diag); + } else { + L[i * n + j] = (A[i * n + j] - sum) / L[j * n + j]; + } + } + } + return true; + } + + static void cholesky_solve(const float* L, const float* b, float* x, int n) { + std::vector y(n); + for (int i = 0; i < n; i++) { + float sum = 0.0f; + for (int j = 0; j < i; j++) + sum += L[i * n + j] * y[j]; + y[i] = (b[i] - sum) / L[i * n + i]; + } + for (int i = n - 1; i >= 0; i--) { + float sum = 0.0f; + for (int j = i + 1; j < n; j++) + sum += L[j * n + i] * x[j]; + x[i] = (y[i] - sum) / L[i * n + i]; + } + } +}; + +#endif // __SPECTRUM_HPP__ diff --git a/otherarch/sdcpp/stable-diffusion.cpp b/otherarch/sdcpp/stable-diffusion.cpp index e1983a23d..edf6f63da 100644 --- a/otherarch/sdcpp/stable-diffusion.cpp +++ b/otherarch/sdcpp/stable-diffusion.cpp @@ -16,6 +16,7 @@ #include "esrgan.hpp" #include "lora.hpp" #include "pmid.hpp" +#include "spectrum.hpp" #include "tae.hpp" #include "ucache.hpp" #include "vae.hpp" @@ -113,6 +114,9 @@ public: bool external_vae_is_invalid = false; bool free_params_immediately = false; + bool circular_x = false; + bool circular_y = false; + std::shared_ptr rng = std::make_shared(); std::shared_ptr sampler_rng = nullptr; int n_threads = -1; @@ -922,12 +926,8 @@ public: if (control_net) { control_net->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); } - if (first_stage_model) { - first_stage_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); - } - if (tae_first_stage) { - tae_first_stage->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); - } + circular_x = sd_ctx_params->circular_x; + circular_y = sd_ctx_params->circular_y; } struct ggml_init_params params; @@ -1664,7 +1664,7 @@ public: sd_progress_cb_t cb = sd_get_progress_callback(); void* cbd = sd_get_progress_callback_data(); sd_set_progress_callback((sd_progress_cb_t)suppress_pp, nullptr); - sd_tiling(input, output, scale, tile_size, tile_overlap_factor, on_processing); + sd_tiling(input, output, scale, tile_size, tile_overlap_factor, circular_x, circular_y, on_processing); sd_set_progress_callback(cb, cbd); } @@ -1873,9 +1873,11 @@ public: EasyCacheState easycache_state; UCacheState ucache_state; CacheDitConditionState cachedit_state; + SpectrumState spectrum_state; bool easycache_enabled = false; bool ucache_enabled = false; bool cachedit_enabled = false; + bool spectrum_enabled = false; if (cache_params != nullptr && cache_params->mode != SD_CACHE_DISABLED) { bool percent_valid = true; @@ -1979,6 +1981,27 @@ public: LOG_WARN("CacheDIT requested but could not be initialized for this run"); } } + } else if (cache_params->mode == SD_CACHE_SPECTRUM) { + bool spectrum_supported = sd_version_is_unet(version); + if (!spectrum_supported) { + LOG_WARN("Spectrum requested but not supported for this model type (only UNET models)"); + } else { + SpectrumConfig spectrum_config; + spectrum_config.w = cache_params->spectrum_w; + spectrum_config.m = cache_params->spectrum_m; + spectrum_config.lam = cache_params->spectrum_lam; + spectrum_config.window_size = cache_params->spectrum_window_size; + spectrum_config.flex_window = cache_params->spectrum_flex_window; + spectrum_config.warmup_steps = cache_params->spectrum_warmup_steps; + spectrum_config.stop_percent = cache_params->spectrum_stop_percent; + size_t total_steps = sigmas.size() > 0 ? sigmas.size() - 1 : 0; + spectrum_state.init(spectrum_config, total_steps); + spectrum_enabled = true; + LOG_INFO("Spectrum enabled - w: %.2f, m: %d, lam: %.2f, window: %d, flex: %.2f, warmup: %d, stop: %.0f%%", + spectrum_config.w, spectrum_config.m, spectrum_config.lam, + spectrum_config.window_size, spectrum_config.flex_window, + spectrum_config.warmup_steps, spectrum_config.stop_percent * 100.0f); + } } } @@ -2201,7 +2224,29 @@ public: timesteps_vec.assign(1, t); } - timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask); + timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask); + + if (spectrum_enabled && spectrum_state.should_predict()) { + spectrum_state.predict(denoised); + + if (denoise_mask != nullptr) { + apply_mask(denoised, init_latent, denoise_mask); + } + + if (sd_preview_cb != nullptr && sd_should_preview_denoised()) { + if (step % sd_get_preview_interval() == 0) { + preview_image(work_ctx, step, denoised, version, sd_preview_mode, preview_tensor, sd_preview_cb, sd_preview_cb_data, false); + } + } + + int64_t t1 = ggml_time_us(); + if (step > 0 || step == -(int)steps) { + int showstep = std::abs(step); + pretty_progress(showstep, (int)steps, (t1 - t0) / 1000000.f / showstep); + } + return denoised; + } + auto timesteps = vector_to_ggml_tensor(work_ctx, timesteps_vec); std::vector guidance_vec(1, guidance.distilled_guidance); auto guidance_tensor = vector_to_ggml_tensor(work_ctx, guidance_vec); @@ -2375,6 +2420,10 @@ public: vec_denoised[i] = latent_result * c_out + vec_input[i] * c_skip; } + if (spectrum_enabled) { + spectrum_state.update(denoised); + } + if (denoise_mask != nullptr) { apply_mask(denoised, init_latent, denoise_mask); } @@ -2466,6 +2515,14 @@ public: } } + if (spectrum_enabled && spectrum_state.total_steps_skipped > 0) { + size_t total_steps = sigmas.size() > 0 ? sigmas.size() - 1 : 0; + double speedup = static_cast(total_steps) / + static_cast(total_steps - spectrum_state.total_steps_skipped); + LOG_INFO("Spectrum skipped %d/%zu steps (%.2fx estimated speedup)", + spectrum_state.total_steps_skipped, total_steps, speedup); + } + if (inverse_noise_scaling) { x = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x); } @@ -2712,14 +2769,14 @@ public: tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y); } - ggml_tensor* vae_encode(ggml_context* work_ctx, ggml_tensor* x, bool encode_video = false) { + ggml_tensor* vae_encode(ggml_context* work_ctx, ggml_tensor* x) { int64_t t0 = ggml_time_ms(); ggml_tensor* result = nullptr; const int vae_scale_factor = get_vae_scale_factor(); int64_t W = x->ne[0] / vae_scale_factor; int64_t H = x->ne[1] / vae_scale_factor; int64_t C = get_latent_channel(); - if (vae_tiling_params.enabled && !encode_video) { + if (vae_tiling_params.enabled) { // TODO wan2.2 vae support? int64_t ne2; int64_t ne3; @@ -2747,7 +2804,7 @@ public: if (!use_tiny_autoencoder) { process_vae_input_tensor(x); - if (vae_tiling_params.enabled && !encode_video) { + if (vae_tiling_params.enabled) { float tile_overlap; int tile_size_x, tile_size_y; // multiply tile size for encode to keep the compute buffer size consistent @@ -2758,18 +2815,18 @@ public: auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { return first_stage_model->compute(n_threads, in, false, &out, work_ctx); }; - sd_tiling_non_square(x, result, vae_scale_factor, tile_size_x, tile_size_y, tile_overlap, on_tiling); + sd_tiling_non_square(x, result, vae_scale_factor, tile_size_x, tile_size_y, tile_overlap, circular_x, circular_y, on_tiling); } else { first_stage_model->compute(n_threads, x, false, &result, work_ctx); } first_stage_model->free_compute_buffer(); } else { - if (vae_tiling_params.enabled && !encode_video) { + if (vae_tiling_params.enabled) { // split latent in 32x32 tiles and compute in several steps auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { return tae_first_stage->compute(n_threads, in, false, &out, nullptr); }; - sd_tiling(x, result, vae_scale_factor, 64, 0.5f, on_tiling); + sd_tiling(x, result, vae_scale_factor, 64, 0.5f, circular_x, circular_y, on_tiling); } else { tae_first_stage->compute(n_threads, x, false, &result, work_ctx); } @@ -2831,7 +2888,7 @@ public: } else { latent = gaussian_latent_sample(work_ctx, vae_output); } - if (!use_tiny_autoencoder) { + if (!use_tiny_autoencoder && version != VERSION_SD1_PIX2PIX) { process_latent_in(latent); } if (sd_version_is_qwen_image(version) || sd_version_is_anima(version)) { @@ -2840,8 +2897,8 @@ public: return latent; } - ggml_tensor* encode_first_stage(ggml_context* work_ctx, ggml_tensor* x, bool encode_video = false) { - ggml_tensor* vae_output = vae_encode(work_ctx, x, encode_video); + ggml_tensor* encode_first_stage(ggml_context* work_ctx, ggml_tensor* x) { + ggml_tensor* vae_output = vae_encode(work_ctx, x); return get_first_stage_encoding(work_ctx, vae_output); } @@ -2888,7 +2945,7 @@ public: auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { return first_stage_model->compute(n_threads, in, true, &out, nullptr); }; - sd_tiling_non_square(x, result, vae_scale_factor, tile_size_x, tile_size_y, tile_overlap, on_tiling); + sd_tiling_non_square(x, result, vae_scale_factor, tile_size_x, tile_size_y, tile_overlap, circular_x, circular_y, on_tiling); } else { if (!first_stage_model->compute(n_threads, x, true, &result, work_ctx)) { LOG_ERROR("Failed to decode latetnts"); @@ -2904,7 +2961,7 @@ public: auto on_tiling = [&](ggml_tensor* in, ggml_tensor* out, bool init) { return tae_first_stage->compute(n_threads, in, true, &out); }; - sd_tiling(x, result, vae_scale_factor, 64, 0.5f, on_tiling); + sd_tiling(x, result, vae_scale_factor, 64, 0.5f, circular_x, circular_y, on_tiling); } else { if (!tae_first_stage->compute(n_threads, x, true, &result)) { LOG_ERROR("Failed to decode latetnts"); @@ -3147,6 +3204,13 @@ void sd_cache_params_init(sd_cache_params_t* cache_params) { cache_params->taylorseer_skip_interval = 1; cache_params->scm_mask = nullptr; cache_params->scm_policy_dynamic = true; + cache_params->spectrum_w = 0.40f; + cache_params->spectrum_m = 3; + cache_params->spectrum_lam = 1.0f; + cache_params->spectrum_window_size = 2; + cache_params->spectrum_flex_window = 0.50f; + cache_params->spectrum_warmup_steps = 4; + cache_params->spectrum_stop_percent = 0.9f; } void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { @@ -3727,8 +3791,9 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx, sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) { sd_ctx->sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; - int width = sd_img_gen_params->width; - int height = sd_img_gen_params->height; + + int width = sd_img_gen_params->width; + int height = sd_img_gen_params->height; int vae_scale_factor = sd_ctx->sd->get_vae_scale_factor(); int diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor(); @@ -3742,6 +3807,40 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g LOG_WARN("align up %dx%d to %dx%d (multiple=%d)", sd_img_gen_params->width, sd_img_gen_params->height, width, height, spatial_multiple); } + bool circular_x = sd_ctx->sd->circular_x; + bool circular_y = sd_ctx->sd->circular_y; + + if (!sd_img_gen_params->vae_tiling_params.enabled) { + if (sd_ctx->sd->first_stage_model) { + sd_ctx->sd->first_stage_model->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); + } + if (sd_ctx->sd->tae_first_stage) { + sd_ctx->sd->tae_first_stage->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); + } + } else { + int tile_size_x, tile_size_y; + float _overlap; + int latent_size_x = width / sd_ctx->sd->get_vae_scale_factor(); + int latent_size_y = height / sd_ctx->sd->get_vae_scale_factor(); + sd_ctx->sd->get_tile_sizes(tile_size_x, tile_size_y, _overlap, sd_img_gen_params->vae_tiling_params, latent_size_x, latent_size_y); + + // force disable circular padding for vae if tiling is enabled unless latent is smaller than tile size + // otherwise it will cause artifacts at the edges of the tiles + sd_ctx->sd->circular_x = sd_ctx->sd->circular_x && (tile_size_x >= latent_size_x); + sd_ctx->sd->circular_y = sd_ctx->sd->circular_y && (tile_size_y >= latent_size_y); + + if (sd_ctx->sd->first_stage_model) { + sd_ctx->sd->first_stage_model->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); + } + if (sd_ctx->sd->tae_first_stage) { + sd_ctx->sd->tae_first_stage->set_circular_axes(sd_ctx->sd->circular_x, sd_ctx->sd->circular_y); + } + + // disable circular tiling if it's enabled for the VAE + sd_ctx->sd->circular_x = circular_x && (tile_size_x < latent_size_x); + sd_ctx->sd->circular_y = circular_y && (tile_size_y < latent_size_y); + } + LOG_DEBUG("generate_image %dx%d", width, height); if (sd_ctx == nullptr || sd_img_gen_params == nullptr) { return nullptr; @@ -4011,6 +4110,10 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g denoise_mask, &sd_img_gen_params->cache); + // restore circular params + sd_ctx->sd->circular_x = circular_x; + sd_ctx->sd->circular_y = circular_y; + size_t t2 = ggml_time_ms(); LOG_INFO("generate_image completed in %.2fs", (t2 - t0) * 1.0f / 1000); diff --git a/otherarch/sdcpp/stable-diffusion.h b/otherarch/sdcpp/stable-diffusion.h index 062afd3c4..3518fcdd6 100644 --- a/otherarch/sdcpp/stable-diffusion.h +++ b/otherarch/sdcpp/stable-diffusion.h @@ -252,6 +252,7 @@ enum sd_cache_mode_t { SD_CACHE_DBCACHE, SD_CACHE_TAYLORSEER, SD_CACHE_CACHE_DIT, + SD_CACHE_SPECTRUM, }; typedef struct { @@ -272,6 +273,13 @@ typedef struct { int taylorseer_skip_interval; const char* scm_mask; bool scm_policy_dynamic; + float spectrum_w; + int spectrum_m; + float spectrum_lam; + int spectrum_window_size; + float spectrum_flex_window; + int spectrum_warmup_steps; + float spectrum_stop_percent; } sd_cache_params_t; typedef struct { diff --git a/otherarch/sdcpp/upscaler.cpp b/otherarch/sdcpp/upscaler.cpp index fd0dc8242..41825ee51 100644 --- a/otherarch/sdcpp/upscaler.cpp +++ b/otherarch/sdcpp/upscaler.cpp @@ -92,7 +92,8 @@ struct UpscalerGGML { return esrgan_upscaler->compute(n_threads, in, &out); }; int64_t t0 = ggml_time_ms(); - sd_tiling(input_image_tensor, upscaled, esrgan_upscaler->scale, esrgan_upscaler->tile_size, 0.25f, on_tiling); + // TODO: circular upscaling? + sd_tiling(input_image_tensor, upscaled, esrgan_upscaler->scale, esrgan_upscaler->tile_size, 0.25f, false, false, on_tiling); esrgan_upscaler->free_compute_buffer(); ggml_ext_tensor_clamp_inplace(upscaled, 0.f, 1.f); uint8_t* upscaled_data = ggml_tensor_to_sd_image(upscaled); diff --git a/tests/test_koboldcpp.py b/tests/test_koboldcpp.py index 11f942b61..d6b643848 100644 --- a/tests/test_koboldcpp.py +++ b/tests/test_koboldcpp.py @@ -85,6 +85,37 @@ def sanitize_lora_multipliers(*args, **kwargs): """ return koboldcpp.sanitize_lora_multipliers(*args, **kwargs) + +def gendefaults_parse_meta_field(*args, **kwargs): + ''' + + >>> [gendefaults_parse_meta_field(x) for x in [{}, None, '', "invalid json", ' ', 4]] + Warning: gendefaults field - not a JSON object. + Warning: couldn't parse gendefaults field. + Warning: gendefaults field - not a JSON object. + [{}, {}, {}, {}, {}, {}] + + >>> [gendefaults_parse_meta_field(x) for x in ['["valid", "json"]', 'but', '1']] + Warning: gendefaults field - not a JSON object. + Warning: couldn't parse gendefaults field. + Warning: gendefaults field - not a JSON object. + [{}, {}, {}] + + >>> gendefaults_parse_meta_field({"key": "value"}) + {'key': 'value'} + + >>> gendefaults_parse_meta_field(' "scheduler": "default", "steps": 10 ') + {'scheduler': 'default', 'steps': 10} + + >>> gendefaults_parse_meta_field('{"cfg-scale": 0.5, "cfg_scale": 0.7}') + {'cfg-scale': 0.5, 'cfg_scale': 0.7} + + >>> gendefaults_parse_meta_field('{"guidance": 1.2, "sampler": "ddim"}') + {'distilled_guidance': 1.2, 'sampler_name': 'ddim', 'guidance': 1.2, 'sampler': 'ddim'} + ''' + return koboldcpp.gendefaults_parse_meta_field(*args, **kwargs) + + if __name__ == '__main__': import doctest failures, _ = doctest.testmod()