diff --git "a/stable-diffusion.cpp/stable-diffusion.cpp" "b/stable-diffusion.cpp/stable-diffusion.cpp" deleted file mode 100644--- "a/stable-diffusion.cpp/stable-diffusion.cpp" +++ /dev/null @@ -1,4388 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "ggml/ggml.h" -#include "rng.h" -#include "rng_philox.h" -#include "stable-diffusion.h" - -static SDLogLevel log_level = SDLogLevel::INFO; - -#define __FILENAME__ "stable-diffusion.cpp" -#define SD_LOG(level, format, ...) \ - do { \ - if (level < log_level) { \ - break; \ - } \ - if (level == SDLogLevel::DEBUG) { \ - printf("[DEBUG] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \ - fflush(stdout); \ - } else if (level == SDLogLevel::INFO) { \ - printf("[INFO] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \ - fflush(stdout); \ - } else if (level == SDLogLevel::WARN) { \ - fprintf(stderr, "[WARN] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \ - fflush(stdout); \ - } else if (level == SDLogLevel::ERROR) { \ - fprintf(stderr, "[ERROR] %s:%-4d - " format "\n", __FILENAME__, __LINE__, ##__VA_ARGS__); \ - fflush(stdout); \ - } \ - } while (0) - -#define LOG_DEBUG(format, ...) SD_LOG(SDLogLevel::DEBUG, format, ##__VA_ARGS__) -#define LOG_INFO(format, ...) SD_LOG(SDLogLevel::INFO, format, ##__VA_ARGS__) -#define LOG_WARN(format, ...) SD_LOG(SDLogLevel::WARN, format, ##__VA_ARGS__) -#define LOG_ERROR(format, ...) SD_LOG(SDLogLevel::ERROR, format, ##__VA_ARGS__) - -#define GGML_FILE_MAGIC 0x67676d6c - -#define TIMESTEPS 1000 - -enum ModelType { - SD1 = 0, - SD2 = 1, - MODEL_TYPE_COUNT, -}; - -const char* model_type_to_str[] = { - "SD1.x", - "SD2.x"}; - -/*================================================== Helper Functions ================================================*/ - -void set_sd_log_level(SDLogLevel level) { - log_level = level; -} - -std::string sd_get_system_info() { - std::stringstream ss; - ss << "System Info: \n"; - ss << " BLAS = " << ggml_cpu_has_blas() << std::endl; - ss << " SSE3 = " << ggml_cpu_has_sse3() << std::endl; - ss << " AVX = " << ggml_cpu_has_avx() << std::endl; - ss << " AVX2 = " << ggml_cpu_has_avx2() << std::endl; - ss << " AVX512 = " << ggml_cpu_has_avx512() << std::endl; - ss << " AVX512_VBMI = " << ggml_cpu_has_avx512_vbmi() << std::endl; - ss << " AVX512_VNNI = " << ggml_cpu_has_avx512_vnni() << std::endl; - ss << " FMA = " << ggml_cpu_has_fma() << std::endl; - ss << " NEON = " << ggml_cpu_has_neon() << std::endl; - ss << " ARM_FMA = " << ggml_cpu_has_arm_fma() << std::endl; - ss << " F16C = " << ggml_cpu_has_f16c() << std::endl; - ss << " FP16_VA = " << ggml_cpu_has_fp16_va() << std::endl; - ss << " WASM_SIMD = " << ggml_cpu_has_wasm_simd() << std::endl; - ss << " VSX = " << ggml_cpu_has_vsx() << std::endl; - return ss.str(); -} - -ggml_tensor* load_tensor_from_file(ggml_context* ctx, const std::string& file_path) { - std::ifstream file(file_path, std::ios::binary); - if (!file.is_open()) { - LOG_ERROR("failed to open '%s'", file_path.c_str()); - return NULL; - } - int32_t n_dims; - int32_t length; - int32_t ttype; - - file.read(reinterpret_cast(&n_dims), sizeof(n_dims)); - file.read(reinterpret_cast(&length), sizeof(length)); - file.read(reinterpret_cast(&ttype), sizeof(ttype)); - - if (file.eof()) { - LOG_ERROR("incomplete file '%s'", file_path.c_str()); - return NULL; - } - - int32_t nelements = 1; - int32_t ne[4] = {1, 1, 1, 1}; - for (int i = 0; i < n_dims; ++i) { - file.read(reinterpret_cast(&ne[i]), sizeof(ne[i])); - nelements *= ne[i]; - } - std::string name(length, 0); - file.read(&name[0], length); - ggml_tensor* tensor = ggml_new_tensor_4d(ctx, (ggml_type)ttype, ne[0], ne[1], ne[2], ne[3]); - const size_t bpe = ggml_type_size(ggml_type(ttype)); - file.read(reinterpret_cast(tensor->data), ggml_nbytes(tensor)); - return tensor; -} - -void ggml_tensor_set_f32_randn(struct ggml_tensor* tensor, std::shared_ptr rng) { - uint32_t n = ggml_nelements(tensor); - std::vector random_numbers = rng->randn(n); - for (int i = 0; i < n; i++) { - ggml_set_f32_1d(tensor, i, random_numbers[i]); - } -} - -// set tensor[i, j, k, l] -// set tensor[l] -// set tensor[k, l] -// set tensor[j, k, l] -void ggml_tensor_set_f32(struct ggml_tensor* tensor, float value, int l, int k = 0, int j = 0, int i = 0) { - GGML_ASSERT(tensor->nb[0] == sizeof(float)); - *(float*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]) = value; -} - -float ggml_tensor_get_f32(const ggml_tensor* tensor, int l, int k = 0, int j = 0, int i = 0) { - GGML_ASSERT(tensor->nb[0] == sizeof(float)); - return *(float*)((char*)(tensor->data) + i * tensor->nb[3] + j * tensor->nb[2] + k * tensor->nb[1] + l * tensor->nb[0]); -} - -void print_ggml_tensor(struct ggml_tensor* tensor, bool shape_only = false) { - printf("shape(%zu, %zu, %zu, %zu)\n", tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3]); - fflush(stdout); - if (shape_only) { - return; - } - int range = 3; - for (int i = 0; i < tensor->ne[3]; i++) { - if (i >= range && i + range < tensor->ne[3]) { - continue; - } - for (int j = 0; j < tensor->ne[2]; j++) { - if (j >= range && j + range < tensor->ne[2]) { - continue; - } - for (int k = 0; k < tensor->ne[1]; k++) { - if (k >= range && k + range < tensor->ne[1]) { - continue; - } - for (int l = 0; l < tensor->ne[0]; l++) { - if (l >= range && l + range < tensor->ne[0]) { - continue; - } - printf(" [%d, %d, %d, %d] = %f\n", i, j, k, l, ggml_tensor_get_f32(tensor, l, k, j, i)); - fflush(stdout); - } - } - } - } -} - -void copy_ggml_tensor( - struct ggml_tensor* dst, - const struct ggml_tensor* src) { - dst->nb[0] = src->nb[0]; - dst->nb[1] = src->nb[1]; - dst->nb[2] = src->nb[2]; - dst->nb[3] = src->nb[3]; - - memcpy(((char*)dst->data), ((char*)src->data), ggml_nbytes(dst)); -} - -// Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 -void set_timestep_embedding(struct ggml_tensor* timesteps, struct ggml_tensor* embedding, int dim, int max_period = 10000) { - // timesteps: [N,] - // embedding: [(dim + 1)/2, N] - int half = dim / 2; - std::vector freqs(half); - for (int i = 0; i < half; ++i) { - freqs[i] = (float)std::exp(-std::log(max_period) * i / half); - } - for (int i = 0; i < timesteps->ne[0]; ++i) { - for (int j = 0; j < half; ++j) { - float arg = ggml_get_f32_1d(timesteps, i) * freqs[j]; - ggml_tensor_set_f32(embedding, std::cos(arg), j, i); - ggml_tensor_set_f32(embedding, std::sin(arg), j + half, i); - } - if (dim % 2 != 0) { - *(float*)((char*)embedding->data + i * embedding->nb[1] + dim * embedding->nb[0]) = 0; - } - } -} - -struct ggml_tensor* new_timestep_embedding(struct ggml_context* ctx, struct ggml_tensor* timesteps, int dim, int max_period = 10000) { - // timesteps: [N,] - // embedding: [(dim + 1)/2, N] - int acutual_dim = dim; - if (dim % 2 != 0) { - acutual_dim = dim + 1; - } - struct ggml_tensor* embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, acutual_dim, timesteps->ne[0]); - if (!ggml_get_no_alloc(ctx)) { - set_timestep_embedding(timesteps, embedding, dim, max_period); - } - return embedding; -} - -std::vector ggml_to_image_vec(struct ggml_tensor* t) { - int64_t w = t->ne[0]; - int64_t h = t->ne[1]; - int64_t c = t->ne[2]; - std::vector vec; - vec.resize(w * h * c); - uint8_t* data = (uint8_t*)vec.data(); - for (int i = 0; i < h; i++) { - for (int j = 0; j < w; j++) { - for (int k = 0; k < c; k++) { - float value = ggml_tensor_get_f32(t, j, i, k); - value = (value + 1.0f) * 0.5f; - if (value < 0) { - value = 0; - } else if (value > 1) { - value = 1; - } - value *= 255.f; - *(data + i * w * c + j * c + k) = (uint8_t)value; - } - } - } - return vec; -} - -void image_vec_to_ggml(const std::vector& vec, - struct ggml_tensor* t) { - int64_t w = t->ne[0]; - int64_t h = t->ne[1]; - int64_t c = t->ne[2]; - uint8_t* data = (uint8_t*)vec.data(); - for (int i = 0; i < h; i++) { - for (int j = 0; j < w; j++) { - for (int k = 0; k < c; k++) { - float value = *(data + i * w * c + j * c + k); - value = value / 255.f; - value = 2 * value - 1; - ggml_tensor_set_f32(t, value, j, i, k); - } - } - } -} - -struct ggml_tensor* ggml_group_norm_32(struct ggml_context* ctx, - struct ggml_tensor* a) { - return ggml_group_norm(ctx, a, 32); -} - -/*================================================== CLIPTokenizer ===================================================*/ - -const std::string UNK_TOKEN = "<|endoftext|>"; -const std::string BOS_TOKEN = "<|startoftext|>"; -const std::string EOS_TOKEN = "<|endoftext|>"; -const std::string PAD_TOEKN = "<|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; - -// Ref: https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py -// TODO: implement bpe -class CLIPTokenizer { - private: - ModelType model_type = SD1; - std::map encoder; - std::regex pat; - - 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; - } - - public: - CLIPTokenizer(ModelType model_type = SD1) - : model_type(model_type){}; - std::string bpe(std::string token) { - std::string word = token + ""; - if (encoder.find(word) != encoder.end()) { - return word; - } else if (encoder.find(token) != encoder.end()) { - return token; - } - return UNK_TOKEN; - } - - void add_token(std::string token, int32_t token_id) { - encoder[token] = token_id; - } - - std::vector tokenize(std::string text, size_t max_length = 0, bool padding = false) { - std::vector tokens = encode(text); - 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) { - int pad_token_id = PAD_TOKEN_ID; - if (model_type == SD2) { - pad_token_id = 0; - } - tokens.insert(tokens.end(), max_length - tokens.size(), pad_token_id); - } - } - } - return tokens; - } - - std::vector encode(std::string text) { - std::string original_text = text; - std::vector bpe_tokens; - text = whitespace_clean(text); - std::transform(text.begin(), text.end(), text.begin(), [](unsigned char c) { return std::tolower(c); }); - - std::regex pat(R"(<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[[:alpha:]]+|[[:digit:]]|[^[:space:][:alpha:][:digit:]]+)", - std::regex::icase); - - std::smatch matches; - std::string str = text; - std::vector token_strs; - while (std::regex_search(str, matches, pat)) { - for (auto& token : matches) { - std::istringstream iss(bpe(token)); - std::vector tokens{std::istream_iterator{iss}, - std::istream_iterator{}}; - for (const auto& bpe_token : tokens) { - bpe_tokens.push_back(encoder[bpe_token]); - token_strs.push_back(bpe_token); - } - } - str = matches.suffix(); - } - 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()); - return bpe_tokens; - } -}; - -// Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/prompt_parser.py#L345 -// -// Parses a string with attention tokens and returns a list of pairs: text and its associated weight. -// Accepted tokens are: -// (abc) - increases attention to abc by a multiplier of 1.1 -// (abc:3.12) - increases attention to abc by a multiplier of 3.12 -// [abc] - decreases attention to abc by a multiplier of 1.1 -// \( - literal character '(' -// \[ - literal character '[' -// \) - literal character ')' -// \] - literal character ']' -// \\ - literal character '\' -// anything else - just text -// -// >>> parse_prompt_attention('normal text') -// [['normal text', 1.0]] -// >>> parse_prompt_attention('an (important) word') -// [['an ', 1.0], ['important', 1.1], [' word', 1.0]] -// >>> parse_prompt_attention('(unbalanced') -// [['unbalanced', 1.1]] -// >>> parse_prompt_attention('\(literal\]') -// [['(literal]', 1.0]] -// >>> parse_prompt_attention('(unnecessary)(parens)') -// [['unnecessaryparens', 1.1]] -// >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).') -// [['a ', 1.0], -// ['house', 1.5730000000000004], -// [' ', 1.1], -// ['on', 1.0], -// [' a ', 1.1], -// ['hill', 0.55], -// [', sun, ', 1.1], -// ['sky', 1.4641000000000006], -// ['.', 1.1]] -std::vector> parse_prompt_attention(const std::string& text) { - std::vector> res; - std::vector round_brackets; - std::vector square_brackets; - - float round_bracket_multiplier = 1.1f; - float square_bracket_multiplier = 1 / 1.1f; - - std::regex re_attention(R"(\\\(|\\\)|\\\[|\\\]|\\\\|\\|\(|\[|:([+-]?[.\d]+)\)|\)|\]|[^\\()\[\]:]+|:)"); - std::regex re_break(R"(\s*\bBREAK\b\s*)"); - - auto multiply_range = [&](int start_position, float multiplier) { - for (int p = start_position; p < res.size(); ++p) { - res[p].second *= multiplier; - } - }; - - std::smatch m; - std::string remaining_text = text; - - while (std::regex_search(remaining_text, m, re_attention)) { - std::string text = m[0]; - std::string weight = m[1]; - - if (text == "(") { - round_brackets.push_back(res.size()); - } else if (text == "[") { - square_brackets.push_back(res.size()); - } else if (!weight.empty()) { - if (!round_brackets.empty()) { - multiply_range(round_brackets.back(), std::stod(weight)); - round_brackets.pop_back(); - } - } else if (text == ")" && !round_brackets.empty()) { - multiply_range(round_brackets.back(), round_bracket_multiplier); - round_brackets.pop_back(); - } else if (text == "]" && !square_brackets.empty()) { - multiply_range(square_brackets.back(), square_bracket_multiplier); - square_brackets.pop_back(); - } else if (text == "\\(") { - res.push_back({text.substr(1), 1.0f}); - } else { - res.push_back({text, 1.0f}); - } - - remaining_text = m.suffix(); - } - - for (int pos : round_brackets) { - multiply_range(pos, round_bracket_multiplier); - } - - for (int pos : square_brackets) { - multiply_range(pos, square_bracket_multiplier); - } - - if (res.empty()) { - res.push_back({"", 1.0f}); - } - - int i = 0; - while (i + 1 < res.size()) { - if (res[i].second == res[i + 1].second) { - res[i].first += res[i + 1].first; - res.erase(res.begin() + i + 1); - } else { - ++i; - } - } - - return res; -} - -/*================================================ FrozenCLIPEmbedder ================================================*/ - -struct ResidualAttentionBlock { - int32_t n_head; - int32_t d_model; - int32_t hidden_size; // n_head * d_model - int32_t intermediate_size; - - // attention - struct ggml_tensor* q_w; // [hidden_size, hidden_size] - struct ggml_tensor* q_b; // [hidden_size, ] - struct ggml_tensor* k_w; // [hidden_size, hidden_size] - struct ggml_tensor* k_b; // [hidden_size, ] - struct ggml_tensor* v_w; // [hidden_size, hidden_size] - struct ggml_tensor* v_b; // [hidden_size, ] - - struct ggml_tensor* out_w; // [hidden_size, hidden_size] - struct ggml_tensor* out_b; // [hidden_size, ] - - // layer norm 1 - struct ggml_tensor* ln1_w; // [hidden_size, ] - struct ggml_tensor* ln1_b; // [hidden_size, ] - - // mlp - struct ggml_tensor* fc1_w; // [intermediate_size, hidden_size] - struct ggml_tensor* fc1_b; // [intermediate_size, ] - - struct ggml_tensor* fc2_w; // [hidden_size, intermediate_size] - struct ggml_tensor* fc2_b; // [hidden_size, ] - - // layer norm 2 - struct ggml_tensor* ln2_w; // [hidden_size, ] - struct ggml_tensor* ln2_b; // [hidden_size, ] - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += 4 * hidden_size * hidden_size * ggml_type_sizef(wtype); // q_w/k_w/v_w/out_w - mem_size += 8 * hidden_size * ggml_type_sizef(GGML_TYPE_F32); // q_b/k_b/v_b/out_b/ln1_w/ln1_b/ln2_w/ln2_b - mem_size += 2 * hidden_size * intermediate_size * ggml_type_sizef(wtype); // fc1_w/fc2_w - mem_size += intermediate_size * ggml_type_sizef(GGML_TYPE_F32); // fc1_b - mem_size += hidden_size * ggml_type_sizef(GGML_TYPE_F32); // fc2_b - mem_size += 16 * ggml_tensor_overhead(); // tensor overhead - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - ln1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - ln1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - - q_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size); - q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - k_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size); - k_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - v_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size); - v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - - out_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size); - out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - - fc1_w = ggml_new_tensor_2d(ctx, wtype, hidden_size, intermediate_size); - fc1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, intermediate_size); - - fc2_w = ggml_new_tensor_2d(ctx, wtype, intermediate_size, hidden_size); - fc2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - - ln2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - ln2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "self_attn.q_proj.weight"] = q_w; - tensors[prefix + "self_attn.q_proj.bias"] = q_b; - tensors[prefix + "self_attn.k_proj.weight"] = k_w; - tensors[prefix + "self_attn.k_proj.bias"] = k_b; - tensors[prefix + "self_attn.v_proj.weight"] = v_w; - tensors[prefix + "self_attn.v_proj.bias"] = v_b; - tensors[prefix + "self_attn.out_proj.weight"] = out_w; - tensors[prefix + "self_attn.out_proj.bias"] = out_b; - - tensors[prefix + "layer_norm1.weight"] = ln1_w; - tensors[prefix + "layer_norm1.bias"] = ln1_b; - - tensors[prefix + "layer_norm2.weight"] = ln2_w; - tensors[prefix + "layer_norm2.bias"] = ln2_b; - - tensors[prefix + "mlp.fc1.weight"] = fc1_w; - tensors[prefix + "mlp.fc1.bias"] = fc1_b; - - tensors[prefix + "mlp.fc2.weight"] = fc2_w; - tensors[prefix + "mlp.fc2.bias"] = fc2_b; - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) { - // x: [N, n_token, hidden_size] - int64_t N = x->ne[2]; - int64_t n_token = x->ne[1]; - int64_t hidden_size = n_head * d_model; - - struct ggml_tensor* r = x; - - // layer norm 1 - { - x = ggml_norm(ctx, x, 1e-6f); - x = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ln1_w, x), x), - ggml_repeat(ctx, ln1_b, x)); - } - // self-attention - { - struct ggml_tensor* q = ggml_add(ctx, - ggml_repeat(ctx, q_b, x), - ggml_mul_mat(ctx, q_w, x)); - q = ggml_scale_inplace(ctx, q, ggml_new_f32(ctx, 1.0f / sqrt((float)d_model))); - q = ggml_reshape_4d(ctx, q, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model] - q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, n_token, d_model] - q = ggml_reshape_3d(ctx, q, d_model, n_token, n_head * N); // [N * n_head, n_token, d_model] - - struct ggml_tensor* k = ggml_add(ctx, - ggml_repeat(ctx, k_b, x), - ggml_mul_mat(ctx, k_w, x)); - k = ggml_reshape_4d(ctx, k, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model] - k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_head, n_token, d_model] - k = ggml_reshape_3d(ctx, k, d_model, n_token, n_head); // [N * n_head, n_token, d_model] - - struct ggml_tensor* v = ggml_add(ctx, - ggml_repeat(ctx, v_b, x), - ggml_mul_mat(ctx, v_w, x)); - v = ggml_reshape_4d(ctx, v, d_model, n_head, n_token, N); // [N, n_token, n_head, d_model] - v = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_head, d_model, n_token] - v = ggml_reshape_3d(ctx, v, n_token, d_model, n_head * N); // [N * n_head, d_model, n_token] - - struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q); // [N * n_head, n_token, n_token] - - kq = ggml_diag_mask_inf_inplace(ctx, kq, 0); - kq = ggml_soft_max_inplace(ctx, kq); - - struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, n_token, d_model] - kqv = ggml_reshape_4d(ctx, kqv, d_model, n_token, n_head, N); - kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); // [N, n_token, n_head, d_model] - - x = ggml_reshape_2d(ctx, kqv, d_model * n_head, n_token * N); // // [N * n_token, d_model * n_head] - } - - // attention output - x = ggml_add(ctx, ggml_repeat(ctx, out_b, x), ggml_mul_mat(ctx, out_w, x)); - - // residual - x = ggml_add(ctx, x, r); - r = x; - - // layer norm 2 - { - x = ggml_norm(ctx, x, 1e-6f); - - x = ggml_add(ctx, ggml_mul(ctx, ggml_repeat(ctx, ln2_w, x), x), - ggml_repeat(ctx, ln2_b, x)); - } - - // mlp - x = ggml_mul_mat(ctx, fc1_w, x); - x = ggml_add(ctx, ggml_repeat(ctx, fc1_b, x), x); - - if (hidden_size == 1024) { // SD 2.x - x = ggml_gelu_inplace(ctx, x); - } else { // SD 1.x - x = ggml_gelu_quick_inplace(ctx, x); - } - - x = ggml_mul_mat(ctx, fc2_w, x); - x = ggml_add(ctx, ggml_repeat(ctx, fc2_b, x), x); - - // residual 2 - x = ggml_add(ctx, x, r); - - return x; - } -}; - -// SD1.x: https://huggingface.co/openai/clip-vit-large-patch14/blob/main/config.json -// SD2.x: https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/blob/main/config.json -struct CLIPTextModel { - ModelType model_type = SD1; - // network hparams - int32_t vocab_size = 49408; - int32_t max_position_embeddings = 77; - int32_t hidden_size = 768; // 1024 for SD 2.x - int32_t intermediate_size = 3072; // 4096 for SD 2.x - int32_t n_head = 12; // num_attention_heads, 16 for SD 2.x - int32_t num_hidden_layers = 12; // 24 for SD 2.x - - // embeddings - struct ggml_tensor* position_ids; - struct ggml_tensor* token_embed_weight; - struct ggml_tensor* position_embed_weight; - // transformer - std::vector resblocks; - struct ggml_tensor* final_ln_w; - struct ggml_tensor* final_ln_b; - - CLIPTextModel(ModelType model_type = SD1) - : model_type(model_type) { - if (model_type == SD2) { - hidden_size = 1024; - intermediate_size = 4096; - n_head = 16; - num_hidden_layers = 24; - } - resblocks.resize(num_hidden_layers); - set_resblocks_hp_params(); - } - - void set_resblocks_hp_params() { - int d_model = hidden_size / n_head; // 64 - for (int i = 0; i < num_hidden_layers; i++) { - resblocks[i].d_model = d_model; - resblocks[i].n_head = n_head; - resblocks[i].hidden_size = hidden_size; - resblocks[i].intermediate_size = intermediate_size; - } - } - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += hidden_size * max_position_embeddings * ggml_type_sizef(GGML_TYPE_I32); // position_ids - mem_size += hidden_size * vocab_size * ggml_type_sizef(wtype); // token_embed_weight - mem_size += hidden_size * max_position_embeddings * ggml_type_sizef(wtype); // position_embed_weight - for (int i = 0; i < num_hidden_layers; i++) { - mem_size += resblocks[i].compute_params_mem_size(wtype); - } - mem_size += 2 * hidden_size * ggml_type_sizef(GGML_TYPE_F32); // final_ln_w/b - mem_size += ggml_tensor_overhead(); // object overhead - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - position_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, max_position_embeddings); - for (int i = 0; i < max_position_embeddings; i++) { - ggml_set_i32_1d(position_ids, i, i); - } - token_embed_weight = ggml_new_tensor_2d(ctx, wtype, hidden_size, vocab_size); - position_embed_weight = ggml_new_tensor_2d(ctx, wtype, hidden_size, max_position_embeddings); - - for (int i = 0; i < num_hidden_layers; i++) { - resblocks[i].init_params(ctx, wtype); - } - - final_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - final_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "embeddings.token_embedding.weight"] = token_embed_weight; - tensors[prefix + "embeddings.position_embedding.weight"] = position_embed_weight; - tensors[prefix + "final_layer_norm.weight"] = final_ln_w; - tensors[prefix + "final_layer_norm.bias"] = final_ln_b; - for (int i = 0; i < num_hidden_layers; i++) { - resblocks[i].map_by_name(tensors, prefix + "encoder.layers." + std::to_string(i) + "."); - } - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* input_ids) { - // input_ids: [N, n_token] - GGML_ASSERT(input_ids->ne[0] <= position_ids->ne[0]); - - // token_embedding + position_embedding - struct ggml_tensor* x; - x = ggml_add(ctx, - ggml_get_rows(ctx, token_embed_weight, input_ids), - ggml_get_rows(ctx, - position_embed_weight, - ggml_view_1d(ctx, position_ids, input_ids->ne[0], 0))); // [N, n_token, hidden_size] - - // transformer - for (int i = 0; i < num_hidden_layers; i++) { - if (model_type == SD2 && i == num_hidden_layers - 1) { // layer: "penultimate" - break; - } - x = resblocks[i].forward(ctx, x); // [N, n_token, hidden_size] - } - - // final layer norm - { - x = ggml_norm(ctx, x, 1e-6f); - - x = ggml_add(ctx, ggml_mul(ctx, ggml_repeat(ctx, final_ln_w, x), x), - ggml_repeat(ctx, final_ln_b, x)); - } - - return x; // [N, n_token, hidden_size] - } -}; - -// ldm.modules.encoders.modules.FrozenCLIPEmbedder -struct FrozenCLIPEmbedder { - CLIPTokenizer tokenizer; - CLIPTextModel text_model; - struct ggml_tensor* forward(struct ggml_context* ctx, const std::string& prompt) { - std::vector tokens = tokenizer.tokenize(prompt, text_model.max_position_embeddings, true); - struct ggml_tensor* input_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens.size()); - memcpy(input_ids->data, tokens.data(), tokens.size() * ggml_element_size(input_ids)); - struct ggml_tensor* hidden_states = text_model.forward(ctx, input_ids); - return hidden_states; - } -}; - -// Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/sd_hijack_clip.py#L283 -struct FrozenCLIPEmbedderWithCustomWords { - ModelType model_type = SD1; - CLIPTokenizer tokenizer; - CLIPTextModel text_model; - - FrozenCLIPEmbedderWithCustomWords(ModelType model_type = SD1) - : model_type(model_type), tokenizer(model_type), text_model(model_type) {} - - std::pair, std::vector> tokenize(std::string text, - size_t max_length = 0, - bool padding = false) { - auto parsed_attention = parse_prompt_attention(text); - - { - std::stringstream ss; - ss << "["; - for (const auto& item : parsed_attention) { - ss << "['" << item.first << "', " << item.second << "], "; - } - ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); - } - - std::vector tokens; - std::vector weights; - for (const auto& item : parsed_attention) { - const std::string& curr_text = item.first; - float curr_weight = item.second; - std::vector 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); - } - tokens.insert(tokens.begin(), BOS_TOKEN_ID); - weights.insert(weights.begin(), 1.0); - - if (max_length > 0) { - if (tokens.size() > max_length - 1) { - tokens.resize(max_length - 1); - weights.resize(max_length - 1); - tokens.push_back(EOS_TOKEN_ID); - weights.push_back(1.0); - } else { - tokens.push_back(EOS_TOKEN_ID); - weights.push_back(1.0); - if (padding) { - int pad_token_id = PAD_TOKEN_ID; - if (model_type == SD2) { - pad_token_id = 0; - } - tokens.insert(tokens.end(), max_length - tokens.size(), pad_token_id); - weights.insert(weights.end(), max_length - weights.size(), 1.0); - } - } - } - - // for (int i = 0; i < tokens.size(); i++) { - // std::cout << tokens[i] << ":" << weights[i] << ", "; - // } - // std::cout << std::endl; - - return {tokens, weights}; - } -}; - -/*==================================================== UnetModel =====================================================*/ - -struct ResBlock { - // network hparams - int channels; // model_channels * (1, 1, 1, 2, 2, 4, 4, 4) - int emb_channels; // time_embed_dim - int out_channels; // mult * model_channels - - // network params - // in_layers - struct ggml_tensor* in_layer_0_w; // [channels, ] - struct ggml_tensor* in_layer_0_b; // [channels, ] - // in_layer_1 is nn.SILU() - struct ggml_tensor* in_layer_2_w; // [out_channels, channels, 3, 3] - struct ggml_tensor* in_layer_2_b; // [out_channels, ] - - // emb_layers - // emb_layer_0 is nn.SILU() - struct ggml_tensor* emb_layer_1_w; // [out_channels, emb_channels] - struct ggml_tensor* emb_layer_1_b; // [out_channels, ] - - // out_layers - struct ggml_tensor* out_layer_0_w; // [out_channels, ] - struct ggml_tensor* out_layer_0_b; // [out_channels, ] - // out_layer_1 is nn.SILU() - // out_layer_2 is nn.Dropout(), p = 0 for inference - struct ggml_tensor* out_layer_3_w; // [out_channels, out_channels, 3, 3] - struct ggml_tensor* out_layer_3_b; // [out_channels, ] - - // skip connection, only if out_channels != channels - struct ggml_tensor* skip_w; // [out_channels, channels, 1, 1] - struct ggml_tensor* skip_b; // [out_channels, ] - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += 2 * channels * ggml_type_sizef(GGML_TYPE_F32); // in_layer_0_w/b - mem_size += out_channels * channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // in_layer_2_w - mem_size += 5 * out_channels * ggml_type_sizef(GGML_TYPE_F32); // in_layer_2_b/emb_layer_1_b/out_layer_0_w/out_layer_0_b/out_layer_3_b - mem_size += out_channels * emb_channels * ggml_type_sizef(wtype); // emb_layer_1_w - mem_size += out_channels * out_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // out_layer_3_w - - mem_size += 10 * ggml_tensor_overhead(); // object overhead - - if (out_channels != channels) { - mem_size += out_channels * channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16); // skip_w - mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32); // skip_b - - mem_size += 2 * ggml_tensor_overhead(); // object overhead - } - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - in_layer_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); - in_layer_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); - in_layer_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels); - in_layer_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - - emb_layer_1_w = ggml_new_tensor_2d(ctx, wtype, emb_channels, out_channels); - emb_layer_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - - out_layer_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - out_layer_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - out_layer_3_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels); - out_layer_3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - - if (out_channels != channels) { - skip_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, channels, out_channels); - skip_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - } - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "in_layers.0.weight"] = in_layer_0_w; - tensors[prefix + "in_layers.0.bias"] = in_layer_0_b; - tensors[prefix + "in_layers.2.weight"] = in_layer_2_w; - tensors[prefix + "in_layers.2.bias"] = in_layer_2_b; - - tensors[prefix + "emb_layers.1.weight"] = emb_layer_1_w; - tensors[prefix + "emb_layers.1.bias"] = emb_layer_1_b; - - tensors[prefix + "out_layers.0.weight"] = out_layer_0_w; - tensors[prefix + "out_layers.0.bias"] = out_layer_0_b; - tensors[prefix + "out_layers.3.weight"] = out_layer_3_w; - tensors[prefix + "out_layers.3.bias"] = out_layer_3_b; - - if (out_channels != channels) { - tensors[prefix + "skip_connection.weight"] = skip_w; - tensors[prefix + "skip_connection.bias"] = skip_b; - } - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* emb) { - // x: [N, channels, h, w] - // emb: [N, emb_channels] - - // in_layers - // group norm 32 - auto h = ggml_group_norm_32(ctx, x); - h = ggml_add(ctx, - ggml_mul(ctx, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, in_layer_0_w, 1, 1, in_layer_0_w->ne[0], 1), - h), - h), - ggml_repeat(ctx, - ggml_reshape_4d(ctx, in_layer_0_b, 1, 1, in_layer_0_b->ne[0], 1), - h)); - // silu - h = ggml_silu_inplace(ctx, h); - // conv2d - h = ggml_conv_2d(ctx, in_layer_2_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, in_layer_2_b, 1, 1, in_layer_2_b->ne[0], 1), - h)); // [N, out_channels, h, w] - - // emb_layers - auto emb_out = ggml_silu(ctx, emb); - emb_out = ggml_mul_mat(ctx, emb_layer_1_w, emb_out); - emb_out = ggml_add(ctx, ggml_repeat(ctx, emb_layer_1_b, emb_out), emb_out); // [N, out_channels] - emb_out = ggml_reshape_4d(ctx, emb_out, 1, 1, emb_out->ne[0], emb_out->ne[1]); // [N, out_channels, 1, 1] - emb_out = ggml_repeat(ctx, emb_out, h); // [N, out_channels, h, w] - - // out_layers - h = ggml_add(ctx, h, emb_out); - // group norm 32 - h = ggml_group_norm_inplace(ctx, h, 32); - h = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, out_layer_0_w, 1, 1, out_layer_0_w->ne[0], 1), h), h), - ggml_repeat(ctx, ggml_reshape_4d(ctx, out_layer_0_b, 1, 1, out_layer_0_b->ne[0], 1), h)); - // silu - h = ggml_silu_inplace(ctx, h); - // dropout, skip for inference - // conv2d - h = ggml_conv_2d(ctx, out_layer_3_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, out_layer_3_b, 1, 1, out_layer_3_b->ne[0], 1), - h)); // [N, out_channels, h, w - - // skip connection - if (out_channels != channels) { - x = ggml_conv_2d(ctx, skip_w, x, 1, 1, 0, 0, 1, 1); - x = ggml_add(ctx, - x, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, skip_b, 1, 1, skip_b->ne[0], 1), - x)); // [N, out_channels, h, w] - } - h = ggml_add(ctx, h, x); - return h; // [N, out_channels, h, w] - } -}; - -struct SpatialTransformer { - int in_channels; // mult * model_channels - int n_head; // num_heads - int d_head; // in_channels // n_heads - int depth = 1; // 1 - int context_dim = 768; // hidden_size, 1024 for SD2.x - - // group norm - struct ggml_tensor* norm_w; // [in_channels,] - struct ggml_tensor* norm_b; // [in_channels,] - - // proj_in - struct ggml_tensor* proj_in_w; // [in_channels, in_channels, 1, 1] - struct ggml_tensor* proj_in_b; // [in_channels,] - - // transformer - struct - { - // layer norm 1 - struct ggml_tensor* norm1_w; // [in_channels, ] - struct ggml_tensor* norm1_b; // [in_channels, ] - - // attn1 - struct ggml_tensor* attn1_q_w; // [in_channels, in_channels] - struct ggml_tensor* attn1_k_w; // [in_channels, in_channels] - struct ggml_tensor* attn1_v_w; // [in_channels, in_channels] - - struct ggml_tensor* attn1_out_w; // [in_channels, in_channels] - struct ggml_tensor* attn1_out_b; // [in_channels, ] - - // layer norm 2 - struct ggml_tensor* norm2_w; // [in_channels, ] - struct ggml_tensor* norm2_b; // [in_channels, ] - - // attn2 - struct ggml_tensor* attn2_q_w; // [in_channels, in_channels] - struct ggml_tensor* attn2_k_w; // [in_channels, context_dim] - struct ggml_tensor* attn2_v_w; // [in_channels, context_dim] - - struct ggml_tensor* attn2_out_w; // [in_channels, in_channels] - struct ggml_tensor* attn2_out_b; // [in_channels, ] - - // layer norm 3 - struct ggml_tensor* norm3_w; // [in_channels, ] - struct ggml_tensor* norm3_b; // [in_channels, ] - - // ff - struct ggml_tensor* ff_0_proj_w; // [in_channels * 4 * 2, in_channels] - struct ggml_tensor* ff_0_proj_b; // [in_channels * 4 * 2] - - struct ggml_tensor* ff_2_w; // [in_channels, in_channels * 4] - struct ggml_tensor* ff_2_b; // [in_channels,] - } transformer; - - // proj_out - struct ggml_tensor* proj_out_w; // [in_channels, in_channels, 1, 1] - struct ggml_tensor* proj_out_b; // [in_channels,] - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32); // norm_w/norm_b - mem_size += 2 * in_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16); // proj_in_w/proj_out_w - mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32); // proj_in_b/proj_out_b - - // transformer - { - mem_size += 6 * in_channels * ggml_type_sizef(GGML_TYPE_F32); // norm1-3_w/b - mem_size += 6 * in_channels * in_channels * ggml_type_sizef(wtype); // attn1_q/k/v/out_w attn2_q/out_w - mem_size += 2 * in_channels * context_dim * ggml_type_sizef(wtype); // attn2_k/v_w - mem_size += in_channels * 4 * 2 * in_channels * ggml_type_sizef(wtype); // ff_0_proj_w - mem_size += in_channels * 4 * 2 * ggml_type_sizef(GGML_TYPE_F32); // ff_0_proj_b - mem_size += in_channels * 4 * in_channels * ggml_type_sizef(wtype); // ff_2_w - mem_size += in_channels * ggml_type_sizef(GGML_TYPE_F32); // ff_2_b - } - mem_size += 26 * ggml_tensor_overhead(); // object overhead - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - norm_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - norm_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - proj_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels); - proj_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels); - proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - // transformer - transformer.norm1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - transformer.norm1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - transformer.attn1_q_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels); - transformer.attn1_k_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels); - transformer.attn1_v_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels); - - transformer.attn1_out_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels); - transformer.attn1_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - transformer.norm2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - transformer.norm2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - transformer.attn2_q_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels); - transformer.attn2_k_w = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels); - transformer.attn2_v_w = ggml_new_tensor_2d(ctx, wtype, context_dim, in_channels); - - transformer.attn2_out_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels); - transformer.attn2_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - transformer.norm3_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - transformer.norm3_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - transformer.ff_0_proj_w = ggml_new_tensor_2d(ctx, wtype, in_channels, in_channels * 4 * 2); - transformer.ff_0_proj_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels * 4 * 2); - - transformer.ff_2_w = ggml_new_tensor_2d(ctx, wtype, in_channels * 4, in_channels); - transformer.ff_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "norm.weight"] = norm_w; - tensors[prefix + "norm.bias"] = norm_b; - tensors[prefix + "proj_in.weight"] = proj_in_w; - tensors[prefix + "proj_in.bias"] = proj_in_b; - - // transformer - { - std::string transformer_prefix = prefix + "transformer_blocks.0."; - tensors[transformer_prefix + "attn1.to_q.weight"] = transformer.attn1_q_w; - tensors[transformer_prefix + "attn1.to_k.weight"] = transformer.attn1_k_w; - tensors[transformer_prefix + "attn1.to_v.weight"] = transformer.attn1_v_w; - - tensors[transformer_prefix + "attn1.to_out.0.weight"] = transformer.attn1_out_w; - tensors[transformer_prefix + "attn1.to_out.0.bias"] = transformer.attn1_out_b; - - tensors[transformer_prefix + "ff.net.0.proj.weight"] = transformer.ff_0_proj_w; - tensors[transformer_prefix + "ff.net.0.proj.bias"] = transformer.ff_0_proj_b; - tensors[transformer_prefix + "ff.net.2.weight"] = transformer.ff_2_w; - tensors[transformer_prefix + "ff.net.2.bias"] = transformer.ff_2_b; - - tensors[transformer_prefix + "attn2.to_q.weight"] = transformer.attn2_q_w; - tensors[transformer_prefix + "attn2.to_k.weight"] = transformer.attn2_k_w; - tensors[transformer_prefix + "attn2.to_v.weight"] = transformer.attn2_v_w; - - tensors[transformer_prefix + "attn2.to_out.0.weight"] = transformer.attn2_out_w; - tensors[transformer_prefix + "attn2.to_out.0.bias"] = transformer.attn2_out_b; - - tensors[transformer_prefix + "norm1.weight"] = transformer.norm1_w; - tensors[transformer_prefix + "norm1.bias"] = transformer.norm1_b; - tensors[transformer_prefix + "norm2.weight"] = transformer.norm2_w; - tensors[transformer_prefix + "norm2.bias"] = transformer.norm2_b; - tensors[transformer_prefix + "norm3.weight"] = transformer.norm3_w; - tensors[transformer_prefix + "norm3.bias"] = transformer.norm3_b; - } - - tensors[prefix + "proj_out.weight"] = proj_out_w; - tensors[prefix + "proj_out.bias"] = proj_out_b; - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* context) { - // x: [N, in_channels, h, w] - // context: [N, max_position, hidden_size(aka context_dim)] - - auto x_in = x; - // group norm 32 - x = ggml_group_norm_32(ctx, x); - x = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_w, 1, 1, norm_w->ne[0], 1), x), x), - ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_b, 1, 1, norm_b->ne[0], 1), x)); - // proj_in - x = ggml_conv_2d(ctx, proj_in_w, x, 1, 1, 0, 0, 1, 1); - x = ggml_add(ctx, - x, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, proj_in_b, 1, 1, proj_in_b->ne[0], 1), - x)); // [N, in_channels, h, w] - - // transformer - const int64_t n = x->ne[3]; - const int64_t c = x->ne[2]; - const int64_t h = x->ne[1]; - const int64_t w = x->ne[0]; - const int64_t max_position = context->ne[1]; - x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3)); // [N, h, w, in_channels] - - { - auto r = x; - // layer norm 1 - { - x = ggml_reshape_2d(ctx, x, c, w * h * n); - x = ggml_norm(ctx, x, 1e-6f); - x = ggml_add(ctx, - ggml_mul(ctx, - ggml_repeat(ctx, transformer.norm1_w, x), - x), - ggml_repeat(ctx, transformer.norm1_b, x)); - } - - // self-attention - { - x = ggml_reshape_2d(ctx, x, c, h * w * n); // [N * h * w, in_channels] - struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn1_q_w, x); // [N * h * w, in_channels] - q = ggml_scale_inplace(ctx, q, ggml_new_f32(ctx, 1.0f / sqrt((float)d_head))); - q = ggml_reshape_4d(ctx, q, d_head, n_head, h * w, n); // [N, h * w, n_head, d_head] - q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, h * w, d_head] - q = ggml_reshape_3d(ctx, q, d_head, h * w, n_head * n); // [N * n_head, h * w, d_head] - - struct ggml_tensor* k = ggml_mul_mat(ctx, transformer.attn1_k_w, x); // [N * h * w, in_channels] - k = ggml_reshape_4d(ctx, k, d_head, n_head, h * w, n); // [N, h * w, n_head, d_head] - k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_head, h * w, d_head] - k = ggml_reshape_3d(ctx, k, d_head, h * w, n_head * n); // [N * n_head, h * w, d_head] - - struct ggml_tensor* v = ggml_mul_mat(ctx, transformer.attn1_v_w, x); // [N * h * w, in_channels] - v = ggml_reshape_4d(ctx, v, d_head, n_head, h * w, n); // [N, h * w, n_head, d_head] - v = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_head, d_head, h * w] - v = ggml_reshape_3d(ctx, v, h * w, d_head, n_head * n); // [N * n_head, d_head, h * w] - - struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q); // [N * n_head, h * w, h * w] - // kq = ggml_diag_mask_inf_inplace(ctx, kq, 0); - kq = ggml_soft_max_inplace(ctx, kq); - - struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, h * w, d_head] - kqv = ggml_reshape_4d(ctx, kqv, d_head, h * w, n_head, n); - kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); // [N, h * w, n_head, d_head] - - // x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n)); - x = ggml_reshape_2d(ctx, kqv, d_head * n_head, h * w * n); - - x = ggml_add(ctx, ggml_repeat(ctx, transformer.attn1_out_b, x), ggml_mul_mat(ctx, transformer.attn1_out_w, x)); - - x = ggml_reshape_4d(ctx, x, c, w, h, n); - } - - x = ggml_add(ctx, x, r); - r = x; - - // layer norm 2 - { - x = ggml_norm(ctx, x, 1e-6f); - x = ggml_add(ctx, - ggml_mul(ctx, - ggml_repeat(ctx, transformer.norm2_w, x), x), - ggml_repeat(ctx, transformer.norm2_b, x)); - } - - // cross-attention - { - x = ggml_reshape_2d(ctx, x, c, h * w * n); // [N * h * w, in_channels] - context = ggml_reshape_2d(ctx, context, context->ne[0], context->ne[1] * context->ne[2]); // [N * max_position, hidden_size] - struct ggml_tensor* q = ggml_mul_mat(ctx, transformer.attn2_q_w, x); // [N * h * w, in_channels] - - q = ggml_scale_inplace(ctx, q, ggml_new_f32(ctx, 1.0f / sqrt((float)d_head))); - q = ggml_reshape_4d(ctx, q, d_head, n_head, h * w, n); // [N, h * w, n_head, d_head] - q = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, h * w, d_head] - q = ggml_reshape_3d(ctx, q, d_head, h * w, n_head * n); // [N * n_head, h * w, d_head] - - struct ggml_tensor* k = ggml_mul_mat(ctx, transformer.attn2_k_w, context); // [N * max_position, in_channels] - k = ggml_reshape_4d(ctx, k, d_head, n_head, max_position, n); // [N, max_position, n_head, d_head] - k = ggml_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_head, max_position, d_head] - k = ggml_reshape_3d(ctx, k, d_head, max_position, n_head * n); // [N * n_head, max_position, d_head] - - struct ggml_tensor* v = ggml_mul_mat(ctx, transformer.attn2_v_w, context); // [N * max_position, in_channels] - v = ggml_reshape_4d(ctx, v, d_head, n_head, max_position, n); // [N, max_position, n_head, d_head] - v = ggml_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_head, d_head, max_position] - v = ggml_reshape_3d(ctx, v, max_position, d_head, n_head * n); // [N * n_head, d_head, max_position] - - struct ggml_tensor* kq = ggml_mul_mat(ctx, k, q); // [N * n_head, h * w, max_position] - // kq = ggml_diag_mask_inf_inplace(ctx, kq, 0); - kq = ggml_soft_max_inplace(ctx, kq); - - struct ggml_tensor* kqv = ggml_mul_mat(ctx, v, kq); // [N * n_head, h * w, d_head] - - kqv = ggml_reshape_4d(ctx, kqv, d_head, h * w, n_head, n); - kqv = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); - - // x = ggml_cpy(ctx, kqv, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_head * n_head, h * w * n)); // [N * h * w, in_channels] - x = ggml_reshape_2d(ctx, kqv, d_head * n_head, h * w * n); // [N * h * w, in_channels] - - x = ggml_add(ctx, ggml_repeat(ctx, transformer.attn2_out_b, x), ggml_mul_mat(ctx, transformer.attn2_out_w, x)); - - x = ggml_reshape_4d(ctx, x, c, w, h, n); - } - - x = ggml_add(ctx, x, r); - r = x; - - // layer norm 3 - { - x = ggml_reshape_2d(ctx, x, c, h * w * n); // [N * h * w, in_channels] - x = ggml_norm(ctx, x, 1e-6f); - x = ggml_add(ctx, - ggml_mul(ctx, - ggml_repeat(ctx, transformer.norm3_w, x), x), - ggml_repeat(ctx, transformer.norm3_b, x)); - } - - // ff - { - // GEGLU - auto x_w = ggml_view_2d(ctx, - transformer.ff_0_proj_w, - transformer.ff_0_proj_w->ne[0], - transformer.ff_0_proj_w->ne[1] / 2, - transformer.ff_0_proj_w->nb[1], - 0); // [in_channels * 4, in_channels] - auto x_b = ggml_view_1d(ctx, - transformer.ff_0_proj_b, - transformer.ff_0_proj_b->ne[0] / 2, - 0); // [in_channels * 4, in_channels] - auto gate_w = ggml_view_2d(ctx, - transformer.ff_0_proj_w, - transformer.ff_0_proj_w->ne[0], - transformer.ff_0_proj_w->ne[1] / 2, - transformer.ff_0_proj_w->nb[1], - transformer.ff_0_proj_w->nb[1] * transformer.ff_0_proj_w->ne[1] / 2); // [in_channels * 4, ] - auto gate_b = ggml_view_1d(ctx, - transformer.ff_0_proj_b, - transformer.ff_0_proj_b->ne[0] / 2, - transformer.ff_0_proj_b->nb[0] * transformer.ff_0_proj_b->ne[0] / 2); // [in_channels * 4, ] - x = ggml_reshape_2d(ctx, x, c, w * h * n); - auto x_in = x; - x = ggml_mul_mat(ctx, x_w, x_in); // [N * h * w, in_channels * 4] - x = ggml_add(ctx, ggml_repeat(ctx, x_b, x), x); - auto gate = ggml_mul_mat(ctx, gate_w, x_in); // [N * h * w, in_channels * 4] - gate = ggml_add(ctx, ggml_repeat(ctx, gate_b, gate), gate); - - gate = ggml_gelu_inplace(ctx, gate); - - x = ggml_mul(ctx, x, gate); // [N * h * w, in_channels * 4] - // fc - x = ggml_mul_mat(ctx, transformer.ff_2_w, x); // [N * h * w, in_channels] - x = ggml_add(ctx, ggml_repeat(ctx, transformer.ff_2_b, x), x); - } - - x = ggml_reshape_4d(ctx, x, c, w, h, n); // [N, h, w, in_channels] - - // residual - x = ggml_add(ctx, x, r); - } - x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3)); // // [N, in_channels, h, w] - - // proj_out - x = ggml_conv_2d(ctx, proj_out_w, x, 1, 1, 0, 0, 1, 1); - x = ggml_add(ctx, - x, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, proj_out_b, 1, 1, proj_out_b->ne[0], 1), - x)); // [N, in_channels, h, w] - x = ggml_add(ctx, x, x_in); - return x; - } -}; - -struct DownSample { - // hparams - int channels; - int out_channels; - - // conv2d params - struct ggml_tensor* op_w; // [out_channels, channels, 3, 3] - struct ggml_tensor* op_b; // [out_channels,] - - bool vae_downsample = false; - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += out_channels * channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // op_w - mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32); // op_b - mem_size += 2 * ggml_tensor_overhead(); // object overhead - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - op_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels); - op_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - if (vae_downsample) { - tensors[prefix + "conv.weight"] = op_w; - tensors[prefix + "conv.bias"] = op_b; - } else { - tensors[prefix + "op.weight"] = op_w; - tensors[prefix + "op.bias"] = op_b; - } - } - - // TODO: making it parallel - static void asymmetric_pad(struct ggml_tensor* dst, - const struct ggml_tensor* a, - const struct ggml_tensor* b, - int ith, - int nth, - void* userdata) { - assert(sizeof(dst->nb[0]) == sizeof(float)); - assert(sizeof(a->nb[0]) == sizeof(float)); - assert(sizeof(b->nb[0]) == sizeof(float)); - float value = 0; - - for (int i = 0; i < dst->ne[3]; i++) { - for (int j = 0; j < dst->ne[2]; j++) { - for (int k = 0; k < dst->ne[1]; k++) { - for (int l = 0; l < dst->ne[0]; l++) { - if (k == dst->ne[1] - 1 || l == dst->ne[0] - 1) { - value = 0; - } else { - value = ggml_tensor_get_f32(b, l, k, j, i); - } - // printf("%d %d %d %d -> %f\n", i, j, k, l, value); - ggml_tensor_set_f32(dst, value, l, k, j, i); - } - } - } - } - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) { - // x: [N, channels, h, w] - if (vae_downsample) { - bool dynamic = ggml_get_dynamic(ctx); - ggml_set_dynamic(ctx, false); - auto pad_x = ggml_new_tensor_4d(ctx, x->type, x->ne[0] + 1, x->ne[1] + 1, x->ne[2], x->ne[3]); - ggml_set_dynamic(ctx, dynamic); - - x = ggml_map_custom2_inplace(ctx, pad_x, x, asymmetric_pad, 1, NULL); - x = ggml_conv_2d(ctx, op_w, x, 2, 2, 0, 0, 1, 1); - } else { - x = ggml_conv_2d(ctx, op_w, x, 2, 2, 1, 1, 1, 1); - } - x = ggml_add(ctx, - x, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, op_b, 1, 1, op_b->ne[0], 1), - x)); // [N, out_channels, h/2, w/2] - return x; - } -}; - -struct UpSample { - // hparams - int channels; - int out_channels; - - // conv2d params - struct ggml_tensor* conv_w; // [out_channels, channels, 3, 3] - struct ggml_tensor* conv_b; // [out_channels,] - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += out_channels * channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // op_w - mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32); // op_b - mem_size += 2 * ggml_tensor_overhead(); // object overhead - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - conv_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, channels, out_channels); - conv_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "conv.weight"] = conv_w; - tensors[prefix + "conv.bias"] = conv_b; - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) { - // x: [N, channels, h, w] - x = ggml_upscale(ctx, x, 2); // [N, channels, h*2, w*2] - x = ggml_conv_2d(ctx, conv_w, x, 1, 1, 1, 1, 1, 1); - - x = ggml_add(ctx, - x, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv_b, 1, 1, conv_b->ne[0], 1), - x)); // [N, out_channels, h*2, w*2] - return x; - } -}; - -// ldm.modules.diffusionmodules.openaimodel.UNetModel -struct UNetModel { - // network hparams - int in_channels = 4; - int model_channels = 320; - int out_channels = 4; - int num_res_blocks = 2; - int attention_resolutions[3] = {4, 2, 1}; - int channel_mult[4] = {1, 2, 4, 4}; - int time_embed_dim = 1280; // model_channels*4 - int num_heads = 8; - int num_head_channels = -1; // channels // num_heads - int context_dim = 768; // 1024 for SD2.x - - // network params - struct ggml_tensor* time_embed_0_w; // [time_embed_dim, model_channels] - struct ggml_tensor* time_embed_0_b; // [time_embed_dim, ] - // time_embed_1 is nn.SILU() - struct ggml_tensor* time_embed_2_w; // [time_embed_dim, time_embed_dim] - struct ggml_tensor* time_embed_2_b; // [time_embed_dim, ] - - struct ggml_tensor* input_block_0_w; // [model_channels, in_channels, 3, 3] - struct ggml_tensor* input_block_0_b; // [model_channels, ] - - // input_blocks - ResBlock input_res_blocks[4][2]; - SpatialTransformer input_transformers[3][2]; - DownSample input_down_samples[3]; - - // middle_block - ResBlock middle_block_0; - SpatialTransformer middle_block_1; - ResBlock middle_block_2; - - // output_blocks - ResBlock output_res_blocks[4][3]; - SpatialTransformer output_transformers[3][3]; - UpSample output_up_samples[3]; - - // out - // group norm 32 - struct ggml_tensor* out_0_w; // [model_channels, ] - struct ggml_tensor* out_0_b; // [model_channels, ] - // out 1 is nn.SILU() - struct ggml_tensor* out_2_w; // [out_channels, model_channels, 3, 3] - struct ggml_tensor* out_2_b; // [out_channels, ] - - UNetModel(ModelType model_type = SD1) { - if (model_type == SD2) { - context_dim = 1024; - num_head_channels = 64; - num_heads = -1; - } - // set up hparams of blocks - - // input_blocks - std::vector input_block_chans; - input_block_chans.push_back(model_channels); - int ch = model_channels; - int ds = 1; - - int len_mults = sizeof(channel_mult) / sizeof(int); - for (int i = 0; i < len_mults; i++) { - int mult = channel_mult[i]; - for (int j = 0; j < num_res_blocks; j++) { - input_res_blocks[i][j].channels = ch; - input_res_blocks[i][j].emb_channels = time_embed_dim; - input_res_blocks[i][j].out_channels = mult * model_channels; - - ch = mult * model_channels; - - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - int n_head = num_heads; - int d_head = ch / num_heads; - if (num_head_channels != -1) { - d_head = num_head_channels; - n_head = ch / d_head; - } - input_transformers[i][j].in_channels = ch; - input_transformers[i][j].n_head = n_head; - input_transformers[i][j].d_head = d_head; - input_transformers[i][j].context_dim = context_dim; - } - input_block_chans.push_back(ch); - } - if (i != len_mults - 1) { - input_down_samples[i].channels = ch; - input_down_samples[i].out_channels = ch; - input_block_chans.push_back(ch); - - ds *= 2; - } - } - - // middle blocks - middle_block_0.channels = ch; - middle_block_0.emb_channels = time_embed_dim; - middle_block_0.out_channels = ch; - - int n_head = num_heads; - int d_head = ch / num_heads; - if (num_head_channels != -1) { - d_head = num_head_channels; - n_head = ch / d_head; - } - middle_block_1.in_channels = ch; - middle_block_1.n_head = n_head; - middle_block_1.d_head = d_head; - middle_block_1.context_dim = context_dim; - - middle_block_2.channels = ch; - middle_block_2.emb_channels = time_embed_dim; - middle_block_2.out_channels = ch; - - // output blocks - for (int i = len_mults - 1; i >= 0; i--) { - int mult = channel_mult[i]; - for (int j = 0; j < num_res_blocks + 1; j++) { - int ich = input_block_chans.back(); - input_block_chans.pop_back(); - - output_res_blocks[i][j].channels = ch + ich; - output_res_blocks[i][j].emb_channels = time_embed_dim; - output_res_blocks[i][j].out_channels = mult * model_channels; - - ch = mult * model_channels; - - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - int n_head = num_heads; - int d_head = ch / num_heads; - if (num_head_channels != -1) { - d_head = num_head_channels; - n_head = ch / d_head; - } - output_transformers[i][j].in_channels = ch; - output_transformers[i][j].n_head = n_head; - output_transformers[i][j].d_head = d_head; - output_transformers[i][j].context_dim = context_dim; - } - - if (i > 0 && j == num_res_blocks) { - output_up_samples[i - 1].channels = ch; - output_up_samples[i - 1].out_channels = ch; - - ds /= 2; - } - } - } - } - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += time_embed_dim * model_channels * ggml_type_sizef(wtype); // time_embed_0_w - mem_size += time_embed_dim * ggml_type_sizef(GGML_TYPE_F32); // time_embed_0_b - mem_size += time_embed_dim * time_embed_dim * ggml_type_sizef(wtype); // time_embed_2_w - mem_size += time_embed_dim * ggml_type_sizef(GGML_TYPE_F32); // time_embed_2_b - - mem_size += model_channels * in_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // input_block_0_w - mem_size += model_channels * ggml_type_sizef(GGML_TYPE_F32); // input_block_0_b - - mem_size += 6 * ggml_tensor_overhead(); // object overhead - - // input_blocks - int ds = 1; - int len_mults = sizeof(channel_mult) / sizeof(int); - for (int i = 0; i < len_mults; i++) { - for (int j = 0; j < num_res_blocks; j++) { - mem_size += input_res_blocks[i][j].compute_params_mem_size(wtype); - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - mem_size += input_transformers[i][j].compute_params_mem_size(wtype); - } - } - if (i != len_mults - 1) { - ds *= 2; - mem_size += input_down_samples[i].compute_params_mem_size(wtype); - } - } - - // middle_block - mem_size += middle_block_0.compute_params_mem_size(wtype); - mem_size += middle_block_1.compute_params_mem_size(wtype); - mem_size += middle_block_2.compute_params_mem_size(wtype); - - // output_blocks - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - mem_size += output_res_blocks[i][j].compute_params_mem_size(wtype); - - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - mem_size += output_transformers[i][j].compute_params_mem_size(wtype); - } - - if (i > 0 && j == num_res_blocks) { - mem_size += output_up_samples[i - 1].compute_params_mem_size(wtype); - - ds /= 2; - } - } - } - - // out - mem_size += 2 * model_channels * ggml_type_sizef(GGML_TYPE_F32); // out_0_w/b - mem_size += out_channels * model_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // out_2_w - mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32); // out_2_b - - mem_size += 4 * ggml_tensor_overhead(); - - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - time_embed_0_w = ggml_new_tensor_2d(ctx, wtype, model_channels, time_embed_dim); - time_embed_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, time_embed_dim); - - time_embed_2_w = ggml_new_tensor_2d(ctx, wtype, time_embed_dim, time_embed_dim); - time_embed_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, time_embed_dim); - - // input_blocks - input_block_0_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, in_channels, model_channels); - input_block_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, model_channels); - int ds = 1; - int len_mults = sizeof(channel_mult) / sizeof(int); - for (int i = 0; i < len_mults; i++) { - for (int j = 0; j < num_res_blocks; j++) { - input_res_blocks[i][j].init_params(ctx, wtype); - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - input_transformers[i][j].init_params(ctx, wtype); - } - } - if (i != len_mults - 1) { - input_down_samples[i].init_params(ctx, wtype); - ds *= 2; - } - } - - // middle_blocks - middle_block_0.init_params(ctx, wtype); - middle_block_1.init_params(ctx, wtype); - middle_block_2.init_params(ctx, wtype); - - // output_blocks - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - output_res_blocks[i][j].init_params(ctx, wtype); - - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - output_transformers[i][j].init_params(ctx, wtype); - } - - if (i > 0 && j == num_res_blocks) { - output_up_samples[i - 1].init_params(ctx, wtype); - - ds /= 2; - } - } - } - - // out - out_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, model_channels); - out_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, model_channels); - - out_2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, model_channels, out_channels); - out_2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "time_embed.0.weight"] = time_embed_0_w; - tensors[prefix + "time_embed.0.bias"] = time_embed_0_b; - - tensors[prefix + "time_embed.2.weight"] = time_embed_2_w; - tensors[prefix + "time_embed.2.bias"] = time_embed_2_b; - - // input_blocks - tensors[prefix + "input_blocks.0.0.weight"] = input_block_0_w; - tensors[prefix + "input_blocks.0.0.bias"] = input_block_0_b; - - int len_mults = sizeof(channel_mult) / sizeof(int); - int input_block_idx = 0; - int ds = 1; - for (int i = 0; i < len_mults; i++) { - for (int j = 0; j < num_res_blocks; j++) { - input_block_idx += 1; - - input_res_blocks[i][j].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".0."); - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - input_transformers[i][j].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".1."); - } - } - if (i != len_mults - 1) { - input_block_idx += 1; - input_down_samples[i].map_by_name(tensors, prefix + "input_blocks." + std::to_string(input_block_idx) + ".0."); - ds *= 2; - } - } - - // middle_blocks - middle_block_0.map_by_name(tensors, prefix + "middle_block.0."); - middle_block_1.map_by_name(tensors, prefix + "middle_block.1."); - middle_block_2.map_by_name(tensors, prefix + "middle_block.2."); - - // output_blocks - int output_block_idx = 0; - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - output_res_blocks[i][j].map_by_name(tensors, prefix + "output_blocks." + std::to_string(output_block_idx) + ".0."); - - int up_sample_idx = 1; - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - output_transformers[i][j].map_by_name(tensors, prefix + "output_blocks." + std::to_string(output_block_idx) + ".1."); - up_sample_idx++; - } - - if (i > 0 && j == num_res_blocks) { - output_up_samples[i - 1].map_by_name(tensors, prefix + "output_blocks." + std::to_string(output_block_idx) + "." + std::to_string(up_sample_idx) + "."); - - ds /= 2; - } - output_block_idx += 1; - } - } - - // out - tensors[prefix + "out.0.weight"] = out_0_w; - tensors[prefix + "out.0.bias"] = out_0_b; - tensors[prefix + "out.2.weight"] = out_2_w; - tensors[prefix + "out.2.bias"] = out_2_b; - } - - struct ggml_tensor* forward(struct ggml_context* ctx, - struct ggml_tensor* x, - struct ggml_tensor* timesteps, - struct ggml_tensor* context, - struct ggml_tensor* t_emb = NULL) { - // x: [N, in_channels, h, w] - // timesteps: [N, ] - // t_emb: [N, model_channels] - // context: [N, max_position, hidden_size]([N, 77, 768]) - if (t_emb == NULL && timesteps != NULL) { - t_emb = new_timestep_embedding(ctx, timesteps, model_channels); // [N, model_channels] - } - - // time_embed - auto emb = ggml_mul_mat(ctx, time_embed_0_w, t_emb); - emb = ggml_add(ctx, ggml_repeat(ctx, time_embed_0_b, emb), emb); - emb = ggml_silu_inplace(ctx, emb); - emb = ggml_mul_mat(ctx, time_embed_2_w, emb); - emb = ggml_add(ctx, ggml_repeat(ctx, time_embed_2_b, emb), emb); // [N, time_embed_dim] - - // input_blocks - std::vector hs; - // input block 0 - auto h = ggml_conv_2d(ctx, input_block_0_w, x, 1, 1, 1, 1, 1, 1); // [N, model_channels, h, w] - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, input_block_0_b, 1, 1, input_block_0_b->ne[0], 1), - h)); // [N, model_channels, h, w] - hs.push_back(h); - // input block 1-11 - int len_mults = sizeof(channel_mult) / sizeof(int); - int ds = 1; - for (int i = 0; i < len_mults; i++) { - int mult = channel_mult[i]; - for (int j = 0; j < num_res_blocks; j++) { - h = input_res_blocks[i][j].forward(ctx, h, emb); // [N, mult*model_channels, h, w] - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - h = input_transformers[i][j].forward(ctx, h, context); // [N, mult*model_channels, h, w] - } - hs.push_back(h); - } - if (i != len_mults - 1) { - ds *= 2; - h = input_down_samples[i].forward(ctx, h); // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))] - hs.push_back(h); - } - } - // [N, 4*model_channels, h/8, w/8] - - // middle_block - h = middle_block_0.forward(ctx, h, emb); // [N, 4*model_channels, h/8, w/8] - h = middle_block_1.forward(ctx, h, context); // [N, 4*model_channels, h/8, w/8] - h = middle_block_2.forward(ctx, h, emb); // [N, 4*model_channels, h/8, w/8] - - // output_blocks - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - auto h_skip = hs.back(); - hs.pop_back(); - - h = ggml_concat(ctx, h, h_skip); - h = output_res_blocks[i][j].forward(ctx, h, emb); - - if (ds == attention_resolutions[0] || ds == attention_resolutions[1] || ds == attention_resolutions[2]) { - h = output_transformers[i][j].forward(ctx, h, context); - } - - if (i > 0 && j == num_res_blocks) { - h = output_up_samples[i - 1].forward(ctx, h); - - ds /= 2; - } - } - } - - // out - // group norm 32 - h = ggml_group_norm_32(ctx, h); - h = ggml_add(ctx, - ggml_mul(ctx, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, out_0_w, 1, 1, out_0_w->ne[0], 1), - h), - h), - ggml_repeat(ctx, - ggml_reshape_4d(ctx, out_0_b, 1, 1, out_0_b->ne[0], 1), - h)); - // silu - h = ggml_silu_inplace(ctx, h); - // conv2d - h = ggml_conv_2d(ctx, out_2_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, out_2_b, 1, 1, out_2_b->ne[0], 1), - h)); // [N, out_channels, h, w] - - return h; - } -}; - -/*================================================== AutoEncoderKL ===================================================*/ - -struct ResnetBlock { - // network hparams - int in_channels; - int out_channels; - - // network params - struct ggml_tensor* norm1_w; // [in_channels, ] - struct ggml_tensor* norm1_b; // [in_channels, ] - - struct ggml_tensor* conv1_w; // [out_channels, in_channels, 3, 3] - struct ggml_tensor* conv1_b; // [out_channels, ] - - struct ggml_tensor* norm2_w; // [out_channels, ] - struct ggml_tensor* norm2_b; // [out_channels, ] - - struct ggml_tensor* conv2_w; // [out_channels, out_channels, 3, 3] - struct ggml_tensor* conv2_b; // [out_channels, ] - - // nin_shortcut, only if out_channels != in_channels - struct ggml_tensor* nin_shortcut_w; // [out_channels, in_channels, 1, 1] - struct ggml_tensor* nin_shortcut_b; // [out_channels, ] - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += 2 * in_channels * ggml_type_sizef(GGML_TYPE_F32); // norm1_w/b - mem_size += out_channels * in_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // conv1_w - mem_size += 4 * out_channels * ggml_type_sizef(GGML_TYPE_F32); // conv1_b/norm2_w/norm2_b/conv2_b - mem_size += out_channels * out_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // conv2_w - - mem_size += 8 * ggml_tensor_overhead(); // object overhead - - if (out_channels != in_channels) { - mem_size += out_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16); // nin_shortcut_w - mem_size += out_channels * ggml_type_sizef(GGML_TYPE_F32); // nin_shortcut_b - - mem_size += 2 * ggml_tensor_overhead(); // object overhead - } - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - norm1_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - norm1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - conv1_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, in_channels, out_channels); - conv1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - - norm2_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - norm2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - conv2_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, out_channels, out_channels); - conv2_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - - if (out_channels != in_channels) { - nin_shortcut_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, out_channels); - nin_shortcut_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_channels); - } - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "norm1.weight"] = norm1_w; - tensors[prefix + "norm1.bias"] = norm1_b; - tensors[prefix + "conv1.weight"] = conv1_w; - tensors[prefix + "conv1.bias"] = conv1_b; - - tensors[prefix + "norm2.weight"] = norm2_w; - tensors[prefix + "norm2.bias"] = norm2_b; - tensors[prefix + "conv2.weight"] = conv2_w; - tensors[prefix + "conv2.bias"] = conv2_b; - - if (out_channels != in_channels) { - tensors[prefix + "nin_shortcut.weight"] = nin_shortcut_w; - tensors[prefix + "nin_shortcut.bias"] = nin_shortcut_b; - } - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* z) { - // z: [N, in_channels, h, w] - - // group norm 32 - auto h = ggml_group_norm_32(ctx, z); - h = ggml_mul(ctx, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, norm1_w, 1, 1, norm1_w->ne[0], 1), - h), - h); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, norm1_b, 1, 1, norm1_b->ne[0], 1), - h)); - // silu - h = ggml_silu_inplace(ctx, h); - // conv2d - h = ggml_conv_2d(ctx, conv1_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv1_b, 1, 1, conv1_b->ne[0], 1), - h)); // [N, out_channels, h, w] - - // group norm 32 - h = ggml_group_norm_32(ctx, h); - h = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm2_w, 1, 1, norm2_w->ne[0], 1), h), h), - ggml_repeat(ctx, ggml_reshape_4d(ctx, norm2_b, 1, 1, norm2_b->ne[0], 1), h)); - // silu - h = ggml_silu_inplace(ctx, h); - // dropout, skip for inference - // conv2d - h = ggml_conv_2d(ctx, conv2_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv2_b, 1, 1, conv2_b->ne[0], 1), - h)); // [N, out_channels, h, w - - // skip connection - if (out_channels != in_channels) { - z = ggml_conv_2d(ctx, nin_shortcut_w, z, 1, 1, 0, 0, 1, 1); - z = ggml_add(ctx, - z, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, nin_shortcut_b, 1, 1, nin_shortcut_b->ne[0], 1), - z)); // [N, out_channels, h, w] - } - h = ggml_add(ctx, h, z); - return h; // [N, out_channels, h, w] - } -}; - -struct AttnBlock { - int in_channels; // mult * model_channels - - // group norm - struct ggml_tensor* norm_w; // [in_channels,] - struct ggml_tensor* norm_b; // [in_channels,] - - // q/k/v - struct ggml_tensor* q_w; // [in_channels, in_channels, 1, 1] - struct ggml_tensor* q_b; // [in_channels,] - struct ggml_tensor* k_w; // [in_channels, in_channels, 1, 1] - struct ggml_tensor* k_b; // [in_channels,] - struct ggml_tensor* v_w; // [in_channels, in_channels, 1, 1] - struct ggml_tensor* v_b; // [in_channels,] - - // proj_out - struct ggml_tensor* proj_out_w; // [in_channels, in_channels, 1, 1] - struct ggml_tensor* proj_out_b; // [in_channels,] - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - mem_size += 6 * in_channels * ggml_type_sizef(GGML_TYPE_F32); // norm_w/norm_b/q_b/k_v/v_b/proj_out_b - mem_size += 4 * in_channels * in_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16); // q_w/k_w/v_w/proj_out_w - mem_size += 10 * ggml_tensor_overhead(); // object overhead - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - norm_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - norm_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - q_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels); - q_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - k_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels); - k_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - v_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels); - v_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - - proj_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, in_channels, in_channels); - proj_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "norm.weight"] = norm_w; - tensors[prefix + "norm.bias"] = norm_b; - tensors[prefix + "q.weight"] = q_w; - tensors[prefix + "q.bias"] = q_b; - tensors[prefix + "k.weight"] = k_w; - tensors[prefix + "k.bias"] = k_b; - tensors[prefix + "v.weight"] = v_w; - tensors[prefix + "v.bias"] = v_b; - tensors[prefix + "proj_out.weight"] = proj_out_w; - tensors[prefix + "proj_out.bias"] = proj_out_b; - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) { - // x: [N, in_channels, h, w] - - // group norm 32 - auto h_ = ggml_group_norm_32(ctx, x); - h_ = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_w, 1, 1, norm_w->ne[0], 1), h_), h_), - ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_b, 1, 1, norm_b->ne[0], 1), h_)); - - const int64_t n = h_->ne[3]; - const int64_t c = h_->ne[2]; - const int64_t h = h_->ne[1]; - const int64_t w = h_->ne[0]; - // q - auto q = ggml_conv_2d(ctx, q_w, h_, 1, 1, 0, 0, 1, 1); - q = ggml_add(ctx, - q, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, q_b, 1, 1, q_b->ne[0], 1), - q)); // [N, in_channels, h, w] - - // k - auto k = ggml_conv_2d(ctx, k_w, h_, 1, 1, 0, 0, 1, 1); - k = ggml_add(ctx, - k, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, k_b, 1, 1, k_b->ne[0], 1), - k)); // [N, in_channels, h, w] - - // v - auto v = ggml_conv_2d(ctx, v_w, h_, 1, 1, 0, 0, 1, 1); - v = ggml_add(ctx, - v, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, v_b, 1, 1, v_b->ne[0], 1), - v)); // [N, in_channels, h, w] - - q = ggml_cont(ctx, ggml_permute(ctx, q, 1, 2, 0, 3)); // [N, h, w, in_channels] - q = ggml_reshape_3d(ctx, q, c, h * w, n); // [N, h * w, in_channels] - - k = ggml_cont(ctx, ggml_permute(ctx, k, 1, 2, 0, 3)); // [N, h, w, in_channels] - k = ggml_reshape_3d(ctx, k, c, h * w, n); // [N, h * w, in_channels] - - auto w_ = ggml_mul_mat(ctx, k, q); // [N, h * w, h * w] - w_ = ggml_scale_inplace(ctx, w_, ggml_new_f32(ctx, 1.0f / sqrt((float)c))); - w_ = ggml_soft_max_inplace(ctx, w_); - - v = ggml_reshape_3d(ctx, v, h * w, c, n); // [N, in_channels, h * w] - h_ = ggml_mul_mat(ctx, v, w_); // [N, h * w, in_channels] - h_ = ggml_cont(ctx, ggml_permute(ctx, h_, 1, 0, 2, 3)); // [N, in_channels, h * w] - h_ = ggml_reshape_4d(ctx, h_, w, h, c, n); // [N, in_channels, h, w] - - // proj_out - h_ = ggml_conv_2d(ctx, proj_out_w, h_, 1, 1, 0, 0, 1, 1); - h_ = ggml_add(ctx, - h_, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, proj_out_b, 1, 1, proj_out_b->ne[0], 1), - h_)); // [N, in_channels, h, w] - h_ = ggml_add(ctx, h_, x); - return h_; - } -}; - -// ldm.modules.diffusionmodules.model.Encoder -struct Encoder { - int embed_dim = 4; - int ch = 128; - int z_channels = 4; - int in_channels = 3; - int num_res_blocks = 2; - int ch_mult[4] = {1, 2, 4, 4}; - - struct ggml_tensor* conv_in_w; // [ch, in_channels, 3, 3] - struct ggml_tensor* conv_in_b; // [ch, ] - - ResnetBlock down_blocks[4][2]; - DownSample down_samples[3]; - - struct - { - ResnetBlock block_1; - AttnBlock attn_1; - ResnetBlock block_2; - } mid; - - // block_in = ch * ch_mult[len_mults - 1] - struct ggml_tensor* norm_out_w; // [block_in, ] - struct ggml_tensor* norm_out_b; // [block_in, ] - - struct ggml_tensor* conv_out_w; // [embed_dim*2, block_in, 3, 3] - struct ggml_tensor* conv_out_b; // [embed_dim*2, ] - - Encoder() { - int len_mults = sizeof(ch_mult) / sizeof(int); - - int block_in = 1; - for (int i = 0; i < len_mults; i++) { - if (i == 0) { - block_in = ch; - } else { - block_in = ch * ch_mult[i - 1]; - } - int block_out = ch * ch_mult[i]; - for (int j = 0; j < num_res_blocks; j++) { - down_blocks[i][j].in_channels = block_in; - down_blocks[i][j].out_channels = block_out; - block_in = block_out; - } - if (i != len_mults - 1) { - down_samples[i].channels = block_in; - down_samples[i].out_channels = block_in; - down_samples[i].vae_downsample = true; - } - } - - mid.block_1.in_channels = block_in; - mid.block_1.out_channels = block_in; - mid.attn_1.in_channels = block_in; - mid.block_2.in_channels = block_in; - mid.block_2.out_channels = block_in; - } - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - int len_mults = sizeof(ch_mult) / sizeof(int); - int block_in = ch * ch_mult[len_mults - 1]; - - mem_size += ch * in_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // conv_in_w - mem_size += ch * ggml_type_sizef(GGML_TYPE_F32); // conv_in_b - - mem_size += 2 * block_in * ggml_type_sizef(GGML_TYPE_F32); // norm_out_w/b - - mem_size += z_channels * 2 * block_in * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // conv_out_w - mem_size += z_channels * 2 * ggml_type_sizef(GGML_TYPE_F32); // conv_out_b - - mem_size += 6 * ggml_tensor_overhead(); // object overhead - - mem_size += mid.block_1.compute_params_mem_size(wtype); - mem_size += mid.attn_1.compute_params_mem_size(wtype); - mem_size += mid.block_2.compute_params_mem_size(wtype); - - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - mem_size += down_blocks[i][j].compute_params_mem_size(wtype); - } - if (i != 0) { - mem_size += down_samples[i - 1].compute_params_mem_size(wtype); - } - } - - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - int len_mults = sizeof(ch_mult) / sizeof(int); - int block_in = ch * ch_mult[len_mults - 1]; - - conv_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, in_channels, ch); - conv_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ch); - - norm_out_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, block_in); - norm_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, block_in); - - conv_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, block_in, z_channels * 2); - conv_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_channels * 2); - - mid.block_1.init_params(ctx, wtype); - mid.attn_1.init_params(ctx, wtype); - mid.block_2.init_params(ctx, wtype); - - for (int i = 0; i < len_mults; i++) { - for (int j = 0; j < num_res_blocks; j++) { - down_blocks[i][j].init_params(ctx, wtype); - } - if (i != len_mults - 1) { - down_samples[i].init_params(ctx, wtype); - } - } - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "norm_out.weight"] = norm_out_w; - tensors[prefix + "norm_out.bias"] = norm_out_b; - tensors[prefix + "conv_in.weight"] = conv_in_w; - tensors[prefix + "conv_in.bias"] = conv_in_b; - tensors[prefix + "conv_out.weight"] = conv_out_w; - tensors[prefix + "conv_out.bias"] = conv_out_b; - - mid.block_1.map_by_name(tensors, prefix + "mid.block_1."); - mid.attn_1.map_by_name(tensors, prefix + "mid.attn_1."); - mid.block_2.map_by_name(tensors, prefix + "mid.block_2."); - - int len_mults = sizeof(ch_mult) / sizeof(int); - for (int i = 0; i < len_mults; i++) { - for (int j = 0; j < num_res_blocks; j++) { - down_blocks[i][j].map_by_name(tensors, prefix + "down." + std::to_string(i) + ".block." + std::to_string(j) + "."); - } - if (i != len_mults - 1) { - down_samples[i].map_by_name(tensors, prefix + "down." + std::to_string(i) + ".downsample."); - } - } - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* x) { - // x: [N, in_channels, h, w] - - // conv_in - auto h = ggml_conv_2d(ctx, conv_in_w, x, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv_in_b, 1, 1, conv_in_b->ne[0], 1), - h)); // [N, ch, h, w] - int len_mults = sizeof(ch_mult) / sizeof(int); - for (int i = 0; i < len_mults; i++) { - for (int j = 0; j < num_res_blocks; j++) { - h = down_blocks[i][j].forward(ctx, h); - } - if (i != len_mults - 1) { - h = down_samples[i].forward(ctx, h); - } - } - - h = mid.block_1.forward(ctx, h); - h = mid.attn_1.forward(ctx, h); - h = mid.block_2.forward(ctx, h); // [N, block_in, h, w] - - // group norm 32 - h = ggml_group_norm_32(ctx, h); - h = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_w, 1, 1, norm_out_w->ne[0], 1), h), h), - ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_b, 1, 1, norm_out_b->ne[0], 1), h)); - - // silu - // silu - h = ggml_silu_inplace(ctx, h); - - // conv_out - h = ggml_conv_2d(ctx, conv_out_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv_out_b, 1, 1, conv_out_b->ne[0], 1), - h)); // [N, z_channels*2, h, w] - - return h; - } -}; - -// ldm.modules.diffusionmodules.model.Decoder -struct Decoder { - int embed_dim = 4; - int ch = 128; - int z_channels = 4; - int out_ch = 3; - int num_res_blocks = 2; - int ch_mult[4] = {1, 2, 4, 4}; - - // block_in = ch * ch_mult[-1], 512 - struct ggml_tensor* conv_in_w; // [block_in, z_channels, 3, 3] - struct ggml_tensor* conv_in_b; // [block_in, ] - - struct - { - ResnetBlock block_1; - AttnBlock attn_1; - ResnetBlock block_2; - } mid; - - ResnetBlock up_blocks[4][3]; - UpSample up_samples[3]; - - struct ggml_tensor* norm_out_w; // [ch * ch_mult[0], ] - struct ggml_tensor* norm_out_b; // [ch * ch_mult[0], ] - - struct ggml_tensor* conv_out_w; // [out_ch, ch * ch_mult[0], 3, 3] - struct ggml_tensor* conv_out_b; // [out_ch, ] - - Decoder() { - int len_mults = sizeof(ch_mult) / sizeof(int); - int block_in = ch * ch_mult[len_mults - 1]; - - mid.block_1.in_channels = block_in; - mid.block_1.out_channels = block_in; - mid.attn_1.in_channels = block_in; - mid.block_2.in_channels = block_in; - mid.block_2.out_channels = block_in; - - for (int i = len_mults - 1; i >= 0; i--) { - int mult = ch_mult[i]; - int block_out = ch * mult; - for (int j = 0; j < num_res_blocks + 1; j++) { - up_blocks[i][j].in_channels = block_in; - up_blocks[i][j].out_channels = block_out; - block_in = block_out; - } - if (i != 0) { - up_samples[i - 1].channels = block_in; - up_samples[i - 1].out_channels = block_in; - } - } - } - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - int len_mults = sizeof(ch_mult) / sizeof(int); - int block_in = ch * ch_mult[len_mults - 1]; - - mem_size += block_in * z_channels * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // conv_in_w - mem_size += block_in * ggml_type_sizef(GGML_TYPE_F32); // conv_in_b - - mem_size += 2 * (ch * ch_mult[0]) * ggml_type_sizef(GGML_TYPE_F32); // norm_out_w/b - - mem_size += (ch * ch_mult[0]) * out_ch * 3 * 3 * ggml_type_sizef(GGML_TYPE_F16); // conv_out_w - mem_size += out_ch * ggml_type_sizef(GGML_TYPE_F32); // conv_out_b - - mem_size += 8 * ggml_tensor_overhead(); // object overhead - - mem_size += mid.block_1.compute_params_mem_size(wtype); - mem_size += mid.attn_1.compute_params_mem_size(wtype); - mem_size += mid.block_2.compute_params_mem_size(wtype); - - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - mem_size += up_blocks[i][j].compute_params_mem_size(wtype); - } - if (i != 0) { - mem_size += up_samples[i - 1].compute_params_mem_size(wtype); - } - } - - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - int len_mults = sizeof(ch_mult) / sizeof(int); - int block_in = ch * ch_mult[len_mults - 1]; - - norm_out_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ch * ch_mult[0]); - norm_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ch * ch_mult[0]); - - conv_in_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, z_channels, block_in); - conv_in_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, block_in); - - conv_out_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 3, 3, ch * ch_mult[0], out_ch); - conv_out_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_ch); - - mid.block_1.init_params(ctx, wtype); - mid.attn_1.init_params(ctx, wtype); - mid.block_2.init_params(ctx, wtype); - - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - up_blocks[i][j].init_params(ctx, wtype); - } - if (i != 0) { - up_samples[i - 1].init_params(ctx, wtype); - } - } - } - - void map_by_name(std::map& tensors, const std::string prefix) { - tensors[prefix + "norm_out.weight"] = norm_out_w; - tensors[prefix + "norm_out.bias"] = norm_out_b; - tensors[prefix + "conv_in.weight"] = conv_in_w; - tensors[prefix + "conv_in.bias"] = conv_in_b; - tensors[prefix + "conv_out.weight"] = conv_out_w; - tensors[prefix + "conv_out.bias"] = conv_out_b; - - mid.block_1.map_by_name(tensors, prefix + "mid.block_1."); - mid.attn_1.map_by_name(tensors, prefix + "mid.attn_1."); - mid.block_2.map_by_name(tensors, prefix + "mid.block_2."); - - int len_mults = sizeof(ch_mult) / sizeof(int); - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - up_blocks[i][j].map_by_name(tensors, prefix + "up." + std::to_string(i) + ".block." + std::to_string(j) + "."); - } - if (i != 0) { - up_samples[i - 1].map_by_name(tensors, prefix + "up." + std::to_string(i) + ".upsample."); - } - } - } - - struct ggml_tensor* forward(struct ggml_context* ctx, struct ggml_tensor* z) { - // z: [N, z_channels, h, w] - - // conv_in - auto h = ggml_conv_2d(ctx, conv_in_w, z, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv_in_b, 1, 1, conv_in_b->ne[0], 1), - h)); // [N, block_in, h, w] - - h = mid.block_1.forward(ctx, h); - h = mid.attn_1.forward(ctx, h); - h = mid.block_2.forward(ctx, h); // [N, block_in, h, w] - - int len_mults = sizeof(ch_mult) / sizeof(int); - for (int i = len_mults - 1; i >= 0; i--) { - for (int j = 0; j < num_res_blocks + 1; j++) { - h = up_blocks[i][j].forward(ctx, h); - } - if (i != 0) { - h = up_samples[i - 1].forward(ctx, h); - } - } - - // group norm 32 - h = ggml_group_norm_32(ctx, h); - h = ggml_add(ctx, - ggml_mul(ctx, ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_w, 1, 1, norm_out_w->ne[0], 1), h), h), - ggml_repeat(ctx, ggml_reshape_4d(ctx, norm_out_b, 1, 1, norm_out_b->ne[0], 1), h)); - - // silu - // silu - h = ggml_silu_inplace(ctx, h); - - // conv_out - h = ggml_conv_2d(ctx, conv_out_w, h, 1, 1, 1, 1, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, conv_out_b, 1, 1, conv_out_b->ne[0], 1), - h)); // [N, out_ch, h, w] - - return h; - } -}; - -// ldm.models.autoencoder.AutoencoderKL -struct AutoEncoderKL { - bool decode_only = true; - int embed_dim = 4; - struct - { - int z_channels = 4; - int resolution = 256; - int in_channels = 3; - int out_ch = 3; - int ch = 128; - int ch_mult[4] = {1, 2, 4, 4}; - int num_res_blocks = 2; - } dd_config; - - struct ggml_tensor* quant_conv_w; // [2*embed_dim, 2*z_channels, 1, 1] - struct ggml_tensor* quant_conv_b; // [2*embed_dim, ] - - struct ggml_tensor* post_quant_conv_w; // [z_channels, embed_dim, 1, 1] - struct ggml_tensor* post_quant_conv_b; // [z_channels, ] - - Encoder encoder; - Decoder decoder; - - AutoEncoderKL(bool decode_only = false) - : decode_only(decode_only) { - assert(sizeof(dd_config.ch_mult) == sizeof(encoder.ch_mult)); - assert(sizeof(dd_config.ch_mult) == sizeof(decoder.ch_mult)); - - encoder.embed_dim = embed_dim; - decoder.embed_dim = embed_dim; - encoder.ch = dd_config.ch; - decoder.ch = dd_config.ch; - encoder.z_channels = dd_config.z_channels; - decoder.z_channels = dd_config.z_channels; - encoder.in_channels = dd_config.in_channels; - decoder.out_ch = dd_config.out_ch; - encoder.num_res_blocks = dd_config.num_res_blocks; - - int len_mults = sizeof(dd_config.ch_mult) / sizeof(int); - for (int i = 0; i < len_mults; i++) { - encoder.ch_mult[i] = dd_config.ch_mult[i]; - decoder.ch_mult[i] = dd_config.ch_mult[i]; - } - } - - size_t compute_params_mem_size(ggml_type wtype) { - double mem_size = 0; - - if (!decode_only) { - mem_size += 2 * embed_dim * 2 * dd_config.z_channels * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16); // quant_conv_w - mem_size += 2 * embed_dim * ggml_type_sizef(GGML_TYPE_F32); // quant_conv_b - mem_size += encoder.compute_params_mem_size(wtype); - } - - mem_size += dd_config.z_channels * embed_dim * 1 * 1 * ggml_type_sizef(GGML_TYPE_F16); // post_quant_conv_w - mem_size += dd_config.z_channels * ggml_type_sizef(GGML_TYPE_F32); // post_quant_conv_b - - mem_size += decoder.compute_params_mem_size(wtype); - return static_cast(mem_size); - } - - void init_params(struct ggml_context* ctx, ggml_type wtype) { - if (!decode_only) { - quant_conv_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, 2 * dd_config.z_channels, 2 * embed_dim); - quant_conv_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2 * embed_dim); - encoder.init_params(ctx, wtype); - } - - post_quant_conv_w = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, embed_dim, dd_config.z_channels); - post_quant_conv_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dd_config.z_channels); - decoder.init_params(ctx, wtype); - } - - void map_by_name(std::map& tensors, const std::string prefix) { - if (!decode_only) { - tensors[prefix + "quant_conv.weight"] = quant_conv_w; - tensors[prefix + "quant_conv.bias"] = quant_conv_b; - encoder.map_by_name(tensors, prefix + "encoder."); - } - - tensors[prefix + "post_quant_conv.weight"] = post_quant_conv_w; - tensors[prefix + "post_quant_conv.bias"] = post_quant_conv_b; - decoder.map_by_name(tensors, prefix + "decoder."); - } - - struct ggml_tensor* decode(struct ggml_context* ctx, struct ggml_tensor* z) { - // z: [N, z_channels, h, w] - - // post_quant_conv - auto h = ggml_conv_2d(ctx, post_quant_conv_w, z, 1, 1, 0, 0, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, post_quant_conv_b, 1, 1, post_quant_conv_b->ne[0], 1), - h)); // [N, z_channels, h, w] - h = decoder.forward(ctx, h); - return h; - } - - struct ggml_tensor* encode(struct ggml_context* ctx, struct ggml_tensor* x) { - // x: [N, in_channels, h, w] - auto h = encoder.forward(ctx, x); // [N, 2*z_channels, h/8, w/8] - // quant_conv - h = ggml_conv_2d(ctx, quant_conv_w, h, 1, 1, 0, 0, 1, 1); - h = ggml_add(ctx, - h, - ggml_repeat(ctx, - ggml_reshape_4d(ctx, quant_conv_b, 1, 1, quant_conv_b->ne[0], 1), - h)); // [N, 2*embed_dim, h/8, w/8] - return h; - } -}; - -/*================================================= CompVisDenoiser ==================================================*/ - -// Ref: https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/external.py - -struct SigmaSchedule { - float alphas_cumprod[TIMESTEPS]; - float sigmas[TIMESTEPS]; - float log_sigmas[TIMESTEPS]; - - virtual std::vector get_sigmas(uint32_t n) = 0; - - float sigma_to_t(float sigma) { - float log_sigma = std::log(sigma); - std::vector dists; - dists.reserve(TIMESTEPS); - for (float log_sigma_val : log_sigmas) { - dists.push_back(log_sigma - log_sigma_val); - } - - int low_idx = 0; - for (size_t i = 0; i < TIMESTEPS; i++) { - if (dists[i] >= 0) { - low_idx++; - } - } - low_idx = std::min(std::max(low_idx - 1, 0), TIMESTEPS - 2); - int high_idx = low_idx + 1; - - float low = log_sigmas[low_idx]; - float high = log_sigmas[high_idx]; - float w = (low - log_sigma) / (low - high); - w = std::max(0.f, std::min(1.f, w)); - float t = (1.0f - w) * low_idx + w * high_idx; - - return t; - } - - float t_to_sigma(float t) { - int low_idx = static_cast(std::floor(t)); - int high_idx = static_cast(std::ceil(t)); - float w = t - static_cast(low_idx); - float log_sigma = (1.0f - w) * log_sigmas[low_idx] + w * log_sigmas[high_idx]; - return std::exp(log_sigma); - } -}; - -struct DiscreteSchedule : SigmaSchedule { - std::vector get_sigmas(uint32_t n) { - std::vector result; - - int t_max = TIMESTEPS - 1; - - if (n == 0) { - return result; - } else if (n == 1) { - result.push_back(t_to_sigma(t_max)); - result.push_back(0); - return result; - } - - float step = static_cast(t_max) / static_cast(n - 1); - for (int i = 0; i < n; ++i) { - float t = t_max - step * i; - result.push_back(t_to_sigma(t)); - } - result.push_back(0); - return result; - } -}; - -struct KarrasSchedule : SigmaSchedule { - std::vector get_sigmas(uint32_t n) { - // These *COULD* be function arguments here, - // but does anybody ever bother to touch them? - float sigma_min = 0.1; - float sigma_max = 10.; - float rho = 7.; - - std::vector result(n + 1); - - float min_inv_rho = pow(sigma_min, (1. / rho)); - float max_inv_rho = pow(sigma_max, (1. / rho)); - for (int i = 0; i < n; i++) { - // Eq. (5) from Karras et al 2022 - result[i] = pow(max_inv_rho + (float)i / ((float)n - 1.) * (min_inv_rho - max_inv_rho), rho); - } - result[n] = 0.; - return result; - } -}; - -struct Denoiser { - std::shared_ptr schedule = std::make_shared(); - virtual std::vector get_scalings(float sigma) = 0; -}; - -struct CompVisDenoiser : public Denoiser { - float sigma_data = 1.0f; - - std::vector get_scalings(float sigma) { - float c_out = -sigma; - float c_in = 1.0f / std::sqrt(sigma * sigma + sigma_data * sigma_data); - return {c_out, c_in}; - } -}; - -struct CompVisVDenoiser : public Denoiser { - float sigma_data = 1.0f; - - std::vector get_scalings(float sigma) { - float c_skip = sigma_data * sigma_data / (sigma * sigma + sigma_data * sigma_data); - float c_out = -sigma * sigma_data / std::sqrt(sigma * sigma + sigma_data * sigma_data); - float c_in = 1.0f / std::sqrt(sigma * sigma + sigma_data * sigma_data); - return {c_skip, c_out, c_in}; - } -}; - -/*=============================================== StableDiffusionGGML ================================================*/ - -class StableDiffusionGGML { - public: - ggml_context* clip_params_ctx = NULL; - ggml_context* unet_params_ctx = NULL; - ggml_context* vae_params_ctx = NULL; - - bool dynamic = true; - bool vae_decode_only = false; - bool free_params_immediately = false; - - std::shared_ptr rng = std::make_shared(); - int32_t ftype = 1; - int n_threads = -1; - float scale_factor = 0.18215f; - size_t max_mem_size = 0; - size_t curr_params_mem_size = 0; - size_t max_params_mem_size = 0; - size_t max_rt_mem_size = 0; - - FrozenCLIPEmbedderWithCustomWords cond_stage_model; - UNetModel diffusion_model; - AutoEncoderKL first_stage_model; - - std::shared_ptr denoiser = std::make_shared(); - - StableDiffusionGGML() = default; - - StableDiffusionGGML(int n_threads, - bool vae_decode_only, - bool free_params_immediately, - RNGType rng_type) - : n_threads(n_threads), - vae_decode_only(vae_decode_only), - free_params_immediately(free_params_immediately) { - first_stage_model.decode_only = vae_decode_only; - if (rng_type == STD_DEFAULT_RNG) { - rng = std::make_shared(); - } else if (rng_type == CUDA_RNG) { - rng = std::make_shared(); - } - } - - ~StableDiffusionGGML() { - if (clip_params_ctx != NULL) { - ggml_free(clip_params_ctx); - clip_params_ctx = NULL; - } - if (unet_params_ctx != NULL) { - ggml_free(unet_params_ctx); - unet_params_ctx = NULL; - } - if (vae_params_ctx != NULL) { - ggml_free(vae_params_ctx); - vae_params_ctx = NULL; - } - } - - bool load_from_file(const std::string& file_path, Schedule schedule) { - LOG_INFO("loading model from '%s'", file_path.c_str()); - - std::ifstream file(file_path, std::ios::binary); - if (!file.is_open()) { - LOG_ERROR("failed to open '%s'", file_path.c_str()); - return false; - } - - LOG_DEBUG("verifying magic"); - // verify magic - { - uint32_t magic; - file.read(reinterpret_cast(&magic), sizeof(magic)); - if (magic != GGML_FILE_MAGIC) { - LOG_ERROR("invalid model file '%s' (bad magic)", file_path.c_str()); - return false; - } - } - - LOG_DEBUG("loading hparams"); - // load hparams - file.read(reinterpret_cast(&ftype), sizeof(ftype)); - - int model_type = (ftype >> 16) & 0xFFFF; - if (model_type >= MODEL_TYPE_COUNT) { - LOG_ERROR("invalid model file '%s' (bad model type value %d)", file_path.c_str(), ftype); - return false; - } - LOG_INFO("model type: %s", model_type_to_str[model_type]); - - if (model_type == SD2) { - cond_stage_model = FrozenCLIPEmbedderWithCustomWords((ModelType)model_type); - diffusion_model = UNetModel((ModelType)model_type); - } - - ggml_type wtype = ggml_ftype_to_ggml_type((ggml_ftype)(ftype & 0xFFFF)); - LOG_INFO("ftype: %s", ggml_type_name(wtype)); - if (wtype == GGML_TYPE_COUNT) { - LOG_ERROR("invalid model file '%s' (bad ftype value %d)", file_path.c_str(), ftype); - return false; - } - - LOG_DEBUG("loading vocab"); - // load vocab - { - int32_t n_vocab = 0; - file.read(reinterpret_cast(&n_vocab), sizeof(n_vocab)); - - if (n_vocab != cond_stage_model.text_model.vocab_size) { - LOG_ERROR("invalid model file '%s' (bad vocab size %d != %d)", - file_path.c_str(), n_vocab, cond_stage_model.text_model.vocab_size); - return false; - } - - std::string word; - std::vector buf(128); - - for (int i = 0; i < n_vocab; i++) { - uint32_t len; - file.read((char*)&len, sizeof(len)); - - buf.resize(len); - file.read((char*)buf.data(), len); - word.assign(buf.data(), len); - - cond_stage_model.tokenizer.add_token(word, i); - } - } - - // create the ggml context for network params - LOG_DEBUG("ggml tensor size = %d bytes", (int)sizeof(ggml_tensor)); - { - // cond_stage_model(FrozenCLIPEmbedder) - double ctx_size = 1 * 1024 * 1024; // 1 MB, for padding - ctx_size += cond_stage_model.text_model.compute_params_mem_size(wtype); - LOG_DEBUG("clip params ctx size = % 6.2f MB", ctx_size / (1024.0 * 1024.0)); - - struct ggml_init_params params; - params.mem_size = static_cast(ctx_size); - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = false; - - clip_params_ctx = ggml_init(params); - if (!clip_params_ctx) { - LOG_ERROR("ggml_init() failed"); - return false; - } - } - - { - // diffusion_model(UNetModel) - double ctx_size = 1 * 1024 * 1024; // 1 MB, for padding - ctx_size += diffusion_model.compute_params_mem_size(wtype); - LOG_DEBUG("unet params ctx size = % 6.2f MB", ctx_size / (1024.0 * 1024.0)); - - struct ggml_init_params params; - params.mem_size = static_cast(ctx_size); - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = false; - - unet_params_ctx = ggml_init(params); - if (!unet_params_ctx) { - LOG_ERROR("ggml_init() failed"); - ggml_free(clip_params_ctx); - clip_params_ctx = NULL; - return false; - } - } - - { - // first_stage_model(AutoEncoderKL) - double ctx_size = 1 * 1024 * 1024; // 1 MB, for padding - ctx_size += first_stage_model.compute_params_mem_size(wtype); - LOG_DEBUG("vae params ctx size = % 6.2f MB", ctx_size / (1024.0 * 1024.0)); - - struct ggml_init_params params; - params.mem_size = static_cast(ctx_size); - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = false; - - vae_params_ctx = ggml_init(params); - if (!vae_params_ctx) { - LOG_ERROR("ggml_init() failed"); - ggml_free(clip_params_ctx); - clip_params_ctx = NULL; - ggml_free(unet_params_ctx); - unet_params_ctx = NULL; - return false; - } - } - - std::map tensors; - - LOG_DEBUG("preparing memory for the weights"); - // prepare memory for the weights - { - // cond_stage_model(FrozenCLIPEmbedder) - cond_stage_model.text_model.init_params(clip_params_ctx, wtype); - cond_stage_model.text_model.map_by_name(tensors, "cond_stage_model.transformer.text_model."); - - // diffusion_model(UNetModel) - diffusion_model.init_params(unet_params_ctx, wtype); - diffusion_model.map_by_name(tensors, "model.diffusion_model."); - - // firest_stage_model(AutoEncoderKL) - first_stage_model.init_params(vae_params_ctx, wtype); - first_stage_model.map_by_name(tensors, "first_stage_model."); - } - - LOG_DEBUG("loading weights"); - std::set tensor_names_in_file; - int64_t t0 = ggml_time_ms(); - // load weights - float alphas_cumprod[TIMESTEPS]; - { - int n_tensors = 0; - size_t total_size = 0; - - while (true) { - int32_t n_dims; - int32_t length; - int32_t ttype; - - file.read(reinterpret_cast(&n_dims), sizeof(n_dims)); - file.read(reinterpret_cast(&length), sizeof(length)); - file.read(reinterpret_cast(&ttype), sizeof(ttype)); - - if (file.eof()) { - break; - } - - int32_t nelements = 1; - int32_t ne[4] = {1, 1, 1, 1}; - for (int i = 0; i < n_dims; ++i) { - file.read(reinterpret_cast(&ne[i]), sizeof(ne[i])); - nelements *= ne[i]; - } - - const size_t num_bytes = nelements / ggml_blck_size(ggml_type(ttype)) * ggml_type_size(ggml_type(ttype)); - - std::string name(length, 0); - file.read(&name[0], length); - - tensor_names_in_file.insert(std::string(name.data())); - - if (std::string(name.data()) == "alphas_cumprod") { - file.read(reinterpret_cast(alphas_cumprod), nelements * ggml_type_size((ggml_type)ttype)); - continue; - } - - struct ggml_tensor* tensor; - if (tensors.find(name.data()) != tensors.end()) { - tensor = tensors[name.data()]; - } else { - if (name.find("quant") == std::string::npos && name.find("first_stage_model.encoder.") == std::string::npos) { - LOG_WARN("unknown tensor '%s' in model file", name.data()); - } else { - if (!vae_decode_only) { - LOG_WARN("unknown tensor '%s' in model file", name.data()); - return false; - } - } - file.ignore(num_bytes); - continue; - } - - if (tensor->ne[0] != ne[0] || tensor->ne[1] != ne[1] || tensor->ne[2] != ne[2] || tensor->ne[3] != ne[3]) { - LOG_ERROR( - "tensor '%s' has wrong shape in model file: " - "got [%d, %d, %d, %d], expected [%d, %d, %d, %d]", - name.data(), - ne[0], ne[1], ne[2], ne[3], - (int)tensor->ne[0], (int)tensor->ne[1], (int)tensor->ne[2], (int)tensor->ne[3]); - return false; - } - - if (ggml_nelements(tensor) != nelements) { - LOG_ERROR( - "tensor '%s' has wrong number of elements in model file: " - "got %u, expert %zu", - name.data(), nelements, ggml_nelements(tensor)); - return false; - } - - if (tensor->type != ttype) { - LOG_ERROR("tensor '%s' has wrong type in model file: got %s, expect %s", - name.data(), ggml_type_name(ggml_type(ttype)), ggml_type_name(tensor->type)); - return false; - } - - file.read(reinterpret_cast(tensor->data), num_bytes); - - total_size += ggml_nbytes(tensor); - } - bool some_tensor_not_init = false; - for (auto pair : tensors) { - if (pair.first.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos) { - continue; - } - if (tensor_names_in_file.find(pair.first) == tensor_names_in_file.end()) { - LOG_ERROR("tensor '%s' not in model file", pair.first.c_str()); - some_tensor_not_init = true; - } - } - if (tensor_names_in_file.find("alphas_cumprod") == tensor_names_in_file.end()) { - LOG_ERROR("tensor alphas_cumprod not in model file"); - some_tensor_not_init = true; - } - if (some_tensor_not_init) { - file.close(); - return false; - } - LOG_DEBUG("model size = %.2fMB", total_size / 1024.0 / 1024.0); - } - max_params_mem_size = ggml_used_mem(clip_params_ctx) + ggml_used_mem(unet_params_ctx) + ggml_used_mem(vae_params_ctx); - max_mem_size = max_params_mem_size; - curr_params_mem_size = max_params_mem_size; - LOG_INFO("total params size = %.2fMB (clip %.2fMB, unet %.2fMB, vae %.2fMB)", - max_params_mem_size / 1024.0 / 1024.0, - ggml_used_mem(clip_params_ctx) / 1024.0 / 1024.0, - ggml_used_mem(unet_params_ctx) / 1024.0 / 1024.0, - ggml_used_mem(vae_params_ctx) / 1024.0 / 1024.0); - int64_t t1 = ggml_time_ms(); - LOG_INFO("loading model from '%s' completed, taking %.2fs", file_path.c_str(), (t1 - t0) * 1.0f / 1000); - file.close(); - - // check is_using_v_parameterization_for_sd2 - bool is_using_v_parameterization = false; - if (model_type == SD2) { - struct ggml_init_params params; - params.mem_size = static_cast(10 * 1024) * 1024; // 10M - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = false; - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return false; - } - if (is_using_v_parameterization_for_sd2(ctx)) { - is_using_v_parameterization = true; - } - } - - if (is_using_v_parameterization) { - denoiser = std::make_shared(); - LOG_INFO("running in v-prediction mode"); - } else { - LOG_INFO("running in eps-prediction mode"); - } - - if (schedule != DEFAULT) { - switch (schedule) { - case DISCRETE: - LOG_INFO("running with discrete schedule"); - denoiser->schedule = std::make_shared(); - break; - case KARRAS: - LOG_INFO("running with Karras schedule"); - denoiser->schedule = std::make_shared(); - break; - case DEFAULT: - // Don't touch anything. - break; - default: - LOG_ERROR("Unknown schedule %i", schedule); - abort(); - } - } - - for (int i = 0; i < TIMESTEPS; i++) { - denoiser->schedule->alphas_cumprod[i] = alphas_cumprod[i]; - denoiser->schedule->sigmas[i] = std::sqrt((1 - denoiser->schedule->alphas_cumprod[i]) / denoiser->schedule->alphas_cumprod[i]); - denoiser->schedule->log_sigmas[i] = std::log(denoiser->schedule->sigmas[i]); - } - - return true; - } - - bool is_using_v_parameterization_for_sd2(ggml_context* res_ctx) { - struct ggml_tensor* x_t = ggml_new_tensor_4d(res_ctx, GGML_TYPE_F32, 8, 8, 4, 1); - ggml_set_f32(x_t, 0.5); - struct ggml_tensor* c = ggml_new_tensor_4d(res_ctx, GGML_TYPE_F32, 1024, 2, 1, 1); - ggml_set_f32(c, 0.5); - - struct ggml_cplan cplan; - - size_t ctx_size = 10 * 1024 * 1024; // 10MB - // calculate the amount of memory required - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = true; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return false; - } - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* timesteps = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); // [N, ] - struct ggml_tensor* t_emb = new_timestep_embedding(ctx, timesteps, diffusion_model.model_channels); // [N, model_channels] - ggml_set_dynamic(ctx, params.dynamic); - - struct ggml_tensor* out = diffusion_model.forward(ctx, x_t, NULL, c, t_emb); - ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx); - - struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out); - cplan = ggml_graph_plan(diffusion_graph, n_threads); - - ctx_size += cplan.work_size; - LOG_DEBUG("diffusion context need %.2fMB static memory, with work_size needing %.2fMB", - ctx_size * 1.0f / 1024 / 1024, - cplan.work_size * 1.0f / 1024 / 1024); - - ggml_free(ctx); - } - - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return false; - } - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* timesteps = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); // [N, ] - struct ggml_tensor* t_emb = new_timestep_embedding(ctx, timesteps, diffusion_model.model_channels); // [N, model_channels] - ggml_set_dynamic(ctx, params.dynamic); - ggml_set_f32(timesteps, 999); - set_timestep_embedding(timesteps, t_emb, diffusion_model.model_channels); - - struct ggml_tensor* out = diffusion_model.forward(ctx, x_t, NULL, c, t_emb); - ggml_hold_dynamic_tensor(out); - - struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out); - cplan = ggml_graph_plan(diffusion_graph, n_threads); - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* buf = ggml_new_tensor_1d(ctx, GGML_TYPE_I8, cplan.work_size); - ggml_set_dynamic(ctx, params.dynamic); - - cplan.work_data = (uint8_t*)buf->data; - - int64_t t0 = ggml_time_ms(); - ggml_graph_compute(diffusion_graph, &cplan); - - double result = 0.f; - - { - float* vec_x = (float*)x_t->data; - float* vec_out = (float*)out->data; - - int64_t n = ggml_nelements(out); - - for (int i = 0; i < n; i++) { - result += ((double)vec_out[i] - (double)vec_x[i]); - } - result /= n; - } - -#ifdef GGML_PERF - ggml_graph_print(&diffusion_graph); -#endif - int64_t t1 = ggml_time_ms(); - LOG_INFO("check is_using_v_parameterization_for_sd2 completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - LOG_DEBUG("diffusion graph use %.2fMB runtime memory: static %.2fMB, dynamic %.2fMB", - (ctx_size + ggml_curr_max_dynamic_size()) * 1.0f / 1024 / 1024, - ctx_size * 1.0f / 1024 / 1024, - ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024); - LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size()); - - return result < -1; - } - - ggml_tensor* get_learned_condition(ggml_context* res_ctx, const std::string& text) { - auto tokens_and_weights = cond_stage_model.tokenize(text, - cond_stage_model.text_model.max_position_embeddings, - true); - std::vector& tokens = tokens_and_weights.first; - std::vector& weights = tokens_and_weights.second; - struct ggml_cplan cplan; - size_t ctx_size = 10 * 1024 * 1024; // 10MB - // calculate the amount of memory required - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = true; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* input_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens.size()); - ggml_set_dynamic(ctx, params.dynamic); - - struct ggml_tensor* hidden_states = cond_stage_model.text_model.forward(ctx, input_ids); - - struct ggml_cgraph* cond_graph = ggml_build_forward_ctx(ctx, hidden_states); - cplan = ggml_graph_plan(cond_graph, n_threads); - ctx_size += cplan.work_size; - - ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx); - LOG_DEBUG("condition context need %.2fMB static memory, with work_size needing %.2fMB", - ctx_size * 1.0f / 1024 / 1024, - cplan.work_size * 1.0f / 1024 / 1024); - ggml_free(ctx); - } - - // allocate the required memory and compute forward - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* input_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens.size()); - ggml_set_dynamic(ctx, params.dynamic); - - struct ggml_tensor* hidden_states = cond_stage_model.text_model.forward(ctx, input_ids); - struct ggml_cgraph* cond_graph = ggml_build_forward_ctx(ctx, hidden_states); - LOG_DEBUG("building condition graph completed: %d nodes, %d leafs", - cond_graph->n_nodes, cond_graph->n_leafs); - - memcpy(input_ids->data, tokens.data(), tokens.size() * ggml_element_size(input_ids)); - - int64_t t0 = ggml_time_ms(); - ggml_graph_compute_with_ctx(ctx, cond_graph, n_threads); - int64_t t1 = ggml_time_ms(); - LOG_DEBUG("computing condition graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - - ggml_tensor* result = ggml_dup_tensor(res_ctx, hidden_states); // [N, n_token, hidden_size] - - { - int64_t nelements = ggml_nelements(hidden_states); - float original_mean = 0.f; - float new_mean = 0.f; - float* vec = (float*)hidden_states->data; - for (int i = 0; i < nelements; i++) { - original_mean += vec[i] / nelements * 1.0f; - } - - for (int i2 = 0; i2 < hidden_states->ne[2]; i2++) { - for (int i1 = 0; i1 < hidden_states->ne[1]; i1++) { - for (int i0 = 0; i0 < hidden_states->ne[0]; i0++) { - float value = ggml_tensor_get_f32(hidden_states, i0, i1, i2); - value *= weights[i1]; - ggml_tensor_set_f32(result, value, i0, i1, i2); - } - } - } - - vec = (float*)result->data; - for (int i = 0; i < nelements; i++) { - new_mean += vec[i] / nelements * 1.0f; - } - - for (int i = 0; i < nelements; i++) { - vec[i] = vec[i] * (original_mean / new_mean); - } - } - - // print_ggml_tensor(result); - - size_t rt_mem_size = ctx_size + ggml_curr_max_dynamic_size(); - if (rt_mem_size > max_rt_mem_size) { - max_rt_mem_size = rt_mem_size; - } - size_t graph_mem_size = ggml_used_mem(clip_params_ctx) + rt_mem_size; - - size_t curr_mem_size = curr_params_mem_size + rt_mem_size; - if (curr_mem_size > max_mem_size) { - max_mem_size = curr_mem_size; - } - - LOG_INFO( - "condition graph use %.2fMB of memory: params %.2fMB, " - "runtime %.2fMB (static %.2fMB, dynamic %.2fMB)", - graph_mem_size * 1.0f / 1024 / 1024, - ggml_used_mem(clip_params_ctx) * 1.0f / 1024 / 1024, - rt_mem_size * 1.0f / 1024 / 1024, - ctx_size * 1.0f / 1024 / 1024, - ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024); - - LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size()); - - ggml_free(ctx); - - return result; // [1, 77, 768] - } - - ggml_tensor* sample(ggml_context* res_ctx, - ggml_tensor* x_t, - ggml_tensor* c, - ggml_tensor* uc, - float cfg_scale, - SampleMethod method, - const std::vector& sigmas) { - size_t steps = sigmas.size() - 1; - // x_t = load_tensor_from_file(res_ctx, "./rand0.bin"); - // print_ggml_tensor(x_t); - struct ggml_tensor* x = ggml_dup_tensor(res_ctx, x_t); - copy_ggml_tensor(x, x_t); - struct ggml_cplan cplan; - - size_t ctx_size = 10 * 1024 * 1024; // 10MB - // calculate the amount of memory required - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = true; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* noised_input = ggml_dup_tensor(ctx, x_t); - struct ggml_tensor* context = ggml_dup_tensor(ctx, c); - struct ggml_tensor* timesteps = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); // [N, ] - struct ggml_tensor* t_emb = new_timestep_embedding(ctx, timesteps, diffusion_model.model_channels); // [N, model_channels] - ggml_set_dynamic(ctx, params.dynamic); - - struct ggml_tensor* out = diffusion_model.forward(ctx, noised_input, NULL, context, t_emb); - ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx); - - struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out); - cplan = ggml_graph_plan(diffusion_graph, n_threads); - - ctx_size += cplan.work_size; - LOG_DEBUG("diffusion context need %.2fMB static memory, with work_size needing %.2fMB", - ctx_size * 1.0f / 1024 / 1024, - cplan.work_size * 1.0f / 1024 / 1024); - - ggml_free(ctx); - } - - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* noised_input = ggml_dup_tensor(ctx, x_t); - struct ggml_tensor* context = ggml_dup_tensor(ctx, c); - struct ggml_tensor* timesteps = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); // [N, ] - struct ggml_tensor* t_emb = new_timestep_embedding(ctx, timesteps, diffusion_model.model_channels); // [N, model_channels] - ggml_set_dynamic(ctx, params.dynamic); - - struct ggml_tensor* out = diffusion_model.forward(ctx, noised_input, NULL, context, t_emb); - ggml_hold_dynamic_tensor(out); - - struct ggml_cgraph* diffusion_graph = ggml_build_forward_ctx(ctx, out); - cplan = ggml_graph_plan(diffusion_graph, n_threads); - - ggml_set_dynamic(ctx, false); - struct ggml_tensor* buf = ggml_new_tensor_1d(ctx, GGML_TYPE_I8, cplan.work_size); - ggml_set_dynamic(ctx, params.dynamic); - - cplan.work_data = (uint8_t*)buf->data; - - // x = x * sigmas[0] - { - float* vec = (float*)x->data; - for (int i = 0; i < ggml_nelements(x); i++) { - vec[i] = vec[i] * sigmas[0]; - } - } - - // denoise wrapper - ggml_set_dynamic(ctx, false); - struct ggml_tensor* out_cond = NULL; - struct ggml_tensor* out_uncond = NULL; - if (cfg_scale != 1.0f && uc != NULL) { - out_uncond = ggml_dup_tensor(ctx, x); - } - struct ggml_tensor* denoised = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - auto denoise = [&](ggml_tensor* input, float sigma, int step) { - int64_t t0 = ggml_time_ms(); - - float c_skip = 1.0f; - float c_out = 1.0f; - float c_in = 1.0f; - std::vector scaling = denoiser->get_scalings(sigma); - if (scaling.size() == 3) { // CompVisVDenoiser - c_skip = scaling[0]; - c_out = scaling[1]; - c_in = scaling[2]; - } else { // CompVisDenoiser - c_out = scaling[0]; - c_in = scaling[1]; - } - - float t = denoiser->schedule->sigma_to_t(sigma); - ggml_set_f32(timesteps, t); - set_timestep_embedding(timesteps, t_emb, diffusion_model.model_channels); - - copy_ggml_tensor(noised_input, input); - // noised_input = noised_input * c_in - { - float* vec = (float*)noised_input->data; - for (int i = 0; i < ggml_nelements(noised_input); i++) { - vec[i] = vec[i] * c_in; - } - } - - if (cfg_scale != 1.0 && uc != NULL) { - // uncond - copy_ggml_tensor(context, uc); - ggml_graph_compute(diffusion_graph, &cplan); - copy_ggml_tensor(out_uncond, out); - - // cond - copy_ggml_tensor(context, c); - ggml_graph_compute(diffusion_graph, &cplan); - - out_cond = out; - - // out_uncond + cfg_scale * (out_cond - out_uncond) - { - float* vec_out = (float*)out->data; - float* vec_out_uncond = (float*)out_uncond->data; - float* vec_out_cond = (float*)out_cond->data; - - for (int i = 0; i < ggml_nelements(out); i++) { - vec_out[i] = vec_out_uncond[i] + cfg_scale * (vec_out_cond[i] - vec_out_uncond[i]); - } - } - } else { - // cond - copy_ggml_tensor(context, c); - ggml_graph_compute(diffusion_graph, &cplan); - } - - // v = out, eps = out - // denoised = (v * c_out + input * c_skip) or (input + eps * c_out) - { - float* vec_denoised = (float*)denoised->data; - float* vec_input = (float*)input->data; - float* vec_out = (float*)out->data; - - for (int i = 0; i < ggml_nelements(denoised); i++) { - vec_denoised[i] = vec_out[i] * c_out + vec_input[i] * c_skip; - } - } - -#ifdef GGML_PERF - ggml_graph_print(&diffusion_graph); -#endif - int64_t t1 = ggml_time_ms(); - if (step > 0) { - LOG_INFO("step %d sampling completed, taking %.2fs", step, (t1 - t0) * 1.0f / 1000); - LOG_DEBUG("diffusion graph use %.2fMB runtime memory: static %.2fMB, dynamic %.2fMB", - (ctx_size + ggml_curr_max_dynamic_size()) * 1.0f / 1024 / 1024, - ctx_size * 1.0f / 1024 / 1024, - ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024); - LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size()); - } - }; - - // sample_euler_ancestral - switch (method) { - case EULER_A: { - LOG_INFO("sampling using Euler A method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* noise = ggml_dup_tensor(ctx, x); - struct ggml_tensor* d = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - for (int i = 0; i < steps; i++) { - float sigma = sigmas[i]; - - // denoise - denoise(x, sigma, i + 1); - - // d = (x - denoised) / sigma - { - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - - for (int i = 0; i < ggml_nelements(d); i++) { - vec_d[i] = (vec_x[i] - vec_denoised[i]) / sigma; - } - } - - // get_ancestral_step - float sigma_up = std::min(sigmas[i + 1], - std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i]))); - float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up); - - // Euler method - float dt = sigma_down - sigmas[i]; - // x = x + d * dt - { - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - - for (int i = 0; i < ggml_nelements(x); i++) { - vec_x[i] = vec_x[i] + vec_d[i] * dt; - } - } - - if (sigmas[i + 1] > 0) { - // x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up - ggml_tensor_set_f32_randn(noise, rng); - // noise = load_tensor_from_file(res_ctx, "./rand" + std::to_string(i+1) + ".bin"); - { - float* vec_x = (float*)x->data; - float* vec_noise = (float*)noise->data; - - for (int i = 0; i < ggml_nelements(x); i++) { - vec_x[i] = vec_x[i] + vec_noise[i] * sigma_up; - } - } - } - } - } break; - case EULER: // Implemented without any sigma churn - { - LOG_INFO("sampling using Euler method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* d = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - for (int i = 0; i < steps; i++) { - float sigma = sigmas[i]; - - // denoise - denoise(x, sigma, i + 1); - - // d = (x - denoised) / sigma - { - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - - for (int j = 0; j < ggml_nelements(d); j++) { - vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigma; - } - } - - float dt = sigmas[i + 1] - sigma; - // x = x + d * dt - { - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x[j] = vec_x[j] + vec_d[j] * dt; - } - } - } - } break; - case HEUN: { - LOG_INFO("sampling using Heun method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* d = ggml_dup_tensor(ctx, x); - struct ggml_tensor* x2 = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - for (int i = 0; i < steps; i++) { - // denoise - denoise(x, sigmas[i], -(i + 1)); - - // d = (x - denoised) / sigma - { - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - - for (int j = 0; j < ggml_nelements(x); j++) { - vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i]; - } - } - - float dt = sigmas[i + 1] - sigmas[i]; - if (sigmas[i + 1] == 0) { - // Euler step - // x = x + d * dt - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x[j] = vec_x[j] + vec_d[j] * dt; - } - } else { - // Heun step - float* vec_d = (float*)d->data; - float* vec_d2 = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_x2 = (float*)x2->data; - - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x2[j] = vec_x[j] + vec_d[j] * dt; - } - - denoise(x2, sigmas[i + 1], i + 1); - float* vec_denoised = (float*)denoised->data; - for (int j = 0; j < ggml_nelements(x); j++) { - float d2 = (vec_x2[j] - vec_denoised[j]) / sigmas[i + 1]; - vec_d[j] = (vec_d[j] + d2) / 2; - vec_x[j] = vec_x[j] + vec_d[j] * dt; - } - } - } - } break; - case DPM2: { - LOG_INFO("sampling using DPM2 method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* d = ggml_dup_tensor(ctx, x); - struct ggml_tensor* x2 = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - for (int i = 0; i < steps; i++) { - // denoise - denoise(x, sigmas[i], i + 1); - - // d = (x - denoised) / sigma - { - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - - for (int j = 0; j < ggml_nelements(x); j++) { - vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i]; - } - } - - if (sigmas[i + 1] == 0) { - // Euler step - // x = x + d * dt - float dt = sigmas[i + 1] - sigmas[i]; - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x[j] = vec_x[j] + vec_d[j] * dt; - } - } else { - // DPM-Solver-2 - float sigma_mid = exp(0.5 * (log(sigmas[i]) + log(sigmas[i + 1]))); - float dt_1 = sigma_mid - sigmas[i]; - float dt_2 = sigmas[i + 1] - sigmas[i]; - - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_x2 = (float*)x2->data; - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x2[j] = vec_x[j] + vec_d[j] * dt_1; - } - - denoise(x2, sigma_mid, i + 1); - float* vec_denoised = (float*)denoised->data; - for (int j = 0; j < ggml_nelements(x); j++) { - float d2 = (vec_x2[j] - vec_denoised[j]) / sigma_mid; - vec_x[j] = vec_x[j] + d2 * dt_2; - } - } - } - - } break; - case DPMPP2S_A: { - LOG_INFO("sampling using DPM++ (2s) a method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* noise = ggml_dup_tensor(ctx, x); - struct ggml_tensor* d = ggml_dup_tensor(ctx, x); - struct ggml_tensor* x2 = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - for (int i = 0; i < steps; i++) { - // denoise - denoise(x, sigmas[i], i + 1); - - // get_ancestral_step - float sigma_up = std::min(sigmas[i + 1], - std::sqrt(sigmas[i + 1] * sigmas[i + 1] * (sigmas[i] * sigmas[i] - sigmas[i + 1] * sigmas[i + 1]) / (sigmas[i] * sigmas[i]))); - float sigma_down = std::sqrt(sigmas[i + 1] * sigmas[i + 1] - sigma_up * sigma_up); - auto t_fn = [](float sigma) -> float { return -log(sigma); }; - auto sigma_fn = [](float t) -> float { return exp(-t); }; - - if (sigma_down == 0) { - // Euler step - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - - for (int j = 0; j < ggml_nelements(d); j++) { - vec_d[j] = (vec_x[j] - vec_denoised[j]) / sigmas[i]; - } - - // TODO: If sigma_down == 0, isn't this wrong? - // But - // https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/sampling.py#L525 - // has this exactly the same way. - float dt = sigma_down - sigmas[i]; - for (int j = 0; j < ggml_nelements(d); j++) { - vec_x[j] = vec_x[j] + vec_d[j] * dt; - } - } else { - // DPM-Solver++(2S) - float t = t_fn(sigmas[i]); - float t_next = t_fn(sigma_down); - float h = t_next - t; - float s = t + 0.5 * h; - - float* vec_d = (float*)d->data; - float* vec_x = (float*)x->data; - float* vec_x2 = (float*)x2->data; - float* vec_denoised = (float*)denoised->data; - - // First half-step - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x2[j] = (sigma_fn(s) / sigma_fn(t)) * vec_x[j] - (exp(-h * 0.5) - 1) * vec_denoised[j]; - } - - denoise(x2, sigmas[i + 1], i + 1); - - // Second half-step - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x[j] = (sigma_fn(t_next) / sigma_fn(t)) * vec_x[j] - (exp(-h) - 1) * vec_denoised[j]; - } - } - - // Noise addition - if (sigmas[i + 1] > 0) { - ggml_tensor_set_f32_randn(noise, rng); - { - float* vec_x = (float*)x->data; - float* vec_noise = (float*)noise->data; - - for (int i = 0; i < ggml_nelements(x); i++) { - vec_x[i] = vec_x[i] + vec_noise[i] * sigma_up; - } - } - } - } - } break; - case DPMPP2M: // DPM++ (2M) from Karras et al (2022) - { - LOG_INFO("sampling using DPM++ (2M) method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* old_denoised = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - auto t_fn = [](float sigma) -> float { return -log(sigma); }; - - for (int i = 0; i < steps; i++) { - // denoise - denoise(x, sigmas[i], i + 1); - - float t = t_fn(sigmas[i]); - float t_next = t_fn(sigmas[i + 1]); - float h = t_next - t; - float a = sigmas[i + 1] / sigmas[i]; - float b = exp(-h) - 1.; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - float* vec_old_denoised = (float*)old_denoised->data; - - if (i == 0 || sigmas[i + 1] == 0) { - // Simpler step for the edge cases - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x[j] = a * vec_x[j] - b * vec_denoised[j]; - } - } else { - float h_last = t - t_fn(sigmas[i - 1]); - float r = h_last / h; - for (int j = 0; j < ggml_nelements(x); j++) { - float denoised_d = (1. + 1. / (2. * r)) * vec_denoised[j] - (1. / (2. * r)) * vec_old_denoised[j]; - vec_x[j] = a * vec_x[j] - b * denoised_d; - } - } - - // old_denoised = denoised - for (int j = 0; j < ggml_nelements(x); j++) { - vec_old_denoised[j] = vec_denoised[j]; - } - } - } break; - case DPMPP2Mv2: // Modified DPM++ (2M) from https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions/8457 - { - LOG_INFO("sampling using modified DPM++ (2M) method"); - ggml_set_dynamic(ctx, false); - struct ggml_tensor* old_denoised = ggml_dup_tensor(ctx, x); - ggml_set_dynamic(ctx, params.dynamic); - - auto t_fn = [](float sigma) -> float { return -log(sigma); }; - - for (int i = 0; i < steps; i++) { - // denoise - denoise(x, sigmas[i], i + 1); - - float t = t_fn(sigmas[i]); - float t_next = t_fn(sigmas[i + 1]); - float h = t_next - t; - float a = sigmas[i + 1] / sigmas[i]; - float* vec_x = (float*)x->data; - float* vec_denoised = (float*)denoised->data; - float* vec_old_denoised = (float*)old_denoised->data; - - if (i == 0 || sigmas[i + 1] == 0) { - // Simpler step for the edge cases - float b = exp(-h) - 1.; - for (int j = 0; j < ggml_nelements(x); j++) { - vec_x[j] = a * vec_x[j] - b * vec_denoised[j]; - } - } else { - float h_last = t - t_fn(sigmas[i - 1]); - float h_min = std::min(h_last, h); - float h_max = std::max(h_last, h); - float r = h_max / h_min; - float h_d = (h_max + h_min) / 2.; - float b = exp(-h_d) - 1.; - for (int j = 0; j < ggml_nelements(x); j++) { - float denoised_d = (1. + 1. / (2. * r)) * vec_denoised[j] - (1. / (2. * r)) * vec_old_denoised[j]; - vec_x[j] = a * vec_x[j] - b * denoised_d; - } - } - - // old_denoised = denoised - for (int j = 0; j < ggml_nelements(x); j++) { - vec_old_denoised[j] = vec_denoised[j]; - } - } - } break; - - default: - LOG_ERROR("Attempting to sample with nonexisting sample method %i", method); - abort(); - } - - size_t rt_mem_size = ctx_size + ggml_curr_max_dynamic_size(); - if (rt_mem_size > max_rt_mem_size) { - max_rt_mem_size = rt_mem_size; - } - size_t graph_mem_size = ggml_used_mem(unet_params_ctx) + rt_mem_size; - - size_t curr_mem_size = curr_params_mem_size + rt_mem_size; - if (curr_mem_size > max_mem_size) { - max_mem_size = curr_mem_size; - } - - LOG_INFO( - "diffusion graph use %.2fMB of memory: params %.2fMB, " - "runtime %.2fMB (static %.2fMB, dynamic %.2fMB)", - graph_mem_size * 1.0f / 1024 / 1024, - ggml_used_mem(unet_params_ctx) * 1.0f / 1024 / 1024, - rt_mem_size * 1.0f / 1024 / 1024, - ctx_size * 1.0f / 1024 / 1024, - ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024); - LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size()); - - ggml_free(ctx); - - return x; - } - - ggml_tensor* encode_first_stage(ggml_context* res_ctx, ggml_tensor* x) { - int64_t W = x->ne[0]; - int64_t H = x->ne[1]; - struct ggml_tensor* result = NULL; - struct ggml_cplan cplan; - - // calculate the amount of memory required - size_t ctx_size = 10 * 1024 * 1024; // 10MB - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = true; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - struct ggml_tensor* moments = first_stage_model.encode(ctx, x); - ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx); - - struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments); - cplan = ggml_graph_plan(vae_graph, n_threads); - - ctx_size += cplan.work_size; - LOG_DEBUG("vae context need %.2fMB static memory, with work_size needing %.2fMB", - ctx_size * 1.0f / 1024 / 1024, - cplan.work_size * 1.0f / 1024 / 1024); - - ggml_free(ctx); - } - - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - struct ggml_tensor* moments = first_stage_model.encode(ctx, x); - struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, moments); - - int64_t t0 = ggml_time_ms(); - ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads); - int64_t t1 = ggml_time_ms(); - -#ifdef GGML_PERF - ggml_graph_print(&vae_graph); -#endif - LOG_DEBUG("computing vae graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - - result = ggml_dup_tensor(res_ctx, moments); - copy_ggml_tensor(result, moments); - - size_t rt_mem_size = ctx_size + ggml_curr_max_dynamic_size(); - if (rt_mem_size > max_rt_mem_size) { - max_rt_mem_size = rt_mem_size; - } - size_t graph_mem_size = ggml_used_mem(vae_params_ctx) + rt_mem_size; - - size_t curr_mem_size = curr_params_mem_size + rt_mem_size; - if (curr_mem_size > max_mem_size) { - max_mem_size = curr_mem_size; - } - - LOG_INFO( - "vae graph use %.2fMB of memory: params %.2fMB, " - "runtime %.2fMB (static %.2fMB, dynamic %.2fMB)", - graph_mem_size * 1.0f / 1024 / 1024, - ggml_used_mem(vae_params_ctx) * 1.0f / 1024 / 1024, - rt_mem_size * 1.0f / 1024 / 1024, - ctx_size * 1.0f / 1024 / 1024, - ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024); - LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size()); - - ggml_free(ctx); - } - - return result; - } - - // ldm.models.diffusion.ddpm.LatentDiffusion.get_first_stage_encoding - ggml_tensor* get_first_stage_encoding(ggml_context* res_ctx, ggml_tensor* moments) { - // ldm.modules.distributions.distributions.DiagonalGaussianDistribution.sample - ggml_tensor* latent = ggml_new_tensor_4d(res_ctx, moments->type, moments->ne[0], - moments->ne[1], moments->ne[2] / 2, moments->ne[3]); - struct ggml_tensor* noise = ggml_dup_tensor(res_ctx, latent); - ggml_tensor_set_f32_randn(noise, rng); - // noise = load_tensor_from_file(res_ctx, "noise.bin"); - { - float mean = 0; - float logvar = 0; - float value = 0; - float std_ = 0; - for (int i = 0; i < latent->ne[3]; i++) { - for (int j = 0; j < latent->ne[2]; j++) { - for (int k = 0; k < latent->ne[1]; k++) { - for (int l = 0; l < latent->ne[0]; l++) { - mean = ggml_tensor_get_f32(moments, l, k, j, i); - logvar = ggml_tensor_get_f32(moments, l, k, j + (int)latent->ne[2], i); - logvar = std::max(-30.0f, std::min(logvar, 20.0f)); - std_ = std::exp(0.5f * logvar); - value = mean + std_ * ggml_tensor_get_f32(noise, l, k, j, i); - value = value * scale_factor; - // printf("%d %d %d %d -> %f\n", i, j, k, l, value); - ggml_tensor_set_f32(latent, value, l, k, j, i); - } - } - } - } - } - return latent; - } - - ggml_tensor* decode_first_stage(ggml_context* res_ctx, ggml_tensor* z) { - int64_t W = z->ne[0]; - int64_t H = z->ne[1]; - struct ggml_tensor* result_img = NULL; - struct ggml_cplan cplan; - - { - float* vec = (float*)z->data; - for (int i = 0; i < ggml_nelements(z); i++) { - vec[i] = 1.0f / scale_factor * vec[i]; - } - } - - // calculate the amount of memory required - size_t ctx_size = 10 * 1024 * 1024; // 10MB - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = true; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - struct ggml_tensor* img = first_stage_model.decoder.forward(ctx, z); - ctx_size += ggml_used_mem(ctx) + ggml_used_mem_of_data(ctx); - - struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img); - cplan = ggml_graph_plan(vae_graph, n_threads); - - ctx_size += cplan.work_size; - LOG_DEBUG("vae context need %.2fMB static memory, with work_size needing %.2fMB", - ctx_size * 1.0f / 1024 / 1024, - cplan.work_size * 1.0f / 1024 / 1024); - - ggml_free(ctx); - } - - { - struct ggml_init_params params; - params.mem_size = ctx_size; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = dynamic; - - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return NULL; - } - - struct ggml_tensor* img = first_stage_model.decode(ctx, z); - struct ggml_cgraph* vae_graph = ggml_build_forward_ctx(ctx, img); - - int64_t t0 = ggml_time_ms(); - ggml_graph_compute_with_ctx(ctx, vae_graph, n_threads); - int64_t t1 = ggml_time_ms(); - -#ifdef GGML_PERF - ggml_graph_print(&vae_graph); -#endif - LOG_DEBUG("computing vae graph completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - - result_img = ggml_dup_tensor(res_ctx, img); - copy_ggml_tensor(result_img, img); - - size_t rt_mem_size = ctx_size + ggml_curr_max_dynamic_size(); - if (rt_mem_size > max_rt_mem_size) { - max_rt_mem_size = rt_mem_size; - } - size_t graph_mem_size = ggml_used_mem(vae_params_ctx) + rt_mem_size; - - size_t curr_mem_size = curr_params_mem_size + rt_mem_size; - if (curr_mem_size > max_mem_size) { - max_mem_size = curr_mem_size; - } - - LOG_INFO( - "vae graph use %.2fMB of memory: params %.2fMB, " - "runtime %.2fMB (static %.2fMB, dynamic %.2fMB)", - graph_mem_size * 1.0f / 1024 / 1024, - ggml_used_mem(vae_params_ctx) * 1.0f / 1024 / 1024, - rt_mem_size * 1.0f / 1024 / 1024, - ctx_size * 1.0f / 1024 / 1024, - ggml_curr_max_dynamic_size() * 1.0f / 1024 / 1024); - LOG_DEBUG("%zu bytes of dynamic memory has not been released yet", ggml_dynamic_size()); - - ggml_free(ctx); - } - - return result_img; - } -}; - -/*================================================= StableDiffusion ==================================================*/ - -StableDiffusion::StableDiffusion(int n_threads, - bool vae_decode_only, - bool free_params_immediately, - RNGType rng_type) { - sd = std::make_shared(n_threads, - vae_decode_only, - free_params_immediately, - rng_type); -} - -bool StableDiffusion::load_from_file(const std::string& file_path, Schedule s) { - return sd->load_from_file(file_path, s); -} - -std::vector StableDiffusion::txt2img(const std::string& prompt, - const std::string& negative_prompt, - float cfg_scale, - int width, - int height, - SampleMethod sample_method, - int sample_steps, - int64_t seed) { - std::vector result; - struct ggml_init_params params; - params.mem_size = static_cast(10 * 1024) * 1024; // 10M - params.mem_size += width * height * 3 * sizeof(float) * 2; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = false; - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return result; - } - - if (seed < 0) { - seed = (int)time(NULL); - } - sd->rng->manual_seed(seed); - - int64_t t0 = ggml_time_ms(); - ggml_tensor* c = sd->get_learned_condition(ctx, prompt); - struct ggml_tensor* uc = NULL; - if (cfg_scale != 1.0) { - uc = sd->get_learned_condition(ctx, negative_prompt); - } - int64_t t1 = ggml_time_ms(); - LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - - if (sd->free_params_immediately) { - sd->curr_params_mem_size -= ggml_used_mem(sd->clip_params_ctx); - ggml_free(sd->clip_params_ctx); - sd->clip_params_ctx = NULL; - } - - int C = 4; - int W = width / 8; - int H = height / 8; - struct ggml_tensor* x_t = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, W, H, C, 1); - ggml_tensor_set_f32_randn(x_t, sd->rng); - - std::vector sigmas = sd->denoiser->schedule->get_sigmas(sample_steps); - - LOG_INFO("start sampling"); - struct ggml_tensor* x_0 = sd->sample(ctx, x_t, c, uc, cfg_scale, sample_method, sigmas); - // struct ggml_tensor* x_0 = load_tensor_from_file(ctx, "samples_ddim.bin"); - // print_ggml_tensor(x_0); - int64_t t2 = ggml_time_ms(); - LOG_INFO("sampling completed, taking %.2fs", (t2 - t1) * 1.0f / 1000); - - if (sd->free_params_immediately) { - sd->curr_params_mem_size -= ggml_used_mem(sd->unet_params_ctx); - ggml_free(sd->unet_params_ctx); - sd->unet_params_ctx = NULL; - } - - struct ggml_tensor* img = sd->decode_first_stage(ctx, x_0); - if (img != NULL) { - result = ggml_to_image_vec(img); - } - int64_t t3 = ggml_time_ms(); - LOG_INFO("decode_first_stage completed, taking %.2fs", (t3 - t2) * 1.0f / 1000); - - if (sd->free_params_immediately) { - sd->curr_params_mem_size -= ggml_used_mem(sd->vae_params_ctx); - ggml_free(sd->vae_params_ctx); - sd->vae_params_ctx = NULL; - } - - LOG_INFO( - "txt2img completed in %.2fs, use %.2fMB of memory: peak params memory %.2fMB, " - "peak runtime memory %.2fMB", - (t3 - t0) * 1.0f / 1000, - sd->max_mem_size * 1.0f / 1024 / 1024, - sd->max_params_mem_size * 1.0f / 1024 / 1024, - sd->max_rt_mem_size * 1.0f / 1024 / 1024); - - ggml_free(ctx); - return result; -} - -std::vector StableDiffusion::img2img(const std::vector& init_img_vec, - const std::string& prompt, - const std::string& negative_prompt, - float cfg_scale, - int width, - int height, - SampleMethod sample_method, - int sample_steps, - float strength, - int64_t seed) { - std::vector result; - if (init_img_vec.size() != width * height * 3) { - return result; - } - LOG_INFO("img2img %dx%d", width, height); - - std::vector sigmas = sd->denoiser->schedule->get_sigmas(sample_steps); - size_t t_enc = static_cast(sample_steps * strength); - LOG_INFO("target t_enc is %zu steps", t_enc); - std::vector sigma_sched; - sigma_sched.assign(sigmas.begin() + sample_steps - t_enc - 1, sigmas.end()); - - struct ggml_init_params params; - params.mem_size = static_cast(10 * 1024) * 1024; // 10M - params.mem_size += width * height * 3 * sizeof(float) * 2; - params.mem_buffer = NULL; - params.no_alloc = false; - params.dynamic = false; - struct ggml_context* ctx = ggml_init(params); - if (!ctx) { - LOG_ERROR("ggml_init() failed"); - return result; - } - - if (seed < 0) { - seed = (int)time(NULL); - } - sd->rng->manual_seed(seed); - - ggml_tensor* init_img = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, height, 3, 1); - image_vec_to_ggml(init_img_vec, init_img); - - int64_t t0 = ggml_time_ms(); - ggml_tensor* moments = sd->encode_first_stage(ctx, init_img); - ggml_tensor* init_latent = sd->get_first_stage_encoding(ctx, moments); - // print_ggml_tensor(init_latent); - int64_t t1 = ggml_time_ms(); - LOG_INFO("encode_first_stage completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - - ggml_reset_curr_max_dynamic_size(); // reset counter - - ggml_tensor* c = sd->get_learned_condition(ctx, prompt); - struct ggml_tensor* uc = NULL; - if (cfg_scale != 1.0) { - uc = sd->get_learned_condition(ctx, negative_prompt); - } - int64_t t2 = ggml_time_ms(); - LOG_INFO("get_learned_condition completed, taking %.2fs", (t2 - t1) * 1.0f / 1000); - if (sd->free_params_immediately) { - sd->curr_params_mem_size -= ggml_used_mem(sd->clip_params_ctx); - ggml_free(sd->clip_params_ctx); - sd->clip_params_ctx = NULL; - } - - LOG_INFO("start sampling"); - struct ggml_tensor* x_0 = sd->sample(ctx, init_latent, c, uc, cfg_scale, sample_method, sigma_sched); - // struct ggml_tensor *x_0 = load_tensor_from_file(ctx, "samples_ddim.bin"); - // print_ggml_tensor(x_0); - int64_t t3 = ggml_time_ms(); - LOG_INFO("sampling completed, taking %.2fs", (t3 - t2) * 1.0f / 1000); - if (sd->free_params_immediately) { - sd->curr_params_mem_size -= ggml_used_mem(sd->unet_params_ctx); - ggml_free(sd->unet_params_ctx); - sd->unet_params_ctx = NULL; - } - - struct ggml_tensor* img = sd->decode_first_stage(ctx, x_0); - if (img != NULL) { - result = ggml_to_image_vec(img); - } - int64_t t4 = ggml_time_ms(); - LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t3) * 1.0f / 1000); - - if (sd->free_params_immediately) { - sd->curr_params_mem_size -= ggml_used_mem(sd->vae_params_ctx); - ggml_free(sd->vae_params_ctx); - sd->vae_params_ctx = NULL; - } - - LOG_INFO( - "img2img completed in %.2fs, use %.2fMB of memory: peak params memory %.2fMB, " - "peak runtime memory %.2fMB", - (t4 - t0) * 1.0f / 1000, - sd->max_mem_size * 1.0f / 1024 / 1024, - sd->max_params_mem_size * 1.0f / 1024 / 1024, - sd->max_rt_mem_size * 1.0f / 1024 / 1024); - - ggml_free(ctx); - - return result; -}