add support for cache modes to accelerate image generation (#2021)

* sd: sync to master-525-d6dd6d7

* sd: add support for cache modes for inference acceleration

* keep gendefaults as a JSON object inside the config file

* covered more invalid cases on gendefaults parsing
This commit is contained in:
Wagner Bruna 2026-03-15 04:27:14 -03:00 committed by GitHub
parent 893b8abc21
commit b437d18319
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 599 additions and 57 deletions

View file

@ -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;

View file

@ -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 "")

View file

@ -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()) {

View file

@ -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<int>(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,

View file

@ -82,6 +82,9 @@ struct SDParams {
std::vector<std::string> lora_paths;
std::vector<float> 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(&params);
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<sd_lora_t> lora_specs;

View file

@ -0,0 +1,195 @@
#ifndef __SPECTRUM_HPP__
#define __SPECTRUM_HPP__
#include <cmath>
#include <cstring>
#include <vector>
#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<std::vector<float>> H_buf;
std::vector<float> 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<float> 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<float> 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<float> 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<float> 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<float> v(M1);
cholesky_solve(L.data(), x_star.data(), v.data(), M1);
// Prediction weights per history entry
std::vector<float> 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<float> 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__

View file

@ -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> rng = std::make_shared<PhiloxRNG>();
std::shared_ptr<RNG> 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<float> 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<double>(total_steps) /
static_cast<double>(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);

View file

@ -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 {

View file

@ -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);

View file

@ -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()