1312 lines
50 KiB
C++
1312 lines
50 KiB
C++
#include <stdarg.h>
|
|
#include <fstream>
|
|
#include <regex>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
#include "model.h"
|
|
#include "stable-diffusion.h"
|
|
#include "util.h"
|
|
#include "vocab.hpp"
|
|
|
|
#include "ggml/ggml-alloc.h"
|
|
#include "ggml/ggml-backend.h"
|
|
#include "ggml/ggml.h"
|
|
|
|
uint64_t read_u64(uint8_t* buffer) {
|
|
// little endian
|
|
uint64_t value = 0;
|
|
value |= static_cast<int64_t>(buffer[7]) << 56;
|
|
value |= static_cast<int64_t>(buffer[6]) << 48;
|
|
value |= static_cast<int64_t>(buffer[5]) << 40;
|
|
value |= static_cast<int64_t>(buffer[4]) << 32;
|
|
value |= static_cast<int64_t>(buffer[3]) << 24;
|
|
value |= static_cast<int64_t>(buffer[2]) << 16;
|
|
value |= static_cast<int64_t>(buffer[1]) << 8;
|
|
value |= static_cast<int64_t>(buffer[0]);
|
|
return value;
|
|
}
|
|
|
|
int32_t read_int(uint8_t* buffer) {
|
|
// little endian
|
|
int value = 0;
|
|
value |= buffer[3] << 24;
|
|
value |= buffer[2] << 16;
|
|
value |= buffer[1] << 8;
|
|
value |= buffer[0];
|
|
return value;
|
|
}
|
|
|
|
uint16_t read_short(uint8_t* buffer) {
|
|
// little endian
|
|
uint16_t value = 0;
|
|
value |= buffer[1] << 8;
|
|
value |= buffer[0];
|
|
return value;
|
|
}
|
|
|
|
/*================================================= Preprocess ==================================================*/
|
|
|
|
std::string self_attn_names[] = {
|
|
"self_attn.q_proj.weight",
|
|
"self_attn.k_proj.weight",
|
|
"self_attn.v_proj.weight",
|
|
|
|
"self_attn.q_proj.bias",
|
|
"self_attn.k_proj.bias",
|
|
"self_attn.v_proj.bias",
|
|
};
|
|
|
|
const char* unused_tensors[] = {
|
|
"betas",
|
|
"alphas_cumprod_prev",
|
|
"sqrt_alphas_cumprod",
|
|
"sqrt_one_minus_alphas_cumprod",
|
|
"log_one_minus_alphas_cumprod",
|
|
"sqrt_recip_alphas_cumprod",
|
|
"sqrt_recipm1_alphas_cumprod",
|
|
"posterior_variance",
|
|
"posterior_log_variance_clipped",
|
|
"posterior_mean_coef1",
|
|
"posterior_mean_coef2",
|
|
"cond_stage_model.transformer.text_model.embeddings.position_ids",
|
|
"cond_stage_model.model.logit_scale",
|
|
"cond_stage_model.model.text_projection",
|
|
"model.diffusion_model.time_embedding.cond_proj.weight",
|
|
"unet.time_embedding.cond_proj.weight",
|
|
"model_ema.decay",
|
|
"model_ema.num_updates",
|
|
"model_ema.diffusion_model",
|
|
"control_model",
|
|
"embedding_manager",
|
|
};
|
|
|
|
bool is_unused_tensor(std::string name) {
|
|
for (int i = 0; i < sizeof(unused_tensors) / sizeof(const char*); i++) {
|
|
if (starts_with(name, unused_tensors[i])) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
std::unordered_map<std::string, std::string> open_clip_to_hf_clip_model = {
|
|
{"cond_stage_model.model.ln_final.bias", "cond_stage_model.transformer.text_model.final_layer_norm.bias"},
|
|
{"cond_stage_model.model.ln_final.weight", "cond_stage_model.transformer.text_model.final_layer_norm.weight"},
|
|
{"cond_stage_model.model.positional_embedding", "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight"},
|
|
{"cond_stage_model.model.token_embedding.weight", "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight"},
|
|
|
|
};
|
|
|
|
std::unordered_map<std::string, std::string> open_clip_to_hk_clip_resblock = {
|
|
{"attn.out_proj.bias", "self_attn.out_proj.bias"},
|
|
{"attn.out_proj.weight", "self_attn.out_proj.weight"},
|
|
{"ln_1.bias", "layer_norm1.bias"},
|
|
{"ln_1.weight", "layer_norm1.weight"},
|
|
{"ln_2.bias", "layer_norm2.bias"},
|
|
{"ln_2.weight", "layer_norm2.weight"},
|
|
{"mlp.c_fc.bias", "mlp.fc1.bias"},
|
|
{"mlp.c_fc.weight", "mlp.fc1.weight"},
|
|
{"mlp.c_proj.bias", "mlp.fc2.bias"},
|
|
{"mlp.c_proj.weight", "mlp.fc2.weight"},
|
|
};
|
|
|
|
std::unordered_map<std::string, std::string> vae_decoder_name_map = {
|
|
{"first_stage_model.decoder.mid.attn_1.to_k.bias", "first_stage_model.decoder.mid.attn_1.k.bias"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_k.weight", "first_stage_model.decoder.mid.attn_1.k.weight"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_out.0.bias", "first_stage_model.decoder.mid.attn_1.proj_out.bias"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_out.0.weight", "first_stage_model.decoder.mid.attn_1.proj_out.weight"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_q.bias", "first_stage_model.decoder.mid.attn_1.q.bias"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_q.weight", "first_stage_model.decoder.mid.attn_1.q.weight"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_v.bias", "first_stage_model.decoder.mid.attn_1.v.bias"},
|
|
{"first_stage_model.decoder.mid.attn_1.to_v.weight", "first_stage_model.decoder.mid.attn_1.v.weight"},
|
|
};
|
|
|
|
std::string convert_open_clip_to_hf_clip(const std::string& name) {
|
|
std::string new_name = name;
|
|
std::string open_clip_resblock_prefix = "cond_stage_model.model.transformer.resblocks.";
|
|
std::string hf_clip_resblock_prefix = "cond_stage_model.transformer.text_model.encoder.layers.";
|
|
|
|
if (open_clip_to_hf_clip_model.find(name) != open_clip_to_hf_clip_model.end()) {
|
|
new_name = open_clip_to_hf_clip_model[name];
|
|
}
|
|
|
|
if (name.find(open_clip_resblock_prefix) == 0) {
|
|
std::string remain = name.substr(open_clip_resblock_prefix.length());
|
|
std::string idx = remain.substr(0, remain.find("."));
|
|
std::string suffix = remain.substr(idx.length() + 1);
|
|
|
|
if (suffix == "attn.in_proj_weight" || suffix == "attn.in_proj_bias") {
|
|
new_name = hf_clip_resblock_prefix + idx + "." + suffix;
|
|
} else if (open_clip_to_hk_clip_resblock.find(suffix) != open_clip_to_hk_clip_resblock.end()) {
|
|
std::string new_suffix = open_clip_to_hk_clip_resblock[suffix];
|
|
new_name = hf_clip_resblock_prefix + idx + "." + new_suffix;
|
|
}
|
|
}
|
|
|
|
return new_name;
|
|
}
|
|
|
|
std::string convert_vae_decoder_name(const std::string& name) {
|
|
if (vae_decoder_name_map.find(name) != vae_decoder_name_map.end()) {
|
|
return vae_decoder_name_map[name];
|
|
}
|
|
return name;
|
|
}
|
|
|
|
std::unordered_map<std::string, std::unordered_map<std::string, std::string>> suffix_conversion_underline = {
|
|
{
|
|
"attentions",
|
|
{
|
|
{"to_k", "k"},
|
|
{"to_q", "q"},
|
|
{"to_v", "v"},
|
|
{"to_out_0", "proj_out"},
|
|
{"group_norm", "norm"},
|
|
},
|
|
},
|
|
{
|
|
"resnets",
|
|
{
|
|
{"conv1", "in_layers_2"},
|
|
{"conv2", "out_layers_3"},
|
|
{"norm1", "in_layers_0"},
|
|
{"norm2", "out_layers_0"},
|
|
{"time_emb_proj", "emb_layers_1"},
|
|
{"conv_shortcut", "skip_connection"},
|
|
},
|
|
},
|
|
};
|
|
|
|
std::unordered_map<std::string, std::unordered_map<std::string, std::string>> suffix_conversion_dot = {
|
|
{
|
|
"attentions",
|
|
{
|
|
{"to_k", "k"},
|
|
{"to_q", "q"},
|
|
{"to_v", "v"},
|
|
{"to_out.0", "proj_out"},
|
|
{"group_norm", "norm"},
|
|
},
|
|
},
|
|
{
|
|
"resnets",
|
|
{
|
|
{"conv1", "in_layers.2"},
|
|
{"conv2", "out_layers.3"},
|
|
{"norm1", "in_layers.0"},
|
|
{"norm2", "out_layers.0"},
|
|
{"time_emb_proj", "emb_layers.1"},
|
|
{"conv_shortcut", "skip_connection"},
|
|
},
|
|
},
|
|
};
|
|
|
|
std::string convert_diffusers_name_to_compvis(const std::string& key, char seq) {
|
|
std::vector<std::string> m;
|
|
|
|
auto match = [](std::vector<std::string>& match_list, const std::regex& regex, const std::string& key) {
|
|
auto r = std::smatch{};
|
|
if (!std::regex_match(key, r, regex)) {
|
|
return false;
|
|
}
|
|
|
|
match_list.clear();
|
|
for (size_t i = 1; i < r.size(); ++i) {
|
|
match_list.push_back(r.str(i));
|
|
}
|
|
return true;
|
|
};
|
|
|
|
std::unordered_map<std::string, std::unordered_map<std::string, std::string>> suffix_conversion;
|
|
if (seq == '_') {
|
|
suffix_conversion = suffix_conversion_underline;
|
|
} else {
|
|
suffix_conversion = suffix_conversion_dot;
|
|
}
|
|
|
|
auto get_converted_suffix = [&suffix_conversion](const std::string& outer_key, const std::string& inner_key) {
|
|
auto outer_iter = suffix_conversion.find(outer_key);
|
|
if (outer_iter != suffix_conversion.end()) {
|
|
auto inner_iter = outer_iter->second.find(inner_key);
|
|
if (inner_iter != outer_iter->second.end()) {
|
|
return inner_iter->second;
|
|
}
|
|
}
|
|
return inner_key;
|
|
};
|
|
|
|
// unet
|
|
if (match(m, std::regex(format("unet%cconv_in(.*)", seq)), key)) {
|
|
return format("model%cdiffusion_model%cinput_blocks%c0%c0", seq, seq, seq, seq) + m[0];
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cconv%cout(.*)", seq, seq)), key)) {
|
|
return format("model%cdiffusion_model%cout%c2", seq, seq, seq) + m[0];
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cconv_norm_out(.*)", seq)), key)) {
|
|
return format("model%cdiffusion_model%cout%c0", seq, seq, seq) + m[0];
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%ctime_embedding%clinear_(\\d+)(.*)", seq, seq)), key)) {
|
|
return format("model%cdiffusion_model%ctime_embed%c", seq, seq, seq) + std::to_string(std::stoi(m[0]) * 2 - 2) + m[1];
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cdown_blocks%c(\\d+)%c(attentions|resnets)%c(\\d+)%c(.+)", seq, seq, seq, seq, seq)), key)) {
|
|
std::string suffix = get_converted_suffix(m[1], m[3]);
|
|
// LOG_DEBUG("%s %s %s %s", m[0].c_str(), m[1].c_str(), m[2].c_str(), m[3].c_str());
|
|
return format("model%cdiffusion_model%cinput_blocks%c", seq, seq, seq) + std::to_string(1 + std::stoi(m[0]) * 3 + std::stoi(m[2])) + seq +
|
|
(m[1] == "attentions" ? "1" : "0") + seq + suffix;
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cmid_block%c(attentions|resnets)%c(\\d+)%c(.+)", seq, seq, seq, seq)), key)) {
|
|
std::string suffix = get_converted_suffix(m[0], m[2]);
|
|
return format("model%cdiffusion_model%cmiddle_block%c", seq, seq, seq) + (m[0] == "attentions" ? "1" : std::to_string(std::stoi(m[1]) * 2)) +
|
|
seq + suffix;
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cup_blocks%c(\\d+)%c(attentions|resnets)%c(\\d+)%c(.+)", seq, seq, seq, seq, seq)), key)) {
|
|
std::string suffix = get_converted_suffix(m[1], m[3]);
|
|
return format("model%cdiffusion_model%coutput_blocks%c", seq, seq, seq) + std::to_string(std::stoi(m[0]) * 3 + std::stoi(m[2])) + seq +
|
|
(m[1] == "attentions" ? "1" : "0") + seq + suffix;
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cdown_blocks%c(\\d+)%cdownsamplers%c0%cconv", seq, seq, seq, seq, seq)), key)) {
|
|
return format("model%cdiffusion_model%cinput_blocks%c", seq, seq, seq) + std::to_string(3 + std::stoi(m[0]) * 3) + seq + "0" + seq + "op";
|
|
}
|
|
|
|
if (match(m, std::regex(format("unet%cup_blocks%c(\\d+)%cupsamplers%c0%cconv", seq, seq, seq, seq, seq)), key)) {
|
|
return format("model%cdiffusion_model%coutput_blocks%c", seq, seq, seq) + std::to_string(2 + std::stoi(m[0]) * 3) + seq +
|
|
(std::stoi(m[0]) > 0 ? "2" : "1") + seq + "conv";
|
|
}
|
|
|
|
// clip
|
|
if (match(m, std::regex(format("te%ctext_model%cencoder%clayers%c(\\d+)%c(.+)", seq, seq, seq, seq, seq)), key)) {
|
|
return format("cond_stage_model%ctransformer%ctext_model%cencoder%clayers%c", seq, seq, seq, seq, seq) + m[0] + seq + m[1];
|
|
}
|
|
|
|
if (match(m, std::regex(format("te%ctext_model(.*)", seq)), key)) {
|
|
return format("cond_stage_model%ctransformer%ctext_model", seq, seq) + m[0];
|
|
}
|
|
|
|
// vae
|
|
if (match(m, std::regex(format("vae%c(.*)%cconv_norm_out(.*)", seq, seq)), key)) {
|
|
return format("first_stage_model%c%s%cnorm_out%s", seq, m[0].c_str(), seq, m[1].c_str());
|
|
}
|
|
|
|
if (match(m, std::regex(format("vae%c(.*)%cmid_block%c(attentions|resnets)%c(\\d+)%c(.+)", seq, seq, seq, seq, seq)), key)) {
|
|
std::string suffix;
|
|
std::string block_name;
|
|
if (m[1] == "attentions") {
|
|
block_name = "attn";
|
|
suffix = get_converted_suffix(m[1], m[3]);
|
|
} else {
|
|
block_name = "block";
|
|
suffix = m[3];
|
|
}
|
|
return format("first_stage_model%c%s%cmid%c%s_%d%c%s",
|
|
seq, m[0].c_str(), seq, seq, block_name.c_str(), std::stoi(m[2]) + 1, seq, suffix.c_str());
|
|
}
|
|
|
|
if (match(m, std::regex(format("vae%c(.*)%cup_blocks%c(\\d+)%cresnets%c(\\d+)%c(.+)", seq, seq, seq, seq, seq, seq)), key)) {
|
|
std::string suffix = m[3];
|
|
if (suffix == "conv_shortcut") {
|
|
suffix = "nin_shortcut";
|
|
}
|
|
return format("first_stage_model%c%s%cup%c%d%cblock%c%s%c%s",
|
|
seq, m[0].c_str(), seq, seq, 3 - std::stoi(m[1]), seq, seq, m[2].c_str(), seq, suffix.c_str());
|
|
}
|
|
|
|
if (match(m, std::regex(format("vae%c(.*)%cdown_blocks%c(\\d+)%cdownsamplers%c0%cconv", seq, seq, seq, seq, seq, seq)), key)) {
|
|
return format("first_stage_model%c%s%cdown%c%d%cdownsample%cconv",
|
|
seq, m[0].c_str(), seq, seq, std::stoi(m[1]), seq, seq);
|
|
}
|
|
|
|
if (match(m, std::regex(format("vae%c(.*)%cdown_blocks%c(\\d+)%cresnets%c(\\d+)%c(.+)", seq, seq, seq, seq, seq, seq)), key)) {
|
|
std::string suffix = m[3];
|
|
if (suffix == "conv_shortcut") {
|
|
suffix = "nin_shortcut";
|
|
}
|
|
return format("first_stage_model%c%s%cdown%c%d%cblock%c%s%c%s",
|
|
seq, m[0].c_str(), seq, seq, std::stoi(m[1]), seq, seq, m[2].c_str(), seq, suffix.c_str());
|
|
}
|
|
|
|
if (match(m, std::regex(format("vae%c(.*)%cup_blocks%c(\\d+)%cupsamplers%c0%cconv", seq, seq, seq, seq, seq, seq)), key)) {
|
|
return format("first_stage_model%c%s%cup%c%d%cupsample%cconv",
|
|
seq, m[0].c_str(), seq, seq, 3 - std::stoi(m[1]), seq, seq);
|
|
}
|
|
|
|
if (match(m, std::regex(format("vae%c(.*)", seq)), key)) {
|
|
return format("first_stage_model%c", seq) + m[0];
|
|
}
|
|
|
|
return key;
|
|
}
|
|
|
|
std::string convert_tensor_name(const std::string& name) {
|
|
std::string new_name;
|
|
if (starts_with(name, "cond_stage_model.model")) {
|
|
new_name = convert_open_clip_to_hf_clip(name);
|
|
} else if (starts_with(name, "first_stage_model.decoder")) {
|
|
new_name = convert_vae_decoder_name(name);
|
|
} else if (starts_with(name, "lora_")) { // for lora
|
|
size_t pos = name.find('.');
|
|
if (pos != std::string::npos) {
|
|
std::string name_without_network_parts = name.substr(5, pos - 5);
|
|
std::string network_part = name.substr(pos + 1);
|
|
// LOG_DEBUG("%s %s", name_without_network_parts.c_str(), network_part.c_str());
|
|
std::string new_key = convert_diffusers_name_to_compvis(name_without_network_parts, '_');
|
|
if (new_key.empty()) {
|
|
new_name = name;
|
|
} else {
|
|
new_name = "lora." + new_key + "." + network_part;
|
|
}
|
|
} else {
|
|
new_name = name;
|
|
}
|
|
} else if (starts_with(name, "unet") || starts_with(name, "vae") || starts_with(name, "te")) { // for diffuser
|
|
size_t pos = name.find_last_of('.');
|
|
if (pos != std::string::npos) {
|
|
std::string name_without_network_parts = name.substr(0, pos);
|
|
std::string network_part = name.substr(pos + 1);
|
|
// LOG_DEBUG("%s %s", name_without_network_parts.c_str(), network_part.c_str());
|
|
std::string new_key = convert_diffusers_name_to_compvis(name_without_network_parts, '.');
|
|
if (new_key.empty()) {
|
|
new_name = name;
|
|
} else {
|
|
new_name = new_key + "." + network_part;
|
|
}
|
|
} else {
|
|
new_name = name;
|
|
}
|
|
} else {
|
|
new_name = name;
|
|
}
|
|
// if (new_name != name) {
|
|
// LOG_DEBUG("%s => %s", name.c_str(), new_name.c_str());
|
|
// }
|
|
return new_name;
|
|
}
|
|
|
|
void preprocess_tensor(TensorStorage tensor_storage,
|
|
std::vector<TensorStorage>& processed_tensor_storages) {
|
|
std::vector<TensorStorage> result;
|
|
std::string new_name = convert_tensor_name(tensor_storage.name);
|
|
|
|
// convert unet transformer linear to conv2d 1x1
|
|
if (starts_with(new_name, "model.diffusion_model.") &&
|
|
(ends_with(new_name, "proj_in.weight") || ends_with(new_name, "proj_out.weight"))) {
|
|
tensor_storage.unsqueeze();
|
|
}
|
|
|
|
// convert vae attn block linear to conv2d 1x1
|
|
if (starts_with(new_name, "first_stage_model.") && new_name.find("attn_1") != std::string::npos) {
|
|
tensor_storage.unsqueeze();
|
|
}
|
|
|
|
tensor_storage.name = new_name;
|
|
|
|
if (starts_with(new_name, "cond_stage_model.transformer.text_model.encoder.layers.") &&
|
|
ends_with(new_name, "attn.in_proj_weight")) {
|
|
size_t prefix_size = new_name.find("attn.in_proj_weight");
|
|
std::string prefix = new_name.substr(0, prefix_size);
|
|
|
|
std::vector<TensorStorage> chunks = tensor_storage.chunk(3);
|
|
chunks[0].name = prefix + "self_attn.q_proj.weight";
|
|
chunks[1].name = prefix + "self_attn.k_proj.weight";
|
|
chunks[2].name = prefix + "self_attn.v_proj.weight";
|
|
|
|
processed_tensor_storages.insert(processed_tensor_storages.end(), chunks.begin(), chunks.end());
|
|
|
|
} else if (starts_with(new_name, "cond_stage_model.transformer.text_model.encoder.layers.") &&
|
|
ends_with(new_name, "attn.in_proj_bias")) {
|
|
size_t prefix_size = new_name.find("attn.in_proj_bias");
|
|
std::string prefix = new_name.substr(0, prefix_size);
|
|
|
|
std::vector<TensorStorage> chunks = tensor_storage.chunk(3);
|
|
chunks[0].name = prefix + "self_attn.q_proj.bias";
|
|
chunks[1].name = prefix + "self_attn.k_proj.bias";
|
|
chunks[2].name = prefix + "self_attn.v_proj.bias";
|
|
|
|
processed_tensor_storages.insert(processed_tensor_storages.end(), chunks.begin(), chunks.end());
|
|
} else {
|
|
processed_tensor_storages.push_back(tensor_storage);
|
|
}
|
|
}
|
|
|
|
float bf16_to_f32(uint16_t bfloat16) {
|
|
uint32_t val_bits = (static_cast<uint32_t>(bfloat16) << 16);
|
|
return *reinterpret_cast<float*>(&val_bits);
|
|
}
|
|
|
|
void bf16_to_f32_vec(uint16_t* src, float* dst, int64_t n) {
|
|
// support inplace op
|
|
for (int64_t i = n - 1; i >= 0; i--) {
|
|
dst[i] = bf16_to_f32(src[i]);
|
|
}
|
|
}
|
|
|
|
void convert_tensor(void* src, ggml_type src_type, void* dst, ggml_type dst_type, int n) {
|
|
if (src_type == dst_type) {
|
|
size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type);
|
|
memcpy(((char*)dst), ((char*)src), nbytes);
|
|
} else if (src_type == GGML_TYPE_F32) {
|
|
if (dst_type == GGML_TYPE_F16) {
|
|
ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n);
|
|
} else {
|
|
int64_t hist[16];
|
|
ggml_quantize_chunk(dst_type, (float*)src, dst, 0, n, hist);
|
|
}
|
|
} else if (dst_type == GGML_TYPE_F32) {
|
|
if (src_type == GGML_TYPE_F16) {
|
|
ggml_fp16_to_fp32_row((ggml_fp16_t*)src, (float*)dst, n);
|
|
} else {
|
|
auto qtype = ggml_internal_get_type_traits(src_type);
|
|
if (qtype.to_float == NULL) {
|
|
throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available",
|
|
ggml_type_name(src_type)));
|
|
}
|
|
qtype.to_float(src, (float*)dst, n);
|
|
}
|
|
} else {
|
|
// src_type == GGML_TYPE_F16 => dst_type is quantized
|
|
// src_type is quantized => dst_type == GGML_TYPE_F16 or dst_type is quantized
|
|
auto qtype = ggml_internal_get_type_traits(src_type);
|
|
if (qtype.to_float == NULL) {
|
|
throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available",
|
|
ggml_type_name(src_type)));
|
|
}
|
|
std::vector<char> buf;
|
|
buf.resize(sizeof(float) * n);
|
|
char* src_data_f32 = buf.data();
|
|
qtype.to_float(src, (float*)src_data_f32, n);
|
|
if (dst_type == GGML_TYPE_F16) {
|
|
ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n);
|
|
} else {
|
|
int64_t hist[16];
|
|
ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, n, hist);
|
|
}
|
|
}
|
|
}
|
|
|
|
/*================================================= ModelLoader ==================================================*/
|
|
|
|
// ported from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py#L16
|
|
std::map<char, int> unicode_to_byte() {
|
|
std::map<int, char> byte_to_unicode;
|
|
|
|
// List of utf-8 byte ranges
|
|
for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
|
|
byte_to_unicode[b] = static_cast<char>(b);
|
|
}
|
|
|
|
for (int b = 49825; b <= 49836; ++b) {
|
|
byte_to_unicode[b] = static_cast<char>(b);
|
|
}
|
|
|
|
for (int b = 49838; b <= 50111; ++b) {
|
|
byte_to_unicode[b] = static_cast<char>(b);
|
|
}
|
|
// printf("%d %d %d %d\n", static_cast<int>('¡'), static_cast<int>('¬'), static_cast<int>('®'), static_cast<int>('ÿ'));
|
|
// exit(1);
|
|
|
|
int n = 0;
|
|
for (int b = 0; b < 256; ++b) {
|
|
if (byte_to_unicode.find(b) == byte_to_unicode.end()) {
|
|
byte_to_unicode[b] = static_cast<char>(256 + n);
|
|
n++;
|
|
}
|
|
}
|
|
|
|
// byte_encoder = bytes_to_unicode()
|
|
// byte_decoder = {v: k for k, v in byte_encoder.items()}
|
|
std::map<char, int> byte_decoder;
|
|
|
|
for (const auto& entry : byte_to_unicode) {
|
|
byte_decoder[entry.second] = entry.first;
|
|
}
|
|
|
|
byte_to_unicode.clear();
|
|
|
|
return byte_decoder;
|
|
}
|
|
|
|
bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
|
|
file_paths_.push_back(file_path);
|
|
return true;
|
|
}
|
|
|
|
bool ModelLoader::init_from_files(const std::vector<std::string>& file_paths) {
|
|
for (auto& file_path : file_paths) {
|
|
if (!init_from_file(file_path)) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
SDVersion ModelLoader::get_sd_version() {
|
|
TensorStorage token_embedding_weight;
|
|
for (auto& tensor_storage : tensor_storages) {
|
|
if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
|
|
tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
|
|
tensor_storage.name == "text_model.embeddings.token_embedding.weight" ||
|
|
tensor_storage.name == "te.text_model.embeddings.token_embedding.weight") {
|
|
token_embedding_weight = tensor_storage;
|
|
break;
|
|
}
|
|
}
|
|
if (token_embedding_weight.ne[0] == 768) {
|
|
return VERSION_1_x;
|
|
} else if (token_embedding_weight.ne[0] == 1024) {
|
|
return VERSION_2_x;
|
|
}
|
|
return VERSION_COUNT;
|
|
}
|
|
|
|
ggml_type ModelLoader::get_sd_wtype() {
|
|
for (auto& tensor_storage : tensor_storages) {
|
|
if (is_unused_tensor(tensor_storage.name)) {
|
|
continue;
|
|
}
|
|
|
|
if (tensor_storage.name.find(".weight") != std::string::npos &&
|
|
tensor_storage.name.find("time_embed") != std::string::npos) {
|
|
return tensor_storage.type;
|
|
}
|
|
}
|
|
return GGML_TYPE_COUNT;
|
|
}
|
|
|
|
bool ModelLoader::load_vocab(on_new_token_cb_t on_new_token_cb) {
|
|
char* vocab_buffer = reinterpret_cast<char*>(vocab_json);
|
|
nlohmann::json vocab = nlohmann::json::parse(vocab_buffer);
|
|
std::map<char, int> decoder = unicode_to_byte();
|
|
for (auto& it : vocab.items()) {
|
|
int token_id = it.value();
|
|
std::string token_str = it.key();
|
|
std::string token = "";
|
|
for (char c : token_str) {
|
|
token += decoder[c];
|
|
}
|
|
on_new_token_cb(token, token_id);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb) {
|
|
bool success = true;
|
|
for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
|
|
std::string file_path = file_paths_[file_index];
|
|
LOG_DEBUG("loading tensors from %s", file_path.c_str());
|
|
|
|
std::ifstream file(file_path, std::ios::binary);
|
|
if (!file.is_open()) {
|
|
LOG_ERROR("failed to open '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
bool is_zip = false;
|
|
for (auto& tensor_storage : tensor_storages) {
|
|
if (tensor_storage.index_in_zip >= 0) {
|
|
is_zip = true;
|
|
break;
|
|
}
|
|
}
|
|
|
|
struct zip_t* zip = NULL;
|
|
if (is_zip) {
|
|
zip = zip_open(file_path.c_str(), 0, 'r');
|
|
if (zip == NULL) {
|
|
LOG_ERROR("failed to open zip '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
}
|
|
|
|
std::vector<uint8_t> read_buffer;
|
|
std::vector<uint8_t> convert_buffer;
|
|
|
|
auto read_data = [&](const TensorStorage& tensor_storage, char* buf, size_t n) {
|
|
if (zip != NULL) {
|
|
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
|
|
size_t entry_size = zip_entry_size(zip);
|
|
if (entry_size != n) {
|
|
read_buffer.resize(entry_size);
|
|
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
|
|
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
|
|
} else {
|
|
zip_entry_noallocread(zip, (void*)buf, n);
|
|
}
|
|
zip_entry_close(zip);
|
|
} else {
|
|
file.seekg(tensor_storage.offset);
|
|
file.read(buf, n);
|
|
if (!file) {
|
|
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
};
|
|
|
|
std::vector<TensorStorage> processed_tensor_storages;
|
|
for (auto& tensor_storage : tensor_storages) {
|
|
if (tensor_storage.file_index != file_index) {
|
|
continue;
|
|
}
|
|
|
|
// LOG_DEBUG("%s", name.c_str());
|
|
|
|
if (is_unused_tensor(tensor_storage.name)) {
|
|
continue;
|
|
}
|
|
|
|
preprocess_tensor(tensor_storage, processed_tensor_storages);
|
|
}
|
|
|
|
for (auto& tensor_storage : processed_tensor_storages) {
|
|
// LOG_DEBUG("%s", name.c_str());
|
|
|
|
ggml_tensor* dst_tensor = NULL;
|
|
|
|
success = on_new_tensor_cb(tensor_storage, &dst_tensor);
|
|
if (!success) {
|
|
LOG_WARN("process tensor failed: '%s'", tensor_storage.name.c_str());
|
|
break;
|
|
}
|
|
|
|
if (dst_tensor == NULL) {
|
|
continue;
|
|
}
|
|
|
|
ggml_backend_t backend = ggml_get_backend(dst_tensor);
|
|
|
|
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
|
|
|
|
if (backend == NULL || ggml_backend_is_cpu(backend)) {
|
|
// for the CPU and Metal backend, we can copy directly into the tensor
|
|
if (tensor_storage.type == dst_tensor->type) {
|
|
GGML_ASSERT(ggml_nbytes(dst_tensor) == nbytes_to_read);
|
|
read_data(tensor_storage, (char*)dst_tensor->data, nbytes_to_read);
|
|
|
|
if (tensor_storage.is_bf16) {
|
|
// inplace op
|
|
bf16_to_f32_vec((uint16_t*)dst_tensor->data, (float*)dst_tensor->data, tensor_storage.nelements());
|
|
}
|
|
} else {
|
|
read_buffer.resize(tensor_storage.nbytes());
|
|
read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
|
|
|
|
if (tensor_storage.is_bf16) {
|
|
// inplace op
|
|
bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
|
|
}
|
|
|
|
convert_tensor((void*)read_buffer.data(), tensor_storage.type, dst_tensor->data,
|
|
dst_tensor->type, (int)tensor_storage.nelements());
|
|
}
|
|
} else {
|
|
read_buffer.resize(tensor_storage.nbytes());
|
|
read_data(tensor_storage, (char*)read_buffer.data(), nbytes_to_read);
|
|
|
|
if (tensor_storage.is_bf16) {
|
|
// inplace op
|
|
bf16_to_f32_vec((uint16_t*)read_buffer.data(), (float*)read_buffer.data(), tensor_storage.nelements());
|
|
}
|
|
|
|
if (tensor_storage.type == dst_tensor->type) {
|
|
// copy to device memory
|
|
ggml_backend_tensor_set(dst_tensor, read_buffer.data(), 0, ggml_nbytes(dst_tensor));
|
|
} else {
|
|
// convert first, then copy to device memory
|
|
convert_buffer.resize(ggml_nbytes(dst_tensor));
|
|
convert_tensor((void*)read_buffer.data(), tensor_storage.type,
|
|
(void*)convert_buffer.data(), dst_tensor->type,
|
|
(int)tensor_storage.nelements());
|
|
ggml_backend_tensor_set(dst_tensor, convert_buffer.data(), 0, ggml_nbytes(dst_tensor));
|
|
}
|
|
}
|
|
}
|
|
|
|
if (zip != NULL) {
|
|
zip_close(zip);
|
|
}
|
|
|
|
if (!success) {
|
|
break;
|
|
}
|
|
}
|
|
return success;
|
|
}
|
|
|
|
int64_t ModelLoader::cal_mem_size() {
|
|
int64_t mem_size = 0;
|
|
for (auto& tensor_storage : tensor_storages) {
|
|
if (is_unused_tensor(tensor_storage.name)) {
|
|
continue;
|
|
}
|
|
|
|
mem_size += tensor_storage.nbytes();
|
|
mem_size += GGML_MEM_ALIGN * 2; // for lora alphas
|
|
}
|
|
|
|
return mem_size + 10 * 1024 * 1024;
|
|
}
|
|
|
|
/*================================================= GGUFModelLoader ==================================================*/
|
|
|
|
bool GGUFModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
|
|
LOG_INFO("loading model from '%s'", file_path.c_str());
|
|
ModelLoader::init_from_file(file_path, prefix);
|
|
size_t file_index = file_paths_.size() - 1;
|
|
|
|
gguf_context* ctx_gguf_ = NULL;
|
|
ggml_context* ctx_meta_ = NULL;
|
|
ctx_gguf_ = gguf_init_from_file(file_path.c_str(), {true, &ctx_meta_});
|
|
if (!ctx_gguf_) {
|
|
LOG_ERROR("failed to open '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
int n_tensors = gguf_get_n_tensors(ctx_gguf_);
|
|
|
|
size_t total_size = 0;
|
|
size_t data_offset = gguf_get_data_offset(ctx_gguf_);
|
|
for (int i = 0; i < n_tensors; i++) {
|
|
std::string name = gguf_get_tensor_name(ctx_gguf_, i);
|
|
struct ggml_tensor* dummy = ggml_get_tensor(ctx_meta_, name.c_str());
|
|
size_t offset = data_offset + gguf_get_tensor_offset(ctx_gguf_, i);
|
|
|
|
// LOG_DEBUG("%s", name.c_str());
|
|
|
|
TensorStorage tensor_storage(prefix + name, dummy->type, dummy->ne, dummy->n_dims, file_index, offset);
|
|
|
|
GGML_ASSERT(ggml_nbytes(dummy) == tensor_storage.nbytes());
|
|
|
|
tensor_storages.push_back(tensor_storage);
|
|
}
|
|
|
|
gguf_free(ctx_gguf_);
|
|
ggml_free(ctx_meta_);
|
|
|
|
return true;
|
|
}
|
|
|
|
/*================================================= SafeTensorsModelLoader ==================================================*/
|
|
|
|
#define ST_HEADER_SIZE_LEN 8
|
|
|
|
ggml_type str_to_ggml_type(const std::string& dtype) {
|
|
ggml_type ttype = GGML_TYPE_COUNT;
|
|
if (dtype == "F16") {
|
|
ttype = GGML_TYPE_F16;
|
|
} else if (dtype == "BF16") {
|
|
ttype = GGML_TYPE_F32;
|
|
} else if (dtype == "F32") {
|
|
ttype = GGML_TYPE_F32;
|
|
}
|
|
return ttype;
|
|
}
|
|
|
|
// https://huggingface.co/docs/safetensors/index
|
|
bool SafeTensorsModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
|
|
ModelLoader::init_from_file(file_path, prefix);
|
|
size_t file_index = file_paths_.size() - 1;
|
|
std::ifstream file(file_path, std::ios::binary);
|
|
if (!file.is_open()) {
|
|
LOG_ERROR("failed to open '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
// get file size
|
|
file.seekg(0, file.end);
|
|
size_t file_size_ = file.tellg();
|
|
file.seekg(0, file.beg);
|
|
|
|
// read header size
|
|
if (file_size_ <= ST_HEADER_SIZE_LEN) {
|
|
LOG_ERROR("invalid safetensor file '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
uint8_t header_size_buf[ST_HEADER_SIZE_LEN];
|
|
file.read((char*)header_size_buf, ST_HEADER_SIZE_LEN);
|
|
if (!file) {
|
|
LOG_ERROR("read safetensors header size failed: '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
size_t header_size_ = read_u64(header_size_buf);
|
|
if (header_size_ >= file_size_) {
|
|
LOG_ERROR("invalid safetensor file '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
// read header
|
|
std::vector<char> header_buf;
|
|
header_buf.resize(header_size_ + 1);
|
|
header_buf[header_size_] = '\0';
|
|
file.read(header_buf.data(), header_size_);
|
|
if (!file) {
|
|
LOG_ERROR("read safetensors header failed: '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
|
|
nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
|
|
|
|
for (auto& item : header_.items()) {
|
|
std::string name = item.key();
|
|
nlohmann::json tensor_info = item.value();
|
|
// LOG_DEBUG("%s %s\n", name.c_str(), tensor_info.dump().c_str());
|
|
|
|
if (name == "__metadata__") {
|
|
continue;
|
|
}
|
|
|
|
if (is_unused_tensor(name)) {
|
|
continue;
|
|
}
|
|
|
|
std::string dtype = tensor_info["dtype"];
|
|
nlohmann::json shape = tensor_info["shape"];
|
|
|
|
size_t begin = tensor_info["data_offsets"][0].get<size_t>();
|
|
size_t end = tensor_info["data_offsets"][1].get<size_t>();
|
|
|
|
ggml_type type = str_to_ggml_type(dtype);
|
|
if (type == GGML_TYPE_COUNT) {
|
|
LOG_ERROR("unsupported dtype '%s'", dtype.c_str());
|
|
return false;
|
|
}
|
|
|
|
if (shape.size() > 4) {
|
|
LOG_ERROR("invalid tensor '%s'", name.c_str());
|
|
return false;
|
|
}
|
|
|
|
int n_dims = (int)shape.size();
|
|
int64_t ne[4] = {1, 1, 1, 1};
|
|
for (int i = 0; i < n_dims; i++) {
|
|
ne[i] = shape[i].get<int64_t>();
|
|
}
|
|
|
|
TensorStorage tensor_storage(prefix + name, type, ne, n_dims, file_index, ST_HEADER_SIZE_LEN + header_size_ + begin);
|
|
|
|
tensor_storage.reverse_ne();
|
|
|
|
size_t tensor_data_size = end - begin;
|
|
|
|
if (dtype == "BF16") {
|
|
tensor_storage.is_bf16 = true;
|
|
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2);
|
|
} else {
|
|
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size);
|
|
}
|
|
|
|
tensor_storages.push_back(tensor_storage);
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/*================================================= DiffusersModelLoader ==================================================*/
|
|
|
|
bool DiffusersModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
|
|
if (!is_directory(file_path)) {
|
|
return SafeTensorsModelLoader::init_from_file(file_path, prefix);
|
|
}
|
|
std::string unet_path = path_join(file_path, "unet/diffusion_pytorch_model.safetensors");
|
|
std::string vae_path = path_join(file_path, "vae/diffusion_pytorch_model.safetensors");
|
|
std::string clip_path = path_join(file_path, "text_encoder/model.safetensors");
|
|
|
|
if (!SafeTensorsModelLoader::init_from_file(unet_path, "unet.")) {
|
|
return false;
|
|
}
|
|
if (!SafeTensorsModelLoader::init_from_file(vae_path, "vae.")) {
|
|
return false;
|
|
}
|
|
if (!SafeTensorsModelLoader::init_from_file(clip_path, "te.")) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
/*================================================= CkptModelLoader ==================================================*/
|
|
|
|
// $ python -m pickletools sd-v1-4/archive/data.pkl | head -n 100
|
|
// 0: \x80 PROTO 2
|
|
// 2: } EMPTY_DICT
|
|
// 3: q BINPUT 0
|
|
// 5: ( MARK
|
|
// 6: X BINUNICODE 'epoch'
|
|
// 16: q BINPUT 1
|
|
// 18: K BININT1 6
|
|
// 20: X BINUNICODE 'global_step'
|
|
// 36: q BINPUT 2
|
|
// 38: J BININT 470000
|
|
// 43: X BINUNICODE 'pytorch-lightning_version'
|
|
// 73: q BINPUT 3
|
|
// 75: X BINUNICODE '1.4.2'
|
|
// 85: q BINPUT 4
|
|
// 87: X BINUNICODE 'state_dict'
|
|
// 102: q BINPUT 5
|
|
// 104: } EMPTY_DICT
|
|
// 105: q BINPUT 6
|
|
// 107: ( MARK
|
|
// 108: X BINUNICODE 'betas'
|
|
// 118: q BINPUT 7
|
|
// 120: c GLOBAL 'torch._utils _rebuild_tensor_v2'
|
|
// 153: q BINPUT 8
|
|
// 155: ( MARK
|
|
// 156: ( MARK
|
|
// 157: X BINUNICODE 'storage'
|
|
// 169: q BINPUT 9
|
|
// 171: c GLOBAL 'torch FloatStorage'
|
|
// 191: q BINPUT 10
|
|
// 193: X BINUNICODE '0'
|
|
// 199: q BINPUT 11
|
|
// 201: X BINUNICODE 'cpu'
|
|
// 209: q BINPUT 12
|
|
// 211: M BININT2 1000
|
|
// 214: t TUPLE (MARK at 156)
|
|
// 215: q BINPUT 13
|
|
// 217: Q BINPERSID
|
|
// 218: K BININT1 0
|
|
// 220: M BININT2 1000
|
|
// ...............................
|
|
// 3201: q BINPUT 250
|
|
// 3203: R REDUCE
|
|
// 3204: q BINPUT 251
|
|
// 3206: X BINUNICODE 'model.diffusion_model.input_blocks.1.1.proj_in.weight'
|
|
// 3264: q BINPUT 252
|
|
// 3266: h BINGET 8
|
|
// 3268: ( MARK
|
|
// 3269: ( MARK
|
|
// 3270: h BINGET 9
|
|
// 3272: h BINGET 10
|
|
// 3274: X BINUNICODE '30'
|
|
// 3281: q BINPUT 253
|
|
// 3283: h BINGET 12
|
|
// 3285: J BININT 102400
|
|
// 3290: t TUPLE (MARK at 3269)
|
|
// 3291: q BINPUT 254
|
|
// 3293: Q BINPERSID
|
|
// 3294: K BININT1 0
|
|
// 3296: ( MARK
|
|
// 3297: M BININT2 320
|
|
// 3300: M BININT2 320
|
|
// 3303: K BININT1 1
|
|
// 3305: K BININT1 1
|
|
// 3307: t TUPLE (MARK at 3296)
|
|
// 3308: q BINPUT 255
|
|
// 3310: ( MARK
|
|
// 3311: M BININT2 320
|
|
// 3314: K BININT1 1
|
|
// 3316: K BININT1 1
|
|
// 3318: K BININT1 1
|
|
// 3320: t TUPLE (MARK at 3310)
|
|
// 3321: r LONG_BINPUT 256
|
|
// 3326: \x89 NEWFALSE
|
|
// 3327: h BINGET 16
|
|
// 3329: ) EMPTY_TUPLE
|
|
// 3330: R REDUCE
|
|
// 3331: r LONG_BINPUT 257
|
|
// 3336: t TUPLE (MARK at 3268)
|
|
// 3337: r LONG_BINPUT 258
|
|
// 3342: R REDUCE
|
|
// 3343: r LONG_BINPUT 259
|
|
// 3348: X BINUNICODE 'model.diffusion_model.input_blocks.1.1.proj_in.bias'
|
|
// 3404: r LONG_BINPUT 260
|
|
// 3409: h BINGET 8
|
|
// 3411: ( MARK
|
|
// 3412: ( MARK
|
|
// 3413: h BINGET 9
|
|
// 3415: h BINGET 10
|
|
// 3417: X BINUNICODE '31'
|
|
|
|
struct PickleTensorReader {
|
|
enum ReadPhase {
|
|
READ_NAME,
|
|
READ_DATA,
|
|
CHECK_SIZE,
|
|
READ_DIMENS
|
|
};
|
|
ReadPhase phase = READ_NAME;
|
|
size_t entry_size = 0;
|
|
int32_t nelements = 0;
|
|
|
|
TensorStorage tensor_storage;
|
|
|
|
static ggml_type global_type; // all pickle_tensors data type
|
|
static bool read_global_type;
|
|
|
|
bool read_int_value(uint32_t value) {
|
|
if (phase == CHECK_SIZE) {
|
|
if (entry_size == value * ggml_type_size(tensor_storage.type)) {
|
|
nelements = value;
|
|
phase = READ_DIMENS;
|
|
return true;
|
|
} else {
|
|
phase = READ_NAME;
|
|
}
|
|
} else if (phase == READ_DIMENS) {
|
|
if (tensor_storage.n_dims + 1 > 4) { // too many dimens
|
|
phase = READ_NAME;
|
|
tensor_storage.n_dims = 0;
|
|
}
|
|
if (nelements % value == 0) {
|
|
tensor_storage.ne[tensor_storage.n_dims] = value;
|
|
tensor_storage.n_dims++;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
void read_global(const std::string& str) {
|
|
if (str == "FloatStorage") {
|
|
if (read_global_type) {
|
|
global_type = GGML_TYPE_F32;
|
|
read_global_type = false;
|
|
}
|
|
tensor_storage.type = GGML_TYPE_F32;
|
|
} else if (str == "HalfStorage") {
|
|
if (read_global_type) {
|
|
global_type = GGML_TYPE_F16;
|
|
read_global_type = false;
|
|
}
|
|
tensor_storage.type = GGML_TYPE_F16;
|
|
}
|
|
}
|
|
|
|
void read_string(const std::string& str, struct zip_t* zip, std::string dir) {
|
|
if (str == "storage") {
|
|
read_global_type = true;
|
|
} else if (str != "state_dict") {
|
|
if (phase == READ_DATA) {
|
|
std::string entry_name = dir + "data/" + std::string(str);
|
|
|
|
size_t i, n = zip_entries_total(zip);
|
|
for (i = 0; i < n; ++i) {
|
|
zip_entry_openbyindex(zip, i);
|
|
{
|
|
std::string name = zip_entry_name(zip);
|
|
if (name == entry_name) {
|
|
tensor_storage.index_in_zip = (int)i;
|
|
entry_size = zip_entry_size(zip);
|
|
zip_entry_close(zip);
|
|
break;
|
|
}
|
|
}
|
|
zip_entry_close(zip);
|
|
}
|
|
|
|
phase = entry_size > 0 ? CHECK_SIZE : READ_NAME;
|
|
}
|
|
if (!read_global_type && phase == READ_NAME) {
|
|
tensor_storage.name = str;
|
|
phase = READ_DATA;
|
|
tensor_storage.type = global_type;
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
ggml_type PickleTensorReader::global_type = GGML_TYPE_F32; // all pickle_tensors data type
|
|
bool PickleTensorReader::read_global_type = false;
|
|
|
|
int find_char(uint8_t* buffer, int len, char c) {
|
|
for (int pos = 0; pos < len; pos++) {
|
|
if (buffer[pos] == c) {
|
|
return pos;
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
#define MAX_STRING_BUFFER 512
|
|
|
|
bool CkptModelLoader::parse_data_pkl(uint8_t* buffer,
|
|
size_t buffer_size,
|
|
zip_t* zip,
|
|
std::string dir,
|
|
size_t file_index,
|
|
const std::string& prefix) {
|
|
uint8_t* buffer_end = buffer + buffer_size;
|
|
if (buffer[0] == 0x80) { // proto
|
|
if (buffer[1] != 2) {
|
|
LOG_ERROR("Unsupported protocol\n");
|
|
return false;
|
|
}
|
|
buffer += 2; // 0x80 and version
|
|
char string_buffer[MAX_STRING_BUFFER];
|
|
bool finish = false;
|
|
PickleTensorReader reader;
|
|
// read pickle binary file
|
|
while (!finish && buffer < buffer_end) {
|
|
uint8_t opcode = *buffer;
|
|
buffer++;
|
|
// https://github.com/python/cpython/blob/3.7/Lib/pickletools.py#L1048
|
|
// https://github.com/python/cpython/blob/main/Lib/pickle.py#L105
|
|
switch (opcode) {
|
|
case '}': // EMPTY_DICT = b'}' # push empty dict
|
|
break;
|
|
case ']': // EMPTY_LIST = b']' # push empty list
|
|
break;
|
|
// skip unused sections
|
|
case 'h': // BINGET = b'h' # " " " " " " ; " " 1-byte arg
|
|
case 'q': // BINPUT = b'q' # " " " " " ; " " 1-byte arg
|
|
case 'Q': // BINPERSID = b'Q' # " " " ; " " " " stack
|
|
buffer++;
|
|
break;
|
|
case 'r': // LONG_BINPUT = b'r' # " " " " " ; " " 4-byte arg
|
|
buffer += 4;
|
|
break;
|
|
case 0x95: // FRAME = b'\x95' # indicate the beginning of a new frame
|
|
buffer += 8;
|
|
break;
|
|
case 0x94: // MEMOIZE = b'\x94' # store top of the stack in memo
|
|
break;
|
|
case '(': // MARK = b'(' # push special markobject on stack
|
|
break;
|
|
case 'K': // BININT1 = b'K' # push 1-byte unsigned int
|
|
{
|
|
uint8_t value = *buffer;
|
|
if (reader.read_int_value(value)) {
|
|
buffer++;
|
|
}
|
|
buffer++;
|
|
} break;
|
|
case 'M': // BININT2 = b'M' # push 2-byte unsigned int
|
|
{
|
|
uint16_t value = read_short(buffer);
|
|
if (reader.read_int_value(value)) {
|
|
buffer++;
|
|
}
|
|
buffer += 2;
|
|
} break;
|
|
case 'J': // BININT = b'J' # push four-byte signed int
|
|
{
|
|
const int32_t value = read_int(buffer);
|
|
if (reader.read_int_value(value)) {
|
|
buffer++; // skip tuple after read num_elements
|
|
}
|
|
buffer += 4;
|
|
} break;
|
|
case 'X': // BINUNICODE = b'X' # " " " ; counted UTF-8 string argument
|
|
{
|
|
const int32_t len = read_int(buffer);
|
|
buffer += 4;
|
|
memset(string_buffer, 0, MAX_STRING_BUFFER);
|
|
if (len > MAX_STRING_BUFFER) {
|
|
LOG_WARN("tensor name very large");
|
|
}
|
|
memcpy(string_buffer, buffer, len < MAX_STRING_BUFFER ? len : (MAX_STRING_BUFFER - 1));
|
|
buffer += len;
|
|
reader.read_string(string_buffer, zip, dir);
|
|
} break;
|
|
case 0x8C: // SHORT_BINUNICODE = b'\x8c' # push short string; UTF-8 length < 256 bytes
|
|
{
|
|
const int8_t len = *buffer;
|
|
buffer++;
|
|
memset(string_buffer, 0, MAX_STRING_BUFFER);
|
|
memcpy(string_buffer, buffer, len);
|
|
buffer += len;
|
|
// printf("String: '%s'\n", string_buffer);
|
|
} break;
|
|
case 'c': // GLOBAL = b'c' # push self.find_class(modname, name); 2 string args
|
|
{
|
|
int len = find_char(buffer, MAX_STRING_BUFFER, '\n');
|
|
|
|
buffer += len + 1;
|
|
len = find_char(buffer, MAX_STRING_BUFFER, '\n');
|
|
|
|
memset(string_buffer, 0, MAX_STRING_BUFFER);
|
|
memcpy(string_buffer, buffer, len);
|
|
buffer += len + 1;
|
|
reader.read_global(string_buffer);
|
|
} break;
|
|
case 0x86: // TUPLE2 = b'\x86' # build 2-tuple from two topmost stack items
|
|
case 0x85: // TUPLE1 = b'\x85' # build 1-tuple from stack top
|
|
case 't': // TUPLE = b't' # build tuple from topmost stack items
|
|
if (reader.phase == PickleTensorReader::READ_DIMENS) {
|
|
reader.tensor_storage.reverse_ne();
|
|
reader.tensor_storage.file_index = file_index;
|
|
reader.tensor_storage.name = prefix + reader.tensor_storage.name;
|
|
tensor_storages.push_back(reader.tensor_storage);
|
|
// reset
|
|
reader = PickleTensorReader();
|
|
}
|
|
break;
|
|
case '.': // STOP = b'.' # every pickle ends with STOP
|
|
finish = true;
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
bool CkptModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
|
|
ModelLoader::init_from_file(file_path, prefix);
|
|
size_t file_index = file_paths_.size() - 1;
|
|
|
|
struct zip_t* zip = zip_open(file_path.c_str(), 0, 'r');
|
|
if (zip == NULL) {
|
|
LOG_ERROR("failed to open '%s'", file_path.c_str());
|
|
return false;
|
|
}
|
|
int n = (int)zip_entries_total(zip);
|
|
for (int i = 0; i < n; ++i) {
|
|
zip_entry_openbyindex(zip, i);
|
|
{
|
|
std::string name = zip_entry_name(zip);
|
|
size_t pos = name.find("data.pkl");
|
|
if (pos != std::string::npos) {
|
|
std::string dir = name.substr(0, pos);
|
|
void* pkl_data = NULL;
|
|
size_t pkl_size;
|
|
zip_entry_read(zip, &pkl_data, &pkl_size);
|
|
|
|
LOG_DEBUG("%lld", pkl_size);
|
|
|
|
parse_data_pkl((uint8_t*)pkl_data, pkl_size, zip, dir, file_index, prefix);
|
|
|
|
free(pkl_data);
|
|
}
|
|
}
|
|
zip_entry_close(zip);
|
|
}
|
|
zip_close(zip);
|
|
return true;
|
|
}
|
|
|
|
/*================================================= init_model_loader_from_file ==================================================*/
|
|
|
|
ModelLoader* init_model_loader_from_file(const std::string& file_path) {
|
|
ModelLoader* model_loader = NULL;
|
|
if (is_directory(file_path)) {
|
|
LOG_DEBUG("load %s using diffusers format", file_path.c_str());
|
|
model_loader = new DiffusersModelLoader();
|
|
} else if (ends_with(file_path, ".gguf")) {
|
|
LOG_DEBUG("load %s using gguf format", file_path.c_str());
|
|
model_loader = new GGUFModelLoader();
|
|
} else if (ends_with(file_path, ".safetensors")) {
|
|
LOG_DEBUG("load %s using safetensors format", file_path.c_str());
|
|
model_loader = new SafeTensorsModelLoader();
|
|
} else if (ends_with(file_path, ".ckpt")) {
|
|
LOG_DEBUG("load %s using checkpoint format", file_path.c_str());
|
|
model_loader = new CkptModelLoader();
|
|
} else {
|
|
LOG_DEBUG("unknown format %s", file_path.c_str());
|
|
return NULL;
|
|
}
|
|
if (!model_loader->init_from_file(file_path)) {
|
|
delete model_loader;
|
|
model_loader = NULL;
|
|
}
|
|
return model_loader;
|
|
} |