Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-08-29 00:20:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

sd: sync to master-678-dfb2390 (#2251)

Browse Source 
* sd: sync to master-666-7948df8

* sd: sync to master-672-1f9ee88

* sd: sync to master-676-b9254dd

* sd: sync to master-678-dfb2390
This commit is contained in:
Wagner Bruna
2026-06-07 11:18:03 -03:00
committed by GitHub
parent 17117ecf1d
commit 635beb6891
51 changed files with 6517 additions and 3077 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -695,7 +695,7 @@ llama-impl.o: src/llama-impl.cpp src/llama-impl.h
budget.o: common/reasoning-budget.cpp common/reasoning-budget.h
$(CXX) $(CXXFLAGS) -c $< -o $@
SDCPP_COMMON_BASENAMES := anima.hpp auto_encoder_kl.hpp avi_writer.h cache_dit.hpp clip.hpp common_block.hpp common_dit.hpp condition_cache_utils.hpp conditioner.hpp control.hpp convert.cpp denoiser.hpp diffusion_model.hpp easycache.hpp ernie_image.hpp esrgan.hpp flux.hpp ggml_extend.hpp ggml_extend_backend.cpp ggml_extend_backend.h ggml_graph_cut.cpp ggml_graph_cut.h gits_noise.inl guidance.cpp guidance.h hidream_o1.hpp image_metadata.cpp image_metadata.h kcpp_sd_extensions.h latent-preview.h llm.hpp lora.hpp ltx_audio_vae.h ltx_latent_upscaler.hpp ltx_vae.hpp ltxv.hpp mmdit.hpp model.cpp model.h model_io/binary_io.h model_io/gguf_io.cpp model_io/gguf_io.h model_io/gguf_reader_ext.h model_io/pickle_io.cpp model_io/pickle_io.h model_io/safetensors_io.cpp model_io/safetensors_io.h model_io/tensor_storage.h model_io/torch_legacy_io.cpp model_io/torch_legacy_io.h model_io/torch_zip_io.cpp model_io/torch_zip_io.h msf_gif.h name_conversion.cpp name_conversion.h ordered_map.hpp pmid.hpp preprocessing.hpp qwen_image.hpp rng.hpp rng_mt19937.hpp rng_philox.hpp rope.hpp sample-cache.cpp sample-cache.h spectrum.hpp stable-diffusion.cpp stable-diffusion.h t5.hpp tae.hpp tensor.hpp tensor_ggml.hpp thirdparty/LICENSE.darts_clone.txt thirdparty/darts.h thirdparty/miniz.h thirdparty/stb_image_resize.h thirdparty/stb_image_write.h thirdparty/zip.c thirdparty/zip.h tokenizers/bpe_tokenizer.cpp tokenizers/bpe_tokenizer.h tokenizers/clip_tokenizer.cpp tokenizers/clip_tokenizer.h tokenizers/gemma_tokenizer.cpp tokenizers/gemma_tokenizer.h tokenizers/gpt_oss_tokenizer.cpp tokenizers/gpt_oss_tokenizer.h tokenizers/mistral_tokenizer.cpp tokenizers/mistral_tokenizer.h tokenizers/qwen2_tokenizer.cpp tokenizers/qwen2_tokenizer.h tokenizers/t5_unigram_tokenizer.cpp tokenizers/t5_unigram_tokenizer.h tokenizers/tokenize_util.cpp tokenizers/tokenize_util.h tokenizers/tokenizer.cpp tokenizers/tokenizer.h tokenizers/vocab/vocab.h ucache.hpp unet.hpp upscaler.cpp upscaler.h util.cpp util.h vae.hpp wan.hpp z_image.hpp
SDCPP_COMMON_BASENAMES := anima.hpp auto_encoder_kl.hpp avi_writer.h cache_dit.hpp clip.hpp common_block.hpp common_dit.hpp condition_cache_utils.hpp conditioner.hpp control.hpp convert.cpp denoiser.hpp diffusion_model.hpp easycache.hpp ernie_image.hpp esrgan.hpp flux.hpp ggml_extend.hpp ggml_extend_backend.cpp ggml_extend_backend.h ggml_graph_cut.cpp ggml_graph_cut.h gits_noise.inl guidance.cpp guidance.h hidream_o1.hpp image_metadata.cpp image_metadata.h kcpp_sd_extensions.h latent-preview.h layer_registry.cpp layer_registry.h llm.hpp lora.hpp ltx_audio_vae.h ltx_latent_upscaler.hpp ltx_vae.hpp ltxv.hpp mmdit.hpp model.cpp model.h model_io/binary_io.h model_io/gguf_io.cpp model_io/gguf_io.h model_io/gguf_reader_ext.h model_io/pickle_io.cpp model_io/pickle_io.h model_io/safetensors_io.cpp model_io/safetensors_io.h model_io/tensor_storage.h model_io/torch_legacy_io.cpp model_io/torch_legacy_io.h model_io/torch_zip_io.cpp model_io/torch_zip_io.h msf_gif.h name_conversion.cpp name_conversion.h ordered_map.hpp pmid.hpp preprocessing.hpp qwen_image.hpp rng.hpp rng_mt19937.hpp rng_philox.hpp rope.hpp sample-cache.cpp sample-cache.h spectrum.hpp stable-diffusion.cpp stable-diffusion.h t5.hpp tae.hpp tensor.hpp tensor_ggml.hpp thirdparty/LICENSE.darts_clone.txt thirdparty/darts.h thirdparty/miniz.h thirdparty/stb_image_resize.h thirdparty/stb_image_write.h thirdparty/zip.c thirdparty/zip.h tokenizers/bpe_tokenizer.cpp tokenizers/bpe_tokenizer.h tokenizers/clip_tokenizer.cpp tokenizers/clip_tokenizer.h tokenizers/gemma_tokenizer.cpp tokenizers/gemma_tokenizer.h tokenizers/gpt_oss_tokenizer.cpp tokenizers/gpt_oss_tokenizer.h tokenizers/mistral_tokenizer.cpp tokenizers/mistral_tokenizer.h tokenizers/qwen2_tokenizer.cpp tokenizers/qwen2_tokenizer.h tokenizers/t5_unigram_tokenizer.cpp tokenizers/t5_unigram_tokenizer.h tokenizers/tokenize_util.cpp tokenizers/tokenize_util.h tokenizers/tokenizer.cpp tokenizers/tokenizer.h tokenizers/vocab/vocab.h ucache.hpp unet.hpp upscaler.cpp upscaler.h util.cpp util.h vae.hpp wan.hpp z_image.hpp
SDCPP_MAIN_BASENAMES := common/common.cpp common/common.h common/log.cpp common/log.h common/media_io.cpp common/media_io.cpp common/media_io.h common/resource_owners.hpp convert.cpp image_metadata.cpp main.cpp tokenizers/vocab/clip_merges.hpp tokenizers/vocab/gemma_merges.hpp tokenizers/vocab/gemma_vocab.hpp tokenizers/vocab/gpt_oss_merges.hpp tokenizers/vocab/gpt_oss_vocab.hpp tokenizers/vocab/mistral_merges.hpp tokenizers/vocab/mistral_vocab.hpp tokenizers/vocab/qwen_merges.hpp tokenizers/vocab/t5.hpp tokenizers/vocab/umt5.hpp tokenizers/vocab/vocab.cpp version.cpp
otherarch/sdcpp/anima.hpp
+69 -55
View File
@@ -1,6 +1,7 @@
#ifndef __ANIMA_HPP__
#define __ANIMA_HPP__
#include <algorithm>
#include <cmath>
#include <memory>
#include <utility>
@@ -14,6 +15,47 @@
namespace Anima {
constexpr int ANIMA_GRAPH_SIZE = 65536;
struct AnimaConfig {
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t hidden_size = 2048;
int64_t text_embed_dim = 1024;
int64_t num_heads = 16;
int64_t head_dim = 128;
int patch_size = 2;
int64_t num_layers = 28;
std::vector<int> axes_dim = {44, 42, 42};
int theta = 10000;
static AnimaConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
AnimaConfig config;
int64_t detected_layers = 0;
std::string layer_tag = prefix.empty() ? "blocks." : prefix + ".blocks.";
for (const auto& [name, _] : tensor_storage_map) {
size_t pos = name.find(layer_tag);
if (pos == std::string::npos) {
continue;
}
size_t start = pos + layer_tag.size();
size_t end = name.find('.', start);
if (end == std::string::npos) {
continue;
}
int64_t layer_id = atoll(name.substr(start, end - start).c_str());
detected_layers = std::max(detected_layers, layer_id + 1);
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
LOG_DEBUG("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64,
config.num_layers,
config.hidden_size,
config.num_heads,
config.head_dim);
}
return config;
}
};
__STATIC_INLINE__ ggml_tensor* apply_gate(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* gate) {
@@ -418,31 +460,22 @@ namespace Anima {
struct AnimaNet : public GGMLBlock {
public:
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t hidden_size = 2048;
int64_t text_embed_dim = 1024;
int64_t num_heads = 16;
int64_t head_dim = 128;
int patch_size = 2;
int64_t num_layers = 28;
std::vector<int> axes_dim = {44, 42, 42};
int theta = 10000;
AnimaConfig config;
public:
AnimaNet() = default;
explicit AnimaNet(int64_t num_layers)
: num_layers(num_layers) {
blocks["x_embedder"] = std::make_shared<XEmbedder>((in_channels + 1) * patch_size * patch_size, hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(hidden_size, hidden_size * 3);
blocks["t_embedding_norm"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
for (int i = 0; i < num_layers; i++) {
blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(hidden_size,
text_embed_dim,
num_heads,
head_dim);
explicit AnimaNet(AnimaConfig config)
: config(config) {
blocks["x_embedder"] = std::make_shared<XEmbedder>((config.in_channels + 1) * config.patch_size * config.patch_size, config.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(config.hidden_size, config.hidden_size * 3);
blocks["t_embedding_norm"] = std::make_shared<RMSNorm>(config.hidden_size, 1e-6f);
for (int i = 0; i < config.num_layers; i++) {
blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(config.hidden_size,
config.text_embed_dim,
config.num_heads,
config.head_dim);
}
blocks["final_layer"] = std::make_shared<FinalLayer>(hidden_size, patch_size, out_channels);
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
}
@@ -469,11 +502,11 @@ namespace Anima {
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W]
x = DiT::pad_and_patchify(ctx, x, patch_size, patch_size); // [N, h*w, (C+1)*ph*pw]
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
x = x_embedder->forward(ctx, x);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(hidden_size));
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
auto temb = t_embedder->forward(ctx, timestep_proj);
auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj);
@@ -505,7 +538,7 @@ namespace Anima {
sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb");
sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context");
for (int i = 0; i < num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
@@ -513,7 +546,7 @@ namespace Anima {
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, false); // [N, C, H, W]
x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false); // [N, C, H, W]
return x;
}
@@ -524,35 +557,16 @@ namespace Anima {
std::vector<float> image_pe_vec;
std::vector<float> adapter_q_pe_vec;
std::vector<float> adapter_k_pe_vec;
AnimaConfig config;
AnimaNet net;
AnimaRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "model.diffusion_model")
: DiffusionModelRunner(backend, params_backend, prefix) {
int64_t num_layers = 0;
std::string layer_tag = prefix + ".net.blocks.";
for (const auto& kv : tensor_storage_map) {
const std::string& tensor_name = kv.first;
size_t pos = tensor_name.find(layer_tag);
if (pos == std::string::npos) {
continue;
}
size_t start = pos + layer_tag.size();
size_t end = tensor_name.find('.', start);
if (end == std::string::npos) {
continue;
}
int64_t layer_id = atoll(tensor_name.substr(start, end - start).c_str());
num_layers = std::max(num_layers, layer_id + 1);
}
if (num_layers <= 0) {
num_layers = 28;
}
LOG_INFO("anima net layers: %" PRId64, num_layers);
net = AnimaNet(num_layers);
: DiffusionModelRunner(backend, params_backend, prefix),
config(AnimaConfig::detect_from_weights(tensor_storage_map, prefix + ".net")) {
net = AnimaNet(config);
net.init(params_ctx, tensor_storage_map, prefix + ".net");
}
@@ -623,22 +637,22 @@ namespace Anima {
GGML_ASSERT(x->ne[3] == 1);
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
int64_t pad_h = (net.patch_size - x->ne[1] % net.patch_size) % net.patch_size;
int64_t pad_w = (net.patch_size - x->ne[0] % net.patch_size) % net.patch_size;
int64_t pad_h = (config.patch_size - x->ne[1] % config.patch_size) % config.patch_size;
int64_t pad_w = (config.patch_size - x->ne[0] % config.patch_size) % config.patch_size;
int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w;
image_pe_vec = gen_anima_image_pe_vec(1,
static_cast<int>(h_pad),
static_cast<int>(w_pad),
static_cast<int>(net.patch_size),
net.theta,
net.axes_dim,
static_cast<int>(config.patch_size),
config.theta,
config.axes_dim,
4.0f,
4.0f,
1.0f);
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (net.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, net.head_dim / 2, image_pos_len);
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
set_backend_tensor_data(image_pe, image_pe_vec.data());
ggml_tensor* adapter_q_pe = nullptr;
otherarch/sdcpp/common/common.cpp
+51 -5
View File
@@ -35,6 +35,22 @@ const char* const modes_str[] = {
"metadata",
};
static sd_vae_format_t str_to_vae_format(const std::string& value) {
if (value == "auto") {
return SD_VAE_FORMAT_AUTO;
}
if (value == "flux") {
return SD_VAE_FORMAT_FLUX;
}
if (value == "sd3") {
return SD_VAE_FORMAT_SD3;
}
if (value == "flux2") {
return SD_VAE_FORMAT_FLUX2;
}
return SD_VAE_FORMAT_COUNT;
}
#if defined(_WIN32)
static std::string utf16_to_utf8(const std::wstring& wstr) {
if (wstr.empty())
@@ -229,6 +245,7 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
return false;
};
bool valid = false;
for (int i = 1; i < argc; i++) {
arg = argv[i];
bool found_arg = false;
@@ -271,7 +288,7 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
break;
if (match_and_apply(options.manual_options, [&](auto& option) {
int ret = option.cb(argc, argv, i);
int ret = option.cb(argc, argv, i, valid);
if (ret < 0) {
invalid_arg = true;
return;
@@ -283,7 +300,9 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
}
if (invalid_arg) {
LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
if (!valid) {
LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
}
return false;
}
if (!found_arg) {
@@ -340,6 +359,10 @@ ArgOptions SDContextParams::get_options() {
"--high-noise-diffusion-model",
"path to the standalone high noise diffusion model",
&high_noise_diffusion_model_path},
{"",
"--uncond-diffusion-model",
"path to the standalone unconditional diffusion model, currently used by Ideogram4 CFG",
&uncond_diffusion_model_path},
{"",
"--embeddings-connectors",
"path to LTXAV embeddings connectors",
@@ -348,6 +371,10 @@ ArgOptions SDContextParams::get_options() {
"--vae",
"path to standalone vae model",
&vae_path},
{"",
"--vae-format",
"VAE latent format override: auto, flux, sd3, or flux2 (default: auto)",
&vae_format},
{"",
"--audio-vae",
"path to standalone LTX audio vae model",
@@ -418,6 +445,10 @@ ArgOptions SDContextParams::get_options() {
};
options.bool_options = {
{"",
"--stream-layers",
"enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)",
true, &stream_layers},
{"",
"--force-sdxl-vae-conv-scale",
"force use of conv scale on sdxl vae",
@@ -639,6 +670,11 @@ bool SDContextParams::validate(SDMode mode) {
}
}
if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) {
LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', or 'flux2'");
return false;
}
return true;
}
@@ -677,8 +713,10 @@ std::string SDContextParams::to_string() const {
<< " llm_vision_path: \"" << llm_vision_path << "\",\n"
<< " diffusion_model_path: \"" << diffusion_model_path << "\",\n"
<< " high_noise_diffusion_model_path: \"" << high_noise_diffusion_model_path << "\",\n"
<< " uncond_diffusion_model_path: \"" << uncond_diffusion_model_path << "\",\n"
<< " embeddings_connectors_path: \"" << embeddings_connectors_path << "\",\n"
<< " vae_path: \"" << vae_path << "\",\n"
<< " vae_format: \"" << vae_format << "\",\n"
<< " audio_vae_path: \"" << audio_vae_path << "\",\n"
<< " taesd_path: \"" << taesd_path << "\",\n"
<< " esrgan_path: \"" << esrgan_path << "\",\n"
@@ -694,6 +732,7 @@ std::string SDContextParams::to_string() const {
<< " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
<< " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
<< " max_vram: " << max_vram << ",\n"
<< " stream_layers: " << (stream_layers ? "true" : "false") << ",\n"
<< " backend: \"" << backend << "\",\n"
<< " params_backend: \"" << params_backend << "\",\n"
<< " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
@@ -738,6 +777,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
llm_vision_path.c_str(),
diffusion_model_path.c_str(),
high_noise_diffusion_model_path.c_str(),
uncond_diffusion_model_path.c_str(),
embeddings_connectors_path.c_str(),
vae_path.c_str(),
audio_vae_path.c_str(),
@@ -772,7 +812,9 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
chroma_use_t5_mask,
chroma_t5_mask_pad,
qwen_image_zero_cond_t,
str_to_vae_format(vae_format),
max_vram,
stream_layers,
backend.c_str(),
params_backend.c_str(),
};
@@ -833,7 +875,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler args, key=value list. lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma",
"extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma",
&extra_sample_args},
{"",
"--extra-tiling-args",
@@ -913,7 +955,7 @@ ArgOptions SDGenerationParams::get_options() {
&sample_params.guidance.txt_cfg},
{"",
"--img-cfg-scale",
"image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)",
"image guidance scale for inpaint or image edit models: (default: same as --cfg-scale)",
&sample_params.guidance.img_cfg},
{"",
"--guidance",
@@ -945,7 +987,7 @@ ArgOptions SDGenerationParams::get_options() {
&high_noise_sample_params.guidance.txt_cfg},
{"",
"--high-noise-img-cfg-scale",
"(high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)",
"(high noise) image guidance scale for inpaint or image edit models (default: same as --cfg-scale)",
&high_noise_sample_params.guidance.img_cfg},
{"",
"--high-noise-guidance",
@@ -2486,6 +2528,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
set_json_basename_if_not_empty(models, "llm_vision", ctx_params.llm_vision_path);
set_json_basename_if_not_empty(models, "diffusion_model", ctx_params.diffusion_model_path);
set_json_basename_if_not_empty(models, "high_noise_diffusion_model", ctx_params.high_noise_diffusion_model_path);
set_json_basename_if_not_empty(models, "uncond_diffusion_model", ctx_params.uncond_diffusion_model_path);
set_json_basename_if_not_empty(models, "vae", ctx_params.vae_path);
set_json_basename_if_not_empty(models, "taesd", ctx_params.taesd_path);
set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path);
@@ -2653,6 +2696,9 @@ std::string get_image_params(const SDContextParams& ctx_params,
if (!ctx_params.diffusion_model_path.empty()) {
parameter_string += "Unet: " + sd_basename(ctx_params.diffusion_model_path) + ", ";
}
if (!ctx_params.uncond_diffusion_model_path.empty()) {
parameter_string += "Uncond Unet: " + sd_basename(ctx_params.uncond_diffusion_model_path) + ", ";
}
if (!ctx_params.vae_path.empty()) {
parameter_string += "VAE: " + sd_basename(ctx_params.vae_path) + ", ";
}
otherarch/sdcpp/common/common.h
+35 -1
View File
@@ -56,11 +56,42 @@ struct BoolOption {
bool* target;
};
struct ManualFunction {
std::function<int(int, const char**, int, bool&)> _func;
ManualFunction() = default;
ManualFunction(std::function<int(int argc, const char** argv, int index, bool& valid)> func)
: _func(std::move(func)) {
}
template <typename F>
ManualFunction(F func)
: _func(make_function(func)) {
}
int operator()(int argc, const char** argv, int index, bool& valid) const {
return _func(argc, argv, index, valid);
}
private:
template <typename F>
static std::function<int(int, const char**, int, bool&)> make_function(F func) {
if constexpr (std::is_invocable_v<F, int, const char**, int, bool&>) {
return func;
} else {
return [func](int argc, const char** argv, int index, bool&) {
return func(argc, argv, index);
};
}
}
};
struct ManualOption {
std::string short_name;
std::string long_name;
std::string desc;
std::function<int(int argc, const char** argv, int index)> cb;
ManualFunction cb;
};
struct ArgOptions {
@@ -92,8 +123,10 @@ struct SDContextParams {
std::string llm_vision_path;
std::string diffusion_model_path;
std::string high_noise_diffusion_model_path;
std::string uncond_diffusion_model_path;
std::string embeddings_connectors_path;
std::string vae_path;
std::string vae_format = "auto";
std::string audio_vae_path;
std::string taesd_path;
std::string esrgan_path;
@@ -112,6 +145,7 @@ struct SDContextParams {
rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
bool offload_params_to_cpu = false;
float max_vram = 0.f;
bool stream_layers = false;
std::string backend;
std::string params_backend;
bool enable_mmap = false;
otherarch/sdcpp/conditioner.hpp
+118 -11
View File
@@ -118,6 +118,7 @@ public:
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual size_t get_params_buffer_size() = 0;
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_flash_attention_enabled(bool enabled) = 0;
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
virtual std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(int n_threads,
@@ -210,6 +211,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
text_model->set_stream_layers_enabled(enabled);
if (sd_version_is_sdxl(version)) {
text_model2->set_stream_layers_enabled(enabled);
}
}
void set_flash_attention_enabled(bool enabled) override {
text_model->set_flash_attention_enabled(enabled);
if (sd_version_is_sdxl(version)) {
@@ -843,6 +851,18 @@ struct SD3CLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_stream_layers_enabled(enabled);
}
if (clip_g) {
clip_g->set_stream_layers_enabled(enabled);
}
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_flash_attention_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_flash_attention_enabled(enabled);
@@ -1171,7 +1191,6 @@ struct FluxCLIPEmbedder : public Conditioner {
return true;
}
void free_params_buffer() override {
if (clip_l) {
clip_l->free_params_buffer();
@@ -1201,6 +1220,15 @@ struct FluxCLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_stream_layers_enabled(enabled);
}
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_flash_attention_enabled(bool enabled) override {
if (clip_l) {
clip_l->set_flash_attention_enabled(enabled);
@@ -1435,6 +1463,12 @@ struct T5CLIPEmbedder : public Conditioner {
}
}
void set_stream_layers_enabled(bool enabled) override {
if (t5) {
t5->set_stream_layers_enabled(enabled);
}
}
void set_flash_attention_enabled(bool enabled) override {
if (t5) {
t5->set_flash_attention_enabled(enabled);
@@ -1601,8 +1635,8 @@ struct AnimaConditioner : public Conditioner {
bool alloc_params_buffer() override {
if (!llm->alloc_params_buffer()) {
return false;
}
return false;
}
return true;
}
@@ -1618,6 +1652,10 @@ struct AnimaConditioner : public Conditioner {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
}
@@ -1719,6 +1757,10 @@ struct LLMEmbedder : public Conditioner {
arch = LLM::LLMArch::MINISTRAL_3_3B;
} else if (sd_version_is_lens(version)) {
arch = LLM::LLMArch::GPT_OSS_20B;
} else if (sd_version_is_pid(version)) {
arch = LLM::LLMArch::GEMMA2_2B;
} else if (sd_version_is_ideogram4(version)) {
arch = LLM::LLMArch::QWEN3_VL;
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
arch = LLM::LLMArch::QWEN3;
}
@@ -1726,6 +1768,8 @@ struct LLMEmbedder : public Conditioner {
tokenizer = std::make_shared<MistralTokenizer>();
} else if (arch == LLM::LLMArch::GPT_OSS_20B) {
tokenizer = std::make_shared<GPTOSSTokenizer>();
} else if (arch == LLM::LLMArch::GEMMA2_2B) {
tokenizer = std::make_shared<Gemma2Tokenizer>();
} else {
tokenizer = std::make_shared<Qwen2Tokenizer>();
}
@@ -1743,7 +1787,7 @@ struct LLMEmbedder : public Conditioner {
bool alloc_params_buffer() override {
if (!llm->alloc_params_buffer()) {
return false;
return false;
}
return true;
}
@@ -1762,6 +1806,10 @@ struct LLMEmbedder : public Conditioner {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
}
@@ -1847,12 +1895,16 @@ struct LLMEmbedder : public Conditioner {
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
sd::Tensor<float> attention_mask;
if (!mask.empty()) {
attention_mask = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
attention_mask = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
const float masked_attention_value = -std::numeric_limits<float>::max() / 4.0f;
for (size_t i1 = 0; i1 < mask.size(); ++i1) {
for (size_t i0 = 0; i0 < mask.size(); ++i0) {
float value = 0.0f;
if (mask[i0] == 0.0f || i0 > i1) {
value = -INFINITY;
if (mask[i0] == 0.0f) {
value += masked_attention_value;
}
if (i0 > i1) {
value += masked_attention_value;
}
attention_mask[static_cast<int64_t>(i0 + mask.size() * i1)] = value;
}
@@ -1919,7 +1971,7 @@ struct LLMEmbedder : public Conditioner {
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size;
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor);
@@ -1990,7 +2042,7 @@ struct LLMEmbedder : public Conditioner {
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size;
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor);
@@ -2051,6 +2103,14 @@ struct LLMEmbedder : public Conditioner {
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "[/INST]";
} else if (sd_version_is_ideogram4(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 36};
prompt = "<|im_start|>user\n";
prompt += conditioner_params.text;
prompt += "<|im_end|>\n<|im_start|>assistant\n";
prompt_attn_range = {0, 0};
} else if (sd_version_is_ernie_image(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {25}; // -2
@@ -2126,6 +2186,53 @@ struct LLMEmbedder : public Conditioner {
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
} else if (sd_version_is_pid(version)) {
constexpr int pixeldit_max_length = 300;
const std::string chi_prompt =
"Given a user prompt, generate an \"Enhanced prompt\" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:\n"
"- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.\n"
"- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.\n"
"Here are examples of how to transform or refine prompts:\n"
"- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.\n"
"- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n"
"Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:\n"
"User Prompt: ";
auto chi_tokens = std::get<0>(tokenize(chi_prompt, {0, 0}));
size_t num_chi_tokens = chi_tokens.size();
max_length = (int)num_chi_tokens + pixeldit_max_length - 2;
min_length = max_length;
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += " " + conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
auto hidden_states = encode_prompt(n_threads,
prompt,
prompt_attn_range,
min_length,
0,
image_embeds,
out_layers,
0,
false,
max_length);
GGML_ASSERT(!hidden_states.empty());
if (hidden_states.shape()[1] > pixeldit_max_length) {
auto bos = sd::ops::slice(hidden_states, 1, 0, 1);
auto tail = sd::ops::slice(hidden_states,
1,
hidden_states.shape()[1] - (pixeldit_max_length - 1),
hidden_states.shape()[1]);
hidden_states = sd::ops::concat(bos, tail, 1);
}
int64_t t1 = ggml_time_ms();
LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
SDCondition result;
result.c_crossattn = std::move(hidden_states);
return result;
} else {
GGML_ABORT("unknown version %d", version);
}
@@ -2268,10 +2375,10 @@ struct LTXAVEmbedder : public Conditioner {
bool alloc_params_buffer() override {
if (!llm->alloc_params_buffer()) {
return false;
return false;
}
if (!projector->alloc_params_buffer()) {
return false;
return false;
}
return true;
}
otherarch/sdcpp/convert.cpp
+2 -2
View File
@@ -8,7 +8,7 @@
#include "model_io/safetensors_io.h"
#include "util.h"
#include "ggml-cpu.h"
#include "ggml_extend_backend.h"
static ggml_type get_export_tensor_type(ModelLoader& model_loader,
const TensorStorage& tensor_storage,
@@ -103,7 +103,7 @@ bool convert(const char* input_path,
bool output_is_safetensors = ends_with(output_path, ".safetensors");
TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
auto backend = ggml_backend_cpu_init();
auto backend = sd_backend_cpu_init();
size_t mem_size = 1 * 1024 * 1024; // for padding
mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
mem_size += model_loader.get_params_mem_size(backend, type);
otherarch/sdcpp/denoiser.hpp
+17 -15
View File
@@ -514,8 +514,6 @@ struct LTX2Scheduler : SigmaScheduler {
if (!parse_strict_bool(value, stretch)) {
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
}
} else {
LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
}
}
}
@@ -1238,20 +1236,26 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
for (const auto& [key, value] : extra_sample_args) {
float parsed = 0.0f;
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
if (key == "noise_clip_std") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
args.noise_clip_std = parsed;
} else if (key == "noise_scale_start") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
args.noise_scale_start = parsed;
noise_scale_start_was_set = true;
} else if (key == "noise_scale_end") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
args.noise_scale_end = parsed;
noise_scale_end_was_set = true;
} else {
LOG_WARN("ignoring unknown lcm extra sample arg '%s'", key.c_str());
}
}
@@ -1795,16 +1799,14 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
float ge_gamma = 2.0f;
for (const auto& [key, value] : extra_sample_args) {
float parsed = 0.0f;
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
if (key == "gamma") {
float parsed = 0.0f;
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
LOG_DEBUG("setting euler_ge gamma to %.2f", parsed);
ge_gamma = parsed;
} else {
LOG_WARN("ignoring unknown euler_ge extra sample arg '%s'", key.c_str());
}
}
otherarch/sdcpp/ernie_image.hpp
+105 -99
View File
@@ -13,6 +13,76 @@
namespace ErnieImage {
constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960;
struct ErnieImageConfig {
int64_t hidden_size = 4096;
int64_t num_heads = 32;
int64_t num_layers = 36;
int64_t ffn_hidden_size = 12288;
int64_t in_channels = 128;
int64_t out_channels = 128;
int patch_size = 1;
int64_t text_in_dim = 3072;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int axes_dim_sum = 128;
float eps = 1e-6f;
static ErnieImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
ErnieImageConfig config;
config.num_layers = 0;
int64_t detected_head_dim = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
config.patch_size = static_cast<int>(tensor_storage.ne[0]);
config.in_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
config.text_in_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
detected_head_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
config.ffn_hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
int64_t out_dim = tensor_storage.ne[1];
int64_t patch_area = config.patch_size * config.patch_size;
config.out_channels = out_dim / patch_area;
}
size_t pos = name.find("layers.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
}
if (config.num_layers == 0) {
config.num_layers = 36;
}
if (detected_head_dim > 0) {
config.num_heads = config.hidden_size / detected_head_dim;
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
LOG_DEBUG("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
config.num_layers,
config.hidden_size,
config.num_heads,
config.ffn_hidden_size,
config.in_channels,
config.out_channels);
return config;
}
};
__STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
ggml_tensor* timesteps,
int dim,
@@ -208,51 +278,36 @@ namespace ErnieImage {
}
};
struct ErnieImageParams {
int64_t hidden_size = 4096;
int64_t num_heads = 32;
int64_t num_layers = 36;
int64_t ffn_hidden_size = 12288;
int64_t in_channels = 128;
int64_t out_channels = 128;
int patch_size = 1;
int64_t text_in_dim = 3072;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int axes_dim_sum = 128;
float eps = 1e-6f;
};
class ErnieImageModel : public GGMLBlock {
public:
ErnieImageParams params;
ErnieImageConfig config;
ErnieImageModel() = default;
ErnieImageModel(ErnieImageParams params)
: params(params) {
blocks["x_embedder.proj"] = std::make_shared<Conv2d>(params.in_channels,
params.hidden_size,
std::pair<int, int>{params.patch_size, params.patch_size},
std::pair<int, int>{params.patch_size, params.patch_size},
ErnieImageModel(ErnieImageConfig config)
: config(config) {
blocks["x_embedder.proj"] = std::make_shared<Conv2d>(config.in_channels,
config.hidden_size,
std::pair<int, int>{config.patch_size, config.patch_size},
std::pair<int, int>{config.patch_size, config.patch_size},
std::pair<int, int>{0, 0},
std::pair<int, int>{1, 1},
true);
if (params.text_in_dim != params.hidden_size) {
blocks["text_proj"] = std::make_shared<Linear>(params.text_in_dim, params.hidden_size, false);
if (config.text_in_dim != config.hidden_size) {
blocks["text_proj"] = std::make_shared<Linear>(config.text_in_dim, config.hidden_size, false);
}
blocks["time_embedding"] = std::make_shared<Qwen::TimestepEmbedding>(params.hidden_size, params.hidden_size);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(params.hidden_size, 6 * params.hidden_size, true);
blocks["time_embedding"] = std::make_shared<Qwen::TimestepEmbedding>(config.hidden_size, config.hidden_size);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
for (int i = 0; i < params.num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(params.hidden_size,
params.num_heads,
params.ffn_hidden_size,
params.eps);
for (int i = 0; i < config.num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(config.hidden_size,
config.num_heads,
config.ffn_hidden_size,
config.eps);
}
blocks["final_norm"] = std::make_shared<ErnieImageAdaLNContinuous>(params.hidden_size, params.eps);
blocks["final_linear"] = std::make_shared<Linear>(params.hidden_size,
params.patch_size * params.patch_size * params.out_channels,
blocks["final_norm"] = std::make_shared<ErnieImageAdaLNContinuous>(config.hidden_size, config.eps);
blocks["final_linear"] = std::make_shared<Linear>(config.hidden_size,
config.patch_size * config.patch_size * config.out_channels,
true);
}
@@ -265,12 +320,12 @@ namespace ErnieImage {
// context: [N, text_tokens, 3072]
// pe: [image_tokens + text_tokens, head_dim/2, 2, 2]
GGML_ASSERT(context != nullptr);
GGML_ASSERT(x->ne[1] % params.patch_size == 0 && x->ne[0] % params.patch_size == 0);
GGML_ASSERT(x->ne[1] % config.patch_size == 0 && x->ne[0] % config.patch_size == 0);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t Hp = H / params.patch_size;
int64_t Wp = W / params.patch_size;
int64_t Hp = H / config.patch_size;
int64_t Wp = W / config.patch_size;
int64_t n_img = Hp * Wp;
int64_t N = x->ne[3];
@@ -292,7 +347,7 @@ namespace ErnieImage {
auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1); // [N, image_tokens + text_tokens, hidden_size]
auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(params.hidden_size));
auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
auto c = time_embedding->forward(ctx, sample); // [N, hidden_size]
auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size]
@@ -305,7 +360,7 @@ namespace ErnieImage {
temb.push_back(ggml_reshape_3d(ctx->ggml_ctx, chunk, chunk->ne[0], 1, chunk->ne[1])); // [N, 1, hidden_size]
}
for (int i = 0; i < params.num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto layer = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
hidden_states = layer->forward(ctx, hidden_states, pe, temb);
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states");
@@ -319,15 +374,15 @@ namespace ErnieImage {
patches,
Hp,
Wp,
params.patch_size,
params.patch_size,
config.patch_size,
config.patch_size,
false); // [N, out_channels, H, W]
return out;
}
};
struct ErnieImageRunner : public DiffusionModelRunner {
ErnieImageParams ernie_params;
ErnieImageConfig config;
ErnieImageModel ernie_image;
std::vector<float> pe_vec;
@@ -335,58 +390,9 @@ namespace ErnieImage {
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix) {
ernie_params.num_layers = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
ernie_params.patch_size = static_cast<int>(tensor_storage.ne[0]);
ernie_params.in_channels = tensor_storage.ne[2];
ernie_params.hidden_size = tensor_storage.ne[3];
} else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
ernie_params.text_in_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
int64_t head_dim = tensor_storage.ne[0];
ernie_params.num_heads = ernie_params.hidden_size / head_dim;
} else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
ernie_params.ffn_hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
int64_t out_dim = tensor_storage.ne[1];
ernie_params.out_channels = out_dim / ernie_params.patch_size / ernie_params.patch_size;
}
size_t pos = name.find("layers.");
if (pos != std::string::npos) {
std::string layer_name = name.substr(pos);
auto items = split_string(layer_name, '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > ernie_params.num_layers) {
ernie_params.num_layers = block_index + 1;
}
}
}
}
if (ernie_params.num_layers == 0) {
ernie_params.num_layers = 36;
}
ernie_params.axes_dim_sum = 0;
for (int axis_dim : ernie_params.axes_dim) {
ernie_params.axes_dim_sum += axis_dim;
}
LOG_INFO("ernie_image: layers = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64
", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
ernie_params.num_layers,
ernie_params.hidden_size,
ernie_params.num_heads,
ernie_params.ffn_hidden_size,
ernie_params.in_channels,
ernie_params.out_channels);
ernie_image = ErnieImageModel(ernie_params);
: DiffusionModelRunner(backend, params_backend, prefix),
config(ErnieImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
ernie_image = ErnieImageModel(config);
ernie_image.init(params_ctx, tensor_storage_map, prefix);
}
@@ -410,15 +416,15 @@ namespace ErnieImage {
pe_vec = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
ernie_params.patch_size,
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ernie_params.theta,
config.theta,
circular_y_enabled,
circular_x_enabled,
ernie_params.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / ernie_params.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, ernie_params.axes_dim_sum, 1, pos_len, 2);
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
otherarch/sdcpp/flux.hpp
+240 -227
View File
@@ -13,6 +13,155 @@
namespace Flux {
struct ChromaRadianceConfig {
int64_t nerf_hidden_size = 64;
int nerf_mlp_ratio = 4;
int nerf_depth = 4;
int nerf_max_freqs = 8;
bool use_x0 = false;
bool fake_patch_size_x2 = false;
};
struct FluxConfig {
SDVersion version = VERSION_FLUX;
bool is_chroma = false;
int patch_size = 2;
int64_t in_channels = 64;
int64_t out_channels = 64;
int64_t vec_in_dim = 768;
int64_t context_in_dim = 4096;
int64_t hidden_size = 3072;
float mlp_ratio = 4.0f;
int num_heads = 24;
int depth = 19;
int depth_single_blocks = 38;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
int theta = 10000;
bool qkv_bias = true;
bool guidance_embed = true;
int64_t in_dim = 64;
bool disable_bias = false;
bool share_modulation = false;
bool semantic_txt_norm = false;
bool use_yak_mlp = false;
bool use_mlp_silu_act = false;
float ref_index_scale = 1.f;
ChromaRadianceConfig chroma_radiance_params;
static FluxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
SDVersion version = VERSION_FLUX) {
FluxConfig config;
config.version = version;
config.guidance_embed = false;
config.depth = 0;
config.depth_single_blocks = 0;
if (version == VERSION_FLUX_FILL) {
config.in_channels = 384;
} else if (version == VERSION_FLUX_CONTROLS) {
config.in_channels = 128;
} else if (version == VERSION_FLEX_2) {
config.in_channels = 196;
} else if (version == VERSION_CHROMA_RADIANCE) {
config.in_channels = 3;
config.patch_size = 16;
} else if (version == VERSION_OVIS_IMAGE) {
config.semantic_txt_norm = true;
config.use_yak_mlp = true;
config.vec_in_dim = 0;
} else if (sd_version_is_flux2(version)) {
config.in_channels = 128;
config.patch_size = 1;
config.out_channels = 128;
config.mlp_ratio = 3.f;
config.theta = 2000;
config.axes_dim = {32, 32, 32, 32};
config.vec_in_dim = 0;
config.qkv_bias = false;
config.disable_bias = true;
config.share_modulation = true;
config.ref_index_scale = 10.f;
config.use_mlp_silu_act = true;
} else if (sd_version_is_longcat(version)) {
config.context_in_dim = 3584;
config.vec_in_dim = 0;
}
int64_t head_dim = 0;
int64_t actual_radiance_patch_size = -1;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (name.find("guidance_in.in_layer.weight") != std::string::npos) {
config.guidance_embed = true;
}
if (name.find("__x0__") != std::string::npos) {
LOG_DEBUG("using x0 prediction");
config.chroma_radiance_params.use_x0 = true;
}
if (name.find("__32x32__") != std::string::npos) {
LOG_DEBUG("using patch size 32");
config.patch_size = 32;
}
if (name.find("img_in_patch.weight") != std::string::npos) {
actual_radiance_patch_size = tensor_storage.ne[0];
LOG_DEBUG("actual radiance patch size: %" PRId64, actual_radiance_patch_size);
}
if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
config.is_chroma = true;
}
size_t db = name.find("double_blocks.");
if (db != std::string::npos) {
std::string block_name = name.substr(db);
int block_depth = atoi(block_name.substr(14, block_name.find(".", 14)).c_str());
if (block_depth + 1 > config.depth) {
config.depth = block_depth + 1;
}
}
size_t sb = name.find("single_blocks.");
if (sb != std::string::npos) {
std::string block_name = name.substr(sb);
int block_depth = atoi(block_name.substr(14, block_name.find(".", 14)).c_str());
if (block_depth + 1 > config.depth_single_blocks) {
config.depth_single_blocks = block_depth + 1;
}
}
if (ends_with(name, "txt_in.weight")) {
config.context_in_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
}
if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) {
head_dim = tensor_storage.ne[0];
}
if (ends_with(name, "double_blocks.0.txt_attn.norm.key_norm.scale")) {
head_dim = tensor_storage.ne[0];
}
}
if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != config.patch_size) {
GGML_ASSERT(config.patch_size == 2 * actual_radiance_patch_size);
LOG_DEBUG("using fake x2 patch size");
config.chroma_radiance_params.fake_patch_size_x2 = true;
}
if (head_dim > 0) {
config.num_heads = static_cast<int>(config.hidden_size / head_dim);
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
LOG_DEBUG("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d",
config.depth,
config.depth_single_blocks,
config.guidance_embed ? "true" : "false",
config.context_in_dim,
config.hidden_size,
config.num_heads);
return config;
}
};
struct MLPEmbedder : public UnaryBlock {
public:
MLPEmbedder(int64_t in_dim, int64_t hidden_dim, bool bias = true) {
@@ -723,127 +872,90 @@ namespace Flux {
}
};
struct ChromaRadianceParams {
int64_t nerf_hidden_size = 64;
int nerf_mlp_ratio = 4;
int nerf_depth = 4;
int nerf_max_freqs = 8;
bool use_x0 = false;
bool fake_patch_size_x2 = false;
};
struct FluxParams {
SDVersion version = VERSION_FLUX;
bool is_chroma = false;
int patch_size = 2;
int64_t in_channels = 64;
int64_t out_channels = 64;
int64_t vec_in_dim = 768;
int64_t context_in_dim = 4096;
int64_t hidden_size = 3072;
float mlp_ratio = 4.0f;
int num_heads = 24;
int depth = 19;
int depth_single_blocks = 38;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
int theta = 10000;
bool qkv_bias = true;
bool guidance_embed = true;
int64_t in_dim = 64;
bool disable_bias = false;
bool share_modulation = false;
bool semantic_txt_norm = false;
bool use_yak_mlp = false;
bool use_mlp_silu_act = false;
float ref_index_scale = 1.f;
ChromaRadianceParams chroma_radiance_params;
};
struct Flux : public GGMLBlock {
public:
FluxParams params;
FluxConfig config;
Flux() {}
Flux(FluxParams params)
: params(params) {
if (params.version == VERSION_CHROMA_RADIANCE) {
std::pair<int, int> kernel_size = {params.patch_size, params.patch_size};
if (params.chroma_radiance_params.fake_patch_size_x2) {
kernel_size = {params.patch_size / 2, params.patch_size / 2};
Flux(FluxConfig config)
: config(config) {
if (config.version == VERSION_CHROMA_RADIANCE) {
std::pair<int, int> kernel_size = {config.patch_size, config.patch_size};
if (config.chroma_radiance_params.fake_patch_size_x2) {
kernel_size = {config.patch_size / 2, config.patch_size / 2};
}
std::pair<int, int> stride = kernel_size;
blocks["img_in_patch"] = std::make_shared<Conv2d>(params.in_channels,
params.hidden_size,
blocks["img_in_patch"] = std::make_shared<Conv2d>(config.in_channels,
config.hidden_size,
kernel_size,
stride);
} else {
blocks["img_in"] = std::make_shared<Linear>(params.in_channels, params.hidden_size, !params.disable_bias);
blocks["img_in"] = std::make_shared<Linear>(config.in_channels, config.hidden_size, !config.disable_bias);
}
if (params.is_chroma) {
blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(params.in_dim, params.hidden_size);
if (config.is_chroma) {
blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(config.in_dim, config.hidden_size);
} else {
blocks["time_in"] = std::make_shared<MLPEmbedder>(256, params.hidden_size, !params.disable_bias);
if (params.vec_in_dim > 0) {
blocks["vector_in"] = std::make_shared<MLPEmbedder>(params.vec_in_dim, params.hidden_size, !params.disable_bias);
blocks["time_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
if (config.vec_in_dim > 0) {
blocks["vector_in"] = std::make_shared<MLPEmbedder>(config.vec_in_dim, config.hidden_size, !config.disable_bias);
}
if (params.guidance_embed) {
blocks["guidance_in"] = std::make_shared<MLPEmbedder>(256, params.hidden_size, !params.disable_bias);
if (config.guidance_embed) {
blocks["guidance_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
}
}
if (params.semantic_txt_norm) {
blocks["txt_norm"] = std::make_shared<RMSNorm>(params.context_in_dim);
if (config.semantic_txt_norm) {
blocks["txt_norm"] = std::make_shared<RMSNorm>(config.context_in_dim);
}
blocks["txt_in"] = std::make_shared<Linear>(params.context_in_dim, params.hidden_size, !params.disable_bias);
blocks["txt_in"] = std::make_shared<Linear>(config.context_in_dim, config.hidden_size, !config.disable_bias);
for (int i = 0; i < params.depth; i++) {
blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamBlock>(params.hidden_size,
params.num_heads,
params.mlp_ratio,
for (int i = 0; i < config.depth; i++) {
blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamBlock>(config.hidden_size,
config.num_heads,
config.mlp_ratio,
i,
params.qkv_bias,
params.is_chroma,
params.share_modulation,
!params.disable_bias,
params.use_yak_mlp,
params.use_mlp_silu_act);
config.qkv_bias,
config.is_chroma,
config.share_modulation,
!config.disable_bias,
config.use_yak_mlp,
config.use_mlp_silu_act);
}
for (int i = 0; i < params.depth_single_blocks; i++) {
blocks["single_blocks." + std::to_string(i)] = std::make_shared<SingleStreamBlock>(params.hidden_size,
params.num_heads,
params.mlp_ratio,
for (int i = 0; i < config.depth_single_blocks; i++) {
blocks["single_blocks." + std::to_string(i)] = std::make_shared<SingleStreamBlock>(config.hidden_size,
config.num_heads,
config.mlp_ratio,
i,
0.f,
params.is_chroma,
params.share_modulation,
!params.disable_bias,
params.use_yak_mlp,
params.use_mlp_silu_act);
config.is_chroma,
config.share_modulation,
!config.disable_bias,
config.use_yak_mlp,
config.use_mlp_silu_act);
}
if (params.version == VERSION_CHROMA_RADIANCE) {
blocks["nerf_image_embedder"] = std::make_shared<NerfEmbedder>(params.in_channels,
params.chroma_radiance_params.nerf_hidden_size,
params.chroma_radiance_params.nerf_max_freqs);
if (config.version == VERSION_CHROMA_RADIANCE) {
blocks["nerf_image_embedder"] = std::make_shared<NerfEmbedder>(config.in_channels,
config.chroma_radiance_params.nerf_hidden_size,
config.chroma_radiance_params.nerf_max_freqs);
for (int i = 0; i < params.chroma_radiance_params.nerf_depth; i++) {
blocks["nerf_blocks." + std::to_string(i)] = std::make_shared<NerfGLUBlock>(params.hidden_size,
params.chroma_radiance_params.nerf_hidden_size,
params.chroma_radiance_params.nerf_mlp_ratio);
for (int i = 0; i < config.chroma_radiance_params.nerf_depth; i++) {
blocks["nerf_blocks." + std::to_string(i)] = std::make_shared<NerfGLUBlock>(config.hidden_size,
config.chroma_radiance_params.nerf_hidden_size,
config.chroma_radiance_params.nerf_mlp_ratio);
}
blocks["nerf_final_layer_conv"] = std::make_shared<NerfFinalLayerConv>(params.chroma_radiance_params.nerf_hidden_size,
params.in_channels);
blocks["nerf_final_layer_conv"] = std::make_shared<NerfFinalLayerConv>(config.chroma_radiance_params.nerf_hidden_size,
config.in_channels);
} else {
blocks["final_layer"] = std::make_shared<LastLayer>(params.hidden_size, 1, params.out_channels, params.is_chroma, !params.disable_bias);
blocks["final_layer"] = std::make_shared<LastLayer>(config.hidden_size, 1, config.out_channels, config.is_chroma, !config.disable_bias);
}
if (params.share_modulation) {
blocks["double_stream_modulation_img"] = std::make_shared<Modulation>(params.hidden_size, true, !params.disable_bias);
blocks["double_stream_modulation_txt"] = std::make_shared<Modulation>(params.hidden_size, true, !params.disable_bias);
blocks["single_stream_modulation"] = std::make_shared<Modulation>(params.hidden_size, false, !params.disable_bias);
if (config.share_modulation) {
blocks["double_stream_modulation_img"] = std::make_shared<Modulation>(config.hidden_size, true, !config.disable_bias);
blocks["double_stream_modulation_txt"] = std::make_shared<Modulation>(config.hidden_size, true, !config.disable_bias);
blocks["single_stream_modulation"] = std::make_shared<Modulation>(config.hidden_size, false, !config.disable_bias);
}
}
@@ -866,7 +978,7 @@ namespace Flux {
ggml_tensor* vec;
ggml_tensor* txt_img_mask = nullptr;
if (params.is_chroma) {
if (config.is_chroma) {
int64_t mod_index_length = 344;
auto approx = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
auto distill_timestep = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 16, 10000, 1000.f);
@@ -894,7 +1006,7 @@ namespace Flux {
} else {
auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f));
if (params.guidance_embed) {
if (config.guidance_embed) {
GGML_ASSERT(guidance != nullptr);
auto guidance_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["guidance_in"]);
// bf16 and fp16 result is different
@@ -902,7 +1014,7 @@ namespace Flux {
vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, g_in));
}
if (params.vec_in_dim > 0) {
if (config.vec_in_dim > 0) {
auto vector_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["vector_in"]);
vec = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y));
}
@@ -911,7 +1023,7 @@ namespace Flux {
std::vector<ModulationOut> ds_img_mods;
std::vector<ModulationOut> ds_txt_mods;
std::vector<ModulationOut> ss_mods;
if (params.share_modulation) {
if (config.share_modulation) {
auto double_stream_modulation_img = std::dynamic_pointer_cast<Modulation>(blocks["double_stream_modulation_img"]);
auto double_stream_modulation_txt = std::dynamic_pointer_cast<Modulation>(blocks["double_stream_modulation_txt"]);
auto single_stream_modulation = std::dynamic_pointer_cast<Modulation>(blocks["single_stream_modulation"]);
@@ -921,7 +1033,7 @@ namespace Flux {
ss_mods = single_stream_modulation->forward(ctx, vec);
}
if (params.semantic_txt_norm) {
if (config.semantic_txt_norm) {
auto semantic_txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
txt = semantic_txt_norm->forward(ctx, txt);
@@ -932,7 +1044,7 @@ namespace Flux {
sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt");
sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec");
for (int i = 0; i < params.depth; i++) {
for (int i = 0; i < config.depth; i++) {
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
continue;
}
@@ -947,8 +1059,8 @@ namespace Flux {
}
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
for (int i = 0; i < params.depth_single_blocks; i++) {
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + params.depth) != skip_layers.end()) {
for (int i = 0; i < config.depth_single_blocks; i++) {
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + config.depth) != skip_layers.end()) {
continue;
}
auto block = std::dynamic_pointer_cast<SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
@@ -999,14 +1111,14 @@ namespace Flux {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int patch_size = params.patch_size;
int patch_size = config.patch_size;
int pad_h = (patch_size - H % patch_size) % patch_size;
int pad_w = (patch_size - W % patch_size) % patch_size;
auto img = DiT::pad_to_patch_size(ctx, x, params.patch_size, params.patch_size);
auto img = DiT::pad_to_patch_size(ctx, x, config.patch_size, config.patch_size);
auto orig_img = img;
if (params.chroma_radiance_params.fake_patch_size_x2) {
if (config.chroma_radiance_params.fake_patch_size_x2) {
// It's supposed to be using GGML_SCALE_MODE_NEAREST, but this seems more stable
// Maybe the implementation of nearest-neighbor interpolation in ggml behaves differently than the one in PyTorch?
// img = F.interpolate(img, size=(H//2, W//2), mode="nearest")
@@ -1037,7 +1149,7 @@ namespace Flux {
auto nerf_hidden = ggml_reshape_2d(ctx->ggml_ctx, out, out->ne[0], out->ne[1] * out->ne[2]); // [N*num_patches, hidden_size]
auto img_dct = nerf_image_embedder->forward(ctx, nerf_pixels, dct); // [N*num_patches, patch_size*patch_size, nerf_hidden_size]
for (int i = 0; i < params.chroma_radiance_params.nerf_depth; i++) {
for (int i = 0; i < config.chroma_radiance_params.nerf_depth; i++) {
auto block = std::dynamic_pointer_cast<NerfGLUBlock>(blocks["nerf_blocks." + std::to_string(i)]);
img_dct = block->forward(ctx, img_dct, nerf_hidden);
@@ -1049,7 +1161,7 @@ namespace Flux {
out = nerf_final_layer_conv->forward(ctx, img_dct); // [N, C, H, W]
if (params.chroma_radiance_params.use_x0) {
if (config.chroma_radiance_params.use_x0) {
out = _apply_x0_residual(ctx, out, orig_img, timestep);
}
@@ -1073,14 +1185,14 @@ namespace Flux {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
int patch_size = params.patch_size;
int patch_size = config.patch_size;
int pad_h = (patch_size - H % patch_size) % patch_size;
int pad_w = (patch_size - W % patch_size) % patch_size;
auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size);
int64_t img_tokens = img->ne[1];
if (params.version == VERSION_FLUX_FILL) {
if (config.version == VERSION_FLUX_FILL) {
GGML_ASSERT(c_concat != nullptr);
ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
@@ -1089,7 +1201,7 @@ namespace Flux {
mask = DiT::pad_and_patchify(ctx, mask, patch_size, patch_size);
img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, masked, mask, 0), 0);
} else if (params.version == VERSION_FLEX_2) {
} else if (config.version == VERSION_FLEX_2) {
GGML_ASSERT(c_concat != nullptr);
ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
@@ -1100,7 +1212,7 @@ namespace Flux {
control = DiT::pad_and_patchify(ctx, control, patch_size, patch_size);
img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, masked, mask, 0), control, 0), 0);
} else if (params.version == VERSION_FLUX_CONTROLS) {
} else if (config.version == VERSION_FLUX_CONTROLS) {
GGML_ASSERT(c_concat != nullptr);
auto control = DiT::pad_and_patchify(ctx, c_concat, patch_size, patch_size);
@@ -1147,7 +1259,7 @@ namespace Flux {
// pe: (L, d_head/2, 2, 2)
// return: (N, C, H, W)
if (params.version == VERSION_CHROMA_RADIANCE) {
if (config.version == VERSION_CHROMA_RADIANCE) {
return forward_chroma_radiance(ctx,
x,
timestep,
@@ -1179,7 +1291,7 @@ namespace Flux {
struct FluxRunner : public DiffusionModelRunner {
public:
FluxParams flux_params;
FluxConfig config;
Flux flux;
std::vector<float> pe_vec;
std::vector<float> mod_index_arange_vec;
@@ -1194,114 +1306,15 @@ namespace Flux {
const std::string prefix = "",
SDVersion version = VERSION_FLUX,
bool use_mask = false)
: DiffusionModelRunner(backend, params_backend, prefix), version(version), use_mask(use_mask) {
flux_params.version = version;
flux_params.guidance_embed = false;
flux_params.depth = 0;
flux_params.depth_single_blocks = 0;
if (version == VERSION_FLUX_FILL) {
flux_params.in_channels = 384;
} else if (version == VERSION_FLUX_CONTROLS) {
flux_params.in_channels = 128;
} else if (version == VERSION_FLEX_2) {
flux_params.in_channels = 196;
} else if (version == VERSION_CHROMA_RADIANCE) {
flux_params.in_channels = 3;
flux_params.patch_size = 16;
} else if (version == VERSION_OVIS_IMAGE) {
flux_params.semantic_txt_norm = true;
flux_params.use_yak_mlp = true;
flux_params.vec_in_dim = 0;
} else if (sd_version_is_flux2(version)) {
flux_params.in_channels = 128;
flux_params.patch_size = 1;
flux_params.out_channels = 128;
flux_params.mlp_ratio = 3.f;
flux_params.theta = 2000;
flux_params.axes_dim = {32, 32, 32, 32};
flux_params.vec_in_dim = 0;
flux_params.qkv_bias = false;
flux_params.disable_bias = true;
flux_params.share_modulation = true;
flux_params.ref_index_scale = 10.f;
flux_params.use_mlp_silu_act = true;
} else if (sd_version_is_longcat(version)) {
flux_params.context_in_dim = 3584;
flux_params.vec_in_dim = 0;
}
int64_t head_dim = 0;
int64_t actual_radiance_patch_size = -1;
for (auto pair : tensor_storage_map) {
std::string tensor_name = pair.first;
if (!starts_with(tensor_name, prefix))
continue;
if (tensor_name.find("guidance_in.in_layer.weight") != std::string::npos) {
flux_params.guidance_embed = true;
}
if (tensor_name.find("__x0__") != std::string::npos) {
LOG_DEBUG("using x0 prediction");
flux_params.chroma_radiance_params.use_x0 = true;
}
if (tensor_name.find("__32x32__") != std::string::npos) {
LOG_DEBUG("using patch size 32");
flux_params.patch_size = 32;
}
if (tensor_name.find("img_in_patch.weight") != std::string::npos) {
actual_radiance_patch_size = pair.second.ne[0];
LOG_DEBUG("actual radiance patch size: %d", actual_radiance_patch_size);
}
if (tensor_name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
// Chroma
flux_params.is_chroma = true;
}
size_t db = tensor_name.find("double_blocks.");
if (db != std::string::npos) {
tensor_name = tensor_name.substr(db); // remove prefix
int block_depth = atoi(tensor_name.substr(14, tensor_name.find(".", 14)).c_str());
if (block_depth + 1 > flux_params.depth) {
flux_params.depth = block_depth + 1;
}
}
size_t sb = tensor_name.find("single_blocks.");
if (sb != std::string::npos) {
tensor_name = tensor_name.substr(sb); // remove prefix
int block_depth = atoi(tensor_name.substr(14, tensor_name.find(".", 14)).c_str());
if (block_depth + 1 > flux_params.depth_single_blocks) {
flux_params.depth_single_blocks = block_depth + 1;
}
}
if (ends_with(tensor_name, "txt_in.weight")) {
flux_params.context_in_dim = pair.second.ne[0];
flux_params.hidden_size = pair.second.ne[1];
}
if (ends_with(tensor_name, "single_blocks.0.norm.key_norm.scale")) {
head_dim = pair.second.ne[0];
}
if (ends_with(tensor_name, "double_blocks.0.txt_attn.norm.key_norm.scale")) {
head_dim = pair.second.ne[0];
}
}
if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != flux_params.patch_size) {
GGML_ASSERT(flux_params.patch_size == 2 * actual_radiance_patch_size);
LOG_DEBUG("using fake x2 patch size");
flux_params.chroma_radiance_params.fake_patch_size_x2 = true;
}
flux_params.num_heads = static_cast<int>(flux_params.hidden_size / head_dim);
LOG_INFO("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64
", hidden_size = %" PRId64 ", num_heads = %d",
flux_params.depth,
flux_params.depth_single_blocks,
flux_params.guidance_embed ? "true" : "false",
flux_params.context_in_dim,
flux_params.hidden_size,
flux_params.num_heads);
if (flux_params.is_chroma) {
: DiffusionModelRunner(backend, params_backend, prefix),
config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)),
version(version),
use_mask(use_mask) {
if (config.is_chroma) {
LOG_INFO("Using pruned modulation (Chroma)");
}
flux = Flux(flux_params);
flux = Flux(config);
flux.init(params_ctx, tensor_storage_map, prefix);
}
@@ -1377,10 +1390,10 @@ namespace Flux {
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
ggml_tensor* y = make_optional_input(y_tensor);
if (flux_params.guidance_embed || flux_params.is_chroma) {
if (config.guidance_embed || config.is_chroma) {
if (!guidance_tensor.empty()) {
this->guidance_tensor = guidance_tensor;
if (flux_params.is_chroma) {
if (config.is_chroma) {
this->guidance_tensor.fill_(0.f);
}
}
@@ -1398,7 +1411,7 @@ namespace Flux {
ggml_tensor* mod_index_arange = nullptr;
ggml_tensor* dct = nullptr; // for chroma radiance
if (flux_params.is_chroma) {
if (config.is_chroma) {
if (!use_mask) {
y = nullptr;
}
@@ -1417,29 +1430,29 @@ namespace Flux {
}
pe_vec = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
flux_params.patch_size,
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
increase_ref_index,
flux_params.ref_index_scale,
flux_params.theta,
config.ref_index_scale,
config.theta,
circular_y_enabled,
circular_x_enabled,
flux_params.axes_dim,
config.axes_dim,
sd_version_is_longcat(version));
int pos_len = static_cast<int>(pe_vec.size() / flux_params.axes_dim_sum / 2);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, flux_params.axes_dim_sum / 2, pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe);
// pe->data = nullptr;
set_backend_tensor_data(pe, pe_vec.data());
if (version == VERSION_CHROMA_RADIANCE) {
int patch_size = flux_params.patch_size;
int nerf_max_freqs = flux_params.chroma_radiance_params.nerf_max_freqs;
int patch_size = config.patch_size;
int nerf_max_freqs = config.chroma_radiance_params.nerf_max_freqs;
dct_vec = fetch_dct_pos(patch_size, nerf_max_freqs);
dct = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nerf_max_freqs * nerf_max_freqs, patch_size * patch_size);
// dct->data = dct_vec.data();
@@ -1567,7 +1580,7 @@ namespace Flux {
static void load_from_file_and_test(const std::string& file_path) {
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_COUNT;
ModelLoader model_loader;
otherarch/sdcpp/ggml_extend.hpp
+439 -51
View File
@@ -28,6 +28,7 @@
#include "ggml.h"
#include "ggml_extend_backend.h"
#include "ggml_graph_cut.h"
#include "layer_registry.h"
#include "model.h"
#include "tensor.hpp"
@@ -1329,13 +1330,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
float scale = (1.0f / sqrt((float)d_head));
int kv_pad = 0;
ggml_tensor* kqv = nullptr;
auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* {
if (kv_pad != 0) {
k_in = ggml_pad(ctx, k_in, 0, kv_pad, 0, 0);
}
if (kv_scale != 1.0f) {
k_in = ggml_ext_scale(ctx, k_in, kv_scale);
}
@@ -1343,9 +1340,6 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
if (kv_pad != 0) {
v_in = ggml_pad(ctx, v_in, 0, kv_pad, 0, 0);
}
if (kv_scale != 1.0f) {
v_in = ggml_ext_scale(ctx, v_in, kv_scale);
}
@@ -1353,26 +1347,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
if (mask_in != nullptr) {
mask_in = ggml_transpose(ctx, mask_in);
} else {
if (kv_pad > 0) {
mask_in = ggml_ext_zeros(ctx, L_k, L_q, 1, 1);
auto pad_tensor = ggml_ext_full(ctx, -INFINITY, kv_pad, L_q, 1, 1);
mask_in = ggml_concat(ctx, mask_in, pad_tensor, 0);
}
}
if (mask_in != nullptr) {
// the need for padding got removed in ggml 4767bda
// ensure we can still use the old version for now
#ifdef GGML_KQ_MASK_PAD
int mask_pad = 0;
if (mask_in->ne[1] % GGML_KQ_MASK_PAD != 0) {
mask_pad = GGML_PAD(L_q, GGML_KQ_MASK_PAD) - mask_in->ne[1];
}
if (mask_pad > 0) {
mask_in = ggml_pad(ctx, mask_in, 0, mask_pad, 0, 0);
}
#endif
mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
}
@@ -1387,10 +1364,6 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
if (flash_attn) {
// LOG_DEBUG("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
bool can_use_flash_attn = true;
if (can_use_flash_attn && L_k % 256 != 0) {
kv_pad = GGML_PAD(L_k, 256) - static_cast<int>(L_k);
}
if (mask != nullptr) {
// TODO: figure out if we can bend t5 to work too
can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1;
@@ -1470,7 +1443,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
__STATIC_INLINE__ void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) {
if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) &&
!ggml_backend_is_cpu(backend)) {
!sd_backend_is_cpu(backend)) {
ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
ggml_backend_synchronize(backend);
return;
@@ -1725,7 +1698,19 @@ protected:
ggml_context* partial_offload_ctx = nullptr;
ggml_backend_buffer_t partial_runtime_params_buffer = nullptr;
std::vector<std::pair<ggml_tensor*, ggml_tensor*>> partial_offload_pairs;
size_t max_graph_vram_bytes = 0;
// Params kept on the runtime backend across streaming segments.
ggml_context* resident_offload_ctx = nullptr;
std::vector<std::pair<ggml_tensor*, ggml_tensor*>> resident_offload_pairs;
ggml_backend_buffer_t resident_runtime_params_buffer = nullptr;
std::unordered_set<ggml_tensor*> resident_param_set;
uint64_t resident_state_token = 0;
size_t max_graph_vram_bytes = 0;
bool stream_layers_enabled = false;
size_t observed_max_effective_budget_ = 0;
sd::layer_registry::LayerRegistry layer_registry_;
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
@@ -1927,7 +1912,7 @@ protected:
LOG_DEBUG("%s compute buffer size: %.2f MB(%s)",
get_desc().c_str(),
compute_buffer_size / 1024.0 / 1024.0,
ggml_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM");
sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM");
return true;
}
@@ -2014,7 +1999,7 @@ protected:
LOG_DEBUG("%s cache backend buffer size = % 6.2f MB(%s) (%i tensors)",
get_desc().c_str(),
cache_buffer_size / (1024.f * 1024.f),
ggml_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM",
sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM",
num_tensors);
if (old_cache_buffer != nullptr) {
ggml_backend_buffer_free(old_cache_buffer);
@@ -2193,6 +2178,9 @@ protected:
if (tensor == nullptr) {
continue;
}
if (resident_param_set.find(tensor) != resident_param_set.end()) {
continue;
}
if (seen_tensors.insert(tensor).second) {
unique_tensors.push_back(tensor);
}
@@ -2315,36 +2303,225 @@ protected:
}
}
bool offload_resident_params(const std::vector<ggml_tensor*>& tensors) {
if (params_backend == runtime_backend) {
return true;
}
if (tensors.empty()) {
return true;
}
GGML_ASSERT(resident_runtime_params_buffer == nullptr);
GGML_ASSERT(resident_offload_ctx == nullptr);
GGML_ASSERT(resident_offload_pairs.empty());
GGML_ASSERT(resident_param_set.empty());
std::vector<ggml_tensor*> unique_tensors;
std::unordered_set<ggml_tensor*> seen;
unique_tensors.reserve(tensors.size());
seen.reserve(tensors.size());
for (ggml_tensor* t : tensors) {
if (t == nullptr)
continue;
if (seen.insert(t).second)
unique_tensors.push_back(t);
}
if (unique_tensors.empty())
return true;
ggml_init_params init = {};
init.mem_size = std::max<size_t>(1, unique_tensors.size()) * ggml_tensor_overhead();
init.mem_buffer = nullptr;
init.no_alloc = true;
resident_offload_ctx = ggml_init(init);
GGML_ASSERT(resident_offload_ctx != nullptr);
resident_offload_pairs.reserve(unique_tensors.size());
for (ggml_tensor* t : unique_tensors) {
GGML_ASSERT(t->view_src == nullptr);
ggml_tensor* twin = ggml_dup_tensor(resident_offload_ctx, t);
ggml_set_name(twin, t->name);
resident_offload_pairs.push_back({t, twin});
}
resident_runtime_params_buffer = ggml_backend_alloc_ctx_tensors(resident_offload_ctx, runtime_backend);
if (resident_runtime_params_buffer == nullptr) {
LOG_ERROR("%s alloc resident runtime params backend buffer failed, num_tensors = %zu",
get_desc().c_str(), resident_offload_pairs.size());
ggml_free(resident_offload_ctx);
resident_offload_ctx = nullptr;
resident_offload_pairs.clear();
return false;
}
ggml_backend_buffer_set_usage(resident_runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
for (auto& pair : resident_offload_pairs) {
ggml_tensor* t = pair.first;
ggml_tensor* twin = pair.second;
ggml_backend_tensor_copy(t, twin);
std::swap(t->buffer, twin->buffer);
std::swap(t->data, twin->data);
std::swap(t->extra, twin->extra);
resident_param_set.insert(t);
}
ggml_backend_synchronize(runtime_backend);
size_t sz = ggml_backend_buffer_get_size(resident_runtime_params_buffer);
LOG_INFO("%s offload resident params (%6.2f MB, %zu tensors) to runtime backend (%s)",
get_desc().c_str(),
sz / (1024.f * 1024.f),
resident_offload_pairs.size(),
ggml_backend_name(runtime_backend));
return true;
}
void restore_resident_params() {
if (resident_offload_pairs.empty()) {
if (resident_runtime_params_buffer != nullptr) {
ggml_backend_buffer_free(resident_runtime_params_buffer);
resident_runtime_params_buffer = nullptr;
}
if (resident_offload_ctx != nullptr) {
ggml_free(resident_offload_ctx);
resident_offload_ctx = nullptr;
}
resident_param_set.clear();
resident_state_token = 0;
return;
}
for (auto& pair : resident_offload_pairs) {
ggml_tensor* t = pair.first;
ggml_tensor* twin = pair.second;
t->buffer = twin->buffer;
t->data = twin->data;
t->extra = twin->extra;
twin->buffer = nullptr;
twin->data = nullptr;
twin->extra = nullptr;
}
if (resident_runtime_params_buffer != nullptr) {
ggml_backend_buffer_free(resident_runtime_params_buffer);
resident_runtime_params_buffer = nullptr;
}
resident_offload_pairs.clear();
if (resident_offload_ctx != nullptr) {
ggml_free(resident_offload_ctx);
resident_offload_ctx = nullptr;
}
resident_param_set.clear();
resident_state_token = 0;
}
bool should_use_graph_cut_segmented_compute(const GraphCutPlan& plan) {
return plan.has_cuts &&
plan.valid &&
max_graph_vram_bytes > 0 &&
plan.segments.size() > 1 &&
params_backend != runtime_backend &&
!ggml_backend_is_cpu(runtime_backend);
!sd_backend_is_cpu(runtime_backend);
}
bool can_attempt_graph_cut_segmented_compute() const {
return max_graph_vram_bytes > 0 &&
params_backend != runtime_backend &&
!ggml_backend_is_cpu(runtime_backend);
!sd_backend_is_cpu(runtime_backend);
}
bool resolve_graph_cut_plan(ggml_cgraph* gf,
GraphCutPlan* plan_out) {
GraphCutPlan* plan_out,
size_t* effective_budget_out = nullptr) {
GGML_ASSERT(plan_out != nullptr);
GGML_ASSERT(gf != nullptr);
// Keep the plan and resident params under the same live-VRAM cap.
// Add back our own resident buffer so we don't see chunk-K's
// allocation as "taken" VRAM and shrink the budget on every step.
size_t effective_budget = max_graph_vram_bytes;
if (stream_layers_enabled && max_graph_vram_bytes > 0 && runtime_backend != nullptr) {
ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend);
if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
size_t free_vram = 0, total_vram = 0;
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
if (resident_runtime_params_buffer != nullptr) {
free_vram += ggml_backend_buffer_get_size(resident_runtime_params_buffer);
}
constexpr size_t safety_margin = 512ull * 1024 * 1024;
size_t free_clamp = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0;
if (free_clamp < effective_budget) {
LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB",
get_desc().c_str(),
free_clamp / (1024.0 * 1024.0),
effective_budget / (1024.0 * 1024.0));
effective_budget = free_clamp;
}
}
}
bool budget_increased = false;
if (stream_layers_enabled) {
if (effective_budget > observed_max_effective_budget_) {
observed_max_effective_budget_ = effective_budget;
budget_increased = true;
} else {
effective_budget = observed_max_effective_budget_;
}
}
if (effective_budget_out != nullptr) {
*effective_budget_out = effective_budget;
}
*plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
gf,
&graph_cut_plan_cache_,
max_graph_vram_bytes,
effective_budget,
params_tensor_set_,
get_desc().c_str());
if (stream_layers_enabled) {
if (budget_increased) {
LOG_INFO("%s streaming budget = %.2f MB",
get_desc().c_str(),
effective_budget / (1024.0 * 1024.0));
} else {
LOG_DEBUG("%s streaming budget = %.2f MB",
get_desc().c_str(),
effective_budget / (1024.0 * 1024.0));
}
}
return true;
}
struct PersistentExternalBinding {
ggml_backend_buffer_t buffer = nullptr;
void* data = nullptr;
void* extra = nullptr;
};
void snapshot_persistent_externals(const sd::ggml_graph_cut::Plan& plan,
ggml_cgraph* gf,
std::unordered_map<ggml_tensor*, PersistentExternalBinding>& out) {
GGML_ASSERT(gf != nullptr);
out.clear();
for (const auto& segment : plan.segments) {
for (const auto& input : segment.input_refs) {
if (input.type != GraphCutSegment::INPUT_EXTERNAL) {
continue;
}
ggml_tensor* tensor = sd::ggml_graph_cut::input_tensor(gf, input);
if (tensor == nullptr || tensor->buffer == nullptr) {
continue;
}
PersistentExternalBinding binding;
binding.buffer = tensor->buffer;
binding.data = tensor->data;
binding.extra = tensor->extra;
out[tensor] = binding;
}
}
}
void reset_segment_runtime_tensors(const GraphCutSegment& segment,
ggml_cgraph* gf) {
ggml_cgraph* gf,
const std::unordered_map<ggml_tensor*, PersistentExternalBinding>* persistent_externals = nullptr) {
GGML_ASSERT(gf != nullptr);
for (const auto& input : segment.input_refs) {
@@ -2354,11 +2531,25 @@ protected:
}
switch (input.type) {
case GraphCutSegment::INPUT_PREVIOUS_CUT:
case GraphCutSegment::INPUT_EXTERNAL:
input_tensor->buffer = nullptr;
input_tensor->data = nullptr;
input_tensor->extra = nullptr;
break;
case GraphCutSegment::INPUT_EXTERNAL: {
if (persistent_externals != nullptr) {
auto it = persistent_externals->find(input_tensor);
if (it != persistent_externals->end()) {
input_tensor->buffer = it->second.buffer;
input_tensor->data = it->second.data;
input_tensor->extra = it->second.extra;
break;
}
}
input_tensor->buffer = nullptr;
input_tensor->data = nullptr;
input_tensor->extra = nullptr;
break;
}
case GraphCutSegment::INPUT_PARAM:
break;
}
@@ -2464,8 +2655,8 @@ protected:
int64_t t_copy_begin = ggml_time_ms();
copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map);
int64_t t_copy_end = ggml_time_ms();
if (ggml_backend_is_cpu(runtime_backend)) {
ggml_backend_cpu_set_n_threads(runtime_backend, n_threads);
if (sd_backend_is_cpu(runtime_backend)) {
sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
}
int64_t t_compute_begin = ggml_time_ms();
@@ -2573,6 +2764,9 @@ protected:
free_compute_buffer();
free_cache_ctx_and_buffer();
std::unordered_map<ggml_tensor*, PersistentExternalBinding> persistent_externals;
snapshot_persistent_externals(plan, gf, persistent_externals);
std::optional<sd::Tensor<T>> output = sd::Tensor<T>();
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) {
int64_t t_segment_begin = ggml_time_ms();
@@ -2584,7 +2778,7 @@ protected:
plan.segments.size(),
segment.group_name.c_str());
reset_segment_runtime_tensors(segment, gf);
reset_segment_runtime_tensors(segment, gf, &persistent_externals);
if (!bind_segment_cached_inputs(gf, segment)) {
free_cache_ctx_and_buffer();
free_compute_buffer();
@@ -2629,6 +2823,150 @@ protected:
return output;
}
public:
void release_streaming_residency() {
restore_resident_params();
}
template <typename T>
std::optional<sd::Tensor<T>> compute_streaming_segments(ggml_cgraph* gf,
const GraphCutPlan& plan,
size_t residency_budget_bytes,
int n_threads,
bool free_compute_buffer_immediately,
bool no_return = false) {
GGML_ASSERT(gf != nullptr);
// Runtime LoRA composes `weight + diff` in the compute graph via
// ggml_add; the resident weight tensor's data is never mutated, so
// chunk-K residency stays valid across sampling steps.
// Reserve room for the worst merged segment so chunk-K can't grow
// large enough to starve later partial-param allocations.
size_t worst_merged_segment_footprint = 0;
for (const auto& seg : plan.segments) {
const size_t fp = seg.input_param_bytes +
seg.compute_buffer_size +
seg.output_bytes +
seg.input_previous_cut_bytes +
seg.input_external_bytes;
if (fp > worst_merged_segment_footprint) {
worst_merged_segment_footprint = fp;
}
}
const size_t residency_budget_for_annotate =
residency_budget_bytes > worst_merged_segment_footprint
? residency_budget_bytes - worst_merged_segment_footprint
: 0;
sd::ggml_graph_cut::Plan& base_plan = graph_cut_plan_cache_.graph_cut_plan;
if (base_plan.available) {
sd::ggml_graph_cut::annotate_residency(base_plan, residency_budget_for_annotate);
std::vector<ggml_tensor*> resident_params;
uint64_t token = 0;
for (const auto& segment : base_plan.segments) {
if (segment.residency != sd::ggml_graph_cut::SegmentResidency::RESIDENT) {
continue;
}
auto seg_params = sd::ggml_graph_cut::param_tensors(gf, segment);
for (ggml_tensor* t : seg_params) {
if (t == nullptr)
continue;
resident_params.push_back(t);
token ^= reinterpret_cast<uintptr_t>(t) * 0x9E3779B97F4A7C15ull;
}
}
if (token != resident_state_token) {
restore_resident_params();
if (!resident_params.empty()) {
if (offload_resident_params(resident_params)) {
resident_state_token = token;
} else {
LOG_ERROR("%s chunk-K: resident offload failed; continuing with per-segment streaming",
get_desc().c_str());
restore_resident_params();
}
}
}
}
free_compute_buffer();
free_cache_ctx_and_buffer();
layer_registry_.move_layer_to_gpu("_global");
std::unordered_map<ggml_tensor*, PersistentExternalBinding> persistent_externals;
snapshot_persistent_externals(plan, gf, persistent_externals);
std::optional<sd::Tensor<T>> output = sd::Tensor<T>();
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) {
int64_t t_segment_begin = ggml_time_ms();
const auto& segment = plan.segments[seg_idx];
const bool is_last = seg_idx + 1 == plan.segments.size();
auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx);
LOG_DEBUG("%s streaming-cut executing segment %zu/%zu: %s (residency=%s)",
get_desc().c_str(),
seg_idx + 1,
plan.segments.size(),
segment.group_name.c_str(),
segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT ? "RESIDENT" : "STREAMED");
if (!layer_registry_.move_layer_to_gpu(segment.group_name)) {
LOG_DEBUG("%s streaming: no registry entry for group '%s' (using upstream offload path)",
get_desc().c_str(),
segment.group_name.c_str());
}
reset_segment_runtime_tensors(segment, gf, &persistent_externals);
if (!bind_segment_cached_inputs(gf, segment)) {
free_cache_ctx_and_buffer();
free_compute_buffer();
free_compute_ctx();
return std::nullopt;
}
if (!is_last) {
for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) {
ggml_tensor* out_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx);
if (out_tensor != nullptr &&
sd::ggml_graph_cut::is_graph_cut_tensor(out_tensor) &&
future_cut_names.find(out_tensor->name) != future_cut_names.end()) {
cache(out_tensor->name, out_tensor);
}
}
}
ggml_context* segment_graph_ctx = nullptr;
ggml_cgraph* segment_graph = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx);
auto segment_output = execute_graph<T>(segment_graph,
n_threads,
/*free_compute_buffer_immediately=*/true,
sd::ggml_graph_cut::runtime_param_tensors(gf, segment, get_desc().c_str()),
/*preserve_backend_tensor_data_map=*/true,
/*no_return=*/!is_last || no_return,
&future_cut_names);
ggml_free(segment_graph_ctx);
if (!segment_output.has_value()) {
free_cache_ctx_and_buffer();
free_compute_buffer();
free_compute_ctx();
return std::nullopt;
}
output = std::move(segment_output);
if (segment.residency == sd::ggml_graph_cut::SegmentResidency::STREAMED) {
layer_registry_.move_layer_to_cpu(segment.group_name);
}
(void)t_segment_begin;
}
backend_tensor_data_map.clear();
free_cache_ctx_and_buffer();
free_compute_ctx();
return output;
}
public:
virtual std::string get_desc() = 0;
@@ -2638,9 +2976,11 @@ public:
GGML_ASSERT(runtime_backend != nullptr);
GGML_ASSERT(params_backend != nullptr);
alloc_params_ctx();
layer_registry_.set_backends(runtime_backend, params_backend);
}
virtual ~GGMLRunner() {
restore_resident_params();
free_params_buffer();
free_compute_buffer();
free_params_ctx();
@@ -2694,7 +3034,18 @@ public:
LOG_DEBUG("%s skipping params allocation (no tensors)", get_desc().c_str());
return true;
}
params_buffer = ggml_backend_alloc_ctx_tensors(params_ctx, params_backend);
// Pinned host buffer when CPU-offloaded for DMA-direct H2D.
ggml_backend_buffer_type_t params_buft = nullptr;
if (params_backend != runtime_backend) {
ggml_backend_dev_t runtime_dev = ggml_backend_get_device(runtime_backend);
if (runtime_dev != nullptr) {
params_buft = ggml_backend_dev_host_buffer_type(runtime_dev);
}
}
if (params_buft == nullptr) {
params_buft = ggml_backend_get_default_buffer_type(params_backend);
}
params_buffer = ggml_backend_alloc_ctx_tensors_from_buft(params_ctx, params_buft);
if (params_buffer == nullptr) {
LOG_ERROR("%s alloc params backend buffer failed, num_tensors = %i",
get_desc().c_str(),
@@ -2707,16 +3058,19 @@ public:
LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)",
get_desc().c_str(),
params_buffer_size / (1024.f * 1024.f),
ggml_backend_is_cpu(params_backend) ? "RAM" : "VRAM",
sd_backend_is_cpu(params_backend) ? "RAM" : "VRAM",
num_tensors);
return true;
}
void free_params_buffer() {
// Restore swapped resident params before freeing their backing buffer.
restore_resident_params();
if (params_buffer != nullptr) {
ggml_backend_buffer_free(params_buffer);
params_buffer = nullptr;
}
observed_max_effective_budget_ = 0;
}
size_t get_params_buffer_size() {
@@ -2774,7 +3128,7 @@ public:
return nullptr;
}
// it's performing a compute, check if backend isn't cpu
if (!ggml_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
// pass input tensors to gpu memory
auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
@@ -2812,11 +3166,20 @@ public:
if (can_attempt_graph_cut_segmented_compute()) {
GraphCutPlan plan;
if (!resolve_graph_cut_plan(gf, &plan)) {
size_t effective_graph_vram_bytes = 0;
if (!resolve_graph_cut_plan(gf, &plan, &effective_graph_vram_bytes)) {
free_compute_ctx();
return std::nullopt;
}
if (should_use_graph_cut_segmented_compute(plan)) {
if (stream_layers_enabled) {
return compute_streaming_segments<T>(gf,
plan,
effective_graph_vram_bytes,
n_threads,
free_compute_buffer_immediately,
no_return);
}
return compute_with_graph_cuts<T>(gf,
plan,
n_threads,
@@ -2857,6 +3220,12 @@ public:
max_graph_vram_bytes = max_vram_bytes;
}
void set_stream_layers_enabled(bool enabled) {
stream_layers_enabled = enabled;
}
sd::layer_registry::LayerRegistry& get_layer_registry() { return layer_registry_; }
ggml_backend_t get_runtime_backend() {
return runtime_backend;
}
@@ -2978,11 +3347,14 @@ protected:
bool bias;
bool force_f32;
bool force_prec_f32;
bool allow_weight_scale;
bool has_weight_scale = false;
float scale;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
this->prefix = prefix;
has_weight_scale = false;
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
wtype = GGML_TYPE_F32;
@@ -2992,20 +3364,26 @@ protected:
enum ggml_type wtype = GGML_TYPE_F32;
params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features);
}
if (allow_weight_scale && tensor_storage_map.find(prefix + "weight_scale") != tensor_storage_map.end()) {
params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
has_weight_scale = true;
}
}
public:
Linear(int64_t in_features,
int64_t out_features,
bool bias = true,
bool force_f32 = false,
bool force_prec_f32 = false,
float scale = 1.f)
bool bias = true,
bool force_f32 = false,
bool force_prec_f32 = false,
float scale = 1.f,
bool allow_weight_scale = false)
: in_features(in_features),
out_features(out_features),
bias(bias),
force_f32(force_f32),
force_prec_f32(force_prec_f32),
allow_weight_scale(allow_weight_scale),
scale(scale) {}
void set_scale(float scale_) {
@@ -3022,14 +3400,24 @@ public:
if (bias) {
b = params["bias"];
}
ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
ggml_tensor* out = nullptr;
if (ctx->weight_adapter) {
WeightAdapter::ForwardParams forward_params;
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
forward_params.linear.force_prec_f32 = force_prec_f32;
forward_params.linear.scale = scale;
return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
out = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params);
} else {
out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
}
return ggml_ext_linear(ctx->ggml_ctx, x, w, b, force_prec_f32, scale);
if (has_weight_scale) {
out = ggml_mul(ctx->ggml_ctx, out, params["weight_scale"]);
if (b != nullptr) {
out = ggml_add_inplace(ctx->ggml_ctx, out, b);
}
}
return out;
}
};
otherarch/sdcpp/ggml_extend_backend.cpp
+61 -5
View File
@@ -8,6 +8,7 @@
#include <stdexcept>
#include <vector>
#include "stable-diffusion.h"
#include "util.h"
static std::string trim_copy(const std::string& value) {
@@ -300,6 +301,61 @@ static ggml_backend_t init_named_backend(const std::string& name) {
return ggml_backend_init_by_name(resolved.c_str(), nullptr);
}
bool sd_backend_is_cpu(ggml_backend_t backend) {
if (backend == nullptr) {
return false;
}
auto dev = ggml_backend_get_device(backend);
return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
}
ggml_backend_t sd_backend_cpu_init() {
ggml_backend_load_all_once();
return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
}
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
if (backend == nullptr) {
return false;
}
auto dev = ggml_backend_get_device(backend);
if (dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
auto reg = ggml_backend_dev_backend_reg(dev);
auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
if (ggml_backend_set_n_threads_fn != nullptr) {
ggml_backend_set_n_threads_fn(backend, n_threads);
return true;
}
}
return false;
}
const char* sd_get_system_info() {
static std::string cache_info = []() -> std::string {
ggml_backend_load_all_once();
std::stringstream ss;
ss << "System Info: \n";
auto dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
if (dev != nullptr) {
auto reg = ggml_backend_dev_backend_reg(dev);
auto ggml_backend_get_features_fn = (ggml_backend_get_features_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features");
if (ggml_backend_get_features_fn != nullptr) {
ggml_backend_feature* feat = ggml_backend_get_features_fn(reg);
while (feat->name && feat->value) {
ss << " " << feat->name << " = " << feat->value << " | ";
feat++;
}
} else {
LOG_WARN("unable to get CPU features");
}
} else {
LOG_WARN("unable to get CPU features");
}
return ss.str();
}();
return cache_info.c_str();
}
static ggml_backend_t sd_get_default_backend() {
ggml_backend_load_all_once();
static std::once_flag once;
@@ -349,10 +405,10 @@ static ggml_backend_t sd_get_default_backend() {
if (!backend) {
LOG_WARN("loading CPU backend");
backend = ggml_backend_cpu_init();
backend = sd_backend_cpu_init();
}
if (ggml_backend_is_cpu(backend)) {
if (sd_backend_is_cpu(backend)) {
LOG_DEBUG("Using CPU backend");
}
@@ -452,11 +508,11 @@ ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
}
bool SDBackendManager::runtime_backend_is_cpu(SDBackendModule module) {
return ggml_backend_is_cpu(runtime_backend(module));
return sd_backend_is_cpu(runtime_backend(module));
}
bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) {
return ggml_backend_is_cpu(params_backend(module));
return sd_backend_is_cpu(params_backend(module));
}
bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
@@ -464,7 +520,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
if (backend == nullptr) {
return false;
}
if (ggml_backend_is_cpu(backend)) {
if (sd_backend_is_cpu(backend)) {
return true;
}
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
otherarch/sdcpp/ggml_extend_backend.h
+3 -1
View File
@@ -8,7 +8,6 @@
#include <unordered_map>
#include "ggml-backend.h"
#include "ggml-cpu.h"
#include "ggml.h"
enum class SDBackendModule {
@@ -72,6 +71,9 @@ private:
};
bool sd_backend_is(ggml_backend_t backend, const std::string& name);
bool sd_backend_is_cpu(ggml_backend_t backend);
ggml_backend_t sd_backend_cpu_init();
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
const char* sd_backend_module_name(SDBackendModule module);
void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
#endif
otherarch/sdcpp/ggml_graph_cut.cpp
+53 -3
View File
@@ -699,9 +699,9 @@ namespace sd::ggml_graph_cut {
}
if (log_desc != nullptr) {
LOG_INFO("%s graph cut max_vram budget merge took %lld ms",
log_desc,
ggml_time_ms() - t_budget_begin);
LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
log_desc,
ggml_time_ms() - t_budget_begin);
}
return merged_plan;
@@ -753,4 +753,54 @@ namespace sd::ggml_graph_cut {
return resolved_plan;
}
void annotate_residency(Plan& plan, size_t max_graph_vram_bytes) {
// Cached plans may be reused with a smaller live budget.
for (auto& seg : plan.segments) {
seg.residency = SegmentResidency::STREAMED;
}
if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) {
return;
}
bool any_param_bearing = false;
for (const auto& seg : plan.segments) {
if (seg.input_param_bytes > 0) {
any_param_bearing = true;
break;
}
}
if (!any_param_bearing) {
return;
}
// Leave room for the largest active streamed segment.
size_t worst_streamed_footprint = 0;
for (const auto& seg : plan.segments) {
const size_t seg_footprint = seg.input_param_bytes +
seg.compute_buffer_size +
seg.output_bytes +
seg.input_previous_cut_bytes +
seg.input_external_bytes;
if (seg_footprint > worst_streamed_footprint) {
worst_streamed_footprint = seg_footprint;
}
}
constexpr size_t safety = 512ull * 1024 * 1024;
const size_t reserved = safety + worst_streamed_footprint;
if (max_graph_vram_bytes <= reserved) {
return;
}
const size_t available = max_graph_vram_bytes - reserved;
size_t cumulative = 0;
for (auto& seg : plan.segments) {
if (cumulative + seg.input_param_bytes > available) {
break;
}
seg.residency = SegmentResidency::RESIDENT;
cumulative += seg.input_param_bytes;
}
}
} // namespace sd::ggml_graph_cut
otherarch/sdcpp/ggml_graph_cut.h
+11
View File
@@ -2,6 +2,7 @@
#define __SD_GGML_GRAPH_CUT_H__
#include <array>
#include <cstdint>
#include <string>
#include <unordered_set>
#include <vector>
@@ -11,6 +12,12 @@
namespace sd::ggml_graph_cut {
// Streaming residency for a segment's params.
enum class SegmentResidency : uint8_t {
STREAMED = 0,
RESIDENT = 1,
};
struct Segment {
enum InputType {
INPUT_EXTERNAL = 0,
@@ -34,6 +41,7 @@ namespace sd::ggml_graph_cut {
std::vector<int> internal_node_indices;
std::vector<int> output_node_indices;
std::vector<InputRef> input_refs;
SegmentResidency residency = SegmentResidency::STREAMED;
};
struct Plan {
@@ -101,6 +109,9 @@ namespace sd::ggml_graph_cut {
size_t max_graph_vram_bytes,
const std::unordered_set<const ggml_tensor*>& params_tensor_set,
const char* log_desc);
// Mark leading segments resident when they fit after streamed-segment headroom.
void annotate_residency(Plan& plan, size_t max_graph_vram_bytes);
} // namespace sd::ggml_graph_cut
#endif
otherarch/sdcpp/guidance.cpp
+180 -10
View File
@@ -1,13 +1,68 @@
#include "guidance.h"
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <string>
#include <utility>
#include "util.h"
namespace sd::guidance {
static bool has_tensor(const sd::Tensor<float>* tensor) {
return tensor != nullptr && !tensor->empty();
}
bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params) {
return params.eta != 1.0f || params.momentum != 0.0f || params.norm_threshold > 0.0f;
}
AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args) {
AdaptiveProjectedGuidanceParams params;
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) {
float parsed = 0.0f;
if (key == "apg_eta") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
params.eta = parsed;
} else if (key == "apg_momentum") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
params.momentum = parsed;
} else if (key == "apg_norm_threshold") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
params.norm_threshold = parsed;
} else if (key == "apg_norm_threshold_smoothing") {
if (!parse_strict_float(value, parsed)) {
LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
params.norm_threshold_smoothing = parsed;
}
}
return params;
}
bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args) {
bool uncond = false;
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) {
if (key == "slg_uncond") {
if (!parse_strict_bool(value, uncond)) {
LOG_WARN("ignoring invalid SLG extra sample arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
return uncond;
}
ClassifierFreeGuidance::ClassifierFreeGuidance(float guidance_scale,
float image_guidance_scale)
: guidance_scale_(guidance_scale),
@@ -27,17 +82,132 @@ namespace sd::guidance {
output.pred = pred_cond;
if (has_tensor(input.pred_uncond)) {
const sd::Tensor<float>& pred_uncond = *input.pred_uncond;
if (has_tensor(input.pred_img_cond)) {
const sd::Tensor<float>& pred_img_cond = *input.pred_img_cond;
output.pred = pred_uncond +
image_guidance_scale_ * (pred_img_cond - pred_uncond) +
guidance_scale_ * (pred_cond - pred_img_cond);
if (has_tensor(input.pred_img_uncond)) {
const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
output.pred = pred_img_uncond +
image_guidance_scale_ * (pred_uncond - pred_img_uncond) +
guidance_scale_ * (pred_cond - pred_uncond);
} else {
output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond);
}
} else if (has_tensor(input.pred_img_cond)) {
const sd::Tensor<float>& pred_img_cond = *input.pred_img_cond;
output.pred = pred_img_cond + guidance_scale_ * (pred_cond - pred_img_cond);
} else if (has_tensor(input.pred_img_uncond)) {
const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
output.pred = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond);
}
return output;
}
AdaptiveProjectedGuidance::AdaptiveProjectedGuidance(float guidance_scale,
float image_guidance_scale,
AdaptiveProjectedGuidanceParams params)
: guidance_scale_(guidance_scale),
image_guidance_scale_(image_guidance_scale),
params_(params) {
}
static sd::Tensor<float> calculate_guidance_delta(const sd::Tensor<float>& pred_cond,
const sd::Tensor<float>* pred_uncond,
const sd::Tensor<float>* pred_img_uncond,
float guidance_scale,
float image_guidance_scale) {
if (pred_img_uncond != nullptr) {
if (pred_uncond != nullptr && guidance_scale == 1.0f) {
return *pred_uncond - *pred_img_uncond;
}
if (pred_uncond != nullptr) {
return pred_cond +
(*pred_uncond * (image_guidance_scale - guidance_scale) +
*pred_img_uncond * (1.0f - image_guidance_scale)) /
(guidance_scale - 1.0f);
}
return pred_cond - *pred_img_uncond;
}
return pred_cond - *pred_uncond;
}
GuiderOutput AdaptiveProjectedGuidance::forward(const GuidanceInput& input,
GuiderOutput previous) const {
(void)previous;
GuiderOutput output;
if (!has_tensor(input.pred_cond)) {
return output;
}
const sd::Tensor<float>& pred_cond = *input.pred_cond;
output.pred = pred_cond;
if (has_tensor(input.pred_uncond)) {
const sd::Tensor<float>& pred_uncond = *input.pred_uncond;
if (has_tensor(input.pred_img_uncond)) {
const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
output.pred = pred_img_uncond +
image_guidance_scale_ * (pred_uncond - pred_img_uncond) +
guidance_scale_ * (pred_cond - pred_uncond);
} else {
output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond);
}
} else if (has_tensor(input.pred_img_uncond)) {
const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
output.pred = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond);
}
if (!has_tensor(input.pred_uncond) && !has_tensor(input.pred_img_uncond)) {
return output;
}
const sd::Tensor<float>* pred_uncond = input.pred_uncond;
const sd::Tensor<float>* pred_img_uncond = input.pred_img_uncond;
sd::Tensor<float> deltas = calculate_guidance_delta(pred_cond,
pred_uncond,
pred_img_uncond,
guidance_scale_,
image_guidance_scale_);
if (params_.momentum != 0.0f) {
if (momentum_buffer_.shape() != deltas.shape()) {
momentum_buffer_ = sd::Tensor<float>::zeros_like(deltas);
}
deltas += params_.momentum * momentum_buffer_;
momentum_buffer_ = deltas;
}
float diff_norm = 0.0f;
if (params_.norm_threshold > 0.0f) {
diff_norm = std::sqrt((deltas * deltas).sum());
}
float apg_scale_factor = 1.0f;
if (params_.norm_threshold > 0.0f) {
if (diff_norm > 0.0f) {
if (params_.norm_threshold_smoothing <= 0.0f) {
apg_scale_factor = std::min(1.0f, params_.norm_threshold / diff_norm);
} else {
float x = params_.norm_threshold / diff_norm;
apg_scale_factor = x / std::pow(1.0f + std::pow(x, 1.0f / params_.norm_threshold_smoothing),
params_.norm_threshold_smoothing);
}
}
}
deltas *= apg_scale_factor;
if (params_.eta != 1.0f) {
float cond_norm_sq = (pred_cond * pred_cond).sum();
if (cond_norm_sq != 0.0f) {
float projection_scale = (pred_cond * deltas).sum() / cond_norm_sq;
deltas += (params_.eta - 1.0f) * (projection_scale * pred_cond);
}
}
output.pred = pred_cond;
if (pred_uncond != nullptr) {
if (guidance_scale_ != 1.0f) {
output.pred = pred_cond + (guidance_scale_ - 1.0f) * deltas;
} else if (pred_img_uncond != nullptr) {
output.pred = pred_cond + (image_guidance_scale_ - 1.0f) * deltas;
}
} else if (pred_img_uncond != nullptr) {
output.pred = *pred_img_uncond + guidance_scale_ * deltas;
}
return output;
@@ -54,7 +224,7 @@ namespace sd::guidance {
}
bool SkipLayerGuidance::is_enabled_for_step(const GuidanceInput& input) const {
if (scale_ == 0.0f || layers_.empty() || input.schedule_size == 0) {
if (layers_.empty() || input.schedule_size == 0) {
return false;
}
@@ -69,7 +239,7 @@ namespace sd::guidance {
GuiderOutput SkipLayerGuidance::forward(const GuidanceInput& input,
GuiderOutput output) const {
if (!is_enabled_for_step(input) || !input.predict_skip_layer) {
if (scale_ == 0.0f || !is_enabled_for_step(input) || !input.predict_skip_layer) {
return output;
}
otherarch/sdcpp/guidance.h
+31 -5
View File
@@ -17,12 +17,23 @@ namespace sd::guidance {
sd::Tensor<float> pred_skip_layer;
};
struct AdaptiveProjectedGuidanceParams {
float eta = 1.0f;
float momentum = 0.0f;
float norm_threshold = 0.0f;
float norm_threshold_smoothing = 0.0f;
};
AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args);
bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params);
bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args);
struct GuidanceInput {
int step = 0;
size_t schedule_size = 0;
const sd::Tensor<float>* pred_cond = nullptr;
const sd::Tensor<float>* pred_uncond = nullptr;
const sd::Tensor<float>* pred_img_cond = nullptr;
int step = 0;
size_t schedule_size = 0;
const sd::Tensor<float>* pred_cond = nullptr;
const sd::Tensor<float>* pred_uncond = nullptr;
const sd::Tensor<float>* pred_img_uncond = nullptr;
std::function<sd::Tensor<float>()> predict_skip_layer;
};
@@ -46,6 +57,21 @@ namespace sd::guidance {
GuiderOutput previous) const override;
};
class AdaptiveProjectedGuidance : public BaseGuidance {
float guidance_scale_ = 1.0f;
float image_guidance_scale_ = 1.0f;
AdaptiveProjectedGuidanceParams params_;
mutable sd::Tensor<float> momentum_buffer_;
public:
AdaptiveProjectedGuidance(float guidance_scale,
float image_guidance_scale,
AdaptiveProjectedGuidanceParams params);
GuiderOutput forward(const GuidanceInput& input,
GuiderOutput previous) const override;
};
class SkipLayerGuidance : public BaseGuidance {
std::vector<int> layers_;
float scale_ = 0.0f;
otherarch/sdcpp/hidream_o1.hpp
+50 -48
View File
@@ -23,6 +23,39 @@ namespace HiDreamO1 {
constexpr int IMAGE_TOKEN_ID = 151655;
constexpr int VISION_START_TOKEN_ID = 151652;
struct HiDreamO1Config {
LLM::LLMConfig llm;
int patch_size = PATCH_SIZE;
static HiDreamO1Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
(void)tensor_storage_map;
(void)prefix;
HiDreamO1Config config;
config.llm.arch = LLM::LLMArch::QWEN3_VL;
config.llm.hidden_size = 4096;
config.llm.intermediate_size = 12288;
config.llm.num_layers = 36;
config.llm.num_heads = 32;
config.llm.num_kv_heads = 8;
config.llm.head_dim = 128;
config.llm.qkv_bias = false;
config.llm.qk_norm = true;
config.llm.vocab_size = 151936;
config.llm.rms_norm_eps = 1e-6f;
config.llm.vision.arch = LLM::LLMVisionArch::QWEN3_VL;
config.llm.vision.num_layers = 27;
config.llm.vision.hidden_size = 1152;
config.llm.vision.intermediate_size = 4304;
config.llm.vision.num_heads = 16;
config.llm.vision.out_hidden_size = 4096;
config.llm.vision.patch_size = 16;
config.llm.vision.spatial_merge_size = 2;
config.llm.vision.temporal_patch_size = 2;
config.llm.vision.num_position_embeddings = 2304;
return config;
}
};
static inline std::string repeat_special_token(const std::string& token, int64_t count) {
std::string out;
out.reserve(static_cast<size_t>(count) * token.size());
@@ -205,50 +238,19 @@ namespace HiDreamO1 {
}
};
struct HiDreamO1Params {
LLM::LLMParams llm;
int patch_size = PATCH_SIZE;
};
static inline HiDreamO1Params make_hidream_o1_params() {
HiDreamO1Params params;
params.llm.arch = LLM::LLMArch::QWEN3_VL;
params.llm.hidden_size = 4096;
params.llm.intermediate_size = 12288;
params.llm.num_layers = 36;
params.llm.num_heads = 32;
params.llm.num_kv_heads = 8;
params.llm.head_dim = 128;
params.llm.qkv_bias = false;
params.llm.qk_norm = true;
params.llm.vocab_size = 151936;
params.llm.rms_norm_eps = 1e-6f;
params.llm.vision.arch = LLM::LLMVisionArch::QWEN3_VL;
params.llm.vision.num_layers = 27;
params.llm.vision.hidden_size = 1152;
params.llm.vision.intermediate_size = 4304;
params.llm.vision.num_heads = 16;
params.llm.vision.out_hidden_size = 4096;
params.llm.vision.patch_size = 16;
params.llm.vision.spatial_merge_size = 2;
params.llm.vision.temporal_patch_size = 2;
params.llm.vision.num_position_embeddings = 2304;
return params;
}
struct HiDreamO1Model : public GGMLBlock {
HiDreamO1Params params;
HiDreamO1Config config;
HiDreamO1Model() = default;
explicit HiDreamO1Model(HiDreamO1Params params)
: params(std::move(params)) {
blocks["language_model"] = std::make_shared<LLM::TextModel>(this->params.llm);
blocks["t_embedder1"] = std::make_shared<TimestepEmbedder>(this->params.llm.hidden_size);
blocks["x_embedder"] = std::make_shared<BottleneckPatchEmbed>(this->params.patch_size * this->params.patch_size * 3,
this->params.llm.hidden_size / 4,
this->params.llm.hidden_size);
blocks["final_layer2"] = std::make_shared<FinalLayer>(this->params.llm.hidden_size,
this->params.patch_size * this->params.patch_size * 3);
explicit HiDreamO1Model(HiDreamO1Config config)
: config(std::move(config)) {
blocks["language_model"] = std::make_shared<LLM::TextModel>(this->config.llm);
blocks["t_embedder1"] = std::make_shared<TimestepEmbedder>(this->config.llm.hidden_size);
blocks["x_embedder"] = std::make_shared<BottleneckPatchEmbed>(this->config.patch_size * this->config.patch_size * 3,
this->config.llm.hidden_size / 4,
this->config.llm.hidden_size);
blocks["final_layer2"] = std::make_shared<FinalLayer>(this->config.llm.hidden_size,
this->config.patch_size * this->config.patch_size * 3);
}
std::shared_ptr<LLM::TextModel> text_model() {
@@ -269,7 +271,7 @@ namespace HiDreamO1 {
};
struct HiDreamO1VisionRunner : public GGMLRunner {
HiDreamO1Params params;
HiDreamO1Config config;
std::shared_ptr<LLM::VisionModel> model;
std::vector<int> window_index_vec;
@@ -284,8 +286,8 @@ namespace HiDreamO1 {
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "model.visual")
: GGMLRunner(backend, params_backend),
params(make_hidream_o1_params()),
model(std::make_shared<LLM::VisionModel>(false, params.llm.vision)) {
config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)),
model(std::make_shared<LLM::VisionModel>(false, config.llm.vision)) {
model->init(params_ctx, tensor_storage_map, prefix);
}
@@ -302,7 +304,7 @@ namespace HiDreamO1 {
compute_ctx,
runner_ctx,
image,
params.llm.vision,
config.llm.vision,
model,
window_index_vec,
window_inverse_index_vec,
@@ -331,7 +333,7 @@ namespace HiDreamO1 {
};
struct HiDreamO1Runner : public DiffusionModelRunner {
HiDreamO1Params params;
HiDreamO1Config config;
HiDreamO1Model model;
std::vector<float> attention_mask_vec;
@@ -341,8 +343,8 @@ namespace HiDreamO1 {
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "model")
: DiffusionModelRunner(backend, params_backend, prefix),
params(make_hidream_o1_params()) {
model = HiDreamO1Model(params);
config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)) {
model = HiDreamO1Model(config);
model.init(params_ctx, tensor_storage_map, prefix);
}
otherarch/sdcpp/ideogram4.hpp
+531
View File
@@ -0,0 +1,531 @@
#ifndef __IDEOGRAM4_HPP__
#define __IDEOGRAM4_HPP__
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include "diffusion_model.hpp"
#include "ggml_extend.hpp"
#include "ggml_graph_cut.h"
#include "rope.hpp"
namespace Ideogram4 {
constexpr int IDEOGRAM4_GRAPH_SIZE = 65536;
constexpr int OUTPUT_IMAGE_INDICATOR = 2;
constexpr int IMAGE_POSITION_OFFSET = 65536;
constexpr int DEFAULT_MROPE_SECTION_T = 24;
constexpr int DEFAULT_MROPE_SECTION_H = 20;
constexpr int DEFAULT_MROPE_SECTION_W = 20;
constexpr int TIMESTEP_MAX_PERIOD = 10000;
constexpr int LLM_HIDDEN_STATE_LAYERS = 13;
struct Ideogram4Config {
int64_t emb_dim = 4608;
int64_t num_layers = 34;
int64_t num_heads = 18;
int64_t intermediate_size = 12288;
int64_t adanln_dim = 512;
int64_t in_channels = 128;
int64_t llm_features_dim = 53248;
int64_t rope_theta = 5000000;
float norm_eps = 1e-5f;
int patch_size = 2;
int ae_channels = 32;
std::vector<int> mrope_section = {DEFAULT_MROPE_SECTION_T,
DEFAULT_MROPE_SECTION_H,
DEFAULT_MROPE_SECTION_W};
static Ideogram4Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
Ideogram4Config config;
int64_t detected_layers = 0;
std::string layer_prefix = prefix.empty() ? "layers." : prefix + ".layers.";
for (const auto& [name, _] : tensor_storage_map) {
if (name.find(layer_prefix) != 0) {
continue;
}
std::string tail = name.substr(layer_prefix.size());
size_t dot = tail.find('.');
if (dot == std::string::npos) {
continue;
}
int layer_idx = std::atoi(tail.substr(0, dot).c_str());
detected_layers = std::max<int64_t>(detected_layers, layer_idx + 1);
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
LOG_DEBUG("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64,
config.num_layers,
config.emb_dim,
config.num_heads,
config.intermediate_size);
}
return config;
}
};
__STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
ggml_tensor* timesteps,
int dim) {
GGML_ASSERT(dim % 2 == 0);
auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
auto chunks = ggml_ext_chunk(ctx, embedding, 2, 0);
return ggml_concat(ctx, chunks[1], chunks[0], 0);
}
__STATIC_INLINE__ ggml_tensor* to_token_modulation(ggml_context* ctx, ggml_tensor* x) {
// [N, C] -> [N, 1, C] in PyTorch layout.
if (ggml_n_dims(x) < 3 || x->ne[1] != 1) {
x = ggml_reshape_3d(ctx, x, x->ne[0], 1, x->ne[1]);
}
return x;
}
__STATIC_INLINE__ ggml_tensor* interleave_hidden_state_layers(ggml_context* ctx, ggml_tensor* x) {
// Match upstream stack(...).permute(1, 2, 3, 0).reshape(...):
// [layers * hidden, tokens, batch] -> [hidden * layers, tokens, batch].
GGML_ASSERT(x->ne[0] % LLM_HIDDEN_STATE_LAYERS == 0);
const int64_t hidden_size = x->ne[0] / LLM_HIDDEN_STATE_LAYERS;
const int64_t token_count = x->ne[1];
const int64_t batch_count = x->ne[2];
x = ggml_reshape_4d(ctx, x, hidden_size, LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
}
__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
scale = to_token_modulation(ctx, scale);
return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
}
__STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
// x: [N, 128, H, W] with channel order [ae, ph, pw].
// return: [N, H*W, 128] with token channel order [ph, pw, ae].
const int64_t W = x->ne[0];
const int64_t H = x->ne[1];
const int64_t C = x->ne[2];
const int64_t N = x->ne[3];
GGML_ASSERT(N == 1);
GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size);
x = ggml_cont(ctx, x);
x = ggml_reshape_4d(ctx, x, W * H, config.patch_size, config.patch_size, config.ae_channels);
x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0));
x = ggml_reshape_3d(ctx, x, C, W * H, N);
return x;
}
__STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx,
ggml_tensor* x,
int64_t H,
int64_t W,
const Ideogram4Config& config) {
const int64_t C = x->ne[0];
const int64_t N = x->ne[2];
GGML_ASSERT(N == 1);
GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size);
GGML_ASSERT(x->ne[1] == H * W);
x = ggml_reshape_4d(ctx, x, config.ae_channels, config.patch_size, config.patch_size, H * W);
x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0));
x = ggml_reshape_4d(ctx, x, W, H, C, N);
return x;
}
__STATIC_INLINE__ std::shared_ptr<Linear> make_linear(int64_t in_features,
int64_t out_features,
bool bias = true) {
return std::make_shared<Linear>(in_features, out_features, bias, false, false, 1.f, true);
}
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
int grid_w,
int bs,
int context_len,
int head_dim,
int rope_theta,
const std::vector<int>& mrope_section) {
GGML_ASSERT(bs == 1);
std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
std::vector<float>(3, 0.f));
for (int i = 0; i < context_len; ++i) {
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
}
int cursor = context_len;
for (int y = 0; y < grid_h; ++y) {
for (int x = 0; x < grid_w; ++x) {
ids[cursor++] = {static_cast<float>(IMAGE_POSITION_OFFSET),
static_cast<float>(IMAGE_POSITION_OFFSET + y),
static_cast<float>(IMAGE_POSITION_OFFSET + x)};
}
}
return Rope::embed_interleaved_mrope(ids, bs, static_cast<float>(rope_theta), head_dim, mrope_section);
}
class Ideogram4Attention : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t num_heads;
int64_t head_dim;
public:
Ideogram4Attention(int64_t hidden_size, int64_t num_heads, float eps)
: hidden_size(hidden_size), num_heads(num_heads), head_dim(hidden_size / num_heads) {
GGML_ASSERT(hidden_size % num_heads == 0);
blocks["qkv"] = make_linear(hidden_size, hidden_size * 3, false);
blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["o"] = make_linear(hidden_size, hidden_size, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* mask = nullptr) {
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
auto qkv = qkv_proj->forward(ctx, x);
auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N);
auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N);
auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N);
q = norm_q->forward(ctx, q);
k = norm_k->forward(ctx, k);
x = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f, false);
x = out_proj->forward(ctx, x);
return x;
}
};
class Ideogram4MLP : public GGMLBlock {
public:
Ideogram4MLP(int64_t dim, int64_t hidden_dim) {
blocks["w1"] = make_linear(dim, hidden_dim, false);
blocks["w2"] = make_linear(hidden_dim, dim, false);
blocks["w3"] = make_linear(dim, hidden_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto x1 = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
auto x3 = w3->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, x1, x3);
x = w2->forward(ctx, x);
return x;
}
};
class Ideogram4TransformerBlock : public GGMLBlock {
public:
Ideogram4TransformerBlock(const Ideogram4Config& config) {
blocks["attention"] = std::make_shared<Ideogram4Attention>(config.emb_dim, config.num_heads, config.norm_eps);
blocks["feed_forward"] = std::make_shared<Ideogram4MLP>(config.emb_dim, config.intermediate_size);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
blocks["adaln_modulation"] = make_linear(config.adanln_dim, 4 * config.emb_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* adaln_input,
ggml_tensor* mask = nullptr) {
auto attention = std::dynamic_pointer_cast<Ideogram4Attention>(blocks["attention"]);
auto feed_forward = std::dynamic_pointer_cast<Ideogram4MLP>(blocks["feed_forward"]);
auto attention_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
auto ffn_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
auto attention_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm2"]);
auto ffn_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
auto mod = adaln_modulation->forward(ctx, adaln_input);
auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0);
auto scale_msa = mods[0];
auto gate_msa = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1]));
auto scale_mlp = mods[2];
auto gate_mlp = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3]));
auto attn_out = attention_norm1->forward(ctx, x);
attn_out = modulate(ctx->ggml_ctx, attn_out, scale_msa);
attn_out = attention->forward(ctx, attn_out, pe, mask);
attn_out = attention_norm2->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
auto ffn_out = ffn_norm1->forward(ctx, x);
ffn_out = modulate(ctx->ggml_ctx, ffn_out, scale_mlp);
ffn_out = feed_forward->forward(ctx, ffn_out);
ffn_out = ffn_norm2->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp));
return x;
}
};
class Ideogram4EmbedScalar : public GGMLBlock {
protected:
int64_t dim;
public:
Ideogram4EmbedScalar(int64_t dim)
: dim(dim) {
blocks["mlp_in"] = make_linear(dim, dim, true);
blocks["mlp_out"] = make_linear(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto mlp_in = std::dynamic_pointer_cast<Linear>(blocks["mlp_in"]);
auto mlp_out = std::dynamic_pointer_cast<Linear>(blocks["mlp_out"]);
x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast<int>(dim));
x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x));
x = mlp_out->forward(ctx, x);
return x;
}
};
class Ideogram4FinalLayer : public GGMLBlock {
public:
Ideogram4FinalLayer(const Ideogram4Config& config) {
blocks["norm_final"] = std::make_shared<LayerNorm>(config.emb_dim, 1e-6f, false);
blocks["linear"] = make_linear(config.emb_dim, config.in_channels, true);
blocks["adaln_modulation"] = make_linear(config.adanln_dim, config.emb_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* c) {
auto norm_final = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
x = norm_final->forward(ctx, x);
x = modulate(ctx->ggml_ctx, x, scale);
x = linear->forward(ctx, x);
return x;
}
};
class Ideogram4Transformer : public GGMLBlock {
protected:
Ideogram4Config config;
public:
Ideogram4Transformer() = default;
explicit Ideogram4Transformer(Ideogram4Config config)
: config(std::move(config)) {
blocks["input_proj"] = make_linear(this->config.in_channels, this->config.emb_dim, true);
blocks["llm_cond_norm"] = std::make_shared<RMSNorm>(this->config.llm_features_dim, 1e-6f);
blocks["llm_cond_proj"] = make_linear(this->config.llm_features_dim, this->config.emb_dim, true);
blocks["t_embedding"] = std::make_shared<Ideogram4EmbedScalar>(this->config.emb_dim);
blocks["adaln_proj"] = make_linear(this->config.emb_dim, this->config.adanln_dim, true);
blocks["embed_image_indicator"] = std::make_shared<Embedding>(2, this->config.emb_dim);
for (int i = 0; i < this->config.num_layers; ++i) {
blocks["layers." + std::to_string(i)] = std::make_shared<Ideogram4TransformerBlock>(this->config);
}
blocks["final_layer"] = std::make_shared<Ideogram4FinalLayer>(this->config);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* image_indicator_ids) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t N = x->ne[3];
GGML_ASSERT(N == 1);
auto input_proj = std::dynamic_pointer_cast<Linear>(blocks["input_proj"]);
auto llm_cond_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["llm_cond_norm"]);
auto llm_cond_proj = std::dynamic_pointer_cast<Linear>(blocks["llm_cond_proj"]);
auto t_embedding = std::dynamic_pointer_cast<Ideogram4EmbedScalar>(blocks["t_embedding"]);
auto adaln_proj = std::dynamic_pointer_cast<Linear>(blocks["adaln_proj"]);
auto embed_image_indicator = std::dynamic_pointer_cast<Embedding>(blocks["embed_image_indicator"]);
auto final_layer = std::dynamic_pointer_cast<Ideogram4FinalLayer>(blocks["final_layer"]);
auto img = patchify(ctx->ggml_ctx, x, config);
img = input_proj->forward(ctx, img);
ggml_tensor* h = img;
int64_t context_len = 0;
if (context != nullptr) {
if (ggml_n_dims(context) < 3) {
context = ggml_reshape_3d(ctx->ggml_ctx, context, context->ne[0], context->ne[1], 1);
}
context = interleave_hidden_state_layers(ctx->ggml_ctx, context);
context_len = context->ne[1];
auto txt = llm_cond_norm->forward(ctx, context);
txt = llm_cond_proj->forward(ctx, txt);
h = ggml_concat(ctx->ggml_ctx, txt, img, 1);
}
auto indicator_embedding = embed_image_indicator->forward(ctx, image_indicator_ids);
h = ggml_add(ctx->ggml_ctx, h, indicator_embedding);
auto t_cond = t_embedding->forward(ctx, timestep);
auto adaln_input = ggml_silu(ctx->ggml_ctx, adaln_proj->forward(ctx, t_cond));
for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<Ideogram4TransformerBlock>(blocks["layers." + std::to_string(i)]);
h = block->forward(ctx, h, pe, adaln_input, nullptr);
sd::ggml_graph_cut::mark_graph_cut(h, "ideogram4.layers." + std::to_string(i), "hidden");
}
h = final_layer->forward(ctx, h, adaln_input);
if (context_len > 0) {
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]);
}
h = unpatchify(ctx->ggml_ctx, h, H, W, config);
h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f);
return h;
}
};
class Ideogram4Runner : public DiffusionModelRunner {
protected:
bool should_use_uncond_model(const DiffusionParams& diffusion_params) const {
return has_uncond_model &&
diffusion_params.context == nullptr &&
diffusion_params.y != nullptr &&
!diffusion_params.y->empty();
}
public:
Ideogram4Config config;
Ideogram4Transformer model;
Ideogram4Transformer uncond_model;
bool has_uncond_model = false;
std::string uncond_prefix;
std::vector<float> pe_vec;
std::vector<int32_t> image_indicator_vec;
Ideogram4Runner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix),
config(Ideogram4Config::detect_from_weights(tensor_storage_map, prefix)),
uncond_prefix(prefix + ".uncond") {
model = Ideogram4Transformer(config);
model.init(params_ctx, tensor_storage_map, prefix);
for (const auto& pair : tensor_storage_map) {
const std::string& name = pair.first;
if (starts_with(name, uncond_prefix)) {
has_uncond_model = true;
break;
}
}
if (has_uncond_model) {
LOG_DEBUG("using uncond model");
uncond_model = Ideogram4Transformer(config);
uncond_model.init(params_ctx, tensor_storage_map, uncond_prefix);
}
}
std::string get_desc() override {
return "ideogram4";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.get_param_tensors(tensors, prefix);
if (has_uncond_model) {
uncond_model.get_param_tensors(tensors, this->uncond_prefix);
}
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
bool use_uncond_model = false) {
ggml_cgraph* gf = new_graph_custom(IDEOGRAM4_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x->ne[3] == 1);
Ideogram4Transformer& active_model = use_uncond_model ? uncond_model : model;
ggml_tensor* context = nullptr;
int64_t context_len = 0;
if (!context_tensor.empty()) {
context = make_input(context_tensor);
context_len = context->ne[1];
}
int64_t grid_w = x->ne[0];
int64_t grid_h = x->ne[1];
int64_t pos_len = context_len + grid_h * grid_w;
int64_t head_dim = config.emb_dim / config.num_heads;
pe_vec = gen_ideogram4_pe(static_cast<int>(grid_h),
static_cast<int>(grid_w),
static_cast<int>(x->ne[3]),
static_cast<int>(context_len),
static_cast<int>(head_dim),
static_cast<int>(config.rope_theta),
config.mrope_section);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
image_indicator_vec.assign(static_cast<size_t>(pos_len), 1);
for (int64_t i = 0; i < context_len; ++i) {
image_indicator_vec[static_cast<size_t>(i)] = 0;
}
auto indicator = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_I32, pos_len, x->ne[3]);
set_backend_tensor_data(indicator, image_indicator_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = active_model.forward(&runner_ctx, x, timesteps, context, pe, indicator);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
bool use_uncond_model = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, use_uncond_model);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
bool use_uncond_model = should_use_uncond_model(diffusion_params);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
use_uncond_model);
}
};
} // namespace Ideogram4
#endif // __IDEOGRAM4_HPP__
otherarch/sdcpp/layer_registry.cpp
+132
View File
@@ -0,0 +1,132 @@
#include "layer_registry.h"
#include <utility>
#include "util.h"
namespace sd::layer_registry {
void LayerRegistry::register_layer(const std::string& name, ggml_tensor* tensor) {
auto& info = layers_[name];
info.tensors.push_back(tensor);
info.bytes += ggml_nbytes(tensor);
}
bool LayerRegistry::move_layer_to_gpu(const std::string& name) {
auto it = layers_.find(name);
if (it == layers_.end())
return false;
LayerInfo& info = it->second;
if (info.on_gpu)
return true;
if (gpu_backend_ == nullptr || cpu_backend_ == nullptr) {
LOG_ERROR("layer_registry: backends not set; cannot move '%s' to GPU",
name.c_str());
return false;
}
if (info.tensors.empty()) {
info.on_gpu = true;
return true;
}
// 1. Build a no_alloc context big enough to hold one twin tensor per CPU
// tensor, plus a little overhead.
const size_t ctx_size = info.tensors.size() * ggml_tensor_overhead() + 1024;
ggml_init_params ctx_params{ctx_size, /*mem_buffer=*/nullptr, /*no_alloc=*/true};
ggml_context* twin_ctx = ggml_init(ctx_params);
if (twin_ctx == nullptr) {
LOG_ERROR("layer_registry: failed to allocate twin context for '%s'",
name.c_str());
return false;
}
// 2. Create one GPU twin per CPU tensor. The twin shares the original
// name so any name-based lookup keeps working.
std::vector<ggml_tensor*> gpu_twins;
gpu_twins.reserve(info.tensors.size());
for (ggml_tensor* cpu_t : info.tensors) {
ggml_tensor* twin = ggml_dup_tensor(twin_ctx, cpu_t);
if (cpu_t->name[0] != '\0') {
ggml_set_name(twin, cpu_t->name);
}
gpu_twins.push_back(twin);
}
// 3. Back the twins with a GPU buffer in one alloc call.
ggml_backend_buffer_t gpu_buffer = ggml_backend_alloc_ctx_tensors(twin_ctx, gpu_backend_);
if (gpu_buffer == nullptr) {
LOG_ERROR("layer_registry: failed to allocate GPU buffer for '%s'",
name.c_str());
ggml_free(twin_ctx);
return false;
}
// 4. H2D copy + sync.
for (size_t i = 0; i < info.tensors.size(); ++i) {
ggml_backend_tensor_copy(info.tensors[i], gpu_twins[i]);
}
ggml_backend_synchronize(gpu_backend_);
// 5. Swap buffer/data/extra so the originals now point at GPU memory.
for (size_t i = 0; i < info.tensors.size(); ++i) {
std::swap(info.tensors[i]->buffer, gpu_twins[i]->buffer);
std::swap(info.tensors[i]->data, gpu_twins[i]->data);
std::swap(info.tensors[i]->extra, gpu_twins[i]->extra);
}
info.gpu_twins = std::move(gpu_twins);
info.twin_ctx = twin_ctx;
info.gpu_buffer = gpu_buffer;
info.on_gpu = true;
return true;
}
bool LayerRegistry::move_layer_to_cpu(const std::string& name) {
auto it = layers_.find(name);
if (it == layers_.end())
return false;
LayerInfo& info = it->second;
if (!info.on_gpu)
return true;
if (info.tensors.size() != info.gpu_twins.size()) {
LOG_ERROR("layer_registry: twin/tensor count mismatch for '%s'",
name.c_str());
return false;
}
// 1. Swap back: originals point at CPU memory again.
for (size_t i = 0; i < info.tensors.size(); ++i) {
if (info.gpu_twins[i] == nullptr)
continue;
std::swap(info.tensors[i]->buffer, info.gpu_twins[i]->buffer);
std::swap(info.tensors[i]->data, info.gpu_twins[i]->data);
std::swap(info.tensors[i]->extra, info.gpu_twins[i]->extra);
}
// 2. Free the GPU buffer + twin context.
if (info.gpu_buffer != nullptr) {
ggml_backend_buffer_free(info.gpu_buffer);
info.gpu_buffer = nullptr;
}
if (info.twin_ctx != nullptr) {
ggml_free(info.twin_ctx);
info.twin_ctx = nullptr;
}
info.gpu_twins.clear();
info.on_gpu = false;
return true;
}
bool LayerRegistry::is_layer_on_gpu(const std::string& name) const {
auto it = layers_.find(name);
return it != layers_.end() && it->second.on_gpu;
}
size_t LayerRegistry::get_layer_size(const std::string& name) const {
auto it = layers_.find(name);
return it != layers_.end() ? it->second.bytes : 0;
}
} // namespace sd::layer_registry
otherarch/sdcpp/layer_registry.h
+50
View File
@@ -0,0 +1,50 @@
#ifndef __LAYER_REGISTRY_H__
#define __LAYER_REGISTRY_H__
#include <map>
#include <set>
#include <string>
#include <vector>
#include "ggml-backend.h"
#include "ggml.h"
namespace sd::layer_registry {
struct LayerInfo {
std::vector<ggml_tensor*> tensors;
std::vector<ggml_tensor*> gpu_twins;
ggml_context* twin_ctx = nullptr;
ggml_backend_buffer_t gpu_buffer = nullptr;
bool on_gpu = false;
size_t bytes = 0;
};
class LayerRegistry {
public:
LayerRegistry() = default;
LayerRegistry(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend)
: gpu_backend_(gpu_backend), cpu_backend_(cpu_backend) {}
void set_backends(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend) {
gpu_backend_ = gpu_backend;
cpu_backend_ = cpu_backend;
}
void register_layer(const std::string& name, ggml_tensor* tensor);
bool move_layer_to_gpu(const std::string& name);
bool move_layer_to_cpu(const std::string& name);
bool is_layer_on_gpu(const std::string& name) const;
size_t get_layer_size(const std::string& name) const;
size_t get_layer_count() const { return layers_.size(); }
const std::map<std::string, LayerInfo>& layers() const { return layers_; }
private:
ggml_backend_t gpu_backend_ = nullptr;
ggml_backend_t cpu_backend_ = nullptr;
std::map<std::string, LayerInfo> layers_;
};
} // namespace sd::layer_registry
#endif
otherarch/sdcpp/lens.hpp
+89 -82
View File
@@ -13,6 +13,71 @@
namespace Lens {
constexpr int LENS_GRAPH_SIZE = 40960;
struct LensConfig {
int patch_size = 2;
int64_t in_channels = 128;
int64_t out_channels = 32;
int num_layers = 48;
int64_t attention_head_dim = 64;
int64_t num_attention_heads = 24;
int64_t joint_attention_dim = 2880;
int selected_layer_count = 4;
int theta = 10000;
std::vector<int> axes_dim = {8, 28, 28};
int axes_dim_sum = 64;
static LensConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
LensConfig config;
config.num_layers = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
config.in_channels = tensor_storage.ne[0];
int64_t inner_dim = tensor_storage.ne[1];
if (config.attention_head_dim > 0) {
config.num_attention_heads = inner_dim / config.attention_head_dim;
}
} else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
config.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / config.joint_attention_dim);
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.out_channels = tensor_storage.ne[1] / patch_area;
} else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
config.attention_head_dim = tensor_storage.ne[0];
}
size_t pos = name.find("transformer_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
}
if (config.num_layers == 0) {
config.num_layers = 48;
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
LOG_DEBUG("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
config.num_layers,
config.selected_layer_count,
config.num_attention_heads * config.attention_head_dim,
config.num_attention_heads,
config.attention_head_dim,
config.in_channels,
config.out_channels);
return config;
}
};
struct LensTimestepProjEmbeddings : public GGMLBlock {
LensTimestepProjEmbeddings(int64_t embedding_dim) {
blocks["timestep_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(256, embedding_dim);
@@ -209,41 +274,27 @@ namespace Lens {
}
};
struct LensParams {
int patch_size = 2;
int64_t in_channels = 128;
int64_t out_channels = 32;
int num_layers = 48;
int64_t attention_head_dim = 64;
int64_t num_attention_heads = 24;
int64_t joint_attention_dim = 2880;
int selected_layer_count = 4;
int theta = 10000;
std::vector<int> axes_dim = {8, 28, 28};
int axes_dim_sum = 64;
};
class LensModel : public GGMLBlock {
public:
LensParams params;
LensConfig config;
LensModel() = default;
LensModel(LensParams params)
: params(params) {
int64_t inner_dim = params.num_attention_heads * params.attention_head_dim;
LensModel(LensConfig config)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::make_shared<LensTimestepProjEmbeddings>(inner_dim);
blocks["img_in"] = std::make_shared<Linear>(params.in_channels, inner_dim, true);
blocks["txt_in"] = std::make_shared<Linear>(params.joint_attention_dim * params.selected_layer_count, inner_dim, true);
for (int i = 0; i < params.selected_layer_count; ++i) {
blocks["txt_norm." + std::to_string(i)] = std::make_shared<RMSNorm>(params.joint_attention_dim, 1e-5f);
blocks["img_in"] = std::make_shared<Linear>(config.in_channels, inner_dim, true);
blocks["txt_in"] = std::make_shared<Linear>(config.joint_attention_dim * config.selected_layer_count, inner_dim, true);
for (int i = 0; i < config.selected_layer_count; ++i) {
blocks["txt_norm." + std::to_string(i)] = std::make_shared<RMSNorm>(config.joint_attention_dim, 1e-5f);
}
for (int i = 0; i < params.num_layers; ++i) {
for (int i = 0; i < config.num_layers; ++i) {
blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<LensTransformerBlock>(inner_dim,
params.num_attention_heads,
params.attention_head_dim);
config.num_attention_heads,
config.attention_head_dim);
}
blocks["norm_out"] = std::make_shared<LensAdaLayerNormContinuous>(inner_dim, 1e-6f);
blocks["proj_out"] = std::make_shared<Linear>(inner_dim, params.patch_size * params.patch_size * params.out_channels, true);
blocks["proj_out"] = std::make_shared<Linear>(inner_dim, config.patch_size * config.patch_size * config.out_channels, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -269,9 +320,9 @@ namespace Lens {
img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
img = img_in->forward(ctx, img);
std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, params.selected_layer_count, 0);
std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0);
ggml_tensor* txt = nullptr;
for (int i = 0; i < params.selected_layer_count; ++i) {
for (int i = 0; i < config.selected_layer_count; ++i) {
auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm." + std::to_string(i)]);
auto chunk = txt_norm->forward(ctx, txt_chunks[i]);
txt = txt == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, txt, chunk, 0);
@@ -281,7 +332,7 @@ namespace Lens {
sd::ggml_graph_cut::mark_graph_cut(img, "lens.prelude", "img");
sd::ggml_graph_cut::mark_graph_cut(txt, "lens.prelude", "txt");
for (int i = 0; i < params.num_layers; ++i) {
for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<LensTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
auto out = block->forward(ctx, img, txt, t_emb, pe);
img = out.first;
@@ -294,13 +345,13 @@ namespace Lens {
img = proj_out->forward(ctx, img);
auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
out = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, params.patch_size * params.patch_size * params.out_channels, N);
out = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N);
return out;
}
};
struct LensRunner : public DiffusionModelRunner {
LensParams lens_params;
LensConfig config;
LensModel lens;
std::vector<float> pe_vec;
@@ -308,53 +359,9 @@ namespace Lens {
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix) {
lens_params.num_layers = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
lens_params.in_channels = tensor_storage.ne[0];
int64_t inner_dim = tensor_storage.ne[1];
lens_params.num_attention_heads = inner_dim / lens_params.attention_head_dim;
} else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
lens_params.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / lens_params.joint_attention_dim);
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
lens_params.out_channels = tensor_storage.ne[1] / lens_params.patch_size / lens_params.patch_size;
} else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
lens_params.attention_head_dim = tensor_storage.ne[0];
}
size_t pos = name.find("transformer_blocks.");
if (pos != std::string::npos) {
std::string layer_name = name.substr(pos);
auto items = split_string(layer_name, '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > lens_params.num_layers) {
lens_params.num_layers = block_index + 1;
}
}
}
}
if (lens_params.num_layers == 0) {
lens_params.num_layers = 48;
}
lens_params.axes_dim_sum = 0;
for (int axis_dim : lens_params.axes_dim) {
lens_params.axes_dim_sum += axis_dim;
}
LOG_INFO("lens: layers = %d, in_channels = %" PRId64 ", out_channels = %" PRId64
", heads = %" PRId64 ", head_dim = %" PRId64,
lens_params.num_layers,
lens_params.in_channels,
lens_params.out_channels,
lens_params.num_attention_heads,
lens_params.attention_head_dim);
lens = LensModel(lens_params);
: DiffusionModelRunner(backend, params_backend, prefix),
config(LensConfig::detect_from_weights(tensor_storage_map, prefix)) {
lens = LensModel(config);
lens.init(params_ctx, tensor_storage_map, prefix);
}
@@ -380,12 +387,12 @@ namespace Lens {
static_cast<int>(x->ne[0]),
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
lens_params.theta,
config.theta,
circular_y_enabled,
circular_x_enabled,
lens_params.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / lens_params.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, lens_params.axes_dim_sum / 2, pos_len);
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
otherarch/sdcpp/llm.hpp
+246 -182
View File
@@ -37,6 +37,7 @@ namespace LLM {
MISTRAL_SMALL_3_2,
MINISTRAL_3_3B,
GEMMA3_12B,
GEMMA2_2B,
GPT_OSS_20B,
ARCH_COUNT,
};
@@ -48,6 +49,7 @@ namespace LLM {
"mistral_small3.2",
"ministral3.3b",
"gemma3_12b",
"gemma2_2b",
"gpt_oss_20b",
};
@@ -61,7 +63,7 @@ namespace LLM {
QWEN3_VL,
};
struct LLMVisionParams {
struct LLMVisionConfig {
LLMVisionArch arch = LLMVisionArch::QWEN2_5_VL;
int num_layers = 32;
int64_t hidden_size = 1280;
@@ -77,7 +79,7 @@ namespace LLM {
std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
};
struct LLMParams {
struct LLMConfig {
LLMArch arch = LLMArch::QWEN2_5_VL;
int64_t num_layers = 28;
int64_t hidden_size = 3584;
@@ -99,7 +101,129 @@ namespace LLM {
std::vector<int> sliding_attention;
int64_t num_experts = 0;
int64_t num_experts_per_tok = 0;
LLMVisionParams vision;
LLMVisionConfig vision;
bool have_vision_weight = false;
bool llama_cpp_style = false;
static LLMConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
LLMArch arch) {
LLMConfig config;
config.arch = arch;
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
config.head_dim = 128;
config.num_heads = 32;
config.num_kv_heads = 8;
config.qkv_bias = false;
config.rms_norm_eps = 1e-5f;
} else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) {
config.head_dim = 128;
config.num_heads = 32;
config.num_kv_heads = 8;
config.qkv_bias = false;
config.qk_norm = true;
config.rms_norm_eps = 1e-6f;
if (arch == LLMArch::QWEN3_VL) {
config.max_position_embeddings = 262144;
config.rope_thetas = {5000000.f};
config.vision.arch = LLMVisionArch::QWEN3_VL;
}
} else if (arch == LLMArch::GEMMA3_12B) {
config.head_dim = 256;
config.num_heads = 16;
config.num_kv_heads = 8;
config.qkv_bias = false;
config.qk_norm = true;
config.rms_norm_eps = 1e-6f;
config.rms_norm_add = false;
config.normalize_input = true;
config.max_position_embeddings = 131072;
config.mlp_activation = MLPActivation::GELU_TANH;
config.rope_thetas = {1000000.f, 10000.f};
config.rope_scales = {8.f, 1.f};
config.sliding_attention = {1024, 1024, 1024, 1024, 1024, 0};
} else if (arch == LLMArch::GEMMA2_2B) {
config.head_dim = 256;
config.num_heads = 8;
config.num_kv_heads = 4;
config.qkv_bias = false;
config.qk_norm = false;
config.rms_norm_eps = 1e-6f;
config.rms_norm_add = true;
config.normalize_input = true;
config.max_position_embeddings = 8192;
config.mlp_activation = MLPActivation::GELU_TANH;
config.hidden_size = 2304;
config.intermediate_size = 9216;
config.num_layers = 26;
config.vocab_size = 256000;
} else if (arch == LLMArch::GPT_OSS_20B) {
config.head_dim = 64;
config.num_heads = 64;
config.num_kv_heads = 8;
config.qkv_bias = true;
config.attention_out_bias = true;
config.qk_norm = false;
config.rms_norm_eps = 1e-5f;
config.hidden_size = 2880;
config.intermediate_size = 2880;
config.num_layers = 24;
config.vocab_size = 201088;
config.max_position_embeddings = 131072;
config.rope_thetas = {150000.f};
config.rope_scales = {32.f};
config.sliding_attention = {128, 0};
config.num_experts = 32;
config.num_experts_per_tok = 4;
}
config.num_layers = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
size_t pos = name.find("visual.");
if (pos != std::string::npos) {
config.have_vision_weight = true;
if (contains(name, "attn.q_proj")) {
config.llama_cpp_style = true;
}
continue;
}
pos = name.find("layers.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
if (contains(name, "embed_tokens.weight")) {
config.hidden_size = tensor_storage.ne[0];
config.vocab_size = tensor_storage.ne[1];
}
if (contains(name, "layers.0.mlp.gate_proj.weight")) {
config.intermediate_size = tensor_storage.ne[1];
}
if (contains(name, "layers.0.mlp.experts.gate_up_proj.weight")) {
config.intermediate_size = tensor_storage.ne[1] / 2;
}
if (contains(name, "layers.0.mlp.experts.gate_proj.weight")) {
config.intermediate_size = tensor_storage.ne[1];
}
}
if (arch == LLMArch::QWEN3 && config.num_layers == 28) {
config.num_heads = 16;
}
LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
config.num_layers,
config.vocab_size,
config.hidden_size,
config.intermediate_size);
return config;
}
};
struct LLMRMSNorm : public UnaryBlock {
@@ -230,11 +354,11 @@ namespace LLM {
}
public:
GPTOSSMLP(const LLMParams& params)
: hidden_size(params.hidden_size),
intermediate_size(params.intermediate_size),
num_experts(params.num_experts),
num_experts_per_tok(params.num_experts_per_tok) {}
GPTOSSMLP(const LLMConfig& config)
: hidden_size(config.hidden_size),
intermediate_size(config.intermediate_size),
num_experts(config.num_experts),
num_experts_per_tok(config.num_experts_per_tok) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
@@ -665,7 +789,7 @@ namespace LLM {
public:
VisionModel(bool llama_cpp_style,
const LLMVisionParams& vision_params,
const LLMVisionConfig& vision_params,
float eps = 1e-6f)
: arch_(vision_params.arch),
num_layers(vision_params.num_layers),
@@ -782,23 +906,23 @@ namespace LLM {
}
public:
Attention(const LLMParams& params)
: arch(params.arch),
num_heads(params.num_heads),
num_kv_heads(params.num_kv_heads),
head_dim(params.head_dim),
qk_norm(params.qk_norm),
max_position_embeddings(params.max_position_embeddings),
rope_thetas(params.rope_thetas),
rope_scales(params.rope_scales),
has_attention_sinks(params.arch == LLMArch::GPT_OSS_20B) {
blocks["q_proj"] = std::make_shared<Linear>(params.hidden_size, num_heads * head_dim, params.qkv_bias);
blocks["k_proj"] = std::make_shared<Linear>(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias);
blocks["v_proj"] = std::make_shared<Linear>(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias);
blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, params.hidden_size, params.attention_out_bias);
if (params.qk_norm) {
blocks["q_norm"] = std::make_shared<LLMRMSNorm>(head_dim, params.rms_norm_eps, params.rms_norm_add);
blocks["k_norm"] = std::make_shared<LLMRMSNorm>(head_dim, params.rms_norm_eps, params.rms_norm_add);
Attention(const LLMConfig& config)
: arch(config.arch),
num_heads(config.num_heads),
num_kv_heads(config.num_kv_heads),
head_dim(config.head_dim),
qk_norm(config.qk_norm),
max_position_embeddings(config.max_position_embeddings),
rope_thetas(config.rope_thetas),
rope_scales(config.rope_scales),
has_attention_sinks(config.arch == LLMArch::GPT_OSS_20B) {
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, config.hidden_size, config.attention_out_bias);
if (config.qk_norm) {
blocks["q_norm"] = std::make_shared<LLMRMSNorm>(head_dim, config.rms_norm_eps, config.rms_norm_add);
blocks["k_norm"] = std::make_shared<LLMRMSNorm>(head_dim, config.rms_norm_eps, config.rms_norm_add);
}
}
@@ -900,6 +1024,33 @@ namespace LLM {
1.f,
32.f,
1.f);
} else if (arch == LLMArch::GEMMA2_2B) {
q = ggml_rope_ext(ctx->ggml_ctx,
q,
input_pos,
nullptr,
head_dim,
GGML_ROPE_TYPE_NEOX,
8192,
10000.f,
1.f,
0.f,
1.f,
32.f,
1.f);
k = ggml_rope_ext(ctx->ggml_ctx,
k,
input_pos,
nullptr,
head_dim,
GGML_ROPE_TYPE_NEOX,
8192,
10000.f,
1.f,
0.f,
1.f,
32.f,
1.f);
} else if (arch == LLMArch::QWEN3_VL) {
int sections[4] = {24, 20, 20, 0};
q = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -953,34 +1104,42 @@ namespace LLM {
std::string post_ffw_norm_name;
public:
TransformerBlock(const LLMParams& params, int layer_index)
: arch(params.arch),
TransformerBlock(const LLMConfig& config, int layer_index)
: arch(config.arch),
sliding_attention(0) {
if (params.arch == LLMArch::GEMMA3_12B) {
post_attention_norm_name = "post_attention_norm";
post_ffw_norm_name = "post_ffw_norm";
}
pre_ffw_norm_name = params.arch == LLMArch::GPT_OSS_20B ? "post_attention_norm" : "post_attention_layernorm";
blocks["self_attn"] = std::make_shared<Attention>(params);
if (params.arch == LLMArch::GPT_OSS_20B) {
blocks["mlp"] = std::make_shared<GPTOSSMLP>(params);
if (config.arch == LLMArch::GEMMA3_12B) {
post_attention_norm_name = "post_attention_norm"; // attn_post_norm
pre_ffw_norm_name = "post_attention_layernorm"; // ffn_norm
post_ffw_norm_name = "post_ffw_norm"; // ffn_post_norm
} else if (config.arch == LLMArch::GEMMA2_2B) {
post_attention_norm_name = "post_attention_layernorm"; // ffn_norm
pre_ffw_norm_name = "pre_feedforward_layernorm";
post_ffw_norm_name = "post_feedforward_layernorm";
} else if (config.arch == LLMArch::GPT_OSS_20B) {
pre_ffw_norm_name = "post_attention_norm"; // attn_post_norm
} else {
blocks["mlp"] = std::make_shared<MLP>(params.hidden_size,
params.intermediate_size,
false,
params.mlp_activation);
pre_ffw_norm_name = "post_attention_layernorm"; // ffn_norm
}
blocks["input_layernorm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
blocks[pre_ffw_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
blocks["self_attn"] = std::make_shared<Attention>(config);
if (config.arch == LLMArch::GPT_OSS_20B) {
blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
} else {
blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
config.intermediate_size,
false,
config.mlp_activation);
}
blocks["input_layernorm"] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
blocks[pre_ffw_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
if (!post_attention_norm_name.empty()) {
blocks[post_attention_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
blocks[post_attention_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
}
if (!post_ffw_norm_name.empty()) {
blocks[post_ffw_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
blocks[post_ffw_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
}
if (!params.sliding_attention.empty()) {
sliding_attention = params.sliding_attention[layer_index % params.sliding_attention.size()];
if (!config.sliding_attention.empty()) {
sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()];
}
}
@@ -1037,16 +1196,16 @@ namespace LLM {
struct TextModel : public GGMLBlock {
protected:
int64_t num_layers;
LLMParams params;
LLMConfig config;
public:
TextModel(const LLMParams& params)
: num_layers(params.num_layers), params(params) {
blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(params.vocab_size, params.hidden_size));
TextModel(const LLMConfig& config)
: num_layers(config.num_layers), config(config) {
blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(config.vocab_size, config.hidden_size));
for (int i = 0; i < num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(params, i));
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
}
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(params.hidden_size, params.rms_norm_eps, params.rms_norm_add));
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
}
ggml_tensor* embed(GGMLRunnerContext* ctx,
@@ -1066,8 +1225,8 @@ namespace LLM {
auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
std::vector<ggml_tensor*> intermediate_outputs;
if (params.normalize_input) {
x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast<float>(params.hidden_size)), true);
if (config.normalize_input) {
x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast<float>(config.hidden_size)), true);
}
if (return_all_hidden_states) {
intermediate_outputs.push_back(x);
@@ -1137,15 +1296,15 @@ namespace LLM {
struct LLM : public GGMLBlock {
bool enable_vision;
LLMParams params;
LLMConfig config;
public:
LLM() = default;
LLM(LLMParams params, bool enable_vision = false, bool llama_cpp_style = false)
: enable_vision(enable_vision), params(params) {
blocks["model"] = std::shared_ptr<GGMLBlock>(new TextModel(params));
LLM(LLMConfig config, bool enable_vision = false, bool llama_cpp_style = false)
: enable_vision(enable_vision), config(config) {
blocks["model"] = std::shared_ptr<GGMLBlock>(new TextModel(config));
if (enable_vision) {
blocks["visual"] = std::shared_ptr<GGMLBlock>(new VisionModel(llama_cpp_style, params.vision));
blocks["visual"] = std::shared_ptr<GGMLBlock>(new VisionModel(llama_cpp_style, config.vision));
}
}
@@ -1189,7 +1348,7 @@ namespace LLM {
};
struct LLMRunner : public GGMLRunner {
LLMParams params;
LLMConfig config;
bool enable_vision;
LLM model;
@@ -1205,7 +1364,7 @@ namespace LLM {
static ggml_tensor* process_image_common(ggml_context* ctx,
ggml_tensor* image,
const LLMVisionParams& vision_params) {
const LLMVisionConfig& vision_params) {
// image: [C, H, W]
// return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
int64_t C = image->ne[2];
@@ -1300,7 +1459,7 @@ namespace LLM {
ggml_context* compute_ctx,
GGMLRunnerContext* runner_ctx,
ggml_tensor* image,
const LLMVisionParams& vision_params,
const LLMVisionConfig& vision_params,
std::shared_ptr<VisionModel> vision_model,
std::vector<int>& window_index_vec,
std::vector<int>& window_inverse_index_vec,
@@ -1415,121 +1574,25 @@ namespace LLM {
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
bool enable_vision_ = false)
: GGMLRunner(backend, params_backend), enable_vision(enable_vision_) {
params.arch = arch;
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
params.head_dim = 128;
params.num_heads = 32;
params.num_kv_heads = 8;
params.qkv_bias = false;
params.rms_norm_eps = 1e-5f;
} else if (arch == LLMArch::QWEN3) {
params.head_dim = 128;
params.num_heads = 32;
params.num_kv_heads = 8;
params.qkv_bias = false;
params.qk_norm = true;
params.rms_norm_eps = 1e-6f;
} else if (arch == LLMArch::GEMMA3_12B) {
params.head_dim = 256;
params.num_heads = 16;
params.num_kv_heads = 8;
params.qkv_bias = false;
params.qk_norm = true;
params.rms_norm_eps = 1e-6f;
// llama.cpp adds +1 to Gemma3 norm.weight when exporting GGUF, so GGUF loading
// must keep rms_norm_add disabled here or the offset gets applied twice.
// Convenient for the converter, less convenient for whoever gets to debug it later.
params.rms_norm_add = false;
params.normalize_input = true;
params.max_position_embeddings = 131072;
params.mlp_activation = MLPActivation::GELU_TANH;
params.rope_thetas = {1000000.f, 10000.f};
params.rope_scales = {8.f, 1.f};
params.sliding_attention = {1024, 1024, 1024, 1024, 1024, 0};
} else if (arch == LLMArch::GPT_OSS_20B) {
params.head_dim = 64;
params.num_heads = 64;
params.num_kv_heads = 8;
params.qkv_bias = true;
params.attention_out_bias = true;
params.qk_norm = false;
params.rms_norm_eps = 1e-5f;
params.hidden_size = 2880;
params.intermediate_size = 2880;
params.num_layers = 24;
params.vocab_size = 201088;
params.max_position_embeddings = 131072;
params.rope_thetas = {150000.f};
params.rope_scales = {32.f};
params.sliding_attention = {128, 0};
params.num_experts = 32;
params.num_experts_per_tok = 4;
}
bool have_vision_weight = false;
bool llama_cpp_style = false;
params.num_layers = 0;
for (auto pair : tensor_storage_map) {
std::string tensor_name = pair.first;
if (tensor_name.find(prefix) == std::string::npos)
continue;
size_t pos = tensor_name.find("visual.");
if (pos != std::string::npos) {
have_vision_weight = true;
if (contains(tensor_name, "attn.q_proj")) {
llama_cpp_style = true;
}
continue;
}
pos = tensor_name.find("layers.");
if (pos != std::string::npos) {
tensor_name = tensor_name.substr(pos); // remove prefix
auto items = split_string(tensor_name, '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > params.num_layers) {
params.num_layers = block_index + 1;
}
}
}
if (contains(tensor_name, "embed_tokens.weight")) {
params.hidden_size = pair.second.ne[0];
params.vocab_size = pair.second.ne[1];
}
if (contains(tensor_name, "layers.0.mlp.gate_proj.weight")) {
params.intermediate_size = pair.second.ne[1];
}
if (contains(tensor_name, "layers.0.mlp.experts.gate_up_proj.weight")) {
params.intermediate_size = pair.second.ne[1] / 2;
}
if (contains(tensor_name, "layers.0.mlp.experts.gate_proj.weight")) {
params.intermediate_size = pair.second.ne[1];
}
}
if (arch == LLMArch::QWEN3 && params.num_layers == 28) { // Qwen3 2B
params.num_heads = 16;
}
LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
params.num_layers,
params.vocab_size,
params.hidden_size,
params.intermediate_size);
if (enable_vision && !have_vision_weight) {
: GGMLRunner(backend, params_backend),
config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch)),
enable_vision(enable_vision_) {
if (enable_vision && !config.have_vision_weight) {
LOG_WARN("no vision weights detected, vision disabled");
enable_vision = false;
}
if (enable_vision) {
LOG_DEBUG("enable llm vision");
if (llama_cpp_style) {
if (config.llama_cpp_style) {
LOG_DEBUG("llama.cpp style vision weight");
}
}
model = LLM(params, enable_vision, llama_cpp_style);
model = LLM(config, enable_vision, config.llama_cpp_style);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return llm_arch_to_str[static_cast<int>(params.arch)];
return llm_arch_to_str[static_cast<int>(config.arch)];
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
@@ -1581,11 +1644,12 @@ namespace LLM {
}
int64_t n_tokens = input_ids->ne[0];
if (params.arch == LLMArch::MISTRAL_SMALL_3_2 ||
params.arch == LLMArch::MINISTRAL_3_3B ||
params.arch == LLMArch::QWEN3 ||
params.arch == LLMArch::GEMMA3_12B ||
params.arch == LLMArch::GPT_OSS_20B) {
if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
config.arch == LLMArch::MINISTRAL_3_3B ||
config.arch == LLMArch::QWEN3 ||
config.arch == LLMArch::GEMMA3_12B ||
config.arch == LLMArch::GEMMA2_2B ||
config.arch == LLMArch::GPT_OSS_20B) {
input_pos_vec.resize(n_tokens);
for (int i = 0; i < n_tokens; ++i) {
input_pos_vec[i] = i;
@@ -1624,9 +1688,9 @@ namespace LLM {
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
}
if (params.arch == LLMArch::GEMMA3_12B || params.arch == LLMArch::GPT_OSS_20B) {
if (config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GPT_OSS_20B) {
int sliding_window = 0;
for (int window : params.sliding_attention) {
for (int window : config.sliding_attention) {
sliding_window = std::max(sliding_window, window);
}
sliding_attention_mask_vec.resize(n_tokens * n_tokens);
@@ -1682,15 +1746,15 @@ namespace LLM {
int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
int64_t grid_t = 1;
int64_t grid_h = h / params.vision.patch_size;
int64_t grid_w = w / params.vision.patch_size;
int64_t llm_grid_h = grid_h / params.vision.spatial_merge_size;
int64_t llm_grid_w = grid_w / params.vision.spatial_merge_size;
int64_t grid_h = h / config.vision.patch_size;
int64_t grid_w = w / config.vision.patch_size;
int64_t llm_grid_h = grid_h / config.vision.spatial_merge_size;
int64_t llm_grid_w = grid_w / config.vision.spatial_merge_size;
return grid_t * grid_h * grid_w;
}
ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
return process_image_common(ctx, image, params.vision);
return process_image_common(ctx, image, config.vision);
}
ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx,
@@ -1712,7 +1776,7 @@ namespace LLM {
compute_ctx,
runner_ctx,
image,
params.vision,
config.vision,
model.vision_model(),
window_index_vec,
window_inverse_index_vec,
@@ -1726,8 +1790,8 @@ namespace LLM {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
GGML_ASSERT(image->ne[1] % (config.vision.patch_size * config.vision.spatial_merge_size) == 0);
GGML_ASSERT(image->ne[0] % (config.vision.patch_size * config.vision.spatial_merge_size) == 0);
auto runnter_ctx = get_context();
ggml_tensor* hidden_states = encode_image(&runnter_ctx, image);
@@ -1988,7 +2052,7 @@ namespace LLM {
static void load_from_file_and_test(const std::string& file_path) {
// cpu f16: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_COUNT;
ModelLoader model_loader;
otherarch/sdcpp/lora.hpp
+2 -3
View File
@@ -91,7 +91,6 @@ struct LoraModel : public GGMLRunner {
return false;
}
dry_run = false;
model_loader.load_tensors(on_new_tensor_cb, n_threads);
@@ -773,7 +772,7 @@ struct LoraModel : public GGMLRunner {
}
ggml_tensor* original_tensor = model_tensor;
if (!ggml_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
if (!sd_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
model_tensor = ggml_dup_tensor(compute_ctx, model_tensor);
set_backend_tensor_data(model_tensor, original_tensor->data);
}
@@ -787,7 +786,7 @@ struct LoraModel : public GGMLRunner {
final_tensor = ggml_add_inplace(compute_ctx, model_tensor, diff);
}
ggml_build_forward_expand(gf, final_tensor);
if (!ggml_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
if (!sd_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
original_tensor_to_final_tensor[original_tensor] = final_tensor;
}
}
otherarch/sdcpp/ltx_audio_vae.h
+12 -5
View File
@@ -58,11 +58,12 @@ namespace LTXV {
return base_output_sample_rate();
}
static LTXAudioVAEConfig detect_from_weights(const String2TensorStorage& tensor_storage_map) {
static LTXAudioVAEConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix = "") {
LTXAudioVAEConfig config;
auto require = [&](const std::string& name) -> const TensorStorage* {
auto iter = tensor_storage_map.find(name);
std::string tensor_name = prefix.empty() ? name : prefix + "." + name;
auto iter = tensor_storage_map.find(tensor_name);
if (iter == tensor_storage_map.end()) {
return nullptr;
}
@@ -168,6 +169,12 @@ namespace LTXV {
if (config.audio_channels != 2 || config.latent_channels != 8 || config.mel_bins != 64) {
return config;
}
LOG_DEBUG("ltx_audio_vae: sample_rate = %d, mel_bins = %d, latent_channels = %d, latent_frequency_bins = %d, has_bwe = %s",
config.sample_rate,
config.mel_bins,
config.latent_channels,
config.latent_frequency_bins,
config.has_bwe ? "true" : "false");
return config;
}
};
@@ -1052,7 +1059,7 @@ namespace LTXV {
static void load_from_file_and_test(const std::string& model_path,
const std::string& input_path,
const std::string& prefix = "") {
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
// ggml_backend_t backend = ggml_backend_cuda_init(0);
LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
@@ -1069,8 +1076,8 @@ namespace LTXV {
prefix);
if (!ltx_audio_vae->alloc_params_buffer()) {
LOG_ERROR("ltx audio vae buffer allocation failed");
return;
LOG_ERROR("ltx audio vae buffer allocation failed");
return;
}
std::map<std::string, ggml_tensor*> tensors;
otherarch/sdcpp/ltx_vae.hpp
+2 -2
View File
@@ -11,7 +11,7 @@
#include "ltxv.hpp"
#include "vae.hpp"
#include "wan.hpp"
#include "wan_vae.hpp"
namespace LTXVAE {
@@ -1517,7 +1517,7 @@ struct LTXVideoVAE : public VAE {
static void load_from_file_and_test(const std::string& model_path,
const std::string& input_path) {
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
ModelLoader model_loader;
otherarch/sdcpp/ltxv.hpp
+346 -320
View File
@@ -72,6 +72,200 @@ namespace LTXV {
return max_block + 1;
}
struct LTXAVConfig {
int64_t in_channels = 128;
int64_t out_channels = 128;
int64_t hidden_size = 3840;
int64_t cross_attention_dim = 4096;
int64_t caption_channels = 3840;
int64_t num_attention_heads = 30;
int64_t attention_head_dim = 128;
int64_t num_layers = 28;
float positional_embedding_theta = 10000.f;
std::vector<int> positional_embedding_max_pos = {20, 2048, 2048};
std::tuple<int, int, int> vae_scale_factors = {8, 32, 32};
bool causal_temporal_positioning = true;
float timestep_scale_multiplier = 1000.f;
int64_t audio_in_channels = 128;
int64_t audio_out_channels = 128;
int64_t audio_hidden_size = 2048;
int64_t audio_cross_attention_dim = 2048;
int64_t audio_num_attention_heads = 32;
int64_t audio_attention_head_dim = 64;
std::vector<int> audio_positional_embedding_max_pos = {20};
float av_ca_timestep_scale_multiplier = 1000.f;
int64_t num_audio_channels = 8;
int64_t audio_frequency_bins = 16;
bool use_connector = false;
int64_t connector_hidden_size = 3840;
int64_t connector_num_heads = 30;
int64_t connector_head_dim = 128;
int64_t connector_num_layers = 2;
int64_t connector_num_registers = 128;
bool connector_rope_interleaved = false;
bool connector_apply_gated_attention = false;
bool use_audio_connector = false;
int64_t audio_connector_hidden_size = 2048;
int64_t audio_connector_num_heads = 32;
int64_t audio_connector_head_dim = 64;
int64_t audio_connector_num_layers = 2;
int64_t audio_connector_num_registers = 128;
bool audio_connector_rope_interleaved = false;
bool audio_connector_apply_gated_attention = false;
bool video_rope_interleaved = false;
bool use_middle_indices_grid = true;
bool cross_attention_adaln = false;
bool use_caption_projection = true;
bool use_audio_caption_projection = true;
bool caption_proj_before_connector = true;
bool caption_projection_first_linear = false;
bool self_attention_gated = false;
bool cross_attention_gated = false;
static std::pair<int64_t, int64_t> infer_attention_layout(int64_t hidden_size,
int64_t preferred_heads = -1) {
if (preferred_heads > 0 && hidden_size % preferred_heads == 0) {
return {preferred_heads, hidden_size / preferred_heads};
}
const int candidates[] = {128, 96, 80, 64, 48, 40, 32};
for (int head_dim : candidates) {
if (hidden_size % head_dim == 0) {
int64_t heads = hidden_size / head_dim;
if (heads >= 8 && heads <= 64) {
return {heads, head_dim};
}
}
}
return {32, hidden_size / 32};
}
static int64_t infer_gate_heads(const String2TensorStorage& tensor_storage_map,
const std::string& bias_name,
int64_t fallback_heads) {
auto it = tensor_storage_map.find(bias_name);
if (it != tensor_storage_map.end()) {
return it->second.ne[0];
}
return fallback_heads;
}
static LTXAVConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
LTXAVConfig config;
auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
if (patchify_proj_iter != tensor_storage_map.end()) {
config.in_channels = patchify_proj_iter->second.ne[0];
config.hidden_size = patchify_proj_iter->second.ne[1];
int64_t video_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32);
auto attn_layout = infer_attention_layout(config.hidden_size, video_heads);
config.num_attention_heads = attn_layout.first;
config.attention_head_dim = attn_layout.second;
}
auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight");
if (audio_patchify_proj_iter != tensor_storage_map.end()) {
config.audio_in_channels = audio_patchify_proj_iter->second.ne[0];
config.audio_hidden_size = audio_patchify_proj_iter->second.ne[1];
config.audio_out_channels = config.audio_in_channels;
int64_t audio_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32);
auto audio_attn_layout = infer_attention_layout(config.audio_hidden_size, audio_heads);
config.audio_num_attention_heads = audio_attn_layout.first;
config.audio_attention_head_dim = audio_attn_layout.second;
}
auto proj_out_iter = tensor_storage_map.find(prefix + ".proj_out.weight");
if (proj_out_iter != tensor_storage_map.end()) {
config.out_channels = proj_out_iter->second.ne[1];
}
auto audio_proj_out_iter = tensor_storage_map.find(prefix + ".audio_proj_out.weight");
if (audio_proj_out_iter != tensor_storage_map.end()) {
config.audio_out_channels = audio_proj_out_iter->second.ne[1];
}
auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight");
if (attn2_iter != tensor_storage_map.end()) {
config.cross_attention_dim = attn2_iter->second.ne[0];
}
auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight");
if (audio_attn2_iter != tensor_storage_map.end()) {
config.audio_cross_attention_dim = audio_attn2_iter->second.ne[0];
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) {
config.cross_attention_adaln = true;
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end() ||
tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
config.self_attention_gated = true;
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_gate_logits.weight") != tensor_storage_map.end() ||
tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) {
config.cross_attention_gated = true;
}
if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) {
config.use_caption_projection = false;
}
if (tensor_storage_map.find(prefix + ".audio_caption_projection.linear_1.weight") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".audio_caption_projection.linear_2.weight") == tensor_storage_map.end()) {
config.use_audio_caption_projection = false;
}
config.num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".", "transformer_blocks.");
auto connector_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
if (connector_iter != tensor_storage_map.end()) {
config.use_connector = true;
config.connector_hidden_size = connector_iter->second.ne[1];
int64_t connector_heads = infer_gate_heads(tensor_storage_map,
prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
32);
auto connector_layout = infer_attention_layout(config.connector_hidden_size, connector_heads);
config.connector_num_heads = connector_layout.first;
config.connector_head_dim = connector_layout.second;
config.connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".video_embeddings_connector.", "transformer_1d_blocks.");
auto register_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.learnable_registers");
if (register_iter != tensor_storage_map.end()) {
config.connector_num_registers = register_iter->second.ne[1];
}
if (tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
config.connector_apply_gated_attention = true;
}
}
auto audio_connector_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
if (audio_connector_iter != tensor_storage_map.end()) {
config.use_audio_connector = true;
config.audio_connector_hidden_size = audio_connector_iter->second.ne[1];
int64_t connector_heads = infer_gate_heads(tensor_storage_map,
prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
32);
auto connector_layout = infer_attention_layout(config.audio_connector_hidden_size, connector_heads);
config.audio_connector_num_heads = connector_layout.first;
config.audio_connector_head_dim = connector_layout.second;
config.audio_connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".audio_embeddings_connector.", "transformer_1d_blocks.");
auto register_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.learnable_registers");
if (register_iter != tensor_storage_map.end()) {
config.audio_connector_num_registers = register_iter->second.ne[1];
}
if (tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
config.audio_connector_apply_gated_attention = true;
}
}
LOG_DEBUG("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64,
config.num_layers,
config.hidden_size,
config.num_attention_heads,
config.audio_hidden_size,
config.audio_num_attention_heads);
return config;
}
};
__STATIC_INLINE__ std::vector<float> generate_freq_grid(float theta,
int positional_dims,
int dim) {
@@ -749,63 +943,6 @@ namespace LTXV {
}
};
struct LTXAVParams {
int64_t in_channels = 128;
int64_t out_channels = 128;
int64_t hidden_size = 3840;
int64_t cross_attention_dim = 4096;
int64_t caption_channels = 3840;
int64_t num_attention_heads = 30;
int64_t attention_head_dim = 128;
int64_t num_layers = 28;
float positional_embedding_theta = 10000.f;
std::vector<int> positional_embedding_max_pos = {20, 2048, 2048};
std::tuple<int, int, int> vae_scale_factors = {8, 32, 32};
bool causal_temporal_positioning = true;
float timestep_scale_multiplier = 1000.f;
int64_t audio_in_channels = 128;
int64_t audio_out_channels = 128;
int64_t audio_hidden_size = 2048;
int64_t audio_cross_attention_dim = 2048;
int64_t audio_num_attention_heads = 32;
int64_t audio_attention_head_dim = 64;
std::vector<int> audio_positional_embedding_max_pos = {20};
float av_ca_timestep_scale_multiplier = 1000.f;
int64_t num_audio_channels = 8;
int64_t audio_frequency_bins = 16;
bool use_connector = false;
int64_t connector_hidden_size = 3840;
int64_t connector_num_heads = 30;
int64_t connector_head_dim = 128;
int64_t connector_num_layers = 2;
int64_t connector_num_registers = 128;
bool connector_rope_interleaved = false;
bool connector_apply_gated_attention = false;
bool use_audio_connector = false;
int64_t audio_connector_hidden_size = 2048;
int64_t audio_connector_num_heads = 32;
int64_t audio_connector_head_dim = 64;
int64_t audio_connector_num_layers = 2;
int64_t audio_connector_num_registers = 128;
bool audio_connector_rope_interleaved = false;
bool audio_connector_apply_gated_attention = false;
bool video_rope_interleaved = false;
bool use_middle_indices_grid = true;
bool cross_attention_adaln = false;
bool use_caption_projection = true;
bool use_audio_caption_projection = true;
bool caption_proj_before_connector = true;
bool caption_projection_first_linear = false;
bool self_attention_gated = false;
bool cross_attention_gated = false;
};
__STATIC_INLINE__ std::pair<int64_t, int64_t> infer_attention_layout(int64_t hidden_size,
int64_t preferred_heads = -1) {
if (preferred_heads > 0 && hidden_size % preferred_heads == 0) {
@@ -1169,92 +1306,92 @@ namespace LTXV {
};
struct LTXAVModelBlock : public GGMLBlock {
LTXAVParams cfg;
LTXAVConfig config;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["scale_shift_table"] = ggml_new_tensor_2d(ctx,
get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32),
cfg.hidden_size,
config.hidden_size,
2);
params["audio_scale_shift_table"] = ggml_new_tensor_2d(ctx,
get_type(prefix + "audio_scale_shift_table", tensor_storage_map, GGML_TYPE_F32),
cfg.audio_hidden_size,
config.audio_hidden_size,
2);
}
LTXAVModelBlock(const LTXAVParams& params)
: cfg(params) {
blocks["patchify_proj"] = std::make_shared<Linear>(cfg.in_channels, cfg.hidden_size, true, true);
blocks["audio_patchify_proj"] = std::make_shared<Linear>(cfg.audio_in_channels, cfg.audio_hidden_size, true, true);
blocks["adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, cfg.cross_attention_adaln ? 9 : 6);
blocks["audio_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, cfg.cross_attention_adaln ? 9 : 6);
if (cfg.cross_attention_adaln) {
blocks["prompt_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, 2);
blocks["audio_prompt_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, 2);
LTXAVModelBlock(const LTXAVConfig& config)
: config(config) {
blocks["patchify_proj"] = std::make_shared<Linear>(config.in_channels, config.hidden_size, true, true);
blocks["audio_patchify_proj"] = std::make_shared<Linear>(config.audio_in_channels, config.audio_hidden_size, true, true);
blocks["adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.hidden_size, config.cross_attention_adaln ? 9 : 6);
blocks["audio_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, config.cross_attention_adaln ? 9 : 6);
if (config.cross_attention_adaln) {
blocks["prompt_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.hidden_size, 2);
blocks["audio_prompt_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, 2);
}
blocks["av_ca_video_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, 4);
blocks["av_ca_a2v_gate_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, 1);
blocks["av_ca_audio_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, 4);
blocks["av_ca_v2a_gate_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, 1);
blocks["av_ca_video_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.hidden_size, 4);
blocks["av_ca_a2v_gate_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.hidden_size, 1);
blocks["av_ca_audio_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, 4);
blocks["av_ca_v2a_gate_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, 1);
if (cfg.use_caption_projection) {
if (cfg.caption_proj_before_connector) {
if (cfg.caption_projection_first_linear) {
blocks["caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(cfg.caption_channels, cfg.hidden_size);
if (config.use_caption_projection) {
if (config.caption_proj_before_connector) {
if (config.caption_projection_first_linear) {
blocks["caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(config.caption_channels, config.hidden_size);
}
} else {
blocks["caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(cfg.caption_channels, cfg.hidden_size, cfg.hidden_size);
blocks["caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(config.caption_channels, config.hidden_size, config.hidden_size);
}
}
if (cfg.use_audio_caption_projection) {
if (cfg.caption_proj_before_connector) {
if (cfg.caption_projection_first_linear) {
blocks["audio_caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(cfg.caption_channels, cfg.audio_hidden_size);
if (config.use_audio_caption_projection) {
if (config.caption_proj_before_connector) {
if (config.caption_projection_first_linear) {
blocks["audio_caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(config.caption_channels, config.audio_hidden_size);
}
} else {
blocks["audio_caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(cfg.caption_channels, cfg.audio_hidden_size, cfg.audio_hidden_size);
blocks["audio_caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(config.caption_channels, config.audio_hidden_size, config.audio_hidden_size);
}
}
if (cfg.use_connector) {
blocks["video_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(cfg.connector_hidden_size,
cfg.connector_num_heads,
cfg.connector_head_dim,
cfg.connector_num_layers,
cfg.connector_num_registers,
cfg.connector_rope_interleaved,
cfg.connector_apply_gated_attention);
if (config.use_connector) {
blocks["video_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(config.connector_hidden_size,
config.connector_num_heads,
config.connector_head_dim,
config.connector_num_layers,
config.connector_num_registers,
config.connector_rope_interleaved,
config.connector_apply_gated_attention);
}
if (cfg.use_audio_connector) {
blocks["audio_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(cfg.audio_connector_hidden_size,
cfg.audio_connector_num_heads,
cfg.audio_connector_head_dim,
cfg.audio_connector_num_layers,
cfg.audio_connector_num_registers,
cfg.audio_connector_rope_interleaved,
cfg.audio_connector_apply_gated_attention);
if (config.use_audio_connector) {
blocks["audio_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(config.audio_connector_hidden_size,
config.audio_connector_num_heads,
config.audio_connector_head_dim,
config.audio_connector_num_layers,
config.audio_connector_num_registers,
config.audio_connector_rope_interleaved,
config.audio_connector_apply_gated_attention);
}
for (int i = 0; i < cfg.num_layers; i++) {
blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<BasicAVTransformerBlock>(cfg.hidden_size,
cfg.audio_hidden_size,
cfg.num_attention_heads,
cfg.audio_num_attention_heads,
cfg.attention_head_dim,
cfg.audio_attention_head_dim,
cfg.cross_attention_dim,
cfg.audio_cross_attention_dim,
cfg.self_attention_gated || cfg.cross_attention_gated,
cfg.cross_attention_adaln,
cfg.video_rope_interleaved);
for (int i = 0; i < config.num_layers; i++) {
blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<BasicAVTransformerBlock>(config.hidden_size,
config.audio_hidden_size,
config.num_attention_heads,
config.audio_num_attention_heads,
config.attention_head_dim,
config.audio_attention_head_dim,
config.cross_attention_dim,
config.audio_cross_attention_dim,
config.self_attention_gated || config.cross_attention_gated,
config.cross_attention_adaln,
config.video_rope_interleaved);
}
blocks["norm_out"] = std::make_shared<LayerNorm>(cfg.hidden_size, 1e-6f, false);
blocks["proj_out"] = std::make_shared<Linear>(cfg.hidden_size, cfg.out_channels, true, true);
blocks["audio_norm_out"] = std::make_shared<LayerNorm>(cfg.audio_hidden_size, 1e-6f, false);
blocks["audio_proj_out"] = std::make_shared<Linear>(cfg.audio_hidden_size, cfg.audio_out_channels, true, true);
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size, config.out_channels, true, true);
blocks["audio_norm_out"] = std::make_shared<LayerNorm>(config.audio_hidden_size, 1e-6f, false);
blocks["audio_proj_out"] = std::make_shared<Linear>(config.audio_hidden_size, config.audio_out_channels, true, true);
}
ggml_tensor* patchify_video(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t n) {
@@ -1293,8 +1430,8 @@ namespace LTXV {
if (ax == nullptr) {
return nullptr;
}
ax = ggml_reshape_4d(ctx->ggml_ctx, ax, cfg.audio_frequency_bins, cfg.num_audio_channels, audio_length, ax->ne[2]); // [b, t, c, f]
ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3)); // [b, c, t, f]
ax = ggml_reshape_4d(ctx->ggml_ctx, ax, config.audio_frequency_bins, config.num_audio_channels, audio_length, ax->ne[2]); // [b, t, c, f]
ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3)); // [b, c, t, f]
return ax;
}
@@ -1308,17 +1445,17 @@ namespace LTXV {
}
bool is_fully_processed_context =
context->ne[0] == cfg.cross_attention_dim + cfg.audio_cross_attention_dim &&
context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim &&
context->ne[1] >= 1024;
bool is_unprocessed_dual_context =
context->ne[0] == cfg.cross_attention_dim + cfg.audio_cross_attention_dim &&
context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim &&
context->ne[1] < 1024;
if (is_fully_processed_context) {
auto v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.cross_attention_dim);
auto v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
ggml_tensor* a_context = nullptr;
if (process_audio_context) {
a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.cross_attention_dim, cfg.cross_attention_dim + cfg.audio_cross_attention_dim);
a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
}
return {v_context, a_context};
}
@@ -1326,32 +1463,32 @@ namespace LTXV {
ggml_tensor* v_context = context;
ggml_tensor* a_context = process_audio_context ? context : nullptr;
if (is_unprocessed_dual_context) {
v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.cross_attention_dim);
v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
if (process_audio_context) {
a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.cross_attention_dim, cfg.cross_attention_dim + cfg.audio_cross_attention_dim);
a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
}
} else if (context->ne[0] == cfg.caption_channels * 2) {
v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.caption_channels);
} else if (context->ne[0] == config.caption_channels * 2) {
v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.caption_channels);
if (process_audio_context) {
a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.caption_channels, cfg.caption_channels * 2);
a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.caption_channels, config.caption_channels * 2);
}
}
if (cfg.caption_proj_before_connector) {
if (cfg.use_caption_projection &&
if (config.caption_proj_before_connector) {
if (config.use_caption_projection &&
blocks.count("caption_projection") > 0 &&
v_context != nullptr &&
v_context->ne[0] == cfg.caption_channels) {
v_context->ne[0] == config.caption_channels) {
auto caption_projection = std::dynamic_pointer_cast<NormSingleLinearTextProjection>(blocks["caption_projection"]);
if (caption_projection != nullptr) {
v_context = caption_projection->forward(ctx, v_context);
}
}
if (process_audio_context &&
cfg.use_audio_caption_projection &&
config.use_audio_caption_projection &&
blocks.count("audio_caption_projection") > 0 &&
a_context != nullptr &&
a_context->ne[0] == cfg.caption_channels) {
a_context->ne[0] == config.caption_channels) {
auto caption_projection = std::dynamic_pointer_cast<NormSingleLinearTextProjection>(blocks["audio_caption_projection"]);
if (caption_projection != nullptr) {
a_context = caption_projection->forward(ctx, a_context);
@@ -1359,34 +1496,34 @@ namespace LTXV {
}
}
if (cfg.use_connector && v_context != nullptr && v_context->ne[0] == cfg.connector_hidden_size) {
if (config.use_connector && v_context != nullptr && v_context->ne[0] == config.connector_hidden_size) {
auto connector = std::dynamic_pointer_cast<Embeddings1DConnector>(blocks["video_embeddings_connector"]);
v_context = connector->forward(ctx, v_context, video_connector_pe);
}
if (process_audio_context &&
cfg.use_audio_connector &&
config.use_audio_connector &&
a_context != nullptr &&
a_context->ne[0] == cfg.audio_connector_hidden_size) {
a_context->ne[0] == config.audio_connector_hidden_size) {
auto connector = std::dynamic_pointer_cast<Embeddings1DConnector>(blocks["audio_embeddings_connector"]);
a_context = connector->forward(ctx, a_context, audio_connector_pe);
}
if (!cfg.caption_proj_before_connector &&
cfg.use_caption_projection &&
if (!config.caption_proj_before_connector &&
config.use_caption_projection &&
blocks.count("caption_projection") > 0 &&
v_context != nullptr &&
v_context->ne[0] == cfg.caption_channels) {
v_context->ne[0] == config.caption_channels) {
auto caption_projection = std::dynamic_pointer_cast<PixArtAlphaTextProjection>(blocks["caption_projection"]);
if (caption_projection != nullptr) {
v_context = caption_projection->forward(ctx, v_context);
}
}
if (process_audio_context &&
!cfg.caption_proj_before_connector &&
cfg.use_audio_caption_projection &&
!config.caption_proj_before_connector &&
config.use_audio_caption_projection &&
blocks.count("audio_caption_projection") > 0 &&
a_context != nullptr &&
a_context->ne[0] == cfg.caption_channels) {
a_context->ne[0] == config.caption_channels) {
auto caption_projection = std::dynamic_pointer_cast<PixArtAlphaTextProjection>(blocks["audio_caption_projection"]);
if (caption_projection != nullptr) {
a_context = caption_projection->forward(ctx, a_context);
@@ -1428,8 +1565,8 @@ namespace LTXV {
auto audio_norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["audio_norm_out"]);
auto audio_proj_out = std::dynamic_pointer_cast<Linear>(blocks["audio_proj_out"]);
GGML_ASSERT(vx->ne[3] % cfg.in_channels == 0);
int64_t n = vx->ne[3] / cfg.in_channels;
GGML_ASSERT(vx->ne[3] % config.in_channels == 0);
int64_t n = vx->ne[3] / config.in_channels;
int64_t width = vx->ne[0];
int64_t height = vx->ne[1];
int64_t frames = vx->ne[2];
@@ -1452,20 +1589,20 @@ namespace LTXV {
a_context = ggml_cont(ctx->ggml_ctx, a_context);
}
auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, cfg.timestep_scale_multiplier);
auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, config.timestep_scale_multiplier);
auto v_pair = adaln_single->forward(ctx, v_timestep_scaled);
auto v_timestep_mod = v_pair.first;
auto v_embedded_time = v_pair.second;
ggml_tensor* effective_audio_timestep = audio_timestep != nullptr ? audio_timestep : timestep;
auto a_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, cfg.timestep_scale_multiplier);
auto a_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, config.timestep_scale_multiplier);
auto a_pair = audio_adaln_single->forward(ctx, a_timestep_scaled);
auto a_timestep_mod = a_pair.first;
auto a_embedded_time = a_pair.second;
ggml_tensor* v_prompt_timestep_mod = nullptr;
ggml_tensor* a_prompt_timestep_mod = nullptr;
if (cfg.cross_attention_adaln) {
if (config.cross_attention_adaln) {
auto prompt_adaln_single = std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["prompt_adaln_single"]);
auto audio_prompt_adaln_single = std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["audio_prompt_adaln_single"]);
v_prompt_timestep_mod = prompt_adaln_single->forward(ctx, a_timestep_scaled).first;
@@ -1474,7 +1611,7 @@ namespace LTXV {
auto av_ca_video_timestep = repeat_scalar_timestep_like(ctx, effective_audio_timestep, timestep);
auto av_ca_audio_timestep = effective_audio_timestep;
auto av_ca_factor = cfg.av_ca_timestep_scale_multiplier / cfg.timestep_scale_multiplier;
auto av_ca_factor = config.av_ca_timestep_scale_multiplier / config.timestep_scale_multiplier;
auto av_ca_video_scale_shift_timestep =
std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_video_scale_shift_adaln_single"])->forward(ctx, av_ca_video_timestep).first;
auto av_ca_a2v_gate_noise_timestep =
@@ -1491,7 +1628,7 @@ namespace LTXV {
sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.prelude", "vx");
sd::ggml_graph_cut::mark_graph_cut(ax, "ltxav.prelude", "ax");
for (int i = 0; i < cfg.num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<BasicAVTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
auto out = block->forward(ctx,
vx,
@@ -1517,14 +1654,14 @@ namespace LTXV {
sd::ggml_graph_cut::mark_graph_cut(ax, "ltxav.transformer_blocks." + std::to_string(i), "ax");
}
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, cfg.hidden_size);
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
vx = norm_out->forward(ctx, vx);
vx = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = proj_out->forward(ctx, vx);
vx = unpatchify_video(ctx, vx, width, height, frames);
if (ax != nullptr && audio_time > 0) {
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, cfg.audio_hidden_size);
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
ax = audio_norm_out->forward(ctx, ax);
ax = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = audio_proj_out->forward(ctx, ax);
@@ -1536,7 +1673,7 @@ namespace LTXV {
};
struct LTXAVRunner : public DiffusionModelRunner {
LTXAVParams params;
LTXAVConfig config;
LTXAVModelBlock model;
std::vector<float> video_pe_vec;
std::vector<float> audio_pe_vec;
@@ -1547,124 +1684,13 @@ namespace LTXV {
sd::Tensor<float> vx_input_cache;
sd::Tensor<float> ax_input_cache;
static int64_t infer_gate_heads(const String2TensorStorage& tensor_storage_map,
const std::string& bias_name,
int64_t fallback_heads) {
auto it = tensor_storage_map.find(bias_name);
if (it != tensor_storage_map.end()) {
return it->second.ne[0];
}
return fallback_heads;
}
LTXAVRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& prefix = "model.diffusion_model")
: DiffusionModelRunner(backend, params_backend, prefix),
params(),
model(params) {
auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
if (patchify_proj_iter != tensor_storage_map.end()) {
params.in_channels = patchify_proj_iter->second.ne[0];
params.hidden_size = patchify_proj_iter->second.ne[1];
int64_t video_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32);
auto attn_layout = infer_attention_layout(params.hidden_size, video_heads);
params.num_attention_heads = attn_layout.first;
params.attention_head_dim = attn_layout.second;
}
auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight");
if (audio_patchify_proj_iter != tensor_storage_map.end()) {
params.audio_in_channels = audio_patchify_proj_iter->second.ne[0];
params.audio_hidden_size = audio_patchify_proj_iter->second.ne[1];
params.audio_out_channels = params.audio_in_channels;
int64_t audio_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32);
auto audio_attn_layout = infer_attention_layout(params.audio_hidden_size, audio_heads);
params.audio_num_attention_heads = audio_attn_layout.first;
params.audio_attention_head_dim = audio_attn_layout.second;
}
auto proj_out_iter = tensor_storage_map.find(prefix + ".proj_out.weight");
if (proj_out_iter != tensor_storage_map.end()) {
params.out_channels = proj_out_iter->second.ne[1];
}
auto audio_proj_out_iter = tensor_storage_map.find(prefix + ".audio_proj_out.weight");
if (audio_proj_out_iter != tensor_storage_map.end()) {
params.audio_out_channels = audio_proj_out_iter->second.ne[1];
}
auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight");
if (attn2_iter != tensor_storage_map.end()) {
params.cross_attention_dim = attn2_iter->second.ne[0];
}
auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight");
if (audio_attn2_iter != tensor_storage_map.end()) {
params.audio_cross_attention_dim = audio_attn2_iter->second.ne[0];
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) {
params.cross_attention_adaln = true;
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end() ||
tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
params.self_attention_gated = true;
}
if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_gate_logits.weight") != tensor_storage_map.end() ||
tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) {
params.cross_attention_gated = true;
}
if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) {
params.use_caption_projection = false;
}
if (tensor_storage_map.find(prefix + ".audio_caption_projection.linear_1.weight") == tensor_storage_map.end() &&
tensor_storage_map.find(prefix + ".audio_caption_projection.linear_2.weight") == tensor_storage_map.end()) {
params.use_audio_caption_projection = false;
}
params.num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".", "transformer_blocks.");
auto connector_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
if (connector_iter != tensor_storage_map.end()) {
params.use_connector = true;
params.connector_hidden_size = connector_iter->second.ne[1];
int64_t connector_heads = infer_gate_heads(tensor_storage_map,
prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
32);
auto connector_layout = infer_attention_layout(params.connector_hidden_size, connector_heads);
params.connector_num_heads = connector_layout.first;
params.connector_head_dim = connector_layout.second;
params.connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".video_embeddings_connector.", "transformer_1d_blocks.");
auto register_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.learnable_registers");
if (register_iter != tensor_storage_map.end()) {
params.connector_num_registers = register_iter->second.ne[1];
}
if (tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
params.connector_apply_gated_attention = true;
}
}
auto audio_connector_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
if (audio_connector_iter != tensor_storage_map.end()) {
params.use_audio_connector = true;
params.audio_connector_hidden_size = audio_connector_iter->second.ne[1];
int64_t connector_heads = infer_gate_heads(tensor_storage_map,
prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
32);
auto connector_layout = infer_attention_layout(params.audio_connector_hidden_size, connector_heads);
params.audio_connector_num_heads = connector_layout.first;
params.audio_connector_head_dim = connector_layout.second;
params.audio_connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".audio_embeddings_connector.", "transformer_1d_blocks.");
auto register_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.learnable_registers");
if (register_iter != tensor_storage_map.end()) {
params.audio_connector_num_registers = register_iter->second.ne[1];
}
if (tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
params.audio_connector_apply_gated_attention = true;
}
}
model = LTXAVModelBlock(params);
config(LTXAVConfig::detect_from_weights(tensor_storage_map, prefix)),
model(config) {
model.init(params_ctx, tensor_storage_map, prefix);
}
@@ -1692,21 +1718,21 @@ namespace LTXV {
int64_t total_channels = x_tensor.shape()[3];
int64_t spatial_size = width * height * frames;
GGML_ASSERT(total_channels >= params.in_channels);
GGML_ASSERT(total_channels >= config.in_channels);
sd::Tensor<float> vx({width, height, frames, params.in_channels});
size_t video_values = static_cast<size_t>(params.in_channels * spatial_size);
sd::Tensor<float> vx({width, height, frames, config.in_channels});
size_t video_values = static_cast<size_t>(config.in_channels * spatial_size);
std::copy_n(x_tensor.data(), video_values, vx.data());
if (audio_length <= 0 || total_channels == params.in_channels) {
if (audio_length <= 0 || total_channels == config.in_channels) {
return {vx, {}};
}
int64_t needed_audio_values = static_cast<int64_t>(audio_length) * params.num_audio_channels * params.audio_frequency_bins;
int64_t packed_audio_values = (total_channels - params.in_channels) * spatial_size;
int64_t needed_audio_values = static_cast<int64_t>(audio_length) * config.num_audio_channels * config.audio_frequency_bins;
int64_t packed_audio_values = (total_channels - config.in_channels) * spatial_size;
GGML_ASSERT(packed_audio_values >= needed_audio_values);
sd::Tensor<float> ax({params.audio_frequency_bins, audio_length, params.num_audio_channels, 1});
sd::Tensor<float> ax({config.audio_frequency_bins, audio_length, config.num_audio_channels, 1});
const float* audio_src = x_tensor.data() + video_values;
std::copy_n(audio_src, static_cast<size_t>(needed_audio_values), ax.data());
return {vx, ax};
@@ -1767,25 +1793,25 @@ namespace LTXV {
if (has_video_positions) {
GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count);
video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor,
static_cast<int>(params.hidden_size),
static_cast<int>(params.num_attention_heads),
params.positional_embedding_theta,
params.positional_embedding_max_pos,
params.use_middle_indices_grid);
static_cast<int>(config.hidden_size),
static_cast<int>(config.num_attention_heads),
config.positional_embedding_theta,
config.positional_embedding_max_pos,
config.use_middle_indices_grid);
} else {
video_pe_vec = build_video_rope_matrix(vx->ne[0],
vx->ne[1],
vx->ne[2],
static_cast<int>(params.hidden_size),
static_cast<int>(params.num_attention_heads),
static_cast<int>(config.hidden_size),
static_cast<int>(config.num_attention_heads),
video_frame_rate,
params.positional_embedding_theta,
params.positional_embedding_max_pos,
params.vae_scale_factors,
params.causal_temporal_positioning,
params.use_middle_indices_grid);
config.positional_embedding_theta,
config.positional_embedding_max_pos,
config.vae_scale_factors,
config.causal_temporal_positioning,
config.use_middle_indices_grid);
}
auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.attention_head_dim / 2, video_token_count * params.num_attention_heads);
auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads);
ggml_set_name(video_pe, "ltxav_video_pe");
set_backend_tensor_data(video_pe, video_pe_vec.data());
@@ -1794,66 +1820,66 @@ namespace LTXV {
ggml_tensor* audio_cross_pe = nullptr;
if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) {
audio_pe_vec = build_audio_rope_matrix(ax->ne[1],
static_cast<int>(params.audio_hidden_size),
static_cast<int>(params.audio_num_attention_heads),
params.positional_embedding_theta,
params.audio_positional_embedding_max_pos[0],
params.use_middle_indices_grid);
audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, ax->ne[1] * params.audio_num_attention_heads);
static_cast<int>(config.audio_hidden_size),
static_cast<int>(config.audio_num_attention_heads),
config.positional_embedding_theta,
config.audio_positional_embedding_max_pos[0],
config.use_middle_indices_grid);
audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
ggml_set_name(audio_pe, "ltxav_audio_pe");
set_backend_tensor_data(audio_pe, audio_pe_vec.data());
int temporal_max_pos = std::max(params.positional_embedding_max_pos[0], params.audio_positional_embedding_max_pos[0]);
int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]);
if (has_video_positions) {
video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor,
static_cast<int>(params.audio_cross_attention_dim),
static_cast<int>(params.audio_num_attention_heads),
params.positional_embedding_theta,
static_cast<int>(config.audio_cross_attention_dim),
static_cast<int>(config.audio_num_attention_heads),
config.positional_embedding_theta,
temporal_max_pos,
true);
} else {
video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0],
vx->ne[1],
vx->ne[2],
static_cast<int>(params.audio_cross_attention_dim),
static_cast<int>(params.audio_num_attention_heads),
static_cast<int>(config.audio_cross_attention_dim),
static_cast<int>(config.audio_num_attention_heads),
video_frame_rate,
params.positional_embedding_theta,
config.positional_embedding_theta,
temporal_max_pos,
std::get<0>(params.vae_scale_factors),
params.causal_temporal_positioning,
std::get<0>(config.vae_scale_factors),
config.causal_temporal_positioning,
true);
}
video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, video_token_count * params.audio_num_attention_heads);
video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads);
ggml_set_name(video_cross_pe, "ltxav_video_cross_pe");
set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data());
audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1],
static_cast<int>(params.audio_cross_attention_dim),
static_cast<int>(params.audio_num_attention_heads),
params.positional_embedding_theta,
static_cast<int>(config.audio_cross_attention_dim),
static_cast<int>(config.audio_num_attention_heads),
config.positional_embedding_theta,
temporal_max_pos,
true);
audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, ax->ne[1] * params.audio_num_attention_heads);
audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe");
set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data());
}
bool needs_video_connector_pe =
params.use_connector &&
config.use_connector &&
context != nullptr &&
(context->ne[0] == params.connector_hidden_size ||
((context->ne[0] == params.cross_attention_dim + params.audio_cross_attention_dim ||
context->ne[0] == params.caption_channels * 2) &&
(context->ne[0] == config.connector_hidden_size ||
((context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim ||
context->ne[0] == config.caption_channels * 2) &&
context->ne[1] < 1024));
ggml_tensor* video_connector_pe = nullptr;
if (needs_video_connector_pe) {
int64_t seq_len = context->ne[1];
int64_t target_len = std::max<int64_t>(1024, seq_len);
int64_t duplications = (target_len + params.connector_num_registers - 1) / params.connector_num_registers;
int64_t full_len = seq_len + duplications * params.connector_num_registers - seq_len;
connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(params.connector_hidden_size), static_cast<int>(params.connector_num_heads), 10000.f, 4096.f, true);
video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.connector_head_dim / 2, full_len * params.connector_num_heads);
int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers;
int64_t full_len = seq_len + duplications * config.connector_num_registers - seq_len;
connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.connector_hidden_size), static_cast<int>(config.connector_num_heads), 10000.f, 4096.f, true);
video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads);
ggml_set_name(video_connector_pe, "ltxav_video_connector_pe");
set_backend_tensor_data(video_connector_pe, connector_pe_vec.data());
}
@@ -1864,20 +1890,20 @@ namespace LTXV {
ax->ne[1] > 0;
bool needs_audio_connector_pe =
run_audio_context &&
params.use_audio_connector &&
config.use_audio_connector &&
context != nullptr &&
(context->ne[0] == params.audio_connector_hidden_size ||
((context->ne[0] == params.cross_attention_dim + params.audio_cross_attention_dim ||
context->ne[0] == params.caption_channels * 2) &&
(context->ne[0] == config.audio_connector_hidden_size ||
((context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim ||
context->ne[0] == config.caption_channels * 2) &&
context->ne[1] < 1024));
ggml_tensor* audio_connector_pe = nullptr;
if (needs_audio_connector_pe) {
int64_t seq_len = context->ne[1];
int64_t target_len = std::max<int64_t>(1024, seq_len);
int64_t duplications = (target_len + params.audio_connector_num_registers - 1) / params.audio_connector_num_registers;
int64_t full_len = seq_len + duplications * params.audio_connector_num_registers - seq_len;
audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(params.audio_connector_hidden_size), static_cast<int>(params.audio_connector_num_heads), 10000.f, 4096.f, true);
audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_connector_head_dim / 2, full_len * params.audio_connector_num_heads);
int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers;
int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len;
audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.audio_connector_hidden_size), static_cast<int>(config.audio_connector_num_heads), 10000.f, 4096.f, true);
audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads);
ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe");
set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data());
}
@@ -1995,7 +2021,7 @@ namespace LTXV {
const std::string& audio_x_path = "",
const std::string& audio_timesteps_path = "") {
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
LOG_INFO("loading ltxav from '%s'", model_path.c_str());
ModelLoader model_loader;
otherarch/sdcpp/main.cpp
+2 -1
View File
@@ -169,8 +169,9 @@ struct SDCliParams {
return 1;
};
auto on_help_arg = [&](int argc, const char** argv, int index) {
auto on_help_arg = [&](int argc, const char** argv, int index, bool& valid) {
normal_exit = true;
valid = true;
return -1;
};
otherarch/sdcpp/mmdit.hpp
+169 -87
View File
@@ -1,7 +1,10 @@
#ifndef __MMDIT_HPP__
#define __MMDIT_HPP__
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "diffusion_model.hpp"
#include "ggml_extend.hpp"
@@ -9,6 +12,128 @@
#define MMDIT_GRAPH_SIZE 10240
struct MMDiTConfig {
int64_t input_size = -1;
int patch_size = 2;
int64_t in_channels = 16;
int64_t d_self = -1; // >=0 for MMdiT-X
int64_t depth = 24;
float mlp_ratio = 4.0f;
int64_t adm_in_channels = 2048;
int64_t out_channels = 16;
int64_t pos_embed_max_size = 192;
int64_t num_patches = 36864; // 192 * 192
int64_t context_size = 4096;
int64_t context_embedder_out_dim = 1536;
int64_t hidden_size = 1536;
std::string qk_norm;
static MMDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
MMDiTConfig config;
bool has_weight_config = false;
bool has_pos_embed = false;
bool has_hidden_size = false;
bool has_context_embed = false;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (name.find("x_embedder.proj.weight") != std::string::npos && tensor_storage.n_dims == 4) {
has_weight_config = true;
has_hidden_size = true;
config.patch_size = static_cast<int>(tensor_storage.ne[0]);
config.in_channels = tensor_storage.ne[2];
config.hidden_size = tensor_storage.ne[3];
} else if (name.find("t_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_hidden_size = true;
config.hidden_size = tensor_storage.ne[1];
} else if (name.find("y_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_hidden_size = true;
config.adm_in_channels = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
} else if (name.find("context_embedder.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_context_embed = true;
config.context_size = tensor_storage.ne[0];
config.context_embedder_out_dim = tensor_storage.ne[1];
} else if (name.find("final_layer.linear.weight") != std::string::npos && tensor_storage.n_dims == 2) {
has_weight_config = true;
has_hidden_size = true;
config.hidden_size = tensor_storage.ne[0];
int64_t patch_area = static_cast<int64_t>(config.patch_size) * config.patch_size;
if (patch_area > 0) {
config.out_channels = tensor_storage.ne[1] / patch_area;
}
} else if (name.find("pos_embed") != std::string::npos && tensor_storage.n_dims == 3) {
has_weight_config = true;
has_pos_embed = true;
has_hidden_size = true;
config.hidden_size = tensor_storage.ne[0];
config.num_patches = tensor_storage.ne[1];
for (int64_t size = 1; size * size <= config.num_patches; size++) {
if (size * size == config.num_patches) {
config.pos_embed_max_size = size;
break;
}
}
}
size_t jb = name.find("joint_blocks.");
if (jb == std::string::npos) {
continue;
}
has_weight_config = true;
std::string block_name = name.substr(jb);
int64_t block_depth = atoi(block_name.substr(13, block_name.find(".", 13)).c_str());
if (block_depth + 1 > config.depth) {
config.depth = block_depth + 1;
}
if (block_name.find("attn.ln") != std::string::npos) {
if (block_name.find(".bias") != std::string::npos) {
config.qk_norm = "ln";
} else {
config.qk_norm = "rms";
}
}
if (block_name.find("attn2") != std::string::npos) {
if (block_depth > config.d_self) {
config.d_self = block_depth;
}
}
}
if (!has_pos_embed && config.d_self >= 0) {
config.pos_embed_max_size *= 2;
config.num_patches *= 4;
}
if (!has_hidden_size || config.hidden_size <= 0) {
config.hidden_size = 64 * config.depth;
}
if (!has_context_embed || config.context_embedder_out_dim <= 0) {
config.context_embedder_out_dim = config.hidden_size;
}
if (has_weight_config) {
LOG_DEBUG("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s",
config.depth,
config.d_self + 1,
config.hidden_size,
config.patch_size,
config.in_channels,
config.out_channels,
config.context_size,
config.adm_in_channels,
config.qk_norm.empty() ? "none" : config.qk_norm.c_str());
}
return config;
}
};
struct Mlp : public GGMLBlock {
public:
Mlp(int64_t in_features,
@@ -612,28 +737,16 @@ public:
struct MMDiT : public GGMLBlock {
// Diffusion model with a Transformer backbone.
protected:
int64_t input_size = -1;
int patch_size = 2;
int64_t in_channels = 16;
int64_t d_self = -1; // >=0 for MMdiT-X
int64_t depth = 24;
float mlp_ratio = 4.0f;
int64_t adm_in_channels = 2048;
int64_t out_channels = 16;
int64_t pos_embed_max_size = 192;
int64_t num_patchs = 36864; // 192 * 192
int64_t context_size = 4096;
int64_t context_embedder_out_dim = 1536;
int64_t hidden_size;
std::string qk_norm;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
enum ggml_type wtype = GGML_TYPE_F32;
params["pos_embed"] = ggml_new_tensor_3d(ctx, wtype, hidden_size, num_patchs, 1);
params["pos_embed"] = ggml_new_tensor_3d(ctx, wtype, config.hidden_size, config.num_patches, 1);
}
public:
MMDiT(const String2TensorStorage& tensor_storage_map = {}) {
MMDiTConfig config;
explicit MMDiT(MMDiTConfig config = {})
: config(config) {
// input_size is always None
// learn_sigma is always False
// register_length is alwalys 0
@@ -646,64 +759,30 @@ public:
// pos_embed_offset is not used
// context_embedder_config is always {'target': 'torch.nn.Linear', 'params': {'in_features': 4096, 'out_features': 1536}}
for (auto pair : tensor_storage_map) {
std::string tensor_name = pair.first;
if (tensor_name.find("model.diffusion_model.") == std::string::npos)
continue;
size_t jb = tensor_name.find("joint_blocks.");
if (jb != std::string::npos) {
tensor_name = tensor_name.substr(jb); // remove prefix
int block_depth = atoi(tensor_name.substr(13, tensor_name.find(".", 13)).c_str());
if (block_depth + 1 > depth) {
depth = block_depth + 1;
}
if (tensor_name.find("attn.ln") != std::string::npos) {
if (tensor_name.find(".bias") != std::string::npos) {
qk_norm = "ln";
} else {
qk_norm = "rms";
}
}
if (tensor_name.find("attn2") != std::string::npos) {
if (block_depth > d_self) {
d_self = block_depth;
}
}
}
blocks["x_embedder"] = std::shared_ptr<GGMLBlock>(new PatchEmbed(config.input_size,
config.patch_size,
config.in_channels,
config.hidden_size,
true));
blocks["t_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedder(config.hidden_size));
if (config.adm_in_channels != -1) {
blocks["y_embedder"] = std::shared_ptr<GGMLBlock>(new VectorEmbedder(config.adm_in_channels, config.hidden_size));
}
if (d_self >= 0) {
pos_embed_max_size *= 2;
num_patchs *= 4;
}
blocks["context_embedder"] = std::shared_ptr<GGMLBlock>(new Linear(config.context_size, config.context_embedder_out_dim, true, true));
LOG_INFO("MMDiT layers: %d (including %d MMDiT-x layers)", depth, d_self + 1);
int64_t default_out_channels = in_channels;
hidden_size = 64 * depth;
context_embedder_out_dim = 64 * depth;
int64_t num_heads = depth;
blocks["x_embedder"] = std::shared_ptr<GGMLBlock>(new PatchEmbed(input_size, patch_size, in_channels, hidden_size, true));
blocks["t_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedder(hidden_size));
if (adm_in_channels != -1) {
blocks["y_embedder"] = std::shared_ptr<GGMLBlock>(new VectorEmbedder(adm_in_channels, hidden_size));
}
blocks["context_embedder"] = std::shared_ptr<GGMLBlock>(new Linear(4096, context_embedder_out_dim, true, true));
for (int i = 0; i < depth; i++) {
blocks["joint_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new JointBlock(hidden_size,
num_heads,
mlp_ratio,
qk_norm,
for (int i = 0; i < config.depth; i++) {
blocks["joint_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new JointBlock(config.hidden_size,
config.depth,
config.mlp_ratio,
config.qk_norm,
true,
i == depth - 1,
i <= d_self));
i == config.depth - 1,
i <= config.d_self));
}
blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(hidden_size, patch_size, out_channels));
blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(config.hidden_size, config.patch_size, config.out_channels));
}
ggml_tensor*
@@ -712,22 +791,22 @@ public:
int64_t w) {
auto pos_embed = params["pos_embed"];
h = (h + 1) / patch_size;
w = (w + 1) / patch_size;
h = (h + 1) / config.patch_size;
w = (w + 1) / config.patch_size;
GGML_ASSERT(h <= pos_embed_max_size && h > 0);
GGML_ASSERT(w <= pos_embed_max_size && w > 0);
GGML_ASSERT(h <= config.pos_embed_max_size && h > 0);
GGML_ASSERT(w <= config.pos_embed_max_size && w > 0);
int64_t top = (pos_embed_max_size - h) / 2;
int64_t left = (pos_embed_max_size - w) / 2;
int64_t top = (config.pos_embed_max_size - h) / 2;
int64_t left = (config.pos_embed_max_size - w) / 2;
auto spatial_pos_embed = ggml_reshape_3d(ctx, pos_embed, hidden_size, pos_embed_max_size, pos_embed_max_size);
auto spatial_pos_embed = ggml_reshape_3d(ctx, pos_embed, config.hidden_size, config.pos_embed_max_size, config.pos_embed_max_size);
// spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :]
spatial_pos_embed = ggml_view_3d(ctx,
spatial_pos_embed,
hidden_size,
pos_embed_max_size,
config.hidden_size,
config.pos_embed_max_size,
h,
spatial_pos_embed->nb[1],
spatial_pos_embed->nb[2],
@@ -735,14 +814,14 @@ public:
spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3)); // [pos_embed_max_size, h, hidden_size]
spatial_pos_embed = ggml_view_3d(ctx,
spatial_pos_embed,
hidden_size,
config.hidden_size,
h,
w,
spatial_pos_embed->nb[1],
spatial_pos_embed->nb[2],
spatial_pos_embed->nb[2] * left); // [w, h, hidden_size]
spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3)); // [h, w, hidden_size]
spatial_pos_embed = ggml_reshape_3d(ctx, spatial_pos_embed, hidden_size, h * w, 1); // [1, h*w, hidden_size]
spatial_pos_embed->nb[2] * left); // [w, h, hidden_size]
spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3)); // [h, w, hidden_size]
spatial_pos_embed = ggml_reshape_3d(ctx, spatial_pos_embed, config.hidden_size, h * w, 1); // [1, h*w, hidden_size]
return spatial_pos_embed;
}
@@ -757,7 +836,7 @@ public:
// return: [N, N*W, patch_size * patch_size * out_channels]
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
for (int i = 0; i < depth; i++) {
for (int i = 0; i < config.depth; i++) {
// skip iteration if i is in skip_layers
if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
continue;
@@ -800,7 +879,7 @@ public:
x = ggml_add(ctx->ggml_ctx, patch_embed, pos_embed); // [N, H*W, hidden_size]
auto c = t_embedder->forward(ctx, t); // [N, hidden_size]
if (y != nullptr && adm_in_channels != -1) {
if (y != nullptr && config.adm_in_channels != -1) {
auto y_embedder = std::dynamic_pointer_cast<VectorEmbedder>(blocks["y_embedder"]);
y = y_embedder->forward(ctx, y); // [N, hidden_size]
@@ -820,19 +899,22 @@ public:
x = forward_core_with_concat(ctx, x, c, context, skip_layers); // (N, H*W, patch_size ** 2 * out_channels)
x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, /*patch_last*/ false); // [N, C, H, W]
x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false); // [N, C, H, W]
return x;
}
};
struct MMDiTRunner : public DiffusionModelRunner {
MMDiTConfig config;
MMDiT mmdit;
MMDiTRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "")
: DiffusionModelRunner(backend, params_backend, prefix), mmdit(tensor_storage_map) {
: DiffusionModelRunner(backend, params_backend, prefix),
config(MMDiTConfig::detect_from_weights(tensor_storage_map, prefix)),
mmdit(config) {
mmdit.init(params_ctx, tensor_storage_map, prefix);
}
@@ -947,7 +1029,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
static void load_from_file_and_test(const std::string& file_path) {
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_F16;
std::shared_ptr<MMDiTRunner> mmdit = std::make_shared<MMDiTRunner>(backend, backend);
{
otherarch/sdcpp/model.cpp
+8
View File
@@ -460,6 +460,12 @@ SDVersion ModelLoader::get_sd_version() {
tensor_storage.name.find("model.diffusion_model.single_transformer_blocks.") != std::string::npos) {
is_flux = true;
}
if (tensor_storage.name.find("model.diffusion_model.net.lq_proj.latent_proj.0.weight") != std::string::npos) {
return VERSION_PID;
}
if (tensor_storage.name.find("embed_image_indicator.weight") != std::string::npos) {
return VERSION_IDEOGRAM4;
}
if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
return VERSION_CHROMA_RADIANCE;
}
@@ -1279,6 +1285,8 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
// Pass, do not convert
} else if (ends_with(name, ".scale")) {
// Pass, do not convert
} else if (ends_with(name, ".weight_scale")) {
// Pass, do not convert
} else if (contains(name, "img_in.") ||
contains(name, "txt_in.") ||
contains(name, "time_in.") ||
otherarch/sdcpp/model.h
+20 -2
View File
@@ -49,6 +49,8 @@ enum SDVersion {
VERSION_ERNIE_IMAGE,
VERSION_LENS,
VERSION_LONGCAT,
VERSION_PID,
VERSION_IDEOGRAM4,
VERSION_COUNT,
};
@@ -164,8 +166,22 @@ static inline bool sd_version_is_lens(SDVersion version) {
return false;
}
static inline bool sd_version_is_pid(SDVersion version) {
if (version == VERSION_PID) {
return true;
}
return false;
}
static inline bool sd_version_is_ideogram4(SDVersion version) {
if (version == VERSION_IDEOGRAM4) {
return true;
}
return false;
}
static inline bool sd_version_uses_flux2_vae(SDVersion version) {
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version)) {
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version)) {
return true;
}
return false;
@@ -194,7 +210,9 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_z_image(version) ||
sd_version_is_ernie_image(version) ||
sd_version_is_lens(version) ||
sd_version_is_longcat(version)) {
sd_version_is_longcat(version) ||
sd_version_is_pid(version) ||
sd_version_is_ideogram4(version)) {
return true;
}
return false;
otherarch/sdcpp/pid.hpp
+847
View File
@@ -0,0 +1,847 @@
#ifndef __SD_PID_HPP__
#define __SD_PID_HPP__
#include <cmath>
#include <cstdlib>
#include <memory>
#include <string>
#include <vector>
#include "common_dit.hpp"
#include "ggml_extend.hpp"
#include "mmdit.hpp"
#include "rope.hpp"
namespace Pid {
constexpr int PID_GRAPH_SIZE = 196608;
constexpr float PID_PI = 3.14159265358979323846f;
struct PixelDiTConfig {
int64_t in_channels = 3;
int64_t hidden_size = 1536;
int64_t num_groups = 24;
int64_t patch_mlp_hidden_dim = 4096;
int64_t pixel_hidden_size = 16;
int64_t pixel_attn_hidden_size = 1152;
int64_t pixel_num_groups = 16;
int64_t patch_depth = 14;
int64_t pixel_depth = 2;
int64_t patch_size = 16;
int64_t txt_embed_dim = 2304;
int64_t txt_max_length = 300;
float text_rope_theta = 10000.f;
int64_t lq_latent_channels = 16;
int64_t lq_hidden_dim = 512;
int64_t lq_num_res_blocks = 4;
int64_t lq_interval = 2;
int64_t lq_sr_scale = 4;
int64_t lq_latent_down_factor = 8;
int64_t rope_ref_grid_h = 64;
int64_t rope_ref_grid_w = 64;
static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
PixelDiTConfig config;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
size_t pos = name.find("patch_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
config.patch_depth = std::max<int64_t>(config.patch_depth, block_index + 1);
}
}
pos = name.find("pixel_blocks.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
}
}
if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
config.lq_latent_channels = tensor_storage.ne[2];
config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
}
if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
config.patch_mlp_hidden_dim = tensor_storage.ne[1];
}
}
LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64,
config.patch_depth,
config.pixel_depth,
config.patch_mlp_hidden_dim,
config.lq_latent_channels,
config.lq_latent_down_factor);
return config;
}
};
inline std::vector<float> make_rope_1d(int length,
int dim,
float theta) {
GGML_ASSERT(dim % 2 == 0);
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
}
inline std::vector<float> make_rope_2d(int height,
int width,
int dim,
float theta = 10000.f,
float scale = 16.f,
int ref_grid_h = 0,
int ref_grid_w = 0) {
GGML_ASSERT(dim % 4 == 0);
return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w);
}
inline std::vector<float> make_pixel_abs_pos(int height,
int width,
int dim) {
GGML_ASSERT(dim % 4 == 0);
int half_dim = dim / 2;
std::vector<float> x_pos;
std::vector<float> y_pos;
x_pos.reserve(static_cast<size_t>(height) * width);
y_pos.reserve(static_cast<size_t>(height) * width);
for (int iy = 0; iy < height; ++iy) {
for (int ix = 0; ix < width; ++ix) {
x_pos.push_back(static_cast<float>(ix));
y_pos.push_back(static_cast<float>(iy));
}
}
auto x_emb = timestep_embedding(x_pos, half_dim, 10000, false);
auto y_emb = timestep_embedding(y_pos, half_dim, 10000, false);
std::vector<float> out(static_cast<size_t>(dim) * height * width);
for (int pos = 0; pos < height * width; ++pos) {
size_t out_base = static_cast<size_t>(pos) * dim;
size_t emb_base = static_cast<size_t>(pos) * half_dim;
for (int i = 0; i < half_dim; ++i) {
out[out_base + i] = x_emb[emb_base + i];
out[out_base + half_dim + i] = y_emb[emb_base + i];
}
}
return out;
}
inline ggml_tensor* apply_adaln(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
}
struct PatchTokenEmbedder : public GGMLBlock {
bool use_rms_norm;
PatchTokenEmbedder(int64_t in_chans,
int64_t embed_dim,
bool use_rms_norm = false,
bool bias = true)
: use_rms_norm(use_rms_norm) {
blocks["proj"] = std::make_shared<Linear>(in_chans, embed_dim, bias);
if (use_rms_norm) {
blocks["norm"] = std::make_shared<RMSNorm>(embed_dim, 1e-6f);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
x = proj->forward(ctx, x);
if (use_rms_norm) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
x = norm->forward(ctx, x);
}
return x;
}
};
struct PixelDiTTimestepEmbedder : public GGMLBlock {
int frequency_embedding_size;
PixelDiTTimestepEmbedder(int64_t hidden_size,
int frequency_embedding_size = 256)
: frequency_embedding_size(frequency_embedding_size) {
blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10);
t_emb = mlp_0->forward(ctx, t_emb);
t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
return mlp_2->forward(ctx, t_emb);
}
};
struct FeedForward : public GGMLBlock {
FeedForward(int64_t dim, int64_t hidden_dim) {
blocks["w1"] = std::make_shared<Linear>(dim, hidden_dim, false);
blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false);
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
auto h = ggml_silu_inplace(ctx->ggml_ctx, w1->forward(ctx, x));
h = ggml_mul_inplace(ctx->ggml_ctx, h, w3->forward(ctx, x));
return w2->forward(ctx, h);
}
};
struct FinalLayer : public GGMLBlock {
FinalLayer(int64_t hidden_size, int64_t out_channels) {
blocks["norm"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["linear"] = std::make_shared<Linear>(hidden_size, out_channels, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
return linear->forward(ctx, norm->forward(ctx, x));
}
};
struct RotaryAttention : public GGMLBlock {
int64_t dim;
int64_t num_heads;
RotaryAttention(int64_t dim, int64_t num_heads)
: dim(dim), num_heads(num_heads) {
int64_t head_dim = dim / num_heads;
blocks["qkv"] = std::make_shared<Linear>(dim, dim * 3, false);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["proj"] = std::make_shared<Linear>(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pos) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto qkv = qkv_proj->forward(ctx, x);
auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
int64_t L = x->ne[1];
int64_t N = x->ne[2];
int64_t head_dim = dim / num_heads;
auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, L, N);
auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, L, N);
auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, L, N);
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
x = Rope::attention(ctx, q, k, v, pos, nullptr, 1.0f / 128.f, true);
return proj->forward(ctx, x);
}
};
struct MMDiTJointAttention : public GGMLBlock {
int64_t dim;
int64_t num_heads;
MMDiTJointAttention(int64_t dim, int64_t num_heads)
: dim(dim), num_heads(num_heads) {
int64_t head_dim = dim / num_heads;
blocks["qkv_x"] = std::make_shared<Linear>(dim, dim * 3, false);
blocks["qkv_y"] = std::make_shared<Linear>(dim, dim * 3, false);
blocks["q_norm_x"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm_x"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["q_norm_y"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["k_norm_y"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
blocks["proj_x"] = std::make_shared<Linear>(dim, dim, true);
blocks["proj_y"] = std::make_shared<Linear>(dim, dim, true);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* y,
ggml_tensor* pos_img,
ggml_tensor* pos_txt) {
auto qkv_x_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_x"]);
auto qkv_y_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_y"]);
auto q_norm_x = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm_x"]);
auto k_norm_x = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm_x"]);
auto q_norm_y = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm_y"]);
auto k_norm_y = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm_y"]);
auto proj_x = std::dynamic_pointer_cast<Linear>(blocks["proj_x"]);
auto proj_y = std::dynamic_pointer_cast<Linear>(blocks["proj_y"]);
int64_t Nx = x->ne[1];
int64_t Ny = y->ne[1];
int64_t N = x->ne[2];
int64_t head_dim = dim / num_heads;
auto qkv_x = split_qkv(ctx->ggml_ctx, qkv_x_proj->forward(ctx, x));
auto qx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N);
auto kx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N);
auto vx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N);
qx = q_norm_x->forward(ctx, qx);
kx = k_norm_x->forward(ctx, kx);
auto qkv_y = split_qkv(ctx->ggml_ctx, qkv_y_proj->forward(ctx, y));
auto qy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N);
auto ky = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N);
auto vy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N);
qy = q_norm_y->forward(ctx, qy);
ky = k_norm_y->forward(ctx, ky);
auto q_joint = ggml_concat(ctx->ggml_ctx, qy, qx, 2);
auto k_joint = ggml_concat(ctx->ggml_ctx, ky, kx, 2);
auto v_joint = ggml_concat(ctx->ggml_ctx, vy, vx, 2);
auto pos_joint = ggml_concat(ctx->ggml_ctx, pos_txt, pos_img, 3);
auto out = Rope::attention(ctx, q_joint, k_joint, v_joint, pos_joint, nullptr, 1.0f, true);
auto out_y = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, Ny);
auto out_x = ggml_ext_slice(ctx->ggml_ctx, out, 1, Ny, Ny + Nx);
return {proj_x->forward(ctx, out_x), proj_y->forward(ctx, out_y)};
}
};
struct MMDiTBlockT2I : public GGMLBlock {
int64_t hidden_size;
MMDiTBlockT2I(int64_t hidden_size, int64_t groups, int64_t mlp_hidden_dim)
: hidden_size(hidden_size) {
blocks["norm_x1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["norm_y1"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["attn"] = std::make_shared<MMDiTJointAttention>(hidden_size, groups);
blocks["norm_x2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["norm_y2"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
blocks["mlp_x"] = std::make_shared<FeedForward>(hidden_size, mlp_hidden_dim);
blocks["mlp_y"] = std::make_shared<FeedForward>(hidden_size, mlp_hidden_dim);
blocks["adaLN_modulation_img.0"] = std::make_shared<Linear>(hidden_size, 6 * hidden_size, true);
blocks["adaLN_modulation_txt.0"] = std::make_shared<Linear>(hidden_size, 6 * hidden_size, true);
}
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* y,
ggml_tensor* c,
ggml_tensor* pos_img,
ggml_tensor* pos_txt) {
auto norm_x1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_x1"]);
auto norm_y1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_y1"]);
auto attn = std::dynamic_pointer_cast<MMDiTJointAttention>(blocks["attn"]);
auto norm_x2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_x2"]);
auto norm_y2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_y2"]);
auto mlp_x = std::dynamic_pointer_cast<FeedForward>(blocks["mlp_x"]);
auto mlp_y = std::dynamic_pointer_cast<FeedForward>(blocks["mlp_y"]);
auto ada_img = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_img.0"]);
auto ada_txt = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_txt.0"]);
auto mx = ggml_ext_chunk(ctx->ggml_ctx, ada_img->forward(ctx, c), 6, 0);
auto my = ggml_ext_chunk(ctx->ggml_ctx, ada_txt->forward(ctx, c), 6, 0);
auto x_norm = apply_adaln(ctx->ggml_ctx, norm_x1->forward(ctx, x), mx[0], mx[1]);
auto y_norm = apply_adaln(ctx->ggml_ctx, norm_y1->forward(ctx, y), my[0], my[1]);
auto attn_out = attn->forward(ctx, x_norm, y_norm, pos_img, pos_txt);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out.first, mx[2]));
y = ggml_add(ctx->ggml_ctx, y, ggml_mul(ctx->ggml_ctx, attn_out.second, my[2]));
auto x_mlp = mlp_x->forward(ctx, apply_adaln(ctx->ggml_ctx, norm_x2->forward(ctx, x), mx[3], mx[4]));
auto y_mlp = mlp_y->forward(ctx, apply_adaln(ctx->ggml_ctx, norm_y2->forward(ctx, y), my[3], my[4]));
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, x_mlp, mx[5]));
y = ggml_add(ctx->ggml_ctx, y, ggml_mul(ctx->ggml_ctx, y_mlp, my[5]));
return {x, y};
}
};
struct PixelTokenEmbedder : public GGMLBlock {
int64_t in_channels;
int64_t hidden_size_output;
PixelTokenEmbedder(int64_t in_channels, int64_t hidden_size_output)
: in_channels(in_channels), hidden_size_output(hidden_size_output) {
blocks["proj"] = std::make_shared<Linear>(in_channels, hidden_size_output, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* inputs,
int64_t patch_size,
ggml_tensor* pos_full) {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
int64_t W = inputs->ne[0];
int64_t H = inputs->ne[1];
int64_t B = inputs->ne[3];
int64_t L = (W / patch_size) * (H / patch_size);
int64_t P2 = patch_size * patch_size;
auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, inputs, 2, 0, 1, 3));
x = ggml_reshape_3d(ctx->ggml_ctx, x, in_channels, W * H, B);
x = proj->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, pos_full);
x = ggml_reshape_4d(ctx->ggml_ctx, x, hidden_size_output, W, H, B);
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
x = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size_output, P2, L * B);
return x;
}
};
struct PiTBlock : public GGMLBlock {
int64_t pixel_dim;
int64_t context_dim;
int64_t attn_dim;
int64_t num_heads;
int64_t patch_size;
PiTBlock(int64_t pixel_dim,
int64_t context_dim,
int64_t patch_size,
int64_t attn_dim,
int64_t num_heads)
: pixel_dim(pixel_dim),
context_dim(context_dim),
attn_dim(attn_dim),
num_heads(num_heads),
patch_size(patch_size) {
int64_t p2 = patch_size * patch_size;
blocks["compress_to_attn"] = std::make_shared<Linear>(p2 * pixel_dim, attn_dim, true);
blocks["expand_from_attn"] = std::make_shared<Linear>(attn_dim, p2 * pixel_dim, true);
blocks["norm1"] = std::make_shared<RMSNorm>(pixel_dim, 1e-6f);
blocks["attn"] = std::make_shared<RotaryAttention>(attn_dim, num_heads);
blocks["norm2"] = std::make_shared<RMSNorm>(pixel_dim, 1e-6f);
blocks["mlp"] = std::make_shared<Mlp>(pixel_dim, pixel_dim * 4);
blocks["adaLN_modulation.0"] = std::make_shared<Linear>(context_dim, 6 * pixel_dim * p2, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* s_cond,
int64_t image_height,
int64_t image_width,
ggml_tensor* pos_comp) {
auto compress = std::dynamic_pointer_cast<Linear>(blocks["compress_to_attn"]);
auto expand = std::dynamic_pointer_cast<Linear>(blocks["expand_from_attn"]);
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<RotaryAttention>(blocks["attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto mlp = std::dynamic_pointer_cast<Mlp>(blocks["mlp"]);
auto ada = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
int64_t Hs = image_height / patch_size;
int64_t Ws = image_width / patch_size;
int64_t L = Hs * Ws;
int64_t BL = x->ne[2];
int64_t B = BL / L;
int64_t P2 = patch_size * patch_size;
auto ada_params = ada->forward(ctx, s_cond);
ada_params = ggml_reshape_3d(ctx->ggml_ctx, ada_params, 6 * pixel_dim, P2, BL);
auto mod = ggml_ext_chunk(ctx->ggml_ctx, ada_params, 6, 0);
auto x_norm = apply_adaln(ctx->ggml_ctx, norm1->forward(ctx, x), mod[0], mod[1]);
auto x_flat = ggml_reshape_2d(ctx->ggml_ctx, x_norm, P2 * pixel_dim, BL);
auto x_comp = compress->forward(ctx, x_flat);
x_comp = ggml_reshape_3d(ctx->ggml_ctx, x_comp, attn_dim, L, B);
auto attn_out = attn->forward(ctx, x_comp, pos_comp);
auto attn_flat = expand->forward(ctx, ggml_reshape_2d(ctx->ggml_ctx, attn_out, attn_dim, BL));
auto attn_exp = ggml_reshape_3d(ctx->ggml_ctx, attn_flat, pixel_dim, P2, BL);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_exp, mod[2]));
auto mlp_out = mlp->forward(ctx, apply_adaln(ctx->ggml_ctx, norm2->forward(ctx, x), mod[3], mod[4]));
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mod[5]));
}
};
struct SigmaAwareGate : public GGMLBlock {
int64_t dim;
SigmaAwareGate(int64_t dim)
: dim(dim) {
blocks["content_proj"] = std::make_shared<Linear>(dim * 2, dim, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
params["log_alpha"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* lq,
ggml_tensor* sigma) {
auto content_proj = std::dynamic_pointer_cast<Linear>(blocks["content_proj"]);
auto content_logit = content_proj->forward(ctx, ggml_concat(ctx->ggml_ctx, x, lq, 0));
sigma = ggml_reshape_3d(ctx->ggml_ctx, sigma, 1, 1, sigma->ne[0]);
auto alpha = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
auto offset = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma));
auto gate = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset));
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq));
}
};
struct PiDResBlock : public GGMLBlock {
PiDResBlock(int64_t channels) {
blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.0"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["block.2"]);
auto norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
h = conv1->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
h = conv2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct LQProjection2D : public GGMLBlock {
PixelDiTConfig config;
LQProjection2D(const PixelDiTConfig& config)
: config(config) {
blocks["latent_proj.0"] = std::make_shared<Conv2d>(config.lq_latent_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim);
}
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
for (int i = 0; i < num_outputs; ++i) {
blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size);
}
}
bool is_gate_active(int block_idx) const {
return block_idx % config.lq_interval == 0;
}
int get_output_index(int block_idx) const {
return block_idx / static_cast<int>(config.lq_interval);
}
ggml_tensor* gate(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* lq,
ggml_tensor* sigma,
int out_idx) {
auto gate_module = std::dynamic_pointer_cast<SigmaAwareGate>(blocks["gate_modules." + std::to_string(out_idx)]);
return gate_module->forward(ctx, x, lq, sigma);
}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx,
ggml_tensor* lq_latent,
int64_t target_pH,
int64_t target_pW) {
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
static_cast<float>(config.patch_size);
GGML_ASSERT(z_to_patch_ratio >= 1.0f);
if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
lq_latent = ggml_interpolate(ctx->ggml_ctx,
lq_latent,
target_pW,
target_pH,
lq_latent->ne[2],
lq_latent->ne[3],
GGML_SCALE_MODE_NEAREST);
}
auto feat = conv0->forward(ctx, lq_latent);
feat = ggml_silu_inplace(ctx->ggml_ctx, feat);
feat = conv2->forward(ctx, feat);
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
feat = block->forward(ctx, feat);
}
int64_t B = feat->ne[3];
int64_t C = feat->ne[2];
int64_t L = target_pH * target_pW;
auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3));
tokens = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B);
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
std::vector<ggml_tensor*> outputs;
outputs.reserve(num_outputs);
for (int i = 0; i < num_outputs; ++i) {
auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
outputs.push_back(head->forward(ctx, tokens));
}
return outputs;
}
};
struct PixelDiT : public GGMLBlock {
PixelDiTConfig config;
PixelDiT() = default;
PixelDiT(const PixelDiTConfig& config)
: config(config) {
blocks["pixel_embedder"] = std::make_shared<PixelTokenEmbedder>(config.in_channels, config.pixel_hidden_size);
blocks["s_embedder"] = std::make_shared<PatchTokenEmbedder>(config.in_channels * config.patch_size * config.patch_size, config.hidden_size, false, true);
blocks["t_embedder"] = std::make_shared<PixelDiTTimestepEmbedder>(config.hidden_size);
blocks["y_embedder"] = std::make_shared<PatchTokenEmbedder>(config.txt_embed_dim, config.hidden_size, true, true);
for (int i = 0; i < config.patch_depth; ++i) {
blocks["patch_blocks." + std::to_string(i)] = std::make_shared<MMDiTBlockT2I>(config.hidden_size, config.num_groups, config.patch_mlp_hidden_dim);
}
for (int i = 0; i < config.pixel_depth; ++i) {
blocks["pixel_blocks." + std::to_string(i)] = std::make_shared<PiTBlock>(config.pixel_hidden_size,
config.hidden_size,
config.patch_size,
config.pixel_attn_hidden_size,
config.pixel_num_groups);
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
blocks["lq_proj"] = std::make_shared<LQProjection2D>(config);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config.hidden_size, config.txt_max_length, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timesteps,
ggml_tensor* context,
ggml_tensor* lq_latent,
ggml_tensor* degrade_sigma,
ggml_tensor* pos_img,
ggml_tensor* pos_txt,
ggml_tensor* pixel_pos_full,
ggml_tensor* pixel_pos_comp) {
auto pixel_embedder = std::dynamic_pointer_cast<PixelTokenEmbedder>(blocks["pixel_embedder"]);
auto s_embedder = std::dynamic_pointer_cast<PatchTokenEmbedder>(blocks["s_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<PixelDiTTimestepEmbedder>(blocks["t_embedder"]);
auto y_embedder = std::dynamic_pointer_cast<PatchTokenEmbedder>(blocks["y_embedder"]);
auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
auto lq_proj = std::dynamic_pointer_cast<LQProjection2D>(blocks["lq_proj"]);
int64_t W_orig = x->ne[0];
int64_t H_orig = x->ne[1];
x = DiT::pad_to_patch_size(ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size));
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t B = x->ne[3];
int64_t Hs = H / config.patch_size;
int64_t Ws = W / config.patch_size;
int64_t L = Hs * Ws;
int64_t P2 = config.patch_size * config.patch_size;
auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
auto t_emb = t_embedder->forward(ctx, timesteps);
auto condition = ggml_silu(ctx->ggml_ctx, t_emb);
GGML_ASSERT(context != nullptr);
int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
auto y = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt);
auto y_emb = y_embedder->forward(ctx, y);
auto y_pos = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt);
y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
auto s = s_embedder->forward(ctx, x_patches);
for (int i = 0; i < config.patch_depth; ++i) {
if (lq_proj->is_gate_active(i)) {
int out_idx = lq_proj->get_output_index(i);
if (out_idx < static_cast<int>(lq_features.size())) {
s = lq_proj->gate(ctx, s, lq_features[out_idx], degrade_sigma, out_idx);
}
}
auto block = std::dynamic_pointer_cast<MMDiTBlockT2I>(blocks["patch_blocks." + std::to_string(i)]);
auto out = block->forward(ctx,
s,
y_emb,
condition,
pos_img,
pos_txt);
s = out.first;
y_emb = out.second;
sd::ggml_graph_cut::mark_graph_cut(s, "pid.patch_blocks." + std::to_string(i), "s");
sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y");
}
s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B);
auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);
for (int i = 0; i < config.pixel_depth; ++i) {
auto block = std::dynamic_pointer_cast<PiTBlock>(blocks["pixel_blocks." + std::to_string(i)]);
pixels = block->forward(ctx, pixels, s_cond, H, W, pixel_pos_comp);
sd::ggml_graph_cut::mark_graph_cut(pixels, "pid.pixel_blocks." + std::to_string(i), "pixels");
}
pixels = final_layer->forward(ctx, pixels);
pixels = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B);
auto out = DiT::unpatchify(ctx->ggml_ctx,
pixels,
Hs,
Ws,
static_cast<int>(config.patch_size),
static_cast<int>(config.patch_size),
false);
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig);
out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig);
return out;
}
};
struct PiDRunner : public DiffusionModelRunner {
PixelDiTConfig config;
PixelDiT model;
std::vector<float> pos_img_vec;
std::vector<float> pos_txt_vec;
std::vector<float> pixel_pos_vec;
std::vector<float> pixel_pos_comp_vec;
PiDRunner(ggml_backend_t backend,
ggml_backend_t params_backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix = "model.diffusion_model")
: DiffusionModelRunner(backend, params_backend, prefix),
config(PixelDiTConfig::detect_from_weights(tensor_storage_map, prefix)) {
model = PixelDiT(config);
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "PiD";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
model.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& lq_latent_tensor,
const sd::Tensor<float>& degrade_sigma_tensor) {
ggml_cgraph* gf = new_graph_custom(PID_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* lq_latent = make_input(lq_latent_tensor);
ggml_tensor* degrade_sigma = make_input(degrade_sigma_tensor);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t B = x->ne[3];
int64_t Wp = align_up(static_cast<int>(W), static_cast<int>(config.patch_size));
int64_t Hp = align_up(static_cast<int>(H), static_cast<int>(config.patch_size));
int64_t Hs = Hp / config.patch_size;
int64_t Ws = Wp / config.patch_size;
pos_img_vec = make_rope_2d(static_cast<int>(Hs),
static_cast<int>(Ws),
static_cast<int>(config.hidden_size / config.num_groups),
10000.f,
16.f,
static_cast<int>(config.rope_ref_grid_h),
static_cast<int>(config.rope_ref_grid_w));
auto pos_img = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
2,
config.hidden_size / config.num_groups / 2,
Hs * Ws);
set_backend_tensor_data(pos_img, pos_img_vec.data());
int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
pos_txt_vec = make_rope_1d(static_cast<int>(Ltxt),
static_cast<int>(config.hidden_size / config.num_groups),
config.text_rope_theta);
auto pos_txt = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
2,
config.hidden_size / config.num_groups / 2,
Ltxt);
set_backend_tensor_data(pos_txt, pos_txt_vec.data());
pixel_pos_vec = make_pixel_abs_pos(static_cast<int>(Hp),
static_cast<int>(Wp),
static_cast<int>(config.pixel_hidden_size));
auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
GGML_TYPE_F32,
config.pixel_hidden_size,
Wp * Hp,
1);
set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
pixel_pos_comp_vec = make_rope_2d(static_cast<int>(Hs),
static_cast<int>(Ws),
static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
10000.f,
16.f,
static_cast<int>(config.rope_ref_grid_h),
static_cast<int>(config.rope_ref_grid_w));
auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
GGML_TYPE_F32,
2,
2,
config.pixel_attn_hidden_size / config.pixel_num_groups / 2,
Hs * Ws);
set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data());
auto runner_ctx = get_context();
auto out = model.forward(&runner_ctx,
x,
timesteps,
context,
lq_latent,
degrade_sigma,
pos_img,
pos_txt,
pixel_pos,
pixel_pos_comp);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const sd::Tensor<float>& lq_latent,
const sd::Tensor<float>& degrade_sigma) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
GGML_ASSERT(diffusion_params.context != nullptr);
GGML_ASSERT(diffusion_params.ref_latents != nullptr);
GGML_ASSERT(!diffusion_params.ref_latents->empty());
auto degrade_sigma = sd::Tensor<float>::from_vector({0.0f});
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*diffusion_params.context,
diffusion_params.ref_latents->front(),
degrade_sigma);
}
};
} // namespace Pid
#endif // __SD_PID_HPP__
otherarch/sdcpp/qwen_image.hpp
+76 -72
View File
@@ -10,6 +10,48 @@
namespace Qwen {
constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480;
struct QwenImageConfig {
int patch_size = 2;
int64_t in_channels = 64;
int64_t out_channels = 16;
int num_layers = 60;
int64_t attention_head_dim = 128;
int64_t num_attention_heads = 24;
int64_t joint_attention_dim = 3584;
int theta = 10000;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
bool zero_cond_t = false;
static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
QwenImageConfig config;
config.num_layers = 0;
for (const auto& [name, _] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (name.find("__index_timestep_zero__") != std::string::npos) {
config.zero_cond_t = true;
}
size_t pos = name.find("transformer_blocks.");
if (pos == std::string::npos) {
continue;
}
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.num_layers) {
config.num_layers = block_index + 1;
}
}
}
LOG_DEBUG("qwen_image: num_layers = %d, zero_cond_t = %s",
config.num_layers,
config.zero_cond_t ? "true" : "false");
return config;
}
};
struct TimestepEmbedding : public GGMLBlock {
public:
TimestepEmbedding(int64_t in_channels,
@@ -350,46 +392,32 @@ namespace Qwen {
}
};
struct QwenImageParams {
int patch_size = 2;
int64_t in_channels = 64;
int64_t out_channels = 16;
int num_layers = 60;
int64_t attention_head_dim = 128;
int64_t num_attention_heads = 24;
int64_t joint_attention_dim = 3584;
int theta = 10000;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
bool zero_cond_t = false;
};
class QwenImageModel : public GGMLBlock {
protected:
QwenImageParams params;
QwenImageConfig config;
public:
QwenImageModel() {}
QwenImageModel(QwenImageParams params)
: params(params) {
int64_t inner_dim = params.num_attention_heads * params.attention_head_dim;
QwenImageModel(QwenImageConfig config)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(params.joint_attention_dim, 1e-6f));
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(params.in_channels, inner_dim));
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(params.joint_attention_dim, inner_dim));
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
// blocks
for (int i = 0; i < params.num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto block = std::shared_ptr<GGMLBlock>(new QwenImageTransformerBlock(inner_dim,
params.num_attention_heads,
params.attention_head_dim,
config.num_attention_heads,
config.attention_head_dim,
1e-6f,
params.zero_cond_t));
config.zero_cond_t));
blocks["transformer_blocks." + std::to_string(i)] = block;
}
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new AdaLayerNormContinuous(inner_dim, inner_dim, false, 1e-6f));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, config.patch_size * config.patch_size * config.out_channels));
}
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
@@ -406,7 +434,7 @@ namespace Qwen {
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto t_emb = time_text_embed->forward(ctx, timestep);
if (params.zero_cond_t) {
if (config.zero_cond_t) {
auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
}
@@ -417,7 +445,7 @@ namespace Qwen {
sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt");
// sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb");
for (int i = 0; i < params.num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
@@ -427,7 +455,7 @@ namespace Qwen {
sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt");
}
if (params.zero_cond_t) {
if (config.zero_cond_t) {
t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
}
@@ -456,12 +484,12 @@ namespace Qwen {
int64_t C = x->ne[2];
int64_t N = x->ne[3];
auto img = DiT::pad_and_patchify(ctx, x, params.patch_size, params.patch_size);
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
int64_t img_tokens = img->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
ref = DiT::pad_and_patchify(ctx, ref, params.patch_size, params.patch_size);
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
}
}
@@ -474,7 +502,7 @@ namespace Qwen {
out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [N, h*w, C * patch_size * patch_size]
}
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, params.patch_size, params.patch_size); // [N, C, H, W]
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W]
return out;
}
@@ -482,7 +510,7 @@ namespace Qwen {
struct QwenImageRunner : public DiffusionModelRunner {
public:
QwenImageParams qwen_image_params;
QwenImageConfig config;
QwenImageModel qwen_image;
std::vector<float> pe_vec;
std::vector<float> modulate_index_vec;
@@ -494,34 +522,10 @@ namespace Qwen {
const std::string prefix = "",
SDVersion version = VERSION_QWEN_IMAGE,
bool zero_cond_t = false)
: DiffusionModelRunner(backend, params_backend, prefix) {
qwen_image_params.num_layers = 0;
qwen_image_params.zero_cond_t = zero_cond_t;
for (auto pair : tensor_storage_map) {
std::string tensor_name = pair.first;
if (tensor_name.find(prefix) == std::string::npos)
continue;
if (tensor_name.find("__index_timestep_zero__") != std::string::npos) {
qwen_image_params.zero_cond_t = true;
}
size_t pos = tensor_name.find("transformer_blocks.");
if (pos != std::string::npos) {
tensor_name = tensor_name.substr(pos); // remove prefix
auto items = split_string(tensor_name, '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > qwen_image_params.num_layers) {
qwen_image_params.num_layers = block_index + 1;
}
}
continue;
}
}
LOG_INFO("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers);
if (qwen_image_params.zero_cond_t) {
LOG_INFO("use zero_cond_t");
}
qwen_image = QwenImageModel(qwen_image_params);
: DiffusionModelRunner(backend, params_backend, prefix),
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
config.zero_cond_t = config.zero_cond_t || zero_cond_t;
qwen_image = QwenImageModel(config);
qwen_image.init(params_ctx, tensor_storage_map, prefix);
}
@@ -552,36 +556,36 @@ namespace Qwen {
pe_vec = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
qwen_image_params.patch_size,
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ref_latents,
increase_ref_index,
qwen_image_params.theta,
config.theta,
circular_y_enabled,
circular_x_enabled,
qwen_image_params.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / qwen_image_params.axes_dim_sum / 2);
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, qwen_image_params.axes_dim_sum / 2, pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe, true, "pe");
// pe->data = nullptr;
set_backend_tensor_data(pe, pe_vec.data());
ggml_tensor* modulate_index = nullptr;
if (qwen_image_params.zero_cond_t) {
if (config.zero_cond_t) {
modulate_index_vec.clear();
int64_t h_len = ((x->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
int64_t w_len = ((x->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
int64_t h_len = ((x->ne[1] + (config.patch_size / 2)) / config.patch_size);
int64_t w_len = ((x->ne[0] + (config.patch_size / 2)) / config.patch_size);
int64_t num_img_tokens = h_len * w_len;
modulate_index_vec.insert(modulate_index_vec.end(), num_img_tokens, 0.f);
int64_t num_ref_img_tokens = 0;
for (ggml_tensor* ref : ref_latents) {
int64_t h_len = ((ref->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
int64_t w_len = ((ref->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
int64_t h_len = ((ref->ne[1] + (config.patch_size / 2)) / config.patch_size);
int64_t w_len = ((ref->ne[0] + (config.patch_size / 2)) / config.patch_size);
num_ref_img_tokens += h_len * w_len;
}
@@ -683,7 +687,7 @@ namespace Qwen {
// cuda q8: pass
// cuda q8 fa: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_Q8_0;
ModelLoader model_loader;
otherarch/sdcpp/rope.hpp
+87
View File
@@ -249,6 +249,93 @@ namespace Rope {
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
}
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
int bs,
float theta,
int head_dim,
const std::vector<int>& mrope_section) {
GGML_ASSERT(bs > 0);
GGML_ASSERT(head_dim % 2 == 0);
GGML_ASSERT(mrope_section.size() >= 3);
std::vector<std::vector<float>> trans_ids = transpose(ids);
size_t pos_len = ids.size() / bs;
int half_dim = head_dim / 2;
std::vector<std::vector<std::vector<float>>> axis_embs;
axis_embs.reserve(3);
for (int axis = 0; axis < 3; ++axis) {
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta));
}
std::vector<std::vector<float>> emb = axis_embs[0];
for (int axis = 1; axis < 3; ++axis) {
int length = std::min<int>(mrope_section[axis] * 3, half_dim);
for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
for (int k = 0; k < 4; ++k) {
emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
}
}
}
}
return flatten(emb);
}
__STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height,
int width,
int dim,
float theta = 10000.f,
float scale = 16.f,
int ref_grid_h = 0,
int ref_grid_w = 0) {
assert(dim % 4 == 0);
int half_dim = dim / 2;
int dim_axis = dim / 2;
int axis_half_dim = dim_axis / 2;
float h_ntk = 1.f;
float w_ntk = 1.f;
if (ref_grid_h > 0 && ref_grid_w > 0 && dim_axis > 2) {
float power = static_cast<float>(dim_axis) / static_cast<float>(dim_axis - 2);
h_ntk = std::pow(static_cast<float>(height) / static_cast<float>(ref_grid_h), power);
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
}
std::vector<float> x_pos;
std::vector<float> y_pos;
x_pos.reserve(static_cast<size_t>(height) * width);
y_pos.reserve(static_cast<size_t>(height) * width);
for (int iy = 0; iy < height; ++iy) {
float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
for (int ix = 0; ix < width; ++ix) {
float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
x_pos.push_back(x);
y_pos.push_back(y);
}
}
auto x_emb = rope(x_pos, dim_axis, theta * w_ntk);
auto y_emb = rope(y_pos, dim_axis, theta * h_ntk);
std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
for (int pos = 0; pos < height * width; ++pos) {
for (int i = 0; i < axis_half_dim; ++i) {
int jx = 2 * i;
int jy = 2 * i + 1;
size_t base_x = static_cast<size_t>(pos) * half_dim * 4 + static_cast<size_t>(jx) * 4;
size_t base_y = static_cast<size_t>(pos) * half_dim * 4 + static_cast<size_t>(jy) * 4;
size_t axis = static_cast<size_t>(i) * 4;
for (int k = 0; k < 4; ++k) {
out[base_x + k] = x_emb[pos][axis + k];
out[base_y + k] = y_emb[pos][axis + k];
}
}
}
return out;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
int bs,
int axes_dim_num,
otherarch/sdcpp/stable-diffusion.cpp
+335 -93
View File
@@ -1,3 +1,8 @@
#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include "ggml_extend.hpp"
#include "ggml_graph_cut.h"
@@ -19,6 +24,7 @@
#include "flux.hpp"
#include "guidance.h"
#include "hidream_o1.hpp"
#include "ideogram4.hpp"
#include "lens.hpp"
#include "lora.hpp"
#include "ltx_audio_vae.h"
@@ -26,6 +32,7 @@
#include "ltx_vae.hpp"
#include "ltxv.hpp"
#include "mmdit.hpp"
#include "pid.hpp"
#include "pmid.hpp"
#include "qwen_image.hpp"
#include "sample-cache.h"
@@ -34,12 +41,14 @@
#include "upscaler.h"
#include "vae.hpp"
#include "wan.hpp"
#include "wan_vae.hpp"
#include "z_image.hpp"
#include "latent-preview.h"
#include "name_conversion.h"
#include <filesystem>
const char* sd_vae_format_name(enum sd_vae_format_t format);
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);
const char* model_version_to_str[] = {
"SD 1.x",
@@ -77,6 +86,8 @@ const char* model_version_to_str[] = {
"Ernie Image",
"Lens",
"Longcat-Image",
"PiD",
"Ideogram 4",
};
const char* sampling_methods_str[] = {
@@ -102,6 +113,19 @@ const char* sampling_methods_str[] = {
/*================================================== Helper Functions ================================================*/
static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
return version == VERSION_FLUX ||
sd_version_is_flux2(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_longcat(version) ||
sd_version_is_z_image(version);
}
static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) {
return sd_version_is_inpaint_or_unet_edit(version) ||
(has_ref_images && sd_version_supports_ref_latent_img_cfg(version));
}
void calculate_alphas_cumprod(float* alphas_cumprod,
float linear_start = 0.00085f,
float linear_end = 0.0120f,
@@ -171,6 +195,7 @@ public:
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
bool offload_params_to_cpu = false;
float max_vram = 0.f;
bool stream_layers = false;
bool use_pmid = false;
std::string backend_spec;
std::string params_backend_spec;
@@ -227,7 +252,7 @@ public:
std::string error;
if (!backend_manager.init(sd_ctx_params->backend,
sd_ctx_params->params_backend,
sd_ctx_params->offload_params_to_cpu,
offload_params_to_cpu,
sd_ctx_params->keep_clip_on_cpu,
sd_ctx_params->keep_vae_on_cpu,
sd_ctx_params->keep_control_net_on_cpu,
@@ -257,8 +282,18 @@ public:
free_params_immediately = sd_ctx_params->free_params_immediately;
offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu;
max_vram = sd_ctx_params->max_vram;
stream_layers = sd_ctx_params->stream_layers;
backend_spec = SAFE_STR(sd_ctx_params->backend);
params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
if (stream_layers && max_vram == 0.f) {
LOG_WARN("--stream-layers has no effect without --max-vram set; ignoring");
stream_layers = false;
}
if (stream_layers && !offload_params_to_cpu && params_backend_spec.empty()) {
// Streaming needs CPU-resident params.
LOG_WARN("--stream-layers has no effect without --offload-to-cpu (or --params-backend); ignoring");
stream_layers = false;
}
bool use_tae = false;
bool use_audio_vae = false;
@@ -310,6 +345,13 @@ public:
}
}
if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) {
LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path);
if (!model_loader.init_from_file(sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) {
LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path);
}
}
bool is_unet = sd_version_is_unet(model_loader.get_sd_version());
// begin kcpp replacements
@@ -658,7 +700,10 @@ public:
}
}
}
if (have_quantized_weight) {
// Avoid full-model LoRA merge buffers on constrained setups.
const bool streaming_constrained = stream_layers ||
sd_ctx_params->offload_params_to_cpu;
if (have_quantized_weight || streaming_constrained) {
apply_lora_immediately = false;
} else {
apply_lora_immediately = true;
@@ -740,6 +785,27 @@ public:
params_backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model");
} else if (sd_version_is_pid(version)) {
vae_decode_only = false;
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
params_backend_for(SDBackendModule::TE),
tensor_storage_map,
version);
diffusion_model = std::make_shared<Pid::PiDRunner>(backend_for(SDBackendModule::DIFFUSION),
params_backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model.net");
} else if (sd_version_is_ideogram4(version)) {
cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
params_backend_for(SDBackendModule::TE),
tensor_storage_map,
version,
"",
false);
diffusion_model = std::make_shared<Ideogram4::Ideogram4Runner>(backend_for(SDBackendModule::DIFFUSION),
params_backend_for(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model");
} else if (sd_version_is_flux(version)) {
bool is_chroma = false;
for (auto pair : tensor_storage_map) {
@@ -944,6 +1010,7 @@ public:
get_param_tensors(cond_stage_model, module_can_mmap(SDBackendModule::TE));
diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
diffusion_model->set_stream_layers_enabled(stream_layers);
get_param_tensors(diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION));
if (sd_version_is_unet_edit(version)) {
@@ -952,6 +1019,7 @@ public:
if (high_noise_diffusion_model) {
high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
high_noise_diffusion_model->set_stream_layers_enabled(stream_layers);
get_param_tensors(high_noise_diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION));
}
@@ -982,6 +1050,16 @@ public:
}
};
sd_vae_format_t vae_format = sd_ctx_params->vae_format;
if (vae_format < SD_VAE_FORMAT_AUTO || vae_format >= SD_VAE_FORMAT_COUNT) {
LOG_WARN("invalid VAE format override, using auto");
vae_format = SD_VAE_FORMAT_AUTO;
}
SDVersion vae_version = version;
if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) {
vae_version = sd_vae_format_to_version(vae_format, vae_version);
}
auto create_vae = [&]() -> std::shared_ptr<VAE> {
if (sd_version_is_ltxav(version)) {
return std::make_shared<LTXVideoVAE>(backend_for(SDBackendModule::VAE),
@@ -1006,7 +1084,7 @@ public:
"first_stage_model",
vae_decode_only,
false,
version);
vae_version);
if (sd_version_is_sdxl(version) &&
(strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || external_vae_is_invalid)) {
float vae_conv_2d_scale = 1.f / 32.f;
@@ -1207,6 +1285,12 @@ public:
ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.gate_up_proj.weight_scale_2");
ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.down_proj.weight_scale_2");
}
if (sd_version_is_ideogram4(version)) {
ignore_tensors.insert("text_encoders.llm.lm_head.");
ignore_tensors.insert("text_encoders.llm.visual.");
ignore_tensors.insert("text_encoders.llm.vision_model.");
ignore_tensors.insert("text_encoders.llm.tokenizer_json");
}
if (version == VERSION_HIDREAM_O1) {
ignore_tensors.insert("lm_head.");
ignore_tensors.insert("model.visual.deepstack_merger_list.");
@@ -1307,7 +1391,7 @@ public:
if (module_backend == nullptr) {
return false;
}
if (ggml_backend_is_cpu(module_backend)) {
if (sd_backend_is_cpu(module_backend)) {
total_params_ram_size += size;
} else {
total_params_vram_size += size;
@@ -1322,7 +1406,7 @@ public:
if (module_backend == nullptr) {
return "N/A";
}
return ggml_backend_is_cpu(module_backend) ? "RAM" : "VRAM";
return sd_backend_is_cpu(module_backend) ? "RAM" : "VRAM";
};
if (!add_params_memory(clip_params_mem_size, SDBackendModule::TE) ||
@@ -1381,12 +1465,18 @@ public:
version == VERSION_HIDREAM_O1 ||
sd_version_is_anima(version) ||
sd_version_is_ernie_image(version) ||
sd_version_is_z_image(version)) {
sd_version_is_z_image(version) ||
sd_version_is_pid(version) ||
sd_version_is_ideogram4(version)) {
pred_type = FLOW_PRED;
if (sd_version_is_wan(version)) {
default_flow_shift = 5.f;
} else if (sd_version_is_ernie_image(version)) {
default_flow_shift = 4.f;
} else if (sd_version_is_pid(version)) {
default_flow_shift = 1.5f;
} else if (sd_version_is_ideogram4(version)) {
default_flow_shift = 1.0f;
} else {
default_flow_shift = 3.f;
}
@@ -1879,12 +1969,15 @@ public:
const sd::Tensor<float>& init_latent,
const sd::Tensor<float>& denoise_mask) {
if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") {
auto new_timesteps = std::vector<float>(static_cast<size_t>(init_latent.shape()[2]), timesteps[0]);
int64_t frame_count = init_latent.shape()[2];
auto new_timesteps = std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
if (!denoise_mask.empty()) {
float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, 0, 0, 0) : denoise_mask.index(0, 0, 0, 0);
if (value == 0.f) {
new_timesteps[0] = 0.f;
if (!denoise_mask.empty() && denoise_mask.dim() >= 4 && denoise_mask.shape()[2] == frame_count) {
for (int64_t frame = 0; frame < frame_count; ++frame) {
float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, frame, 0, 0) : denoise_mask.index(0, 0, frame, 0);
if (value == 0.f) {
new_timesteps[static_cast<size_t>(frame)] = 0.f;
}
}
}
return new_timesteps;
@@ -2065,7 +2158,7 @@ public:
if (version == VERSION_HIDREAM_O1) {
return std::vector<float>{1.0f - (t / static_cast<float>(TIMESTEPS))};
}
if (sd_version_is_z_image(version)) {
if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) {
return std::vector<float>{1000.f - t};
}
return std::vector<float>{t};
@@ -2144,7 +2237,7 @@ public:
sd::Tensor<float> noise,
const SDCondition& cond,
const SDCondition& uncond,
const SDCondition& img_cond,
const SDCondition& img_uncond,
const SDCondition& id_cond,
const sd::Tensor<float>& control_image,
float control_strength,
@@ -2169,6 +2262,7 @@ public:
float cfg_scale = guidance.txt_cfg;
float img_cfg_scale = guidance.img_cfg;
float slg_scale = guidance.slg.scale;
bool slg_uncond = sd::guidance::parse_skip_layer_guidance_uncond_arg(extra_sample_args);
sd_sample::SampleCacheRuntime cache_runtime = sd_sample::init_sample_cache_runtime(version,
cache_params,
@@ -2185,12 +2279,21 @@ public:
}
size_t steps = sigmas.size() - 1;
bool has_skiplayer = slg_scale != 0.0f && !skip_layers.empty();
bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
if (has_skiplayer && !sd_version_is_dit(version)) {
has_skiplayer = false;
LOG_WARN("SLG is incompatible with this model type");
}
sd::guidance::AdaptiveProjectedGuidanceParams apg_params = sd::guidance::parse_adaptive_projected_guidance_args(extra_sample_args);
bool use_apg_guidance = sd::guidance::is_adaptive_projected_guidance_enabled(apg_params);
if (use_apg_guidance) {
LOG_INFO("using Adaptive Projected Guidance (APG)");
}
sd::guidance::ClassifierFreeGuidance classifier_free_guidance(cfg_scale, img_cfg_scale);
sd::guidance::AdaptiveProjectedGuidance adaptive_projected_guidance(cfg_scale, img_cfg_scale, apg_params);
const sd::guidance::BaseGuidance& primary_guidance = use_apg_guidance
? static_cast<const sd::guidance::BaseGuidance&>(adaptive_projected_guidance)
: static_cast<const sd::guidance::BaseGuidance&>(classifier_free_guidance);
sd::guidance::SkipLayerGuidance skip_layer_guidance(has_skiplayer ? skip_layers : std::vector<int>(),
has_skiplayer ? slg_scale : 0.0f,
guidance.slg.layer_start,
@@ -2259,13 +2362,17 @@ public:
sd::Tensor<float> cond_out;
sd::Tensor<float> uncond_out;
sd::Tensor<float> img_cond_out;
sd::Tensor<float> img_uncond_out;
sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
std::vector<sd::Tensor<float>> controls;
DiffusionParams diffusion_params;
diffusion_params.x = &noised_input;
diffusion_params.timesteps = &timesteps_tensor;
diffusion_params.increase_ref_index = increase_ref_index;
sd::guidance::GuidanceInput step_guidance_input;
step_guidance_input.step = step;
step_guidance_input.schedule_size = sigmas.size();
bool is_skiplayer_step = skip_layer_guidance.is_enabled_for_step(step_guidance_input);
compute_sample_controls(control_image,
noised_input,
@@ -2273,13 +2380,19 @@ public:
cond,
&controls);
static const std::vector<sd::Tensor<float>> empty_ref_latents;
bool uncond_without_ref_latents = !img_uncond.empty() &&
!ref_latents.empty() &&
sd_version_supports_ref_latent_img_cfg(version);
auto run_condition = [&](const SDCondition& condition,
const sd::Tensor<float>* c_concat_override = nullptr,
const std::vector<int>* local_skip_layers = nullptr) -> sd::Tensor<float> {
const sd::Tensor<float>* c_concat_override = nullptr,
const std::vector<int>* local_skip_layers = nullptr,
const std::vector<sd::Tensor<float>>* ref_latents_override = nullptr) -> sd::Tensor<float> {
diffusion_params.context = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn;
diffusion_params.c_concat = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat);
diffusion_params.y = condition.c_vector.empty() ? nullptr : &condition.c_vector;
diffusion_params.ref_latents = condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images;
diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images);
if (sd_version_is_unet(version)) {
diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength};
@@ -2349,31 +2462,40 @@ public:
uncond,
&controls);
}
uncond_out = run_condition(uncond);
const std::vector<int>* uncond_skip_layers = nullptr;
if (is_skiplayer_step && slg_uncond) {
LOG_DEBUG("Skipping layers at uncond step %d\n", step);
uncond_skip_layers = &skip_layer_guidance.layers();
}
uncond_out = run_condition(uncond,
uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
uncond_skip_layers);
if (uncond_out.empty()) {
return {};
}
}
if (!img_cond.empty()) {
img_cond_out = run_condition(img_cond,
cond.c_concat.empty() ? nullptr : &cond.c_concat);
if (img_cond_out.empty()) {
if (!img_uncond.empty()) {
img_uncond_out = run_condition(img_uncond,
img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
nullptr,
uncond_without_ref_latents ? &empty_ref_latents : nullptr);
if (img_uncond_out.empty()) {
return {};
}
}
sd::guidance::GuidanceInput guidance_input;
guidance_input.step = step;
guidance_input.schedule_size = sigmas.size();
guidance_input.pred_cond = &cond_out;
guidance_input.pred_uncond = uncond_out.empty() ? nullptr : &uncond_out;
guidance_input.pred_img_cond = img_cond_out.empty() ? nullptr : &img_cond_out;
guidance_input.step = step;
guidance_input.schedule_size = sigmas.size();
guidance_input.pred_cond = &cond_out;
guidance_input.pred_uncond = uncond_out.empty() ? nullptr : &uncond_out;
guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out;
sd::guidance::GuiderOutput guided = classifier_free_guidance.forward(guidance_input, {});
sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {});
if (guided.pred.empty()) {
return {};
}
if (skip_layer_guidance.is_enabled_for_step(guidance_input)) {
if (is_skiplayer_step && slg_scale != 0.0f) {
LOG_DEBUG("Skipping layers at step %d\n", step);
if (!step_cache.is_step_skipped()) {
guidance_input.predict_skip_layer = [&]() -> sd::Tensor<float> {
@@ -2393,7 +2515,9 @@ public:
sd::guidance::GuiderOutput output;
output.pred = denoised;
if (needs_uncond_denoised) {
const sd::Tensor<float>& base_uncond = !uncond_out.empty() ? uncond_out : cond_out;
const sd::Tensor<float>& base_uncond = !img_uncond_out.empty()
? img_uncond_out
: (!uncond_out.empty() ? uncond_out : cond_out);
output.pred_uncond = base_uncond * c_out + x * c_skip;
}
if (cache_runtime.spectrum_enabled) {
@@ -2440,6 +2564,9 @@ public:
}
int get_vae_scale_factor() {
if (sd_version_is_pid(version)) {
return 1;
}
return first_stage_model->get_scale_factor();
}
@@ -2466,6 +2593,8 @@ public:
latent_channel = 3;
} else if (version == VERSION_CHROMA_RADIANCE) {
latent_channel = 3;
} else if (sd_version_is_pid(version)) {
latent_channel = 3;
} else if (sd_version_uses_flux2_vae(version)) {
latent_channel = 128;
} else {
@@ -2543,6 +2672,18 @@ public:
}
sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false) {
if (sd_version_is_pid(version)) {
return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
}
// Free resident diffusion params before VAE allocates its compute buffer.
if (stream_layers) {
if (diffusion_model) {
diffusion_model->release_streaming_residency();
}
if (high_noise_diffusion_model) {
high_noise_diffusion_model->release_streaming_residency();
}
}
auto latents = first_stage_model->diffusion_to_vae_latents(x);
first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
@@ -2820,6 +2961,35 @@ enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) {
return SD_HIRES_UPSCALER_COUNT;
}
const char* sd_vae_format_name(enum sd_vae_format_t format) {
switch (format) {
case SD_VAE_FORMAT_AUTO:
return "auto";
case SD_VAE_FORMAT_FLUX:
return "flux";
case SD_VAE_FORMAT_SD3:
return "sd3";
case SD_VAE_FORMAT_FLUX2:
return "flux2";
default:
return NONE_STR;
}
}
static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback) {
switch (format) {
case SD_VAE_FORMAT_FLUX:
return VERSION_FLUX;
case SD_VAE_FORMAT_SD3:
return VERSION_SD3;
case SD_VAE_FORMAT_FLUX2:
return VERSION_FLUX2;
case SD_VAE_FORMAT_AUTO:
default:
return fallback;
}
}
void sd_cache_params_init(sd_cache_params_t* cache_params) {
*cache_params = {};
cache_params->mode = SD_CACHE_DISABLED;
@@ -2875,6 +3045,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO;
sd_ctx_params->offload_params_to_cpu = false;
sd_ctx_params->max_vram = 0.f;
sd_ctx_params->stream_layers = false;
sd_ctx_params->enable_mmap = false;
sd_ctx_params->keep_clip_on_cpu = false;
sd_ctx_params->keep_control_net_on_cpu = false;
@@ -2885,6 +3056,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
sd_ctx_params->chroma_use_dit_mask = true;
sd_ctx_params->chroma_use_t5_mask = false;
sd_ctx_params->chroma_t5_mask_pad = 1;
sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO;
sd_ctx_params->backend = nullptr;
sd_ctx_params->params_backend = nullptr;
}
@@ -2905,6 +3077,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"llm_vision_path: %s\n"
"diffusion_model_path: %s\n"
"high_noise_diffusion_model_path: %s\n"
"uncond_diffusion_model_path: %s\n"
"embeddings_connectors_path: %s\n"
"vae_path: %s\n"
"audio_vae_path: %s\n"
@@ -2921,6 +3094,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"prediction: %s\n"
"offload_params_to_cpu: %s\n"
"max_vram: %.3f\n"
"stream_layers: %s\n"
"backend: %s\n"
"params_backend: %s\n"
"keep_clip_on_cpu: %s\n"
@@ -2932,7 +3106,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"circular_y: %s\n"
"chroma_use_dit_mask: %s\n"
"chroma_use_t5_mask: %s\n"
"chroma_t5_mask_pad: %d\n",
"chroma_t5_mask_pad: %d\n"
"vae_format: %s\n",
SAFE_STR(sd_ctx_params->model_path),
SAFE_STR(sd_ctx_params->clip_l_path),
SAFE_STR(sd_ctx_params->clip_g_path),
@@ -2942,6 +3117,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
SAFE_STR(sd_ctx_params->llm_vision_path),
SAFE_STR(sd_ctx_params->diffusion_model_path),
SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path),
SAFE_STR(sd_ctx_params->uncond_diffusion_model_path),
SAFE_STR(sd_ctx_params->embeddings_connectors_path),
SAFE_STR(sd_ctx_params->vae_path),
SAFE_STR(sd_ctx_params->audio_vae_path),
@@ -2958,6 +3134,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
sd_prediction_name(sd_ctx_params->prediction),
BOOL_STR(sd_ctx_params->offload_params_to_cpu),
sd_ctx_params->max_vram,
BOOL_STR(sd_ctx_params->stream_layers),
SAFE_STR(sd_ctx_params->backend),
SAFE_STR(sd_ctx_params->params_backend),
BOOL_STR(sd_ctx_params->keep_clip_on_cpu),
@@ -2969,7 +3146,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
BOOL_STR(sd_ctx_params->circular_y),
BOOL_STR(sd_ctx_params->chroma_use_dit_mask),
BOOL_STR(sd_ctx_params->chroma_use_t5_mask),
sd_ctx_params->chroma_t5_mask_pad);
sd_ctx_params->chroma_t5_mask_pad,
sd_vae_format_name(sd_ctx_params->vae_format));
return buf;
}
@@ -3246,6 +3424,9 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
if (sd_ctx != nullptr && sd_ctx->sd != nullptr) {
if (sd_version_is_pid(sd_ctx->sd->version)) {
return LCM_SAMPLE_METHOD;
}
if (sd_version_is_dit(sd_ctx->sd->version)) {
return EULER_SAMPLE_METHOD;
}
@@ -3329,9 +3510,10 @@ struct GenerationRequest {
int diffusion_model_down_factor = -1;
int64_t seed = -1;
bool use_uncond = false;
bool use_img_cond = false;
bool use_img_uncond = false;
bool use_high_noise_uncond = false;
bool use_high_noise_img_cond = false;
bool use_high_noise_img_uncond = false;
bool has_ref_images = false;
const sd_cache_params_t* cache_params = nullptr;
int batch_count = 1;
int shifted_timestep = 0;
@@ -3365,6 +3547,7 @@ struct GenerationRequest {
eta = sd_img_gen_params->sample_params.eta;
increase_ref_index = sd_img_gen_params->increase_ref_index;
auto_resize_ref_image = sd_img_gen_params->auto_resize_ref_image;
has_ref_images = sd_img_gen_params->ref_images_count > 0;
guidance = sd_img_gen_params->sample_params.guidance;
pm_params = sd_img_gen_params->pm_params;
hires = sd_img_gen_params->hires;
@@ -3487,28 +3670,36 @@ struct GenerationRequest {
static void resolve_guidance(sd_ctx_t* sd_ctx,
sd_guidance_params_t* guidance,
bool* use_uncond,
bool* use_img_cond,
bool* use_img_uncond,
bool has_ref_images,
const char* stage_name = nullptr) {
GGML_ASSERT(guidance != nullptr);
GGML_ASSERT(use_uncond != nullptr);
GGML_ASSERT(use_img_cond != nullptr);
// out_uncond + text_cfg_scale * (out_cond - out_img_cond) + image_cfg_scale * (out_img_cond - out_uncond)
// img_cfg == txt_cfg means that img_cfg is not used
if (!std::isfinite(guidance->img_cfg)) {
guidance->img_cfg = guidance->txt_cfg;
GGML_ASSERT(use_img_uncond != nullptr);
// out_img_uncond + text_cfg_scale * (out_cond - out_uncond) + image_cfg_scale * (out_uncond - out_img_uncond)
// -> text_cfg_scale * out_cond + (image_cfg_scale - text_cfg_scale) * out_uncond + (1 - image_cfg_scale) * out_img_uncond
// out_cond : prompt, image latent
// out_uncond : negative prompt, image latent
// out_img_uncond : negative prompt, zero image latent
// image_cfg_scale == 1 reduces 3-cond CFG to 2-cond CFG.
bool img_cfg_was_set = std::isfinite(guidance->img_cfg);
if (!img_cfg_was_set) {
guidance->img_cfg = 1.f;
}
if (!sd_version_is_inpaint_or_unet_edit(sd_ctx->sd->version)) {
guidance->img_cfg = guidance->txt_cfg;
}
if (guidance->txt_cfg != 1.f) {
*use_uncond = true;
if (!sd_version_supports_img_cfg(sd_ctx->sd->version, has_ref_images)) {
if (img_cfg_was_set && guidance->img_cfg != 1.f) {
LOG_WARN("3-conditioning CFG is not supported with this model, disabling it for better performance");
}
guidance->img_cfg = 1.f;
}
if (guidance->img_cfg != guidance->txt_cfg) {
*use_img_cond = true;
*use_uncond = true;
*use_uncond = true;
}
if (guidance->img_cfg != 1.f) {
*use_img_uncond = true;
}
if (guidance->txt_cfg < 1.f) {
@@ -3527,12 +3718,13 @@ struct GenerationRequest {
resolve_hires();
seed = resolve_seed(seed);
resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_cond);
resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_uncond, has_ref_images);
if (sd_ctx->sd->high_noise_diffusion_model) {
resolve_guidance(sd_ctx,
&high_noise_guidance,
&use_high_noise_uncond,
&use_high_noise_img_cond,
&use_high_noise_img_uncond,
has_ref_images,
"high noise: ");
}
@@ -3650,7 +3842,7 @@ struct SamplePlan {
struct ImageGenerationLatents {
sd::Tensor<float> init_latent;
sd::Tensor<float> concat_latent;
sd::Tensor<float> uncond_concat_latent;
sd::Tensor<float> img_uncond_concat_latent;
sd::Tensor<float> audio_latent;
sd::Tensor<float> video_positions;
sd::Tensor<float> control_image;
@@ -3973,7 +4165,7 @@ static int get_ltxav_num_audio_latents(int frames, int fps) {
struct ImageGenerationEmbeds {
SDCondition cond;
SDCondition uncond;
SDCondition img_cond;
SDCondition img_uncond;
SDCondition id_cond;
};
@@ -4132,6 +4324,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
LOG_WARN("This model needs at least one reference image; using an empty reference");
ref_images.push_back(sd::zeros<float>({request->width, request->height, 3, 1}));
request->guidance.img_cfg = request->guidance.txt_cfg;
request->use_img_uncond = false;
}
if (!ref_images.empty()) {
@@ -4144,7 +4337,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
continue;
}
sd::Tensor<float> ref_latent;
if (request->auto_resize_ref_image) {
if (request->auto_resize_ref_image && !sd_version_is_pid(sd_ctx->sd->version)) {
LOG_DEBUG("auto resize ref images");
int vae_image_size = std::min(1024 * 1024, request->width * request->height);
double vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
@@ -4176,8 +4369,15 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
ref_latents.push_back(std::move(ref_latent));
}
if (sd_version_is_pid(sd_ctx->sd->version)) {
if (ref_latents.empty()) {
LOG_ERROR("PiD requires a reference image");
return std::nullopt;
}
}
sd::Tensor<float> concat_latent;
sd::Tensor<float> uncond_concat_latent;
sd::Tensor<float> img_uncond_concat_latent;
if (sd_version_is_inpaint(sd_ctx->sd->version)) {
sd::Tensor<float> masked_init_latent;
@@ -4205,8 +4405,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
request->height / request->vae_scale_factor});
mask = mask.permute({1, 3, 0, 2}).reshape({request->width / request->vae_scale_factor, request->height / request->vae_scale_factor, request->vae_scale_factor * request->vae_scale_factor, 1});
concat_latent = sd::ops::concat(masked_init_latent, mask, 2);
uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, mask, 2);
concat_latent = sd::ops::concat(masked_init_latent, mask, 2);
img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, mask, 2);
} else if (sd_ctx->sd->version == VERSION_FLEX_2) {
concat_latent = sd::ops::concat(masked_init_latent, latent_mask, 2);
if (!control_latent.empty()) {
@@ -4215,16 +4415,16 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
concat_latent = sd::ops::concat(concat_latent, sd::Tensor<float>::zeros_like(masked_init_latent), 2);
}
uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, latent_mask, 2);
uncond_concat_latent = sd::ops::concat(uncond_concat_latent, sd::Tensor<float>::zeros_like(masked_init_latent), 2);
img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, latent_mask, 2);
img_uncond_concat_latent = sd::ops::concat(img_uncond_concat_latent, sd::Tensor<float>::zeros_like(masked_init_latent), 2);
} else { // SD1.x SD2.x SDXL inpaint
concat_latent = sd::ops::concat(latent_mask, masked_init_latent, 2);
uncond_concat_latent = sd::ops::concat(latent_mask, uncond_masked_init_latent, 2);
concat_latent = sd::ops::concat(latent_mask, masked_init_latent, 2);
img_uncond_concat_latent = sd::ops::concat(latent_mask, uncond_masked_init_latent, 2);
}
}
if (sd_version_is_unet_edit(sd_ctx->sd->version)) {
concat_latent = sd::ops::interpolate<float>(ref_latents[0], init_latent.shape());
uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
concat_latent = sd::ops::interpolate<float>(ref_latents[0], init_latent.shape());
img_uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
}
if (sd_ctx->sd->version == VERSION_FLUX_CONTROLS) {
if (!control_latent.empty()) {
@@ -4232,7 +4432,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
} else {
concat_latent = sd::Tensor<float>::zeros_like(init_latent);
}
uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
img_uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
}
if (sd_img_gen_params->init_image.data != nullptr || sd_img_gen_params->ref_images_count > 0) {
@@ -4241,12 +4441,12 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
}
ImageGenerationLatents latents;
latents.init_latent = std::move(init_latent);
latents.concat_latent = std::move(concat_latent);
latents.uncond_concat_latent = std::move(uncond_concat_latent);
latents.control_image = std::move(control_image_tensor);
latents.ref_images = std::move(ref_images);
latents.ref_latents = std::move(ref_latents);
latents.init_latent = std::move(init_latent);
latents.concat_latent = std::move(concat_latent);
latents.img_uncond_concat_latent = std::move(img_uncond_concat_latent);
latents.control_image = std::move(control_image_tensor);
latents.ref_images = std::move(ref_images);
latents.ref_latents = std::move(ref_latents);
if (sd_version_is_inpaint(sd_ctx->sd->version)) {
latent_mask = sd::ops::max_pool_2d(latent_mask,
@@ -4280,20 +4480,53 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
cond.c_concat = latents->concat_latent; // TODO: optimize
}
bool use_ref_latent_img_cfg = request->use_img_uncond &&
!latents->ref_images.empty() &&
sd_version_supports_ref_latent_img_cfg(sd_ctx->sd->version);
SDCondition uncond;
if (request->use_uncond || request->use_high_noise_uncond) {
bool zero_out_masked = false;
if (sd_version_is_sdxl(sd_ctx->sd->version) &&
request->negative_prompt.empty() &&
!sd_ctx->sd->is_using_edm_v_parameterization) {
zero_out_masked = true;
if (sd_version_is_ideogram4(sd_ctx->sd->version)) {
uncond.c_vector = sd::Tensor<float>::from_vector({1.0f});
} else {
bool zero_out_masked = false;
if (sd_version_is_sdxl(sd_ctx->sd->version) &&
request->negative_prompt.empty() &&
!sd_ctx->sd->is_using_edm_v_parameterization) {
zero_out_masked = true;
}
condition_params.text = request->negative_prompt;
condition_params.zero_out_masked = zero_out_masked;
uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
}
condition_params.text = request->negative_prompt;
condition_params.zero_out_masked = zero_out_masked;
uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
if (uncond.c_concat.empty()) {
uncond.c_concat = latents->uncond_concat_latent; // TODO: optimize
uncond.c_concat = latents->concat_latent; // TODO: optimize
}
}
SDCondition img_uncond;
if (request->use_img_uncond) {
if ((request->use_uncond || request->use_high_noise_uncond) && (latents->ref_images.empty() || !use_ref_latent_img_cfg)) {
img_uncond = SDCondition(uncond.c_crossattn, uncond.c_vector, latents->img_uncond_concat_latent);
} else {
bool zero_out_masked = false;
if (sd_version_is_sdxl(sd_ctx->sd->version) &&
request->negative_prompt.empty() &&
!sd_ctx->sd->is_using_edm_v_parameterization) {
zero_out_masked = true;
}
condition_params.text = request->negative_prompt;
condition_params.zero_out_masked = zero_out_masked;
if (use_ref_latent_img_cfg) {
std::vector<sd::Tensor<float>> empty_ref_images;
condition_params.ref_images = &empty_ref_images;
}
img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
condition_params);
if (img_uncond.c_concat.empty()) {
img_uncond.c_concat = latents->img_uncond_concat_latent; // TODO: optimize
}
}
}
@@ -4305,12 +4538,10 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
}
ImageGenerationEmbeds embeds;
if (request->use_img_cond) {
embeds.img_cond = SDCondition(uncond.c_crossattn, uncond.c_vector, cond.c_concat);
}
embeds.cond = std::move(cond);
embeds.uncond = std::move(uncond);
embeds.id_cond = std::move(id_cond);
embeds.img_uncond = std::move(img_uncond);
embeds.cond = std::move(cond);
embeds.uncond = std::move(uncond);
embeds.id_cond = std::move(id_cond);
return embeds;
}
@@ -4590,7 +4821,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
std::move(noise),
embeds.cond,
embeds.uncond,
embeds.img_cond,
embeds.img_uncond,
embeds.id_cond,
latents.control_image,
request.control_strength,
@@ -4710,7 +4941,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
std::move(noise),
embeds.cond,
embeds.uncond,
embeds.img_cond,
embeds.img_uncond,
embeds.id_cond,
latents.control_image,
request.control_strength,
@@ -4978,6 +5209,17 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
latents.denoise_mask = sd::full<float>({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f);
sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f);
if (!end_image.empty()) {
auto end_img = end_image.reshape({end_image.shape()[0], end_image.shape()[1], 1, end_image.shape()[2], 1});
auto end_image_latent = sd_ctx->sd->encode_first_stage(end_img); // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor]
if (end_image_latent.empty()) {
LOG_ERROR("failed to encode end video frame");
return std::nullopt;
}
sd::ops::slice_assign(&latents.init_latent, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], end_image_latent);
sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f);
}
int64_t t2 = ggml_time_ms();
LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
} else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" ||
@@ -5414,7 +5656,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
std::move(noise),
embeds.cond,
request.use_high_noise_uncond ? embeds.uncond : SDCondition(),
embeds.img_cond,
embeds.img_uncond,
embeds.id_cond,
sd::Tensor<float>(),
0.f,
@@ -5460,7 +5702,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
std::move(noise),
embeds.cond,
request.use_uncond ? embeds.uncond : SDCondition(),
embeds.img_cond,
embeds.img_uncond,
embeds.id_cond,
sd::Tensor<float>(),
0.f,
@@ -5604,7 +5846,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
std::move(noise),
embeds.cond,
hires_request.use_uncond ? embeds.uncond : SDCondition(),
embeds.img_cond,
embeds.img_uncond,
embeds.id_cond,
sd::Tensor<float>(),
0.f,
@@ -5716,7 +5958,7 @@ namespace kcpp_sd {
if (ctx != nullptr && ctx->sd != nullptr) {
auto maybe_flux = std::dynamic_pointer_cast<Flux::FluxRunner>(ctx->sd->diffusion_model);
if (maybe_flux != nullptr) {
return maybe_flux->flux_params.is_chroma;
return maybe_flux->config.is_chroma;
}
}
return false;
otherarch/sdcpp/stable-diffusion.h
+11
View File
@@ -168,6 +168,14 @@ typedef struct {
const char* path;
} sd_embedding_t;
enum sd_vae_format_t {
SD_VAE_FORMAT_AUTO = -1,
SD_VAE_FORMAT_FLUX,
SD_VAE_FORMAT_SD3,
SD_VAE_FORMAT_FLUX2,
SD_VAE_FORMAT_COUNT,
};
typedef struct {
const char* model_path;
const char* clip_l_path;
@@ -178,6 +186,7 @@ typedef struct {
const char* llm_vision_path;
const char* diffusion_model_path;
const char* high_noise_diffusion_model_path;
const char* uncond_diffusion_model_path;
const char* embeddings_connectors_path;
const char* vae_path;
const char* audio_vae_path;
@@ -212,7 +221,9 @@ typedef struct {
bool chroma_use_t5_mask;
int chroma_t5_mask_pad;
bool qwen_image_zero_cond_t;
enum sd_vae_format_t vae_format;
float max_vram; // GiB budget for graph-cut segmented param offload (0 = disabled, -1 = auto free VRAM minus 1 GiB)
bool stream_layers; // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
const char* backend;
const char* params_backend;
} sd_ctx_params_t;
otherarch/sdcpp/t5.hpp
+39 -29
View File
@@ -14,6 +14,28 @@
#include "model.h"
#include "tokenizers/t5_unigram_tokenizer.h"
struct T5Config {
int64_t num_layers = 24;
int64_t model_dim = 4096;
int64_t ff_dim = 10240;
int64_t num_heads = 64;
int64_t vocab_size = 32128;
bool relative_attention = true;
static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
bool is_umt5 = false) {
(void)tensor_storage_map;
(void)prefix;
T5Config config;
if (is_umt5) {
config.vocab_size = 256384;
config.relative_attention = false;
}
return config;
}
};
class T5LayerNorm : public UnaryBlock {
protected:
int64_t hidden_size;
@@ -272,30 +294,21 @@ public:
}
};
struct T5Params {
int64_t num_layers = 24;
int64_t model_dim = 4096;
int64_t ff_dim = 10240;
int64_t num_heads = 64;
int64_t vocab_size = 32128;
bool relative_attention = true;
};
struct T5 : public GGMLBlock {
T5Params params;
T5Config config;
public:
T5() {}
T5(T5Params params)
: params(params) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(params.num_layers,
params.model_dim,
params.model_dim,
params.ff_dim,
params.num_heads,
params.relative_attention));
blocks["shared"] = std::shared_ptr<GGMLBlock>(new Embedding(params.vocab_size,
params.model_dim));
T5(T5Config config)
: config(config) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
config.model_dim,
config.model_dim,
config.ff_dim,
config.num_heads,
config.relative_attention));
blocks["shared"] = std::shared_ptr<GGMLBlock>(new Embedding(config.vocab_size,
config.model_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -316,7 +329,7 @@ public:
};
struct T5Runner : public GGMLRunner {
T5Params params;
T5Config config;
T5 model;
std::vector<int> relative_position_bucket_vec;
@@ -325,12 +338,9 @@ struct T5Runner : public GGMLRunner {
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
bool is_umt5 = false)
: GGMLRunner(backend, params_backend) {
if (is_umt5) {
params.vocab_size = 256384;
params.relative_attention = false;
}
model = T5(params);
: GGMLRunner(backend, params_backend),
config(T5Config::detect_from_weights(tensor_storage_map, prefix, is_umt5)) {
model = T5(config);
model.init(params_ctx, tensor_storage_map, prefix);
}
@@ -478,7 +488,7 @@ struct T5Embedder {
bool alloc_params_buffer() {
if (!model.alloc_params_buffer()) {
return false;
}
}
return true;
}
@@ -563,7 +573,7 @@ struct T5Embedder {
// cuda f32: pass
// cuda q8_0: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_F16;
ModelLoader model_loader;
otherarch/sdcpp/tensor.hpp
+19
View File
@@ -235,6 +235,7 @@ namespace sd {
Tensor& masked_fill_(const Tensor<uint8_t>& mask, const T& value);
T sum() const;
T mean() const;
static Tensor zeros(std::vector<int64_t> shape) {
@@ -327,6 +328,24 @@ namespace sd {
std::vector<int64_t> shape_;
};
template <typename T>
inline T Tensor<T>::sum() const {
T total = T{};
for (const T& value : data_) {
total += value;
}
return total;
}
template <>
inline float Tensor<float>::sum() const {
double total = 0.0;
for (float value : data_) {
total += static_cast<double>(value);
}
return static_cast<float>(total);
}
template <typename T>
inline T Tensor<T>::mean() const {
if (empty()) {
otherarch/sdcpp/tokenizers/bpe_tokenizer.cpp
+2 -1
View File
@@ -182,7 +182,8 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
unsigned char b = utf8_token_str[i];
char hex_buf[16];
snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", b);
iter = encoder.find(utf8_to_utf32(hex_buf));
iter = encoder.find(utf8_to_utf32(hex_buf));
token_id = iter != encoder.end() ? iter->second : UNK_TOKEN_ID;
bpe_tokens.push_back(token_id);
token_strs.push_back(hex_buf);
}
otherarch/sdcpp/tokenizers/gemma_tokenizer.cpp
+161
View File
@@ -189,3 +189,164 @@ GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::st
load_from_merges(load_gemma_merges(), load_gemma_vocab_json());
}
}
std::string Gemma2Tokenizer::normalize(const std::string& text) const {
std::string normalized = text;
size_t pos = 0;
while ((pos = normalized.find(' ', pos)) != std::string::npos) {
normalized.replace(pos, 1, "\xE2\x96\x81");
pos += 3;
}
return normalized;
}
void Gemma2Tokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
nlohmann::json vocab;
try {
vocab = nlohmann::json::parse(vocab_utf8_str);
} catch (const nlohmann::json::parse_error&) {
GGML_ABORT("invalid vocab json str");
}
for (const auto& [key, value] : vocab.items()) {
std::u32string token = utf8_to_utf32(key);
int i = value;
encoder[token] = i;
decoder[i] = token;
}
encoder_len = static_cast<int>(vocab.size());
LOG_DEBUG("vocab size: %d", encoder_len);
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
for (const auto& merge : merges) {
size_t space_pos = merge.find(' ');
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
}
LOG_DEBUG("merges size %zu", merge_pairs.size());
int rank = 0;
for (const auto& merge : merge_pairs) {
bpe_ranks[merge] = rank++;
}
bpe_len = rank;
}
Gemma2Tokenizer::Gemma2Tokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
byte_level_bpe = false;
byte_fallback = true;
add_bos_token = true;
PAD_TOKEN = "<pad>";
EOS_TOKEN = "<eos>";
BOS_TOKEN = "<bos>";
UNK_TOKEN = "<unk>";
PAD_TOKEN_ID = 0;
EOS_TOKEN_ID = 1;
BOS_TOKEN_ID = 2;
UNK_TOKEN_ID = 3;
std::vector<std::string> special_tokens_before_merge = {
PAD_TOKEN,
EOS_TOKEN,
BOS_TOKEN,
UNK_TOKEN,
"<mask>",
"<2mass>",
"[@BOS@]",
};
for (int i = 0; i <= 98; i++) {
special_tokens_before_merge.push_back("<unused" + std::to_string(i) + ">");
}
special_tokens_before_merge.push_back("<start_of_turn>");
special_tokens_before_merge.push_back("<end_of_turn>");
for (int i = 1; i <= 31; i++) {
special_tokens_before_merge.push_back(std::string(i, '\n'));
}
for (int i = 2; i <= 31; i++) {
std::string whitespace_token;
for (int j = 0; j < i; j++) {
whitespace_token += "\xE2\x96\x81";
}
special_tokens_before_merge.push_back(whitespace_token);
}
std::vector<std::string> html_tokens = {
"<table>",
"<caption>",
"<thead>",
"<tbody>",
"<tfoot>",
"<tr>",
"<th>",
"<td>",
"</table>",
"</caption>",
"</thead>",
"</tbody>",
"</tfoot>",
"</tr>",
"</th>",
"</td>",
"<h1>",
"<h2>",
"<h3>",
"<h4>",
"<h5>",
"<h6>",
"<blockquote>",
"</h1>",
"</h2>",
"</h3>",
"</h4>",
"</h5>",
"</h6>",
"</blockquote>",
"<strong>",
"<em>",
"<b>",
"<i>",
"<u>",
"<s>",
"<sub>",
"<sup>",
"<code>",
"</strong>",
"</em>",
"</b>",
"</i>",
"</u>",
"</s>",
"</sub>",
"</sup>",
"</code>",
};
special_tokens_before_merge.insert(special_tokens_before_merge.end(),
html_tokens.begin(),
html_tokens.end());
for (int i = 0; i <= 0xFF; i++) {
char hex_buf[16];
snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", i);
special_tokens_before_merge.push_back(hex_buf);
}
std::vector<std::string> special_tokens_after_merge = {
"[toxicity=0]",
};
for (int i = 1; i <= 31; i++) {
special_tokens_after_merge.insert(special_tokens_after_merge.begin() + i - 1,
std::string(i, '\t'));
}
for (int i = 99; i <= 99; i++) {
special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
}
special_tokens = special_tokens_before_merge;
special_tokens.insert(special_tokens.end(),
special_tokens_after_merge.begin(),
special_tokens_after_merge.end());
if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
load_from_merges(merges_utf8_str, vocab_utf8_str);
} else {
load_from_merges(load_gemma2_merges(), load_gemma2_vocab_json());
}
}
otherarch/sdcpp/tokenizers/gemma_tokenizer.h
+9
View File
@@ -14,4 +14,13 @@ public:
explicit GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
};
class Gemma2Tokenizer : public BPETokenizer {
protected:
void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
std::string normalize(const std::string& text) const override;
public:
explicit Gemma2Tokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
};
#endif // __SD_TOKENIZERS_GEMMA_TOKENIZER_H__
otherarch/sdcpp/tokenizers/vocab/gemma2_merges.hpp
+3
View File
File diff suppressed because one or more lines are too long
otherarch/sdcpp/tokenizers/vocab/gemma2_vocab.hpp
+3
View File
File diff suppressed because one or more lines are too long
otherarch/sdcpp/tokenizers/vocab/vocab.cpp
+12
View File
@@ -1,5 +1,7 @@
#include "vocab.h"
#include "clip_merges.hpp"
#include "gemma2_merges.hpp"
#include "gemma2_vocab.hpp"
#include "gemma_merges.hpp"
#include "gemma_vocab.hpp"
#include "gpt_oss_merges.hpp"
@@ -50,6 +52,16 @@ std::string load_gemma_vocab_json() {
return json_str;
}
std::string load_gemma2_merges() {
std::string merges_utf8_str(reinterpret_cast<const char*>(gemma2_merges_utf8_c_str), sizeof(gemma2_merges_utf8_c_str));
return merges_utf8_str;
}
std::string load_gemma2_vocab_json() {
std::string json_str(reinterpret_cast<const char*>(gemma2_vocab_json_utf8_c_str), sizeof(gemma2_vocab_json_utf8_c_str));
return json_str;
}
std::string load_gpt_oss_merges() {
std::string merges_utf8_str(reinterpret_cast<const char*>(gpt_oss_merges_utf8_c_str), sizeof(gpt_oss_merges_utf8_c_str));
return merges_utf8_str;
otherarch/sdcpp/tokenizers/vocab/vocab.h
+2
View File
@@ -11,6 +11,8 @@ std::string load_t5_tokenizer_json();
std::string load_umt5_tokenizer_json();
std::string load_gemma_merges();
std::string load_gemma_vocab_json();
std::string load_gemma2_merges();
std::string load_gemma2_vocab_json();
std::string load_gpt_oss_merges();
std::string load_gpt_oss_vocab_json();
otherarch/sdcpp/unet.hpp
+153 -61
View File
@@ -1,6 +1,9 @@
#ifndef __UNET_HPP__
#define __UNET_HPP__
#include <algorithm>
#include <vector>
#include "common_block.hpp"
#include "diffusion_model.hpp"
#include "model.h"
@@ -9,6 +12,125 @@
#define UNET_GRAPH_SIZE 102400
struct UNetConfig {
SDVersion version = VERSION_SD1;
// network hparams
int in_channels = 4;
int out_channels = 4;
int num_res_blocks = 2;
std::vector<int> attention_resolutions = {4, 2, 1};
std::vector<int> channel_mult = {1, 2, 4, 4};
std::vector<int> transformer_depth = {1, 1, 1, 1};
int time_embed_dim = 1280; // model_channels*4
int num_heads = 8;
int num_head_channels = -1; // channels // num_heads
int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
bool use_linear_projection = false;
bool tiny_unet = false;
int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
SDVersion version = VERSION_SD1) {
UNetConfig config;
config.version = version;
if (sd_version_is_sd2(version)) {
config.context_dim = 1024;
config.num_head_channels = 64;
config.num_heads = -1;
config.use_linear_projection = true;
} else if (sd_version_is_sdxl(version)) {
config.context_dim = 2048;
config.attention_resolutions = {4, 2};
config.channel_mult = {1, 2, 4};
config.transformer_depth = {1, 2, 10};
config.num_head_channels = 64;
config.num_heads = -1;
config.use_linear_projection = true;
if (version == VERSION_SDXL_VEGA) {
config.transformer_depth = {1, 1, 2};
}
} else if (version == VERSION_SVD) {
config.in_channels = 8;
config.out_channels = 4;
config.context_dim = 1024;
config.adm_in_channels = 768;
config.num_head_channels = 64;
config.num_heads = -1;
config.use_linear_projection = true;
}
if (sd_version_is_inpaint(version)) {
config.in_channels = 9;
} else if (sd_version_is_unet_edit(version)) {
config.in_channels = 8;
}
if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
config.num_res_blocks = 1;
config.channel_mult = {1, 2, 4};
config.tiny_unet = true;
if (version == VERSION_SDXS_512_DS) {
config.attention_resolutions = {4, 2}; // here just like SDXL
}
}
auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
std::string name = prefix.empty() ? suffix : prefix + "." + suffix;
auto it = tensor_storage_map.find(name);
if (it == tensor_storage_map.end()) {
return nullptr;
}
return &it->second;
};
if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
if (input->n_dims == 4) {
config.in_channels = static_cast<int>(input->ne[2]);
config.model_channels = static_cast<int>(input->ne[3]);
config.time_embed_dim = config.model_channels * 4;
}
}
if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) {
if (time_embed->n_dims == 2) {
config.model_channels = static_cast<int>(time_embed->ne[0]);
config.time_embed_dim = static_cast<int>(time_embed->ne[1]);
}
}
if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) {
if (label_emb->n_dims == 2) {
config.adm_in_channels = static_cast<int>(label_emb->ne[0]);
config.time_embed_dim = static_cast<int>(label_emb->ne[1]);
}
}
if (const TensorStorage* out = find_weight("out.2.weight")) {
if (out->n_dims == 4) {
config.out_channels = static_cast<int>(out->ne[3]);
}
}
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) {
config.context_dim = static_cast<int>(tensor_storage.ne[0]);
break;
}
}
LOG_DEBUG("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s",
config.in_channels,
config.out_channels,
config.model_channels,
config.time_embed_dim,
config.context_dim,
config.adm_in_channels,
config.num_res_blocks,
config.tiny_unet ? "true" : "false");
return config;
}
};
class SpatialVideoTransformer : public SpatialTransformer {
protected:
int64_t time_depth;
@@ -166,66 +288,26 @@ public:
// ldm.modules.diffusionmodules.openaimodel.UNetModel
class UnetModelBlock : public GGMLBlock {
protected:
SDVersion version = VERSION_SD1;
// network hparams
int in_channels = 4;
int out_channels = 4;
int num_res_blocks = 2;
std::vector<int> attention_resolutions = {4, 2, 1};
std::vector<int> channel_mult = {1, 2, 4, 4};
std::vector<int> transformer_depth = {1, 1, 1, 1};
int time_embed_dim = 1280; // model_channels*4
int num_heads = 8;
int num_head_channels = -1; // channels // num_heads
int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
bool use_linear_projection = false;
bool tiny_unet = false;
public:
int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
UNetConfig config;
UnetModelBlock(SDVersion version = VERSION_SD1, const String2TensorStorage& tensor_storage_map = {})
: version(version) {
if (sd_version_is_sd2(version)) {
context_dim = 1024;
num_head_channels = 64;
num_heads = -1;
use_linear_projection = true;
} else if (sd_version_is_sdxl(version)) {
context_dim = 2048;
attention_resolutions = {4, 2};
channel_mult = {1, 2, 4};
transformer_depth = {1, 2, 10};
num_head_channels = 64;
num_heads = -1;
use_linear_projection = true;
if (version == VERSION_SDXL_VEGA) {
transformer_depth = {1, 1, 2};
}
} else if (version == VERSION_SVD) {
in_channels = 8;
out_channels = 4;
context_dim = 1024;
adm_in_channels = 768;
num_head_channels = 64;
num_heads = -1;
use_linear_projection = true;
}
if (sd_version_is_inpaint(version)) {
in_channels = 9;
} else if (sd_version_is_unet_edit(version)) {
in_channels = 8;
}
if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
num_res_blocks = 1;
channel_mult = {1, 2, 4};
tiny_unet = true;
if (version == VERSION_SDXS_512_DS) {
attention_resolutions = {4, 2}; // here just like SDXL
}
}
explicit UnetModelBlock(UNetConfig config = {})
: config(config) {
const SDVersion version = this->config.version;
const int in_channels = this->config.in_channels;
const int out_channels = this->config.out_channels;
const int num_res_blocks = this->config.num_res_blocks;
const auto& attention_resolutions = this->config.attention_resolutions;
const auto& channel_mult = this->config.channel_mult;
const auto& transformer_depth = this->config.transformer_depth;
const int time_embed_dim = this->config.time_embed_dim;
const int num_heads = this->config.num_heads;
const int num_head_channels = this->config.num_head_channels;
const int context_dim = this->config.context_dim;
const bool use_linear_projection = this->config.use_linear_projection;
const bool tiny_unet = this->config.tiny_unet;
const int model_channels = this->config.model_channels;
const int adm_in_channels = this->config.adm_in_channels;
// dims is always 2
// use_temporal_attention is always True for SVD
@@ -398,7 +480,7 @@ public:
ggml_tensor* x,
ggml_tensor* emb,
int num_video_frames) {
if (version == VERSION_SVD) {
if (config.version == VERSION_SVD) {
auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
return block->forward(ctx, x, emb, num_video_frames);
@@ -414,7 +496,7 @@ public:
ggml_tensor* x,
ggml_tensor* context,
int timesteps) {
if (version == VERSION_SVD) {
if (config.version == VERSION_SVD) {
auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
return block->forward(ctx, x, context, timesteps);
@@ -440,6 +522,13 @@ public:
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
// y: [N, adm_in_channels] or [1, adm_in_channels]
// return: [N, out_channels, h, w]
const SDVersion version = config.version;
const int model_channels = config.model_channels;
const int num_res_blocks = config.num_res_blocks;
const auto& attention_resolutions = config.attention_resolutions;
const auto& channel_mult = config.channel_mult;
const bool tiny_unet = config.tiny_unet;
if (context != nullptr) {
if (context->ne[2] != x->ne[3]) {
context = ggml_repeat(ctx->ggml_ctx, context, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, context->ne[0], context->ne[1], x->ne[3]));
@@ -601,6 +690,7 @@ public:
};
struct UNetModelRunner : public DiffusionModelRunner {
UNetConfig config;
UnetModelBlock unet;
UNetModelRunner(ggml_backend_t backend,
@@ -608,7 +698,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
SDVersion version = VERSION_SD1)
: DiffusionModelRunner(backend, params_backend, prefix), unet(version, tensor_storage_map) {
: DiffusionModelRunner(backend, params_backend, prefix),
config(UNetConfig::detect_from_weights(tensor_storage_map, prefix, version)),
unet(config) {
unet.init(params_ctx, tensor_storage_map, prefix);
}
otherarch/sdcpp/upscaler.cpp
+8
View File
@@ -25,6 +25,13 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
}
}
void UpscalerGGML::set_stream_layers_enabled(bool enabled) {
stream_layers_enabled = enabled;
if (esrgan_upscaler) {
esrgan_upscaler->set_stream_layers_enabled(enabled);
}
}
bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
bool offload_params_to_cpu,
int n_threads) {
@@ -76,6 +83,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
tile_size,
model_loader.get_tensor_storage_map());
esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled);
if (direct) {
esrgan_upscaler->set_conv2d_direct_enabled(true);
}
otherarch/sdcpp/upscaler.h
+2
View File
@@ -18,6 +18,7 @@ struct UpscalerGGML {
bool direct = false;
int tile_size = 128;
size_t max_graph_vram_bytes = 0;
bool stream_layers_enabled = false;
std::string backend_spec;
std::string params_backend_spec;
@@ -31,6 +32,7 @@ struct UpscalerGGML {
bool offload_params_to_cpu,
int n_threads);
void set_max_graph_vram_bytes(size_t max_vram_bytes);
void set_stream_layers_enabled(bool enabled);
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
};
otherarch/sdcpp/util.cpp
-29
View File
@@ -29,9 +29,7 @@
#include <unistd.h>
#endif
#include "ggml-backend.h"
#include "ggml.h"
#include "ggml_extend_backend.h"
#include "stable-diffusion.h"
bool ends_with(const std::string& str, const std::string& ending) {
@@ -1019,30 +1017,3 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
}
return result;
}
// namespace is needed to avoid conflicts with ggml_backend_extend.hpp
namespace ggml_cpu {
#include "ggml-cpu.h"
}
const char* sd_get_system_info() {
using namespace ggml_cpu;
static char buffer[1024];
std::stringstream ss;
ss << "System Info: \n";
ss << " SSE3 = " << ggml_cpu_has_sse3() << " | ";
ss << " AVX = " << ggml_cpu_has_avx() << " | ";
ss << " AVX2 = " << ggml_cpu_has_avx2() << " | ";
ss << " AVX512 = " << ggml_cpu_has_avx512() << " | ";
ss << " AVX512_VBMI = " << ggml_cpu_has_avx512_vbmi() << " | ";
ss << " AVX512_VNNI = " << ggml_cpu_has_avx512_vnni() << " | ";
ss << " FMA = " << ggml_cpu_has_fma() << " | ";
ss << " NEON = " << ggml_cpu_has_neon() << " | ";
ss << " ARM_FMA = " << ggml_cpu_has_arm_fma() << " | ";
ss << " F16C = " << ggml_cpu_has_f16c() << " | ";
ss << " FP16_VA = " << ggml_cpu_has_fp16_va() << " | ";
ss << " WASM_SIMD = " << ggml_cpu_has_wasm_simd() << " | ";
ss << " VSX = " << ggml_cpu_has_vsx() << " | ";
snprintf(buffer, sizeof(buffer), "%s", ss.str().c_str());
return buffer;
}
otherarch/sdcpp/wan.hpp
+163 -1499
View File
File diff suppressed because it is too large Load Diff
otherarch/sdcpp/wan_vae.hpp
+1361
View File
File diff suppressed because it is too large Load Diff
otherarch/sdcpp/z_image.hpp
+149 -70
View File
@@ -20,6 +20,104 @@ namespace ZImage {
constexpr int ADALN_EMBED_DIM = 256;
constexpr int SEQ_MULTI_OF = 32;
struct ZImageConfig {
int patch_size = 2;
int64_t hidden_size = 3840;
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t num_layers = 30;
int64_t num_refiner_layers = 2;
int64_t head_dim = 128;
int64_t num_heads = 30;
int64_t num_kv_heads = 30;
int64_t multiple_of = 256;
float ffn_dim_multiplier = 8.0f / 3.0f;
float norm_eps = 1e-5f;
bool qk_norm = true;
int64_t cap_feat_dim = 2560;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int64_t axes_dim_sum = 128;
static ZImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
ZImageConfig config;
int64_t detected_layers = 0;
int64_t detected_refiner_layers = 0;
int64_t detected_context_refiner = 0;
int64_t detected_head_dim = 0;
int64_t detected_qkv_dim = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.in_channels = tensor_storage.ne[0] / patch_area;
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.cap_feat_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
} else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) {
detected_head_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.attention.qkv.weight") && tensor_storage.n_dims == 2) {
detected_qkv_dim = tensor_storage.ne[1];
} else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
int64_t patch_area = config.patch_size * config.patch_size;
config.out_channels = tensor_storage.ne[1] / patch_area;
}
size_t pos = name.find("layers.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
detected_layers = std::max<int64_t>(detected_layers, block_index + 1);
}
}
pos = name.find("noise_refiner.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
detected_refiner_layers = std::max<int64_t>(detected_refiner_layers, block_index + 1);
}
}
pos = name.find("context_refiner.");
if (pos != std::string::npos) {
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
detected_context_refiner = std::max<int64_t>(detected_context_refiner, block_index + 1);
}
}
}
if (detected_layers > 0) {
config.num_layers = detected_layers;
}
if (detected_refiner_layers > 0 || detected_context_refiner > 0) {
config.num_refiner_layers = std::max(detected_refiner_layers, detected_context_refiner);
}
if (detected_head_dim > 0) {
config.head_dim = detected_head_dim;
config.num_heads = config.hidden_size / config.head_dim;
if (detected_qkv_dim > 0) {
int64_t qkv_heads = detected_qkv_dim / config.head_dim;
config.num_kv_heads = std::max<int64_t>(1, (qkv_heads - config.num_heads) / 2);
}
}
LOG_DEBUG("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
config.num_layers,
config.num_refiner_layers,
config.hidden_size,
config.num_heads,
config.num_kv_heads,
config.in_channels,
config.out_channels);
return config;
}
};
struct JointAttention : public GGMLBlock {
protected:
int64_t head_dim;
@@ -263,90 +361,70 @@ namespace ZImage {
}
};
struct ZImageParams {
int patch_size = 2;
int64_t hidden_size = 3840;
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t num_layers = 30;
int64_t num_refiner_layers = 2;
int64_t head_dim = 128;
int64_t num_heads = 30;
int64_t num_kv_heads = 30;
int64_t multiple_of = 256;
float ffn_dim_multiplier = 8.0f / 3.0f;
float norm_eps = 1e-5f;
bool qk_norm = true;
int64_t cap_feat_dim = 2560;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int64_t axes_dim_sum = 128;
};
class ZImageModel : public GGMLBlock {
protected:
ZImageParams z_image_params;
ZImageConfig config;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
}
public:
ZImageModel() = default;
ZImageModel(ZImageParams z_image_params)
: z_image_params(z_image_params) {
blocks["x_embedder"] = std::make_shared<Linear>(z_image_params.patch_size * z_image_params.patch_size * z_image_params.in_channels, z_image_params.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(z_image_params.hidden_size, 1024), 256, 256);
blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(z_image_params.cap_feat_dim, z_image_params.norm_eps);
blocks["cap_embedder.1"] = std::make_shared<Linear>(z_image_params.cap_feat_dim, z_image_params.hidden_size);
ZImageModel(ZImageConfig config)
: config(config) {
blocks["x_embedder"] = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, 256);
blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps);
blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::make_shared<JointTransformerBlock>(i,
z_image_params.hidden_size,
z_image_params.head_dim,
z_image_params.num_heads,
z_image_params.num_kv_heads,
z_image_params.multiple_of,
z_image_params.ffn_dim_multiplier,
z_image_params.norm_eps,
z_image_params.qk_norm,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
true);
blocks["noise_refiner." + std::to_string(i)] = block;
}
for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::make_shared<JointTransformerBlock>(i,
z_image_params.hidden_size,
z_image_params.head_dim,
z_image_params.num_heads,
z_image_params.num_kv_heads,
z_image_params.multiple_of,
z_image_params.ffn_dim_multiplier,
z_image_params.norm_eps,
z_image_params.qk_norm,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
false);
blocks["context_refiner." + std::to_string(i)] = block;
}
for (int i = 0; i < z_image_params.num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto block = std::make_shared<JointTransformerBlock>(i,
z_image_params.hidden_size,
z_image_params.head_dim,
z_image_params.num_heads,
z_image_params.num_kv_heads,
z_image_params.multiple_of,
z_image_params.ffn_dim_multiplier,
z_image_params.norm_eps,
z_image_params.qk_norm,
config.hidden_size,
config.head_dim,
config.num_heads,
config.num_kv_heads,
config.multiple_of,
config.ffn_dim_multiplier,
config.norm_eps,
config.qk_norm,
true);
blocks["layers." + std::to_string(i)] = block;
}
blocks["final_layer"] = std::make_shared<FinalLayer>(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels);
blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
}
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
@@ -393,14 +471,14 @@ namespace ZImage {
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt");
}
for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
@@ -410,7 +488,7 @@ namespace ZImage {
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img");
for (int i = 0; i < z_image_params.num_layers; i++) {
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
@@ -442,7 +520,7 @@ namespace ZImage {
int64_t C = x->ne[2];
int64_t N = x->ne[3];
int patch_size = z_image_params.patch_size;
int patch_size = config.patch_size;
auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
uint64_t n_img_token = img->ne[1];
@@ -467,7 +545,7 @@ namespace ZImage {
struct ZImageRunner : public DiffusionModelRunner {
public:
ZImageParams z_image_params;
ZImageConfig config;
ZImageModel z_image;
std::vector<float> pe_vec;
std::vector<float> timestep_vec;
@@ -478,8 +556,9 @@ namespace ZImage {
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_Z_IMAGE)
: DiffusionModelRunner(backend, params_backend, prefix) {
z_image = ZImageModel(z_image_params);
: DiffusionModelRunner(backend, params_backend, prefix),
config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
z_image = ZImageModel(config);
z_image.init(params_ctx, tensor_storage_map, prefix);
}
@@ -510,19 +589,19 @@ namespace ZImage {
pe_vec = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
z_image_params.patch_size,
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
SEQ_MULTI_OF,
ref_latents,
increase_ref_index,
z_image_params.theta,
config.theta,
circular_y_enabled,
circular_x_enabled,
z_image_params.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / z_image_params.axes_dim_sum / 2);
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, z_image_params.axes_dim_sum / 2, pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe, true, "pe");
// pe->data = nullptr;
@@ -615,7 +694,7 @@ namespace ZImage {
// cuda q8: pass
// cuda q8 fa: pass
// ggml_backend_t backend = ggml_backend_cuda_init(0);
ggml_backend_t backend = ggml_backend_cpu_init();
ggml_backend_t backend = sd_backend_cpu_init();
ggml_type model_data_type = GGML_TYPE_Q8_0;
ModelLoader model_loader;