mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-08-02 00:00:41 -05:00
feat: add boogu image support (#1688)
This commit is contained in:
leejet
835
src/model/diffusion/boogu.hpp
Normal file
835
src/model/diffusion/boogu.hpp
Normal file
@@ -0,0 +1,835 @@
|
||||
#ifndef __SD_MODEL_DIFFUSION_BOOGU_HPP__
|
||||
#define __SD_MODEL_DIFFUSION_BOOGU_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
#include "core/ggml_extend.hpp"
|
||||
#include "model/common/rope.hpp"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/diffusion/qwen_image.hpp"
|
||||
#include "model_loader.h"
|
||||
|
||||
namespace Boogu {
|
||||
constexpr int BOOGU_GRAPH_SIZE = 65536;
|
||||
|
||||
struct BooguConfig {
|
||||
int patch_size = 2;
|
||||
int64_t in_channels = 16;
|
||||
int64_t out_channels = 16;
|
||||
int64_t hidden_size = 3360;
|
||||
int64_t num_layers = 32;
|
||||
int64_t num_double_stream_layers = 8;
|
||||
int64_t num_refiner_layers = 2;
|
||||
int64_t num_attention_heads = 28;
|
||||
int64_t num_kv_heads = 7;
|
||||
int64_t head_dim = 120;
|
||||
int64_t multiple_of = 256;
|
||||
int64_t instruction_feat_dim = 4096;
|
||||
int64_t timestep_embed_dim = 1024;
|
||||
int theta = 10000;
|
||||
float timestep_scale = 1000.0f;
|
||||
float norm_eps = 1e-5f;
|
||||
std::vector<int> axes_dim = {40, 40, 40};
|
||||
int64_t axes_dim_sum = 120;
|
||||
|
||||
static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix,
|
||||
const std::string& block_prefix) {
|
||||
int64_t count = 0;
|
||||
for (const auto& [name, _] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
}
|
||||
size_t pos = name.find(block_prefix);
|
||||
if (pos == std::string::npos) {
|
||||
continue;
|
||||
}
|
||||
auto items = split_string(name.substr(pos), '.');
|
||||
if (items.size() > 1) {
|
||||
count = std::max<int64_t>(count, atoi(items[1].c_str()) + 1);
|
||||
}
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
static BooguConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
|
||||
BooguConfig config;
|
||||
int64_t detected_head_dim = 0;
|
||||
int64_t detected_kv_dim = 0;
|
||||
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
}
|
||||
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
|
||||
int64_t patch_area = config.patch_size * config.patch_size;
|
||||
config.in_channels = tensor_storage.ne[0] / patch_area;
|
||||
config.hidden_size = tensor_storage.ne[1];
|
||||
} else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) {
|
||||
config.instruction_feat_dim = tensor_storage.ne[0];
|
||||
config.hidden_size = tensor_storage.ne[1];
|
||||
} else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
|
||||
detected_head_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "double_stream_layers.0.img_self_attn.norm_q.weight") && tensor_storage.n_dims == 1) {
|
||||
detected_head_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "single_stream_layers.0.attn.to_k.weight") && tensor_storage.n_dims == 2) {
|
||||
detected_kv_dim = tensor_storage.ne[1];
|
||||
} else if (ends_with(name, "double_stream_layers.0.img_instruct_attn.processor.img_to_k.weight") && tensor_storage.n_dims == 2) {
|
||||
detected_kv_dim = tensor_storage.ne[1];
|
||||
} else if (ends_with(name, "norm_out.linear_2.weight") && tensor_storage.n_dims == 2) {
|
||||
int64_t patch_area = config.patch_size * config.patch_size;
|
||||
config.out_channels = tensor_storage.ne[1] / patch_area;
|
||||
}
|
||||
}
|
||||
|
||||
config.num_layers = std::max<int64_t>(1, count_blocks(tensor_storage_map, prefix, "single_stream_layers."));
|
||||
config.num_double_stream_layers = std::max<int64_t>(0, count_blocks(tensor_storage_map, prefix, "double_stream_layers."));
|
||||
int64_t noise_refiner_layers = count_blocks(tensor_storage_map, prefix, "noise_refiner.");
|
||||
int64_t ref_refiner_layers = count_blocks(tensor_storage_map, prefix, "ref_image_refiner.");
|
||||
int64_t context_refiner_layers = count_blocks(tensor_storage_map, prefix, "context_refiner.");
|
||||
config.num_refiner_layers = std::max<int64_t>(1, std::max(noise_refiner_layers, std::max(ref_refiner_layers, context_refiner_layers)));
|
||||
|
||||
if (detected_head_dim > 0) {
|
||||
config.head_dim = detected_head_dim;
|
||||
config.num_attention_heads = config.hidden_size / config.head_dim;
|
||||
config.axes_dim_sum = config.head_dim;
|
||||
if (detected_kv_dim > 0) {
|
||||
config.num_kv_heads = detected_kv_dim / config.head_dim;
|
||||
}
|
||||
if (config.axes_dim_sum == 120) {
|
||||
config.axes_dim = {40, 40, 40};
|
||||
} else if (config.axes_dim_sum % 3 == 0) {
|
||||
int axis = static_cast<int>(config.axes_dim_sum / 3);
|
||||
config.axes_dim = {axis, axis, axis};
|
||||
}
|
||||
}
|
||||
config.timestep_embed_dim = std::min<int64_t>(config.hidden_size, 1024);
|
||||
|
||||
LOG_DEBUG("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64,
|
||||
config.num_layers,
|
||||
config.num_double_stream_layers,
|
||||
config.num_refiner_layers,
|
||||
config.hidden_size,
|
||||
config.num_attention_heads,
|
||||
config.num_kv_heads,
|
||||
config.head_dim,
|
||||
config.in_channels,
|
||||
config.out_channels);
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* scale_modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
|
||||
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]);
|
||||
return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* gate_residual(ggml_context* ctx, ggml_tensor* residual, ggml_tensor* x, ggml_tensor* gate) {
|
||||
gate = ggml_tanh(ctx, gate);
|
||||
gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]);
|
||||
x = ggml_mul(ctx, x, gate);
|
||||
return ggml_add(ctx, residual, x);
|
||||
}
|
||||
|
||||
struct LuminaCombinedTimestepCaptionEmbedding : public GGMLBlock {
|
||||
int64_t frequency_embedding_size;
|
||||
float timestep_scale;
|
||||
|
||||
LuminaCombinedTimestepCaptionEmbedding(int64_t hidden_size,
|
||||
int64_t instruction_feat_dim,
|
||||
int64_t frequency_embedding_size,
|
||||
float norm_eps,
|
||||
float timestep_scale)
|
||||
: frequency_embedding_size(frequency_embedding_size),
|
||||
timestep_scale(timestep_scale) {
|
||||
blocks["timestep_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(frequency_embedding_size, std::min<int64_t>(hidden_size, 1024));
|
||||
blocks["caption_embedder.0"] = std::make_shared<RMSNorm>(instruction_feat_dim, norm_eps);
|
||||
blocks["caption_embedder.1"] = std::make_shared<Linear>(instruction_feat_dim, hidden_size, true);
|
||||
}
|
||||
|
||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* timestep, ggml_tensor* text_hidden_states) {
|
||||
auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
|
||||
auto caption_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["caption_embedder.0"]);
|
||||
auto caption_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["caption_embedder.1"]);
|
||||
|
||||
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(frequency_embedding_size), 10000, timestep_scale);
|
||||
auto time_embed = timestep_embedder->forward(ctx, timestep_proj);
|
||||
auto caption_embed = caption_embedder_1->forward(ctx, caption_embedder_0->forward(ctx, text_hidden_states));
|
||||
return {time_embed, caption_embed};
|
||||
}
|
||||
};
|
||||
|
||||
struct LuminaRMSNormZero : public GGMLBlock {
|
||||
LuminaRMSNormZero(int64_t embedding_dim, int64_t conditioning_embedding_dim, float norm_eps) {
|
||||
blocks["linear"] = std::make_shared<Linear>(conditioning_embedding_dim, 4 * embedding_dim, true);
|
||||
blocks["norm"] = std::make_shared<RMSNorm>(embedding_dim, norm_eps);
|
||||
}
|
||||
|
||||
std::tuple<ggml_tensor*, ggml_tensor*, ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb) {
|
||||
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
||||
|
||||
emb = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, emb));
|
||||
auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 4, 0);
|
||||
|
||||
auto scale_msa = mods[0];
|
||||
auto gate_msa = mods[1];
|
||||
auto scale_mlp = mods[2];
|
||||
auto gate_mlp = mods[3];
|
||||
|
||||
x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), scale_msa);
|
||||
return {x, gate_msa, scale_mlp, gate_mlp};
|
||||
}
|
||||
};
|
||||
|
||||
struct LuminaFeedForward : public GGMLBlock {
|
||||
LuminaFeedForward(int64_t dim, int64_t inner_dim, int64_t multiple_of) {
|
||||
inner_dim = multiple_of * ((inner_dim + multiple_of - 1) / multiple_of);
|
||||
blocks["linear_1"] = std::make_shared<Linear>(dim, inner_dim, false);
|
||||
blocks["linear_2"] = std::make_shared<Linear>(inner_dim, dim, false);
|
||||
blocks["linear_3"] = std::make_shared<Linear>(dim, inner_dim, false);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
|
||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
|
||||
auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
linear_2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
auto h1 = linear_1->forward(ctx, x);
|
||||
auto h2 = linear_3->forward(ctx, x);
|
||||
x = ggml_swiglu_split(ctx->ggml_ctx, h1, h2);
|
||||
x = linear_2->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LuminaLayerNormContinuous : public GGMLBlock {
|
||||
LuminaLayerNormContinuous(int64_t embedding_dim,
|
||||
int64_t conditioning_embedding_dim,
|
||||
int64_t out_dim) {
|
||||
blocks["linear_1"] = std::make_shared<Linear>(conditioning_embedding_dim, embedding_dim, true);
|
||||
blocks["norm"] = std::make_shared<LayerNorm>(embedding_dim, 1e-6f, false);
|
||||
blocks["linear_2"] = std::make_shared<Linear>(embedding_dim, out_dim, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning_embedding) {
|
||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
|
||||
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
|
||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
|
||||
|
||||
auto emb = linear_1->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning_embedding));
|
||||
x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), emb);
|
||||
x = linear_2->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct Attention : public GGMLBlock {
|
||||
int64_t dim_head;
|
||||
int64_t heads;
|
||||
int64_t kv_heads;
|
||||
|
||||
Attention(int64_t query_dim, int64_t dim_head, int64_t heads, int64_t kv_heads, float eps = 1e-5f)
|
||||
: dim_head(dim_head), heads(heads), kv_heads(kv_heads) {
|
||||
blocks["to_q"] = std::make_shared<Linear>(query_dim, heads * dim_head, false);
|
||||
blocks["to_k"] = std::make_shared<Linear>(query_dim, kv_heads * dim_head, false);
|
||||
blocks["to_v"] = std::make_shared<Linear>(query_dim, kv_heads * dim_head, false);
|
||||
blocks["norm_q"] = std::make_shared<RMSNorm>(dim_head, eps);
|
||||
blocks["norm_k"] = std::make_shared<RMSNorm>(dim_head, eps);
|
||||
blocks["to_out.0"] = std::make_shared<Linear>(heads * dim_head, query_dim, false);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* hidden_states,
|
||||
ggml_tensor* encoder_hidden_states,
|
||||
ggml_tensor* rotary_emb,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
|
||||
auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
|
||||
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
|
||||
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
int64_t N = hidden_states->ne[2];
|
||||
int64_t Lq = hidden_states->ne[1];
|
||||
int64_t Lk = encoder_hidden_states->ne[1];
|
||||
|
||||
auto q = to_q->forward(ctx, hidden_states);
|
||||
q = ggml_reshape_4d(ctx->ggml_ctx, q, dim_head, heads, Lq, N);
|
||||
auto k = to_k->forward(ctx, encoder_hidden_states);
|
||||
k = ggml_reshape_4d(ctx->ggml_ctx, k, dim_head, kv_heads, Lk, N);
|
||||
auto v = to_v->forward(ctx, encoder_hidden_states);
|
||||
v = ggml_reshape_4d(ctx->ggml_ctx, v, dim_head, kv_heads, Lk, N);
|
||||
|
||||
q = norm_q->forward(ctx, q);
|
||||
k = norm_k->forward(ctx, k);
|
||||
|
||||
auto out = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask);
|
||||
out = to_out_0->forward(ctx, out);
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct BooguImageTransformerBlock : public GGMLBlock {
|
||||
bool modulation;
|
||||
|
||||
BooguImageTransformerBlock(int64_t dim,
|
||||
int64_t num_attention_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t multiple_of,
|
||||
float norm_eps,
|
||||
bool modulation)
|
||||
: modulation(modulation) {
|
||||
int64_t head_dim = dim / num_attention_heads;
|
||||
blocks["attn"] = std::make_shared<Attention>(dim, head_dim, num_attention_heads, num_kv_heads, 1e-5f);
|
||||
blocks["feed_forward"] = std::make_shared<LuminaFeedForward>(dim, 4 * dim, multiple_of);
|
||||
if (modulation) {
|
||||
blocks["norm1"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
|
||||
} else {
|
||||
blocks["norm1"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
}
|
||||
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["norm2"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* hidden_states,
|
||||
ggml_tensor* rotary_emb,
|
||||
ggml_tensor* temb = nullptr,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
auto attn = std::dynamic_pointer_cast<Attention>(blocks["attn"]);
|
||||
auto feed_forward = std::dynamic_pointer_cast<LuminaFeedForward>(blocks["feed_forward"]);
|
||||
auto ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
|
||||
auto ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
|
||||
|
||||
if (modulation) {
|
||||
auto norm1 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["norm1"]);
|
||||
auto mods = norm1->forward(ctx, hidden_states, temb);
|
||||
|
||||
auto norm_hidden_states = std::get<0>(mods);
|
||||
auto gate_msa = std::get<1>(mods);
|
||||
auto scale_mlp = std::get<2>(mods);
|
||||
auto gate_mlp = std::get<3>(mods);
|
||||
|
||||
auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask);
|
||||
hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output), gate_msa);
|
||||
|
||||
auto mlp_input = scale_modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp);
|
||||
auto mlp_output = feed_forward->forward(ctx, mlp_input);
|
||||
hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output), gate_mlp);
|
||||
} else {
|
||||
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
|
||||
|
||||
auto norm_hidden_states = norm1->forward(ctx, hidden_states);
|
||||
auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask);
|
||||
hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output));
|
||||
|
||||
auto mlp_output = feed_forward->forward(ctx, ffn_norm1->forward(ctx, hidden_states));
|
||||
hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output));
|
||||
}
|
||||
return hidden_states;
|
||||
}
|
||||
};
|
||||
|
||||
struct BooguImageJointAttention : public GGMLBlock {
|
||||
int64_t dim_head;
|
||||
int64_t heads;
|
||||
int64_t kv_heads;
|
||||
|
||||
BooguImageJointAttention(int64_t dim, int64_t dim_head, int64_t heads, int64_t kv_heads)
|
||||
: dim_head(dim_head), heads(heads), kv_heads(kv_heads) {
|
||||
blocks["norm_q"] = std::make_shared<RMSNorm>(dim_head, 1e-5f);
|
||||
blocks["norm_k"] = std::make_shared<RMSNorm>(dim_head, 1e-5f);
|
||||
blocks["to_out.0"] = std::make_shared<Linear>(heads * dim_head, dim, false);
|
||||
blocks["processor.img_to_q"] = std::make_shared<Linear>(dim, heads * dim_head, false);
|
||||
blocks["processor.img_to_k"] = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
|
||||
blocks["processor.img_to_v"] = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
|
||||
blocks["processor.instruct_to_q"] = std::make_shared<Linear>(dim, heads * dim_head, false);
|
||||
blocks["processor.instruct_to_k"] = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
|
||||
blocks["processor.instruct_to_v"] = std::make_shared<Linear>(dim, kv_heads * dim_head, false);
|
||||
blocks["processor.instruct_out"] = std::make_shared<Linear>(heads * dim_head, dim, false);
|
||||
blocks["processor.img_out"] = std::make_shared<Linear>(heads * dim_head, dim, false);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* img_hidden_states,
|
||||
ggml_tensor* instruct_hidden_states,
|
||||
ggml_tensor* rotary_emb,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
auto img_to_q = std::dynamic_pointer_cast<Linear>(blocks["processor.img_to_q"]);
|
||||
auto img_to_k = std::dynamic_pointer_cast<Linear>(blocks["processor.img_to_k"]);
|
||||
auto img_to_v = std::dynamic_pointer_cast<Linear>(blocks["processor.img_to_v"]);
|
||||
auto instruct_to_q = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_to_q"]);
|
||||
auto instruct_to_k = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_to_k"]);
|
||||
auto instruct_to_v = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_to_v"]);
|
||||
auto instruct_out = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
|
||||
auto img_out = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
int64_t N = img_hidden_states->ne[2];
|
||||
int64_t L_img = img_hidden_states->ne[1];
|
||||
int64_t L_instruct = instruct_hidden_states->ne[1];
|
||||
|
||||
auto img_q = img_to_q->forward(ctx, img_hidden_states);
|
||||
img_q = ggml_reshape_4d(ctx->ggml_ctx, img_q, dim_head, heads, L_img, N);
|
||||
auto img_k = img_to_k->forward(ctx, img_hidden_states);
|
||||
img_k = ggml_reshape_4d(ctx->ggml_ctx, img_k, dim_head, kv_heads, L_img, N);
|
||||
auto img_v = img_to_v->forward(ctx, img_hidden_states);
|
||||
img_v = ggml_reshape_4d(ctx->ggml_ctx, img_v, dim_head, kv_heads, L_img, N);
|
||||
|
||||
auto instruct_q = instruct_to_q->forward(ctx, instruct_hidden_states);
|
||||
instruct_q = ggml_reshape_4d(ctx->ggml_ctx, instruct_q, dim_head, heads, L_instruct, N);
|
||||
auto instruct_k = instruct_to_k->forward(ctx, instruct_hidden_states);
|
||||
instruct_k = ggml_reshape_4d(ctx->ggml_ctx, instruct_k, dim_head, kv_heads, L_instruct, N);
|
||||
auto instruct_v = instruct_to_v->forward(ctx, instruct_hidden_states);
|
||||
instruct_v = ggml_reshape_4d(ctx->ggml_ctx, instruct_v, dim_head, kv_heads, L_instruct, N);
|
||||
|
||||
auto q = ggml_concat(ctx->ggml_ctx, instruct_q, img_q, 2);
|
||||
auto k = ggml_concat(ctx->ggml_ctx, instruct_k, img_k, 2);
|
||||
auto v = ggml_concat(ctx->ggml_ctx, instruct_v, img_v, 2);
|
||||
q = norm_q->forward(ctx, q);
|
||||
k = norm_k->forward(ctx, k);
|
||||
|
||||
auto hidden_states = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask);
|
||||
auto instruct_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, L_instruct);
|
||||
auto img_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, L_instruct, L_instruct + L_img);
|
||||
|
||||
instruct_attn = instruct_out->forward(ctx, instruct_attn);
|
||||
img_attn = img_out->forward(ctx, img_attn);
|
||||
hidden_states = ggml_concat(ctx->ggml_ctx, instruct_attn, img_attn, 1);
|
||||
hidden_states = to_out_0->forward(ctx, hidden_states);
|
||||
return hidden_states;
|
||||
}
|
||||
};
|
||||
|
||||
struct BooguImageDoubleStreamBlock : public GGMLBlock {
|
||||
BooguImageDoubleStreamBlock(int64_t dim,
|
||||
int64_t num_attention_heads,
|
||||
int64_t num_kv_heads,
|
||||
int64_t multiple_of,
|
||||
float norm_eps) {
|
||||
int64_t head_dim = dim / num_attention_heads;
|
||||
blocks["img_instruct_attn"] = std::make_shared<BooguImageJointAttention>(dim, head_dim, num_attention_heads, num_kv_heads);
|
||||
blocks["img_self_attn"] = std::make_shared<Attention>(dim, head_dim, num_attention_heads, num_kv_heads, 1e-5f);
|
||||
blocks["img_feed_forward"] = std::make_shared<LuminaFeedForward>(dim, 4 * dim, multiple_of);
|
||||
blocks["instruct_feed_forward"] = std::make_shared<LuminaFeedForward>(dim, 4 * dim, multiple_of);
|
||||
blocks["img_norm1"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
|
||||
blocks["img_norm2"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
|
||||
blocks["img_norm3"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
|
||||
blocks["instruct_norm1"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
|
||||
blocks["instruct_norm2"] = std::make_shared<LuminaRMSNormZero>(dim, std::min<int64_t>(dim, 1024), norm_eps);
|
||||
blocks["img_attn_norm"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["img_self_attn_norm"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["img_ffn_norm1"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["img_ffn_norm2"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["instruct_attn_norm"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["instruct_ffn_norm1"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
blocks["instruct_ffn_norm2"] = std::make_shared<RMSNorm>(dim, norm_eps);
|
||||
}
|
||||
|
||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* img_hidden_states,
|
||||
ggml_tensor* instruct_hidden_states,
|
||||
ggml_tensor* joint_rotary_emb,
|
||||
ggml_tensor* img_rotary_emb,
|
||||
ggml_tensor* temb) {
|
||||
auto img_instruct_attn = std::dynamic_pointer_cast<BooguImageJointAttention>(blocks["img_instruct_attn"]);
|
||||
auto img_self_attn = std::dynamic_pointer_cast<Attention>(blocks["img_self_attn"]);
|
||||
auto img_feed_forward = std::dynamic_pointer_cast<LuminaFeedForward>(blocks["img_feed_forward"]);
|
||||
auto instruct_feed_forward = std::dynamic_pointer_cast<LuminaFeedForward>(blocks["instruct_feed_forward"]);
|
||||
auto img_norm1 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["img_norm1"]);
|
||||
auto img_norm2 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["img_norm2"]);
|
||||
auto img_norm3 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["img_norm3"]);
|
||||
auto instruct_norm1 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["instruct_norm1"]);
|
||||
auto instruct_norm2 = std::dynamic_pointer_cast<LuminaRMSNormZero>(blocks["instruct_norm2"]);
|
||||
auto img_attn_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["img_attn_norm"]);
|
||||
auto img_self_attn_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["img_self_attn_norm"]);
|
||||
auto img_ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_ffn_norm1"]);
|
||||
auto img_ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_ffn_norm2"]);
|
||||
auto instruct_attn_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["instruct_attn_norm"]);
|
||||
auto instruct_ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["instruct_ffn_norm1"]);
|
||||
auto instruct_ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["instruct_ffn_norm2"]);
|
||||
|
||||
int64_t L_instruct = instruct_hidden_states->ne[1];
|
||||
|
||||
auto img_norm1_out_vec = img_norm1->forward(ctx, img_hidden_states, temb);
|
||||
auto img_norm2_out_vec = img_norm2->forward(ctx, img_hidden_states, temb);
|
||||
auto img_norm3_out_vec = img_norm3->forward(ctx, img_hidden_states, temb);
|
||||
auto instruct_norm1_out_vec = instruct_norm1->forward(ctx, instruct_hidden_states, temb);
|
||||
auto instruct_norm2_out_vec = instruct_norm2->forward(ctx, instruct_hidden_states, temb);
|
||||
|
||||
auto img_norm1_out = std::get<0>(img_norm1_out_vec);
|
||||
auto img_gate_msa = std::get<1>(img_norm1_out_vec);
|
||||
auto img_scale_mlp = std::get<2>(img_norm1_out_vec);
|
||||
auto img_gate_mlp = std::get<3>(img_norm1_out_vec);
|
||||
|
||||
auto img_norm2_out = std::get<0>(img_norm2_out_vec);
|
||||
auto img_shift_mlp = std::get<1>(img_norm2_out_vec);
|
||||
|
||||
auto img_norm3_out = std::get<0>(img_norm3_out_vec);
|
||||
auto img_gate_self = std::get<1>(img_norm3_out_vec);
|
||||
|
||||
auto instruct_norm1_out = std::get<0>(instruct_norm1_out_vec);
|
||||
auto instruct_gate_msa = std::get<1>(instruct_norm1_out_vec);
|
||||
auto instruct_scale_mlp = std::get<2>(instruct_norm1_out_vec);
|
||||
auto instruct_gate_mlp = std::get<3>(instruct_norm1_out_vec);
|
||||
|
||||
auto instruct_norm2_out = std::get<0>(instruct_norm2_out_vec);
|
||||
auto instruct_shift_mlp = std::get<1>(instruct_norm2_out_vec);
|
||||
|
||||
auto joint_attn_out = img_instruct_attn->forward(ctx, img_norm1_out, instruct_norm1_out, joint_rotary_emb);
|
||||
auto instruct_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, 0, L_instruct);
|
||||
auto img_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]);
|
||||
|
||||
auto img_self_attn_out = img_self_attn->forward(ctx, img_norm3_out, img_norm3_out, img_rotary_emb);
|
||||
|
||||
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_attn_norm->forward(ctx, img_attn_out), img_gate_msa);
|
||||
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_self_attn_norm->forward(ctx, img_self_attn_out), img_gate_self);
|
||||
|
||||
auto img_mlp_input = scale_modulate(ctx->ggml_ctx, img_norm2_out, img_scale_mlp);
|
||||
img_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, img_shift_mlp, img_shift_mlp->ne[0], 1, img_shift_mlp->ne[1]);
|
||||
img_mlp_input = ggml_add(ctx->ggml_ctx, img_mlp_input, img_shift_mlp);
|
||||
auto img_mlp_out = img_feed_forward->forward(ctx, img_ffn_norm1->forward(ctx, img_mlp_input));
|
||||
img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_ffn_norm2->forward(ctx, img_mlp_out), img_gate_mlp);
|
||||
|
||||
instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_attn_norm->forward(ctx, instruct_attn_out), instruct_gate_msa);
|
||||
auto instruct_mlp_input = scale_modulate(ctx->ggml_ctx, instruct_norm2_out, instruct_scale_mlp);
|
||||
instruct_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, instruct_shift_mlp, instruct_shift_mlp->ne[0], 1, instruct_shift_mlp->ne[1]);
|
||||
instruct_mlp_input = ggml_add(ctx->ggml_ctx, instruct_mlp_input, instruct_shift_mlp);
|
||||
auto instruct_mlp_out = instruct_feed_forward->forward(ctx, instruct_ffn_norm1->forward(ctx, instruct_mlp_input));
|
||||
instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_ffn_norm2->forward(ctx, instruct_mlp_out), instruct_gate_mlp);
|
||||
|
||||
return {img_hidden_states, instruct_hidden_states};
|
||||
}
|
||||
};
|
||||
|
||||
struct BooguImageModel : public GGMLBlock {
|
||||
BooguConfig config;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
GGML_UNUSED(tensor_storage_map);
|
||||
GGML_UNUSED(prefix);
|
||||
params["image_index_embedding"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, 5);
|
||||
}
|
||||
|
||||
BooguImageModel() = default;
|
||||
BooguImageModel(BooguConfig config)
|
||||
: config(std::move(config)) {
|
||||
blocks["x_embedder"] = std::make_shared<Linear>(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true);
|
||||
blocks["ref_image_patch_embedder"] = std::make_shared<Linear>(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true);
|
||||
blocks["time_caption_embed"] = std::make_shared<LuminaCombinedTimestepCaptionEmbedding>(this->config.hidden_size,
|
||||
this->config.instruction_feat_dim,
|
||||
256,
|
||||
this->config.norm_eps,
|
||||
this->config.timestep_scale);
|
||||
|
||||
for (int i = 0; i < this->config.num_refiner_layers; i++) {
|
||||
blocks["noise_refiner." + std::to_string(i)] = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
|
||||
this->config.num_attention_heads,
|
||||
this->config.num_kv_heads,
|
||||
this->config.multiple_of,
|
||||
this->config.norm_eps,
|
||||
true);
|
||||
blocks["ref_image_refiner." + std::to_string(i)] = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
|
||||
this->config.num_attention_heads,
|
||||
this->config.num_kv_heads,
|
||||
this->config.multiple_of,
|
||||
this->config.norm_eps,
|
||||
true);
|
||||
blocks["context_refiner." + std::to_string(i)] = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
|
||||
this->config.num_attention_heads,
|
||||
this->config.num_kv_heads,
|
||||
this->config.multiple_of,
|
||||
this->config.norm_eps,
|
||||
false);
|
||||
}
|
||||
|
||||
for (int i = 0; i < this->config.num_double_stream_layers; i++) {
|
||||
blocks["double_stream_layers." + std::to_string(i)] = std::make_shared<BooguImageDoubleStreamBlock>(this->config.hidden_size,
|
||||
this->config.num_attention_heads,
|
||||
this->config.num_kv_heads,
|
||||
this->config.multiple_of,
|
||||
this->config.norm_eps);
|
||||
}
|
||||
|
||||
for (int i = 0; i < this->config.num_layers; i++) {
|
||||
blocks["single_stream_layers." + std::to_string(i)] = std::make_shared<BooguImageTransformerBlock>(this->config.hidden_size,
|
||||
this->config.num_attention_heads,
|
||||
this->config.num_kv_heads,
|
||||
this->config.multiple_of,
|
||||
this->config.norm_eps,
|
||||
true);
|
||||
}
|
||||
|
||||
blocks["norm_out"] = std::make_shared<LuminaLayerNormContinuous>(this->config.hidden_size,
|
||||
this->config.timestep_embed_dim,
|
||||
this->config.patch_size * this->config.patch_size * this->config.out_channels);
|
||||
}
|
||||
|
||||
ggml_tensor* image_index_embedding(GGMLRunnerContext* ctx, int index) {
|
||||
GGML_ASSERT(index >= 0 && index < 5);
|
||||
auto embedding = params["image_index_embedding"];
|
||||
auto out = ggml_view_1d(ctx->ggml_ctx,
|
||||
embedding,
|
||||
config.hidden_size,
|
||||
index * config.hidden_size * ggml_element_size(embedding));
|
||||
out = ggml_reshape_3d(ctx->ggml_ctx, out, config.hidden_size, 1, 1);
|
||||
return out;
|
||||
}
|
||||
|
||||
ggml_tensor* embed_refs(GGMLRunnerContext* ctx, const std::vector<ggml_tensor*>& ref_latents) {
|
||||
if (ref_latents.empty()) {
|
||||
return nullptr;
|
||||
}
|
||||
auto ref_image_patch_embedder = std::dynamic_pointer_cast<Linear>(blocks["ref_image_patch_embedder"]);
|
||||
|
||||
ggml_tensor* ref_img = nullptr;
|
||||
for (int i = 0; i < static_cast<int>(ref_latents.size()); i++) {
|
||||
auto ref = DiT::pad_and_patchify(ctx, ref_latents[i], config.patch_size, config.patch_size, false);
|
||||
ref = ref_image_patch_embedder->forward(ctx, ref);
|
||||
ref = ggml_add(ctx->ggml_ctx, ref, image_index_embedding(ctx, std::min(i, 4)));
|
||||
ref_img = ref_img == nullptr ? ref : ggml_concat(ctx->ggml_ctx, ref_img, ref, 1);
|
||||
}
|
||||
return ref_img;
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* timesteps,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe,
|
||||
std::vector<ggml_tensor*> ref_latents = {}) {
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t N = x->ne[3];
|
||||
GGML_ASSERT(N == 1);
|
||||
|
||||
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
|
||||
auto time_caption_embed = std::dynamic_pointer_cast<LuminaCombinedTimestepCaptionEmbedding>(blocks["time_caption_embed"]);
|
||||
auto norm_out = std::dynamic_pointer_cast<LuminaLayerNormContinuous>(blocks["norm_out"]);
|
||||
|
||||
auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, timesteps), timesteps);
|
||||
auto embeds = time_caption_embed->forward(ctx, timestep, context);
|
||||
auto temb = embeds.first;
|
||||
auto txt = embeds.second;
|
||||
|
||||
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, false);
|
||||
int64_t img_len = img->ne[1];
|
||||
img = x_embedder->forward(ctx, img);
|
||||
auto ref_img = embed_refs(ctx, ref_latents);
|
||||
int64_t ref_len = ref_img != nullptr ? ref_img->ne[1] : 0;
|
||||
int64_t txt_len = txt->ne[1];
|
||||
|
||||
GGML_ASSERT(pe->ne[3] == txt_len + ref_len + img_len);
|
||||
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt_len);
|
||||
auto noise_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len);
|
||||
|
||||
for (int i = 0; i < config.num_refiner_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
|
||||
txt = block->forward(ctx, txt, txt_pe);
|
||||
sd::ggml_graph_cut::mark_graph_cut(txt, "boogu.context_refiner." + std::to_string(i), "txt");
|
||||
}
|
||||
|
||||
for (int i = 0; i < config.num_refiner_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
|
||||
img = block->forward(ctx, img, noise_pe, temb);
|
||||
sd::ggml_graph_cut::mark_graph_cut(img, "boogu.noise_refiner." + std::to_string(i), "img");
|
||||
}
|
||||
|
||||
ggml_tensor* combined_img = img;
|
||||
if (ref_img != nullptr) {
|
||||
auto ref_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + ref_len);
|
||||
for (int i = 0; i < config.num_refiner_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["ref_image_refiner." + std::to_string(i)]);
|
||||
ref_img = block->forward(ctx, ref_img, ref_pe, temb);
|
||||
sd::ggml_graph_cut::mark_graph_cut(ref_img, "boogu.ref_image_refiner." + std::to_string(i), "ref_img");
|
||||
}
|
||||
combined_img = ggml_concat(ctx->ggml_ctx, ref_img, img, 1);
|
||||
}
|
||||
|
||||
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]);
|
||||
for (int i = 0; i < config.num_double_stream_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<BooguImageDoubleStreamBlock>(blocks["double_stream_layers." + std::to_string(i)]);
|
||||
auto result = block->forward(ctx, combined_img, txt, pe, img_pe, temb);
|
||||
combined_img = result.first;
|
||||
txt = result.second;
|
||||
sd::ggml_graph_cut::mark_graph_cut(combined_img, "boogu.double_stream_layers." + std::to_string(i), "img");
|
||||
sd::ggml_graph_cut::mark_graph_cut(txt, "boogu.double_stream_layers." + std::to_string(i), "txt");
|
||||
}
|
||||
|
||||
auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, combined_img, 1);
|
||||
for (int i = 0; i < config.num_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<BooguImageTransformerBlock>(blocks["single_stream_layers." + std::to_string(i)]);
|
||||
hidden_states = block->forward(ctx, hidden_states, pe, temb);
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "boogu.single_stream_layers." + std::to_string(i), "hidden_states");
|
||||
}
|
||||
|
||||
hidden_states = norm_out->forward(ctx, hidden_states, temb);
|
||||
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]);
|
||||
hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, false);
|
||||
hidden_states = ggml_ext_scale(ctx->ggml_ctx, hidden_states, -1.f);
|
||||
return hidden_states;
|
||||
}
|
||||
};
|
||||
|
||||
__STATIC_INLINE__ int patched_token_count(int64_t size, int patch_size) {
|
||||
int pad = (patch_size - (static_cast<int>(size) % patch_size)) % patch_size;
|
||||
return (static_cast<int>(size) + pad) / patch_size;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ void append_spatial_ids(std::vector<std::vector<float>>& ids,
|
||||
int bs,
|
||||
int pe_shift,
|
||||
int h_tokens,
|
||||
int w_tokens) {
|
||||
std::vector<std::vector<float>> image_ids(h_tokens * w_tokens, std::vector<float>(3, 0.0f));
|
||||
for (int h = 0; h < h_tokens; h++) {
|
||||
for (int w = 0; w < w_tokens; w++) {
|
||||
image_ids[h * w_tokens + w][0] = static_cast<float>(pe_shift);
|
||||
image_ids[h * w_tokens + w][1] = static_cast<float>(h);
|
||||
image_ids[h * w_tokens + w][2] = static_cast<float>(w);
|
||||
}
|
||||
}
|
||||
for (int b = 0; b < bs; b++) {
|
||||
ids.insert(ids.end(), image_ids.begin(), image_ids.end());
|
||||
}
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids;
|
||||
ids.reserve(static_cast<size_t>(bs) * context_len);
|
||||
for (int b = 0; b < bs; b++) {
|
||||
for (int i = 0; i < context_len; i++) {
|
||||
float pos = static_cast<float>(i);
|
||||
ids.push_back({pos, pos, pos});
|
||||
}
|
||||
}
|
||||
|
||||
int pe_shift = context_len;
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
|
||||
int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
|
||||
append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
|
||||
pe_shift += std::max(ref_h_tokens, ref_w_tokens);
|
||||
}
|
||||
|
||||
int h_tokens = patched_token_count(h, patch_size);
|
||||
int w_tokens = patched_token_count(w, patch_size);
|
||||
append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
|
||||
|
||||
return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
struct BooguImageRunner : public DiffusionModelRunner {
|
||||
BooguConfig config;
|
||||
BooguImageModel boogu;
|
||||
std::vector<float> pe_vec;
|
||||
|
||||
BooguImageRunner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "",
|
||||
SDVersion version = VERSION_BOOGU_IMAGE,
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
config(BooguConfig::detect_from_weights(tensor_storage_map, prefix)) {
|
||||
boogu = BooguImageModel(config);
|
||||
boogu.init(params_ctx, tensor_storage_map, prefix);
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "boogu_image";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
|
||||
boogu.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
const sd::Tensor<float>& timesteps_tensor,
|
||||
const sd::Tensor<float>& context_tensor,
|
||||
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
|
||||
ggml_cgraph* gf = new_graph_custom(BOOGU_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
GGML_ASSERT(!context_tensor.empty());
|
||||
ggml_tensor* context = make_input(context_tensor);
|
||||
|
||||
std::vector<ggml_tensor*> ref_latents;
|
||||
ref_latents.reserve(ref_latents_tensor.size());
|
||||
for (const auto& ref_latent_tensor : ref_latents_tensor) {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
|
||||
pe_vec = gen_boogu_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
ref_latents,
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
const sd::Tensor<float>& timesteps,
|
||||
const sd::Tensor<float>& context,
|
||||
const std::vector<sd::Tensor<float>>& ref_latents = {}) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context, ref_latents);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const DiffusionParams& diffusion_params) override {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
static const std::vector<sd::Tensor<float>> empty_ref_latents;
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context),
|
||||
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
|
||||
}
|
||||
};
|
||||
} // namespace Boogu
|
||||
|
||||
#endif // __SD_MODEL_DIFFUSION_BOOGU_HPP__
|
||||
Reference in New Issue
Block a user