upstream-archive/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-08-05 01:30:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

feat: add lingbot video support (#1770)

Browse Source
This commit is contained in:
leejet
2026-07-11 17:50:59 +08:00
committed by GitHub
parent c79d24bdbd
commit b5d812008e
10 changed files with 934 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
src/model/diffusion/dit.hpp
View File

@@ -104,14 +104,15 @@ namespace DiT {
return x;
}
inline ggml_tensor* patchify(ggml_context* ctx,
ggml_tensor* x,
int pt,
int ph,
int pw,
int64_t N = 1) {
inline ggml_tensor* patchify_3d(ggml_context* ctx,
ggml_tensor* x,
int pt,
int ph,
int pw,
int64_t N = 1,
bool patch_last = true) {
// x: [N*C, T, H, W]
// return: [N, h*w, C*pt*ph*pw]
// return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
int64_t C = x->ne[3] / N;
int64_t T = x->ne[2];
int64_t H = x->ne[1];
@@ -123,15 +124,20 @@ namespace DiT {
GGML_ASSERT(C * N == x->ne[3]);
GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); // [N, t_len*h_len*w_len, C*pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
if (patch_last) {
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
} else {
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C]
}
// [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
return x;
}
@@ -142,16 +148,23 @@ namespace DiT {
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
int pw,
bool patch_last = true) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C]
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
int64_t N = x->ne[2];
int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
if (patch_last) {
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
} else {
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
}
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]

692
src/model/diffusion/lingbot_video.hpp Normal file
View File

@@ -0,0 +1,692 @@
#ifndef __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
#define __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdlib>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "core/util.h"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/qwen_image.hpp"
namespace LingBotVideo {
constexpr int LINGBOT_VIDEO_GRAPH_SIZE = 65536;
struct LingBotVideoConfig {
int patch_t = 1;
int patch_h = 2;
int patch_w = 2;
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t hidden_size = 2048;
int64_t num_attention_heads = 16;
int64_t depth = 24;
int64_t intermediate_size = 6144;
int64_t text_dim = 2560;
int64_t freq_dim = 256;
float norm_eps = 1e-6f;
int rope_theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int axes_dim_sum = 128;
bool qkv_bias = false;
bool out_bias = true;
bool patch_embed_bias = true;
bool timestep_mlp_bias = true;
int64_t num_experts = 0;
int64_t num_experts_per_tok = 8;
int64_t moe_intermediate_size = 512;
int64_t decoder_sparse_step = 1;
int64_t n_shared_experts = 0;
bool norm_topk_prob = true;
float routed_scaling_factor = 1.0f;
int64_t n_group = 0;
int64_t topk_group = 0;
std::set<int> sparse_layers;
static LingBotVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
LingBotVideoConfig config;
config.depth = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
if (patch_dim % patch_volume == 0) {
config.in_channels = patch_dim / patch_volume;
}
} else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.text_dim = tensor_storage.ne[0];
} else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.freq_dim = tensor_storage.ne[0];
} else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
int64_t head_dim = tensor_storage.ne[0];
if (head_dim > 0) {
config.num_attention_heads = config.hidden_size / head_dim;
}
} else if (name.find(".attn.to_q.bias") != std::string::npos) {
config.qkv_bias = true;
} else if (name.find(".ffn.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
config.intermediate_size = tensor_storage.ne[1];
} else if (name.find(".ffn.experts.w1") != std::string::npos && tensor_storage.n_dims == 3) {
config.num_experts = tensor_storage.ne[2];
config.moe_intermediate_size = tensor_storage.ne[1];
} else if (name.find(".ffn.shared_experts.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
if (config.moe_intermediate_size > 0) {
config.n_shared_experts = tensor_storage.ne[1] / config.moe_intermediate_size;
}
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
int64_t out_dim = tensor_storage.ne[1];
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
config.out_channels = patch_volume > 0 ? out_dim / patch_volume : config.out_channels;
}
size_t block_pos = name.find("blocks.");
if (block_pos != std::string::npos) {
auto items = split_string(name.substr(block_pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.depth) {
config.depth = block_index + 1;
}
if (name.find("blocks." + std::to_string(block_index) + ".ffn.experts.w1") != std::string::npos) {
config.sparse_layers.insert(block_index);
}
}
}
}
if (config.depth == 0) {
config.depth = 24;
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
if (!config.sparse_layers.empty()) {
config.num_experts = 128;
config.num_experts_per_tok = 8;
config.moe_intermediate_size = 768;
config.decoder_sparse_step = 1;
config.n_shared_experts = 1;
config.norm_topk_prob = true;
config.n_group = 4;
config.topk_group = 2;
config.routed_scaling_factor = 2.5f;
}
LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
config.depth,
config.hidden_size,
config.num_attention_heads,
config.text_dim,
config.num_experts,
config.num_experts_per_tok,
config.n_group,
config.topk_group,
config.routed_scaling_factor,
config.sparse_layers.size());
return config;
}
};
struct LingBotVideoTextEmbedder : public GGMLBlock {
LingBotVideoTextEmbedder(int64_t text_dim,
int64_t hidden_size,
float eps = 1e-6f) {
blocks["norm"] = std::make_shared<RMSNorm>(text_dim, eps);
blocks["linear_1"] = std::make_shared<Linear>(text_dim, hidden_size, true);
blocks["linear_2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
x = norm->forward(ctx, x);
x = linear_1->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
x = linear_2->forward(ctx, x);
return x;
}
};
struct LingBotVideoAttention : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
LingBotVideoAttention(int64_t hidden_size,
int64_t num_heads,
bool qkv_bias = false,
bool out_bias = true,
float eps = 1e-6f)
: num_heads(num_heads),
head_dim(hidden_size / num_heads) {
int64_t inner_dim = num_heads * head_dim;
blocks["to_q"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
blocks["to_k"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
blocks["to_v"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["to_out"] = std::make_shared<Linear>(inner_dim, hidden_size, out_bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* attention_mask = nullptr) {
// x: [N, video_tokens + text_tokens, hidden_size]
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
int64_t S = x->ne[1];
int64_t N = x->ne[2];
auto q = to_q->forward(ctx, x);
auto k = to_k->forward(ctx, x);
auto v = to_v->forward(ctx, x);
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N);
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);
v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);
q = norm_q->forward(ctx, q);
k = norm_k->forward(ctx, k);
x = Rope::attention(ctx, q, k, v, pe, attention_mask);
x = to_out->forward(ctx, x);
return x;
}
};
struct LingBotVideoMLP : public UnaryBlock {
LingBotVideoMLP(int64_t hidden_size,
int64_t intermediate_size) {
blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
blocks["up_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
auto gate = gate_proj->forward(ctx, x);
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
auto up = up_proj->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, gate, up);
x = down_proj->forward(ctx, x);
return x;
}
};
struct LingBotVideoSparseMoeBlock : public UnaryBlock {
int64_t hidden_size;
int64_t intermediate_size;
int64_t num_experts;
int64_t num_experts_per_tok;
bool has_shared_experts;
bool norm_topk_prob;
float routed_scaling_factor;
int64_t n_group;
int64_t topk_group;
std::vector<float> group_expert_mask_vec;
bool has_correction_bias = false;
LingBotVideoSparseMoeBlock(const LingBotVideoConfig& config)
: hidden_size(config.hidden_size),
intermediate_size(config.moe_intermediate_size),
num_experts(config.num_experts),
num_experts_per_tok(config.num_experts_per_tok),
has_shared_experts(config.n_shared_experts > 0),
norm_topk_prob(config.norm_topk_prob),
routed_scaling_factor(config.routed_scaling_factor),
n_group(config.n_group),
topk_group(config.topk_group) {
if (n_group > 1) {
GGML_ASSERT(num_experts % n_group == 0);
int64_t experts_per_group = num_experts / n_group;
group_expert_mask_vec.assign(static_cast<size_t>(num_experts * n_group), 0.f);
for (int64_t group = 0; group < n_group; ++group) {
int64_t expert_begin = group * experts_per_group;
int64_t expert_end = expert_begin + experts_per_group;
for (int64_t expert = expert_begin; expert < expert_end; ++expert) {
group_expert_mask_vec[static_cast<size_t>(group * num_experts + expert)] = 1.f;
}
}
}
if (has_shared_experts) {
blocks["shared_experts"] = std::make_shared<LingBotVideoMLP>(hidden_size,
intermediate_size * config.n_shared_experts);
}
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
ggml_type router_type = get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32);
ggml_type w1_type = get_type(prefix + "experts.w1", tensor_storage_map, GGML_TYPE_F32);
ggml_type w2_type = get_type(prefix + "experts.w2", tensor_storage_map, GGML_TYPE_F32);
ggml_type w3_type = get_type(prefix + "experts.w3", tensor_storage_map, GGML_TYPE_F32);
params["router.weight"] = ggml_new_tensor_2d(ctx, router_type, hidden_size, num_experts);
if (tensor_storage_map.find(prefix + "router.e_score_correction_bias") != tensor_storage_map.end()) {
params["router.e_score_correction_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
has_correction_bias = true;
}
params["experts.w1"] = ggml_new_tensor_3d(ctx, w1_type, hidden_size, intermediate_size, num_experts);
params["experts.w2"] = ggml_new_tensor_3d(ctx, w2_type, intermediate_size, hidden_size, num_experts);
params["experts.w3"] = ggml_new_tensor_3d(ctx, w3_type, hidden_size, intermediate_size, num_experts);
}
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
const std::string& weight_name,
ggml_tensor* x,
ggml_tensor* selected_experts) {
return ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
}
ggml_tensor* select_experts(GGMLRunnerContext* ctx, ggml_tensor* choice_scores) {
ggml_context* gctx = ctx->ggml_ctx;
if (n_group <= 1 || topk_group <= 0) {
return ggml_argsort_top_k(gctx, choice_scores, static_cast<int>(num_experts_per_tok));
}
GGML_ASSERT(choice_scores->ne[0] == num_experts);
GGML_ASSERT(num_experts % n_group == 0);
GGML_ASSERT(topk_group > 0 && topk_group <= n_group);
GGML_ASSERT(!group_expert_mask_vec.empty());
const int64_t n_token_total = choice_scores->ne[1];
const int64_t experts_per_group = num_experts / n_group;
const int group_score_k = 2;
GGML_ASSERT(experts_per_group >= group_score_k);
ggml_tensor* grouped_scores = ggml_reshape_3d(gctx, choice_scores, experts_per_group, n_group, n_token_total);
ggml_tensor* group_top_ids = ggml_argsort_top_k(gctx, grouped_scores, group_score_k);
grouped_scores = ggml_reshape_3d(gctx, grouped_scores, 1, experts_per_group, n_group * n_token_total);
group_top_ids = ggml_cont(gctx, group_top_ids);
group_top_ids = ggml_reshape_2d(gctx, group_top_ids, group_score_k, n_group * n_token_total);
ggml_tensor* group_top_values = ggml_get_rows(gctx, grouped_scores, group_top_ids);
group_top_values = ggml_reshape_3d(gctx, group_top_values, group_score_k, n_group, n_token_total);
ggml_tensor* group_scores = nullptr;
for (int rank = 0; rank < group_score_k; ++rank) {
ggml_tensor* value = ggml_view_3d(gctx,
group_top_values,
1,
n_group,
n_token_total,
group_top_values->nb[1],
group_top_values->nb[2],
rank * group_top_values->nb[0]);
group_scores = group_scores == nullptr ? value : ggml_add(gctx, group_scores, value);
}
group_scores = ggml_reshape_2d(gctx, group_scores, n_group, n_token_total);
ggml_tensor* selected_groups = ggml_argsort_top_k(gctx, group_scores, static_cast<int>(topk_group));
selected_groups = ggml_cont(gctx, selected_groups);
ggml_tensor* group_expert_mask = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, 1);
ctx->bind_backend_tensor_data(group_expert_mask, group_expert_mask_vec.data());
ggml_tensor* group_expert_mask_template = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, n_token_total);
group_expert_mask = ggml_repeat(gctx, group_expert_mask, group_expert_mask_template);
ggml_tensor* selected_group_masks = ggml_get_rows(gctx, group_expert_mask, selected_groups);
ggml_tensor* selected_mask = nullptr;
for (int64_t rank = 0; rank < topk_group; ++rank) {
ggml_tensor* mask = ggml_view_3d(gctx,
selected_group_masks,
num_experts,
1,
n_token_total,
selected_group_masks->nb[1],
selected_group_masks->nb[2],
rank * selected_group_masks->nb[1]);
selected_mask = selected_mask == nullptr ? mask : ggml_add(gctx, selected_mask, mask);
}
selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
ggml_tensor* masked_scores = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, tokens, hidden_size]
GGML_ASSERT(num_experts > 0);
GGML_ASSERT(num_experts_per_tok > 0 && num_experts_per_tok <= num_experts);
ggml_tensor* residual = x;
const int64_t n_token = x->ne[1];
const int64_t N = x->ne[2];
const int64_t n_token_total = n_token * N;
ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, params["router.weight"], x);
router_logits = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
ggml_tensor* probs = ggml_sigmoid(ctx->ggml_ctx, router_logits);
ggml_tensor* choice_scores = probs;
if (has_correction_bias) {
choice_scores = ggml_add(ctx->ggml_ctx, choice_scores, params["router.e_score_correction_bias"]);
}
ggml_tensor* selected_experts = select_experts(ctx, choice_scores);
ggml_tensor* weights = ggml_get_rows(ctx->ggml_ctx,
ggml_reshape_3d(ctx->ggml_ctx, probs, 1, num_experts, n_token_total),
selected_experts);
weights = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
if (norm_topk_prob && num_experts_per_tok > 1) {
auto weights_sum = ggml_sum_rows(ctx->ggml_ctx, weights);
weights_sum = ggml_clamp(ctx->ggml_ctx, weights_sum, 6.103515625e-5f, INFINITY);
weights = ggml_div(ctx->ggml_ctx, weights, weights_sum);
}
if (routed_scaling_factor != 1.0f) {
weights = ggml_scale(ctx->ggml_ctx, weights, routed_scaling_factor);
}
weights = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total);
x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
auto gate = expert_linear(ctx, "experts.w1", x, selected_experts);
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
auto up = expert_linear(ctx, "experts.w3", x, selected_experts);
auto act = ggml_mul(ctx->ggml_ctx, gate, up);
auto out = expert_linear(ctx, "experts.w2", act, selected_experts);
out = ggml_mul(ctx->ggml_ctx, out, weights);
ggml_tensor* summed = nullptr;
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
auto expert_out = ggml_view_2d(ctx->ggml_ctx,
out,
hidden_size,
n_token_total,
out->nb[2],
i * out->nb[1]);
summed = summed == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, summed, expert_out);
}
if (num_experts_per_tok == 1) {
summed = ggml_cont(ctx->ggml_ctx, summed);
}
summed = ggml_reshape_3d(ctx->ggml_ctx, summed, hidden_size, n_token, N);
if (has_shared_experts) {
auto shared_experts = std::dynamic_pointer_cast<LingBotVideoMLP>(blocks["shared_experts"]);
summed = ggml_add(ctx->ggml_ctx, summed, shared_experts->forward(ctx, residual));
}
return summed;
}
};
struct LingBotVideoBlock : public GGMLBlock {
int64_t hidden_size;
LingBotVideoBlock(const LingBotVideoConfig& config,
bool sparse)
: hidden_size(config.hidden_size) {
blocks["norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
blocks["attn"] = std::make_shared<LingBotVideoAttention>(config.hidden_size,
config.num_attention_heads,
config.qkv_bias,
config.out_bias,
config.norm_eps);
blocks["norm_post_attn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
blocks["norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
if (sparse) {
blocks["ffn"] = std::make_shared<LingBotVideoSparseMoeBlock>(config);
} else {
blocks["ffn"] = std::make_shared<LingBotVideoMLP>(config.hidden_size, config.intermediate_size);
}
blocks["norm_post_ffn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["scale_shift_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * 6, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* temb6,
ggml_tensor* pe,
ggml_tensor* attention_mask = nullptr) {
// x: [N, tokens, hidden_size], temb6: [N, tokens, 6 * hidden_size]
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<LingBotVideoAttention>(blocks["attn"]);
auto norm_post_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto ffn = std::dynamic_pointer_cast<UnaryBlock>(blocks["ffn"]);
auto norm_post_ffn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_ffn"]);
auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
auto mod = ggml_add(ctx->ggml_ctx, temb6, table);
auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
auto shift_msa = mods[0];
auto scale_msa = mods[1];
auto gate_msa = ggml_tanh(ctx->ggml_ctx, mods[2]);
auto shift_mlp = mods[3];
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_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_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));
return x;
}
};
class LingBotVideoModel : public GGMLBlock {
public:
LingBotVideoConfig config;
LingBotVideoModel() = default;
LingBotVideoModel(LingBotVideoConfig config)
: config(config) {
int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels;
blocks["patch_embedder"] = std::make_shared<Linear>(patch_dim, config.hidden_size, config.patch_embed_bias);
blocks["time_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(config.freq_dim,
config.hidden_size,
config.hidden_size,
0,
config.timestep_mlp_bias);
blocks["time_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
blocks["text_embedder"] = std::make_shared<LingBotVideoTextEmbedder>(config.text_dim,
config.hidden_size,
config.norm_eps);
for (int i = 0; i < config.depth; i++) {
bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end();
blocks["blocks." + std::to_string(i)] = std::make_shared<LingBotVideoBlock>(config, sparse);
}
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["norm_out_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 2 * config.hidden_size, true);
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size,
config.patch_t * config.patch_h * config.patch_w * config.out_channels,
true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
// x: [N*C, T, H, W], context: [N, text_tokens, text_dim]
GGML_ASSERT(context != nullptr);
GGML_ASSERT(x->ne[3] == config.in_channels);
GGML_ASSERT(x->ne[2] % config.patch_t == 0);
GGML_ASSERT(x->ne[1] % config.patch_h == 0);
GGML_ASSERT(x->ne[0] % config.patch_w == 0);
auto patch_embedder = std::dynamic_pointer_cast<Linear>(blocks["patch_embedder"]);
auto time_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedder"]);
auto time_modulation = std::dynamic_pointer_cast<Linear>(blocks["time_modulation.1"]);
auto text_embedder = std::dynamic_pointer_cast<LingBotVideoTextEmbedder>(blocks["text_embedder"]);
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
auto norm_out_modulation = std::dynamic_pointer_cast<Linear>(blocks["norm_out_modulation.1"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int64_t N = 1;
int64_t t_len = T / config.patch_t;
int64_t h_len = H / config.patch_h;
int64_t w_len = W / config.patch_w;
int64_t n_img = t_len * h_len * w_len;
auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
img = patch_embedder->forward(ctx, img);
auto txt = text_embedder->forward(ctx, context);
auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
timestep,
static_cast<int>(config.freq_dim),
10000,
1.0f);
auto t_emb = time_embedder->forward(ctx, timestep_proj);
auto temb_template = ggml_new_tensor_3d(ctx->ggml_ctx, t_emb->type, t_emb->ne[0], hidden_states->ne[1], t_emb->ne[1]);
auto temb_tokens = ggml_repeat(ctx->ggml_ctx,
ggml_reshape_3d(ctx->ggml_ctx, t_emb, t_emb->ne[0], 1, t_emb->ne[1]),
temb_template);
auto temb6 = time_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens));
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.prelude", "hidden_states");
for (int i = 0; i < config.depth; i++) {
auto block = std::dynamic_pointer_cast<LingBotVideoBlock>(blocks["blocks." + std::to_string(i)]);
hidden_states = block->forward(ctx, hidden_states, temb6, pe);
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
}
auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
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 = proj_out->forward(ctx, hidden_states);
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
hidden_states,
t_len,
h_len,
w_len,
config.patch_t,
config.patch_h,
config.patch_w,
false);
return out;
}
};
struct LingBotVideoRunner : public DiffusionModelRunner {
LingBotVideoConfig config;
LingBotVideoModel lingbot_video;
std::vector<float> pe_vec;
LingBotVideoRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(LingBotVideoConfig::detect_from_weights(tensor_storage_map, prefix)) {
SD_UNUSED(model_args);
lingbot_video = LingBotVideoModel(config);
lingbot_video.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "lingbot_video";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
lingbot_video.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor) {
ggml_cgraph* gf = new_graph_custom(LINGBOT_VIDEO_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x_tensor.dim() == 5);
GGML_ASSERT(x->ne[3] == config.in_channels);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int64_t N = x_tensor.shape()[4];
GGML_ASSERT(N == 1);
pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
static_cast<int>(H),
static_cast<int>(W),
config.patch_t,
config.patch_h,
config.patch_w,
static_cast<int>(N),
static_cast<int>(context->ne[1]),
config.rope_theta,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context));
}
};
} // namespace LingBotVideo
#endif // __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__

2
src/model/diffusion/qwen_image.hpp
View File

@@ -517,7 +517,7 @@ namespace Qwen {
if (input->ne[3] == 1) {
input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
}
return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
};
auto img = patchify_input(x);