mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-03-24 18:28:57 +00:00
Compare commits
4 Commits
master-533
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
545fac4f3f | ||
|
|
5265a5efa1 | ||
|
|
84cbd88df1 | ||
|
|
997bb11fb6 |
130
src/anima.hpp
130
src/anima.hpp
@ -13,9 +13,9 @@
|
|||||||
namespace Anima {
|
namespace Anima {
|
||||||
constexpr int ANIMA_GRAPH_SIZE = 65536;
|
constexpr int ANIMA_GRAPH_SIZE = 65536;
|
||||||
|
|
||||||
__STATIC_INLINE__ struct ggml_tensor* apply_gate(struct ggml_context* ctx,
|
__STATIC_INLINE__ ggml_tensor* apply_gate(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* gate) {
|
ggml_tensor* gate) {
|
||||||
gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]); // [N, 1, C]
|
gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]); // [N, 1, C]
|
||||||
return ggml_mul(ctx, x, gate);
|
return ggml_mul(ctx, x, gate);
|
||||||
}
|
}
|
||||||
@ -26,7 +26,7 @@ namespace Anima {
|
|||||||
blocks["proj.1"] = std::make_shared<Linear>(in_dim, out_dim, false);
|
blocks["proj.1"] = std::make_shared<Linear>(in_dim, out_dim, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj.1"]);
|
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj.1"]);
|
||||||
return proj->forward(ctx, x);
|
return proj->forward(ctx, x);
|
||||||
}
|
}
|
||||||
@ -39,7 +39,7 @@ namespace Anima {
|
|||||||
blocks["1.linear_2"] = std::make_shared<Linear>(in_dim, out_dim, false);
|
blocks["1.linear_2"] = std::make_shared<Linear>(in_dim, out_dim, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_1"]);
|
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_1"]);
|
||||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_2"]);
|
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_2"]);
|
||||||
|
|
||||||
@ -62,10 +62,10 @@ namespace Anima {
|
|||||||
blocks["2"] = std::make_shared<Linear>(hidden_features, 3 * in_features, false);
|
blocks["2"] = std::make_shared<Linear>(hidden_features, 3 * in_features, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* hidden_states,
|
ggml_tensor* hidden_states,
|
||||||
struct ggml_tensor* embedded_timestep,
|
ggml_tensor* embedded_timestep,
|
||||||
struct ggml_tensor* temb = nullptr) {
|
ggml_tensor* temb = nullptr) {
|
||||||
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
|
||||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
|
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
|
||||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
||||||
@ -102,10 +102,10 @@ namespace Anima {
|
|||||||
blocks["2"] = std::make_shared<Linear>(hidden_features, 2 * in_features, false);
|
blocks["2"] = std::make_shared<Linear>(hidden_features, 2 * in_features, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* hidden_states,
|
ggml_tensor* hidden_states,
|
||||||
struct ggml_tensor* embedded_timestep,
|
ggml_tensor* embedded_timestep,
|
||||||
struct ggml_tensor* temb = nullptr) {
|
ggml_tensor* temb = nullptr) {
|
||||||
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
|
||||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
|
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
|
||||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
||||||
@ -152,11 +152,11 @@ namespace Anima {
|
|||||||
blocks[this->out_proj_name] = std::make_shared<Linear>(inner_dim, query_dim, false);
|
blocks[this->out_proj_name] = std::make_shared<Linear>(inner_dim, query_dim, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* hidden_states,
|
ggml_tensor* hidden_states,
|
||||||
struct ggml_tensor* encoder_hidden_states = nullptr,
|
ggml_tensor* encoder_hidden_states = nullptr,
|
||||||
struct ggml_tensor* pe_q = nullptr,
|
ggml_tensor* pe_q = nullptr,
|
||||||
struct ggml_tensor* pe_k = nullptr) {
|
ggml_tensor* pe_k = nullptr) {
|
||||||
if (encoder_hidden_states == nullptr) {
|
if (encoder_hidden_states == nullptr) {
|
||||||
encoder_hidden_states = hidden_states;
|
encoder_hidden_states = hidden_states;
|
||||||
}
|
}
|
||||||
@ -183,7 +183,7 @@ namespace Anima {
|
|||||||
q4 = q_norm->forward(ctx, q4);
|
q4 = q_norm->forward(ctx, q4);
|
||||||
k4 = k_norm->forward(ctx, k4);
|
k4 = k_norm->forward(ctx, k4);
|
||||||
|
|
||||||
struct ggml_tensor* attn_out = nullptr;
|
ggml_tensor* attn_out = nullptr;
|
||||||
if (pe_q != nullptr || pe_k != nullptr) {
|
if (pe_q != nullptr || pe_k != nullptr) {
|
||||||
if (pe_q == nullptr) {
|
if (pe_q == nullptr) {
|
||||||
pe_q = pe_k;
|
pe_q = pe_k;
|
||||||
@ -227,7 +227,7 @@ namespace Anima {
|
|||||||
blocks["layer2"] = std::make_shared<Linear>(hidden_dim, dim, false);
|
blocks["layer2"] = std::make_shared<Linear>(hidden_dim, dim, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto layer1 = std::dynamic_pointer_cast<Linear>(blocks["layer1"]);
|
auto layer1 = std::dynamic_pointer_cast<Linear>(blocks["layer1"]);
|
||||||
auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["layer2"]);
|
auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["layer2"]);
|
||||||
|
|
||||||
@ -245,7 +245,7 @@ namespace Anima {
|
|||||||
blocks["2"] = std::make_shared<Linear>(hidden_dim, dim, true);
|
blocks["2"] = std::make_shared<Linear>(hidden_dim, dim, true);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto layer0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
|
auto layer0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
|
||||||
auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
||||||
|
|
||||||
@ -267,11 +267,11 @@ namespace Anima {
|
|||||||
blocks["mlp"] = std::make_shared<AdapterMLP>(model_dim, model_dim * 4);
|
blocks["mlp"] = std::make_shared<AdapterMLP>(model_dim, model_dim * 4);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* target_pe,
|
ggml_tensor* target_pe,
|
||||||
struct ggml_tensor* context_pe) {
|
ggml_tensor* context_pe) {
|
||||||
auto norm_self_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_self_attn"]);
|
auto norm_self_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_self_attn"]);
|
||||||
auto self_attn = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
|
auto self_attn = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
|
||||||
auto norm_cross_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_cross_attn"]);
|
auto norm_cross_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_cross_attn"]);
|
||||||
@ -317,11 +317,11 @@ namespace Anima {
|
|||||||
blocks["norm"] = std::make_shared<RMSNorm>(target_dim, 1e-6f);
|
blocks["norm"] = std::make_shared<RMSNorm>(target_dim, 1e-6f);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* source_hidden_states,
|
ggml_tensor* source_hidden_states,
|
||||||
struct ggml_tensor* target_input_ids,
|
ggml_tensor* target_input_ids,
|
||||||
struct ggml_tensor* target_pe,
|
ggml_tensor* target_pe,
|
||||||
struct ggml_tensor* source_pe) {
|
ggml_tensor* source_pe) {
|
||||||
GGML_ASSERT(target_input_ids != nullptr);
|
GGML_ASSERT(target_input_ids != nullptr);
|
||||||
if (ggml_n_dims(target_input_ids) == 1) {
|
if (ggml_n_dims(target_input_ids) == 1) {
|
||||||
target_input_ids = ggml_reshape_2d(ctx->ggml_ctx, target_input_ids, target_input_ids->ne[0], 1);
|
target_input_ids = ggml_reshape_2d(ctx->ggml_ctx, target_input_ids, target_input_ids->ne[0], 1);
|
||||||
@ -360,12 +360,12 @@ namespace Anima {
|
|||||||
blocks["mlp"] = std::make_shared<AnimaMLP>(hidden_size, hidden_size * mlp_ratio);
|
blocks["mlp"] = std::make_shared<AnimaMLP>(hidden_size, hidden_size * mlp_ratio);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* hidden_states,
|
ggml_tensor* hidden_states,
|
||||||
struct ggml_tensor* encoder_hidden_states,
|
ggml_tensor* encoder_hidden_states,
|
||||||
struct ggml_tensor* embedded_timestep,
|
ggml_tensor* embedded_timestep,
|
||||||
struct ggml_tensor* temb,
|
ggml_tensor* temb,
|
||||||
struct ggml_tensor* image_pe) {
|
ggml_tensor* image_pe) {
|
||||||
auto norm1 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_self_attn"]);
|
auto norm1 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_self_attn"]);
|
||||||
auto attn1 = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
|
auto attn1 = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
|
||||||
auto norm2 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_cross_attn"]);
|
auto norm2 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_cross_attn"]);
|
||||||
@ -402,10 +402,10 @@ namespace Anima {
|
|||||||
blocks["linear"] = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, false);
|
blocks["linear"] = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* hidden_states,
|
ggml_tensor* hidden_states,
|
||||||
struct ggml_tensor* embedded_timestep,
|
ggml_tensor* embedded_timestep,
|
||||||
struct ggml_tensor* temb) {
|
ggml_tensor* temb) {
|
||||||
auto adaln = std::dynamic_pointer_cast<AdaLayerNorm>(blocks["adaln_modulation"]);
|
auto adaln = std::dynamic_pointer_cast<AdaLayerNorm>(blocks["adaln_modulation"]);
|
||||||
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
||||||
|
|
||||||
@ -445,15 +445,15 @@ namespace Anima {
|
|||||||
blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
|
blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* encoder_hidden_states,
|
ggml_tensor* encoder_hidden_states,
|
||||||
struct ggml_tensor* image_pe,
|
ggml_tensor* image_pe,
|
||||||
struct ggml_tensor* t5_ids = nullptr,
|
ggml_tensor* t5_ids = nullptr,
|
||||||
struct ggml_tensor* t5_weights = nullptr,
|
ggml_tensor* t5_weights = nullptr,
|
||||||
struct ggml_tensor* adapter_q_pe = nullptr,
|
ggml_tensor* adapter_q_pe = nullptr,
|
||||||
struct ggml_tensor* adapter_k_pe = nullptr) {
|
ggml_tensor* adapter_k_pe = nullptr) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
|
|
||||||
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
|
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
|
||||||
@ -553,7 +553,7 @@ namespace Anima {
|
|||||||
return "anima";
|
return "anima";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
net.get_param_tensors(tensors, prefix + ".net");
|
net.get_param_tensors(tensors, prefix + ".net");
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -602,13 +602,13 @@ namespace Anima {
|
|||||||
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
|
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* t5_ids = nullptr,
|
ggml_tensor* t5_ids = nullptr,
|
||||||
struct ggml_tensor* t5_weights = nullptr) {
|
ggml_tensor* t5_weights = nullptr) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
struct ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
timesteps = to_backend(timesteps);
|
timesteps = to_backend(timesteps);
|
||||||
@ -668,14 +668,14 @@ namespace Anima {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* t5_ids = nullptr,
|
ggml_tensor* t5_ids = nullptr,
|
||||||
struct ggml_tensor* t5_weights = nullptr,
|
ggml_tensor* t5_weights = nullptr,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, t5_ids, t5_weights);
|
return build_graph(x, timesteps, context, t5_ids, t5_weights);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
|
||||||
|
|||||||
@ -29,7 +29,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// t_emb is always None
|
// t_emb is always None
|
||||||
auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
|
auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
|
||||||
@ -65,7 +65,7 @@ protected:
|
|||||||
int64_t in_channels;
|
int64_t in_channels;
|
||||||
bool use_linear;
|
bool use_linear;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
|
||||||
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
|
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
|
||||||
if (iter != tensor_storage_map.end()) {
|
if (iter != tensor_storage_map.end()) {
|
||||||
if (iter->second.n_dims == 4 && use_linear) {
|
if (iter->second.n_dims == 4 && use_linear) {
|
||||||
@ -101,7 +101,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
|
||||||
auto q_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
|
auto q_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
|
||||||
@ -178,8 +178,8 @@ public:
|
|||||||
{kernel_padding, 0, 0}));
|
{kernel_padding, 0, 0}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x) override {
|
ggml_tensor* x) override {
|
||||||
// timesteps always None
|
// timesteps always None
|
||||||
// skip_video always False
|
// skip_video always False
|
||||||
// x: [N, IC, IH, IW]
|
// x: [N, IC, IH, IW]
|
||||||
@ -208,7 +208,7 @@ public:
|
|||||||
|
|
||||||
class VideoResnetBlock : public ResnetBlock {
|
class VideoResnetBlock : public ResnetBlock {
|
||||||
protected:
|
protected:
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type wtype = get_type(prefix + "mix_factor", tensor_storage_map, GGML_TYPE_F32);
|
enum ggml_type wtype = get_type(prefix + "mix_factor", tensor_storage_map, GGML_TYPE_F32);
|
||||||
params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
|
params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
|
||||||
}
|
}
|
||||||
@ -227,7 +227,7 @@ public:
|
|||||||
blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, 0, out_channels, {video_kernel_size, 1}, 3, false, true));
|
blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, 0, out_channels, {video_kernel_size, 1}, 3, false, true));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N, in_channels, h, w] aka [b*t, in_channels, h, w]
|
// x: [N, in_channels, h, w] aka [b*t, in_channels, h, w]
|
||||||
// return: [N, out_channels, h, w] aka [b*t, out_channels, h, w]
|
// return: [N, out_channels, h, w] aka [b*t, out_channels, h, w]
|
||||||
// t_emb is always None
|
// t_emb is always None
|
||||||
@ -317,7 +317,7 @@ public:
|
|||||||
blocks["conv_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(block_in, double_z ? z_channels * 2 : z_channels, {3, 3}, {1, 1}, {1, 1}));
|
blocks["conv_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(block_in, double_z ? z_channels * 2 : z_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
|
|
||||||
auto conv_in = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"]);
|
auto conv_in = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"]);
|
||||||
@ -435,7 +435,7 @@ public:
|
|||||||
blocks["conv_out"] = get_conv_out(block_in, out_ch, {3, 3}, {1, 1}, {1, 1});
|
blocks["conv_out"] = get_conv_out(block_in, out_ch, {3, 3}, {1, 1}, {1, 1});
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||||
// z: [N, z_channels, h, w]
|
// z: [N, z_channels, h, w]
|
||||||
// alpha is always 0
|
// alpha is always 0
|
||||||
// merge_strategy is always learned
|
// merge_strategy is always learned
|
||||||
@ -549,7 +549,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
|
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||||
// z: [N, z_channels, h, w]
|
// z: [N, z_channels, h, w]
|
||||||
if (sd_version_is_flux2(version)) {
|
if (sd_version_is_flux2(version)) {
|
||||||
// [N, C*p*p, h, w] -> [N, C, h*p, w*p]
|
// [N, C*p*p, h, w] -> [N, C, h*p, w*p]
|
||||||
@ -581,7 +581,7 @@ public:
|
|||||||
return h;
|
return h;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
|
auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
|
||||||
|
|
||||||
@ -613,6 +613,9 @@ public:
|
|||||||
|
|
||||||
int get_encoder_output_channels() {
|
int get_encoder_output_channels() {
|
||||||
int factor = dd_config.double_z ? 2 : 1;
|
int factor = dd_config.double_z ? 2 : 1;
|
||||||
|
if (sd_version_is_flux2(version)) {
|
||||||
|
return dd_config.z_channels * 4;
|
||||||
|
}
|
||||||
return dd_config.z_channels * factor;
|
return dd_config.z_channels * factor;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -678,18 +681,18 @@ struct AutoEncoderKL : public VAE {
|
|||||||
return "vae";
|
return "vae";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
|
||||||
ae.get_param_tensors(tensors, prefix);
|
ae.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
|
ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
|
|
||||||
z = to_backend(z);
|
z = to_backend(z);
|
||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
|
ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -697,12 +700,12 @@ struct AutoEncoderKL : public VAE {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool _compute(const int n_threads,
|
bool _compute(const int n_threads,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
bool decode_graph,
|
bool decode_graph,
|
||||||
struct ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
GGML_ASSERT(!decode_only || decode_graph);
|
GGML_ASSERT(!decode_only || decode_graph);
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(z, decode_graph);
|
return build_graph(z, decode_graph);
|
||||||
};
|
};
|
||||||
// ggml_set_f32(z, 0.5f);
|
// ggml_set_f32(z, 0.5f);
|
||||||
@ -712,8 +715,8 @@ struct AutoEncoderKL : public VAE {
|
|||||||
|
|
||||||
ggml_tensor* gaussian_latent_sample(ggml_context* work_ctx, ggml_tensor* moments, std::shared_ptr<RNG> rng) {
|
ggml_tensor* gaussian_latent_sample(ggml_context* work_ctx, ggml_tensor* moments, std::shared_ptr<RNG> rng) {
|
||||||
// ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample
|
// ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample
|
||||||
ggml_tensor* latents = ggml_new_tensor_4d(work_ctx, moments->type, moments->ne[0], moments->ne[1], moments->ne[2] / 2, moments->ne[3]);
|
ggml_tensor* latents = ggml_new_tensor_4d(work_ctx, moments->type, moments->ne[0], moments->ne[1], moments->ne[2] / 2, moments->ne[3]);
|
||||||
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, latents);
|
ggml_tensor* noise = ggml_dup_tensor(work_ctx, latents);
|
||||||
ggml_ext_im_set_randn_f32(noise, rng);
|
ggml_ext_im_set_randn_f32(noise, rng);
|
||||||
{
|
{
|
||||||
float mean = 0;
|
float mean = 0;
|
||||||
@ -881,12 +884,12 @@ struct AutoEncoderKL : public VAE {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -897,7 +900,7 @@ struct AutoEncoderKL : public VAE {
|
|||||||
auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 64, 64, 3, 2);
|
auto x = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 64, 64, 3, 2);
|
||||||
ggml_set_f32(x, 0.5f);
|
ggml_set_f32(x, 0.5f);
|
||||||
print_ggml_tensor(x);
|
print_ggml_tensor(x);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
_compute(8, x, false, &out, work_ctx);
|
_compute(8, x, false, &out, work_ctx);
|
||||||
@ -915,7 +918,7 @@ struct AutoEncoderKL : public VAE {
|
|||||||
auto z = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 4, 1);
|
auto z = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 8, 8, 4, 1);
|
||||||
ggml_set_f32(z, 0.5f);
|
ggml_set_f32(z, 0.5f);
|
||||||
print_ggml_tensor(z);
|
print_ggml_tensor(z);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
_compute(8, z, true, &out, work_ctx);
|
_compute(8, z, true, &out, work_ctx);
|
||||||
|
|||||||
@ -799,7 +799,7 @@ struct CacheDitConditionState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
bool before_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output, float sigma, int step_index) {
|
bool before_condition(const void* cond, ggml_tensor* input, ggml_tensor* output, float sigma, int step_index) {
|
||||||
if (!enabled() || step_index < 0)
|
if (!enabled() || step_index < 0)
|
||||||
return false;
|
return false;
|
||||||
|
|
||||||
@ -867,7 +867,7 @@ struct CacheDitConditionState {
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
void after_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output) {
|
void after_condition(const void* cond, ggml_tensor* input, ggml_tensor* output) {
|
||||||
if (!step_is_active())
|
if (!step_is_active())
|
||||||
return;
|
return;
|
||||||
|
|
||||||
|
|||||||
118
src/clip.hpp
118
src/clip.hpp
@ -473,7 +473,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, n_token, d_model]
|
// x: [N, n_token, d_model]
|
||||||
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
||||||
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
||||||
@ -511,7 +511,7 @@ public:
|
|||||||
blocks["mlp"] = std::shared_ptr<GGMLBlock>(new CLIPMLP(d_model, intermediate_size));
|
blocks["mlp"] = std::shared_ptr<GGMLBlock>(new CLIPMLP(d_model, intermediate_size));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* mask = nullptr) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask = nullptr) {
|
||||||
// x: [N, n_token, d_model]
|
// x: [N, n_token, d_model]
|
||||||
auto self_attn = std::dynamic_pointer_cast<MultiheadAttention>(blocks["self_attn"]);
|
auto self_attn = std::dynamic_pointer_cast<MultiheadAttention>(blocks["self_attn"]);
|
||||||
auto layer_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm1"]);
|
auto layer_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm1"]);
|
||||||
@ -541,10 +541,10 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
int clip_skip = -1) {
|
int clip_skip = -1) {
|
||||||
// x: [N, n_token, d_model]
|
// x: [N, n_token, d_model]
|
||||||
int layer_idx = n_layer - 1;
|
int layer_idx = n_layer - 1;
|
||||||
// LOG_DEBUG("clip_skip %d", clip_skip);
|
// LOG_DEBUG("clip_skip %d", clip_skip);
|
||||||
@ -573,7 +573,7 @@ protected:
|
|||||||
int64_t num_positions;
|
int64_t num_positions;
|
||||||
bool force_clip_f32;
|
bool force_clip_f32;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type token_wtype = GGML_TYPE_F32;
|
enum ggml_type token_wtype = GGML_TYPE_F32;
|
||||||
if (!force_clip_f32) {
|
if (!force_clip_f32) {
|
||||||
token_wtype = get_type(prefix + "token_embedding.weight", tensor_storage_map, GGML_TYPE_F32);
|
token_wtype = get_type(prefix + "token_embedding.weight", tensor_storage_map, GGML_TYPE_F32);
|
||||||
@ -597,13 +597,13 @@ public:
|
|||||||
force_clip_f32(force_clip_f32) {
|
force_clip_f32(force_clip_f32) {
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* get_token_embed_weight() {
|
ggml_tensor* get_token_embed_weight() {
|
||||||
return params["token_embedding.weight"];
|
return params["token_embedding.weight"];
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* custom_embed_weight) {
|
ggml_tensor* custom_embed_weight) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
auto token_embed_weight = params["token_embedding.weight"];
|
auto token_embed_weight = params["token_embedding.weight"];
|
||||||
auto position_embed_weight = params["position_embedding.weight"];
|
auto position_embed_weight = params["position_embedding.weight"];
|
||||||
@ -630,7 +630,7 @@ protected:
|
|||||||
int num_patches;
|
int num_patches;
|
||||||
int64_t num_positions;
|
int64_t num_positions;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type patch_wtype = GGML_TYPE_F16;
|
enum ggml_type patch_wtype = GGML_TYPE_F16;
|
||||||
enum ggml_type class_wtype = GGML_TYPE_F32;
|
enum ggml_type class_wtype = GGML_TYPE_F32;
|
||||||
enum ggml_type position_wtype = GGML_TYPE_F32;
|
enum ggml_type position_wtype = GGML_TYPE_F32;
|
||||||
@ -653,7 +653,7 @@ public:
|
|||||||
num_positions = num_patches + 1;
|
num_positions = num_patches + 1;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* pixel_values) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* pixel_values) {
|
||||||
// pixel_values: [N, num_channels, image_size, image_size]
|
// pixel_values: [N, num_channels, image_size, image_size]
|
||||||
// return: [N, num_positions, embed_dim]
|
// return: [N, num_positions, embed_dim]
|
||||||
GGML_ASSERT(pixel_values->ne[0] == image_size && pixel_values->ne[1] == image_size && pixel_values->ne[2] == num_channels);
|
GGML_ASSERT(pixel_values->ne[0] == image_size && pixel_values->ne[1] == image_size && pixel_values->ne[2] == num_channels);
|
||||||
@ -663,20 +663,20 @@ public:
|
|||||||
auto position_embed_weight = params["position_embedding.weight"];
|
auto position_embed_weight = params["position_embedding.weight"];
|
||||||
|
|
||||||
// concat(patch_embedding, class_embedding) + position_embedding
|
// concat(patch_embedding, class_embedding) + position_embedding
|
||||||
struct ggml_tensor* patch_embedding;
|
ggml_tensor* patch_embedding;
|
||||||
int64_t N = pixel_values->ne[3];
|
int64_t N = pixel_values->ne[3];
|
||||||
patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size); // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
|
patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size); // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
|
||||||
patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N); // [N, embed_dim, num_patches]
|
patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N); // [N, embed_dim, num_patches]
|
||||||
patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3)); // [N, num_patches, embed_dim]
|
patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3)); // [N, num_patches, embed_dim]
|
||||||
patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N); // [N, num_patches, embed_dim, 1]
|
patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N); // [N, num_patches, embed_dim, 1]
|
||||||
|
|
||||||
struct ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
|
ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
|
||||||
class_embedding = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding); // [N, embed_dim]
|
class_embedding = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding); // [N, embed_dim]
|
||||||
class_embedding = ggml_reshape_4d(ctx->ggml_ctx, class_embedding, 1, embed_dim, 1, N); // [N, 1, embed_dim, 1]
|
class_embedding = ggml_reshape_4d(ctx->ggml_ctx, class_embedding, 1, embed_dim, 1, N); // [N, 1, embed_dim, 1]
|
||||||
|
|
||||||
struct ggml_tensor* x = ggml_concat(ctx->ggml_ctx, class_embedding, patch_embedding, 2); // [N, num_positions, embed_dim, 1]
|
ggml_tensor* x = ggml_concat(ctx->ggml_ctx, class_embedding, patch_embedding, 2); // [N, num_positions, embed_dim, 1]
|
||||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, embed_dim, num_positions, N); // [N, num_positions, embed_dim]
|
x = ggml_reshape_3d(ctx->ggml_ctx, x, embed_dim, num_positions, N); // [N, num_positions, embed_dim]
|
||||||
x = ggml_add(ctx->ggml_ctx, x, position_embed_weight);
|
x = ggml_add(ctx->ggml_ctx, x, position_embed_weight);
|
||||||
return x; // [N, num_positions, embed_dim]
|
return x; // [N, num_positions, embed_dim]
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -693,7 +693,7 @@ enum CLIPVersion {
|
|||||||
|
|
||||||
class CLIPTextModel : public GGMLBlock {
|
class CLIPTextModel : public GGMLBlock {
|
||||||
protected:
|
protected:
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
if (version == OPEN_CLIP_VIT_BIGG_14) {
|
if (version == OPEN_CLIP_VIT_BIGG_14) {
|
||||||
enum ggml_type wtype = GGML_TYPE_F32;
|
enum ggml_type wtype = GGML_TYPE_F32;
|
||||||
params["text_projection"] = ggml_new_tensor_2d(ctx, wtype, projection_dim, hidden_size);
|
params["text_projection"] = ggml_new_tensor_2d(ctx, wtype, projection_dim, hidden_size);
|
||||||
@ -734,18 +734,18 @@ public:
|
|||||||
blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
|
blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* get_token_embed_weight() {
|
ggml_tensor* get_token_embed_weight() {
|
||||||
auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
|
auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
|
||||||
return embeddings->get_token_embed_weight();
|
return embeddings->get_token_embed_weight();
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* tkn_embeddings,
|
ggml_tensor* tkn_embeddings,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
size_t max_token_idx = 0,
|
size_t max_token_idx = 0,
|
||||||
bool return_pooled = false,
|
bool return_pooled = false,
|
||||||
int clip_skip = -1) {
|
int clip_skip = -1) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
|
auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
|
||||||
auto encoder = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
|
auto encoder = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
|
||||||
@ -804,10 +804,10 @@ public:
|
|||||||
blocks["post_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
|
blocks["post_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* pixel_values,
|
ggml_tensor* pixel_values,
|
||||||
bool return_pooled = true,
|
bool return_pooled = true,
|
||||||
int clip_skip = -1) {
|
int clip_skip = -1) {
|
||||||
// pixel_values: [N, num_channels, image_size, image_size]
|
// pixel_values: [N, num_channels, image_size, image_size]
|
||||||
auto embeddings = std::dynamic_pointer_cast<CLIPVisionEmbeddings>(blocks["embeddings"]);
|
auto embeddings = std::dynamic_pointer_cast<CLIPVisionEmbeddings>(blocks["embeddings"]);
|
||||||
auto pre_layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["pre_layernorm"]);
|
auto pre_layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["pre_layernorm"]);
|
||||||
@ -839,7 +839,7 @@ protected:
|
|||||||
int64_t out_features;
|
int64_t out_features;
|
||||||
bool transpose_weight;
|
bool transpose_weight;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
||||||
if (transpose_weight) {
|
if (transpose_weight) {
|
||||||
params["weight"] = ggml_new_tensor_2d(ctx, wtype, out_features, in_features);
|
params["weight"] = ggml_new_tensor_2d(ctx, wtype, out_features, in_features);
|
||||||
@ -856,8 +856,8 @@ public:
|
|||||||
out_features(out_features),
|
out_features(out_features),
|
||||||
transpose_weight(transpose_weight) {}
|
transpose_weight(transpose_weight) {}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
struct ggml_tensor* w = params["weight"];
|
ggml_tensor* w = params["weight"];
|
||||||
if (transpose_weight) {
|
if (transpose_weight) {
|
||||||
w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
|
w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
|
||||||
}
|
}
|
||||||
@ -886,10 +886,10 @@ public:
|
|||||||
blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new CLIPProjection(hidden_size, projection_dim, transpose_proj_w));
|
blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new CLIPProjection(hidden_size, projection_dim, transpose_proj_w));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* pixel_values,
|
ggml_tensor* pixel_values,
|
||||||
bool return_pooled = true,
|
bool return_pooled = true,
|
||||||
int clip_skip = -1) {
|
int clip_skip = -1) {
|
||||||
// pixel_values: [N, num_channels, image_size, image_size]
|
// pixel_values: [N, num_channels, image_size, image_size]
|
||||||
// return: [N, projection_dim] if return_pooled else [N, n_token, hidden_size]
|
// return: [N, projection_dim] if return_pooled else [N, n_token, hidden_size]
|
||||||
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
|
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
|
||||||
@ -936,17 +936,17 @@ struct CLIPTextModelRunner : public GGMLRunner {
|
|||||||
return "clip";
|
return "clip";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
model.get_param_tensors(tensors, prefix);
|
model.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* embeddings,
|
ggml_tensor* embeddings,
|
||||||
struct ggml_tensor* mask,
|
ggml_tensor* mask,
|
||||||
size_t max_token_idx = 0,
|
size_t max_token_idx = 0,
|
||||||
bool return_pooled = false,
|
bool return_pooled = false,
|
||||||
int clip_skip = -1) {
|
int clip_skip = -1) {
|
||||||
size_t N = input_ids->ne[1];
|
size_t N = input_ids->ne[1];
|
||||||
size_t n_token = input_ids->ne[0];
|
size_t n_token = input_ids->ne[0];
|
||||||
if (input_ids->ne[0] > model.n_token) {
|
if (input_ids->ne[0] > model.n_token) {
|
||||||
@ -957,17 +957,17 @@ struct CLIPTextModelRunner : public GGMLRunner {
|
|||||||
return model.forward(ctx, input_ids, embeddings, mask, max_token_idx, return_pooled, clip_skip);
|
return model.forward(ctx, input_ids, embeddings, mask, max_token_idx, return_pooled, clip_skip);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
|
ggml_cgraph* build_graph(ggml_tensor* input_ids,
|
||||||
int num_custom_embeddings = 0,
|
int num_custom_embeddings = 0,
|
||||||
void* custom_embeddings_data = nullptr,
|
void* custom_embeddings_data = nullptr,
|
||||||
size_t max_token_idx = 0,
|
size_t max_token_idx = 0,
|
||||||
bool return_pooled = false,
|
bool return_pooled = false,
|
||||||
int clip_skip = -1) {
|
int clip_skip = -1) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(2048);
|
ggml_cgraph* gf = new_graph_custom(2048);
|
||||||
|
|
||||||
input_ids = to_backend(input_ids);
|
input_ids = to_backend(input_ids);
|
||||||
|
|
||||||
struct ggml_tensor* embeddings = nullptr;
|
ggml_tensor* embeddings = nullptr;
|
||||||
|
|
||||||
if (num_custom_embeddings > 0 && custom_embeddings_data != nullptr) {
|
if (num_custom_embeddings > 0 && custom_embeddings_data != nullptr) {
|
||||||
auto token_embed_weight = model.get_token_embed_weight();
|
auto token_embed_weight = model.get_token_embed_weight();
|
||||||
@ -997,7 +997,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, attention_mask, max_token_idx, return_pooled, clip_skip);
|
ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, attention_mask, max_token_idx, return_pooled, clip_skip);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, hidden_states);
|
ggml_build_forward_expand(gf, hidden_states);
|
||||||
|
|
||||||
@ -1005,7 +1005,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(const int n_threads,
|
bool compute(const int n_threads,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
int num_custom_embeddings,
|
int num_custom_embeddings,
|
||||||
void* custom_embeddings_data,
|
void* custom_embeddings_data,
|
||||||
size_t max_token_idx,
|
size_t max_token_idx,
|
||||||
@ -1013,7 +1013,7 @@ struct CLIPTextModelRunner : public GGMLRunner {
|
|||||||
int clip_skip,
|
int clip_skip,
|
||||||
ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
|
return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
||||||
|
|||||||
@ -23,7 +23,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
if (vae_downsample) {
|
if (vae_downsample) {
|
||||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||||
@ -52,7 +52,7 @@ public:
|
|||||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||||
|
|
||||||
@ -121,7 +121,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* emb = nullptr) {
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb = nullptr) {
|
||||||
// For dims==3, we reduce dimension from 5d to 4d by merging h and w, in order not to change ggml
|
// For dims==3, we reduce dimension from 5d to 4d by merging h and w, in order not to change ggml
|
||||||
// [N, c, t, h, w] => [N, c, t, h * w]
|
// [N, c, t, h, w] => [N, c, t, h * w]
|
||||||
// x: [N, channels, h, w] if dims == 2 else [N, channels, t, h, w]
|
// x: [N, channels, h, w] if dims == 2 else [N, channels, t, h, w]
|
||||||
@ -188,7 +188,7 @@ public:
|
|||||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out * 2));
|
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out * 2));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [ne3, ne2, ne1, dim_in]
|
// x: [ne3, ne2, ne1, dim_in]
|
||||||
// return: [ne3, ne2, ne1, dim_out]
|
// return: [ne3, ne2, ne1, dim_out]
|
||||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
||||||
@ -214,7 +214,7 @@ public:
|
|||||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out, bias));
|
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out, bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [ne3, ne2, ne1, dim_in]
|
// x: [ne3, ne2, ne1, dim_in]
|
||||||
// return: [ne3, ne2, ne1, dim_out]
|
// return: [ne3, ne2, ne1, dim_out]
|
||||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
||||||
@ -258,7 +258,7 @@ public:
|
|||||||
blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
|
blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [ne3, ne2, ne1, dim]
|
// x: [ne3, ne2, ne1, dim]
|
||||||
// return: [ne3, ne2, ne1, dim_out]
|
// return: [ne3, ne2, ne1, dim_out]
|
||||||
|
|
||||||
@ -297,9 +297,9 @@ public:
|
|||||||
// to_out_1 is nn.Dropout(), skip for inference
|
// to_out_1 is nn.Dropout(), skip for inference
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context) {
|
ggml_tensor* context) {
|
||||||
// x: [N, n_token, query_dim]
|
// x: [N, n_token, query_dim]
|
||||||
// context: [N, n_context, context_dim]
|
// context: [N, n_context, context_dim]
|
||||||
// return: [N, n_token, query_dim]
|
// return: [N, n_token, query_dim]
|
||||||
@ -355,9 +355,9 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context) {
|
ggml_tensor* context) {
|
||||||
// x: [N, n_token, query_dim]
|
// x: [N, n_token, query_dim]
|
||||||
// context: [N, n_context, context_dim]
|
// context: [N, n_context, context_dim]
|
||||||
// return: [N, n_token, query_dim]
|
// return: [N, n_token, query_dim]
|
||||||
@ -406,7 +406,7 @@ protected:
|
|||||||
int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2
|
int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2
|
||||||
bool use_linear = false;
|
bool use_linear = false;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
|
||||||
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
|
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
|
||||||
if (iter != tensor_storage_map.end()) {
|
if (iter != tensor_storage_map.end()) {
|
||||||
int64_t inner_dim = n_head * d_head;
|
int64_t inner_dim = n_head * d_head;
|
||||||
@ -456,9 +456,9 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context) {
|
ggml_tensor* context) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// context: [N, max_position(aka n_token), hidden_size(aka context_dim)]
|
// context: [N, max_position(aka n_token), hidden_size(aka context_dim)]
|
||||||
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
|
||||||
@ -510,7 +510,7 @@ public:
|
|||||||
|
|
||||||
class AlphaBlender : public GGMLBlock {
|
class AlphaBlender : public GGMLBlock {
|
||||||
protected:
|
protected:
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
|
||||||
// Get the type of the "mix_factor" tensor from the input tensors map with the specified prefix
|
// Get the type of the "mix_factor" tensor from the input tensors map with the specified prefix
|
||||||
enum ggml_type wtype = GGML_TYPE_F32;
|
enum ggml_type wtype = GGML_TYPE_F32;
|
||||||
params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
|
params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
|
||||||
@ -530,9 +530,9 @@ public:
|
|||||||
// since mix_factor.shape is [1,], we don't need rearrange using rearrange_pattern
|
// since mix_factor.shape is [1,], we don't need rearrange using rearrange_pattern
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x_spatial,
|
ggml_tensor* x_spatial,
|
||||||
struct ggml_tensor* x_temporal) {
|
ggml_tensor* x_temporal) {
|
||||||
// image_only_indicator is always tensor([0.])
|
// image_only_indicator is always tensor([0.])
|
||||||
float alpha = get_alpha();
|
float alpha = get_alpha();
|
||||||
auto x = ggml_add(ctx->ggml_ctx,
|
auto x = ggml_add(ctx->ggml_ctx,
|
||||||
@ -555,10 +555,10 @@ public:
|
|||||||
blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
|
blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* emb,
|
ggml_tensor* emb,
|
||||||
int num_video_frames) {
|
int num_video_frames) {
|
||||||
// x: [N, channels, h, w] aka [b*t, channels, h, w]
|
// x: [N, channels, h, w] aka [b*t, channels, h, w]
|
||||||
// emb: [N, emb_channels] aka [b*t, emb_channels]
|
// emb: [N, emb_channels] aka [b*t, emb_channels]
|
||||||
// image_only_indicator is always tensor([0.])
|
// image_only_indicator is always tensor([0.])
|
||||||
|
|||||||
@ -6,17 +6,17 @@
|
|||||||
#include "t5.hpp"
|
#include "t5.hpp"
|
||||||
|
|
||||||
struct SDCondition {
|
struct SDCondition {
|
||||||
struct ggml_tensor* c_crossattn = nullptr; // aka context
|
ggml_tensor* c_crossattn = nullptr; // aka context
|
||||||
struct ggml_tensor* c_vector = nullptr; // aka y
|
ggml_tensor* c_vector = nullptr; // aka y
|
||||||
struct ggml_tensor* c_concat = nullptr;
|
ggml_tensor* c_concat = nullptr;
|
||||||
|
|
||||||
std::vector<struct ggml_tensor*> extra_c_crossattns;
|
std::vector<ggml_tensor*> extra_c_crossattns;
|
||||||
|
|
||||||
SDCondition() = default;
|
SDCondition() = default;
|
||||||
SDCondition(struct ggml_tensor* c_crossattn,
|
SDCondition(ggml_tensor* c_crossattn,
|
||||||
struct ggml_tensor* c_vector,
|
ggml_tensor* c_vector,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
const std::vector<struct ggml_tensor*>& extra_c_crossattns = {})
|
const std::vector<ggml_tensor*>& extra_c_crossattns = {})
|
||||||
: c_crossattn(c_crossattn), c_vector(c_vector), c_concat(c_concat), extra_c_crossattns(extra_c_crossattns) {}
|
: c_crossattn(c_crossattn), c_vector(c_vector), c_concat(c_concat), extra_c_crossattns(extra_c_crossattns) {}
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -37,7 +37,7 @@ struct Conditioner {
|
|||||||
const ConditionerParams& conditioner_params) = 0;
|
const ConditionerParams& conditioner_params) = 0;
|
||||||
virtual void alloc_params_buffer() = 0;
|
virtual void alloc_params_buffer() = 0;
|
||||||
virtual void free_params_buffer() = 0;
|
virtual void free_params_buffer() = 0;
|
||||||
virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) = 0;
|
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
|
||||||
virtual size_t get_params_buffer_size() = 0;
|
virtual size_t get_params_buffer_size() = 0;
|
||||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
||||||
@ -92,7 +92,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
text_model->get_param_tensors(tensors, "cond_stage_model.transformer.text_model");
|
text_model->get_param_tensors(tensors, "cond_stage_model.transformer.text_model");
|
||||||
if (sd_version_is_sdxl(version)) {
|
if (sd_version_is_sdxl(version)) {
|
||||||
text_model2->get_param_tensors(tensors, "cond_stage_model.1.transformer.text_model");
|
text_model2->get_param_tensors(tensors, "cond_stage_model.1.transformer.text_model");
|
||||||
@ -149,14 +149,14 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
|||||||
}
|
}
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = 100 * 1024 * 1024; // max for custom embeddings 100 MB
|
params.mem_size = 100 * 1024 * 1024; // max for custom embeddings 100 MB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
struct ggml_context* embd_ctx = ggml_init(params);
|
ggml_context* embd_ctx = ggml_init(params);
|
||||||
struct ggml_tensor* embd = nullptr;
|
ggml_tensor* embd = nullptr;
|
||||||
struct ggml_tensor* embd2 = nullptr;
|
ggml_tensor* embd2 = nullptr;
|
||||||
auto on_load = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
|
auto on_load = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) {
|
||||||
if (tensor_storage.ne[0] != text_model->model.hidden_size) {
|
if (tensor_storage.ne[0] != text_model->model.hidden_size) {
|
||||||
if (text_model2) {
|
if (text_model2) {
|
||||||
if (tensor_storage.ne[0] == text_model2->model.hidden_size) {
|
if (tensor_storage.ne[0] == text_model2->model.hidden_size) {
|
||||||
@ -435,12 +435,12 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
|||||||
int height,
|
int height,
|
||||||
int adm_in_channels = -1,
|
int adm_in_channels = -1,
|
||||||
bool zero_out_masked = false) {
|
bool zero_out_masked = false) {
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
struct ggml_tensor* hidden_states = nullptr; // [N, n_token, hidden_size]
|
ggml_tensor* hidden_states = nullptr; // [N, n_token, hidden_size]
|
||||||
struct ggml_tensor* chunk_hidden_states = nullptr; // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2]
|
ggml_tensor* chunk_hidden_states = nullptr; // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2]
|
||||||
struct ggml_tensor* chunk_hidden_states1 = nullptr; // [n_token, hidden_size]
|
ggml_tensor* chunk_hidden_states1 = nullptr; // [n_token, hidden_size]
|
||||||
struct ggml_tensor* chunk_hidden_states2 = nullptr; // [n_token, hidden_size2]
|
ggml_tensor* chunk_hidden_states2 = nullptr; // [n_token, hidden_size2]
|
||||||
struct ggml_tensor* pooled = nullptr;
|
ggml_tensor* pooled = nullptr;
|
||||||
std::vector<float> hidden_states_vec;
|
std::vector<float> hidden_states_vec;
|
||||||
|
|
||||||
if (clip_skip <= 0) {
|
if (clip_skip <= 0) {
|
||||||
@ -455,9 +455,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
|||||||
std::vector<float> chunk_weights(weights.begin() + chunk_idx * chunk_len,
|
std::vector<float> chunk_weights(weights.begin() + chunk_idx * chunk_len,
|
||||||
weights.begin() + (chunk_idx + 1) * chunk_len);
|
weights.begin() + (chunk_idx + 1) * chunk_len);
|
||||||
|
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, chunk_tokens);
|
||||||
struct ggml_tensor* input_ids2 = nullptr;
|
ggml_tensor* input_ids2 = nullptr;
|
||||||
size_t max_token_idx = 0;
|
size_t max_token_idx = 0;
|
||||||
if (sd_version_is_sdxl(version)) {
|
if (sd_version_is_sdxl(version)) {
|
||||||
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer.EOS_TOKEN_ID);
|
auto it = std::find(chunk_tokens.begin(), chunk_tokens.end(), tokenizer.EOS_TOKEN_ID);
|
||||||
if (it != chunk_tokens.end()) {
|
if (it != chunk_tokens.end()) {
|
||||||
@ -676,18 +676,18 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
|
|||||||
return "clip_vision";
|
return "clip_vision";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
|
||||||
vision_model.get_param_tensors(tensors, "cond_stage_model.transformer");
|
vision_model.get_param_tensors(tensors, "cond_stage_model.transformer");
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* pixel_values, bool return_pooled, int clip_skip) {
|
ggml_cgraph* build_graph(ggml_tensor* pixel_values, bool return_pooled, int clip_skip) {
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
|
|
||||||
pixel_values = to_backend(pixel_values);
|
pixel_values = to_backend(pixel_values);
|
||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* hidden_states = vision_model.forward(&runner_ctx, pixel_values, return_pooled, clip_skip);
|
ggml_tensor* hidden_states = vision_model.forward(&runner_ctx, pixel_values, return_pooled, clip_skip);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, hidden_states);
|
ggml_build_forward_expand(gf, hidden_states);
|
||||||
|
|
||||||
@ -700,7 +700,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
|
|||||||
int clip_skip,
|
int clip_skip,
|
||||||
ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
ggml_context* output_ctx) {
|
ggml_context* output_ctx) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(pixel_values, return_pooled, clip_skip);
|
return build_graph(pixel_values, return_pooled, clip_skip);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
||||||
@ -746,7 +746,7 @@ struct SD3CLIPEmbedder : public Conditioner {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
if (clip_l) {
|
if (clip_l) {
|
||||||
clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
|
clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
|
||||||
}
|
}
|
||||||
@ -909,15 +909,15 @@ struct SD3CLIPEmbedder : public Conditioner {
|
|||||||
clip_skip = 2;
|
clip_skip = 2;
|
||||||
}
|
}
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
struct ggml_tensor* hidden_states = nullptr; // [N, n_token*2, 4096]
|
ggml_tensor* hidden_states = nullptr; // [N, n_token*2, 4096]
|
||||||
struct ggml_tensor* chunk_hidden_states = nullptr; // [n_token*2, 4096]
|
ggml_tensor* chunk_hidden_states = nullptr; // [n_token*2, 4096]
|
||||||
struct ggml_tensor* chunk_hidden_states_l = nullptr; // [n_token, hidden_size_l]
|
ggml_tensor* chunk_hidden_states_l = nullptr; // [n_token, hidden_size_l]
|
||||||
struct ggml_tensor* chunk_hidden_states_g = nullptr; // [n_token, hidden_size_g]
|
ggml_tensor* chunk_hidden_states_g = nullptr; // [n_token, hidden_size_g]
|
||||||
struct ggml_tensor* chunk_hidden_states_t5 = nullptr; // [n_token, hidden_size_t5]
|
ggml_tensor* chunk_hidden_states_t5 = nullptr; // [n_token, hidden_size_t5]
|
||||||
struct ggml_tensor* pooled = nullptr;
|
ggml_tensor* pooled = nullptr;
|
||||||
struct ggml_tensor* pooled_l = nullptr; // [768,]
|
ggml_tensor* pooled_l = nullptr; // [768,]
|
||||||
struct ggml_tensor* pooled_g = nullptr; // [1280,]
|
ggml_tensor* pooled_g = nullptr; // [1280,]
|
||||||
std::vector<float> hidden_states_vec;
|
std::vector<float> hidden_states_vec;
|
||||||
|
|
||||||
size_t chunk_len = 77;
|
size_t chunk_len = 77;
|
||||||
@ -1178,7 +1178,7 @@ struct FluxCLIPEmbedder : public Conditioner {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
if (clip_l) {
|
if (clip_l) {
|
||||||
clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
|
clip_l->get_param_tensors(tensors, "text_encoders.clip_l.transformer.text_model");
|
||||||
}
|
}
|
||||||
@ -1306,10 +1306,10 @@ struct FluxCLIPEmbedder : public Conditioner {
|
|||||||
clip_skip = 2;
|
clip_skip = 2;
|
||||||
}
|
}
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
struct ggml_tensor* hidden_states = nullptr; // [N, n_token, 4096]
|
ggml_tensor* hidden_states = nullptr; // [N, n_token, 4096]
|
||||||
struct ggml_tensor* chunk_hidden_states = nullptr; // [n_token, 4096]
|
ggml_tensor* chunk_hidden_states = nullptr; // [n_token, 4096]
|
||||||
struct ggml_tensor* pooled = nullptr; // [768,]
|
ggml_tensor* pooled = nullptr; // [768,]
|
||||||
std::vector<float> hidden_states_vec;
|
std::vector<float> hidden_states_vec;
|
||||||
|
|
||||||
size_t chunk_count = std::max(clip_l_tokens.size() > 0 ? chunk_len : 0, t5_tokens.size()) / chunk_len;
|
size_t chunk_count = std::max(clip_l_tokens.size() > 0 ? chunk_len : 0, t5_tokens.size()) / chunk_len;
|
||||||
@ -1448,7 +1448,7 @@ struct T5CLIPEmbedder : public Conditioner {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
if (t5) {
|
if (t5) {
|
||||||
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||||
}
|
}
|
||||||
@ -1523,7 +1523,7 @@ struct T5CLIPEmbedder : public Conditioner {
|
|||||||
return {t5_tokens, t5_weights, t5_mask};
|
return {t5_tokens, t5_weights, t5_mask};
|
||||||
}
|
}
|
||||||
|
|
||||||
void modify_mask_to_attend_padding(struct ggml_tensor* mask, int max_seq_length, int num_extra_padding = 8) {
|
void modify_mask_to_attend_padding(ggml_tensor* mask, int max_seq_length, int num_extra_padding = 8) {
|
||||||
float* mask_data = (float*)mask->data;
|
float* mask_data = (float*)mask->data;
|
||||||
int num_pad = 0;
|
int num_pad = 0;
|
||||||
for (int64_t i = 0; i < max_seq_length; i++) {
|
for (int64_t i = 0; i < max_seq_length; i++) {
|
||||||
@ -1554,11 +1554,11 @@ struct T5CLIPEmbedder : public Conditioner {
|
|||||||
auto& t5_weights = std::get<1>(token_and_weights);
|
auto& t5_weights = std::get<1>(token_and_weights);
|
||||||
auto& t5_attn_mask_vec = std::get<2>(token_and_weights);
|
auto& t5_attn_mask_vec = std::get<2>(token_and_weights);
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
struct ggml_tensor* hidden_states = nullptr; // [N, n_token, 4096]
|
ggml_tensor* hidden_states = nullptr; // [N, n_token, 4096]
|
||||||
struct ggml_tensor* chunk_hidden_states = nullptr; // [n_token, 4096]
|
ggml_tensor* chunk_hidden_states = nullptr; // [n_token, 4096]
|
||||||
struct ggml_tensor* pooled = nullptr;
|
ggml_tensor* pooled = nullptr;
|
||||||
struct ggml_tensor* t5_attn_mask = vector_to_ggml_tensor(work_ctx, t5_attn_mask_vec); // [n_token]
|
ggml_tensor* t5_attn_mask = vector_to_ggml_tensor(work_ctx, t5_attn_mask_vec); // [n_token]
|
||||||
|
|
||||||
std::vector<float> hidden_states_vec;
|
std::vector<float> hidden_states_vec;
|
||||||
|
|
||||||
@ -1658,7 +1658,7 @@ struct AnimaConditioner : public Conditioner {
|
|||||||
false);
|
false);
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1736,7 +1736,7 @@ struct AnimaConditioner : public Conditioner {
|
|||||||
|
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, qwen_tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, qwen_tokens);
|
||||||
|
|
||||||
struct ggml_tensor* hidden_states = nullptr; // [N, n_token, 1024]
|
ggml_tensor* hidden_states = nullptr; // [N, n_token, 1024]
|
||||||
llm->compute(n_threads,
|
llm->compute(n_threads,
|
||||||
input_ids,
|
input_ids,
|
||||||
nullptr,
|
nullptr,
|
||||||
@ -1763,8 +1763,8 @@ struct AnimaConditioner : public Conditioner {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* t5_ids_tensor = nullptr;
|
ggml_tensor* t5_ids_tensor = nullptr;
|
||||||
struct ggml_tensor* t5_weight_tensor = nullptr;
|
ggml_tensor* t5_weight_tensor = nullptr;
|
||||||
if (!t5_tokens.empty()) {
|
if (!t5_tokens.empty()) {
|
||||||
t5_ids_tensor = vector_to_ggml_tensor_i32(work_ctx, t5_tokens);
|
t5_ids_tensor = vector_to_ggml_tensor_i32(work_ctx, t5_tokens);
|
||||||
t5_weight_tensor = vector_to_ggml_tensor(work_ctx, t5_weights);
|
t5_weight_tensor = vector_to_ggml_tensor(work_ctx, t5_weights);
|
||||||
@ -1808,7 +1808,7 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
enable_vision);
|
enable_vision);
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1904,7 +1904,7 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
tokenizer->pad_tokens(tokens, weights, max_length, true);
|
tokenizer->pad_tokens(tokens, weights, max_length, true);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* hidden_states = nullptr; // [N, n_token, hidden_size]
|
ggml_tensor* hidden_states = nullptr; // [N, n_token, hidden_size]
|
||||||
|
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
||||||
|
|
||||||
|
|||||||
@ -164,26 +164,26 @@ public:
|
|||||||
blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
|
blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* resblock_forward(std::string name,
|
ggml_tensor* resblock_forward(std::string name,
|
||||||
GGMLRunnerContext* ctx,
|
GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* emb) {
|
ggml_tensor* emb) {
|
||||||
auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
|
auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
|
||||||
return block->forward(ctx, x, emb);
|
return block->forward(ctx, x, emb);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* attention_layer_forward(std::string name,
|
ggml_tensor* attention_layer_forward(std::string name,
|
||||||
GGMLRunnerContext* ctx,
|
GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context) {
|
ggml_tensor* context) {
|
||||||
auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
|
auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
|
||||||
return block->forward(ctx, x, context);
|
return block->forward(ctx, x, context);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
|
ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* hint,
|
ggml_tensor* hint,
|
||||||
struct ggml_tensor* emb,
|
ggml_tensor* emb,
|
||||||
struct ggml_tensor* context) {
|
ggml_tensor* context) {
|
||||||
int num_input_blocks = 15;
|
int num_input_blocks = 15;
|
||||||
auto h = hint;
|
auto h = hint;
|
||||||
for (int i = 0; i < num_input_blocks; i++) {
|
for (int i = 0; i < num_input_blocks; i++) {
|
||||||
@ -198,13 +198,13 @@ public:
|
|||||||
return h;
|
return h;
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* hint,
|
ggml_tensor* hint,
|
||||||
struct ggml_tensor* guided_hint,
|
ggml_tensor* guided_hint,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* y = nullptr) {
|
ggml_tensor* y = nullptr) {
|
||||||
// x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
|
// x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
|
||||||
// timesteps: [N,]
|
// timesteps: [N,]
|
||||||
// context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
|
// context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
|
||||||
@ -246,7 +246,7 @@ public:
|
|||||||
emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim]
|
emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim]
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<struct ggml_tensor*> outs;
|
std::vector<ggml_tensor*> outs;
|
||||||
|
|
||||||
if (guided_hint == nullptr) {
|
if (guided_hint == nullptr) {
|
||||||
guided_hint = input_hint_block_forward(ctx, hint, emb, context);
|
guided_hint = input_hint_block_forward(ctx, hint, emb, context);
|
||||||
@ -312,9 +312,9 @@ struct ControlNet : public GGMLRunner {
|
|||||||
|
|
||||||
ggml_backend_buffer_t control_buffer = nullptr; // keep control output tensors in backend memory
|
ggml_backend_buffer_t control_buffer = nullptr; // keep control output tensors in backend memory
|
||||||
ggml_context* control_ctx = nullptr;
|
ggml_context* control_ctx = nullptr;
|
||||||
std::vector<struct ggml_tensor*> controls; // (12 input block outputs, 1 middle block output) SD 1.5
|
std::vector<ggml_tensor*> controls; // (12 input block outputs, 1 middle block output) SD 1.5
|
||||||
struct ggml_tensor* guided_hint = nullptr; // guided_hint cache, for faster inference
|
ggml_tensor* guided_hint = nullptr; // guided_hint cache, for faster inference
|
||||||
bool guided_hint_cached = false;
|
bool guided_hint_cached = false;
|
||||||
|
|
||||||
ControlNet(ggml_backend_t backend,
|
ControlNet(ggml_backend_t backend,
|
||||||
bool offload_params_to_cpu,
|
bool offload_params_to_cpu,
|
||||||
@ -328,8 +328,8 @@ struct ControlNet : public GGMLRunner {
|
|||||||
free_control_ctx();
|
free_control_ctx();
|
||||||
}
|
}
|
||||||
|
|
||||||
void alloc_control_ctx(std::vector<struct ggml_tensor*> outs) {
|
void alloc_control_ctx(std::vector<ggml_tensor*> outs) {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(outs.size() * ggml_tensor_overhead()) + 1024 * 1024;
|
params.mem_size = static_cast<size_t>(outs.size() * ggml_tensor_overhead()) + 1024 * 1024;
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = true;
|
params.no_alloc = true;
|
||||||
@ -370,16 +370,16 @@ struct ControlNet : public GGMLRunner {
|
|||||||
return "control_net";
|
return "control_net";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
control_net.get_param_tensors(tensors, prefix);
|
control_net.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* hint,
|
ggml_tensor* hint,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* y = nullptr) {
|
ggml_tensor* y = nullptr) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
if (guided_hint_cached) {
|
if (guided_hint_cached) {
|
||||||
@ -414,18 +414,18 @@ struct ControlNet : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* hint,
|
ggml_tensor* hint,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
|
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
|
||||||
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, hint, timesteps, context, y);
|
return build_graph(x, hint, timesteps, context, y);
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@ -773,8 +773,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
// sample_euler_ancestral
|
// sample_euler_ancestral
|
||||||
switch (method) {
|
switch (method) {
|
||||||
case EULER_A_SAMPLE_METHOD: {
|
case EULER_A_SAMPLE_METHOD: {
|
||||||
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
float sigma = sigmas[i];
|
float sigma = sigmas[i];
|
||||||
@ -830,7 +830,7 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
} break;
|
} break;
|
||||||
case EULER_SAMPLE_METHOD: // Implemented without any sigma churn
|
case EULER_SAMPLE_METHOD: // Implemented without any sigma churn
|
||||||
{
|
{
|
||||||
struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
float sigma = sigmas[i];
|
float sigma = sigmas[i];
|
||||||
@ -865,8 +865,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
}
|
}
|
||||||
} break;
|
} break;
|
||||||
case HEUN_SAMPLE_METHOD: {
|
case HEUN_SAMPLE_METHOD: {
|
||||||
struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
// denoise
|
// denoise
|
||||||
@ -921,8 +921,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
}
|
}
|
||||||
} break;
|
} break;
|
||||||
case DPM2_SAMPLE_METHOD: {
|
case DPM2_SAMPLE_METHOD: {
|
||||||
struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
// denoise
|
// denoise
|
||||||
@ -979,8 +979,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
|
|
||||||
} break;
|
} break;
|
||||||
case DPMPP2S_A_SAMPLE_METHOD: {
|
case DPMPP2S_A_SAMPLE_METHOD: {
|
||||||
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
// denoise
|
// denoise
|
||||||
@ -1050,7 +1050,7 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
} break;
|
} break;
|
||||||
case DPMPP2M_SAMPLE_METHOD: // DPM++ (2M) from Karras et al (2022)
|
case DPMPP2M_SAMPLE_METHOD: // DPM++ (2M) from Karras et al (2022)
|
||||||
{
|
{
|
||||||
struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
auto t_fn = [](float sigma) -> float { return -log(sigma); };
|
auto t_fn = [](float sigma) -> float { return -log(sigma); };
|
||||||
|
|
||||||
@ -1092,7 +1092,7 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
} break;
|
} break;
|
||||||
case DPMPP2Mv2_SAMPLE_METHOD: // Modified DPM++ (2M) from https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457
|
case DPMPP2Mv2_SAMPLE_METHOD: // Modified DPM++ (2M) from https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457
|
||||||
{
|
{
|
||||||
struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
auto t_fn = [](float sigma) -> float { return -log(sigma); };
|
auto t_fn = [](float sigma) -> float { return -log(sigma); };
|
||||||
|
|
||||||
@ -1157,8 +1157,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
}
|
}
|
||||||
float* vec_denoised = (float*)denoised->data;
|
float* vec_denoised = (float*)denoised->data;
|
||||||
// d_cur = (x_cur - denoised) / sigma
|
// d_cur = (x_cur - denoised) / sigma
|
||||||
struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x_cur);
|
ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x_cur);
|
||||||
float* vec_d_cur = (float*)d_cur->data;
|
float* vec_d_cur = (float*)d_cur->data;
|
||||||
|
|
||||||
for (int j = 0; j < ggml_nelements(d_cur); j++) {
|
for (int j = 0; j < ggml_nelements(d_cur); j++) {
|
||||||
vec_d_cur[j] = (vec_x_cur[j] - vec_denoised[j]) / sigma;
|
vec_d_cur[j] = (vec_x_cur[j] - vec_denoised[j]) / sigma;
|
||||||
@ -1225,11 +1225,11 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
float t_next = sigmas[i + 1];
|
float t_next = sigmas[i + 1];
|
||||||
|
|
||||||
// Denoising step
|
// Denoising step
|
||||||
ggml_tensor* denoised = model(x, sigma, i + 1);
|
ggml_tensor* denoised = model(x, sigma, i + 1);
|
||||||
float* vec_denoised = (float*)denoised->data;
|
float* vec_denoised = (float*)denoised->data;
|
||||||
struct ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* d_cur = ggml_dup_tensor(work_ctx, x);
|
||||||
float* vec_d_cur = (float*)d_cur->data;
|
float* vec_d_cur = (float*)d_cur->data;
|
||||||
float* vec_x = (float*)x->data;
|
float* vec_x = (float*)x->data;
|
||||||
|
|
||||||
// d_cur = (x - denoised) / sigma
|
// d_cur = (x - denoised) / sigma
|
||||||
for (int j = 0; j < ggml_nelements(d_cur); j++) {
|
for (int j = 0; j < ggml_nelements(d_cur); j++) {
|
||||||
@ -1290,8 +1290,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
} break;
|
} break;
|
||||||
case LCM_SAMPLE_METHOD: // Latent Consistency Models
|
case LCM_SAMPLE_METHOD: // Latent Consistency Models
|
||||||
{
|
{
|
||||||
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* d = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
float sigma = sigmas[i];
|
float sigma = sigmas[i];
|
||||||
@ -1358,9 +1358,9 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
alphas_cumprod[i]);
|
alphas_cumprod[i]);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* pred_original_sample =
|
ggml_tensor* pred_original_sample =
|
||||||
ggml_dup_tensor(work_ctx, x);
|
ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* variance_noise =
|
ggml_tensor* variance_noise =
|
||||||
ggml_dup_tensor(work_ctx, x);
|
ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
@ -1422,7 +1422,7 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
// model_output = model() is the D(x, sigma) as
|
// model_output = model() is the D(x, sigma) as
|
||||||
// defined in Karras et al. (2022), p. 3, Table 1 and
|
// defined in Karras et al. (2022), p. 3, Table 1 and
|
||||||
// p. 8 (7), compare also p. 38 (226) therein.
|
// p. 8 (7), compare also p. 38 (226) therein.
|
||||||
struct ggml_tensor* model_output =
|
ggml_tensor* model_output =
|
||||||
model(x, sigma, i + 1);
|
model(x, sigma, i + 1);
|
||||||
// Here model_output is still the k-diffusion denoiser
|
// Here model_output is still the k-diffusion denoiser
|
||||||
// output, not the U-net output F_theta(c_in(sigma) x;
|
// output, not the U-net output F_theta(c_in(sigma) x;
|
||||||
@ -1545,9 +1545,9 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
}
|
}
|
||||||
int original_steps = 50;
|
int original_steps = 50;
|
||||||
|
|
||||||
struct ggml_tensor* pred_original_sample =
|
ggml_tensor* pred_original_sample =
|
||||||
ggml_dup_tensor(work_ctx, x);
|
ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* noise =
|
ggml_tensor* noise =
|
||||||
ggml_dup_tensor(work_ctx, x);
|
ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
for (int i = 0; i < steps; i++) {
|
for (int i = 0; i < steps; i++) {
|
||||||
@ -1581,7 +1581,7 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
vec_x[j] *= std::sqrt(sigma * sigma + 1);
|
vec_x[j] *= std::sqrt(sigma * sigma + 1);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
struct ggml_tensor* model_output =
|
ggml_tensor* model_output =
|
||||||
model(x, sigma, i + 1);
|
model(x, sigma, i + 1);
|
||||||
{
|
{
|
||||||
float* vec_x = (float*)x->data;
|
float* vec_x = (float*)x->data;
|
||||||
@ -1689,8 +1689,8 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
} break;
|
} break;
|
||||||
case RES_MULTISTEP_SAMPLE_METHOD: // Res Multistep sampler
|
case RES_MULTISTEP_SAMPLE_METHOD: // Res Multistep sampler
|
||||||
{
|
{
|
||||||
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* old_denoised = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
bool have_old_sigma = false;
|
bool have_old_sigma = false;
|
||||||
float old_sigma_down = 0.0f;
|
float old_sigma_down = 0.0f;
|
||||||
@ -1797,9 +1797,9 @@ static bool sample_k_diffusion(sample_method_t method,
|
|||||||
} break;
|
} break;
|
||||||
case RES_2S_SAMPLE_METHOD: // Res 2s sampler
|
case RES_2S_SAMPLE_METHOD: // Res 2s sampler
|
||||||
{
|
{
|
||||||
struct ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* noise = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* x0 = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* x0 = ggml_dup_tensor(work_ctx, x);
|
||||||
struct ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
ggml_tensor* x2 = ggml_dup_tensor(work_ctx, x);
|
||||||
|
|
||||||
const float c2 = 0.5f;
|
const float c2 = 0.5f;
|
||||||
auto t_fn = [](float sigma) -> float { return -logf(sigma); };
|
auto t_fn = [](float sigma) -> float { return -logf(sigma); };
|
||||||
|
|||||||
@ -10,33 +10,33 @@
|
|||||||
#include "z_image.hpp"
|
#include "z_image.hpp"
|
||||||
|
|
||||||
struct DiffusionParams {
|
struct DiffusionParams {
|
||||||
struct ggml_tensor* x = nullptr;
|
ggml_tensor* x = nullptr;
|
||||||
struct ggml_tensor* timesteps = nullptr;
|
ggml_tensor* timesteps = nullptr;
|
||||||
struct ggml_tensor* context = nullptr;
|
ggml_tensor* context = nullptr;
|
||||||
struct ggml_tensor* c_concat = nullptr;
|
ggml_tensor* c_concat = nullptr;
|
||||||
struct ggml_tensor* y = nullptr;
|
ggml_tensor* y = nullptr;
|
||||||
struct ggml_tensor* guidance = nullptr;
|
ggml_tensor* guidance = nullptr;
|
||||||
std::vector<ggml_tensor*> ref_latents = {};
|
std::vector<ggml_tensor*> ref_latents = {};
|
||||||
bool increase_ref_index = false;
|
bool increase_ref_index = false;
|
||||||
int num_video_frames = -1;
|
int num_video_frames = -1;
|
||||||
std::vector<struct ggml_tensor*> controls = {};
|
std::vector<ggml_tensor*> controls = {};
|
||||||
float control_strength = 0.f;
|
float control_strength = 0.f;
|
||||||
struct ggml_tensor* vace_context = nullptr;
|
ggml_tensor* vace_context = nullptr;
|
||||||
float vace_strength = 1.f;
|
float vace_strength = 1.f;
|
||||||
std::vector<int> skip_layers = {};
|
std::vector<int> skip_layers = {};
|
||||||
};
|
};
|
||||||
|
|
||||||
struct DiffusionModel {
|
struct DiffusionModel {
|
||||||
virtual std::string get_desc() = 0;
|
virtual std::string get_desc() = 0;
|
||||||
virtual bool compute(int n_threads,
|
virtual bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) = 0;
|
ggml_context* output_ctx = nullptr) = 0;
|
||||||
virtual void alloc_params_buffer() = 0;
|
virtual void alloc_params_buffer() = 0;
|
||||||
virtual void free_params_buffer() = 0;
|
virtual void free_params_buffer() = 0;
|
||||||
virtual void free_compute_buffer() = 0;
|
virtual void free_compute_buffer() = 0;
|
||||||
virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) = 0;
|
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
|
||||||
virtual size_t get_params_buffer_size() = 0;
|
virtual size_t get_params_buffer_size() = 0;
|
||||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
|
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
|
||||||
virtual int64_t get_adm_in_channels() = 0;
|
virtual int64_t get_adm_in_channels() = 0;
|
||||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||||
@ -69,7 +69,7 @@ struct UNetModel : public DiffusionModel {
|
|||||||
unet.free_compute_buffer();
|
unet.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
unet.get_param_tensors(tensors, "model.diffusion_model");
|
unet.get_param_tensors(tensors, "model.diffusion_model");
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -95,8 +95,8 @@ struct UNetModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return unet.compute(n_threads,
|
return unet.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
@ -134,7 +134,7 @@ struct MMDiTModel : public DiffusionModel {
|
|||||||
mmdit.free_compute_buffer();
|
mmdit.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
mmdit.get_param_tensors(tensors, "model.diffusion_model");
|
mmdit.get_param_tensors(tensors, "model.diffusion_model");
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -160,8 +160,8 @@ struct MMDiTModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return mmdit.compute(n_threads,
|
return mmdit.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
@ -200,7 +200,7 @@ struct FluxModel : public DiffusionModel {
|
|||||||
flux.free_compute_buffer();
|
flux.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
flux.get_param_tensors(tensors, "model.diffusion_model");
|
flux.get_param_tensors(tensors, "model.diffusion_model");
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -226,8 +226,8 @@ struct FluxModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return flux.compute(n_threads,
|
return flux.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
@ -270,7 +270,7 @@ struct AnimaModel : public DiffusionModel {
|
|||||||
anima.free_compute_buffer();
|
anima.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
anima.get_param_tensors(tensors, prefix);
|
anima.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -296,8 +296,8 @@ struct AnimaModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return anima.compute(n_threads,
|
return anima.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
@ -337,7 +337,7 @@ struct WanModel : public DiffusionModel {
|
|||||||
wan.free_compute_buffer();
|
wan.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
wan.get_param_tensors(tensors, prefix);
|
wan.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -363,8 +363,8 @@ struct WanModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return wan.compute(n_threads,
|
return wan.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
@ -408,7 +408,7 @@ struct QwenImageModel : public DiffusionModel {
|
|||||||
qwen_image.free_compute_buffer();
|
qwen_image.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
qwen_image.get_param_tensors(tensors, prefix);
|
qwen_image.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -434,8 +434,8 @@ struct QwenImageModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return qwen_image.compute(n_threads,
|
return qwen_image.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
@ -475,7 +475,7 @@ struct ZImageModel : public DiffusionModel {
|
|||||||
z_image.free_compute_buffer();
|
z_image.free_compute_buffer();
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
z_image.get_param_tensors(tensors, prefix);
|
z_image.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -501,8 +501,8 @@ struct ZImageModel : public DiffusionModel {
|
|||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
DiffusionParams diffusion_params,
|
DiffusionParams diffusion_params,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
return z_image.compute(n_threads,
|
return z_image.compute(n_threads,
|
||||||
diffusion_params.x,
|
diffusion_params.x,
|
||||||
diffusion_params.timesteps,
|
diffusion_params.timesteps,
|
||||||
|
|||||||
@ -27,11 +27,11 @@ public:
|
|||||||
blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* lrelu(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
|
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [n, num_feat, h, w]
|
// x: [n, num_feat, h, w]
|
||||||
// return: [n, num_feat, h, w]
|
// return: [n, num_feat, h, w]
|
||||||
|
|
||||||
@ -64,7 +64,7 @@ public:
|
|||||||
blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
|
blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [n, num_feat, h, w]
|
// x: [n, num_feat, h, w]
|
||||||
// return: [n, num_feat, h, w]
|
// return: [n, num_feat, h, w]
|
||||||
|
|
||||||
@ -112,11 +112,11 @@ public:
|
|||||||
int get_scale() { return scale; }
|
int get_scale() { return scale; }
|
||||||
int get_num_block() { return num_block; }
|
int get_num_block() { return num_block; }
|
||||||
|
|
||||||
struct ggml_tensor* lrelu(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
|
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [n, num_in_ch, h, w]
|
// x: [n, num_in_ch, h, w]
|
||||||
// return: [n, num_out_ch, h*scale, w*scale]
|
// return: [n, num_out_ch, h*scale, w*scale]
|
||||||
auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
|
auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
|
||||||
@ -341,24 +341,24 @@ struct ESRGAN : public GGMLRunner {
|
|||||||
return success;
|
return success;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x) {
|
ggml_cgraph* build_graph(ggml_tensor* x) {
|
||||||
if (!rrdb_net)
|
if (!rrdb_net)
|
||||||
return nullptr;
|
return nullptr;
|
||||||
constexpr int kGraphNodes = 1 << 16; // 65k
|
constexpr int kGraphNodes = 1 << 16; // 65k
|
||||||
struct ggml_cgraph* gf = new_graph_custom(kGraphNodes);
|
ggml_cgraph* gf = new_graph_custom(kGraphNodes);
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
struct ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
|
ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
bool compute(const int n_threads,
|
bool compute(const int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x);
|
return build_graph(x);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
|
||||||
|
|||||||
286
src/flux.hpp
286
src/flux.hpp
@ -19,7 +19,7 @@ namespace Flux {
|
|||||||
blocks["out_layer"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_dim, hidden_dim, bias));
|
blocks["out_layer"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_dim, hidden_dim, bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [..., in_dim]
|
// x: [..., in_dim]
|
||||||
// return: [..., hidden_dim]
|
// return: [..., hidden_dim]
|
||||||
auto in_layer = std::dynamic_pointer_cast<Linear>(blocks["in_layer"]);
|
auto in_layer = std::dynamic_pointer_cast<Linear>(blocks["in_layer"]);
|
||||||
@ -37,7 +37,7 @@ namespace Flux {
|
|||||||
int64_t hidden_size;
|
int64_t hidden_size;
|
||||||
float eps;
|
float eps;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
ggml_type wtype = GGML_TYPE_F32;
|
ggml_type wtype = GGML_TYPE_F32;
|
||||||
params["scale"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
|
params["scale"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
|
||||||
}
|
}
|
||||||
@ -48,10 +48,10 @@ namespace Flux {
|
|||||||
: hidden_size(hidden_size),
|
: hidden_size(hidden_size),
|
||||||
eps(eps) {}
|
eps(eps) {}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
struct ggml_tensor* w = params["scale"];
|
ggml_tensor* w = params["scale"];
|
||||||
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
|
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
|
||||||
x = ggml_mul(ctx->ggml_ctx, x, w);
|
x = ggml_mul(ctx->ggml_ctx, x, w);
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -63,7 +63,7 @@ namespace Flux {
|
|||||||
blocks["key_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim));
|
blocks["key_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* query_norm(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* query_norm(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [..., dim]
|
// x: [..., dim]
|
||||||
// return: [..., dim]
|
// return: [..., dim]
|
||||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["query_norm"]);
|
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["query_norm"]);
|
||||||
@ -72,7 +72,7 @@ namespace Flux {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* key_norm(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* key_norm(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [..., dim]
|
// x: [..., dim]
|
||||||
// return: [..., dim]
|
// return: [..., dim]
|
||||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["key_norm"]);
|
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["key_norm"]);
|
||||||
@ -98,7 +98,7 @@ namespace Flux {
|
|||||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, proj_bias));
|
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, proj_bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<struct ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
|
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
|
||||||
auto norm = std::dynamic_pointer_cast<QKNorm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<QKNorm>(blocks["norm"]);
|
||||||
|
|
||||||
@ -115,17 +115,17 @@ namespace Flux {
|
|||||||
return {q, k, v};
|
return {q, k, v};
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* post_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* post_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
||||||
|
|
||||||
x = proj->forward(ctx, x); // [N, n_token, dim]
|
x = proj->forward(ctx, x); // [N, n_token, dim]
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask) {
|
ggml_tensor* mask) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// pe: [n_token, d_head/2, 2, 2]
|
// pe: [n_token, d_head/2, 2, 2]
|
||||||
// return [N, n_token, dim]
|
// return [N, n_token, dim]
|
||||||
@ -147,7 +147,7 @@ namespace Flux {
|
|||||||
blocks["2"] = std::make_shared<Linear>(intermediate_size, hidden_size, bias);
|
blocks["2"] = std::make_shared<Linear>(intermediate_size, hidden_size, bias);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
|
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
|
||||||
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
|
||||||
|
|
||||||
@ -170,7 +170,7 @@ namespace Flux {
|
|||||||
blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, bias);
|
blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, bias);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
|
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
|
||||||
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
|
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
|
||||||
auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
|
auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
|
||||||
@ -212,7 +212,7 @@ namespace Flux {
|
|||||||
blocks["lin"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * multiplier, bias));
|
blocks["lin"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * multiplier, bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<ModulationOut> forward(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
|
std::vector<ModulationOut> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
||||||
// x: [N, dim]
|
// x: [N, dim]
|
||||||
// return: [ModulationOut, ModulationOut]
|
// return: [ModulationOut, ModulationOut]
|
||||||
auto lin = std::dynamic_pointer_cast<Linear>(blocks["lin"]);
|
auto lin = std::dynamic_pointer_cast<Linear>(blocks["lin"]);
|
||||||
@ -232,11 +232,11 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
__STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
|
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* shift,
|
ggml_tensor* shift,
|
||||||
struct ggml_tensor* scale,
|
ggml_tensor* scale,
|
||||||
bool skip_reshape = false) {
|
bool skip_reshape = false) {
|
||||||
// x: [N, L, C]
|
// x: [N, L, C]
|
||||||
// scale: [N, C]
|
// scale: [N, C]
|
||||||
// shift: [N, C]
|
// shift: [N, C]
|
||||||
@ -294,7 +294,7 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<ModulationOut> get_distil_img_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
|
std::vector<ModulationOut> get_distil_img_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
||||||
// TODO: not hardcoded?
|
// TODO: not hardcoded?
|
||||||
const int single_blocks_count = 38;
|
const int single_blocks_count = 38;
|
||||||
const int double_blocks_count = 19;
|
const int double_blocks_count = 19;
|
||||||
@ -303,7 +303,7 @@ namespace Flux {
|
|||||||
return {ModulationOut(ctx, vec, offset), ModulationOut(ctx, vec, offset + 3)};
|
return {ModulationOut(ctx, vec, offset), ModulationOut(ctx, vec, offset + 3)};
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<ModulationOut> get_distil_txt_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
|
std::vector<ModulationOut> get_distil_txt_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
||||||
// TODO: not hardcoded?
|
// TODO: not hardcoded?
|
||||||
const int single_blocks_count = 38;
|
const int single_blocks_count = 38;
|
||||||
const int double_blocks_count = 19;
|
const int double_blocks_count = 19;
|
||||||
@ -312,14 +312,14 @@ namespace Flux {
|
|||||||
return {ModulationOut(ctx, vec, offset), ModulationOut(ctx, vec, offset + 3)};
|
return {ModulationOut(ctx, vec, offset), ModulationOut(ctx, vec, offset + 3)};
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* img,
|
ggml_tensor* img,
|
||||||
struct ggml_tensor* txt,
|
ggml_tensor* txt,
|
||||||
struct ggml_tensor* vec,
|
ggml_tensor* vec,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
std::vector<ModulationOut> img_mods = {},
|
std::vector<ModulationOut> img_mods = {},
|
||||||
std::vector<ModulationOut> txt_mods = {}) {
|
std::vector<ModulationOut> txt_mods = {}) {
|
||||||
// img: [N, n_img_token, hidden_size]
|
// img: [N, n_img_token, hidden_size]
|
||||||
// txt: [N, n_txt_token, hidden_size]
|
// txt: [N, n_txt_token, hidden_size]
|
||||||
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
|
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
|
||||||
@ -457,17 +457,17 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
ModulationOut get_distil_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
|
ModulationOut get_distil_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
||||||
int64_t offset = 3 * idx;
|
int64_t offset = 3 * idx;
|
||||||
return ModulationOut(ctx, vec, offset);
|
return ModulationOut(ctx, vec, offset);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* vec,
|
ggml_tensor* vec,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
std::vector<ModulationOut> mods = {}) {
|
std::vector<ModulationOut> mods = {}) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// pe: [n_token, d_head/2, 2, 2]
|
// pe: [n_token, d_head/2, 2, 2]
|
||||||
// return: [N, n_token, hidden_size]
|
// return: [N, n_token, hidden_size]
|
||||||
@ -539,7 +539,7 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
ModulationOut get_distil_mod(GGMLRunnerContext* ctx, struct ggml_tensor* vec) {
|
ModulationOut get_distil_mod(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
||||||
int64_t offset = vec->ne[2] - 2;
|
int64_t offset = vec->ne[2] - 2;
|
||||||
int64_t stride = vec->nb[1] * vec->ne[1];
|
int64_t stride = vec->nb[1] * vec->ne[1];
|
||||||
auto shift = ggml_view_2d(ctx->ggml_ctx, vec, vec->ne[0], vec->ne[1], vec->nb[1], stride * (offset + 0)); // [N, dim]
|
auto shift = ggml_view_2d(ctx->ggml_ctx, vec, vec->ne[0], vec->ne[1], vec->nb[1], stride * (offset + 0)); // [N, dim]
|
||||||
@ -548,15 +548,15 @@ namespace Flux {
|
|||||||
return {shift, scale, nullptr};
|
return {shift, scale, nullptr};
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
// return: [N, n_token, patch_size * patch_size * out_channels]
|
// return: [N, n_token, patch_size * patch_size * out_channels]
|
||||||
auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
|
auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
|
||||||
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
||||||
struct ggml_tensor *shift, *scale;
|
ggml_tensor *shift, *scale;
|
||||||
if (prune_mod) {
|
if (prune_mod) {
|
||||||
auto mod = get_distil_mod(ctx, c);
|
auto mod = get_distil_mod(ctx, c);
|
||||||
shift = mod.shift;
|
shift = mod.shift;
|
||||||
@ -589,7 +589,7 @@ namespace Flux {
|
|||||||
blocks["out_proj"] = std::shared_ptr<GGMLBlock>(new Linear(inner_size, hidden_size, true));
|
blocks["out_proj"] = std::shared_ptr<GGMLBlock>(new Linear(inner_size, hidden_size, true));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto in_proj = std::dynamic_pointer_cast<Linear>(blocks["in_proj"]);
|
auto in_proj = std::dynamic_pointer_cast<Linear>(blocks["in_proj"]);
|
||||||
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
|
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
|
||||||
|
|
||||||
@ -612,9 +612,9 @@ namespace Flux {
|
|||||||
blocks["embedder.0"] = std::make_shared<Linear>(in_channels + max_freqs * max_freqs, hidden_size_input);
|
blocks["embedder.0"] = std::make_shared<Linear>(in_channels + max_freqs * max_freqs, hidden_size_input);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* dct) {
|
ggml_tensor* dct) {
|
||||||
// x: (B, P^2, C)
|
// x: (B, P^2, C)
|
||||||
// dct: (1, P^2, max_freqs^2)
|
// dct: (1, P^2, max_freqs^2)
|
||||||
// return: (B, P^2, hidden_size_input)
|
// return: (B, P^2, hidden_size_input)
|
||||||
@ -639,9 +639,9 @@ namespace Flux {
|
|||||||
blocks["norm"] = std::make_shared<RMSNorm>(hidden_size_x);
|
blocks["norm"] = std::make_shared<RMSNorm>(hidden_size_x);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* s) {
|
ggml_tensor* s) {
|
||||||
// x: (batch_size, n_token, hidden_size_x)
|
// x: (batch_size, n_token, hidden_size_x)
|
||||||
// s: (batch_size, hidden_size_s)
|
// s: (batch_size, hidden_size_s)
|
||||||
// return: (batch_size, n_token, hidden_size_x)
|
// return: (batch_size, n_token, hidden_size_x)
|
||||||
@ -689,8 +689,8 @@ namespace Flux {
|
|||||||
blocks["linear"] = std::make_shared<Linear>(hidden_size, out_channels);
|
blocks["linear"] = std::make_shared<Linear>(hidden_size, out_channels);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
||||||
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
||||||
|
|
||||||
@ -708,8 +708,8 @@ namespace Flux {
|
|||||||
blocks["conv"] = std::make_shared<Conv2d>(hidden_size, out_channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
|
blocks["conv"] = std::make_shared<Conv2d>(hidden_size, out_channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
// x: [N, C, H, W]
|
// x: [N, C, H, W]
|
||||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
||||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||||
@ -847,15 +847,15 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* img,
|
ggml_tensor* img,
|
||||||
struct ggml_tensor* txt,
|
ggml_tensor* txt,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor* guidance,
|
ggml_tensor* guidance,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mod_index_arange = nullptr,
|
ggml_tensor* mod_index_arange = nullptr,
|
||||||
std::vector<int> skip_layers = {}) {
|
std::vector<int> skip_layers = {}) {
|
||||||
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
|
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
|
||||||
auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
|
auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
|
||||||
auto final_layer = std::dynamic_pointer_cast<LastLayer>(blocks["final_layer"]);
|
auto final_layer = std::dynamic_pointer_cast<LastLayer>(blocks["final_layer"]);
|
||||||
@ -864,8 +864,8 @@ namespace Flux {
|
|||||||
img = img_in->forward(ctx, img);
|
img = img_in->forward(ctx, img);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* vec;
|
ggml_tensor* vec;
|
||||||
struct ggml_tensor* txt_img_mask = nullptr;
|
ggml_tensor* txt_img_mask = nullptr;
|
||||||
if (params.is_chroma) {
|
if (params.is_chroma) {
|
||||||
int64_t mod_index_length = 344;
|
int64_t mod_index_length = 344;
|
||||||
auto approx = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
|
auto approx = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
|
||||||
@ -967,27 +967,27 @@ namespace Flux {
|
|||||||
return img;
|
return img;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* _apply_x0_residual(GGMLRunnerContext* ctx,
|
ggml_tensor* _apply_x0_residual(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* predicted,
|
ggml_tensor* predicted,
|
||||||
struct ggml_tensor* noisy,
|
ggml_tensor* noisy,
|
||||||
struct ggml_tensor* timesteps) {
|
ggml_tensor* timesteps) {
|
||||||
auto x = ggml_sub(ctx->ggml_ctx, noisy, predicted);
|
auto x = ggml_sub(ctx->ggml_ctx, noisy, predicted);
|
||||||
x = ggml_div(ctx->ggml_ctx, x, timesteps);
|
x = ggml_div(ctx->ggml_ctx, x, timesteps);
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_chroma_radiance(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_chroma_radiance(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor* guidance,
|
ggml_tensor* guidance,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mod_index_arange = nullptr,
|
ggml_tensor* mod_index_arange = nullptr,
|
||||||
struct ggml_tensor* dct = nullptr,
|
ggml_tensor* dct = nullptr,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
std::vector<int> skip_layers = {}) {
|
std::vector<int> skip_layers = {}) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
|
|
||||||
int64_t W = x->ne[0];
|
int64_t W = x->ne[0];
|
||||||
@ -1050,18 +1050,18 @@ namespace Flux {
|
|||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_flux_chroma(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_flux_chroma(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor* guidance,
|
ggml_tensor* guidance,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mod_index_arange = nullptr,
|
ggml_tensor* mod_index_arange = nullptr,
|
||||||
struct ggml_tensor* dct = nullptr,
|
ggml_tensor* dct = nullptr,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
std::vector<int> skip_layers = {}) {
|
std::vector<int> skip_layers = {}) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
|
|
||||||
int64_t W = x->ne[0];
|
int64_t W = x->ne[0];
|
||||||
@ -1119,18 +1119,18 @@ namespace Flux {
|
|||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor* guidance,
|
ggml_tensor* guidance,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mod_index_arange = nullptr,
|
ggml_tensor* mod_index_arange = nullptr,
|
||||||
struct ggml_tensor* dct = nullptr,
|
ggml_tensor* dct = nullptr,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
std::vector<int> skip_layers = {}) {
|
std::vector<int> skip_layers = {}) {
|
||||||
// Forward pass of DiT.
|
// Forward pass of DiT.
|
||||||
// x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
// x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
||||||
// timestep: (N,) tensor of diffusion timesteps
|
// timestep: (N,) tensor of diffusion timesteps
|
||||||
@ -1299,7 +1299,7 @@ namespace Flux {
|
|||||||
return "flux";
|
return "flux";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
flux.get_param_tensors(tensors, prefix);
|
flux.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1353,20 +1353,20 @@ namespace Flux {
|
|||||||
return dct;
|
return dct;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor* guidance,
|
ggml_tensor* guidance,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
bool increase_ref_index = false,
|
bool increase_ref_index = false,
|
||||||
std::vector<int> skip_layers = {}) {
|
std::vector<int> skip_layers = {}) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
struct ggml_cgraph* gf = new_graph_custom(FLUX_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(FLUX_GRAPH_SIZE);
|
||||||
|
|
||||||
struct ggml_tensor* mod_index_arange = nullptr;
|
ggml_tensor* mod_index_arange = nullptr;
|
||||||
struct ggml_tensor* dct = nullptr; // for chroma radiance
|
ggml_tensor* dct = nullptr; // for chroma radiance
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
context = to_backend(context);
|
context = to_backend(context);
|
||||||
@ -1437,18 +1437,18 @@ namespace Flux {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = flux.forward(&runner_ctx,
|
ggml_tensor* out = flux.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
timesteps,
|
timesteps,
|
||||||
context,
|
context,
|
||||||
c_concat,
|
c_concat,
|
||||||
y,
|
y,
|
||||||
guidance,
|
guidance,
|
||||||
pe,
|
pe,
|
||||||
mod_index_arange,
|
mod_index_arange,
|
||||||
dct,
|
dct,
|
||||||
ref_latents,
|
ref_latents,
|
||||||
skip_layers);
|
skip_layers);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -1456,23 +1456,23 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor* guidance,
|
ggml_tensor* guidance,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
bool increase_ref_index = false,
|
bool increase_ref_index = false,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr,
|
ggml_context* output_ctx = nullptr,
|
||||||
std::vector<int> skip_layers = std::vector<int>()) {
|
std::vector<int> skip_layers = std::vector<int>()) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]
|
// context: [N, max_position, hidden_size]
|
||||||
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
||||||
// guidance: [N, ]
|
// guidance: [N, ]
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
|
return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -1480,12 +1480,12 @@ namespace Flux {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -1513,7 +1513,7 @@ namespace Flux {
|
|||||||
auto y = nullptr;
|
auto y = nullptr;
|
||||||
// print_ggml_tensor(y);
|
// print_ggml_tensor(y);
|
||||||
|
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
compute(8, x, timesteps, context, nullptr, y, guidance, {}, false, &out, work_ctx);
|
compute(8, x, timesteps, context, nullptr, y, guidance, {}, false, &out, work_ctx);
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@ -163,7 +163,7 @@ const float sd_latent_rgb_proj[4][3] = {
|
|||||||
{-0.178022f, -0.200862f, -0.678514f}};
|
{-0.178022f, -0.200862f, -0.678514f}};
|
||||||
float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f};
|
float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f};
|
||||||
|
|
||||||
void preview_latent_video(uint8_t* buffer, struct ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
|
void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
|
||||||
size_t buffer_head = 0;
|
size_t buffer_head = 0;
|
||||||
|
|
||||||
uint32_t latent_width = static_cast<uint32_t>(latents->ne[0]);
|
uint32_t latent_width = static_cast<uint32_t>(latents->ne[0]);
|
||||||
|
|||||||
180
src/llm.hpp
180
src/llm.hpp
@ -522,7 +522,7 @@ namespace LLM {
|
|||||||
blocks["down_proj"] = std::shared_ptr<GGMLBlock>(new Linear(intermediate_size, hidden_size, bias));
|
blocks["down_proj"] = std::shared_ptr<GGMLBlock>(new Linear(intermediate_size, hidden_size, bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
|
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
|
||||||
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
|
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
|
||||||
@ -582,7 +582,7 @@ namespace LLM {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N*grid_t*grid_h*grid_w, in_channels, temporal_patch_size*patch_size*patch_size]
|
// x: [N*grid_t*grid_h*grid_w, in_channels, temporal_patch_size*patch_size*patch_size]
|
||||||
// return: [N*grid_t*grid_h*grid_w, embed_dim]
|
// return: [N*grid_t*grid_h*grid_w, embed_dim]
|
||||||
x = ggml_reshape_4d(ctx->ggml_ctx,
|
x = ggml_reshape_4d(ctx->ggml_ctx,
|
||||||
@ -631,7 +631,7 @@ namespace LLM {
|
|||||||
blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, dim));
|
blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto ln_q = std::dynamic_pointer_cast<RMSNorm>(blocks["ln_q"]);
|
auto ln_q = std::dynamic_pointer_cast<RMSNorm>(blocks["ln_q"]);
|
||||||
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
||||||
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
|
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
|
||||||
@ -668,10 +668,10 @@ namespace LLM {
|
|||||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size));
|
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr) {
|
ggml_tensor* mask = nullptr) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
int64_t n_token = x->ne[1];
|
int64_t n_token = x->ne[1];
|
||||||
int64_t N = x->ne[2];
|
int64_t N = x->ne[2];
|
||||||
@ -718,10 +718,10 @@ namespace LLM {
|
|||||||
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_size, eps));
|
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_size, eps));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr) {
|
ggml_tensor* mask = nullptr) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
auto attn = std::dynamic_pointer_cast<VisionAttention>(blocks["attn"]);
|
auto attn = std::dynamic_pointer_cast<VisionAttention>(blocks["attn"]);
|
||||||
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
|
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
|
||||||
@ -778,12 +778,12 @@ namespace LLM {
|
|||||||
blocks["merger"] = std::shared_ptr<GGMLBlock>(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size));
|
blocks["merger"] = std::shared_ptr<GGMLBlock>(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* pixel_values,
|
ggml_tensor* pixel_values,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* window_index,
|
ggml_tensor* window_index,
|
||||||
struct ggml_tensor* window_inverse_index,
|
ggml_tensor* window_inverse_index,
|
||||||
struct ggml_tensor* window_mask) {
|
ggml_tensor* window_mask) {
|
||||||
// pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
|
// pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
|
||||||
// window_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
// window_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
||||||
// window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
// window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
||||||
@ -836,10 +836,10 @@ namespace LLM {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
struct ggml_tensor* attention_mask = nullptr) {
|
ggml_tensor* attention_mask = nullptr) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
int64_t n_token = x->ne[1];
|
int64_t n_token = x->ne[1];
|
||||||
int64_t N = x->ne[2];
|
int64_t N = x->ne[2];
|
||||||
@ -898,10 +898,10 @@ namespace LLM {
|
|||||||
blocks["post_attention_layernorm"] = std::make_shared<RMSNorm>(params.hidden_size, params.rms_norm_eps);
|
blocks["post_attention_layernorm"] = std::make_shared<RMSNorm>(params.hidden_size, params.rms_norm_eps);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
struct ggml_tensor* attention_mask = nullptr) {
|
ggml_tensor* attention_mask = nullptr) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
auto self_attn = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
|
auto self_attn = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
|
||||||
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
|
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
|
||||||
@ -936,12 +936,12 @@ namespace LLM {
|
|||||||
blocks["norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.hidden_size, params.rms_norm_eps));
|
blocks["norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.hidden_size, params.rms_norm_eps));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
struct ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
||||||
std::set<int> out_layers) {
|
std::set<int> out_layers) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
// return: [N, n_token, hidden_size]
|
// return: [N, n_token, hidden_size]
|
||||||
|
|
||||||
@ -1037,12 +1037,12 @@ namespace LLM {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
struct ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
||||||
std::set<int> out_layers) {
|
std::set<int> out_layers) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
|
auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
|
||||||
|
|
||||||
@ -1050,12 +1050,12 @@ namespace LLM {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
|
ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* pixel_values,
|
ggml_tensor* pixel_values,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* window_index,
|
ggml_tensor* window_index,
|
||||||
struct ggml_tensor* window_inverse_index,
|
ggml_tensor* window_inverse_index,
|
||||||
struct ggml_tensor* window_mask) {
|
ggml_tensor* window_mask) {
|
||||||
GGML_ASSERT(enable_vision);
|
GGML_ASSERT(enable_vision);
|
||||||
auto vision_model = std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
|
auto vision_model = std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
|
||||||
return vision_model->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
|
return vision_model->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
|
||||||
@ -1156,35 +1156,35 @@ namespace LLM {
|
|||||||
return llm_arch_to_str[static_cast<int>(params.arch)];
|
return llm_arch_to_str[static_cast<int>(params.arch)];
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
model.get_param_tensors(tensors, prefix);
|
model.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
struct ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
||||||
std::set<int> out_layers) {
|
std::set<int> out_layers) {
|
||||||
auto hidden_states = model.forward(ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers); // [N, n_token, hidden_size]
|
auto hidden_states = model.forward(ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers); // [N, n_token, hidden_size]
|
||||||
return hidden_states;
|
return hidden_states;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
|
ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* pixel_values,
|
ggml_tensor* pixel_values,
|
||||||
struct ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
struct ggml_tensor* window_index,
|
ggml_tensor* window_index,
|
||||||
struct ggml_tensor* window_inverse_index,
|
ggml_tensor* window_inverse_index,
|
||||||
struct ggml_tensor* window_mask) {
|
ggml_tensor* window_mask) {
|
||||||
auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask);
|
auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask);
|
||||||
return hidden_states;
|
return hidden_states;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
|
ggml_cgraph* build_graph(ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
||||||
std::set<int> out_layers) {
|
std::set<int> out_layers) {
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
|
|
||||||
input_ids = to_backend(input_ids);
|
input_ids = to_backend(input_ids);
|
||||||
|
|
||||||
@ -1232,7 +1232,7 @@ namespace LLM {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);
|
ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, input_pos, attention_mask, image_embeds, out_layers);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, hidden_states);
|
ggml_build_forward_expand(gf, hidden_states);
|
||||||
|
|
||||||
@ -1240,13 +1240,13 @@ namespace LLM {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(const int n_threads,
|
bool compute(const int n_threads,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
||||||
std::set<int> out_layers,
|
std::set<int> out_layers,
|
||||||
ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(input_ids, attention_mask, image_embeds, out_layers);
|
return build_graph(input_ids, attention_mask, image_embeds, out_layers);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
||||||
@ -1261,7 +1261,7 @@ namespace LLM {
|
|||||||
return grid_t * grid_h * grid_w;
|
return grid_t * grid_h * grid_w;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* process_image(struct ggml_context* ctx, struct ggml_tensor* image) {
|
ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
|
||||||
// image: [C, H, W]
|
// image: [C, H, W]
|
||||||
// return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
|
// return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
|
||||||
int64_t C = image->ne[2];
|
int64_t C = image->ne[2];
|
||||||
@ -1288,8 +1288,8 @@ namespace LLM {
|
|||||||
return image;
|
return image;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_encode_image_graph(struct ggml_tensor* image) {
|
ggml_cgraph* build_encode_image_graph(ggml_tensor* image) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
|
||||||
|
|
||||||
GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
|
GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
|
||||||
GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
|
GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
|
||||||
@ -1399,23 +1399,23 @@ namespace LLM {
|
|||||||
// pe->data = nullptr;
|
// pe->data = nullptr;
|
||||||
set_backend_tensor_data(pe, pe_vec.data());
|
set_backend_tensor_data(pe, pe_vec.data());
|
||||||
|
|
||||||
auto runnter_ctx = get_context();
|
auto runnter_ctx = get_context();
|
||||||
struct ggml_tensor* hidden_states = vision_forward(&runnter_ctx,
|
ggml_tensor* hidden_states = vision_forward(&runnter_ctx,
|
||||||
pixel_values,
|
pixel_values,
|
||||||
pe,
|
pe,
|
||||||
window_index,
|
window_index,
|
||||||
window_inverse_index,
|
window_inverse_index,
|
||||||
window_mask);
|
window_mask);
|
||||||
ggml_build_forward_expand(gf, hidden_states);
|
ggml_build_forward_expand(gf, hidden_states);
|
||||||
|
|
||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
void encode_image(const int n_threads,
|
void encode_image(const int n_threads,
|
||||||
struct ggml_tensor* image,
|
ggml_tensor* image,
|
||||||
ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_encode_image_graph(image);
|
return build_encode_image_graph(image);
|
||||||
};
|
};
|
||||||
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
|
GGMLRunner::compute(get_graph, n_threads, false, output, output_ctx);
|
||||||
@ -1440,7 +1440,7 @@ namespace LLM {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
model.get_param_tensors(tensors, prefix);
|
model.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1492,12 +1492,12 @@ namespace LLM {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
bool test_mistral = false;
|
bool test_mistral = false;
|
||||||
bool test_qwen3 = true;
|
bool test_qwen3 = true;
|
||||||
@ -1509,7 +1509,7 @@ namespace LLM {
|
|||||||
{
|
{
|
||||||
auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
|
auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
|
||||||
print_ggml_tensor(image, false, "image");
|
print_ggml_tensor(image, false, "image");
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.encode_image(8, image, &out, work_ctx);
|
model.encode_image(8, image, &out, work_ctx);
|
||||||
@ -1547,8 +1547,8 @@ namespace LLM {
|
|||||||
printf("%d ", token);
|
printf("%d ", token);
|
||||||
}
|
}
|
||||||
printf("\n");
|
printf("\n");
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.compute(8, input_ids, nullptr, image_embeds, {}, &out, work_ctx);
|
model.compute(8, input_ids, nullptr, image_embeds, {}, &out, work_ctx);
|
||||||
@ -1561,7 +1561,7 @@ namespace LLM {
|
|||||||
// ggml_set_f32(image, 0.f);
|
// ggml_set_f32(image, 0.f);
|
||||||
auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
|
auto image = load_tensor_from_file(work_ctx, "qwen2vl_normalized.bin");
|
||||||
print_ggml_tensor(image, false, "image");
|
print_ggml_tensor(image, false, "image");
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.encode_image(8, image, &out, work_ctx);
|
model.encode_image(8, image, &out, work_ctx);
|
||||||
@ -1587,8 +1587,8 @@ namespace LLM {
|
|||||||
printf("%d ", token);
|
printf("%d ", token);
|
||||||
}
|
}
|
||||||
printf("\n");
|
printf("\n");
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.compute(8, input_ids, nullptr, {}, {10, 20, 30}, &out, work_ctx);
|
model.compute(8, input_ids, nullptr, {}, {10, 20, 30}, &out, work_ctx);
|
||||||
@ -1610,8 +1610,8 @@ namespace LLM {
|
|||||||
printf("%d ", token);
|
printf("%d ", token);
|
||||||
}
|
}
|
||||||
printf("\n");
|
printf("\n");
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.compute(8, input_ids, nullptr, {}, {35}, &out, work_ctx);
|
model.compute(8, input_ids, nullptr, {}, {35}, &out, work_ctx);
|
||||||
@ -1633,8 +1633,8 @@ namespace LLM {
|
|||||||
printf("%d ", token);
|
printf("%d ", token);
|
||||||
}
|
}
|
||||||
printf("\n");
|
printf("\n");
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.compute(8, input_ids, nullptr, {}, {}, &out, work_ctx);
|
model.compute(8, input_ids, nullptr, {}, {}, &out, work_ctx);
|
||||||
|
|||||||
32
src/lora.hpp
32
src/lora.hpp
@ -9,7 +9,7 @@
|
|||||||
struct LoraModel : public GGMLRunner {
|
struct LoraModel : public GGMLRunner {
|
||||||
std::string lora_id;
|
std::string lora_id;
|
||||||
float multiplier = 1.0f;
|
float multiplier = 1.0f;
|
||||||
std::unordered_map<std::string, struct ggml_tensor*> lora_tensors;
|
std::unordered_map<std::string, ggml_tensor*> lora_tensors;
|
||||||
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
|
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
|
||||||
std::set<std::string> applied_lora_tensors;
|
std::set<std::string> applied_lora_tensors;
|
||||||
std::string file_path;
|
std::string file_path;
|
||||||
@ -76,13 +76,13 @@ struct LoraModel : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
for (const auto& pair : tensors_to_create) {
|
for (const auto& pair : tensors_to_create) {
|
||||||
const auto& name = pair.first;
|
const auto& name = pair.first;
|
||||||
const auto& ts = pair.second;
|
const auto& ts = pair.second;
|
||||||
struct ggml_tensor* real = ggml_new_tensor(params_ctx,
|
ggml_tensor* real = ggml_new_tensor(params_ctx,
|
||||||
ts.type,
|
ts.type,
|
||||||
ts.n_dims,
|
ts.n_dims,
|
||||||
ts.ne);
|
ts.ne);
|
||||||
lora_tensors[name] = real;
|
lora_tensors[name] = real;
|
||||||
}
|
}
|
||||||
|
|
||||||
alloc_params_buffer();
|
alloc_params_buffer();
|
||||||
@ -337,10 +337,10 @@ struct LoraModel : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
scale_value *= multiplier;
|
scale_value *= multiplier;
|
||||||
|
|
||||||
struct ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
|
ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
|
||||||
struct ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
|
ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
|
||||||
auto curr_updown = ggml_mul_inplace(ctx, updown_1, updown_2);
|
auto curr_updown = ggml_mul_inplace(ctx, updown_1, updown_2);
|
||||||
curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
|
curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
|
||||||
if (updown == nullptr) {
|
if (updown == nullptr) {
|
||||||
updown = curr_updown;
|
updown = curr_updown;
|
||||||
} else {
|
} else {
|
||||||
@ -747,9 +747,9 @@ struct LoraModel : public GGMLRunner {
|
|||||||
return out_diff;
|
return out_diff;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
|
ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
|
||||||
size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
|
size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
|
||||||
struct ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
|
ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
|
||||||
|
|
||||||
preprocess_lora_tensors(model_tensors);
|
preprocess_lora_tensors(model_tensors);
|
||||||
|
|
||||||
@ -788,8 +788,8 @@ struct LoraModel : public GGMLRunner {
|
|||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
void apply(std::map<std::string, struct ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
|
void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_lora_graph(model_tensors, version);
|
return build_lora_graph(model_tensors, version);
|
||||||
};
|
};
|
||||||
GGMLRunner::compute(get_graph, n_threads, false);
|
GGMLRunner::compute(get_graph, n_threads, false);
|
||||||
|
|||||||
@ -26,9 +26,9 @@ namespace LTXV {
|
|||||||
bias));
|
bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
bool causal = true) {
|
bool causal = true) {
|
||||||
// x: [N*IC, ID, IH, IW]
|
// x: [N*IC, ID, IH, IW]
|
||||||
// result: [N*OC, OD, OH, OW]
|
// result: [N*OC, OD, OH, OW]
|
||||||
auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
|
auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
|
||||||
|
|||||||
174
src/mmdit.hpp
174
src/mmdit.hpp
@ -27,7 +27,7 @@ public:
|
|||||||
blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
|
blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, n_token, in_features]
|
// x: [N, n_token, in_features]
|
||||||
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
||||||
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
||||||
@ -72,7 +72,7 @@ public:
|
|||||||
bias));
|
bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, C, H, W]
|
// x: [N, C, H, W]
|
||||||
// return: [N, H*W, embed_dim]
|
// return: [N, H*W, embed_dim]
|
||||||
auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
|
auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
|
||||||
@ -111,7 +111,7 @@ public:
|
|||||||
blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_channels, true, true));
|
blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_channels, true, true));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* t) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
|
||||||
// t: [N, ]
|
// t: [N, ]
|
||||||
// return: [N, hidden_size]
|
// return: [N, hidden_size]
|
||||||
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
||||||
@ -135,7 +135,7 @@ public:
|
|||||||
blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size, true, true));
|
blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size, true, true));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, input_dim]
|
// x: [N, input_dim]
|
||||||
// return: [N, hidden_size]
|
// return: [N, hidden_size]
|
||||||
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
|
||||||
@ -175,7 +175,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<struct ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
|
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
|
||||||
|
|
||||||
auto qkv = qkv_proj->forward(ctx, x);
|
auto qkv = qkv_proj->forward(ctx, x);
|
||||||
@ -198,7 +198,7 @@ public:
|
|||||||
return {q, k, v};
|
return {q, k, v};
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* post_attention(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* post_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
GGML_ASSERT(!pre_only);
|
GGML_ASSERT(!pre_only);
|
||||||
|
|
||||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
||||||
@ -208,8 +208,8 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
auto qkv = pre_attention(ctx, 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 = 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 = post_attention(ctx, x); // [N, n_token, dim]
|
||||||
@ -217,10 +217,10 @@ public:
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
__STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
|
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* shift,
|
ggml_tensor* shift,
|
||||||
struct ggml_tensor* scale) {
|
ggml_tensor* scale) {
|
||||||
// x: [N, L, C]
|
// x: [N, L, C]
|
||||||
// scale: [N, C]
|
// scale: [N, C]
|
||||||
// shift: [N, C]
|
// shift: [N, C]
|
||||||
@ -274,8 +274,8 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
std::tuple<std::vector<ggml_tensor*>, std::vector<ggml_tensor*>, std::vector<ggml_tensor*>> pre_attention_x(GGMLRunnerContext* ctx,
|
std::tuple<std::vector<ggml_tensor*>, std::vector<ggml_tensor*>, std::vector<ggml_tensor*>> pre_attention_x(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
GGML_ASSERT(self_attn);
|
GGML_ASSERT(self_attn);
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
@ -309,9 +309,9 @@ public:
|
|||||||
return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}};
|
return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}};
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<std::vector<struct ggml_tensor*>, std::vector<struct ggml_tensor*>> pre_attention(GGMLRunnerContext* ctx,
|
std::pair<std::vector<ggml_tensor*>, std::vector<ggml_tensor*>> pre_attention(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
|
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
|
||||||
@ -346,15 +346,15 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* post_attention_x(GGMLRunnerContext* ctx,
|
ggml_tensor* post_attention_x(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* attn_out,
|
ggml_tensor* attn_out,
|
||||||
struct ggml_tensor* attn2_out,
|
ggml_tensor* attn2_out,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* gate_msa,
|
ggml_tensor* gate_msa,
|
||||||
struct ggml_tensor* shift_mlp,
|
ggml_tensor* shift_mlp,
|
||||||
struct ggml_tensor* scale_mlp,
|
ggml_tensor* scale_mlp,
|
||||||
struct ggml_tensor* gate_mlp,
|
ggml_tensor* gate_mlp,
|
||||||
struct ggml_tensor* gate_msa2) {
|
ggml_tensor* gate_msa2) {
|
||||||
// attn_out: [N, n_token, hidden_size]
|
// attn_out: [N, n_token, hidden_size]
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// gate_msa: [N, hidden_size]
|
// gate_msa: [N, hidden_size]
|
||||||
@ -384,13 +384,13 @@ public:
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* post_attention(GGMLRunnerContext* ctx,
|
ggml_tensor* post_attention(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* attn_out,
|
ggml_tensor* attn_out,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* gate_msa,
|
ggml_tensor* gate_msa,
|
||||||
struct ggml_tensor* shift_mlp,
|
ggml_tensor* shift_mlp,
|
||||||
struct ggml_tensor* scale_mlp,
|
ggml_tensor* scale_mlp,
|
||||||
struct ggml_tensor* gate_mlp) {
|
ggml_tensor* gate_mlp) {
|
||||||
// attn_out: [N, n_token, hidden_size]
|
// attn_out: [N, n_token, hidden_size]
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// gate_msa: [N, hidden_size]
|
// gate_msa: [N, hidden_size]
|
||||||
@ -416,9 +416,9 @@ public:
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
// return: [N, n_token, hidden_size]
|
// return: [N, n_token, hidden_size]
|
||||||
@ -463,11 +463,11 @@ public:
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
__STATIC_INLINE__ std::pair<struct ggml_tensor*, struct ggml_tensor*>
|
__STATIC_INLINE__ std::pair<ggml_tensor*, ggml_tensor*>
|
||||||
block_mixing(GGMLRunnerContext* ctx,
|
block_mixing(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c,
|
ggml_tensor* c,
|
||||||
std::shared_ptr<DismantledBlock> context_block,
|
std::shared_ptr<DismantledBlock> context_block,
|
||||||
std::shared_ptr<DismantledBlock> x_block) {
|
std::shared_ptr<DismantledBlock> x_block) {
|
||||||
// context: [N, n_context, hidden_size]
|
// context: [N, n_context, hidden_size]
|
||||||
@ -489,7 +489,7 @@ block_mixing(GGMLRunnerContext* ctx,
|
|||||||
x_qkv = x_qkv_intermediates.first;
|
x_qkv = x_qkv_intermediates.first;
|
||||||
x_intermediates = x_qkv_intermediates.second;
|
x_intermediates = x_qkv_intermediates.second;
|
||||||
}
|
}
|
||||||
std::vector<struct ggml_tensor*> qkv;
|
std::vector<ggml_tensor*> qkv;
|
||||||
for (int i = 0; i < 3; i++) {
|
for (int i = 0; i < 3; i++) {
|
||||||
qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
|
qkv.push_back(ggml_concat(ctx->ggml_ctx, context_qkv[i], x_qkv[i], 1));
|
||||||
}
|
}
|
||||||
@ -563,10 +563,10 @@ public:
|
|||||||
blocks["x_block"] = std::shared_ptr<GGMLBlock>(new DismantledBlock(hidden_size, num_heads, mlp_ratio, qk_norm, qkv_bias, false, self_attn_x));
|
blocks["x_block"] = std::shared_ptr<GGMLBlock>(new DismantledBlock(hidden_size, num_heads, mlp_ratio, qk_norm, qkv_bias, false, self_attn_x));
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
auto context_block = std::dynamic_pointer_cast<DismantledBlock>(blocks["context_block"]);
|
auto context_block = std::dynamic_pointer_cast<DismantledBlock>(blocks["context_block"]);
|
||||||
auto x_block = std::dynamic_pointer_cast<DismantledBlock>(blocks["x_block"]);
|
auto x_block = std::dynamic_pointer_cast<DismantledBlock>(blocks["x_block"]);
|
||||||
|
|
||||||
@ -586,9 +586,9 @@ public:
|
|||||||
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size));
|
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
// return: [N, n_token, patch_size * patch_size * out_channels]
|
// return: [N, n_token, patch_size * patch_size * out_channels]
|
||||||
@ -626,7 +626,7 @@ protected:
|
|||||||
int64_t hidden_size;
|
int64_t hidden_size;
|
||||||
std::string qk_norm;
|
std::string qk_norm;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
|
||||||
enum ggml_type wtype = GGML_TYPE_F32;
|
enum ggml_type wtype = GGML_TYPE_F32;
|
||||||
params["pos_embed"] = ggml_new_tensor_3d(ctx, wtype, hidden_size, num_patchs, 1);
|
params["pos_embed"] = ggml_new_tensor_3d(ctx, wtype, hidden_size, num_patchs, 1);
|
||||||
}
|
}
|
||||||
@ -705,8 +705,8 @@ public:
|
|||||||
blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(hidden_size, patch_size, out_channels));
|
blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(hidden_size, patch_size, out_channels));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor*
|
ggml_tensor*
|
||||||
cropped_pos_embed(struct ggml_context* ctx,
|
cropped_pos_embed(ggml_context* ctx,
|
||||||
int64_t h,
|
int64_t h,
|
||||||
int64_t w) {
|
int64_t w) {
|
||||||
auto pos_embed = params["pos_embed"];
|
auto pos_embed = params["pos_embed"];
|
||||||
@ -745,11 +745,11 @@ public:
|
|||||||
return spatial_pos_embed;
|
return spatial_pos_embed;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_core_with_concat(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_core_with_concat(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c_mod,
|
ggml_tensor* c_mod,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
std::vector<int> skip_layers = std::vector<int>()) {
|
std::vector<int> skip_layers = std::vector<int>()) {
|
||||||
// x: [N, H*W, hidden_size]
|
// x: [N, H*W, hidden_size]
|
||||||
// context: [N, n_context, d_context]
|
// context: [N, n_context, d_context]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
@ -774,12 +774,12 @@ public:
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* t,
|
ggml_tensor* t,
|
||||||
struct ggml_tensor* y = nullptr,
|
ggml_tensor* y = nullptr,
|
||||||
struct ggml_tensor* context = nullptr,
|
ggml_tensor* context = nullptr,
|
||||||
std::vector<int> skip_layers = std::vector<int>()) {
|
std::vector<int> skip_layers = std::vector<int>()) {
|
||||||
// Forward pass of DiT.
|
// Forward pass of DiT.
|
||||||
// x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
// x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
|
||||||
// t: (N,) tensor of diffusion timesteps
|
// t: (N,) tensor of diffusion timesteps
|
||||||
@ -832,29 +832,29 @@ struct MMDiTRunner : public GGMLRunner {
|
|||||||
return "mmdit";
|
return "mmdit";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
mmdit.get_param_tensors(tensors, prefix);
|
mmdit.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
std::vector<int> skip_layers = std::vector<int>()) {
|
std::vector<int> skip_layers = std::vector<int>()) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(MMDIT_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(MMDIT_GRAPH_SIZE);
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
context = to_backend(context);
|
context = to_backend(context);
|
||||||
y = to_backend(y);
|
y = to_backend(y);
|
||||||
timesteps = to_backend(timesteps);
|
timesteps = to_backend(timesteps);
|
||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
struct ggml_tensor* out = mmdit.forward(&runner_ctx,
|
ggml_tensor* out = mmdit.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
timesteps,
|
timesteps,
|
||||||
y,
|
y,
|
||||||
context,
|
context,
|
||||||
skip_layers);
|
skip_layers);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -862,18 +862,18 @@ struct MMDiTRunner : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr,
|
ggml_context* output_ctx = nullptr,
|
||||||
std::vector<int> skip_layers = std::vector<int>()) {
|
std::vector<int> skip_layers = std::vector<int>()) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]([N, 154, 4096]) or [1, max_position, hidden_size]
|
// context: [N, max_position, hidden_size]([N, 154, 4096]) or [1, max_position, hidden_size]
|
||||||
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, y, skip_layers);
|
return build_graph(x, timesteps, context, y, skip_layers);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -881,12 +881,12 @@ struct MMDiTRunner : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -908,7 +908,7 @@ struct MMDiTRunner : public GGMLRunner {
|
|||||||
ggml_set_f32(y, 0.01f);
|
ggml_set_f32(y, 0.01f);
|
||||||
// print_ggml_tensor(y);
|
// print_ggml_tensor(y);
|
||||||
|
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
compute(8, x, timesteps, context, y, &out, work_ctx);
|
compute(8, x, timesteps, context, y, &out, work_ctx);
|
||||||
|
|||||||
@ -287,7 +287,7 @@ void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool is_zip_file(const std::string& file_path) {
|
bool is_zip_file(const std::string& file_path) {
|
||||||
struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
||||||
if (zip == nullptr) {
|
if (zip == nullptr) {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
@ -453,9 +453,9 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
|
|||||||
size_t total_size = 0;
|
size_t total_size = 0;
|
||||||
size_t data_offset = gguf_get_data_offset(ctx_gguf_);
|
size_t data_offset = gguf_get_data_offset(ctx_gguf_);
|
||||||
for (int i = 0; i < n_tensors; i++) {
|
for (int i = 0; i < n_tensors; i++) {
|
||||||
std::string name = gguf_get_tensor_name(ctx_gguf_, i);
|
std::string name = gguf_get_tensor_name(ctx_gguf_, i);
|
||||||
struct ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
|
ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
|
||||||
size_t offset = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
|
size_t offset = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
|
||||||
|
|
||||||
// LOG_DEBUG("%s", name.c_str());
|
// LOG_DEBUG("%s", name.c_str());
|
||||||
|
|
||||||
@ -812,7 +812,7 @@ struct PickleTensorReader {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void read_string(const std::string& str, struct zip_t* zip, std::string dir) {
|
void read_string(const std::string& str, zip_t* zip, std::string dir) {
|
||||||
if (str == "storage") {
|
if (str == "storage") {
|
||||||
read_global_type = true;
|
read_global_type = true;
|
||||||
} else if (str != "state_dict") {
|
} else if (str != "state_dict") {
|
||||||
@ -995,7 +995,7 @@ bool ModelLoader::init_from_ckpt_file(const std::string& file_path, const std::s
|
|||||||
file_paths_.push_back(file_path);
|
file_paths_.push_back(file_path);
|
||||||
size_t file_index = file_paths_.size() - 1;
|
size_t file_index = file_paths_.size() - 1;
|
||||||
|
|
||||||
struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
||||||
if (zip == nullptr) {
|
if (zip == nullptr) {
|
||||||
LOG_ERROR("failed to open '%s'", file_path.c_str());
|
LOG_ERROR("failed to open '%s'", file_path.c_str());
|
||||||
return false;
|
return false;
|
||||||
@ -1413,7 +1413,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
|
|||||||
for (int i = 0; i < n_threads; ++i) {
|
for (int i = 0; i < n_threads; ++i) {
|
||||||
workers.emplace_back([&, file_path, is_zip]() {
|
workers.emplace_back([&, file_path, is_zip]() {
|
||||||
std::ifstream file;
|
std::ifstream file;
|
||||||
struct zip_t* zip = nullptr;
|
zip_t* zip = nullptr;
|
||||||
if (is_zip) {
|
if (is_zip) {
|
||||||
zip = zip_open(file_path.c_str(), 0, 'r');
|
zip = zip_open(file_path.c_str(), 0, 'r');
|
||||||
if (zip == nullptr) {
|
if (zip == nullptr) {
|
||||||
@ -1601,7 +1601,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
|
|||||||
return success;
|
return success;
|
||||||
}
|
}
|
||||||
|
|
||||||
bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
|
bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
|
||||||
std::set<std::string> ignore_tensors,
|
std::set<std::string> ignore_tensors,
|
||||||
int n_threads,
|
int n_threads,
|
||||||
bool enable_mmap) {
|
bool enable_mmap) {
|
||||||
@ -1615,7 +1615,7 @@ bool ModelLoader::load_tensors(std::map<std::string, struct ggml_tensor*>& tenso
|
|||||||
tensor_names_in_file.insert(name);
|
tensor_names_in_file.insert(name);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* real;
|
ggml_tensor* real;
|
||||||
if (tensors.find(name) != tensors.end()) {
|
if (tensors.find(name) != tensors.end()) {
|
||||||
real = tensors[name];
|
real = tensors[name];
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
@ -323,7 +323,7 @@ public:
|
|||||||
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
|
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
|
||||||
void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
|
void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
|
||||||
bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
|
bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
|
||||||
bool load_tensors(std::map<std::string, struct ggml_tensor*>& tensors,
|
bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
|
||||||
std::set<std::string> ignore_tensors = {},
|
std::set<std::string> ignore_tensors = {},
|
||||||
int n_threads = 0,
|
int n_threads = 0,
|
||||||
bool use_mmap = false);
|
bool use_mmap = false);
|
||||||
|
|||||||
204
src/pmid.hpp
204
src/pmid.hpp
@ -21,14 +21,14 @@ public:
|
|||||||
blocks["layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(in_dim));
|
blocks["layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(in_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
|
|
||||||
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
||||||
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
||||||
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
|
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
|
||||||
|
|
||||||
struct ggml_tensor* r = x;
|
ggml_tensor* r = x;
|
||||||
// x = ggml_ext_layer_norm(ctx, x, ln_w, ln_b);
|
// x = ggml_ext_layer_norm(ctx, x, ln_w, ln_b);
|
||||||
x = layer_norm->forward(ctx, x);
|
x = layer_norm->forward(ctx, x);
|
||||||
// x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x), fc1_b);
|
// x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x), fc1_b);
|
||||||
@ -54,8 +54,8 @@ public:
|
|||||||
blocks["1"] = std::shared_ptr<GGMLBlock>(new Mlp(dim, inner_dim, dim, false));
|
blocks["1"] = std::shared_ptr<GGMLBlock>(new Mlp(dim, inner_dim, dim, false));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["0"]);
|
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["0"]);
|
||||||
auto ff = std::dynamic_pointer_cast<Mlp>(blocks["1"]);
|
auto ff = std::dynamic_pointer_cast<Mlp>(blocks["1"]);
|
||||||
|
|
||||||
@ -81,9 +81,9 @@ public:
|
|||||||
blocks["to_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim, false));
|
blocks["to_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim, false));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* reshape_tensor(struct ggml_context* ctx,
|
ggml_tensor* reshape_tensor(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int heads) {
|
int heads) {
|
||||||
int64_t ne[4];
|
int64_t ne[4];
|
||||||
for (int i = 0; i < 4; ++i)
|
for (int i = 0; i < 4; ++i)
|
||||||
ne[i] = x->ne[i];
|
ne[i] = x->ne[i];
|
||||||
@ -92,17 +92,17 @@ public:
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
std::vector<struct ggml_tensor*> chunk_half(struct ggml_context* ctx,
|
std::vector<ggml_tensor*> chunk_half(ggml_context* ctx,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
auto tlo = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0);
|
auto tlo = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0);
|
||||||
auto tli = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], x->nb[0] * x->ne[0] / 2);
|
auto tli = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], x->nb[0] * x->ne[0] / 2);
|
||||||
return {ggml_cont(ctx, tlo),
|
return {ggml_cont(ctx, tlo),
|
||||||
ggml_cont(ctx, tli)};
|
ggml_cont(ctx, tli)};
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* latents) {
|
ggml_tensor* latents) {
|
||||||
// x (torch.Tensor): image features
|
// x (torch.Tensor): image features
|
||||||
// shape (b, n1, D)
|
// shape (b, n1, D)
|
||||||
// latent (torch.Tensor): latent features
|
// latent (torch.Tensor): latent features
|
||||||
@ -176,9 +176,9 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* latents,
|
ggml_tensor* latents,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
|
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
|
||||||
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
|
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
|
||||||
@ -225,19 +225,19 @@ public:
|
|||||||
4));
|
4));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* last_hidden_state) {
|
ggml_tensor* last_hidden_state) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
auto token_proj = std::dynamic_pointer_cast<Mlp>(blocks["token_proj"]);
|
auto token_proj = std::dynamic_pointer_cast<Mlp>(blocks["token_proj"]);
|
||||||
auto token_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["token_norm"]);
|
auto token_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["token_norm"]);
|
||||||
auto perceiver_resampler = std::dynamic_pointer_cast<FacePerceiverResampler>(blocks["perceiver_resampler"]);
|
auto perceiver_resampler = std::dynamic_pointer_cast<FacePerceiverResampler>(blocks["perceiver_resampler"]);
|
||||||
|
|
||||||
x = token_proj->forward(ctx, x);
|
x = token_proj->forward(ctx, x);
|
||||||
int64_t nel = ggml_nelements(x);
|
int64_t nel = ggml_nelements(x);
|
||||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, cross_attention_dim, num_tokens, nel / (cross_attention_dim * num_tokens));
|
x = ggml_reshape_3d(ctx->ggml_ctx, x, cross_attention_dim, num_tokens, nel / (cross_attention_dim * num_tokens));
|
||||||
x = token_norm->forward(ctx, x);
|
x = token_norm->forward(ctx, x);
|
||||||
struct ggml_tensor* out = perceiver_resampler->forward(ctx, x, last_hidden_state);
|
ggml_tensor* out = perceiver_resampler->forward(ctx, x, last_hidden_state);
|
||||||
if (use_residul)
|
if (use_residul)
|
||||||
out = ggml_add(ctx->ggml_ctx, x, out);
|
out = ggml_add(ctx->ggml_ctx, x, out);
|
||||||
return out;
|
return out;
|
||||||
@ -256,9 +256,9 @@ public:
|
|||||||
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(embed_dim));
|
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(embed_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
|
ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* prompt_embeds,
|
ggml_tensor* prompt_embeds,
|
||||||
struct ggml_tensor* id_embeds) {
|
ggml_tensor* id_embeds) {
|
||||||
auto mlp1 = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp1"]);
|
auto mlp1 = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp1"]);
|
||||||
auto mlp2 = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp2"]);
|
auto mlp2 = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp2"]);
|
||||||
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
|
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
|
||||||
@ -273,24 +273,24 @@ public:
|
|||||||
return stacked_id_embeds;
|
return stacked_id_embeds;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* prompt_embeds,
|
ggml_tensor* prompt_embeds,
|
||||||
struct ggml_tensor* id_embeds,
|
ggml_tensor* id_embeds,
|
||||||
struct ggml_tensor* class_tokens_mask,
|
ggml_tensor* class_tokens_mask,
|
||||||
struct ggml_tensor* class_tokens_mask_pos,
|
ggml_tensor* class_tokens_mask_pos,
|
||||||
struct ggml_tensor* left,
|
ggml_tensor* left,
|
||||||
struct ggml_tensor* right) {
|
ggml_tensor* right) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
|
|
||||||
struct ggml_tensor* valid_id_embeds = id_embeds;
|
ggml_tensor* valid_id_embeds = id_embeds;
|
||||||
// # slice out the image token embeddings
|
// # slice out the image token embeddings
|
||||||
ggml_set_name(class_tokens_mask_pos, "class_tokens_mask_pos");
|
ggml_set_name(class_tokens_mask_pos, "class_tokens_mask_pos");
|
||||||
ggml_set_name(prompt_embeds, "prompt_embeds");
|
ggml_set_name(prompt_embeds, "prompt_embeds");
|
||||||
struct ggml_tensor* image_token_embeds = ggml_get_rows(ctx->ggml_ctx, prompt_embeds, class_tokens_mask_pos);
|
ggml_tensor* image_token_embeds = ggml_get_rows(ctx->ggml_ctx, prompt_embeds, class_tokens_mask_pos);
|
||||||
ggml_set_name(image_token_embeds, "image_token_embeds");
|
ggml_set_name(image_token_embeds, "image_token_embeds");
|
||||||
valid_id_embeds = ggml_reshape_2d(ctx->ggml_ctx, valid_id_embeds, valid_id_embeds->ne[0],
|
valid_id_embeds = ggml_reshape_2d(ctx->ggml_ctx, valid_id_embeds, valid_id_embeds->ne[0],
|
||||||
ggml_nelements(valid_id_embeds) / valid_id_embeds->ne[0]);
|
ggml_nelements(valid_id_embeds) / valid_id_embeds->ne[0]);
|
||||||
struct ggml_tensor* stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds);
|
ggml_tensor* stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds);
|
||||||
|
|
||||||
if (left && right) {
|
if (left && right) {
|
||||||
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, left, stacked_id_embeds, 1);
|
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, left, stacked_id_embeds, 1);
|
||||||
@ -301,10 +301,10 @@ public:
|
|||||||
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, stacked_id_embeds, right, 1);
|
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, stacked_id_embeds, right, 1);
|
||||||
}
|
}
|
||||||
|
|
||||||
class_tokens_mask = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, class_tokens_mask));
|
class_tokens_mask = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, class_tokens_mask));
|
||||||
class_tokens_mask = ggml_repeat(ctx->ggml_ctx, class_tokens_mask, prompt_embeds);
|
class_tokens_mask = ggml_repeat(ctx->ggml_ctx, class_tokens_mask, prompt_embeds);
|
||||||
prompt_embeds = ggml_mul(ctx->ggml_ctx, prompt_embeds, class_tokens_mask);
|
prompt_embeds = ggml_mul(ctx->ggml_ctx, prompt_embeds, class_tokens_mask);
|
||||||
struct ggml_tensor* updated_prompt_embeds = ggml_add(ctx->ggml_ctx, prompt_embeds, stacked_id_embeds);
|
ggml_tensor* updated_prompt_embeds = ggml_add(ctx->ggml_ctx, prompt_embeds, stacked_id_embeds);
|
||||||
ggml_set_name(updated_prompt_embeds, "updated_prompt_embeds");
|
ggml_set_name(updated_prompt_embeds, "updated_prompt_embeds");
|
||||||
return updated_prompt_embeds;
|
return updated_prompt_embeds;
|
||||||
}
|
}
|
||||||
@ -317,22 +317,22 @@ struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
|
|||||||
blocks["fuse_module"] = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
|
blocks["fuse_module"] = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* id_pixel_values,
|
ggml_tensor* id_pixel_values,
|
||||||
struct ggml_tensor* prompt_embeds,
|
ggml_tensor* prompt_embeds,
|
||||||
struct ggml_tensor* class_tokens_mask,
|
ggml_tensor* class_tokens_mask,
|
||||||
struct ggml_tensor* class_tokens_mask_pos,
|
ggml_tensor* class_tokens_mask_pos,
|
||||||
struct ggml_tensor* left,
|
ggml_tensor* left,
|
||||||
struct ggml_tensor* right) {
|
ggml_tensor* right) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
|
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
|
||||||
auto visual_projection = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);
|
auto visual_projection = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);
|
||||||
auto visual_projection_2 = std::dynamic_pointer_cast<Linear>(blocks["visual_projection_2"]);
|
auto visual_projection_2 = std::dynamic_pointer_cast<Linear>(blocks["visual_projection_2"]);
|
||||||
auto fuse_module = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
|
auto fuse_module = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
|
||||||
|
|
||||||
struct ggml_tensor* shared_id_embeds = vision_model->forward(ctx, id_pixel_values); // [N, hidden_size]
|
ggml_tensor* shared_id_embeds = vision_model->forward(ctx, id_pixel_values); // [N, hidden_size]
|
||||||
struct ggml_tensor* id_embeds = visual_projection->forward(ctx, shared_id_embeds); // [N, proj_dim(768)]
|
ggml_tensor* id_embeds = visual_projection->forward(ctx, shared_id_embeds); // [N, proj_dim(768)]
|
||||||
struct ggml_tensor* id_embeds_2 = visual_projection_2->forward(ctx, shared_id_embeds); // [N, 1280]
|
ggml_tensor* id_embeds_2 = visual_projection_2->forward(ctx, shared_id_embeds); // [N, 1280]
|
||||||
|
|
||||||
id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 2, 0, 1, 3));
|
id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 2, 0, 1, 3));
|
||||||
id_embeds_2 = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds_2, 2, 0, 1, 3));
|
id_embeds_2 = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds_2, 2, 0, 1, 3));
|
||||||
@ -340,12 +340,12 @@ struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
|
|||||||
id_embeds = ggml_concat(ctx->ggml_ctx, id_embeds, id_embeds_2, 2); // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
|
id_embeds = ggml_concat(ctx->ggml_ctx, id_embeds, id_embeds_2, 2); // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
|
||||||
id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 1, 2, 0, 3));
|
id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 1, 2, 0, 3));
|
||||||
|
|
||||||
struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
|
ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
|
||||||
prompt_embeds,
|
prompt_embeds,
|
||||||
id_embeds,
|
id_embeds,
|
||||||
class_tokens_mask,
|
class_tokens_mask,
|
||||||
class_tokens_mask_pos,
|
class_tokens_mask_pos,
|
||||||
left, right);
|
left, right);
|
||||||
return updated_prompt_embeds;
|
return updated_prompt_embeds;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -365,29 +365,29 @@ struct PhotoMakerIDEncoder_CLIPInsightfaceExtendtokenBlock : public CLIPVisionMo
|
|||||||
num_tokens));
|
num_tokens));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* id_pixel_values,
|
ggml_tensor* id_pixel_values,
|
||||||
struct ggml_tensor* prompt_embeds,
|
ggml_tensor* prompt_embeds,
|
||||||
struct ggml_tensor* class_tokens_mask,
|
ggml_tensor* class_tokens_mask,
|
||||||
struct ggml_tensor* class_tokens_mask_pos,
|
ggml_tensor* class_tokens_mask_pos,
|
||||||
struct ggml_tensor* id_embeds,
|
ggml_tensor* id_embeds,
|
||||||
struct ggml_tensor* left,
|
ggml_tensor* left,
|
||||||
struct ggml_tensor* right) {
|
ggml_tensor* right) {
|
||||||
// x: [N, channels, h, w]
|
// x: [N, channels, h, w]
|
||||||
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
|
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
|
||||||
auto fuse_module = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
|
auto fuse_module = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
|
||||||
auto qformer_perceiver = std::dynamic_pointer_cast<QFormerPerceiver>(blocks["qformer_perceiver"]);
|
auto qformer_perceiver = std::dynamic_pointer_cast<QFormerPerceiver>(blocks["qformer_perceiver"]);
|
||||||
|
|
||||||
// struct ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values); // [N, hidden_size]
|
// ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values); // [N, hidden_size]
|
||||||
struct ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values, false); // [N, hidden_size]
|
ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values, false); // [N, hidden_size]
|
||||||
id_embeds = qformer_perceiver->forward(ctx, id_embeds, last_hidden_state);
|
id_embeds = qformer_perceiver->forward(ctx, id_embeds, last_hidden_state);
|
||||||
|
|
||||||
struct ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
|
ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
|
||||||
prompt_embeds,
|
prompt_embeds,
|
||||||
id_embeds,
|
id_embeds,
|
||||||
class_tokens_mask,
|
class_tokens_mask,
|
||||||
class_tokens_mask_pos,
|
class_tokens_mask_pos,
|
||||||
left, right);
|
left, right);
|
||||||
return updated_prompt_embeds;
|
return updated_prompt_embeds;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -436,18 +436,18 @@ public:
|
|||||||
return pm_version;
|
return pm_version;
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
if (pm_version == PM_VERSION_1)
|
if (pm_version == PM_VERSION_1)
|
||||||
id_encoder.get_param_tensors(tensors, prefix);
|
id_encoder.get_param_tensors(tensors, prefix);
|
||||||
else if (pm_version == PM_VERSION_2)
|
else if (pm_version == PM_VERSION_2)
|
||||||
id_encoder2.get_param_tensors(tensors, prefix);
|
id_encoder2.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph( // struct ggml_allocr* allocr,
|
ggml_cgraph* build_graph( // ggml_allocr* allocr,
|
||||||
struct ggml_tensor* id_pixel_values,
|
ggml_tensor* id_pixel_values,
|
||||||
struct ggml_tensor* prompt_embeds,
|
ggml_tensor* prompt_embeds,
|
||||||
std::vector<bool>& class_tokens_mask,
|
std::vector<bool>& class_tokens_mask,
|
||||||
struct ggml_tensor* id_embeds) {
|
ggml_tensor* id_embeds) {
|
||||||
ctm.clear();
|
ctm.clear();
|
||||||
ctmf16.clear();
|
ctmf16.clear();
|
||||||
ctmpos.clear();
|
ctmpos.clear();
|
||||||
@ -458,20 +458,20 @@ public:
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
|
|
||||||
int64_t hidden_size = prompt_embeds->ne[0];
|
int64_t hidden_size = prompt_embeds->ne[0];
|
||||||
int64_t seq_length = prompt_embeds->ne[1];
|
int64_t seq_length = prompt_embeds->ne[1];
|
||||||
ggml_type type = GGML_TYPE_F32;
|
ggml_type type = GGML_TYPE_F32;
|
||||||
|
|
||||||
struct ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
|
ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
|
||||||
|
|
||||||
struct ggml_tensor* id_pixel_values_d = to_backend(id_pixel_values);
|
ggml_tensor* id_pixel_values_d = to_backend(id_pixel_values);
|
||||||
struct ggml_tensor* prompt_embeds_d = to_backend(prompt_embeds);
|
ggml_tensor* prompt_embeds_d = to_backend(prompt_embeds);
|
||||||
struct ggml_tensor* id_embeds_d = to_backend(id_embeds);
|
ggml_tensor* id_embeds_d = to_backend(id_embeds);
|
||||||
|
|
||||||
struct ggml_tensor* left = nullptr;
|
ggml_tensor* left = nullptr;
|
||||||
struct ggml_tensor* right = nullptr;
|
ggml_tensor* right = nullptr;
|
||||||
for (int i = 0; i < class_tokens_mask.size(); i++) {
|
for (int i = 0; i < class_tokens_mask.size(); i++) {
|
||||||
if (class_tokens_mask[i]) {
|
if (class_tokens_mask[i]) {
|
||||||
// printf(" 1,");
|
// printf(" 1,");
|
||||||
@ -495,7 +495,7 @@ public:
|
|||||||
right = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type,
|
right = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type,
|
||||||
hidden_size, seq_length - ctmpos[ctmpos.size() - 1] - 1, 1);
|
hidden_size, seq_length - ctmpos[ctmpos.size() - 1] - 1, 1);
|
||||||
}
|
}
|
||||||
struct ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
|
ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
|
||||||
|
|
||||||
{
|
{
|
||||||
if (type == GGML_TYPE_F16)
|
if (type == GGML_TYPE_F16)
|
||||||
@ -526,7 +526,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
struct ggml_tensor* updated_prompt_embeds = nullptr;
|
ggml_tensor* updated_prompt_embeds = nullptr;
|
||||||
if (pm_version == PM_VERSION_1)
|
if (pm_version == PM_VERSION_1)
|
||||||
updated_prompt_embeds = id_encoder.forward(&runner_ctx,
|
updated_prompt_embeds = id_encoder.forward(&runner_ctx,
|
||||||
id_pixel_values_d,
|
id_pixel_values_d,
|
||||||
@ -549,13 +549,13 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(const int n_threads,
|
bool compute(const int n_threads,
|
||||||
struct ggml_tensor* id_pixel_values,
|
ggml_tensor* id_pixel_values,
|
||||||
struct ggml_tensor* prompt_embeds,
|
ggml_tensor* prompt_embeds,
|
||||||
struct ggml_tensor* id_embeds,
|
ggml_tensor* id_embeds,
|
||||||
std::vector<bool>& class_tokens_mask,
|
std::vector<bool>& class_tokens_mask,
|
||||||
struct ggml_tensor** updated_prompt_embeds,
|
ggml_tensor** updated_prompt_embeds,
|
||||||
ggml_context* output_ctx) {
|
ggml_context* output_ctx) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
// return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask);
|
// return build_graph(compute_allocr, id_pixel_values, prompt_embeds, class_tokens_mask);
|
||||||
return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
|
return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
|
||||||
};
|
};
|
||||||
@ -566,7 +566,7 @@ public:
|
|||||||
};
|
};
|
||||||
|
|
||||||
struct PhotoMakerIDEmbed : public GGMLRunner {
|
struct PhotoMakerIDEmbed : public GGMLRunner {
|
||||||
std::map<std::string, struct ggml_tensor*> tensors;
|
std::map<std::string, ggml_tensor*> tensors;
|
||||||
std::string file_path;
|
std::string file_path;
|
||||||
ModelLoader* model_loader;
|
ModelLoader* model_loader;
|
||||||
bool load_failed = false;
|
bool load_failed = false;
|
||||||
@ -606,11 +606,11 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
if (dry_run) {
|
if (dry_run) {
|
||||||
std::lock_guard<std::mutex> lock(tensor_mutex);
|
std::lock_guard<std::mutex> lock(tensor_mutex);
|
||||||
struct ggml_tensor* real = ggml_new_tensor(params_ctx,
|
ggml_tensor* real = ggml_new_tensor(params_ctx,
|
||||||
tensor_storage.type,
|
tensor_storage.type,
|
||||||
tensor_storage.n_dims,
|
tensor_storage.n_dims,
|
||||||
tensor_storage.ne);
|
tensor_storage.ne);
|
||||||
tensors[name] = real;
|
tensors[name] = real;
|
||||||
} else {
|
} else {
|
||||||
auto real = tensors[name];
|
auto real = tensors[name];
|
||||||
*dst_tensor = real;
|
*dst_tensor = real;
|
||||||
@ -629,8 +629,8 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
|
|||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* get() {
|
ggml_tensor* get() {
|
||||||
std::map<std::string, struct ggml_tensor*>::iterator pos;
|
std::map<std::string, ggml_tensor*>::iterator pos;
|
||||||
pos = tensors.find("pmid.id_embeds");
|
pos = tensors.find("pmid.id_embeds");
|
||||||
if (pos != tensors.end())
|
if (pos != tensors.end())
|
||||||
return pos->second;
|
return pos->second;
|
||||||
|
|||||||
@ -4,13 +4,13 @@
|
|||||||
#include "ggml_extend.hpp"
|
#include "ggml_extend.hpp"
|
||||||
#define M_PI_ 3.14159265358979323846f
|
#define M_PI_ 3.14159265358979323846f
|
||||||
|
|
||||||
void convolve(struct ggml_tensor* input, struct ggml_tensor* output, struct ggml_tensor* kernel, int padding) {
|
void convolve(ggml_tensor* input, ggml_tensor* output, ggml_tensor* kernel, int padding) {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = 80 * input->ne[0] * input->ne[1]; // 20M for 512x512
|
params.mem_size = 80 * input->ne[0] * input->ne[1]; // 20M for 512x512
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
struct ggml_context* ctx0 = ggml_init(params);
|
ggml_context* ctx0 = ggml_init(params);
|
||||||
struct ggml_tensor* kernel_fp16 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, kernel->ne[0], kernel->ne[1], 1, 1);
|
ggml_tensor* kernel_fp16 = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, kernel->ne[0], kernel->ne[1], 1, 1);
|
||||||
ggml_fp32_to_fp16_row((float*)kernel->data, (ggml_fp16_t*)kernel_fp16->data, ggml_nelements(kernel));
|
ggml_fp32_to_fp16_row((float*)kernel->data, (ggml_fp16_t*)kernel_fp16->data, ggml_nelements(kernel));
|
||||||
ggml_tensor* h = ggml_conv_2d(ctx0, kernel_fp16, input, 1, 1, padding, padding, 1, 1);
|
ggml_tensor* h = ggml_conv_2d(ctx0, kernel_fp16, input, 1, 1, padding, padding, 1, 1);
|
||||||
ggml_cgraph* gf = ggml_new_graph(ctx0);
|
ggml_cgraph* gf = ggml_new_graph(ctx0);
|
||||||
@ -19,7 +19,7 @@ void convolve(struct ggml_tensor* input, struct ggml_tensor* output, struct ggml
|
|||||||
ggml_free(ctx0);
|
ggml_free(ctx0);
|
||||||
}
|
}
|
||||||
|
|
||||||
void gaussian_kernel(struct ggml_tensor* kernel) {
|
void gaussian_kernel(ggml_tensor* kernel) {
|
||||||
int ks_mid = static_cast<int>(kernel->ne[0] / 2);
|
int ks_mid = static_cast<int>(kernel->ne[0] / 2);
|
||||||
float sigma = 1.4f;
|
float sigma = 1.4f;
|
||||||
float normal = 1.f / (2.0f * M_PI_ * powf(sigma, 2.0f));
|
float normal = 1.f / (2.0f * M_PI_ * powf(sigma, 2.0f));
|
||||||
@ -33,7 +33,7 @@ void gaussian_kernel(struct ggml_tensor* kernel) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
|
void grayscale(ggml_tensor* rgb_img, ggml_tensor* grayscale) {
|
||||||
for (int iy = 0; iy < rgb_img->ne[1]; iy++) {
|
for (int iy = 0; iy < rgb_img->ne[1]; iy++) {
|
||||||
for (int ix = 0; ix < rgb_img->ne[0]; ix++) {
|
for (int ix = 0; ix < rgb_img->ne[0]; ix++) {
|
||||||
float r = ggml_ext_tensor_get_f32(rgb_img, ix, iy);
|
float r = ggml_ext_tensor_get_f32(rgb_img, ix, iy);
|
||||||
@ -45,7 +45,7 @@ void grayscale(struct ggml_tensor* rgb_img, struct ggml_tensor* grayscale) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
|
void prop_hypot(ggml_tensor* x, ggml_tensor* y, ggml_tensor* h) {
|
||||||
int n_elements = static_cast<int>(ggml_nelements(h));
|
int n_elements = static_cast<int>(ggml_nelements(h));
|
||||||
float* dx = (float*)x->data;
|
float* dx = (float*)x->data;
|
||||||
float* dy = (float*)y->data;
|
float* dy = (float*)y->data;
|
||||||
@ -55,7 +55,7 @@ void prop_hypot(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tensor* h) {
|
void prop_arctan2(ggml_tensor* x, ggml_tensor* y, ggml_tensor* h) {
|
||||||
int n_elements = static_cast<int>(ggml_nelements(h));
|
int n_elements = static_cast<int>(ggml_nelements(h));
|
||||||
float* dx = (float*)x->data;
|
float* dx = (float*)x->data;
|
||||||
float* dy = (float*)y->data;
|
float* dy = (float*)y->data;
|
||||||
@ -65,7 +65,7 @@ void prop_arctan2(struct ggml_tensor* x, struct ggml_tensor* y, struct ggml_tens
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void normalize_tensor(struct ggml_tensor* g) {
|
void normalize_tensor(ggml_tensor* g) {
|
||||||
int n_elements = static_cast<int>(ggml_nelements(g));
|
int n_elements = static_cast<int>(ggml_nelements(g));
|
||||||
float* dg = (float*)g->data;
|
float* dg = (float*)g->data;
|
||||||
float max = -INFINITY;
|
float max = -INFINITY;
|
||||||
@ -78,7 +78,7 @@ void normalize_tensor(struct ggml_tensor* g) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struct ggml_tensor* D) {
|
void non_max_supression(ggml_tensor* result, ggml_tensor* G, ggml_tensor* D) {
|
||||||
for (int iy = 1; iy < result->ne[1] - 1; iy++) {
|
for (int iy = 1; iy < result->ne[1] - 1; iy++) {
|
||||||
for (int ix = 1; ix < result->ne[0] - 1; ix++) {
|
for (int ix = 1; ix < result->ne[0] - 1; ix++) {
|
||||||
float angle = ggml_ext_tensor_get_f32(D, ix, iy) * 180.0f / M_PI_;
|
float angle = ggml_ext_tensor_get_f32(D, ix, iy) * 180.0f / M_PI_;
|
||||||
@ -117,7 +117,7 @@ void non_max_supression(struct ggml_tensor* result, struct ggml_tensor* G, struc
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float low_threshold, float weak, float strong) {
|
void threshold_hystersis(ggml_tensor* img, float high_threshold, float low_threshold, float weak, float strong) {
|
||||||
int n_elements = static_cast<int>(ggml_nelements(img));
|
int n_elements = static_cast<int>(ggml_nelements(img));
|
||||||
float* imd = (float*)img->data;
|
float* imd = (float*)img->data;
|
||||||
float max = -INFINITY;
|
float max = -INFINITY;
|
||||||
@ -163,11 +163,11 @@ void threshold_hystersis(struct ggml_tensor* img, float high_threshold, float lo
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
|
bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, float weak, float strong, bool inverse) {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(40 * img.width * img.height); // 10MB for 512x512
|
params.mem_size = static_cast<size_t>(40 * img.width * img.height); // 10MB for 512x512
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
|
|
||||||
if (!work_ctx) {
|
if (!work_ctx) {
|
||||||
LOG_ERROR("ggml_init() failed");
|
LOG_ERROR("ggml_init() failed");
|
||||||
@ -185,19 +185,19 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
|
|||||||
-1, -2, -1};
|
-1, -2, -1};
|
||||||
|
|
||||||
// generate kernel
|
// generate kernel
|
||||||
int kernel_size = 5;
|
int kernel_size = 5;
|
||||||
struct ggml_tensor* gkernel = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, 1);
|
ggml_tensor* gkernel = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, 1);
|
||||||
struct ggml_tensor* sf_kx = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
|
ggml_tensor* sf_kx = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
|
||||||
memcpy(sf_kx->data, kX, ggml_nbytes(sf_kx));
|
memcpy(sf_kx->data, kX, ggml_nbytes(sf_kx));
|
||||||
struct ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
|
ggml_tensor* sf_ky = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, 3, 3, 1, 1);
|
||||||
memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
|
memcpy(sf_ky->data, kY, ggml_nbytes(sf_ky));
|
||||||
gaussian_kernel(gkernel);
|
gaussian_kernel(gkernel);
|
||||||
struct ggml_tensor* image = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
|
ggml_tensor* image = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 3, 1);
|
||||||
struct ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
|
ggml_tensor* image_gray = ggml_new_tensor_4d(work_ctx, GGML_TYPE_F32, img.width, img.height, 1, 1);
|
||||||
struct ggml_tensor* iX = ggml_dup_tensor(work_ctx, image_gray);
|
ggml_tensor* iX = ggml_dup_tensor(work_ctx, image_gray);
|
||||||
struct ggml_tensor* iY = ggml_dup_tensor(work_ctx, image_gray);
|
ggml_tensor* iY = ggml_dup_tensor(work_ctx, image_gray);
|
||||||
struct ggml_tensor* G = ggml_dup_tensor(work_ctx, image_gray);
|
ggml_tensor* G = ggml_dup_tensor(work_ctx, image_gray);
|
||||||
struct ggml_tensor* tetha = ggml_dup_tensor(work_ctx, image_gray);
|
ggml_tensor* tetha = ggml_dup_tensor(work_ctx, image_gray);
|
||||||
sd_image_to_ggml_tensor(img, image);
|
sd_image_to_ggml_tensor(img, image);
|
||||||
grayscale(image, image_gray);
|
grayscale(image, image_gray);
|
||||||
convolve(image_gray, image_gray, gkernel, 2);
|
convolve(image_gray, image_gray, gkernel, 2);
|
||||||
|
|||||||
@ -26,9 +26,9 @@ namespace Qwen {
|
|||||||
blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, out_dim, sample_proj_bias));
|
blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, out_dim, sample_proj_bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* sample,
|
ggml_tensor* sample,
|
||||||
struct ggml_tensor* condition = nullptr) {
|
ggml_tensor* condition = nullptr) {
|
||||||
if (condition != nullptr) {
|
if (condition != nullptr) {
|
||||||
auto cond_proj = std::dynamic_pointer_cast<Linear>(blocks["cond_proj"]);
|
auto cond_proj = std::dynamic_pointer_cast<Linear>(blocks["cond_proj"]);
|
||||||
sample = ggml_add(ctx->ggml_ctx, sample, cond_proj->forward(ctx, condition));
|
sample = ggml_add(ctx->ggml_ctx, sample, cond_proj->forward(ctx, condition));
|
||||||
@ -49,8 +49,8 @@ namespace Qwen {
|
|||||||
blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
|
blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* timesteps) {
|
ggml_tensor* timesteps) {
|
||||||
// timesteps: [N,]
|
// timesteps: [N,]
|
||||||
// return: [N, embedding_dim]
|
// return: [N, embedding_dim]
|
||||||
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
|
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
|
||||||
@ -107,10 +107,10 @@ namespace Qwen {
|
|||||||
}
|
}
|
||||||
|
|
||||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* img,
|
ggml_tensor* img,
|
||||||
struct ggml_tensor* txt,
|
ggml_tensor* txt,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr) {
|
ggml_tensor* mask = nullptr) {
|
||||||
// img: [N, n_img_token, hidden_size]
|
// img: [N, n_img_token, hidden_size]
|
||||||
// txt: [N, n_txt_token, hidden_size]
|
// txt: [N, n_txt_token, hidden_size]
|
||||||
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
|
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
|
||||||
@ -249,11 +249,11 @@ namespace Qwen {
|
|||||||
}
|
}
|
||||||
|
|
||||||
virtual std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
virtual std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* img,
|
ggml_tensor* img,
|
||||||
struct ggml_tensor* txt,
|
ggml_tensor* txt,
|
||||||
struct ggml_tensor* t_emb,
|
ggml_tensor* t_emb,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* modulate_index = nullptr) {
|
ggml_tensor* modulate_index = nullptr) {
|
||||||
// img: [N, n_img_token, hidden_size]
|
// img: [N, n_img_token, hidden_size]
|
||||||
// txt: [N, n_txt_token, hidden_size]
|
// txt: [N, n_txt_token, hidden_size]
|
||||||
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
|
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
|
||||||
@ -325,9 +325,9 @@ namespace Qwen {
|
|||||||
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(conditioning_embedding_dim, embedding_dim * 2, bias));
|
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(conditioning_embedding_dim, embedding_dim * 2, bias));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
// return: [N, n_token, patch_size * patch_size * out_channels]
|
// return: [N, n_token, patch_size * patch_size * out_channels]
|
||||||
@ -389,12 +389,12 @@ namespace Qwen {
|
|||||||
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
|
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* modulate_index = nullptr) {
|
ggml_tensor* modulate_index = nullptr) {
|
||||||
auto time_text_embed = std::dynamic_pointer_cast<QwenTimestepProjEmbeddings>(blocks["time_text_embed"]);
|
auto time_text_embed = std::dynamic_pointer_cast<QwenTimestepProjEmbeddings>(blocks["time_text_embed"]);
|
||||||
auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
|
auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
|
||||||
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
|
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
|
||||||
@ -429,13 +429,13 @@ namespace Qwen {
|
|||||||
return img;
|
return img;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
struct ggml_tensor* modulate_index = nullptr) {
|
ggml_tensor* modulate_index = nullptr) {
|
||||||
// Forward pass of DiT.
|
// Forward pass of DiT.
|
||||||
// x: [N, C, H, W]
|
// x: [N, C, H, W]
|
||||||
// timestep: [N,]
|
// timestep: [N,]
|
||||||
@ -521,17 +521,17 @@ namespace Qwen {
|
|||||||
return "qwen_image";
|
return "qwen_image";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
qwen_image.get_param_tensors(tensors, prefix);
|
qwen_image.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
bool increase_ref_index = false) {
|
bool increase_ref_index = false) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
struct ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
context = to_backend(context);
|
context = to_backend(context);
|
||||||
@ -587,13 +587,13 @@ namespace Qwen {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = qwen_image.forward(&runner_ctx,
|
ggml_tensor* out = qwen_image.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
timesteps,
|
timesteps,
|
||||||
context,
|
context,
|
||||||
pe,
|
pe,
|
||||||
ref_latents,
|
ref_latents,
|
||||||
modulate_index);
|
modulate_index);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -601,17 +601,17 @@ namespace Qwen {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
bool increase_ref_index = false,
|
bool increase_ref_index = false,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]
|
// context: [N, max_position, hidden_size]
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
|
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -619,12 +619,12 @@ namespace Qwen {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -641,7 +641,7 @@ namespace Qwen {
|
|||||||
auto context = load_tensor_from_file(work_ctx, "./qwen_image_context.bin");
|
auto context = load_tensor_from_file(work_ctx, "./qwen_image_context.bin");
|
||||||
print_ggml_tensor(context);
|
print_ggml_tensor(context);
|
||||||
|
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
compute(8, x, timesteps, context, {}, false, &out, work_ctx);
|
compute(8, x, timesteps, context, {}, false, &out, work_ctx);
|
||||||
|
|||||||
24
src/rope.hpp
24
src/rope.hpp
@ -600,10 +600,10 @@ namespace Rope {
|
|||||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||||
}
|
}
|
||||||
|
|
||||||
__STATIC_INLINE__ struct ggml_tensor* apply_rope(struct ggml_context* ctx,
|
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
bool rope_interleaved = true) {
|
bool rope_interleaved = true) {
|
||||||
// x: [N, L, n_head, d_head]
|
// x: [N, L, n_head, d_head]
|
||||||
// pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
|
// pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
|
||||||
int64_t d_head = x->ne[0];
|
int64_t d_head = x->ne[0];
|
||||||
@ -641,14 +641,14 @@ namespace Rope {
|
|||||||
return x_out;
|
return x_out;
|
||||||
}
|
}
|
||||||
|
|
||||||
__STATIC_INLINE__ struct ggml_tensor* attention(GGMLRunnerContext* ctx,
|
__STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* q,
|
ggml_tensor* q,
|
||||||
struct ggml_tensor* k,
|
ggml_tensor* k,
|
||||||
struct ggml_tensor* v,
|
ggml_tensor* v,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask,
|
ggml_tensor* mask,
|
||||||
float kv_scale = 1.0f,
|
float kv_scale = 1.0f,
|
||||||
bool rope_interleaved = true) {
|
bool rope_interleaved = true) {
|
||||||
// q,k,v: [N, L, n_head, d_head]
|
// q,k,v: [N, L, n_head, d_head]
|
||||||
// pe: [L, d_head/2, 2, 2]
|
// pe: [L, d_head/2, 2, 2]
|
||||||
// return: [N, L, n_head*d_head]
|
// return: [N, L, n_head*d_head]
|
||||||
|
|||||||
@ -57,7 +57,7 @@ struct SpectrumState {
|
|||||||
return (num_cached + 1) % ws != 0;
|
return (num_cached + 1) % ws != 0;
|
||||||
}
|
}
|
||||||
|
|
||||||
void update(const struct ggml_tensor* denoised) {
|
void update(const ggml_tensor* denoised) {
|
||||||
int64_t ne = ggml_nelements(denoised);
|
int64_t ne = ggml_nelements(denoised);
|
||||||
const float* data = (const float*)denoised->data;
|
const float* data = (const float*)denoised->data;
|
||||||
|
|
||||||
@ -76,7 +76,7 @@ struct SpectrumState {
|
|||||||
cnt++;
|
cnt++;
|
||||||
}
|
}
|
||||||
|
|
||||||
void predict(struct ggml_tensor* denoised) {
|
void predict(ggml_tensor* denoised) {
|
||||||
int64_t F = (int64_t)H_buf[0].size();
|
int64_t F = (int64_t)H_buf[0].size();
|
||||||
int K_curr = (int)H_buf.size();
|
int K_curr = (int)H_buf.size();
|
||||||
int M1 = config.m + 1;
|
int M1 = config.m + 1;
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
114
src/t5.hpp
114
src/t5.hpp
@ -211,9 +211,9 @@ protected:
|
|||||||
// implementation. It's based on the following three ideas:
|
// implementation. It's based on the following three ideas:
|
||||||
//
|
//
|
||||||
// 1. Because it uses the *unigram* model:
|
// 1. Because it uses the *unigram* model:
|
||||||
// best_score(x1, x2, …, xt) = best_score(x1, x2, …, x{t-1}) + score(xt)
|
// best_score(x1, x2, <EFBFBD>? xt) = best_score(x1, x2, <20>? x{t-1}) + score(xt)
|
||||||
// Deciding the best path (and score) can be decoupled into two isolated
|
// Deciding the best path (and score) can be decoupled into two isolated
|
||||||
// terms: (a) the best path ended before the last token `best_score(x1, x2, …,
|
// terms: (a) the best path ended before the last token `best_score(x1, x2, <EFBFBD>?
|
||||||
// x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
|
// x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
|
||||||
// not related to each other at all.
|
// not related to each other at all.
|
||||||
//
|
//
|
||||||
@ -462,7 +462,7 @@ protected:
|
|||||||
int64_t hidden_size;
|
int64_t hidden_size;
|
||||||
float eps;
|
float eps;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type wtype = GGML_TYPE_F32;
|
enum ggml_type wtype = GGML_TYPE_F32;
|
||||||
params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
|
params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
|
||||||
}
|
}
|
||||||
@ -473,10 +473,10 @@ public:
|
|||||||
: hidden_size(hidden_size),
|
: hidden_size(hidden_size),
|
||||||
eps(eps) {}
|
eps(eps) {}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
struct ggml_tensor* w = params["weight"];
|
ggml_tensor* w = params["weight"];
|
||||||
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
|
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
|
||||||
x = ggml_mul(ctx->ggml_ctx, x, w);
|
x = ggml_mul(ctx->ggml_ctx, x, w);
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -488,7 +488,7 @@ public:
|
|||||||
blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false));
|
blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
auto wi = std::dynamic_pointer_cast<Linear>(blocks["wi"]);
|
auto wi = std::dynamic_pointer_cast<Linear>(blocks["wi"]);
|
||||||
auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
|
auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
|
||||||
@ -510,7 +510,7 @@ public:
|
|||||||
blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false, false, false, scale));
|
blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false, false, false, scale));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
auto wi_0 = std::dynamic_pointer_cast<Linear>(blocks["wi_0"]);
|
auto wi_0 = std::dynamic_pointer_cast<Linear>(blocks["wi_0"]);
|
||||||
auto wi_1 = std::dynamic_pointer_cast<Linear>(blocks["wi_1"]);
|
auto wi_1 = std::dynamic_pointer_cast<Linear>(blocks["wi_1"]);
|
||||||
@ -531,7 +531,7 @@ public:
|
|||||||
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
|
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
auto DenseReluDense = std::dynamic_pointer_cast<T5DenseGatedActDense>(blocks["DenseReluDense"]);
|
auto DenseReluDense = std::dynamic_pointer_cast<T5DenseGatedActDense>(blocks["DenseReluDense"]);
|
||||||
auto layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
|
auto layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
|
||||||
@ -570,8 +570,8 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
|
ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* relative_position_bucket) {
|
ggml_tensor* relative_position_bucket) {
|
||||||
auto relative_attention_bias = std::dynamic_pointer_cast<Embedding>(blocks["relative_attention_bias"]);
|
auto relative_attention_bias = std::dynamic_pointer_cast<Embedding>(blocks["relative_attention_bias"]);
|
||||||
|
|
||||||
auto values = relative_attention_bias->forward(ctx, relative_position_bucket); // shape (query_length, key_length, num_heads)
|
auto values = relative_attention_bias->forward(ctx, relative_position_bucket); // shape (query_length, key_length, num_heads)
|
||||||
@ -580,11 +580,11 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* past_bias = nullptr,
|
ggml_tensor* past_bias = nullptr,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
struct ggml_tensor* relative_position_bucket = nullptr) {
|
ggml_tensor* relative_position_bucket = nullptr) {
|
||||||
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q"]);
|
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q"]);
|
||||||
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k"]);
|
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k"]);
|
||||||
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v"]);
|
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v"]);
|
||||||
@ -629,11 +629,11 @@ public:
|
|||||||
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
|
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* past_bias = nullptr,
|
ggml_tensor* past_bias = nullptr,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
struct ggml_tensor* relative_position_bucket = nullptr) {
|
ggml_tensor* relative_position_bucket = nullptr) {
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
auto SelfAttention = std::dynamic_pointer_cast<T5Attention>(blocks["SelfAttention"]);
|
auto SelfAttention = std::dynamic_pointer_cast<T5Attention>(blocks["SelfAttention"]);
|
||||||
auto layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
|
auto layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
|
||||||
@ -655,11 +655,11 @@ public:
|
|||||||
blocks["layer.1"] = std::shared_ptr<GGMLBlock>(new T5LayerFF(model_dim, ff_dim));
|
blocks["layer.1"] = std::shared_ptr<GGMLBlock>(new T5LayerFF(model_dim, ff_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
std::pair<struct ggml_tensor*, struct ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* past_bias = nullptr,
|
ggml_tensor* past_bias = nullptr,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
struct ggml_tensor* relative_position_bucket = nullptr) {
|
ggml_tensor* relative_position_bucket = nullptr) {
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
auto layer_0 = std::dynamic_pointer_cast<T5LayerSelfAttention>(blocks["layer.0"]);
|
auto layer_0 = std::dynamic_pointer_cast<T5LayerSelfAttention>(blocks["layer.0"]);
|
||||||
auto layer_1 = std::dynamic_pointer_cast<T5LayerFF>(blocks["layer.1"]);
|
auto layer_1 = std::dynamic_pointer_cast<T5LayerFF>(blocks["layer.1"]);
|
||||||
@ -690,11 +690,11 @@ public:
|
|||||||
blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
|
blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* past_bias = nullptr,
|
ggml_tensor* past_bias = nullptr,
|
||||||
struct ggml_tensor* attention_mask = nullptr,
|
ggml_tensor* attention_mask = nullptr,
|
||||||
struct ggml_tensor* relative_position_bucket = nullptr) {
|
ggml_tensor* relative_position_bucket = nullptr) {
|
||||||
// x: [N, n_token, model_dim]
|
// x: [N, n_token, model_dim]
|
||||||
for (int i = 0; i < num_layers; i++) {
|
for (int i = 0; i < num_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
|
auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
|
||||||
@ -737,11 +737,11 @@ public:
|
|||||||
params.model_dim));
|
params.model_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* past_bias = nullptr,
|
ggml_tensor* past_bias = nullptr,
|
||||||
struct ggml_tensor* attention_mask = nullptr,
|
ggml_tensor* attention_mask = nullptr,
|
||||||
struct ggml_tensor* relative_position_bucket = nullptr) {
|
ggml_tensor* relative_position_bucket = nullptr) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
|
|
||||||
auto shared = std::dynamic_pointer_cast<Embedding>(blocks["shared"]);
|
auto shared = std::dynamic_pointer_cast<Embedding>(blocks["shared"]);
|
||||||
@ -776,14 +776,14 @@ struct T5Runner : public GGMLRunner {
|
|||||||
return "t5";
|
return "t5";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
model.get_param_tensors(tensors, prefix);
|
model.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* relative_position_bucket,
|
ggml_tensor* relative_position_bucket,
|
||||||
struct ggml_tensor* attention_mask = nullptr) {
|
ggml_tensor* attention_mask = nullptr) {
|
||||||
size_t N = input_ids->ne[1];
|
size_t N = input_ids->ne[1];
|
||||||
size_t n_token = input_ids->ne[0];
|
size_t n_token = input_ids->ne[0];
|
||||||
|
|
||||||
@ -791,9 +791,9 @@ struct T5Runner : public GGMLRunner {
|
|||||||
return hidden_states;
|
return hidden_states;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* input_ids,
|
ggml_cgraph* build_graph(ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* attention_mask = nullptr) {
|
ggml_tensor* attention_mask = nullptr) {
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
|
|
||||||
input_ids = to_backend(input_ids);
|
input_ids = to_backend(input_ids);
|
||||||
attention_mask = to_backend(attention_mask);
|
attention_mask = to_backend(attention_mask);
|
||||||
@ -813,8 +813,8 @@ struct T5Runner : public GGMLRunner {
|
|||||||
input_ids->ne[0]);
|
input_ids->ne[0]);
|
||||||
set_backend_tensor_data(relative_position_bucket, relative_position_bucket_vec.data());
|
set_backend_tensor_data(relative_position_bucket, relative_position_bucket_vec.data());
|
||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
struct ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
|
ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, hidden_states);
|
ggml_build_forward_expand(gf, hidden_states);
|
||||||
|
|
||||||
@ -822,11 +822,11 @@ struct T5Runner : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(const int n_threads,
|
bool compute(const int n_threads,
|
||||||
struct ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
struct ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(input_ids, attention_mask);
|
return build_graph(input_ids, attention_mask);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
||||||
@ -912,7 +912,7 @@ struct T5Embedder {
|
|||||||
: model(backend, offload_params_to_cpu, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
|
: model(backend, offload_params_to_cpu, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
model.get_param_tensors(tensors, prefix);
|
model.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -962,17 +962,17 @@ struct T5Embedder {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
std::string text("a lovely cat");
|
std::string text("a lovely cat");
|
||||||
// std::string text("一只可爱的猫"); // umt5 chinease test
|
// std::string text("一只可爱的<EFBFBD>?); // umt5 chinease test
|
||||||
auto tokens_and_weights = tokenize(text, 512, true);
|
auto tokens_and_weights = tokenize(text, 512, true);
|
||||||
std::vector<int>& tokens = std::get<0>(tokens_and_weights);
|
std::vector<int>& tokens = std::get<0>(tokens_and_weights);
|
||||||
std::vector<float>& weights = std::get<1>(tokens_and_weights);
|
std::vector<float>& weights = std::get<1>(tokens_and_weights);
|
||||||
@ -981,9 +981,9 @@ struct T5Embedder {
|
|||||||
printf("%d ", token);
|
printf("%d ", token);
|
||||||
}
|
}
|
||||||
printf("\n");
|
printf("\n");
|
||||||
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
auto input_ids = vector_to_ggml_tensor_i32(work_ctx, tokens);
|
||||||
auto attention_mask = vector_to_ggml_tensor(work_ctx, masks);
|
auto attention_mask = vector_to_ggml_tensor(work_ctx, masks);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
model.compute(8, input_ids, attention_mask, &out, work_ctx);
|
model.compute(8, input_ids, attention_mask, &out, work_ctx);
|
||||||
|
|||||||
80
src/tae.hpp
80
src/tae.hpp
@ -37,7 +37,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [n, n_in, h, w]
|
// x: [n, n_in, h, w]
|
||||||
// return: [n, n_out, h, w]
|
// return: [n, n_out, h, w]
|
||||||
|
|
||||||
@ -107,7 +107,7 @@ public:
|
|||||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
|
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [n, in_channels, h, w]
|
// x: [n, in_channels, h, w]
|
||||||
// return: [n, z_channels, h/8, w/8]
|
// return: [n, z_channels, h/8, w/8]
|
||||||
|
|
||||||
@ -157,7 +157,7 @@ public:
|
|||||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
|
||||||
// z: [n, z_channels, h, w]
|
// z: [n, z_channels, h, w]
|
||||||
// return: [n, out_channels, h*8, w*8]
|
// return: [n, out_channels, h*8, w*8]
|
||||||
|
|
||||||
@ -192,7 +192,7 @@ public:
|
|||||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * stride, channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
|
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * stride, channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
|
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
|
||||||
auto h = x;
|
auto h = x;
|
||||||
if (stride != 1) {
|
if (stride != 1) {
|
||||||
@ -212,7 +212,7 @@ public:
|
|||||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels * stride, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
|
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels * stride, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
|
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
|
||||||
auto h = conv->forward(ctx, x);
|
auto h = conv->forward(ctx, x);
|
||||||
if (stride != 1) {
|
if (stride != 1) {
|
||||||
@ -236,7 +236,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* past) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* past) {
|
||||||
// x: [n, channels, h, w]
|
// x: [n, channels, h, w]
|
||||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
|
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
|
||||||
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
|
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
|
||||||
@ -260,10 +260,10 @@ public:
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
struct ggml_tensor* patchify(struct ggml_context* ctx,
|
ggml_tensor* patchify(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t patch_size,
|
int64_t patch_size,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// x: [f, b*c, h*q, w*r]
|
// x: [f, b*c, h*q, w*r]
|
||||||
// return: [f, b*c*r*q, h, w]
|
// return: [f, b*c*r*q, h, w]
|
||||||
if (patch_size == 1) {
|
if (patch_size == 1) {
|
||||||
@ -289,10 +289,10 @@ struct ggml_tensor* patchify(struct ggml_context* ctx,
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* unpatchify(struct ggml_context* ctx,
|
ggml_tensor* unpatchify(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t patch_size,
|
int64_t patch_size,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// x: [f, b*c*r*q, h, w]
|
// x: [f, b*c*r*q, h, w]
|
||||||
// return: [f, b*c, h*q, w*r]
|
// return: [f, b*c, h*q, w*r]
|
||||||
if (patch_size == 1) {
|
if (patch_size == 1) {
|
||||||
@ -339,7 +339,7 @@ public:
|
|||||||
blocks[std::to_string(index)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, z_channels, {3, 3}, {1, 1}, {1, 1}));
|
blocks[std::to_string(index)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, z_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
|
||||||
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
|
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
|
||||||
|
|
||||||
if (patch_size > 1) {
|
if (patch_size > 1) {
|
||||||
@ -396,7 +396,7 @@ public:
|
|||||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels[num_layers], out_channels * patch_size * patch_size, {3, 3}, {1, 1}, {1, 1}));
|
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels[num_layers], out_channels * patch_size * patch_size, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
|
||||||
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
|
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
|
||||||
|
|
||||||
// Clamp()
|
// Clamp()
|
||||||
@ -459,7 +459,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
|
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||||
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
|
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
|
||||||
if (sd_version_is_wan(version)) {
|
if (sd_version_is_wan(version)) {
|
||||||
// (W, H, C, T) -> (W, H, T, C)
|
// (W, H, C, T) -> (W, H, T, C)
|
||||||
@ -473,7 +473,7 @@ public:
|
|||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
|
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
|
||||||
// (W, H, T, C) -> (W, H, C, T)
|
// (W, H, T, C) -> (W, H, C, T)
|
||||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||||
@ -519,7 +519,7 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* decode(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
|
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||||
auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
|
auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
|
||||||
if (taef2) {
|
if (taef2) {
|
||||||
z = unpatchify(ctx->ggml_ctx, z, 2);
|
z = unpatchify(ctx->ggml_ctx, z, 2);
|
||||||
@ -527,7 +527,7 @@ public:
|
|||||||
return decoder->forward(ctx, z);
|
return decoder->forward(ctx, z);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* encode(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
|
auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
|
||||||
auto z = encoder->forward(ctx, x);
|
auto z = encoder->forward(ctx, x);
|
||||||
if (taef2) {
|
if (taef2) {
|
||||||
@ -558,7 +558,7 @@ struct TinyImageAutoEncoder : public VAE {
|
|||||||
return "taesd";
|
return "taesd";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
taesd.get_param_tensors(tensors, prefix);
|
taesd.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -578,21 +578,21 @@ struct TinyImageAutoEncoder : public VAE {
|
|||||||
return taesd.z_channels;
|
return taesd.z_channels;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
|
ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
z = to_backend(z);
|
z = to_backend(z);
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
struct ggml_tensor* out = decode_graph ? taesd.decode(&runner_ctx, z) : taesd.encode(&runner_ctx, z);
|
ggml_tensor* out = decode_graph ? taesd.decode(&runner_ctx, z) : taesd.encode(&runner_ctx, z);
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
bool _compute(const int n_threads,
|
bool _compute(const int n_threads,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
bool decode_graph,
|
bool decode_graph,
|
||||||
struct ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(z, decode_graph);
|
return build_graph(z, decode_graph);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -621,7 +621,7 @@ struct TinyVideoAutoEncoder : public VAE {
|
|||||||
return "taehv";
|
return "taehv";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
taehv.get_param_tensors(tensors, prefix);
|
taehv.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -641,21 +641,21 @@ struct TinyVideoAutoEncoder : public VAE {
|
|||||||
return taehv.z_channels;
|
return taehv.z_channels;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
|
ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
|
||||||
struct ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||||
z = to_backend(z);
|
z = to_backend(z);
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
struct ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
|
ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
bool _compute(const int n_threads,
|
bool _compute(const int n_threads,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
bool decode_graph,
|
bool decode_graph,
|
||||||
struct ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(z, decode_graph);
|
return build_graph(z, decode_graph);
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
114
src/unet.hpp
114
src/unet.hpp
@ -60,10 +60,10 @@ public:
|
|||||||
blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
|
blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int timesteps) {
|
int timesteps) {
|
||||||
// x: [N, in_channels, h, w] aka [b*t, in_channels, h, w], t == timesteps
|
// x: [N, in_channels, h, w] aka [b*t, in_channels, h, w], t == timesteps
|
||||||
// context: [N, max_position(aka n_context), hidden_size(aka context_dim)] aka [b*t, n_context, context_dim], t == timesteps
|
// context: [N, max_position(aka n_context), hidden_size(aka context_dim)] aka [b*t, n_context, context_dim], t == timesteps
|
||||||
// t_emb: [N, in_channels] aka [b*t, in_channels]
|
// t_emb: [N, in_channels] aka [b*t, in_channels]
|
||||||
@ -388,11 +388,11 @@ public:
|
|||||||
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* resblock_forward(std::string name,
|
ggml_tensor* resblock_forward(std::string name,
|
||||||
GGMLRunnerContext* ctx,
|
GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* emb,
|
ggml_tensor* emb,
|
||||||
int num_video_frames) {
|
int num_video_frames) {
|
||||||
if (version == VERSION_SVD) {
|
if (version == VERSION_SVD) {
|
||||||
auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
|
auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
|
||||||
|
|
||||||
@ -404,11 +404,11 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* attention_layer_forward(std::string name,
|
ggml_tensor* attention_layer_forward(std::string name,
|
||||||
GGMLRunnerContext* ctx,
|
GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int timesteps) {
|
int timesteps) {
|
||||||
if (version == VERSION_SVD) {
|
if (version == VERSION_SVD) {
|
||||||
auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
|
auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
|
||||||
|
|
||||||
@ -420,15 +420,15 @@ public:
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat = nullptr,
|
ggml_tensor* c_concat = nullptr,
|
||||||
struct ggml_tensor* y = nullptr,
|
ggml_tensor* y = nullptr,
|
||||||
int num_video_frames = -1,
|
int num_video_frames = -1,
|
||||||
std::vector<struct ggml_tensor*> controls = {},
|
std::vector<ggml_tensor*> controls = {},
|
||||||
float control_strength = 0.f) {
|
float control_strength = 0.f) {
|
||||||
// x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
|
// x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
|
||||||
// timesteps: [N,]
|
// timesteps: [N,]
|
||||||
// context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
|
// context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
|
||||||
@ -480,7 +480,7 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
// input_blocks
|
// input_blocks
|
||||||
std::vector<struct ggml_tensor*> hs;
|
std::vector<ggml_tensor*> hs;
|
||||||
|
|
||||||
// input block 0
|
// input block 0
|
||||||
auto h = input_blocks_0_0->forward(ctx, x);
|
auto h = input_blocks_0_0->forward(ctx, x);
|
||||||
@ -605,19 +605,19 @@ struct UNetModelRunner : public GGMLRunner {
|
|||||||
return "unet";
|
return "unet";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
unet.get_param_tensors(tensors, prefix);
|
unet.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat = nullptr,
|
ggml_tensor* c_concat = nullptr,
|
||||||
struct ggml_tensor* y = nullptr,
|
ggml_tensor* y = nullptr,
|
||||||
int num_video_frames = -1,
|
int num_video_frames = -1,
|
||||||
std::vector<struct ggml_tensor*> controls = {},
|
std::vector<ggml_tensor*> controls = {},
|
||||||
float control_strength = 0.f) {
|
float control_strength = 0.f) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
|
||||||
|
|
||||||
if (num_video_frames == -1) {
|
if (num_video_frames == -1) {
|
||||||
num_video_frames = static_cast<int>(x->ne[3]);
|
num_video_frames = static_cast<int>(x->ne[3]);
|
||||||
@ -635,15 +635,15 @@ struct UNetModelRunner : public GGMLRunner {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = unet.forward(&runner_ctx,
|
ggml_tensor* out = unet.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
timesteps,
|
timesteps,
|
||||||
context,
|
context,
|
||||||
c_concat,
|
c_concat,
|
||||||
y,
|
y,
|
||||||
num_video_frames,
|
num_video_frames,
|
||||||
controls,
|
controls,
|
||||||
control_strength);
|
control_strength);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -651,22 +651,22 @@ struct UNetModelRunner : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* c_concat,
|
ggml_tensor* c_concat,
|
||||||
struct ggml_tensor* y,
|
ggml_tensor* y,
|
||||||
int num_video_frames = -1,
|
int num_video_frames = -1,
|
||||||
std::vector<struct ggml_tensor*> controls = {},
|
std::vector<ggml_tensor*> controls = {},
|
||||||
float control_strength = 0.f,
|
float control_strength = 0.f,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
|
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
|
||||||
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
|
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
|
||||||
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
|
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -674,12 +674,12 @@ struct UNetModelRunner : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -703,7 +703,7 @@ struct UNetModelRunner : public GGMLRunner {
|
|||||||
ggml_set_f32(y, 0.5f);
|
ggml_set_f32(y, 0.5f);
|
||||||
// print_ggml_tensor(y);
|
// print_ggml_tensor(y);
|
||||||
|
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
compute(8, x, timesteps, context, nullptr, y, num_video_frames, {}, 0.f, &out, work_ctx);
|
compute(8, x, timesteps, context, nullptr, y, num_video_frames, {}, 0.f, &out, work_ctx);
|
||||||
|
|||||||
@ -72,13 +72,13 @@ struct UpscalerGGML {
|
|||||||
LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
|
LOG_INFO("upscaling from (%i x %i) to (%i x %i)",
|
||||||
input_image.width, input_image.height, output_width, output_height);
|
input_image.width, input_image.height, output_width, output_height);
|
||||||
|
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1G
|
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1G
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
// draft context
|
// draft context
|
||||||
struct ggml_context* upscale_ctx = ggml_init(params);
|
ggml_context* upscale_ctx = ggml_init(params);
|
||||||
if (!upscale_ctx) {
|
if (!upscale_ctx) {
|
||||||
LOG_ERROR("ggml_init() failed");
|
LOG_ERROR("ggml_init() failed");
|
||||||
return upscaled_image;
|
return upscaled_image;
|
||||||
|
|||||||
18
src/vae.hpp
18
src/vae.hpp
@ -6,12 +6,12 @@
|
|||||||
struct VAE : public GGMLRunner {
|
struct VAE : public GGMLRunner {
|
||||||
protected:
|
protected:
|
||||||
SDVersion version;
|
SDVersion version;
|
||||||
bool scale_input = true;
|
bool scale_input = true;
|
||||||
virtual bool _compute(const int n_threads,
|
virtual bool _compute(const int n_threads,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
bool decode_graph,
|
bool decode_graph,
|
||||||
struct ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
struct ggml_context* output_ctx) = 0;
|
ggml_context* output_ctx) = 0;
|
||||||
|
|
||||||
public:
|
public:
|
||||||
VAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
|
VAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
|
||||||
@ -186,7 +186,7 @@ public:
|
|||||||
virtual ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) = 0;
|
virtual ggml_tensor* vae_output_to_latents(ggml_context* work_ctx, ggml_tensor* vae_output, std::shared_ptr<RNG> rng) = 0;
|
||||||
virtual ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) = 0;
|
virtual ggml_tensor* diffusion_to_vae_latents(ggml_context* work_ctx, ggml_tensor* latents) = 0;
|
||||||
virtual ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) = 0;
|
virtual ggml_tensor* vae_to_diffuison_latents(ggml_context* work_ctx, ggml_tensor* latents) = 0;
|
||||||
virtual void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) = 0;
|
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) = 0;
|
||||||
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
|
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -199,10 +199,10 @@ struct FakeVAE : public VAE {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool _compute(const int n_threads,
|
bool _compute(const int n_threads,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
bool decode_graph,
|
bool decode_graph,
|
||||||
struct ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
struct ggml_context* output_ctx) override {
|
ggml_context* output_ctx) override {
|
||||||
if (*output == nullptr && output_ctx != nullptr) {
|
if (*output == nullptr && output_ctx != nullptr) {
|
||||||
*output = ggml_dup_tensor(output_ctx, z);
|
*output = ggml_dup_tensor(output_ctx, z);
|
||||||
}
|
}
|
||||||
@ -225,7 +225,7 @@ struct FakeVAE : public VAE {
|
|||||||
return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
|
return ggml_ext_dup_and_cpy_tensor(work_ctx, latents);
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {}
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {}
|
||||||
|
|
||||||
std::string get_desc() override {
|
std::string get_desc() override {
|
||||||
return "fake_vae";
|
return "fake_vae";
|
||||||
|
|||||||
396
src/wan.hpp
396
src/wan.hpp
@ -25,7 +25,7 @@ namespace WAN {
|
|||||||
std::tuple<int, int, int> dilation;
|
std::tuple<int, int, int> dilation;
|
||||||
bool bias;
|
bool bias;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
params["weight"] = ggml_new_tensor_4d(ctx,
|
params["weight"] = ggml_new_tensor_4d(ctx,
|
||||||
GGML_TYPE_F16,
|
GGML_TYPE_F16,
|
||||||
std::get<2>(kernel_size),
|
std::get<2>(kernel_size),
|
||||||
@ -53,11 +53,11 @@ namespace WAN {
|
|||||||
dilation(std::move(dilation)),
|
dilation(std::move(dilation)),
|
||||||
bias(bias) {}
|
bias(bias) {}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x, struct ggml_tensor* cache_x = nullptr) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) {
|
||||||
// x: [N*IC, ID, IH, IW]
|
// x: [N*IC, ID, IH, IW]
|
||||||
// result: x: [N*OC, ID, IH, IW]
|
// result: x: [N*OC, ID, IH, IW]
|
||||||
struct ggml_tensor* w = params["weight"];
|
ggml_tensor* w = params["weight"];
|
||||||
struct ggml_tensor* b = nullptr;
|
ggml_tensor* b = nullptr;
|
||||||
if (bias) {
|
if (bias) {
|
||||||
b = params["bias"];
|
b = params["bias"];
|
||||||
}
|
}
|
||||||
@ -86,7 +86,7 @@ namespace WAN {
|
|||||||
protected:
|
protected:
|
||||||
int64_t dim;
|
int64_t dim;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
ggml_type wtype = GGML_TYPE_F32;
|
ggml_type wtype = GGML_TYPE_F32;
|
||||||
auto iter = tensor_storage_map.find(prefix + "gamma");
|
auto iter = tensor_storage_map.find(prefix + "gamma");
|
||||||
if (iter != tensor_storage_map.end()) {
|
if (iter != tensor_storage_map.end()) {
|
||||||
@ -100,16 +100,16 @@ namespace WAN {
|
|||||||
RMS_norm(int64_t dim)
|
RMS_norm(int64_t dim)
|
||||||
: dim(dim) {}
|
: dim(dim) {}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) override {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
// x: [N*IC, ID, IH, IW], IC == dim
|
// x: [N*IC, ID, IH, IW], IC == dim
|
||||||
// assert N == 1
|
// assert N == 1
|
||||||
|
|
||||||
struct ggml_tensor* w = params["gamma"];
|
ggml_tensor* w = params["gamma"];
|
||||||
w = ggml_reshape_1d(ctx->ggml_ctx, w, ggml_nelements(w));
|
w = ggml_reshape_1d(ctx->ggml_ctx, w, ggml_nelements(w));
|
||||||
auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2)); // [ID, IH, IW, N*IC]
|
auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2)); // [ID, IH, IW, N*IC]
|
||||||
h = ggml_rms_norm(ctx->ggml_ctx, h, 1e-12f);
|
h = ggml_rms_norm(ctx->ggml_ctx, h, 1e-12f);
|
||||||
h = ggml_mul(ctx->ggml_ctx, h, w);
|
h = ggml_mul(ctx->ggml_ctx, h, w);
|
||||||
h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
|
h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0));
|
||||||
|
|
||||||
return h;
|
return h;
|
||||||
}
|
}
|
||||||
@ -148,12 +148,12 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b,
|
int64_t b,
|
||||||
std::vector<struct ggml_tensor*>& feat_cache,
|
std::vector<ggml_tensor*>& feat_cache,
|
||||||
int& feat_idx,
|
int& feat_idx,
|
||||||
int chunk_idx) {
|
int chunk_idx) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
int64_t c = x->ne[3] / b;
|
int64_t c = x->ne[3] / b;
|
||||||
@ -254,9 +254,9 @@ namespace WAN {
|
|||||||
GGML_ASSERT(in_channels * factor % out_channels == 0);
|
GGML_ASSERT(in_channels * factor % out_channels == 0);
|
||||||
group_size = in_channels * factor / out_channels;
|
group_size = in_channels * factor / out_channels;
|
||||||
}
|
}
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t B = 1) {
|
int64_t B = 1) {
|
||||||
// x: [B*IC, T, H, W]
|
// x: [B*IC, T, H, W]
|
||||||
// return: [B*OC, T/factor_t, H/factor_s, W/factor_s]
|
// return: [B*OC, T/factor_t, H/factor_s, W/factor_s]
|
||||||
GGML_ASSERT(B == 1);
|
GGML_ASSERT(B == 1);
|
||||||
@ -301,10 +301,10 @@ namespace WAN {
|
|||||||
GGML_ASSERT(out_channels * factor % in_channels == 0);
|
GGML_ASSERT(out_channels * factor % in_channels == 0);
|
||||||
repeats = out_channels * factor / in_channels;
|
repeats = out_channels * factor / in_channels;
|
||||||
}
|
}
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
bool first_chunk = false,
|
bool first_chunk = false,
|
||||||
int64_t B = 1) {
|
int64_t B = 1) {
|
||||||
// x: [B*IC, T, H, W]
|
// x: [B*IC, T, H, W]
|
||||||
// return: [B*OC, T/factor_t, H/factor_s, W/factor_s]
|
// return: [B*OC, T/factor_t, H/factor_s, W/factor_s]
|
||||||
GGML_ASSERT(B == 1);
|
GGML_ASSERT(B == 1);
|
||||||
@ -356,14 +356,14 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b,
|
int64_t b,
|
||||||
std::vector<struct ggml_tensor*>& feat_cache,
|
std::vector<ggml_tensor*>& feat_cache,
|
||||||
int& feat_idx) {
|
int& feat_idx) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
struct ggml_tensor* h = x;
|
ggml_tensor* h = x;
|
||||||
if (in_dim != out_dim) {
|
if (in_dim != out_dim) {
|
||||||
auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
|
auto shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["shortcut"]);
|
||||||
|
|
||||||
@ -430,15 +430,15 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b,
|
int64_t b,
|
||||||
std::vector<struct ggml_tensor*>& feat_cache,
|
std::vector<ggml_tensor*>& feat_cache,
|
||||||
int& feat_idx,
|
int& feat_idx,
|
||||||
int chunk_idx) {
|
int chunk_idx) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
struct ggml_tensor* x_copy = x;
|
ggml_tensor* x_copy = x;
|
||||||
|
|
||||||
auto avg_shortcut = std::dynamic_pointer_cast<AvgDown3D>(blocks["avg_shortcut"]);
|
auto avg_shortcut = std::dynamic_pointer_cast<AvgDown3D>(blocks["avg_shortcut"]);
|
||||||
|
|
||||||
@ -492,15 +492,15 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b,
|
int64_t b,
|
||||||
std::vector<struct ggml_tensor*>& feat_cache,
|
std::vector<ggml_tensor*>& feat_cache,
|
||||||
int& feat_idx,
|
int& feat_idx,
|
||||||
int chunk_idx) {
|
int chunk_idx) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
struct ggml_tensor* x_copy = x;
|
ggml_tensor* x_copy = x;
|
||||||
|
|
||||||
int i = 0;
|
int i = 0;
|
||||||
for (; i < mult; i++) {
|
for (; i < mult; i++) {
|
||||||
@ -537,9 +537,9 @@ namespace WAN {
|
|||||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {1, 1}));
|
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b) {
|
int64_t b) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
auto norm = std::dynamic_pointer_cast<RMS_norm>(blocks["norm"]);
|
auto norm = std::dynamic_pointer_cast<RMS_norm>(blocks["norm"]);
|
||||||
@ -659,12 +659,12 @@ namespace WAN {
|
|||||||
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
|
blocks["head.2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b,
|
int64_t b,
|
||||||
std::vector<struct ggml_tensor*>& feat_cache,
|
std::vector<ggml_tensor*>& feat_cache,
|
||||||
int& feat_idx,
|
int& feat_idx,
|
||||||
int chunk_idx) {
|
int chunk_idx) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
||||||
@ -830,12 +830,12 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b,
|
int64_t b,
|
||||||
std::vector<struct ggml_tensor*>& feat_cache,
|
std::vector<ggml_tensor*>& feat_cache,
|
||||||
int& feat_idx,
|
int& feat_idx,
|
||||||
int chunk_idx) {
|
int chunk_idx) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
||||||
@ -934,16 +934,16 @@ namespace WAN {
|
|||||||
|
|
||||||
int _conv_num = 33;
|
int _conv_num = 33;
|
||||||
int _conv_idx = 0;
|
int _conv_idx = 0;
|
||||||
std::vector<struct ggml_tensor*> _feat_map;
|
std::vector<ggml_tensor*> _feat_map;
|
||||||
int _enc_conv_num = 28;
|
int _enc_conv_num = 28;
|
||||||
int _enc_conv_idx = 0;
|
int _enc_conv_idx = 0;
|
||||||
std::vector<struct ggml_tensor*> _enc_feat_map;
|
std::vector<ggml_tensor*> _enc_feat_map;
|
||||||
|
|
||||||
void clear_cache() {
|
void clear_cache() {
|
||||||
_conv_idx = 0;
|
_conv_idx = 0;
|
||||||
_feat_map = std::vector<struct ggml_tensor*>(_conv_num, nullptr);
|
_feat_map = std::vector<ggml_tensor*>(_conv_num, nullptr);
|
||||||
_enc_conv_idx = 0;
|
_enc_conv_idx = 0;
|
||||||
_enc_feat_map = std::vector<struct ggml_tensor*>(_enc_conv_num, nullptr);
|
_enc_feat_map = std::vector<ggml_tensor*>(_enc_conv_num, nullptr);
|
||||||
}
|
}
|
||||||
|
|
||||||
public:
|
public:
|
||||||
@ -966,10 +966,10 @@ namespace WAN {
|
|||||||
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
|
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* patchify(struct ggml_context* ctx,
|
ggml_tensor* patchify(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t patch_size,
|
int64_t patch_size,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// x: [b*c, f, h*q, w*r]
|
// x: [b*c, f, h*q, w*r]
|
||||||
// return: [b*c*r*q, f, h, w]
|
// return: [b*c*r*q, f, h, w]
|
||||||
if (patch_size == 1) {
|
if (patch_size == 1) {
|
||||||
@ -993,10 +993,10 @@ namespace WAN {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* unpatchify(struct ggml_context* ctx,
|
ggml_tensor* unpatchify(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t patch_size,
|
int64_t patch_size,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// x: [b*c*r*q, f, h, w]
|
// x: [b*c*r*q, f, h, w]
|
||||||
// return: [b*c, f, h*q, w*r]
|
// return: [b*c, f, h*q, w*r]
|
||||||
if (patch_size == 1) {
|
if (patch_size == 1) {
|
||||||
@ -1019,9 +1019,9 @@ namespace WAN {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* encode(GGMLRunnerContext* ctx,
|
ggml_tensor* encode(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// x: [b*c, t, h, w]
|
// x: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
GGML_ASSERT(decode_only == false);
|
GGML_ASSERT(decode_only == false);
|
||||||
@ -1037,7 +1037,7 @@ namespace WAN {
|
|||||||
|
|
||||||
int64_t t = x->ne[2];
|
int64_t t = x->ne[2];
|
||||||
int64_t iter_ = 1 + (t - 1) / 4;
|
int64_t iter_ = 1 + (t - 1) / 4;
|
||||||
struct ggml_tensor* out;
|
ggml_tensor* out;
|
||||||
for (int i = 0; i < iter_; i++) {
|
for (int i = 0; i < iter_; i++) {
|
||||||
_enc_conv_idx = 0;
|
_enc_conv_idx = 0;
|
||||||
if (i == 0) {
|
if (i == 0) {
|
||||||
@ -1055,9 +1055,9 @@ namespace WAN {
|
|||||||
return mu;
|
return mu;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* decode(GGMLRunnerContext* ctx,
|
ggml_tensor* decode(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// z: [b*c, t, h, w]
|
// z: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
|
|
||||||
@ -1068,7 +1068,7 @@ namespace WAN {
|
|||||||
|
|
||||||
int64_t iter_ = z->ne[2];
|
int64_t iter_ = z->ne[2];
|
||||||
auto x = conv2->forward(ctx, z);
|
auto x = conv2->forward(ctx, z);
|
||||||
struct ggml_tensor* out;
|
ggml_tensor* out;
|
||||||
for (int i = 0; i < iter_; i++) {
|
for (int i = 0; i < iter_; i++) {
|
||||||
_conv_idx = 0;
|
_conv_idx = 0;
|
||||||
if (i == 0) {
|
if (i == 0) {
|
||||||
@ -1087,10 +1087,10 @@ namespace WAN {
|
|||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* decode_partial(GGMLRunnerContext* ctx,
|
ggml_tensor* decode_partial(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
int i,
|
int i,
|
||||||
int64_t b = 1) {
|
int64_t b = 1) {
|
||||||
// z: [b*c, t, h, w]
|
// z: [b*c, t, h, w]
|
||||||
GGML_ASSERT(b == 1);
|
GGML_ASSERT(b == 1);
|
||||||
|
|
||||||
@ -1127,7 +1127,7 @@ namespace WAN {
|
|||||||
return "wan_vae";
|
return "wan_vae";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) override {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
|
||||||
ae.get_param_tensors(tensors, prefix);
|
ae.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1226,22 +1226,22 @@ namespace WAN {
|
|||||||
return static_cast<int>(ae.z_dim);
|
return static_cast<int>(ae.z_dim);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* z, bool decode_graph) {
|
ggml_cgraph* build_graph(ggml_tensor* z, bool decode_graph) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(10240 * z->ne[2]);
|
ggml_cgraph* gf = new_graph_custom(10240 * z->ne[2]);
|
||||||
|
|
||||||
z = to_backend(z);
|
z = to_backend(z);
|
||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
|
ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
return gf;
|
return gf;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph_partial(struct ggml_tensor* z, bool decode_graph, int i) {
|
ggml_cgraph* build_graph_partial(ggml_tensor* z, bool decode_graph, int i) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(20480);
|
ggml_cgraph* gf = new_graph_custom(20480);
|
||||||
|
|
||||||
ae.clear_cache();
|
ae.clear_cache();
|
||||||
|
|
||||||
@ -1254,7 +1254,7 @@ namespace WAN {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = decode_graph ? ae.decode_partial(&runner_ctx, z, i) : ae.encode(&runner_ctx, z);
|
ggml_tensor* out = decode_graph ? ae.decode_partial(&runner_ctx, z, i) : ae.encode(&runner_ctx, z);
|
||||||
|
|
||||||
for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
|
for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
|
||||||
ggml_tensor* feat_cache = ae._feat_map[feat_idx];
|
ggml_tensor* feat_cache = ae._feat_map[feat_idx];
|
||||||
@ -1270,12 +1270,12 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool _compute(const int n_threads,
|
bool _compute(const int n_threads,
|
||||||
struct ggml_tensor* z,
|
ggml_tensor* z,
|
||||||
bool decode_graph,
|
bool decode_graph,
|
||||||
struct ggml_tensor** output,
|
ggml_tensor** output,
|
||||||
struct ggml_context* output_ctx = nullptr) override {
|
ggml_context* output_ctx = nullptr) override {
|
||||||
if (true) {
|
if (true) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(z, decode_graph);
|
return build_graph(z, decode_graph);
|
||||||
};
|
};
|
||||||
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
return GGMLRunner::compute(get_graph, n_threads, true, output, output_ctx);
|
||||||
@ -1283,11 +1283,11 @@ namespace WAN {
|
|||||||
ae.clear_cache();
|
ae.clear_cache();
|
||||||
int64_t t = z->ne[2];
|
int64_t t = z->ne[2];
|
||||||
int i = 0;
|
int i = 0;
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph_partial(z, decode_graph, i);
|
return build_graph_partial(z, decode_graph, i);
|
||||||
};
|
};
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
bool res = GGMLRunner::compute(get_graph, n_threads, true, &out, output_ctx);
|
bool res = GGMLRunner::compute(get_graph, n_threads, true, &out, output_ctx);
|
||||||
ae.clear_cache();
|
ae.clear_cache();
|
||||||
if (t == 1) {
|
if (t == 1) {
|
||||||
*output = out;
|
*output = out;
|
||||||
@ -1325,12 +1325,12 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1G
|
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1G
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
if (true) {
|
if (true) {
|
||||||
@ -1342,7 +1342,7 @@ namespace WAN {
|
|||||||
ggml_set_f32(z, 0.5f);
|
ggml_set_f32(z, 0.5f);
|
||||||
z = load_tensor_from_file(work_ctx, "wan_vae_z.bin");
|
z = load_tensor_from_file(work_ctx, "wan_vae_z.bin");
|
||||||
print_ggml_tensor(z);
|
print_ggml_tensor(z);
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
_compute(8, z, true, &out, work_ctx);
|
_compute(8, z, true, &out, work_ctx);
|
||||||
@ -1410,10 +1410,10 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr) {
|
ggml_tensor* mask = nullptr) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// pe: [n_token, d_head/2, 2, 2]
|
// pe: [n_token, d_head/2, 2, 2]
|
||||||
// return [N, n_token, dim]
|
// return [N, n_token, dim]
|
||||||
@ -1451,10 +1451,10 @@ namespace WAN {
|
|||||||
bool qk_norm = true,
|
bool qk_norm = true,
|
||||||
float eps = 1e-6)
|
float eps = 1e-6)
|
||||||
: WanSelfAttention(dim, num_heads, qk_norm, eps) {}
|
: WanSelfAttention(dim, num_heads, qk_norm, eps) {}
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int64_t context_img_len) = 0;
|
int64_t context_img_len) = 0;
|
||||||
};
|
};
|
||||||
|
|
||||||
class WanT2VCrossAttention : public WanCrossAttention {
|
class WanT2VCrossAttention : public WanCrossAttention {
|
||||||
@ -1464,10 +1464,10 @@ namespace WAN {
|
|||||||
bool qk_norm = true,
|
bool qk_norm = true,
|
||||||
float eps = 1e-6)
|
float eps = 1e-6)
|
||||||
: WanCrossAttention(dim, num_heads, qk_norm, eps) {}
|
: WanCrossAttention(dim, num_heads, qk_norm, eps) {}
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int64_t context_img_len) override {
|
int64_t context_img_len) override {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// context: [N, n_context, dim]
|
// context: [N, n_context, dim]
|
||||||
// context_img_len: unused
|
// context_img_len: unused
|
||||||
@ -1512,10 +1512,10 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int64_t context_img_len) override {
|
int64_t context_img_len) override {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// context: [N, context_img_len + context_txt_len, dim]
|
// context: [N, context_img_len + context_txt_len, dim]
|
||||||
// return [N, n_token, dim]
|
// return [N, n_token, dim]
|
||||||
@ -1560,7 +1560,7 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
static struct ggml_tensor* modulate_add(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* e) {
|
static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// e: [N, 1, dim] or [N, T, 1, dim]
|
// e: [N, 1, dim] or [N, T, 1, dim]
|
||||||
if (ggml_n_dims(e) == 3) {
|
if (ggml_n_dims(e) == 3) {
|
||||||
@ -1574,7 +1574,7 @@ namespace WAN {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
static struct ggml_tensor* modulate_mul(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* e) {
|
static ggml_tensor* modulate_mul(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// e: [N, 1, dim] or [N, T, 1, dim]
|
// e: [N, 1, dim] or [N, T, 1, dim]
|
||||||
if (ggml_n_dims(e) == 3) {
|
if (ggml_n_dims(e) == 3) {
|
||||||
@ -1592,7 +1592,7 @@ namespace WAN {
|
|||||||
protected:
|
protected:
|
||||||
int64_t dim;
|
int64_t dim;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
||||||
params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
|
params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
|
||||||
}
|
}
|
||||||
@ -1626,12 +1626,12 @@ namespace WAN {
|
|||||||
blocks["ffn.2"] = std::shared_ptr<GGMLBlock>(new Linear(ffn_dim, dim));
|
blocks["ffn.2"] = std::shared_ptr<GGMLBlock>(new Linear(ffn_dim, dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* e,
|
ggml_tensor* e,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int64_t context_img_len = 257) {
|
int64_t context_img_len = 257) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// e: [N, 6, dim] or [N, T, 6, dim]
|
// e: [N, 6, dim] or [N, T, 6, dim]
|
||||||
// context: [N, context_img_len + context_txt_len, dim]
|
// context: [N, context_img_len + context_txt_len, dim]
|
||||||
@ -1680,7 +1680,7 @@ namespace WAN {
|
|||||||
class VaceWanAttentionBlock : public WanAttentionBlock {
|
class VaceWanAttentionBlock : public WanAttentionBlock {
|
||||||
protected:
|
protected:
|
||||||
int block_id;
|
int block_id;
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
||||||
params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
|
params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1);
|
||||||
}
|
}
|
||||||
@ -1702,11 +1702,11 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
|
|
||||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* c,
|
ggml_tensor* c,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* e,
|
ggml_tensor* e,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
int64_t context_img_len = 257) {
|
int64_t context_img_len = 257) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// e: [N, 6, dim] or [N, T, 6, dim]
|
// e: [N, 6, dim] or [N, T, 6, dim]
|
||||||
@ -1732,7 +1732,7 @@ namespace WAN {
|
|||||||
protected:
|
protected:
|
||||||
int64_t dim;
|
int64_t dim;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
||||||
params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 2, 1);
|
params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 2, 1);
|
||||||
}
|
}
|
||||||
@ -1749,9 +1749,9 @@ namespace WAN {
|
|||||||
blocks["head"] = std::shared_ptr<GGMLBlock>(new Linear(dim, out_dim));
|
blocks["head"] = std::shared_ptr<GGMLBlock>(new Linear(dim, out_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* e) {
|
ggml_tensor* e) {
|
||||||
// x: [N, n_token, dim]
|
// x: [N, n_token, dim]
|
||||||
// e: [N, dim] or [N, T, dim]
|
// e: [N, dim] or [N, T, dim]
|
||||||
// return [N, n_token, out_dim]
|
// return [N, n_token, out_dim]
|
||||||
@ -1779,7 +1779,7 @@ namespace WAN {
|
|||||||
int64_t in_dim;
|
int64_t in_dim;
|
||||||
int64_t flf_pos_embed_token_number;
|
int64_t flf_pos_embed_token_number;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
if (flf_pos_embed_token_number > 0) {
|
if (flf_pos_embed_token_number > 0) {
|
||||||
params["emb_pos"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, in_dim, flf_pos_embed_token_number, 1);
|
params["emb_pos"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, in_dim, flf_pos_embed_token_number, 1);
|
||||||
}
|
}
|
||||||
@ -1797,8 +1797,8 @@ namespace WAN {
|
|||||||
blocks["proj.4"] = std::shared_ptr<GGMLBlock>(new LayerNorm(out_dim));
|
blocks["proj.4"] = std::shared_ptr<GGMLBlock>(new LayerNorm(out_dim));
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* image_embeds) {
|
ggml_tensor* image_embeds) {
|
||||||
if (flf_pos_embed_token_number > 0) {
|
if (flf_pos_embed_token_number > 0) {
|
||||||
auto emb_pos = params["emb_pos"];
|
auto emb_pos = params["emb_pos"];
|
||||||
|
|
||||||
@ -1917,8 +1917,8 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
|
ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x) {
|
ggml_tensor* x) {
|
||||||
int64_t W = x->ne[0];
|
int64_t W = x->ne[0];
|
||||||
int64_t H = x->ne[1];
|
int64_t H = x->ne[1];
|
||||||
int64_t T = x->ne[2];
|
int64_t T = x->ne[2];
|
||||||
@ -1930,11 +1930,11 @@ namespace WAN {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* unpatchify(struct ggml_context* ctx,
|
ggml_tensor* unpatchify(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
int64_t t_len,
|
int64_t t_len,
|
||||||
int64_t h_len,
|
int64_t h_len,
|
||||||
int64_t w_len) {
|
int64_t w_len) {
|
||||||
// x: [N, t_len*h_len*w_len, pt*ph*pw*C]
|
// x: [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||||
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
|
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
|
||||||
int64_t N = x->ne[3];
|
int64_t N = x->ne[3];
|
||||||
@ -1957,15 +1957,15 @@ namespace WAN {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* clip_fea = nullptr,
|
ggml_tensor* clip_fea = nullptr,
|
||||||
struct ggml_tensor* vace_context = nullptr,
|
ggml_tensor* vace_context = nullptr,
|
||||||
float vace_strength = 1.f,
|
float vace_strength = 1.f,
|
||||||
int64_t N = 1) {
|
int64_t N = 1) {
|
||||||
// x: [N*C, T, H, W], C => in_dim
|
// x: [N*C, T, H, W], C => in_dim
|
||||||
// vace_context: [N*vace_in_dim, T, H, W]
|
// vace_context: [N*vace_in_dim, T, H, W]
|
||||||
// timestep: [N,] or [T]
|
// timestep: [N,] or [T]
|
||||||
@ -2051,16 +2051,16 @@ namespace WAN {
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* clip_fea = nullptr,
|
ggml_tensor* clip_fea = nullptr,
|
||||||
struct ggml_tensor* time_dim_concat = nullptr,
|
ggml_tensor* time_dim_concat = nullptr,
|
||||||
struct ggml_tensor* vace_context = nullptr,
|
ggml_tensor* vace_context = nullptr,
|
||||||
float vace_strength = 1.f,
|
float vace_strength = 1.f,
|
||||||
int64_t N = 1) {
|
int64_t N = 1) {
|
||||||
// Forward pass of DiT.
|
// Forward pass of DiT.
|
||||||
// x: [N*C, T, H, W]
|
// x: [N*C, T, H, W]
|
||||||
// timestep: [N,]
|
// timestep: [N,]
|
||||||
@ -2225,19 +2225,19 @@ namespace WAN {
|
|||||||
return desc;
|
return desc;
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
wan.get_param_tensors(tensors, prefix);
|
wan.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* clip_fea = nullptr,
|
ggml_tensor* clip_fea = nullptr,
|
||||||
struct ggml_tensor* c_concat = nullptr,
|
ggml_tensor* c_concat = nullptr,
|
||||||
struct ggml_tensor* time_dim_concat = nullptr,
|
ggml_tensor* time_dim_concat = nullptr,
|
||||||
struct ggml_tensor* vace_context = nullptr,
|
ggml_tensor* vace_context = nullptr,
|
||||||
float vace_strength = 1.f) {
|
float vace_strength = 1.f) {
|
||||||
struct ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE);
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
timesteps = to_backend(timesteps);
|
timesteps = to_backend(timesteps);
|
||||||
@ -2270,15 +2270,15 @@ namespace WAN {
|
|||||||
|
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = wan.forward(&runner_ctx,
|
ggml_tensor* out = wan.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
timesteps,
|
timesteps,
|
||||||
context,
|
context,
|
||||||
pe,
|
pe,
|
||||||
clip_fea,
|
clip_fea,
|
||||||
time_dim_concat,
|
time_dim_concat,
|
||||||
vace_context,
|
vace_context,
|
||||||
vace_strength);
|
vace_strength);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -2286,17 +2286,17 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* clip_fea = nullptr,
|
ggml_tensor* clip_fea = nullptr,
|
||||||
struct ggml_tensor* c_concat = nullptr,
|
ggml_tensor* c_concat = nullptr,
|
||||||
struct ggml_tensor* time_dim_concat = nullptr,
|
ggml_tensor* time_dim_concat = nullptr,
|
||||||
struct ggml_tensor* vace_context = nullptr,
|
ggml_tensor* vace_context = nullptr,
|
||||||
float vace_strength = 1.f,
|
float vace_strength = 1.f,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
|
return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -2304,12 +2304,12 @@ namespace WAN {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(200 * 1024 * 1024); // 200 MB
|
params.mem_size = static_cast<size_t>(200 * 1024 * 1024); // 200 MB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -2332,7 +2332,7 @@ namespace WAN {
|
|||||||
// auto clip_fea = load_tensor_from_file(work_ctx, "wan_dit_clip_fea.bin");
|
// auto clip_fea = load_tensor_from_file(work_ctx, "wan_dit_clip_fea.bin");
|
||||||
// print_ggml_tensor(clip_fea);
|
// print_ggml_tensor(clip_fea);
|
||||||
|
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
compute(8, x, timesteps, context, nullptr, nullptr, nullptr, nullptr, 1.f, &out, work_ctx);
|
compute(8, x, timesteps, context, nullptr, nullptr, nullptr, nullptr, 1.f, &out, work_ctx);
|
||||||
|
|||||||
102
src/z_image.hpp
102
src/z_image.hpp
@ -42,10 +42,10 @@ namespace ZImage {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr) {
|
ggml_tensor* mask = nullptr) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
int64_t n_token = x->ne[1];
|
int64_t n_token = x->ne[1];
|
||||||
int64_t N = x->ne[2];
|
int64_t N = x->ne[2];
|
||||||
@ -124,23 +124,23 @@ namespace ZImage {
|
|||||||
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
|
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
|
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
|
||||||
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
|
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
|
||||||
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
|
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
|
||||||
|
|
||||||
auto x1 = w1->forward(ctx, x);
|
auto x1 = w1->forward(ctx, x);
|
||||||
auto x3 = w3->forward(ctx, x);
|
auto x3 = w3->forward(ctx, x);
|
||||||
x = ggml_mul(ctx->ggml_ctx, ggml_silu(ctx->ggml_ctx, x1), x3);
|
x = ggml_swiglu_split(ctx->ggml_ctx, x1, x3);
|
||||||
x = w2->forward(ctx, x);
|
x = w2->forward(ctx, x);
|
||||||
|
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
__STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx,
|
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* scale) {
|
ggml_tensor* scale) {
|
||||||
// x: [N, L, C]
|
// x: [N, L, C]
|
||||||
// scale: [N, C]
|
// scale: [N, C]
|
||||||
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
|
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
|
||||||
@ -175,11 +175,11 @@ namespace ZImage {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
struct ggml_tensor* mask = nullptr,
|
ggml_tensor* mask = nullptr,
|
||||||
struct ggml_tensor* adaln_input = nullptr) {
|
ggml_tensor* adaln_input = nullptr) {
|
||||||
auto attention = std::dynamic_pointer_cast<JointAttention>(blocks["attention"]);
|
auto attention = std::dynamic_pointer_cast<JointAttention>(blocks["attention"]);
|
||||||
auto feed_forward = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"]);
|
auto feed_forward = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"]);
|
||||||
auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
|
auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
|
||||||
@ -241,9 +241,9 @@ namespace ZImage {
|
|||||||
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), hidden_size);
|
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), hidden_size);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* c) {
|
ggml_tensor* c) {
|
||||||
// x: [N, n_token, hidden_size]
|
// x: [N, n_token, hidden_size]
|
||||||
// c: [N, hidden_size]
|
// c: [N, hidden_size]
|
||||||
// return: [N, n_token, patch_size * patch_size * out_channels]
|
// return: [N, n_token, patch_size * patch_size * out_channels]
|
||||||
@ -284,7 +284,7 @@ namespace ZImage {
|
|||||||
protected:
|
protected:
|
||||||
ZImageParams z_image_params;
|
ZImageParams z_image_params;
|
||||||
|
|
||||||
void init_params(struct ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
|
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
|
||||||
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
|
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
|
||||||
}
|
}
|
||||||
@ -346,11 +346,11 @@ namespace ZImage {
|
|||||||
blocks["final_layer"] = std::make_shared<FinalLayer>(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels);
|
blocks["final_layer"] = std::make_shared<FinalLayer>(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward_core(GGMLRunnerContext* ctx,
|
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* pe) {
|
ggml_tensor* pe) {
|
||||||
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
|
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
|
||||||
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
|
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
|
||||||
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
|
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
|
||||||
@ -414,12 +414,12 @@ namespace ZImage {
|
|||||||
return img;
|
return img;
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timestep,
|
ggml_tensor* timestep,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
struct ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
std::vector<ggml_tensor*> ref_latents = {}) {
|
std::vector<ggml_tensor*> ref_latents = {}) {
|
||||||
// Forward pass of DiT.
|
// Forward pass of DiT.
|
||||||
// x: [N, C, H, W]
|
// x: [N, C, H, W]
|
||||||
// timestep: [N,]
|
// timestep: [N,]
|
||||||
@ -477,17 +477,17 @@ namespace ZImage {
|
|||||||
return "z_image";
|
return "z_image";
|
||||||
}
|
}
|
||||||
|
|
||||||
void get_param_tensors(std::map<std::string, struct ggml_tensor*>& tensors, const std::string prefix) {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||||
z_image.get_param_tensors(tensors, prefix);
|
z_image.get_param_tensors(tensors, prefix);
|
||||||
}
|
}
|
||||||
|
|
||||||
struct ggml_cgraph* build_graph(struct ggml_tensor* x,
|
ggml_cgraph* build_graph(ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
bool increase_ref_index = false) {
|
bool increase_ref_index = false) {
|
||||||
GGML_ASSERT(x->ne[3] == 1);
|
GGML_ASSERT(x->ne[3] == 1);
|
||||||
struct ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
|
||||||
|
|
||||||
x = to_backend(x);
|
x = to_backend(x);
|
||||||
context = to_backend(context);
|
context = to_backend(context);
|
||||||
@ -518,12 +518,12 @@ namespace ZImage {
|
|||||||
set_backend_tensor_data(pe, pe_vec.data());
|
set_backend_tensor_data(pe, pe_vec.data());
|
||||||
auto runner_ctx = get_context();
|
auto runner_ctx = get_context();
|
||||||
|
|
||||||
struct ggml_tensor* out = z_image.forward(&runner_ctx,
|
ggml_tensor* out = z_image.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
timesteps,
|
timesteps,
|
||||||
context,
|
context,
|
||||||
pe,
|
pe,
|
||||||
ref_latents);
|
ref_latents);
|
||||||
|
|
||||||
ggml_build_forward_expand(gf, out);
|
ggml_build_forward_expand(gf, out);
|
||||||
|
|
||||||
@ -531,17 +531,17 @@ namespace ZImage {
|
|||||||
}
|
}
|
||||||
|
|
||||||
bool compute(int n_threads,
|
bool compute(int n_threads,
|
||||||
struct ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
struct ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
struct ggml_tensor* context,
|
ggml_tensor* context,
|
||||||
std::vector<ggml_tensor*> ref_latents = {},
|
std::vector<ggml_tensor*> ref_latents = {},
|
||||||
bool increase_ref_index = false,
|
bool increase_ref_index = false,
|
||||||
struct ggml_tensor** output = nullptr,
|
ggml_tensor** output = nullptr,
|
||||||
struct ggml_context* output_ctx = nullptr) {
|
ggml_context* output_ctx = nullptr) {
|
||||||
// x: [N, in_channels, h, w]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]
|
// context: [N, max_position, hidden_size]
|
||||||
auto get_graph = [&]() -> struct ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
|
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -549,12 +549,12 @@ namespace ZImage {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void test() {
|
void test() {
|
||||||
struct ggml_init_params params;
|
ggml_init_params params;
|
||||||
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
|
||||||
params.mem_buffer = nullptr;
|
params.mem_buffer = nullptr;
|
||||||
params.no_alloc = false;
|
params.no_alloc = false;
|
||||||
|
|
||||||
struct ggml_context* work_ctx = ggml_init(params);
|
ggml_context* work_ctx = ggml_init(params);
|
||||||
GGML_ASSERT(work_ctx != nullptr);
|
GGML_ASSERT(work_ctx != nullptr);
|
||||||
|
|
||||||
{
|
{
|
||||||
@ -571,7 +571,7 @@ namespace ZImage {
|
|||||||
auto context = load_tensor_from_file(work_ctx, "./z_image_context.bin");
|
auto context = load_tensor_from_file(work_ctx, "./z_image_context.bin");
|
||||||
print_ggml_tensor(context);
|
print_ggml_tensor(context);
|
||||||
|
|
||||||
struct ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
compute(8, x, timesteps, context, {}, false, &out, work_ctx);
|
compute(8, x, timesteps, context, {}, false, &out, work_ctx);
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user