DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 19:36:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add custom modulate operator and ConvRot fusion (#2258)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-11-02 01:19:28 +08:00
1 parent 133087deb8
commit 290148d2e3
36 files changed
+655 -161
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/diffusion/anima.hpp
+2 -2
View File
@@ -129,7 +129,7 @@ namespace Anima {
auto gate = emb_chunks[2];
auto x = norm->forward(ctx, hidden_states);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
x = Flux::modulate(ctx, x, shift, scale);
return {x, gate};
}
@@ -169,7 +169,7 @@ namespace Anima {
auto scale = emb_chunks[1];
auto x = norm->forward(ctx, hidden_states);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
x = Flux::modulate(ctx, x, shift, scale);
return x;
}
};
src/model/diffusion/boogu.hpp
+8 -8
View File
@@ -128,9 +128,9 @@ namespace Boogu {
}
};
__STATIC_INLINE__ ggml_tensor* scale_modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);
return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
__STATIC_INLINE__ ggml_tensor* scale_modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) {
scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]);
return ggml_ext_modulate(ctx, x, nullptr, scale);
}
__STATIC_INLINE__ ggml_tensor* gate_residual(ggml_context* ctx, ggml_tensor* residual, ggml_tensor* x, ggml_tensor* gate) {
@@ -186,7 +186,7 @@ namespace Boogu {
auto scale_mlp = mods[2];
auto gate_mlp = mods[3];
x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), scale_msa);
x = scale_modulate(ctx, norm->forward(ctx, x), scale_msa);
return {x, gate_msa, scale_mlp, gate_mlp};
}
};
@@ -231,7 +231,7 @@ namespace Boogu {
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
auto emb = linear_1->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning_embedding));
x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), emb);
x = scale_modulate(ctx, norm->forward(ctx, x), emb);
x = linear_2->forward(ctx, x);
return x;
}
@@ -334,7 +334,7 @@ namespace Boogu {
auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask);
hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output), gate_msa);
auto mlp_input = scale_modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp);
auto mlp_input = scale_modulate(ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp);
auto mlp_output = feed_forward->forward(ctx, mlp_input);
hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output), gate_mlp);
} else {
@@ -512,14 +512,14 @@ namespace Boogu {
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_attn_norm->forward(ctx, img_attn_out), img_gate_msa);
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_self_attn_norm->forward(ctx, img_self_attn_out), img_gate_self);
auto img_mlp_input = scale_modulate(ctx->ggml_ctx, img_norm2_out, img_scale_mlp);
auto img_mlp_input = scale_modulate(ctx, img_norm2_out, img_scale_mlp);
img_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, img_shift_mlp, img_shift_mlp->ne[0], 1, img_shift_mlp->ne[1]);
img_mlp_input = ggml_add(ctx->ggml_ctx, img_mlp_input, img_shift_mlp);
auto img_mlp_out = img_feed_forward->forward(ctx, img_ffn_norm1->forward(ctx, img_mlp_input));
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_ffn_norm2->forward(ctx, img_mlp_out), img_gate_mlp);
instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_attn_norm->forward(ctx, instruct_attn_out), instruct_gate_msa);
auto instruct_mlp_input = scale_modulate(ctx->ggml_ctx, instruct_norm2_out, instruct_scale_mlp);
auto instruct_mlp_input = scale_modulate(ctx, instruct_norm2_out, instruct_scale_mlp);
instruct_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, instruct_shift_mlp, instruct_shift_mlp->ne[0], 1, instruct_shift_mlp->ne[1]);
instruct_mlp_input = ggml_add(ctx->ggml_ctx, instruct_mlp_input, instruct_shift_mlp);
auto instruct_mlp_out = instruct_feed_forward->forward(ctx, instruct_ffn_norm1->forward(ctx, instruct_mlp_input));
src/model/diffusion/ernie_image.hpp
+3 -3
View File
@@ -246,13 +246,13 @@ namespace ErnieImage {
auto residual = x;
x = adaLN_sa_ln->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift_msa, scale_msa, true);
x = Flux::modulate(ctx, x, shift_msa, scale_msa, true);
auto attn_out = self_attention->forward(ctx, x, pe, attention_mask);
x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
residual = x;
x = adaLN_mlp_ln->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift_mlp, scale_mlp, true);
x = Flux::modulate(ctx, x, shift_mlp, scale_mlp, true);
x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, x), gate_mlp));
return x;
}
@@ -274,7 +274,7 @@ namespace ErnieImage {
auto shift = mods[1];
x = norm->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
x = Flux::modulate(ctx, x, shift, scale);
return x;
}
};
src/model/diffusion/fibo2.hpp
+2 -2
View File
@@ -262,7 +262,7 @@ namespace Fibo2 {
}
return m;
};
auto h = ZImage::modulate(c, norm("attention_norm1", x), select(0, false));
auto h = ZImage::modulate(ctx, norm("attention_norm1", x), select(0, false));
h = std::dynamic_pointer_cast<Attention>(blocks["attention"])->forward(ctx, h, h, pe);
x = ggml_add(c, x, ggml_mul(c, norm("attention_norm2", h), select(1, true)));
if (injection) {
@@ -270,7 +270,7 @@ namespace Fibo2 {
h = ggml_mul(c, norm("cross_attn_norm2", h), ggml_tanh(c, params["cross_attn_gate"]));
x = ggml_add(c, x, h);
}
h = ZImage::modulate(c, norm("ffn_norm1", x), select(2, false));
h = ZImage::modulate(ctx, norm("ffn_norm1", x), select(2, false));
h = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"])->forward(ctx, h);
return ggml_add(c, x, ggml_mul(c, norm("ffn_norm2", h), select(3, true)));
}
src/model/diffusion/flux.hpp
+10 -12
View File
@@ -413,7 +413,7 @@ namespace Flux {
}
};
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale,
@@ -422,12 +422,10 @@ namespace Flux {
// scale: [N, C]
// shift: [N, C]
if (!skip_reshape) {
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
shift = ggml_reshape_3d(ctx->ggml_ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
}
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
x = ggml_add(ctx, x, shift);
return x;
return ggml_ext_modulate(ctx, x, shift, scale);
}
struct DoubleStreamBlock : public GGMLBlock {
@@ -540,7 +538,7 @@ namespace Flux {
// prepare image for attention
auto img_modulated = img_norm1->forward(ctx, img);
img_modulated = Flux::modulate(ctx->ggml_ctx, img_modulated, img_mod1.shift, img_mod1.scale);
img_modulated = Flux::modulate(ctx, img_modulated, img_mod1.shift, img_mod1.scale);
auto img_qkv = img_attn->pre_attention(ctx, img_modulated); // q,k,v: [N, n_img_token, n_head, d_head]
auto img_q = img_qkv[0];
auto img_k = img_qkv[1];
@@ -548,7 +546,7 @@ namespace Flux {
// prepare txt for attention
auto txt_modulated = txt_norm1->forward(ctx, txt);
txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_modulated, txt_mod1.shift, txt_mod1.scale);
txt_modulated = Flux::modulate(ctx, txt_modulated, txt_mod1.shift, txt_mod1.scale);
auto txt_qkv = txt_attn->pre_attention(ctx, txt_modulated); // q,k,v: [N, n_txt_token, n_head, d_head]
auto txt_q = txt_qkv[0];
auto txt_k = txt_qkv[1];
@@ -580,14 +578,14 @@ namespace Flux {
// calculate the img bloks
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn->post_attention(ctx, img_attn_out), img_mod1.gate));
auto img_mlp_out = img_mlp->forward(ctx, Flux::modulate(ctx->ggml_ctx, img_norm2->forward(ctx, img), img_mod2.shift, img_mod2.scale));
auto img_mlp_out = img_mlp->forward(ctx, Flux::modulate(ctx, img_norm2->forward(ctx, img), img_mod2.shift, img_mod2.scale));
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp_out, img_mod2.gate));
// calculate the txt bloks
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn->post_attention(ctx, txt_attn_out), txt_mod1.gate));
auto txt_mlp_out = txt_mlp->forward(ctx, Flux::modulate(ctx->ggml_ctx, txt_norm2->forward(ctx, txt), txt_mod2.shift, txt_mod2.scale));
auto txt_mlp_out = txt_mlp->forward(ctx, Flux::modulate(ctx, txt_norm2->forward(ctx, txt), txt_mod2.shift, txt_mod2.scale));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp_out, txt_mod2.gate));
return {img, txt};
@@ -670,7 +668,7 @@ namespace Flux {
}
}
auto x_mod = Flux::modulate(ctx->ggml_ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale);
auto x_mod = Flux::modulate(ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale);
auto qkv_mlp = linear1->forward(ctx, x_mod); // [N, n_token, hidden_size * 3 + mlp_hidden_dim*mlp_mult_factor]
int64_t head_dim = hidden_size / num_heads;
@@ -752,7 +750,7 @@ namespace Flux {
scale = m_vec[1];
}
x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
x = Flux::modulate(ctx, norm_final->forward(ctx, x), shift, scale);
x = linear->forward(ctx, x);
return x;
src/model/diffusion/flux3.hpp
+1 -1
View File
@@ -83,7 +83,7 @@ namespace Flux3 {
int64_t offset = 0;
for (size_t i = 0; i < lengths.size(); ++i) {
auto part = token_slice(gc, x, offset, lengths[i]);
part = Flux::modulate(gc, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale);
part = Flux::modulate(ctx, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale);
x_mod = x_mod ? ggml_concat(gc, x_mod, part, 1) : part;
offset += lengths[i];
}
src/model/diffusion/ideogram4.hpp
+6 -6
View File
@@ -103,9 +103,9 @@ namespace Ideogram4 {
return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
}
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
scale = to_token_modulation(ctx, scale);
return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) {
scale = to_token_modulation(ctx->ggml_ctx, scale);
return ggml_ext_modulate(ctx, x, nullptr, scale);
}
__STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
@@ -285,13 +285,13 @@ namespace Ideogram4 {
auto gate_mlp = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3]));
auto attn_out = attention_norm1->forward(ctx, x);
attn_out = modulate(ctx->ggml_ctx, attn_out, scale_msa);
attn_out = modulate(ctx, attn_out, scale_msa);
attn_out = attention->forward(ctx, attn_out, pe, mask);
attn_out = attention_norm2->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
auto ffn_out = ffn_norm1->forward(ctx, x);
ffn_out = modulate(ctx->ggml_ctx, ffn_out, scale_mlp);
ffn_out = modulate(ctx, ffn_out, scale_mlp);
ffn_out = feed_forward->forward(ctx, ffn_out);
ffn_out = ffn_norm2->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp));
@@ -337,7 +337,7 @@ namespace Ideogram4 {
auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
x = norm_final->forward(ctx, x);
x = modulate(ctx->ggml_ctx, x, scale);
x = modulate(ctx, x, scale);
x = linear->forward(ctx, x);
return x;
}
src/model/diffusion/krea2.hpp
+7 -7
View File
@@ -465,8 +465,8 @@ namespace Krea2 {
auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0);
auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1);
auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true);
auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
auto attn_in_main = Flux::modulate(ctx, pre_x_main, mods_main[1], mods_main[0], true);
auto attn_in_refs = Flux::modulate(ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1);
@@ -487,8 +487,8 @@ namespace Krea2 {
auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0);
auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]);
auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true);
auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
auto mlp_in_main = Flux::modulate(ctx, post_x_main, mods_main[4], mods_main[3], true);
auto mlp_in_refs = Flux::modulate(ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1);
auto mlp_out = mlp->forward(ctx, mlp_input);
@@ -503,7 +503,7 @@ namespace Krea2 {
x = ggml_add(ctx->ggml_ctx, x, mlp_res);
} else {
auto mods = mod->forward(ctx, vec);
auto attn_input = Flux::modulate(ctx->ggml_ctx,
auto attn_input = Flux::modulate(ctx,
prenorm->forward(ctx, x),
mods[1],
mods[0],
@@ -511,7 +511,7 @@ namespace Krea2 {
auto attn_out = attn->forward(ctx, attn_input, pe);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
auto mlp_input = Flux::modulate(ctx,
postnorm->forward(ctx, x),
mods[4],
mods[3],
@@ -595,7 +595,7 @@ namespace Krea2 {
auto modulation = std::dynamic_pointer_cast<KreaFinalModulation>(blocks["modulation"]);
auto mods = modulation->forward(ctx, vec);
x = Flux::modulate(ctx->ggml_ctx,
x = Flux::modulate(ctx,
norm->forward(ctx, x),
mods[1],
mods[0],
src/model/diffusion/lens.hpp
+5 -5
View File
@@ -237,9 +237,9 @@ namespace Lens {
auto txt_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
auto img_normed = img_norm1->forward(ctx, img);
auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
auto txt_normed = txt_norm1->forward(ctx, txt);
auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
@@ -247,9 +247,9 @@ namespace Lens {
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_mod_vec[2]));
auto img_normed2 = img_norm2->forward(ctx, img);
auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
auto txt_normed2 = txt_norm2->forward(ctx, txt);
auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp->forward(ctx, img_modulated2), img_mod_vec[5]));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp->forward(ctx, txt_modulated2), txt_mod_vec[5]));
@@ -272,7 +272,7 @@ namespace Lens {
auto scale = mods[0];
auto shift = mods[1];
x = ggml_norm(ctx->ggml_ctx, x, eps);
return Flux::modulate(ctx->ggml_ctx, x, shift, scale);
return Flux::modulate(ctx, x, shift, scale);
}
};
src/model/diffusion/lingbot_video.hpp
+3 -3
View File
@@ -485,12 +485,12 @@ namespace LingBotVideo {
auto scale_mlp = mods[4];
auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]);
auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
auto attn_in = Flux::modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask);
attn_out = norm_post_attn->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out));
auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
auto ffn_in = Flux::modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
auto ffn_out = ffn->forward(ctx, ffn_in);
ffn_out = norm_post_ffn->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out));
@@ -587,7 +587,7 @@ namespace LingBotVideo {
2,
0);
hidden_states = norm_out->forward(ctx, hidden_states);
hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true);
hidden_states = Flux::modulate(ctx, hidden_states, final_mods[0], final_mods[1], true);
hidden_states = proj_out->forward(ctx, hidden_states);
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
src/model/diffusion/ltxv.hpp
+19 -19
View File
@@ -38,12 +38,12 @@ namespace LTXV {
return mod;
}
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
shift = align_token_modulation(ctx, x, shift);
scale = align_token_modulation(ctx, x, scale);
shift = align_token_modulation(ctx->ggml_ctx, x, shift);
scale = align_token_modulation(ctx->ggml_ctx, x, scale);
return Flux::modulate(ctx, x, shift, scale, true);
}
@@ -821,7 +821,7 @@ namespace LTXV {
auto gate_mlp = mods[5];
auto x_norm = rms_norm(ctx->ggml_ctx, x);
x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
x_norm = LTXV::modulate(ctx, x_norm, shift_msa, scale_msa);
auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa));
@@ -831,12 +831,12 @@ namespace LTXV {
auto gate_q = mods[8];
auto q = rms_norm(ctx->ggml_ctx, x);
q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q);
q = LTXV::modulate(ctx, q, shift_q, scale_q);
auto context_mod = context;
if (prompt_timestep != nullptr) {
auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
context_mod = LTXV::modulate(ctx, context_mod, prompt_mods[0], prompt_mods[1]);
}
auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
@@ -847,7 +847,7 @@ namespace LTXV {
}
auto y = rms_norm(ctx->ggml_ctx, x);
y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
y = LTXV::modulate(ctx, y, shift_mlp, scale_mlp);
auto mlp_out = ff->forward(ctx, y);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -1199,11 +1199,11 @@ namespace LTXV {
if (cross_attention_adaln) {
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
auto q = rms_norm(ctx->ggml_ctx, x);
q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
q = LTXV::modulate(ctx, q, q_mods[0], q_mods[1]);
auto context_mod = context;
if (prompt_timestep != nullptr && prompt_table != nullptr) {
auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
context_mod = LTXV::modulate(ctx, context_mod, p_mods[0], p_mods[1]);
}
auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
@@ -1250,7 +1250,7 @@ namespace LTXV {
auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6);
auto v_norm = rms_norm(ctx->ggml_ctx, vx);
v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
v_norm = LTXV::modulate(ctx, v_norm, v_mods[0], v_mods[1]);
auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2]));
auto v_txt = apply_text_cross_attention(ctx,
@@ -1268,7 +1268,7 @@ namespace LTXV {
if (run_ax) {
auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
auto a_norm = rms_norm(ctx->ggml_ctx, ax);
a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
a_norm = LTXV::modulate(ctx, a_norm, a_mods[0], a_mods[1]);
auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2]));
auto a_txt = apply_text_cross_attention(ctx,
@@ -1291,8 +1291,8 @@ namespace LTXV {
auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto vx_scaled = LTXV::modulate(ctx, vx_norm3, a2v_video[1], a2v_video[0]);
auto ax_scaled = LTXV::modulate(ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
@@ -1304,8 +1304,8 @@ namespace LTXV {
auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto ax_scaled = LTXV::modulate(ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
auto vx_scaled = LTXV::modulate(ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
@@ -1313,14 +1313,14 @@ namespace LTXV {
}
auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto ax_scaled = rms_norm(ctx->ggml_ctx, ax);
ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
ax_scaled = LTXV::modulate(ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
auto a_ff_out = audio_ff->forward(ctx, ax_scaled);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2]));
}
auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto vx_scaled = rms_norm(ctx->ggml_ctx, vx);
vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
vx_scaled = LTXV::modulate(ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
auto v_ff_out = ff->forward(ctx, vx_scaled);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2]));
@@ -1720,14 +1720,14 @@ namespace LTXV {
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
vx = norm_out->forward(ctx, vx);
vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = LTXV::modulate(ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = proj_out->forward(ctx, vx);
vx = unpatchify_video(ctx, vx, width, height, frames);
if (ax != nullptr && audio_time > 0) {
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
ax = audio_norm_out->forward(ctx, ax);
ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = LTXV::modulate(ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = audio_proj_out->forward(ctx, ax);
ax = unpatchify_audio(ctx, ax, audio_time);
}
src/model/diffusion/minimax_h3.hpp
+1 -3
View File
@@ -445,9 +445,7 @@ namespace MiniMaxH3 {
inputs.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end, false));
}
auto modulate = [&](ggml_tensor* value, const std::vector<ggml_tensor*>& row, int shift, int scale) {
return ggml_add(ctx->ggml_ctx,
ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[scale])),
row[shift]);
return ggml_ext_modulate(ctx, value, row[shift], row[scale]);
};
ggml_tensor* qkv = nullptr;
for (size_t i = 0; i < segments.size(); ++i) {
src/model/diffusion/mmdit.hpp
+11 -13
View File
@@ -372,18 +372,16 @@ public:
}
};
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C]
// shift: [N, C]
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
x = ggml_add(ctx, x, shift);
return x;
scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
shift = ggml_reshape_3d(ctx->ggml_ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C]
return ggml_ext_modulate(ctx, x, shift, scale);
}
struct DismantledBlock : public GGMLBlock {
@@ -455,10 +453,10 @@ public:
auto x_norm = norm1->forward(ctx, x);
auto attn_in = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
auto attn_in = modulate(ctx, x_norm, shift_msa, scale_msa);
auto qkv = attn->pre_attention(ctx, attn_in);
auto attn2_in = modulate(ctx->ggml_ctx, x_norm, shift_msa2, scale_msa2);
auto attn2_in = modulate(ctx, x_norm, shift_msa2, scale_msa2);
auto qkv2 = attn2->pre_attention(ctx, attn2_in);
return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}};
@@ -488,13 +486,13 @@ public:
auto scale_mlp = m_vec[4]; // [N, hidden_size]
auto gate_mlp = m_vec[5]; // [N, hidden_size]
auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
auto attn_in = modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
auto qkv = attn->pre_attention(ctx, attn_in);
return {qkv, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp}};
} else {
auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
auto attn_in = modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa);
auto qkv = attn->pre_attention(ctx, attn_in);
return {qkv, {nullptr, nullptr, nullptr, nullptr, nullptr}};
@@ -533,7 +531,7 @@ public:
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn2_out, gate_msa2));
auto mlp_out = mlp->forward(ctx, modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
auto mlp_out = mlp->forward(ctx, modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -565,7 +563,7 @@ public:
attn_out = attn->post_attention(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
auto mlp_out = mlp->forward(ctx, modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
auto mlp_out = mlp->forward(ctx, modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@@ -756,7 +754,7 @@ public:
auto shift = m_vec[0]; // [N, hidden_size]
auto scale = m_vec[1]; // [N, hidden_size]
x = modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale);
x = modulate(ctx, norm_final->forward(ctx, x), shift, scale);
x = linear->forward(ctx, x);
return x;
src/model/diffusion/pixart.hpp
+3 -3
View File
@@ -198,11 +198,11 @@ namespace PixArt {
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, y, g));
};
auto h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]);
auto h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]);
x = gate(attn1->forward(ctx, h, h), mv[2]);
// ada_norm_single: no norm before cross-attention (PixArtMS.py)
x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask));
h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]);
h = proj->forward(ctx, h);
h = ggml_ext_gelu(ctx->ggml_ctx, h, true);
h = fc2->forward(ctx, h);
@@ -288,7 +288,7 @@ namespace PixArt {
ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N),
2, 0);
h = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer.norm_final"])->forward(ctx, h);
h = modulate(ctx->ggml_ctx, h, parts[0], parts[1]);
h = modulate(ctx, h, parts[0], parts[1]);
h = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch]
h = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast<int>(p), static_cast<int>(p), false);
return h; // [N, out_channels, H, W]
src/model/diffusion/qwen_image.hpp
+5 -5
View File
@@ -350,11 +350,11 @@ namespace Qwen {
auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params);
auto img_normed = img_norm1->forward(ctx, img);
auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
auto img_gate1 = img_mod_param_vec[2];
auto txt_normed = txt_norm1->forward(ctx, txt);
auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]);
auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]);
auto txt_gate1 = txt_mod_param_vec[2];
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
@@ -363,11 +363,11 @@ namespace Qwen {
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_gate1));
auto img_normed2 = img_norm2->forward(ctx, img);
auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr);
auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr);
auto img_gate2 = img_mod_param_vec[5];
auto txt_normed2 = txt_norm2->forward(ctx, txt);
auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]);
auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]);
auto txt_gate2 = txt_mod_param_vec[5];
auto img_mlp_out = img_mlp->forward(ctx, img_modulated2);
@@ -407,7 +407,7 @@ namespace Qwen {
auto shift = mods[1];
x = norm->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
x = Flux::modulate(ctx, x, shift, scale);
return x;
}
src/model/diffusion/z_image.hpp
+8 -9
View File
@@ -294,17 +294,16 @@ namespace ZImage {
}
};
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing
// feeds a per-token timestep embedding so each segment carries its own modulation).
if (scale->ne[1] != x->ne[1]) {
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
}
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
return x;
return ggml_ext_modulate(ctx, x, nullptr, scale);
}
struct JointTransformerBlock : public GGMLBlock {
@@ -361,14 +360,14 @@ namespace ZImage {
auto gate_mlp = mods[3];
auto residual = x;
x = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, x), scale_msa);
x = modulate(ctx, attention_norm1->forward(ctx, x), scale_msa);
x = attention->forward(ctx, x, pe, mask);
x = attention_norm2->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_msa));
x = ggml_add(ctx->ggml_ctx, x, residual);
residual = x;
x = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, x), scale_mlp);
x = modulate(ctx, ffn_norm1->forward(ctx, x), scale_mlp);
x = feed_forward->forward(ctx, x);
x = ffn_norm2->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_mlp));
@@ -505,14 +504,14 @@ namespace ZImage {
auto gate_mlp = mods[3];
auto residual = c;
c = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, c), scale_msa);
c = modulate(ctx, attention_norm1->forward(ctx, c), scale_msa);
c = forward_attention(ctx, c, pe);
c = attention_norm2->forward(ctx, c);
c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_msa));
c = ggml_add(ctx->ggml_ctx, c, residual);
residual = c;
c = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, c), scale_mlp);
c = modulate(ctx, ffn_norm1->forward(ctx, c), scale_mlp);
c = feed_forward->forward(ctx, c);
c = ffn_norm2->forward(ctx, c);
c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_mlp));
@@ -544,7 +543,7 @@ namespace ZImage {
auto scale = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, hidden_size]
x = norm_final->forward(ctx, x);
x = modulate(ctx->ggml_ctx, x, scale);
x = modulate(ctx, x, scale);
x = linear->forward(ctx, x);
return x;