Skip to content

Commit 65fa646

Browse files
stduhpfleejet
stduhpf
and
leejet
authored
feat: add sd3.5 medium and skip layer guidance support (#451)
* mmdit-x * add support for sd3.5 medium * add skip layer guidance support (mmdit only) * ignore slg if slg_scale is zero (optimization) * init out_skip once * slg support for flux (expermiental) * warn if version doesn't support slg * refactor slg cli args * set default slg_scale to 0 (oops) * format code --------- Co-authored-by: leejet <leejet714@gmail.com>
1 parent ac54e00 commit 65fa646

File tree

9 files changed

+416
-81
lines changed

9 files changed

+416
-81
lines changed

‎diffusion_model.hpp‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ struct DiffusionModel {
1717
std::vector<struct ggml_tensor*> controls = {},
1818
float control_strength = 0.f,
1919
struct ggml_tensor** output = NULL,
20-
struct ggml_context* output_ctx = NULL) = 0;
20+
struct ggml_context* output_ctx = NULL,
21+
std::vector<int> skip_layers = std::vector<int>()) = 0;
2122
virtual void alloc_params_buffer() = 0;
2223
virtual void free_params_buffer() = 0;
2324
virtual void free_compute_buffer() = 0;
@@ -70,7 +71,9 @@ struct UNetModel : public DiffusionModel {
7071
std::vector<struct ggml_tensor*> controls = {},
7172
float control_strength = 0.f,
7273
struct ggml_tensor** output = NULL,
73-
struct ggml_context* output_ctx = NULL) {
74+
struct ggml_context* output_ctx = NULL,
75+
std::vector<int> skip_layers = std::vector<int>()) {
76+
(void)skip_layers; // SLG doesn't work with UNet models
7477
return unet.compute(n_threads, x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, output, output_ctx);
7578
}
7679
};
@@ -119,8 +122,9 @@ struct MMDiTModel : public DiffusionModel {
119122
std::vector<struct ggml_tensor*> controls = {},
120123
float control_strength = 0.f,
121124
struct ggml_tensor** output = NULL,
122-
struct ggml_context* output_ctx = NULL) {
123-
return mmdit.compute(n_threads, x, timesteps, context, y, output, output_ctx);
125+
struct ggml_context* output_ctx = NULL,
126+
std::vector<int> skip_layers = std::vector<int>()) {
127+
return mmdit.compute(n_threads, x, timesteps, context, y, output, output_ctx, skip_layers);
124128
}
125129
};
126130

@@ -168,8 +172,9 @@ struct FluxModel : public DiffusionModel {
168172
std::vector<struct ggml_tensor*> controls = {},
169173
float control_strength = 0.f,
170174
struct ggml_tensor** output = NULL,
171-
struct ggml_context* output_ctx = NULL) {
172-
return flux.compute(n_threads, x, timesteps, context, y, guidance, output, output_ctx);
175+
struct ggml_context* output_ctx = NULL,
176+
std::vector<int> skip_layers = std::vector<int>()) {
177+
return flux.compute(n_threads, x, timesteps, context, y, guidance, output, output_ctx, skip_layers);
173178
}
174179
};
175180

‎examples/cli/main.cpp‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,11 @@ struct SDParams {
119119
bool canny_preprocess = false;
120120
bool color = false;
121121
int upscale_repeats = 1;
122+
123+
std::vector<int> skip_layers = {7, 8, 9};
124+
float slg_scale = 0.;
125+
float skip_layer_start = 0.01;
126+
float skip_layer_end = 0.2;
122127
};
123128

124129
void print_params(SDParams params) {
@@ -151,6 +156,7 @@ void print_params(SDParams params) {
151156
printf(" negative_prompt: %s\n", params.negative_prompt.c_str());
152157
printf(" min_cfg: %.2f\n", params.min_cfg);
153158
printf(" cfg_scale: %.2f\n", params.cfg_scale);
159+
printf(" slg_scale: %.2f\n", params.slg_scale);
154160
printf(" guidance: %.2f\n", params.guidance);
155161
printf(" clip_skip: %d\n", params.clip_skip);
156162
printf(" width: %d\n", params.width);
@@ -197,6 +203,12 @@ void print_usage(int argc, const char* argv[]) {
197203
printf(" -p, --prompt [PROMPT] the prompt to render\n");
198204
printf(" -n, --negative-prompt PROMPT the negative prompt (default: \"\")\n");
199205
printf(" --cfg-scale SCALE unconditional guidance scale: (default: 7.0)\n");
206+
printf(" --slg-scale SCALE skip layer guidance (SLG) scale, only for DiT models: (default: 0)\n");
207+
printf(" 0 means disabled, a value of 2.5 is nice for sd3.5 medium\n");
208+
printf(" --skip_layers LAYERS Layers to skip for SLG steps: (default: [7,8,9])\n");
209+
printf(" --skip_layer_start START SLG enabling point: (default: 0.01)\n");
210+
printf(" --skip_layer_end END SLG disabling point: (default: 0.2)\n");
211+
printf(" SLG will be enabled at step int([STEPS]*[START]) and disabled at int([STEPS]*[END])\n");
200212
printf(" --strength STRENGTH strength for noising/unnoising (default: 0.75)\n");
201213
printf(" --style-ratio STYLE-RATIO strength for keeping input identity (default: 20%%)\n");
202214
printf(" --control-strength STRENGTH strength to apply Control Net (default: 0.9)\n");
@@ -534,6 +546,61 @@ void parse_args(int argc, const char** argv, SDParams& params) {
534546
params.verbose = true;
535547
} else if (arg == "--color") {
536548
params.color = true;
549+
} else if (arg == "--slg-scale") {
550+
if (++i >= argc) {
551+
invalid_arg = true;
552+
break;
553+
}
554+
params.slg_scale = std::stof(argv[i]);
555+
} else if (arg == "--skip-layers") {
556+
if (++i >= argc) {
557+
invalid_arg = true;
558+
break;
559+
}
560+
if (argv[i][0] != '[') {
561+
invalid_arg = true;
562+
break;
563+
}
564+
std::string layers_str = argv[i];
565+
while (layers_str.back() != ']') {
566+
if (++i >= argc) {
567+
invalid_arg = true;
568+
break;
569+
}
570+
layers_str += " " + std::string(argv[i]);
571+
}
572+
layers_str = layers_str.substr(1, layers_str.size() - 2);
573+
574+
std::regex regex("[, ]+");
575+
std::sregex_token_iterator iter(layers_str.begin(), layers_str.end(), regex, -1);
576+
std::sregex_token_iterator end;
577+
std::vector<std::string> tokens(iter, end);
578+
std::vector<int> layers;
579+
for (const auto& token : tokens) {
580+
try {
581+
layers.push_back(std::stoi(token));
582+
} catch (const std::invalid_argument& e) {
583+
invalid_arg = true;
584+
break;
585+
}
586+
}
587+
params.skip_layers = layers;
588+
589+
if (invalid_arg) {
590+
break;
591+
}
592+
} else if (arg == "--skip-layer-start") {
593+
if (++i >= argc) {
594+
invalid_arg = true;
595+
break;
596+
}
597+
params.skip_layer_start = std::stof(argv[i]);
598+
} else if (arg == "--skip-layer-end") {
599+
if (++i >= argc) {
600+
invalid_arg = true;
601+
break;
602+
}
603+
params.skip_layer_end = std::stof(argv[i]);
537604
} else {
538605
fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
539606
print_usage(argc, argv);
@@ -624,6 +691,16 @@ std::string get_image_params(SDParams params, int64_t seed) {
624691
}
625692
parameter_string += "Steps: " + std::to_string(params.sample_steps) + ", ";
626693
parameter_string += "CFG scale: " + std::to_string(params.cfg_scale) + ", ";
694+
if (params.slg_scale != 0 && params.skip_layers.size() != 0) {
695+
parameter_string += "SLG scale: " + std::to_string(params.cfg_scale) + ", ";
696+
parameter_string += "Skip layers: [";
697+
for (const auto& layer : params.skip_layers) {
698+
parameter_string += std::to_string(layer) + ", ";
699+
}
700+
parameter_string += "], ";
701+
parameter_string += "Skip layer start: " + std::to_string(params.skip_layer_start) + ", ";
702+
parameter_string += "Skip layer end: " + std::to_string(params.skip_layer_end) + ", ";
703+
}
627704
parameter_string += "Guidance: " + std::to_string(params.guidance) + ", ";
628705
parameter_string += "Seed: " + std::to_string(seed) + ", ";
629706
parameter_string += "Size: " + std::to_string(params.width) + "x" + std::to_string(params.height) + ", ";
@@ -840,7 +917,11 @@ int main(int argc, const char* argv[]) {
840917
params.control_strength,
841918
params.style_ratio,
842919
params.normalize_input,
843-
params.input_id_images_path.c_str());
920+
params.input_id_images_path.c_str(),
921+
params.skip_layers,
922+
params.slg_scale,
923+
params.skip_layer_start,
924+
params.skip_layer_end);
844925
} else {
845926
sd_image_t input_image = {(uint32_t)params.width,
846927
(uint32_t)params.height,

‎flux.hpp‎

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -711,7 +711,8 @@ namespace Flux {
711711
struct ggml_tensor* timesteps,
712712
struct ggml_tensor* y,
713713
struct ggml_tensor* guidance,
714-
struct ggml_tensor* pe) {
714+
struct ggml_tensor* pe,
715+
std::vector<int> skip_layers = std::vector<int>()) {
715716
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
716717
auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
717718
auto vector_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["vector_in"]);
@@ -733,6 +734,10 @@ namespace Flux {
733734
txt = txt_in->forward(ctx, txt);
734735

735736
for (int i = 0; i < params.depth; i++) {
737+
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
738+
continue;
739+
}
740+
736741
auto block = std::dynamic_pointer_cast<DoubleStreamBlock>(blocks["double_blocks." + std::to_string(i)]);
737742

738743
auto img_txt = block->forward(ctx, img, txt, vec, pe);
@@ -742,6 +747,9 @@ namespace Flux {
742747

743748
auto txt_img = ggml_concat(ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
744749
for (int i = 0; i < params.depth_single_blocks; i++) {
750+
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + params.depth) != skip_layers.end()) {
751+
continue;
752+
}
745753
auto block = std::dynamic_pointer_cast<SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
746754

747755
txt_img = block->forward(ctx, txt_img, vec, pe);
@@ -769,7 +777,8 @@ namespace Flux {
769777
struct ggml_tensor* context,
770778
struct ggml_tensor* y,
771779
struct ggml_tensor* guidance,
772-
struct ggml_tensor* pe) {
780+
struct ggml_tensor* pe,
781+
std::vector<int> skip_layers = std::vector<int>()) {
773782
// Forward pass of DiT.
774783
// x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
775784
// timestep: (N,) tensor of diffusion timesteps
@@ -791,7 +800,7 @@ namespace Flux {
791800
// img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
792801
auto img = patchify(ctx, x, patch_size); // [N, h*w, C * patch_size * patch_size]
793802

794-
auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe); // [N, h*w, C * patch_size * patch_size]
803+
auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, skip_layers); // [N, h*w, C * patch_size * patch_size]
795804

796805
// rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)
797806
out = unpatchify(ctx, out, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size); // [N, C, H + pad_h, W + pad_w]
@@ -829,7 +838,8 @@ namespace Flux {
829838
struct ggml_tensor* timesteps,
830839
struct ggml_tensor* context,
831840
struct ggml_tensor* y,
832-
struct ggml_tensor* guidance) {
841+
struct ggml_tensor* guidance,
842+
std::vector<int> skip_layers = std::vector<int>()) {
833843
GGML_ASSERT(x->ne[3] == 1);
834844
struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, FLUX_GRAPH_SIZE, false);
835845

@@ -856,7 +866,8 @@ namespace Flux {
856866
context,
857867
y,
858868
guidance,
859-
pe);
869+
pe,
870+
skip_layers);
860871

861872
ggml_build_forward_expand(gf, out);
862873

@@ -870,14 +881,15 @@ namespace Flux {
870881
struct ggml_tensor* y,
871882
struct ggml_tensor* guidance,
872883
struct ggml_tensor** output = NULL,
873-
struct ggml_context* output_ctx = NULL) {
884+
struct ggml_context* output_ctx = NULL,
885+
std::vector<int> skip_layers = std::vector<int>()) {
874886
// x: [N, in_channels, h, w]
875887
// timesteps: [N, ]
876888
// context: [N, max_position, hidden_size]
877889
// y: [N, adm_in_channels] or [1, adm_in_channels]
878890
// guidance: [N, ]
879891
auto get_graph = [&]() -> struct ggml_cgraph* {
880-
return build_graph(x, timesteps, context, y, guidance);
892+
return build_graph(x, timesteps, context, y, guidance, skip_layers);
881893
};
882894

883895
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);

0 commit comments

Comments
 (0)