From 5c40f07d4a30dab1b44cf18d62f54efdb3ad02e5 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Sat, 28 Feb 2026 01:22:32 -0300 Subject: [PATCH] sd: sync to 0752cc9 (master-507-b314d80 +1) (#1999) * sd: sync to 0752cc9 (master-507-b314d80 +1) * sd: add flow-shift support to gendefaults --- expose.h | 1 + koboldcpp.py | 7 ++++ otherarch/sdcpp/common/common.hpp | 11 +++-- otherarch/sdcpp/denoiser.hpp | 60 +++++----------------------- otherarch/sdcpp/sdtype_adapter.cpp | 8 +++- otherarch/sdcpp/stable-diffusion.cpp | 49 ++++++++++++++--------- otherarch/sdcpp/stable-diffusion.h | 2 +- otherarch/sdcpp/vae.hpp | 2 +- otherarch/sdcpp/wan.hpp | 4 +- 9 files changed, 64 insertions(+), 80 deletions(-) diff --git a/expose.h b/expose.h index 939a13ba7..0847392ae 100644 --- a/expose.h +++ b/expose.h @@ -213,6 +213,7 @@ struct sd_generation_inputs const float cfg_scale = 0.0f; const float distilled_guidance = -1.0f; const int shifted_timestep = 0; + const float flow_shift = 0.0f; const int sample_steps = 0; const int width = 0; const int height = 0; diff --git a/koboldcpp.py b/koboldcpp.py index 03f670d4b..4c7b43ed9 100755 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -343,6 +343,7 @@ class sd_generation_inputs(ctypes.Structure): ("cfg_scale", ctypes.c_float), ("distilled_guidance", ctypes.c_float), ("shifted_timestep", ctypes.c_int), + ("flow_shift", ctypes.c_float), ("sample_steps", ctypes.c_int), ("width", ctypes.c_int), ("height", ctypes.c_int), @@ -2073,6 +2074,7 @@ def gendefaults_parse_meta_field(input_str): 'sampler': 'sampler_name', 'sampling-method': 'sampler_name', 'timestep-shift': 'shifted_timestep', + 'flow-shift': 'flow_shift', } if not isinstance(input_str, str) or not input_str.strip(): return {} @@ -2141,6 +2143,7 @@ def sd_generate(genparams): cfg_scale = tryparsefloat(genparams.get("cfg_scale", 5),5) distilled_guidance = tryparsefloat(genparams.get("distilled_guidance", None), None) shifted_timestep = tryparseint(genparams.get("shifted_timestep", None), None) + flow_shift = tryparsefloat(genparams.get("flow_shift", None), None) sample_steps = tryparseint(genparams.get("steps", 20),20) width = tryparseint(genparams.get("width", 512),512) height = tryparseint(genparams.get("height", 512),512) @@ -2164,6 +2167,8 @@ def sd_generate(genparams): distilled_guidance = None # fall back to the default if shifted_timestep is not None and (shifted_timestep < 0 or shifted_timestep > 1000): shifted_timestep = None # fall back to the default + if flow_shift is not None and flow_shift < 0: + flow_shift = None # fall back to the default sample_steps = (1 if sample_steps < 1 else (forced_steplimit if sample_steps > forced_steplimit else sample_steps)) vid_req_frames = (1 if vid_req_frames < 1 else (100 if vid_req_frames > 100 else vid_req_frames)) @@ -2189,6 +2194,8 @@ def sd_generate(genparams): inputs.denoising_strength = denoising_strength if shifted_timestep is not None: inputs.shifted_timestep = shifted_timestep + if flow_shift is not None: + inputs.flow_shift = flow_shift inputs.sample_steps = sample_steps inputs.width = width inputs.height = height diff --git a/otherarch/sdcpp/common/common.hpp b/otherarch/sdcpp/common/common.hpp index 749ab2b86..4150bf9d7 100644 --- a/otherarch/sdcpp/common/common.hpp +++ b/otherarch/sdcpp/common/common.hpp @@ -586,10 +586,6 @@ struct SDContextParams { "--vae-tile-overlap", "tile overlap for vae tiling, in fraction of tile size (default: 0.5)", &vae_tiling_params.target_overlap}, - {"", - "--flow-shift", - "shift value for Flow models like SD3.x or WAN (default: auto)", - &flow_shift}, }; options.bool_options = { @@ -908,7 +904,6 @@ struct SDContextParams { << " photo_maker_path: \"" << photo_maker_path << "\",\n" << " rng_type: " << sd_rng_type_name(rng_type) << ",\n" << " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n" - << " flow_shift: " << (std::isinf(flow_shift) ? "INF" : std::to_string(flow_shift)) << "\n" << " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n" << " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n" << " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n" @@ -991,7 +986,6 @@ struct SDContextParams { chroma_use_t5_mask, chroma_t5_mask_pad, qwen_image_zero_cond_t, - flow_shift, }; return sd_ctx_params; } @@ -1211,6 +1205,10 @@ struct SDGenerationParams { "--eta", "eta in DDIM, only for DDIM and TCD (default: 0)", &sample_params.eta}, + {"", + "--flow-shift", + "shift value for Flow models like SD3.x or WAN (default: auto)", + &sample_params.flow_shift}, {"", "--high-noise-cfg-scale", "(high noise) unconditional guidance scale: (default: 7.0)", @@ -1611,6 +1609,7 @@ struct SDGenerationParams { load_if_exists("cfg_scale", sample_params.guidance.txt_cfg); load_if_exists("img_cfg_scale", sample_params.guidance.img_cfg); load_if_exists("guidance", sample_params.guidance.distilled_guidance); + load_if_exists("flow_shift", sample_params.flow_shift); auto load_sampler_if_exists = [&](const char* key, enum sample_method_t& out) { if (j.contains(key) && j[key].is_string()) { diff --git a/otherarch/sdcpp/denoiser.hpp b/otherarch/sdcpp/denoiser.hpp index 7e99b84a8..40bd7cb7f 100644 --- a/otherarch/sdcpp/denoiser.hpp +++ b/otherarch/sdcpp/denoiser.hpp @@ -657,9 +657,8 @@ struct DiscreteFlowDenoiser : public Denoiser { float sigma_data = 1.0f; - DiscreteFlowDenoiser(float shift = 3.0f) - : shift(shift) { - set_parameters(); + DiscreteFlowDenoiser(float shift = 3.0f) { + set_shift(shift); } void set_parameters() { @@ -668,6 +667,11 @@ struct DiscreteFlowDenoiser : public Denoiser { } } + void set_shift(float shift) { + this->shift = shift; + set_parameters(); + } + float sigma_min() override { return sigmas[0]; } @@ -710,34 +714,8 @@ float flux_time_shift(float mu, float sigma, float t) { return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma)); } -struct FluxFlowDenoiser : public Denoiser { - float sigmas[TIMESTEPS]; - float shift = 1.15f; - - float sigma_data = 1.0f; - - FluxFlowDenoiser(float shift = 1.15f) { - set_parameters(shift); - } - - void set_shift(float shift) { - this->shift = shift; - } - - void set_parameters(float shift) { - set_shift(shift); - for (int i = 0; i < TIMESTEPS; i++) { - sigmas[i] = t_to_sigma(static_cast(i)); - } - } - - float sigma_min() override { - return sigmas[0]; - } - - float sigma_max() override { - return sigmas[TIMESTEPS - 1]; - } +struct FluxFlowDenoiser : public DiscreteFlowDenoiser { + FluxFlowDenoiser() = default; float sigma_to_t(float sigma) override { return sigma; @@ -747,26 +725,6 @@ struct FluxFlowDenoiser : public Denoiser { t = t + 1; return flux_time_shift(shift, 1.0f, t / TIMESTEPS); } - - std::vector get_scalings(float sigma) override { - float c_skip = 1.0f; - float c_out = -sigma; - float c_in = 1.0f; - return {c_skip, c_out, c_in}; - } - - // this function will modify noise/latent - ggml_tensor* noise_scaling(float sigma, ggml_tensor* noise, ggml_tensor* latent) override { - ggml_ext_tensor_scale_inplace(noise, sigma); - ggml_ext_tensor_scale_inplace(latent, 1.0f - sigma); - ggml_ext_tensor_add_inplace(latent, noise); - return latent; - } - - ggml_tensor* inverse_noise_scaling(float sigma, ggml_tensor* latent) override { - ggml_ext_tensor_scale_inplace(latent, 1.0f / (1.0f - sigma)); - return latent; - } }; struct Flux2FlowDenoiser : public FluxFlowDenoiser { diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index a65811a6c..e88809420 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -69,6 +69,7 @@ struct SDParams { int sample_steps = 20; float distilled_guidance = -1.0f; float shifted_timestep = 0; + float flow_shift = -1.0f; float strength = 0.75f; int64_t seed = 42; bool clip_on_cpu = false; @@ -370,7 +371,6 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { params.keep_vae_on_cpu = inputs.vae_cpu; params.keep_clip_on_cpu = inputs.clip_cpu; params.lora_apply_mode = (lora_apply_mode_t)lora_apply_mode; - // params.flow_shift = 5.0f; // also switches flash attn for the vae and conditioner params.flash_attn = params.diffusion_flash_attn; @@ -492,6 +492,8 @@ static std::string get_image_params(const sd_img_gen_params_t & params) { << get_scheduler_name(params.sample_params.scheduler, true); if (params.sample_params.shifted_timestep != 0) ss << "| Timestep Shift: " << params.sample_params.shifted_timestep; + if (params.sample_params.flow_shift > 0.f && params.sample_params.flow_shift != INFINITY) + ss << "| Flow Shift: " << params.sample_params.flow_shift; ss << " | Clip skip: " << params.clip_skip << " | Model: " << sdmodelfilename << " | Version: KoboldCpp"; @@ -785,6 +787,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) sd_params->distilled_guidance = inputs.distilled_guidance; sd_params->sample_steps = inputs.sample_steps; sd_params->shifted_timestep = inputs.shifted_timestep; + sd_params->flow_shift = inputs.flow_shift; sd_params->seed = inputs.seed; sd_params->width = inputs.width; sd_params->height = inputs.height; @@ -1022,6 +1025,9 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs) params.sample_params.scheduler = sd_params->scheduler; params.sample_params.sample_steps = sd_params->sample_steps; params.sample_params.shifted_timestep = sd_params->shifted_timestep; + if (sd_params->flow_shift > 0.f && sd_params->flow_shift != INFINITY) { + params.sample_params.flow_shift = sd_params->flow_shift; + } params.seed = sd_params->seed; params.strength = sd_params->strength; params.vae_tiling_params.enabled = dotile; diff --git a/otherarch/sdcpp/stable-diffusion.cpp b/otherarch/sdcpp/stable-diffusion.cpp index 8176e1dca..773ff731d 100644 --- a/otherarch/sdcpp/stable-diffusion.cpp +++ b/otherarch/sdcpp/stable-diffusion.cpp @@ -117,6 +117,7 @@ public: int n_threads = -1; float scale_factor = 0.18215f; float shift_factor = 0.f; + float default_flow_shift = INFINITY; std::shared_ptr cond_stage_model; std::shared_ptr clip_vision; // for svd or wan2.1 i2v @@ -1028,7 +1029,6 @@ public: // init denoiser { prediction_t pred_type = sd_ctx_params->prediction; - float flow_shift = sd_ctx_params->flow_shift; if (pred_type == PREDICTION_COUNT) { if (sd_version_is_sd2(version)) { @@ -1053,22 +1053,19 @@ public: sd_version_is_qwen_image(version) || sd_version_is_z_image(version)) { pred_type = FLOW_PRED; - if (flow_shift == INFINITY) { - if (sd_version_is_wan(version)) { - flow_shift = 5.f; - } else { - flow_shift = 3.f; - } + if (sd_version_is_wan(version)) { + default_flow_shift = 5.f; + } else { + default_flow_shift = 3.f; } } else if (sd_version_is_flux(version)) { pred_type = FLUX_FLOW_PRED; - if (flow_shift == INFINITY) { - flow_shift = 1.0f; // TODO: validate - for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { - flow_shift = 1.15f; - } + default_flow_shift = 1.0f; // TODO: validate + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { + default_flow_shift = 1.15f; + break; } } } else if (sd_version_is_flux2(version)) { @@ -1092,12 +1089,12 @@ public: break; case FLOW_PRED: { LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(flow_shift); + denoiser = std::make_shared(); break; } case FLUX_FLOW_PRED: { LOG_INFO("running in Flux FLOW mode"); - denoiser = std::make_shared(flow_shift); + denoiser = std::make_shared(); break; } case FLUX2_FLOW_PRED: { @@ -2859,6 +2856,16 @@ public: return result; } + void set_flow_shift(float flow_shift = INFINITY) { + auto flow_denoiser = std::dynamic_pointer_cast(denoiser); + if (flow_denoiser) { + if (flow_shift == INFINITY) { + flow_shift = default_flow_shift; + } + flow_denoiser->set_shift(flow_shift); + } + } + //added for kcpp void SetCircularAxesAll(bool circular_x, bool circular_y) { @@ -3098,7 +3105,6 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->chroma_use_dit_mask = true; sd_ctx_params->chroma_use_t5_mask = false; sd_ctx_params->chroma_t5_mask_pad = 1; - sd_ctx_params->flow_shift = INFINITY; } char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { @@ -3190,6 +3196,7 @@ void sd_sample_params_init(sd_sample_params_t* sample_params) { sample_params->sample_steps = 20; sample_params->custom_sigmas = nullptr; sample_params->custom_sigmas_count = 0; + sample_params->flow_shift = INFINITY; } char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { @@ -3210,7 +3217,8 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { "sample_method: %s, " "sample_steps: %d, " "eta: %.2f, " - "shifted_timestep: %d)", + "shifted_timestep: %d, " + "flow_shift: %.2f)", sample_params->guidance.txt_cfg, std::isfinite(sample_params->guidance.img_cfg) ? sample_params->guidance.img_cfg @@ -3224,7 +3232,8 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) { sd_sample_method_name(sample_params->sample_method), sample_params->sample_steps, sample_params->eta, - sample_params->shifted_timestep); + sample_params->shifted_timestep, + sample_params->flow_shift); return buf; } @@ -3695,6 +3704,8 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g size_t t0 = ggml_time_ms(); + sd_ctx->sd->set_flow_shift(sd_img_gen_params->sample_params.flow_shift); + // Apply lora sd_ctx->sd->apply_loras(sd_img_gen_params->loras, sd_img_gen_params->lora_count); @@ -3970,6 +3981,8 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s } LOG_INFO("generate_video %dx%dx%d", width, height, frames); + sd_ctx->sd->set_flow_shift(sd_vid_gen_params->sample_params.flow_shift); + enum sample_method_t sample_method = sd_vid_gen_params->sample_params.sample_method; if (sample_method == SAMPLE_METHOD_COUNT) { sample_method = sd_get_default_sample_method(sd_ctx); diff --git a/otherarch/sdcpp/stable-diffusion.h b/otherarch/sdcpp/stable-diffusion.h index ac5e7c9a9..8b3545295 100644 --- a/otherarch/sdcpp/stable-diffusion.h +++ b/otherarch/sdcpp/stable-diffusion.h @@ -201,7 +201,6 @@ typedef struct { bool chroma_use_t5_mask; int chroma_t5_mask_pad; bool qwen_image_zero_cond_t; - float flow_shift; } sd_ctx_params_t; typedef struct { @@ -235,6 +234,7 @@ typedef struct { int shifted_timestep; float* custom_sigmas; int custom_sigmas_count; + float flow_shift; } sd_sample_params_t; typedef struct { diff --git a/otherarch/sdcpp/vae.hpp b/otherarch/sdcpp/vae.hpp index 01081343b..c627616c2 100644 --- a/otherarch/sdcpp/vae.hpp +++ b/otherarch/sdcpp/vae.hpp @@ -141,7 +141,7 @@ public: v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n); // [N, h * w, in_channels] } - h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, true, ctx->flash_attn_enabled); + h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); if (use_linear) { h_ = proj_out->forward(ctx, h_); // [N, h * w, in_channels] diff --git a/otherarch/sdcpp/wan.hpp b/otherarch/sdcpp/wan.hpp index 7b1059785..90de3bdd1 100644 --- a/otherarch/sdcpp/wan.hpp +++ b/otherarch/sdcpp/wan.hpp @@ -572,8 +572,8 @@ namespace WAN { auto v = qkv_vec[2]; v = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n); // [t, c, h * w] - v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c] - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, true, ctx->flash_attn_enabled); // [t, h * w, c] + v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c] + x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [t, h * w, c] x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [t, c, h * w] x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n); // [t, c, h, w]