Skip to content

Commit edbad43

Browse files
committed
Add flow shift parameter (for SD3 and Wan)
1 parent 797d2f9 commit edbad43

File tree

4 files changed

+20
-3
lines changed

4 files changed

+20
-3
lines changed

denoiser.hpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -382,7 +382,7 @@ struct DiscreteFlowDenoiser : public Denoiser {
382382

383383
float sigma_data = 1.0f;
384384

385-
DiscreteFlowDenoiser() {
385+
DiscreteFlowDenoiser(float shift = 3.0f) : shift(shift) {
386386
set_parameters();
387387
}
388388

examples/cli/main.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@ struct SDParams {
112112
bool chroma_use_dit_mask = true;
113113
bool chroma_use_t5_mask = false;
114114
int chroma_t5_mask_pad = 1;
115+
float flow_shift = INFINITY;
115116

116117
SDParams() {
117118
sd_sample_params_init(&sample_params);
@@ -271,6 +272,7 @@ void print_usage(int argc, const char* argv[]) {
271272
printf(" --chroma-t5-mask-pad PAD_SIZE t5 mask pad size of chroma\n");
272273
printf(" --video-frames video frames (default: 1)\n");
273274
printf(" --fps fps (default: 24)\n");
275+
printf(" --flow-shift SHIFT Shift value for Flow models like SD3.x or WAN (default: auto)");
274276
printf(" -v, --verbose print extra info\n");
275277
}
276278

@@ -493,6 +495,7 @@ void parse_args(int argc, const char** argv, SDParams& params) {
493495
{"", "--strength", "", &params.strength},
494496
{"", "--style-ratio", "", &params.style_ratio},
495497
{"", "--control-strength", "", &params.control_strength},
498+
{"", "--flow-shift", "", &params.flow_shift},
496499
};
497500

498501
options.bool_options = {
@@ -1092,6 +1095,7 @@ int main(int argc, const char* argv[]) {
10921095
params.chroma_use_dit_mask,
10931096
params.chroma_use_t5_mask,
10941097
params.chroma_t5_mask_pad,
1098+
params.flow_shift,
10951099
};
10961100

10971101
sd_ctx_t* sd_ctx = new_sd_ctx(&sd_ctx_params);

stable-diffusion.cpp

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -681,7 +681,11 @@ class StableDiffusionGGML {
681681

682682
if (sd_version_is_sd3(version)) {
683683
LOG_INFO("running in FLOW mode");
684-
denoiser = std::make_shared<DiscreteFlowDenoiser>();
684+
float shift = sd_ctx_params->flow_shift;
685+
if( shift == INFINITY){
686+
shift = 3.0;
687+
}
688+
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
685689
} else if (sd_version_is_flux(version)) {
686690
LOG_INFO("running in Flux FLOW mode");
687691
float shift = 1.0f; // TODO: validate
@@ -694,7 +698,14 @@ class StableDiffusionGGML {
694698
denoiser = std::make_shared<FluxFlowDenoiser>(shift);
695699
} else if (sd_version_is_wan(version)) {
696700
LOG_INFO("running in FLOW mode");
697-
denoiser = std::make_shared<DiscreteFlowDenoiser>();
701+
float shift = sd_ctx_params->flow_shift;
702+
if(shift == INFINITY) {
703+
shift = 5.0;
704+
if (version == VERSION_WAN2){
705+
shift = 12.0;
706+
}
707+
}
708+
denoiser = std::make_shared<DiscreteFlowDenoiser>(shift);
698709
} else if (is_using_v_parameterization) {
699710
LOG_INFO("running in v-prediction mode");
700711
denoiser = std::make_shared<CompVisVDenoiser>();
@@ -1553,6 +1564,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
15531564
sd_ctx_params->chroma_use_dit_mask = true;
15541565
sd_ctx_params->chroma_use_t5_mask = false;
15551566
sd_ctx_params->chroma_t5_mask_pad = 1;
1567+
sd_ctx_params->flow_shift = INFINITY;
15561568
}
15571569

15581570
char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {

stable-diffusion.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ typedef struct {
142142
bool chroma_use_dit_mask;
143143
bool chroma_use_t5_mask;
144144
int chroma_t5_mask_pad;
145+
float flow_shift;
145146
} sd_ctx_params_t;
146147

147148
typedef struct {

0 commit comments

Comments
 (0)