mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-10-03 11:27:43 -05:00
feat: add linear and attention scale overrides (#1964)
This commit is contained in:
@@ -237,8 +237,7 @@ namespace Anima {
|
||||
}
|
||||
auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
|
||||
auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
|
||||
attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
attn_out = ggml_ext_attention_ext(ctx,
|
||||
q_rope,
|
||||
k_rope,
|
||||
v4,
|
||||
@@ -249,8 +248,7 @@ namespace Anima {
|
||||
} else {
|
||||
auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
|
||||
auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
|
||||
attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
attn_out = ggml_ext_attention_ext(ctx,
|
||||
q_flat,
|
||||
k_flat,
|
||||
v,
|
||||
|
||||
@@ -61,7 +61,7 @@ namespace AnimateDiff {
|
||||
auto k = to_k->forward(ctx, x_pe);
|
||||
auto v = to_v->forward(ctx, x_pe);
|
||||
|
||||
auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false);
|
||||
auto a = ggml_ext_attention_ext(ctx, q, k, v, (int)num_heads, nullptr, false);
|
||||
return to_out->forward(ctx, a);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -183,7 +183,7 @@ namespace ErnieImage {
|
||||
k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, heads, S, head_dim]
|
||||
k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size]
|
||||
x = to_out_0->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -504,6 +504,10 @@ namespace HiDreamO1 {
|
||||
vision_runner->set_flash_attention_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_scale_overrides(float linear_scale, float attn_scale) override {
|
||||
vision_runner->set_scale_overrides(linear_scale, attn_scale);
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
vision_runner->set_weight_adapter(adapter);
|
||||
}
|
||||
|
||||
@@ -54,7 +54,7 @@ namespace Hunyuan {
|
||||
auto k = qkv_vec[1];
|
||||
auto v = qkv_vec[2];
|
||||
|
||||
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
|
||||
auto attn_out = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
|
||||
attn_out = self_attn_proj->forward(ctx, attn_out);
|
||||
|
||||
// adaLN_modulation
|
||||
|
||||
@@ -232,8 +232,7 @@ namespace Krea2 {
|
||||
q = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim_ * heads, Lq, N);
|
||||
k = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim_ * kv_heads, Lk, N);
|
||||
v = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim_ * kv_heads, Lk, N);
|
||||
return ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
return ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -709,8 +709,7 @@ namespace LTXV {
|
||||
k = apply_hidden_rope(ctx->ggml_ctx, k, k_pe, heads, dim_head, rope_interleaved);
|
||||
}
|
||||
|
||||
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
auto out = ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -215,8 +215,7 @@ namespace MiniMaxH3 {
|
||||
q = attention_layout(ctx->ggml_ctx, q);
|
||||
k = attention_layout(ctx->ggml_ctx, k);
|
||||
}
|
||||
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
|
||||
ctx->backend,
|
||||
auto out = ggml_ext_attention_ext(ctx,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
|
||||
@@ -365,8 +365,8 @@ public:
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x) {
|
||||
auto qkv = pre_attention(ctx, x);
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention(ctx, x); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention(ctx, x); // [N, n_token, dim]
|
||||
return x;
|
||||
}
|
||||
};
|
||||
@@ -587,8 +587,8 @@ public:
|
||||
auto qkv2 = std::get<1>(qkv_intermediates);
|
||||
auto intermediates = std::get<2>(qkv_intermediates);
|
||||
|
||||
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn2_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn2_out = ggml_ext_attention_ext(ctx, qkv2[0], qkv2[1], qkv2[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention_x(ctx,
|
||||
attn_out,
|
||||
attn2_out,
|
||||
@@ -604,7 +604,7 @@ public:
|
||||
auto qkv = qkv_intermediates.first;
|
||||
auto intermediates = qkv_intermediates.second;
|
||||
|
||||
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto attn_out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = post_attention(ctx,
|
||||
attn_out,
|
||||
intermediates[0],
|
||||
@@ -648,7 +648,7 @@ block_mixing(GGMLRunnerContext* ctx,
|
||||
qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
|
||||
}
|
||||
|
||||
auto attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size]
|
||||
auto attn = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_context + n_token, hidden_size]
|
||||
|
||||
auto context_attn = ggml_view_3d(ctx->ggml_ctx,
|
||||
attn,
|
||||
@@ -680,7 +680,7 @@ block_mixing(GGMLRunnerContext* ctx,
|
||||
}
|
||||
|
||||
if (x_block->self_attn) {
|
||||
auto attn2 = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size]
|
||||
auto attn2 = ggml_ext_attention_ext(ctx, x_qkv2[0], x_qkv2[1], x_qkv2[2], x_block->num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, hidden_size]
|
||||
|
||||
x = x_block->post_attention_x(ctx,
|
||||
x_attn,
|
||||
|
||||
@@ -193,7 +193,7 @@ namespace WAN {
|
||||
k = norm_k->forward(ctx, k);
|
||||
auto v = v_proj->forward(ctx, context); // [N, n_context, dim]
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
|
||||
x = o_proj->forward(ctx, x); // [N, n_token, dim]
|
||||
return x;
|
||||
@@ -255,8 +255,8 @@ namespace WAN {
|
||||
k_img = norm_k_img->forward(ctx, k_img);
|
||||
auto v_img = v_img_proj->forward(ctx, context_img); // [N, context_img_len, dim]
|
||||
|
||||
auto img_x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
auto img_x = ggml_ext_attention_ext(ctx, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim]
|
||||
|
||||
x = ggml_add(ctx->ggml_ctx, x, img_x);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user