mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 11:56:36 +00:00
refactor: reorganize src model layout (#1615)
This commit is contained in:
leejet authored and GitHub committed
1 parent
dfb2390dd4
commit
f3fd359b58
81 files changed
+407
-385
No files matched your search
@@ -0,0 +1,637 @@
|
||||
#ifndef __SD_MODEL_COMMON_BLOCK_HPP__
|
||||
#define __SD_MODEL_COMMON_BLOCK_HPP__
|
||||
|
||||
#include "core/ggml_extend.hpp"
|
||||
#include "core/util.h"
|
||||
#include "ggml-backend.h"
|
||||
|
||||
class DownSampleBlock : public GGMLBlock {
|
||||
protected:
|
||||
int channels;
|
||||
int out_channels;
|
||||
bool vae_downsample;
|
||||
|
||||
public:
|
||||
DownSampleBlock(int channels,
|
||||
int out_channels,
|
||||
bool vae_downsample = false)
|
||||
: channels(channels),
|
||||
out_channels(out_channels),
|
||||
vae_downsample(vae_downsample) {
|
||||
if (vae_downsample) {
|
||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {2, 2}, {0, 0}));
|
||||
} else {
|
||||
blocks["op"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {2, 2}, {1, 1}));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, channels, h, w]
|
||||
if (vae_downsample) {
|
||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||
|
||||
x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
|
||||
x = conv->forward(ctx, x);
|
||||
} else {
|
||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["op"]);
|
||||
|
||||
x = conv->forward(ctx, x);
|
||||
}
|
||||
return x; // [N, out_channels, h/2, w/2]
|
||||
}
|
||||
};
|
||||
|
||||
class UpSampleBlock : public GGMLBlock {
|
||||
protected:
|
||||
int channels;
|
||||
int out_channels;
|
||||
|
||||
public:
|
||||
UpSampleBlock(int channels,
|
||||
int out_channels)
|
||||
: channels(channels),
|
||||
out_channels(out_channels) {
|
||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, channels, h, w]
|
||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||
|
||||
x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST); // [N, channels, h*2, w*2]
|
||||
x = conv->forward(ctx, x); // [N, out_channels, h*2, w*2]
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class ResBlock : public GGMLBlock {
|
||||
protected:
|
||||
// network hparams
|
||||
int64_t channels; // model_channels * (1, 1, 1, 2, 2, 4, 4, 4)
|
||||
int64_t emb_channels; // time_embed_dim
|
||||
int64_t out_channels; // mult * model_channels
|
||||
std::pair<int, int> kernel_size;
|
||||
int dims;
|
||||
bool skip_t_emb;
|
||||
bool exchange_temb_dims;
|
||||
|
||||
std::shared_ptr<GGMLBlock> conv_nd(int dims,
|
||||
int64_t in_channels,
|
||||
int64_t out_channels,
|
||||
std::pair<int, int> kernel_size,
|
||||
std::pair<int, int> padding) {
|
||||
GGML_ASSERT(dims == 2 || dims == 3);
|
||||
if (dims == 3) {
|
||||
return std::shared_ptr<GGMLBlock>(new Conv3d(in_channels, out_channels, {kernel_size.first, 1, 1}, {1, 1, 1}, {padding.first, 0, 0}));
|
||||
} else {
|
||||
return std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, out_channels, kernel_size, {1, 1}, padding));
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
ResBlock(int64_t channels,
|
||||
int64_t emb_channels,
|
||||
int64_t out_channels,
|
||||
std::pair<int, int> kernel_size = {3, 3},
|
||||
int dims = 2,
|
||||
bool exchange_temb_dims = false,
|
||||
bool skip_t_emb = false)
|
||||
: channels(channels),
|
||||
emb_channels(emb_channels),
|
||||
out_channels(out_channels),
|
||||
kernel_size(kernel_size),
|
||||
dims(dims),
|
||||
skip_t_emb(skip_t_emb),
|
||||
exchange_temb_dims(exchange_temb_dims) {
|
||||
std::pair<int, int> padding = {kernel_size.first / 2, kernel_size.second / 2};
|
||||
blocks["in_layers.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(channels));
|
||||
// in_layer_1 is nn.SILU()
|
||||
blocks["in_layers.2"] = conv_nd(dims, channels, out_channels, kernel_size, padding);
|
||||
|
||||
if (!skip_t_emb) {
|
||||
// emb_layer_0 is nn.SILU()
|
||||
blocks["emb_layers.1"] = std::shared_ptr<GGMLBlock>(new Linear(emb_channels, out_channels));
|
||||
}
|
||||
|
||||
blocks["out_layers.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(out_channels));
|
||||
// out_layer_1 is nn.SILU()
|
||||
// out_layer_2 is nn.Dropout(), skip for inference
|
||||
blocks["out_layers.3"] = conv_nd(dims, out_channels, out_channels, kernel_size, padding);
|
||||
|
||||
if (out_channels != channels) {
|
||||
blocks["skip_connection"] = conv_nd(dims, channels, out_channels, {1, 1}, {0, 0});
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
// [N, c, t, h, w] => [N, c, t, h * w]
|
||||
// x: [N, channels, h, w] if dims == 2 else [N, channels, t, h, w]
|
||||
// emb: [N, emb_channels] if dims == 2 else [N, t, emb_channels]
|
||||
auto in_layers_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["in_layers.0"]);
|
||||
auto in_layers_2 = std::dynamic_pointer_cast<UnaryBlock>(blocks["in_layers.2"]);
|
||||
auto out_layers_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["out_layers.0"]);
|
||||
auto out_layers_3 = std::dynamic_pointer_cast<UnaryBlock>(blocks["out_layers.3"]);
|
||||
|
||||
if (emb == nullptr) {
|
||||
GGML_ASSERT(skip_t_emb);
|
||||
}
|
||||
|
||||
// in_layers
|
||||
auto h = in_layers_0->forward(ctx, x);
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, h);
|
||||
h = in_layers_2->forward(ctx, h); // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
|
||||
|
||||
// emb_layers
|
||||
if (!skip_t_emb) {
|
||||
auto emb_layer_1 = std::dynamic_pointer_cast<Linear>(blocks["emb_layers.1"]);
|
||||
|
||||
auto emb_out = ggml_silu(ctx->ggml_ctx, emb);
|
||||
emb_out = emb_layer_1->forward(ctx, emb_out); // [N, out_channels] if dims == 2 else [N, t, out_channels]
|
||||
|
||||
if (dims == 2) {
|
||||
emb_out = ggml_reshape_4d(ctx->ggml_ctx, emb_out, 1, 1, emb_out->ne[0], emb_out->ne[1]); // [N, out_channels, 1, 1]
|
||||
} else {
|
||||
emb_out = ggml_reshape_4d(ctx->ggml_ctx, emb_out, 1, emb_out->ne[0], emb_out->ne[1], emb_out->ne[2]); // [N, t, out_channels, 1]
|
||||
if (exchange_temb_dims) {
|
||||
// emb_out = rearrange(emb_out, "b t c ... -> b c t ...")
|
||||
emb_out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, emb_out, 0, 2, 1, 3)); // [N, out_channels, t, 1]
|
||||
}
|
||||
}
|
||||
|
||||
h = ggml_add(ctx->ggml_ctx, h, emb_out); // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
|
||||
}
|
||||
|
||||
// out_layers
|
||||
h = out_layers_0->forward(ctx, h);
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, h);
|
||||
// dropout, skip for inference
|
||||
h = out_layers_3->forward(ctx, h);
|
||||
|
||||
// skip connection
|
||||
if (out_channels != channels) {
|
||||
auto skip_connection = std::dynamic_pointer_cast<UnaryBlock>(blocks["skip_connection"]);
|
||||
x = skip_connection->forward(ctx, x); // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
|
||||
}
|
||||
|
||||
h = ggml_add(ctx->ggml_ctx, h, x);
|
||||
return h; // [N, out_channels, h, w] if dims == 2 else [N, out_channels, t, h, w]
|
||||
}
|
||||
};
|
||||
|
||||
class GEGLU : public UnaryBlock {
|
||||
protected:
|
||||
int64_t dim_in;
|
||||
int64_t dim_out;
|
||||
|
||||
public:
|
||||
GEGLU(int64_t dim_in, int64_t dim_out)
|
||||
: dim_in(dim_in), dim_out(dim_out) {
|
||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out * 2));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
// x: [ne3, ne2, ne1, dim_in]
|
||||
// return: [ne3, ne2, ne1, dim_out]
|
||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
||||
|
||||
x = proj->forward(ctx, x); // [ne3, ne2, ne1, dim_out*2]
|
||||
auto x_vec = ggml_ext_chunk(ctx->ggml_ctx, x, 2, 0, false);
|
||||
x = x_vec[0]; // [ne3, ne2, ne1, dim_out]
|
||||
auto gate = x_vec[1]; // [ne3, ne2, ne1, dim_out]
|
||||
|
||||
gate = ggml_cont(ctx->ggml_ctx, gate);
|
||||
|
||||
gate = ggml_ext_gelu(ctx->ggml_ctx, gate, true);
|
||||
|
||||
x = ggml_mul(ctx->ggml_ctx, x, gate); // [ne3, ne2, ne1, dim_out]
|
||||
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class GELU : public UnaryBlock {
|
||||
public:
|
||||
GELU(int64_t dim_in, int64_t dim_out, bool bias = true) {
|
||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim_in, dim_out, bias));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
// x: [ne3, ne2, ne1, dim_in]
|
||||
// return: [ne3, ne2, ne1, dim_out]
|
||||
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
|
||||
|
||||
x = proj->forward(ctx, x);
|
||||
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct Mlp : public GGMLBlock {
|
||||
public:
|
||||
Mlp(int64_t in_features,
|
||||
int64_t hidden_features = -1,
|
||||
int64_t out_features = -1,
|
||||
bool bias = true) {
|
||||
// act_layer is always lambda: nn.GELU(approximate="tanh")
|
||||
// norm_layer is always None
|
||||
// use_conv is always False
|
||||
if (hidden_features == -1) {
|
||||
hidden_features = in_features;
|
||||
}
|
||||
if (out_features == -1) {
|
||||
out_features = in_features;
|
||||
}
|
||||
blocks["fc1"] = std::shared_ptr<GGMLBlock>(new Linear(in_features, hidden_features, bias));
|
||||
blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, n_token, in_features]
|
||||
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
|
||||
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
|
||||
|
||||
x = fc1->forward(ctx, x);
|
||||
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
|
||||
x = fc2->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class FeedForward : public GGMLBlock {
|
||||
public:
|
||||
enum class Activation {
|
||||
GEGLU,
|
||||
GELU
|
||||
};
|
||||
FeedForward(int64_t dim,
|
||||
int64_t dim_out,
|
||||
int64_t mult = 4,
|
||||
Activation activation = Activation::GEGLU,
|
||||
bool precision_fix = false) {
|
||||
int64_t inner_dim = dim * mult;
|
||||
if (activation == Activation::GELU) {
|
||||
blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GELU(dim, inner_dim));
|
||||
} else {
|
||||
blocks["net.0"] = std::shared_ptr<GGMLBlock>(new GEGLU(dim, inner_dim));
|
||||
}
|
||||
|
||||
// net_1 is nn.Dropout(), skip for inference
|
||||
bool force_prec_f32 = false;
|
||||
float scale = 1.f;
|
||||
if (precision_fix) {
|
||||
scale = 1.f / 128.f;
|
||||
}
|
||||
// The purpose of the scale here is to prevent NaN issues in certain situations.
|
||||
// For example, when using Vulkan without enabling force_prec_f32,
|
||||
// or when using CUDA but the weights are k-quants.
|
||||
blocks["net.2"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim_out, true, false, force_prec_f32, scale));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [ne3, ne2, ne1, dim]
|
||||
// return: [ne3, ne2, ne1, dim_out]
|
||||
|
||||
auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
|
||||
auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
net_2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
x = net_0->forward(ctx, x); // [ne3, ne2, ne1, inner_dim]
|
||||
x = net_2->forward(ctx, x); // [ne3, ne2, ne1, dim_out]
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class CrossAttention : public GGMLBlock {
|
||||
protected:
|
||||
int64_t query_dim;
|
||||
int64_t context_dim;
|
||||
int64_t n_head;
|
||||
int64_t d_head;
|
||||
bool xtra_dim = false;
|
||||
|
||||
public:
|
||||
CrossAttention(int64_t query_dim,
|
||||
int64_t context_dim,
|
||||
int64_t n_head,
|
||||
int64_t d_head)
|
||||
: n_head(n_head),
|
||||
d_head(d_head),
|
||||
query_dim(query_dim),
|
||||
context_dim(context_dim) {
|
||||
int64_t inner_dim = d_head * n_head;
|
||||
if (context_dim == 320 && d_head == 320) {
|
||||
// LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
|
||||
xtra_dim = true;
|
||||
context_dim = 1024;
|
||||
}
|
||||
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, false));
|
||||
blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(context_dim, inner_dim, false));
|
||||
blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(context_dim, inner_dim, false));
|
||||
|
||||
blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, query_dim));
|
||||
// to_out_1 is nn.Dropout(), skip for inference
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* context) {
|
||||
// x: [N, n_token, query_dim]
|
||||
// context: [N, n_context, context_dim]
|
||||
// return: [N, n_token, query_dim]
|
||||
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 to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
|
||||
int64_t n = x->ne[2];
|
||||
int64_t n_token = x->ne[1];
|
||||
int64_t n_context = context->ne[1];
|
||||
int64_t inner_dim = d_head * n_head;
|
||||
|
||||
auto q = to_q->forward(ctx, x); // [N, n_token, inner_dim]
|
||||
if (xtra_dim) {
|
||||
// LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
|
||||
context->ne[0] = 1024; // patch dim
|
||||
}
|
||||
auto k = to_k->forward(ctx, context); // [N, n_context, inner_dim]
|
||||
auto v = to_v->forward(ctx, context); // [N, n_context, inner_dim]
|
||||
if (xtra_dim) {
|
||||
context->ne[0] = 320; // reset dim to orig
|
||||
}
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim]
|
||||
|
||||
x = to_out_0->forward(ctx, x); // [N, n_token, query_dim]
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class BasicTransformerBlock : public GGMLBlock {
|
||||
protected:
|
||||
int64_t n_head;
|
||||
int64_t d_head;
|
||||
bool ff_in;
|
||||
|
||||
public:
|
||||
BasicTransformerBlock(int64_t dim,
|
||||
int64_t n_head,
|
||||
int64_t d_head,
|
||||
int64_t context_dim,
|
||||
bool ff_in = false)
|
||||
: n_head(n_head), d_head(d_head), ff_in(ff_in) {
|
||||
// disable_self_attn is always False
|
||||
// disable_temporal_crossattention is always False
|
||||
// switch_temporal_ca_to_sa is always False
|
||||
// inner_dim is always None or equal to dim
|
||||
// gated_ff is always True
|
||||
blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
|
||||
blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head));
|
||||
blocks["ff"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
|
||||
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||
blocks["norm3"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||
|
||||
if (ff_in) {
|
||||
blocks["norm_in"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||
blocks["ff_in"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* context) {
|
||||
// x: [N, n_token, query_dim]
|
||||
// context: [N, n_context, context_dim]
|
||||
// return: [N, n_token, query_dim]
|
||||
|
||||
auto attn1 = std::dynamic_pointer_cast<CrossAttention>(blocks["attn1"]);
|
||||
auto attn2 = std::dynamic_pointer_cast<CrossAttention>(blocks["attn2"]);
|
||||
auto ff = std::dynamic_pointer_cast<FeedForward>(blocks["ff"]);
|
||||
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
|
||||
auto norm3 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm3"]);
|
||||
|
||||
if (ff_in) {
|
||||
auto norm_in = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_in"]);
|
||||
auto ff_in = std::dynamic_pointer_cast<FeedForward>(blocks["ff_in"]);
|
||||
|
||||
auto x_skip = x;
|
||||
x = norm_in->forward(ctx, x);
|
||||
x = ff_in->forward(ctx, x);
|
||||
// self.is_res is always True
|
||||
x = ggml_add(ctx->ggml_ctx, x, x_skip);
|
||||
}
|
||||
|
||||
auto r = x;
|
||||
x = norm1->forward(ctx, x);
|
||||
x = attn1->forward(ctx, x, x); // self-attention
|
||||
x = ggml_add(ctx->ggml_ctx, x, r);
|
||||
r = x;
|
||||
x = norm2->forward(ctx, x);
|
||||
x = attn2->forward(ctx, x, context); // cross-attention
|
||||
x = ggml_add(ctx->ggml_ctx, x, r);
|
||||
r = x;
|
||||
x = norm3->forward(ctx, x);
|
||||
x = ff->forward(ctx, x);
|
||||
x = ggml_add(ctx->ggml_ctx, x, r);
|
||||
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class SpatialTransformer : public GGMLBlock {
|
||||
protected:
|
||||
int64_t in_channels; // mult * model_channels
|
||||
int64_t n_head;
|
||||
int64_t d_head;
|
||||
int64_t depth = 1; // 1
|
||||
int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2
|
||||
bool use_linear = false;
|
||||
|
||||
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");
|
||||
if (iter != tensor_storage_map.end()) {
|
||||
int64_t inner_dim = n_head * d_head;
|
||||
if (iter->second.n_dims == 4 && use_linear) {
|
||||
use_linear = false;
|
||||
blocks["proj_in"] = std::make_shared<Conv2d>(in_channels, inner_dim, std::pair{1, 1});
|
||||
blocks["proj_out"] = std::make_shared<Conv2d>(inner_dim, in_channels, std::pair{1, 1});
|
||||
} else if (iter->second.n_dims == 2 && !use_linear) {
|
||||
use_linear = true;
|
||||
blocks["proj_in"] = std::make_shared<Linear>(in_channels, inner_dim);
|
||||
blocks["proj_out"] = std::make_shared<Linear>(inner_dim, in_channels);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
SpatialTransformer(int64_t in_channels,
|
||||
int64_t n_head,
|
||||
int64_t d_head,
|
||||
int64_t depth,
|
||||
int64_t context_dim,
|
||||
bool use_linear)
|
||||
: in_channels(in_channels),
|
||||
n_head(n_head),
|
||||
d_head(d_head),
|
||||
depth(depth),
|
||||
context_dim(context_dim),
|
||||
use_linear(use_linear) {
|
||||
// disable_self_attn is always False
|
||||
int64_t inner_dim = n_head * d_head; // in_channels
|
||||
blocks["norm"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(in_channels));
|
||||
if (use_linear) {
|
||||
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, inner_dim));
|
||||
} else {
|
||||
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, inner_dim, {1, 1}));
|
||||
}
|
||||
|
||||
for (int i = 0; i < depth; i++) {
|
||||
std::string name = "transformer_blocks." + std::to_string(i);
|
||||
blocks[name] = std::shared_ptr<GGMLBlock>(new BasicTransformerBlock(inner_dim, n_head, d_head, context_dim, false));
|
||||
}
|
||||
|
||||
if (use_linear) {
|
||||
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, in_channels));
|
||||
} else {
|
||||
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(inner_dim, in_channels, {1, 1}));
|
||||
}
|
||||
}
|
||||
|
||||
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* context) {
|
||||
// x: [N, in_channels, h, w]
|
||||
// context: [N, max_position(aka n_token), hidden_size(aka context_dim)]
|
||||
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
|
||||
auto proj_in = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj_in"]);
|
||||
auto proj_out = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj_out"]);
|
||||
|
||||
auto x_in = x;
|
||||
int64_t n = x->ne[3];
|
||||
int64_t h = x->ne[1];
|
||||
int64_t w = x->ne[0];
|
||||
int64_t inner_dim = n_head * d_head;
|
||||
|
||||
x = norm->forward(ctx, x);
|
||||
if (use_linear) {
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); // [N, h, w, inner_dim]
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, inner_dim, w * h, n); // [N, h * w, inner_dim]
|
||||
x = proj_in->forward(ctx, x); // [N, inner_dim, h, w]
|
||||
} else {
|
||||
x = proj_in->forward(ctx, x); // [N, inner_dim, h, w]
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); // [N, h, w, inner_dim]
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, inner_dim, w * h, n); // [N, h * w, inner_dim]
|
||||
}
|
||||
|
||||
for (int i = 0; i < depth; i++) {
|
||||
std::string name = "transformer_blocks." + std::to_string(i);
|
||||
auto transformer_block = std::dynamic_pointer_cast<BasicTransformerBlock>(blocks[name]);
|
||||
|
||||
x = transformer_block->forward(ctx, x, context);
|
||||
}
|
||||
|
||||
if (use_linear) {
|
||||
// proj_out
|
||||
x = proj_out->forward(ctx, x); // [N, in_channels, h, w]
|
||||
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, inner_dim, h * w]
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, inner_dim, n); // [N, inner_dim, h, w]
|
||||
} else {
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, inner_dim, h * w]
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, inner_dim, n); // [N, inner_dim, h, w]
|
||||
|
||||
// proj_out
|
||||
x = proj_out->forward(ctx, x); // [N, in_channels, h, w]
|
||||
}
|
||||
|
||||
x = ggml_add(ctx->ggml_ctx, x, x_in);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class AlphaBlender : public GGMLBlock {
|
||||
protected:
|
||||
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
|
||||
enum ggml_type wtype = GGML_TYPE_F32;
|
||||
params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
|
||||
}
|
||||
|
||||
float get_alpha() {
|
||||
// image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,]
|
||||
// so learned_with_images is same as learned
|
||||
float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]);
|
||||
return sigmoid(alpha);
|
||||
}
|
||||
|
||||
public:
|
||||
AlphaBlender() {
|
||||
// merge_strategy is always learned_with_images
|
||||
// for inference, we don't need to set alpha
|
||||
// since mix_factor.shape is [1,], we don't need rearrange using rearrange_pattern
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x_spatial,
|
||||
ggml_tensor* x_temporal) {
|
||||
// image_only_indicator is always tensor([0.])
|
||||
float alpha = get_alpha();
|
||||
auto x = ggml_add(ctx->ggml_ctx,
|
||||
ggml_ext_scale(ctx->ggml_ctx, x_spatial, alpha),
|
||||
ggml_ext_scale(ctx->ggml_ctx, x_temporal, 1.0f - alpha));
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class VideoResBlock : public ResBlock {
|
||||
public:
|
||||
VideoResBlock(int64_t channels,
|
||||
int64_t emb_channels,
|
||||
int64_t out_channels,
|
||||
std::pair<int, int> kernel_size = {3, 3},
|
||||
int64_t video_kernel_size = 3,
|
||||
int dims = 2) // always 2
|
||||
: ResBlock(channels, emb_channels, out_channels, kernel_size, dims) {
|
||||
blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, emb_channels, out_channels, kernel_size, 3, true));
|
||||
blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* emb,
|
||||
int num_video_frames) {
|
||||
// x: [N, channels, h, w] aka [b*t, channels, h, w]
|
||||
// emb: [N, emb_channels] aka [b*t, emb_channels]
|
||||
// image_only_indicator is always tensor([0.])
|
||||
auto time_stack = std::dynamic_pointer_cast<ResBlock>(blocks["time_stack"]);
|
||||
auto time_mixer = std::dynamic_pointer_cast<AlphaBlender>(blocks["time_mixer"]);
|
||||
|
||||
x = ResBlock::forward(ctx, x, emb);
|
||||
|
||||
int64_t T = num_video_frames;
|
||||
int64_t B = x->ne[3] / T;
|
||||
int64_t C = x->ne[2];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t W = x->ne[0];
|
||||
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, C, T, B); // (b t) c h w -> b t c (h w)
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // b t c (h w) -> b c t (h w)
|
||||
auto x_mix = x;
|
||||
|
||||
emb = ggml_reshape_4d(ctx->ggml_ctx, emb, emb->ne[0], T, B, emb->ne[3]); // (b t) ... -> b t ...
|
||||
|
||||
x = time_stack->forward(ctx, x, emb); // b t c (h w)
|
||||
|
||||
x = time_mixer->forward(ctx, x_mix, x); // b t c (h w)
|
||||
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // b c t (h w) -> b t c (h w)
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, C, T * B); // b t c (h w) -> (b t) c h w
|
||||
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_MODEL_COMMON_BLOCK_HPP__
|
||||
@@ -0,0 +1,905 @@
|
||||
#ifndef __SD_MODEL_COMMON_ROPE_HPP__
|
||||
#define __SD_MODEL_COMMON_ROPE_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
#include "core/ggml_extend.hpp"
|
||||
|
||||
namespace Rope {
|
||||
enum class EmbedNDLayout {
|
||||
Matrix,
|
||||
ErnieImage,
|
||||
};
|
||||
|
||||
template <class T>
|
||||
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
|
||||
std::vector<T> result(num);
|
||||
if (num == 1) {
|
||||
result[0] = start;
|
||||
return result;
|
||||
}
|
||||
T step = (end - start) / (num - 1);
|
||||
for (int i = 0; i < num; ++i) {
|
||||
result[i] = start + i * step;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> transpose(const std::vector<std::vector<float>>& mat) {
|
||||
size_t rows = mat.size();
|
||||
size_t cols = mat[0].size();
|
||||
std::vector<std::vector<float>> transposed(cols, std::vector<float>(rows));
|
||||
for (size_t i = 0; i < rows; ++i) {
|
||||
for (size_t j = 0; j < cols; ++j) {
|
||||
transposed[j][i] = mat[i][j];
|
||||
}
|
||||
}
|
||||
return transposed;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> flatten(const std::vector<std::vector<float>>& vec) {
|
||||
std::vector<float> flat_vec;
|
||||
for (const auto& sub_vec : vec) {
|
||||
flat_vec.insert(flat_vec.end(), sub_vec.begin(), sub_vec.end());
|
||||
}
|
||||
return flat_vec;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
|
||||
int dim,
|
||||
float theta,
|
||||
const std::vector<int>& axis_wrap_dims = {}) {
|
||||
assert(dim % 2 == 0);
|
||||
int half_dim = dim / 2;
|
||||
|
||||
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
|
||||
|
||||
std::vector<float> omega(half_dim);
|
||||
for (int i = 0; i < half_dim; ++i) {
|
||||
omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
|
||||
}
|
||||
|
||||
size_t pos_size = pos.size();
|
||||
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
|
||||
for (size_t i = 0; i < pos_size; ++i) {
|
||||
for (size_t j = 0; j < half_dim; ++j) {
|
||||
float angle = pos[i] * omega[j];
|
||||
if (!axis_wrap_dims.empty()) {
|
||||
size_t wrap_size = axis_wrap_dims.size();
|
||||
// mod batch size since we only store this for one item in the batch
|
||||
size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
|
||||
int wrap_dim = axis_wrap_dims[wrap_idx];
|
||||
if (wrap_dim > 0) {
|
||||
constexpr float TWO_PI = 6.28318530717958647692f;
|
||||
float cycles = omega[j] * wrap_dim / TWO_PI;
|
||||
// closest periodic harmonic, necessary to ensure things neatly tile
|
||||
// without this round, things don't tile at the boundaries and you end up
|
||||
// with the model knowing what is "center"
|
||||
float rounded = std::round(cycles);
|
||||
angle = pos[i] * TWO_PI * rounded / wrap_dim;
|
||||
}
|
||||
}
|
||||
|
||||
out[i][j] = angle;
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> result(pos_size, std::vector<float>(half_dim * 4));
|
||||
for (int i = 0; i < pos_size; ++i) {
|
||||
for (int j = 0; j < half_dim; ++j) {
|
||||
result[i][4 * j] = std::cos(out[i][j]);
|
||||
result[i][4 * j + 1] = -std::sin(out[i][j]);
|
||||
result[i][4 * j + 2] = std::sin(out[i][j]);
|
||||
result[i][4 * j + 3] = std::cos(out[i][j]);
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// Generate IDs for image patches and text
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
|
||||
for (int dim = 0; dim < axes_dim_num; dim++) {
|
||||
if (arange_dims.find(dim) != arange_dims.end()) {
|
||||
for (int i = 0; i < bs * context_len; i++) {
|
||||
txt_ids[i][dim] = 1.f * (i % context_len);
|
||||
}
|
||||
}
|
||||
}
|
||||
return txt_ids;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_longcat_txt_ids(int bs, int context_len, int axes_dim_num) {
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
|
||||
for (int i = 0; i < bs * context_len; i++) {
|
||||
float token_index = static_cast<float>(i % context_len);
|
||||
txt_ids[i][1] = token_index;
|
||||
txt_ids[i][2] = token_index;
|
||||
}
|
||||
return txt_ids;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_img_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int axes_dim_num,
|
||||
int index = 0,
|
||||
int h_offset = 0,
|
||||
int w_offset = 0,
|
||||
bool scale_rope = false) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
|
||||
|
||||
int h_start = h_offset;
|
||||
int w_start = w_offset;
|
||||
|
||||
if (scale_rope) {
|
||||
h_start -= h_len / 2;
|
||||
w_start -= w_len / 2;
|
||||
}
|
||||
|
||||
std::vector<float> row_ids = linspace<float>(1.f * h_start, 1.f * h_start + h_len - 1, h_len);
|
||||
std::vector<float> col_ids = linspace<float>(1.f * w_start, 1.f * w_start + w_len - 1, w_len);
|
||||
for (int i = 0; i < h_len; ++i) {
|
||||
for (int j = 0; j < w_len; ++j) {
|
||||
img_ids[i * w_len + j][0] = 1.f * index;
|
||||
img_ids[i * w_len + j][1] = row_ids[i];
|
||||
img_ids[i * w_len + j][2] = col_ids[j];
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> img_ids_repeated(bs * img_ids.size(), std::vector<float>(3));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < img_ids.size(); ++j) {
|
||||
img_ids_repeated[i * img_ids.size() + j] = img_ids[j];
|
||||
}
|
||||
}
|
||||
return img_ids_repeated;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> concat_ids(const std::vector<std::vector<float>>& a,
|
||||
const std::vector<std::vector<float>>& b,
|
||||
int bs) {
|
||||
size_t a_len = a.size() / bs;
|
||||
size_t b_len = b.size() / bs;
|
||||
std::vector<std::vector<float>> ids(a.size() + b.size(), std::vector<float>(3));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < a_len; ++j) {
|
||||
ids[i * (a_len + b_len) + j] = a[i * a_len + j];
|
||||
}
|
||||
for (int j = 0; j < b_len; ++j) {
|
||||
ids[i * (a_len + b_len) + a_len + j] = b[i * b_len + j];
|
||||
}
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
|
||||
int bs,
|
||||
const std::vector<float>& axis_thetas,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {},
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
|
||||
std::vector<std::vector<float>> trans_ids = transpose(ids);
|
||||
size_t pos_len = ids.size() / bs;
|
||||
size_t num_axes = axes_dim.size();
|
||||
// for (int i = 0; i < pos_len; i++) {
|
||||
// std::cout << trans_ids[0][i] << " " << trans_ids[1][i] << " " << trans_ids[2][i] << std::endl;
|
||||
// }
|
||||
|
||||
int emb_dim = 0;
|
||||
for (int d : axes_dim)
|
||||
emb_dim += d / 2;
|
||||
|
||||
std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
|
||||
size_t offset = 0;
|
||||
for (size_t i = 0; i < num_axes; ++i) {
|
||||
std::vector<int> axis_wrap_dims;
|
||||
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
|
||||
axis_wrap_dims = wrap_dims[i];
|
||||
}
|
||||
float axis_theta = 10000.0f;
|
||||
if (!axis_thetas.empty()) {
|
||||
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
|
||||
}
|
||||
std::vector<std::vector<float>> rope_emb =
|
||||
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
|
||||
for (int b = 0; b < bs; ++b) {
|
||||
for (int j = 0; j < pos_len; ++j) {
|
||||
for (int k = 0; k < rope_emb[0].size(); ++k) {
|
||||
emb[b * pos_len + j][offset + k] = rope_emb[j][k];
|
||||
}
|
||||
}
|
||||
}
|
||||
offset += rope_emb[0].size();
|
||||
}
|
||||
|
||||
if (layout == EmbedNDLayout::ErnieImage) {
|
||||
int head_dim = emb_dim * 2;
|
||||
std::vector<float> ernie_emb(bs * pos_len * head_dim * 2, 0.0f);
|
||||
for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
|
||||
for (int i = 0; i < emb_dim; ++i) {
|
||||
float cos_val = emb[pos_idx][4 * i];
|
||||
float sin_val = emb[pos_idx][4 * i + 2];
|
||||
size_t cos_offset = pos_idx * head_dim + 2 * i;
|
||||
size_t sin_offset = bs * pos_len * head_dim + cos_offset;
|
||||
ernie_emb[cos_offset] = cos_val;
|
||||
ernie_emb[cos_offset + 1] = cos_val;
|
||||
ernie_emb[sin_offset] = sin_val;
|
||||
ernie_emb[sin_offset + 1] = sin_val;
|
||||
}
|
||||
}
|
||||
return ernie_emb;
|
||||
}
|
||||
|
||||
return flatten(emb);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_nd(const std::vector<std::vector<float>>& ids,
|
||||
int bs,
|
||||
float theta,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {},
|
||||
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
|
||||
std::vector<float> axis_thetas(axes_dim.size(), theta);
|
||||
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
|
||||
int bs,
|
||||
float theta,
|
||||
int head_dim,
|
||||
const std::vector<int>& mrope_section) {
|
||||
GGML_ASSERT(bs > 0);
|
||||
GGML_ASSERT(head_dim % 2 == 0);
|
||||
GGML_ASSERT(mrope_section.size() >= 3);
|
||||
|
||||
std::vector<std::vector<float>> trans_ids = transpose(ids);
|
||||
size_t pos_len = ids.size() / bs;
|
||||
int half_dim = head_dim / 2;
|
||||
|
||||
std::vector<std::vector<std::vector<float>>> axis_embs;
|
||||
axis_embs.reserve(3);
|
||||
for (int axis = 0; axis < 3; ++axis) {
|
||||
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta));
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> emb = axis_embs[0];
|
||||
for (int axis = 1; axis < 3; ++axis) {
|
||||
int length = std::min<int>(mrope_section[axis] * 3, half_dim);
|
||||
for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
|
||||
for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
|
||||
for (int k = 0; k < 4; ++k) {
|
||||
emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return flatten(emb);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height,
|
||||
int width,
|
||||
int dim,
|
||||
float theta = 10000.f,
|
||||
float scale = 16.f,
|
||||
int ref_grid_h = 0,
|
||||
int ref_grid_w = 0) {
|
||||
assert(dim % 4 == 0);
|
||||
int half_dim = dim / 2;
|
||||
int dim_axis = dim / 2;
|
||||
int axis_half_dim = dim_axis / 2;
|
||||
|
||||
float h_ntk = 1.f;
|
||||
float w_ntk = 1.f;
|
||||
if (ref_grid_h > 0 && ref_grid_w > 0 && dim_axis > 2) {
|
||||
float power = static_cast<float>(dim_axis) / static_cast<float>(dim_axis - 2);
|
||||
h_ntk = std::pow(static_cast<float>(height) / static_cast<float>(ref_grid_h), power);
|
||||
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
|
||||
}
|
||||
|
||||
std::vector<float> x_pos;
|
||||
std::vector<float> y_pos;
|
||||
x_pos.reserve(static_cast<size_t>(height) * width);
|
||||
y_pos.reserve(static_cast<size_t>(height) * width);
|
||||
for (int iy = 0; iy < height; ++iy) {
|
||||
float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
|
||||
for (int ix = 0; ix < width; ++ix) {
|
||||
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
|
||||
x_pos.push_back(x);
|
||||
y_pos.push_back(y);
|
||||
}
|
||||
}
|
||||
|
||||
auto x_emb = rope(x_pos, dim_axis, theta * w_ntk);
|
||||
auto y_emb = rope(y_pos, dim_axis, theta * h_ntk);
|
||||
|
||||
std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
|
||||
for (int pos = 0; pos < height * width; ++pos) {
|
||||
for (int i = 0; i < axis_half_dim; ++i) {
|
||||
int jx = 2 * i;
|
||||
int jy = 2 * i + 1;
|
||||
size_t base_x = static_cast<size_t>(pos) * half_dim * 4 + static_cast<size_t>(jx) * 4;
|
||||
size_t base_y = static_cast<size_t>(pos) * half_dim * 4 + static_cast<size_t>(jy) * 4;
|
||||
size_t axis = static_cast<size_t>(i) * 4;
|
||||
for (int k = 0; k < 4; ++k) {
|
||||
out[base_x + k] = x_emb[pos][axis + k];
|
||||
out[base_y + k] = y_emb[pos][axis + k];
|
||||
}
|
||||
}
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
|
||||
int bs,
|
||||
int axes_dim_num,
|
||||
int start_index,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
float ref_index_scale,
|
||||
bool scale_rope,
|
||||
int base_offset = 0) {
|
||||
std::vector<std::vector<float>> ids;
|
||||
int curr_h_offset = 0;
|
||||
int curr_w_offset = 0;
|
||||
int index = start_index;
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
int h_offset = 0;
|
||||
int w_offset = 0;
|
||||
if (!increase_ref_index) {
|
||||
if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
|
||||
w_offset = curr_w_offset;
|
||||
} else {
|
||||
h_offset = curr_h_offset;
|
||||
}
|
||||
scale_rope = false;
|
||||
}
|
||||
|
||||
auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
|
||||
static_cast<int>(ref->ne[0]),
|
||||
patch_size,
|
||||
bs,
|
||||
axes_dim_num,
|
||||
static_cast<int>(index * ref_index_scale),
|
||||
h_offset + base_offset,
|
||||
w_offset + base_offset,
|
||||
scale_rope);
|
||||
ids = concat_ids(ids, ref_ids, bs);
|
||||
|
||||
if (increase_ref_index) {
|
||||
index++;
|
||||
}
|
||||
|
||||
curr_h_offset = std::max(curr_h_offset, static_cast<int>(ref->ne[1]) + h_offset);
|
||||
curr_w_offset = std::max(curr_w_offset, static_cast<int>(ref->ne[0]) + w_offset);
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int axes_dim_num,
|
||||
int context_len,
|
||||
std::set<int> txt_arange_dims,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
float ref_index_scale,
|
||||
bool is_longcat) {
|
||||
int x_index = is_longcat ? 1 : 0;
|
||||
|
||||
auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
|
||||
int offset = is_longcat ? context_len : 0;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
|
||||
|
||||
auto ids = concat_ids(txt_ids, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
|
||||
ids = concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
// Generate flux positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_flux_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
std::set<int> txt_arange_dims,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
float ref_index_scale,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim,
|
||||
bool is_longcat) {
|
||||
std::vector<std::vector<float>> ids = gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
bs,
|
||||
static_cast<int>(axes_dim.size()),
|
||||
context_len,
|
||||
txt_arange_dims,
|
||||
ref_latents,
|
||||
increase_ref_index,
|
||||
ref_index_scale,
|
||||
is_longcat);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
size_t cursor = context_len; // text first
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
cursor += img_tokens;
|
||||
// reference latents
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
if (ref == nullptr) {
|
||||
continue;
|
||||
}
|
||||
int ref_h = static_cast<int>(ref->ne[1]);
|
||||
int ref_w = static_cast<int>(ref->ne[0]);
|
||||
int ref_h_l = (ref_h + (patch_size / 2)) / patch_size;
|
||||
int ref_w_l = (ref_w + (patch_size / 2)) / patch_size;
|
||||
size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
|
||||
for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = ref_h_l;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = ref_w_l;
|
||||
}
|
||||
}
|
||||
cursor += ref_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
int txt_id_start = std::max(h_len, w_len);
|
||||
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
|
||||
std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < txt_ids.size(); ++j) {
|
||||
txt_ids_repeated[i * txt_ids.size() + j] = {txt_ids[j], txt_ids[j], txt_ids[j]};
|
||||
}
|
||||
}
|
||||
int axes_dim_num = 3;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
|
||||
auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
|
||||
ids = concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
// Generate qwen_image positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int pad_h = (patch_size - (h % patch_size)) % patch_size;
|
||||
int pad_w = (patch_size - (w % patch_size)) % patch_size;
|
||||
int h_len = (h + pad_h) / patch_size;
|
||||
int w_len = (w + pad_w) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
const size_t total_tokens = ids.size();
|
||||
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
|
||||
size_t cursor = context_len; // ignore text tokens
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
cursor += img_tokens;
|
||||
// For each reference image, store wrap sizes as well
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
if (ref == nullptr) {
|
||||
continue;
|
||||
}
|
||||
int ref_h = static_cast<int>(ref->ne[1]);
|
||||
int ref_w = static_cast<int>(ref->ne[0]);
|
||||
int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size;
|
||||
int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size;
|
||||
int ref_h_len = (ref_h + ref_pad_h) / patch_size;
|
||||
int ref_w_len = (ref_w + ref_pad_w) / patch_size;
|
||||
size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
|
||||
for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = ref_h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = ref_w_len;
|
||||
}
|
||||
}
|
||||
cursor += ref_n_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
bool scale_rope = true) {
|
||||
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
|
||||
|
||||
int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
|
||||
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
|
||||
std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < txt_ids.size(); ++j) {
|
||||
txt_ids_repeated[i * txt_ids.size() + j] = {txt_ids[j], txt_ids[j], txt_ids[j]};
|
||||
}
|
||||
}
|
||||
|
||||
return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len) {
|
||||
int h_len = h / patch_size;
|
||||
int w_len = w / patch_size;
|
||||
|
||||
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
|
||||
std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
|
||||
std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
|
||||
for (int i = 0; i < h_len; ++i) {
|
||||
for (int j = 0; j < w_len; ++j) {
|
||||
img_ids[i * w_len + j][0] = static_cast<float>(context_len);
|
||||
img_ids[i * w_len + j][1] = h_ids[i];
|
||||
img_ids[i * w_len + j][2] = w_ids[j];
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> img_ids_repeated(bs * img_ids.size(), std::vector<float>(3, 0.0f));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < static_cast<int>(img_ids.size()); ++j) {
|
||||
img_ids_repeated[i * img_ids.size() + j] = img_ids[j];
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < context_len; ++j) {
|
||||
txt_ids[i * context_len + j][0] = static_cast<float>(j);
|
||||
}
|
||||
}
|
||||
|
||||
return concat_ids(img_ids_repeated, txt_ids, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int h_len = h / patch_size;
|
||||
int w_len = w / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
|
||||
int h,
|
||||
int w,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int bs,
|
||||
int t_offset = 0,
|
||||
int h_offset = 0,
|
||||
int w_offset = 0) {
|
||||
int t_len = (t + (pt / 2)) / pt;
|
||||
int h_len = (h + (ph / 2)) / ph;
|
||||
int w_len = (w + (pw / 2)) / pw;
|
||||
|
||||
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
|
||||
|
||||
std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
|
||||
std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
|
||||
std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
|
||||
|
||||
for (int i = 0; i < t_len; ++i) {
|
||||
for (int j = 0; j < h_len; ++j) {
|
||||
for (int k = 0; k < w_len; ++k) {
|
||||
int idx = i * h_len * w_len + j * w_len + k;
|
||||
vid_ids[idx][0] = t_ids[i];
|
||||
vid_ids[idx][1] = h_ids[j];
|
||||
vid_ids[idx][2] = w_ids[k];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < vid_ids.size(); ++j) {
|
||||
vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
|
||||
}
|
||||
}
|
||||
return vid_ids_repeated;
|
||||
}
|
||||
|
||||
// Generate wan positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int bs,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_vid_ids(t, h, w, pt, ph, pw, bs);
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
|
||||
int grid_w,
|
||||
int merge_size,
|
||||
const std::vector<int>& window_index) {
|
||||
std::vector<std::vector<float>> ids(grid_h * grid_w, std::vector<float>(2, 0.0));
|
||||
int index = 0;
|
||||
for (int ih = 0; ih < grid_h; ih += merge_size) {
|
||||
for (int iw = 0; iw < grid_w; iw += merge_size) {
|
||||
for (int iy = 0; iy < merge_size; iy++) {
|
||||
for (int ix = 0; ix < merge_size; ix++) {
|
||||
int inverse_index = window_index[index / (merge_size * merge_size)];
|
||||
int i = inverse_index * (merge_size * merge_size) + index % (merge_size * merge_size);
|
||||
|
||||
GGML_ASSERT(i < grid_h * grid_w);
|
||||
|
||||
ids[i][0] = static_cast<float>(ih + iy);
|
||||
ids[i][1] = static_cast<float>(iw + ix);
|
||||
index++;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return ids;
|
||||
}
|
||||
|
||||
// Generate qwen2vl positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_qwen2vl_pe(int grid_h,
|
||||
int grid_w,
|
||||
int merge_size,
|
||||
const std::vector<int>& window_index,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_qwen2vl_ids(grid_h, grid_w, merge_size, window_index);
|
||||
return embed_nd(ids, 1, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ int bound_mod(int a, int m) {
|
||||
return (m - (a % m)) % m;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_z_image_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int seq_multi_of,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index) {
|
||||
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
|
||||
for (int i = 0; i < bs * padded_context_len; i++) {
|
||||
txt_ids[i][0] = (i % padded_context_len) + 1.f;
|
||||
}
|
||||
|
||||
int axes_dim_num = 3;
|
||||
int index = padded_context_len + 1;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
|
||||
|
||||
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
|
||||
if (img_pad_len > 0) {
|
||||
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
|
||||
img_ids = concat_ids(img_ids, img_pad_ids, bs);
|
||||
}
|
||||
|
||||
auto ids = concat_ids(txt_ids, img_ids, bs);
|
||||
|
||||
// ignore ref_latents for now
|
||||
return ids;
|
||||
}
|
||||
|
||||
// Generate z_image positional embeddings
|
||||
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
int bs,
|
||||
int context_len,
|
||||
int seq_multi_of,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int pad_h = (patch_size - (h % patch_size)) % patch_size;
|
||||
int pad_w = (patch_size - (w % patch_size)) % patch_size;
|
||||
int h_len = (h + pad_h) / patch_size;
|
||||
int w_len = (w + pad_w) / patch_size;
|
||||
if (h_len > 0 && w_len > 0) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding)
|
||||
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][cursor + token_i] = h_len;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][cursor + token_i] = w_len;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* pe,
|
||||
bool rope_interleaved = true) {
|
||||
// x: [N, L, n_head, d_head]
|
||||
// pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]]
|
||||
int64_t d_head = x->ne[0];
|
||||
int64_t n_head = x->ne[1];
|
||||
int64_t L = x->ne[2];
|
||||
int64_t N = x->ne[3];
|
||||
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N, n_head, L, d_head]
|
||||
if (rope_interleaved) {
|
||||
x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N); // [N * n_head, L, d_head/2, 2]
|
||||
x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2)); // [2, N * n_head, L, d_head/2]
|
||||
} else {
|
||||
x = ggml_reshape_4d(ctx, x, d_head / 2, 2, L, n_head * N); // [N * n_head, L, 2, d_head/2]
|
||||
x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 3, 1)); // [2, N * n_head, L, d_head/2]
|
||||
}
|
||||
|
||||
int64_t offset = x->nb[2] * x->ne[2];
|
||||
auto x_0 = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0); // [N * n_head, L, d_head/2]
|
||||
auto x_1 = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1); // [N * n_head, L, d_head/2]
|
||||
x_0 = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]); // [N * n_head, L, d_head/2, 1]
|
||||
x_1 = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]); // [N * n_head, L, d_head/2, 1]
|
||||
auto temp_x = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]);
|
||||
x_0 = ggml_repeat(ctx, x_0, temp_x); // [N * n_head, L, d_head/2, 2]
|
||||
x_1 = ggml_repeat(ctx, x_1, temp_x); // [N * n_head, L, d_head/2, 2]
|
||||
|
||||
pe = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2)); // [2, L, d_head/2, 2]
|
||||
offset = pe->nb[2] * pe->ne[2];
|
||||
auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0); // [L, d_head/2, 2]
|
||||
auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1); // [L, d_head/2, 2]
|
||||
|
||||
auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1)); // [N * n_head, L, d_head/2, 2]
|
||||
if (!rope_interleaved) {
|
||||
x_out = ggml_cont(ctx, ggml_permute(ctx, x_out, 1, 0, 2, 3)); // [N * n_head, L, x, d_head/2]
|
||||
}
|
||||
x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N); // [N*n_head, L, d_head]
|
||||
return x_out;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* q,
|
||||
ggml_tensor* k,
|
||||
ggml_tensor* v,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* mask,
|
||||
float kv_scale = 1.0f,
|
||||
bool rope_interleaved = true) {
|
||||
// q,k,v: [N, L, n_head, d_head]
|
||||
// pe: [L, d_head/2, 2, 2]
|
||||
// return: [N, L, n_head*d_head]
|
||||
q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head]
|
||||
k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head]
|
||||
|
||||
auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, v->ne[1], mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head]
|
||||
return x;
|
||||
}
|
||||
}; // namespace Rope
|
||||
|
||||
#endif // __SD_MODEL_COMMON_ROPE_HPP__
|
||||
Reference in new issue
Block a user