mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-08-26 15:15:16 +02:00
sd: sync to master-582-7023fc4 (#2150)
* sd: remove sampler alias handling from the C++ layer It's already handled at the Python layer. * sd: sync to master-580-7d33d4b * sd: sync to master-582-7023fc4
This commit is contained in:
@@ -679,7 +679,7 @@ llama-impl.o: src/llama-impl.cpp src/llama-impl.h
|
||||
budget.o: common/reasoning-budget.cpp common/reasoning-budget.h
|
||||
$(CXX) $(CXXFLAGS) -c $< -o $@
|
||||
|
||||
SDCPP_COMMON_BASENAMES := stable-diffusion.h stable-diffusion.cpp sample-cache.h sample-cache.cpp util.cpp upscaler.cpp model.cpp name_conversion.cpp tokenize_util.cpp thirdparty/zip.c
|
||||
SDCPP_COMMON_BASENAMES := stable-diffusion.h stable-diffusion.cpp sample-cache.h sample-cache.cpp util.cpp upscaler.cpp model.cpp name_conversion.cpp tokenizers/bpe_tokenizer.cpp tokenizers/bpe_tokenizer.h tokenizers/clip_tokenizer.cpp tokenizers/clip_tokenizer.h tokenizers/mistral_tokenizer.cpp tokenizers/mistral_tokenizer.h tokenizers/qwen2_tokenizer.cpp tokenizers/qwen2_tokenizer.h tokenizers/t5_unigram_tokenizer.cpp tokenizers/t5_unigram_tokenizer.h tokenizers/tokenizer.cpp tokenizers/tokenizer.h tokenizers/tokenize_util.cpp tokenizers/tokenize_util.h thirdparty/zip.c
|
||||
SDCPP_COMMON_SOURCES := $(foreach f,$(SDCPP_COMMON_BASENAMES),otherarch/sdcpp/$(f))
|
||||
SDCPP_FLAGS := -I./vendor/nlohmann
|
||||
|
||||
@@ -736,7 +736,7 @@ mainvk: tools/completion/completion.cpp common/arg.cpp common/speculative.cpp co
|
||||
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
fitparams: tools/fit-params/fit-params.cpp common/arg.cpp common/speculative.cpp common/ngram-cache.cpp common/ngram-map.cpp common/ngram-mod.cpp common/chat.cpp common/preset.cpp common/download.cpp build-info.h ggml_v4_vulkan.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_vulkan.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_vulkan.o ggml-vulkan.o ggml-vulkan-shaders.o ggml-repack.o $(OBJS_FULL) $(OBJS) lib/vulkan-1.lib
|
||||
$(CXX) $(CXXFLAGS) -DGGML_USE_VULKAN -DSD_USE_VULKAN $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
sdmain: $(SDCPP_COMMON_SOURCES) otherarch/sdcpp/main.cpp otherarch/sdcpp/image_metadata.cpp otherarch/sdcpp/common/log.cpp otherarch/sdcpp/common/media_io.cpp otherarch/sdcpp/common/common.cpp otherarch/sdcpp/version.cpp otherarch/sdcpp/vocab/vocab.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
sdmain: $(SDCPP_COMMON_SOURCES) otherarch/sdcpp/main.cpp otherarch/sdcpp/image_metadata.cpp otherarch/sdcpp/common/log.cpp otherarch/sdcpp/common/media_io.cpp otherarch/sdcpp/common/common.cpp otherarch/sdcpp/version.cpp otherarch/sdcpp/tokenizers/vocab/vocab.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(SDCPP_FLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
whispermain: otherarch/whispercpp/main.cpp otherarch/whispercpp/whisper.cpp build-info.h ggml.o ggml-cpu.o ggml-ops.o ggml-vec.o ggml-binops.o ggml-unops.o llama.o console.o llavaclip_default.o llava.o ggml-backend.o ggml-backend-meta.o ggml-backend-reg_default.o ggml-repack.o $(OBJS_FULL) $(OBJS)
|
||||
$(CXX) $(CXXFLAGS) $(filter-out %.h,$^) -o $@ $(LDFLAGS)
|
||||
|
||||
@@ -533,7 +533,7 @@ public:
|
||||
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_is_flux2(version)) {
|
||||
if (sd_version_uses_flux2_vae(version)) {
|
||||
dd_config.z_channels = 32;
|
||||
embed_dim = 32;
|
||||
} else {
|
||||
@@ -578,7 +578,7 @@ public:
|
||||
|
||||
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||
// z: [N, z_channels, h, w]
|
||||
if (sd_version_is_flux2(version)) {
|
||||
if (sd_version_uses_flux2_vae(version)) {
|
||||
// [N, C*p*p, h, w] -> [N, C, h*p, w*p]
|
||||
int64_t p = 2;
|
||||
|
||||
@@ -617,7 +617,7 @@ public:
|
||||
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]
|
||||
}
|
||||
if (sd_version_is_flux2(version)) {
|
||||
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]
|
||||
@@ -640,7 +640,7 @@ public:
|
||||
|
||||
int get_encoder_output_channels() {
|
||||
int factor = dd_config.double_z ? 2 : 1;
|
||||
if (sd_version_is_flux2(version)) {
|
||||
if (sd_version_uses_flux2_vae(version)) {
|
||||
return dd_config.z_channels * 4;
|
||||
}
|
||||
return dd_config.z_channels * factor;
|
||||
@@ -673,7 +673,7 @@ struct AutoEncoderKL : public VAE {
|
||||
} else if (sd_version_is_flux(version) || sd_version_is_z_image(version)) {
|
||||
scale_factor = 0.3611f;
|
||||
shift_factor = 0.1159f;
|
||||
} else if (sd_version_is_flux2(version)) {
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
scale_factor = 1.0f;
|
||||
shift_factor = 0.f;
|
||||
}
|
||||
@@ -747,7 +747,7 @@ struct AutoEncoderKL : public VAE {
|
||||
}
|
||||
|
||||
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
|
||||
if (sd_version_is_flux2(version)) {
|
||||
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];
|
||||
@@ -758,7 +758,7 @@ struct AutoEncoderKL : public VAE {
|
||||
|
||||
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_is_flux2(version)) {
|
||||
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];
|
||||
@@ -804,7 +804,7 @@ struct AutoEncoderKL : public VAE {
|
||||
}
|
||||
|
||||
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
|
||||
if (sd_version_is_flux2(version)) {
|
||||
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;
|
||||
@@ -813,7 +813,7 @@ struct AutoEncoderKL : public VAE {
|
||||
}
|
||||
|
||||
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
|
||||
if (sd_version_is_flux2(version)) {
|
||||
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;
|
||||
|
||||
+1
-449
@@ -3,455 +3,7 @@
|
||||
|
||||
#include "ggml_extend.hpp"
|
||||
#include "model.h"
|
||||
#include "tokenize_util.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
/*================================================== CLIPTokenizer ===================================================*/
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::pair<int, std::u32string>> bytes_to_unicode() {
|
||||
std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;
|
||||
std::set<int> byte_set;
|
||||
for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
|
||||
byte_set.insert(b);
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
|
||||
}
|
||||
for (int b = 161; b <= 172; ++b) {
|
||||
byte_set.insert(b);
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
|
||||
}
|
||||
for (int b = 174; b <= 255; ++b) {
|
||||
byte_set.insert(b);
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
|
||||
}
|
||||
int n = 0;
|
||||
for (int b = 0; b < 256; ++b) {
|
||||
if (byte_set.find(b) == byte_set.end()) {
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
|
||||
++n;
|
||||
}
|
||||
}
|
||||
// LOG_DEBUG("byte_unicode_pairs %d", byte_unicode_pairs.size());
|
||||
return byte_unicode_pairs;
|
||||
}
|
||||
|
||||
// Ref: https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py
|
||||
|
||||
typedef std::function<bool(std::string&, std::vector<int32_t>&)> on_new_token_cb_t;
|
||||
|
||||
class CLIPTokenizer {
|
||||
private:
|
||||
std::map<int, std::u32string> byte_encoder;
|
||||
std::map<std::u32string, int> byte_decoder;
|
||||
std::map<std::u32string, int> encoder;
|
||||
std::map<int, std::u32string> decoder;
|
||||
std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
|
||||
std::regex pat;
|
||||
int encoder_len;
|
||||
int bpe_len;
|
||||
|
||||
std::vector<std::string> special_tokens;
|
||||
|
||||
public:
|
||||
const std::string UNK_TOKEN = "<|endoftext|>";
|
||||
const std::string BOS_TOKEN = "<|startoftext|>";
|
||||
const std::string EOS_TOKEN = "<|endoftext|>";
|
||||
const std::string PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
const int UNK_TOKEN_ID = 49407;
|
||||
const int BOS_TOKEN_ID = 49406;
|
||||
const int EOS_TOKEN_ID = 49407;
|
||||
const int PAD_TOKEN_ID = 49407;
|
||||
|
||||
private:
|
||||
static std::string strip(const std::string& str) {
|
||||
std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
|
||||
std::string::size_type end = str.find_last_not_of(" \t\n\r\v\f");
|
||||
|
||||
if (start == std::string::npos) {
|
||||
// String contains only whitespace characters
|
||||
return "";
|
||||
}
|
||||
|
||||
return str.substr(start, end - start + 1);
|
||||
}
|
||||
|
||||
static std::string whitespace_clean(std::string text) {
|
||||
text = std::regex_replace(text, std::regex(R"(\s+)"), " ");
|
||||
text = strip(text);
|
||||
return text;
|
||||
}
|
||||
|
||||
static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
|
||||
std::set<std::pair<std::u32string, std::u32string>> pairs;
|
||||
if (subwords.size() == 0) {
|
||||
return pairs;
|
||||
}
|
||||
std::u32string prev_subword = subwords[0];
|
||||
for (int i = 1; i < subwords.size(); i++) {
|
||||
std::u32string subword = subwords[i];
|
||||
std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
|
||||
pairs.insert(pair);
|
||||
prev_subword = subword;
|
||||
}
|
||||
return pairs;
|
||||
}
|
||||
|
||||
bool is_special_token(const std::string& token) {
|
||||
for (auto& special_token : special_tokens) {
|
||||
if (special_token == token) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
public:
|
||||
CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "")
|
||||
: PAD_TOKEN_ID(pad_token_id) {
|
||||
if (merges_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_clip_merges());
|
||||
}
|
||||
add_special_token("<|startoftext|>");
|
||||
add_special_token("<|endoftext|>");
|
||||
}
|
||||
|
||||
void load_from_merges(const std::string& merges_utf8_str) {
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
// printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size());
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
// for (auto & pair: byte_unicode_pairs) {
|
||||
// std::cout << pair.first << ": " << pair.second << std::endl;
|
||||
// }
|
||||
std::vector<std::u32string> merges;
|
||||
size_t start = 0;
|
||||
size_t pos;
|
||||
std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
|
||||
while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
|
||||
merges.push_back(merges_utf32_str.substr(start, pos - start));
|
||||
start = pos + 1;
|
||||
}
|
||||
// LOG_DEBUG("merges size %llu", merges.size());
|
||||
GGML_ASSERT(merges.size() == 48895);
|
||||
merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
// LOG_DEBUG("%s", utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
|
||||
// printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
|
||||
// utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
|
||||
}
|
||||
std::vector<std::u32string> vocab;
|
||||
for (const auto& pair : byte_unicode_pairs) {
|
||||
vocab.push_back(pair.second);
|
||||
}
|
||||
for (const auto& pair : byte_unicode_pairs) {
|
||||
vocab.push_back(pair.second + utf8_to_utf32("</w>"));
|
||||
}
|
||||
for (const auto& merge : merge_pairs) {
|
||||
vocab.push_back(merge.first + merge.second);
|
||||
}
|
||||
vocab.push_back(utf8_to_utf32("<|startoftext|>"));
|
||||
vocab.push_back(utf8_to_utf32("<|endoftext|>"));
|
||||
LOG_DEBUG("vocab size: %llu", vocab.size());
|
||||
int i = 0;
|
||||
for (const auto& token : vocab) {
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
i++;
|
||||
}
|
||||
encoder_len = i;
|
||||
|
||||
auto it = encoder.find(utf8_to_utf32("img</w>"));
|
||||
if (it != encoder.end()) {
|
||||
LOG_DEBUG("trigger word img already in vocab");
|
||||
} else {
|
||||
LOG_DEBUG("trigger word img not in vocab yet");
|
||||
}
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
};
|
||||
|
||||
void add_token(const std::string& text) {
|
||||
std::u32string token = utf8_to_utf32(text);
|
||||
auto it = encoder.find(token);
|
||||
if (it != encoder.end()) {
|
||||
encoder[token] = encoder_len;
|
||||
decoder[encoder_len] = token;
|
||||
encoder_len++;
|
||||
}
|
||||
}
|
||||
|
||||
void add_special_token(const std::string& token) {
|
||||
special_tokens.push_back(token);
|
||||
}
|
||||
|
||||
std::u32string bpe(const std::u32string& token) {
|
||||
std::vector<std::u32string> word;
|
||||
|
||||
for (int i = 0; i < token.size() - 1; i++) {
|
||||
word.emplace_back(1, token[i]);
|
||||
}
|
||||
word.push_back(token.substr(token.size() - 1) + utf8_to_utf32("</w>"));
|
||||
|
||||
std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
|
||||
|
||||
if (pairs.empty()) {
|
||||
return token + utf8_to_utf32("</w>");
|
||||
}
|
||||
|
||||
while (true) {
|
||||
auto min_pair_iter = std::min_element(pairs.begin(),
|
||||
pairs.end(),
|
||||
[&](const std::pair<std::u32string, std::u32string>& a,
|
||||
const std::pair<std::u32string, std::u32string>& b) {
|
||||
if (bpe_ranks.find(a) == bpe_ranks.end()) {
|
||||
return false;
|
||||
} else if (bpe_ranks.find(b) == bpe_ranks.end()) {
|
||||
return true;
|
||||
}
|
||||
return bpe_ranks.at(a) < bpe_ranks.at(b);
|
||||
});
|
||||
|
||||
const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
|
||||
|
||||
if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
|
||||
break;
|
||||
}
|
||||
|
||||
std::u32string first = bigram.first;
|
||||
std::u32string second = bigram.second;
|
||||
std::vector<std::u32string> new_word;
|
||||
int32_t i = 0;
|
||||
|
||||
while (i < word.size()) {
|
||||
auto it = std::find(word.begin() + i, word.end(), first);
|
||||
if (it == word.end()) {
|
||||
new_word.insert(new_word.end(), word.begin() + i, word.end());
|
||||
break;
|
||||
}
|
||||
new_word.insert(new_word.end(), word.begin() + i, it);
|
||||
i = static_cast<int32_t>(std::distance(word.begin(), it));
|
||||
|
||||
if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
|
||||
new_word.push_back(first + second);
|
||||
i += 2;
|
||||
} else {
|
||||
new_word.push_back(word[i]);
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
|
||||
word = new_word;
|
||||
|
||||
if (word.size() == 1) {
|
||||
break;
|
||||
}
|
||||
pairs = get_pairs(word);
|
||||
}
|
||||
|
||||
std::u32string result;
|
||||
for (int i = 0; i < word.size(); i++) {
|
||||
result += word[i];
|
||||
if (i != word.size() - 1) {
|
||||
result += utf8_to_utf32(" ");
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<int> tokenize(std::string text,
|
||||
on_new_token_cb_t on_new_token_cb,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
std::vector<int32_t> tokens = encode(text, on_new_token_cb);
|
||||
|
||||
tokens.insert(tokens.begin(), BOS_TOKEN_ID);
|
||||
if (max_length > 0) {
|
||||
if (tokens.size() > max_length - 1) {
|
||||
tokens.resize(max_length - 1);
|
||||
tokens.push_back(EOS_TOKEN_ID);
|
||||
} else {
|
||||
tokens.push_back(EOS_TOKEN_ID);
|
||||
if (padding) {
|
||||
tokens.insert(tokens.end(), max_length - tokens.size(), PAD_TOKEN_ID);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return tokens;
|
||||
}
|
||||
|
||||
void pad_tokens(std::vector<int>& tokens,
|
||||
std::vector<float>& weights,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
if (max_length > 0 && padding) {
|
||||
size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.0 / (max_length - 2)));
|
||||
if (n == 0) {
|
||||
n = 1;
|
||||
}
|
||||
size_t length = max_length * n;
|
||||
LOG_DEBUG("token length: %llu", length);
|
||||
std::vector<int> new_tokens;
|
||||
std::vector<float> new_weights;
|
||||
new_tokens.push_back(BOS_TOKEN_ID);
|
||||
new_weights.push_back(1.0);
|
||||
int token_idx = 0;
|
||||
for (int i = 1; i < length; i++) {
|
||||
if (token_idx >= tokens.size()) {
|
||||
break;
|
||||
}
|
||||
if (i % max_length == 0) {
|
||||
new_tokens.push_back(BOS_TOKEN_ID);
|
||||
new_weights.push_back(1.0);
|
||||
} else if (i % max_length == max_length - 1) {
|
||||
new_tokens.push_back(EOS_TOKEN_ID);
|
||||
new_weights.push_back(1.0);
|
||||
} else {
|
||||
new_tokens.push_back(tokens[token_idx]);
|
||||
new_weights.push_back(weights[token_idx]);
|
||||
token_idx++;
|
||||
}
|
||||
}
|
||||
|
||||
new_tokens.push_back(EOS_TOKEN_ID);
|
||||
new_weights.push_back(1.0);
|
||||
tokens = new_tokens;
|
||||
weights = new_weights;
|
||||
|
||||
if (padding) {
|
||||
tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
|
||||
weights.insert(weights.end(), length - weights.size(), 1.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::string clean_up_tokenization(std::string& text) {
|
||||
std::regex pattern(R"( ,)");
|
||||
// Replace " ," with ","
|
||||
std::string result = std::regex_replace(text, pattern, ",");
|
||||
return result;
|
||||
}
|
||||
|
||||
std::string decode(const std::vector<int>& tokens) {
|
||||
std::string text = "";
|
||||
for (int t : tokens) {
|
||||
if (t == 49406 || t == 49407)
|
||||
continue;
|
||||
std::u32string ts = decoder[t];
|
||||
// printf("%d, %s \n", t, utf32_to_utf8(ts).c_str());
|
||||
std::string s = utf32_to_utf8(ts);
|
||||
if (s.length() >= 4) {
|
||||
if (ends_with(s, "</w>")) {
|
||||
text += s.replace(s.length() - 4, s.length() - 1, "") + " ";
|
||||
} else {
|
||||
text += s;
|
||||
}
|
||||
} else {
|
||||
text += " " + s;
|
||||
}
|
||||
}
|
||||
// std::vector<unsigned char> bytes;
|
||||
// for (auto c : text){
|
||||
// bytes.push_back(byte_decoder[c]);
|
||||
// }
|
||||
|
||||
// std::string s((char *)bytes.data());
|
||||
// std::string s = "";
|
||||
text = clean_up_tokenization(text);
|
||||
return trim(text);
|
||||
}
|
||||
|
||||
std::vector<std::string> token_split(const std::string& text) {
|
||||
std::regex pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
|
||||
std::regex::icase);
|
||||
std::sregex_iterator iter(text.begin(), text.end(), pat);
|
||||
std::sregex_iterator end;
|
||||
|
||||
std::vector<std::string> result;
|
||||
for (; iter != end; ++iter) {
|
||||
result.emplace_back(iter->str());
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb) {
|
||||
std::string original_text = text;
|
||||
std::vector<int32_t> bpe_tokens;
|
||||
text = whitespace_clean(text);
|
||||
std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return std::tolower(c); });
|
||||
|
||||
std::string str = text;
|
||||
std::vector<std::string> token_strs;
|
||||
|
||||
auto splited_texts = split_with_special_tokens(text, special_tokens);
|
||||
|
||||
for (auto& splited_text : splited_texts) {
|
||||
LOG_DEBUG("token %s", splited_text.c_str());
|
||||
if (is_special_token(splited_text)) {
|
||||
LOG_DEBUG("special %s", splited_text.c_str());
|
||||
bool skip = on_new_token_cb(splited_text, bpe_tokens);
|
||||
if (skip) {
|
||||
token_strs.push_back(splited_text);
|
||||
continue;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
auto tokens = token_split(splited_text);
|
||||
for (auto& token : tokens) {
|
||||
if (on_new_token_cb != nullptr) {
|
||||
bool skip = on_new_token_cb(token, bpe_tokens);
|
||||
if (skip) {
|
||||
token_strs.push_back(token);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
std::string token_str = token;
|
||||
std::u32string utf32_token;
|
||||
for (int i = 0; i < token_str.length(); i++) {
|
||||
unsigned char b = token_str[i];
|
||||
utf32_token += byte_encoder[b];
|
||||
}
|
||||
auto bpe_strs = bpe(utf32_token);
|
||||
size_t start = 0;
|
||||
size_t pos;
|
||||
while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {
|
||||
auto bpe_str = bpe_strs.substr(start, pos - start);
|
||||
bpe_tokens.push_back(encoder[bpe_str]);
|
||||
token_strs.push_back(utf32_to_utf8(bpe_str));
|
||||
|
||||
start = pos + 1;
|
||||
}
|
||||
auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);
|
||||
bpe_tokens.push_back(encoder[bpe_str]);
|
||||
token_strs.push_back(utf32_to_utf8(bpe_str));
|
||||
}
|
||||
}
|
||||
// std::stringstream ss;
|
||||
// ss << "[";
|
||||
// for (auto token : token_strs) {
|
||||
// ss << "\"" << token << "\", ";
|
||||
// }
|
||||
// ss << "]";
|
||||
// LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
|
||||
// printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str());
|
||||
return bpe_tokens;
|
||||
}
|
||||
};
|
||||
#include "tokenizers/clip_tokenizer.h"
|
||||
|
||||
/*================================================ FrozenCLIPEmbedder ================================================*/
|
||||
|
||||
|
||||
@@ -21,6 +21,7 @@
|
||||
#endif // _WIN32
|
||||
|
||||
#include "log.h"
|
||||
#include "media_io.h"
|
||||
#include "resource_owners.hpp"
|
||||
|
||||
using json = nlohmann::json;
|
||||
@@ -578,7 +579,17 @@ void SDContextParams::build_embedding_map() {
|
||||
}
|
||||
}
|
||||
|
||||
bool SDContextParams::process_and_check(SDMode mode) {
|
||||
bool SDContextParams::resolve(SDMode mode) {
|
||||
if (n_threads <= 0) {
|
||||
n_threads = sd_get_num_physical_cores();
|
||||
}
|
||||
|
||||
build_embedding_map();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDContextParams::validate(SDMode mode) {
|
||||
if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) {
|
||||
LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n");
|
||||
return false;
|
||||
@@ -591,12 +602,16 @@ bool SDContextParams::process_and_check(SDMode mode) {
|
||||
}
|
||||
}
|
||||
|
||||
if (n_threads <= 0) {
|
||||
n_threads = sd_get_num_physical_cores();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDContextParams::resolve_and_validate(SDMode mode) {
|
||||
if (!resolve(mode)) {
|
||||
return false;
|
||||
}
|
||||
if (!validate(mode)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
build_embedding_map();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -840,7 +855,7 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
&sample_params.guidance.slg.layer_end},
|
||||
{"",
|
||||
"--eta",
|
||||
"noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)",
|
||||
"noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
|
||||
&sample_params.eta},
|
||||
{"",
|
||||
"--flow-shift",
|
||||
@@ -872,7 +887,7 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
&high_noise_sample_params.guidance.slg.layer_end},
|
||||
{"",
|
||||
"--high-noise-eta",
|
||||
"(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a and dpm++2s_a)",
|
||||
"(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
|
||||
&high_noise_sample_params.eta},
|
||||
{"",
|
||||
"--strength",
|
||||
@@ -1170,12 +1185,12 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
on_seed_arg},
|
||||
{"",
|
||||
"--sampling-method",
|
||||
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s] "
|
||||
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde] "
|
||||
"(default: euler for Flux/SD3/Wan, euler_a otherwise)",
|
||||
on_sample_method_arg},
|
||||
{"",
|
||||
"--high-noise-sampling-method",
|
||||
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s]"
|
||||
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde]"
|
||||
" default: euler for Flux/SD3/Wan, euler_a otherwise",
|
||||
on_high_noise_sample_method_arg},
|
||||
{"",
|
||||
@@ -1228,7 +1243,190 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
return options;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::from_json_str(const std::string& json_str) {
|
||||
static const std::string k_base64_chars =
|
||||
"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
|
||||
"abcdefghijklmnopqrstuvwxyz"
|
||||
"0123456789+/";
|
||||
|
||||
static bool is_base64(unsigned char c) {
|
||||
return std::isalnum(c) || c == '+' || c == '/';
|
||||
}
|
||||
|
||||
static std::vector<uint8_t> decode_base64_bytes(const std::string& encoded_string) {
|
||||
int in_len = static_cast<int>(encoded_string.size());
|
||||
int i = 0;
|
||||
int j = 0;
|
||||
int in_ = 0;
|
||||
uint8_t char_array_4[4];
|
||||
uint8_t char_array_3[3];
|
||||
std::vector<uint8_t> ret;
|
||||
|
||||
while (in_len-- && encoded_string[in_] != '=' && is_base64(encoded_string[in_])) {
|
||||
char_array_4[i++] = encoded_string[in_];
|
||||
in_++;
|
||||
if (i == 4) {
|
||||
for (i = 0; i < 4; i++) {
|
||||
char_array_4[i] = static_cast<uint8_t>(k_base64_chars.find(char_array_4[i]));
|
||||
}
|
||||
|
||||
char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
|
||||
char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
|
||||
char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
|
||||
|
||||
for (i = 0; i < 3; i++) {
|
||||
ret.push_back(char_array_3[i]);
|
||||
}
|
||||
i = 0;
|
||||
}
|
||||
}
|
||||
|
||||
if (i) {
|
||||
for (j = i; j < 4; j++) {
|
||||
char_array_4[j] = 0;
|
||||
}
|
||||
|
||||
for (j = 0; j < 4; j++) {
|
||||
char_array_4[j] = static_cast<uint8_t>(k_base64_chars.find(char_array_4[j]));
|
||||
}
|
||||
|
||||
char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
|
||||
char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
|
||||
char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
|
||||
|
||||
for (j = 0; j < i - 1; j++) {
|
||||
ret.push_back(char_array_3[j]);
|
||||
}
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
|
||||
bool decode_base64_image(const std::string& encoded_input,
|
||||
int target_channels,
|
||||
int expected_width,
|
||||
int expected_height,
|
||||
SDImageOwner& out_image) {
|
||||
std::string encoded = encoded_input;
|
||||
auto comma_pos = encoded.find(',');
|
||||
if (comma_pos != std::string::npos) {
|
||||
encoded = encoded.substr(comma_pos + 1);
|
||||
}
|
||||
|
||||
std::vector<uint8_t> image_bytes = decode_base64_bytes(encoded);
|
||||
if (image_bytes.empty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
int decoded_width = 0;
|
||||
int decoded_height = 0;
|
||||
uint8_t* raw_data = load_image_from_memory(reinterpret_cast<const char*>(image_bytes.data()),
|
||||
static_cast<int>(image_bytes.size()),
|
||||
decoded_width,
|
||||
decoded_height,
|
||||
expected_width,
|
||||
expected_height,
|
||||
target_channels);
|
||||
if (raw_data == nullptr) {
|
||||
return false;
|
||||
}
|
||||
|
||||
out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)target_channels, raw_data});
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool parse_image_json_field(const json& parent,
|
||||
const char* key,
|
||||
int channels,
|
||||
int expected_width,
|
||||
int expected_height,
|
||||
SDImageOwner& out_image) {
|
||||
if (!parent.contains(key)) {
|
||||
return true;
|
||||
}
|
||||
if (parent.at(key).is_null()) {
|
||||
out_image.reset({0, 0, (uint32_t)channels, nullptr});
|
||||
return true;
|
||||
}
|
||||
if (!parent.at(key).is_string()) {
|
||||
return false;
|
||||
}
|
||||
return decode_base64_image(parent.at(key).get<std::string>(), channels, expected_width, expected_height, out_image);
|
||||
}
|
||||
|
||||
static bool parse_image_array_json_field(const json& parent,
|
||||
const char* key,
|
||||
int channels,
|
||||
int expected_width,
|
||||
int expected_height,
|
||||
std::vector<SDImageOwner>& out_images) {
|
||||
if (!parent.contains(key)) {
|
||||
return true;
|
||||
}
|
||||
if (parent.at(key).is_null()) {
|
||||
out_images.clear();
|
||||
return true;
|
||||
}
|
||||
if (!parent.at(key).is_array()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
out_images.clear();
|
||||
for (const auto& item : parent.at(key)) {
|
||||
if (!item.is_string()) {
|
||||
return false;
|
||||
}
|
||||
SDImageOwner image;
|
||||
if (!decode_base64_image(item.get<std::string>(), channels, expected_width, expected_height, image)) {
|
||||
return false;
|
||||
}
|
||||
out_images.push_back(std::move(image));
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool parse_lora_json_field(const json& parent,
|
||||
const std::function<std::string(const std::string&)>& lora_path_resolver,
|
||||
std::map<std::string, float>& lora_map,
|
||||
std::map<std::string, float>& high_noise_lora_map) {
|
||||
if (!parent.contains("lora")) {
|
||||
return true;
|
||||
}
|
||||
if (!parent.at("lora").is_array()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
lora_map.clear();
|
||||
high_noise_lora_map.clear();
|
||||
for (const auto& item : parent.at("lora")) {
|
||||
if (!item.is_object()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::string path = item.value("path", "");
|
||||
if (path.empty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
std::string resolved_path = lora_path_resolver ? lora_path_resolver(path) : path;
|
||||
if (resolved_path.empty()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const float multiplier = item.value("multiplier", 1.0f);
|
||||
const bool is_high_noise = item.value("is_high_noise", false);
|
||||
if (is_high_noise) {
|
||||
high_noise_lora_map[resolved_path] += multiplier;
|
||||
} else {
|
||||
lora_map[resolved_path] += multiplier;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::from_json_str(
|
||||
const std::string& json_str,
|
||||
const std::function<std::string(const std::string&)>& lora_path_resolver) {
|
||||
json j;
|
||||
try {
|
||||
j = json::parse(json_str);
|
||||
@@ -1255,6 +1453,9 @@ bool SDGenerationParams::from_json_str(const std::string& json_str) {
|
||||
} else if constexpr (std::is_same_v<T, std::vector<int>>) {
|
||||
if (j[key].is_array())
|
||||
out = j[key].get<std::vector<int>>();
|
||||
} else if constexpr (std::is_same_v<T, std::vector<float>>) {
|
||||
if (j[key].is_array())
|
||||
out = j[key].get<std::vector<float>>();
|
||||
} else if constexpr (std::is_same_v<T, std::vector<std::string>>) {
|
||||
if (j[key].is_array())
|
||||
out = j[key].get<std::vector<std::string>>();
|
||||
@@ -1279,7 +1480,6 @@ bool SDGenerationParams::from_json_str(const std::string& json_str) {
|
||||
|
||||
load_if_exists("strength", strength);
|
||||
load_if_exists("control_strength", control_strength);
|
||||
load_if_exists("pm_style_strength", pm_style_strength);
|
||||
load_if_exists("moe_boundary", moe_boundary);
|
||||
load_if_exists("vace_strength", vace_strength);
|
||||
|
||||
@@ -1287,32 +1487,127 @@ bool SDGenerationParams::from_json_str(const std::string& json_str) {
|
||||
load_if_exists("increase_ref_index", increase_ref_index);
|
||||
load_if_exists("embed_image_metadata", embed_image_metadata);
|
||||
|
||||
load_if_exists("skip_layers", skip_layers);
|
||||
load_if_exists("high_noise_skip_layers", high_noise_skip_layers);
|
||||
|
||||
load_if_exists("steps", sample_params.sample_steps);
|
||||
load_if_exists("high_noise_steps", high_noise_sample_params.sample_steps);
|
||||
load_if_exists("cfg_scale", sample_params.guidance.txt_cfg);
|
||||
load_if_exists("img_cfg_scale", sample_params.guidance.img_cfg);
|
||||
load_if_exists("guidance", sample_params.guidance.distilled_guidance);
|
||||
load_if_exists("flow_shift", sample_params.flow_shift);
|
||||
|
||||
auto load_sampler_if_exists = [&](const char* key, enum sample_method_t& out) {
|
||||
if (j.contains(key) && j[key].is_string()) {
|
||||
enum sample_method_t tmp = str_to_sample_method(j[key].get<std::string>().c_str());
|
||||
auto parse_sample_params_json = [&](const json& sample_json,
|
||||
sd_sample_params_t& target_params,
|
||||
std::vector<int>& target_skip_layers,
|
||||
std::vector<float>* target_custom_sigmas) {
|
||||
if (sample_json.contains("sample_steps") && sample_json["sample_steps"].is_number_integer()) {
|
||||
target_params.sample_steps = sample_json["sample_steps"];
|
||||
}
|
||||
if (sample_json.contains("eta") && sample_json["eta"].is_number()) {
|
||||
target_params.eta = sample_json["eta"];
|
||||
}
|
||||
if (sample_json.contains("shifted_timestep") && sample_json["shifted_timestep"].is_number_integer()) {
|
||||
target_params.shifted_timestep = sample_json["shifted_timestep"];
|
||||
}
|
||||
if (sample_json.contains("flow_shift") && sample_json["flow_shift"].is_number()) {
|
||||
target_params.flow_shift = sample_json["flow_shift"];
|
||||
}
|
||||
if (target_custom_sigmas != nullptr &&
|
||||
sample_json.contains("custom_sigmas") &&
|
||||
sample_json["custom_sigmas"].is_array()) {
|
||||
*target_custom_sigmas = sample_json["custom_sigmas"].get<std::vector<float>>();
|
||||
}
|
||||
if (sample_json.contains("sample_method") && sample_json["sample_method"].is_string()) {
|
||||
enum sample_method_t tmp = str_to_sample_method(sample_json["sample_method"].get<std::string>().c_str());
|
||||
if (tmp != SAMPLE_METHOD_COUNT) {
|
||||
out = tmp;
|
||||
target_params.sample_method = tmp;
|
||||
}
|
||||
}
|
||||
if (sample_json.contains("scheduler") && sample_json["scheduler"].is_string()) {
|
||||
enum scheduler_t tmp = str_to_scheduler(sample_json["scheduler"].get<std::string>().c_str());
|
||||
if (tmp != SCHEDULER_COUNT) {
|
||||
target_params.scheduler = tmp;
|
||||
}
|
||||
}
|
||||
if (sample_json.contains("guidance") && sample_json["guidance"].is_object()) {
|
||||
const json& guidance_json = sample_json["guidance"];
|
||||
if (guidance_json.contains("txt_cfg") && guidance_json["txt_cfg"].is_number()) {
|
||||
target_params.guidance.txt_cfg = guidance_json["txt_cfg"];
|
||||
}
|
||||
if (guidance_json.contains("img_cfg") && guidance_json["img_cfg"].is_number()) {
|
||||
target_params.guidance.img_cfg = guidance_json["img_cfg"];
|
||||
}
|
||||
if (guidance_json.contains("distilled_guidance") && guidance_json["distilled_guidance"].is_number()) {
|
||||
target_params.guidance.distilled_guidance = guidance_json["distilled_guidance"];
|
||||
}
|
||||
if (guidance_json.contains("slg") && guidance_json["slg"].is_object()) {
|
||||
const json& slg_json = guidance_json["slg"];
|
||||
if (slg_json.contains("layers") && slg_json["layers"].is_array()) {
|
||||
target_skip_layers = slg_json["layers"].get<std::vector<int>>();
|
||||
}
|
||||
if (slg_json.contains("layer_start") && slg_json["layer_start"].is_number()) {
|
||||
target_params.guidance.slg.layer_start = slg_json["layer_start"];
|
||||
}
|
||||
if (slg_json.contains("layer_end") && slg_json["layer_end"].is_number()) {
|
||||
target_params.guidance.slg.layer_end = slg_json["layer_end"];
|
||||
}
|
||||
if (slg_json.contains("scale") && slg_json["scale"].is_number()) {
|
||||
target_params.guidance.slg.scale = slg_json["scale"];
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
load_sampler_if_exists("sample_method", sample_params.sample_method);
|
||||
load_sampler_if_exists("high_noise_sample_method", high_noise_sample_params.sample_method);
|
||||
|
||||
if (j.contains("scheduler") && j["scheduler"].is_string()) {
|
||||
enum scheduler_t tmp = str_to_scheduler(j["scheduler"].get<std::string>().c_str());
|
||||
if (tmp != SCHEDULER_COUNT) {
|
||||
sample_params.scheduler = tmp;
|
||||
if (j.contains("sample_params") && j["sample_params"].is_object()) {
|
||||
parse_sample_params_json(j["sample_params"], sample_params, skip_layers, &custom_sigmas);
|
||||
}
|
||||
if (j.contains("high_noise_sample_params") && j["high_noise_sample_params"].is_object()) {
|
||||
parse_sample_params_json(j["high_noise_sample_params"],
|
||||
high_noise_sample_params,
|
||||
high_noise_skip_layers,
|
||||
nullptr);
|
||||
}
|
||||
|
||||
if (j.contains("vae_tiling_params") && j["vae_tiling_params"].is_object()) {
|
||||
const json& tiling_json = j["vae_tiling_params"];
|
||||
if (tiling_json.contains("enabled") && tiling_json["enabled"].is_boolean()) {
|
||||
vae_tiling_params.enabled = tiling_json["enabled"];
|
||||
}
|
||||
if (tiling_json.contains("tile_size_x") && tiling_json["tile_size_x"].is_number_integer()) {
|
||||
vae_tiling_params.tile_size_x = tiling_json["tile_size_x"];
|
||||
}
|
||||
if (tiling_json.contains("tile_size_y") && tiling_json["tile_size_y"].is_number_integer()) {
|
||||
vae_tiling_params.tile_size_y = tiling_json["tile_size_y"];
|
||||
}
|
||||
if (tiling_json.contains("target_overlap") && tiling_json["target_overlap"].is_number()) {
|
||||
vae_tiling_params.target_overlap = tiling_json["target_overlap"];
|
||||
}
|
||||
if (tiling_json.contains("rel_size_x") && tiling_json["rel_size_x"].is_number()) {
|
||||
vae_tiling_params.rel_size_x = tiling_json["rel_size_x"];
|
||||
}
|
||||
if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) {
|
||||
vae_tiling_params.rel_size_y = tiling_json["rel_size_y"];
|
||||
}
|
||||
}
|
||||
|
||||
if (!parse_lora_json_field(j, lora_path_resolver, lora_map, high_noise_lora_map)) {
|
||||
LOG_ERROR("invalid lora");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_json_field(j, "init_image", 3, width, height, init_image)) {
|
||||
LOG_ERROR("invalid init_image");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) {
|
||||
LOG_ERROR("invalid end_image");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_array_json_field(j, "ref_images", 3, width, height, ref_images)) {
|
||||
LOG_ERROR("invalid ref_images");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_array_json_field(j, "control_frames", 3, width, height, control_frames)) {
|
||||
LOG_ERROR("invalid control_frames");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_json_field(j, "mask_image", 1, width, height, mask_image)) {
|
||||
LOG_ERROR("invalid mask_image");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_json_field(j, "control_image", 3, width, height, control_image)) {
|
||||
LOG_ERROR("invalid control_image");
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
@@ -1384,22 +1679,6 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
|
||||
|
||||
tmp = m.suffix().str();
|
||||
}
|
||||
|
||||
for (const auto& kv : lora_map) {
|
||||
sd_lora_t item;
|
||||
item.is_high_noise = false;
|
||||
item.path = kv.first.c_str();
|
||||
item.multiplier = kv.second;
|
||||
lora_vec.emplace_back(item);
|
||||
}
|
||||
|
||||
for (const auto& kv : high_noise_lora_map) {
|
||||
sd_lora_t item;
|
||||
item.is_high_noise = true;
|
||||
item.path = kv.first.c_str();
|
||||
item.multiplier = kv.second;
|
||||
lora_vec.emplace_back(item);
|
||||
}
|
||||
}
|
||||
|
||||
bool SDGenerationParams::width_and_height_are_set() const {
|
||||
@@ -1422,23 +1701,7 @@ int SDGenerationParams::get_resolved_height() const {
|
||||
return (height > 0) ? height : 512;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::process_and_check(SDMode mode, const std::string& lora_model_dir) {
|
||||
prompt_with_lora = prompt;
|
||||
|
||||
if (sample_params.sample_steps <= 0) {
|
||||
LOG_ERROR("error: the sample_steps must be greater than 0\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (high_noise_sample_params.sample_steps <= 0) {
|
||||
high_noise_sample_params.sample_steps = -1;
|
||||
}
|
||||
|
||||
if (strength < 0.f || strength > 1.f) {
|
||||
LOG_ERROR("error: can only work with strength in [0.0, 1.0]\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::initialize_cache_params() {
|
||||
sd_cache_params_init(&cache_params);
|
||||
|
||||
auto parse_named_params = [&](const std::string& opt_str) -> bool {
|
||||
@@ -1504,7 +1767,9 @@ bool SDGenerationParams::process_and_check(SDMode mode, const std::string& lora_
|
||||
};
|
||||
|
||||
if (!cache_mode.empty()) {
|
||||
if (cache_mode == "easycache") {
|
||||
if (cache_mode == "disabled") {
|
||||
cache_params.mode = SD_CACHE_DISABLED;
|
||||
} else if (cache_mode == "easycache") {
|
||||
cache_params.mode = SD_CACHE_EASYCACHE;
|
||||
} else if (cache_mode == "ucache") {
|
||||
cache_params.mode = SD_CACHE_UCACHE;
|
||||
@@ -1516,14 +1781,73 @@ bool SDGenerationParams::process_and_check(SDMode mode, const std::string& lora_
|
||||
cache_params.mode = SD_CACHE_CACHE_DIT;
|
||||
} else if (cache_mode == "spectrum") {
|
||||
cache_params.mode = SD_CACHE_SPECTRUM;
|
||||
} else {
|
||||
LOG_ERROR("error: invalid cache mode '%s'", cache_mode.c_str());
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if (!cache_option.empty()) {
|
||||
if (!parse_named_params(cache_option)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if (!cache_option.empty() && !parse_named_params(cache_option)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (cache_params.mode == SD_CACHE_DBCACHE ||
|
||||
cache_params.mode == SD_CACHE_TAYLORSEER ||
|
||||
cache_params.mode == SD_CACHE_CACHE_DIT) {
|
||||
cache_params.scm_policy_dynamic = scm_policy_dynamic;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::resolve(const std::string& lora_model_dir, bool strict) {
|
||||
if (high_noise_sample_params.sample_steps <= 0) {
|
||||
high_noise_sample_params.sample_steps = -1;
|
||||
}
|
||||
|
||||
if (!initialize_cache_params()) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (seed < 0) {
|
||||
srand((int)time(nullptr));
|
||||
seed = rand();
|
||||
}
|
||||
|
||||
if (strict) {
|
||||
batch_count = std::clamp(batch_count, 1, 8);
|
||||
sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100);
|
||||
}
|
||||
|
||||
prompt_with_lora = prompt;
|
||||
if (!lora_model_dir.empty()) {
|
||||
extract_and_remove_lora(lora_model_dir);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::validate(SDMode mode) {
|
||||
if (batch_count <= 0) {
|
||||
LOG_ERROR("error: batch_count must be greater than 0");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (sample_params.sample_steps <= 0) {
|
||||
LOG_ERROR("error: the sample_steps must be greater than 0\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (strength < 0.f || strength > 1.f) {
|
||||
LOG_ERROR("error: can only work with strength in [0.0, 1.0]\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (sample_params.guidance.txt_cfg < 0.f) {
|
||||
LOG_ERROR("error: cfg_scale must be positive");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!cache_mode.empty()) {
|
||||
if (cache_mode == "easycache" || cache_mode == "ucache") {
|
||||
if (cache_params.reuse_threshold < 0.0f) {
|
||||
LOG_ERROR("error: cache threshold must be non-negative");
|
||||
@@ -1538,22 +1862,6 @@ bool SDGenerationParams::process_and_check(SDMode mode, const std::string& lora_
|
||||
}
|
||||
}
|
||||
|
||||
if (cache_params.mode == SD_CACHE_DBCACHE ||
|
||||
cache_params.mode == SD_CACHE_TAYLORSEER ||
|
||||
cache_params.mode == SD_CACHE_CACHE_DIT) {
|
||||
if (!scm_mask.empty()) {
|
||||
cache_params.scm_mask = scm_mask.c_str();
|
||||
}
|
||||
cache_params.scm_policy_dynamic = scm_policy_dynamic;
|
||||
}
|
||||
|
||||
sample_params.guidance.slg.layers = skip_layers.data();
|
||||
sample_params.guidance.slg.layer_count = skip_layers.size();
|
||||
sample_params.custom_sigmas = custom_sigmas.data();
|
||||
sample_params.custom_sigmas_count = static_cast<int>(custom_sigmas.size());
|
||||
high_noise_sample_params.guidance.slg.layers = high_noise_skip_layers.data();
|
||||
high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
|
||||
|
||||
if (mode == VID_GEN && video_frames <= 0) {
|
||||
return false;
|
||||
}
|
||||
@@ -1563,6 +1871,7 @@ bool SDGenerationParams::process_and_check(SDMode mode, const std::string& lora_
|
||||
}
|
||||
|
||||
if (sample_params.shifted_timestep < 0 || sample_params.shifted_timestep > 1000) {
|
||||
LOG_ERROR("error: shifted_timestep must be in range [0, 1000]");
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -1581,14 +1890,132 @@ bool SDGenerationParams::process_and_check(SDMode mode, const std::string& lora_
|
||||
}
|
||||
}
|
||||
|
||||
if (seed < 0) {
|
||||
srand((int)time(nullptr));
|
||||
seed = rand();
|
||||
return true;
|
||||
}
|
||||
|
||||
bool SDGenerationParams::resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict) {
|
||||
if (!resolve(lora_model_dir, strict)) {
|
||||
return false;
|
||||
}
|
||||
if (!validate(mode)) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
||||
sd_img_gen_params_t params;
|
||||
sd_img_gen_params_init(¶ms);
|
||||
|
||||
lora_vec.clear();
|
||||
lora_vec.reserve(lora_map.size() + high_noise_lora_map.size());
|
||||
for (const auto& kv : lora_map) {
|
||||
lora_vec.push_back({false, kv.second, kv.first.c_str()});
|
||||
}
|
||||
for (const auto& kv : high_noise_lora_map) {
|
||||
lora_vec.push_back({true, kv.second, kv.first.c_str()});
|
||||
}
|
||||
|
||||
extract_and_remove_lora(lora_model_dir);
|
||||
ref_image_views.clear();
|
||||
ref_image_views.reserve(ref_images.size());
|
||||
for (auto& ref_image : ref_images) {
|
||||
ref_image_views.push_back(ref_image.get());
|
||||
}
|
||||
|
||||
return true;
|
||||
pm_id_image_views.clear();
|
||||
pm_id_image_views.reserve(pm_id_images.size());
|
||||
for (auto& image : pm_id_images) {
|
||||
pm_id_image_views.push_back(image.get());
|
||||
}
|
||||
|
||||
sample_params.guidance.slg.layers = skip_layers.empty() ? nullptr : skip_layers.data();
|
||||
sample_params.guidance.slg.layer_count = skip_layers.size();
|
||||
high_noise_sample_params.guidance.slg.layers = high_noise_skip_layers.empty() ? nullptr : high_noise_skip_layers.data();
|
||||
high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
|
||||
sample_params.custom_sigmas = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
|
||||
sample_params.custom_sigmas_count = static_cast<int>(custom_sigmas.size());
|
||||
cache_params.scm_mask = scm_mask.empty() ? nullptr : scm_mask.c_str();
|
||||
|
||||
sd_pm_params_t pm_params = {
|
||||
pm_id_image_views.empty() ? nullptr : pm_id_image_views.data(),
|
||||
static_cast<int>(pm_id_image_views.size()),
|
||||
pm_id_embed_path.empty() ? nullptr : pm_id_embed_path.c_str(),
|
||||
pm_style_strength,
|
||||
};
|
||||
|
||||
params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||
params.prompt = prompt.c_str();
|
||||
params.negative_prompt = negative_prompt.c_str();
|
||||
params.clip_skip = clip_skip;
|
||||
params.init_image = init_image.get();
|
||||
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
|
||||
params.ref_images_count = static_cast<int>(ref_image_views.size());
|
||||
params.auto_resize_ref_image = auto_resize_ref_image;
|
||||
params.increase_ref_index = increase_ref_index;
|
||||
params.mask_image = mask_image.get();
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.batch_count = batch_count;
|
||||
params.control_image = control_image.get();
|
||||
params.control_strength = control_strength;
|
||||
params.pm_params = pm_params;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
return params;
|
||||
}
|
||||
|
||||
sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
|
||||
sd_vid_gen_params_t params;
|
||||
sd_vid_gen_params_init(¶ms);
|
||||
|
||||
lora_vec.clear();
|
||||
lora_vec.reserve(lora_map.size() + high_noise_lora_map.size());
|
||||
for (const auto& kv : lora_map) {
|
||||
lora_vec.push_back({false, kv.second, kv.first.c_str()});
|
||||
}
|
||||
for (const auto& kv : high_noise_lora_map) {
|
||||
lora_vec.push_back({true, kv.second, kv.first.c_str()});
|
||||
}
|
||||
|
||||
control_frame_views.clear();
|
||||
control_frame_views.reserve(control_frames.size());
|
||||
for (auto& frame : control_frames) {
|
||||
control_frame_views.push_back(frame.get());
|
||||
}
|
||||
|
||||
sample_params.guidance.slg.layers = skip_layers.empty() ? nullptr : skip_layers.data();
|
||||
sample_params.guidance.slg.layer_count = skip_layers.size();
|
||||
high_noise_sample_params.guidance.slg.layers = high_noise_skip_layers.empty() ? nullptr : high_noise_skip_layers.data();
|
||||
high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
|
||||
sample_params.custom_sigmas = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
|
||||
sample_params.custom_sigmas_count = static_cast<int>(custom_sigmas.size());
|
||||
cache_params.scm_mask = scm_mask.empty() ? nullptr : scm_mask.c_str();
|
||||
|
||||
params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||
params.prompt = prompt.c_str();
|
||||
params.negative_prompt = negative_prompt.c_str();
|
||||
params.clip_skip = clip_skip;
|
||||
params.init_image = init_image.get();
|
||||
params.end_image = end_image.get();
|
||||
params.control_frames = control_frame_views.empty() ? nullptr : control_frame_views.data();
|
||||
params.control_frames_size = static_cast<int>(control_frame_views.size());
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.high_noise_sample_params = high_noise_sample_params;
|
||||
params.moe_boundary = moe_boundary;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.video_frames = video_frames;
|
||||
params.vace_strength = vace_strength;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
return params;
|
||||
}
|
||||
|
||||
std::string SDGenerationParams::to_string() const {
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
#include <vector>
|
||||
|
||||
#include "log.h"
|
||||
#include "resource_owners.hpp"
|
||||
#include "stable-diffusion.h"
|
||||
|
||||
#define SAFE_STR(s) ((s) ? (s) : "")
|
||||
@@ -74,6 +75,11 @@ struct ArgOptions {
|
||||
};
|
||||
|
||||
bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& options_list);
|
||||
bool decode_base64_image(const std::string& encoded_input,
|
||||
int target_channels,
|
||||
int expected_width,
|
||||
int expected_height,
|
||||
SDImageOwner& out_image);
|
||||
|
||||
struct SDContextParams {
|
||||
int n_threads = -1;
|
||||
@@ -129,34 +135,39 @@ struct SDContextParams {
|
||||
float flow_shift = INFINITY;
|
||||
ArgOptions get_options();
|
||||
void build_embedding_map();
|
||||
bool process_and_check(SDMode mode);
|
||||
bool resolve(SDMode mode);
|
||||
bool validate(SDMode mode);
|
||||
bool resolve_and_validate(SDMode mode);
|
||||
std::string to_string() const;
|
||||
sd_ctx_params_t to_sd_ctx_params_t(bool vae_decode_only, bool free_params_immediately, bool taesd_preview);
|
||||
};
|
||||
|
||||
struct SDGenerationParams {
|
||||
// User-facing input fields.
|
||||
std::string prompt;
|
||||
std::string prompt_with_lora; // for metadata record only
|
||||
std::string negative_prompt;
|
||||
int clip_skip = -1; // <= 0 represents unspecified
|
||||
int width = -1;
|
||||
int height = -1;
|
||||
int batch_count = 1;
|
||||
int clip_skip = -1; // <= 0 represents unspecified
|
||||
int width = -1;
|
||||
int height = -1;
|
||||
int batch_count = 1;
|
||||
int64_t seed = 42;
|
||||
float strength = 0.75f;
|
||||
float control_strength = 0.9f;
|
||||
bool auto_resize_ref_image = true;
|
||||
bool increase_ref_index = false;
|
||||
bool embed_image_metadata = true;
|
||||
|
||||
std::string init_image_path;
|
||||
std::string end_image_path;
|
||||
std::string mask_image_path;
|
||||
std::string control_image_path;
|
||||
std::vector<std::string> ref_image_paths;
|
||||
std::string control_video_path;
|
||||
bool auto_resize_ref_image = true;
|
||||
bool increase_ref_index = false;
|
||||
bool embed_image_metadata = true;
|
||||
|
||||
std::vector<int> skip_layers = {7, 8, 9};
|
||||
sd_sample_params_t sample_params;
|
||||
|
||||
std::vector<int> high_noise_skip_layers = {7, 8, 9};
|
||||
sd_sample_params_t high_noise_sample_params;
|
||||
std::vector<int> skip_layers = {7, 8, 9};
|
||||
std::vector<int> high_noise_skip_layers = {7, 8, 9};
|
||||
|
||||
std::vector<float> custom_sigmas;
|
||||
|
||||
@@ -166,19 +177,12 @@ struct SDGenerationParams {
|
||||
bool scm_policy_dynamic = true;
|
||||
sd_cache_params_t cache_params{};
|
||||
|
||||
float moe_boundary = 0.875f;
|
||||
int video_frames = 1;
|
||||
int fps = 16;
|
||||
float vace_strength = 1.f;
|
||||
|
||||
float strength = 0.75f;
|
||||
float control_strength = 0.9f;
|
||||
|
||||
int64_t seed = 42;
|
||||
|
||||
float moe_boundary = 0.875f;
|
||||
int video_frames = 1;
|
||||
int fps = 16;
|
||||
float vace_strength = 1.f;
|
||||
sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
|
||||
|
||||
// Photo Maker
|
||||
std::string pm_id_images_dir;
|
||||
std::string pm_id_embed_path;
|
||||
float pm_style_strength = 20.f;
|
||||
@@ -188,16 +192,44 @@ struct SDGenerationParams {
|
||||
|
||||
std::map<std::string, float> lora_map;
|
||||
std::map<std::string, float> high_noise_lora_map;
|
||||
|
||||
// Derived and normalized fields.
|
||||
std::string prompt_with_lora; // for metadata record only
|
||||
std::vector<sd_lora_t> lora_vec;
|
||||
|
||||
// Owned execution payload.
|
||||
SDImageOwner init_image;
|
||||
SDImageOwner end_image;
|
||||
std::vector<SDImageOwner> ref_images;
|
||||
SDImageOwner mask_image;
|
||||
SDImageOwner control_image;
|
||||
std::vector<SDImageOwner> pm_id_images;
|
||||
std::vector<SDImageOwner> control_frames;
|
||||
|
||||
// Backing storage for sd_img_gen_params_t view fields.
|
||||
std::vector<sd_image_t> ref_image_views;
|
||||
std::vector<sd_image_t> pm_id_image_views;
|
||||
std::vector<sd_image_t> control_frame_views;
|
||||
|
||||
SDGenerationParams();
|
||||
SDGenerationParams(const SDGenerationParams& other) = default;
|
||||
SDGenerationParams& operator=(const SDGenerationParams& other) = default;
|
||||
SDGenerationParams(SDGenerationParams&& other) noexcept = default;
|
||||
SDGenerationParams& operator=(SDGenerationParams&& other) noexcept = default;
|
||||
ArgOptions get_options();
|
||||
bool from_json_str(const std::string& json_str);
|
||||
bool from_json_str(const std::string& json_str,
|
||||
const std::function<std::string(const std::string&)>& lora_path_resolver = {});
|
||||
bool initialize_cache_params();
|
||||
void extract_and_remove_lora(const std::string& lora_model_dir);
|
||||
bool width_and_height_are_set() const;
|
||||
void set_width_and_height_if_unset(int w, int h);
|
||||
int get_resolved_width() const;
|
||||
int get_resolved_height() const;
|
||||
bool process_and_check(SDMode mode, const std::string& lora_model_dir);
|
||||
bool resolve(const std::string& lora_model_dir, bool strict = false);
|
||||
bool validate(SDMode mode);
|
||||
bool resolve_and_validate(SDMode mode, const std::string& lora_model_dir, bool strict = false);
|
||||
sd_img_gen_params_t to_sd_img_gen_params_t();
|
||||
sd_vid_gen_params_t to_sd_vid_gen_params_t();
|
||||
std::string to_string() const;
|
||||
};
|
||||
|
||||
|
||||
+257
-140
@@ -95,6 +95,57 @@ using WebPMuxPtr = std::unique_ptr<WebPMux, WebPMuxDeleter>;
|
||||
using WebPAnimEncoderPtr = std::unique_ptr<WebPAnimEncoder, WebPAnimEncoderDeleter>;
|
||||
#endif
|
||||
|
||||
#ifdef SD_USE_WEBM
|
||||
class MemoryMkvWriter : public mkvmuxer::IMkvWriter {
|
||||
public:
|
||||
mkvmuxer::int32 Write(const void* buf, mkvmuxer::uint32 len) override {
|
||||
if (buf == nullptr && len > 0) {
|
||||
return -1;
|
||||
}
|
||||
const size_t end_pos = position_ + static_cast<size_t>(len);
|
||||
if (end_pos > data_.size()) {
|
||||
data_.resize(end_pos);
|
||||
}
|
||||
if (len > 0) {
|
||||
memcpy(data_.data() + position_, buf, len);
|
||||
}
|
||||
position_ = end_pos;
|
||||
return 0;
|
||||
}
|
||||
|
||||
mkvmuxer::int64 Position() const override {
|
||||
return static_cast<mkvmuxer::int64>(position_);
|
||||
}
|
||||
|
||||
mkvmuxer::int32 Position(mkvmuxer::int64 position) override {
|
||||
if (position < 0) {
|
||||
return -1;
|
||||
}
|
||||
const size_t target = static_cast<size_t>(position);
|
||||
if (target > data_.size()) {
|
||||
data_.resize(target);
|
||||
}
|
||||
position_ = target;
|
||||
return 0;
|
||||
}
|
||||
|
||||
bool Seekable() const override {
|
||||
return true;
|
||||
}
|
||||
|
||||
void ElementStartNotify(mkvmuxer::uint64, mkvmuxer::int64) override {
|
||||
}
|
||||
|
||||
const std::vector<uint8_t>& data() const {
|
||||
return data_;
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<uint8_t> data_;
|
||||
size_t position_ = 0;
|
||||
};
|
||||
#endif
|
||||
|
||||
bool read_binary_file_bytes(const char* path, std::vector<uint8_t>& data) {
|
||||
std::ifstream fin(fs::path(path), std::ios::binary);
|
||||
if (!fin) {
|
||||
@@ -569,6 +620,33 @@ void write_u32_le(FILE* f, uint32_t val) {
|
||||
void write_u16_le(FILE* f, uint16_t val) {
|
||||
fwrite(&val, 2, 1, f);
|
||||
}
|
||||
|
||||
void write_u32_le(std::vector<uint8_t>& data, uint32_t val) {
|
||||
data.push_back(static_cast<uint8_t>(val & 0xFF));
|
||||
data.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
|
||||
data.push_back(static_cast<uint8_t>((val >> 16) & 0xFF));
|
||||
data.push_back(static_cast<uint8_t>((val >> 24) & 0xFF));
|
||||
}
|
||||
|
||||
void write_u16_le(std::vector<uint8_t>& data, uint16_t val) {
|
||||
data.push_back(static_cast<uint8_t>(val & 0xFF));
|
||||
data.push_back(static_cast<uint8_t>((val >> 8) & 0xFF));
|
||||
}
|
||||
|
||||
void patch_u32_le(std::vector<uint8_t>& data, size_t offset, uint32_t val) {
|
||||
if (offset + 4 > data.size()) {
|
||||
return;
|
||||
}
|
||||
data[offset + 0] = static_cast<uint8_t>(val & 0xFF);
|
||||
data[offset + 1] = static_cast<uint8_t>((val >> 8) & 0xFF);
|
||||
data[offset + 2] = static_cast<uint8_t>((val >> 16) & 0xFF);
|
||||
data[offset + 3] = static_cast<uint8_t>((val >> 24) & 0xFF);
|
||||
}
|
||||
|
||||
void write_fourcc(std::vector<uint8_t>& data, const char* fourcc) {
|
||||
data.insert(data.end(), fourcc, fourcc + 4);
|
||||
}
|
||||
|
||||
EncodedImageFormat encoded_image_format_from_path(const std::string& path) {
|
||||
std::string ext = fs::path(path).extension().string();
|
||||
std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
|
||||
@@ -682,8 +760,9 @@ bool load_sd_image_from_file(sd_image_t* image,
|
||||
if (image->data == nullptr) {
|
||||
return false;
|
||||
}
|
||||
image->width = width;
|
||||
image->height = height;
|
||||
image->width = width;
|
||||
image->height = height;
|
||||
image->channel = expected_channel;
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -697,95 +776,96 @@ uint8_t* load_image_from_memory(const char* image_bytes,
|
||||
return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
|
||||
}
|
||||
|
||||
int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
|
||||
if (num_images == 0) {
|
||||
fprintf(stderr, "Error: Image array is empty.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
FilePtr file(fopen(filename, "wb"));
|
||||
if (!file) {
|
||||
perror("Error opening file for writing");
|
||||
return -1;
|
||||
}
|
||||
FILE* f = file.get();
|
||||
|
||||
uint32_t width = images[0].width;
|
||||
uint32_t height = images[0].height;
|
||||
uint32_t channels = images[0].channel;
|
||||
if (channels != 3 && channels != 4) {
|
||||
fprintf(stderr, "Error: Unsupported channel count: %u\n", channels);
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
fwrite("RIFF", 4, 1, f);
|
||||
long riff_size_pos = ftell(f);
|
||||
write_u32_le(f, 0);
|
||||
fwrite("AVI ", 4, 1, f);
|
||||
// stb_image_write changes JPEG sampling behavior above quality 90.
|
||||
// MJPG AVI playback is more compatible when we keep the encoder on the
|
||||
// <= 90 path.
|
||||
const int mjpg_quality = std::clamp(quality, 1, 90);
|
||||
|
||||
fwrite("LIST", 4, 1, f);
|
||||
write_u32_le(f, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
|
||||
fwrite("hdrl", 4, 1, f);
|
||||
std::vector<uint8_t> avi_data;
|
||||
avi_data.reserve(static_cast<size_t>(num_images) * 1024);
|
||||
|
||||
fwrite("avih", 4, 1, f);
|
||||
write_u32_le(f, 56);
|
||||
write_u32_le(f, 1000000 / fps);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0x110);
|
||||
write_u32_le(f, num_images);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 1);
|
||||
write_u32_le(f, width * height * 3);
|
||||
write_u32_le(f, width);
|
||||
write_u32_le(f, height);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_fourcc(avi_data, "RIFF");
|
||||
const size_t riff_size_pos = avi_data.size();
|
||||
write_u32_le(avi_data, 0);
|
||||
write_fourcc(avi_data, "AVI ");
|
||||
|
||||
fwrite("LIST", 4, 1, f);
|
||||
write_u32_le(f, 4 + 8 + 56 + 8 + 40);
|
||||
fwrite("strl", 4, 1, f);
|
||||
write_fourcc(avi_data, "LIST");
|
||||
write_u32_le(avi_data, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
|
||||
write_fourcc(avi_data, "hdrl");
|
||||
|
||||
fwrite("strh", 4, 1, f);
|
||||
write_u32_le(f, 56);
|
||||
fwrite("vids", 4, 1, f);
|
||||
fwrite("MJPG", 4, 1, f);
|
||||
write_u32_le(f, 0);
|
||||
write_u16_le(f, 0);
|
||||
write_u16_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 1);
|
||||
write_u32_le(f, fps);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, num_images);
|
||||
write_u32_le(f, width * height * 3);
|
||||
write_u32_le(f, (uint32_t)-1);
|
||||
write_u32_le(f, 0);
|
||||
write_u16_le(f, 0);
|
||||
write_u16_le(f, 0);
|
||||
write_u16_le(f, 0);
|
||||
write_u16_le(f, 0);
|
||||
write_fourcc(avi_data, "avih");
|
||||
write_u32_le(avi_data, 56);
|
||||
write_u32_le(avi_data, 1000000 / fps);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0x110);
|
||||
write_u32_le(avi_data, num_images);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 1);
|
||||
write_u32_le(avi_data, width * height * 3);
|
||||
write_u32_le(avi_data, width);
|
||||
write_u32_le(avi_data, height);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
|
||||
fwrite("strf", 4, 1, f);
|
||||
write_u32_le(f, 40);
|
||||
write_u32_le(f, 40);
|
||||
write_u32_le(f, width);
|
||||
write_u32_le(f, height);
|
||||
write_u16_le(f, 1);
|
||||
write_u16_le(f, 24);
|
||||
fwrite("MJPG", 4, 1, f);
|
||||
write_u32_le(f, width * height * 3);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_u32_le(f, 0);
|
||||
write_fourcc(avi_data, "LIST");
|
||||
write_u32_le(avi_data, 4 + 8 + 56 + 8 + 40);
|
||||
write_fourcc(avi_data, "strl");
|
||||
|
||||
fwrite("LIST", 4, 1, f);
|
||||
long movi_size_pos = ftell(f);
|
||||
write_u32_le(f, 0);
|
||||
fwrite("movi", 4, 1, f);
|
||||
write_fourcc(avi_data, "strh");
|
||||
write_u32_le(avi_data, 56);
|
||||
write_fourcc(avi_data, "vids");
|
||||
write_fourcc(avi_data, "MJPG");
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u16_le(avi_data, 0);
|
||||
write_u16_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 1);
|
||||
write_u32_le(avi_data, fps);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, num_images);
|
||||
write_u32_le(avi_data, width * height * 3);
|
||||
write_u32_le(avi_data, static_cast<uint32_t>(-1));
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u16_le(avi_data, 0);
|
||||
write_u16_le(avi_data, 0);
|
||||
write_u16_le(avi_data, 0);
|
||||
write_u16_le(avi_data, 0);
|
||||
|
||||
write_fourcc(avi_data, "strf");
|
||||
write_u32_le(avi_data, 40);
|
||||
write_u32_le(avi_data, 40);
|
||||
write_u32_le(avi_data, width);
|
||||
write_u32_le(avi_data, height);
|
||||
write_u16_le(avi_data, 1);
|
||||
write_u16_le(avi_data, 24);
|
||||
write_fourcc(avi_data, "MJPG");
|
||||
write_u32_le(avi_data, width * height * 3);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
write_u32_le(avi_data, 0);
|
||||
|
||||
write_fourcc(avi_data, "LIST");
|
||||
const size_t movi_size_pos = avi_data.size();
|
||||
write_u32_le(avi_data, 0);
|
||||
write_fourcc(avi_data, "movi");
|
||||
|
||||
std::vector<avi_index_entry> index(static_cast<size_t>(num_images));
|
||||
std::vector<uint8_t> jpeg_data;
|
||||
@@ -799,55 +879,61 @@ int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int
|
||||
buffer->insert(buffer->end(), src, src + size);
|
||||
};
|
||||
|
||||
if (!stbi_write_jpg_to_func(write_to_buf, &jpeg_data, images[i].width, images[i].height, channels, images[i].data, quality)) {
|
||||
if (!stbi_write_jpg_to_func(write_to_buf, &jpeg_data, images[i].width, images[i].height, channels, images[i].data, mjpg_quality)) {
|
||||
fprintf(stderr, "Error: Failed to encode JPEG frame.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
fwrite("00dc", 4, 1, f);
|
||||
write_u32_le(f, (uint32_t)jpeg_data.size());
|
||||
index[i].offset = ftell(f) - 8;
|
||||
index[i].size = (uint32_t)jpeg_data.size();
|
||||
fwrite(jpeg_data.data(), 1, jpeg_data.size(), f);
|
||||
index[i].offset = static_cast<uint32_t>(avi_data.size());
|
||||
write_fourcc(avi_data, "00dc");
|
||||
write_u32_le(avi_data, static_cast<uint32_t>(jpeg_data.size()));
|
||||
index[i].size = (uint32_t)jpeg_data.size();
|
||||
avi_data.insert(avi_data.end(), jpeg_data.begin(), jpeg_data.end());
|
||||
|
||||
if (jpeg_data.size() % 2) {
|
||||
fputc(0, f);
|
||||
avi_data.push_back(0);
|
||||
}
|
||||
}
|
||||
|
||||
long cur_pos = ftell(f);
|
||||
long movi_size = cur_pos - movi_size_pos - 4;
|
||||
fseek(f, movi_size_pos, SEEK_SET);
|
||||
write_u32_le(f, movi_size);
|
||||
fseek(f, cur_pos, SEEK_SET);
|
||||
const size_t movi_size = avi_data.size() - movi_size_pos - 4;
|
||||
patch_u32_le(avi_data, movi_size_pos, static_cast<uint32_t>(movi_size));
|
||||
|
||||
fwrite("idx1", 4, 1, f);
|
||||
write_u32_le(f, num_images * 16);
|
||||
write_fourcc(avi_data, "idx1");
|
||||
write_u32_le(avi_data, num_images * 16);
|
||||
for (int i = 0; i < num_images; i++) {
|
||||
fwrite("00dc", 4, 1, f);
|
||||
write_u32_le(f, 0x10);
|
||||
write_u32_le(f, index[i].offset);
|
||||
write_u32_le(f, index[i].size);
|
||||
write_fourcc(avi_data, "00dc");
|
||||
write_u32_le(avi_data, 0x10);
|
||||
write_u32_le(avi_data, index[i].offset);
|
||||
write_u32_le(avi_data, index[i].size);
|
||||
}
|
||||
|
||||
cur_pos = ftell(f);
|
||||
long file_size = cur_pos - riff_size_pos - 4;
|
||||
fseek(f, riff_size_pos, SEEK_SET);
|
||||
write_u32_le(f, file_size);
|
||||
fseek(f, cur_pos, SEEK_SET);
|
||||
const size_t file_size = avi_data.size() - riff_size_pos - 4;
|
||||
patch_u32_le(avi_data, riff_size_pos, static_cast<uint32_t>(file_size));
|
||||
|
||||
return avi_data;
|
||||
}
|
||||
|
||||
int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality);
|
||||
if (avi_data.empty()) {
|
||||
return -1;
|
||||
}
|
||||
if (!write_binary_file_bytes(filename, avi_data)) {
|
||||
perror("Error opening file for writing");
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
#ifdef SD_USE_WEBP
|
||||
int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::vector<uint8_t> create_animated_webp_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
|
||||
if (num_images == 0) {
|
||||
fprintf(stderr, "Error: Image array is empty.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
if (fps <= 0) {
|
||||
fprintf(stderr, "Error: FPS must be positive.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
const int width = static_cast<int>(images[0].width);
|
||||
@@ -855,14 +941,14 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
|
||||
const int channels = static_cast<int>(images[0].channel);
|
||||
if (channels != 1 && channels != 3 && channels != 4) {
|
||||
fprintf(stderr, "Error: Unsupported channel count: %d\n", channels);
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
WebPAnimEncoderOptions anim_options;
|
||||
WebPConfig config;
|
||||
if (!WebPAnimEncoderOptionsInit(&anim_options) || !WebPConfigInit(&config)) {
|
||||
fprintf(stderr, "Error: Failed to initialize WebP animation encoder.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
config.quality = static_cast<float>(quality);
|
||||
@@ -873,13 +959,13 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
|
||||
}
|
||||
if (!WebPValidateConfig(&config)) {
|
||||
fprintf(stderr, "Error: Invalid WebP encoder configuration.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
WebPAnimEncoderPtr enc(WebPAnimEncoderNew(width, height, &anim_options));
|
||||
if (enc == nullptr) {
|
||||
fprintf(stderr, "Error: Could not create WebPAnimEncoder object.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
const int frame_duration_ms = std::max(1, static_cast<int>(std::lround(1000.0 / static_cast<double>(fps))));
|
||||
@@ -889,13 +975,13 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
|
||||
const sd_image_t& image = images[i];
|
||||
if (static_cast<int>(image.width) != width || static_cast<int>(image.height) != height) {
|
||||
fprintf(stderr, "Error: Frame dimensions do not match.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
WebPPictureGuard picture;
|
||||
if (!picture.initialized) {
|
||||
fprintf(stderr, "Error: Failed to initialize WebPPicture.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
picture.picture.use_argb = 1;
|
||||
picture.picture.width = width;
|
||||
@@ -919,12 +1005,12 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
|
||||
|
||||
if (!picture_ok) {
|
||||
fprintf(stderr, "Error: Failed to import frame into WebPPicture.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
if (!WebPAnimEncoderAdd(enc.get(), &picture.picture, timestamp_ms, &config)) {
|
||||
fprintf(stderr, "Error: Failed to add frame to animated WebP: %s\n", WebPAnimEncoderGetError(enc.get()));
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
timestamp_ms += frame_duration_ms;
|
||||
@@ -932,52 +1018,50 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
|
||||
|
||||
if (!WebPAnimEncoderAdd(enc.get(), nullptr, timestamp_ms, nullptr)) {
|
||||
fprintf(stderr, "Error: Failed to finalize animated WebP frames: %s\n", WebPAnimEncoderGetError(enc.get()));
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
WebPDataGuard webp_data;
|
||||
if (!WebPAnimEncoderAssemble(enc.get(), &webp_data.data)) {
|
||||
fprintf(stderr, "Error: Failed to assemble animated WebP: %s\n", WebPAnimEncoderGetError(enc.get()));
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
FilePtr f(fopen(filename, "wb"));
|
||||
if (!f) {
|
||||
return std::vector<uint8_t>(webp_data.data.bytes, webp_data.data.bytes + webp_data.data.size);
|
||||
}
|
||||
|
||||
int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::vector<uint8_t> webp_data = create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality);
|
||||
if (webp_data.empty()) {
|
||||
return -1;
|
||||
}
|
||||
if (!write_binary_file_bytes(filename, webp_data)) {
|
||||
perror("Error opening file for writing");
|
||||
return -1;
|
||||
}
|
||||
if (webp_data.data.size > 0 && fwrite(webp_data.data.bytes, 1, webp_data.data.size, f.get()) != webp_data.data.size) {
|
||||
fprintf(stderr, "Error: Failed to write animated WebP file.\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
#endif
|
||||
|
||||
#ifdef SD_USE_WEBM
|
||||
int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
|
||||
if (num_images == 0) {
|
||||
fprintf(stderr, "Error: Image array is empty.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
if (fps <= 0) {
|
||||
fprintf(stderr, "Error: FPS must be positive.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
const int width = static_cast<int>(images[0].width);
|
||||
const int height = static_cast<int>(images[0].height);
|
||||
if (width <= 0 || height <= 0) {
|
||||
fprintf(stderr, "Error: Invalid frame dimensions.\n");
|
||||
return -1;
|
||||
return {};
|
||||
}
|
||||
|
||||
mkvmuxer::MkvWriter writer;
|
||||
if (!writer.Open(filename)) {
|
||||
fprintf(stderr, "Error: Could not open WebM file for writing.\n");
|
||||
return -1;
|
||||
}
|
||||
MemoryMkvWriter writer;
|
||||
|
||||
const int ret = [&]() -> int {
|
||||
mkvmuxer::Segment segment;
|
||||
@@ -1043,30 +1127,63 @@ int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num
|
||||
}
|
||||
return 0;
|
||||
}();
|
||||
writer.Close();
|
||||
return ret;
|
||||
if (ret != 0) {
|
||||
return {};
|
||||
}
|
||||
return writer.data();
|
||||
}
|
||||
|
||||
int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality);
|
||||
if (webm_data.empty()) {
|
||||
return -1;
|
||||
}
|
||||
if (!write_binary_file_bytes(filename, webm_data)) {
|
||||
perror("Error opening file for writing");
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
#endif
|
||||
|
||||
int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::string path = filename ? filename : "";
|
||||
auto pos = path.find_last_of('.');
|
||||
std::string ext = pos == std::string::npos ? "" : path.substr(pos);
|
||||
for (char& ch : ext) {
|
||||
ch = static_cast<char>(tolower(static_cast<unsigned char>(ch)));
|
||||
std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
|
||||
sd_image_t* images,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality) {
|
||||
std::string format = output_format;
|
||||
std::transform(format.begin(), format.end(), format.begin(),
|
||||
[](unsigned char c) { return static_cast<char>(tolower(c)); });
|
||||
if (!format.empty() && format[0] == '.') {
|
||||
format.erase(format.begin());
|
||||
}
|
||||
|
||||
#ifdef SD_USE_WEBM
|
||||
if (ext == ".webm") {
|
||||
return create_webm_from_sd_images(filename, images, num_images, fps, quality);
|
||||
if (format == "webm") {
|
||||
return create_webm_from_sd_images_to_vector(images, num_images, fps, quality);
|
||||
}
|
||||
#endif
|
||||
|
||||
#ifdef SD_USE_WEBP
|
||||
if (ext == ".webp") {
|
||||
return create_animated_webp_from_sd_images(filename, images, num_images, fps, quality);
|
||||
if (format == "webp") {
|
||||
return create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality);
|
||||
}
|
||||
#endif
|
||||
|
||||
return create_mjpg_avi_from_sd_images(filename, images, num_images, fps, quality);
|
||||
return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality);
|
||||
}
|
||||
|
||||
int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
|
||||
std::string path = filename ? filename : "";
|
||||
auto pos = path.find_last_of('.');
|
||||
std::string ext = pos == std::string::npos ? "" : path.substr(pos);
|
||||
std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality);
|
||||
if (video_data.empty()) {
|
||||
return -1;
|
||||
}
|
||||
if (!write_binary_file_bytes(filename, video_data)) {
|
||||
perror("Error opening file for writing");
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -58,6 +58,10 @@ int create_mjpg_avi_from_sd_images(const char* filename,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
|
||||
#ifdef SD_USE_WEBP
|
||||
int create_animated_webp_from_sd_images(const char* filename,
|
||||
@@ -65,6 +69,10 @@ int create_animated_webp_from_sd_images(const char* filename,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
std::vector<uint8_t> create_animated_webp_from_sd_images_to_vector(sd_image_t* images,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
#endif
|
||||
|
||||
#ifdef SD_USE_WEBM
|
||||
@@ -73,6 +81,10 @@ int create_webm_from_sd_images(const char* filename,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
#endif
|
||||
|
||||
int create_video_from_sd_images(const char* filename,
|
||||
@@ -80,5 +92,10 @@ int create_video_from_sd_images(const char* filename,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
|
||||
sd_image_t* images,
|
||||
int num_images,
|
||||
int fps,
|
||||
int quality = 90);
|
||||
|
||||
#endif // __MEDIA_IO_H__
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
@@ -47,14 +48,40 @@ using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
|
||||
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
|
||||
|
||||
class SDImageOwner {
|
||||
private:
|
||||
static sd_image_t copy_image(const sd_image_t& image) {
|
||||
if (image.data == nullptr) {
|
||||
return {image.width, image.height, image.channel, nullptr};
|
||||
}
|
||||
|
||||
const size_t byte_count = static_cast<size_t>(image.width) * image.height * image.channel;
|
||||
uint8_t* raw_copy = static_cast<uint8_t*>(malloc(byte_count));
|
||||
if (raw_copy == nullptr) {
|
||||
return {0, 0, 0, nullptr};
|
||||
}
|
||||
|
||||
std::memcpy(raw_copy, image.data, byte_count);
|
||||
return {image.width, image.height, image.channel, raw_copy};
|
||||
}
|
||||
|
||||
sd_image_t image_ = {0, 0, 0, nullptr};
|
||||
|
||||
public:
|
||||
SDImageOwner() = default;
|
||||
explicit SDImageOwner(sd_image_t image)
|
||||
: image_(image) {
|
||||
}
|
||||
|
||||
SDImageOwner(const SDImageOwner&) = delete;
|
||||
SDImageOwner& operator=(const SDImageOwner&) = delete;
|
||||
SDImageOwner(const SDImageOwner& other)
|
||||
: image_(copy_image(other.image_)) {
|
||||
}
|
||||
|
||||
SDImageOwner& operator=(const SDImageOwner& other) {
|
||||
if (this != &other) {
|
||||
reset(copy_image(other.image_));
|
||||
}
|
||||
return *this;
|
||||
}
|
||||
|
||||
SDImageOwner(SDImageOwner&& other) noexcept
|
||||
: image_(other.release()) {
|
||||
@@ -77,8 +104,9 @@ public:
|
||||
free(image_.data);
|
||||
image_.data = nullptr;
|
||||
}
|
||||
image_.width = 0;
|
||||
image_.height = 0;
|
||||
image_.width = 0;
|
||||
image_.height = 0;
|
||||
image_.channel = 0;
|
||||
return &image_;
|
||||
}
|
||||
|
||||
@@ -102,12 +130,12 @@ public:
|
||||
}
|
||||
image_ = image;
|
||||
}
|
||||
|
||||
private:
|
||||
sd_image_t image_ = {0, 0, 0, nullptr};
|
||||
};
|
||||
|
||||
class SDImageVec {
|
||||
private:
|
||||
std::vector<sd_image_t> images_;
|
||||
|
||||
public:
|
||||
SDImageVec() = default;
|
||||
|
||||
@@ -164,6 +192,10 @@ public:
|
||||
return images_.empty();
|
||||
}
|
||||
|
||||
int count() const {
|
||||
return static_cast<int>(images_.size());
|
||||
}
|
||||
|
||||
explicit operator bool() const {
|
||||
return !images_.empty();
|
||||
}
|
||||
@@ -199,9 +231,6 @@ public:
|
||||
}
|
||||
images_.clear();
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<sd_image_t> images_;
|
||||
};
|
||||
|
||||
#endif // __EXAMPLE_RESOURCE_OWNERS_H__
|
||||
|
||||
@@ -277,6 +277,7 @@ protected:
|
||||
int64_t context_dim;
|
||||
int64_t n_head;
|
||||
int64_t d_head;
|
||||
bool xtra_dim = false;
|
||||
|
||||
public:
|
||||
CrossAttention(int64_t query_dim,
|
||||
@@ -288,7 +289,11 @@ public:
|
||||
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));
|
||||
@@ -313,10 +318,16 @@ public:
|
||||
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]
|
||||
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]
|
||||
|
||||
@@ -256,15 +256,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
return true;
|
||||
}
|
||||
|
||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
|
||||
tokenize_with_trigger_token(std::string text,
|
||||
int num_input_imgs,
|
||||
int32_t image_token,
|
||||
bool padding = false) {
|
||||
return tokenize_with_trigger_token(text, num_input_imgs, image_token,
|
||||
text_model->model.n_token, padding);
|
||||
}
|
||||
|
||||
std::vector<int> convert_token_to_id(std::string text) {
|
||||
auto on_new_token_cb = [&](std::string& str, std::vector<int32_t>& bpe_tokens) -> bool {
|
||||
auto iter = embedding_map.find(str);
|
||||
@@ -288,9 +279,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
|
||||
tokenize_with_trigger_token(std::string text,
|
||||
int num_input_imgs,
|
||||
int32_t image_token,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
int32_t image_token) {
|
||||
auto parsed_attention = parse_prompt_attention(text);
|
||||
|
||||
{
|
||||
@@ -377,7 +366,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
// tokens.insert(tokens.begin(), tokenizer.BOS_TOKEN_ID);
|
||||
// weights.insert(weights.begin(), 1.0);
|
||||
|
||||
tokenizer.pad_tokens(tokens, weights, max_length, padding);
|
||||
tokenizer.pad_tokens(tokens, &weights, nullptr, text_model->model.n_token, text_model->model.n_token, true);
|
||||
int offset = pm_version == PM_VERSION_2 ? 2 * num_input_imgs : num_input_imgs;
|
||||
for (int i = 0; i < tokens.size(); i++) {
|
||||
// if (class_idx + 1 <= i && i < class_idx + 1 + 2*num_input_imgs) // photomaker V2 has num_tokens(=2)*num_input_imgs
|
||||
@@ -403,13 +392,9 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
}
|
||||
|
||||
std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
|
||||
bool padding = false) {
|
||||
return tokenize(text, text_model->model.n_token, padding);
|
||||
}
|
||||
|
||||
std::pair<std::vector<int>, std::vector<float>> tokenize(std::string text,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 0,
|
||||
bool allow_overflow_expand = true) {
|
||||
auto parsed_attention = parse_prompt_attention(text);
|
||||
|
||||
{
|
||||
@@ -460,7 +445,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
|
||||
tokenizer.pad_tokens(tokens, weights, max_length, padding);
|
||||
tokenizer.pad_tokens(tokens, &weights, nullptr, min_length, max_length, allow_overflow_expand);
|
||||
|
||||
// for (int i = 0; i < tokens.size(); i++) {
|
||||
// std::cout << tokens[i] << ":" << weights[i] << ", ";
|
||||
@@ -603,8 +588,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
GGML_ASSERT(image_tokens.size() == 1);
|
||||
auto tokens_and_weights = tokenize_with_trigger_token(conditioner_params.text,
|
||||
conditioner_params.num_input_imgs,
|
||||
image_tokens[0],
|
||||
true);
|
||||
image_tokens[0]);
|
||||
std::vector<int>& tokens = std::get<0>(tokens_and_weights);
|
||||
std::vector<float>& weights = std::get<1>(tokens_and_weights);
|
||||
std::vector<bool>& clsm = std::get<2>(tokens_and_weights);
|
||||
@@ -630,7 +614,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
std::string remove_trigger_from_prompt(const std::string& prompt) override {
|
||||
auto image_tokens = convert_token_to_id(trigger_word);
|
||||
GGML_ASSERT(image_tokens.size() == 1);
|
||||
auto tokens_and_weights = tokenize(prompt, false);
|
||||
auto tokens_and_weights = tokenize(prompt);
|
||||
std::vector<int>& tokens = tokens_and_weights.first;
|
||||
auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
|
||||
GGML_ASSERT(it != tokens.end()); // prompt must have trigger word
|
||||
@@ -640,7 +624,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
|
||||
SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) override {
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, true);
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true);
|
||||
std::vector<int>& tokens = tokens_and_weights.first;
|
||||
std::vector<float>& weights = tokens_and_weights.second;
|
||||
return get_learned_condition_common(n_threads,
|
||||
@@ -822,8 +806,9 @@ struct SD3CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
|
||||
std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 0,
|
||||
bool allow_overflow_expand = true) {
|
||||
auto parsed_attention = parse_prompt_attention(text);
|
||||
|
||||
{
|
||||
@@ -860,20 +845,20 @@ struct SD3CLIPEmbedder : public Conditioner {
|
||||
clip_g_weights.insert(clip_g_weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
if (t5) {
|
||||
std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
|
||||
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
|
||||
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
|
||||
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
}
|
||||
|
||||
if (clip_l) {
|
||||
clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, max_length, padding);
|
||||
clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, min_length, max_length, allow_overflow_expand);
|
||||
}
|
||||
if (clip_g) {
|
||||
clip_g_tokenizer.pad_tokens(clip_g_tokens, clip_g_weights, max_length, padding);
|
||||
clip_g_tokenizer.pad_tokens(clip_g_tokens, &clip_g_weights, nullptr, min_length, max_length, allow_overflow_expand);
|
||||
}
|
||||
if (t5) {
|
||||
t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding);
|
||||
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
|
||||
}
|
||||
|
||||
// for (int i = 0; i < clip_l_tokens.size(); i++) {
|
||||
@@ -1056,7 +1041,7 @@ struct SD3CLIPEmbedder : public Conditioner {
|
||||
|
||||
SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) override {
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, 77, true);
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, 77, 77, true);
|
||||
return get_learned_condition_common(n_threads,
|
||||
tokens_and_weights,
|
||||
conditioner_params.clip_skip,
|
||||
@@ -1158,8 +1143,8 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
}
|
||||
|
||||
std::vector<std::pair<std::vector<int>, std::vector<float>>> tokenize(std::string text,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 0) {
|
||||
auto parsed_attention = parse_prompt_attention(text);
|
||||
|
||||
{
|
||||
@@ -1189,17 +1174,17 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
clip_l_weights.insert(clip_l_weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
if (t5) {
|
||||
std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
|
||||
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
|
||||
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
|
||||
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
}
|
||||
|
||||
if (clip_l) {
|
||||
clip_l_tokenizer.pad_tokens(clip_l_tokens, clip_l_weights, 77, padding);
|
||||
clip_l_tokenizer.pad_tokens(clip_l_tokens, &clip_l_weights, nullptr, 77, 77, true);
|
||||
}
|
||||
if (t5) {
|
||||
t5_tokenizer.pad_tokens(t5_tokens, t5_weights, nullptr, max_length, padding);
|
||||
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr, min_length, max_length, true);
|
||||
}
|
||||
|
||||
// for (int i = 0; i < clip_l_tokens.size(); i++) {
|
||||
@@ -1300,7 +1285,7 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
|
||||
SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) override {
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
|
||||
return get_learned_condition_common(n_threads,
|
||||
tokens_and_weights,
|
||||
conditioner_params.clip_skip,
|
||||
@@ -1377,8 +1362,8 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
|
||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 0) {
|
||||
auto parsed_attention = parse_prompt_attention(text);
|
||||
|
||||
{
|
||||
@@ -1403,12 +1388,15 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
const std::string& curr_text = item.first;
|
||||
float curr_weight = item.second;
|
||||
|
||||
std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
|
||||
std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
|
||||
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
|
||||
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
|
||||
t5_tokenizer.pad_tokens(t5_tokens, t5_weights, &t5_mask, max_length, padding);
|
||||
t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, &t5_mask, min_length, max_length, true);
|
||||
for (auto& mask_value : t5_mask) {
|
||||
mask_value = mask_value > 0.0f ? 0.0f : -HUGE_VALF;
|
||||
}
|
||||
}
|
||||
return {t5_tokens, t5_weights, t5_mask};
|
||||
}
|
||||
@@ -1496,7 +1484,7 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
|
||||
SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) override {
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, true);
|
||||
auto tokens_and_weights = tokenize(conditioner_params.text, chunk_len, chunk_len);
|
||||
return get_learned_condition_common(n_threads,
|
||||
tokens_and_weights,
|
||||
conditioner_params.clip_skip,
|
||||
@@ -1505,14 +1493,14 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
};
|
||||
|
||||
struct AnimaConditioner : public Conditioner {
|
||||
std::shared_ptr<LLM::BPETokenizer> qwen_tokenizer;
|
||||
std::shared_ptr<BPETokenizer> qwen_tokenizer;
|
||||
T5UniGramTokenizer t5_tokenizer;
|
||||
std::shared_ptr<LLM::LLMRunner> llm;
|
||||
|
||||
AnimaConditioner(ggml_backend_t backend,
|
||||
bool offload_params_to_cpu,
|
||||
const String2TensorStorage& tensor_storage_map = {}) {
|
||||
qwen_tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
|
||||
qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
|
||||
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
|
||||
backend,
|
||||
offload_params_to_cpu,
|
||||
@@ -1578,7 +1566,7 @@ struct AnimaConditioner : public Conditioner {
|
||||
for (const auto& item : parsed_attention) {
|
||||
const std::string& curr_text = item.first;
|
||||
float curr_weight = item.second;
|
||||
std::vector<int> curr_tokens = t5_tokenizer.Encode(curr_text, true);
|
||||
std::vector<int> curr_tokens = t5_tokenizer.tokenize(curr_text, nullptr, true);
|
||||
t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
|
||||
t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
@@ -1620,7 +1608,7 @@ struct AnimaConditioner : public Conditioner {
|
||||
|
||||
struct LLMEmbedder : public Conditioner {
|
||||
SDVersion version;
|
||||
std::shared_ptr<LLM::BPETokenizer> tokenizer;
|
||||
std::shared_ptr<BPETokenizer> tokenizer;
|
||||
std::shared_ptr<LLM::LLMRunner> llm;
|
||||
|
||||
LLMEmbedder(ggml_backend_t backend,
|
||||
@@ -1633,13 +1621,15 @@ struct LLMEmbedder : public Conditioner {
|
||||
LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL;
|
||||
if (version == VERSION_FLUX2) {
|
||||
arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
|
||||
} else if (sd_version_is_ernie_image(version)) {
|
||||
arch = LLM::LLMArch::MINISTRAL_3_3B;
|
||||
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
|
||||
arch = LLM::LLMArch::QWEN3;
|
||||
}
|
||||
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2) {
|
||||
tokenizer = std::make_shared<LLM::MistralTokenizer>();
|
||||
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
|
||||
tokenizer = std::make_shared<MistralTokenizer>();
|
||||
} else {
|
||||
tokenizer = std::make_shared<LLM::Qwen2Tokenizer>();
|
||||
tokenizer = std::make_shared<Qwen2Tokenizer>();
|
||||
}
|
||||
llm = std::make_shared<LLM::LLMRunner>(arch,
|
||||
backend,
|
||||
@@ -1677,20 +1667,24 @@ struct LLMEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<std::vector<int>, std::vector<float>> tokenize(std::string text,
|
||||
const std::pair<int, int>& attn_range,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
|
||||
const std::pair<int, int>& attn_range,
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 100000000) {
|
||||
std::vector<std::pair<std::string, float>> parsed_attention;
|
||||
if (attn_range.first >= 0 && attn_range.second > 0) {
|
||||
parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
|
||||
if (attn_range.first > 0) {
|
||||
parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
|
||||
}
|
||||
if (attn_range.second - attn_range.first > 0) {
|
||||
auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
|
||||
parsed_attention.insert(parsed_attention.end(),
|
||||
new_parsed_attention.begin(),
|
||||
new_parsed_attention.end());
|
||||
}
|
||||
parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
|
||||
if (attn_range.second < text.size()) {
|
||||
parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
|
||||
}
|
||||
} else {
|
||||
parsed_attention.emplace_back(text, 1.f);
|
||||
}
|
||||
@@ -1710,39 +1704,34 @@ struct LLMEmbedder : public Conditioner {
|
||||
for (const auto& item : parsed_attention) {
|
||||
const std::string& curr_text = item.first;
|
||||
float curr_weight = item.second;
|
||||
std::vector<int> curr_tokens = tokenizer->tokenize(curr_text, nullptr);
|
||||
std::vector<int> curr_tokens = tokenizer->encode(curr_text, nullptr);
|
||||
tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
|
||||
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
|
||||
tokenizer->pad_tokens(tokens, weights, max_length, padding);
|
||||
std::vector<float> mask;
|
||||
tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
|
||||
|
||||
// for (int i = 0; i < tokens.size(); i++) {
|
||||
// std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl;
|
||||
// }
|
||||
// std::cout << std::endl;
|
||||
|
||||
return {tokens, weights};
|
||||
return {tokens, weights, mask};
|
||||
}
|
||||
|
||||
sd::Tensor<float> encode_prompt(int n_threads,
|
||||
const std::string prompt,
|
||||
const std::pair<int, int>& prompt_attn_range,
|
||||
int max_length,
|
||||
int min_length,
|
||||
int hidden_states_min_length,
|
||||
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
|
||||
const std::set<int>& out_layers,
|
||||
int prompt_template_encode_start_idx) {
|
||||
auto tokens_and_weights = tokenize(prompt, prompt_attn_range);
|
||||
auto& tokens = std::get<0>(tokens_and_weights);
|
||||
auto& weights = std::get<1>(tokens_and_weights);
|
||||
std::vector<float> mask;
|
||||
|
||||
if (max_length > 0 && tokens.size() < max_length) {
|
||||
mask.insert(mask.end(), tokens.size(), 1.f);
|
||||
mask.insert(mask.end(), max_length - tokens.size(), 0.f);
|
||||
tokenizer->pad_tokens(tokens, weights, max_length, true);
|
||||
}
|
||||
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length);
|
||||
auto& tokens = std::get<0>(tokens_weights_mask);
|
||||
auto& weights = std::get<1>(tokens_weights_mask);
|
||||
auto& mask = std::get<2>(tokens_weights_mask);
|
||||
|
||||
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
|
||||
sd::Tensor<float> attention_mask;
|
||||
@@ -1769,9 +1758,9 @@ struct LLMEmbedder : public Conditioner {
|
||||
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
|
||||
|
||||
int64_t zero_pad_len = 0;
|
||||
if (min_length > 0) {
|
||||
if (hidden_states.shape()[1] - prompt_template_encode_start_idx < min_length) {
|
||||
zero_pad_len = min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
|
||||
if (hidden_states_min_length > 0) {
|
||||
if (hidden_states.shape()[1] - prompt_template_encode_start_idx < hidden_states_min_length) {
|
||||
zero_pad_len = hidden_states_min_length - hidden_states.shape()[1] + prompt_template_encode_start_idx;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1798,8 +1787,8 @@ struct LLMEmbedder : public Conditioner {
|
||||
std::vector<std::pair<int, int>> extra_prompts_attn_range;
|
||||
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
|
||||
int prompt_template_encode_start_idx = 34;
|
||||
int max_length = 0; // pad tokens
|
||||
int min_length = 0; // zero pad hidden_states
|
||||
int min_length = 0; // pad tokens
|
||||
int hidden_states_min_length = 0; // zero pad hidden_states
|
||||
std::set<int> out_layers;
|
||||
|
||||
int64_t t0 = ggml_time_ms();
|
||||
@@ -1874,7 +1863,7 @@ struct LLMEmbedder : public Conditioner {
|
||||
}
|
||||
} else if (version == VERSION_FLUX2) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
min_length = 512;
|
||||
hidden_states_min_length = 512;
|
||||
out_layers = {10, 20, 30};
|
||||
|
||||
prompt = "[SYSTEM_PROMPT]You are an AI that reasons about image descriptions. You give structured responses focusing on object relationships, object\nattribution and actions without speculation.[/SYSTEM_PROMPT][INST]";
|
||||
@@ -1884,6 +1873,13 @@ struct LLMEmbedder : public Conditioner {
|
||||
prompt_attn_range.second = static_cast<int>(prompt.size());
|
||||
|
||||
prompt += "[/INST]";
|
||||
} else if (sd_version_is_ernie_image(version)) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
out_layers = {25}; // -2
|
||||
|
||||
prompt_attn_range.first = 0;
|
||||
prompt += conditioner_params.text;
|
||||
prompt_attn_range.second = static_cast<int>(prompt.size());
|
||||
} else if (sd_version_is_z_image(version)) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
out_layers = {35}; // -2
|
||||
@@ -1907,7 +1903,7 @@ struct LLMEmbedder : public Conditioner {
|
||||
}
|
||||
} else if (version == VERSION_FLUX2_KLEIN) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
max_length = 512;
|
||||
min_length = 512;
|
||||
out_layers = {9, 18, 27};
|
||||
|
||||
prompt = "<|im_start|>user\n";
|
||||
@@ -1919,7 +1915,7 @@ struct LLMEmbedder : public Conditioner {
|
||||
prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
|
||||
} else if (version == VERSION_OVIS_IMAGE) {
|
||||
prompt_template_encode_start_idx = 28;
|
||||
max_length = prompt_template_encode_start_idx + 256;
|
||||
min_length = prompt_template_encode_start_idx + 256;
|
||||
|
||||
prompt = "<|im_start|>user\nDescribe the image by detailing the color, quantity, text, shape, size, texture, spatial relationships of the objects and background:";
|
||||
|
||||
@@ -1935,8 +1931,8 @@ struct LLMEmbedder : public Conditioner {
|
||||
auto hidden_states = encode_prompt(n_threads,
|
||||
prompt,
|
||||
prompt_attn_range,
|
||||
max_length,
|
||||
min_length,
|
||||
hidden_states_min_length,
|
||||
image_embeds,
|
||||
out_layers,
|
||||
prompt_template_encode_start_idx);
|
||||
@@ -1945,8 +1941,8 @@ struct LLMEmbedder : public Conditioner {
|
||||
auto extra_hidden_states = encode_prompt(n_threads,
|
||||
extra_prompts[i],
|
||||
extra_prompts_attn_range[i],
|
||||
max_length,
|
||||
min_length,
|
||||
hidden_states_min_length,
|
||||
image_embeds,
|
||||
out_layers,
|
||||
prompt_template_encode_start_idx);
|
||||
|
||||
+279
-63
@@ -658,32 +658,22 @@ inline float time_snr_shift(float alpha, float t) {
|
||||
}
|
||||
|
||||
struct DiscreteFlowDenoiser : public Denoiser {
|
||||
float sigmas[TIMESTEPS];
|
||||
float shift = 3.0f;
|
||||
|
||||
float sigma_data = 1.0f;
|
||||
|
||||
DiscreteFlowDenoiser(float shift = 3.0f) {
|
||||
set_shift(shift);
|
||||
}
|
||||
|
||||
void set_parameters() {
|
||||
for (int i = 0; i < TIMESTEPS; i++) {
|
||||
sigmas[i] = t_to_sigma(static_cast<float>(i));
|
||||
}
|
||||
}
|
||||
|
||||
void set_shift(float shift) {
|
||||
this->shift = shift;
|
||||
set_parameters();
|
||||
}
|
||||
|
||||
float sigma_min() override {
|
||||
return sigmas[0];
|
||||
return t_to_sigma(0);
|
||||
}
|
||||
|
||||
float sigma_max() override {
|
||||
return sigmas[TIMESTEPS - 1];
|
||||
return t_to_sigma(TIMESTEPS - 1);
|
||||
}
|
||||
|
||||
float sigma_to_t(float sigma) override {
|
||||
@@ -963,8 +953,9 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
|
||||
float t_next = t_fn(sigma_down);
|
||||
float h = t_next - t;
|
||||
float s = t + 0.5f * h;
|
||||
sd::Tensor<float> x2 = (sigma_fn(s) / sigma_fn(t)) * x - (exp(-h * 0.5f) - 1) * denoised;
|
||||
auto denoised2_opt = model(x2, sigmas[i + 1], i + 1);
|
||||
float sigma_s = sigma_fn(s);
|
||||
sd::Tensor<float> x2 = (sigma_s / sigma_fn(t)) * x - (exp(-h * 0.5f) - 1) * denoised;
|
||||
auto denoised2_opt = model(x2, sigma_s, i + 1);
|
||||
if (denoised2_opt.empty()) {
|
||||
return {};
|
||||
}
|
||||
@@ -979,6 +970,100 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
|
||||
return x;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
|
||||
sd::Tensor<float> x,
|
||||
const std::vector<float>& sigmas,
|
||||
std::shared_ptr<RNG> rng,
|
||||
float eta = 1.0f) {
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
float sigma = sigmas[i];
|
||||
float sigma_to = sigmas[i + 1];
|
||||
|
||||
bool opt_first_step = (1.0 - sigma < 1e-6);
|
||||
|
||||
auto denoised_opt = model(x, sigma, (opt_first_step ? 1 : -1) * (i + 1));
|
||||
if (denoised_opt.empty()) {
|
||||
return {};
|
||||
}
|
||||
sd::Tensor<float> denoised = std::move(denoised_opt);
|
||||
|
||||
if (sigma_to == 0.0f) {
|
||||
// Euler method (final step, no noise)
|
||||
// sigma_to == 0 --> sigma_down = 0, so:
|
||||
// x + d * (sigma_down - sigma)
|
||||
// = x + ((x - denoised) / sigma) * (sigma_down - sigma)
|
||||
// = x + ((x - denoised) / sigma) * ( 0 - sigma)
|
||||
// = x + ((x - denoised) ) * -1
|
||||
// = x -x + denoised
|
||||
x = denoised;
|
||||
|
||||
} else {
|
||||
auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step_flow(sigma, sigma_to, eta);
|
||||
sd::Tensor<float> D_i;
|
||||
|
||||
if (opt_first_step) {
|
||||
// the reformulated exp_s calc already accounts for this, but we can avoid
|
||||
// a redundant model call for the typical sigma 1 at the first step:
|
||||
// exp_s = sqrt((1-sigma)/sigma * (1-sigma_down)/sigma_down)
|
||||
// = sqrt((1- 1)/ 1 * (1-sigma_down)/sigma_down)
|
||||
// = 0
|
||||
// so sigma_s = 1 = sigma, and sigma_s_i_ratio = sigma_s / sigma = 1
|
||||
// u = (x*sigma_s_i_ratio)+(denoised*(1.0f-sigma_s_i_ratio))
|
||||
// = (x*1)+(denoised*0) = x
|
||||
// so D_i = model(u, sigma_s, i + 1)
|
||||
// = model(x, sigma, i + 1)
|
||||
// = denoised
|
||||
D_i = denoised;
|
||||
|
||||
} else {
|
||||
float sigma_s;
|
||||
|
||||
// ref implementation would be:
|
||||
// auto lambda_fn = [](float sigma) -> float {
|
||||
// return std::log((1.0f - sigma) / sigma); };
|
||||
// auto sigma_fn = [](float lbda) -> float {
|
||||
// return 1.0f / (std::exp(lbda) + 1.0f); };
|
||||
// t_i = lambda_fn(sigma);
|
||||
// t_down = lambda_fn(sigma_down);
|
||||
// float r = 0.5f;
|
||||
// h = t_down - t_i;
|
||||
// s = t_i + r * h;
|
||||
// sigma_s = sigma_fn(s);
|
||||
|
||||
// assuming r is constant, we sidestep the singularity at sigma -> 1 by:
|
||||
// s = 0.5 * (lambda_fn(sigma) + lambda_fn(sigma_down))
|
||||
// = 0.5 * (log((1-sigma)/sigma) + log((1-sigma_down)/sigma_down))
|
||||
// = 0.5 * log(((1-sigma)/sigma) * ((1-sigma_down)/sigma_down))
|
||||
// = log(sqrt (((1-sigma)/sigma) * ((1-sigma_down)/sigma_down)))
|
||||
// so exp(s) = sqrt((1-sigma)/sigma * (1-sigma_down)/sigma_down)
|
||||
// and sigma_s = sigma_fn(s) = 1.0f / (exp(s) + 1.0f)
|
||||
|
||||
float exp_s = std::sqrt(((1 - sigma) / sigma) * ((1 - sigma_down) / sigma_down));
|
||||
sigma_s = 1.0f / (exp_s + 1.0f);
|
||||
|
||||
float sigma_s_i_ratio = sigma_s / sigma;
|
||||
sd::Tensor<float> u = (x * sigma_s_i_ratio) + (denoised * (1.0f - sigma_s_i_ratio));
|
||||
|
||||
auto denoised2_opt = model(u, sigma_s, i + 1);
|
||||
if (denoised2_opt.empty()) {
|
||||
return {};
|
||||
}
|
||||
D_i = std::move(denoised2_opt);
|
||||
}
|
||||
|
||||
float sigma_down_i_ratio = sigma_down / sigma;
|
||||
x = (x * sigma_down_i_ratio) + (D_i * (1.0f - sigma_down_i_ratio));
|
||||
|
||||
if (sigma_to > 0.0f && eta > 0.0f) {
|
||||
x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return x;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> sample_dpmpp_2m(denoise_cb_t model,
|
||||
sd::Tensor<float> x,
|
||||
const std::vector<float>& sigmas) {
|
||||
@@ -1050,7 +1135,8 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
|
||||
static sd::Tensor<float> sample_lcm(denoise_cb_t model,
|
||||
sd::Tensor<float> x,
|
||||
const std::vector<float>& sigmas,
|
||||
std::shared_ptr<RNG> rng) {
|
||||
std::shared_ptr<RNG> rng,
|
||||
bool is_flow_denoiser) {
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
auto denoised_opt = model(x, sigmas[i], i + 1);
|
||||
@@ -1059,6 +1145,9 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
|
||||
}
|
||||
x = std::move(denoised_opt);
|
||||
if (sigmas[i + 1] > 0) {
|
||||
if (is_flow_denoiser) {
|
||||
x *= (1 - sigmas[i + 1]);
|
||||
}
|
||||
x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
|
||||
}
|
||||
}
|
||||
@@ -1295,37 +1384,151 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
|
||||
return x;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
|
||||
sd::Tensor<float> x,
|
||||
std::vector<float> sigmas,
|
||||
std::shared_ptr<RNG> rng,
|
||||
bool is_flow_denoiser,
|
||||
float eta) {
|
||||
constexpr int max_stage = 3;
|
||||
constexpr int num_integration_points = 200;
|
||||
constexpr float num_integration_points_f = 200.0f;
|
||||
float s_noise = eta;
|
||||
|
||||
auto er_sde_flow_sigma = [](float sigma) -> float {
|
||||
sigma = std::max(sigma, 1e-6f);
|
||||
sigma = std::min(sigma, 1.0f - 1e-4f);
|
||||
return sigma;
|
||||
};
|
||||
|
||||
auto sigma_to_er_sde_lambda = [&](float sigma, bool is_flow_denoiser) -> float {
|
||||
if (is_flow_denoiser) {
|
||||
sigma = er_sde_flow_sigma(sigma);
|
||||
return sigma / std::max(1.0f - sigma, 1e-6f);
|
||||
}
|
||||
return std::max(sigma, 1e-6f);
|
||||
};
|
||||
|
||||
auto sigma_to_er_sde_alpha = [&](float sigma, bool is_flow_denoiser) -> float {
|
||||
if (is_flow_denoiser) {
|
||||
sigma = er_sde_flow_sigma(sigma);
|
||||
return 1.0f - sigma;
|
||||
}
|
||||
return 1.0f;
|
||||
};
|
||||
|
||||
auto er_sde_noise_scaler = [](float x) -> float {
|
||||
x = std::max(x, 0.0f);
|
||||
return x * (std::exp(std::pow(x, 0.3f)) + 10.0f);
|
||||
};
|
||||
|
||||
if (is_flow_denoiser) {
|
||||
for (size_t i = 0; i + 1 < sigmas.size(); ++i) {
|
||||
if (sigmas[i] > 1.0f) {
|
||||
sigmas[i] = er_sde_flow_sigma(sigmas[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<float> er_lambdas(sigmas.size(), 0.0f);
|
||||
for (size_t i = 0; i < sigmas.size(); ++i) {
|
||||
er_lambdas[i] = sigma_to_er_sde_lambda(sigmas[i], is_flow_denoiser);
|
||||
}
|
||||
|
||||
sd::Tensor<float> old_denoised = x;
|
||||
sd::Tensor<float> old_denoised_d = x;
|
||||
bool have_old_denoised = false;
|
||||
bool have_old_denoised_d = false;
|
||||
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
sd::Tensor<float> denoised = model(x, sigmas[i], i + 1);
|
||||
if (denoised.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
int stage_used = std::min(max_stage, i + 1);
|
||||
|
||||
if (sigmas[i + 1] == 0.0f) {
|
||||
x = denoised;
|
||||
} else {
|
||||
float er_lambda_s = er_lambdas[i];
|
||||
float er_lambda_t = er_lambdas[i + 1];
|
||||
float alpha_s = sigma_to_er_sde_alpha(sigmas[i], is_flow_denoiser);
|
||||
float alpha_t = sigma_to_er_sde_alpha(sigmas[i + 1], is_flow_denoiser);
|
||||
float scaled_s = er_sde_noise_scaler(er_lambda_s);
|
||||
float scaled_t = er_sde_noise_scaler(er_lambda_t);
|
||||
float r_alpha = alpha_s > 0.0f ? alpha_t / alpha_s : 0.0f;
|
||||
float r = scaled_s > 0.0f ? scaled_t / scaled_s : 0.0f;
|
||||
|
||||
x = r_alpha * r * x + alpha_t * (1.0f - r) * denoised;
|
||||
|
||||
if (stage_used >= 2 && have_old_denoised) {
|
||||
float dt = er_lambda_t - er_lambda_s;
|
||||
float lambda_step_size = -dt / num_integration_points_f;
|
||||
float s = 0.0f;
|
||||
float s_u = 0.0f;
|
||||
|
||||
for (int p = 0; p < num_integration_points; ++p) {
|
||||
float lambda_pos = er_lambda_t + p * lambda_step_size;
|
||||
float scaled_pos = er_sde_noise_scaler(lambda_pos);
|
||||
if (scaled_pos <= 0.0f) {
|
||||
continue;
|
||||
}
|
||||
|
||||
s += 1.0f / scaled_pos;
|
||||
if (stage_used >= 3 && have_old_denoised_d) {
|
||||
s_u += (lambda_pos - er_lambda_s) / scaled_pos;
|
||||
}
|
||||
}
|
||||
|
||||
s *= lambda_step_size;
|
||||
|
||||
float denom_d = er_lambda_s - er_lambdas[i - 1];
|
||||
if (std::fabs(denom_d) > 1e-12f) {
|
||||
float coeff_d = alpha_t * (dt + s * scaled_t);
|
||||
sd::Tensor<float> denoised_d = (denoised - old_denoised) / denom_d;
|
||||
x += coeff_d * denoised_d;
|
||||
|
||||
if (stage_used >= 3 && have_old_denoised_d) {
|
||||
float denom_u = (er_lambda_s - er_lambdas[i - 2]) * 0.5f;
|
||||
if (std::fabs(denom_u) > 1e-12f) {
|
||||
s_u *= lambda_step_size;
|
||||
float coeff_u = alpha_t * (0.5f * dt * dt + s_u * scaled_t);
|
||||
sd::Tensor<float> denoised_u = (denoised_d - old_denoised_d) / denom_u;
|
||||
x += coeff_u * denoised_u;
|
||||
}
|
||||
}
|
||||
|
||||
old_denoised_d = denoised_d;
|
||||
have_old_denoised_d = true;
|
||||
}
|
||||
}
|
||||
|
||||
float noise_scale_sq = er_lambda_t * er_lambda_t - er_lambda_s * er_lambda_s * r * r;
|
||||
if (s_noise > 0.0f && noise_scale_sq > 0.0f) {
|
||||
float noise_scale = alpha_t * std::sqrt(std::max(noise_scale_sq, 0.0f));
|
||||
x += sd::Tensor<float>::randn_like(x, rng) * noise_scale;
|
||||
}
|
||||
}
|
||||
|
||||
old_denoised = denoised;
|
||||
have_old_denoised = true;
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
|
||||
sd::Tensor<float> x,
|
||||
const std::vector<float>& sigmas,
|
||||
std::shared_ptr<RNG> rng,
|
||||
float eta) {
|
||||
float beta_start = 0.00085f;
|
||||
float beta_end = 0.0120f;
|
||||
std::vector<double> alphas_cumprod(TIMESTEPS);
|
||||
std::vector<double> compvis_sigmas(TIMESTEPS);
|
||||
for (int i = 0; i < TIMESTEPS; i++) {
|
||||
alphas_cumprod[i] =
|
||||
(i == 0 ? 1.0f : alphas_cumprod[i - 1]) *
|
||||
(1.0f -
|
||||
std::pow(sqrtf(beta_start) +
|
||||
(sqrtf(beta_end) - sqrtf(beta_start)) *
|
||||
((float)i / (TIMESTEPS - 1)),
|
||||
2));
|
||||
compvis_sigmas[i] =
|
||||
std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
|
||||
}
|
||||
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
int timestep = static_cast<int>(roundf(TIMESTEPS - i * ((float)TIMESTEPS / steps))) - 1;
|
||||
int prev_timestep = timestep - TIMESTEPS / steps;
|
||||
float sigma = static_cast<float>(compvis_sigmas[timestep]);
|
||||
if (i == 0) {
|
||||
x *= std::sqrt(sigma * sigma + 1) / sigma;
|
||||
} else {
|
||||
x *= std::sqrt(sigma * sigma + 1);
|
||||
}
|
||||
|
||||
float sigma = sigmas[i];
|
||||
float sigma_to = sigmas[i + 1];
|
||||
|
||||
auto model_output_opt = model(x, sigma, i + 1);
|
||||
if (model_output_opt.empty()) {
|
||||
@@ -1334,8 +1537,8 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
|
||||
sd::Tensor<float> model_output = std::move(model_output_opt);
|
||||
model_output = (x - model_output) * (1.0f / sigma);
|
||||
|
||||
float alpha_prod_t = static_cast<float>(alphas_cumprod[timestep]);
|
||||
float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
|
||||
float alpha_prod_t = 1.0f / (sigma * sigma + 1.0f);
|
||||
float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
|
||||
float beta_prod_t = 1.0f - alpha_prod_t;
|
||||
|
||||
sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
|
||||
@@ -1347,12 +1550,13 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
|
||||
(1.0f - alpha_prod_t / alpha_prod_t_prev);
|
||||
float std_dev_t = eta * std::sqrt(variance);
|
||||
|
||||
x = std::sqrt(alpha_prod_t_prev) * pred_original_sample +
|
||||
std::sqrt(1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) * model_output;
|
||||
x = pred_original_sample +
|
||||
std::sqrt((1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2))/ alpha_prod_t_prev) * model_output;
|
||||
|
||||
if (eta > 0) {
|
||||
x += std_dev_t * sd::Tensor<float>::randn_like(x, rng);
|
||||
x+= std_dev_t / std::sqrt(alpha_prod_t_prev) * sd::Tensor<float>::randn_like(x, rng);
|
||||
}
|
||||
|
||||
}
|
||||
return x;
|
||||
}
|
||||
@@ -1378,19 +1582,25 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
||||
std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]);
|
||||
}
|
||||
|
||||
int original_steps = 50;
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
int timestep = TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor(i * ((float)original_steps / steps));
|
||||
int prev_timestep = i >= steps - 1 ? 0 : TIMESTEPS - 1 - (TIMESTEPS / original_steps) * (int)floor((i + 1) * ((float)original_steps / steps));
|
||||
int timestep_s = (int)floor((1 - eta) * prev_timestep);
|
||||
float sigma = static_cast<float>(compvis_sigmas[timestep]);
|
||||
|
||||
if (i == 0) {
|
||||
x *= std::sqrt(sigma * sigma + 1) / sigma;
|
||||
} else {
|
||||
x *= std::sqrt(sigma * sigma + 1);
|
||||
auto get_timestep_from_sigma = [&](float s) -> int {
|
||||
auto it = std::lower_bound(compvis_sigmas.begin(), compvis_sigmas.end(), s);
|
||||
if (it == compvis_sigmas.begin()) return 0;
|
||||
if (it == compvis_sigmas.end()) return TIMESTEPS - 1;
|
||||
int idx_high = static_cast<int>(std::distance(compvis_sigmas.begin(), it));
|
||||
int idx_low = idx_high - 1;
|
||||
if (std::abs(compvis_sigmas[idx_high] - s) < std::abs(compvis_sigmas[idx_low] - s)) {
|
||||
return idx_high;
|
||||
}
|
||||
return idx_low;
|
||||
};
|
||||
|
||||
int steps = static_cast<int>(sigmas.size()) - 1;
|
||||
for (int i = 0; i < steps; i++) {
|
||||
|
||||
float sigma_to = sigmas[i + 1];
|
||||
int prev_timestep = get_timestep_from_sigma(sigma_to);
|
||||
int timestep_s = (int)floor((1 - eta) * prev_timestep);
|
||||
float sigma = sigmas[i];
|
||||
|
||||
auto model_output_opt = model(x, sigma, i + 1);
|
||||
if (model_output_opt.empty()) {
|
||||
@@ -1399,9 +1609,9 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
||||
sd::Tensor<float> model_output = std::move(model_output_opt);
|
||||
model_output = (x - model_output) * (1.0f / sigma);
|
||||
|
||||
float alpha_prod_t = static_cast<float>(alphas_cumprod[timestep]);
|
||||
float alpha_prod_t = 1.0f / (sigma * sigma + 1.0f);
|
||||
float beta_prod_t = 1.0f - alpha_prod_t;
|
||||
float alpha_prod_t_prev = static_cast<float>(prev_timestep >= 0 ? alphas_cumprod[prev_timestep] : alphas_cumprod[0]);
|
||||
float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
|
||||
float alpha_prod_s = static_cast<float>(alphas_cumprod[timestep_s]);
|
||||
float beta_prod_s = 1.0f - alpha_prod_s;
|
||||
|
||||
@@ -1409,13 +1619,14 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
||||
std::sqrt(beta_prod_t) * model_output) *
|
||||
(1.0f / std::sqrt(alpha_prod_t));
|
||||
|
||||
x = std::sqrt(alpha_prod_s) * pred_original_sample +
|
||||
std::sqrt(beta_prod_s) * model_output;
|
||||
x = std::sqrt(alpha_prod_s / alpha_prod_t_prev) * pred_original_sample +
|
||||
std::sqrt(beta_prod_s / alpha_prod_t_prev) * model_output;
|
||||
|
||||
if (eta > 0 && i != steps - 1) {
|
||||
if (eta > 0 && sigma_to > 0.0f) {
|
||||
x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
|
||||
std::sqrt(1.0f - alpha_prod_t_prev / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
|
||||
std::sqrt(1.0f / alpha_prod_t_prev - 1.0f / alpha_prod_s) * sd::Tensor<float>::randn_like(x, rng);
|
||||
}
|
||||
|
||||
}
|
||||
return x;
|
||||
}
|
||||
@@ -1441,13 +1652,16 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
|
||||
case DPM2_SAMPLE_METHOD:
|
||||
return sample_dpm2(model, std::move(x), sigmas);
|
||||
case DPMPP2S_A_SAMPLE_METHOD:
|
||||
return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta);
|
||||
if (is_flow_denoiser)
|
||||
return sample_dpmpp_2s_ancestral_flow(model, std::move(x), sigmas, rng, eta);
|
||||
else
|
||||
return sample_dpmpp_2s_ancestral(model, std::move(x), sigmas, rng, eta);
|
||||
case DPMPP2M_SAMPLE_METHOD:
|
||||
return sample_dpmpp_2m(model, std::move(x), sigmas);
|
||||
case DPMPP2Mv2_SAMPLE_METHOD:
|
||||
return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
|
||||
case LCM_SAMPLE_METHOD:
|
||||
return sample_lcm(model, std::move(x), sigmas, rng);
|
||||
return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser);
|
||||
case IPNDM_SAMPLE_METHOD:
|
||||
return sample_ipndm(model, std::move(x), sigmas);
|
||||
case IPNDM_V_SAMPLE_METHOD:
|
||||
@@ -1456,6 +1670,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
|
||||
return sample_res_multistep(model, std::move(x), sigmas, rng, eta);
|
||||
case RES_2S_SAMPLE_METHOD:
|
||||
return sample_res_2s(model, std::move(x), sigmas, rng, eta);
|
||||
case ER_SDE_SAMPLE_METHOD:
|
||||
return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
||||
case DDIM_TRAILING_SAMPLE_METHOD:
|
||||
return sample_ddim_trailing(model, std::move(x), sigmas, rng, eta);
|
||||
case TCD_SAMPLE_METHOD:
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
#include <optional>
|
||||
#include "anima.hpp"
|
||||
#include "ernie_image.hpp"
|
||||
#include "flux.hpp"
|
||||
#include "mmdit.hpp"
|
||||
#include "qwen_image.hpp"
|
||||
@@ -516,4 +517,66 @@ struct ZImageModel : public DiffusionModel {
|
||||
}
|
||||
};
|
||||
|
||||
struct ErnieImageModel : public DiffusionModel {
|
||||
std::string prefix;
|
||||
ErnieImage::ErnieImageRunner ernie_image;
|
||||
|
||||
ErnieImageModel(ggml_backend_t backend,
|
||||
bool offload_params_to_cpu,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "model.diffusion_model")
|
||||
: prefix(prefix), ernie_image(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return ernie_image.get_desc();
|
||||
}
|
||||
|
||||
void alloc_params_buffer() override {
|
||||
ernie_image.alloc_params_buffer();
|
||||
}
|
||||
|
||||
void free_params_buffer() override {
|
||||
ernie_image.free_params_buffer();
|
||||
}
|
||||
|
||||
void free_compute_buffer() override {
|
||||
ernie_image.free_compute_buffer();
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
ernie_image.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
size_t get_params_buffer_size() override {
|
||||
return ernie_image.get_params_buffer_size();
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
ernie_image.set_weight_adapter(adapter);
|
||||
}
|
||||
|
||||
int64_t get_adm_in_channels() override {
|
||||
return 768;
|
||||
}
|
||||
|
||||
void set_flash_attention_enabled(bool enabled) {
|
||||
ernie_image.set_flash_attention_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||
ernie_image.set_circular_axes(circular_x, circular_y);
|
||||
}
|
||||
|
||||
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 ernie_image.compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context));
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,438 @@
|
||||
#ifndef __SD_ERNIE_IMAGE_HPP__
|
||||
#define __SD_ERNIE_IMAGE_HPP__
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "common_dit.hpp"
|
||||
#include "flux.hpp"
|
||||
#include "qwen_image.hpp"
|
||||
#include "rope.hpp"
|
||||
|
||||
namespace ErnieImage {
|
||||
constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960;
|
||||
|
||||
__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;
|
||||
}
|
||||
};
|
||||
|
||||
struct ErnieImageParams {
|
||||
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;
|
||||
};
|
||||
|
||||
class ErnieImageModel : public GGMLBlock {
|
||||
public:
|
||||
ErnieImageParams params;
|
||||
|
||||
ErnieImageModel() = default;
|
||||
ErnieImageModel(ErnieImageParams params)
|
||||
: params(params) {
|
||||
blocks["x_embedder.proj"] = std::make_shared<Conv2d>(params.in_channels,
|
||||
params.hidden_size,
|
||||
std::pair<int, int>{params.patch_size, params.patch_size},
|
||||
std::pair<int, int>{params.patch_size, params.patch_size},
|
||||
std::pair<int, int>{0, 0},
|
||||
std::pair<int, int>{1, 1},
|
||||
true);
|
||||
if (params.text_in_dim != params.hidden_size) {
|
||||
blocks["text_proj"] = std::make_shared<Linear>(params.text_in_dim, params.hidden_size, false);
|
||||
}
|
||||
blocks["time_embedding"] = std::make_shared<Qwen::TimestepEmbedding>(params.hidden_size, params.hidden_size);
|
||||
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(params.hidden_size, 6 * params.hidden_size, true);
|
||||
|
||||
for (int i = 0; i < params.num_layers; i++) {
|
||||
blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(params.hidden_size,
|
||||
params.num_heads,
|
||||
params.ffn_hidden_size,
|
||||
params.eps);
|
||||
}
|
||||
|
||||
blocks["final_norm"] = std::make_shared<ErnieImageAdaLNContinuous>(params.hidden_size, params.eps);
|
||||
blocks["final_linear"] = std::make_shared<Linear>(params.hidden_size,
|
||||
params.patch_size * params.patch_size * params.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] % params.patch_size == 0 && x->ne[0] % params.patch_size == 0);
|
||||
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t Hp = H / params.patch_size;
|
||||
int64_t Wp = W / params.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>(params.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]
|
||||
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 < params.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);
|
||||
}
|
||||
|
||||
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,
|
||||
params.patch_size,
|
||||
params.patch_size,
|
||||
false); // [N, out_channels, H, W]
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct ErnieImageRunner : public GGMLRunner {
|
||||
ErnieImageParams ernie_params;
|
||||
ErnieImageModel ernie_image;
|
||||
std::vector<float> pe_vec;
|
||||
|
||||
ErnieImageRunner(ggml_backend_t backend,
|
||||
bool offload_params_to_cpu,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "")
|
||||
: GGMLRunner(backend, offload_params_to_cpu) {
|
||||
ernie_params.num_layers = 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) {
|
||||
ernie_params.patch_size = static_cast<int>(tensor_storage.ne[0]);
|
||||
ernie_params.in_channels = tensor_storage.ne[2];
|
||||
ernie_params.hidden_size = tensor_storage.ne[3];
|
||||
} else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
|
||||
ernie_params.text_in_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
|
||||
int64_t head_dim = tensor_storage.ne[0];
|
||||
ernie_params.num_heads = ernie_params.hidden_size / head_dim;
|
||||
} else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
|
||||
ernie_params.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];
|
||||
ernie_params.out_channels = out_dim / ernie_params.patch_size / ernie_params.patch_size;
|
||||
}
|
||||
|
||||
size_t pos = name.find("layers.");
|
||||
if (pos != std::string::npos) {
|
||||
std::string layer_name = name.substr(pos);
|
||||
auto items = split_string(layer_name, '.');
|
||||
if (items.size() > 1) {
|
||||
int block_index = atoi(items[1].c_str());
|
||||
if (block_index + 1 > ernie_params.num_layers) {
|
||||
ernie_params.num_layers = block_index + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (ernie_params.num_layers == 0) {
|
||||
ernie_params.num_layers = 36;
|
||||
}
|
||||
ernie_params.axes_dim_sum = 0;
|
||||
for (int axis_dim : ernie_params.axes_dim) {
|
||||
ernie_params.axes_dim_sum += axis_dim;
|
||||
}
|
||||
|
||||
LOG_INFO("ernie_image: layers = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64
|
||||
", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
|
||||
ernie_params.num_layers,
|
||||
ernie_params.hidden_size,
|
||||
ernie_params.num_heads,
|
||||
ernie_params.ffn_hidden_size,
|
||||
ernie_params.in_channels,
|
||||
ernie_params.out_channels);
|
||||
|
||||
ernie_image = ErnieImageModel(ernie_params);
|
||||
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) {
|
||||
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]),
|
||||
ernie_params.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
ernie_params.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
ernie_params.axes_dim);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / ernie_params.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, ernie_params.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());
|
||||
}
|
||||
};
|
||||
} // namespace ErnieImage
|
||||
|
||||
#endif // __SD_ERNIE_IMAGE_HPP__
|
||||
@@ -59,6 +59,9 @@ private:
|
||||
if (!safe_read(fin, key_len))
|
||||
return false;
|
||||
|
||||
if (key_len > 4096)
|
||||
return false;
|
||||
|
||||
std::string key(key_len, '\0');
|
||||
if (!safe_read(fin, (char*)key.data(), key_len))
|
||||
return false;
|
||||
|
||||
@@ -1,61 +0,0 @@
|
||||
/*================================================= LCMSchedule =================================================*/
|
||||
|
||||
struct LCMSchedule : SigmaSchedule {
|
||||
// The number of steps in the original, reference schedule that LCM was trained on.
|
||||
// This is a default value from the Diffusers implementation.
|
||||
const uint32_t original_inference_steps = 50;
|
||||
// The 'rho' parameter for the Karras schedule used to generate the reference sigmas.
|
||||
const float rho = 7.0f;
|
||||
|
||||
std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t /* t_to_sigma */) override {
|
||||
// Note: The t_to_sigma function is not used here, as LCM's schedule is defined
|
||||
// directly in sigma-space, not based on a model's timestep-to-sigma conversion.
|
||||
|
||||
if (n == 0) {
|
||||
// Return an empty vector if no steps are requested.
|
||||
return {};
|
||||
}
|
||||
|
||||
// 1. Generate the "original" Karras schedule.
|
||||
// This is a full reference schedule of `original_inference_steps` that LCM was
|
||||
// trained to condense. The sigmas are ordered from high to low.
|
||||
std::vector<float> original_sigmas(original_inference_steps);
|
||||
float min_inv_rho = std::pow(sigma_min, (1.0f / rho));
|
||||
float max_inv_rho = std::pow(sigma_max, (1.0f / rho));
|
||||
for (uint32_t i = 0; i < original_inference_steps; i++) {
|
||||
// Formula for Karras schedule: sigma = (sigma_max^(1/rho) + i/(N-1) * (sigma_min^(1/rho) - sigma_max^(1/rho)))^rho
|
||||
// This correctly produces a list of sigmas from sigma_max (at i=0) to sigma_min (at i=N-1).
|
||||
original_sigmas[i] = std::pow(max_inv_rho + (static_cast<float>(i) / (original_inference_steps - 1.0f)) * (min_inv_rho - max_inv_rho), rho);
|
||||
}
|
||||
|
||||
std::vector<float> result;
|
||||
result.reserve(n + 1);
|
||||
|
||||
// 2. Select `n` evenly spaced points from the original schedule.
|
||||
// We sample indices from 0 to `original_inference_steps - 1` and pick the
|
||||
// corresponding sigmas. This ensures we start with sigma_max and end with sigma_min.
|
||||
if (n == 1) {
|
||||
// Special case for a single step: just take the start and end.
|
||||
result.push_back(original_sigmas.front()); // sigma_max
|
||||
result.push_back(0.0f);
|
||||
return result;
|
||||
}
|
||||
|
||||
float step_size = static_cast<float>(original_inference_steps - 1) / static_cast<float>(n - 1);
|
||||
for (uint32_t i = 0; i < n; ++i) {
|
||||
// Calculate the index into the original_sigmas array.
|
||||
// We use round for a more even distribution of indices.
|
||||
int idx = static_cast<int>(std::round(step_size * static_cast<float>(i)));
|
||||
|
||||
// Clamp index to be safe against floating point precision issues.
|
||||
idx = std::max(0, std::min(static_cast<int>(original_inference_steps - 1), idx));
|
||||
|
||||
result.push_back(original_sigmas[idx]);
|
||||
}
|
||||
|
||||
// 3. Append the final zero sigma, representing a fully denoised latent.
|
||||
result.push_back(0.0f);
|
||||
|
||||
return result;
|
||||
}
|
||||
};
|
||||
@@ -1,38 +0,0 @@
|
||||
/* LCM (Latent Consistency Model) Schedule
|
||||
* Reference: https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_lcm.py
|
||||
*
|
||||
* LCM uses a simple linear spacing of timesteps from training timesteps down to 0,
|
||||
* matching the original training distribution for few-step inference.
|
||||
*/
|
||||
struct LCMSchedule : SigmaSchedule {
|
||||
std::vector<float> get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override {
|
||||
std::vector<float> result;
|
||||
|
||||
if (n == 0) {
|
||||
result.push_back(0.0f);
|
||||
return result;
|
||||
}
|
||||
|
||||
result.reserve(n + 1);
|
||||
|
||||
// Handle n == 1 as a special case to avoid division by zero in linear_space
|
||||
if (n == 1) {
|
||||
result.push_back(t_to_sigma(static_cast<float>(TIMESTEPS - 1)));
|
||||
result.push_back(0.0f);
|
||||
return result;
|
||||
}
|
||||
|
||||
// LCM uses linearly spaced timesteps from TIMESTEPS-1 down to 0
|
||||
// This matches the Python implementation's behavior:
|
||||
// timesteps = np.linspace(0, train_timesteps - 1, num_inference_steps)[::-1]
|
||||
std::vector<float> timesteps = linear_space(static_cast<float>(TIMESTEPS - 1), 0.0f, n);
|
||||
|
||||
for (uint32_t i = 0; i < n; ++i) {
|
||||
result.push_back(t_to_sigma(timesteps[i]));
|
||||
}
|
||||
|
||||
// Append the final sigma of 0 as required by the sampling loop
|
||||
result.push_back(0.0f);
|
||||
return result;
|
||||
}
|
||||
};
|
||||
+12
-456
@@ -14,469 +14,21 @@
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "clip.hpp"
|
||||
#include "ggml_extend.hpp"
|
||||
#include "json.hpp"
|
||||
#include "rope.hpp"
|
||||
#include "tokenize_util.h"
|
||||
#include "vocab/vocab.h"
|
||||
#include "tokenizers/bpe_tokenizer.h"
|
||||
#include "tokenizers/mistral_tokenizer.h"
|
||||
#include "tokenizers/qwen2_tokenizer.h"
|
||||
|
||||
namespace LLM {
|
||||
constexpr int LLM_GRAPH_SIZE = 10240;
|
||||
|
||||
class BPETokenizer {
|
||||
protected:
|
||||
std::map<int, std::u32string> byte_encoder;
|
||||
std::map<std::u32string, int> byte_decoder;
|
||||
std::map<std::u32string, int> encoder;
|
||||
std::map<int, std::u32string> decoder;
|
||||
std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
|
||||
std::regex pat;
|
||||
int encoder_len;
|
||||
int bpe_len;
|
||||
|
||||
std::string UNK_TOKEN;
|
||||
std::string BOS_TOKEN;
|
||||
std::string EOS_TOKEN;
|
||||
std::string PAD_TOKEN;
|
||||
|
||||
int UNK_TOKEN_ID;
|
||||
int BOS_TOKEN_ID;
|
||||
int EOS_TOKEN_ID;
|
||||
int PAD_TOKEN_ID;
|
||||
|
||||
std::vector<std::string> special_tokens;
|
||||
|
||||
bool add_bos_token = false;
|
||||
|
||||
protected:
|
||||
static std::string strip(const std::string& str) {
|
||||
std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
|
||||
std::string::size_type end = str.find_last_not_of(" \t\n\r\v\f");
|
||||
|
||||
if (start == std::string::npos) {
|
||||
// String contains only whitespace characters
|
||||
return "";
|
||||
}
|
||||
|
||||
return str.substr(start, end - start + 1);
|
||||
}
|
||||
|
||||
static std::string whitespace_clean(std::string text) {
|
||||
text = std::regex_replace(text, std::regex(R"(\s+)"), " ");
|
||||
text = strip(text);
|
||||
return text;
|
||||
}
|
||||
|
||||
static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
|
||||
std::set<std::pair<std::u32string, std::u32string>> pairs;
|
||||
if (subwords.size() == 0) {
|
||||
return pairs;
|
||||
}
|
||||
std::u32string prev_subword = subwords[0];
|
||||
for (int i = 1; i < subwords.size(); i++) {
|
||||
std::u32string subword = subwords[i];
|
||||
std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
|
||||
pairs.insert(pair);
|
||||
prev_subword = subword;
|
||||
}
|
||||
return pairs;
|
||||
}
|
||||
|
||||
bool is_special_token(const std::string& token) {
|
||||
for (auto& special_token : special_tokens) {
|
||||
if (special_token == token) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
public:
|
||||
BPETokenizer() = default;
|
||||
|
||||
std::u32string bpe(const std::u32string& token) {
|
||||
std::vector<std::u32string> word;
|
||||
|
||||
for (int i = 0; i < token.size(); i++) {
|
||||
word.emplace_back(1, token[i]);
|
||||
}
|
||||
|
||||
std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
|
||||
|
||||
if (pairs.empty()) {
|
||||
return token;
|
||||
}
|
||||
|
||||
while (true) {
|
||||
auto min_pair_iter = std::min_element(pairs.begin(),
|
||||
pairs.end(),
|
||||
[&](const std::pair<std::u32string, std::u32string>& a,
|
||||
const std::pair<std::u32string, std::u32string>& b) {
|
||||
if (bpe_ranks.find(a) == bpe_ranks.end()) {
|
||||
return false;
|
||||
} else if (bpe_ranks.find(b) == bpe_ranks.end()) {
|
||||
return true;
|
||||
}
|
||||
return bpe_ranks.at(a) < bpe_ranks.at(b);
|
||||
});
|
||||
|
||||
const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
|
||||
|
||||
if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
|
||||
break;
|
||||
}
|
||||
|
||||
std::u32string first = bigram.first;
|
||||
std::u32string second = bigram.second;
|
||||
std::vector<std::u32string> new_word;
|
||||
int32_t i = 0;
|
||||
|
||||
while (i < word.size()) {
|
||||
auto it = std::find(word.begin() + i, word.end(), first);
|
||||
if (it == word.end()) {
|
||||
new_word.insert(new_word.end(), word.begin() + i, word.end());
|
||||
break;
|
||||
}
|
||||
new_word.insert(new_word.end(), word.begin() + i, it);
|
||||
i = static_cast<int32_t>(std::distance(word.begin(), it));
|
||||
|
||||
if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
|
||||
new_word.push_back(first + second);
|
||||
i += 2;
|
||||
} else {
|
||||
new_word.push_back(word[i]);
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
|
||||
word = new_word;
|
||||
|
||||
if (word.size() == 1) {
|
||||
break;
|
||||
}
|
||||
pairs = get_pairs(word);
|
||||
}
|
||||
|
||||
std::u32string result;
|
||||
for (int i = 0; i < word.size(); i++) {
|
||||
result += word[i];
|
||||
if (i != word.size() - 1) {
|
||||
result += utf8_to_utf32(" ");
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<int> tokenize(std::string text,
|
||||
on_new_token_cb_t on_new_token_cb = nullptr,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
std::vector<int32_t> tokens = encode(text, on_new_token_cb);
|
||||
|
||||
if (max_length > 0) {
|
||||
if (tokens.size() < max_length) {
|
||||
tokens.resize(max_length);
|
||||
} else {
|
||||
if (padding) {
|
||||
tokens.insert(tokens.end(), max_length - tokens.size(), PAD_TOKEN_ID);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return tokens;
|
||||
}
|
||||
|
||||
void pad_tokens(std::vector<int>& tokens,
|
||||
std::vector<float>& weights,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
if (add_bos_token) {
|
||||
tokens.insert(tokens.begin(), BOS_TOKEN_ID);
|
||||
weights.insert(weights.begin(), 1.f);
|
||||
}
|
||||
if (max_length > 0 && padding) {
|
||||
size_t n = static_cast<size_t>(std::ceil(tokens.size() * 1.f / max_length));
|
||||
if (n == 0) {
|
||||
n = 1;
|
||||
}
|
||||
size_t length = max_length * n;
|
||||
LOG_DEBUG("token length: %llu", length);
|
||||
tokens.insert(tokens.end(), length - tokens.size(), PAD_TOKEN_ID);
|
||||
weights.insert(weights.end(), length - weights.size(), 1.f);
|
||||
}
|
||||
}
|
||||
|
||||
std::vector<int> encode(std::string text, on_new_token_cb_t on_new_token_cb = nullptr) {
|
||||
std::string original_text = text;
|
||||
std::vector<int32_t> bpe_tokens;
|
||||
std::vector<std::string> token_strs;
|
||||
|
||||
auto splited_texts = split_with_special_tokens(text, special_tokens);
|
||||
|
||||
for (auto& splited_text : splited_texts) {
|
||||
if (is_special_token(splited_text)) {
|
||||
bpe_tokens.push_back(encoder[utf8_to_utf32(splited_text)]);
|
||||
token_strs.push_back(splited_text);
|
||||
continue;
|
||||
}
|
||||
auto tokens = token_split(splited_text);
|
||||
for (auto& token : tokens) {
|
||||
if (on_new_token_cb != nullptr) {
|
||||
bool skip = on_new_token_cb(token, bpe_tokens);
|
||||
if (skip) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
std::string token_str = token;
|
||||
std::u32string utf32_token;
|
||||
for (int i = 0; i < token_str.length(); i++) {
|
||||
unsigned char b = token_str[i];
|
||||
utf32_token += byte_encoder[b];
|
||||
}
|
||||
auto bpe_strs = bpe(utf32_token);
|
||||
size_t start = 0;
|
||||
size_t pos;
|
||||
while ((pos = bpe_strs.find(' ', start)) != std::u32string::npos) {
|
||||
auto bpe_str = bpe_strs.substr(start, pos - start);
|
||||
bpe_tokens.push_back(encoder[bpe_str]);
|
||||
token_strs.push_back(utf32_to_utf8(bpe_str));
|
||||
|
||||
start = pos + 1;
|
||||
}
|
||||
auto bpe_str = bpe_strs.substr(start, bpe_strs.size() - start);
|
||||
bpe_tokens.push_back(encoder[bpe_str]);
|
||||
token_strs.push_back(utf32_to_utf8(bpe_str));
|
||||
}
|
||||
}
|
||||
|
||||
std::stringstream ss;
|
||||
ss << "[";
|
||||
for (auto token : token_strs) {
|
||||
ss << "\"" << token << "\", ";
|
||||
}
|
||||
ss << "]";
|
||||
LOG_DEBUG("split prompt \"%s\" to tokens %s", original_text.c_str(), ss.str().c_str());
|
||||
// printf("split prompt \"%s\" to tokens %s \n", original_text.c_str(), ss.str().c_str());
|
||||
return bpe_tokens;
|
||||
}
|
||||
};
|
||||
|
||||
class Qwen2Tokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str) {
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
// printf("byte_unicode_pairs have %lu pairs \n", byte_unicode_pairs.size());
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
// for (auto & pair: byte_unicode_pairs) {
|
||||
// std::cout << pair.first << ": " << pair.second << std::endl;
|
||||
// }
|
||||
std::vector<std::u32string> merges;
|
||||
size_t start = 0;
|
||||
size_t pos;
|
||||
std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
|
||||
while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
|
||||
merges.push_back(merges_utf32_str.substr(start, pos - start));
|
||||
start = pos + 1;
|
||||
}
|
||||
LOG_DEBUG("merges size %llu", merges.size());
|
||||
merges = std::vector<std::u32string>(merges.begin(), merges.end());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
// int print_num = 10;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
// if (print_num > 0) {
|
||||
// print_num--;
|
||||
// printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
|
||||
// utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
|
||||
// }
|
||||
}
|
||||
|
||||
std::vector<std::u32string> tokens;
|
||||
for (const auto& pair : byte_unicode_pairs) {
|
||||
tokens.push_back(pair.second);
|
||||
}
|
||||
for (const auto& merge : merge_pairs) {
|
||||
tokens.push_back(merge.first + merge.second);
|
||||
}
|
||||
for (auto& special_token : special_tokens) {
|
||||
tokens.push_back(utf8_to_utf32(special_token));
|
||||
}
|
||||
|
||||
int i = 0;
|
||||
for (const auto& token : tokens) {
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
i++;
|
||||
}
|
||||
encoder_len = i;
|
||||
LOG_DEBUG("vocab size: %d", encoder_len);
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
};
|
||||
|
||||
public:
|
||||
explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "") {
|
||||
UNK_TOKEN = "<|endoftext|>";
|
||||
EOS_TOKEN = "<|endoftext|>";
|
||||
PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
UNK_TOKEN_ID = 151643;
|
||||
EOS_TOKEN_ID = 151643;
|
||||
PAD_TOKEN_ID = 151643;
|
||||
|
||||
special_tokens = {
|
||||
"<|endoftext|>",
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>",
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
"<|fim_prefix|>",
|
||||
"<|fim_middle|>",
|
||||
"<|fim_suffix|>",
|
||||
"<|fim_pad|>",
|
||||
"<|repo_name|>",
|
||||
"<|file_sep|>",
|
||||
"<tool_response>",
|
||||
"</tool_response>",
|
||||
"<think>",
|
||||
"</think>",
|
||||
};
|
||||
|
||||
if (merges_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_qwen2_merges());
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
class MistralTokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
|
||||
nlohmann::json vocab;
|
||||
|
||||
try {
|
||||
vocab = nlohmann::json::parse(vocab_utf8_str);
|
||||
} catch (const nlohmann::json::parse_error&) {
|
||||
GGML_ABORT("invalid vocab json str");
|
||||
}
|
||||
for (const auto& [key, value] : vocab.items()) {
|
||||
std::u32string token = utf8_to_utf32(key);
|
||||
int i = value;
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
}
|
||||
encoder_len = static_cast<int>(vocab.size());
|
||||
LOG_DEBUG("vocab size: %d", encoder_len);
|
||||
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
std::vector<std::u32string> merges;
|
||||
size_t start = 0;
|
||||
size_t pos;
|
||||
std::u32string merges_utf32_str = utf8_to_utf32(merges_utf8_str);
|
||||
while ((pos = merges_utf32_str.find('\n', start)) != std::string::npos) {
|
||||
merges.push_back(merges_utf32_str.substr(start, pos - start));
|
||||
start = pos + 1;
|
||||
}
|
||||
LOG_DEBUG("merges size %llu", merges.size());
|
||||
merges = std::vector<std::u32string>(merges.begin(), merges.end());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
// int print_num = 10;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
// if (print_num > 0) {
|
||||
// print_num--;
|
||||
// printf("%s :: %s | %s \n", utf32_to_utf8(merge).c_str(), utf32_to_utf8(merge.substr(0, space_pos)).c_str(),
|
||||
// utf32_to_utf8(merge.substr(space_pos + 1)).c_str());
|
||||
// }
|
||||
}
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
};
|
||||
|
||||
public:
|
||||
explicit MistralTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "") {
|
||||
add_bos_token = true;
|
||||
|
||||
UNK_TOKEN = "<unk>";
|
||||
BOS_TOKEN = "<s>";
|
||||
EOS_TOKEN = "</s>";
|
||||
PAD_TOKEN = "<pad>";
|
||||
|
||||
UNK_TOKEN_ID = 0;
|
||||
BOS_TOKEN_ID = 1;
|
||||
EOS_TOKEN_ID = 2;
|
||||
PAD_TOKEN_ID = 11;
|
||||
|
||||
special_tokens = {
|
||||
"<unk>",
|
||||
"<s>",
|
||||
"</s>",
|
||||
"[INST]",
|
||||
"[/INST]",
|
||||
"[AVAILABLE_TOOLS]",
|
||||
"[/AVAILABLE_TOOLS]",
|
||||
"[TOOL_RESULTS]",
|
||||
"[/TOOL_RESULTS]",
|
||||
"[TOOL_CALLS]",
|
||||
"[IMG]",
|
||||
"<pad>",
|
||||
"[IMG_BREAK]",
|
||||
"[IMG_END]",
|
||||
"[PREFIX]",
|
||||
"[MIDDLE]",
|
||||
"[SUFFIX]",
|
||||
"[SYSTEM_PROMPT]",
|
||||
"[/SYSTEM_PROMPT]",
|
||||
"[TOOL_CONTENT]",
|
||||
};
|
||||
for (int i = 20; i < 1000; i++) {
|
||||
special_tokens.push_back("<SPECIAL_" + std::to_string(i) + ">");
|
||||
}
|
||||
|
||||
if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str, vocab_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
enum class LLMArch {
|
||||
QWEN2_5_VL,
|
||||
QWEN3,
|
||||
MISTRAL_SMALL_3_2,
|
||||
MINISTRAL_3_3B,
|
||||
ARCH_COUNT,
|
||||
};
|
||||
|
||||
@@ -484,6 +36,7 @@ namespace LLM {
|
||||
"qwen2.5vl",
|
||||
"qwen3",
|
||||
"mistral_small3.2",
|
||||
"ministral3.3b",
|
||||
};
|
||||
|
||||
struct LLMVisionParams {
|
||||
@@ -868,6 +421,9 @@ namespace LLM {
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2) {
|
||||
q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NORMAL, 8192, 1000000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
} else if (arch == LLMArch::MINISTRAL_3_3B) {
|
||||
q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 262144, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
} else if (arch == LLMArch::QWEN3) {
|
||||
q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
@@ -1083,7 +639,7 @@ namespace LLM {
|
||||
bool enable_vision_ = false)
|
||||
: GGMLRunner(backend, offload_params_to_cpu), enable_vision(enable_vision_) {
|
||||
params.arch = arch;
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2) {
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
|
||||
params.head_dim = 128;
|
||||
params.num_heads = 32;
|
||||
params.num_kv_heads = 8;
|
||||
@@ -1195,7 +751,7 @@ namespace LLM {
|
||||
}
|
||||
|
||||
int64_t n_tokens = input_ids->ne[0];
|
||||
if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::QWEN3) {
|
||||
if (params.arch == LLMArch::MISTRAL_SMALL_3_2 || params.arch == LLMArch::MINISTRAL_3_3B || params.arch == LLMArch::QWEN3) {
|
||||
input_pos_vec.resize(n_tokens);
|
||||
for (int i = 0; i < n_tokens; ++i) {
|
||||
input_pos_vec[i] = i;
|
||||
@@ -1431,7 +987,7 @@ namespace LLM {
|
||||
const std::string prefix = "",
|
||||
bool enable_vision = false)
|
||||
: model(arch, backend, offload_params_to_cpu, tensor_storage_map, prefix, enable_vision) {
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2) {
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
|
||||
tokenizer = std::make_shared<MistralTokenizer>();
|
||||
} else {
|
||||
tokenizer = std::make_shared<Qwen2Tokenizer>();
|
||||
@@ -1479,7 +1035,7 @@ namespace LLM {
|
||||
weights.insert(weights.end(), curr_tokens.size(), curr_weight);
|
||||
}
|
||||
|
||||
tokenizer->pad_tokens(tokens, weights, max_length, padding);
|
||||
tokenizer->pad_tokens(tokens, &weights, nullptr, padding ? max_length : 0, padding ? max_length : 100000000, padding);
|
||||
|
||||
// for (int i = 0; i < tokens.size(); i++) {
|
||||
// std::cout << tokens[i] << ":" << weights[i] << ", ";
|
||||
|
||||
+46
-85
@@ -192,17 +192,22 @@ struct SDCliParams {
|
||||
return options;
|
||||
};
|
||||
|
||||
bool process_and_check() {
|
||||
if (mode != METADATA && output_path.length() == 0) {
|
||||
LOG_ERROR("error: the following arguments are required: output_path");
|
||||
return false;
|
||||
}
|
||||
|
||||
bool resolve() {
|
||||
if (mode == CONVERT) {
|
||||
if (output_path == "output.png") {
|
||||
output_path = "output.gguf";
|
||||
}
|
||||
} else if (mode == METADATA) {
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool validate() {
|
||||
if (mode != METADATA) {
|
||||
if (output_path.length() == 0) {
|
||||
LOG_ERROR("error: the following arguments are required: output_path");
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
if (image_path.empty()) {
|
||||
LOG_ERROR("error: metadata mode needs an image path (--image)");
|
||||
return false;
|
||||
@@ -216,6 +221,16 @@ struct SDCliParams {
|
||||
return true;
|
||||
}
|
||||
|
||||
bool resolve_and_validate() {
|
||||
if (!resolve()) {
|
||||
return false;
|
||||
}
|
||||
if (!validate()) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::string to_string() const {
|
||||
std::ostringstream oss;
|
||||
oss << "SDCliParams {\n"
|
||||
@@ -260,10 +275,10 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP
|
||||
exit(cli_params.normal_exit ? 0 : 1);
|
||||
}
|
||||
|
||||
bool valid = cli_params.process_and_check();
|
||||
bool valid = cli_params.resolve_and_validate();
|
||||
if (valid && cli_params.mode != METADATA) {
|
||||
valid = ctx_params.process_and_check(cli_params.mode) &&
|
||||
gen_params.process_and_check(cli_params.mode, ctx_params.lora_model_dir);
|
||||
valid = ctx_params.resolve_and_validate(cli_params.mode) &&
|
||||
gen_params.resolve_and_validate(cli_params.mode, ctx_params.lora_model_dir);
|
||||
}
|
||||
|
||||
if (!valid) {
|
||||
@@ -278,7 +293,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
|
||||
}
|
||||
|
||||
bool load_images_from_dir(const std::string dir,
|
||||
SDImageVec& images,
|
||||
std::vector<SDImageOwner>& images,
|
||||
int expected_width = 0,
|
||||
int expected_height = 0,
|
||||
int max_image_num = 0,
|
||||
@@ -315,10 +330,10 @@ bool load_images_from_dir(const std::string dir,
|
||||
return false;
|
||||
}
|
||||
|
||||
images.push_back({(uint32_t)width,
|
||||
(uint32_t)height,
|
||||
3,
|
||||
image_buffer});
|
||||
images.emplace_back(sd_image_t{(uint32_t)width,
|
||||
(uint32_t)height,
|
||||
3,
|
||||
image_buffer});
|
||||
|
||||
if (max_image_num > 0 && static_cast<int>(images.size()) >= max_image_num) {
|
||||
break;
|
||||
@@ -558,13 +573,6 @@ int main(int argc, const char* argv[]) {
|
||||
}
|
||||
|
||||
bool vae_decode_only = true;
|
||||
SDImageOwner init_image({0, 0, 3, nullptr});
|
||||
SDImageOwner end_image({0, 0, 3, nullptr});
|
||||
SDImageOwner control_image({0, 0, 3, nullptr});
|
||||
SDImageOwner mask_image({0, 0, 1, nullptr});
|
||||
SDImageVec ref_images;
|
||||
SDImageVec pmid_images;
|
||||
SDImageVec control_frames;
|
||||
|
||||
auto load_image_and_update_size = [&](const std::string& path,
|
||||
SDImageOwner& image,
|
||||
@@ -588,31 +596,32 @@ int main(int argc, const char* argv[]) {
|
||||
|
||||
if (gen_params.init_image_path.size() > 0) {
|
||||
vae_decode_only = false;
|
||||
if (!load_image_and_update_size(gen_params.init_image_path, init_image)) {
|
||||
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
if (gen_params.end_image_path.size() > 0) {
|
||||
vae_decode_only = false;
|
||||
if (!load_image_and_update_size(gen_params.end_image_path, end_image)) {
|
||||
if (!load_image_and_update_size(gen_params.end_image_path, gen_params.end_image)) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
if (gen_params.ref_image_paths.size() > 0) {
|
||||
vae_decode_only = false;
|
||||
gen_params.ref_images.clear();
|
||||
for (auto& path : gen_params.ref_image_paths) {
|
||||
SDImageOwner ref_image({0, 0, 3, nullptr});
|
||||
if (!load_image_and_update_size(path, ref_image, false)) {
|
||||
return 1;
|
||||
}
|
||||
ref_images.push_back(std::move(ref_image));
|
||||
gen_params.ref_images.push_back(std::move(ref_image));
|
||||
}
|
||||
}
|
||||
|
||||
if (gen_params.mask_image_path.size() > 0) {
|
||||
if (!load_sd_image_from_file(mask_image.put(),
|
||||
if (!load_sd_image_from_file(gen_params.mask_image.put(),
|
||||
gen_params.mask_image_path.c_str(),
|
||||
gen_params.get_resolved_width(),
|
||||
gen_params.get_resolved_height(),
|
||||
@@ -630,11 +639,11 @@ int main(int argc, const char* argv[]) {
|
||||
generated_mask.width = gen_params.get_resolved_width();
|
||||
generated_mask.height = gen_params.get_resolved_height();
|
||||
memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
|
||||
mask_image.reset(generated_mask);
|
||||
gen_params.mask_image.reset(generated_mask);
|
||||
}
|
||||
|
||||
if (gen_params.control_image_path.size() > 0) {
|
||||
if (!load_sd_image_from_file(control_image.put(),
|
||||
if (!load_sd_image_from_file(gen_params.control_image.put(),
|
||||
gen_params.control_image_path.c_str(),
|
||||
gen_params.get_resolved_width(),
|
||||
gen_params.get_resolved_height())) {
|
||||
@@ -642,7 +651,7 @@ int main(int argc, const char* argv[]) {
|
||||
return 1;
|
||||
}
|
||||
if (cli_params.canny_preprocess) { // apply preprocessor
|
||||
preprocess_canny(control_image.get(),
|
||||
preprocess_canny(gen_params.control_image.get(),
|
||||
0.08f,
|
||||
0.08f,
|
||||
0.8f,
|
||||
@@ -652,8 +661,9 @@ int main(int argc, const char* argv[]) {
|
||||
}
|
||||
|
||||
if (!gen_params.control_video_path.empty()) {
|
||||
gen_params.control_frames.clear();
|
||||
if (!load_images_from_dir(gen_params.control_video_path,
|
||||
control_frames,
|
||||
gen_params.control_frames,
|
||||
gen_params.get_resolved_width(),
|
||||
gen_params.get_resolved_height(),
|
||||
gen_params.video_frames,
|
||||
@@ -663,8 +673,9 @@ int main(int argc, const char* argv[]) {
|
||||
}
|
||||
|
||||
if (!gen_params.pm_id_images_dir.empty()) {
|
||||
gen_params.pm_id_images.clear();
|
||||
if (!load_images_from_dir(gen_params.pm_id_images_dir,
|
||||
pmid_images,
|
||||
gen_params.pm_id_images,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
@@ -684,7 +695,7 @@ int main(int argc, const char* argv[]) {
|
||||
|
||||
if (cli_params.mode == UPSCALE) {
|
||||
num_results = 1;
|
||||
results.push_back(init_image.release());
|
||||
results.push_back(gen_params.init_image.release());
|
||||
} else {
|
||||
SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
|
||||
|
||||
@@ -706,63 +717,13 @@ int main(int argc, const char* argv[]) {
|
||||
}
|
||||
|
||||
if (cli_params.mode == IMG_GEN) {
|
||||
sd_img_gen_params_t img_gen_params = {
|
||||
gen_params.lora_vec.data(),
|
||||
static_cast<uint32_t>(gen_params.lora_vec.size()),
|
||||
gen_params.prompt.c_str(),
|
||||
gen_params.negative_prompt.c_str(),
|
||||
gen_params.clip_skip,
|
||||
init_image.get(),
|
||||
ref_images.data(),
|
||||
(int)ref_images.size(),
|
||||
gen_params.auto_resize_ref_image,
|
||||
gen_params.increase_ref_index,
|
||||
mask_image.get(),
|
||||
gen_params.get_resolved_width(),
|
||||
gen_params.get_resolved_height(),
|
||||
gen_params.sample_params,
|
||||
gen_params.strength,
|
||||
gen_params.seed,
|
||||
gen_params.batch_count,
|
||||
control_image.get(),
|
||||
gen_params.control_strength,
|
||||
{
|
||||
pmid_images.data(),
|
||||
(int)pmid_images.size(),
|
||||
gen_params.pm_id_embed_path.c_str(),
|
||||
gen_params.pm_style_strength,
|
||||
}, // pm_params
|
||||
gen_params.vae_tiling_params,
|
||||
gen_params.cache_params,
|
||||
};
|
||||
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
|
||||
|
||||
num_results = gen_params.batch_count;
|
||||
results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
|
||||
} else if (cli_params.mode == VID_GEN) {
|
||||
sd_vid_gen_params_t vid_gen_params = {
|
||||
gen_params.lora_vec.data(),
|
||||
static_cast<uint32_t>(gen_params.lora_vec.size()),
|
||||
gen_params.prompt.c_str(),
|
||||
gen_params.negative_prompt.c_str(),
|
||||
gen_params.clip_skip,
|
||||
init_image.get(),
|
||||
end_image.get(),
|
||||
control_frames.data(),
|
||||
(int)control_frames.size(),
|
||||
gen_params.get_resolved_width(),
|
||||
gen_params.get_resolved_height(),
|
||||
gen_params.sample_params,
|
||||
gen_params.high_noise_sample_params,
|
||||
gen_params.moe_boundary,
|
||||
gen_params.strength,
|
||||
gen_params.seed,
|
||||
gen_params.video_frames,
|
||||
gen_params.vace_strength,
|
||||
gen_params.vae_tiling_params,
|
||||
gen_params.cache_params,
|
||||
};
|
||||
|
||||
sd_image_t* generated_video = generate_video(sd_ctx.get(), &vid_gen_params, &num_results);
|
||||
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
|
||||
sd_image_t* generated_video = generate_video(sd_ctx.get(), &vid_gen_params, &num_results);
|
||||
results.adopt(generated_video, num_results);
|
||||
}
|
||||
|
||||
|
||||
+82
-63
@@ -329,8 +329,9 @@ bool is_safetensors_file(const std::string& file_path) {
|
||||
if (!file) {
|
||||
return false;
|
||||
}
|
||||
nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
|
||||
if (header_.is_discarded()) {
|
||||
try {
|
||||
nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
@@ -525,7 +526,14 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
|
||||
return false;
|
||||
}
|
||||
|
||||
nlohmann::json header_ = nlohmann::json::parse(header_buf.data());
|
||||
nlohmann::json header_;
|
||||
try {
|
||||
header_ = nlohmann::json::parse(header_buf.data());
|
||||
} catch (const std::exception&) {
|
||||
LOG_ERROR("parsing safetensors header failed", file_path.c_str());
|
||||
file_paths_.pop_back();
|
||||
return false;
|
||||
}
|
||||
|
||||
for (auto& item : header_.items()) {
|
||||
std::string name = item.key();
|
||||
@@ -591,24 +599,29 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
|
||||
|
||||
size_t tensor_data_size = end - begin;
|
||||
|
||||
bool tensor_size_ok;
|
||||
if (dtype == "F8_E4M3") {
|
||||
tensor_storage.is_f8_e4m3 = true;
|
||||
// f8 -> f16
|
||||
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2);
|
||||
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
|
||||
} else if (dtype == "F8_E5M2") {
|
||||
tensor_storage.is_f8_e5m2 = true;
|
||||
// f8 -> f16
|
||||
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size * 2);
|
||||
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
|
||||
} else if (dtype == "F64") {
|
||||
tensor_storage.is_f64 = true;
|
||||
// f64 -> f32
|
||||
GGML_ASSERT(tensor_storage.nbytes() * 2 == tensor_data_size);
|
||||
tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
|
||||
} else if (dtype == "I64") {
|
||||
tensor_storage.is_i64 = true;
|
||||
// i64 -> i32
|
||||
GGML_ASSERT(tensor_storage.nbytes() * 2 == tensor_data_size);
|
||||
tensor_size_ok = (tensor_storage.nbytes() * 2 == tensor_data_size);
|
||||
} else {
|
||||
GGML_ASSERT(tensor_storage.nbytes() == tensor_data_size);
|
||||
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
|
||||
}
|
||||
if (!tensor_size_ok) {
|
||||
LOG_ERROR("size mismatch for tensor '%s' (%s)\n", name.c_str(), dtype.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
add_tensor_storage(tensor_storage);
|
||||
@@ -1032,64 +1045,66 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
bool has_middle_block_1 = false;
|
||||
bool has_output_block_311 = false;
|
||||
bool has_output_block_71 = false;
|
||||
bool has_attn_1024 = false;
|
||||
|
||||
for (auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!(is_xl)) {
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos) {
|
||||
is_flux = true;
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos) {
|
||||
is_flux = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
|
||||
return VERSION_CHROMA_RADIANCE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) {
|
||||
return VERSION_SD3;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
|
||||
return VERSION_QWEN_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
|
||||
return VERSION_ANIMA;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) {
|
||||
is_flux2 = true;
|
||||
}
|
||||
if (tensor_storage.name.find("single_blocks.47.linear1.weight") != std::string::npos) {
|
||||
has_single_block_47 = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) {
|
||||
return VERSION_OVIS_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) {
|
||||
return VERSION_Z_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
|
||||
return VERSION_ERNIE_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
|
||||
is_wan = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) {
|
||||
patch_embedding_channels = tensor_storage.ne[3];
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.img_emb") != std::string::npos) {
|
||||
has_img_emb = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.input_blocks.") != std::string::npos ||
|
||||
tensor_storage.name.find("unet.down_blocks.") != std::string::npos) {
|
||||
is_unet = true;
|
||||
if (has_multiple_encoders) {
|
||||
is_xl = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
|
||||
return VERSION_CHROMA_RADIANCE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) {
|
||||
return VERSION_SD3;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
|
||||
return VERSION_QWEN_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
|
||||
return VERSION_ANIMA;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) {
|
||||
is_flux2 = true;
|
||||
}
|
||||
if (tensor_storage.name.find("single_blocks.47.linear1.weight") != std::string::npos) {
|
||||
has_single_block_47 = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) {
|
||||
return VERSION_OVIS_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) {
|
||||
return VERSION_Z_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
|
||||
is_wan = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) {
|
||||
patch_embedding_channels = tensor_storage.ne[3];
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.img_emb") != std::string::npos) {
|
||||
has_img_emb = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.input_blocks.") != std::string::npos ||
|
||||
tensor_storage.name.find("unet.down_blocks.") != std::string::npos) {
|
||||
is_unet = true;
|
||||
if (has_multiple_encoders) {
|
||||
is_xl = true;
|
||||
}
|
||||
}
|
||||
if (tensor_storage.name.find("conditioner.embedders.1") != std::string::npos ||
|
||||
tensor_storage.name.find("cond_stage_model.1") != std::string::npos ||
|
||||
tensor_storage.name.find("te.1") != std::string::npos) {
|
||||
has_multiple_encoders = true;
|
||||
if (is_unet) {
|
||||
is_xl = true;
|
||||
}
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.input_blocks.8.0.time_mixer.mix_factor") != std::string::npos) {
|
||||
return VERSION_SVD;
|
||||
}
|
||||
if (tensor_storage.name.find("conditioner.embedders.1") != std::string::npos ||
|
||||
tensor_storage.name.find("cond_stage_model.1") != std::string::npos ||
|
||||
tensor_storage.name.find("te.1") != std::string::npos) {
|
||||
has_multiple_encoders = true;
|
||||
if (is_unet) {
|
||||
is_xl = true;
|
||||
}
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.input_blocks.8.0.time_mixer.mix_factor") != std::string::npos) {
|
||||
return VERSION_SVD;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.middle_block.1.") != std::string::npos ||
|
||||
tensor_storage.name.find("unet.mid_block.resnets.1.") != std::string::npos) {
|
||||
has_middle_block_1 = true;
|
||||
@@ -1101,6 +1116,10 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1") != std::string::npos ||
|
||||
tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) {
|
||||
has_output_block_71 = true;
|
||||
if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight") != std::string::npos) {
|
||||
if (tensor_storage.ne[0] == 1024)
|
||||
has_attn_1024 = true;
|
||||
}
|
||||
}
|
||||
if (tensor_storage.name == "cond_stage_model.transformer.text_model.embeddings.token_embedding.weight" ||
|
||||
tensor_storage.name == "cond_stage_model.model.token_embedding.weight" ||
|
||||
@@ -1174,7 +1193,7 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
}
|
||||
if (!has_middle_block_1) {
|
||||
if (!has_output_block_71) {
|
||||
return VERSION_SDXS;
|
||||
return VERSION_SDXS_512_DS;
|
||||
}
|
||||
return VERSION_SD1_TINY_UNET;
|
||||
}
|
||||
@@ -1184,7 +1203,7 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
return VERSION_SD2_INPAINT;
|
||||
}
|
||||
if (!has_middle_block_1) {
|
||||
return VERSION_SD2_TINY_UNET;
|
||||
return has_attn_1024 ? VERSION_SDXS_09 : VERSION_SD2_TINY_UNET;
|
||||
}
|
||||
return VERSION_SD2;
|
||||
}
|
||||
|
||||
+21
-4
@@ -28,7 +28,8 @@ enum SDVersion {
|
||||
VERSION_SD2,
|
||||
VERSION_SD2_INPAINT,
|
||||
VERSION_SD2_TINY_UNET,
|
||||
VERSION_SDXS,
|
||||
VERSION_SDXS_512_DS,
|
||||
VERSION_SDXS_09,
|
||||
VERSION_SDXL,
|
||||
VERSION_SDXL_INPAINT,
|
||||
VERSION_SDXL_PIX2PIX,
|
||||
@@ -50,18 +51,19 @@ enum SDVersion {
|
||||
VERSION_FLUX2_KLEIN,
|
||||
VERSION_Z_IMAGE,
|
||||
VERSION_OVIS_IMAGE,
|
||||
VERSION_ERNIE_IMAGE,
|
||||
VERSION_COUNT,
|
||||
};
|
||||
|
||||
static inline bool sd_version_is_sd1(SDVersion version) {
|
||||
if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS) {
|
||||
if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS_512_DS) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_sd2(SDVersion version) {
|
||||
if (version == VERSION_SD2 || version == VERSION_SD2_INPAINT || version == VERSION_SD2_TINY_UNET) {
|
||||
if (version == VERSION_SD2 || version == VERSION_SD2_INPAINT || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_09) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
@@ -137,6 +139,20 @@ static inline bool sd_version_is_z_image(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_ernie_image(SDVersion version) {
|
||||
if (version == VERSION_ERNIE_IMAGE) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_flux2_vae(SDVersion version) {
|
||||
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_inpaint(SDVersion version) {
|
||||
if (version == VERSION_SD1_INPAINT ||
|
||||
version == VERSION_SD2_INPAINT ||
|
||||
@@ -155,7 +171,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
||||
sd_version_is_wan(version) ||
|
||||
sd_version_is_qwen_image(version) ||
|
||||
sd_version_is_anima(version) ||
|
||||
sd_version_is_z_image(version)) {
|
||||
sd_version_is_z_image(version) ||
|
||||
sd_version_is_ernie_image(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
|
||||
@@ -1120,7 +1120,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
|
||||
for (const auto& prefix : first_stage_model_prefix_vec) {
|
||||
if (starts_with(name, prefix)) {
|
||||
name = convert_first_stage_model_name(name.substr(prefix.size()), prefix);
|
||||
if (version == VERSION_SDXS) {
|
||||
if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
|
||||
name = "tae." + name;
|
||||
} else {
|
||||
name = prefix + name;
|
||||
|
||||
@@ -7,6 +7,11 @@
|
||||
#include "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);
|
||||
@@ -169,7 +174,8 @@ namespace Rope {
|
||||
int bs,
|
||||
const std::vector<float>& axis_thetas,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {}) {
|
||||
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();
|
||||
@@ -204,6 +210,24 @@ namespace Rope {
|
||||
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);
|
||||
}
|
||||
|
||||
@@ -211,9 +235,10 @@ namespace Rope {
|
||||
int bs,
|
||||
float theta,
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<std::vector<int>>& wrap_dims = {}) {
|
||||
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);
|
||||
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
|
||||
@@ -437,6 +462,74 @@ namespace Rope {
|
||||
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,
|
||||
|
||||
@@ -15,14 +15,20 @@
|
||||
#include <filesystem>
|
||||
|
||||
#include "model_adapter.h"
|
||||
#include "vocab/vocab.h"
|
||||
#include "tokenizers/vocab/vocab.h"
|
||||
#include "flux.hpp"
|
||||
#include "sample-cache.cpp"
|
||||
#include "util.cpp"
|
||||
#include "name_conversion.cpp"
|
||||
#include "upscaler.cpp"
|
||||
#include "model.cpp"
|
||||
#include "tokenize_util.cpp"
|
||||
#include "tokenizers/bpe_tokenizer.cpp"
|
||||
#include "tokenizers/clip_tokenizer.cpp"
|
||||
#include "tokenizers/mistral_tokenizer.cpp"
|
||||
#include "tokenizers/qwen2_tokenizer.cpp"
|
||||
#include "tokenizers/t5_unigram_tokenizer.cpp"
|
||||
#include "tokenizers/tokenizer.cpp"
|
||||
#include "tokenizers/tokenize_util.cpp"
|
||||
#include "zip.c"
|
||||
|
||||
#include "otherarch/utils.h"
|
||||
@@ -706,48 +712,9 @@ static enum sample_method_t sampler_from_name(const std::string& sampler)
|
||||
{
|
||||
// all lowercase
|
||||
enum sample_method_t result = str_to_sample_method(sampler.c_str());
|
||||
if (result != sample_method_t::SAMPLE_METHOD_COUNT)
|
||||
{
|
||||
if (result != sample_method_t::SAMPLE_METHOD_COUNT) {
|
||||
return result;
|
||||
}
|
||||
else if(sampler=="euler a"||sampler=="k_euler_a")
|
||||
{
|
||||
return sample_method_t::EULER_A_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="k_euler")
|
||||
{
|
||||
return sample_method_t::EULER_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="k_heun")
|
||||
{
|
||||
return sample_method_t::HEUN_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="k_dpm_2")
|
||||
{
|
||||
return sample_method_t::DPM2_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="k_lcm")
|
||||
{
|
||||
return sample_method_t::LCM_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="ddim")
|
||||
{
|
||||
return sample_method_t::DDIM_TRAILING_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="dpm++ 2m karras" || sampler=="dpm++ 2m" || sampler=="k_dpmpp_2m")
|
||||
{
|
||||
return sample_method_t::DPMPP2M_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="res multistep" || sampler=="k_res_multistep")
|
||||
{
|
||||
return sample_method_t::RES_MULTISTEP_SAMPLE_METHOD;
|
||||
}
|
||||
else if(sampler=="res 2s" || sampler=="k_res_2s")
|
||||
{
|
||||
return sample_method_t::RES_2S_SAMPLE_METHOD;
|
||||
}
|
||||
else
|
||||
{
|
||||
} else {
|
||||
return sample_method_t::SAMPLE_METHOD_COUNT;
|
||||
}
|
||||
}
|
||||
@@ -1058,7 +1025,7 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
|
||||
}
|
||||
}
|
||||
|
||||
if(loadedsdver == SDVersion::VERSION_SDXS)
|
||||
if(loadedsdver == SDVersion::VERSION_SDXS_512_DS || loadedsdver == SDVersion::VERSION_SDXS_09)
|
||||
{
|
||||
if(sd_params->cfg_scale > 1.0f || sd_params->sample_steps > 1)
|
||||
{
|
||||
|
||||
@@ -32,7 +32,8 @@ const char* model_version_to_str[] = {
|
||||
"SD 2.x",
|
||||
"SD 2.x Inpaint",
|
||||
"SD 2.x Tiny UNet",
|
||||
"SDXS",
|
||||
"SDXS (512-DS)",
|
||||
"SDXS (09)",
|
||||
"SDXL",
|
||||
"SDXL Inpaint",
|
||||
"SDXL Instruct-Pix2Pix",
|
||||
@@ -54,6 +55,7 @@ const char* model_version_to_str[] = {
|
||||
"Flux.2 klein",
|
||||
"Z-Image",
|
||||
"Ovis Image",
|
||||
"Ernie Image",
|
||||
};
|
||||
|
||||
const char* sampling_methods_str[] = {
|
||||
@@ -71,6 +73,7 @@ const char* sampling_methods_str[] = {
|
||||
"TCD",
|
||||
"Res Multistep",
|
||||
"Res 2s",
|
||||
"ER-SDE",
|
||||
};
|
||||
|
||||
/*================================================== Helper Functions ================================================*/
|
||||
@@ -366,7 +369,8 @@ public:
|
||||
bool isflux2 = sd_version_is_flux2(tempver);
|
||||
bool is_ovis = (tempver==VERSION_OVIS_IMAGE);
|
||||
bool is_anima = sd_version_is_anima(tempver);
|
||||
bool conditioner_is_llm = (is_qwenimg || iszimg || isflux2 || is_ovis || is_anima);
|
||||
bool is_ernie = sd_version_is_ernie_image(tempver);
|
||||
bool conditioner_is_llm = (is_qwenimg || iszimg || isflux2 || is_ovis || is_anima || is_ernie);
|
||||
|
||||
//kcpp qol fallback: if a llm was loaded as t5 by mistake
|
||||
if(conditioner_is_llm && t5_path_fixed!="")
|
||||
@@ -457,7 +461,7 @@ public:
|
||||
{
|
||||
to_replace = "taesd_3.embd";
|
||||
}
|
||||
else if(sd_version_is_flux2(tempver))
|
||||
else if(sd_version_uses_flux2_vae(tempver))
|
||||
{
|
||||
to_replace = "taesd_f2.embd";
|
||||
}
|
||||
@@ -634,7 +638,7 @@ public:
|
||||
}
|
||||
|
||||
bool tae_preview_only = sd_ctx_params->tae_preview_only;
|
||||
if (version == VERSION_SDXS) {
|
||||
if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
|
||||
tae_preview_only = false;
|
||||
use_tae = true;
|
||||
}
|
||||
@@ -772,6 +776,15 @@ public:
|
||||
tensor_storage_map,
|
||||
"model.diffusion_model",
|
||||
version);
|
||||
} else if (sd_version_is_ernie_image(version)) {
|
||||
cond_stage_model = std::make_shared<LLMEmbedder>(clip_backend,
|
||||
offload_params_to_cpu,
|
||||
tensor_storage_map,
|
||||
version);
|
||||
diffusion_model = std::make_shared<ErnieImageModel>(backend,
|
||||
offload_params_to_cpu,
|
||||
tensor_storage_map,
|
||||
"model.diffusion_model");
|
||||
} else { // SD1.x SD2.x SDXL
|
||||
std::map<std::string, std::string> embbeding_map;
|
||||
for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
|
||||
@@ -1047,6 +1060,10 @@ public:
|
||||
if (version == VERSION_SVD) {
|
||||
ignore_tensors.insert("conditioner.embedders.3");
|
||||
}
|
||||
if (sd_version_is_ernie_image(version)) {
|
||||
ignore_tensors.insert("text_encoders.llm.vision_tower.");
|
||||
ignore_tensors.insert("text_encoders.llm.multi_modal_projector.");
|
||||
}
|
||||
bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap);
|
||||
if (!success) {
|
||||
LOG_ERROR("load tensors from model loader failed");
|
||||
@@ -1150,10 +1167,13 @@ public:
|
||||
sd_version_is_wan(version) ||
|
||||
sd_version_is_qwen_image(version) ||
|
||||
sd_version_is_anima(version) ||
|
||||
sd_version_is_ernie_image(version) ||
|
||||
sd_version_is_z_image(version)) {
|
||||
pred_type = FLOW_PRED;
|
||||
if (sd_version_is_wan(version)) {
|
||||
default_flow_shift = 5.f;
|
||||
} else if (sd_version_is_ernie_image(version)) {
|
||||
default_flow_shift = 4.f;
|
||||
} else {
|
||||
default_flow_shift = 3.f;
|
||||
}
|
||||
@@ -1644,7 +1664,7 @@ public:
|
||||
uint32_t dim = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
|
||||
|
||||
if (dim == 128) {
|
||||
if (sd_version_is_flux2(version)) {
|
||||
if (sd_version_uses_flux2_vae(version)) {
|
||||
latent_rgb_proj = flux2_latent_rgb_proj;
|
||||
latent_rgb_bias = flux2_latent_rgb_bias;
|
||||
patch_sz = 2;
|
||||
@@ -2093,7 +2113,7 @@ public:
|
||||
latent_channel = 48;
|
||||
} else if (version == VERSION_CHROMA_RADIANCE) {
|
||||
latent_channel = 3;
|
||||
} else if (sd_version_is_flux2(version)) {
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
latent_channel = 128;
|
||||
} else {
|
||||
latent_channel = 16;
|
||||
@@ -2241,6 +2261,7 @@ const char* sample_method_to_str[] = {
|
||||
"tcd",
|
||||
"res_multistep",
|
||||
"res_2s",
|
||||
"er_sde",
|
||||
};
|
||||
|
||||
const char* sd_sample_method_name(enum sample_method_t sample_method) {
|
||||
@@ -2636,6 +2657,14 @@ struct sd_ctx_t {
|
||||
StableDiffusionGGML* sd = nullptr;
|
||||
};
|
||||
|
||||
static bool sd_version_supports_video_generation(SDVersion version) {
|
||||
return version == VERSION_SVD || sd_version_is_wan(version);
|
||||
}
|
||||
|
||||
static bool sd_version_supports_image_generation(SDVersion version) {
|
||||
return !sd_version_supports_video_generation(version);
|
||||
}
|
||||
|
||||
sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) {
|
||||
sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
|
||||
if (sd_ctx == nullptr) {
|
||||
@@ -2665,6 +2694,20 @@ void free_sd_ctx(sd_ctx_t* sd_ctx) {
|
||||
free(sd_ctx);
|
||||
}
|
||||
|
||||
SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return sd_version_supports_image_generation(sd_ctx->sd->version);
|
||||
}
|
||||
|
||||
SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
|
||||
return false;
|
||||
}
|
||||
return sd_version_supports_video_generation(sd_ctx->sd->version);
|
||||
}
|
||||
|
||||
enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
|
||||
if (sd_ctx != nullptr && sd_ctx->sd != nullptr) {
|
||||
if (sd_version_is_dit(sd_ctx->sd->version)) {
|
||||
@@ -2723,6 +2766,7 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
|
||||
return 0.0f;
|
||||
case EULER_A_SAMPLE_METHOD:
|
||||
case DPMPP2S_A_SAMPLE_METHOD:
|
||||
case ER_SDE_SAMPLE_METHOD:
|
||||
return 1.0f;
|
||||
default:;
|
||||
}
|
||||
|
||||
@@ -50,6 +50,7 @@ enum sample_method_t {
|
||||
TCD_SAMPLE_METHOD,
|
||||
RES_MULTISTEP_SAMPLE_METHOD,
|
||||
RES_2S_SAMPLE_METHOD,
|
||||
ER_SDE_SAMPLE_METHOD,
|
||||
SAMPLE_METHOD_COUNT
|
||||
};
|
||||
|
||||
@@ -348,6 +349,8 @@ SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data);
|
||||
SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data);
|
||||
SD_API int32_t sd_get_num_physical_cores();
|
||||
SD_API const char* sd_get_system_info();
|
||||
SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
|
||||
SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx);
|
||||
|
||||
SD_API const char* sd_type_name(enum sd_type_t type);
|
||||
SD_API enum sd_type_t str_to_sd_type(const char* str);
|
||||
|
||||
+6
-450
@@ -10,452 +10,9 @@
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
|
||||
#include "darts.h"
|
||||
#include "ggml_extend.hpp"
|
||||
#include "json.hpp"
|
||||
#include "model.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
// Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
|
||||
// and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
|
||||
// Original License: https://github.com/google/sentencepiece/blob/master/LICENSE
|
||||
//
|
||||
// Since tokenization is not the bottleneck in SD, performance was not a major consideration
|
||||
// during the migration.
|
||||
class MetaspacePreTokenizer {
|
||||
private:
|
||||
std::string replacement;
|
||||
bool add_prefix_space;
|
||||
|
||||
public:
|
||||
MetaspacePreTokenizer(const std::string replacement = " ", bool add_prefix_space = true)
|
||||
: replacement(replacement), add_prefix_space(add_prefix_space) {}
|
||||
|
||||
std::string tokenize(const std::string& input) const {
|
||||
std::string tokens;
|
||||
std::stringstream ss(input);
|
||||
|
||||
if (add_prefix_space) {
|
||||
tokens += replacement;
|
||||
}
|
||||
|
||||
std::string token;
|
||||
bool firstToken = true;
|
||||
while (std::getline(ss, token, ' ')) {
|
||||
if (!firstToken)
|
||||
tokens += replacement + token;
|
||||
else
|
||||
tokens += token;
|
||||
|
||||
firstToken = false;
|
||||
}
|
||||
|
||||
return tokens;
|
||||
}
|
||||
};
|
||||
|
||||
using EncodeResult = std::vector<std::pair<std::string, int>>;
|
||||
class T5UniGramTokenizer {
|
||||
public:
|
||||
enum Status {
|
||||
OK,
|
||||
NO_PIECES_LOADED,
|
||||
NO_ENTRY_FOUND,
|
||||
BUILD_DOUBLE_ARRAY_FAILED,
|
||||
PIECE_ALREADY_DEFINED,
|
||||
INVLIAD_JSON
|
||||
};
|
||||
|
||||
protected:
|
||||
MetaspacePreTokenizer pre_tokenizer;
|
||||
|
||||
// all <piece, score> pairs
|
||||
std::vector<std::pair<std::string, float>> piece_score_pairs;
|
||||
|
||||
float min_score_ = 0.0;
|
||||
float max_score_ = 0.0;
|
||||
std::unique_ptr<Darts::DoubleArray> trie_;
|
||||
|
||||
// Maximum size of the return value of Trie, which corresponds
|
||||
// to the maximum size of shared common prefix in the sentence pieces.
|
||||
int trie_results_size_;
|
||||
// unknown id.
|
||||
int unk_id_ = 2;
|
||||
std::string eos_token_ = "</s>";
|
||||
int eos_id_ = 1;
|
||||
int pad_id_ = 0;
|
||||
// status.
|
||||
Status status_ = OK;
|
||||
|
||||
float kUnkPenalty = 10.0;
|
||||
|
||||
std::string replacement;
|
||||
bool add_prefix_space = true;
|
||||
|
||||
void InitializePieces(const std::string& json_str) {
|
||||
nlohmann::json data;
|
||||
|
||||
try {
|
||||
data = nlohmann::json::parse(json_str);
|
||||
} catch (const nlohmann::json::parse_error&) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
if (!data.contains("model")) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
nlohmann::json model = data["model"];
|
||||
if (!model.contains("vocab")) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
if (model.contains("unk_id")) {
|
||||
unk_id_ = model["unk_id"];
|
||||
}
|
||||
|
||||
replacement = data["pre_tokenizer"]["replacement"];
|
||||
add_prefix_space = data["pre_tokenizer"]["add_prefix_space"];
|
||||
|
||||
pre_tokenizer = MetaspacePreTokenizer(replacement, add_prefix_space);
|
||||
|
||||
for (const auto& item : model["vocab"]) {
|
||||
if (item.size() != 2 || !item[0].is_string() || !item[1].is_number_float()) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
std::string piece = item[0];
|
||||
if (piece.empty()) {
|
||||
piece = "<empty_token>";
|
||||
}
|
||||
float score = item[1];
|
||||
piece_score_pairs.emplace_back(piece, score);
|
||||
}
|
||||
}
|
||||
|
||||
// Builds a Trie index.
|
||||
void BuildTrie(std::vector<std::pair<std::string, int>>* pieces) {
|
||||
if (status_ != OK)
|
||||
return;
|
||||
|
||||
if (pieces->empty()) {
|
||||
status_ = NO_PIECES_LOADED;
|
||||
return;
|
||||
}
|
||||
|
||||
// sort by sentencepiece since DoubleArray::build()
|
||||
// only accepts sorted strings.
|
||||
sort(pieces->begin(), pieces->end());
|
||||
|
||||
// Makes key/value set for DoubleArrayTrie.
|
||||
std::vector<const char*> key(pieces->size());
|
||||
std::vector<int> value(pieces->size());
|
||||
for (size_t i = 0; i < pieces->size(); ++i) {
|
||||
// LOG_DEBUG("%s %d", (*pieces)[i].first.c_str(), (*pieces)[i].second);
|
||||
key[i] = (*pieces)[i].first.data(); // sorted piece.
|
||||
value[i] = (*pieces)[i].second; // vocab_id
|
||||
}
|
||||
|
||||
trie_ = std::unique_ptr<Darts::DoubleArray>(new Darts::DoubleArray());
|
||||
if (trie_->build(key.size(), const_cast<char**>(&key[0]), nullptr,
|
||||
&value[0]) != 0) {
|
||||
status_ = BUILD_DOUBLE_ARRAY_FAILED;
|
||||
return;
|
||||
}
|
||||
|
||||
// Computes the maximum number of shared prefixes in the trie.
|
||||
const int kMaxTrieResultsSize = 1024;
|
||||
std::vector<Darts::DoubleArray::result_pair_type> results(
|
||||
kMaxTrieResultsSize);
|
||||
trie_results_size_ = 0;
|
||||
for (const auto& p : *pieces) {
|
||||
const size_t num_nodes = trie_->commonPrefixSearch(
|
||||
p.first.data(), results.data(), results.size(), p.first.size());
|
||||
trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
|
||||
}
|
||||
|
||||
if (trie_results_size_ == 0)
|
||||
status_ = NO_ENTRY_FOUND;
|
||||
}
|
||||
|
||||
// Non-virtual (inlined) implementation for faster execution.
|
||||
inline float GetScoreInlined(int id) const {
|
||||
return piece_score_pairs[id].second;
|
||||
}
|
||||
|
||||
inline bool IsUnusedInlined(int id) const {
|
||||
return false; // TODO
|
||||
}
|
||||
|
||||
inline bool IsUserDefinedInlined(int id) const {
|
||||
return false; // TODO
|
||||
}
|
||||
|
||||
inline size_t OneCharLen(const char* src) const {
|
||||
return "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"[(*src & 0xFF) >> 4];
|
||||
}
|
||||
|
||||
// The optimized Viterbi encode.
|
||||
// Main differences from the original function:
|
||||
// 1. Memorizes the best path at each postion so far,
|
||||
// 2. No need to store the Lattice nodes,
|
||||
// 3. Works in utf-8 directly,
|
||||
// 4. Defines a new struct with fewer fields than Lattice,
|
||||
// 5. Does not depend on `class Lattice` nor call `SetSentence()`,
|
||||
// `PopulateNodes()`, or `Viterbi()`. It does everything in one function.
|
||||
// For detailed explanations please see the comments inside the function body.
|
||||
EncodeResult EncodeOptimized(const std::string& normalized) const {
|
||||
// An optimized Viterbi algorithm for unigram language models. Benchmarking
|
||||
// results show that it generates almost identical outputs and achieves 2.1x
|
||||
// speedup on average for 102 languages compared to the original
|
||||
// implementation. It's based on the following three ideas:
|
||||
//
|
||||
// 1. Because it uses the *unigram* model:
|
||||
// best_score(x1, x2, ... xt) = best_score(x1, x2, ... x{t-1}) + score(xt)
|
||||
// Deciding the best path (and score) can be decoupled into two isolated
|
||||
// terms: (a) the best path ended before the last token `best_score(x1, x2, ...)`
|
||||
// x{t-1})`, and (b) the last token and its `score(xt)`. The two terms are
|
||||
// not related to each other at all.
|
||||
//
|
||||
// Therefore, we can compute once and store the *best_path ending at
|
||||
// each character position*. In this way, when we know best_path_ends_at[M],
|
||||
// we can reuse it to compute all the best_path_ends_at_[...] where the last
|
||||
// token starts at the same character position M.
|
||||
//
|
||||
// This improves the time complexity from O(n*k*k) to O(n*k) because it
|
||||
// eliminates the extra loop of recomputing the best path ending at the same
|
||||
// position, where n is the input length and k is the maximum number of tokens
|
||||
// that can be recognized starting at each position.
|
||||
//
|
||||
// 2. Again, because it uses the *unigram* model, we don't need to actually
|
||||
// store the lattice nodes. We still recognize all the tokens and lattice
|
||||
// nodes from the input, but along identifying them, we use and discard them
|
||||
// on the fly. There is no need to actually store them for best path Viterbi
|
||||
// decoding. The only thing we need to store is the best_path ending at
|
||||
// each character position.
|
||||
//
|
||||
// This improvement reduces the things needed to store in memory from O(n*k)
|
||||
// to O(n), where n is the input length and k is the maximum number of tokens
|
||||
// that can be recognized starting at each position.
|
||||
//
|
||||
// It also avoids the need of dynamic-size lattice node pool, because the
|
||||
// number of things to store is fixed as n.
|
||||
//
|
||||
// 3. SentencePiece is designed to work with unicode, taking utf-8 encoding
|
||||
// inputs. In the original implementation, the lattice positions are based on
|
||||
// unicode positions. A mapping from unicode position to the utf-8 position is
|
||||
// maintained to recover the utf-8 string piece.
|
||||
//
|
||||
// We found that it is sufficient and beneficial to directly work with utf-8
|
||||
// positions:
|
||||
//
|
||||
// Firstly, it saves the conversion and mapping between unicode positions and
|
||||
// utf-8 positions.
|
||||
//
|
||||
// Secondly, it reduces the number of fields we need to maintain in the
|
||||
// node/path structure. Specifically, there are 8 fields defined in
|
||||
// `Lattice::Node` used by the original encoder, but here in the optimized
|
||||
// encoder we only need to define 3 fields in `BestPathNode`.
|
||||
|
||||
if (status() != OK || normalized.empty()) {
|
||||
return {};
|
||||
}
|
||||
// Represents the last node of the best path.
|
||||
struct BestPathNode {
|
||||
int id = -1; // The vocab id. (maybe -1 for UNK)
|
||||
float best_path_score =
|
||||
0; // The total score of the best path ending at this node.
|
||||
int starts_at =
|
||||
-1; // The starting position (in utf-8) of this node. The entire best
|
||||
// path can be constructed by backtracking along this link.
|
||||
};
|
||||
const int size = static_cast<int>(normalized.size());
|
||||
const float unk_score = min_score() - kUnkPenalty;
|
||||
// The ends are exclusive.
|
||||
std::vector<BestPathNode> best_path_ends_at(size + 1);
|
||||
// Generate lattice on-the-fly (not stored) and update best_path_ends_at.
|
||||
int starts_at = 0;
|
||||
while (starts_at < size) {
|
||||
std::size_t node_pos = 0;
|
||||
std::size_t key_pos = starts_at;
|
||||
const auto best_path_score_till_here =
|
||||
best_path_ends_at[starts_at].best_path_score;
|
||||
bool has_single_node = false;
|
||||
const int mblen =
|
||||
std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)),
|
||||
size - starts_at);
|
||||
while (key_pos < size) {
|
||||
const int ret =
|
||||
trie_->traverse(normalized.data(), node_pos, key_pos, key_pos + 1);
|
||||
if (ret == -2)
|
||||
break;
|
||||
if (ret >= 0) {
|
||||
if (IsUnusedInlined(ret))
|
||||
continue;
|
||||
// Update the best path node.
|
||||
auto& target_node = best_path_ends_at[key_pos];
|
||||
const auto length = (key_pos - starts_at);
|
||||
// User defined symbol receives extra bonus to always be selected.
|
||||
const auto score = IsUserDefinedInlined(ret)
|
||||
? (length * max_score_ - 0.1)
|
||||
: GetScoreInlined(ret);
|
||||
const auto candidate_best_path_score =
|
||||
score + best_path_score_till_here;
|
||||
if (target_node.starts_at == -1 ||
|
||||
candidate_best_path_score > target_node.best_path_score) {
|
||||
target_node.best_path_score = static_cast<float>(candidate_best_path_score);
|
||||
target_node.starts_at = starts_at;
|
||||
target_node.id = ret;
|
||||
}
|
||||
if (!has_single_node && length == mblen) {
|
||||
has_single_node = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!has_single_node) {
|
||||
auto& target_node = best_path_ends_at[starts_at + mblen];
|
||||
const auto candidate_best_path_score =
|
||||
unk_score + best_path_score_till_here;
|
||||
if (target_node.starts_at == -1 ||
|
||||
candidate_best_path_score > target_node.best_path_score) {
|
||||
target_node.best_path_score = candidate_best_path_score;
|
||||
target_node.starts_at = starts_at;
|
||||
target_node.id = unk_id_;
|
||||
}
|
||||
}
|
||||
// Move by one unicode character.
|
||||
starts_at += mblen;
|
||||
}
|
||||
// Backtrack to identify the best path.
|
||||
EncodeResult results;
|
||||
int ends_at = size;
|
||||
while (ends_at > 0) {
|
||||
const auto& node = best_path_ends_at[ends_at];
|
||||
results.emplace_back(
|
||||
normalized.substr(node.starts_at, ends_at - node.starts_at), node.id);
|
||||
ends_at = node.starts_at;
|
||||
}
|
||||
std::reverse(results.begin(), results.end());
|
||||
return results;
|
||||
}
|
||||
|
||||
public:
|
||||
explicit T5UniGramTokenizer(bool is_umt5 = false) {
|
||||
if (is_umt5) {
|
||||
InitializePieces(load_umt5_tokenizer_json());
|
||||
} else {
|
||||
InitializePieces(load_t5_tokenizer_json());
|
||||
}
|
||||
|
||||
min_score_ = FLT_MAX;
|
||||
max_score_ = FLT_MIN;
|
||||
|
||||
std::vector<std::pair<std::string, int>> pieces;
|
||||
for (int i = 0; i < piece_score_pairs.size(); i++) {
|
||||
const auto& sp = piece_score_pairs[i];
|
||||
|
||||
min_score_ = std::min(min_score_, sp.second);
|
||||
max_score_ = std::max(max_score_, sp.second);
|
||||
|
||||
pieces.emplace_back(sp.first, i);
|
||||
}
|
||||
|
||||
BuildTrie(&pieces);
|
||||
}
|
||||
~T5UniGramTokenizer(){};
|
||||
|
||||
std::string Normalize(const std::string& input) const {
|
||||
// Ref: https://github.com/huggingface/tokenizers/blob/1ff56c0c70b045f0cd82da1af9ac08cd4c7a6f9f/bindings/python/py_src/tokenizers/implementations/sentencepiece_unigram.py#L29
|
||||
// TODO: nmt-nfkc
|
||||
std::string normalized = std::regex_replace(input, std::regex(" {2,}"), " ");
|
||||
return normalized;
|
||||
}
|
||||
|
||||
std::vector<int> Encode(const std::string& input, bool append_eos_if_not_present = true) const {
|
||||
std::string normalized = Normalize(input);
|
||||
normalized = pre_tokenizer.tokenize(normalized);
|
||||
EncodeResult result = EncodeOptimized(normalized);
|
||||
if (result.size() > 0 && append_eos_if_not_present) {
|
||||
auto item = result[result.size() - 1];
|
||||
if (item.first != eos_token_) {
|
||||
result.emplace_back(eos_token_, eos_id_);
|
||||
}
|
||||
}
|
||||
std::vector<int> tokens;
|
||||
for (auto item : result) {
|
||||
tokens.push_back(item.second);
|
||||
}
|
||||
return tokens;
|
||||
}
|
||||
|
||||
void pad_tokens(std::vector<int>& tokens,
|
||||
std::vector<float>& weights,
|
||||
std::vector<float>* attention_mask,
|
||||
size_t max_length = 0,
|
||||
bool padding = false) {
|
||||
if (max_length > 0 && padding) {
|
||||
size_t orig_token_num = tokens.size() - 1;
|
||||
size_t n = static_cast<size_t>(std::ceil(orig_token_num * 1.0 / (max_length - 1)));
|
||||
if (n == 0) {
|
||||
n = 1;
|
||||
}
|
||||
size_t length = max_length * n;
|
||||
LOG_DEBUG("token length: %llu", length);
|
||||
std::vector<int> new_tokens;
|
||||
std::vector<float> new_weights;
|
||||
std::vector<float> new_attention_mask;
|
||||
int token_idx = 0;
|
||||
for (int i = 0; i < length; i++) {
|
||||
if (token_idx >= orig_token_num) {
|
||||
break;
|
||||
}
|
||||
if (attention_mask != nullptr) {
|
||||
new_attention_mask.push_back(0.0);
|
||||
}
|
||||
if (i % max_length == max_length - 1) {
|
||||
new_tokens.push_back(eos_id_);
|
||||
new_weights.push_back(1.0);
|
||||
} else {
|
||||
new_tokens.push_back(tokens[token_idx]);
|
||||
new_weights.push_back(weights[token_idx]);
|
||||
token_idx++;
|
||||
}
|
||||
}
|
||||
|
||||
new_tokens.push_back(eos_id_);
|
||||
new_weights.push_back(1.0);
|
||||
if (attention_mask != nullptr) {
|
||||
new_attention_mask.push_back(0.0);
|
||||
}
|
||||
|
||||
tokens = new_tokens;
|
||||
weights = new_weights;
|
||||
if (attention_mask != nullptr) {
|
||||
*attention_mask = new_attention_mask;
|
||||
}
|
||||
|
||||
if (padding) {
|
||||
int pad_token_id = pad_id_;
|
||||
tokens.insert(tokens.end(), length - tokens.size(), pad_token_id);
|
||||
weights.insert(weights.end(), length - weights.size(), 1.0);
|
||||
if (attention_mask != nullptr) {
|
||||
// maybe keep some padding tokens unmasked?
|
||||
attention_mask->insert(attention_mask->end(), length - attention_mask->size(), -HUGE_VALF);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Returns the minimum score in sentence pieces.
|
||||
// min_score() - 10 is used for the cost of unknown sentence.
|
||||
float min_score() const { return min_score_; }
|
||||
|
||||
// Returns the maximum score in sentence pieces.
|
||||
// max_score() is used for the cost of user defined symbols.
|
||||
float max_score() const { return max_score_; }
|
||||
|
||||
Status status() const { return status_; }
|
||||
};
|
||||
#include "tokenizers/t5_unigram_tokenizer.h"
|
||||
|
||||
class T5LayerNorm : public UnaryBlock {
|
||||
protected:
|
||||
@@ -937,18 +494,17 @@ struct T5Embedder {
|
||||
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, false);
|
||||
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);
|
||||
}
|
||||
|
||||
int EOS_TOKEN_ID = 1;
|
||||
tokens.push_back(EOS_TOKEN_ID);
|
||||
weights.push_back(1.0);
|
||||
|
||||
std::vector<float> attention_mask;
|
||||
|
||||
tokenizer.pad_tokens(tokens, weights, &attention_mask, max_length, padding);
|
||||
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] << ", ";
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
#include "bpe_tokenizer.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <sstream>
|
||||
|
||||
#include "tokenize_util.h"
|
||||
#include "util.h"
|
||||
|
||||
std::vector<std::pair<int, std::u32string>> BPETokenizer::bytes_to_unicode() {
|
||||
std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;
|
||||
std::set<int> byte_set;
|
||||
for (int b = static_cast<int>('!'); b <= static_cast<int>('~'); ++b) {
|
||||
byte_set.insert(b);
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
|
||||
}
|
||||
for (int b = 161; b <= 172; ++b) {
|
||||
byte_set.insert(b);
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
|
||||
}
|
||||
for (int b = 174; b <= 255; ++b) {
|
||||
byte_set.insert(b);
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(b)));
|
||||
}
|
||||
int n = 0;
|
||||
for (int b = 0; b < 256; ++b) {
|
||||
if (byte_set.find(b) == byte_set.end()) {
|
||||
byte_unicode_pairs.push_back(std::pair<int, std::u32string>(b, unicode_value_to_utf32(n + 256)));
|
||||
++n;
|
||||
}
|
||||
}
|
||||
return byte_unicode_pairs;
|
||||
}
|
||||
|
||||
std::vector<std::string> BPETokenizer::token_split(const std::string& text) const {
|
||||
return ::token_split(text);
|
||||
}
|
||||
|
||||
std::vector<std::u32string> BPETokenizer::split_utf32(const std::string& text, char32_t delimiter) {
|
||||
std::vector<std::u32string> result;
|
||||
size_t start = 0;
|
||||
size_t pos = 0;
|
||||
std::u32string utf32_text = utf8_to_utf32(text);
|
||||
while ((pos = utf32_text.find(delimiter, start)) != std::u32string::npos) {
|
||||
result.push_back(utf32_text.substr(start, pos - start));
|
||||
start = pos + 1;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
static std::set<std::pair<std::u32string, std::u32string>> get_pairs(const std::vector<std::u32string>& subwords) {
|
||||
std::set<std::pair<std::u32string, std::u32string>> pairs;
|
||||
if (subwords.empty()) {
|
||||
return pairs;
|
||||
}
|
||||
|
||||
std::u32string prev_subword = subwords[0];
|
||||
for (int i = 1; i < static_cast<int>(subwords.size()); i++) {
|
||||
std::u32string subword = subwords[i];
|
||||
std::pair<std::u32string, std::u32string> pair(prev_subword, subword);
|
||||
pairs.insert(pair);
|
||||
prev_subword = subword;
|
||||
}
|
||||
return pairs;
|
||||
}
|
||||
|
||||
std::vector<std::u32string> BPETokenizer::bpe(const std::u32string& token) const {
|
||||
std::vector<std::u32string> word;
|
||||
|
||||
for (int i = 0; i < static_cast<int>(token.size()) - 1; i++) {
|
||||
word.emplace_back(1, token[i]);
|
||||
}
|
||||
word.push_back(token.substr(token.size() - 1) + utf8_to_utf32(end_of_word_suffix));
|
||||
|
||||
std::set<std::pair<std::u32string, std::u32string>> pairs = get_pairs(word);
|
||||
|
||||
if (pairs.empty()) {
|
||||
return {token + utf8_to_utf32(end_of_word_suffix)};
|
||||
}
|
||||
|
||||
while (true) {
|
||||
auto min_pair_iter = std::min_element(pairs.begin(),
|
||||
pairs.end(),
|
||||
[&](const std::pair<std::u32string, std::u32string>& a,
|
||||
const std::pair<std::u32string, std::u32string>& b) {
|
||||
if (bpe_ranks.find(a) == bpe_ranks.end()) {
|
||||
return false;
|
||||
} else if (bpe_ranks.find(b) == bpe_ranks.end()) {
|
||||
return true;
|
||||
}
|
||||
return bpe_ranks.at(a) < bpe_ranks.at(b);
|
||||
});
|
||||
|
||||
const std::pair<std::u32string, std::u32string>& bigram = *min_pair_iter;
|
||||
|
||||
if (bpe_ranks.find(bigram) == bpe_ranks.end()) {
|
||||
break;
|
||||
}
|
||||
|
||||
std::u32string first = bigram.first;
|
||||
std::u32string second = bigram.second;
|
||||
std::vector<std::u32string> new_word;
|
||||
int32_t i = 0;
|
||||
|
||||
while (i < static_cast<int32_t>(word.size())) {
|
||||
auto it = std::find(word.begin() + i, word.end(), first);
|
||||
if (it == word.end()) {
|
||||
new_word.insert(new_word.end(), word.begin() + i, word.end());
|
||||
break;
|
||||
}
|
||||
new_word.insert(new_word.end(), word.begin() + i, it);
|
||||
i = static_cast<int32_t>(std::distance(word.begin(), it));
|
||||
|
||||
if (word[i] == first && i < static_cast<int32_t>(word.size()) - 1 && word[i + 1] == second) {
|
||||
new_word.push_back(first + second);
|
||||
i += 2;
|
||||
} else {
|
||||
new_word.push_back(word[i]);
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
|
||||
word = new_word;
|
||||
|
||||
if (word.size() == 1) {
|
||||
break;
|
||||
}
|
||||
pairs = get_pairs(word);
|
||||
}
|
||||
|
||||
return word;
|
||||
}
|
||||
|
||||
std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t on_new_token_cb) {
|
||||
std::string normalized_text = normalize(text);
|
||||
std::vector<int32_t> bpe_tokens;
|
||||
std::vector<std::string> token_strs;
|
||||
|
||||
auto splited_texts = split_with_special_tokens(normalized_text, special_tokens);
|
||||
|
||||
for (auto& splited_text : splited_texts) {
|
||||
if (is_special_token(splited_text)) {
|
||||
if (on_new_token_cb != nullptr) {
|
||||
bool skip = on_new_token_cb(splited_text, bpe_tokens);
|
||||
if (skip) {
|
||||
token_strs.push_back(splited_text);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
bpe_tokens.push_back(encoder[utf8_to_utf32(splited_text)]);
|
||||
token_strs.push_back(splited_text);
|
||||
continue;
|
||||
}
|
||||
auto tokens = token_split(splited_text);
|
||||
for (auto& token : tokens) {
|
||||
if (on_new_token_cb != nullptr) {
|
||||
bool skip = on_new_token_cb(token, bpe_tokens);
|
||||
if (skip) {
|
||||
token_strs.push_back(splited_text);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
std::string token_str = token;
|
||||
std::u32string utf32_token;
|
||||
for (int i = 0; i < static_cast<int>(token_str.length()); i++) {
|
||||
unsigned char b = token_str[i];
|
||||
utf32_token += byte_encoder[b];
|
||||
}
|
||||
auto bpe_strs = bpe(utf32_token);
|
||||
for (auto bpe_str : bpe_strs) {
|
||||
bpe_tokens.push_back(encoder[bpe_str]);
|
||||
token_strs.push_back(utf32_to_utf8(bpe_str));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
std::stringstream ss;
|
||||
ss << "[";
|
||||
for (auto token : token_strs) {
|
||||
ss << "\"" << token << "\", ";
|
||||
}
|
||||
ss << "]";
|
||||
LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
|
||||
return bpe_tokens;
|
||||
}
|
||||
|
||||
std::string BPETokenizer::decode_token(int token_id) const {
|
||||
return utf32_to_utf8(decoder.at(token_id));
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
#ifndef __SD_TOKENIZERS_BPE_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_BPE_TOKENIZER_H__
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <map>
|
||||
#include <regex>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "tokenizer.h"
|
||||
|
||||
class BPETokenizer : public Tokenizer {
|
||||
protected:
|
||||
std::map<int, std::u32string> byte_encoder;
|
||||
std::map<std::u32string, int> byte_decoder;
|
||||
std::map<std::u32string, int> encoder;
|
||||
std::map<int, std::u32string> decoder;
|
||||
std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
|
||||
int encoder_len = 0;
|
||||
int bpe_len = 0;
|
||||
|
||||
protected:
|
||||
static std::vector<std::pair<int, std::u32string>> bytes_to_unicode();
|
||||
static std::vector<std::u32string> split_utf32(const std::string& text, char32_t delimiter = U'\n');
|
||||
virtual std::vector<std::string> token_split(const std::string& text) const;
|
||||
std::vector<std::u32string> bpe(const std::u32string& token) const;
|
||||
std::string decode_token(int token_id) const override;
|
||||
|
||||
public:
|
||||
BPETokenizer() = default;
|
||||
virtual ~BPETokenizer() = default;
|
||||
|
||||
std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) override;
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_BPE_TOKENIZER_H__
|
||||
@@ -0,0 +1,116 @@
|
||||
#include "clip_tokenizer.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <regex>
|
||||
#include <set>
|
||||
|
||||
#include "ggml.h"
|
||||
#include "tokenize_util.h"
|
||||
#include "util.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str) {
|
||||
UNK_TOKEN = "<|endoftext|>";
|
||||
BOS_TOKEN = "<|startoftext|>";
|
||||
EOS_TOKEN = "<|endoftext|>";
|
||||
PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
UNK_TOKEN_ID = 49407;
|
||||
BOS_TOKEN_ID = 49406;
|
||||
EOS_TOKEN_ID = 49407;
|
||||
PAD_TOKEN_ID = pad_token_id;
|
||||
|
||||
end_of_word_suffix = "</w>";
|
||||
add_bos_token = true;
|
||||
add_eos_token = true;
|
||||
|
||||
if (merges_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_clip_merges());
|
||||
}
|
||||
add_special_token("<|startoftext|>");
|
||||
add_special_token("<|endoftext|>");
|
||||
}
|
||||
|
||||
void CLIPTokenizer::load_from_merges(const std::string& merges_utf8_str) {
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
|
||||
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
|
||||
GGML_ASSERT(merges.size() == 48895);
|
||||
merges = std::vector<std::u32string>(merges.begin() + 1, merges.end());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
}
|
||||
std::vector<std::u32string> vocab;
|
||||
for (const auto& pair : byte_unicode_pairs) {
|
||||
vocab.push_back(pair.second);
|
||||
}
|
||||
for (const auto& pair : byte_unicode_pairs) {
|
||||
vocab.push_back(pair.second + utf8_to_utf32("</w>"));
|
||||
}
|
||||
for (const auto& merge : merge_pairs) {
|
||||
vocab.push_back(merge.first + merge.second);
|
||||
}
|
||||
vocab.push_back(utf8_to_utf32("<|startoftext|>"));
|
||||
vocab.push_back(utf8_to_utf32("<|endoftext|>"));
|
||||
LOG_DEBUG("vocab size: %llu", vocab.size());
|
||||
int i = 0;
|
||||
for (const auto& token : vocab) {
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
i++;
|
||||
}
|
||||
encoder_len = i;
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
}
|
||||
|
||||
static std::string strip(const std::string& str) {
|
||||
std::string::size_type start = str.find_first_not_of(" \t\n\r\v\f");
|
||||
std::string::size_type end = str.find_last_not_of(" \t\n\r\v\f");
|
||||
|
||||
if (start == std::string::npos) {
|
||||
return "";
|
||||
}
|
||||
|
||||
return str.substr(start, end - start + 1);
|
||||
}
|
||||
|
||||
static std::string whitespace_clean(const std::string& text) {
|
||||
auto result = std::regex_replace(text, std::regex(R"(\s+)"), " ");
|
||||
result = strip(result);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::string CLIPTokenizer::normalize(const std::string& text) const {
|
||||
auto normalized_text = whitespace_clean(text);
|
||||
std::transform(normalized_text.begin(), normalized_text.end(), normalized_text.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
||||
return normalized_text;
|
||||
}
|
||||
|
||||
std::vector<std::string> CLIPTokenizer::token_split(const std::string& text) const {
|
||||
std::regex clip_pat(R"('s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)",
|
||||
std::regex::icase);
|
||||
std::sregex_iterator iter(text.begin(), text.end(), clip_pat);
|
||||
std::sregex_iterator end;
|
||||
|
||||
std::vector<std::string> result;
|
||||
for (; iter != end; ++iter) {
|
||||
result.emplace_back(iter->str());
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
#ifndef __SD_TOKENIZERS_CLIP_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_CLIP_TOKENIZER_H__
|
||||
|
||||
#include <cstddef>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "bpe_tokenizer.h"
|
||||
|
||||
class CLIPTokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str);
|
||||
std::string normalize(const std::string& text) const override;
|
||||
std::vector<std::string> token_split(const std::string& text) const override;
|
||||
|
||||
public:
|
||||
explicit CLIPTokenizer(int pad_token_id = 49407, const std::string& merges_utf8_str = "");
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_CLIP_TOKENIZER_H__
|
||||
@@ -0,0 +1,89 @@
|
||||
#include "mistral_tokenizer.h"
|
||||
|
||||
#include "ggml.h"
|
||||
#include "json.hpp"
|
||||
#include "util.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
|
||||
nlohmann::json vocab;
|
||||
|
||||
try {
|
||||
vocab = nlohmann::json::parse(vocab_utf8_str);
|
||||
} catch (const nlohmann::json::parse_error&) {
|
||||
GGML_ABORT("invalid vocab json str");
|
||||
}
|
||||
for (const auto& [key, value] : vocab.items()) {
|
||||
std::u32string token = utf8_to_utf32(key);
|
||||
int i = value;
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
}
|
||||
encoder_len = static_cast<int>(vocab.size());
|
||||
LOG_DEBUG("vocab size: %d", encoder_len);
|
||||
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
|
||||
LOG_DEBUG("merges size %llu", merges.size());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
}
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
}
|
||||
|
||||
MistralTokenizer::MistralTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
|
||||
add_bos_token = true;
|
||||
|
||||
UNK_TOKEN = "<unk>";
|
||||
BOS_TOKEN = "<s>";
|
||||
EOS_TOKEN = "</s>";
|
||||
PAD_TOKEN = "<pad>";
|
||||
|
||||
UNK_TOKEN_ID = 0;
|
||||
BOS_TOKEN_ID = 1;
|
||||
EOS_TOKEN_ID = 2;
|
||||
PAD_TOKEN_ID = 11;
|
||||
|
||||
special_tokens = {
|
||||
"<unk>",
|
||||
"<s>",
|
||||
"</s>",
|
||||
"[INST]",
|
||||
"[/INST]",
|
||||
"[AVAILABLE_TOOLS]",
|
||||
"[/AVAILABLE_TOOLS]",
|
||||
"[TOOL_RESULTS]",
|
||||
"[/TOOL_RESULTS]",
|
||||
"[TOOL_CALLS]",
|
||||
"[IMG]",
|
||||
"<pad>",
|
||||
"[IMG_BREAK]",
|
||||
"[IMG_END]",
|
||||
"[PREFIX]",
|
||||
"[MIDDLE]",
|
||||
"[SUFFIX]",
|
||||
"[SYSTEM_PROMPT]",
|
||||
"[/SYSTEM_PROMPT]",
|
||||
"[TOOL_CONTENT]",
|
||||
};
|
||||
for (int i = 20; i < 1000; i++) {
|
||||
special_tokens.push_back("<SPECIAL_" + std::to_string(i) + ">");
|
||||
}
|
||||
|
||||
if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str, vocab_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_mistral_merges(), load_mistral_vocab_json());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
#ifndef __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "bpe_tokenizer.h"
|
||||
|
||||
class MistralTokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
|
||||
|
||||
public:
|
||||
explicit MistralTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_MISTRAL_TOKENIZER_H__
|
||||
@@ -0,0 +1,91 @@
|
||||
#include "qwen2_tokenizer.h"
|
||||
|
||||
#include "util.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) {
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
|
||||
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
|
||||
LOG_DEBUG("merges size %llu", merges.size());
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
}
|
||||
|
||||
std::vector<std::u32string> tokens;
|
||||
for (const auto& pair : byte_unicode_pairs) {
|
||||
tokens.push_back(pair.second);
|
||||
}
|
||||
for (const auto& merge : merge_pairs) {
|
||||
tokens.push_back(merge.first + merge.second);
|
||||
}
|
||||
for (auto& special_token : special_tokens) {
|
||||
tokens.push_back(utf8_to_utf32(special_token));
|
||||
}
|
||||
|
||||
int i = 0;
|
||||
for (const auto& token : tokens) {
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
i++;
|
||||
}
|
||||
encoder_len = i;
|
||||
LOG_DEBUG("vocab size: %d", encoder_len);
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
}
|
||||
|
||||
Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
|
||||
UNK_TOKEN = "<|endoftext|>";
|
||||
EOS_TOKEN = "<|endoftext|>";
|
||||
PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
UNK_TOKEN_ID = 151643;
|
||||
EOS_TOKEN_ID = 151643;
|
||||
PAD_TOKEN_ID = 151643;
|
||||
|
||||
special_tokens = {
|
||||
"<|endoftext|>",
|
||||
"<|im_start|>",
|
||||
"<|im_end|>",
|
||||
"<|object_ref_start|>",
|
||||
"<|object_ref_end|>",
|
||||
"<|box_start|>",
|
||||
"<|box_end|>",
|
||||
"<|quad_start|>",
|
||||
"<|quad_end|>",
|
||||
"<|vision_start|>",
|
||||
"<|vision_end|>",
|
||||
"<|vision_pad|>",
|
||||
"<|image_pad|>",
|
||||
"<|video_pad|>",
|
||||
"<tool_call>",
|
||||
"</tool_call>",
|
||||
"<|fim_prefix|>",
|
||||
"<|fim_middle|>",
|
||||
"<|fim_suffix|>",
|
||||
"<|fim_pad|>",
|
||||
"<|repo_name|>",
|
||||
"<|file_sep|>",
|
||||
"<tool_response>",
|
||||
"</tool_response>",
|
||||
"<think>",
|
||||
"</think>",
|
||||
};
|
||||
|
||||
if (merges_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_qwen2_merges());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
#ifndef __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "bpe_tokenizer.h"
|
||||
|
||||
class Qwen2Tokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str);
|
||||
|
||||
public:
|
||||
explicit Qwen2Tokenizer(const std::string& merges_utf8_str = "");
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_QWEN2_TOKENIZER_H__
|
||||
@@ -0,0 +1,339 @@
|
||||
#include "t5_unigram_tokenizer.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cfloat>
|
||||
#include <cmath>
|
||||
#include <regex>
|
||||
#include <sstream>
|
||||
|
||||
#include "json.hpp"
|
||||
#include "tokenize_util.h"
|
||||
#include "util.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
// Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h
|
||||
// and https://github.com/google/sentencepiece/blob/master/src/unigram_model.h.
|
||||
// Original License: https://github.com/google/sentencepiece/blob/master/LICENSE
|
||||
//
|
||||
// Since tokenization is not the bottleneck in SD, performance was not a major consideration
|
||||
// during the migration.
|
||||
|
||||
MetaspacePreTokenizer::MetaspacePreTokenizer(const std::string replacement, bool add_prefix_space)
|
||||
: replacement(replacement), add_prefix_space(add_prefix_space) {}
|
||||
|
||||
std::string MetaspacePreTokenizer::tokenize(const std::string& input) const {
|
||||
std::string tokens;
|
||||
std::stringstream ss(input);
|
||||
|
||||
if (add_prefix_space) {
|
||||
tokens += replacement;
|
||||
}
|
||||
|
||||
std::string token;
|
||||
bool first_token = true;
|
||||
while (std::getline(ss, token, ' ')) {
|
||||
if (!first_token) {
|
||||
tokens += replacement + token;
|
||||
} else {
|
||||
tokens += token;
|
||||
}
|
||||
|
||||
first_token = false;
|
||||
}
|
||||
|
||||
return tokens;
|
||||
}
|
||||
|
||||
void T5UniGramTokenizer::InitializePieces(const std::string& json_str) {
|
||||
nlohmann::json data;
|
||||
|
||||
try {
|
||||
data = nlohmann::json::parse(json_str);
|
||||
} catch (const nlohmann::json::parse_error&) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
if (!data.contains("model")) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
nlohmann::json model = data["model"];
|
||||
if (!model.contains("vocab")) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
if (model.contains("unk_id")) {
|
||||
UNK_TOKEN_ID = model["unk_id"];
|
||||
}
|
||||
|
||||
replacement = data["pre_tokenizer"]["replacement"];
|
||||
add_prefix_space = data["pre_tokenizer"]["add_prefix_space"];
|
||||
|
||||
pre_tokenizer = MetaspacePreTokenizer(replacement, add_prefix_space);
|
||||
|
||||
for (const auto& item : model["vocab"]) {
|
||||
if (item.size() != 2 || !item[0].is_string() || !item[1].is_number_float()) {
|
||||
status_ = INVLIAD_JSON;
|
||||
return;
|
||||
}
|
||||
std::string piece = item[0];
|
||||
if (piece.empty()) {
|
||||
piece = "<empty_token>";
|
||||
}
|
||||
float score = item[1];
|
||||
piece_score_pairs.emplace_back(piece, score);
|
||||
}
|
||||
}
|
||||
|
||||
void T5UniGramTokenizer::BuildTrie(std::vector<std::pair<std::string, int>>* pieces) {
|
||||
if (status_ != OK) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (pieces->empty()) {
|
||||
status_ = NO_PIECES_LOADED;
|
||||
return;
|
||||
}
|
||||
|
||||
std::sort(pieces->begin(), pieces->end());
|
||||
|
||||
std::vector<const char*> key(pieces->size());
|
||||
std::vector<int> value(pieces->size());
|
||||
for (size_t i = 0; i < pieces->size(); ++i) {
|
||||
key[i] = (*pieces)[i].first.data();
|
||||
value[i] = (*pieces)[i].second;
|
||||
}
|
||||
|
||||
trie_ = std::unique_ptr<Darts::DoubleArray>(new Darts::DoubleArray());
|
||||
if (trie_->build(key.size(), const_cast<char**>(&key[0]), nullptr, &value[0]) != 0) {
|
||||
status_ = BUILD_DOUBLE_ARRAY_FAILED;
|
||||
return;
|
||||
}
|
||||
|
||||
const int kMaxTrieResultsSize = 1024;
|
||||
std::vector<Darts::DoubleArray::result_pair_type> results(kMaxTrieResultsSize);
|
||||
trie_results_size_ = 0;
|
||||
for (const auto& p : *pieces) {
|
||||
const size_t num_nodes = trie_->commonPrefixSearch(
|
||||
p.first.data(), results.data(), results.size(), p.first.size());
|
||||
trie_results_size_ = std::max(trie_results_size_, static_cast<int>(num_nodes));
|
||||
}
|
||||
|
||||
if (trie_results_size_ == 0) {
|
||||
status_ = NO_ENTRY_FOUND;
|
||||
}
|
||||
}
|
||||
|
||||
float T5UniGramTokenizer::GetScoreInlined(int id) const {
|
||||
return piece_score_pairs[id].second;
|
||||
}
|
||||
|
||||
bool T5UniGramTokenizer::IsUnusedInlined(int id) const {
|
||||
(void)id;
|
||||
return false;
|
||||
}
|
||||
|
||||
bool T5UniGramTokenizer::IsUserDefinedInlined(int id) const {
|
||||
(void)id;
|
||||
return false;
|
||||
}
|
||||
|
||||
size_t T5UniGramTokenizer::OneCharLen(const char* src) const {
|
||||
return "\1\1\1\1\1\1\1\1\1\1\1\1\2\2\3\4"[(*src & 0xFF) >> 4];
|
||||
}
|
||||
|
||||
EncodeResult T5UniGramTokenizer::EncodeOptimized(const std::string& normalized) const {
|
||||
if (status() != OK || normalized.empty()) {
|
||||
return {};
|
||||
}
|
||||
|
||||
struct BestPathNode {
|
||||
int id = -1;
|
||||
float best_path_score = 0;
|
||||
int starts_at = -1;
|
||||
};
|
||||
|
||||
const int size = static_cast<int>(normalized.size());
|
||||
const float unk_score = min_score() - kUnkPenalty;
|
||||
std::vector<BestPathNode> best_path_ends_at(size + 1);
|
||||
|
||||
int starts_at = 0;
|
||||
while (starts_at < size) {
|
||||
std::size_t node_pos = 0;
|
||||
std::size_t key_pos = starts_at;
|
||||
const auto best_path_score_till_here = best_path_ends_at[starts_at].best_path_score;
|
||||
bool has_single_node = false;
|
||||
const int mblen = std::min<int>(static_cast<int>(OneCharLen(normalized.data() + starts_at)), size - starts_at);
|
||||
while (key_pos < static_cast<size_t>(size)) {
|
||||
const int ret = trie_->traverse(normalized.data(), node_pos, key_pos, key_pos + 1);
|
||||
if (ret == -2) {
|
||||
break;
|
||||
}
|
||||
if (ret >= 0) {
|
||||
if (IsUnusedInlined(ret)) {
|
||||
continue;
|
||||
}
|
||||
auto& target_node = best_path_ends_at[key_pos];
|
||||
const auto length = static_cast<int>(key_pos - starts_at);
|
||||
const auto score = IsUserDefinedInlined(ret) ? (length * max_score_ - 0.1f) : GetScoreInlined(ret);
|
||||
const auto candidate_best_path_score = score + best_path_score_till_here;
|
||||
if (target_node.starts_at == -1 || candidate_best_path_score > target_node.best_path_score) {
|
||||
target_node.best_path_score = static_cast<float>(candidate_best_path_score);
|
||||
target_node.starts_at = starts_at;
|
||||
target_node.id = ret;
|
||||
}
|
||||
if (!has_single_node && length == mblen) {
|
||||
has_single_node = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (!has_single_node) {
|
||||
auto& target_node = best_path_ends_at[starts_at + mblen];
|
||||
const auto candidate_best_path_score = unk_score + best_path_score_till_here;
|
||||
if (target_node.starts_at == -1 || candidate_best_path_score > target_node.best_path_score) {
|
||||
target_node.best_path_score = candidate_best_path_score;
|
||||
target_node.starts_at = starts_at;
|
||||
target_node.id = UNK_TOKEN_ID;
|
||||
}
|
||||
}
|
||||
starts_at += mblen;
|
||||
}
|
||||
|
||||
EncodeResult results;
|
||||
int ends_at = size;
|
||||
while (ends_at > 0) {
|
||||
const auto& node = best_path_ends_at[ends_at];
|
||||
results.emplace_back(normalized.substr(node.starts_at, ends_at - node.starts_at), node.id);
|
||||
ends_at = node.starts_at;
|
||||
}
|
||||
std::reverse(results.begin(), results.end());
|
||||
return results;
|
||||
}
|
||||
|
||||
T5UniGramTokenizer::T5UniGramTokenizer(bool is_umt5) {
|
||||
add_bos_token = false;
|
||||
add_eos_token = true;
|
||||
|
||||
if (is_umt5) {
|
||||
PAD_TOKEN_ID = 0;
|
||||
EOS_TOKEN_ID = 1;
|
||||
BOS_TOKEN_ID = 2;
|
||||
UNK_TOKEN_ID = 3;
|
||||
|
||||
PAD_TOKEN = "<pad>";
|
||||
EOS_TOKEN = "</s>";
|
||||
BOS_TOKEN = "<s>";
|
||||
UNK_TOKEN = "<unk>";
|
||||
} else {
|
||||
PAD_TOKEN_ID = 0;
|
||||
EOS_TOKEN_ID = 1;
|
||||
UNK_TOKEN_ID = 2;
|
||||
|
||||
PAD_TOKEN = "<pad>";
|
||||
EOS_TOKEN = "</s>";
|
||||
UNK_TOKEN = "<unk>";
|
||||
}
|
||||
|
||||
special_tokens = {
|
||||
"<pad>",
|
||||
"</s>",
|
||||
"<unk>",
|
||||
};
|
||||
|
||||
if (is_umt5) {
|
||||
special_tokens.push_back("<s>");
|
||||
}
|
||||
|
||||
if (is_umt5) {
|
||||
InitializePieces(load_umt5_tokenizer_json());
|
||||
} else {
|
||||
InitializePieces(load_t5_tokenizer_json());
|
||||
}
|
||||
|
||||
min_score_ = FLT_MAX;
|
||||
max_score_ = FLT_MIN;
|
||||
|
||||
std::vector<std::pair<std::string, int>> pieces;
|
||||
for (int i = 0; i < static_cast<int>(piece_score_pairs.size()); i++) {
|
||||
const auto& sp = piece_score_pairs[i];
|
||||
|
||||
min_score_ = std::min(min_score_, sp.second);
|
||||
max_score_ = std::max(max_score_, sp.second);
|
||||
|
||||
pieces.emplace_back(sp.first, i);
|
||||
}
|
||||
|
||||
BuildTrie(&pieces);
|
||||
}
|
||||
|
||||
T5UniGramTokenizer::~T5UniGramTokenizer() = default;
|
||||
|
||||
std::string T5UniGramTokenizer::decode_token(int token_id) const {
|
||||
if (token_id < 0 || token_id >= static_cast<int>(piece_score_pairs.size())) {
|
||||
return "";
|
||||
}
|
||||
|
||||
const std::string& piece = piece_score_pairs[token_id].first;
|
||||
if (piece == "<empty_token>") {
|
||||
return "";
|
||||
}
|
||||
return piece;
|
||||
}
|
||||
|
||||
std::string T5UniGramTokenizer::normalize(const std::string& input) const {
|
||||
// Ref: https://github.com/huggingface/tokenizers/blob/1ff56c0c70b045f0cd82da1af9ac08cd4c7a6f9f/bindings/python/py_src/tokenizers/implementations/sentencepiece_unigram.py#L29
|
||||
// TODO: nmt-nfkc
|
||||
std::string normalized = std::regex_replace(input, std::regex(" {2,}"), " ");
|
||||
return normalized;
|
||||
}
|
||||
|
||||
std::vector<int> T5UniGramTokenizer::encode(const std::string& input, on_new_token_cb_t on_new_token_cb) {
|
||||
std::vector<int32_t> tokens;
|
||||
std::vector<std::string> token_strs;
|
||||
std::string normalized = normalize(input);
|
||||
auto splited_texts = split_with_special_tokens(normalized, special_tokens);
|
||||
if (splited_texts.empty()) {
|
||||
splited_texts.push_back(normalized); // for empty string
|
||||
}
|
||||
|
||||
for (auto& splited_text : splited_texts) {
|
||||
if (is_special_token(splited_text)) {
|
||||
if (on_new_token_cb != nullptr) {
|
||||
bool skip = on_new_token_cb(splited_text, tokens);
|
||||
if (skip) {
|
||||
token_strs.push_back(splited_text);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if (splited_text == UNK_TOKEN) {
|
||||
tokens.push_back(UNK_TOKEN_ID);
|
||||
token_strs.push_back(UNK_TOKEN);
|
||||
} else if (splited_text == EOS_TOKEN) {
|
||||
tokens.push_back(EOS_TOKEN_ID);
|
||||
token_strs.push_back(EOS_TOKEN);
|
||||
} else if (splited_text == PAD_TOKEN) {
|
||||
tokens.push_back(PAD_TOKEN_ID);
|
||||
token_strs.push_back(PAD_TOKEN);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
std::string pretokenized = pre_tokenizer.tokenize(splited_text);
|
||||
EncodeResult result = EncodeOptimized(pretokenized);
|
||||
for (const auto& item : result) {
|
||||
tokens.push_back(item.second);
|
||||
token_strs.push_back(item.first);
|
||||
}
|
||||
}
|
||||
|
||||
std::stringstream ss;
|
||||
ss << "[";
|
||||
for (const auto& token_str : token_strs) {
|
||||
ss << "\"" << token_str << "\", ";
|
||||
}
|
||||
ss << "]";
|
||||
LOG_DEBUG("split prompt \"%s\" to tokens %s", input.c_str(), ss.str().c_str());
|
||||
|
||||
return tokens;
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
#ifndef __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
|
||||
|
||||
#include <cstddef>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "darts.h"
|
||||
#include "tokenizer.h"
|
||||
|
||||
class MetaspacePreTokenizer {
|
||||
private:
|
||||
std::string replacement;
|
||||
bool add_prefix_space;
|
||||
|
||||
public:
|
||||
MetaspacePreTokenizer(const std::string replacement = " ", bool add_prefix_space = true);
|
||||
|
||||
std::string tokenize(const std::string& input) const;
|
||||
};
|
||||
|
||||
using EncodeResult = std::vector<std::pair<std::string, int>>;
|
||||
|
||||
class T5UniGramTokenizer : public Tokenizer {
|
||||
public:
|
||||
enum Status {
|
||||
OK,
|
||||
NO_PIECES_LOADED,
|
||||
NO_ENTRY_FOUND,
|
||||
BUILD_DOUBLE_ARRAY_FAILED,
|
||||
PIECE_ALREADY_DEFINED,
|
||||
INVLIAD_JSON
|
||||
};
|
||||
|
||||
protected:
|
||||
MetaspacePreTokenizer pre_tokenizer;
|
||||
std::vector<std::pair<std::string, float>> piece_score_pairs;
|
||||
float min_score_ = 0.0f;
|
||||
float max_score_ = 0.0f;
|
||||
std::unique_ptr<Darts::DoubleArray> trie_;
|
||||
int trie_results_size_ = 0;
|
||||
Status status_ = OK;
|
||||
float kUnkPenalty = 10.0f;
|
||||
std::string replacement;
|
||||
bool add_prefix_space = true;
|
||||
|
||||
void InitializePieces(const std::string& json_str);
|
||||
void BuildTrie(std::vector<std::pair<std::string, int>>* pieces);
|
||||
float GetScoreInlined(int id) const;
|
||||
bool IsUnusedInlined(int id) const;
|
||||
bool IsUserDefinedInlined(int id) const;
|
||||
size_t OneCharLen(const char* src) const;
|
||||
EncodeResult EncodeOptimized(const std::string& normalized) const;
|
||||
|
||||
float min_score() const { return min_score_; }
|
||||
float max_score() const { return max_score_; }
|
||||
Status status() const { return status_; }
|
||||
std::string decode_token(int token_id) const override;
|
||||
std::string normalize(const std::string& input) const override;
|
||||
|
||||
public:
|
||||
explicit T5UniGramTokenizer(bool is_umt5 = false);
|
||||
~T5UniGramTokenizer();
|
||||
|
||||
std::vector<int> encode(const std::string& input, on_new_token_cb_t on_new_token_cb = nullptr) override;
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_T5_UNIGRAM_TOKENIZER_H__
|
||||
@@ -1,5 +1,5 @@
|
||||
#ifndef __TOKENIZE_UTIL__
|
||||
#define __TOKENIZE_UTIL__
|
||||
#ifndef __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
|
||||
#define __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
|
||||
|
||||
#include <string>
|
||||
#include <vector>
|
||||
@@ -7,4 +7,4 @@
|
||||
std::vector<std::string> token_split(const std::string& text);
|
||||
std::vector<std::string> split_with_special_tokens(const std::string& text, const std::vector<std::string>& special_tokens);
|
||||
|
||||
#endif // __TOKENIZE_UTIL__
|
||||
#endif // __SD_TOKENIZERS_BPE_TOKENIZE_UTIL_H__
|
||||
@@ -0,0 +1,222 @@
|
||||
#include "tokenizer.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <regex>
|
||||
|
||||
#include "util.h"
|
||||
|
||||
void Tokenizer::add_special_token(const std::string& token) {
|
||||
special_tokens.push_back(token);
|
||||
}
|
||||
|
||||
bool Tokenizer::is_special_token(const std::string& token) const {
|
||||
for (const auto& special_token : special_tokens) {
|
||||
if (special_token == token) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
std::string Tokenizer::normalize(const std::string& text) const {
|
||||
return text;
|
||||
}
|
||||
|
||||
std::vector<int> Tokenizer::tokenize(const std::string& text,
|
||||
on_new_token_cb_t on_new_token_cb,
|
||||
bool padding,
|
||||
size_t min_length,
|
||||
size_t max_length,
|
||||
bool allow_overflow_expand) {
|
||||
std::vector<int> tokens = encode(text, on_new_token_cb);
|
||||
if (padding) {
|
||||
pad_tokens(tokens, nullptr, nullptr, min_length, max_length, allow_overflow_expand);
|
||||
}
|
||||
return tokens;
|
||||
}
|
||||
|
||||
void Tokenizer::pad_tokens(std::vector<int>& tokens,
|
||||
std::vector<float>* weights,
|
||||
std::vector<float>* mask,
|
||||
size_t min_length,
|
||||
size_t max_length,
|
||||
bool allow_overflow_expand) {
|
||||
const bool use_weights = weights != nullptr;
|
||||
const bool use_mask = mask != nullptr;
|
||||
|
||||
if (use_weights && tokens.size() != weights->size()) {
|
||||
LOG_ERROR("tokens size != weights size");
|
||||
return;
|
||||
}
|
||||
|
||||
const size_t bos_count = add_bos_token ? 1 : 0;
|
||||
const size_t eos_count = add_eos_token ? 1 : 0;
|
||||
const size_t special_token_count = bos_count + eos_count;
|
||||
|
||||
auto build_sequence = [&](size_t begin,
|
||||
size_t count,
|
||||
size_t target_length,
|
||||
std::vector<int>& out_tokens,
|
||||
std::vector<float>& out_weights,
|
||||
std::vector<float>& out_mask) {
|
||||
const size_t base_length = count + special_token_count;
|
||||
const size_t final_length = std::max(target_length, base_length);
|
||||
|
||||
out_tokens.clear();
|
||||
out_weights.clear();
|
||||
out_mask.clear();
|
||||
|
||||
out_tokens.reserve(final_length);
|
||||
if (use_weights) {
|
||||
out_weights.reserve(final_length);
|
||||
}
|
||||
if (use_mask) {
|
||||
out_mask.reserve(final_length);
|
||||
}
|
||||
|
||||
if (add_bos_token) {
|
||||
out_tokens.push_back(BOS_TOKEN_ID);
|
||||
if (use_weights) {
|
||||
out_weights.push_back(1.0f);
|
||||
}
|
||||
if (use_mask) {
|
||||
out_mask.push_back(1.0f);
|
||||
}
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < count; ++i) {
|
||||
out_tokens.push_back(tokens[begin + i]);
|
||||
if (use_weights) {
|
||||
out_weights.push_back((*weights)[begin + i]);
|
||||
}
|
||||
if (use_mask) {
|
||||
out_mask.push_back(1.0f);
|
||||
}
|
||||
}
|
||||
|
||||
if (add_eos_token) {
|
||||
out_tokens.push_back(EOS_TOKEN_ID);
|
||||
if (use_weights) {
|
||||
out_weights.push_back(1.0f);
|
||||
}
|
||||
if (use_mask) {
|
||||
out_mask.push_back(1.0f);
|
||||
}
|
||||
}
|
||||
|
||||
if (final_length > out_tokens.size()) {
|
||||
const size_t pad_count = final_length - out_tokens.size();
|
||||
if (pad_left) {
|
||||
out_tokens.insert(out_tokens.begin(), pad_count, PAD_TOKEN_ID);
|
||||
|
||||
if (use_weights) {
|
||||
out_weights.insert(out_weights.begin(), pad_count, 1.0f);
|
||||
}
|
||||
if (use_mask) {
|
||||
out_mask.insert(out_mask.begin(), pad_count, 0.0f);
|
||||
}
|
||||
} else {
|
||||
out_tokens.insert(out_tokens.end(), pad_count, PAD_TOKEN_ID);
|
||||
|
||||
if (use_weights) {
|
||||
out_weights.insert(out_weights.end(), pad_count, 1.0f);
|
||||
}
|
||||
if (use_mask) {
|
||||
out_mask.insert(out_mask.end(), pad_count, 0.0f);
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const size_t single_length = std::max(min_length, tokens.size() + special_token_count);
|
||||
const bool exceeds_max_length = max_length > 0 && single_length > max_length;
|
||||
|
||||
std::vector<int> new_tokens;
|
||||
std::vector<float> new_weights;
|
||||
std::vector<float> new_mask;
|
||||
|
||||
if (!exceeds_max_length) {
|
||||
build_sequence(0, tokens.size(), min_length, new_tokens, new_weights, new_mask);
|
||||
} else if (!allow_overflow_expand) {
|
||||
build_sequence(0, tokens.size(), 0, new_tokens, new_weights, new_mask);
|
||||
|
||||
new_tokens.resize(max_length);
|
||||
if (use_weights) {
|
||||
new_weights.resize(max_length);
|
||||
}
|
||||
if (use_mask) {
|
||||
new_mask.resize(max_length);
|
||||
}
|
||||
|
||||
if (add_eos_token && !new_tokens.empty()) {
|
||||
new_tokens.back() = EOS_TOKEN_ID;
|
||||
if (use_weights) {
|
||||
new_weights.back() = 1.0f;
|
||||
}
|
||||
if (use_mask) {
|
||||
new_mask.back() = 1.0f;
|
||||
}
|
||||
}
|
||||
} else if (min_length > special_token_count) {
|
||||
const size_t tokens_per_chunk = min_length - special_token_count;
|
||||
size_t offset = 0;
|
||||
|
||||
while (offset < tokens.size()) {
|
||||
const size_t remaining = tokens.size() - offset;
|
||||
const size_t take = std::min(tokens_per_chunk, remaining);
|
||||
|
||||
std::vector<int> chunk_tokens;
|
||||
std::vector<float> chunk_weights;
|
||||
std::vector<float> chunk_mask;
|
||||
|
||||
build_sequence(offset, take, min_length, chunk_tokens, chunk_weights, chunk_mask);
|
||||
|
||||
new_tokens.insert(new_tokens.end(), chunk_tokens.begin(), chunk_tokens.end());
|
||||
if (use_weights) {
|
||||
new_weights.insert(new_weights.end(), chunk_weights.begin(), chunk_weights.end());
|
||||
}
|
||||
if (use_mask) {
|
||||
new_mask.insert(new_mask.end(), chunk_mask.begin(), chunk_mask.end());
|
||||
}
|
||||
|
||||
offset += take;
|
||||
}
|
||||
} else {
|
||||
build_sequence(0, tokens.size(), min_length, new_tokens, new_weights, new_mask);
|
||||
}
|
||||
|
||||
tokens = std::move(new_tokens);
|
||||
if (use_weights) {
|
||||
*weights = std::move(new_weights);
|
||||
}
|
||||
if (use_mask) {
|
||||
*mask = std::move(new_mask);
|
||||
}
|
||||
}
|
||||
|
||||
static std::string clean_up_tokenization(std::string& text) {
|
||||
std::regex pattern(R"( ,)");
|
||||
return std::regex_replace(text, pattern, ",");
|
||||
}
|
||||
|
||||
std::string Tokenizer::decode(const std::vector<int>& tokens) const {
|
||||
std::string text;
|
||||
|
||||
for (int token_id : tokens) {
|
||||
if (token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID || token_id == PAD_TOKEN_ID) {
|
||||
continue;
|
||||
}
|
||||
|
||||
std::string piece = decode_token(token_id);
|
||||
if (!end_of_word_suffix.empty() && ends_with(piece, end_of_word_suffix)) {
|
||||
piece.erase(piece.size() - end_of_word_suffix.size());
|
||||
text += piece + " ";
|
||||
} else {
|
||||
text += piece;
|
||||
}
|
||||
}
|
||||
|
||||
text = clean_up_tokenization(text);
|
||||
return trim(text);
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
#ifndef __SD_TOKENIZERS_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_TOKENIZER_H__
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
using on_new_token_cb_t = std::function<bool(std::string&, std::vector<int32_t>&)>;
|
||||
|
||||
class Tokenizer {
|
||||
protected:
|
||||
std::vector<std::string> special_tokens;
|
||||
bool add_bos_token = false;
|
||||
bool add_eos_token = false;
|
||||
bool pad_left = false;
|
||||
std::string end_of_word_suffix;
|
||||
|
||||
virtual std::string decode_token(int token_id) const = 0;
|
||||
virtual std::string normalize(const std::string& text) const;
|
||||
|
||||
public:
|
||||
std::string UNK_TOKEN;
|
||||
std::string BOS_TOKEN;
|
||||
std::string EOS_TOKEN;
|
||||
std::string PAD_TOKEN;
|
||||
int UNK_TOKEN_ID = 0;
|
||||
int BOS_TOKEN_ID = 0;
|
||||
int EOS_TOKEN_ID = 0;
|
||||
int PAD_TOKEN_ID = 0;
|
||||
|
||||
virtual ~Tokenizer() = default;
|
||||
|
||||
void add_special_token(const std::string& token);
|
||||
bool is_special_token(const std::string& token) const;
|
||||
virtual std::vector<int> encode(const std::string& text, on_new_token_cb_t on_new_token_cb = nullptr) = 0;
|
||||
std::vector<int> tokenize(const std::string& text,
|
||||
on_new_token_cb_t on_new_token_cb = nullptr,
|
||||
bool padding = false,
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 100000000,
|
||||
bool allow_overflow_expand = false);
|
||||
void pad_tokens(std::vector<int>& tokens,
|
||||
std::vector<float>* weights,
|
||||
std::vector<float>* mask,
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 100000000,
|
||||
bool allow_overflow_expand = false);
|
||||
std::string decode(const std::vector<int>& tokens) const;
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_TOKENIZER_H__
|
||||
@@ -1,5 +1,5 @@
|
||||
#ifndef __VOCAB_H__
|
||||
#define __VOCAB_H__
|
||||
#ifndef __SD_TOKENIZERS_VOCAB_VOCAB_H__
|
||||
#define __SD_TOKENIZERS_VOCAB_VOCAB_H__
|
||||
|
||||
#include <string>
|
||||
|
||||
@@ -10,4 +10,4 @@ std::string load_mistral_vocab_json();
|
||||
std::string load_t5_tokenizer_json();
|
||||
std::string load_umt5_tokenizer_json();
|
||||
|
||||
#endif // __VOCAB_H__
|
||||
#endif // __SD_TOKENIZERS_VOCAB_VOCAB_H__
|
||||
@@ -217,11 +217,11 @@ public:
|
||||
} else if (sd_version_is_unet_edit(version)) {
|
||||
in_channels = 8;
|
||||
}
|
||||
if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS) {
|
||||
if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
|
||||
num_res_blocks = 1;
|
||||
channel_mult = {1, 2, 4};
|
||||
tiny_unet = true;
|
||||
if (version == VERSION_SDXS) {
|
||||
if (version == VERSION_SDXS_512_DS) {
|
||||
attention_resolutions = {4, 2}; // here just like SDXL
|
||||
}
|
||||
}
|
||||
@@ -264,6 +264,10 @@ public:
|
||||
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);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -69,7 +69,7 @@ public:
|
||||
int scale_factor = 8;
|
||||
if (version == VERSION_WAN2_2_TI2V) {
|
||||
scale_factor = 16;
|
||||
} else if (sd_version_is_flux2(version)) {
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
scale_factor = 16;
|
||||
} else if (version == VERSION_CHROMA_RADIANCE) {
|
||||
scale_factor = 1;
|
||||
|
||||
Reference in New Issue
Block a user