DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-11 11:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: reorganize src model layout (#1615)

Browse Source
This commit is contained in:
leejet authored and GitHub committed 2026-06-07 03:21:12 +08:00
1 parent dfb2390dd4
commit f3fd359b58
81 files changed
+407 -385
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/model/adapter/lora.hpp
+917
View File
@@ -0,0 +1,917 @@
#ifndef __SD_MODEL_ADAPTER_LORA_HPP__
#define __SD_MODEL_ADAPTER_LORA_HPP__
#include <mutex>
#include "core/ggml_extend.hpp"
#define LORA_GRAPH_BASE_SIZE 10240
struct LoraModel : public GGMLRunner {
std::string lora_id;
float multiplier = 1.0f;
std::unordered_map<std::string, ggml_tensor*> lora_tensors;
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
std::set<std::string> applied_lora_tensors;
std::string file_path;
ModelLoader model_loader;
bool load_failed = false;
bool applied = false;
bool tensor_preprocessed = false;
typedef std::function<bool(const std::string&)> filter_t;
LoraModel(const std::string& lora_id,
ggml_backend_t backend,
ggml_backend_t params_backend,
const std::string& file_path = "",
std::string prefix = "",
SDVersion version = VERSION_COUNT)
: lora_id(lora_id), file_path(file_path), GGMLRunner(backend, params_backend) {
prefix = "lora." + prefix;
if (!model_loader.init_from_file_and_convert_name(file_path, prefix, version)) {
load_failed = true;
}
}
std::string get_desc() override {
return "lora";
}
bool load_from_file(int n_threads, filter_t filter = nullptr) {
LOG_INFO("loading LoRA from '%s'", file_path.c_str());
if (load_failed) {
LOG_ERROR("init lora model loader from file failed: '%s'", file_path.c_str());
return false;
}
std::unordered_map<std::string, TensorStorage> tensors_to_create;
std::mutex lora_mutex;
bool dry_run = true;
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
if (dry_run) {
const std::string& name = tensor_storage.name;
if (filter && !filter(name)) {
return true;
}
{
std::lock_guard<std::mutex> lock(lora_mutex);
tensors_to_create[name] = tensor_storage;
}
} else {
const std::string& name = tensor_storage.name;
auto iter = lora_tensors.find(name);
if (iter != lora_tensors.end()) {
*dst_tensor = iter->second;
}
}
return true;
};
model_loader.load_tensors(on_new_tensor_cb, n_threads);
if (tensors_to_create.empty()) {
return true;
}
for (const auto& pair : tensors_to_create) {
const auto& name = pair.first;
const auto& ts = pair.second;
ggml_tensor* real = ggml_new_tensor(params_ctx,
ts.type,
ts.n_dims,
ts.ne);
lora_tensors[name] = real;
}
if (!alloc_params_buffer()) {
LOG_ERROR("lora model buffer allocation failed");
return false;
}
dry_run = false;
model_loader.load_tensors(on_new_tensor_cb, n_threads);
LOG_DEBUG("finished loaded lora");
return true;
}
void preprocess_lora_tensors(const std::map<std::string, ggml_tensor*>& model_tensors) {
if (tensor_preprocessed) {
return;
}
tensor_preprocessed = true;
// I really hate these hardcoded processes.
if (model_tensors.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensors.end()) {
std::unordered_map<std::string, ggml_tensor*> new_lora_tensors;
for (auto& [old_name, tensor] : lora_tensors) {
std::string new_name = old_name;
if (contains(new_name, "cond_stage_model.1.transformer.text_model.encoder.layers")) {
std::vector<std::pair<std::string, std::string>> qkv_name_map = {
{"self_attn.q_proj.weight", "self_attn.in_proj.weight"},
{"self_attn.q_proj.bias", "self_attn.in_proj.bias"},
{"self_attn.k_proj.weight", "self_attn.in_proj.weight.1"},
{"self_attn.k_proj.bias", "self_attn.in_proj.bias.1"},
{"self_attn.v_proj.weight", "self_attn.in_proj.weight.2"},
{"self_attn.v_proj.bias", "self_attn.in_proj.bias.2"},
};
for (auto kv : qkv_name_map) {
size_t pos = new_name.find(kv.first);
if (pos != std::string::npos) {
new_name.replace(pos, kv.first.size(), kv.second);
}
}
}
new_lora_tensors[new_name] = tensor;
}
lora_tensors = std::move(new_lora_tensors);
}
}
ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
ggml_tensor* updown = nullptr;
int index = 0;
while (true) {
std::string key;
if (index == 0) {
key = model_tensor_name;
} else {
key = model_tensor_name + "." + std::to_string(index);
}
std::string lora_down_name = "lora." + key + ".lora_down";
std::string lora_up_name = "lora." + key + ".lora_up";
std::string lora_mid_name = "lora." + key + ".lora_mid";
std::string scale_name = "lora." + key + ".scale";
std::string alpha_name = "lora." + key + ".alpha";
ggml_tensor* lora_up = nullptr;
ggml_tensor* lora_mid = nullptr;
ggml_tensor* lora_down = nullptr;
auto iter = lora_tensors.find(lora_up_name);
if (iter != lora_tensors.end()) {
lora_up = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(lora_mid_name);
if (iter != lora_tensors.end()) {
lora_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(lora_down_name);
if (iter != lora_tensors.end()) {
lora_down = ggml_ext_cast_f32(ctx, backend, iter->second);
}
if (lora_up == nullptr || lora_down == nullptr) {
break;
}
applied_lora_tensors.insert(lora_up_name);
applied_lora_tensors.insert(lora_down_name);
if (lora_mid) {
applied_lora_tensors.insert(lora_mid_name);
}
float scale_value = 1.0f;
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
iter = lora_tensors.find(scale_name);
if (iter != lora_tensors.end()) {
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
applied_lora_tensors.insert(scale_name);
} else {
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
// LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
applied_lora_tensors.insert(alpha_name);
}
}
scale_value *= multiplier;
auto curr_updown = ggml_ext_merge_lora(ctx, lora_down, lora_up, lora_mid);
curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
if (updown == nullptr) {
updown = curr_updown;
} else {
updown = ggml_concat(ctx, updown, curr_updown, ggml_n_dims(updown) - 1);
}
index++;
}
return updown;
}
ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
ggml_tensor* updown = nullptr;
int index = 0;
while (true) {
std::string key;
if (index == 0) {
key = model_tensor_name;
} else {
key = model_tensor_name + "." + std::to_string(index);
}
std::string diff_name = "lora." + key + ".diff";
ggml_tensor* curr_updown = nullptr;
auto iter = lora_tensors.find(diff_name);
if (iter != lora_tensors.end()) {
curr_updown = ggml_ext_cast_f32(ctx, backend, iter->second);
} else {
break;
}
applied_lora_tensors.insert(diff_name);
float scale_value = 1.0f;
scale_value *= multiplier;
curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
if (updown == nullptr) {
updown = curr_updown;
} else {
updown = ggml_concat(ctx, updown, curr_updown, ggml_n_dims(updown) - 1);
}
index++;
}
return updown;
}
ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
ggml_tensor* updown = nullptr;
int index = 0;
while (true) {
std::string key;
if (index == 0) {
key = model_tensor_name;
} else {
key = model_tensor_name + "." + std::to_string(index);
}
std::string hada_1_down_name = "lora." + key + ".hada_w1_b";
std::string hada_1_mid_name = "lora." + key + ".hada_t1";
std::string hada_1_up_name = "lora." + key + ".hada_w1_a";
std::string hada_2_down_name = "lora." + key + ".hada_w2_b";
std::string hada_2_mid_name = "lora." + key + ".hada_t2";
std::string hada_2_up_name = "lora." + key + ".hada_w2_a";
std::string alpha_name = "lora." + key + ".alpha";
ggml_tensor* hada_1_mid = nullptr; // tau for tucker decomposition
ggml_tensor* hada_1_up = nullptr;
ggml_tensor* hada_1_down = nullptr;
ggml_tensor* hada_2_mid = nullptr; // tau for tucker decomposition
ggml_tensor* hada_2_up = nullptr;
ggml_tensor* hada_2_down = nullptr;
auto iter = lora_tensors.find(hada_1_down_name);
if (iter != lora_tensors.end()) {
hada_1_down = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(hada_1_up_name);
if (iter != lora_tensors.end()) {
hada_1_up = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(hada_1_mid_name);
if (iter != lora_tensors.end()) {
hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
}
iter = lora_tensors.find(hada_2_down_name);
if (iter != lora_tensors.end()) {
hada_2_down = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(hada_2_up_name);
if (iter != lora_tensors.end()) {
hada_2_up = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(hada_2_mid_name);
if (iter != lora_tensors.end()) {
hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
}
if (hada_1_up == nullptr || hada_1_down == nullptr || hada_2_up == nullptr || hada_2_down == nullptr) {
break;
}
applied_lora_tensors.insert(hada_1_down_name);
applied_lora_tensors.insert(hada_1_up_name);
applied_lora_tensors.insert(hada_2_down_name);
applied_lora_tensors.insert(hada_2_up_name);
applied_lora_tensors.insert(alpha_name);
if (hada_1_mid) {
applied_lora_tensors.insert(hada_1_mid_name);
}
if (hada_2_mid) {
applied_lora_tensors.insert(hada_2_mid_name);
}
float scale_value = 1.0f;
// calc_scale
// TODO: .dora_scale?
int64_t rank = hada_1_down->ne[ggml_n_dims(hada_1_down) - 1];
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
scale_value *= multiplier;
ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid);
ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid);
auto curr_updown = ggml_mul_inplace(ctx, updown_1, updown_2);
curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
if (updown == nullptr) {
updown = curr_updown;
} else {
updown = ggml_concat(ctx, updown, curr_updown, ggml_n_dims(updown) - 1);
}
index++;
}
return updown;
}
ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
ggml_tensor* updown = nullptr;
int index = 0;
while (true) {
std::string key;
if (index == 0) {
key = model_tensor_name;
} else {
key = model_tensor_name + "." + std::to_string(index);
}
std::string lokr_w1_name = "lora." + key + ".lokr_w1";
std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
std::string lokr_w1_b_name = "lora." + key + ".lokr_w1_b";
std::string lokr_w2_name = "lora." + key + ".lokr_w2";
std::string lokr_w2_a_name = "lora." + key + ".lokr_w2_a";
std::string lokr_w2_b_name = "lora." + key + ".lokr_w2_b";
std::string alpha_name = "lora." + key + ".alpha";
ggml_tensor* lokr_w1 = nullptr;
ggml_tensor* lokr_w1_a = nullptr;
ggml_tensor* lokr_w1_b = nullptr;
ggml_tensor* lokr_w2 = nullptr;
ggml_tensor* lokr_w2_a = nullptr;
ggml_tensor* lokr_w2_b = nullptr;
auto iter = lora_tensors.find(lokr_w1_name);
if (iter != lora_tensors.end()) {
lokr_w1 = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(lokr_w2_name);
if (iter != lora_tensors.end()) {
lokr_w2 = ggml_ext_cast_f32(ctx, backend, iter->second);
}
int64_t rank = 1;
if (lokr_w1 == nullptr) {
iter = lora_tensors.find(lokr_w1_a_name);
if (iter != lora_tensors.end()) {
lokr_w1_a = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(lokr_w1_b_name);
if (iter != lora_tensors.end()) {
lokr_w1_b = ggml_ext_cast_f32(ctx, backend, iter->second);
}
if (lokr_w1_a == nullptr || lokr_w1_b == nullptr) {
break;
}
rank = lokr_w1_b->ne[ggml_n_dims(lokr_w1_b) - 1];
lokr_w1 = ggml_ext_merge_lora(ctx, lokr_w1_b, lokr_w1_a);
}
if (lokr_w2 == nullptr) {
iter = lora_tensors.find(lokr_w2_a_name);
if (iter != lora_tensors.end()) {
lokr_w2_a = ggml_ext_cast_f32(ctx, backend, iter->second);
}
iter = lora_tensors.find(lokr_w2_b_name);
if (iter != lora_tensors.end()) {
lokr_w2_b = ggml_ext_cast_f32(ctx, backend, iter->second);
}
if (lokr_w2_a == nullptr || lokr_w2_b == nullptr) {
break;
}
rank = lokr_w2_b->ne[ggml_n_dims(lokr_w2_b) - 1];
lokr_w2 = ggml_ext_merge_lora(ctx, lokr_w2_b, lokr_w2_a);
}
if (!lokr_w1_a) {
applied_lora_tensors.insert(lokr_w1_name);
} else {
applied_lora_tensors.insert(lokr_w1_a_name);
applied_lora_tensors.insert(lokr_w1_b_name);
}
if (!lokr_w2_a) {
applied_lora_tensors.insert(lokr_w2_name);
} else {
applied_lora_tensors.insert(lokr_w2_a_name);
applied_lora_tensors.insert(lokr_w2_b_name);
}
float scale_value = 1.0f;
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
if (rank == 1) {
scale_value = 1.0f;
}
scale_value *= multiplier;
auto curr_updown = ggml_ext_kronecker(ctx, lokr_w1, lokr_w2);
curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true);
if (updown == nullptr) {
updown = curr_updown;
} else {
updown = ggml_concat(ctx, updown, curr_updown, ggml_n_dims(updown) - 1);
}
index++;
}
return updown;
}
ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_backend_t backend, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) {
// lora
ggml_tensor* diff = nullptr;
if (with_lora_and_lokr) {
diff = get_lora_weight_diff(model_tensor_name, ctx, backend);
}
// diff
if (diff == nullptr) {
diff = get_raw_weight_diff(model_tensor_name, ctx, backend);
}
// loha
if (diff == nullptr) {
diff = get_loha_weight_diff(model_tensor_name, ctx, backend);
}
// lokr
if (diff == nullptr && with_lora_and_lokr) {
diff = get_lokr_weight_diff(model_tensor_name, ctx, backend);
}
if (diff != nullptr) {
if (ggml_nelements(diff) < ggml_nelements(model_tensor)) {
if (ggml_n_dims(diff) == 2 && ggml_n_dims(model_tensor) == 2 && diff->ne[0] == model_tensor->ne[0]) {
LOG_WARN("pad for %s", model_tensor_name.c_str());
auto pad_tensor = ggml_ext_zeros(ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1);
diff = ggml_concat(ctx, diff, pad_tensor, 1);
}
}
GGML_ASSERT(ggml_nelements(diff) == ggml_nelements(model_tensor));
diff = ggml_reshape(ctx, diff, model_tensor);
}
return diff;
}
ggml_tensor* get_out_diff(ggml_context* ctx,
ggml_backend_t backend,
ggml_tensor* x,
WeightAdapter::ForwardParams forward_params,
const std::string& model_tensor_name) {
ggml_tensor* out_diff = nullptr;
int index = 0;
while (true) {
std::string key;
if (index == 0) {
key = model_tensor_name;
} else {
key = model_tensor_name + "." + std::to_string(index);
}
bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
std::string lokr_w1_name = "lora." + key + ".lokr_w1";
std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
// if either of these is found, then we have a lokr lora
auto iter = lora_tensors.find(lokr_w1_name);
auto iter_a = lora_tensors.find(lokr_w1_a_name);
if (iter != lora_tensors.end() || iter_a != lora_tensors.end()) {
std::string lokr_w1_b_name = "lora." + key + ".lokr_w1_b";
std::string lokr_w2_name = "lora." + key + ".lokr_w2";
std::string lokr_w2_a_name = "lora." + key + ".lokr_w2_a";
std::string lokr_w2_b_name = "lora." + key + ".lokr_w2_b";
std::string alpha_name = "lora." + key + ".alpha";
ggml_tensor* lokr_w1 = nullptr;
ggml_tensor* lokr_w1_a = nullptr;
ggml_tensor* lokr_w1_b = nullptr;
ggml_tensor* lokr_w2 = nullptr;
ggml_tensor* lokr_w2_a = nullptr;
ggml_tensor* lokr_w2_b = nullptr;
if (iter != lora_tensors.end()) {
lokr_w1 = iter->second;
}
iter = iter_a;
if (iter != lora_tensors.end()) {
lokr_w1_a = iter->second;
}
iter = lora_tensors.find(lokr_w1_b_name);
if (iter != lora_tensors.end()) {
lokr_w1_b = iter->second;
}
iter = lora_tensors.find(lokr_w2_name);
if (iter != lora_tensors.end()) {
lokr_w2 = iter->second;
if (is_conv2d && lokr_w2->type != GGML_TYPE_F16) {
lokr_w2 = ggml_cast(ctx, lokr_w2, GGML_TYPE_F16);
}
}
iter = lora_tensors.find(lokr_w2_a_name);
if (iter != lora_tensors.end()) {
lokr_w2_a = iter->second;
if (is_conv2d && lokr_w2_a->type != GGML_TYPE_F16) {
lokr_w2_a = ggml_cast(ctx, lokr_w2_a, GGML_TYPE_F16);
}
}
iter = lora_tensors.find(lokr_w2_b_name);
if (iter != lora_tensors.end()) {
lokr_w2_b = iter->second;
if (is_conv2d && lokr_w2_b->type != GGML_TYPE_F16) {
lokr_w2_b = ggml_cast(ctx, lokr_w2_b, GGML_TYPE_F16);
}
}
int rank = 1;
if (lokr_w1_b) {
rank = (int)lokr_w1_b->ne[ggml_n_dims(lokr_w1_b) - 1];
}
if (lokr_w2_b) {
rank = (int)lokr_w2_b->ne[ggml_n_dims(lokr_w2_b) - 1];
}
float scale_value = 1.0f;
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
if (rank == 1) {
scale_value = 1.0f;
}
scale_value *= multiplier;
auto curr_out_diff = ggml_ext_lokr_forward(ctx, backend, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
if (out_diff == nullptr) {
out_diff = curr_out_diff;
} else {
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
}
if (lokr_w1)
applied_lora_tensors.insert(lokr_w1_name);
if (lokr_w1_a)
applied_lora_tensors.insert(lokr_w1_a_name);
if (lokr_w1_b)
applied_lora_tensors.insert(lokr_w1_b_name);
if (lokr_w2)
applied_lora_tensors.insert(lokr_w2_name);
if (lokr_w2_a)
applied_lora_tensors.insert(lokr_w2_name);
if (lokr_w2_b)
applied_lora_tensors.insert(lokr_w2_b_name);
applied_lora_tensors.insert(alpha_name);
index++;
continue;
}
// not a lokr, normal lora path
std::string lora_down_name = "lora." + key + ".lora_down";
std::string lora_up_name = "lora." + key + ".lora_up";
std::string lora_mid_name = "lora." + key + ".lora_mid";
std::string scale_name = "lora." + key + ".scale";
std::string alpha_name = "lora." + key + ".alpha";
ggml_tensor* lora_up = nullptr;
ggml_tensor* lora_mid = nullptr;
ggml_tensor* lora_down = nullptr;
iter = lora_tensors.find(lora_up_name);
if (iter != lora_tensors.end()) {
lora_up = iter->second;
if (is_conv2d && lora_up->type != GGML_TYPE_F16) {
lora_up = ggml_cast(ctx, lora_up, GGML_TYPE_F16);
}
}
iter = lora_tensors.find(lora_mid_name);
if (iter != lora_tensors.end()) {
lora_mid = iter->second;
if (is_conv2d && lora_mid->type != GGML_TYPE_F16) {
lora_mid = ggml_cast(ctx, lora_mid, GGML_TYPE_F16);
}
}
iter = lora_tensors.find(lora_down_name);
if (iter != lora_tensors.end()) {
lora_down = iter->second;
if (is_conv2d && lora_down->type != GGML_TYPE_F16) {
lora_down = ggml_cast(ctx, lora_down, GGML_TYPE_F16);
}
}
if (lora_up == nullptr || lora_down == nullptr) {
break;
}
applied_lora_tensors.insert(lora_up_name);
applied_lora_tensors.insert(lora_down_name);
if (lora_mid) {
applied_lora_tensors.insert(lora_mid_name);
}
float scale_value = 1.0f;
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
iter = lora_tensors.find(scale_name);
if (iter != lora_tensors.end()) {
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
applied_lora_tensors.insert(scale_name);
} else {
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
// LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
applied_lora_tensors.insert(alpha_name);
}
}
scale_value *= multiplier;
ggml_tensor* lx;
if (!is_conv2d) {
lx = ggml_ext_linear(ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
if (lora_mid) {
lx = ggml_ext_linear(ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
}
lx = ggml_ext_linear(ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale);
} else { // OP_CONV2D
lx = ggml_ext_conv_2d(ctx,
x,
lora_down,
nullptr,
forward_params.conv2d.s0,
forward_params.conv2d.s1,
forward_params.conv2d.p0,
forward_params.conv2d.p1,
forward_params.conv2d.d0,
forward_params.conv2d.d1,
forward_params.conv2d.direct,
forward_params.conv2d.circular_x,
forward_params.conv2d.circular_y,
forward_params.conv2d.scale);
if (lora_mid) {
lx = ggml_ext_conv_2d(ctx,
lx,
lora_mid,
nullptr,
1,
1,
0,
0,
1,
1,
forward_params.conv2d.direct,
forward_params.conv2d.circular_x,
forward_params.conv2d.circular_y,
forward_params.conv2d.scale);
}
lx = ggml_ext_conv_2d(ctx,
lx,
lora_up,
nullptr,
1,
1,
0,
0,
1,
1,
forward_params.conv2d.direct,
forward_params.conv2d.circular_x,
forward_params.conv2d.circular_y,
forward_params.conv2d.scale);
}
auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
if (out_diff == nullptr) {
out_diff = curr_out_diff;
} else {
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
}
index++;
}
return out_diff;
}
ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
preprocess_lora_tensors(model_tensors);
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
for (auto it : model_tensors) {
std::string model_tensor_name = it.first;
ggml_tensor* model_tensor = it.second;
// lora
ggml_tensor* diff = get_weight_diff(model_tensor_name, runtime_backend, compute_ctx, model_tensor);
if (diff == nullptr) {
continue;
}
ggml_tensor* original_tensor = model_tensor;
if (!sd_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
model_tensor = ggml_dup_tensor(compute_ctx, model_tensor);
set_backend_tensor_data(model_tensor, original_tensor->data);
}
ggml_tensor* final_tensor;
if (model_tensor->type != GGML_TYPE_F32 && model_tensor->type != GGML_TYPE_F16) {
final_tensor = ggml_ext_cast_f32(compute_ctx, runtime_backend, model_tensor);
final_tensor = ggml_add_inplace(compute_ctx, final_tensor, diff);
final_tensor = ggml_cpy(compute_ctx, final_tensor, model_tensor);
} else {
final_tensor = ggml_add_inplace(compute_ctx, model_tensor, diff);
}
ggml_build_forward_expand(gf, final_tensor);
if (!sd_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
original_tensor_to_final_tensor[original_tensor] = final_tensor;
}
}
return gf;
}
void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_lora_graph(model_tensors, version);
};
GGMLRunner::compute<float>(get_graph, n_threads, false, true);
stat();
for (auto item : original_tensor_to_final_tensor) {
ggml_tensor* original_tensor = item.first;
ggml_tensor* final_tensor = item.second;
ggml_backend_tensor_copy(final_tensor, original_tensor);
}
original_tensor_to_final_tensor.clear();
GGMLRunner::free_compute_buffer();
}
void stat(bool at_runntime = false) {
size_t total_lora_tensors_count = 0;
size_t applied_lora_tensors_count = 0;
for (auto& kv : lora_tensors) {
total_lora_tensors_count++;
if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
if (!at_runntime) {
LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
print_ggml_tensor(kv.second, true);
}
} else {
applied_lora_tensors_count++;
}
}
/* Don't worry if this message shows up twice in the logs per LoRA,
* this function is called once to calculate the required buffer size
* and then again to actually generate a graph to be used */
if (!at_runntime && applied_lora_tensors_count != total_lora_tensors_count) {
LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
} else {
LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
}
}
};
struct MultiLoraAdapter : public WeightAdapter {
protected:
std::vector<std::shared_ptr<LoraModel>> lora_models;
public:
explicit MultiLoraAdapter(const std::vector<std::shared_ptr<LoraModel>>& lora_models)
: lora_models(lora_models) {
}
ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) {
for (auto& lora_model : lora_models) {
ggml_tensor* diff = lora_model->get_weight_diff(weight_name, backend, ctx, weight, with_lora_and_lokr);
if (diff == nullptr) {
continue;
}
if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) {
weight = ggml_ext_cast_f32(ctx, backend, weight);
}
weight = ggml_add(ctx, weight, diff);
}
return weight;
}
ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) override {
return patch_weight(ctx, backend, weight, weight_name, true);
}
ggml_tensor* forward_with_lora(ggml_context* ctx,
ggml_backend_t backend,
ggml_tensor* x,
ggml_tensor* w,
ggml_tensor* b,
const std::string& prefix,
WeightAdapter::ForwardParams forward_params) override {
w = patch_weight(ctx, backend, w, prefix + "weight", false);
if (b) {
b = patch_weight(ctx, backend, b, prefix + "bias", false);
}
ggml_tensor* out;
if (forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR) {
out = ggml_ext_linear(ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale);
} else { // OP_CONV2D
out = ggml_ext_conv_2d(ctx,
x,
w,
b,
forward_params.conv2d.s0,
forward_params.conv2d.s1,
forward_params.conv2d.p0,
forward_params.conv2d.p1,
forward_params.conv2d.d0,
forward_params.conv2d.d1,
forward_params.conv2d.direct,
forward_params.conv2d.circular_x,
forward_params.conv2d.circular_y,
forward_params.conv2d.scale);
}
for (auto& lora_model : lora_models) {
ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, forward_params, prefix + "weight");
if (out_diff == nullptr) {
continue;
}
out = ggml_add_inplace(ctx, out, out_diff);
}
return out;
}
size_t get_extra_graph_size() override {
size_t lora_tensor_num = 0;
for (auto& lora_model : lora_models) {
lora_tensor_num += lora_model->lora_tensors.size();
}
return LORA_GRAPH_BASE_SIZE + lora_tensor_num * 10;
}
};
#endif // __SD_MODEL_ADAPTER_LORA_HPP__
src/model/adapter/pmid.hpp
+640
View File
@@ -0,0 +1,640 @@
#ifndef __SD_MODEL_ADAPTER_PMID_HPP__
#define __SD_MODEL_ADAPTER_PMID_HPP__
#include "core/ggml_extend.hpp"
#include "model/adapter/lora.hpp"
#include "model/common/block.hpp"
#include "model/te/clip.hpp"
struct FuseBlock : public GGMLBlock {
// network hparams
int in_dim;
int out_dim;
int hidden_dim;
bool use_residue;
public:
FuseBlock(int i_d, int o_d, int h_d, bool use_residue = true)
: in_dim(i_d), out_dim(o_d), hidden_dim(h_d), use_residue(use_residue) {
blocks["fc1"] = std::shared_ptr<GGMLBlock>(new Linear(in_dim, hidden_dim, true));
blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_dim, out_dim, true));
blocks["layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(in_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, channels, h, w]
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
ggml_tensor* r = x;
// x = ggml_ext_layer_norm(ctx, x, ln_w, ln_b);
x = layer_norm->forward(ctx, x);
// x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x), fc1_b);
x = fc1->forward(ctx, x);
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
x = fc2->forward(ctx, x);
// x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x), fc2_b);
if (use_residue)
x = ggml_add(ctx->ggml_ctx, x, r);
return x;
}
};
struct PMFeedForward : public GGMLBlock {
// network hparams
int dim;
public:
PMFeedForward(int d, int multi = 4)
: dim(d) {
int inner_dim = dim * multi;
blocks["0"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks["1"] = std::shared_ptr<GGMLBlock>(new Mlp(dim, inner_dim, dim, false));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["0"]);
auto ff = std::dynamic_pointer_cast<Mlp>(blocks["1"]);
x = norm->forward(ctx, x);
x = ff->forward(ctx, x);
return x;
}
};
struct PerceiverAttention : public GGMLBlock {
// network hparams
float scale; // = dim_head**-0.5
int dim_head; // = dim_head
int heads; // = heads
public:
PerceiverAttention(int dim, int dim_h = 64, int h = 8)
: scale(powf(static_cast<float>(dim_h), -0.5f)), dim_head(dim_h), heads(h) {
int inner_dim = dim_head * heads;
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(dim, inner_dim, false));
blocks["to_kv"] = std::shared_ptr<GGMLBlock>(new Linear(dim, inner_dim * 2, false));
blocks["to_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, dim, false));
}
ggml_tensor* reshape_tensor(ggml_context* ctx,
ggml_tensor* x,
int heads) {
int64_t ne[4];
for (int i = 0; i < 4; ++i)
ne[i] = x->ne[i];
x = ggml_reshape_4d(ctx, x, x->ne[0] / heads, heads, x->ne[1], x->ne[2]);
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));
return x;
}
std::vector<ggml_tensor*> chunk_half(ggml_context* ctx,
ggml_tensor* x) {
auto tlo = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], 0);
auto tli = ggml_view_4d(ctx, x, x->ne[0] / 2, x->ne[1], x->ne[2], x->ne[3], x->nb[1], x->nb[2], x->nb[3], x->nb[0] * x->ne[0] / 2);
return {ggml_cont(ctx, tlo),
ggml_cont(ctx, tli)};
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* latents) {
// x (torch.Tensor): image features
// shape (b, n1, D)
// latent (torch.Tensor): latent features
// shape (b, n2, D)
int64_t ne[4];
for (int i = 0; i < 4; ++i)
ne[i] = latents->ne[i];
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
x = norm1->forward(ctx, x);
latents = norm2->forward(ctx, latents);
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto q = to_q->forward(ctx, latents);
auto kv_input = ggml_concat(ctx->ggml_ctx, x, latents, 1);
auto to_kv = std::dynamic_pointer_cast<Linear>(blocks["to_kv"]);
auto kv = to_kv->forward(ctx, kv_input);
auto k = ggml_view_4d(ctx->ggml_ctx, kv, kv->ne[0] / 2, kv->ne[1], kv->ne[2], kv->ne[3], kv->nb[1] / 2, kv->nb[2] / 2, kv->nb[3] / 2, 0);
auto v = ggml_view_4d(ctx->ggml_ctx, kv, kv->ne[0] / 2, kv->ne[1], kv->ne[2], kv->ne[3], kv->nb[1] / 2, kv->nb[2] / 2, kv->nb[3] / 2, kv->nb[0] * (kv->ne[0] / 2));
k = ggml_cont(ctx->ggml_ctx, k);
v = ggml_cont(ctx->ggml_ctx, v);
q = reshape_tensor(ctx->ggml_ctx, q, heads);
k = reshape_tensor(ctx->ggml_ctx, k, heads);
v = reshape_tensor(ctx->ggml_ctx, v, heads);
scale = 1.f / sqrt(sqrt((float)dim_head));
k = ggml_ext_scale(ctx->ggml_ctx, k, scale, true);
q = ggml_ext_scale(ctx->ggml_ctx, q, scale, true);
// auto weight = ggml_mul_mat(ctx, q, k);
auto weight = ggml_mul_mat(ctx->ggml_ctx, k, q); // NOTE order of mul is opposite to pytorch
// GGML's softmax() is equivalent to pytorch's softmax(x, dim=-1)
// in this case, dimension along which Softmax will be computed is the last dim
// in torch and the first dim in GGML, consistent with the convention that pytorch's
// last dimension (varying most rapidly) corresponds to GGML's first (varying most rapidly).
// weight = ggml_soft_max(ctx, weight);
weight = ggml_soft_max_inplace(ctx->ggml_ctx, weight);
v = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, v));
// auto out = ggml_mul_mat(ctx, weight, v);
auto out = ggml_mul_mat(ctx->ggml_ctx, v, weight); // NOTE order of mul is opposite to pytorch
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));
out = ggml_reshape_3d(ctx->ggml_ctx, out, ne[0], ne[1], ggml_nelements(out) / (ne[0] * ne[1]));
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
out = to_out->forward(ctx, out);
return out;
}
};
struct FacePerceiverResampler : public GGMLBlock {
// network hparams
int depth;
public:
FacePerceiverResampler(int dim = 768,
int d = 4,
int dim_head = 64,
int heads = 16,
int embedding_dim = 1280,
int output_dim = 768,
int ff_mult = 4)
: depth(d) {
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(embedding_dim, dim, true));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, output_dim, true));
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new LayerNorm(output_dim));
for (int i = 0; i < depth; i++) {
std::string name = "layers." + std::to_string(i) + ".0";
blocks[name] = std::shared_ptr<GGMLBlock>(new PerceiverAttention(dim, dim_head, heads));
name = "layers." + std::to_string(i) + ".1";
blocks[name] = std::shared_ptr<GGMLBlock>(new PMFeedForward(dim, ff_mult));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* latents,
ggml_tensor* x) {
// x: [N, channels, h, w]
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
x = proj_in->forward(ctx, x);
for (int i = 0; i < depth; i++) {
std::string name = "layers." + std::to_string(i) + ".0";
auto attn = std::dynamic_pointer_cast<PerceiverAttention>(blocks[name]);
name = "layers." + std::to_string(i) + ".1";
auto ff = std::dynamic_pointer_cast<PMFeedForward>(blocks[name]);
auto t = attn->forward(ctx, x, latents);
latents = ggml_add(ctx->ggml_ctx, t, latents);
t = ff->forward(ctx, latents);
latents = ggml_add(ctx->ggml_ctx, t, latents);
}
latents = proj_out->forward(ctx, latents);
latents = norm_out->forward(ctx, latents);
return latents;
}
};
struct QFormerPerceiver : public GGMLBlock {
// network hparams
int num_tokens;
int cross_attention_dim;
bool use_residul;
public:
QFormerPerceiver(int id_embeddings_dim, int cross_attention_d, int num_t, int embedding_dim = 1024, bool use_r = true, int ratio = 4)
: cross_attention_dim(cross_attention_d), num_tokens(num_t), use_residul(use_r) {
blocks["token_proj"] = std::shared_ptr<GGMLBlock>(new Mlp(id_embeddings_dim,
id_embeddings_dim * ratio,
cross_attention_dim * num_tokens,
true));
blocks["token_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(cross_attention_d));
blocks["perceiver_resampler"] = std::shared_ptr<GGMLBlock>(new FacePerceiverResampler(
cross_attention_dim,
4,
128,
cross_attention_dim / 128,
embedding_dim,
cross_attention_dim,
4));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* last_hidden_state) {
// x: [N, channels, h, w]
auto token_proj = std::dynamic_pointer_cast<Mlp>(blocks["token_proj"]);
auto token_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["token_norm"]);
auto perceiver_resampler = std::dynamic_pointer_cast<FacePerceiverResampler>(blocks["perceiver_resampler"]);
x = token_proj->forward(ctx, x);
int64_t nel = ggml_nelements(x);
x = ggml_reshape_3d(ctx->ggml_ctx, x, cross_attention_dim, num_tokens, nel / (cross_attention_dim * num_tokens));
x = token_norm->forward(ctx, x);
ggml_tensor* out = perceiver_resampler->forward(ctx, x, last_hidden_state);
if (use_residul)
out = ggml_add(ctx->ggml_ctx, x, out);
return out;
}
};
struct FuseModule : public GGMLBlock {
// network hparams
int embed_dim;
public:
FuseModule(int imb_d)
: embed_dim(imb_d) {
blocks["mlp1"] = std::shared_ptr<GGMLBlock>(new FuseBlock(imb_d * 2, imb_d, imb_d, false));
blocks["mlp2"] = std::shared_ptr<GGMLBlock>(new FuseBlock(imb_d, imb_d, imb_d, true));
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(embed_dim));
}
ggml_tensor* fuse_fn(GGMLRunnerContext* ctx,
ggml_tensor* prompt_embeds,
ggml_tensor* id_embeds) {
auto mlp1 = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp1"]);
auto mlp2 = std::dynamic_pointer_cast<FuseBlock>(blocks["mlp2"]);
auto layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm"]);
auto stacked_id_embeds = ggml_concat(ctx->ggml_ctx, prompt_embeds, id_embeds, 0);
stacked_id_embeds = mlp1->forward(ctx, stacked_id_embeds);
stacked_id_embeds = ggml_add(ctx->ggml_ctx, stacked_id_embeds, prompt_embeds);
stacked_id_embeds = mlp2->forward(ctx, stacked_id_embeds);
stacked_id_embeds = layer_norm->forward(ctx, stacked_id_embeds);
return stacked_id_embeds;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* prompt_embeds,
ggml_tensor* id_embeds,
ggml_tensor* class_tokens_mask,
ggml_tensor* class_tokens_mask_pos,
ggml_tensor* left,
ggml_tensor* right) {
// x: [N, channels, h, w]
ggml_tensor* valid_id_embeds = id_embeds;
// # slice out the image token embeddings
ggml_set_name(class_tokens_mask_pos, "class_tokens_mask_pos");
ggml_set_name(prompt_embeds, "prompt_embeds");
ggml_tensor* image_token_embeds = ggml_get_rows(ctx->ggml_ctx, prompt_embeds, class_tokens_mask_pos);
ggml_set_name(image_token_embeds, "image_token_embeds");
valid_id_embeds = ggml_reshape_2d(ctx->ggml_ctx, valid_id_embeds, valid_id_embeds->ne[0],
ggml_nelements(valid_id_embeds) / valid_id_embeds->ne[0]);
ggml_tensor* stacked_id_embeds = fuse_fn(ctx, image_token_embeds, valid_id_embeds);
if (left && right) {
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, left, stacked_id_embeds, 1);
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, stacked_id_embeds, right, 1);
} else if (left) {
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, left, stacked_id_embeds, 1);
} else if (right) {
stacked_id_embeds = ggml_concat(ctx->ggml_ctx, stacked_id_embeds, right, 1);
}
class_tokens_mask = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, class_tokens_mask));
class_tokens_mask = ggml_repeat(ctx->ggml_ctx, class_tokens_mask, prompt_embeds);
prompt_embeds = ggml_mul(ctx->ggml_ctx, prompt_embeds, class_tokens_mask);
ggml_tensor* updated_prompt_embeds = ggml_add(ctx->ggml_ctx, prompt_embeds, stacked_id_embeds);
ggml_set_name(updated_prompt_embeds, "updated_prompt_embeds");
return updated_prompt_embeds;
}
};
struct PhotoMakerIDEncoderBlock : public CLIPVisionModelProjection {
PhotoMakerIDEncoderBlock()
: CLIPVisionModelProjection(OPENAI_CLIP_VIT_L_14) {
blocks["visual_projection_2"] = std::shared_ptr<GGMLBlock>(new Linear(1024, 1280, false));
blocks["fuse_module"] = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* id_pixel_values,
ggml_tensor* prompt_embeds,
ggml_tensor* class_tokens_mask,
ggml_tensor* class_tokens_mask_pos,
ggml_tensor* left,
ggml_tensor* right) {
// x: [N, channels, h, w]
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
auto visual_projection = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);
auto visual_projection_2 = std::dynamic_pointer_cast<Linear>(blocks["visual_projection_2"]);
auto fuse_module = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
ggml_tensor* shared_id_embeds = vision_model->forward(ctx, id_pixel_values); // [N, hidden_size]
ggml_tensor* id_embeds = visual_projection->forward(ctx, shared_id_embeds); // [N, proj_dim(768)]
ggml_tensor* id_embeds_2 = visual_projection_2->forward(ctx, shared_id_embeds); // [N, 1280]
id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 2, 0, 1, 3));
id_embeds_2 = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds_2, 2, 0, 1, 3));
id_embeds = ggml_concat(ctx->ggml_ctx, id_embeds, id_embeds_2, 2); // [batch_size, seq_length, 1, 2048] check whether concat at dim 2 is right
id_embeds = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, id_embeds, 1, 2, 0, 3));
ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
prompt_embeds,
id_embeds,
class_tokens_mask,
class_tokens_mask_pos,
left, right);
return updated_prompt_embeds;
}
};
struct PhotoMakerIDEncoder_CLIPInsightfaceExtendtokenBlock : public CLIPVisionModelProjection {
int cross_attention_dim;
int num_tokens;
PhotoMakerIDEncoder_CLIPInsightfaceExtendtokenBlock(int id_embeddings_dim = 512)
: CLIPVisionModelProjection(OPENAI_CLIP_VIT_L_14),
cross_attention_dim(2048),
num_tokens(2) {
blocks["visual_projection_2"] = std::shared_ptr<GGMLBlock>(new Linear(1024, 1280, false));
blocks["fuse_module"] = std::shared_ptr<GGMLBlock>(new FuseModule(2048));
blocks["qformer_perceiver"] = std::shared_ptr<GGMLBlock>(new QFormerPerceiver(id_embeddings_dim,
cross_attention_dim,
num_tokens));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* id_pixel_values,
ggml_tensor* prompt_embeds,
ggml_tensor* class_tokens_mask,
ggml_tensor* class_tokens_mask_pos,
ggml_tensor* id_embeds,
ggml_tensor* left,
ggml_tensor* right) {
// x: [N, channels, h, w]
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
auto fuse_module = std::dynamic_pointer_cast<FuseModule>(blocks["fuse_module"]);
auto qformer_perceiver = std::dynamic_pointer_cast<QFormerPerceiver>(blocks["qformer_perceiver"]);
// ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values); // [N, hidden_size]
ggml_tensor* last_hidden_state = vision_model->forward(ctx, id_pixel_values, false); // [N, hidden_size]
id_embeds = qformer_perceiver->forward(ctx, id_embeds, last_hidden_state);
ggml_tensor* updated_prompt_embeds = fuse_module->forward(ctx,
prompt_embeds,
id_embeds,
class_tokens_mask,
class_tokens_mask_pos,
left, right);
return updated_prompt_embeds;
}
};
struct PhotoMakerIDEncoder : public GGMLRunner {
public:
SDVersion version = VERSION_SDXL;
PMVersion pm_version = PM_VERSION_1;
PhotoMakerIDEncoderBlock id_encoder;
PhotoMakerIDEncoder_CLIPInsightfaceExtendtokenBlock id_encoder2;
float style_strength;
std::vector<float> ctm;
std::vector<ggml_fp16_t> ctmf16;
std::vector<int> ctmpos;
std::vector<ggml_fp16_t> zeros_left_16;
std::vector<float> zeros_left;
std::vector<ggml_fp16_t> zeros_right_16;
std::vector<float> zeros_right;
public:
PhotoMakerIDEncoder(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
SDVersion version = VERSION_SDXL,
PMVersion pm_v = PM_VERSION_1,
float sty = 20.f)
: GGMLRunner(backend, params_backend),
version(version),
pm_version(pm_v),
style_strength(sty) {
if (pm_version == PM_VERSION_1) {
id_encoder.init(params_ctx, tensor_storage_map, prefix);
} else if (pm_version == PM_VERSION_2) {
id_encoder2.init(params_ctx, tensor_storage_map, prefix);
}
}
std::string get_desc() {
return "pmid";
}
PMVersion get_version() const {
return pm_version;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
if (pm_version == PM_VERSION_1)
id_encoder.get_param_tensors(tensors, prefix);
else if (pm_version == PM_VERSION_2)
id_encoder2.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& id_pixel_values_tensor,
const sd::Tensor<float>& prompt_embeds_tensor,
std::vector<bool>& class_tokens_mask,
const sd::Tensor<float>& id_embeds_tensor = {}) {
ctm.clear();
ctmf16.clear();
ctmpos.clear();
zeros_left.clear();
zeros_left_16.clear();
zeros_right.clear();
zeros_right_16.clear();
auto runner_ctx = get_context();
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
ggml_tensor* id_pixel_values = make_input(id_pixel_values_tensor);
ggml_tensor* prompt_embeds = make_input(prompt_embeds_tensor);
ggml_tensor* id_embeds = make_optional_input(id_embeds_tensor);
int64_t hidden_size = prompt_embeds->ne[0];
int64_t seq_length = prompt_embeds->ne[1];
ggml_type type = GGML_TYPE_F32;
ggml_tensor* class_tokens_mask_d = ggml_new_tensor_1d(runner_ctx.ggml_ctx, type, class_tokens_mask.size());
ggml_tensor* left = nullptr;
ggml_tensor* right = nullptr;
for (int i = 0; i < class_tokens_mask.size(); i++) {
if (class_tokens_mask[i]) {
// printf(" 1,");
ctm.push_back(0.f); // here use 0.f instead of 1.f to make a scale mask
ctmf16.push_back(ggml_fp32_to_fp16(0.f)); // here use 0.f instead of 1.f to make a scale mask
ctmpos.push_back(i);
} else {
// printf(" 0,");
ctm.push_back(1.f); // here use 1.f instead of 0.f to make a scale mask
ctmf16.push_back(ggml_fp32_to_fp16(1.f)); // here use 0.f instead of 1.f to make a scale mask
}
}
// printf("\n");
if (ctmpos[0] > 0) {
// left = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type, hidden_size, 1, ctmpos[0]);
left = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type, hidden_size, ctmpos[0], 1);
}
if (ctmpos[ctmpos.size() - 1] < seq_length - 1) {
// right = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type,
// hidden_size, 1, seq_length - ctmpos[ctmpos.size() - 1] - 1);
right = ggml_new_tensor_3d(runner_ctx.ggml_ctx, type,
hidden_size, seq_length - ctmpos[ctmpos.size() - 1] - 1, 1);
}
ggml_tensor* class_tokens_mask_pos = ggml_new_tensor_1d(runner_ctx.ggml_ctx, GGML_TYPE_I32, ctmpos.size());
{
if (type == GGML_TYPE_F16)
set_backend_tensor_data(class_tokens_mask_d, ctmf16.data());
else
set_backend_tensor_data(class_tokens_mask_d, ctm.data());
set_backend_tensor_data(class_tokens_mask_pos, ctmpos.data());
if (left) {
if (type == GGML_TYPE_F16) {
for (int i = 0; i < ggml_nelements(left); ++i)
zeros_left_16.push_back(ggml_fp32_to_fp16(0.f));
set_backend_tensor_data(left, zeros_left_16.data());
} else {
for (int i = 0; i < ggml_nelements(left); ++i)
zeros_left.push_back(0.f);
set_backend_tensor_data(left, zeros_left.data());
}
}
if (right) {
if (type == GGML_TYPE_F16) {
for (int i = 0; i < ggml_nelements(right); ++i)
zeros_right_16.push_back(ggml_fp32_to_fp16(0.f));
set_backend_tensor_data(right, zeros_right_16.data());
} else {
for (int i = 0; i < ggml_nelements(right); ++i)
zeros_right.push_back(0.f);
set_backend_tensor_data(right, zeros_right.data());
}
}
}
ggml_tensor* updated_prompt_embeds = nullptr;
if (pm_version == PM_VERSION_1)
updated_prompt_embeds = id_encoder.forward(&runner_ctx,
id_pixel_values,
prompt_embeds,
class_tokens_mask_d,
class_tokens_mask_pos,
left, right);
else if (pm_version == PM_VERSION_2)
updated_prompt_embeds = id_encoder2.forward(&runner_ctx,
id_pixel_values,
prompt_embeds,
class_tokens_mask_d,
class_tokens_mask_pos,
id_embeds,
left, right);
ggml_build_forward_expand(gf, updated_prompt_embeds);
return gf;
}
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& id_pixel_values,
const sd::Tensor<float>& prompt_embeds,
const sd::Tensor<float>& id_embeds,
std::vector<bool>& class_tokens_mask) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
}
};
struct PhotoMakerIDEmbed : public GGMLRunner {
std::map<std::string, ggml_tensor*> tensors;
std::string file_path;
ModelLoader* model_loader;
bool load_failed = false;
bool applied = false;
PhotoMakerIDEmbed(ggml_backend_t backend,
ggml_backend_t params_backend,
ModelLoader* ml,
const std::string& file_path = "",
const std::string& prefix = "")
: file_path(file_path), GGMLRunner(backend, params_backend), model_loader(ml) {
if (!model_loader->init_from_file_and_convert_name(file_path, prefix)) {
load_failed = true;
}
}
std::string get_desc() {
return "id_embeds";
}
bool load_from_file(bool filter_tensor, int n_threads) {
LOG_INFO("loading PhotoMaker ID Embeds from '%s'", file_path.c_str());
if (load_failed) {
LOG_ERROR("init photomaker id embed from file failed: '%s'", file_path.c_str());
return false;
}
bool dry_run = true;
std::mutex tensor_mutex;
auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
const std::string& name = tensor_storage.name;
if (filter_tensor && !contains(name, "pmid.id_embeds")) {
// LOG_INFO("skipping LoRA tesnor '%s'", name.c_str());
return true;
}
if (dry_run) {
std::lock_guard<std::mutex> lock(tensor_mutex);
ggml_tensor* real = ggml_new_tensor(params_ctx,
tensor_storage.type,
tensor_storage.n_dims,
tensor_storage.ne);
tensors[name] = real;
} else {
auto real = tensors[name];
*dst_tensor = real;
}
return true;
};
model_loader->load_tensors(on_new_tensor_cb, n_threads);
if (!alloc_params_buffer()) {
LOG_ERROR("PhotoMaker ID embeds buffer allocation failed");
return false;
}
dry_run = false;
model_loader->load_tensors(on_new_tensor_cb, n_threads);
LOG_DEBUG("finished loading PhotoMaker ID Embeds ");
return true;
}
ggml_tensor* get() {
std::map<std::string, ggml_tensor*>::iterator pos;
pos = tensors.find("pmid.id_embeds");
if (pos != tensors.end())
return pos->second;
return nullptr;
}
};
#endif // __SD_MODEL_ADAPTER_PMID_HPP__
src/model/common/block.hpp
+637
View File
@@ -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__
src/model/common/rope.hpp
+905
View File
@@ -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__
src/model/diffusion/anima.hpp
+718
View File
@@ -0,0 +1,718 @@
#ifndef __SD_MODEL_DIFFUSION_ANIMA_HPP__
#define __SD_MODEL_DIFFUSION_ANIMA_HPP__
#include <algorithm>
#include <cmath>
#include <memory>
#include <utility>
#include <vector>
#include "model/common/block.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
namespace Anima {
constexpr int ANIMA_GRAPH_SIZE = 65536;
struct AnimaConfig {
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t hidden_size = 2048;
int64_t text_embed_dim = 1024;
int64_t num_heads = 16;
int64_t head_dim = 128;
int patch_size = 2;
int64_t num_layers = 28;
std::vector<int> axes_dim = {44, 42, 42};
int theta = 10000;
static AnimaConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
AnimaConfig config;
int64_t detected_layers = 0;
std::string layer_tag = prefix.empty() ? "blocks." : prefix + ".blocks.";
for (const auto& [name, _] : tensor_storage_map) {
size_t pos = name.find(layer_tag);
if (pos == std::string::npos) {
continue;
}
size_t start = pos + layer_tag.size();
size_t end = name.find('.', start);
if (end == std::string::npos) {
continue;
}
int64_t layer_id = atoll(name.substr(start, end - start).c_str());
detected_layers = std::max(detected_layers, layer_id + 1);
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
LOG_DEBUG("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64,
config.num_layers,
config.hidden_size,
config.num_heads,
config.head_dim);
}
return config;
}
};
__STATIC_INLINE__ ggml_tensor* apply_gate(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* gate) {
gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]); // [N, 1, C]
return ggml_mul(ctx, x, gate);
}
struct XEmbedder : public GGMLBlock {
public:
XEmbedder(int64_t in_dim, int64_t out_dim) {
blocks["proj.1"] = std::make_shared<Linear>(in_dim, out_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj.1"]);
return proj->forward(ctx, x);
}
};
struct TimestepEmbedder : public GGMLBlock {
public:
TimestepEmbedder(int64_t in_dim, int64_t out_dim) {
blocks["1.linear_1"] = std::make_shared<Linear>(in_dim, in_dim, false);
blocks["1.linear_2"] = std::make_shared<Linear>(in_dim, out_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["1.linear_2"]);
x = linear_1->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
x = linear_2->forward(ctx, x);
return x;
}
};
struct AdaLayerNormZero : public GGMLBlock {
protected:
int64_t in_features;
public:
AdaLayerNormZero(int64_t in_features, int64_t hidden_features = 256)
: in_features(in_features) {
blocks["norm"] = std::make_shared<LayerNorm>(in_features, 1e-6f, false, false);
blocks["1"] = std::make_shared<Linear>(in_features, hidden_features, false);
blocks["2"] = std::make_shared<Linear>(hidden_features, 3 * in_features, false);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* hidden_states,
ggml_tensor* embedded_timestep,
ggml_tensor* temb = nullptr) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
auto emb = ggml_silu(ctx->ggml_ctx, embedded_timestep);
emb = linear_1->forward(ctx, emb);
emb = linear_2->forward(ctx, emb); // [N, 3*C]
if (temb != nullptr) {
emb = ggml_add(ctx->ggml_ctx, emb, temb);
}
auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 3, 0);
auto shift = emb_chunks[0];
auto scale = emb_chunks[1];
auto gate = emb_chunks[2];
auto x = norm->forward(ctx, hidden_states);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
return {x, gate};
}
};
struct AdaLayerNorm : public GGMLBlock {
protected:
int64_t embedding_dim;
public:
AdaLayerNorm(int64_t in_features, int64_t hidden_features = 256)
: embedding_dim(in_features) {
blocks["norm"] = std::make_shared<LayerNorm>(in_features, 1e-6f, false, false);
blocks["1"] = std::make_shared<Linear>(in_features, hidden_features, false);
blocks["2"] = std::make_shared<Linear>(hidden_features, 2 * in_features, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* hidden_states,
ggml_tensor* embedded_timestep,
ggml_tensor* temb = nullptr) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
auto emb = ggml_silu(ctx->ggml_ctx, embedded_timestep);
emb = linear_1->forward(ctx, emb);
emb = linear_2->forward(ctx, emb); // [N, 2*C]
if (temb != nullptr) {
auto temb_2c = ggml_view_2d(ctx->ggml_ctx, temb, 2 * embedding_dim, temb->ne[1], temb->nb[1], 0);
emb = ggml_add(ctx->ggml_ctx, emb, temb_2c);
}
auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
auto shift = emb_chunks[0];
auto scale = emb_chunks[1];
auto x = norm->forward(ctx, hidden_states);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
return x;
}
};
struct AnimaAttention : public GGMLBlock {
protected:
int64_t num_heads;
int64_t head_dim;
std::string out_proj_name;
public:
AnimaAttention(int64_t query_dim,
int64_t context_dim,
int64_t num_heads,
int64_t head_dim,
const std::string& out_proj_name = "output_proj")
: num_heads(num_heads), head_dim(head_dim), out_proj_name(out_proj_name) {
int64_t inner_dim = num_heads * head_dim;
blocks["q_proj"] = std::make_shared<Linear>(query_dim, inner_dim, false);
blocks["k_proj"] = std::make_shared<Linear>(context_dim, inner_dim, false);
blocks["v_proj"] = std::make_shared<Linear>(context_dim, inner_dim, false);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks[this->out_proj_name] = std::make_shared<Linear>(inner_dim, query_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* hidden_states,
ggml_tensor* encoder_hidden_states = nullptr,
ggml_tensor* pe_q = nullptr,
ggml_tensor* pe_k = nullptr) {
if (encoder_hidden_states == nullptr) {
encoder_hidden_states = hidden_states;
}
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[out_proj_name]);
auto q = q_proj->forward(ctx, hidden_states);
auto k = k_proj->forward(ctx, encoder_hidden_states);
auto v = v_proj->forward(ctx, encoder_hidden_states);
int64_t N = q->ne[2];
int64_t L_q = q->ne[1];
int64_t L_k = k->ne[1];
auto q4 = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, L_q, N); // [N, L_q, H, D]
auto k4 = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, L_k, N); // [N, L_k, H, D]
auto v4 = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, L_k, N); // [N, L_k, H, D]
q4 = q_norm->forward(ctx, q4);
k4 = k_norm->forward(ctx, k4);
ggml_tensor* attn_out = nullptr;
if (pe_q != nullptr || pe_k != nullptr) {
if (pe_q == nullptr) {
pe_q = pe_k;
}
if (pe_k == nullptr) {
pe_k = pe_q;
}
auto q_rope = Rope::apply_rope(ctx->ggml_ctx, q4, pe_q, false);
auto k_rope = Rope::apply_rope(ctx->ggml_ctx, k4, pe_k, false);
attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q_rope,
k_rope,
v4,
num_heads,
nullptr,
true,
ctx->flash_attn_enabled);
} else {
auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q_flat,
k_flat,
v,
num_heads,
nullptr,
false,
ctx->flash_attn_enabled);
}
return out_proj->forward(ctx, attn_out);
}
};
struct AnimaMLP : public GGMLBlock {
public:
AnimaMLP(int64_t dim, int64_t hidden_dim) {
blocks["layer1"] = std::make_shared<Linear>(dim, hidden_dim, false);
blocks["layer2"] = std::make_shared<Linear>(hidden_dim, dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto layer1 = std::dynamic_pointer_cast<Linear>(blocks["layer1"]);
auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["layer2"]);
x = layer1->forward(ctx, x);
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
x = layer2->forward(ctx, x);
return x;
}
};
struct AdapterMLP : public GGMLBlock {
public:
AdapterMLP(int64_t dim, int64_t hidden_dim) {
blocks["0"] = std::make_shared<Linear>(dim, hidden_dim, true);
blocks["2"] = std::make_shared<Linear>(hidden_dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto layer0 = std::dynamic_pointer_cast<Linear>(blocks["0"]);
auto layer2 = std::dynamic_pointer_cast<Linear>(blocks["2"]);
x = layer0->forward(ctx, x);
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
x = layer2->forward(ctx, x);
return x;
}
};
struct LLMAdapterBlock : public GGMLBlock {
public:
LLMAdapterBlock(int64_t model_dim = 1024, int64_t source_dim = 1024, int64_t num_heads = 16, int64_t head_dim = 64) {
blocks["norm_self_attn"] = std::make_shared<RMSNorm>(model_dim, 1e-6f);
blocks["self_attn"] = std::make_shared<AnimaAttention>(model_dim, model_dim, num_heads, head_dim, "o_proj");
blocks["norm_cross_attn"] = std::make_shared<RMSNorm>(model_dim, 1e-6f);
blocks["cross_attn"] = std::make_shared<AnimaAttention>(model_dim, source_dim, num_heads, head_dim, "o_proj");
blocks["norm_mlp"] = std::make_shared<RMSNorm>(model_dim, 1e-6f);
blocks["mlp"] = std::make_shared<AdapterMLP>(model_dim, model_dim * 4);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* context,
ggml_tensor* target_pe,
ggml_tensor* context_pe) {
auto norm_self_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_self_attn"]);
auto self_attn = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
auto norm_cross_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_cross_attn"]);
auto cross_attn = std::dynamic_pointer_cast<AnimaAttention>(blocks["cross_attn"]);
auto norm_mlp = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_mlp"]);
auto mlp = std::dynamic_pointer_cast<AdapterMLP>(blocks["mlp"]);
auto h = norm_self_attn->forward(ctx, x);
h = self_attn->forward(ctx, h, nullptr, target_pe, target_pe);
x = ggml_add(ctx->ggml_ctx, x, h);
h = norm_cross_attn->forward(ctx, x);
h = cross_attn->forward(ctx, h, context, target_pe, context_pe);
x = ggml_add(ctx->ggml_ctx, x, h);
h = norm_mlp->forward(ctx, x);
h = mlp->forward(ctx, h);
x = ggml_add(ctx->ggml_ctx, x, h);
return x;
}
};
struct LLMAdapter : public GGMLBlock {
protected:
int num_layers;
public:
LLMAdapter(int64_t source_dim = 1024,
int64_t target_dim = 1024,
int64_t model_dim = 1024,
int num_layers = 6,
int num_heads = 16)
: num_layers(num_layers) {
int64_t head_dim = model_dim / num_heads;
blocks["embed"] = std::make_shared<Embedding>(32128, target_dim);
for (int i = 0; i < num_layers; i++) {
blocks["blocks." + std::to_string(i)] =
std::make_shared<LLMAdapterBlock>(model_dim, source_dim, num_heads, head_dim);
}
blocks["out_proj"] = std::make_shared<Linear>(model_dim, target_dim, true);
blocks["norm"] = std::make_shared<RMSNorm>(target_dim, 1e-6f);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* source_hidden_states,
ggml_tensor* target_input_ids,
ggml_tensor* target_pe,
ggml_tensor* source_pe) {
GGML_ASSERT(target_input_ids != nullptr);
if (ggml_n_dims(target_input_ids) == 1) {
target_input_ids = ggml_reshape_2d(ctx->ggml_ctx, target_input_ids, target_input_ids->ne[0], 1);
}
auto embed = std::dynamic_pointer_cast<Embedding>(blocks["embed"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
auto x = embed->forward(ctx, target_input_ids); // [N, target_len, target_dim]
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<LLMAdapterBlock>(blocks["blocks." + std::to_string(i)]);
x = block->forward(ctx, x, source_hidden_states, target_pe, source_pe);
}
x = out_proj->forward(ctx, x);
x = norm->forward(ctx, x);
return x;
}
};
struct TransformerBlock : public GGMLBlock {
public:
TransformerBlock(int64_t hidden_size,
int64_t text_embed_dim,
int64_t num_heads,
int64_t head_dim,
int64_t mlp_ratio = 4,
int64_t adaln_lora_dim = 256) {
blocks["adaln_modulation_self_attn"] = std::make_shared<AdaLayerNormZero>(hidden_size, adaln_lora_dim);
blocks["self_attn"] = std::make_shared<AnimaAttention>(hidden_size, hidden_size, num_heads, head_dim);
blocks["adaln_modulation_cross_attn"] = std::make_shared<AdaLayerNormZero>(hidden_size, adaln_lora_dim);
blocks["cross_attn"] = std::make_shared<AnimaAttention>(hidden_size, text_embed_dim, num_heads, head_dim);
blocks["adaln_modulation_mlp"] = std::make_shared<AdaLayerNormZero>(hidden_size, adaln_lora_dim);
blocks["mlp"] = std::make_shared<AnimaMLP>(hidden_size, hidden_size * mlp_ratio);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* hidden_states,
ggml_tensor* encoder_hidden_states,
ggml_tensor* embedded_timestep,
ggml_tensor* temb,
ggml_tensor* image_pe) {
auto norm1 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_self_attn"]);
auto attn1 = std::dynamic_pointer_cast<AnimaAttention>(blocks["self_attn"]);
auto norm2 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_cross_attn"]);
auto attn2 = std::dynamic_pointer_cast<AnimaAttention>(blocks["cross_attn"]);
auto norm3 = std::dynamic_pointer_cast<AdaLayerNormZero>(blocks["adaln_modulation_mlp"]);
auto mlp = std::dynamic_pointer_cast<AnimaMLP>(blocks["mlp"]);
auto [normed1, gate1] = norm1->forward(ctx, hidden_states, embedded_timestep, temb);
auto h = attn1->forward(ctx, normed1, nullptr, image_pe, image_pe);
hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, apply_gate(ctx->ggml_ctx, h, gate1));
auto [normed2, gate2] = norm2->forward(ctx, hidden_states, embedded_timestep, temb);
h = attn2->forward(ctx, normed2, encoder_hidden_states, nullptr, nullptr);
hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, apply_gate(ctx->ggml_ctx, h, gate2));
auto [normed3, gate3] = norm3->forward(ctx, hidden_states, embedded_timestep, temb);
h = mlp->forward(ctx, normed3);
hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, apply_gate(ctx->ggml_ctx, h, gate3));
return hidden_states;
}
};
struct FinalLayer : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t patch_size;
int64_t out_channels;
public:
FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels)
: hidden_size(hidden_size), patch_size(patch_size), out_channels(out_channels) {
blocks["adaln_modulation"] = std::make_shared<AdaLayerNorm>(hidden_size, 256);
blocks["linear"] = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* hidden_states,
ggml_tensor* embedded_timestep,
ggml_tensor* temb) {
auto adaln = std::dynamic_pointer_cast<AdaLayerNorm>(blocks["adaln_modulation"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
hidden_states = adaln->forward(ctx, hidden_states, embedded_timestep, temb);
hidden_states = linear->forward(ctx, hidden_states);
return hidden_states;
}
};
struct AnimaNet : public GGMLBlock {
public:
AnimaConfig config;
public:
AnimaNet() = default;
explicit AnimaNet(AnimaConfig config)
: config(config) {
blocks["x_embedder"] = std::make_shared<XEmbedder>((config.in_channels + 1) * config.patch_size * config.patch_size, config.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(config.hidden_size, config.hidden_size * 3);
blocks["t_embedding_norm"] = std::make_shared<RMSNorm>(config.hidden_size, 1e-6f);
for (int i = 0; i < config.num_layers; i++) {
blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(config.hidden_size,
config.text_embed_dim,
config.num_heads,
config.head_dim);
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* encoder_hidden_states,
ggml_tensor* image_pe,
ggml_tensor* t5_ids = nullptr,
ggml_tensor* t5_weights = nullptr,
ggml_tensor* adapter_q_pe = nullptr,
ggml_tensor* adapter_k_pe = nullptr) {
GGML_ASSERT(x->ne[3] == 1);
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto t_embedding_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["t_embedding_norm"]);
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
auto llm_adapter = std::dynamic_pointer_cast<LLMAdapter>(blocks["llm_adapter"]);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W]
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
x = x_embedder->forward(ctx, x);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
auto temb = t_embedder->forward(ctx, timestep_proj);
auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj);
if (t5_ids != nullptr) {
auto adapted_context = llm_adapter->forward(ctx, encoder_hidden_states, t5_ids, adapter_q_pe, adapter_k_pe);
if (t5_weights != nullptr) {
auto w = t5_weights;
if (ggml_n_dims(w) == 1) {
w = ggml_reshape_3d(ctx->ggml_ctx, w, 1, w->ne[0], 1);
}
w = ggml_repeat_4d(ctx->ggml_ctx, w, adapted_context->ne[0], adapted_context->ne[1], adapted_context->ne[2], 1);
adapted_context = ggml_mul(ctx->ggml_ctx, adapted_context, w);
}
if (adapted_context->ne[1] < 512) {
auto pad_ctx = ggml_ext_zeros(ctx->ggml_ctx,
adapted_context->ne[0],
512 - adapted_context->ne[1],
adapted_context->ne[2],
1);
adapted_context = ggml_concat(ctx->ggml_ctx, adapted_context, pad_ctx, 1);
} else if (adapted_context->ne[1] > 512) {
adapted_context = ggml_ext_slice(ctx->ggml_ctx, adapted_context, 1, 0, 512);
}
encoder_hidden_states = adapted_context;
}
sd::ggml_graph_cut::mark_graph_cut(x, "anima.prelude", "x");
sd::ggml_graph_cut::mark_graph_cut(embedded_timestep, "anima.prelude", "embedded_timestep");
sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb");
sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context");
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
}
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false); // [N, C, H, W]
return x;
}
};
struct AnimaRunner : public DiffusionModelRunner {
public:
std::vector<float> image_pe_vec;
std::vector<float> adapter_q_pe_vec;
std::vector<float> adapter_k_pe_vec;
AnimaConfig config;
AnimaNet net;
AnimaRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "model.diffusion_model")
: DiffusionModelRunner(backend, params_backend, prefix),
config(AnimaConfig::detect_from_weights(tensor_storage_map, prefix + ".net")) {
net = AnimaNet(config);
net.init(params_ctx, tensor_storage_map, prefix + ".net");
}
std::string get_desc() override {
return "anima";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
net.get_param_tensors(tensors, prefix + ".net");
}
static std::vector<float> gen_1d_rope_pe_vec(int64_t seq_len, int dim, float theta = 10000.f) {
std::vector<float> pos(seq_len);
for (int64_t i = 0; i < seq_len; i++) {
pos[i] = static_cast<float>(i);
}
auto rope_emb = Rope::rope(pos, dim, theta);
return Rope::flatten(rope_emb);
}
static float calc_ntk_factor(float extrapolation_ratio, int axis_dim) {
if (extrapolation_ratio == 1.0f || axis_dim <= 2) {
return 1.0f;
}
return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
}
static std::vector<float> gen_anima_image_pe_vec(int bs,
int h,
int w,
int patch_size,
int theta,
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio) {
static const std::vector<ggml_tensor*> empty_ref_latents;
auto ids = Rope::gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
0,
{},
empty_ref_latents,
false,
1.0f,
false);
std::vector<float> axis_thetas = {
static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
};
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor = {},
const sd::Tensor<int32_t>& t5_ids_tensor = {},
const sd::Tensor<float>& t5_weights_tensor = {}) {
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor);
ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
GGML_ASSERT(x->ne[3] == 1);
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
int64_t pad_h = (config.patch_size - x->ne[1] % config.patch_size) % config.patch_size;
int64_t pad_w = (config.patch_size - x->ne[0] % config.patch_size) % config.patch_size;
int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w;
image_pe_vec = gen_anima_image_pe_vec(1,
static_cast<int>(h_pad),
static_cast<int>(w_pad),
static_cast<int>(config.patch_size),
config.theta,
config.axes_dim,
4.0f,
4.0f,
1.0f);
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
set_backend_tensor_data(image_pe, image_pe_vec.data());
ggml_tensor* adapter_q_pe = nullptr;
ggml_tensor* adapter_k_pe = nullptr;
if (t5_ids != nullptr) {
int64_t target_len = t5_ids->ne[0];
int64_t source_len = context->ne[1];
adapter_q_pe_vec = gen_1d_rope_pe_vec(target_len, 64, 10000.f);
adapter_k_pe_vec = gen_1d_rope_pe_vec(source_len, 64, 10000.f);
int64_t target_pos_len = static_cast<int64_t>(adapter_q_pe_vec.size()) / (2 * 2 * 32);
int64_t source_pos_len = static_cast<int64_t>(adapter_k_pe_vec.size()) / (2 * 2 * 32);
adapter_q_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, target_pos_len);
adapter_k_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, source_pos_len);
set_backend_tensor_data(adapter_q_pe, adapter_q_pe_vec.data());
set_backend_tensor_data(adapter_k_pe, adapter_k_pe_vec.data());
}
auto runner_ctx = get_context();
auto out = net.forward(&runner_ctx,
x,
timesteps,
context,
image_pe,
t5_ids,
t5_weights,
adapter_q_pe,
adapter_k_pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context = {},
const sd::Tensor<int32_t>& t5_ids = {},
const sd::Tensor<float>& t5_weights = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, t5_ids, t5_weights);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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);
const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
tensor_or_empty(extra->t5_ids),
tensor_or_empty(extra->t5_weights));
}
};
} // namespace Anima
#endif // __SD_MODEL_DIFFUSION_ANIMA_HPP__
src/model/diffusion/control.hpp
+487
View File
@@ -0,0 +1,487 @@
#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
#define __SD_MODEL_DIFFUSION_CONTROL_HPP__
#include "model.h"
#include "model/common/block.hpp"
#define CONTROL_NET_GRAPH_SIZE 1536
/*
=================================== ControlNet ===================================
Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/cldm/cldm.py
*/
class ControlNetBlock : public GGMLBlock {
protected:
SDVersion version = VERSION_SD1;
// network hparams
int in_channels = 4;
int out_channels = 4;
int hint_channels = 3;
int num_res_blocks = 2;
std::vector<int> attention_resolutions = {4, 2, 1};
std::vector<int> channel_mult = {1, 2, 4, 4};
std::vector<int> transformer_depth = {1, 1, 1, 1};
int time_embed_dim = 1280; // model_channels*4
int num_heads = 8;
int num_head_channels = -1; // channels // num_heads
int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
bool use_linear_projection = false;
public:
int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL
ControlNetBlock(SDVersion version = VERSION_SD1)
: version(version) {
if (sd_version_is_sd2(version)) {
context_dim = 1024;
num_head_channels = 64;
num_heads = -1;
} else if (sd_version_is_sdxl(version)) {
context_dim = 2048;
attention_resolutions = {4, 2};
channel_mult = {1, 2, 4};
transformer_depth = {1, 2, 10};
num_head_channels = 64;
num_heads = -1;
} else if (version == VERSION_SVD) {
in_channels = 8;
out_channels = 4;
context_dim = 1024;
adm_in_channels = 768;
num_head_channels = 64;
num_heads = -1;
}
blocks["time_embed.0"] = std::shared_ptr<GGMLBlock>(new Linear(model_channels, time_embed_dim));
// time_embed_1 is nn.SiLU()
blocks["time_embed.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
if (sd_version_is_sdxl(version) || version == VERSION_SVD) {
blocks["label_emb.0.0"] = std::shared_ptr<GGMLBlock>(new Linear(adm_in_channels, time_embed_dim));
// label_emb_1 is nn.SiLU()
blocks["label_emb.0.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
}
// input_blocks
blocks["input_blocks.0.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, model_channels, {3, 3}, {1, 1}, {1, 1}));
std::vector<int> input_block_chans;
input_block_chans.push_back(model_channels);
int ch = model_channels;
int input_block_idx = 0;
int ds = 1;
auto get_resblock = [&](int64_t channels, int64_t emb_channels, int64_t out_channels) -> ResBlock* {
return new ResBlock(channels, emb_channels, out_channels);
};
auto get_attention_layer = [&](int64_t in_channels,
int64_t n_head,
int64_t d_head,
int64_t depth,
int64_t context_dim) -> SpatialTransformer* {
return new SpatialTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear_projection);
};
auto make_zero_conv = [&](int64_t channels) {
return new Conv2d(channels, channels, {1, 1});
};
blocks["zero_convs.0.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(model_channels));
blocks["input_hint_block.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(hint_channels, 16, {3, 3}, {1, 1}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(16, 16, {3, 3}, {1, 1}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(16, 32, {3, 3}, {2, 2}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.6"] = std::shared_ptr<GGMLBlock>(new Conv2d(32, 32, {3, 3}, {1, 1}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.8"] = std::shared_ptr<GGMLBlock>(new Conv2d(32, 96, {3, 3}, {2, 2}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.10"] = std::shared_ptr<GGMLBlock>(new Conv2d(96, 96, {3, 3}, {1, 1}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.12"] = std::shared_ptr<GGMLBlock>(new Conv2d(96, 256, {3, 3}, {2, 2}, {1, 1}));
// nn.SiLU()
blocks["input_hint_block.14"] = std::shared_ptr<GGMLBlock>(new Conv2d(256, model_channels, {3, 3}, {1, 1}, {1, 1}));
size_t len_mults = channel_mult.size();
for (int i = 0; i < len_mults; i++) {
int mult = channel_mult[i];
for (int j = 0; j < num_res_blocks; j++) {
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
blocks[name] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, mult * model_channels));
ch = mult * model_channels;
if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
int n_head = num_heads;
int d_head = ch / num_heads;
if (num_head_channels != -1) {
d_head = num_head_channels;
n_head = ch / d_head;
}
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
blocks[name] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
n_head,
d_head,
transformer_depth[i],
context_dim));
}
blocks["zero_convs." + std::to_string(input_block_idx) + ".0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
input_block_chans.push_back(ch);
}
if (i != len_mults - 1) {
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
blocks[name] = std::shared_ptr<GGMLBlock>(new DownSampleBlock(ch, ch));
blocks["zero_convs." + std::to_string(input_block_idx) + ".0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
input_block_chans.push_back(ch);
ds *= 2;
}
}
// middle blocks
int n_head = num_heads;
int d_head = ch / num_heads;
if (num_head_channels != -1) {
d_head = num_head_channels;
n_head = ch / d_head;
}
blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
n_head,
d_head,
transformer_depth[transformer_depth.size() - 1],
context_dim));
blocks["middle_block.2"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
// middle_block_out
blocks["middle_block_out.0"] = std::shared_ptr<GGMLBlock>(make_zero_conv(ch));
}
ggml_tensor* resblock_forward(std::string name,
GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* emb) {
auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
return block->forward(ctx, x, emb);
}
ggml_tensor* attention_layer_forward(std::string name,
GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* context) {
auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
return block->forward(ctx, x, context);
}
ggml_tensor* input_hint_block_forward(GGMLRunnerContext* ctx,
ggml_tensor* hint,
ggml_tensor* emb,
ggml_tensor* context) {
int num_input_blocks = 15;
auto h = hint;
for (int i = 0; i < num_input_blocks; i++) {
if (i % 2 == 0) {
auto block = std::dynamic_pointer_cast<Conv2d>(blocks["input_hint_block." + std::to_string(i)]);
h = block->forward(ctx, h);
} else {
h = ggml_silu_inplace(ctx->ggml_ctx, h);
}
}
return h;
}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* hint,
ggml_tensor* guided_hint,
ggml_tensor* timesteps,
ggml_tensor* context,
ggml_tensor* y = nullptr) {
// x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
// timesteps: [N,]
// context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
// y: [N, adm_in_channels] or [1, adm_in_channels]
if (context != nullptr) {
if (context->ne[2] != x->ne[3]) {
context = ggml_repeat(ctx->ggml_ctx, context, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, context->ne[0], context->ne[1], x->ne[3]));
}
}
if (y != nullptr) {
if (y->ne[1] != x->ne[3]) {
y = ggml_repeat(ctx->ggml_ctx, y, ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, y->ne[0], x->ne[3]));
}
}
auto time_embed_0 = std::dynamic_pointer_cast<Linear>(blocks["time_embed.0"]);
auto time_embed_2 = std::dynamic_pointer_cast<Linear>(blocks["time_embed.2"]);
auto input_blocks_0_0 = std::dynamic_pointer_cast<Conv2d>(blocks["input_blocks.0.0"]);
auto zero_convs_0 = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs.0.0"]);
auto middle_block_out = std::dynamic_pointer_cast<Conv2d>(blocks["middle_block_out.0"]);
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels); // [N, model_channels]
auto emb = time_embed_0->forward(ctx, t_emb);
emb = ggml_silu_inplace(ctx->ggml_ctx, emb);
emb = time_embed_2->forward(ctx, emb); // [N, time_embed_dim]
// SDXL/SVD
if (y != nullptr) {
auto label_embed_0 = std::dynamic_pointer_cast<Linear>(blocks["label_emb.0.0"]);
auto label_embed_2 = std::dynamic_pointer_cast<Linear>(blocks["label_emb.0.2"]);
auto label_emb = label_embed_0->forward(ctx, y);
label_emb = ggml_silu_inplace(ctx->ggml_ctx, label_emb);
label_emb = label_embed_2->forward(ctx, label_emb); // [N, time_embed_dim]
emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim]
}
std::vector<ggml_tensor*> outs;
if (guided_hint == nullptr) {
guided_hint = input_hint_block_forward(ctx, hint, emb, context);
}
outs.push_back(guided_hint);
// input_blocks
// input block 0
auto h = input_blocks_0_0->forward(ctx, x);
h = ggml_add(ctx->ggml_ctx, h, guided_hint);
outs.push_back(zero_convs_0->forward(ctx, h));
// input block 1-11
size_t len_mults = channel_mult.size();
int input_block_idx = 0;
int ds = 1;
for (int i = 0; i < len_mults; i++) {
int mult = channel_mult[i];
for (int j = 0; j < num_res_blocks; j++) {
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
h = resblock_forward(name, ctx, h, emb); // [N, mult*model_channels, h, w]
if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
h = attention_layer_forward(name, ctx, h, context); // [N, mult*model_channels, h, w]
}
auto zero_conv = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs." + std::to_string(input_block_idx) + ".0"]);
outs.push_back(zero_conv->forward(ctx, h));
}
if (i != len_mults - 1) {
ds *= 2;
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
auto block = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
h = block->forward(ctx, h); // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
auto zero_conv = std::dynamic_pointer_cast<Conv2d>(blocks["zero_convs." + std::to_string(input_block_idx) + ".0"]);
outs.push_back(zero_conv->forward(ctx, h));
}
}
// [N, 4*model_channels, h/8, w/8]
// middle_block
h = resblock_forward("middle_block.0", ctx, h, emb); // [N, 4*model_channels, h/8, w/8]
h = attention_layer_forward("middle_block.1", ctx, h, context); // [N, 4*model_channels, h/8, w/8]
h = resblock_forward("middle_block.2", ctx, h, emb); // [N, 4*model_channels, h/8, w/8]
// out
outs.push_back(middle_block_out->forward(ctx, h));
return outs;
}
};
struct ControlNet : public GGMLRunner {
SDVersion version = VERSION_SD1;
ControlNetBlock control_net;
ggml_backend_buffer_t control_buffer = nullptr;
ggml_context* control_ctx = nullptr;
std::vector<ggml_tensor*> control_outputs_ggml;
ggml_tensor* guided_hint_output_ggml = nullptr;
std::vector<sd::Tensor<float>> controls;
sd::Tensor<float> guided_hint;
bool guided_hint_cached = false;
ControlNet(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
SDVersion version = VERSION_SD1)
: GGMLRunner(backend, params_backend), control_net(version) {
control_net.init(params_ctx, tensor_storage_map, "");
}
~ControlNet() override {
free_control_ctx();
}
void alloc_control_ctx(std::vector<ggml_tensor*> outs) {
ggml_init_params params;
params.mem_size = static_cast<size_t>(outs.size() * ggml_tensor_overhead()) + 1024 * 1024;
params.mem_buffer = nullptr;
params.no_alloc = true;
control_ctx = ggml_init(params);
control_outputs_ggml.resize(outs.size() - 1);
size_t control_buffer_size = 0;
guided_hint_output_ggml = ggml_dup_tensor(control_ctx, outs[0]);
control_buffer_size += ggml_nbytes(guided_hint_output_ggml);
for (int i = 0; i < outs.size() - 1; i++) {
control_outputs_ggml[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
control_buffer_size += ggml_nbytes(control_outputs_ggml[i]);
}
control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, runtime_backend);
LOG_DEBUG("control buffer size %.2fMB", control_buffer_size * 1.f / 1024.f / 1024.f);
}
void free_control_ctx() {
if (control_buffer != nullptr) {
ggml_backend_buffer_free(control_buffer);
control_buffer = nullptr;
}
if (control_ctx != nullptr) {
ggml_free(control_ctx);
control_ctx = nullptr;
}
guided_hint_output_ggml = nullptr;
guided_hint_cached = false;
guided_hint = {};
control_outputs_ggml.clear();
controls.clear();
}
std::string get_desc() override {
return "control_net";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
control_net.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& hint_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor = {},
const sd::Tensor<float>& y_tensor = {}) {
ggml_cgraph* gf = new_graph_custom(CONTROL_NET_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* hint = nullptr;
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* y = make_optional_input(y_tensor);
ggml_tensor* guided_hint_input = nullptr;
if (guided_hint_cached && !guided_hint.empty()) {
guided_hint_input = make_input(guided_hint);
hint = nullptr;
} else {
hint = make_input(hint_tensor);
}
auto runner_ctx = get_context();
auto outs = control_net.forward(&runner_ctx,
x,
hint,
guided_hint_input,
timesteps,
context,
y);
if (control_ctx == nullptr) {
alloc_control_ctx(outs);
}
ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint_output_ggml));
for (int i = 0; i < outs.size() - 1; i++) {
ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], control_outputs_ggml[i]));
}
return gf;
}
std::optional<std::vector<sd::Tensor<float>>> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& hint,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context = {},
const sd::Tensor<float>& y = {}) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
// y: [N, adm_in_channels] or [1, adm_in_channels]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, hint, timesteps, context, y);
};
auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false);
if (!compute_result.has_value()) {
return std::nullopt;
}
if (guided_hint_output_ggml != nullptr) {
guided_hint = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(guided_hint_output_ggml),
4);
}
controls.clear();
controls.reserve(control_outputs_ggml.size());
for (ggml_tensor* control : control_outputs_ggml) {
auto control_host = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(control), 4);
GGML_ASSERT(!control_host.empty());
controls.push_back(std::move(control_host));
}
guided_hint_cached = true;
return controls;
}
bool load_from_file(const std::string& file_path, int n_threads) {
LOG_INFO("loading control net from '%s'", file_path.c_str());
if (!alloc_params_buffer()) {
LOG_ERROR("control net model buffer allocation failed");
return false;
}
std::map<std::string, ggml_tensor*> tensors;
control_net.get_param_tensors(tensors);
std::set<std::string> ignore_tensors;
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str());
return false;
}
bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);
if (!success) {
LOG_ERROR("load control net tensors from model loader failed");
return false;
}
LOG_INFO("control net model loaded");
return success;
}
};
#endif // __SD_MODEL_DIFFUSION_CONTROL_HPP__
src/model/diffusion/dit.hpp
+166
View File
@@ -0,0 +1,166 @@
#ifndef __SD_MODEL_DIFFUSION_DIT_HPP__
#define __SD_MODEL_DIFFUSION_DIT_HPP__
#include "core/ggml_extend.hpp"
namespace DiT {
inline ggml_tensor* patchify(ggml_context* ctx,
ggml_tensor* x,
int pw,
int ph,
bool patch_last = true) {
// x: [N, C, H, W]
// return: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
int64_t N = x->ne[3];
int64_t C = x->ne[2];
int64_t H = x->ne[1];
int64_t W = x->ne[0];
int64_t h = H / ph;
int64_t w = W / pw;
GGML_ASSERT(h * ph == H && w * pw == W);
x = ggml_reshape_4d(ctx, x, pw, w, ph, h * C * N); // [N*C*h, ph, w, pw]
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N*C*h, w, ph, pw]
x = ggml_reshape_4d(ctx, x, pw * ph, w * h, C, N); // [N, C, h*w, ph*pw]
if (patch_last) {
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N, h*w, C, ph*pw]
x = ggml_reshape_3d(ctx, x, pw * ph * C, w * h, N); // [N, h*w, C*ph*pw]
} else {
x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, h*w, C, ph*pw]
x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N); // [N, h*w, ph*pw*C]
}
return x;
}
inline ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t h,
int64_t w,
int ph,
int pw,
bool patch_last = true) {
// x: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C]
// return: [N, C, H, W]
int64_t N = x->ne[2];
int64_t C = x->ne[0] / ph / pw;
int64_t H = h * ph;
int64_t W = w * pw;
GGML_ASSERT(C * ph * pw == x->ne[0]);
if (patch_last) {
x = ggml_reshape_4d(ctx, x, pw * ph, C, w * h, N); // [N, h*w, C, ph*pw]
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N, C, h*w, ph*pw]
} else {
x = ggml_reshape_4d(ctx, x, C, pw * ph, w * h, N); // [N, h*w, ph*pw, C]
x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3)); // [N, C, h*w, ph*pw]
}
x = ggml_reshape_4d(ctx, x, pw, ph, w, h * C * N); // [N*C*h, w, ph, pw]
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N*C*h, ph, w, pw]
x = ggml_reshape_4d(ctx, x, W, H, C, N); // [N, C, h*ph, w*pw]
return x;
}
inline ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx,
ggml_tensor* x,
int ph,
int pw) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int pad_h = (ph - H % ph) % ph;
int pad_w = (pw - W % pw) % pw;
x = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
return x;
}
inline ggml_tensor* pad_and_patchify(GGMLRunnerContext* ctx,
ggml_tensor* x,
int ph,
int pw,
bool patch_last = true) {
x = pad_to_patch_size(ctx, x, ph, pw);
x = patchify(ctx->ggml_ctx, x, ph, pw, patch_last);
return x;
}
inline ggml_tensor* unpatchify_and_crop(ggml_context* ctx,
ggml_tensor* x,
int64_t H,
int64_t W,
int ph,
int pw,
bool patch_last = true) {
int pad_h = (ph - H % ph) % ph;
int pad_w = (pw - W % pw) % pw;
int64_t h = ((H + pad_h) / ph);
int64_t w = ((W + pad_w) / pw);
x = unpatchify(ctx, x, h, w, ph, pw, patch_last); // [N, C, H + pad_h, W + pad_w]
x = ggml_ext_slice(ctx, x, 1, 0, H); // [N, C, H, W + pad_w]
x = ggml_ext_slice(ctx, x, 0, 0, W); // [N, C, H, W]
return x;
}
inline ggml_tensor* patchify(ggml_context* ctx,
ggml_tensor* x,
int pt,
int ph,
int pw,
int64_t N = 1) {
// x: [N*C, T, H, W]
// return: [N, h*w, C*pt*ph*pw]
int64_t C = x->ne[3] / N;
int64_t T = x->ne[2];
int64_t H = x->ne[1];
int64_t W = x->ne[0];
int64_t t_len = T / pt;
int64_t h_len = H / ph;
int64_t w_len = W / pw;
GGML_ASSERT(C * N == x->ne[3]);
GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); // [N, t_len*h_len*w_len, C*pt*ph*pw]
return x;
}
inline ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, pt*ph*pw*C]
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
int64_t N = x->ne[3];
int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw]
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw]
return x;
}
} // namespace DiT
#endif // __SD_MODEL_DIFFUSION_DIT_HPP__
src/model/diffusion/ernie_image.hpp
+458
View File
@@ -0,0 +1,458 @@
#ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
#include <memory>
#include <vector>
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/qwen_image.hpp"
namespace ErnieImage {
constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960;
struct ErnieImageConfig {
int64_t hidden_size = 4096;
int64_t num_heads = 32;
int64_t num_layers = 36;
int64_t ffn_hidden_size = 12288;
int64_t in_channels = 128;
int64_t out_channels = 128;
int patch_size = 1;
int64_t text_in_dim = 3072;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int axes_dim_sum = 128;
float eps = 1e-6f;
static ErnieImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
ErnieImageConfig config;
config.num_layers = 0;
int64_t detected_head_dim = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
config.patch_size = static_cast<int>(tensor_storage.ne[0]);
config.in_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
config.text_in_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
detected_head_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
config.ffn_hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
int64_t out_dim = tensor_storage.ne[1];
int64_t patch_area = config.patch_size * config.patch_size;
config.out_channels = out_dim / patch_area;
}
size_t pos = name.find("layers.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
}
if (config.num_layers == 0) {
config.num_layers = 36;
}
if (detected_head_dim > 0) {
config.num_heads = config.hidden_size / detected_head_dim;
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
LOG_DEBUG("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
config.num_layers,
config.hidden_size,
config.num_heads,
config.ffn_hidden_size,
config.in_channels,
config.out_channels);
return config;
}
};
__STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
ggml_tensor* timesteps,
int dim,
int max_period = 10000) {
auto emb = ggml_ext_timestep_embedding(ctx, timesteps, dim, max_period, 1.0f);
int64_t half = dim / 2;
auto cos_part = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], 0);
auto sin_part = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], half * emb->nb[0]);
auto sin_first = ggml_concat(ctx, sin_part, cos_part, 0);
return sin_first;
}
__STATIC_INLINE__ ggml_tensor* apply_rotary_emb(ggml_context* ctx, ggml_tensor* x, ggml_tensor* pe) {
// x: [N, S, heads, head_dim]
// pe: [2, S, 1, head_dim], stored as ggml [head_dim, 1, S, 2].
int64_t head_dim = x->ne[0];
int64_t heads = x->ne[1];
int64_t S = x->ne[2];
int64_t N = x->ne[3];
int64_t rot_dim = pe->ne[0];
GGML_ASSERT(rot_dim <= head_dim);
GGML_ASSERT(rot_dim % 2 == 0);
GGML_ASSERT(pe->ne[1] == 1 && pe->ne[2] == S && pe->ne[3] == 2);
x = ggml_cont(ctx, x);
auto x_rot = ggml_ext_slice(ctx, x, 0, 0, rot_dim, false);
auto x_pass = rot_dim < head_dim ? ggml_ext_slice(ctx, x, 0, rot_dim, head_dim, false) : nullptr;
int64_t half = rot_dim / 2;
auto x1 = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], 0);
auto x2 = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], half * x_rot->nb[0]);
x1 = ggml_cont(ctx, x1);
x2 = ggml_cont(ctx, x2);
auto rotated = ggml_concat(ctx, ggml_neg(ctx, x2), x1, 0);
auto cos_emb = ggml_ext_slice(ctx, pe, 3, 0, 1, false);
auto sin_emb = ggml_ext_slice(ctx, pe, 3, 1, 2, false);
auto out = ggml_add(ctx, ggml_mul(ctx, x_rot, cos_emb), ggml_mul(ctx, rotated, sin_emb));
if (x_pass != nullptr) {
out = ggml_concat(ctx, out, x_pass, 0);
}
return out;
}
struct ErnieImageAttention : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
ErnieImageAttention(int64_t query_dim,
int64_t heads,
int64_t dim_head,
float eps = 1e-6f)
: num_heads(heads), head_dim(dim_head) {
int64_t inner_dim = heads * dim_head;
blocks["to_q"] = std::make_shared<Linear>(query_dim, inner_dim, false);
blocks["to_k"] = std::make_shared<Linear>(query_dim, inner_dim, false);
blocks["to_v"] = std::make_shared<Linear>(query_dim, inner_dim, 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>(inner_dim, query_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* attention_mask = nullptr) {
// x: [N, S, hidden_size]
// pe: [S, head_dim/2, 2, 2], generated in image-token-first order.
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"]);
int64_t S = x->ne[1];
int64_t N = x->ne[2];
auto q = to_q->forward(ctx, x);
auto k = to_k->forward(ctx, x);
auto v = to_v->forward(ctx, x);
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N); // [N, S, heads, head_dim]
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N); // [N, S, heads, head_dim]
v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N); // [N, S, heads, head_dim]
q = norm_q->forward(ctx, q);
k = norm_k->forward(ctx, k);
q = apply_rotary_emb(ctx->ggml_ctx, q, pe);
k = apply_rotary_emb(ctx->ggml_ctx, k, pe);
q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)); // [N, heads, S, head_dim]
q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]);
k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, heads, S, head_dim]
k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, S, hidden_size]
x = to_out_0->forward(ctx, x);
return x;
}
};
struct ErnieImageFeedForward : public GGMLBlock {
public:
ErnieImageFeedForward(int64_t hidden_size, int64_t ffn_hidden_size) {
blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, ffn_hidden_size, false);
blocks["up_proj"] = std::make_shared<Linear>(hidden_size, ffn_hidden_size, false);
blocks["linear_fc2"] = std::make_shared<Linear>(ffn_hidden_size, hidden_size, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
auto gate = gate_proj->forward(ctx, x);
gate = ggml_ext_gelu(ctx->ggml_ctx, gate);
x = up_proj->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, gate);
x = linear_fc2->forward(ctx, x);
return x;
}
};
struct ErnieImageSharedAdaLNBlock : public GGMLBlock {
public:
ErnieImageSharedAdaLNBlock(int64_t hidden_size,
int64_t num_heads,
int64_t ffn_hidden_size,
float eps = 1e-6f) {
blocks["adaLN_sa_ln"] = std::make_shared<RMSNorm>(hidden_size, eps);
blocks["self_attention"] = std::make_shared<ErnieImageAttention>(hidden_size,
num_heads,
hidden_size / num_heads,
eps);
blocks["adaLN_mlp_ln"] = std::make_shared<RMSNorm>(hidden_size, eps);
blocks["mlp"] = std::make_shared<ErnieImageFeedForward>(hidden_size, ffn_hidden_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
const std::vector<ggml_tensor*>& temb,
ggml_tensor* attention_mask = nullptr) {
// x: [N, image_tokens + text_tokens, hidden_size]
auto adaLN_sa_ln = std::dynamic_pointer_cast<RMSNorm>(blocks["adaLN_sa_ln"]);
auto self_attention = std::dynamic_pointer_cast<ErnieImageAttention>(blocks["self_attention"]);
auto adaLN_mlp_ln = std::dynamic_pointer_cast<RMSNorm>(blocks["adaLN_mlp_ln"]);
auto mlp = std::dynamic_pointer_cast<ErnieImageFeedForward>(blocks["mlp"]);
auto shift_msa = temb[0];
auto scale_msa = temb[1];
auto gate_msa = temb[2];
auto shift_mlp = temb[3];
auto scale_mlp = temb[4];
auto gate_mlp = temb[5];
auto residual = x;
x = adaLN_sa_ln->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift_msa, scale_msa, true);
auto attn_out = self_attention->forward(ctx, x, pe, attention_mask);
x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
residual = x;
x = adaLN_mlp_ln->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift_mlp, scale_mlp, true);
x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, x), gate_mlp));
return x;
}
};
struct ErnieImageAdaLNContinuous : public GGMLBlock {
public:
ErnieImageAdaLNContinuous(int64_t hidden_size, float eps = 1e-6f) {
blocks["norm"] = std::make_shared<LayerNorm>(hidden_size, eps, false);
blocks["linear"] = std::make_shared<Linear>(hidden_size, hidden_size * 2, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
auto mods = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, conditioning), 2, 0);
auto scale = mods[0];
auto shift = mods[1];
x = norm->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
return x;
}
};
class ErnieImageModel : public GGMLBlock {
public:
ErnieImageConfig config;
ErnieImageModel() = default;
ErnieImageModel(ErnieImageConfig config)
: config(config) {
blocks["x_embedder.proj"] = std::make_shared<Conv2d>(config.in_channels,
config.hidden_size,
std::pair<int, int>{config.patch_size, config.patch_size},
std::pair<int, int>{config.patch_size, config.patch_size},
std::pair<int, int>{0, 0},
std::pair<int, int>{1, 1},
true);
if (config.text_in_dim != config.hidden_size) {
blocks["text_proj"] = std::make_shared<Linear>(config.text_in_dim, config.hidden_size, false);
}
blocks["time_embedding"] = std::make_shared<Qwen::TimestepEmbedding>(config.hidden_size, config.hidden_size);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
for (int i = 0; i < config.num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(config.hidden_size,
config.num_heads,
config.ffn_hidden_size,
config.eps);
}
blocks["final_norm"] = std::make_shared<ErnieImageAdaLNContinuous>(config.hidden_size, config.eps);
blocks["final_linear"] = std::make_shared<Linear>(config.hidden_size,
config.patch_size * config.patch_size * config.out_channels,
true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
// x: [N, C, H, W]
// context: [N, text_tokens, 3072]
// pe: [image_tokens + text_tokens, head_dim/2, 2, 2]
GGML_ASSERT(context != nullptr);
GGML_ASSERT(x->ne[1] % config.patch_size == 0 && x->ne[0] % config.patch_size == 0);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t Hp = H / config.patch_size;
int64_t Wp = W / config.patch_size;
int64_t n_img = Hp * Wp;
int64_t N = x->ne[3];
auto x_embedder_proj = std::dynamic_pointer_cast<Conv2d>(blocks["x_embedder.proj"]);
auto time_embedding = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedding"]);
auto adaLN_mod = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
auto final_norm = std::dynamic_pointer_cast<ErnieImageAdaLNContinuous>(blocks["final_norm"]);
auto final_linear = std::dynamic_pointer_cast<Linear>(blocks["final_linear"]);
auto img = x_embedder_proj->forward(ctx, x); // [N, hidden_size, Hp, Wp]
img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N); // [N, hidden_size, image_tokens]
img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); // [N, image_tokens, hidden_size]
auto txt = context;
auto text_proj = std::dynamic_pointer_cast<Linear>(blocks["text_proj"]);
if (text_proj) {
txt = text_proj->forward(ctx, txt);
}
auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1); // [N, image_tokens + text_tokens, hidden_size]
auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
auto c = time_embedding->forward(ctx, sample); // [N, hidden_size]
auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size]
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states");
// sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params");
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
std::vector<ggml_tensor*> temb;
temb.reserve(6);
for (auto chunk : chunks) {
temb.push_back(ggml_reshape_3d(ctx->ggml_ctx, chunk, chunk->ne[0], 1, chunk->ne[1])); // [N, 1, hidden_size]
}
for (int i = 0; i < config.num_layers; i++) {
auto layer = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
hidden_states = layer->forward(ctx, hidden_states, pe, temb);
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states");
}
hidden_states = final_norm->forward(ctx, hidden_states, c);
hidden_states = final_linear->forward(ctx, hidden_states); // [N, image_tokens, p*p*out_channels]
auto patches = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img); // [N, image_tokens, hidden_size]
auto out = DiT::unpatchify(ctx->ggml_ctx,
patches,
Hp,
Wp,
config.patch_size,
config.patch_size,
false); // [N, out_channels, H, W]
return out;
}
};
struct ErnieImageRunner : public DiffusionModelRunner {
ErnieImageConfig config;
ErnieImageModel ernie_image;
std::vector<float> pe_vec;
ErnieImageRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix),
config(ErnieImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
ernie_image = ErnieImageModel(config);
ernie_image.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "ernie_image";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
ernie_image.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor) {
ggml_cgraph* gf = new_graph_custom(ERNIE_IMAGE_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);
pe_vec = Rope::gen_ernie_image_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]),
config.theta,
circular_y_enabled,
circular_x_enabled,
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, config.axes_dim_sum, 1, pos_len, 2);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context));
}
};
} // namespace ErnieImage
#endif // __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
src/model/diffusion/flux.hpp
+1630
View File
File diff suppressed because it is too large. Load diff
src/model/diffusion/hidream_o1.hpp
+681
View File
@@ -0,0 +1,681 @@
#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
#define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
#include <algorithm>
#include <array>
#include <cmath>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "conditioning/conditioner.hpp"
#include "core/util.h"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/model.hpp"
#include "model/te/llm.hpp"
namespace HiDreamO1 {
constexpr int HIDREAM_O1_GRAPH_SIZE = 32768;
constexpr int PATCH_SIZE = 32;
constexpr int TIMESTEP_TOKEN_NUM = 1;
constexpr int IMAGE_TOKEN_ID = 151655;
constexpr int VISION_START_TOKEN_ID = 151652;
struct HiDreamO1Config {
LLM::LLMConfig llm;
int patch_size = PATCH_SIZE;
static HiDreamO1Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
(void)tensor_storage_map;
(void)prefix;
HiDreamO1Config config;
config.llm.arch = LLM::LLMArch::QWEN3_VL;
config.llm.hidden_size = 4096;
config.llm.intermediate_size = 12288;
config.llm.num_layers = 36;
config.llm.num_heads = 32;
config.llm.num_kv_heads = 8;
config.llm.head_dim = 128;
config.llm.qkv_bias = false;
config.llm.qk_norm = true;
config.llm.vocab_size = 151936;
config.llm.rms_norm_eps = 1e-6f;
config.llm.vision.arch = LLM::LLMVisionArch::QWEN3_VL;
config.llm.vision.num_layers = 27;
config.llm.vision.hidden_size = 1152;
config.llm.vision.intermediate_size = 4304;
config.llm.vision.num_heads = 16;
config.llm.vision.out_hidden_size = 4096;
config.llm.vision.patch_size = 16;
config.llm.vision.spatial_merge_size = 2;
config.llm.vision.temporal_patch_size = 2;
config.llm.vision.num_position_embeddings = 2304;
return config;
}
};
static inline std::string repeat_special_token(const std::string& token, int64_t count) {
std::string out;
out.reserve(static_cast<size_t>(count) * token.size());
for (int64_t i = 0; i < count; ++i) {
out += token;
}
return out;
}
static inline std::pair<int, int> calculate_dimensions(int max_size, double ratio) {
int width = static_cast<int>(std::sqrt(max_size * max_size * ratio));
int height = static_cast<int>(width / ratio);
width = (width / PATCH_SIZE) * PATCH_SIZE;
height = (height / PATCH_SIZE) * PATCH_SIZE;
width = std::max(width, PATCH_SIZE);
height = std::max(height, PATCH_SIZE);
return {width, height};
}
static inline sd::Tensor<float> resize_to_area(const sd::Tensor<float>& image, int image_size) {
int64_t width = image.shape()[0];
int64_t height = image.shape()[1];
int64_t s_max = static_cast<int64_t>(image_size) * image_size;
double scale = std::sqrt(static_cast<double>(s_max) / static_cast<double>(width * height));
std::vector<std::pair<int64_t, int64_t>> sizes = {
{(static_cast<int64_t>(std::llround(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::llround(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
{(static_cast<int64_t>(std::llround(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::floor(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
{(static_cast<int64_t>(std::floor(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::llround(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
{(static_cast<int64_t>(std::floor(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::floor(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
};
std::sort(sizes.begin(), sizes.end(), [](const auto& a, const auto& b) {
return a.first * a.second > b.first * b.second;
});
std::pair<int64_t, int64_t> new_size = sizes.back();
for (const auto& size : sizes) {
if (size.first > 0 && size.second > 0 && size.first * size.second <= s_max) {
new_size = size;
break;
}
}
double s1 = static_cast<double>(width) / static_cast<double>(new_size.first);
double s2 = static_cast<double>(height) / static_cast<double>(new_size.second);
sd::Tensor<float> resized;
if (s1 < s2) {
int64_t resized_h = static_cast<int64_t>(std::llround(height / s1));
resized = sd::ops::interpolate(image,
{new_size.first, resized_h, image.shape()[2], image.shape()[3]},
sd::ops::InterpolateMode::Bicubic);
int64_t top = (resized_h - new_size.second) / 2;
resized = sd::ops::slice(resized, 1, top, top + new_size.second);
} else {
int64_t resized_w = static_cast<int64_t>(std::llround(width / s2));
resized = sd::ops::interpolate(image,
{resized_w, new_size.second, image.shape()[2], image.shape()[3]},
sd::ops::InterpolateMode::Bicubic);
int64_t left = (resized_w - new_size.first) / 2;
resized = sd::ops::slice(resized, 0, left, left + new_size.first);
}
return resized;
}
static inline std::vector<int32_t> build_position_ids(const std::vector<int32_t>& input_ids,
const std::vector<std::array<int32_t, 3>>& image_grids,
const std::vector<int32_t>& skip_vision_start_token) {
std::vector<int32_t> position_ids(4 * input_ids.size(), 0);
int image_index = 0;
int st = 0;
int fix_point = 4096;
std::vector<int32_t> out_t;
std::vector<int32_t> out_h;
std::vector<int32_t> out_w;
while (st < static_cast<int>(input_ids.size())) {
int ed = st;
while (ed < static_cast<int>(input_ids.size()) && input_ids[ed] != IMAGE_TOKEN_ID) {
ed++;
}
if (ed >= static_cast<int>(input_ids.size())) {
int st_idx = out_t.empty() ? 0 : (*std::max_element(out_t.begin(), out_t.end()) + 1);
for (int i = 0; i < static_cast<int>(input_ids.size()) - st; ++i) {
out_t.push_back(st_idx + i);
out_h.push_back(st_idx + i);
out_w.push_back(st_idx + i);
}
break;
}
int text_len = std::max(0, ed - st - skip_vision_start_token[image_index]);
int st_idx = out_t.empty() ? 0 : (*std::max_element(out_t.begin(), out_t.end()) + 1);
for (int i = 0; i < text_len; ++i) {
out_t.push_back(st_idx + i);
out_h.push_back(st_idx + i);
out_w.push_back(st_idx + i);
}
auto grid = image_grids[image_index];
int base;
if (skip_vision_start_token[image_index]) {
if (fix_point > 0) {
base = fix_point;
fix_point = 0;
} else {
base = st_idx;
}
} else {
base = text_len + st_idx;
}
for (int32_t ti = 0; ti < grid[0]; ++ti) {
for (int32_t hi = 0; hi < grid[1]; ++hi) {
for (int32_t wi = 0; wi < grid[2]; ++wi) {
out_t.push_back(base + ti);
out_h.push_back(base + hi);
out_w.push_back(base + wi);
}
}
}
st = ed + grid[0] * grid[1] * grid[2];
image_index++;
}
GGML_ASSERT(out_t.size() == input_ids.size());
for (size_t i = 0; i < input_ids.size(); ++i) {
// ggml IMROPE consumes 4 flattened position streams:
// [t, h, w, e]
// llama.cpp's generic Qwen-VL fallback expands text positions as
// [pos, pos, pos, 0]. Keep the extra stream zeroed here too.
position_ids[i] = out_t[i];
position_ids[input_ids.size() + i] = out_h[i];
position_ids[input_ids.size() * 2 + i] = out_w[i];
position_ids[input_ids.size() * 3 + i] = 0;
}
return position_ids;
}
struct TimestepEmbedder : public GGMLBlock {
int frequency_embedding_size = 256;
TimestepEmbedder(int64_t hidden_size) {
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1000.0f);
emb = mlp_0->forward(ctx, emb);
emb = ggml_silu_inplace(ctx->ggml_ctx, emb);
emb = mlp_2->forward(ctx, emb);
return emb;
}
};
struct BottleneckPatchEmbed : public GGMLBlock {
BottleneckPatchEmbed(int64_t in_dim, int64_t pca_dim, int64_t embed_dim) {
blocks["proj1"] = std::make_shared<Linear>(in_dim, pca_dim, false);
blocks["proj2"] = std::make_shared<Linear>(pca_dim, embed_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj1 = std::dynamic_pointer_cast<Linear>(blocks["proj1"]);
auto proj2 = std::dynamic_pointer_cast<Linear>(blocks["proj2"]);
return proj2->forward(ctx, proj1->forward(ctx, x));
}
};
struct FinalLayer : public GGMLBlock {
FinalLayer(int64_t hidden_size, int64_t out_dim) {
blocks["linear"] = std::make_shared<Linear>(hidden_size, out_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
return linear->forward(ctx, x);
}
};
struct HiDreamO1Model : public GGMLBlock {
HiDreamO1Config config;
HiDreamO1Model() = default;
explicit HiDreamO1Model(HiDreamO1Config config)
: config(std::move(config)) {
blocks["language_model"] = std::make_shared<LLM::TextModel>(this->config.llm);
blocks["t_embedder1"] = std::make_shared<TimestepEmbedder>(this->config.llm.hidden_size);
blocks["x_embedder"] = std::make_shared<BottleneckPatchEmbed>(this->config.patch_size * this->config.patch_size * 3,
this->config.llm.hidden_size / 4,
this->config.llm.hidden_size);
blocks["final_layer2"] = std::make_shared<FinalLayer>(this->config.llm.hidden_size,
this->config.patch_size * this->config.patch_size * 3);
}
std::shared_ptr<LLM::TextModel> text_model() {
return std::dynamic_pointer_cast<LLM::TextModel>(blocks["language_model"]);
}
std::shared_ptr<TimestepEmbedder> timestep_embedder() {
return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder1"]);
}
std::shared_ptr<BottleneckPatchEmbed> patch_embedder() {
return std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["x_embedder"]);
}
std::shared_ptr<FinalLayer> final_layer() {
return std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer2"]);
}
};
struct HiDreamO1VisionRunner : public GGMLRunner {
HiDreamO1Config config;
std::shared_ptr<LLM::VisionModel> model;
std::vector<int> window_index_vec;
std::vector<int> window_inverse_index_vec;
std::vector<float> window_mask_vec;
std::vector<float> pe_vec;
std::array<std::vector<int32_t>, 4> pos_embed_idx_data_;
std::array<std::vector<float>, 4> pos_embed_weight_data_;
HiDreamO1VisionRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "model.visual")
: GGMLRunner(backend, params_backend),
config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)),
model(std::make_shared<LLM::VisionModel>(false, config.llm.vision)) {
model->init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "hidream_o1_vision";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix = "model.visual") {
model->get_param_tensors(tensors, prefix);
}
ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
return LLM::LLMRunner::encode_image_common(this,
compute_ctx,
runner_ctx,
image,
config.llm.vision,
model,
window_index_vec,
window_inverse_index_vec,
window_mask_vec,
pe_vec,
pos_embed_idx_data_,
pos_embed_weight_data_);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
auto runner_ctx = get_context();
auto image_embeds = encode_image(&runner_ctx, image);
ggml_build_forward_expand(gf, image_embeds);
return gf;
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image) {
auto get_graph = [&]() {
return build_graph(image);
};
auto output = GGMLRunner::compute<float>(get_graph, n_threads, false);
return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
}
};
struct HiDreamO1Runner : public DiffusionModelRunner {
HiDreamO1Config config;
HiDreamO1Model model;
std::vector<float> attention_mask_vec;
HiDreamO1Runner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "model")
: DiffusionModelRunner(backend, params_backend, prefix),
config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)) {
model = HiDreamO1Model(config);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "hidream_o1";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timestep_tensor,
const sd::Tensor<int32_t>& input_ids_tensor,
const sd::Tensor<int32_t>& input_pos_tensor,
const sd::Tensor<int32_t>& token_types_tensor,
const sd::Tensor<int32_t>& vinput_mask_tensor,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
const std::vector<sd::Tensor<float>>& ref_images) {
ggml_cgraph* gf = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timestep = make_input(timestep_tensor);
ggml_tensor* input_ids = make_input(input_ids_tensor);
ggml_tensor* input_pos = make_input(input_pos_tensor);
auto text_model = model.text_model();
auto t_embedder1 = model.timestep_embedder();
auto x_embedder = model.patch_embedder();
auto final_layer2 = model.final_layer();
std::vector<ggml_tensor*> ref_image_tensors;
for (const auto& image : ref_images) {
ref_image_tensors.push_back(make_input(image));
}
attention_mask_vec = std::vector<float>(static_cast<size_t>(token_types_tensor.shape()[0] * token_types_tensor.shape()[0]), 0.0f);
int64_t total_seq_len = token_types_tensor.shape()[0];
for (int64_t query = 0; query < total_seq_len; ++query) {
bool is_gen = token_types_tensor.values()[static_cast<size_t>(query)] > 0;
for (int64_t key = 0; key < total_seq_len; ++key) {
if (!is_gen && key > query) {
attention_mask_vec[static_cast<size_t>(query * total_seq_len + key)] = -INFINITY;
}
}
}
auto attention_mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, total_seq_len, total_seq_len);
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
auto runner_ctx = get_context();
auto txt = text_model->embed(&runner_ctx, input_ids);
std::vector<std::pair<int, ggml_tensor*>> image_embeds;
image_embeds.reserve(image_embeds_tensor.size());
for (const auto& image_embed : image_embeds_tensor) {
image_embeds.emplace_back(image_embed.first, make_input(image_embed.second));
}
txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds);
auto t_emb = t_embedder1->forward(&runner_ctx, timestep);
int64_t txt_seq_len = input_ids->ne[0];
if (txt_seq_len > 1) {
auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1);
txt = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1);
} else {
txt = ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1);
}
auto vinputs = DiT::pad_and_patchify(&runner_ctx, x, PATCH_SIZE, PATCH_SIZE);
int64_t target_tokens = vinputs->ne[1];
for (ggml_tensor* ref_image : ref_image_tensors) {
auto ref = DiT::pad_and_patchify(&runner_ctx, ref_image, PATCH_SIZE, PATCH_SIZE);
vinputs = ggml_concat(compute_ctx, vinputs, ref, 1);
}
auto vis = x_embedder->forward(&runner_ctx, vinputs);
auto inputs_embeds = ggml_concat(compute_ctx, txt, vis, 1);
auto hidden_states = text_model->forward_embeds(&runner_ctx, inputs_embeds, input_pos, attention_mask, {});
auto x_pred_all = final_layer2->forward(&runner_ctx, hidden_states);
int64_t x_pred_start = txt_seq_len;
if (!vinput_mask_tensor.empty()) {
int64_t seq_len = static_cast<int64_t>(vinput_mask_tensor.shape()[0]);
int64_t first_vinput = 0;
while (first_vinput < seq_len && vinput_mask_tensor.values()[static_cast<size_t>(first_vinput)] == 0) {
first_vinput++;
}
x_pred_start = first_vinput;
}
auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
x_pred = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
float sigma = 1.0f - timestep_tensor.values()[0];
sigma = std::max(1e-6f, sigma);
auto out = ggml_scale(compute_ctx, ggml_sub(compute_ctx, x, x_pred), 1.0f / sigma);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timestep,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<int32_t>& input_pos,
const sd::Tensor<int32_t>& token_types,
const sd::Tensor<int32_t>& vinput_mask,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
const std::vector<sd::Tensor<float>>& ref_images) {
auto get_graph = [&]() {
return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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);
const auto* extra = diffusion_extra_as<HiDreamO1DiffusionExtra>(diffusion_params);
GGML_ASSERT(extra != nullptr);
GGML_ASSERT(extra->input_ids != nullptr);
GGML_ASSERT(extra->input_pos != nullptr);
GGML_ASSERT(extra->token_types != nullptr);
static const std::vector<sd::Tensor<float>> empty_images;
static const std::vector<std::pair<int, sd::Tensor<float>>> empty_image_embeds;
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*extra->input_ids,
*extra->input_pos,
*extra->token_types,
tensor_or_empty(extra->vinput_mask),
extra->image_embeds ? *extra->image_embeds : empty_image_embeds,
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_images);
}
};
struct HiDreamO1Conditioner : public Conditioner {
Qwen2Tokenizer tokenizer;
std::shared_ptr<HiDreamO1VisionRunner> vision_runner;
HiDreamO1Conditioner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {})
: vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, params_backend, tensor_storage_map)) {}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
vision_runner->get_param_tensors(tensors);
}
bool alloc_params_buffer() override {
if (!vision_runner->alloc_params_buffer()) {
return false;
}
return true;
}
void free_params_buffer() override {
vision_runner->free_params_buffer();
}
size_t get_params_buffer_size() override {
return vision_runner->get_params_buffer_size();
}
void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override {
vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes);
}
void set_flash_attention_enabled(bool enabled) override {
vision_runner->set_flash_attention_enabled(enabled);
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
vision_runner->set_weight_adapter(adapter);
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
SDCondition result;
int width = conditioner_params.width;
int height = conditioner_params.height;
int64_t target_image_len = static_cast<int64_t>(width / PATCH_SIZE) * static_cast<int64_t>(height / PATCH_SIZE);
std::vector<sd::Tensor<float>> ref_images;
if (conditioner_params.ref_images != nullptr) {
ref_images = *conditioner_params.ref_images;
}
std::vector<std::pair<int, sd::Tensor<float>>> vlm_images;
std::vector<std::array<int32_t, 3>> image_grids;
std::vector<int32_t> skip_vision_start;
std::string prompt = "<|im_start|>user\n";
if (ref_images.empty()) {
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
auto input_ids = tokenizer.encode(prompt, nullptr);
std::vector<int32_t> input_ids_pad = input_ids;
input_ids_pad.push_back(VISION_START_TOKEN_ID);
input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
skip_vision_start.push_back(1);
std::vector<int32_t> token_types(input_ids_pad.size(), 0);
int txt_seq_len = static_cast<int>(input_ids.size());
int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM;
for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
token_types[i] = 1;
}
auto position_ids = build_position_ids(input_ids_pad, image_grids, skip_vision_start);
std::vector<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
std::vector<int32_t> vinput_mask(token_types.size(), 0);
for (int64_t i = txt_seq_len; i < static_cast<int64_t>(vinput_mask.size()); ++i) {
vinput_mask[static_cast<size_t>(i)] = 1;
}
std::vector<int64_t> vinput_mask_shape{static_cast<int64_t>(vinput_mask.size())};
result.c_input_ids = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
result.c_position_ids = sd::Tensor<int32_t>(position_shape, position_ids);
result.c_token_types = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
result.c_vinput_mask = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
return result;
}
int K = static_cast<int>(ref_images.size());
int max_size;
if (K == 1) {
max_size = std::max(height, width);
} else if (K == 2) {
max_size = std::max(height, width) * 48 / 64;
} else if (K <= 4) {
max_size = std::max(height, width) / 2;
} else if (K <= 8) {
max_size = std::max(height, width) * 24 / 64;
} else {
max_size = std::max(height, width) / 4;
}
int cond_img_size;
if (K <= 4) {
cond_img_size = 384;
} else if (K <= 8) {
cond_img_size = 384 * 48 / 64;
} else {
cond_img_size = 384 / 2;
}
for (const auto& ref_image : ref_images) {
auto resized_ref = resize_to_area(ref_image, max_size);
resized_ref = sd::ops::clamp(resized_ref, 0.0f, 1.0f);
// VLM image: Qwen3-VL expects mean=[0.5]/std=[0.5] (i.e. range [-1,1]),
// not CLIP normalization. Resize the already-resized ref directly to
// (cond_w, cond_h) to match the Python pipeline's pil_r.resize().
auto dims = calculate_dimensions(cond_img_size,
static_cast<double>(resized_ref.shape()[0]) / static_cast<double>(resized_ref.shape()[1]));
sd::Tensor<float> vlm_image = sd::ops::interpolate(
resized_ref,
{dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]});
vlm_image = vlm_image * 2.0f - 1.0f;
int64_t image_tokens = static_cast<int64_t>(dims.first / PATCH_SIZE) * static_cast<int64_t>(dims.second / PATCH_SIZE);
auto patch_img = resized_ref * 2.0f - 1.0f;
result.c_ref_images.push_back(std::move(patch_img));
int64_t prompt_start = static_cast<int64_t>(tokenizer.encode(prompt + "<|vision_start|>", nullptr).size());
prompt += "<|vision_start|>";
prompt += repeat_special_token("<|image_pad|>", image_tokens);
prompt += "<|vision_end|>";
vlm_images.emplace_back(static_cast<int>(prompt_start), std::move(vlm_image));
image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE});
skip_vision_start.push_back(0);
}
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
auto input_ids = tokenizer.encode(prompt, nullptr);
std::vector<int32_t> input_ids_pad = input_ids;
input_ids_pad.push_back(VISION_START_TOKEN_ID);
input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
skip_vision_start.push_back(1);
for (const auto& ref_image : result.c_ref_images) {
int64_t ref_len = static_cast<int64_t>(ref_image.shape()[0] / PATCH_SIZE) * static_cast<int64_t>(ref_image.shape()[1] / PATCH_SIZE);
input_ids_pad.push_back(VISION_START_TOKEN_ID);
input_ids_pad.insert(input_ids_pad.end(), ref_len - 1, IMAGE_TOKEN_ID);
image_grids.push_back({1, static_cast<int32_t>(ref_image.shape()[1] / PATCH_SIZE), static_cast<int32_t>(ref_image.shape()[0] / PATCH_SIZE)});
skip_vision_start.push_back(1);
}
std::vector<int32_t> token_types(input_ids_pad.size(), 0);
int txt_seq_len = static_cast<int>(input_ids.size());
int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM;
for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
token_types[i] = 1;
}
std::vector<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
std::vector<int32_t> vinput_mask(token_types.size(), 0);
for (int i = txt_seq_len; i < static_cast<int>(vinput_mask.size()); ++i) {
vinput_mask[static_cast<size_t>(i)] = 1;
}
std::vector<int64_t> vinput_mask_shape{static_cast<int64_t>(vinput_mask.size())};
result.c_input_ids = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
result.c_position_ids = sd::Tensor<int32_t>(position_shape, build_position_ids(input_ids_pad, image_grids, skip_vision_start));
result.c_token_types = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
result.c_vinput_mask = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
result.c_image_embeds.reserve(vlm_images.size());
for (const auto& vlm_image : vlm_images) {
auto image_embed = vision_runner->compute(n_threads, vlm_image.second);
if (image_embed.empty()) {
LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
return SDCondition();
}
result.c_image_embeds.emplace_back(vlm_image.first, std::move(image_embed));
}
return result;
}
};
} // namespace HiDreamO1
#endif // __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
src/model/diffusion/ideogram4.hpp
+531
View File
@@ -0,0 +1,531 @@
#ifndef __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
#define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.hpp"
#include "core/ggml_graph_cut.h"
#include "model/common/rope.hpp"
#include "model/diffusion/model.hpp"
namespace Ideogram4 {
constexpr int IDEOGRAM4_GRAPH_SIZE = 65536;
constexpr int OUTPUT_IMAGE_INDICATOR = 2;
constexpr int IMAGE_POSITION_OFFSET = 65536;
constexpr int DEFAULT_MROPE_SECTION_T = 24;
constexpr int DEFAULT_MROPE_SECTION_H = 20;
constexpr int DEFAULT_MROPE_SECTION_W = 20;
constexpr int TIMESTEP_MAX_PERIOD = 10000;
constexpr int LLM_HIDDEN_STATE_LAYERS = 13;
struct Ideogram4Config {
int64_t emb_dim = 4608;
int64_t num_layers = 34;
int64_t num_heads = 18;
int64_t intermediate_size = 12288;
int64_t adanln_dim = 512;
int64_t in_channels = 128;
int64_t llm_features_dim = 53248;
int64_t rope_theta = 5000000;
float norm_eps = 1e-5f;
int patch_size = 2;
int ae_channels = 32;
std::vector<int> mrope_section = {DEFAULT_MROPE_SECTION_T,
DEFAULT_MROPE_SECTION_H,
DEFAULT_MROPE_SECTION_W};
static Ideogram4Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
Ideogram4Config config;
int64_t detected_layers = 0;
std::string layer_prefix = prefix.empty() ? "layers." : prefix + ".layers.";
for (const auto& [name, _] : tensor_storage_map) {
if (name.find(layer_prefix) != 0) {
continue;
}
std::string tail = name.substr(layer_prefix.size());
size_t dot = tail.find('.');
if (dot == std::string::npos) {
continue;
}
int layer_idx = std::atoi(tail.substr(0, dot).c_str());
detected_layers = std::max<int64_t>(detected_layers, layer_idx + 1);
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
LOG_DEBUG("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64,
config.num_layers,
config.emb_dim,
config.num_heads,
config.intermediate_size);
}
return config;
}
};
__STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
ggml_tensor* timesteps,
int dim) {
GGML_ASSERT(dim % 2 == 0);
auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
auto chunks = ggml_ext_chunk(ctx, embedding, 2, 0);
return ggml_concat(ctx, chunks[1], chunks[0], 0);
}
__STATIC_INLINE__ ggml_tensor* to_token_modulation(ggml_context* ctx, ggml_tensor* x) {
// [N, C] -> [N, 1, C] in PyTorch layout.
if (ggml_n_dims(x) < 3 || x->ne[1] != 1) {
x = ggml_reshape_3d(ctx, x, x->ne[0], 1, x->ne[1]);
}
return x;
}
__STATIC_INLINE__ ggml_tensor* interleave_hidden_state_layers(ggml_context* ctx, ggml_tensor* x) {
// Match upstream stack(...).permute(1, 2, 3, 0).reshape(...):
// [layers * hidden, tokens, batch] -> [hidden * layers, tokens, batch].
GGML_ASSERT(x->ne[0] % LLM_HIDDEN_STATE_LAYERS == 0);
const int64_t hidden_size = x->ne[0] / LLM_HIDDEN_STATE_LAYERS;
const int64_t token_count = x->ne[1];
const int64_t batch_count = x->ne[2];
x = ggml_reshape_4d(ctx, x, hidden_size, LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
}
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
scale = to_token_modulation(ctx, scale);
return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
}
__STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
// x: [N, 128, H, W] with channel order [ae, ph, pw].
// return: [N, H*W, 128] with token channel order [ph, pw, ae].
const int64_t W = x->ne[0];
const int64_t H = x->ne[1];
const int64_t C = x->ne[2];
const int64_t N = x->ne[3];
GGML_ASSERT(N == 1);
GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size);
x = ggml_cont(ctx, x);
x = ggml_reshape_4d(ctx, x, W * H, config.patch_size, config.patch_size, config.ae_channels);
x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0));
x = ggml_reshape_3d(ctx, x, C, W * H, N);
return x;
}
__STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t H,
int64_t W,
const Ideogram4Config& config) {
const int64_t C = x->ne[0];
const int64_t N = x->ne[2];
GGML_ASSERT(N == 1);
GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size);
GGML_ASSERT(x->ne[1] == H * W);
x = ggml_reshape_4d(ctx, x, config.ae_channels, config.patch_size, config.patch_size, H * W);
x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0));
x = ggml_reshape_4d(ctx, x, W, H, C, N);
return x;
}
__STATIC_INLINE__ std::shared_ptr<Linear> make_linear(int64_t in_features,
int64_t out_features,
bool bias = true) {
return std::make_shared<Linear>(in_features, out_features, bias, false, false, 1.f, true);
}
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
int grid_w,
int bs,
int context_len,
int head_dim,
int rope_theta,
const std::vector<int>& mrope_section) {
GGML_ASSERT(bs == 1);
std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
std::vector<float>(3, 0.f));
for (int i = 0; i < context_len; ++i) {
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
}
int cursor = context_len;
for (int y = 0; y < grid_h; ++y) {
for (int x = 0; x < grid_w; ++x) {
ids[cursor++] = {static_cast<float>(IMAGE_POSITION_OFFSET),
static_cast<float>(IMAGE_POSITION_OFFSET + y),
static_cast<float>(IMAGE_POSITION_OFFSET + x)};
}
}
return Rope::embed_interleaved_mrope(ids, bs, static_cast<float>(rope_theta), head_dim, mrope_section);
}
class Ideogram4Attention : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t num_heads;
int64_t head_dim;
public:
Ideogram4Attention(int64_t hidden_size, int64_t num_heads, float eps)
: hidden_size(hidden_size), num_heads(num_heads), head_dim(hidden_size / num_heads) {
GGML_ASSERT(hidden_size % num_heads == 0);
blocks["qkv"] = make_linear(hidden_size, hidden_size * 3, false);
blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["o"] = make_linear(hidden_size, hidden_size, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
auto qkv = qkv_proj->forward(ctx, x);
auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N);
auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N);
auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N);
q = norm_q->forward(ctx, q);
k = norm_k->forward(ctx, k);
x = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f, false);
x = out_proj->forward(ctx, x);
return x;
}
};
class Ideogram4MLP : public GGMLBlock {
public:
Ideogram4MLP(int64_t dim, int64_t hidden_dim) {
blocks["w1"] = make_linear(dim, hidden_dim, false);
blocks["w2"] = make_linear(hidden_dim, dim, false);
blocks["w3"] = make_linear(dim, hidden_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto x1 = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
auto x3 = w3->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x1, x3);
x = w2->forward(ctx, x);
return x;
}
};
class Ideogram4TransformerBlock : public GGMLBlock {
public:
Ideogram4TransformerBlock(const Ideogram4Config& config) {
blocks["attention"] = std::make_shared<Ideogram4Attention>(config.emb_dim, config.num_heads, config.norm_eps);
blocks["feed_forward"] = std::make_shared<Ideogram4MLP>(config.emb_dim, config.intermediate_size);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["adaln_modulation"] = make_linear(config.adanln_dim, 4 * config.emb_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* adaln_input,
ggml_tensor* mask = nullptr) {
auto attention = std::dynamic_pointer_cast<Ideogram4Attention>(blocks["attention"]);
auto feed_forward = std::dynamic_pointer_cast<Ideogram4MLP>(blocks["feed_forward"]);
auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
auto ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
auto attention_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm2"]);
auto ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
auto mod = adaln_modulation->forward(ctx, adaln_input);
auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0);
auto scale_msa = mods[0];
auto gate_msa = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1]));
auto scale_mlp = mods[2];
auto gate_mlp = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3]));
auto attn_out = attention_norm1->forward(ctx, x);
attn_out = modulate(ctx->ggml_ctx, attn_out, scale_msa);
attn_out = attention->forward(ctx, attn_out, pe, mask);
attn_out = attention_norm2->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
auto ffn_out = ffn_norm1->forward(ctx, x);
ffn_out = modulate(ctx->ggml_ctx, ffn_out, scale_mlp);
ffn_out = feed_forward->forward(ctx, ffn_out);
ffn_out = ffn_norm2->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp));
return x;
}
};
class Ideogram4EmbedScalar : public GGMLBlock {
protected:
int64_t dim;
public:
Ideogram4EmbedScalar(int64_t dim)
: dim(dim) {
blocks["mlp_in"] = make_linear(dim, dim, true);
blocks["mlp_out"] = make_linear(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto mlp_in = std::dynamic_pointer_cast<Linear>(blocks["mlp_in"]);
auto mlp_out = std::dynamic_pointer_cast<Linear>(blocks["mlp_out"]);
x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast<int>(dim));
x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x));
x = mlp_out->forward(ctx, x);
return x;
}
};
class Ideogram4FinalLayer : public GGMLBlock {
public:
Ideogram4FinalLayer(const Ideogram4Config& config) {
blocks["norm_final"] = std::make_shared<LayerNorm>(config.emb_dim, 1e-6f, false);
blocks["linear"] = make_linear(config.emb_dim, config.in_channels, true);
blocks["adaln_modulation"] = make_linear(config.adanln_dim, config.emb_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* c) {
auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
x = norm_final->forward(ctx, x);
x = modulate(ctx->ggml_ctx, x, scale);
x = linear->forward(ctx, x);
return x;
}
};
class Ideogram4Transformer : public GGMLBlock {
protected:
Ideogram4Config config;
public:
Ideogram4Transformer() = default;
explicit Ideogram4Transformer(Ideogram4Config config)
: config(std::move(config)) {
blocks["input_proj"] = make_linear(this->config.in_channels, this->config.emb_dim, true);
blocks["llm_cond_norm"] = std::make_shared<RMSNorm>(this->config.llm_features_dim, 1e-6f);
blocks["llm_cond_proj"] = make_linear(this->config.llm_features_dim, this->config.emb_dim, true);
blocks["t_embedding"] = std::make_shared<Ideogram4EmbedScalar>(this->config.emb_dim);
blocks["adaln_proj"] = make_linear(this->config.emb_dim, this->config.adanln_dim, true);
blocks["embed_image_indicator"] = std::make_shared<Embedding>(2, this->config.emb_dim);
for (int i = 0; i < this->config.num_layers; ++i) {
blocks["layers." + std::to_string(i)] = std::make_shared<Ideogram4TransformerBlock>(this->config);
}
blocks["final_layer"] = std::make_shared<Ideogram4FinalLayer>(this->config);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* image_indicator_ids) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t N = x->ne[3];
GGML_ASSERT(N == 1);
auto input_proj = std::dynamic_pointer_cast<Linear>(blocks["input_proj"]);
auto llm_cond_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["llm_cond_norm"]);
auto llm_cond_proj = std::dynamic_pointer_cast<Linear>(blocks["llm_cond_proj"]);
auto t_embedding = std::dynamic_pointer_cast<Ideogram4EmbedScalar>(blocks["t_embedding"]);
auto adaln_proj = std::dynamic_pointer_cast<Linear>(blocks["adaln_proj"]);
auto embed_image_indicator = std::dynamic_pointer_cast<Embedding>(blocks["embed_image_indicator"]);
auto final_layer = std::dynamic_pointer_cast<Ideogram4FinalLayer>(blocks["final_layer"]);
auto img = patchify(ctx->ggml_ctx, x, config);
img = input_proj->forward(ctx, img);
ggml_tensor* h = img;
int64_t context_len = 0;
if (context != nullptr) {
if (ggml_n_dims(context) < 3) {
context = ggml_reshape_3d(ctx->ggml_ctx, context, context->ne[0], context->ne[1], 1);
}
context = interleave_hidden_state_layers(ctx->ggml_ctx, context);
context_len = context->ne[1];
auto txt = llm_cond_norm->forward(ctx, context);
txt = llm_cond_proj->forward(ctx, txt);
h = ggml_concat(ctx->ggml_ctx, txt, img, 1);
}
auto indicator_embedding = embed_image_indicator->forward(ctx, image_indicator_ids);
h = ggml_add(ctx->ggml_ctx, h, indicator_embedding);
auto t_cond = t_embedding->forward(ctx, timestep);
auto adaln_input = ggml_silu(ctx->ggml_ctx, adaln_proj->forward(ctx, t_cond));
for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<Ideogram4TransformerBlock>(blocks["layers." + std::to_string(i)]);
h = block->forward(ctx, h, pe, adaln_input, nullptr);
sd::ggml_graph_cut::mark_graph_cut(h, "ideogram4.layers." + std::to_string(i), "hidden");
}
h = final_layer->forward(ctx, h, adaln_input);
if (context_len > 0) {
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]);
}
h = unpatchify(ctx->ggml_ctx, h, H, W, config);
h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f);
return h;
}
};
class Ideogram4Runner : public DiffusionModelRunner {
protected:
bool should_use_uncond_model(const DiffusionParams& diffusion_params) const {
return has_uncond_model &&
diffusion_params.context == nullptr &&
diffusion_params.y != nullptr &&
!diffusion_params.y->empty();
}
public:
Ideogram4Config config;
Ideogram4Transformer model;
Ideogram4Transformer uncond_model;
bool has_uncond_model = false;
std::string uncond_prefix;
std::vector<float> pe_vec;
std::vector<int32_t> image_indicator_vec;
Ideogram4Runner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix),
config(Ideogram4Config::detect_from_weights(tensor_storage_map, prefix)),
uncond_prefix(prefix + ".uncond") {
model = Ideogram4Transformer(config);
model.init(params_ctx, tensor_storage_map, prefix);
for (const auto& pair : tensor_storage_map) {
const std::string& name = pair.first;
if (starts_with(name, uncond_prefix)) {
has_uncond_model = true;
break;
}
}
if (has_uncond_model) {
LOG_DEBUG("using uncond model");
uncond_model = Ideogram4Transformer(config);
uncond_model.init(params_ctx, tensor_storage_map, uncond_prefix);
}
}
std::string get_desc() override {
return "ideogram4";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.get_param_tensors(tensors, prefix);
if (has_uncond_model) {
uncond_model.get_param_tensors(tensors, this->uncond_prefix);
}
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
bool use_uncond_model = false) {
ggml_cgraph* gf = new_graph_custom(IDEOGRAM4_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x->ne[3] == 1);
Ideogram4Transformer& active_model = use_uncond_model ? uncond_model : model;
ggml_tensor* context = nullptr;
int64_t context_len = 0;
if (!context_tensor.empty()) {
context = make_input(context_tensor);
context_len = context->ne[1];
}
int64_t grid_w = x->ne[0];
int64_t grid_h = x->ne[1];
int64_t pos_len = context_len + grid_h * grid_w;
int64_t head_dim = config.emb_dim / config.num_heads;
pe_vec = gen_ideogram4_pe(static_cast<int>(grid_h),
static_cast<int>(grid_w),
static_cast<int>(x->ne[3]),
static_cast<int>(context_len),
static_cast<int>(head_dim),
static_cast<int>(config.rope_theta),
config.mrope_section);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
image_indicator_vec.assign(static_cast<size_t>(pos_len), 1);
for (int64_t i = 0; i < context_len; ++i) {
image_indicator_vec[static_cast<size_t>(i)] = 0;
}
auto indicator = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_I32, pos_len, x->ne[3]);
set_backend_tensor_data(indicator, image_indicator_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = active_model.forward(&runner_ctx, x, timesteps, context, pe, indicator);
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,
bool use_uncond_model = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, use_uncond_model);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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);
bool use_uncond_model = should_use_uncond_model(diffusion_params);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
use_uncond_model);
}
};
} // namespace Ideogram4
#endif // __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
src/model/diffusion/lens.hpp
+426
View File
@@ -0,0 +1,426 @@
#ifndef __SD_MODEL_DIFFUSION_LENS_HPP__
#define __SD_MODEL_DIFFUSION_LENS_HPP__
#include <memory>
#include <vector>
#include "model/common/block.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/qwen_image.hpp"
namespace Lens {
constexpr int LENS_GRAPH_SIZE = 40960;
struct LensConfig {
int patch_size = 2;
int64_t in_channels = 128;
int64_t out_channels = 32;
int num_layers = 48;
int64_t attention_head_dim = 64;
int64_t num_attention_heads = 24;
int64_t joint_attention_dim = 2880;
int selected_layer_count = 4;
int theta = 10000;
std::vector<int> axes_dim = {8, 28, 28};
int axes_dim_sum = 64;
static LensConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
LensConfig config;
config.num_layers = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
config.in_channels = tensor_storage.ne[0];
int64_t inner_dim = tensor_storage.ne[1];
if (config.attention_head_dim > 0) {
config.num_attention_heads = inner_dim / config.attention_head_dim;
}
} else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
config.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / config.joint_attention_dim);
} else if (ends_with(name, "proj_out.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;
} else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
config.attention_head_dim = tensor_storage.ne[0];
}
size_t pos = name.find("transformer_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
}
if (config.num_layers == 0) {
config.num_layers = 48;
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
LOG_DEBUG("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
config.num_layers,
config.selected_layer_count,
config.num_attention_heads * config.attention_head_dim,
config.num_attention_heads,
config.attention_head_dim,
config.in_channels,
config.out_channels);
return config;
}
};
struct LensTimestepProjEmbeddings : public GGMLBlock {
LensTimestepProjEmbeddings(int64_t embedding_dim) {
blocks["timestep_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(256, embedding_dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f);
return timestep_embedder->forward(ctx, timesteps_proj);
}
};
struct LensGateMLP : public GGMLBlock {
LensGateMLP(int64_t dim, int64_t hidden_dim) {
blocks["w1"] = std::make_shared<Linear>(dim, hidden_dim, false);
blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false);
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
auto up = w3->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, gate, up);
return w2->forward(ctx, x);
}
};
struct LensJointAttention : public GGMLBlock {
int64_t dim_head;
int64_t num_heads;
LensJointAttention(int64_t query_dim,
int64_t dim_head,
int64_t num_heads,
float eps = 1e-5f)
: dim_head(dim_head), num_heads(num_heads) {
int64_t inner_dim = dim_head * num_heads;
blocks["img_qkv"] = std::make_shared<Linear>(query_dim, inner_dim * 3, true);
blocks["txt_qkv"] = std::make_shared<Linear>(query_dim, inner_dim * 3, true);
blocks["norm_q"] = std::make_shared<RMSNorm>(dim_head, eps);
blocks["norm_k"] = std::make_shared<RMSNorm>(dim_head, eps);
blocks["norm_added_q"] = std::make_shared<RMSNorm>(dim_head, eps);
blocks["norm_added_k"] = std::make_shared<RMSNorm>(dim_head, eps);
blocks["to_out.0"] = std::make_shared<Linear>(inner_dim, query_dim, true);
blocks["to_add_out"] = std::make_shared<Linear>(inner_dim, query_dim, true);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
auto img_qkv = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
auto txt_qkv = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
auto norm_add_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_added_q"]);
auto norm_add_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_added_k"]);
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
auto to_add_out = std::dynamic_pointer_cast<Linear>(blocks["to_add_out"]);
int64_t n_img = img->ne[1];
int64_t n_txt = txt->ne[1];
int64_t N = img->ne[2];
int64_t inner = dim_head * num_heads;
auto img_qkv_vec = split_qkv(ctx->ggml_ctx, img_qkv->forward(ctx, img));
auto txt_qkv_vec = split_qkv(ctx->ggml_ctx, txt_qkv->forward(ctx, txt));
auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N);
auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N);
auto img_v = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[2], dim_head, num_heads, n_img, N);
img_q = norm_q->forward(ctx, img_q);
img_k = norm_k->forward(ctx, img_k);
auto txt_q = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[0], dim_head, num_heads, n_txt, N);
auto txt_k = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[1], dim_head, num_heads, n_txt, N);
auto txt_v = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[2], dim_head, num_heads, n_txt, N);
txt_q = norm_add_q->forward(ctx, txt_q);
txt_k = norm_add_k->forward(ctx, txt_k);
auto q = ggml_concat(ctx->ggml_ctx, img_q, txt_q, 2);
auto k = ggml_concat(ctx->ggml_ctx, img_k, txt_k, 2);
auto v = ggml_concat(ctx->ggml_ctx, img_v, txt_v, 2);
auto attn = Rope::attention(ctx, q, k, v, pe, mask, (1.0f / 128.f));
auto img_attn_out = ggml_view_3d(ctx->ggml_ctx,
attn,
inner,
n_img,
N,
attn->nb[1],
attn->nb[2],
0);
auto txt_attn_out = ggml_view_3d(ctx->ggml_ctx,
attn,
inner,
n_txt,
N,
attn->nb[1],
attn->nb[2],
n_img * attn->nb[1]);
img_attn_out = to_out_0->forward(ctx, ggml_cont(ctx->ggml_ctx, img_attn_out));
txt_attn_out = to_add_out->forward(ctx, ggml_cont(ctx->ggml_ctx, txt_attn_out));
return {img_attn_out, txt_attn_out};
}
};
struct LensTransformerBlock : public GGMLBlock {
LensTransformerBlock(int64_t dim,
int64_t num_attention_heads,
int64_t attention_head_dim,
float eps = 1e-6f) {
int64_t mlp_hidden_dim = dim / 3 * 8;
blocks["img_mod.1"] = std::make_shared<Linear>(dim, 6 * dim, true);
blocks["txt_mod.1"] = std::make_shared<Linear>(dim, 6 * dim, true);
blocks["img_norm1"] = std::make_shared<RMSNorm>(dim, eps);
blocks["img_norm2"] = std::make_shared<RMSNorm>(dim, eps);
blocks["txt_norm1"] = std::make_shared<RMSNorm>(dim, eps);
blocks["txt_norm2"] = std::make_shared<RMSNorm>(dim, eps);
blocks["img_mlp"] = std::make_shared<LensGateMLP>(dim, mlp_hidden_dim);
blocks["txt_mlp"] = std::make_shared<LensGateMLP>(dim, mlp_hidden_dim);
blocks["attn"] = std::make_shared<LensJointAttention>(dim, attention_head_dim, num_attention_heads);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* t_emb,
ggml_tensor* pe) {
auto img_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["img_mod.1"]);
auto txt_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["txt_mod.1"]);
auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
auto img_mlp = std::dynamic_pointer_cast<LensGateMLP>(blocks["img_mlp"]);
auto txt_mlp = std::dynamic_pointer_cast<LensGateMLP>(blocks["txt_mlp"]);
auto attn = std::dynamic_pointer_cast<LensJointAttention>(blocks["attn"]);
auto temb = ggml_silu(ctx->ggml_ctx, t_emb);
auto img_mod_params = img_mod_1->forward(ctx, temb);
auto img_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0);
auto txt_mod_params = txt_mod_1->forward(ctx, temb);
auto txt_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
auto img_normed = img_norm1->forward(ctx, img);
auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
auto txt_normed = txt_norm1->forward(ctx, txt);
auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn_output, img_mod_vec[2]));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_mod_vec[2]));
auto img_normed2 = img_norm2->forward(ctx, img);
auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
auto txt_normed2 = txt_norm2->forward(ctx, txt);
auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp->forward(ctx, img_modulated2), img_mod_vec[5]));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp->forward(ctx, txt_modulated2), txt_mod_vec[5]));
return {img, txt};
}
};
struct LensAdaLayerNormContinuous : public GGMLBlock {
int64_t hidden_size;
float eps;
LensAdaLayerNormContinuous(int64_t hidden_size, float eps = 1e-6f)
: hidden_size(hidden_size), eps(eps) {
blocks["linear"] = std::make_shared<Linear>(hidden_size, hidden_size * 2, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
auto mods = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
auto scale = mods[0];
auto shift = mods[1];
x = ggml_norm(ctx->ggml_ctx, x, eps);
return Flux::modulate(ctx->ggml_ctx, x, shift, scale);
}
};
class LensModel : public GGMLBlock {
public:
LensConfig config;
LensModel() = default;
LensModel(LensConfig config)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::make_shared<LensTimestepProjEmbeddings>(inner_dim);
blocks["img_in"] = std::make_shared<Linear>(config.in_channels, inner_dim, true);
blocks["txt_in"] = std::make_shared<Linear>(config.joint_attention_dim * config.selected_layer_count, inner_dim, true);
for (int i = 0; i < config.selected_layer_count; ++i) {
blocks["txt_norm." + std::to_string(i)] = std::make_shared<RMSNorm>(config.joint_attention_dim, 1e-5f);
}
for (int i = 0; i < config.num_layers; ++i) {
blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<LensTransformerBlock>(inner_dim,
config.num_attention_heads,
config.attention_head_dim);
}
blocks["norm_out"] = std::make_shared<LensAdaLayerNormContinuous>(inner_dim, 1e-6f);
blocks["proj_out"] = std::make_shared<Linear>(inner_dim, config.patch_size * config.patch_size * config.out_channels, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
GGML_ASSERT(context != nullptr);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int64_t N = x->ne[3];
auto time_text_embed = std::dynamic_pointer_cast<LensTimestepProjEmbeddings>(blocks["time_text_embed"]);
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
auto norm_out = std::dynamic_pointer_cast<LensAdaLayerNormContinuous>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto t_emb = time_text_embed->forward(ctx, timestep);
auto img = ggml_reshape_3d(ctx->ggml_ctx, x, W * H, C, N);
img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
img = img_in->forward(ctx, img);
std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0);
ggml_tensor* txt = nullptr;
for (int i = 0; i < config.selected_layer_count; ++i) {
auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm." + std::to_string(i)]);
auto chunk = txt_norm->forward(ctx, txt_chunks[i]);
txt = txt == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, txt, chunk, 0);
}
txt = txt_in->forward(ctx, txt);
sd::ggml_graph_cut::mark_graph_cut(img, "lens.prelude", "img");
sd::ggml_graph_cut::mark_graph_cut(txt, "lens.prelude", "txt");
for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<LensTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
auto out = block->forward(ctx, img, txt, t_emb, pe);
img = out.first;
txt = out.second;
sd::ggml_graph_cut::mark_graph_cut(img, "lens.transformer_blocks." + std::to_string(i), "img");
sd::ggml_graph_cut::mark_graph_cut(txt, "lens.transformer_blocks." + std::to_string(i), "txt");
}
img = norm_out->forward(ctx, img, t_emb);
img = proj_out->forward(ctx, img);
auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
out = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N);
return out;
}
};
struct LensRunner : public DiffusionModelRunner {
LensConfig config;
LensModel lens;
std::vector<float> pe_vec;
LensRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix),
config(LensConfig::detect_from_weights(tensor_storage_map, prefix)) {
lens = LensModel(config);
lens.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "lens";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
lens.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor) {
ggml_cgraph* gf = new_graph_custom(LENS_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);
pe_vec = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
circular_y_enabled,
circular_x_enabled,
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 = lens.forward(&runner_ctx, x, timesteps, context, pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context));
}
};
} // namespace Lens
#endif // __SD_MODEL_DIFFUSION_LENS_HPP__
src/model/diffusion/ltxv.hpp
+2065
View File
File diff suppressed because it is too large. Load diff
src/model/diffusion/mmdit.hpp
+1035
View File
File diff suppressed because it is too large. Load diff
src/model/diffusion/model.hpp
+107
View File
@@ -0,0 +1,107 @@
#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
#define __SD_MODEL_DIFFUSION_MODEL_HPP__
#include <string>
#include <utility>
#include <variant>
#include "core/ggml_extend.hpp"
#include "core/tensor_ggml.hpp"
struct UNetDiffusionExtra {
int num_video_frames = -1;
const std::vector<sd::Tensor<float>>* controls = nullptr;
float control_strength = 0.f;
};
struct SkipLayerDiffusionExtra {
const std::vector<int>* skip_layers = nullptr;
};
struct FluxDiffusionExtra {
const sd::Tensor<float>* guidance = nullptr;
const std::vector<int>* skip_layers = nullptr;
};
struct AnimaDiffusionExtra {
const sd::Tensor<int32_t>* t5_ids = nullptr;
const sd::Tensor<float>* t5_weights = nullptr;
};
struct WanDiffusionExtra {
const sd::Tensor<float>* vace_context = nullptr;
float vace_strength = 1.f;
};
struct HiDreamO1DiffusionExtra {
const sd::Tensor<int32_t>* input_ids = nullptr;
const sd::Tensor<int32_t>* input_pos = nullptr;
const sd::Tensor<int32_t>* token_types = nullptr;
const sd::Tensor<int32_t>* vinput_mask = nullptr;
const std::vector<std::pair<int, sd::Tensor<float>>>* image_embeds = nullptr;
};
struct LTXAVDiffusionExtra {
const sd::Tensor<float>* audio_x = nullptr;
const sd::Tensor<float>* audio_timesteps = nullptr;
int audio_length = 0;
float frame_rate = 24.f;
const sd::Tensor<float>* video_positions = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
FluxDiffusionExtra,
AnimaDiffusionExtra,
WanDiffusionExtra,
HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;
const sd::Tensor<float>* timesteps = nullptr;
const sd::Tensor<float>* context = nullptr;
const sd::Tensor<float>* c_concat = nullptr;
const sd::Tensor<float>* y = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
bool increase_ref_index = false;
DiffusionExtraParams extra = std::monostate{};
};
template <typename T>
static inline const T* diffusion_extra_as(const DiffusionParams& params) {
const auto* extra = std::get_if<T>(&params.extra);
GGML_ASSERT(extra != nullptr);
return extra;
}
template <typename T>
static inline const sd::Tensor<T>& tensor_or_empty(const sd::Tensor<T>* tensor) {
static const sd::Tensor<T> kEmpty;
return tensor != nullptr ? *tensor : kEmpty;
}
struct DiffusionModelRunner : public GGMLRunner {
protected:
std::string prefix;
public:
DiffusionModelRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const std::string& prefix)
: GGMLRunner(backend, params_backend),
prefix(prefix) {}
virtual sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) = 0;
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
get_param_tensors(tensors, prefix);
}
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors,
const std::string& prefix) = 0;
};
#endif // __SD_MODEL_DIFFUSION_MODEL_HPP__
src/model/diffusion/pid.hpp
+847
View File
@@ -0,0 +1,847 @@
#ifndef __SD_MODEL_DIFFUSION_PID_HPP__
#define __SD_MODEL_DIFFUSION_PID_HPP__
#include <cmath>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.hpp"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/mmdit.hpp"
namespace Pid {
constexpr int PID_GRAPH_SIZE = 196608;
constexpr float PID_PI = 3.14159265358979323846f;
struct PixelDiTConfig {
int64_t in_channels = 3;
int64_t hidden_size = 1536;
int64_t num_groups = 24;
int64_t patch_mlp_hidden_dim = 4096;
int64_t pixel_hidden_size = 16;
int64_t pixel_attn_hidden_size = 1152;
int64_t pixel_num_groups = 16;
int64_t patch_depth = 14;
int64_t pixel_depth = 2;
int64_t patch_size = 16;
int64_t txt_embed_dim = 2304;
int64_t txt_max_length = 300;
float text_rope_theta = 10000.f;
int64_t lq_latent_channels = 16;
int64_t lq_hidden_dim = 512;
int64_t lq_num_res_blocks = 4;
int64_t lq_interval = 2;
int64_t lq_sr_scale = 4;
int64_t lq_latent_down_factor = 8;
int64_t rope_ref_grid_h = 64;
int64_t rope_ref_grid_w = 64;
static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
PixelDiTConfig config;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
size_t pos = name.find("patch_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
config.patch_depth = std::max<int64_t>(config.patch_depth, block_index + 1);
}
}
pos = name.find("pixel_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
}
}
if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
config.lq_latent_channels = tensor_storage.ne[2];
config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
}
if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
config.patch_mlp_hidden_dim = tensor_storage.ne[1];
}
}
LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64,
config.patch_depth,
config.pixel_depth,
config.patch_mlp_hidden_dim,
config.lq_latent_channels,
config.lq_latent_down_factor);
return config;
}
};
inline std::vector<float> make_rope_1d(int length,
int dim,
float theta) {
GGML_ASSERT(dim % 2 == 0);
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
}
inline std::vector<float> make_rope_2d(int height,
int width,
int dim,
float theta = 10000.f,
float scale = 16.f,
int ref_grid_h = 0,
int ref_grid_w = 0) {
GGML_ASSERT(dim % 4 == 0);
return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w);
}
inline std::vector<float> make_pixel_abs_pos(int height,
int width,
int dim) {
GGML_ASSERT(dim % 4 == 0);
int half_dim = dim / 2;
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) {
for (int ix = 0; ix < width; ++ix) {
x_pos.push_back(static_cast<float>(ix));
y_pos.push_back(static_cast<float>(iy));
}
}
auto x_emb = timestep_embedding(x_pos, half_dim, 10000, false);
auto y_emb = timestep_embedding(y_pos, half_dim, 10000, false);
std::vector<float> out(static_cast<size_t>(dim) * height * width);
for (int pos = 0; pos < height * width; ++pos) {
size_t out_base = static_cast<size_t>(pos) * dim;
size_t emb_base = static_cast<size_t>(pos) * half_dim;
for (int i = 0; i < half_dim; ++i) {
out[out_base + i] = x_emb[emb_base + i];
out[out_base + half_dim + i] = y_emb[emb_base + i];
}
}
return out;
}
inline ggml_tensor* apply_adaln(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
}
struct PatchTokenEmbedder : public GGMLBlock {
bool use_rms_norm;
PatchTokenEmbedder(int64_t in_chans,
int64_t embed_dim,
bool use_rms_norm = false,
bool bias = true)
: use_rms_norm(use_rms_norm) {
blocks["proj"] = std::make_shared<Linear>(in_chans, embed_dim, bias);
if (use_rms_norm) {
blocks["norm"] = std::make_shared<RMSNorm>(embed_dim, 1e-6f);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
x = proj->forward(ctx, x);
if (use_rms_norm) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
x = norm->forward(ctx, x);
}
return x;
}
};
struct PixelDiTTimestepEmbedder : public GGMLBlock {
int frequency_embedding_size;
PixelDiTTimestepEmbedder(int64_t hidden_size,
int frequency_embedding_size = 256)
: frequency_embedding_size(frequency_embedding_size) {
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10);
t_emb = mlp_0->forward(ctx, t_emb);
t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
return mlp_2->forward(ctx, t_emb);
}
};
struct FeedForward : public GGMLBlock {
FeedForward(int64_t dim, int64_t hidden_dim) {
blocks["w1"] = std::make_shared<Linear>(dim, hidden_dim, false);
blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false);
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto h = ggml_silu_inplace(ctx->ggml_ctx, w1->forward(ctx, x));
h = ggml_mul_inplace(ctx->ggml_ctx, h, w3->forward(ctx, x));
return w2->forward(ctx, h);
}
};
struct FinalLayer : public GGMLBlock {
FinalLayer(int64_t hidden_size, int64_t out_channels) {
blocks["norm"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["linear"] = std::make_shared<Linear>(hidden_size, out_channels, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
return linear->forward(ctx, norm->forward(ctx, x));
}
};
struct RotaryAttention : public GGMLBlock {
int64_t dim;
int64_t num_heads;
RotaryAttention(int64_t dim, int64_t num_heads)
: dim(dim), num_heads(num_heads) {
int64_t head_dim = dim / num_heads;
blocks["qkv"] = std::make_shared<Linear>(dim, dim * 3, false);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["proj"] = std::make_shared<Linear>(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pos) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto qkv = qkv_proj->forward(ctx, x);
auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
int64_t L = x->ne[1];
int64_t N = x->ne[2];
int64_t head_dim = dim / num_heads;
auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, L, N);
auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, L, N);
auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, L, N);
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
x = Rope::attention(ctx, q, k, v, pos, nullptr, 1.0f / 128.f, true);
return proj->forward(ctx, x);
}
};
struct MMDiTJointAttention : public GGMLBlock {
int64_t dim;
int64_t num_heads;
MMDiTJointAttention(int64_t dim, int64_t num_heads)
: dim(dim), num_heads(num_heads) {
int64_t head_dim = dim / num_heads;
blocks["qkv_x"] = std::make_shared<Linear>(dim, dim * 3, false);
blocks["qkv_y"] = std::make_shared<Linear>(dim, dim * 3, false);
blocks["q_norm_x"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm_x"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["q_norm_y"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm_y"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["proj_x"] = std::make_shared<Linear>(dim, dim, true);
blocks["proj_y"] = std::make_shared<Linear>(dim, dim, true);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* y,
ggml_tensor* pos_img,
ggml_tensor* pos_txt) {
auto qkv_x_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_x"]);
auto qkv_y_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_y"]);
auto q_norm_x = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm_x"]);
auto k_norm_x = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm_x"]);
auto q_norm_y = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm_y"]);
auto k_norm_y = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm_y"]);
auto proj_x = std::dynamic_pointer_cast<Linear>(blocks["proj_x"]);
auto proj_y = std::dynamic_pointer_cast<Linear>(blocks["proj_y"]);
int64_t Nx = x->ne[1];
int64_t Ny = y->ne[1];
int64_t N = x->ne[2];
int64_t head_dim = dim / num_heads;
auto qkv_x = split_qkv(ctx->ggml_ctx, qkv_x_proj->forward(ctx, x));
auto qx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N);
auto kx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N);
auto vx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N);
qx = q_norm_x->forward(ctx, qx);
kx = k_norm_x->forward(ctx, kx);
auto qkv_y = split_qkv(ctx->ggml_ctx, qkv_y_proj->forward(ctx, y));
auto qy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N);
auto ky = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N);
auto vy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N);
qy = q_norm_y->forward(ctx, qy);
ky = k_norm_y->forward(ctx, ky);
auto q_joint = ggml_concat(ctx->ggml_ctx, qy, qx, 2);
auto k_joint = ggml_concat(ctx->ggml_ctx, ky, kx, 2);
auto v_joint = ggml_concat(ctx->ggml_ctx, vy, vx, 2);
auto pos_joint = ggml_concat(ctx->ggml_ctx, pos_txt, pos_img, 3);
auto out = Rope::attention(ctx, q_joint, k_joint, v_joint, pos_joint, nullptr, 1.0f, true);
auto out_y = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, Ny);
auto out_x = ggml_ext_slice(ctx->ggml_ctx, out, 1, Ny, Ny + Nx);
return {proj_x->forward(ctx, out_x), proj_y->forward(ctx, out_y)};
}
};
struct MMDiTBlockT2I : public GGMLBlock {
int64_t hidden_size;
MMDiTBlockT2I(int64_t hidden_size, int64_t groups, int64_t mlp_hidden_dim)
: hidden_size(hidden_size) {
blocks["norm_x1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["norm_y1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["attn"] = std::make_shared<MMDiTJointAttention>(hidden_size, groups);
blocks["norm_x2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["norm_y2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["mlp_x"] = std::make_shared<FeedForward>(hidden_size, mlp_hidden_dim);
blocks["mlp_y"] = std::make_shared<FeedForward>(hidden_size, mlp_hidden_dim);
blocks["adaLN_modulation_img.0"] = std::make_shared<Linear>(hidden_size, 6 * hidden_size, true);
blocks["adaLN_modulation_txt.0"] = std::make_shared<Linear>(hidden_size, 6 * hidden_size, true);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* y,
ggml_tensor* c,
ggml_tensor* pos_img,
ggml_tensor* pos_txt) {
auto norm_x1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_x1"]);
auto norm_y1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_y1"]);
auto attn = std::dynamic_pointer_cast<MMDiTJointAttention>(blocks["attn"]);
auto norm_x2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_x2"]);
auto norm_y2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_y2"]);
auto mlp_x = std::dynamic_pointer_cast<FeedForward>(blocks["mlp_x"]);
auto mlp_y = std::dynamic_pointer_cast<FeedForward>(blocks["mlp_y"]);
auto ada_img = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_img.0"]);
auto ada_txt = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_txt.0"]);
auto mx = ggml_ext_chunk(ctx->ggml_ctx, ada_img->forward(ctx, c), 6, 0);
auto my = ggml_ext_chunk(ctx->ggml_ctx, ada_txt->forward(ctx, c), 6, 0);
auto x_norm = apply_adaln(ctx->ggml_ctx, norm_x1->forward(ctx, x), mx[0], mx[1]);
auto y_norm = apply_adaln(ctx->ggml_ctx, norm_y1->forward(ctx, y), my[0], my[1]);
auto attn_out = attn->forward(ctx, x_norm, y_norm, pos_img, pos_txt);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out.first, mx[2]));
y = ggml_add(ctx->ggml_ctx, y, ggml_mul(ctx->ggml_ctx, attn_out.second, my[2]));
auto x_mlp = mlp_x->forward(ctx, apply_adaln(ctx->ggml_ctx, norm_x2->forward(ctx, x), mx[3], mx[4]));
auto y_mlp = mlp_y->forward(ctx, apply_adaln(ctx->ggml_ctx, norm_y2->forward(ctx, y), my[3], my[4]));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, x_mlp, mx[5]));
y = ggml_add(ctx->ggml_ctx, y, ggml_mul(ctx->ggml_ctx, y_mlp, my[5]));
return {x, y};
}
};
struct PixelTokenEmbedder : public GGMLBlock {
int64_t in_channels;
int64_t hidden_size_output;
PixelTokenEmbedder(int64_t in_channels, int64_t hidden_size_output)
: in_channels(in_channels), hidden_size_output(hidden_size_output) {
blocks["proj"] = std::make_shared<Linear>(in_channels, hidden_size_output, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* inputs,
int64_t patch_size,
ggml_tensor* pos_full) {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
int64_t W = inputs->ne[0];
int64_t H = inputs->ne[1];
int64_t B = inputs->ne[3];
int64_t L = (W / patch_size) * (H / patch_size);
int64_t P2 = patch_size * patch_size;
auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, inputs, 2, 0, 1, 3));
x = ggml_reshape_3d(ctx->ggml_ctx, x, in_channels, W * H, B);
x = proj->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, pos_full);
x = ggml_reshape_4d(ctx->ggml_ctx, x, hidden_size_output, W, H, B);
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
x = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size_output, P2, L * B);
return x;
}
};
struct PiTBlock : public GGMLBlock {
int64_t pixel_dim;
int64_t context_dim;
int64_t attn_dim;
int64_t num_heads;
int64_t patch_size;
PiTBlock(int64_t pixel_dim,
int64_t context_dim,
int64_t patch_size,
int64_t attn_dim,
int64_t num_heads)
: pixel_dim(pixel_dim),
context_dim(context_dim),
attn_dim(attn_dim),
num_heads(num_heads),
patch_size(patch_size) {
int64_t p2 = patch_size * patch_size;
blocks["compress_to_attn"] = std::make_shared<Linear>(p2 * pixel_dim, attn_dim, true);
blocks["expand_from_attn"] = std::make_shared<Linear>(attn_dim, p2 * pixel_dim, true);
blocks["norm1"] = std::make_shared<RMSNorm>(pixel_dim, 1e-6f);
blocks["attn"] = std::make_shared<RotaryAttention>(attn_dim, num_heads);
blocks["norm2"] = std::make_shared<RMSNorm>(pixel_dim, 1e-6f);
blocks["mlp"] = std::make_shared<Mlp>(pixel_dim, pixel_dim * 4);
blocks["adaLN_modulation.0"] = std::make_shared<Linear>(context_dim, 6 * pixel_dim * p2, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* s_cond,
int64_t image_height,
int64_t image_width,
ggml_tensor* pos_comp) {
auto compress = std::dynamic_pointer_cast<Linear>(blocks["compress_to_attn"]);
auto expand = std::dynamic_pointer_cast<Linear>(blocks["expand_from_attn"]);
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<RotaryAttention>(blocks["attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto mlp = std::dynamic_pointer_cast<Mlp>(blocks["mlp"]);
auto ada = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
int64_t Hs = image_height / patch_size;
int64_t Ws = image_width / patch_size;
int64_t L = Hs * Ws;
int64_t BL = x->ne[2];
int64_t B = BL / L;
int64_t P2 = patch_size * patch_size;
auto ada_params = ada->forward(ctx, s_cond);
ada_params = ggml_reshape_3d(ctx->ggml_ctx, ada_params, 6 * pixel_dim, P2, BL);
auto mod = ggml_ext_chunk(ctx->ggml_ctx, ada_params, 6, 0);
auto x_norm = apply_adaln(ctx->ggml_ctx, norm1->forward(ctx, x), mod[0], mod[1]);
auto x_flat = ggml_reshape_2d(ctx->ggml_ctx, x_norm, P2 * pixel_dim, BL);
auto x_comp = compress->forward(ctx, x_flat);
x_comp = ggml_reshape_3d(ctx->ggml_ctx, x_comp, attn_dim, L, B);
auto attn_out = attn->forward(ctx, x_comp, pos_comp);
auto attn_flat = expand->forward(ctx, ggml_reshape_2d(ctx->ggml_ctx, attn_out, attn_dim, BL));
auto attn_exp = ggml_reshape_3d(ctx->ggml_ctx, attn_flat, pixel_dim, P2, BL);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_exp, mod[2]));
auto mlp_out = mlp->forward(ctx, apply_adaln(ctx->ggml_ctx, norm2->forward(ctx, x), mod[3], mod[4]));
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mod[5]));
}
};
struct SigmaAwareGate : public GGMLBlock {
int64_t dim;
SigmaAwareGate(int64_t dim)
: dim(dim) {
blocks["content_proj"] = std::make_shared<Linear>(dim * 2, dim, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
params["log_alpha"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* lq,
ggml_tensor* sigma) {
auto content_proj = std::dynamic_pointer_cast<Linear>(blocks["content_proj"]);
auto content_logit = content_proj->forward(ctx, ggml_concat(ctx->ggml_ctx, x, lq, 0));
sigma = ggml_reshape_3d(ctx->ggml_ctx, sigma, 1, 1, sigma->ne[0]);
auto alpha = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
auto offset = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma));
auto gate = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset));
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq));
}
};
struct PiDResBlock : public GGMLBlock {
PiDResBlock(int64_t channels) {
blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.0"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["block.2"]);
auto norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
h = conv1->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
h = conv2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct LQProjection2D : public GGMLBlock {
PixelDiTConfig config;
LQProjection2D(const PixelDiTConfig& config)
: config(config) {
blocks["latent_proj.0"] = std::make_shared<Conv2d>(config.lq_latent_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim);
}
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
for (int i = 0; i < num_outputs; ++i) {
blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size);
}
}
bool is_gate_active(int block_idx) const {
return block_idx % config.lq_interval == 0;
}
int get_output_index(int block_idx) const {
return block_idx / static_cast<int>(config.lq_interval);
}
ggml_tensor* gate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* lq,
ggml_tensor* sigma,
int out_idx) {
auto gate_module = std::dynamic_pointer_cast<SigmaAwareGate>(blocks["gate_modules." + std::to_string(out_idx)]);
return gate_module->forward(ctx, x, lq, sigma);
}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* lq_latent,
int64_t target_pH,
int64_t target_pW) {
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
static_cast<float>(config.patch_size);
GGML_ASSERT(z_to_patch_ratio >= 1.0f);
if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
lq_latent = ggml_interpolate(ctx->ggml_ctx,
lq_latent,
target_pW,
target_pH,
lq_latent->ne[2],
lq_latent->ne[3],
GGML_SCALE_MODE_NEAREST);
}
auto feat = conv0->forward(ctx, lq_latent);
feat = ggml_silu_inplace(ctx->ggml_ctx, feat);
feat = conv2->forward(ctx, feat);
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
feat = block->forward(ctx, feat);
}
int64_t B = feat->ne[3];
int64_t C = feat->ne[2];
int64_t L = target_pH * target_pW;
auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3));
tokens = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B);
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
std::vector<ggml_tensor*> outputs;
outputs.reserve(num_outputs);
for (int i = 0; i < num_outputs; ++i) {
auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
outputs.push_back(head->forward(ctx, tokens));
}
return outputs;
}
};
struct PixelDiT : public GGMLBlock {
PixelDiTConfig config;
PixelDiT() = default;
PixelDiT(const PixelDiTConfig& config)
: config(config) {
blocks["pixel_embedder"] = std::make_shared<PixelTokenEmbedder>(config.in_channels, config.pixel_hidden_size);
blocks["s_embedder"] = std::make_shared<PatchTokenEmbedder>(config.in_channels * config.patch_size * config.patch_size, config.hidden_size, false, true);
blocks["t_embedder"] = std::make_shared<PixelDiTTimestepEmbedder>(config.hidden_size);
blocks["y_embedder"] = std::make_shared<PatchTokenEmbedder>(config.txt_embed_dim, config.hidden_size, true, true);
for (int i = 0; i < config.patch_depth; ++i) {
blocks["patch_blocks." + std::to_string(i)] = std::make_shared<MMDiTBlockT2I>(config.hidden_size, config.num_groups, config.patch_mlp_hidden_dim);
}
for (int i = 0; i < config.pixel_depth; ++i) {
blocks["pixel_blocks." + std::to_string(i)] = std::make_shared<PiTBlock>(config.pixel_hidden_size,
config.hidden_size,
config.patch_size,
config.pixel_attn_hidden_size,
config.pixel_num_groups);
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
blocks["lq_proj"] = std::make_shared<LQProjection2D>(config);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config.hidden_size, config.txt_max_length, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timesteps,
ggml_tensor* context,
ggml_tensor* lq_latent,
ggml_tensor* degrade_sigma,
ggml_tensor* pos_img,
ggml_tensor* pos_txt,
ggml_tensor* pixel_pos_full,
ggml_tensor* pixel_pos_comp) {
auto pixel_embedder = std::dynamic_pointer_cast<PixelTokenEmbedder>(blocks["pixel_embedder"]);
auto s_embedder = std::dynamic_pointer_cast<PatchTokenEmbedder>(blocks["s_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<PixelDiTTimestepEmbedder>(blocks["t_embedder"]);
auto y_embedder = std::dynamic_pointer_cast<PatchTokenEmbedder>(blocks["y_embedder"]);
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
auto lq_proj = std::dynamic_pointer_cast<LQProjection2D>(blocks["lq_proj"]);
int64_t W_orig = x->ne[0];
int64_t H_orig = x->ne[1];
x = DiT::pad_to_patch_size(ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size));
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t B = x->ne[3];
int64_t Hs = H / config.patch_size;
int64_t Ws = W / config.patch_size;
int64_t L = Hs * Ws;
int64_t P2 = config.patch_size * config.patch_size;
auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
auto t_emb = t_embedder->forward(ctx, timesteps);
auto condition = ggml_silu(ctx->ggml_ctx, t_emb);
GGML_ASSERT(context != nullptr);
int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
auto y = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt);
auto y_emb = y_embedder->forward(ctx, y);
auto y_pos = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt);
y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
auto s = s_embedder->forward(ctx, x_patches);
for (int i = 0; i < config.patch_depth; ++i) {
if (lq_proj->is_gate_active(i)) {
int out_idx = lq_proj->get_output_index(i);
if (out_idx < static_cast<int>(lq_features.size())) {
s = lq_proj->gate(ctx, s, lq_features[out_idx], degrade_sigma, out_idx);
}
}
auto block = std::dynamic_pointer_cast<MMDiTBlockT2I>(blocks["patch_blocks." + std::to_string(i)]);
auto out = block->forward(ctx,
s,
y_emb,
condition,
pos_img,
pos_txt);
s = out.first;
y_emb = out.second;
sd::ggml_graph_cut::mark_graph_cut(s, "pid.patch_blocks." + std::to_string(i), "s");
sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y");
}
s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B);
auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);
for (int i = 0; i < config.pixel_depth; ++i) {
auto block = std::dynamic_pointer_cast<PiTBlock>(blocks["pixel_blocks." + std::to_string(i)]);
pixels = block->forward(ctx, pixels, s_cond, H, W, pixel_pos_comp);
sd::ggml_graph_cut::mark_graph_cut(pixels, "pid.pixel_blocks." + std::to_string(i), "pixels");
}
pixels = final_layer->forward(ctx, pixels);
pixels = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B);
auto out = DiT::unpatchify(ctx->ggml_ctx,
pixels,
Hs,
Ws,
static_cast<int>(config.patch_size),
static_cast<int>(config.patch_size),
false);
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig);
out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig);
return out;
}
};
struct PiDRunner : public DiffusionModelRunner {
PixelDiTConfig config;
PixelDiT model;
std::vector<float> pos_img_vec;
std::vector<float> pos_txt_vec;
std::vector<float> pixel_pos_vec;
std::vector<float> pixel_pos_comp_vec;
PiDRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix = "model.diffusion_model")
: DiffusionModelRunner(backend, params_backend, prefix),
config(PixelDiTConfig::detect_from_weights(tensor_storage_map, prefix)) {
model = PixelDiT(config);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "PiD";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.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 sd::Tensor<float>& lq_latent_tensor,
const sd::Tensor<float>& degrade_sigma_tensor) {
ggml_cgraph* gf = new_graph_custom(PID_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* lq_latent = make_input(lq_latent_tensor);
ggml_tensor* degrade_sigma = make_input(degrade_sigma_tensor);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t B = x->ne[3];
int64_t Wp = align_up(static_cast<int>(W), static_cast<int>(config.patch_size));
int64_t Hp = align_up(static_cast<int>(H), static_cast<int>(config.patch_size));
int64_t Hs = Hp / config.patch_size;
int64_t Ws = Wp / config.patch_size;
pos_img_vec = make_rope_2d(static_cast<int>(Hs),
static_cast<int>(Ws),
static_cast<int>(config.hidden_size / config.num_groups),
10000.f,
16.f,
static_cast<int>(config.rope_ref_grid_h),
static_cast<int>(config.rope_ref_grid_w));
auto pos_img = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
2,
config.hidden_size / config.num_groups / 2,
Hs * Ws);
set_backend_tensor_data(pos_img, pos_img_vec.data());
int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
pos_txt_vec = make_rope_1d(static_cast<int>(Ltxt),
static_cast<int>(config.hidden_size / config.num_groups),
config.text_rope_theta);
auto pos_txt = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
2,
config.hidden_size / config.num_groups / 2,
Ltxt);
set_backend_tensor_data(pos_txt, pos_txt_vec.data());
pixel_pos_vec = make_pixel_abs_pos(static_cast<int>(Hp),
static_cast<int>(Wp),
static_cast<int>(config.pixel_hidden_size));
auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
GGML_TYPE_F32,
config.pixel_hidden_size,
Wp * Hp,
1);
set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
pixel_pos_comp_vec = make_rope_2d(static_cast<int>(Hs),
static_cast<int>(Ws),
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
10000.f,
16.f,
static_cast<int>(config.rope_ref_grid_h),
static_cast<int>(config.rope_ref_grid_w));
auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
2,
config.pixel_attn_hidden_size / config.pixel_num_groups / 2,
Hs * Ws);
set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data());
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx,
x,
timesteps,
context,
lq_latent,
degrade_sigma,
pos_img,
pos_txt,
pixel_pos,
pixel_pos_comp);
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 sd::Tensor<float>& lq_latent,
const sd::Tensor<float>& degrade_sigma) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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);
GGML_ASSERT(diffusion_params.context != nullptr);
GGML_ASSERT(diffusion_params.ref_latents != nullptr);
GGML_ASSERT(!diffusion_params.ref_latents->empty());
auto degrade_sigma = sd::Tensor<float>::from_vector({0.0f});
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*diffusion_params.context,
diffusion_params.ref_latents->front(),
degrade_sigma);
}
};
} // namespace Pid
#endif // __SD_MODEL_DIFFUSION_PID_HPP__
src/model/diffusion/qwen_image.hpp
+734
View File
@@ -0,0 +1,734 @@
#ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
#include <memory>
#include "model/common/block.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
namespace Qwen {
constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480;
struct QwenImageConfig {
int patch_size = 2;
int64_t in_channels = 64;
int64_t out_channels = 16;
int num_layers = 60;
int64_t attention_head_dim = 128;
int64_t num_attention_heads = 24;
int64_t joint_attention_dim = 3584;
int theta = 10000;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
bool zero_cond_t = false;
static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
QwenImageConfig config;
config.num_layers = 0;
for (const auto& [name, _] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (name.find("__index_timestep_zero__") != std::string::npos) {
config.zero_cond_t = true;
}
size_t pos = name.find("transformer_blocks.");
if (pos == std::string::npos) {
continue;
}
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
LOG_DEBUG("qwen_image: num_layers = %d, zero_cond_t = %s",
config.num_layers,
config.zero_cond_t ? "true" : "false");
return config;
}
};
struct TimestepEmbedding : public GGMLBlock {
public:
TimestepEmbedding(int64_t in_channels,
int64_t time_embed_dim,
int64_t out_dim = 0,
int64_t cond_proj_dim = 0,
bool sample_proj_bias = true) {
blocks["linear_1"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, time_embed_dim, sample_proj_bias));
if (cond_proj_dim > 0) {
blocks["cond_proj"] = std::shared_ptr<GGMLBlock>(new Linear(cond_proj_dim, in_channels, false));
}
if (out_dim <= 0) {
out_dim = time_embed_dim;
}
blocks["linear_2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, out_dim, sample_proj_bias));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* sample,
ggml_tensor* condition = nullptr) {
if (condition != nullptr) {
auto cond_proj = std::dynamic_pointer_cast<Linear>(blocks["cond_proj"]);
sample = ggml_add(ctx->ggml_ctx, sample, cond_proj->forward(ctx, condition));
}
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
sample = linear_1->forward(ctx, sample);
sample = ggml_silu_inplace(ctx->ggml_ctx, sample);
sample = linear_2->forward(ctx, sample);
return sample;
}
};
struct QwenTimestepProjEmbeddings : public GGMLBlock {
public:
QwenTimestepProjEmbeddings(int64_t embedding_dim) {
blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps) {
// timesteps: [N,]
// return: [N, embedding_dim]
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
auto timesteps_emb = timestep_embedder->forward(ctx, timesteps_proj);
return timesteps_emb;
}
};
struct QwenImageAttention : public GGMLBlock {
protected:
int64_t dim_head;
public:
QwenImageAttention(int64_t query_dim,
int64_t dim_head,
int64_t num_heads,
int64_t out_dim = 0,
int64_t out_context_dim = 0,
bool bias = true,
bool out_bias = true,
float eps = 1e-6)
: dim_head(dim_head) {
int64_t inner_dim = out_dim > 0 ? out_dim : dim_head * num_heads;
out_dim = out_dim > 0 ? out_dim : query_dim;
out_context_dim = out_context_dim > 0 ? out_context_dim : query_dim;
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
blocks["norm_q"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
blocks["norm_k"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
blocks["add_q_proj"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
blocks["add_k_proj"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
blocks["add_v_proj"] = std::shared_ptr<GGMLBlock>(new Linear(query_dim, inner_dim, bias));
blocks["norm_added_q"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
blocks["norm_added_k"] = std::shared_ptr<GGMLBlock>(new RMSNorm(dim_head, eps));
float scale = 1.f / 32.f;
bool force_prec_f32 = false;
// The purpose of the scale here is to prevent NaN issues in certain situations.
// For example when using CUDA but the weights are k-quants (not all prompts).
blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, out_dim, out_bias, false, force_prec_f32, scale));
// to_out.1 is nn.Dropout
blocks["to_add_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, out_context_dim, out_bias, false, false, scale));
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
// img: [N, n_img_token, hidden_size]
// txt: [N, n_txt_token, hidden_size]
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
// return: ([N, n_img_token, hidden_size], [N, n_txt_token, hidden_size])
auto norm_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_k"]);
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"]);
if (sd_backend_is(ctx->backend, "Vulkan")) {
to_out_0->set_force_prec_f32(true);
}
auto norm_added_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_q"]);
auto norm_added_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_k"]);
auto add_q_proj = std::dynamic_pointer_cast<Linear>(blocks["add_q_proj"]);
auto add_k_proj = std::dynamic_pointer_cast<Linear>(blocks["add_k_proj"]);
auto add_v_proj = std::dynamic_pointer_cast<Linear>(blocks["add_v_proj"]);
auto to_add_out = std::dynamic_pointer_cast<Linear>(blocks["to_add_out"]);
int64_t N = img->ne[2];
int64_t n_img_token = img->ne[1];
int64_t n_txt_token = txt->ne[1];
auto img_q = to_q->forward(ctx, img);
int64_t num_heads = img_q->ne[0] / dim_head;
img_q = ggml_reshape_4d(ctx->ggml_ctx, img_q, dim_head, num_heads, n_img_token, N); // [N, n_img_token, n_head, d_head]
auto img_k = to_k->forward(ctx, img);
img_k = ggml_reshape_4d(ctx->ggml_ctx, img_k, dim_head, num_heads, n_img_token, N); // [N, n_img_token, n_head, d_head]
auto img_v = to_v->forward(ctx, img);
img_v = ggml_reshape_4d(ctx->ggml_ctx, img_v, dim_head, num_heads, n_img_token, N); // [N, n_img_token, n_head, d_head]
img_q = norm_q->forward(ctx, img_q);
img_k = norm_k->forward(ctx, img_k);
auto txt_q = add_q_proj->forward(ctx, txt);
txt_q = ggml_reshape_4d(ctx->ggml_ctx, txt_q, dim_head, num_heads, n_txt_token, N); // [N, n_txt_token, n_head, d_head]
auto txt_k = add_k_proj->forward(ctx, txt);
txt_k = ggml_reshape_4d(ctx->ggml_ctx, txt_k, dim_head, num_heads, n_txt_token, N); // [N, n_txt_token, n_head, d_head]
auto txt_v = add_v_proj->forward(ctx, txt);
txt_v = ggml_reshape_4d(ctx->ggml_ctx, txt_v, dim_head, num_heads, n_txt_token, N); // [N, n_txt_token, n_head, d_head]
txt_q = norm_added_q->forward(ctx, txt_q);
txt_k = norm_added_k->forward(ctx, txt_k);
auto q = ggml_concat(ctx->ggml_ctx, txt_q, img_q, 2); // [N, n_txt_token + n_img_token, n_head, d_head]
auto k = ggml_concat(ctx->ggml_ctx, txt_k, img_k, 2); // [N, n_txt_token + n_img_token, n_head, d_head]
auto v = ggml_concat(ctx->ggml_ctx, txt_v, img_v, 2); // [N, n_txt_token + n_img_token, n_head, d_head]
auto attn = Rope::attention(ctx, q, k, v, pe, mask, (1.0f / 128.f)); // [N, n_txt_token + n_img_token, n_head*d_head]
auto txt_attn_out = ggml_view_3d(ctx->ggml_ctx,
attn,
attn->ne[0],
txt->ne[1],
attn->ne[2],
attn->nb[1],
attn->nb[2],
0); // [N, n_txt_token, n_head*d_head]
auto img_attn_out = ggml_view_3d(ctx->ggml_ctx,
attn,
attn->ne[0],
img->ne[1],
attn->ne[2],
attn->nb[1],
attn->nb[2],
txt->ne[1] * attn->nb[1]); // [N, n_img_token, n_head*d_head]
img_attn_out = ggml_cont(ctx->ggml_ctx, img_attn_out);
txt_attn_out = ggml_cont(ctx->ggml_ctx, txt_attn_out);
img_attn_out = to_out_0->forward(ctx, img_attn_out);
txt_attn_out = to_add_out->forward(ctx, txt_attn_out);
return {img_attn_out, txt_attn_out};
}
};
class QwenImageTransformerBlock : public GGMLBlock {
protected:
bool zero_cond_t;
public:
QwenImageTransformerBlock(int64_t dim,
int64_t num_attention_heads,
int64_t attention_head_dim,
float eps = 1e-6,
bool zero_cond_t = false)
: zero_cond_t(zero_cond_t) {
// img_mod.0 is nn.SiLU()
blocks["img_mod.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, 6 * dim, true));
blocks["img_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
blocks["img_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
blocks["img_mlp"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim, 4, FeedForward::Activation::GELU, true));
// txt_mod.0 is nn.SiLU()
blocks["txt_mod.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, 6 * dim, true));
blocks["txt_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
blocks["txt_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim, eps, false));
blocks["txt_mlp"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim, 4, FeedForward::Activation::GELU, true));
blocks["attn"] = std::shared_ptr<GGMLBlock>(new QwenImageAttention(dim,
attention_head_dim,
num_attention_heads,
0, // out_dim
0, // out_context-dim
true, // bias
true, // out_bias
eps));
}
std::vector<ggml_tensor*> get_mod_params_vec(ggml_context* ctx, ggml_tensor* mod_params, ggml_tensor* index = nullptr) {
// index: [N, n_img_token]
// mod_params: [N, hidden_size * 12]
if (index == nullptr) {
return ggml_ext_chunk(ctx, mod_params, 6, 0);
}
mod_params = ggml_reshape_1d(ctx, mod_params, ggml_nelements(mod_params));
auto mod_params_vec = ggml_ext_chunk(ctx, mod_params, 12, 0);
index = ggml_reshape_3d(ctx, index, 1, index->ne[0], index->ne[1]); // [N, n_img_token, 1]
index = ggml_repeat_4d(ctx, index, mod_params_vec[0]->ne[0], index->ne[1], index->ne[2], index->ne[3]); // [N, n_img_token, hidden_size]
std::vector<ggml_tensor*> mod_results;
for (int i = 0; i < 6; i++) {
auto mod_0 = mod_params_vec[i];
auto mod_1 = mod_params_vec[i + 6];
// mod_result = torch.where(index == 0, mod_0, mod_1)
// mod_result = (1 - index)*mod_0 + index*mod_1
mod_0 = ggml_sub(ctx, ggml_repeat(ctx, mod_0, index), ggml_mul(ctx, index, mod_0)); // [N, n_img_token, hidden_size]
mod_1 = ggml_mul(ctx, index, mod_1); // [N, n_img_token, hidden_size]
auto mod_result = ggml_add(ctx, mod_0, mod_1);
mod_results.push_back(mod_result);
}
return mod_results;
}
virtual std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* t_emb,
ggml_tensor* pe,
ggml_tensor* modulate_index = nullptr) {
// img: [N, n_img_token, hidden_size]
// txt: [N, n_txt_token, hidden_size]
// pe: [n_img_token + n_txt_token, d_head/2, 2, 2]
// return: ([N, n_img_token, hidden_size], [N, n_txt_token, hidden_size])
auto img_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["img_mod.1"]);
auto img_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"]);
auto img_norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"]);
auto img_mlp = std::dynamic_pointer_cast<FeedForward>(blocks["img_mlp"]);
auto txt_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["txt_mod.1"]);
auto txt_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["txt_norm1"]);
auto txt_norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["txt_norm2"]);
auto txt_mlp = std::dynamic_pointer_cast<FeedForward>(blocks["txt_mlp"]);
auto attn = std::dynamic_pointer_cast<QwenImageAttention>(blocks["attn"]);
auto img_mod_params = ggml_silu(ctx->ggml_ctx, t_emb);
img_mod_params = img_mod_1->forward(ctx, img_mod_params);
auto img_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, img_mod_params, modulate_index);
if (zero_cond_t) {
t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
}
auto txt_mod_params = ggml_silu(ctx->ggml_ctx, t_emb);
txt_mod_params = txt_mod_1->forward(ctx, txt_mod_params);
auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params);
auto img_normed = img_norm1->forward(ctx, img);
auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr);
auto img_gate1 = img_mod_param_vec[2];
auto txt_normed = txt_norm1->forward(ctx, txt);
auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]);
auto txt_gate1 = txt_mod_param_vec[2];
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn_output, img_gate1));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_gate1));
auto img_normed2 = img_norm2->forward(ctx, img);
auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr);
auto img_gate2 = img_mod_param_vec[5];
auto txt_normed2 = txt_norm2->forward(ctx, txt);
auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]);
auto txt_gate2 = txt_mod_param_vec[5];
auto img_mlp_out = img_mlp->forward(ctx, img_modulated2);
auto txt_mlp_out = txt_mlp->forward(ctx, txt_modulated2);
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp_out, img_gate2));
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp_out, txt_gate2));
return {img, txt};
}
};
struct AdaLayerNormContinuous : public GGMLBlock {
public:
AdaLayerNormContinuous(int64_t embedding_dim,
int64_t conditioning_embedding_dim,
bool elementwise_affine = true,
float eps = 1e-5f,
bool bias = true) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(conditioning_embedding_dim, eps, elementwise_affine, bias));
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(conditioning_embedding_dim, embedding_dim * 2, bias));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* c) {
// x: [N, n_token, hidden_size]
// c: [N, hidden_size]
// return: [N, n_token, patch_size * patch_size * out_channels]
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
auto emb = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
auto scale = mods[0];
auto shift = mods[1];
x = norm->forward(ctx, x);
x = Flux::modulate(ctx->ggml_ctx, x, shift, scale);
return x;
}
};
class QwenImageModel : public GGMLBlock {
protected:
QwenImageConfig config;
public:
QwenImageModel() {}
QwenImageModel(QwenImageConfig config)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
// blocks
for (int i = 0; i < config.num_layers; i++) {
auto block = std::shared_ptr<GGMLBlock>(new QwenImageTransformerBlock(inner_dim,
config.num_attention_heads,
config.attention_head_dim,
1e-6f,
config.zero_cond_t));
blocks["transformer_blocks." + std::to_string(i)] = block;
}
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new AdaLayerNormContinuous(inner_dim, inner_dim, false, 1e-6f));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, config.patch_size * config.patch_size * config.out_channels));
}
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* modulate_index = nullptr) {
auto time_text_embed = std::dynamic_pointer_cast<QwenTimestepProjEmbeddings>(blocks["time_text_embed"]);
auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
auto norm_out = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto t_emb = time_text_embed->forward(ctx, timestep);
if (config.zero_cond_t) {
auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
}
auto img = img_in->forward(ctx, x);
auto txt = txt_norm->forward(ctx, context);
txt = txt_in->forward(ctx, txt);
sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.prelude", "img");
sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt");
// sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb");
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
img = result.first;
txt = result.second;
sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.transformer_blocks." + std::to_string(i), "img");
sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt");
}
if (config.zero_cond_t) {
t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
}
img = norm_out->forward(ctx, img, t_emb);
img = proj_out->forward(ctx, img);
return img;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {},
ggml_tensor* modulate_index = nullptr) {
// Forward pass of DiT.
// x: [N, C, H, W]
// timestep: [N,]
// context: [N, L, D]
// pe: [L, d_head/2, 2, 2]
// return: [N, C, H, W]
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int64_t N = x->ne[3];
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
int64_t img_tokens = img->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
}
}
auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index); // [N, h_len*w_len, ph*pw*C]
if (out->ne[1] > img_tokens) {
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [num_tokens, N, C * patch_size * patch_size]
out = ggml_view_3d(ctx->ggml_ctx, out, out->ne[0], out->ne[1], img_tokens, out->nb[1], out->nb[2], 0);
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [N, h*w, C * patch_size * patch_size]
}
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W]
return out;
}
};
struct QwenImageRunner : public DiffusionModelRunner {
public:
QwenImageConfig config;
QwenImageModel qwen_image;
std::vector<float> pe_vec;
std::vector<float> modulate_index_vec;
SDVersion version;
QwenImageRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_QWEN_IMAGE,
bool zero_cond_t = false)
: DiffusionModelRunner(backend, params_backend, prefix),
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
config.zero_cond_t = config.zero_cond_t || zero_cond_t;
qwen_image = QwenImageModel(config);
qwen_image.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "qwen_image";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
qwen_image.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 = {},
bool increase_ref_index = false) {
ggml_cgraph* gf = new_graph_custom(QWEN_IMAGE_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 = Rope::gen_qwen_image_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,
increase_ref_index,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe, true, "pe");
// pe->data = nullptr;
set_backend_tensor_data(pe, pe_vec.data());
ggml_tensor* modulate_index = nullptr;
if (config.zero_cond_t) {
modulate_index_vec.clear();
int64_t h_len = ((x->ne[1] + (config.patch_size / 2)) / config.patch_size);
int64_t w_len = ((x->ne[0] + (config.patch_size / 2)) / config.patch_size);
int64_t num_img_tokens = h_len * w_len;
modulate_index_vec.insert(modulate_index_vec.end(), num_img_tokens, 0.f);
int64_t num_ref_img_tokens = 0;
for (ggml_tensor* ref : ref_latents) {
int64_t h_len = ((ref->ne[1] + (config.patch_size / 2)) / config.patch_size);
int64_t w_len = ((ref->ne[0] + (config.patch_size / 2)) / config.patch_size);
num_ref_img_tokens += h_len * w_len;
}
if (num_ref_img_tokens > 0) {
modulate_index_vec.insert(modulate_index_vec.end(), num_ref_img_tokens, 1.f);
}
modulate_index = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, modulate_index_vec.size());
set_backend_tensor_data(modulate_index, modulate_index_vec.data());
}
auto runner_ctx = get_context();
ggml_tensor* out = qwen_image.forward(&runner_ctx,
x,
timesteps,
context,
pe,
ref_latents,
modulate_index);
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 = {},
bool increase_ref_index = false) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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,
diffusion_params.increase_ref_index);
}
void test() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
params.mem_buffer = nullptr;
params.no_alloc = false;
ggml_context* ctx = ggml_init(params);
GGML_ASSERT(ctx != nullptr);
{
// auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 16, 16, 16, 1);
// ggml_set_f32(x, 0.01f);
auto x = sd::load_tensor_from_file_as_tensor<float>("./qwen_image_x.bin");
print_sd_tensor(x);
std::vector<float> timesteps_vec(1, 1000.f);
auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
// auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3584, 256, 1);
// ggml_set_f32(context, 0.01f);
auto context = sd::load_tensor_from_file_as_tensor<float>("./qwen_image_context.bin");
print_sd_tensor(context);
sd::Tensor<float> out;
int64_t t0 = ggml_time_ms();
auto out_opt = compute(8,
x,
timesteps,
context,
{},
false);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
out = std::move(out_opt);
print_sd_tensor(out);
LOG_DEBUG("qwen_image test done in %lldms", t1 - t0);
}
}
static void load_from_file_and_test(const std::string& file_path) {
// cuda q8: pass
// cuda q8 fa: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_Q8_0;
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
}
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (ends_with(name, "weight")) {
tensor_storage.expected_type = model_data_type;
}
}
std::shared_ptr<QwenImageRunner> qwen_image = std::make_shared<QwenImageRunner>(backend,
backend,
tensor_storage_map,
"model.diffusion_model",
VERSION_QWEN_IMAGE);
if (!qwen_image->alloc_params_buffer()) {
LOG_ERROR("qwen_image buffer allocation failed");
return;
}
std::map<std::string, ggml_tensor*> tensors;
qwen_image->get_param_tensors(tensors, "model.diffusion_model");
bool success = model_loader.load_tensors(tensors);
if (!success) {
LOG_ERROR("load tensors from model loader failed");
return;
}
LOG_INFO("qwen_image model loaded");
qwen_image->test();
}
};
} // namespace name
#endif // __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
src/model/diffusion/unet.hpp
+847
View File
@@ -0,0 +1,847 @@
#ifndef __SD_MODEL_DIFFUSION_UNET_HPP__
#define __SD_MODEL_DIFFUSION_UNET_HPP__
#include <algorithm>
#include <vector>
#include "model.h"
#include "model/common/block.hpp"
#include "model/diffusion/model.hpp"
/*==================================================== UnetModel =====================================================*/
#define UNET_GRAPH_SIZE 102400
struct UNetConfig {
SDVersion version = VERSION_SD1;
// network hparams
int in_channels = 4;
int out_channels = 4;
int num_res_blocks = 2;
std::vector<int> attention_resolutions = {4, 2, 1};
std::vector<int> channel_mult = {1, 2, 4, 4};
std::vector<int> transformer_depth = {1, 1, 1, 1};
int time_embed_dim = 1280; // model_channels*4
int num_heads = 8;
int num_head_channels = -1; // channels // num_heads
int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
bool use_linear_projection = false;
bool tiny_unet = false;
int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
SDVersion version = VERSION_SD1) {
UNetConfig config;
config.version = version;
if (sd_version_is_sd2(version)) {
config.context_dim = 1024;
config.num_head_channels = 64;
config.num_heads = -1;
config.use_linear_projection = true;
} else if (sd_version_is_sdxl(version)) {
config.context_dim = 2048;
config.attention_resolutions = {4, 2};
config.channel_mult = {1, 2, 4};
config.transformer_depth = {1, 2, 10};
config.num_head_channels = 64;
config.num_heads = -1;
config.use_linear_projection = true;
if (version == VERSION_SDXL_VEGA) {
config.transformer_depth = {1, 1, 2};
}
} else if (version == VERSION_SVD) {
config.in_channels = 8;
config.out_channels = 4;
config.context_dim = 1024;
config.adm_in_channels = 768;
config.num_head_channels = 64;
config.num_heads = -1;
config.use_linear_projection = true;
}
if (sd_version_is_inpaint(version)) {
config.in_channels = 9;
} else if (sd_version_is_unet_edit(version)) {
config.in_channels = 8;
}
if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
config.num_res_blocks = 1;
config.channel_mult = {1, 2, 4};
config.tiny_unet = true;
if (version == VERSION_SDXS_512_DS) {
config.attention_resolutions = {4, 2}; // here just like SDXL
}
}
auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
std::string name = prefix.empty() ? suffix : prefix + "." + suffix;
auto it = tensor_storage_map.find(name);
if (it == tensor_storage_map.end()) {
return nullptr;
}
return &it->second;
};
if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
if (input->n_dims == 4) {
config.in_channels = static_cast<int>(input->ne[2]);
config.model_channels = static_cast<int>(input->ne[3]);
config.time_embed_dim = config.model_channels * 4;
}
}
if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) {
if (time_embed->n_dims == 2) {
config.model_channels = static_cast<int>(time_embed->ne[0]);
config.time_embed_dim = static_cast<int>(time_embed->ne[1]);
}
}
if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) {
if (label_emb->n_dims == 2) {
config.adm_in_channels = static_cast<int>(label_emb->ne[0]);
config.time_embed_dim = static_cast<int>(label_emb->ne[1]);
}
}
if (const TensorStorage* out = find_weight("out.2.weight")) {
if (out->n_dims == 4) {
config.out_channels = static_cast<int>(out->ne[3]);
}
}
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) {
config.context_dim = static_cast<int>(tensor_storage.ne[0]);
break;
}
}
LOG_DEBUG("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s",
config.in_channels,
config.out_channels,
config.model_channels,
config.time_embed_dim,
config.context_dim,
config.adm_in_channels,
config.num_res_blocks,
config.tiny_unet ? "true" : "false");
return config;
}
};
class SpatialVideoTransformer : public SpatialTransformer {
protected:
int64_t time_depth;
int max_time_embed_period;
public:
SpatialVideoTransformer(int64_t in_channels,
int64_t n_head,
int64_t d_head,
int64_t depth,
int64_t context_dim,
bool use_linear,
int64_t time_depth = 1,
int max_time_embed_period = 10000)
: SpatialTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear),
max_time_embed_period(max_time_embed_period) {
// We will convert unet transformer linear to conv2d 1x1 when loading the weights, so use_linear is always False
// use_spatial_context is always True
// merge_strategy is always learned_with_images
// merge_factor is loaded from weights
// time_context_dim is always None
// ff_in is always True
// disable_self_attn is always False
// disable_temporal_crossattention is always False
int64_t inner_dim = n_head * d_head;
GGML_ASSERT(depth == time_depth);
GGML_ASSERT(in_channels == inner_dim);
int64_t time_mix_d_head = d_head;
int64_t n_time_mix_heads = n_head;
int64_t time_mix_inner_dim = time_mix_d_head * n_time_mix_heads; // equal to inner_dim
int64_t time_context_dim = context_dim;
for (int i = 0; i < time_depth; i++) {
std::string name = "time_stack." + std::to_string(i);
blocks[name] = std::shared_ptr<GGMLBlock>(new BasicTransformerBlock(inner_dim,
n_time_mix_heads,
time_mix_d_head,
time_context_dim,
true));
}
int64_t time_embed_dim = in_channels * 4;
blocks["time_pos_embed.0"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, time_embed_dim));
// time_pos_embed.1 is nn.SiLU()
blocks["time_pos_embed.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, in_channels));
blocks["time_mixer"] = std::shared_ptr<GGMLBlock>(new AlphaBlender());
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* context,
int timesteps) {
// x: [N, in_channels, h, w] aka [b*t, in_channels, h, w], t == timesteps
// context: [N, max_position(aka n_context), hidden_size(aka context_dim)] aka [b*t, n_context, context_dim], t == timesteps
// t_emb: [N, in_channels] aka [b*t, in_channels]
// timesteps is num_frames
// time_context is always None
// image_only_indicator is always tensor([0.])
// transformer_options is not used
// GGML_ASSERT(ggml_n_dims(context) == 3);
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
auto proj_in = std::dynamic_pointer_cast<Conv2d>(blocks["proj_in"]);
auto proj_out = std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"]);
auto time_pos_embed_0 = std::dynamic_pointer_cast<Linear>(blocks["time_pos_embed.0"]);
auto time_pos_embed_2 = std::dynamic_pointer_cast<Linear>(blocks["time_pos_embed.2"]);
auto time_mixer = std::dynamic_pointer_cast<AlphaBlender>(blocks["time_mixer"]);
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;
GGML_ASSERT(n == timesteps); // We compute cond and uncond separately, so batch_size==1
auto time_context = context; // [b*t, n_context, context_dim]
auto spatial_context = context;
// time_context_first_timestep = time_context[::timesteps]
auto time_context_first_timestep = ggml_view_3d(ctx->ggml_ctx,
time_context,
time_context->ne[0],
time_context->ne[1],
1,
time_context->nb[1],
time_context->nb[2],
0); // [b, n_context, context_dim]
time_context = ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32,
time_context_first_timestep->ne[0],
time_context_first_timestep->ne[1],
time_context_first_timestep->ne[2] * h * w);
time_context = ggml_repeat(ctx->ggml_ctx, time_context_first_timestep, time_context); // [b*h*w, n_context, context_dim]
x = norm->forward(ctx, x);
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]
auto num_frames = ggml_arange(ctx->ggml_ctx, 0.f, static_cast<float>(timesteps), 1.f);
// since b is 1, no need to do repeat
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, num_frames, static_cast<int>(in_channels), max_time_embed_period); // [N, in_channels]
auto emb = time_pos_embed_0->forward(ctx, t_emb);
emb = ggml_silu_inplace(ctx->ggml_ctx, emb);
emb = time_pos_embed_2->forward(ctx, emb); // [N, in_channels]
emb = ggml_reshape_3d(ctx->ggml_ctx, emb, emb->ne[0], 1, emb->ne[1]); // [N, 1, in_channels]
for (int i = 0; i < depth; i++) {
std::string transformer_name = "transformer_blocks." + std::to_string(i);
std::string time_stack_name = "time_stack." + std::to_string(i);
auto block = std::dynamic_pointer_cast<BasicTransformerBlock>(blocks[transformer_name]);
auto mix_block = std::dynamic_pointer_cast<BasicTransformerBlock>(blocks[time_stack_name]);
x = block->forward(ctx, x, spatial_context); // [N, h * w, inner_dim]
// in_channels == inner_dim
auto x_mix = x;
x_mix = ggml_add(ctx->ggml_ctx, x_mix, emb); // [N, h * w, inner_dim]
int64_t N = x_mix->ne[2];
int64_t T = timesteps;
int64_t B = N / T;
int64_t S = x_mix->ne[1];
int64_t C = x_mix->ne[0];
x_mix = ggml_reshape_4d(ctx->ggml_ctx, x_mix, C, S, T, B); // (b t) s c -> b t s c
x_mix = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x_mix, 0, 2, 1, 3)); // b t s c -> b s t c
x_mix = ggml_reshape_3d(ctx->ggml_ctx, x_mix, C, T, S * B); // b s t c -> (b s) t c
x_mix = mix_block->forward(ctx, x_mix, time_context); // [B * h * w, T, inner_dim]
x_mix = ggml_reshape_4d(ctx->ggml_ctx, x_mix, C, T, S, B); // (b s) t c -> b s t c
x_mix = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x_mix, 0, 2, 1, 3)); // b s t c -> b t s c
x_mix = ggml_reshape_3d(ctx->ggml_ctx, x_mix, C, S, T * B); // b t s c -> (b t) s c
x = time_mixer->forward(ctx, x, x_mix); // [N, h * w, inner_dim]
}
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;
}
};
// ldm.modules.diffusionmodules.openaimodel.UNetModel
class UnetModelBlock : public GGMLBlock {
public:
UNetConfig config;
explicit UnetModelBlock(UNetConfig config = {})
: config(config) {
const SDVersion version = this->config.version;
const int in_channels = this->config.in_channels;
const int out_channels = this->config.out_channels;
const int num_res_blocks = this->config.num_res_blocks;
const auto& attention_resolutions = this->config.attention_resolutions;
const auto& channel_mult = this->config.channel_mult;
const auto& transformer_depth = this->config.transformer_depth;
const int time_embed_dim = this->config.time_embed_dim;
const int num_heads = this->config.num_heads;
const int num_head_channels = this->config.num_head_channels;
const int context_dim = this->config.context_dim;
const bool use_linear_projection = this->config.use_linear_projection;
const bool tiny_unet = this->config.tiny_unet;
const int model_channels = this->config.model_channels;
const int adm_in_channels = this->config.adm_in_channels;
// dims is always 2
// use_temporal_attention is always True for SVD
blocks["time_embed.0"] = std::shared_ptr<GGMLBlock>(new Linear(model_channels, time_embed_dim));
// time_embed_1 is nn.SiLU()
blocks["time_embed.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
if (sd_version_is_sdxl(version) || version == VERSION_SVD) {
blocks["label_emb.0.0"] = std::shared_ptr<GGMLBlock>(new Linear(adm_in_channels, time_embed_dim));
// label_emb_1 is nn.SiLU()
blocks["label_emb.0.2"] = std::shared_ptr<GGMLBlock>(new Linear(time_embed_dim, time_embed_dim));
}
// input_blocks
blocks["input_blocks.0.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, model_channels, {3, 3}, {1, 1}, {1, 1}));
std::vector<int> input_block_chans;
input_block_chans.push_back(model_channels);
int ch = model_channels;
int input_block_idx = 0;
int ds = 1;
auto get_resblock = [&](int64_t channels, int64_t emb_channels, int64_t out_channels) -> ResBlock* {
if (version == VERSION_SVD) {
return new VideoResBlock(channels, emb_channels, out_channels);
} else {
return new ResBlock(channels, emb_channels, out_channels);
}
};
auto get_attention_layer = [&](int64_t in_channels,
int64_t n_head,
int64_t d_head,
int64_t depth,
int64_t context_dim) -> SpatialTransformer* {
if (version == VERSION_SVD) {
return new SpatialVideoTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear_projection);
} else {
if (version == VERSION_SDXS_09 && n_head == 5) {
n_head = 1; // to carry a special case of sdxs_09 into CrossAttentionLayer,
d_head = 320; // works as long the product remains equal (5*64 == 1*320)
}
return new SpatialTransformer(in_channels, n_head, d_head, depth, context_dim, use_linear_projection);
}
};
size_t len_mults = channel_mult.size();
for (int i = 0; i < len_mults; i++) {
int mult = channel_mult[i];
for (int j = 0; j < num_res_blocks; j++) {
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
blocks[name] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, mult * model_channels));
ch = mult * model_channels;
if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
int n_head = num_heads;
int d_head = ch / num_heads;
if (num_head_channels != -1) {
d_head = num_head_channels;
n_head = ch / d_head;
}
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
int td = transformer_depth[i];
if (version == VERSION_SDXL_SSD1B) {
if (i == 2) {
td = 4;
}
}
blocks[name] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
n_head,
d_head,
td,
context_dim));
}
input_block_chans.push_back(ch);
if (tiny_unet) {
input_block_idx++;
}
}
if (i != len_mults - 1) {
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
blocks[name] = std::shared_ptr<GGMLBlock>(new DownSampleBlock(ch, ch));
input_block_chans.push_back(ch);
ds *= 2;
}
}
// middle blocks
int n_head = num_heads;
int d_head = ch / num_heads;
if (num_head_channels != -1) {
d_head = num_head_channels;
n_head = ch / d_head;
}
if (!tiny_unet) {
blocks["middle_block.0"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
if (version != VERSION_SDXL_SSD1B && version != VERSION_SDXL_VEGA) {
blocks["middle_block.1"] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch,
n_head,
d_head,
transformer_depth[transformer_depth.size() - 1],
context_dim));
blocks["middle_block.2"] = std::shared_ptr<GGMLBlock>(get_resblock(ch, time_embed_dim, ch));
}
}
// output_blocks
int output_block_idx = 0;
for (int i = (int)len_mults - 1; i >= 0; i--) {
int mult = channel_mult[i];
for (int j = 0; j < num_res_blocks + 1; j++) {
int ich = input_block_chans.back();
input_block_chans.pop_back();
std::string name = "output_blocks." + std::to_string(output_block_idx) + ".0";
blocks[name] = std::shared_ptr<GGMLBlock>(get_resblock(ch + ich, time_embed_dim, mult * model_channels));
ch = mult * model_channels;
int up_sample_idx = 1;
if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
int n_head = num_heads;
int d_head = ch / num_heads;
if (num_head_channels != -1) {
d_head = num_head_channels;
n_head = ch / d_head;
}
std::string name = "output_blocks." + std::to_string(output_block_idx) + ".1";
int td = transformer_depth[i];
if (version == VERSION_SDXL_SSD1B) {
if (i == 2 && (j == 0 || j == 1)) {
td = 4;
}
if (i == 1 && (j == 1 || j == 2)) {
td = 1;
}
}
blocks[name] = std::shared_ptr<GGMLBlock>(get_attention_layer(ch, n_head, d_head, td, context_dim));
up_sample_idx++;
}
if (i > 0 && j == num_res_blocks) {
if (tiny_unet) {
output_block_idx++;
if (output_block_idx == 2) {
up_sample_idx = 1;
}
}
std::string name = "output_blocks." + std::to_string(output_block_idx) + "." + std::to_string(up_sample_idx);
blocks[name] = std::shared_ptr<GGMLBlock>(new UpSampleBlock(ch, ch));
ds /= 2;
}
output_block_idx += 1;
}
}
// out
blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels
// out_1 is nn.SiLU()
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
}
ggml_tensor* resblock_forward(std::string name,
GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* emb,
int num_video_frames) {
if (config.version == VERSION_SVD) {
auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
return block->forward(ctx, x, emb, num_video_frames);
} else {
auto block = std::dynamic_pointer_cast<ResBlock>(blocks[name]);
return block->forward(ctx, x, emb);
}
}
ggml_tensor* attention_layer_forward(std::string name,
GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* context,
int timesteps) {
if (config.version == VERSION_SVD) {
auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
return block->forward(ctx, x, context, timesteps);
} else {
auto block = std::dynamic_pointer_cast<SpatialTransformer>(blocks[name]);
return block->forward(ctx, x, context);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timesteps,
ggml_tensor* context,
ggml_tensor* c_concat = nullptr,
ggml_tensor* y = nullptr,
int num_video_frames = -1,
std::vector<ggml_tensor*> controls = {},
float control_strength = 0.f) {
// x: [N, in_channels, h, w] or [N, in_channels/2, h, w]
// timesteps: [N,]
// context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768]
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
// y: [N, adm_in_channels] or [1, adm_in_channels]
// return: [N, out_channels, h, w]
const SDVersion version = config.version;
const int model_channels = config.model_channels;
const int num_res_blocks = config.num_res_blocks;
const auto& attention_resolutions = config.attention_resolutions;
const auto& channel_mult = config.channel_mult;
const bool tiny_unet = config.tiny_unet;
if (context != nullptr) {
if (context->ne[2] != x->ne[3]) {
context = ggml_repeat(ctx->ggml_ctx, context, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, context->ne[0], context->ne[1], x->ne[3]));
}
}
if (c_concat != nullptr) {
if (c_concat->ne[3] != x->ne[3]) {
c_concat = ggml_repeat(ctx->ggml_ctx, c_concat, x);
}
x = ggml_concat(ctx->ggml_ctx, x, c_concat, 2);
}
if (y != nullptr) {
if (y->ne[1] != x->ne[3]) {
y = ggml_repeat(ctx->ggml_ctx, y, ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, y->ne[0], x->ne[3]));
}
}
auto time_embed_0 = std::dynamic_pointer_cast<Linear>(blocks["time_embed.0"]);
auto time_embed_2 = std::dynamic_pointer_cast<Linear>(blocks["time_embed.2"]);
auto input_blocks_0_0 = std::dynamic_pointer_cast<Conv2d>(blocks["input_blocks.0.0"]);
auto out_0 = std::dynamic_pointer_cast<GroupNorm32>(blocks["out.0"]);
auto out_2 = std::dynamic_pointer_cast<Conv2d>(blocks["out.2"]);
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels); // [N, model_channels]
auto emb = time_embed_0->forward(ctx, t_emb);
emb = ggml_silu_inplace(ctx->ggml_ctx, emb);
emb = time_embed_2->forward(ctx, emb); // [N, time_embed_dim]
// SDXL/SVD
if (y != nullptr) {
auto label_embed_0 = std::dynamic_pointer_cast<Linear>(blocks["label_emb.0.0"]);
auto label_embed_2 = std::dynamic_pointer_cast<Linear>(blocks["label_emb.0.2"]);
auto label_emb = label_embed_0->forward(ctx, y);
label_emb = ggml_silu_inplace(ctx->ggml_ctx, label_emb);
label_emb = label_embed_2->forward(ctx, label_emb); // [N, time_embed_dim]
emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim]
}
// sd::ggml_graph_cut::mark_graph_cut(emb, "unet.prelude", "emb");
// input_blocks
std::vector<ggml_tensor*> hs;
// input block 0
auto h = input_blocks_0_0->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks.0", "h");
ggml_set_name(h, "bench-start");
hs.push_back(h);
// input block 1-11
size_t len_mults = channel_mult.size();
int input_block_idx = 0;
int ds = 1;
for (int i = 0; i < len_mults; i++) {
int mult = channel_mult[i];
for (int j = 0; j < num_res_blocks; j++) {
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
h = resblock_forward(name, ctx, h, emb, num_video_frames); // [N, mult*model_channels, h, w]
if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w]
}
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
hs.push_back(h);
}
if (tiny_unet) {
input_block_idx++;
}
if (i != len_mults - 1) {
ds *= 2;
input_block_idx += 1;
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".0";
auto block = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
h = block->forward(ctx, h); // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
// sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
hs.push_back(h);
}
}
// [N, 4*model_channels, h/8, w/8]
// middle_block
if (!tiny_unet) {
h = resblock_forward("middle_block.0", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
if (version != VERSION_SDXL_SSD1B && version != VERSION_SDXL_VEGA) {
h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8]
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
}
}
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
if (controls.size() > 0) {
auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true);
h = ggml_add(ctx->ggml_ctx, h, cs); // middle control
}
int control_offset = static_cast<int>(controls.size() - 2);
// output_blocks
int output_block_idx = 0;
for (int i = (int)len_mults - 1; i >= 0; i--) {
for (int j = 0; j < num_res_blocks + 1; j++) {
auto h_skip = hs.back();
hs.pop_back();
if (controls.size() > 0) {
auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, true);
h_skip = ggml_add(ctx->ggml_ctx, h_skip, cs); // control net condition
control_offset--;
}
h = ggml_concat(ctx->ggml_ctx, h, h_skip, 2);
std::string name = "output_blocks." + std::to_string(output_block_idx) + ".0";
h = resblock_forward(name, ctx, h, emb, num_video_frames);
int up_sample_idx = 1;
if (std::find(attention_resolutions.begin(), attention_resolutions.end(), ds) != attention_resolutions.end()) {
std::string name = "output_blocks." + std::to_string(output_block_idx) + ".1";
h = attention_layer_forward(name, ctx, h, context, num_video_frames);
up_sample_idx++;
}
if (i > 0 && j == num_res_blocks) {
if (tiny_unet) {
output_block_idx++;
if (output_block_idx == 2) {
up_sample_idx = 1;
}
}
std::string name = "output_blocks." + std::to_string(output_block_idx) + "." + std::to_string(up_sample_idx);
auto block = std::dynamic_pointer_cast<UpSampleBlock>(blocks[name]);
h = block->forward(ctx, h);
ds /= 2;
}
output_block_idx += 1;
sd::ggml_graph_cut::mark_graph_cut(h, "unet.output_blocks." + std::to_string(output_block_idx - 1), "h");
}
}
// out
h = out_0->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h);
h = out_2->forward(ctx, h);
ggml_set_name(h, "bench-end");
return h; // [N, out_channels, h, w]
}
};
struct UNetModelRunner : public DiffusionModelRunner {
UNetConfig config;
UnetModelBlock unet;
UNetModelRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
SDVersion version = VERSION_SD1)
: DiffusionModelRunner(backend, params_backend, prefix),
config(UNetConfig::detect_from_weights(tensor_storage_map, prefix, version)),
unet(config) {
unet.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "unet";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
unet.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 sd::Tensor<float>& c_concat_tensor = {},
const sd::Tensor<float>& y_tensor = {},
int num_video_frames = -1,
const std::vector<sd::Tensor<float>>& controls_tensor = {},
float control_strength = 0.f) {
ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
ggml_tensor* y = make_optional_input(y_tensor);
std::vector<ggml_tensor*> controls;
controls.reserve(controls_tensor.size());
for (const auto& control_tensor : controls_tensor) {
controls.push_back(make_input(control_tensor));
}
if (num_video_frames == -1) {
num_video_frames = static_cast<int>(x->ne[3]);
}
auto runner_ctx = get_context();
ggml_tensor* out = unet.forward(&runner_ctx,
x,
timesteps,
context,
c_concat,
y,
num_video_frames,
controls,
control_strength);
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 sd::Tensor<float>& c_concat = {},
const sd::Tensor<float>& y = {},
int num_video_frames = -1,
const std::vector<sd::Tensor<float>>& controls = {},
float control_strength = 0.f) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
// y: [N, adm_in_channels] or [1, adm_in_channels]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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);
const auto* extra = diffusion_extra_as<UNetDiffusionExtra>(diffusion_params);
static const std::vector<sd::Tensor<float>> empty_controls;
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
tensor_or_empty(diffusion_params.c_concat),
tensor_or_empty(diffusion_params.y),
extra->num_video_frames,
extra->controls ? *extra->controls : empty_controls,
extra->control_strength);
}
void test() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
params.mem_buffer = nullptr;
params.no_alloc = false;
ggml_context* ctx = ggml_init(params);
GGML_ASSERT(ctx != nullptr);
{
// CPU, num_video_frames = 1, x{num_video_frames, 8, 8, 8}: Pass
// CUDA, num_video_frames = 1, x{num_video_frames, 8, 8, 8}: Pass
// CPU, num_video_frames = 3, x{num_video_frames, 8, 8, 8}: Wrong result
// CUDA, num_video_frames = 3, x{num_video_frames, 8, 8, 8}: nan
int num_video_frames = 3;
sd::Tensor<float> x({8, 8, 8, num_video_frames});
std::vector<float> timesteps_vec(num_video_frames, 999.f);
auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
x.fill_(0.5f);
// print_ggml_tensor(x);
sd::Tensor<float> context({1024, 1, num_video_frames});
context.fill_(0.5f);
// print_ggml_tensor(context);
sd::Tensor<float> y({768, num_video_frames});
y.fill_(0.5f);
// print_ggml_tensor(y);
sd::Tensor<float> out;
int64_t t0 = ggml_time_ms();
auto out_opt = compute(8,
x,
timesteps,
context,
{},
y,
num_video_frames,
{},
0.f);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
out = std::move(out_opt);
print_sd_tensor(out);
LOG_DEBUG("unet test done in %lldms", t1 - t0);
}
}
};
#endif // __SD_MODEL_DIFFUSION_UNET_HPP__
src/model/diffusion/wan.hpp
+1061
View File
File diff suppressed because it is too large. Load diff
src/model/diffusion/z_image.hpp
+742
View File
@@ -0,0 +1,742 @@
#ifndef __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
#include <algorithm>
#include "core/ggml_extend.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/mmdit.hpp"
#include "model/diffusion/model.hpp"
// Ref: https://github.com/Alpha-VLLM/Lumina-Image-2.0/blob/main/model/model.py
// Ref: https://github.com/huggingface/diffusers/pull/12703
#ifndef MIN
#define MIN(a, b) ((a) < (b) ? (a) : (b))
#endif
namespace ZImage {
constexpr int Z_IMAGE_GRAPH_SIZE = 20480;
constexpr int ADALN_EMBED_DIM = 256;
constexpr int SEQ_MULTI_OF = 32;
struct ZImageConfig {
int patch_size = 2;
int64_t hidden_size = 3840;
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t num_layers = 30;
int64_t num_refiner_layers = 2;
int64_t head_dim = 128;
int64_t num_heads = 30;
int64_t num_kv_heads = 30;
int64_t multiple_of = 256;
float ffn_dim_multiplier = 8.0f / 3.0f;
float norm_eps = 1e-5f;
bool qk_norm = true;
int64_t cap_feat_dim = 2560;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int64_t axes_dim_sum = 128;
static ZImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
ZImageConfig config;
int64_t detected_layers = 0;
int64_t detected_refiner_layers = 0;
int64_t detected_context_refiner = 0;
int64_t detected_head_dim = 0;
int64_t detected_qkv_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, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.cap_feat_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) {
detected_head_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.attention.qkv.weight") && tensor_storage.n_dims == 2) {
detected_qkv_dim = tensor_storage.ne[1];
} else if (ends_with(name, "final_layer.linear.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;
}
size_t pos = name.find("layers.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
detected_layers = std::max<int64_t>(detected_layers, block_index + 1);
}
}
pos = name.find("noise_refiner.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
detected_refiner_layers = std::max<int64_t>(detected_refiner_layers, block_index + 1);
}
}
pos = name.find("context_refiner.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
detected_context_refiner = std::max<int64_t>(detected_context_refiner, block_index + 1);
}
}
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
}
if (detected_refiner_layers > 0 || detected_context_refiner > 0) {
config.num_refiner_layers = std::max(detected_refiner_layers, detected_context_refiner);
}
if (detected_head_dim > 0) {
config.head_dim = detected_head_dim;
config.num_heads = config.hidden_size / config.head_dim;
if (detected_qkv_dim > 0) {
int64_t qkv_heads = detected_qkv_dim / config.head_dim;
config.num_kv_heads = std::max<int64_t>(1, (qkv_heads - config.num_heads) / 2);
}
}
LOG_DEBUG("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
config.num_layers,
config.num_refiner_layers,
config.hidden_size,
config.num_heads,
config.num_kv_heads,
config.in_channels,
config.out_channels);
return config;
}
};
struct JointAttention : public GGMLBlock {
protected:
int64_t head_dim;
int64_t num_heads;
int64_t num_kv_heads;
bool qk_norm;
public:
JointAttention(int64_t hidden_size, int64_t head_dim, int64_t num_heads, int64_t num_kv_heads, bool qk_norm)
: head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) {
blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
float scale = 1.f;
blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
if (qk_norm) {
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
// x: [N, n_token, hidden_size]
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out"]);
if (sd_backend_is(ctx->backend, "ROCm")) {
out_proj->set_scale(1.f / 16.f);
}
auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
qkv = ggml_reshape_4d(ctx->ggml_ctx, qkv, head_dim, num_heads + num_kv_heads * 2, qkv->ne[1], qkv->ne[2]); // [N, n_token, num_heads + num_kv_heads*2, head_dim]
auto q = ggml_view_4d(ctx->ggml_ctx,
qkv,
qkv->ne[0],
num_heads,
qkv->ne[2],
qkv->ne[3],
qkv->nb[1],
qkv->nb[2],
qkv->nb[3],
0); // [N, n_token, num_heads, head_dim]
auto k = ggml_view_4d(ctx->ggml_ctx,
qkv,
qkv->ne[0],
num_kv_heads,
qkv->ne[2],
qkv->ne[3],
qkv->nb[1],
qkv->nb[2],
qkv->nb[3],
num_heads * qkv->nb[1]); // [N, n_token, num_kv_heads, head_dim]
auto v = ggml_view_4d(ctx->ggml_ctx,
qkv,
qkv->ne[0],
num_kv_heads,
qkv->ne[2],
qkv->ne[3],
qkv->nb[1],
qkv->nb[2],
qkv->nb[3],
(num_heads + num_kv_heads) * qkv->nb[1]); // [N, n_token, num_kv_heads, head_dim]
if (qk_norm) {
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
}
x = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f); // [N, n_token, num_heads * head_dim]
x = out_proj->forward(ctx, x); // [N, n_token, hidden_size]
return x;
}
};
class FeedForward : public GGMLBlock {
public:
FeedForward(int64_t dim,
int64_t hidden_dim,
int64_t multiple_of,
float ffn_dim_multiplier = 0.f) {
if (ffn_dim_multiplier > 0.f) {
hidden_dim = static_cast<int64_t>(ffn_dim_multiplier * hidden_dim);
}
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) / multiple_of);
blocks["w1"] = std::make_shared<Linear>(dim, hidden_dim, false);
bool force_prec_f32 = false;
float scale = 1.f / 128.f;
// The purpose of the scale here is to prevent NaN issues in certain situations.
// For example, when using CUDA but the weights are k-quants.
blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false, false, force_prec_f32, scale);
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
if (sd_backend_is(ctx->backend, "Vulkan")) {
w2->set_force_prec_f32(true);
}
auto x1 = w1->forward(ctx, x);
auto x3 = w3->forward(ctx, x);
x = ggml_swiglu_split(ctx->ggml_ctx, x1, x3);
x = w2->forward(ctx, x);
return x;
}
};
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C]
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
return x;
}
struct JointTransformerBlock : public GGMLBlock {
protected:
bool modulation;
public:
JointTransformerBlock(int layer_id,
int64_t hidden_size,
int64_t head_dim,
int64_t num_heads,
int64_t num_kv_heads,
int64_t multiple_of,
float ffn_dim_multiplier,
float norm_eps,
bool qk_norm,
bool modulation = true)
: modulation(modulation) {
blocks["attention"] = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm);
blocks["feed_forward"] = std::make_shared<FeedForward>(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
if (modulation) {
blocks["adaLN_modulation.0"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* mask = nullptr,
ggml_tensor* adaln_input = nullptr) {
auto attention = std::dynamic_pointer_cast<JointAttention>(blocks["attention"]);
auto feed_forward = std::dynamic_pointer_cast<FeedForward>(blocks["feed_forward"]);
auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
auto ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
auto attention_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm2"]);
auto ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
if (modulation) {
GGML_ASSERT(adaln_input != nullptr);
auto adaLN_modulation_0 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
auto m = adaLN_modulation_0->forward(ctx, adaln_input); // [N, 4 * hidden_size]
auto mods = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0);
auto scale_msa = mods[0];
auto gate_msa = mods[1];
auto scale_mlp = mods[2];
auto gate_mlp = mods[3];
auto residual = x;
x = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, x), scale_msa);
x = attention->forward(ctx, x, pe, mask);
x = attention_norm2->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_msa));
x = ggml_add(ctx->ggml_ctx, x, residual);
residual = x;
x = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, x), scale_mlp);
x = feed_forward->forward(ctx, x);
x = ffn_norm2->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_mlp));
x = ggml_add(ctx->ggml_ctx, x, residual);
} else {
GGML_ASSERT(adaln_input == nullptr);
auto residual = x;
x = attention_norm1->forward(ctx, x);
x = attention->forward(ctx, x, pe, mask);
x = attention_norm2->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, residual);
residual = x;
x = ffn_norm1->forward(ctx, x);
x = feed_forward->forward(ctx, x);
x = ffn_norm2->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, residual);
}
return x;
}
};
struct FinalLayer : public GGMLBlock {
public:
FinalLayer(int64_t hidden_size,
int64_t patch_size,
int64_t out_channels) {
blocks["norm_final"] = std::make_shared<LayerNorm>(hidden_size, 1e-06f, false);
blocks["linear"] = std::make_shared<Linear>(hidden_size, patch_size * patch_size * out_channels, true, true);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), hidden_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* c) {
// x: [N, n_token, hidden_size]
// c: [N, hidden_size]
// return: [N, n_token, patch_size * patch_size * out_channels]
auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
auto scale = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, hidden_size]
x = norm_final->forward(ctx, x);
x = modulate(ctx->ggml_ctx, x, scale);
x = linear->forward(ctx, x);
return x;
}
};
class ZImageModel : public GGMLBlock {
protected:
ZImageConfig config;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
}
public:
ZImageModel() = default;
ZImageModel(ZImageConfig config)
: config(config) {
blocks["x_embedder"] = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, 256);
blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps);
blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::make_shared<JointTransformerBlock>(i,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
true);
blocks["noise_refiner." + std::to_string(i)] = block;
}
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::make_shared<JointTransformerBlock>(i,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
false);
blocks["context_refiner." + std::to_string(i)] = block;
}
for (int i = 0; i < config.num_layers; i++) {
auto block = std::make_shared<JointTransformerBlock>(i,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
true);
blocks["layers." + std::to_string(i)] = block;
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
}
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
auto norm_final = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_final"]);
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
auto txt_pad_token = params["cap_pad_token"];
auto img_pad_token = params["x_pad_token"];
int64_t N = x->ne[2];
int64_t n_img_token = x->ne[1];
int64_t n_txt_token = context->ne[1];
auto t_emb = t_embedder->forward(ctx, timestep);
auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size]
auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.prelude", "txt");
sd::ggml_graph_cut::mark_graph_cut(img, "z_image.prelude", "img");
sd::ggml_graph_cut::mark_graph_cut(t_emb, "z_image.prelude", "t_emb");
int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), SEQ_MULTI_OF);
if (n_txt_pad_token > 0) {
auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1);
txt = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1); // [N, n_txt_token + n_txt_pad_token, hidden_size]
}
int64_t n_img_pad_token = Rope::bound_mod(static_cast<int>(n_img_token), SEQ_MULTI_OF);
if (n_img_pad_token > 0) {
auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1);
img = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1); // [N, n_img_token + n_img_pad_token, hidden_size]
}
GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt");
}
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img");
}
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img");
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img");
}
txt_img = final_layer->forward(ctx, txt_img, t_emb); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]
img = ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); // [N, n_img_token, ph*pw*C]
return img;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {}) {
// Forward pass of DiT.
// x: [N, C, H, W]
// timestep: [N,]
// context: [N, L, D]
// pe: [L, d_head/2, 2, 2]
// return: [N, C, H, W]
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int64_t N = x->ne[3];
int patch_size = config.patch_size;
auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
uint64_t n_img_token = img->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
ref = DiT::pad_and_patchify(ctx, ref, patch_size, patch_size, false);
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
}
}
auto out = forward_core(ctx, img, timestep, context, pe);
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token); // [N, n_img_token, ph*pw*C]
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); // [N, C, H, W]
out = ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
return out;
}
};
struct ZImageRunner : public DiffusionModelRunner {
public:
ZImageConfig config;
ZImageModel z_image;
std::vector<float> pe_vec;
std::vector<float> timestep_vec;
SDVersion version;
ZImageRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_Z_IMAGE)
: DiffusionModelRunner(backend, params_backend, prefix),
config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
z_image = ZImageModel(config);
z_image.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "z_image";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
z_image.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 = {},
bool increase_ref_index = false) {
ggml_cgraph* gf = new_graph_custom(Z_IMAGE_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 = Rope::gen_z_image_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]),
SEQ_MULTI_OF,
ref_latents,
increase_ref_index,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe, true, "pe");
// pe->data = nullptr;
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = z_image.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 = {},
bool increase_ref_index = false) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, 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,
diffusion_params.increase_ref_index);
}
void test() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(1024 * 1024) * 1024; // 1GB
params.mem_buffer = nullptr;
params.no_alloc = false;
ggml_context* ctx = ggml_init(params);
GGML_ASSERT(ctx != nullptr);
{
// auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 16, 16, 16, 1);
// ggml_set_f32(x, 0.01f);
auto x = sd::load_tensor_from_file_as_tensor<float>("./z_image_x.bin");
print_sd_tensor(x);
std::vector<float> timesteps_vec(1, 0.f);
auto timesteps = sd::Tensor<float>::from_vector(timesteps_vec);
// auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2560, 256, 1);
// ggml_set_f32(context, 0.01f);
auto context = sd::load_tensor_from_file_as_tensor<float>("./z_image_context.bin");
print_sd_tensor(context);
sd::Tensor<float> out;
int64_t t0 = ggml_time_ms();
auto out_opt = compute(8,
x,
timesteps,
context,
{},
false);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
out = std::move(out_opt);
print_sd_tensor(out);
LOG_DEBUG("z_image test done in %lldms", t1 - t0);
}
}
static void load_from_file_and_test(const std::string& file_path) {
// cuda q8: pass
// cuda q8 fa: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_Q8_0;
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
}
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
if (model_data_type != GGML_TYPE_COUNT) {
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (ends_with(name, "weight")) {
tensor_storage.expected_type = model_data_type;
}
}
}
std::shared_ptr<ZImageRunner> z_image = std::make_shared<ZImageRunner>(backend,
backend,
tensor_storage_map,
"model.diffusion_model",
VERSION_QWEN_IMAGE);
if (!z_image->alloc_params_buffer()) {
LOG_ERROR("z_image buffer allocation failed");
return;
}
std::map<std::string, ggml_tensor*> tensors;
z_image->get_param_tensors(tensors, "model.diffusion_model");
bool success = model_loader.load_tensors(tensors);
if (!success) {
LOG_ERROR("load tensors from model loader failed");
return;
}
LOG_INFO("z_image model loaded");
z_image->test();
}
};
} // namespace ZImage
#endif // __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
src/model/te/clip.hpp
+582
View File
@@ -0,0 +1,582 @@
#ifndef __SD_MODEL_TE_CLIP_HPP__
#define __SD_MODEL_TE_CLIP_HPP__
#include "core/ggml_extend.hpp"
#include "model.h"
#include "tokenizers/clip_tokenizer.h"
/*================================================ FrozenCLIPEmbedder ================================================*/
// Ref: https://github.com/huggingface/transformers/blob/main/src/transformers/model/clip/modeling_clip.py
struct CLIPMLP : public GGMLBlock {
protected:
bool use_gelu;
public:
CLIPMLP(int64_t d_model, int64_t intermediate_size) {
blocks["fc1"] = std::shared_ptr<GGMLBlock>(new Linear(d_model, intermediate_size));
blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(intermediate_size, d_model));
if (d_model == 1024 || d_model == 1280) { // SD 2.x
use_gelu = true;
} else { // SD 1.x
use_gelu = false;
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, d_model]
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
x = fc1->forward(ctx, x);
if (use_gelu) {
x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
} else {
x = ggml_ext_gelu_quick(ctx->ggml_ctx, x, true);
}
x = fc2->forward(ctx, x);
return x;
}
};
struct CLIPLayer : public GGMLBlock {
protected:
int64_t d_model; // hidden_size/embed_dim
int64_t n_head;
int64_t intermediate_size;
public:
CLIPLayer(int64_t d_model,
int64_t n_head,
int64_t intermediate_size,
bool proj_in = false)
: d_model(d_model),
n_head(n_head),
intermediate_size(intermediate_size) {
blocks["self_attn"] = std::shared_ptr<GGMLBlock>(new MultiheadAttention(d_model, n_head, true, true, proj_in));
blocks["layer_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_model));
blocks["layer_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(d_model));
blocks["mlp"] = std::shared_ptr<GGMLBlock>(new CLIPMLP(d_model, intermediate_size));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* mask = nullptr) {
// x: [N, n_token, d_model]
auto self_attn = std::dynamic_pointer_cast<MultiheadAttention>(blocks["self_attn"]);
auto layer_norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm1"]);
auto layer_norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["layer_norm2"]);
auto mlp = std::dynamic_pointer_cast<CLIPMLP>(blocks["mlp"]);
x = ggml_add(ctx->ggml_ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x), mask));
x = ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, layer_norm2->forward(ctx, x)));
return x;
}
};
struct CLIPEncoder : public GGMLBlock {
protected:
int n_layer;
public:
CLIPEncoder(int n_layer,
int64_t d_model,
int64_t n_head,
int64_t intermediate_size,
bool proj_in = false)
: n_layer(n_layer) {
for (int i = 0; i < n_layer; i++) {
std::string name = "layers." + std::to_string(i);
blocks[name] = std::shared_ptr<GGMLBlock>(new CLIPLayer(d_model, n_head, intermediate_size, proj_in));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* mask = nullptr,
int clip_skip = -1,
const std::string& graph_cut_prefix = "") {
// x: [N, n_token, d_model]
int layer_idx = n_layer - 1;
// LOG_DEBUG("clip_skip %d", clip_skip);
if (clip_skip > 0) {
layer_idx = n_layer - clip_skip;
}
for (int i = 0; i < n_layer; i++) {
// LOG_DEBUG("layer %d", i);
if (i == layer_idx + 1) {
break;
}
std::string name = "layers." + std::to_string(i);
auto layer = std::dynamic_pointer_cast<CLIPLayer>(blocks[name]);
x = layer->forward(ctx, x, mask); // [N, n_token, d_model]
if (!graph_cut_prefix.empty()) {
sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x");
}
// LOG_DEBUG("layer %d", i);
}
return x;
}
};
class CLIPEmbeddings : public GGMLBlock {
protected:
int64_t embed_dim;
int64_t vocab_size;
int64_t num_positions;
bool force_clip_f32;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
enum ggml_type token_wtype = GGML_TYPE_F32;
if (!force_clip_f32) {
token_wtype = get_type(prefix + "token_embedding.weight", tensor_storage_map, GGML_TYPE_F32);
if (!support_get_rows(token_wtype)) {
token_wtype = GGML_TYPE_F32;
}
}
enum ggml_type position_wtype = GGML_TYPE_F32;
params["token_embedding.weight"] = ggml_new_tensor_2d(ctx, token_wtype, embed_dim, vocab_size);
params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions);
}
public:
CLIPEmbeddings(int64_t embed_dim,
int64_t vocab_size = 49408,
int64_t num_positions = 77,
bool force_clip_f32 = false)
: embed_dim(embed_dim),
vocab_size(vocab_size),
num_positions(num_positions),
force_clip_f32(force_clip_f32) {
}
ggml_tensor* get_token_embed_weight() {
return params["token_embedding.weight"];
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* custom_embed_weight) {
// input_ids: [N, n_token]
auto token_embed_weight = params["token_embedding.weight"];
auto position_embed_weight = params["position_embedding.weight"];
GGML_ASSERT(input_ids->ne[0] == position_embed_weight->ne[1]);
input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]);
auto token_embedding = ggml_get_rows(ctx->ggml_ctx, custom_embed_weight != nullptr ? custom_embed_weight : token_embed_weight, input_ids);
token_embedding = ggml_reshape_3d(ctx->ggml_ctx, token_embedding, token_embedding->ne[0], token_embedding->ne[1], token_embedding->ne[3]);
// token_embedding + position_embedding
auto x = ggml_add(ctx->ggml_ctx,
token_embedding,
position_embed_weight); // [N, n_token, embed_dim]
return x;
}
};
class CLIPVisionEmbeddings : public GGMLBlock {
protected:
int64_t embed_dim;
int num_channels;
int patch_size;
int image_size;
int num_patches;
int64_t num_positions;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
enum ggml_type patch_wtype = GGML_TYPE_F16;
enum ggml_type class_wtype = GGML_TYPE_F32;
enum ggml_type position_wtype = GGML_TYPE_F32;
params["patch_embedding.weight"] = ggml_new_tensor_4d(ctx, patch_wtype, patch_size, patch_size, num_channels, embed_dim);
params["class_embedding"] = ggml_new_tensor_1d(ctx, class_wtype, embed_dim);
params["position_embedding.weight"] = ggml_new_tensor_2d(ctx, position_wtype, embed_dim, num_positions);
}
public:
CLIPVisionEmbeddings(int64_t embed_dim,
int num_channels = 3,
int patch_size = 14,
int image_size = 224)
: embed_dim(embed_dim),
num_channels(num_channels),
patch_size(patch_size),
image_size(image_size) {
num_patches = (image_size / patch_size) * (image_size / patch_size);
num_positions = num_patches + 1;
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* pixel_values) {
// pixel_values: [N, num_channels, image_size, image_size]
// return: [N, num_positions, embed_dim]
GGML_ASSERT(pixel_values->ne[0] == image_size && pixel_values->ne[1] == image_size && pixel_values->ne[2] == num_channels);
auto patch_embed_weight = params["patch_embedding.weight"];
auto class_embed_weight = params["class_embedding"];
auto position_embed_weight = params["position_embedding.weight"];
// concat(patch_embedding, class_embedding) + position_embedding
ggml_tensor* patch_embedding;
int64_t N = pixel_values->ne[3];
patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size); // [N, embed_dim, image_size // pacht_size, image_size // pacht_size]
patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N); // [N, embed_dim, num_patches]
patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3)); // [N, num_patches, embed_dim]
patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N); // [N, num_patches, embed_dim, 1]
ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N);
class_embedding = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding); // [N, embed_dim]
class_embedding = ggml_reshape_4d(ctx->ggml_ctx, class_embedding, 1, embed_dim, 1, N); // [N, 1, embed_dim, 1]
ggml_tensor* x = ggml_concat(ctx->ggml_ctx, class_embedding, patch_embedding, 2); // [N, num_positions, embed_dim, 1]
x = ggml_reshape_3d(ctx->ggml_ctx, x, embed_dim, num_positions, N); // [N, num_positions, embed_dim]
x = ggml_add(ctx->ggml_ctx, x, position_embed_weight);
return x; // [N, num_positions, embed_dim]
}
};
// OPENAI_CLIP_VIT_L_14: https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json
// OPEN_CLIP_VIT_H_14: https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/blob/main/config.json
// OPEN_CLIP_VIT_BIGG_14: https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/blob/main/config.json (CLIPTextModelWithProjection)
enum CLIPVersion {
OPENAI_CLIP_VIT_L_14, // SD 1.x and SDXL
OPEN_CLIP_VIT_H_14, // SD 2.x
OPEN_CLIP_VIT_BIGG_14, // SDXL
};
class CLIPTextModel : public GGMLBlock {
protected:
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
if (version == OPEN_CLIP_VIT_BIGG_14) {
enum ggml_type wtype = GGML_TYPE_F32;
params["text_projection"] = ggml_new_tensor_2d(ctx, wtype, projection_dim, hidden_size);
}
}
public:
CLIPVersion version = OPENAI_CLIP_VIT_L_14;
// network hparams
int32_t vocab_size = 49408;
int32_t n_token = 77; // max_position_embeddings
int32_t hidden_size = 768;
int32_t intermediate_size = 3072;
int32_t n_head = 12;
int32_t n_layer = 12; // num_hidden_layers
int32_t projection_dim = 1280; // only for OPEN_CLIP_VIT_BIGG_14
bool with_final_ln = true;
CLIPTextModel(CLIPVersion version = OPENAI_CLIP_VIT_L_14,
bool with_final_ln = true,
bool force_clip_f32 = false,
bool proj_in = false)
: version(version), with_final_ln(with_final_ln) {
if (version == OPEN_CLIP_VIT_H_14) {
hidden_size = 1024;
intermediate_size = 4096;
n_head = 16;
n_layer = 24;
} else if (version == OPEN_CLIP_VIT_BIGG_14) { // CLIPTextModelWithProjection
hidden_size = 1280;
intermediate_size = 5120;
n_head = 20;
n_layer = 32;
}
blocks["embeddings"] = std::shared_ptr<GGMLBlock>(new CLIPEmbeddings(hidden_size, vocab_size, n_token, force_clip_f32));
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new CLIPEncoder(n_layer, hidden_size, n_head, intermediate_size, proj_in));
blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
}
ggml_tensor* get_token_embed_weight() {
auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
return embeddings->get_token_embed_weight();
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* tkn_embeddings,
ggml_tensor* mask = nullptr,
size_t max_token_idx = 0,
bool return_pooled = false,
int clip_skip = -1) {
// input_ids: [N, n_token]
auto embeddings = std::dynamic_pointer_cast<CLIPEmbeddings>(blocks["embeddings"]);
auto encoder = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
auto x = embeddings->forward(ctx, input_ids, tkn_embeddings); // [N, n_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(x, "clip_text.prelude", "x");
x = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip, "clip_text");
if (return_pooled || with_final_ln) {
x = final_layer_norm->forward(ctx, x);
}
if (return_pooled) {
auto text_projection = params["text_projection"];
ggml_tensor* pooled = ggml_view_1d(ctx->ggml_ctx, x, hidden_size, x->nb[1] * max_token_idx);
if (text_projection != nullptr) {
pooled = ggml_ext_linear(ctx->ggml_ctx, pooled, text_projection, nullptr);
} else {
LOG_DEBUG("identity projection");
}
return pooled; // [hidden_size, 1, 1]
}
return x; // [N, n_token, hidden_size]
}
};
class CLIPVisionModel : public GGMLBlock {
public:
// network hparams
int32_t num_channels = 3;
int32_t patch_size = 14;
int32_t image_size = 224;
int32_t num_positions = 257; // (image_size / patch_size)^2 + 1
int32_t hidden_size = 1024;
int32_t intermediate_size = 4096;
int32_t n_head = 16;
int32_t n_layer = 24;
public:
CLIPVisionModel(CLIPVersion version = OPENAI_CLIP_VIT_L_14, bool proj_in = false) {
if (version == OPEN_CLIP_VIT_H_14) {
hidden_size = 1280;
intermediate_size = 5120;
n_head = 16;
n_layer = 32;
} else if (version == OPEN_CLIP_VIT_BIGG_14) {
hidden_size = 1664;
intermediate_size = 8192;
n_head = 16;
n_layer = 48;
}
blocks["embeddings"] = std::shared_ptr<GGMLBlock>(new CLIPVisionEmbeddings(hidden_size, num_channels, patch_size, image_size));
blocks["pre_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new CLIPEncoder(n_layer, hidden_size, n_head, intermediate_size, proj_in));
blocks["post_layernorm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
bool return_pooled = true,
int clip_skip = -1) {
// pixel_values: [N, num_channels, image_size, image_size]
auto embeddings = std::dynamic_pointer_cast<CLIPVisionEmbeddings>(blocks["embeddings"]);
auto pre_layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["pre_layernorm"]);
auto encoder = std::dynamic_pointer_cast<CLIPEncoder>(blocks["encoder"]);
auto post_layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["post_layernorm"]);
auto x = embeddings->forward(ctx, pixel_values); // [N, num_positions, embed_dim]
x = pre_layernorm->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x, "clip_vision.prelude", "x");
x = encoder->forward(ctx, x, nullptr, clip_skip, "clip_vision");
auto last_hidden_state = x;
x = post_layernorm->forward(ctx, x); // [N, n_token, hidden_size]
GGML_ASSERT(x->ne[3] == 1);
if (return_pooled) {
ggml_tensor* pooled = ggml_cont(ctx->ggml_ctx, ggml_view_2d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->nb[2], 0));
return pooled; // [N, hidden_size]
} else {
// return x; // [N, n_token, hidden_size]
return last_hidden_state; // [N, n_token, hidden_size]
}
}
};
class CLIPProjection : public UnaryBlock {
protected:
int64_t in_features;
int64_t out_features;
bool transpose_weight;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
if (transpose_weight) {
params["weight"] = ggml_new_tensor_2d(ctx, wtype, out_features, in_features);
} else {
params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
}
}
public:
CLIPProjection(int64_t in_features,
int64_t out_features,
bool transpose_weight = false)
: in_features(in_features),
out_features(out_features),
transpose_weight(transpose_weight) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
if (transpose_weight) {
w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w));
}
return ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr);
}
};
class CLIPVisionModelProjection : public GGMLBlock {
public:
int32_t hidden_size = 1024;
int32_t projection_dim = 768;
int32_t image_size = 224;
public:
CLIPVisionModelProjection(CLIPVersion version = OPENAI_CLIP_VIT_L_14,
bool transpose_proj_w = false,
bool proj_in = false) {
if (version == OPEN_CLIP_VIT_H_14) {
hidden_size = 1280;
projection_dim = 1024;
} else if (version == OPEN_CLIP_VIT_BIGG_14) {
hidden_size = 1664;
}
blocks["vision_model"] = std::shared_ptr<GGMLBlock>(new CLIPVisionModel(version, proj_in));
blocks["visual_projection"] = std::shared_ptr<GGMLBlock>(new CLIPProjection(hidden_size, projection_dim, transpose_proj_w));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
bool return_pooled = true,
int clip_skip = -1) {
// pixel_values: [N, num_channels, image_size, image_size]
// return: [N, projection_dim] if return_pooled else [N, n_token, hidden_size]
auto vision_model = std::dynamic_pointer_cast<CLIPVisionModel>(blocks["vision_model"]);
auto visual_projection = std::dynamic_pointer_cast<CLIPProjection>(blocks["visual_projection"]);
auto x = vision_model->forward(ctx, pixel_values, return_pooled, clip_skip); // [N, hidden_size] or [N, n_token, hidden_size]
if (return_pooled) {
x = visual_projection->forward(ctx, x); // [N, projection_dim]
}
return x;
}
};
struct CLIPTextModelRunner : public GGMLRunner {
CLIPTextModel model;
std::vector<float> attention_mask_vec;
CLIPTextModelRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
CLIPVersion version = OPENAI_CLIP_VIT_L_14,
bool with_final_ln = true,
bool force_clip_f32 = false)
: GGMLRunner(backend, params_backend) {
bool proj_in = false;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (contains(name, "self_attn.in_proj")) {
proj_in = true;
break;
}
}
model = CLIPTextModel(version, with_final_ln, force_clip_f32, proj_in);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "clip";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
model.get_param_tensors(tensors, prefix);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* embeddings,
ggml_tensor* mask,
size_t max_token_idx = 0,
bool return_pooled = false,
int clip_skip = -1) {
size_t N = input_ids->ne[1];
size_t n_token = input_ids->ne[0];
if (input_ids->ne[0] > model.n_token) {
GGML_ASSERT(input_ids->ne[0] % model.n_token == 0);
input_ids = ggml_reshape_2d(ctx->ggml_ctx, input_ids, model.n_token, input_ids->ne[0] / model.n_token);
}
return model.forward(ctx, input_ids, embeddings, mask, max_token_idx, return_pooled, clip_skip);
}
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
int num_custom_embeddings = 0,
void* custom_embeddings_data = nullptr,
size_t max_token_idx = 0,
bool return_pooled = false,
int clip_skip = -1) {
ggml_cgraph* gf = new_graph_custom(2048);
ggml_tensor* input_ids = make_input(input_ids_tensor);
ggml_tensor* embeddings = nullptr;
if (num_custom_embeddings > 0 && custom_embeddings_data != nullptr) {
auto token_embed_weight = model.get_token_embed_weight();
auto custom_embeddings = ggml_new_tensor_2d(compute_ctx,
token_embed_weight->type,
model.hidden_size,
num_custom_embeddings);
set_backend_tensor_data(custom_embeddings, custom_embeddings_data);
// concatenate custom embeddings
embeddings = ggml_concat(compute_ctx, token_embed_weight, custom_embeddings, 1);
}
int n_tokens = static_cast<int>(input_ids->ne[0]);
attention_mask_vec.resize(n_tokens * n_tokens);
for (int i0 = 0; i0 < n_tokens; i0++) {
for (int i1 = 0; i1 < n_tokens; i1++) {
float value = 0.f;
if (i0 > i1) {
value = -INFINITY;
}
attention_mask_vec[i1 * n_tokens + i0] = value;
}
}
auto attention_mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, n_tokens, n_tokens);
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
auto runner_ctx = get_context();
ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, embeddings, attention_mask, max_token_idx, return_pooled, clip_skip);
ggml_build_forward_expand(gf, hidden_states);
return gf;
}
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
int num_custom_embeddings,
void* custom_embeddings_data,
size_t max_token_idx,
bool return_pooled,
int clip_skip) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
};
auto result = GGMLRunner::compute<float>(get_graph, n_threads, true);
if (return_pooled) {
return take_or_empty(std::move(result));
}
return restore_trailing_singleton_dims(std::move(result), 3);
}
};
#endif // __SD_MODEL_TE_CLIP_HPP__
src/model/te/llm.hpp
+2103
View File
File diff suppressed because it is too large. Load diff
src/model/te/t5.hpp
+613
View File
@@ -0,0 +1,613 @@
#ifndef __SD_MODEL_TE_T5_HPP__
#define __SD_MODEL_TE_T5_HPP__
#include <cfloat>
#include <limits>
#include <map>
#include <memory>
#include <regex>
#include <sstream>
#include <string>
#include <unordered_map>
#include "core/ggml_extend.hpp"
#include "model.h"
#include "tokenizers/t5_unigram_tokenizer.h"
struct T5Config {
int64_t num_layers = 24;
int64_t model_dim = 4096;
int64_t ff_dim = 10240;
int64_t num_heads = 64;
int64_t vocab_size = 32128;
bool relative_attention = true;
static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
bool is_umt5 = false) {
(void)tensor_storage_map;
(void)prefix;
T5Config config;
if (is_umt5) {
config.vocab_size = 256384;
config.relative_attention = false;
}
return config;
}
};
class T5LayerNorm : public UnaryBlock {
protected:
int64_t hidden_size;
float eps;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
enum ggml_type wtype = GGML_TYPE_F32;
params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
}
public:
T5LayerNorm(int64_t hidden_size,
float eps = 1e-06f)
: hidden_size(hidden_size),
eps(eps) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
x = ggml_mul(ctx->ggml_ctx, x, w);
return x;
}
};
struct T5DenseActDense : public UnaryBlock {
public:
T5DenseActDense(int64_t model_dim, int64_t ff_dim) {
blocks["wi"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, ff_dim, false));
blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [N, n_token, model_dim]
auto wi = std::dynamic_pointer_cast<Linear>(blocks["wi"]);
auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
x = wi->forward(ctx, x);
x = ggml_relu_inplace(ctx->ggml_ctx, x);
x = wo->forward(ctx, x);
return x;
}
};
struct T5DenseGatedActDense : public UnaryBlock {
public:
T5DenseGatedActDense(int64_t model_dim, int64_t ff_dim) {
blocks["wi_0"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, ff_dim, false));
blocks["wi_1"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, ff_dim, false));
float scale = 1.f / 32.f;
// The purpose of the scale here is to prevent NaN issues on some backends(CUDA, ...).
blocks["wo"] = std::shared_ptr<GGMLBlock>(new Linear(ff_dim, model_dim, false, false, false, scale));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [N, n_token, model_dim]
auto wi_0 = std::dynamic_pointer_cast<Linear>(blocks["wi_0"]);
auto wi_1 = std::dynamic_pointer_cast<Linear>(blocks["wi_1"]);
auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
auto hidden_gelu = ggml_ext_gelu(ctx->ggml_ctx, wi_0->forward(ctx, x), true);
auto hidden_linear = wi_1->forward(ctx, x);
x = ggml_mul_inplace(ctx->ggml_ctx, hidden_gelu, hidden_linear);
x = wo->forward(ctx, x);
return x;
}
};
struct T5LayerFF : public UnaryBlock {
public:
T5LayerFF(int64_t model_dim, int64_t ff_dim) {
blocks["DenseReluDense"] = std::shared_ptr<GGMLBlock>(new T5DenseGatedActDense(model_dim, ff_dim));
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [N, n_token, model_dim]
auto DenseReluDense = std::dynamic_pointer_cast<T5DenseGatedActDense>(blocks["DenseReluDense"]);
auto layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
auto forwarded_states = layer_norm->forward(ctx, x);
forwarded_states = DenseReluDense->forward(ctx, forwarded_states);
x = ggml_add_inplace(ctx->ggml_ctx, forwarded_states, x);
return x;
}
};
class T5Attention : public GGMLBlock {
protected:
int64_t model_dim;
int64_t inner_dim;
int64_t num_heads;
bool using_relative_attention_bias;
int64_t relative_attention_num_buckets = 32;
int64_t relative_attention_max_distance = 128;
public:
T5Attention(int64_t model_dim,
int64_t inner_dim,
int64_t num_heads,
bool using_relative_attention_bias = false)
: model_dim(model_dim),
inner_dim(inner_dim),
num_heads(num_heads),
using_relative_attention_bias(using_relative_attention_bias) {
blocks["q"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, inner_dim, false));
blocks["k"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, inner_dim, false));
blocks["v"] = std::shared_ptr<GGMLBlock>(new Linear(model_dim, inner_dim, false));
blocks["o"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, model_dim, false));
if (using_relative_attention_bias) {
blocks["relative_attention_bias"] = std::shared_ptr<GGMLBlock>(new Embedding(relative_attention_num_buckets, num_heads));
}
}
ggml_tensor* compute_bias(GGMLRunnerContext* ctx,
ggml_tensor* relative_position_bucket) {
auto relative_attention_bias = std::dynamic_pointer_cast<Embedding>(blocks["relative_attention_bias"]);
auto values = relative_attention_bias->forward(ctx, relative_position_bucket); // shape (query_length, key_length, num_heads)
values = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, values, 2, 0, 1, 3)); // shape (1, num_heads, query_length, key_length)
return values;
}
// x: [N, n_token, model_dim]
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* past_bias = nullptr,
ggml_tensor* mask = nullptr,
ggml_tensor* relative_position_bucket = nullptr) {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
int64_t n_head = num_heads;
int64_t d_head = inner_dim / n_head;
auto q = q_proj->forward(ctx, x);
auto k = k_proj->forward(ctx, x);
auto v = v_proj->forward(ctx, x);
if (using_relative_attention_bias && relative_position_bucket != nullptr) {
past_bias = compute_bias(ctx, relative_position_bucket);
}
if (past_bias != nullptr) {
if (mask != nullptr) {
mask = ggml_repeat(ctx->ggml_ctx, mask, past_bias);
mask = ggml_add(ctx->ggml_ctx, mask, past_bias);
} else {
mask = past_bias;
}
}
k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast<float>(d_head)), true);
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask); // [N, n_token, d_head * n_head]
x = out_proj->forward(ctx, x); // [N, n_token, model_dim]
return {x, past_bias};
}
};
struct T5LayerSelfAttention : public GGMLBlock {
public:
T5LayerSelfAttention(int64_t model_dim,
int64_t inner_dim,
int64_t ff_dim,
int64_t num_heads,
bool using_relative_attention_bias) {
blocks["SelfAttention"] = std::shared_ptr<GGMLBlock>(new T5Attention(model_dim, inner_dim, num_heads, using_relative_attention_bias));
blocks["layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* past_bias = nullptr,
ggml_tensor* mask = nullptr,
ggml_tensor* relative_position_bucket = nullptr) {
// x: [N, n_token, model_dim]
auto SelfAttention = std::dynamic_pointer_cast<T5Attention>(blocks["SelfAttention"]);
auto layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["layer_norm"]);
auto normed_hidden_state = layer_norm->forward(ctx, x);
auto ret = SelfAttention->forward(ctx, normed_hidden_state, past_bias, mask, relative_position_bucket);
auto output = ret.first;
past_bias = ret.second;
x = ggml_add_inplace(ctx->ggml_ctx, output, x);
return {x, past_bias};
}
};
struct T5Block : public GGMLBlock {
public:
T5Block(int64_t model_dim, int64_t inner_dim, int64_t ff_dim, int64_t num_heads, bool using_relative_attention_bias) {
blocks["layer.0"] = std::shared_ptr<GGMLBlock>(new T5LayerSelfAttention(model_dim, inner_dim, ff_dim, num_heads, using_relative_attention_bias));
blocks["layer.1"] = std::shared_ptr<GGMLBlock>(new T5LayerFF(model_dim, ff_dim));
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* past_bias = nullptr,
ggml_tensor* mask = nullptr,
ggml_tensor* relative_position_bucket = nullptr) {
// x: [N, n_token, model_dim]
auto layer_0 = std::dynamic_pointer_cast<T5LayerSelfAttention>(blocks["layer.0"]);
auto layer_1 = std::dynamic_pointer_cast<T5LayerFF>(blocks["layer.1"]);
auto ret = layer_0->forward(ctx, x, past_bias, mask, relative_position_bucket);
x = ret.first;
past_bias = ret.second;
x = layer_1->forward(ctx, x);
return {x, past_bias};
}
};
struct T5Stack : public GGMLBlock {
int64_t num_layers;
public:
T5Stack(int64_t num_layers,
int64_t model_dim,
int64_t inner_dim,
int64_t ff_dim,
int64_t num_heads,
bool relative_attention = true)
: num_layers(num_layers) {
for (int i = 0; i < num_layers; i++) {
blocks["block." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new T5Block(model_dim, inner_dim, ff_dim, num_heads, (!relative_attention || i == 0)));
}
blocks["final_layer_norm"] = std::shared_ptr<GGMLBlock>(new T5LayerNorm(model_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* past_bias = nullptr,
ggml_tensor* attention_mask = nullptr,
ggml_tensor* relative_position_bucket = nullptr,
const std::string& graph_cut_prefix = "") {
// x: [N, n_token, model_dim]
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
auto ret = block->forward(ctx, x, past_bias, attention_mask, relative_position_bucket);
x = ret.first;
past_bias = ret.second;
if (!graph_cut_prefix.empty()) {
sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".block." + std::to_string(i), "x");
}
}
auto final_layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["final_layer_norm"]);
x = final_layer_norm->forward(ctx, x);
return x;
}
};
struct T5 : public GGMLBlock {
T5Config config;
public:
T5() {}
T5(T5Config config)
: config(config) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
config.model_dim,
config.model_dim,
config.ff_dim,
config.num_heads,
config.relative_attention));
blocks["shared"] = std::shared_ptr<GGMLBlock>(new Embedding(config.vocab_size,
config.model_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* past_bias = nullptr,
ggml_tensor* attention_mask = nullptr,
ggml_tensor* relative_position_bucket = nullptr) {
// input_ids: [N, n_token]
auto shared = std::dynamic_pointer_cast<Embedding>(blocks["shared"]);
auto encoder = std::dynamic_pointer_cast<T5Stack>(blocks["encoder"]);
auto x = shared->forward(ctx, input_ids);
sd::ggml_graph_cut::mark_graph_cut(x, "t5.prelude", "x");
x = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket, "t5");
return x;
}
};
struct T5Runner : public GGMLRunner {
T5Config config;
T5 model;
std::vector<int> relative_position_bucket_vec;
T5Runner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
bool is_umt5 = false)
: GGMLRunner(backend, params_backend),
config(T5Config::detect_from_weights(tensor_storage_map, prefix, is_umt5)) {
model = T5(config);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "t5";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
model.get_param_tensors(tensors, prefix);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* relative_position_bucket,
ggml_tensor* attention_mask = nullptr) {
size_t N = input_ids->ne[1];
size_t n_token = input_ids->ne[0];
auto hidden_states = model.forward(ctx, input_ids, nullptr, attention_mask, relative_position_bucket); // [N, n_token, model_dim]
return hidden_states;
}
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
const sd::Tensor<float>& attention_mask_tensor = {}) {
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
ggml_tensor* input_ids = make_input(input_ids_tensor);
ggml_tensor* attention_mask = attention_mask_tensor.empty() ? nullptr : make_input(attention_mask_tensor);
relative_position_bucket_vec = compute_relative_position_bucket(static_cast<int>(input_ids->ne[0]), static_cast<int>(input_ids->ne[0]));
// for (int i = 0; i < relative_position_bucket_vec.size(); i++) {
// if (i % 77 == 0) {
// printf("\n");
// }
// printf("%d ", relative_position_bucket_vec[i]);
// }
auto relative_position_bucket = ggml_new_tensor_2d(compute_ctx,
GGML_TYPE_I32,
input_ids->ne[0],
input_ids->ne[0]);
set_backend_tensor_data(relative_position_bucket, relative_position_bucket_vec.data());
auto runner_ctx = get_context();
ggml_tensor* hidden_states = forward(&runner_ctx, input_ids, relative_position_bucket, attention_mask);
ggml_build_forward_expand(gf, hidden_states);
return gf;
}
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<float>& attention_mask) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids, attention_mask);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true), 3);
}
static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
bool bidirectional = true,
int num_buckets = 32,
int max_distance = 128) {
std::vector<int> relative_buckets(relative_position.size(), 0);
std::vector<int> abs_relative_position = relative_position;
if (bidirectional) {
num_buckets = num_buckets / 2;
for (size_t i = 0; i < relative_position.size(); ++i) {
if (relative_position[i] > 0) {
relative_buckets[i] += num_buckets;
}
abs_relative_position[i] = std::abs(relative_position[i]);
}
} else {
for (size_t i = 0; i < relative_position.size(); ++i) {
abs_relative_position[i] = std::max(-relative_position[i], 0);
}
}
int max_exact = num_buckets / 2;
std::vector<int> relative_position_if_large(relative_position.size(), 0);
for (size_t i = 0; i < relative_position.size(); ++i) {
if (abs_relative_position[i] < max_exact) {
relative_buckets[i] += abs_relative_position[i];
} else {
float log_pos = std::log(static_cast<float>(abs_relative_position[i]) / max_exact);
float log_base = std::log(static_cast<float>(max_distance) / max_exact);
relative_position_if_large[i] = max_exact + static_cast<int>((log_pos / log_base) * (num_buckets - max_exact));
relative_position_if_large[i] = std::min(relative_position_if_large[i], num_buckets - 1);
relative_buckets[i] += relative_position_if_large[i];
}
}
return relative_buckets;
}
std::vector<int> compute_relative_position_bucket(int query_length,
int key_length) {
std::vector<int> context_position(query_length);
std::vector<int> memory_position(key_length);
for (int i = 0; i < query_length; ++i) {
context_position[i] = i;
}
for (int i = 0; i < key_length; ++i) {
memory_position[i] = i;
}
std::vector<std::vector<int>> relative_position(query_length, std::vector<int>(key_length, 0));
for (int i = 0; i < query_length; ++i) {
for (int j = 0; j < key_length; ++j) {
relative_position[i][j] = memory_position[j] - context_position[i];
}
}
std::vector<int> relative_position_bucket;
for (int i = 0; i < query_length; ++i) {
std::vector<int> result = _relative_position_bucket(relative_position[i], true);
relative_position_bucket.insert(relative_position_bucket.end(), result.begin(), result.end());
}
return relative_position_bucket;
}
};
struct T5Embedder {
T5UniGramTokenizer tokenizer;
T5Runner model;
T5Embedder(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
bool is_umt5 = false)
: model(backend, params_backend, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
model.get_param_tensors(tensors, prefix);
}
bool alloc_params_buffer() {
if (!model.alloc_params_buffer()) {
return false;
}
return true;
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
size_t max_length = 0,
bool padding = false) {
auto parsed_attention = parse_prompt_attention(text);
{
std::stringstream ss;
ss << "[";
for (const auto& item : parsed_attention) {
ss << "['" << item.first << "', " << item.second << "], ";
}
ss << "]";
LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str());
}
std::vector<int> tokens;
std::vector<float> weights;
for (const auto& item : parsed_attention) {
const std::string& curr_text = item.first;
float curr_weight = item.second;
std::vector<int> curr_tokens = tokenizer.encode(curr_text);
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
}
std::vector<float> attention_mask;
tokenizer.pad_tokens(tokens, &weights, &attention_mask, padding ? max_length : 0, padding ? max_length : 100000000, padding);
for (auto& mask_value : attention_mask) {
mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
}
// for (int i = 0; i < tokens.size(); i++) {
// std::cout << tokens[i] << ":" << weights[i] << ", ";
// }
// std::cout << std::endl;
return {tokens, weights, attention_mask};
}
void test() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
params.mem_buffer = nullptr;
params.no_alloc = false;
ggml_context* ctx = ggml_init(params);
GGML_ASSERT(ctx != nullptr);
{
std::string text("a lovely cat");
auto tokens_and_weights = tokenize(text, 512, true);
std::vector<int>& tokens = std::get<0>(tokens_and_weights);
std::vector<float>& weights = std::get<1>(tokens_and_weights);
std::vector<float>& masks = std::get<2>(tokens_and_weights);
for (auto token : tokens) {
printf("%d ", token);
}
printf("\n");
auto input_ids = sd::Tensor<int32_t>::from_vector(tokens);
auto attention_mask = sd::Tensor<float>::from_vector(masks);
sd::Tensor<float> out;
int64_t t0 = ggml_time_ms();
auto out_opt = model.compute(8, input_ids, attention_mask);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
out = std::move(out_opt);
print_sd_tensor(out);
LOG_DEBUG("t5 test done in %lldms", t1 - t0);
}
}
static void load_from_file_and_test(const std::string& file_path) {
// cpu f16: pass
// cpu f32: pass
// cuda f16: pass
// cuda f32: pass
// cuda q8_0: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_F16;
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
return;
}
auto& tensor_storage_map = model_loader.get_tensor_storage_map();
for (auto& [name, tensor_storage] : tensor_storage_map) {
if (ends_with(name, "weight")) {
tensor_storage.expected_type = model_data_type;
}
}
std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, backend, tensor_storage_map, "", true);
if (!t5->alloc_params_buffer()) {
LOG_ERROR("t5 params buffer allocation failed");
return;
}
std::map<std::string, ggml_tensor*> tensors;
t5->get_param_tensors(tensors, "");
bool success = model_loader.load_tensors(tensors);
if (!success) {
LOG_ERROR("load tensors from model loader failed");
return;
}
LOG_INFO("t5 model loaded");
t5->test();
}
};
#endif // __SD_MODEL_TE_T5_HPP__
src/model/upscaler/esrgan.hpp
+375
View File
@@ -0,0 +1,375 @@
#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
#define __SD_MODEL_UPSCALER_ESRGAN_HPP__
#include "core/ggml_extend.hpp"
#include "model.h"
/*
=================================== ESRGAN ===================================
References:
https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py
https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py
*/
class ResidualDenseBlock : public GGMLBlock {
protected:
int num_feat;
int num_grow_ch;
public:
ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32)
: num_feat(num_feat), num_grow_ch(num_grow_ch) {
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
blocks["conv3"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
blocks["conv4"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
}
ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [n, num_feat, h, w]
// return: [n, num_feat, h, w]
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
auto x1 = lrelu(ctx, conv1->forward(ctx, x));
auto x_cat = ggml_concat(ctx->ggml_ctx, x, x1, 2);
auto x2 = lrelu(ctx, conv2->forward(ctx, x_cat));
x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x2, 2);
auto x3 = lrelu(ctx, conv3->forward(ctx, x_cat));
x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x3, 2);
auto x4 = lrelu(ctx, conv4->forward(ctx, x_cat));
x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
auto x5 = conv5->forward(ctx, x_cat);
x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
return x5;
}
};
class RRDB : public GGMLBlock {
public:
RRDB(int num_feat, int num_grow_ch = 32) {
blocks["rdb1"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
blocks["rdb2"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [n, num_feat, h, w]
// return: [n, num_feat, h, w]
auto rdb1 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb1"]);
auto rdb2 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb2"]);
auto rdb3 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb3"]);
auto out = rdb1->forward(ctx, x);
out = rdb2->forward(ctx, out);
out = rdb3->forward(ctx, out);
out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
return out;
}
};
class RRDBNet : public GGMLBlock {
protected:
int scale = 4;
int num_block = 23;
int num_in_ch = 3;
int num_out_ch = 3;
int num_feat = 64;
int num_grow_ch = 32;
public:
RRDBNet(int scale, int num_block, int num_in_ch, int num_out_ch, int num_feat, int num_grow_ch)
: scale(scale), num_block(num_block), num_in_ch(num_in_ch), num_out_ch(num_out_ch), num_feat(num_feat), num_grow_ch(num_grow_ch) {
blocks["conv_first"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_in_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
for (int i = 0; i < num_block; i++) {
std::string name = "body." + std::to_string(i);
blocks[name] = std::shared_ptr<GGMLBlock>(new RRDB(num_feat, num_grow_ch));
}
blocks["conv_body"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
if (scale >= 2) {
blocks["conv_up1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
}
if (scale == 4) {
blocks["conv_up2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
}
blocks["conv_hr"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
blocks["conv_last"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_out_ch, {3, 3}, {1, 1}, {1, 1}));
}
int get_scale() { return scale; }
int get_num_block() { return num_block; }
ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [n, num_in_ch, h, w]
// return: [n, num_out_ch, h*scale, w*scale]
auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
auto conv_body = std::dynamic_pointer_cast<Conv2d>(blocks["conv_body"]);
auto conv_hr = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
auto conv_last = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
auto feat = conv_first->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat");
auto body_feat = feat;
for (int i = 0; i < num_block; i++) {
std::string name = "body." + std::to_string(i);
auto block = std::dynamic_pointer_cast<RRDB>(blocks[name]);
body_feat = block->forward(ctx, body_feat);
sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat");
}
body_feat = conv_body->forward(ctx, body_feat);
feat = ggml_add(ctx->ggml_ctx, feat, body_feat);
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat");
// upsample
if (scale >= 2) {
auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat");
if (scale == 4) {
auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat");
}
}
// for all scales
auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out");
return out;
}
};
struct ESRGAN : public GGMLRunner {
std::unique_ptr<RRDBNet> rrdb_net;
int scale = 4;
int tile_size = 128; // avoid cuda OOM for 4gb VRAM
ESRGAN(ggml_backend_t backend,
ggml_backend_t params_backend,
int tile_size = 128,
const String2TensorStorage& tensor_storage_map = {})
: GGMLRunner(backend, params_backend) {
this->tile_size = tile_size;
}
std::string get_desc() override {
return "esrgan";
}
bool load_from_file(const std::string& file_path, int n_threads) {
LOG_INFO("loading esrgan from '%s'", file_path.c_str());
ModelLoader model_loader;
if (!model_loader.init_from_file_and_convert_name(file_path)) {
LOG_ERROR("init esrgan model loader from file failed: '%s'", file_path.c_str());
return false;
}
// Get tensor names
auto tensor_names = model_loader.get_tensor_names();
// Detect if it's ESRGAN format
bool is_ESRGAN = std::find(tensor_names.begin(), tensor_names.end(), "model.0.weight") != tensor_names.end();
// Detect parameters from tensor names
int detected_num_block = 0;
if (is_ESRGAN) {
for (const auto& name : tensor_names) {
if (name.find("model.1.sub.") == 0) {
size_t first_dot = name.find('.', 12);
if (first_dot != std::string::npos) {
size_t second_dot = name.find('.', first_dot + 1);
if (second_dot != std::string::npos && name.substr(first_dot + 1, 3) == "RDB") {
try {
int idx = std::stoi(name.substr(12, first_dot - 12));
detected_num_block = std::max(detected_num_block, idx + 1);
} catch (...) {
}
}
}
}
}
} else {
// Original format
for (const auto& name : tensor_names) {
if (name.find("body.") == 0) {
size_t pos = name.find('.', 5);
if (pos != std::string::npos) {
try {
int idx = std::stoi(name.substr(5, pos - 5));
detected_num_block = std::max(detected_num_block, idx + 1);
} catch (...) {
}
}
}
}
}
int detected_scale = 4; // default
if (is_ESRGAN) {
// For ESRGAN format, detect scale by highest model number
int max_model_num = 0;
for (const auto& name : tensor_names) {
if (name.find("model.") == 0) {
size_t dot_pos = name.find('.', 6);
if (dot_pos != std::string::npos) {
try {
int num = std::stoi(name.substr(6, dot_pos - 6));
max_model_num = std::max(max_model_num, num);
} catch (...) {
}
}
}
}
if (max_model_num <= 4) {
detected_scale = 1;
} else if (max_model_num <= 7) {
detected_scale = 2;
} else {
detected_scale = 4;
}
} else {
// Original format
bool has_conv_up2 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
return name == "conv_up2.weight";
});
bool has_conv_up1 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
return name == "conv_up1.weight";
});
if (has_conv_up2) {
detected_scale = 4;
} else if (has_conv_up1) {
detected_scale = 2;
} else {
detected_scale = 1;
}
}
int detected_num_in_ch = 3;
int detected_num_out_ch = 3;
int detected_num_feat = 64;
int detected_num_grow_ch = 32;
// Create RRDBNet with detected parameters
rrdb_net = std::make_unique<RRDBNet>(detected_scale, detected_num_block, detected_num_in_ch, detected_num_out_ch, detected_num_feat, detected_num_grow_ch);
rrdb_net->init(params_ctx, {}, "");
if (!alloc_params_buffer()) {
LOG_ERROR("esrgan model buffer allocation failed");
return false;
}
std::map<std::string, ggml_tensor*> esrgan_tensors;
rrdb_net->get_param_tensors(esrgan_tensors);
bool success;
if (is_ESRGAN) {
// Build name mapping for ESRGAN format
std::map<std::string, std::string> expected_to_model;
expected_to_model["conv_first.weight"] = "model.0.weight";
expected_to_model["conv_first.bias"] = "model.0.bias";
for (int i = 0; i < detected_num_block; i++) {
for (int j = 1; j <= 3; j++) {
for (int k = 1; k <= 5; k++) {
std::string expected_weight = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".weight";
std::string model_weight = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.weight";
expected_to_model[expected_weight] = model_weight;
std::string expected_bias = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".bias";
std::string model_bias = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.bias";
expected_to_model[expected_bias] = model_bias;
}
}
}
if (detected_scale == 1) {
expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
expected_to_model["conv_body.bias"] = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
expected_to_model["conv_hr.weight"] = "model.2.weight";
expected_to_model["conv_hr.bias"] = "model.2.bias";
expected_to_model["conv_last.weight"] = "model.4.weight";
expected_to_model["conv_last.bias"] = "model.4.bias";
} else {
expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
expected_to_model["conv_body.bias"] = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
if (detected_scale >= 2) {
expected_to_model["conv_up1.weight"] = "model.3.weight";
expected_to_model["conv_up1.bias"] = "model.3.bias";
}
if (detected_scale == 4) {
expected_to_model["conv_up2.weight"] = "model.6.weight";
expected_to_model["conv_up2.bias"] = "model.6.bias";
expected_to_model["conv_hr.weight"] = "model.8.weight";
expected_to_model["conv_hr.bias"] = "model.8.bias";
expected_to_model["conv_last.weight"] = "model.10.weight";
expected_to_model["conv_last.bias"] = "model.10.bias";
} else if (detected_scale == 2) {
expected_to_model["conv_hr.weight"] = "model.5.weight";
expected_to_model["conv_hr.bias"] = "model.5.bias";
expected_to_model["conv_last.weight"] = "model.7.weight";
expected_to_model["conv_last.bias"] = "model.7.bias";
}
}
std::map<std::string, ggml_tensor*> model_tensors;
for (auto& p : esrgan_tensors) {
auto it = expected_to_model.find(p.first);
if (it != expected_to_model.end()) {
model_tensors[it->second] = p.second;
}
}
success = model_loader.load_tensors(model_tensors, {}, n_threads);
} else {
success = model_loader.load_tensors(esrgan_tensors, {}, n_threads);
}
if (!success) {
LOG_ERROR("load esrgan tensors from model loader failed");
return false;
}
scale = rrdb_net->get_scale();
LOG_INFO("esrgan model loaded with scale=%d, num_block=%d", scale, detected_num_block);
return success;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
if (!rrdb_net)
return nullptr;
constexpr int kGraphNodes = 1 << 16; // 65k
ggml_cgraph* gf = new_graph_custom(kGraphNodes);
ggml_tensor* x = make_input(x_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& x) {
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
return result;
}
};
#endif // __SD_MODEL_UPSCALER_ESRGAN_HPP__
src/model/upscaler/ltx_latent_upscaler.hpp
+551
View File
@@ -0,0 +1,551 @@
#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
#define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
#include <cinttypes>
#include <cmath>
#include <cstdlib>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include "core/ggml_extend.hpp"
#include "core/ggml_graph_cut.h"
#include "core/util.h"
#include "model.h"
#include "model/diffusion/dit.hpp"
namespace LTXVUpsampler {
constexpr int LTX_UPSAMPLER_GRAPH_SIZE = 10240;
struct LatentUpsamplerConfig {
int64_t in_channels = 128;
int64_t mid_channels = 1024;
int num_blocks_per_stage = 4;
int dims = 3;
bool spatial_upsample = true;
bool temporal_upsample = false;
bool rational_resampler = false;
float spatial_scale = 2.f;
int spatial_up_num = 2;
int spatial_down_den = 1;
int temporal_up_factor = 1;
};
static inline bool has_tensor(const String2TensorStorage& tensor_storage_map,
const std::string& name) {
return tensor_storage_map.find(name) != tensor_storage_map.end();
}
static inline int64_t get_tensor_ne(const String2TensorStorage& tensor_storage_map,
const std::string& name,
int axis,
int64_t fallback) {
auto it = tensor_storage_map.find(name);
if (it == tensor_storage_map.end() || axis < 0 || axis >= GGML_MAX_DIMS) {
return fallback;
}
return it->second.ne[axis];
}
static inline int64_t get_tensor_ne0(const String2TensorStorage& tensor_storage_map,
const std::string& name,
int64_t fallback) {
return get_tensor_ne(tensor_storage_map, name, 0, fallback);
}
static inline int count_module_blocks(const String2TensorStorage& tensor_storage_map,
const std::string& module_name) {
int max_block = -1;
const std::string prefix = module_name + ".";
for (const auto& pair : tensor_storage_map) {
const std::string& name = pair.first;
if (name.find(prefix) != 0) {
continue;
}
size_t begin = prefix.size();
size_t end = name.find('.', begin);
if (end == std::string::npos) {
continue;
}
int index = atoi(name.substr(begin, end - begin).c_str());
max_block = std::max(max_block, index);
}
return max_block + 1;
}
static inline LatentUpsamplerConfig detect_config_from_weights(const String2TensorStorage& tensor_storage_map) {
LatentUpsamplerConfig config;
config.mid_channels = get_tensor_ne0(tensor_storage_map, "initial_norm.weight", config.mid_channels);
config.in_channels = get_tensor_ne0(tensor_storage_map, "final_conv.bias", config.in_channels);
int detected_blocks = count_module_blocks(tensor_storage_map, "res_blocks");
if (detected_blocks > 0) {
config.num_blocks_per_stage = detected_blocks;
}
config.rational_resampler = has_tensor(tensor_storage_map, "upsampler.conv.weight");
int64_t upsampler_out_channels = get_tensor_ne0(tensor_storage_map, "upsampler.0.bias", 0);
config.spatial_upsample = config.rational_resampler || upsampler_out_channels == 4 * config.mid_channels;
config.temporal_upsample = upsampler_out_channels == 2 * config.mid_channels;
if (config.temporal_upsample) {
config.temporal_up_factor = 2;
}
if (config.rational_resampler) {
int64_t out_channels = get_tensor_ne(tensor_storage_map,
"upsampler.conv.weight",
3,
config.mid_channels * 9);
if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
int64_t ratio = out_channels / config.mid_channels;
int num = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
config.spatial_up_num = num;
}
}
if (config.spatial_up_num == 3) {
config.spatial_down_den = 2;
config.spatial_scale = 1.5f;
} else if (config.spatial_up_num == 4) {
config.spatial_down_den = 1;
config.spatial_scale = 4.f;
} else {
config.spatial_down_den = 1;
config.spatial_scale = static_cast<float>(config.spatial_up_num);
}
}
return config;
}
class VideoGroupNorm : public GGMLBlock {
protected:
int num_groups;
int64_t num_channels;
float eps;
std::string prefix;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
this->prefix = prefix;
params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
}
public:
VideoGroupNorm(int num_groups, int64_t num_channels, float eps = 1e-05f)
: num_groups(num_groups),
num_channels(num_channels),
eps(eps) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// LTX video latent layout is [W, H, T, C]. ggml_group_norm treats ne[2]
// as channels, so fold only H/T internally and restore the same layout.
GGML_ASSERT(x->ne[3] == num_channels);
const int64_t W = x->ne[0];
const int64_t H = x->ne[1];
const int64_t T = x->ne[2];
x = ggml_ext_cont(ctx->ggml_ctx, x);
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, num_channels, 1);
x = ggml_group_norm(ctx->ggml_ctx, x, num_groups, eps);
ggml_tensor* weight = params["weight"];
ggml_tensor* bias = params["bias"];
if (ctx->weight_adapter) {
weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
}
weight = ggml_reshape_4d(ctx->ggml_ctx, weight, 1, 1, num_channels, 1);
bias = ggml_reshape_4d(ctx->ggml_ctx, bias, 1, 1, num_channels, 1);
x = ggml_mul_inplace(ctx->ggml_ctx, x, weight);
x = ggml_add_inplace(ctx->ggml_ctx, x, bias);
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, num_channels);
}
};
class ResBlock : public GGMLBlock {
public:
ResBlock(int64_t channels, int dims = 3) {
GGML_ASSERT(dims == 3);
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels, channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, channels));
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels, channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, channels));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv1 = std::dynamic_pointer_cast<Conv3d>(blocks["conv1"]);
auto norm1 = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["norm1"]);
auto conv2 = std::dynamic_pointer_cast<Conv3d>(blocks["conv2"]);
auto norm2 = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["norm2"]);
ggml_tensor* residual = x;
x = conv1->forward(ctx, x);
x = norm1->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
x = conv2->forward(ctx, x);
x = norm2->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, residual);
return ggml_silu(ctx->ggml_ctx, x);
}
};
class PixelShuffleND : public UnaryBlock {
protected:
int upscale_factor;
public:
explicit PixelShuffleND(int upscale_factor)
: upscale_factor(upscale_factor) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
GGML_ASSERT(upscale_factor > 0);
int64_t h = x->ne[1];
int64_t w = x->ne[0];
GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
// x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [b*f, h, w, c*p1*p2]
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [b*f, h*w, c*p1*p2]
return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
}
};
class TemporalPixelShuffleND : public UnaryBlock {
protected:
int upscale_factor;
public:
explicit TemporalPixelShuffleND(int upscale_factor)
: upscale_factor(upscale_factor) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
GGML_ASSERT(upscale_factor > 0);
GGML_ASSERT(x->ne[3] % upscale_factor == 0);
const int64_t W = x->ne[0];
const int64_t H = x->ne[1];
const int64_t F = x->ne[2];
const int64_t C = x->ne[3] / upscale_factor;
// x: [b, c*p, f, h, w] -> [b, c, f*p, h, w]
x = ggml_ext_cont(ctx->ggml_ctx, x);
x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C);
}
};
class BlurDownsample : public GGMLBlock {
protected:
int64_t channels;
int stride;
ggml_tensor* kernel = nullptr;
std::vector<float> kernel_data;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
if (stride == 1) {
return;
}
kernel = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
std::string name = prefix + "kernel";
ggml_set_name(kernel, name.c_str());
static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
for (int64_t c = 0; c < channels; ++c) {
for (int y = 0; y < 5; ++y) {
for (int x = 0; x < 5; ++x) {
kernel_data[static_cast<size_t>(x + 5 * (y + 5 * c))] =
binomial[y] * binomial[x] / 256.f;
}
}
}
}
public:
BlurDownsample(int64_t channels, int stride)
: channels(channels),
stride(stride) {
GGML_ASSERT(stride >= 1);
}
void load_fixed_tensors() {
if (kernel == nullptr || kernel_data.empty()) {
return;
}
ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (stride == 1) {
return x;
}
GGML_ASSERT(kernel != nullptr);
GGML_ASSERT(x->ne[2] == channels);
if (ctx->conv2d_direct_enabled) {
return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
}
return ggml_conv_2d_dw(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
}
};
class SpatialRationalResampler : public GGMLBlock {
protected:
int64_t mid_channels;
int num;
int den;
public:
SpatialRationalResampler(int64_t mid_channels, int num, int den)
: mid_channels(mid_channels),
num(num),
den(den) {
GGML_ASSERT(num >= 1);
GGML_ASSERT(den >= 1);
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(mid_channels, num * num * mid_channels, {3, 3}, {1, 1}, {1, 1}));
blocks["pixel_shuffle"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(num));
blocks["blur_down"] = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
}
void load_fixed_tensors() {
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
blur_down->load_fixed_tensors();
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
// rearrange(x, "b c f h w -> (b f) c h w")
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
x = conv->forward(ctx, x);
x = pixel_shuffle->forward(ctx, x);
x = blur_down->forward(ctx, x);
return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}
};
class LatentUpsampler : public GGMLBlock {
public:
LatentUpsamplerConfig config;
explicit LatentUpsampler(LatentUpsamplerConfig config)
: config(std::move(config)) {
GGML_ASSERT(this->config.dims == 3);
GGML_ASSERT(this->config.spatial_upsample || this->config.temporal_upsample);
blocks["initial_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.in_channels,
this->config.mid_channels,
{3, 3, 3},
{1, 1, 1},
{1, 1, 1}));
blocks["initial_norm"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, this->config.mid_channels));
for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
blocks["res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
}
if (this->config.rational_resampler) {
blocks["upsampler"] = std::shared_ptr<GGMLBlock>(new SpatialRationalResampler(this->config.mid_channels,
this->config.spatial_up_num,
this->config.spatial_down_den));
} else if (this->config.temporal_upsample) {
blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
this->config.temporal_up_factor * this->config.mid_channels,
{3, 3, 3},
{1, 1, 1},
{1, 1, 1}));
blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new TemporalPixelShuffleND(this->config.temporal_up_factor));
} else {
blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
4 * this->config.mid_channels,
{3, 3},
{1, 1},
{1, 1}));
blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
}
for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
blocks["post_upsample_res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
}
blocks["final_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
this->config.in_channels,
{3, 3, 3},
{1, 1, 1},
{1, 1, 1}));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [b, c, f, h, w]
// return: [b, c, scaled_f, scaled_h, scaled_w]
auto initial_conv = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
auto initial_norm = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
auto final_conv = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
x = initial_conv->forward(ctx, x);
x = initial_norm->forward(ctx, x);
x = ggml_silu(ctx->ggml_ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.initial", "x");
for (int i = 0; i < config.num_blocks_per_stage; ++i) {
auto block = std::dynamic_pointer_cast<ResBlock>(blocks["res_blocks." + std::to_string(i)]);
x = block->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.res_blocks." + std::to_string(i), "x");
}
if (config.rational_resampler) {
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
x = upsampler->forward(ctx, x);
} else if (config.temporal_upsample) {
auto upsample_conv = std::dynamic_pointer_cast<Conv3d>(blocks["upsampler.0"]);
auto pixel_shuffle = std::dynamic_pointer_cast<TemporalPixelShuffleND>(blocks["upsampler.1"]);
x = upsample_conv->forward(ctx, x); // [b, c*2, f, h, w]
x = pixel_shuffle->forward(ctx, x); // [b, c, f*2, h, w]
x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); // x[:, :, 1:, :, :]
} else {
auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
// rearrange(x, "b c f h w -> (b f) c h w"),
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*f, c, h, w]
x = upsample_conv->forward(ctx, x); // [b*f, c*4, h, w]
x = pixel_shuffle->forward(ctx, x); // [b*f, c, h*2, w*2]
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*c, f, h, w]
}
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
for (int i = 0; i < config.num_blocks_per_stage; ++i) {
auto block = std::dynamic_pointer_cast<ResBlock>(blocks["post_upsample_res_blocks." + std::to_string(i)]);
x = block->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.post_blocks." + std::to_string(i), "x");
}
x = final_conv->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
return x;
}
void load_fixed_tensors() {
if (!config.rational_resampler) {
return;
}
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
upsampler->load_fixed_tensors();
}
};
struct LatentUpsamplerRunner : public GGMLRunner {
std::unique_ptr<LatentUpsampler> model;
LatentUpsamplerRunner(ggml_backend_t backend,
ggml_backend_t params_backend)
: GGMLRunner(backend, params_backend) {}
std::string get_desc() override {
return "ltx_latent_upsampler";
}
bool load_from_file(const std::string& file_path, int n_threads) {
LOG_INFO("loading LTX latent upsampler from '%s'", file_path.c_str());
ModelLoader model_loader;
if (!model_loader.init_from_file(file_path)) {
LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", file_path.c_str());
return false;
}
const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
bool has_regular_upsampler = has_tensor(tensor_storage_map, "upsampler.0.weight");
bool has_rational_spatial = has_tensor(tensor_storage_map, "upsampler.conv.weight");
if (!has_tensor(tensor_storage_map, "post_upsample_res_blocks.0.conv2.bias") ||
(!has_regular_upsampler && !has_rational_spatial)) {
LOG_ERROR("unsupported LTX latent upsampler weights: expected upsampler tensors");
return false;
}
LatentUpsamplerConfig config = detect_config_from_weights(tensor_storage_map);
if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
config.dims,
config.spatial_upsample,
config.temporal_upsample,
config.rational_resampler,
config.spatial_scale,
config.temporal_up_factor);
return false;
}
model = std::make_unique<LatentUpsampler>(config);
model->init(params_ctx, tensor_storage_map, "");
if (!alloc_params_buffer()) {
LOG_ERROR("LTX latent upsampler params buffer allocation failed");
return false;
}
std::map<std::string, ggml_tensor*> tensors;
model->get_param_tensors(tensors);
std::set<std::string> ignore_tensors;
if (config.rational_resampler) {
ignore_tensors.insert("upsampler.blur_down.kernel");
}
if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
LOG_ERROR("load LTX latent upsampler tensors failed");
return false;
}
model->load_fixed_tensors();
LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, temporal_factor=%d, rational=%d",
config.in_channels,
config.mid_channels,
config.num_blocks_per_stage,
config.spatial_scale,
config.temporal_up_factor,
config.rational_resampler);
return true;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
if (!model) {
return nullptr;
}
ggml_cgraph* gf = new_graph_custom(LTX_UPSAMPLER_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = model->forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<float>& x) {
if (!model) {
LOG_ERROR("LTX latent upsampler is not loaded");
return {};
}
if (x.dim() != 4 && x.dim() != 5) {
LOG_ERROR("LTX latent upsampler expects 4D or 5D video latent, got dim=%lld",
(long long)x.dim());
return {};
}
if (x.dim() == 5 && x.shape()[4] != 1) {
LOG_ERROR("LTX latent upsampler currently supports batch size 1, got batch=%lld",
(long long)x.shape()[4]);
return {};
}
if (x.shape()[3] != model->config.in_channels) {
LOG_ERROR("LTX latent upsampler expected %" PRId64 " channels, got %lld",
model->config.in_channels,
(long long)x.shape()[3]);
return {};
}
size_t expected_dim = static_cast<size_t>(x.dim());
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
}
};
} // namespace LTXVUpsampler
#endif // __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
src/model/vae/auto_encoder_kl.hpp
+889
View File
@@ -0,0 +1,889 @@
#ifndef __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
#define __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
#include "model/vae/vae.hpp"
/*================================================== AutoEncoderKL ===================================================*/
#define VAE_GRAPH_SIZE 20480
class ResnetBlock : public UnaryBlock {
protected:
int64_t in_channels;
int64_t out_channels;
public:
ResnetBlock(int64_t in_channels,
int64_t out_channels)
: in_channels(in_channels),
out_channels(out_channels) {
// temb_channels is always 0
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(in_channels));
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(out_channels));
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
if (out_channels != in_channels) {
blocks["nin_shortcut"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, out_channels, {1, 1}));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [N, in_channels, h, w]
// t_emb is always None
auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto norm2 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm2"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
auto h = x;
h = norm1->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish
h = conv1->forward(ctx, h);
// return h;
h = norm2->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish
// dropout, skip for inference
h = conv2->forward(ctx, h);
// skip connection
if (out_channels != in_channels) {
auto nin_shortcut = std::dynamic_pointer_cast<Conv2d>(blocks["nin_shortcut"]);
x = nin_shortcut->forward(ctx, x); // [N, out_channels, h, w]
}
h = ggml_add(ctx->ggml_ctx, h, x);
return h; // [N, out_channels, h, w]
}
};
class AttnBlock : public UnaryBlock {
protected:
int64_t in_channels;
bool use_linear;
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()) {
if (iter->second.n_dims == 4 && use_linear) {
use_linear = false;
blocks["q"] = std::make_shared<Conv2d>(in_channels, in_channels, std::pair{1, 1});
blocks["k"] = std::make_shared<Conv2d>(in_channels, in_channels, std::pair{1, 1});
blocks["v"] = std::make_shared<Conv2d>(in_channels, in_channels, std::pair{1, 1});
blocks["proj_out"] = std::make_shared<Conv2d>(in_channels, in_channels, std::pair{1, 1});
} else if (iter->second.n_dims == 2 && !use_linear) {
use_linear = true;
blocks["q"] = std::make_shared<Linear>(in_channels, in_channels);
blocks["k"] = std::make_shared<Linear>(in_channels, in_channels);
blocks["v"] = std::make_shared<Linear>(in_channels, in_channels);
blocks["proj_out"] = std::make_shared<Linear>(in_channels, in_channels);
}
}
}
public:
AttnBlock(int64_t in_channels, bool use_linear)
: in_channels(in_channels), use_linear(use_linear) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(in_channels));
if (use_linear) {
blocks["q"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, in_channels));
blocks["k"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, in_channels));
blocks["v"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, in_channels));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(in_channels, in_channels));
} else {
blocks["q"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, in_channels, {1, 1}));
blocks["k"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, in_channels, {1, 1}));
blocks["v"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, in_channels, {1, 1}));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, in_channels, {1, 1}));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [N, in_channels, h, w]
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
auto q_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
auto k_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["k"]);
auto v_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["v"]);
auto proj_out = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj_out"]);
auto h_ = norm->forward(ctx, x);
const int64_t n = h_->ne[3];
const int64_t c = h_->ne[2];
const int64_t h = h_->ne[1];
const int64_t w = h_->ne[0];
ggml_tensor* q;
ggml_tensor* k;
ggml_tensor* v;
if (use_linear) {
h_ = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h_, 1, 2, 0, 3)); // [N, h, w, in_channels]
h_ = ggml_reshape_3d(ctx->ggml_ctx, h_, c, h * w, n); // [N, h * w, in_channels]
q = q_proj->forward(ctx, h_); // [N, h * w, in_channels]
k = k_proj->forward(ctx, h_); // [N, h * w, in_channels]
v = v_proj->forward(ctx, h_); // [N, h * w, in_channels]
} else {
q = q_proj->forward(ctx, h_); // [N, in_channels, h, w]
q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 1, 2, 0, 3)); // [N, h, w, in_channels]
q = ggml_reshape_3d(ctx->ggml_ctx, q, c, h * w, n); // [N, h * w, in_channels]
k = k_proj->forward(ctx, h_); // [N, in_channels, h, w]
k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 1, 2, 0, 3)); // [N, h, w, in_channels]
k = ggml_reshape_3d(ctx->ggml_ctx, k, c, h * w, n); // [N, h * w, in_channels]
v = v_proj->forward(ctx, h_); // [N, in_channels, h, w]
v = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3)); // [N, h, w, in_channels]
v = ggml_reshape_3d(ctx->ggml_ctx, v, c, h * w, n); // [N, h * w, in_channels]
}
h_ = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
if (use_linear) {
h_ = proj_out->forward(ctx, h_); // [N, h * w, in_channels]
h_ = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h_, 1, 0, 2, 3)); // [N, in_channels, h * w]
h_ = ggml_reshape_4d(ctx->ggml_ctx, h_, w, h, c, n); // [N, in_channels, h, w]
} else {
h_ = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h_, 1, 0, 2, 3)); // [N, in_channels, h * w]
h_ = ggml_reshape_4d(ctx->ggml_ctx, h_, w, h, c, n); // [N, in_channels, h, w]
h_ = proj_out->forward(ctx, h_); // [N, in_channels, h, w]
}
h_ = ggml_add(ctx->ggml_ctx, h_, x);
return h_;
}
};
class AE3DConv : public Conv2d {
public:
AE3DConv(int64_t in_channels,
int64_t out_channels,
std::pair<int, int> kernel_size,
int video_kernel_size = 3,
std::pair<int, int> stride = {1, 1},
std::pair<int, int> padding = {0, 0},
std::pair<int, int> dilation = {1, 1},
bool bias = true)
: Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias) {
int kernel_padding = video_kernel_size / 2;
blocks["time_mix_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(out_channels,
out_channels,
{video_kernel_size, 1, 1},
{1, 1, 1},
{kernel_padding, 0, 0}));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x) override {
// timesteps always None
// skip_video always False
// x: [N, IC, IH, IW]
// result: [N, OC, OH, OW]
auto time_mix_conv = std::dynamic_pointer_cast<Conv3d>(blocks["time_mix_conv"]);
x = Conv2d::forward(ctx, x);
// timesteps = x.shape[0]
// x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
// x = conv3d(x)
// return rearrange(x, "b c t h w -> (b t) c h w")
int64_t T = x->ne[3];
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)
x = time_mix_conv->forward(ctx, x); // [B, OC, T, OH * OW]
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; // [B*T, OC, OH, OW]
}
};
class VideoResnetBlock : public ResnetBlock {
protected:
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
enum ggml_type wtype = get_type(prefix + "mix_factor", tensor_storage_map, GGML_TYPE_F32);
params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
}
float get_alpha() {
float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]);
return sigmoid(alpha);
}
public:
VideoResnetBlock(int64_t in_channels,
int64_t out_channels,
int video_kernel_size = 3)
: ResnetBlock(in_channels, out_channels) {
// merge_strategy is always learned
blocks["time_stack"] = std::shared_ptr<GGMLBlock>(new ResBlock(out_channels, 0, out_channels, {video_kernel_size, 1}, 3, false, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [N, in_channels, h, w] aka [b*t, in_channels, h, w]
// return: [N, out_channels, h, w] aka [b*t, out_channels, h, w]
// t_emb is always None
// skip_video is always False
// timesteps is always None
auto time_stack = std::dynamic_pointer_cast<ResBlock>(blocks["time_stack"]);
x = ResnetBlock::forward(ctx, x); // [N, out_channels, h, w]
// return x;
int64_t T = x->ne[3];
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;
x = time_stack->forward(ctx, x); // b t c (h w)
float alpha = get_alpha();
x = ggml_add(ctx->ggml_ctx,
ggml_ext_scale(ctx->ggml_ctx, x, alpha),
ggml_ext_scale(ctx->ggml_ctx, x_mix, 1.0f - alpha));
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;
}
};
// ldm.modules.diffusionmodules.model.Encoder
class Encoder : public GGMLBlock {
protected:
int ch = 128;
std::vector<int> ch_mult = {1, 2, 4, 4};
int num_res_blocks = 2;
int in_channels = 3;
int z_channels = 4;
bool double_z = true;
public:
Encoder(int ch,
std::vector<int> ch_mult,
int num_res_blocks,
int in_channels,
int z_channels,
bool double_z = true,
bool use_linear_projection = false)
: ch(ch),
ch_mult(ch_mult),
num_res_blocks(num_res_blocks),
in_channels(in_channels),
z_channels(z_channels),
double_z(double_z) {
blocks["conv_in"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, ch, {3, 3}, {1, 1}, {1, 1}));
size_t num_resolutions = ch_mult.size();
int block_in = 1;
for (int i = 0; i < num_resolutions; i++) {
if (i == 0) {
block_in = ch;
} else {
block_in = ch * ch_mult[i - 1];
}
int block_out = ch * ch_mult[i];
for (int j = 0; j < num_res_blocks; j++) {
std::string name = "down." + std::to_string(i) + ".block." + std::to_string(j);
blocks[name] = std::shared_ptr<GGMLBlock>(new ResnetBlock(block_in, block_out));
block_in = block_out;
}
if (i != num_resolutions - 1) {
std::string name = "down." + std::to_string(i) + ".downsample";
blocks[name] = std::shared_ptr<GGMLBlock>(new DownSampleBlock(block_in, block_in, true));
}
}
blocks["mid.block_1"] = std::shared_ptr<GGMLBlock>(new ResnetBlock(block_in, block_in));
blocks["mid.attn_1"] = std::shared_ptr<GGMLBlock>(new AttnBlock(block_in, use_linear_projection));
blocks["mid.block_2"] = std::shared_ptr<GGMLBlock>(new ResnetBlock(block_in, block_in));
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(block_in));
blocks["conv_out"] = std::shared_ptr<GGMLBlock>(new Conv2d(block_in, double_z ? z_channels * 2 : z_channels, {3, 3}, {1, 1}, {1, 1}));
}
virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, in_channels, h, w]
auto conv_in = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"]);
auto mid_block_1 = std::dynamic_pointer_cast<ResnetBlock>(blocks["mid.block_1"]);
auto mid_attn_1 = std::dynamic_pointer_cast<AttnBlock>(blocks["mid.attn_1"]);
auto mid_block_2 = std::dynamic_pointer_cast<ResnetBlock>(blocks["mid.block_2"]);
auto norm_out = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm_out"]);
auto conv_out = std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"]);
auto h = conv_in->forward(ctx, x); // [N, ch, h, w]
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.prelude", "h");
// downsampling
size_t num_resolutions = ch_mult.size();
for (int i = 0; i < num_resolutions; i++) {
for (int j = 0; j < num_res_blocks; j++) {
std::string name = "down." + std::to_string(i) + ".block." + std::to_string(j);
auto down_block = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
h = down_block->forward(ctx, h);
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".block." + std::to_string(j), "h");
}
if (i != num_resolutions - 1) {
std::string name = "down." + std::to_string(i) + ".downsample";
auto down_sample = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
h = down_sample->forward(ctx, h);
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".downsample", "h");
}
}
// middle
h = mid_block_1->forward(ctx, h);
h = mid_attn_1->forward(ctx, h);
h = mid_block_2->forward(ctx, h); // [N, block_in, h, w]
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.mid", "h");
// end
h = norm_out->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // nonlinearity/swish
h = conv_out->forward(ctx, h); // [N, z_channels*2, h, w]
return h;
}
};
// ldm.modules.diffusionmodules.model.Decoder
class Decoder : public GGMLBlock {
protected:
int ch = 128;
int out_ch = 3;
std::vector<int> ch_mult = {1, 2, 4, 4};
int num_res_blocks = 2;
int z_channels = 4;
bool video_decoder = false;
int video_kernel_size = 3;
virtual std::shared_ptr<GGMLBlock> get_conv_out(int64_t in_channels,
int64_t out_channels,
std::pair<int, int> kernel_size,
std::pair<int, int> stride = {1, 1},
std::pair<int, int> padding = {0, 0}) {
if (video_decoder) {
return std::shared_ptr<GGMLBlock>(new AE3DConv(in_channels, out_channels, kernel_size, video_kernel_size, stride, padding));
} else {
return std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, out_channels, kernel_size, stride, padding));
}
}
virtual std::shared_ptr<GGMLBlock> get_resnet_block(int64_t in_channels,
int64_t out_channels) {
if (video_decoder) {
return std::shared_ptr<GGMLBlock>(new VideoResnetBlock(in_channels, out_channels, video_kernel_size));
} else {
return std::shared_ptr<GGMLBlock>(new ResnetBlock(in_channels, out_channels));
}
}
public:
Decoder(int ch,
int out_ch,
std::vector<int> ch_mult,
int num_res_blocks,
int z_channels,
bool use_linear_projection = false,
bool video_decoder = false,
int video_kernel_size = 3)
: ch(ch),
out_ch(out_ch),
ch_mult(ch_mult),
num_res_blocks(num_res_blocks),
z_channels(z_channels),
video_decoder(video_decoder),
video_kernel_size(video_kernel_size) {
int num_resolutions = static_cast<int>(ch_mult.size());
int block_in = ch * ch_mult[num_resolutions - 1];
blocks["conv_in"] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, block_in, {3, 3}, {1, 1}, {1, 1}));
blocks["mid.block_1"] = get_resnet_block(block_in, block_in);
blocks["mid.attn_1"] = std::shared_ptr<GGMLBlock>(new AttnBlock(block_in, use_linear_projection));
blocks["mid.block_2"] = get_resnet_block(block_in, block_in);
for (int i = num_resolutions - 1; i >= 0; i--) {
int mult = ch_mult[i];
int block_out = ch * mult;
for (int j = 0; j < num_res_blocks + 1; j++) {
std::string name = "up." + std::to_string(i) + ".block." + std::to_string(j);
blocks[name] = get_resnet_block(block_in, block_out);
block_in = block_out;
}
if (i != 0) {
std::string name = "up." + std::to_string(i) + ".upsample";
blocks[name] = std::shared_ptr<GGMLBlock>(new UpSampleBlock(block_in, block_in));
}
}
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(block_in));
blocks["conv_out"] = get_conv_out(block_in, out_ch, {3, 3}, {1, 1}, {1, 1});
}
virtual ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) {
// z: [N, z_channels, h, w]
// alpha is always 0
// merge_strategy is always learned
// time_mode is always conv-only, so we need to replace conv_out_op/resnet_op to AE3DConv/VideoResBlock
// AttnVideoBlock will not be used
auto conv_in = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"]);
auto mid_block_1 = std::dynamic_pointer_cast<ResnetBlock>(blocks["mid.block_1"]);
auto mid_attn_1 = std::dynamic_pointer_cast<AttnBlock>(blocks["mid.attn_1"]);
auto mid_block_2 = std::dynamic_pointer_cast<ResnetBlock>(blocks["mid.block_2"]);
auto norm_out = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm_out"]);
auto conv_out = std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"]);
// conv_in
auto h = conv_in->forward(ctx, z); // [N, block_in, h, w]
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.prelude", "h");
// middle
h = mid_block_1->forward(ctx, h);
// return h;
h = mid_attn_1->forward(ctx, h);
h = mid_block_2->forward(ctx, h); // [N, block_in, h, w]
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.mid", "h");
// upsampling
int num_resolutions = static_cast<int>(ch_mult.size());
for (int i = num_resolutions - 1; i >= 0; i--) {
for (int j = 0; j < num_res_blocks + 1; j++) {
std::string name = "up." + std::to_string(i) + ".block." + std::to_string(j);
auto up_block = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
h = up_block->forward(ctx, h);
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".block." + std::to_string(j), "h");
}
if (i != 0) {
std::string name = "up." + std::to_string(i) + ".upsample";
auto up_sample = std::dynamic_pointer_cast<UpSampleBlock>(blocks[name]);
h = up_sample->forward(ctx, h);
// sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".upsample", "h");
}
}
h = norm_out->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // nonlinearity/swish
h = conv_out->forward(ctx, h); // [N, out_ch, h*8, w*8]
return h;
}
};
// ldm.models.autoencoder.AutoencoderKL
class AutoEncoderKLModel : public GGMLBlock {
protected:
SDVersion version;
bool decode_only = true;
bool use_video_decoder = false;
bool use_quant = true;
int embed_dim = 4;
struct {
int z_channels = 4;
int resolution = 256;
int in_channels = 3;
int out_ch = 3;
int ch = 128;
std::vector<int> ch_mult = {1, 2, 4, 4};
int num_res_blocks = 2;
bool double_z = true;
} dd_config;
static std::string get_tensor_name(const std::string& prefix, const std::string& name) {
return prefix.empty() ? name : prefix + "." + name;
}
void detect_decoder_ch(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
int& decoder_ch) {
auto conv_in_iter = tensor_storage_map.find(get_tensor_name(prefix, "decoder.conv_in.weight"));
if (conv_in_iter != tensor_storage_map.end() && conv_in_iter->second.n_dims >= 4 && conv_in_iter->second.ne[3] > 0) {
int last_ch_mult = dd_config.ch_mult.back();
int64_t conv_in_out_channels = conv_in_iter->second.ne[3];
if (last_ch_mult > 0 && conv_in_out_channels % last_ch_mult == 0) {
decoder_ch = static_cast<int>(conv_in_out_channels / last_ch_mult);
LOG_INFO("vae decoder: ch = %d", decoder_ch);
} else {
LOG_WARN("vae decoder: failed to infer ch from %s (%" PRId64 " / %d)",
get_tensor_name(prefix, "decoder.conv_in.weight").c_str(),
conv_in_out_channels,
last_ch_mult);
}
}
}
public:
AutoEncoderKLModel(SDVersion version = VERSION_SD1,
bool decode_only = true,
bool use_linear_projection = false,
bool use_video_decoder = false,
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "")
: version(version), decode_only(decode_only), use_video_decoder(use_video_decoder) {
if (sd_version_is_dit(version)) {
if (sd_version_uses_flux2_vae(version)) {
dd_config.z_channels = 32;
embed_dim = 32;
} else {
use_quant = false;
dd_config.z_channels = 16;
}
}
if (use_video_decoder) {
use_quant = false;
}
int decoder_ch = dd_config.ch;
detect_decoder_ch(tensor_storage_map, prefix, decoder_ch);
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder(decoder_ch,
dd_config.out_ch,
dd_config.ch_mult,
dd_config.num_res_blocks,
dd_config.z_channels,
use_linear_projection,
use_video_decoder));
if (use_quant) {
blocks["post_quant_conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(dd_config.z_channels,
embed_dim,
{1, 1}));
}
if (!decode_only) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder(dd_config.ch,
dd_config.ch_mult,
dd_config.num_res_blocks,
dd_config.in_channels,
dd_config.z_channels,
dd_config.double_z,
use_linear_projection));
if (use_quant) {
int factor = dd_config.double_z ? 2 : 1;
blocks["quant_conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(embed_dim * factor,
dd_config.z_channels * factor,
{1, 1}));
}
}
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
// z: [N, z_channels, h, w]
if (sd_version_uses_flux2_vae(version)) {
// [N, C*p*p, h, w] -> [N, C, h*p, w*p]
int64_t p = 2;
int64_t N = z->ne[3];
int64_t C = z->ne[2] / p / p;
int64_t h = z->ne[1];
int64_t w = z->ne[0];
int64_t H = h * p;
int64_t W = w * p;
z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N); // [N, C, p*p, h*w]
z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 1, 0, 2, 3)); // [N, C, h*w, p*p]
z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N); // [N*C*h, w, p, p]
z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 0, 2, 1, 3)); // [N*C*h, p, w, p]
z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N); // [N, C, h*p, w*p]
}
if (use_quant) {
auto post_quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["post_quant_conv"]);
z = post_quant_conv->forward(ctx, z); // [N, z_channels, h, w]
// sd::ggml_graph_cut::mark_graph_cut(z, "vae.decode.prelude", "z");
}
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
ggml_set_name(z, "bench-start");
auto h = decoder->forward(ctx, z);
ggml_set_name(h, "bench-end");
return h;
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, in_channels, h, w]
auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
auto z = encoder->forward(ctx, x); // [N, 2*z_channels, h/8, w/8]
if (use_quant) {
auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["quant_conv"]);
z = quant_conv->forward(ctx, z); // [N, 2*embed_dim, h/8, w/8]
// sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z");
}
if (sd_version_uses_flux2_vae(version)) {
z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
// [N, C, H, W] -> [N, C*p*p, H/p, W/p]
int64_t p = 2;
int64_t N = z->ne[3];
int64_t C = z->ne[2];
int64_t H = z->ne[1];
int64_t W = z->ne[0];
int64_t h = H / p;
int64_t w = W / p;
z = ggml_reshape_4d(ctx->ggml_ctx, z, p, w, p, h * C * N); // [N*C*h, p, w, p]
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3)); // [N*C*h, w, p, p]
z = ggml_reshape_4d(ctx->ggml_ctx, z, p * p, w * h, C, N); // [N, C, h*w, p*p]
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 1, 0, 2, 3)); // [N, C, p*p, h*w]
z = ggml_reshape_4d(ctx->ggml_ctx, z, w, h, p * p * C, N); // [N, C*p*p, h*w]
}
return z;
}
int get_encoder_output_channels() {
int factor = dd_config.double_z ? 2 : 1;
if (sd_version_uses_flux2_vae(version)) {
return dd_config.z_channels * 4;
}
return dd_config.z_channels * factor;
}
};
struct AutoEncoderKL : public VAE {
float scale_factor = 1.f;
float shift_factor = 0.f;
bool decode_only = true;
AutoEncoderKLModel ae;
AutoEncoderKL(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
bool decode_only = false,
bool use_video_decoder = false,
SDVersion version = VERSION_SD1)
: decode_only(decode_only), VAE(version, backend, params_backend) {
if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
scale_factor = 0.18215f;
shift_factor = 0.f;
} else if (sd_version_is_sdxl(version)) {
scale_factor = 0.13025f;
shift_factor = 0.f;
} else if (sd_version_is_sd3(version)) {
scale_factor = 1.5305f;
shift_factor = 0.0609f;
} else if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_longcat(version)) {
scale_factor = 0.3611f;
shift_factor = 0.1159f;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 1.0f;
shift_factor = 0.f;
}
bool use_linear_projection = false;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "attn_1.proj_out.weight")) {
if (tensor_storage.n_dims == 2) {
use_linear_projection = true;
}
break;
}
}
ae = AutoEncoderKLModel(version, decode_only, use_linear_projection, use_video_decoder, tensor_storage_map, prefix);
ae.init(params_ctx, tensor_storage_map, prefix);
}
void set_conv2d_scale(float scale) override {
std::vector<GGMLBlock*> blocks;
ae.get_all_blocks(blocks);
for (auto block : blocks) {
if (block->get_desc() == "Conv2d") {
auto conv_block = (Conv2d*)block;
conv_block->set_scale(scale);
}
}
}
std::string get_desc() override {
return "vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
ae.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
ggml_tensor* z = make_input(z_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = decode_graph ? ae.decode(&runner_ctx, z) : ae.encode(&runner_ctx, z);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
bool decode_graph) override {
GGML_ASSERT(!decode_only || decode_graph);
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(z, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
}
sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {
// ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample
auto chunks = sd::ops::chunk(moments, 2, 2);
const auto& mean = chunks[0];
const auto& logvar = chunks[1];
sd::Tensor<float> stddev = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -30.0f, 20.0f));
sd::Tensor<float> noise = sd::Tensor<float>::randn_like(mean, rng);
sd::Tensor<float> latents = mean + stddev * noise;
return latents;
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
if (sd_version_uses_flux2_vae(version)) {
return vae_output;
} else if (version == VERSION_SD1_PIX2PIX) {
return sd::ops::chunk(vae_output, 2, 2)[0];
} else {
return gaussian_latent_sample(vae_output, rng);
}
}
std::pair<sd::Tensor<float>, sd::Tensor<float>> get_latents_mean_std(const sd::Tensor<float>& latents, int channel_dim) {
GGML_ASSERT(channel_dim >= 0 && static_cast<size_t>(channel_dim) < static_cast<size_t>(latents.dim()));
if (sd_version_uses_flux2_vae(version)) {
GGML_ASSERT(latents.shape()[channel_dim] == 128);
std::vector<int64_t> stats_shape(static_cast<size_t>(latents.dim()), 1);
stats_shape[static_cast<size_t>(channel_dim)] = latents.shape()[channel_dim];
auto mean_tensor = sd::Tensor<float>::from_vector({-0.0676f, -0.0715f, -0.0753f, -0.0745f, 0.0223f, 0.0180f, 0.0142f, 0.0184f,
-0.0001f, -0.0063f, -0.0002f, -0.0031f, -0.0272f, -0.0281f, -0.0276f, -0.0290f,
-0.0769f, -0.0672f, -0.0902f, -0.0892f, 0.0168f, 0.0152f, 0.0079f, 0.0086f,
0.0083f, 0.0015f, 0.0003f, -0.0043f, -0.0439f, -0.0419f, -0.0438f, -0.0431f,
-0.0102f, -0.0132f, -0.0066f, -0.0048f, -0.0311f, -0.0306f, -0.0279f, -0.0180f,
0.0030f, 0.0015f, 0.0126f, 0.0145f, 0.0347f, 0.0338f, 0.0337f, 0.0283f,
0.0020f, 0.0047f, 0.0047f, 0.0050f, 0.0123f, 0.0081f, 0.0081f, 0.0146f,
0.0681f, 0.0679f, 0.0767f, 0.0732f, -0.0462f, -0.0474f, -0.0392f, -0.0511f,
-0.0528f, -0.0477f, -0.0470f, -0.0517f, -0.0317f, -0.0316f, -0.0345f, -0.0283f,
0.0510f, 0.0445f, 0.0578f, 0.0458f, -0.0412f, -0.0458f, -0.0487f, -0.0467f,
-0.0088f, -0.0106f, -0.0088f, -0.0046f, -0.0376f, -0.0432f, -0.0436f, -0.0499f,
0.0118f, 0.0166f, 0.0203f, 0.0279f, 0.0113f, 0.0129f, 0.0016f, 0.0072f,
-0.0118f, -0.0018f, -0.0141f, -0.0054f, -0.0091f, -0.0138f, -0.0145f, -0.0187f,
0.0323f, 0.0305f, 0.0259f, 0.0300f, 0.0540f, 0.0614f, 0.0495f, 0.0590f,
-0.0511f, -0.0603f, -0.0478f, -0.0524f, -0.0227f, -0.0274f, -0.0154f, -0.0255f,
-0.0572f, -0.0565f, -0.0518f, -0.0496f, 0.0116f, 0.0054f, 0.0163f, 0.0104f});
mean_tensor.reshape_(stats_shape);
auto std_tensor = sd::Tensor<float>::from_vector({1.8029f, 1.7786f, 1.7868f, 1.7837f, 1.7717f, 1.7590f, 1.7610f, 1.7479f,
1.7336f, 1.7373f, 1.7340f, 1.7343f, 1.8626f, 1.8527f, 1.8629f, 1.8589f,
1.7593f, 1.7526f, 1.7556f, 1.7583f, 1.7363f, 1.7400f, 1.7355f, 1.7394f,
1.7342f, 1.7246f, 1.7392f, 1.7304f, 1.7551f, 1.7513f, 1.7559f, 1.7488f,
1.8449f, 1.8454f, 1.8550f, 1.8535f, 1.8240f, 1.7813f, 1.7854f, 1.7945f,
1.8047f, 1.7876f, 1.7695f, 1.7676f, 1.7782f, 1.7667f, 1.7925f, 1.7848f,
1.7579f, 1.7407f, 1.7483f, 1.7368f, 1.7961f, 1.7998f, 1.7920f, 1.7925f,
1.7780f, 1.7747f, 1.7727f, 1.7749f, 1.7526f, 1.7447f, 1.7657f, 1.7495f,
1.7775f, 1.7720f, 1.7813f, 1.7813f, 1.8162f, 1.8013f, 1.8023f, 1.8033f,
1.7527f, 1.7331f, 1.7563f, 1.7482f, 1.7610f, 1.7507f, 1.7681f, 1.7613f,
1.7665f, 1.7545f, 1.7828f, 1.7726f, 1.7896f, 1.7999f, 1.7864f, 1.7760f,
1.7613f, 1.7625f, 1.7560f, 1.7577f, 1.7783f, 1.7671f, 1.7810f, 1.7799f,
1.7201f, 1.7068f, 1.7265f, 1.7091f, 1.7793f, 1.7578f, 1.7502f, 1.7455f,
1.7587f, 1.7500f, 1.7525f, 1.7362f, 1.7616f, 1.7572f, 1.7444f, 1.7430f,
1.7509f, 1.7610f, 1.7634f, 1.7612f, 1.7254f, 1.7135f, 1.7321f, 1.7226f,
1.7664f, 1.7624f, 1.7718f, 1.7664f, 1.7457f, 1.7441f, 1.7569f, 1.7530f});
std_tensor.reshape_(stats_shape);
return {std::move(mean_tensor), std::move(std_tensor)};
} else {
GGML_ABORT("unknown version %d", version);
}
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
if (sd_version_uses_flux2_vae(version)) {
int channel_dim = 2;
auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
return (latents * std_tensor) / scale_factor + mean_tensor;
}
return (latents / scale_factor) + shift_factor;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
if (sd_version_uses_flux2_vae(version)) {
int channel_dim = 2;
auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim);
return ((latents - mean_tensor) * scale_factor) / std_tensor;
}
return (latents - shift_factor) * scale_factor;
}
int get_encoder_output_channels(int input_channels) {
return ae.get_encoder_output_channels();
}
void test() {
ggml_init_params params;
params.mem_size = static_cast<size_t>(10 * 1024 * 1024); // 10 MB
params.mem_buffer = nullptr;
params.no_alloc = false;
ggml_context* ctx = ggml_init(params);
GGML_ASSERT(ctx != nullptr);
{
// CPU, x{1, 3, 64, 64}: Pass
// CUDA, x{1, 3, 64, 64}: Pass, but sill get wrong result for some image, may be due to interlnal nan
// CPU, x{2, 3, 64, 64}: Wrong result
// CUDA, x{2, 3, 64, 64}: Wrong result, and different from CPU result
sd::Tensor<float> x({64, 64, 3, 2});
x.fill_(0.5f);
print_sd_tensor(x);
sd::Tensor<float> out;
int64_t t0 = ggml_time_ms();
auto out_opt = _compute(8, x, false);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
out = std::move(out_opt);
print_sd_tensor(out);
LOG_DEBUG("encode test done in %lldms", t1 - t0);
}
if (false) {
// CPU, z{1, 4, 8, 8}: Pass
// CUDA, z{1, 4, 8, 8}: Pass
// CPU, z{3, 4, 8, 8}: Wrong result
// CUDA, z{3, 4, 8, 8}: Wrong result, and different from CPU result
sd::Tensor<float> z({8, 8, 4, 1});
z.fill_(0.5f);
print_sd_tensor(z);
sd::Tensor<float> out;
int64_t t0 = ggml_time_ms();
auto out_opt = _compute(8, z, true);
int64_t t1 = ggml_time_ms();
GGML_ASSERT(!out_opt.empty());
out = std::move(out_opt);
print_sd_tensor(out);
LOG_DEBUG("decode test done in %lldms", t1 - t0);
}
};
};
#endif // __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
src/model/vae/ltx_audio_vae.hpp
+1098
View File
File diff suppressed because it is too large. Load diff
src/model/vae/ltx_vae.hpp
+1555
View File
File diff suppressed because it is too large. Load diff
src/model/vae/tae.hpp
+753
View File
@@ -0,0 +1,753 @@
#ifndef __SD_MODEL_VAE_TAE_HPP__
#define __SD_MODEL_VAE_TAE_HPP__
#include "core/ggml_extend.hpp"
#include "model.h"
/*
=================================== TinyAutoEncoder ===================================
References:
https://github.com/huggingface/diffusers/blob/main/src/diffusers/model/autoencoders/vae.py
https://github.com/madebyollin/taesd/blob/main/taesd.py
*/
class TAEBlock : public UnaryBlock {
protected:
int n_in;
int n_out;
bool use_midblock_gn;
public:
TAEBlock(int n_in, int n_out, bool use_midblock_gn = false)
: n_in(n_in), n_out(n_out), use_midblock_gn(use_midblock_gn) {
blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {3, 3}, {1, 1}, {1, 1}));
blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_out, n_out, {3, 3}, {1, 1}, {1, 1}));
if (n_in != n_out) {
blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_out, {1, 1}, {1, 1}, {1, 1}, {1, 1}, false));
}
if (use_midblock_gn) {
int n_gn = n_in * 4;
blocks["pool.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_in, n_gn, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
blocks["pool.1"] = std::shared_ptr<GGMLBlock>(new GroupNorm(4, n_gn));
// pool.2 is ReLU, handled in forward
blocks["pool.3"] = std::shared_ptr<GGMLBlock>(new Conv2d(n_gn, n_in, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [n, n_in, h, w]
// return: [n, n_out, h, w]
if (use_midblock_gn) {
auto pool_0 = std::dynamic_pointer_cast<Conv2d>(blocks["pool.0"]);
auto pool_1 = std::dynamic_pointer_cast<GroupNorm>(blocks["pool.1"]);
auto pool_3 = std::dynamic_pointer_cast<Conv2d>(blocks["pool.3"]);
auto p = pool_0->forward(ctx, x);
p = pool_1->forward(ctx, p);
p = ggml_relu_inplace(ctx->ggml_ctx, p);
p = pool_3->forward(ctx, p);
x = ggml_add(ctx->ggml_ctx, x, p);
}
auto conv_0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
auto conv_2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
auto conv_4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
auto h = conv_0->forward(ctx, x);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv_2->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv_4->forward(ctx, h);
if (n_in != n_out) {
auto skip = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
LOG_DEBUG("skip");
x = skip->forward(ctx, x);
}
h = ggml_add(ctx->ggml_ctx, h, x);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
return h;
}
};
class TinyEncoder : public UnaryBlock {
int in_channels = 3;
int channels = 64;
int z_channels = 4;
int num_blocks = 3;
public:
TinyEncoder(int z_channels = 4, bool use_midblock_gn = false)
: z_channels(z_channels) {
int index = 0;
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels, channels, {3, 3}, {1, 1}, {1, 1}));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
for (int i = 0; i < num_blocks; i++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
}
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
for (int i = 0; i < num_blocks; i++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
}
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
for (int i = 0; i < num_blocks; i++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
}
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, z_channels, {3, 3}, {1, 1}, {1, 1}));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [n, in_channels, h, w]
// return: [n, z_channels, h/8, w/8]
for (int i = 0; i < num_blocks * 3 + 6; i++) {
auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
x = block->forward(ctx, x);
}
return x;
}
};
class TinyDecoder : public UnaryBlock {
int z_channels = 4;
int channels = 64;
int out_channels = 3;
int num_blocks = 3;
public:
TinyDecoder(int z_channels = 4, bool use_midblock_gn = false)
: z_channels(z_channels) {
int index = 0;
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels, {3, 3}, {1, 1}, {1, 1}));
index++; // nn.ReLU()
for (int i = 0; i < num_blocks; i++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels, use_midblock_gn));
}
index++; // nn.Upsample()
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
for (int i = 0; i < num_blocks; i++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
}
index++; // nn.Upsample()
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
for (int i = 0; i < num_blocks; i++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
}
index++; // nn.Upsample()
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels, {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TAEBlock(channels, channels));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
// z: [n, z_channels, h, w]
// return: [n, out_channels, h*8, w*8]
auto h = ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f);
h = ggml_tanh_inplace(ctx->ggml_ctx, h);
h = ggml_ext_scale(ctx->ggml_ctx, h, 3.0f);
for (int i = 0; i < num_blocks * 3 + 10; i++) {
if (blocks.find(std::to_string(i)) == blocks.end()) {
if (i == 1) {
h = ggml_relu_inplace(ctx->ggml_ctx, h);
} else {
h = ggml_upscale(ctx->ggml_ctx, h, 2, GGML_SCALE_MODE_NEAREST);
}
continue;
}
auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(i)]);
h = block->forward(ctx, h);
}
return h;
}
};
class TPool : public UnaryBlock {
int stride;
public:
TPool(int channels, int stride)
: stride(stride) {
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * stride, channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
auto h = x;
if (stride != 1) {
h = ggml_reshape_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2] * stride, h->ne[3] / stride);
}
h = conv->forward(ctx, h);
return h;
}
};
class TGrow : public UnaryBlock {
int stride;
public:
TGrow(int channels, int stride)
: stride(stride) {
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, channels * stride, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks["conv"]);
auto h = conv->forward(ctx, x);
if (stride != 1) {
h = ggml_reshape_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2] / stride, h->ne[3] * stride);
}
return h;
}
};
class MemBlock : public GGMLBlock {
bool has_skip_conv = false;
public:
MemBlock(int channels, int out_channels)
: has_skip_conv(channels != out_channels) {
blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * 2, out_channels, {3, 3}, {1, 1}, {1, 1}));
blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
if (has_skip_conv) {
blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* past) {
// x: [n, channels, h, w]
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.2"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
auto h = ggml_concat(ctx->ggml_ctx, x, past, 2);
h = conv0->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv1->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv2->forward(ctx, h);
auto skip = x;
if (has_skip_conv) {
auto skip_conv = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
skip = skip_conv->forward(ctx, x);
}
h = ggml_add_inplace(ctx->ggml_ctx, h, skip);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
return h;
}
};
class WideMemBlock : public GGMLBlock {
bool has_skip_conv = false;
public:
WideMemBlock(int channels, int out_channels)
: has_skip_conv(channels != out_channels) {
int groups = std::max(1, out_channels / 64);
blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * 2, out_channels, {1, 1}, {1, 1}));
blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d_grouped(out_channels, out_channels, groups, {3, 3}, {1, 1}, {1, 1}));
blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {1, 1}, {1, 1}));
blocks["conv.6"] = std::shared_ptr<GGMLBlock>(new Conv2d_grouped(out_channels, out_channels, groups, {3, 3}, {1, 1}, {1, 1}));
if (has_skip_conv) {
blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* past) {
// x: [n, channels, h, w]
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv.2"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
auto conv3 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv.6"]);
auto h = ggml_concat(ctx->ggml_ctx, x, past, 2);
h = conv0->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv1->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv2->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv3->forward(ctx, h);
auto skip = x;
if (has_skip_conv) {
auto skip_conv = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
skip = skip_conv->forward(ctx, x);
}
h = ggml_add_inplace(ctx->ggml_ctx, h, skip);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
return h;
}
};
ggml_tensor*
patchify(ggml_context* ctx,
ggml_tensor* x,
int64_t patch_size,
int64_t b = 1) {
// x: [f, b*c, h*q, w*r]
// return: [f, b*c*r*q, h, w]
if (patch_size == 1) {
return x;
}
int64_t r = patch_size;
int64_t q = patch_size;
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int64_t f = x->ne[3];
int64_t w = W / r;
int64_t h = H / q;
x = ggml_reshape_4d(ctx, x, W, q, h, C * f); // [W, q, h, C*f]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [W, h, q, C*f]
x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f); // [r, w, h, q*C*f]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [w, h, r, q*C*f]
x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f); // [f, b*c*r*q, h, w]
return x;
}
ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t patch_size,
int64_t b = 1) {
// x: [f, b*c*r*q, h, w]
// return: [f, b*c, h*q, w*r]
if (patch_size == 1) {
return x;
}
int64_t r = patch_size;
int64_t q = patch_size;
int64_t c = x->ne[2] / b / q / r;
int64_t f = x->ne[3];
int64_t h = x->ne[1];
int64_t w = x->ne[0];
x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f); // [q*c*b*f, r, h, w]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [r, w, h, q*c*b*f]
x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f); // [c*b*f, q, h, r*w]
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [r*w, q, h, c*b*f]
x = ggml_reshape_4d(ctx, x, r * w, q * h, c * b, f);
return x;
}
class TinyVideoEncoder : public UnaryBlock {
int in_channels = 3;
int hidden = 64;
int z_channels = 4;
int num_blocks = 3;
int num_layers = 3;
int patch_size = 1;
public:
int t_downscale = 1;
TinyVideoEncoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_downscale = {true, true, false})
: z_channels(z_channels), patch_size(patch_size) {
t_downscale = 1;
for (bool downscale : time_downscale) {
if (downscale) {
t_downscale *= 2;
}
}
int index = 0;
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels * patch_size * patch_size, hidden, {3, 3}, {1, 1}, {1, 1}));
index++; // nn.ReLU()
for (int i = 0; i < num_layers; i++) {
int stride = time_downscale[i] ? 2 : 1;
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TPool(hidden, stride));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, hidden, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
for (int j = 0; j < num_blocks; j++) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(hidden, hidden));
}
}
blocks[std::to_string(index)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, z_channels, {3, 3}, {1, 1}, {1, 1}));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["0"]);
if (patch_size > 1) {
z = patchify(ctx->ggml_ctx, z, patch_size, 1);
}
auto h = first_conv->forward(ctx, z);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
int index = 2;
for (int i = 0; i < num_layers; i++) {
auto pool = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
auto conv = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
h = pool->forward(ctx, h);
h = conv->forward(ctx, h);
for (int j = 0; j < num_blocks; j++) {
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
h = block->forward(ctx, h, mem);
}
}
auto last_conv = std::dynamic_pointer_cast<Conv2d>(blocks[std::to_string(index)]);
h = last_conv->forward(ctx, h);
return h;
}
};
class TinyVideoDecoder : public UnaryBlock {
int z_channels = 4;
int out_channels = 3;
int num_blocks = 3;
static const int num_layers = 3;
int channels[num_layers + 1] = {256, 128, 64, 64};
int patch_size = 1;
int t_upscale = 1;
bool is_wide = false;
public:
TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_upscale = {false, true, true}, bool is_wide = false)
: z_channels(z_channels), patch_size(patch_size), is_wide(is_wide) {
t_upscale = 1;
if (is_wide) {
channels[0] = 1024;
channels[1] = 512;
channels[2] = 256;
}
for (bool upscale : time_upscale) {
if (upscale) {
t_upscale *= 2;
}
}
int index = 1; // Clamp()
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels[0], {3, 3}, {1, 1}, {1, 1}));
index++; // nn.ReLU()
for (int i = 0; i < num_layers; i++) {
int stride = time_upscale[i] ? 2 : 1;
for (int j = 0; j < num_blocks; j++) {
if (is_wide) {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new WideMemBlock(channels[i], channels[i]));
} else {
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(channels[i], channels[i]));
}
}
index++; // nn.Upsample()
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TGrow(channels[i], stride));
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels[i], channels[i + 1], {3, 3}, {1, 1}, {1, 1}, {1, 1}, false));
}
index++; // nn.ReLU()
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(channels[num_layers], out_channels * patch_size * patch_size, {3, 3}, {1, 1}, {1, 1}));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z) override {
auto first_conv = std::dynamic_pointer_cast<Conv2d>(blocks["1"]);
// Clamp()
auto h = ggml_ext_scale(ctx->ggml_ctx,
ggml_tanh_inplace(ctx->ggml_ctx,
ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f)),
3.0f,
true);
h = first_conv->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
int index = 3;
for (int i = 0; i < num_layers; i++) {
for (int j = 0; j < num_blocks; j++) {
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
if (is_wide) {
auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
h = block->forward(ctx, h, mem);
} else {
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
h = block->forward(ctx, h, mem);
}
}
// upsample
index++;
h = ggml_upscale(ctx->ggml_ctx, h, 2, GGML_SCALE_MODE_NEAREST);
auto block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
h = block->forward(ctx, h);
block = std::dynamic_pointer_cast<UnaryBlock>(blocks[std::to_string(index++)]);
h = block->forward(ctx, h);
}
h = ggml_relu_inplace(ctx->ggml_ctx, h);
auto last_conv = std::dynamic_pointer_cast<Conv2d>(blocks[std::to_string(++index)]);
h = last_conv->forward(ctx, h);
if (patch_size > 1) {
h = unpatchify(ctx->ggml_ctx, h, patch_size, 1);
}
// shape(W, H, 3, (t_upscale - 1) + T) => shape(W, H, 3, T)
h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - (t_upscale - 1), h->nb[1], h->nb[2], h->nb[3], (t_upscale - 1) * h->nb[3]);
return h;
}
};
class TAEHV : public GGMLBlock {
protected:
bool decode_only;
SDVersion version;
bool is_wide;
public:
int z_channels = 16;
std::vector<bool> time_downscale = {true, true, false};
std::vector<bool> time_upscale = {false, true, true};
public:
TAEHV(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_wide = false)
: decode_only(decode_only), version(version), is_wide(is_wide) {
int patch = 1;
if (version == VERSION_WAN2_2_TI2V) {
z_channels = 48;
patch = 2;
} else if (sd_version_is_ltxav(version)) {
z_channels = 128;
patch = 4;
time_downscale = {true, true, true};
time_upscale = {true, true, true};
}
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch, time_upscale, is_wide));
if (!decode_only) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoEncoder(z_channels, patch, time_downscale));
}
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
}
auto result = decoder->forward(ctx, z);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
}
return result;
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
// (W, H, T, C) -> (W, H, C, T)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
int64_t num_frames = x->ne[3];
if (num_frames % encoder->t_downscale) {
// pad to multiple of encoder->t_downscale at the end
auto last_frame = ggml_view_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], x->ne[2], 1, x->nb[1], x->nb[2], x->nb[3], (num_frames - 1) * x->nb[3]);
for (int i = 0; i < encoder->t_downscale - num_frames % encoder->t_downscale; i++) {
x = ggml_concat(ctx->ggml_ctx, x, last_frame, 3);
}
}
x = encoder->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
return x;
}
};
class TAESD : public GGMLBlock {
protected:
bool decode_only;
bool taef2 = false;
public:
int z_channels = 4;
public:
TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
: decode_only(decode_only) {
bool use_midblock_gn = false;
taef2 = sd_version_uses_flux2_vae(version);
if (sd_version_is_dit(version)) {
z_channels = 16;
}
if (taef2) {
z_channels = 32;
use_midblock_gn = true;
}
blocks["decoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyDecoder(z_channels, use_midblock_gn));
if (!decode_only) {
blocks["encoder.layers"] = std::shared_ptr<GGMLBlock>(new TinyEncoder(z_channels, use_midblock_gn));
}
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
auto decoder = std::dynamic_pointer_cast<TinyDecoder>(blocks["decoder.layers"]);
if (taef2) {
z = unpatchify(ctx->ggml_ctx, z, 2);
}
return decoder->forward(ctx, z);
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto encoder = std::dynamic_pointer_cast<TinyEncoder>(blocks["encoder.layers"]);
auto z = encoder->forward(ctx, x);
if (taef2) {
z = patchify(ctx->ggml_ctx, z, 2);
}
return z;
}
};
struct TinyImageAutoEncoder : public VAE {
TAESD taesd;
bool decode_only = false;
TinyImageAutoEncoder(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
bool decoder_only = true,
SDVersion version = VERSION_SD1)
: decode_only(decoder_only),
taesd(decoder_only, version),
VAE(version, backend, params_backend) {
scale_input = false;
taesd.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "taesd";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
taesd.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents;
}
int get_encoder_output_channels(int input_channels) {
return taesd.z_channels;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
ggml_tensor* z = make_input(z_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = decode_graph ? taesd.decode(&runner_ctx, z) : taesd.encode(&runner_ctx, z);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z_tensor,
bool decode_graph) override {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
}
};
struct TinyVideoAutoEncoder : public VAE {
TAEHV taehv;
bool decode_only = false;
bool is_wide = false;
TinyVideoAutoEncoder(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
bool decoder_only = true,
SDVersion version = VERSION_WAN2)
: decode_only(decoder_only),
VAE(version, backend, params_backend) {
for (auto tensor_storage : tensor_storage_map) {
if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) {
is_wide = true;
break;
}
}
taehv = TAEHV(decoder_only, version, is_wide);
scale_input = false;
taehv.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "taehv";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
taehv.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents;
}
int get_encoder_output_channels(int input_channels) {
return taehv.z_channels;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
ggml_cgraph* gf = decode_graph && is_wide ? ggml_new_graph_custom(compute_ctx, 4096, false)
: ggml_new_graph(compute_ctx);
ggml_tensor* z = make_input(z_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z_tensor,
bool decode_graph) override {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(z_tensor, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
}
};
#endif // __SD_MODEL_VAE_TAE_HPP__
src/model/vae/vae.hpp
+261
View File
@@ -0,0 +1,261 @@
#ifndef __SD_MODEL_VAE_VAE_HPP__
#define __SD_MODEL_VAE_VAE_HPP__
#include "core/tensor_ggml.hpp"
#include "model/common/block.hpp"
struct VAE : public GGMLRunner {
protected:
SDVersion version;
bool scale_input = true;
virtual sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
bool decode_graph) = 0;
static inline void scale_tensor_to_minus1_1(sd::Tensor<float>* tensor) {
GGML_ASSERT(tensor != nullptr);
for (int64_t i = 0; i < tensor->numel(); ++i) {
(*tensor)[i] = (*tensor)[i] * 2.0f - 1.0f;
}
}
static inline void scale_tensor_to_0_1(sd::Tensor<float>* tensor) {
GGML_ASSERT(tensor != nullptr);
for (int64_t i = 0; i < tensor->numel(); ++i) {
float value = ((*tensor)[i] + 1.0f) * 0.5f;
(*tensor)[i] = std::max(0.0f, std::min(1.0f, value));
}
}
sd::Tensor<float> tiled_compute(const sd::Tensor<float>& input,
int n_threads,
int output_width,
int output_height,
int scale,
int p_tile_size_x,
int p_tile_size_y,
float tile_overlap_factor,
bool circular_x,
bool circular_y,
bool decode_graph,
const char* error_message,
bool silent = false) {
auto on_processing = [&](const sd::Tensor<float>& input_tile) {
auto output_tile = _compute(n_threads, input_tile, decode_graph);
if (output_tile.empty()) {
LOG_ERROR("%s", error_message);
return sd::Tensor<float>();
}
return output_tile;
};
return ::process_tiles_2d(input,
output_width,
output_height,
scale,
p_tile_size_x,
p_tile_size_y,
tile_overlap_factor,
circular_x,
circular_y,
on_processing,
silent);
}
public:
VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
: version(version), GGMLRunner(backend, params_backend) {}
int get_scale_factor() {
int scale_factor = 8;
if (version == VERSION_LTXAV) {
scale_factor = 32;
} else if (version == VERSION_WAN2_2_TI2V) {
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;
} else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
scale_factor = 1;
}
return scale_factor;
}
virtual int get_encoder_output_channels(int input_channels) = 0;
void get_tile_sizes(int& tile_size_x,
int& tile_size_y,
float& tile_overlap,
const sd_tiling_params_t& params,
int64_t latent_x,
int64_t latent_y,
float encoding_factor = 1.0f) {
tile_overlap = std::max(std::min(params.target_overlap, 0.5f), 0.0f);
auto get_tile_size = [&](int requested_size, float factor, int64_t latent_size) {
const int default_tile_size = 32;
const int min_tile_dimension = 4;
int tile_size = default_tile_size;
// factor <= 1 means simple fraction of the latent dimension
// factor > 1 means number of tiles across that dimension
if (factor > 0.f) {
if (factor > 1.0)
factor = 1 / (factor - factor * tile_overlap + tile_overlap);
tile_size = static_cast<int>(std::round(latent_size * factor));
} else if (requested_size >= min_tile_dimension) {
tile_size = requested_size;
}
tile_size = static_cast<int>(tile_size * encoding_factor);
return std::max(std::min(tile_size, static_cast<int>(latent_size)), min_tile_dimension);
};
tile_size_x = get_tile_size(params.tile_size_x, params.rel_size_x, latent_x);
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;
if (scale_input) {
scale_tensor_to_minus1_1(&input);
}
if (tiling_params.enabled) {
const int scale_factor = get_scale_factor();
int64_t W = input.shape()[0] / scale_factor;
int64_t H = input.shape()[1] / scale_factor;
float tile_overlap;
int tile_size_x, tile_size_y;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, 1.30539f);
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,
n_threads,
static_cast<int>(W),
static_cast<int>(H),
scale_factor,
tile_size_x,
tile_size_y,
tile_overlap,
circular_x,
circular_y,
false,
"vae encode compute failed while processing a tile");
} else {
output = _compute(n_threads, input, false);
}
free_compute_buffer();
if (output.empty()) {
LOG_ERROR("vae encode compute failed");
return {};
}
int64_t t1 = ggml_time_ms();
LOG_DEBUG("computing vae encode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
return std::move(output);
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) {
int64_t t0 = ggml_time_ms();
sd::Tensor<float> input = x;
sd::Tensor<float> output;
set_tiling_params(tiling_params);
if (tiling_params.enabled) {
const int scale_factor = get_scale_factor();
int64_t W = input.shape()[0] * scale_factor;
int64_t H = input.shape()[1] * scale_factor;
float tile_overlap;
int tile_size_x, tile_size_y;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, input.shape()[0], input.shape()[1]);
if (!silent) {
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
}
output = tiled_compute(
input,
n_threads,
static_cast<int>(W),
static_cast<int>(H),
scale_factor,
tile_size_x,
tile_size_y,
tile_overlap,
circular_x,
circular_y,
true,
"vae decode compute failed while processing a tile",
silent);
} else {
output = _compute(n_threads, input, true);
}
free_compute_buffer();
if (output.empty()) {
LOG_ERROR("vae decode compute failed");
return {};
}
if (scale_input) {
scale_tensor_to_0_1(&output);
}
int64_t t1 = ggml_time_ms();
LOG_DEBUG("computing vae decode graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
return std::move(output);
}
virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) = 0;
virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) = 0;
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
virtual void set_tiling_params(const sd_tiling_params_t& params) {
set_temporal_tiling_enabled(params.temporal_tiling);
};
};
struct FakeVAE : public VAE {
FakeVAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
: VAE(version, backend, params_backend) {}
int get_encoder_output_channels(int input_channels) {
return input_channels;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
bool decode_graph) override {
SD_UNUSED(n_threads);
SD_UNUSED(decode_graph);
return z;
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {}
std::string get_desc() override {
return "fake_vae";
}
};
#endif // __SD_MODEL_VAE_VAE_HPP__
src/model/vae/wan_vae.hpp
+1361
View File
File diff suppressed because it is too large. Load diff