vrr/koboldcpp
2
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-10-21 13:51:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

sd: sync with master-917-4c3cf75 (#2496)

Browse source 
* sd: sync with master-897-ac45422

* sd: sync to master-907-500ef5f

* sd: sync with master-913-b167b94

* sd: sync with master-917-4c3cf75

* sd: fix Qwen Image 2.1 edit mode detection

* sd: disable TAE for Qwen Image 2.1

* sd: warn about sdvramlimit without sdoffloadcpu
This commit is contained in:
Wagner Bruna 2026-10-03 12:52:58 -03:00 • committed by GitHub
parent d15a496ecb
commit 40fbf48a16
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
69 changed files with 4586 additions and 993 deletions
Download patch file Download diff file Expand all files Collapse all files

2
Makefile
View file

File diff suppressed because one or more lines are too long

69
otherarch/sdcpp/examples/cli/main.cpp
View file

@ -41,7 +41,6 @@ struct SDCliParams {
std::string metadata_format = "text";
sd_log_level_t log_level = SD_LOG_INFO;
bool canny_preprocess = false;
bool convert_name = false;
preview_t preview_method = PREVIEW_NONE;
@ -107,10 +106,6 @@ struct SDCliParams {
};
options.bool_options = {
{"",
"--canny",
"apply canny preprocessor (edge detection)",
true, &canny_preprocess},
{"",
"--convert-name",
"convert tensor name (for convert mode)",
@ -268,7 +263,6 @@ struct SDCliParams {
<< " metadata_format: \"" << metadata_format << "\",\n"
<< " log_level: " << log_level_name(log_level) << ",\n"
<< " color: " << (color ? "true" : "false") << ",\n"
<< " canny_preprocess: " << (canny_preprocess ? "true" : "false") << ",\n"
<< " convert_name: " << (convert_name ? "true" : "false") << ",\n"
<< " preview_method: " << previews_str[preview_method] << ",\n"
<< " preview_interval: " << preview_interval << ",\n"
@ -328,9 +322,7 @@ void sd_log_cb(enum sd_log_level_t level, const char* log, void* data) {
bool load_images_from_dir(const std::string dir,
std::vector<SDImageOwner>& images,
int expected_width = 0,
int expected_height = 0,
int max_image_num = 0) {
int max_image_num = 0) {
if (!fs::exists(dir) || !fs::is_directory(dir)) {
LOG_ERROR("'%s' is not a valid directory\n", dir.c_str());
return false;
@ -357,7 +349,8 @@ bool load_images_from_dir(const std::string dir,
LOG_VERBOSE("load image %zu from '%s'", images.size(), path.c_str());
int width = 0;
int height = 0;
uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, expected_width, expected_height);
int loaded_channel = 0;
uint8_t* image_buffer = load_image_from_file(path.c_str(), width, height, loaded_channel, 0, 0);
if (image_buffer == nullptr) {
LOG_ERROR("load image from '%s' failed", path.c_str());
return false;
@ -365,7 +358,7 @@ bool load_images_from_dir(const std::string dir,
images.emplace_back(sd_image_t{(uint32_t)width,
(uint32_t)height,
3,
(uint32_t)loaded_channel,
image_buffer});
if (max_image_num > 0 && static_cast<int>(images.size()) >= max_image_num) {
@ -651,10 +644,11 @@ int main(int argc, const char* argv[]) {
SDCliParams cli_params;
SDContextParams ctx_params;
ctx_params.conditioning_cache_size = 0;
SDGenerationParams gen_params;
parse_args(argc, argv, cli_params, ctx_params, gen_params);
sd_set_log_callback(sd_log_cb, (void*)&cli_params);
parse_args(argc, argv, cli_params, ctx_params, gen_params);
if (cli_params.mode == METADATA) {
MetadataReadOptions options;
@ -750,16 +744,8 @@ int main(int argc, const char* argv[]) {
auto load_image_and_update_size = [&](const std::string& path,
SDImageOwner& image,
bool resize_image = true,
int expected_channel = 3) -> bool {
int expected_width = 0;
int expected_height = 0;
if (resize_image && gen_params.width_and_height_are_set()) {
expected_width = gen_params.width;
expected_height = gen_params.height;
}
if (!load_sd_image_from_file(image.put(), path.c_str(), expected_width, expected_height, expected_channel)) {
if (!load_sd_image_from_file(image.put(), path.c_str(), 0, 0, expected_channel)) {
LOG_ERROR("load image from '%s' failed", path.c_str());
return false;
}
@ -781,7 +767,8 @@ int main(int argc, const char* argv[]) {
};
if (gen_params.init_image_path.size() > 0) {
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
const bool native_init = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image, native_init ? 0 : 3)) {
return 1;
}
}
@ -795,8 +782,8 @@ int main(int argc, const char* argv[]) {
if (gen_params.ref_image_paths.size() > 0) {
gen_params.ref_images.clear();
for (auto& path : gen_params.ref_image_paths) {
SDImageOwner ref_image({0, 0, 3, nullptr});
if (!load_image_and_update_size(path, ref_image, false)) {
SDImageOwner ref_image({0, 0, 0, nullptr});
if (!load_image_and_update_size(path, ref_image, 0)) {
return 1;
}
gen_params.ref_images.push_back(std::move(ref_image));
@ -837,41 +824,22 @@ int main(int argc, const char* argv[]) {
if (gen_params.mask_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.mask_image.put(),
gen_params.mask_image_path.c_str(),
gen_params.get_resolved_width(),
gen_params.get_resolved_height(),
0,
0,
1)) {
LOG_ERROR("load image from '%s' failed", gen_params.mask_image_path.c_str());
return 1;
}
} else {
sd_image_t generated_mask = {0, 0, 1, nullptr};
generated_mask.data = (uint8_t*)malloc(gen_params.get_resolved_width() * gen_params.get_resolved_height());
if (generated_mask.data == nullptr) {
LOG_ERROR("malloc mask image failed");
return 1;
}
generated_mask.width = gen_params.get_resolved_width();
generated_mask.height = gen_params.get_resolved_height();
memset(generated_mask.data, 255, gen_params.get_resolved_width() * gen_params.get_resolved_height());
gen_params.mask_image.reset(generated_mask);
}
if (gen_params.control_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.control_image.put(),
gen_params.control_image_path.c_str(),
gen_params.get_resolved_width(),
gen_params.get_resolved_height())) {
0,
0)) {
LOG_ERROR("load image from '%s' failed", gen_params.control_image_path.c_str());
return 1;
}
if (cli_params.canny_preprocess) { // apply preprocessor
preprocess_canny(gen_params.control_image.get(),
0.08f,
0.08f,
0.8f,
1.0f,
false);
}
}
if (gen_params.ip_adapter_image_path.size() > 0) {
@ -888,8 +856,6 @@ int main(int argc, const char* argv[]) {
gen_params.control_frames.clear();
if (!load_images_from_dir(gen_params.control_video_path,
gen_params.control_frames,
gen_params.get_resolved_width(),
gen_params.get_resolved_height(),
gen_params.video_frames)) {
return 1;
}
@ -898,10 +864,7 @@ int main(int argc, const char* argv[]) {
if (!gen_params.pm_id_images_dir.empty()) {
gen_params.pm_id_images.clear();
if (!load_images_from_dir(gen_params.pm_id_images_dir,
gen_params.pm_id_images,
0,
0,
0)) {
gen_params.pm_id_images)) {
return 1;
}
}

158
otherarch/sdcpp/examples/common/common.cpp
View file

@ -518,7 +518,8 @@ ArgOptions SDContextParams::get_options() {
{"",
"--model-args",
"extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
"chroma_t5_mask_pad, qwen_image_zero_cond_t",
"chroma_t5_mask_pad, qwen_image_zero_cond_t, qwen_image_2_1_prefix_cache, "
"qwen_image_2_1_prefix_cache_type (auto or a type name from --type)",
(int)',',
&model_args},
{"",
@ -571,6 +572,10 @@ ArgOptions SDContextParams::get_options() {
"number of threads to use during computation (default: -1). "
"If threads <= 0, then threads will be set to the number of CPU physical cores",
&n_threads},
{"",
"--conditioning-cache-size",
"maximum number of conditioning results cached per model context (default: " + std::to_string(conditioning_cache_size) + ", 0 disables caching)",
&conditioning_cache_size},
};
options.bool_options = {
@ -618,6 +623,10 @@ ArgOptions SDContextParams::get_options() {
"--diffusion-fa",
"use flash attention in the diffusion model only",
true, &diffusion_flash_attn},
{"",
"--sage-attn",
"use native CUDA SageAttention in the diffusion model, with flash/default attention fallback",
true, &sage_attn},
{"",
"--diffusion-conv-direct",
"use ggml_conv2d_direct in the diffusion model",
@ -818,6 +827,10 @@ bool SDContextParams::resolve(SDMode mode) {
}
bool SDContextParams::validate(SDMode mode) {
if (conditioning_cache_size < 0) {
LOG_ERROR("error: conditioning-cache-size must be non-negative");
return false;
}
if (mode == CONVERT) {
const bool has_convert_input = model_path.length() != 0 ||
clip_l_path.length() != 0 ||
@ -894,6 +907,7 @@ std::string SDContextParams::to_string() const {
std::ostringstream oss;
oss << "SDContextParams {\n"
<< " n_threads: " << n_threads << ",\n"
<< " conditioning_cache_size: " << conditioning_cache_size << ",\n"
<< " model_path: \"" << model_path << "\",\n"
<< " clip_l_path: \"" << clip_l_path << "\",\n"
<< " clip_g_path: \"" << clip_g_path << "\",\n"
@ -938,6 +952,7 @@ std::string SDContextParams::to_string() const {
<< " vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n"
<< " flash_attn: " << (flash_attn ? "true" : "false") << ",\n"
<< " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
<< " sage_attn: " << (sage_attn ? "true" : "false") << ",\n"
<< " linear_scale: " << linear_scale << ",\n"
<< " attn_scale: " << attn_scale << ",\n"
<< " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
@ -987,6 +1002,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.pulid_weights_path = pulid_weights_path.c_str();
sd_ctx_params.tensor_type_rules = tensor_type_rules.c_str();
sd_ctx_params.n_threads = n_threads;
sd_ctx_params.conditioning_cache_size = conditioning_cache_size;
sd_ctx_params.wtype = wtype;
sd_ctx_params.rng_type = rng_type;
sd_ctx_params.sampler_rng_type = sampler_rng_type;
@ -995,6 +1011,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.enable_mmap = enable_mmap;
sd_ctx_params.flash_attn = flash_attn;
sd_ctx_params.diffusion_flash_attn = diffusion_flash_attn;
sd_ctx_params.sage_attn = sage_attn;
sd_ctx_params.linear_scale = linear_scale;
sd_ctx_params.attn_scale = attn_scale;
sd_ctx_params.tae_preview_only = taesd_preview;
@ -1109,7 +1126,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; 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; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; 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; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng",
(int)',',
&extra_sample_args},
{"",
@ -1122,6 +1139,9 @@ ArgOptions SDGenerationParams::get_options() {
"Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
(int)',',
&ref_image_args},
{"", "--image-preprocess",
"Image preprocessing rule: target=init|end|mask|control|ref|ip-adapter|id|control-frame,index=N,mode=auto|none|stretch|crop|crop-resize|fit-pad,filter=auto|nearest|nearest-exact|bilinear|bicubic|lanczos,antialias=auto|true|false,width=W,height=H,anchor=center|top|bottom|left|right,pad_color=#RRGGBB[AA],canny=true|false. Repeat for multiple rules.",
(int)';', &image_preprocess},
};
options.int_options = {
@ -1302,11 +1322,6 @@ ArgOptions SDGenerationParams::get_options() {
"automatically increase the indices of references images based on the order they are listed (starting with 1).",
true,
&increase_ref_index},
{"",
"--disable-auto-resize-ref-image",
"disable auto resize of ref images",
false,
&auto_resize_ref_image},
{"",
"--circular",
"enable circular padding on both axes for tileable output",
@ -1748,7 +1763,7 @@ ArgOptions SDGenerationParams::get_options() {
on_scm_policy_arg},
{"",
"--vae-tile-size",
"tile size for vae tiling, format [X]x[Y] (default: 32x32)",
"tile size for vae tiling in latent units, not image pixels, format [X]x[Y] (default: 32x32)",
on_tile_size_arg},
{"",
"--vae-relative-tile-size",
@ -1842,28 +1857,28 @@ bool decode_base64_image(const std::string& encoded_input,
return false;
}
int decoded_width = 0;
int decoded_height = 0;
uint8_t* raw_data = load_image_from_memory(reinterpret_cast<const char*>(image_bytes.data()),
static_cast<int>(image_bytes.size()),
decoded_width,
decoded_height,
expected_width,
expected_height,
target_channels);
int decoded_width = 0;
int decoded_height = 0;
int resolved_channel = target_channels;
uint8_t* raw_data = load_image_from_memory(reinterpret_cast<const char*>(image_bytes.data()),
static_cast<int>(image_bytes.size()),
decoded_width,
decoded_height,
resolved_channel,
expected_width,
expected_height,
target_channels);
if (raw_data == nullptr) {
return false;
}
out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)target_channels, raw_data});
out_image.reset({(uint32_t)decoded_width, (uint32_t)decoded_height, (uint32_t)resolved_channel, raw_data});
return true;
}
static bool parse_image_json_field(const json& parent,
const char* key,
int channels,
int expected_width,
int expected_height,
SDImageOwner& out_image) {
if (!parent.contains(key)) {
return true;
@ -1875,14 +1890,12 @@ static bool parse_image_json_field(const json& parent,
if (!parent.at(key).is_string()) {
return false;
}
return decode_base64_image(parent.at(key).get<std::string>(), channels, expected_width, expected_height, out_image);
return decode_base64_image(parent.at(key).get<std::string>(), channels, 0, 0, out_image);
}
static bool parse_image_array_json_field(const json& parent,
const char* key,
int channels,
int expected_width,
int expected_height,
std::vector<SDImageOwner>& out_images) {
if (!parent.contains(key)) {
return true;
@ -1901,7 +1914,7 @@ static bool parse_image_array_json_field(const json& parent,
return false;
}
SDImageOwner image;
if (!decode_base64_image(item.get<std::string>(), channels, expected_width, expected_height, image)) {
if (!decode_base64_image(item.get<std::string>(), channels, 0, 0, image)) {
return false;
}
out_images.push_back(std::move(image));
@ -2000,6 +2013,29 @@ static bool resolve_model_file_from_dir(const std::string& model_name,
return false;
}
bool SDGenerationParams::parse_image_preprocess_json(const std::string& json_str) {
const auto value = json::parse(json_str, nullptr, false);
std::string rules;
if (value.is_string()) {
rules = value.get<std::string>();
} else if (value.is_array()) {
for (const auto& item : value) {
if (!item.is_string()) {
LOG_ERROR("image_preprocess must contain rule strings");
return false;
}
if (!rules.empty())
rules += ";";
rules += item.get<std::string>();
}
} else {
LOG_ERROR("image_preprocess must be a string or array of strings");
return false;
}
image_preprocess = std::move(rules);
return true;
}
bool SDGenerationParams::from_json_str(
const std::string& json_str,
const std::function<std::string(const std::string&)>& lora_path_resolver) {
@ -2011,6 +2047,9 @@ bool SDGenerationParams::from_json_str(
return false;
}
if (j.contains("image_preprocess") && !parse_image_preprocess_json(j["image_preprocess"].dump()))
return false;
auto load_if_exists = [&](const char* key, auto& out) {
if (j.contains(key)) {
using T = std::decay_t<decltype(out)>;
@ -2048,6 +2087,7 @@ bool SDGenerationParams::from_json_str(
load_if_exists("cache_mode", cache_mode);
load_if_exists("cache_option", cache_option);
load_if_exists("scm_mask", scm_mask);
load_if_exists("ref_image_args", ref_image_args);
load_if_exists("clip_skip", clip_skip);
load_if_exists("width", width);
@ -2065,7 +2105,6 @@ bool SDGenerationParams::from_json_str(
load_if_exists("moe_boundary", moe_boundary);
load_if_exists("vace_strength", vace_strength);
load_if_exists("auto_resize_ref_image", auto_resize_ref_image);
load_if_exists("increase_ref_index", increase_ref_index);
load_if_exists("embed_image_metadata", embed_image_metadata);
@ -2209,32 +2248,23 @@ bool SDGenerationParams::from_json_str(
LOG_ERROR("invalid lora");
return false;
}
if (!parse_image_json_field(j, "init_image", 3, width, height, init_image)) {
LOG_ERROR("invalid init_image");
auto load_image = [&](const char* key, int channels, SDImageOwner& image) {
if (!parse_image_json_field(j, key, channels, image)) {
LOG_ERROR("invalid %s", key);
return false;
}
return true;
};
if (!load_image("init_image", 0, init_image) ||
!load_image("end_image", 3, end_image) ||
!load_image("mask_image", 1, mask_image) ||
!load_image("control_image", 3, control_image) ||
!load_image("ip_adapter_image", 3, ip_adapter_image)) {
return false;
}
if (!parse_image_json_field(j, "end_image", 3, width, height, end_image)) {
LOG_ERROR("invalid end_image");
return false;
}
if (!parse_image_array_json_field(j, "ref_images", 3, width, height, ref_images)) {
LOG_ERROR("invalid ref_images");
return false;
}
if (!parse_image_array_json_field(j, "control_frames", 3, width, height, control_frames)) {
LOG_ERROR("invalid control_frames");
return false;
}
if (!parse_image_json_field(j, "mask_image", 1, width, height, mask_image)) {
LOG_ERROR("invalid mask_image");
return false;
}
if (!parse_image_json_field(j, "control_image", 3, width, height, control_image)) {
LOG_ERROR("invalid control_image");
return false;
}
if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
LOG_ERROR("invalid ip_adapter_image");
if (!parse_image_array_json_field(j, "ref_images", 0, ref_images) ||
!parse_image_array_json_field(j, "control_frames", 3, control_frames)) {
LOG_ERROR("invalid input image array");
return false;
}
@ -2478,6 +2508,10 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s
}
bool SDGenerationParams::validate(SDMode mode) {
if (!image_preprocess.empty() && mode != IMG_GEN && mode != VID_GEN) {
LOG_ERROR("--image-preprocess requires img_gen or vid_gen mode");
return false;
}
if (batch_count <= 0) {
LOG_ERROR("error: batch_count must be greater than 0");
return false;
@ -2653,14 +2687,6 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
pulid_id_weight,
};
if (!auto_resize_ref_image) {
if (!ref_image_args.empty()) {
ref_image_args += ",";
}
ref_image_args += "resize_before_vae=0";
LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
}
if (increase_ref_index) {
if (!ref_image_args.empty()) {
ref_image_args += ",";
@ -2708,6 +2734,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
params.circular_x = circular || circular_x;
params.circular_y = circular || circular_y;
params.image_preprocess = {image_preprocess.c_str()};
return params;
}
@ -2810,6 +2837,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
params.circular_x = circular || circular_x;
params.circular_y = circular || circular_y;
params.image_preprocess = {image_preprocess.c_str()};
return params;
}
@ -2866,7 +2894,8 @@ std::string SDGenerationParams::to_string() const {
<< " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n"
<< " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n"
<< " control_video_path: \"" << control_video_path << "\",\n"
<< " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n"
<< " image_preprocess: " << image_preprocess << ",\n"
<< " ref_image_args: " << ref_image_args << ",\n"
<< " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
<< " pm_id_images_dir: \"" << pm_id_images_dir << "\",\n"
<< " pm_id_embed_path: \"" << pm_id_embed_path << "\",\n"
@ -3017,12 +3046,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path);
root["models"] = std::move(models);
root["clip_skip"] = gen_params.clip_skip;
root["strength"] = gen_params.strength;
root["control_strength"] = gen_params.control_strength;
root["ip_adapter_strength"] = gen_params.ip_adapter_strength;
root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
root["increase_ref_index"] = gen_params.increase_ref_index;
root["clip_skip"] = gen_params.clip_skip;
root["strength"] = gen_params.strength;
root["control_strength"] = gen_params.control_strength;
root["ip_adapter_strength"] = gen_params.ip_adapter_strength;
root["ref_image_args"] = gen_params.ref_image_args;
root["image_preprocess"] = gen_params.image_preprocess;
root["increase_ref_index"] = gen_params.increase_ref_index;
if (mode == VID_GEN) {
root["video"] = {
{"frame_count", gen_params.video_frames},

29
otherarch/sdcpp/examples/common/common.h
View file

@ -116,7 +116,8 @@ bool decode_base64_image(const std::string& encoded_input,
SDImageOwner& out_image);
struct SDContextParams {
int n_threads = -1;
int n_threads = -1;
int conditioning_cache_size = 4;
std::string model_path;
std::string clip_l_path;
std::string clip_g_path;
@ -170,6 +171,7 @@ struct SDContextParams {
bool vae_on_cpu = false;
bool flash_attn = false;
bool diffusion_flash_attn = false;
bool sage_attn = false;
bool diffusion_conv_direct = false;
bool vae_conv_direct = false;
@ -199,18 +201,17 @@ struct SDGenerationParams {
std::string ad_prompt;
std::string ad_negative_prompt;
std::string extra_ad_args;
int clip_skip = -1; // <= 0 represents unspecified
int width = -1;
int height = -1;
int batch_count = 1;
int qwen_image_layers = 3;
int64_t seed = 42;
float strength = 0.75f;
float control_strength = 0.9f;
float ip_adapter_strength = 1.0f;
bool auto_resize_ref_image = true;
bool increase_ref_index = false;
bool embed_image_metadata = true;
int clip_skip = -1; // <= 0 represents unspecified
int width = -1;
int height = -1;
int batch_count = 1;
int qwen_image_layers = 3;
int64_t seed = 42;
float strength = 0.75f;
float control_strength = 0.9f;
float ip_adapter_strength = 1.0f;
bool increase_ref_index = false;
bool embed_image_metadata = true;
std::string init_image_path;
std::string end_image_path;
@ -246,6 +247,7 @@ struct SDGenerationParams {
std::string extra_tiling_args;
std::string ref_image_args;
std::string image_preprocess;
std::string pm_id_images_dir;
std::string pm_id_embed_path;
@ -309,6 +311,7 @@ struct SDGenerationParams {
ArgOptions get_options();
bool from_json_str(const std::string& json_str,
const std::function<std::string(const std::string&)>& lora_path_resolver = {});
bool parse_image_preprocess_json(const std::string& json_str);
bool initialize_cache_params();
void extract_and_remove_lora(const std::string& lora_model_dir);
bool width_and_height_are_set() const;

47
otherarch/sdcpp/examples/common/media_io.cpp
View file

@ -275,6 +275,10 @@ uint8_t* decode_webp_image_to_buffer(const uint8_t* data,
height = features.height;
source_channel_count = features.has_alpha ? 4 : 3;
if (expected_channel == 0) {
expected_channel = source_channel_count;
}
const size_t pixel_count = static_cast<size_t>(width) * static_cast<size_t>(height);
if (expected_channel == 1) {
@ -495,7 +499,8 @@ uint8_t* load_image_common(bool from_memory,
int& height,
int expected_width,
int expected_height,
int expected_channel) {
int expected_channel,
int& out_channel) {
const char* image_path;
FreeUniquePtr<uint8_t> image_buffer;
int source_channel_count = 0;
@ -552,6 +557,32 @@ uint8_t* load_image_common(bool from_memory,
LOG_ERROR("load image from '%s' failed", image_path);
return nullptr;
}
if (expected_channel == 0) {
expected_channel = source_channel_count == 2 ? 4 : (source_channel_count == 1 ? 3 : source_channel_count);
if (expected_channel != source_channel_count) {
FreeUniquePtr<uint8_t> promoted((uint8_t*)malloc((size_t)width * height * expected_channel));
if (promoted == nullptr) {
LOG_ERROR("error: allocate memory for channel promotion, image_path = %s", image_path);
return nullptr;
}
const size_t pixel_count = (size_t)width * (size_t)height;
for (size_t i = 0; i < pixel_count; ++i) {
if (source_channel_count == 1) {
promoted.get()[i * 3 + 0] = image_buffer.get()[i];
promoted.get()[i * 3 + 1] = image_buffer.get()[i];
promoted.get()[i * 3 + 2] = image_buffer.get()[i];
} else {
promoted.get()[i * 4 + 0] = image_buffer.get()[i * 2];
promoted.get()[i * 4 + 1] = image_buffer.get()[i * 2];
promoted.get()[i * 4 + 2] = image_buffer.get()[i * 2];
promoted.get()[i * 4 + 3] = image_buffer.get()[i * 2 + 1];
}
}
image_buffer = std::move(promoted);
source_channel_count = expected_channel;
}
}
// stb reports the source channel count even when it converts the output.
if (source_channel_count < expected_channel) {
fprintf(stderr,
"the number of channels for the input image must be >= %d,"
@ -611,7 +642,7 @@ uint8_t* load_image_common(bool from_memory,
}
stbir_resize(image_buffer.get(), width, height, 0,
resized_image_buffer.get(), expected_width, expected_height, 0, STBIR_TYPE_UINT8,
expected_channel, STBIR_ALPHA_CHANNEL_NONE, 0,
expected_channel, expected_channel == 4 ? 3 : STBIR_ALPHA_CHANNEL_NONE, 0,
STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
STBIR_FILTER_BOX, STBIR_FILTER_BOX,
STBIR_COLORSPACE_SRGB, nullptr);
@ -619,6 +650,7 @@ uint8_t* load_image_common(bool from_memory,
height = expected_height;
image_buffer = std::move(resized_image_buffer);
}
out_channel = expected_channel;
return image_buffer.release();
}
@ -791,10 +823,11 @@ bool write_image_to_file(const std::string& path,
uint8_t* load_image_from_file(const char* image_path,
int& width,
int& height,
int& out_channel,
int expected_width,
int expected_height,
int expected_channel) {
return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel);
return load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, out_channel);
}
bool load_sd_image_from_file(sd_image_t* image,
@ -804,13 +837,14 @@ bool load_sd_image_from_file(sd_image_t* image,
int expected_channel) {
int width;
int height;
image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel);
int resolved_channel = expected_channel;
image->data = load_image_common(false, image_path, 0, width, height, expected_width, expected_height, expected_channel, resolved_channel);
if (image->data == nullptr) {
return false;
}
image->width = width;
image->height = height;
image->channel = expected_channel;
image->channel = resolved_channel;
return true;
}
@ -818,10 +852,11 @@ uint8_t* load_image_from_memory(const char* image_bytes,
int len,
int& width,
int& height,
int& out_channel,
int expected_width,
int expected_height,
int expected_channel) {
return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel, out_channel);
}
static void append_avi_metadata(std::vector<uint8_t>& data, const std::string& parameters) {

4
otherarch/sdcpp/examples/common/media_io.h
View file

@ -32,9 +32,12 @@ bool write_image_to_file(const std::string& path,
const std::string& parameters = "",
int quality = 90);
// expected_channel == 0 preserves native channels (grayscale -> RGB, gray+alpha -> RGBA).
// out_channel receives the output channel count.
uint8_t* load_image_from_file(const char* image_path,
int& width,
int& height,
int& out_channel,
int expected_width = 0,
int expected_height = 0,
int expected_channel = 3);
@ -49,6 +52,7 @@ uint8_t* load_image_from_memory(const char* image_bytes,
int len,
int& width,
int& height,
int& out_channel,
int expected_width = 0,
int expected_height = 0,
int expected_channel = 3);

12
otherarch/sdcpp/include/stable-diffusion.h
View file

@ -79,6 +79,7 @@ enum scheduler_t {
FLUX2_SCHEDULER,
FLUX_SCHEDULER,
BETA_SCHEDULER,
LLADA_IMAGE_SCHEDULER,
SCHEDULER_COUNT
};
@ -245,6 +246,8 @@ typedef struct {
float linear_scale; // Override linear input scaling; 0 keeps the model default
float attn_scale; // Override flash-attention K/V scaling; 0 keeps the model default
const char* tokenizer; // tokenizer.json path or main=FILE,clip-l=FILE,clip-g=FILE assignments; required for PiD and Lens
bool sage_attn;
int conditioning_cache_size; // Maximum cached conditioning entries per context; 0 disables caching (default: 4)
} sd_ctx_params_t;
typedef struct {
@ -261,6 +264,11 @@ typedef struct {
uint8_t* data;
} sd_image_t;
typedef struct {
// Semicolon-separated target=...,key=value rules. NULL preserves defaults.
const char* rules;
} sd_image_preprocess_params_t;
typedef struct {
sd_image_t* frames;
int frame_count;
@ -408,6 +416,7 @@ typedef struct {
int qwen_image_layers;
bool circular_x;
bool circular_y;
sd_image_preprocess_params_t image_preprocess;
} sd_img_gen_params_t;
typedef struct {
@ -442,6 +451,7 @@ typedef struct {
sd_hires_params_t hires;
bool circular_x;
bool circular_y;
sd_image_preprocess_params_t image_preprocess;
} sd_vid_gen_params_t;
typedef struct sd_ctx_t sd_ctx_t;
@ -549,6 +559,8 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
int* num_images_out);
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
// Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model.
SD_API int get_upscaler_model_scale(const char* model_path);
typedef struct adetailer_ctx_t adetailer_ctx_t;

3
otherarch/sdcpp/sdtype_adapter.cpp
View file

@ -488,6 +488,9 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
if(inputs.max_vram && *inputs.max_vram) {
max_vram = inputs.max_vram;
printf("Using max VRAM = %s GB\n", max_vram.c_str());
if (params_backend == "") {
printf("Note: a VRAM limit may not be effective without offloading!\n");
}
}
if(inputs.quant > 0)
{

215
otherarch/sdcpp/src/conditioning/conditioner.hpp
View file

@ -14,6 +14,7 @@
#include "core/util.h"
#include "model/diffusion/model.hpp"
#include "model/te/clip.hpp"
#include "model/te/llada_image_te.hpp"
#include "model/te/llm.hpp"
#include "model/te/t5.hpp"
#include "model_loader.h"
@ -2218,7 +2219,10 @@ struct LLMEmbedder : public Conditioner {
false,
deepstack_image_embeds,
image_grids);
GGML_ASSERT(!hidden_states.empty());
if (hidden_states.empty()) {
LOG_ERROR("LLM prompt encoding failed");
return {};
}
hidden_states = apply_token_weights(std::move(hidden_states), weights);
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
@ -3162,6 +3166,7 @@ struct LLMEmbedder : public Conditioner {
int64_t tag_count = static_cast<int64_t>(tags.size());
result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags));
}
return result;
}
};
@ -3236,6 +3241,214 @@ struct LTXAVTextProjectionRunner : public GGMLRunner {
}
};
// LLaDA-Image's text path is a three-stage pipeline rather than a single encoder pass:
// the token embeddings feed a QueryFormer whose 256 queries are appended to the backbone
// input, and the backbone's final hidden states are projected to the denoiser's caption dim.
// Ref: LLaDAImagePipeline._encode_text.
struct LLaDAImageEmbedder : public Conditioner {
std::shared_ptr<Tokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<LLaDAImageTE::QueryFormerRunner> query_former;
std::shared_ptr<LLaDAImageTE::TextProjectionRunner> text_projection;
std::shared_ptr<LLaDAImageTE::SigVQRunner> sigvq;
std::string llm_prefix;
std::string query_former_prefix;
std::string text_projection_prefix;
std::string sigvq_prefix;
LLaDAImageEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string& llm_prefix = "text_encoders.llm",
const std::string& query_former_prefix = "queryformer",
const std::string& text_projection_prefix = "text_projection",
const std::string& sigvq_prefix = "sigvq",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const TokenizerConfig& tokenizers = {})
: llm_prefix(llm_prefix),
query_former_prefix(query_former_prefix),
text_projection_prefix(text_projection_prefix),
sigvq_prefix(sigvq_prefix) {
if (!tokenizers.has(TokenizerConfig::MAIN)) {
throw std::runtime_error("LLaDA-Image requires an external LLaDA2 tokenizer.json; pass --tokenizer FILE or set sd_ctx_params_t::tokenizer");
}
llm = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::LLADA2_MOE,
backend,
tensor_storage_map,
llm_prefix,
false,
weight_manager);
// <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json.
tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892);
query_former = std::make_shared<LLaDAImageTE::QueryFormerRunner>(backend,
tensor_storage_map,
query_former_prefix,
weight_manager);
text_projection = std::make_shared<LLaDAImageTE::TextProjectionRunner>(backend,
tensor_storage_map,
text_projection_prefix,
weight_manager);
// SigVQ is only present when the user supplies the editing weights.
for (const auto& [name, _] : tensor_storage_map) {
if (starts_with(name, sigvq_prefix + ".")) {
sigvq = std::make_shared<LLaDAImageTE::SigVQRunner>(backend,
tensor_storage_map,
sigvq_prefix,
weight_manager);
break;
}
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, llm_prefix);
query_former->get_param_tensors(tensors, query_former_prefix);
text_projection->get_param_tensors(tensors, text_projection_prefix);
if (sigvq != nullptr) {
sigvq->get_param_tensors(tensors, sigvq_prefix);
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
query_former->set_flash_attention_enabled(enabled);
text_projection->set_flash_attention_enabled(enabled);
if (sigvq != nullptr) {
sigvq->set_flash_attention_enabled(enabled);
}
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
query_former->set_max_graph_vram_bytes(max_vram_bytes);
text_projection->set_max_graph_vram_bytes(max_vram_bytes);
if (sigvq != nullptr) {
sigvq->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
llm->set_graph_cut_layer_split_enabled(enabled);
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, llm_prefix);
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
llm->set_weight_adapter(adapter);
query_former->set_weight_adapter(adapter);
text_projection->set_weight_adapter(adapter);
if (sigvq != nullptr) {
sigvq->set_weight_adapter(adapter);
}
}
void runner_end() override {
llm->runner_end();
query_former->runner_end();
text_projection->runner_end();
if (sigvq != nullptr) {
sigvq->runner_end();
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
const int64_t num_queries = 256;
const bool has_ref_images = conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty();
if (has_ref_images && sigvq == nullptr) {
LOG_ERROR("LLaDA-Image editing requires connectors with SigVQ weights");
return {};
}
std::string text = conditioner_params.text;
while (!text.empty() && std::isspace(static_cast<unsigned char>(text.front()))) {
text.erase(text.begin());
}
while (!text.empty() && std::isspace(static_cast<unsigned char>(text.back()))) {
text.pop_back();
}
std::string prompt = text.empty()
? "<role>HUMAN</role> Generate an image.\n<role>ASSISTANT</role>\n<IMAGE1>"
: "<role>HUMAN</role> Generate an image: " + text + "\n<role>ASSISTANT</role>\n<IMAGE1>";
std::vector<int> tokens;
if (!tokenizer->encode(prompt, tokens, nullptr)) {
return {};
}
int64_t n_text = static_cast<int64_t>(tokens.size());
GGML_ASSERT(n_text > 0);
sd::Tensor<int32_t> text_ids({n_text}, std::vector<int32_t>(tokens.begin(), tokens.end()));
auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids);
auto query_embeds = query_former->compute(n_threads, inputs_embeds);
// splice_image_embeds() replaces tokens in place, so the query slots have to exist in
// input_ids; their ids are irrelevant because the embeddings are overwritten.
std::vector<int32_t> padded(tokens.begin(), tokens.end());
padded.resize(static_cast<size_t>(n_text + num_queries), tokenizer->PAD_TOKEN_ID);
int64_t n_total = static_cast<int64_t>(padded.size());
sd::Tensor<int32_t> input_ids({n_total}, padded);
// Bidirectional everywhere except that the text tokens must not see the appended
// queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min.
const float mask_min = std::numeric_limits<float>::lowest() / 4.0f;
sd::Tensor<float> attention_mask({n_total, n_total});
for (int64_t i1 = 0; i1 < n_total; ++i1) {
for (int64_t i0 = 0; i0 < n_total; ++i0) {
float value = (i1 < n_text && i0 >= n_text) ? mask_min : 0.0f;
attention_mask[i0 + n_total * i1] = value;
}
}
LLM::ImageEmbeds image_embeds;
image_embeds.emplace_back(static_cast<int>(n_text), query_embeds);
std::set<int> out_layers = {static_cast<int>(llm->config.num_layers) + 1};
auto hidden_states = llm->compute(n_threads,
input_ids,
attention_mask,
image_embeds,
out_layers);
SDCondition result;
result.c_crossattn = text_projection->compute(n_threads, hidden_states);
// Editing: SigVQ sees the reference at half the output resolution, as in
// LLaDAImagePipeline._encode_source_image.
if (has_ref_images) {
const auto& ref = conditioner_params.ref_images->front();
auto resized = sd::ops::interpolate(ref,
{conditioner_params.width / 2,
conditioner_params.height / 2,
ref.shape()[2],
ref.shape()[3]},
sd::ops::InterpolateMode::Bilinear);
resized = resized * 2.f - 1.f;
auto semantic = sigvq->compute(n_threads, resized);
if (semantic.empty()) {
return {};
}
result.extra_c_crossattns.push_back(std::move(semantic));
}
return result;
}
};
struct LTXAVEmbedder : public Conditioner {
static constexpr int64_t kHiddenSize = 3840;
static constexpr int64_t kNumStates = 49;

107
otherarch/sdcpp/src/conditioning/conditioning_cache.h Normal file
View file

@ -0,0 +1,107 @@
#ifndef __SD_CONDITIONING_CONDITIONING_CACHE_H__
#define __SD_CONDITIONING_CONDITIONING_CACHE_H__
#include <algorithm>
#include <list>
#include <tuple>
#include "conditioning/conditioner.hpp"
class ConditioningCache {
struct Entry {
ConditionerParams params;
std::vector<sd::Tensor<float>> ref_images;
std::vector<MiniMaxH3PresentationItem> references;
SDCondition condition;
Entry(const ConditionerParams& input, const SDCondition& output)
: params(input), condition(output) {
// Request-owned reference pointers must not outlive the request.
if (input.ref_images != nullptr) {
ref_images = *input.ref_images;
params.ref_images = &ref_images;
}
if (input.minimax_h3_references != nullptr) {
references = *input.minimax_h3_references;
params.minimax_h3_references = &references;
}
}
Entry(const Entry&) = delete;
Entry& operator=(const Entry&) = delete;
};
size_t capacity_ = 4;
std::list<Entry> entries_;
static bool same_images(const std::vector<sd::Tensor<float>>& a,
const std::vector<sd::Tensor<float>>& b) {
return std::equal(a.begin(), a.end(), b.begin(), b.end(),
[](const sd::Tensor<float>& x, const sd::Tensor<float>& y) {
return x.shape() == y.shape() && x.values() == y.values();
});
}
static bool same_params(const ConditionerParams& a, const ConditionerParams& b) {
const auto fields = [](const ConditionerParams& p) {
const auto& r = p.ref_image_params;
return std::tie(p.text, p.clip_skip, p.width, p.height, p.zero_out_masked,
r.pass_to_vlm, r.pass_to_dit, r.ref_index_mode,
r.force_ref_timestep_zero, r.resize_before_vae, r.vae_input_max_pixels,
r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target);
};
if (fields(a) != fields(b) ||
(a.ref_images == nullptr) != (b.ref_images == nullptr) ||
(a.minimax_h3_references == nullptr) != (b.minimax_h3_references == nullptr)) {
return false;
}
if (a.ref_images != nullptr && !same_images(*a.ref_images, *b.ref_images)) {
return false;
}
if (a.minimax_h3_references != nullptr &&
!std::equal(a.minimax_h3_references->begin(), a.minimax_h3_references->end(),
b.minimax_h3_references->begin(), b.minimax_h3_references->end(),
[](const MiniMaxH3PresentationItem& x, const MiniMaxH3PresentationItem& y) {
return x.kind == y.kind && x.timestamps == y.timestamps && same_images(x.frames, y.frames);
})) {
return false;
}
return true;
}
public:
void set_capacity(size_t capacity) {
capacity_ = capacity;
while (entries_.size() > capacity_) {
entries_.pop_back();
}
}
void clear() {
entries_.clear();
}
SDCondition get(Conditioner& conditioner, int n_threads, const ConditionerParams& params) {
if (capacity_ == 0) {
return conditioner.get_learned_condition(n_threads, params);
}
for (auto it = entries_.begin(); it != entries_.end(); ++it) {
if (same_params(it->params, params)) {
entries_.splice(entries_.begin(), entries_, it);
LOG_INFO("conditioning cache hit");
return entries_.front().condition;
}
}
auto condition = conditioner.get_learned_condition(n_threads, params);
if (!condition.empty()) {
if (entries_.size() == capacity_) {
entries_.pop_back();
}
entries_.emplace_front(params, condition);
LOG_VERBOSE("conditioning cache stored (%zu/%zu)", entries_.size(), capacity_);
}
return condition;
}
};
#endif // __SD_CONDITIONING_CONDITIONING_CACHE_H__

8
otherarch/sdcpp/src/core/backend_fit.cpp
View file

@ -478,7 +478,11 @@ namespace sd::backend_fit {
return true;
}
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling, ggml_status status) {
// Execution failures can leave the device unusable; tiling only helps with allocation failures.
if (status != GGML_STATUS_ALLOC_FAILED) {
return false;
}
const char* retry_mode = nullptr;
if (prefer_temporal_tiling && !tiling_params.temporal_tiling) {
tiling_params.temporal_tiling = true;
@ -498,7 +502,7 @@ namespace sd::backend_fit {
return false;
}
LOG_WARN("VAE decode failed (likely out of memory); retrying with %s tiling",
LOG_WARN("VAE decode ran out of memory; retrying with %s tiling",
retry_mode);
return true;
}

3
otherarch/sdcpp/src/core/backend_fit.h
View file

@ -16,7 +16,8 @@ namespace sd::backend_fit {
std::string& params_spec);
bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
bool prefer_temporal_tiling);
bool prefer_temporal_tiling,
ggml_status status);
} // namespace sd::backend_fit

51
otherarch/sdcpp/src/core/ggml_extend.cpp
View file

@ -465,7 +465,11 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
int d1,
int d2,
bool force_prec_f32,
bool direct) {
bool direct,
float scale) {
if (scale != 1.f) {
x = ggml_ext_scale(ctx, x, scale);
}
if (direct) {
int64_t OC = w->ne[3] / IC;
int64_t N = x->ne[3] / IC;
@ -502,6 +506,9 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
}
}
if (scale != 1.f) {
x = ggml_ext_scale(ctx, x, 1.f / scale);
}
if (b != nullptr) {
b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1]
x = ggml_add_inplace(ctx, x, b);
@ -622,7 +629,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
ggml_tensor* mask,
bool skip_reshape,
bool flash_attn,
float kv_scale) { // avoid overflow
float kv_scale,
bool sage_attn,
bool* used_flash_attn) { // avoid overflow
if (used_flash_attn != nullptr) {
*used_flash_attn = false;
}
int64_t L_q;
int64_t L_k;
int64_t C;
@ -713,7 +725,37 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
return out;
};
if (flash_attn) {
#ifndef SD_USE_UPSTREAM_GGML
if (sage_attn && mask == nullptr && d_head > 0 && d_head <= 128) {
auto q_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, q->type == GGML_TYPE_F32 ? q : ggml_cast(ctx, q, GGML_TYPE_F32)), d_head, L_q, n_head, N);
auto k_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, k->type == GGML_TYPE_F32 ? k : ggml_cast(ctx, k, GGML_TYPE_F32)), d_head, L_k, n_kv_head, N);
auto v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3));
const int64_t padded_head = d_head <= 64 ? 64 : 128;
if ((padded_head != d_head || kv_scale != 1.0f) && v_in->type != GGML_TYPE_F32) {
v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32);
}
if (padded_head != d_head) {
// Keep the original head's softmax scale when padding for the CUDA kernel.
q_in = ggml_pad(ctx, q_in, padded_head - d_head, 0, 0, 0);
k_in = ggml_pad(ctx, k_in, padded_head - d_head, 0, 0, 0);
v_in = ggml_pad(ctx, v_in, padded_head - d_head, 0, 0, 0);
}
if (kv_scale != 1.0f) {
k_in = ggml_ext_scale(ctx, k_in, kv_scale);
v_in = ggml_ext_scale(ctx, v_in, kv_scale);
}
v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16);
auto out = ggml_sage_attn(ctx, q_in, k_in, v_in, scale / kv_scale, GGML_SAGE_ATTN_AUTO);
if (ggml_backend_supports_op(backend, out)) {
kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out;
if (padded_head != d_head) {
kqv = ggml_ext_slice(ctx, kqv, 0, 0, d_head);
}
}
}
#endif
if (kqv == nullptr && (flash_attn || sage_attn)) {
// LOG_VERBOSE("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 (mask != nullptr) {
@ -724,6 +766,9 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
if (can_use_flash_attn) {
kqv = build_kqv(q, k, v, mask);
if (kqv != nullptr) {
if (used_flash_attn != nullptr) {
*used_flash_attn = true;
}
kqv = ggml_view_4d(ctx,
kqv,
d_head,

13
otherarch/sdcpp/src/core/ggml_extend.h
View file

@ -154,7 +154,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
int d1 = 1,
int d2 = 1,
bool force_prec_f32 = false,
bool direct = false);
bool direct = false,
float scale = 1.f);
// w: [OC，IC, KD, 1 * 1]
// x: [N, IC, ID, IH*IW]
@ -217,10 +218,12 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
ggml_tensor* k,
ggml_tensor* v,
int64_t n_head,
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.0f);
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.0f,
bool sage_attn = false,
bool* used_flash_attn = nullptr);
ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx,
ggml_tensor* x,

19
otherarch/sdcpp/src/core/ggml_extend_backend.cpp
View file

@ -13,6 +13,7 @@
#endif
#include "core/util.h"
#include "ggml-backend-impl.h"
#include "ggml-impl.h"
#include "stable-diffusion.h"
@ -433,6 +434,24 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
}
ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
void* ptr,
size_t size,
size_t max_tensor_size) {
ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
if (buffer != nullptr && buffer->context == nullptr) {
ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
// Metal can wrap a failed mapping in a non-null buffer. Its free callback also
// dereferences the missing context, so only release the outer buffer.
buffer->iface.free_buffer = nullptr;
ggml_backend_buffer_free(buffer);
return nullptr;
}
}
return buffer;
}
bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
#ifdef SD_USE_CUDA
if (!sd_backend_is(backend, "CUDA")) {

4
otherarch/sdcpp/src/core/ggml_extend_backend.h
View file

@ -88,6 +88,10 @@ private:
bool sd_backend_is(ggml_backend_t backend, const std::string& name);
bool sd_backend_is_cpu(ggml_backend_t backend);
bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
void* ptr,
size_t size,
size_t max_tensor_size);
ggml_backend_t sd_backend_cpu_init();
bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,

76
otherarch/sdcpp/src/core/ggml_runner.cpp
View file

@ -21,11 +21,12 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* mask,
bool skip_reshape,
bool flash_attn,
float kv_scale) {
float kv_scale,
bool* used_flash_attn) {
if (ctx->attn_scale > 0.f) {
kv_scale = ctx->attn_scale;
}
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale);
return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn);
}
void GGMLRunner::alloc_params_ctx() {
@ -515,11 +516,13 @@ GGMLRunner::~GGMLRunner() {
free_params_ctx();
}
GGMLRunnerContext GGMLRunner::get_context() {
GGMLRunnerContext GGMLRunner::get_context(ggml_cgraph* graph) {
GGMLRunnerContext runner_ctx;
runner_ctx.ggml_ctx = compute_ctx;
runner_ctx.graph = graph;
runner_ctx.backend = runtime_backend;
runner_ctx.flash_attn_enabled = flash_attn_enabled;
runner_ctx.sage_attn_enabled = sage_attn_enabled;
runner_ctx.linear_scale = linear_scale;
runner_ctx.attn_scale = attn_scale;
runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
@ -531,8 +534,8 @@ GGMLRunnerContext GGMLRunner::get_context() {
runner_ctx.get_cache_tensor = [this](const std::string& name) {
return this->get_cache_tensor_by_name(name);
};
runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) {
this->cache(name, tensor);
runner_ctx.cache_tensor = [this, graph](const std::string& name, ggml_tensor* tensor) {
this->cache(name, tensor, graph);
};
runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) {
this->set_backend_tensor_data(tensor, data);
@ -574,7 +577,7 @@ ggml_tensor* GGMLRunner::to_backend(ggml_tensor* tensor) {
}
}
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph) {
if (tensor != nullptr && tensor->view_src != nullptr) {
tensor = ggml_cont(compute_ctx, tensor);
}
@ -582,6 +585,10 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor) {
ggml_set_output(tensor);
}
cache_.stage(name, tensor);
if (graph != nullptr && tensor != nullptr) {
// Schedule the cache output here so its source can be reused before graph end.
ggml_build_forward_expand(graph, tensor);
}
}
std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
@ -589,6 +596,7 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
bool auto_runner_end,
bool no_return,
const std::function<bool()>& read_outputs) {
last_compute_status_ = GGML_STATUS_FAILED;
if (graph_active_) {
LOG_ERROR("%s does not support reentrant graph execution", get_desc().c_str());
return std::nullopt;
@ -612,7 +620,9 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
GGMLRunner& runner;
const bool& success;
~GraphEndGuard() {
runner.workspace_.segment_end();
if (!runner.workspace_.segment_end()) {
runner.last_compute_status_ = GGML_STATUS_FAILED;
}
runner.cache_.graph_end(false);
runner.cut_cache_.clear();
runner.free_compute_ctx();
@ -640,7 +650,12 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
std::optional<sd::Tensor<float>> output;
try {
output = execute_graph(graph, n_threads, no_return, read_outputs);
} catch (const std::bad_alloc&) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
LOG_ERROR("%s graph allocation failed", get_desc().c_str());
return std::nullopt;
} catch (const std::exception& error) {
last_compute_status_ = GGML_STATUS_FAILED;
LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(),
ggml_backend_name(runtime_backend), error.what());
return std::nullopt;
@ -648,6 +663,7 @@ std::optional<sd::Tensor<float>> GGMLRunner::compute(get_graph_cb_t get_graph,
success = output.has_value();
if (success) {
cache_.graph_end(true);
last_compute_status_ = GGML_STATUS_SUCCESS;
}
return output;
}
@ -765,6 +781,7 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) {
}
workspace_.synchronize();
if (status != GGML_STATUS_SUCCESS) {
last_compute_status_ = status;
LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
return false;
}
@ -817,13 +834,23 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
const auto& cached_plan = resolve_graph_cut_plan(graph);
const auto full_measurement = measure(graph, cached_plan.compute_buffer_size);
if (full_measurement.buffers.empty()) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return std::nullopt;
}
auto fits_monolithic = [&]() {
// Planning headroom absorbs allocation estimate drift; execution keeps the normal limits.
constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL;
auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph));
for (auto& request : requests) {
request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom);
}
return fits(requests, params);
};
auto manager = residency_manager.lock();
const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) &&
manager != nullptr && manager->segmented_compute_enabled() &&
cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 &&
!fits(memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)), params);
!fits_monolithic();
ggml_graph_cut::Plan monolithic_plan;
if (!segmented) {
monolithic_plan.segments.emplace_back();
@ -887,7 +914,9 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
SegmentGraphBindings& bindings;
ggml_context* context;
~SegmentCleanup() {
runner.workspace_.segment_end();
if (!runner.workspace_.segment_end()) {
runner.last_compute_status_ = GGML_STATUS_FAILED;
}
bindings.restore();
weights.segment_end();
ggml_free(context);
@ -897,6 +926,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement;
if (!workspace_.prepare(measurement)) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return fail_segment("workspace preparation");
}
const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment);
@ -911,7 +941,11 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
sync_runtime_residency();
requests = memory_requests(measurement.buffers, new_cache_bytes);
}
return weights.ensure_segment_capacity(index, requests);
const bool ready = weights.ensure_segment_capacity(index, requests);
if (!ready && manager != nullptr) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
}
return ready;
};
if (!weights.segment_start(index, ensure_capacity)) {
return fail_segment("weight preparation");
@ -920,12 +954,17 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
if (!workspace_.measurement_matches(segment_graph, measurement)) {
measurement = measure(segment_graph, segment.compute_buffer_size);
}
if (!workspace_.prepare(measurement) || !ensure_capacity()) {
if (!workspace_.prepare(measurement)) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return fail_segment("workspace preparation");
}
if (!ensure_capacity()) {
return fail_segment("workspace capacity check");
}
if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) {
pin_multi_device_nodes(scheduler, current);
})) {
last_compute_status_ = GGML_STATUS_ALLOC_FAILED;
return fail_segment("workspace allocation");
}
for (const auto& size : measurement.buffers) {
@ -944,10 +983,16 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
}
LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(),
index + 1, plan.segments.size(), segment.group_name.c_str());
if (!execute_segment(segment_graph, n_threads) ||
!cache_.capture(segment_graph) ||
!cut_cache_.capture(graph, segment, get_desc().c_str())) {
return fail_segment("execution or output caching");
if (!execute_segment(segment_graph, n_threads)) {
return fail_segment("execution");
}
auto cache_status = cache_.capture(segment_graph);
if (cache_status == GGML_STATUS_SUCCESS) {
cache_status = cut_cache_.capture(graph, segment, get_desc().c_str());
}
if (cache_status != GGML_STATUS_SUCCESS) {
last_compute_status_ = cache_status;
return fail_segment("output caching");
}
sync_runtime_residency();
if (last) {
@ -963,6 +1008,7 @@ std::optional<Tensor<float>> GGMLRunner::execute_graph(ggml_cgraph* graph, int n
}
}
if (!workspace_.segment_end()) {
last_compute_status_ = GGML_STATUS_FAILED;
return fail_segment("workspace synchronization");
}
// Final outputs and their callbacks may still be views of consumed cuts.

36
otherarch/sdcpp/src/core/ggml_runner.h
View file

@ -67,7 +67,9 @@ struct WeightAdapter {
struct GGMLRunnerContext {
ggml_backend_t backend = nullptr;
ggml_context* ggml_ctx = nullptr;
ggml_cgraph* graph = nullptr;
bool flash_attn_enabled = false;
bool sage_attn_enabled = false;
float linear_scale = 0.f;
float attn_scale = 0.f;
bool conv2d_direct_enabled = false;
@ -101,6 +103,12 @@ struct GGMLRunnerContext {
return get_cache_tensor(name);
}
void expand_graph(ggml_tensor* tensor) const {
if (graph != nullptr && tensor != nullptr) {
ggml_build_forward_expand(graph, tensor);
}
}
void persist_cache_tensor(const std::string& name, ggml_tensor* tensor) const {
if (!cache_tensor || tensor == nullptr) {
return;
@ -121,15 +129,17 @@ ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx,
ggml_tensor* k,
ggml_tensor* v,
int64_t n_head,
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.f);
ggml_tensor* mask = nullptr,
bool skip_reshape = false,
bool flash_attn = false,
float kv_scale = 1.f,
bool* used_flash_attn = nullptr);
struct GGMLRunner {
private:
std::map<ggml_backend_t, size_t> logged_compute_bytes_;
size_t logged_segment_count_ = 0;
size_t logged_segment_count_ = 0;
ggml_status last_compute_status_ = GGML_STATUS_SUCCESS;
sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes);
std::vector<DeviceMemoryRequest> memory_requests(const std::vector<sd::BackendBufferSize>& sizes,
@ -176,6 +186,7 @@ protected:
const std::string final_result_name = "ggml_runner_final_result_tensor";
bool flash_attn_enabled = false;
bool sage_attn_enabled = false;
float linear_scale = 0.f;
float attn_scale = 0.f;
bool conv2d_direct_enabled = false;
@ -286,7 +297,8 @@ public:
virtual ~GGMLRunner();
virtual GGMLRunnerContext get_context();
// Binding a graph schedules cache outputs at registration instead of graph end.
virtual GGMLRunnerContext get_context(ggml_cgraph* graph = nullptr);
void reset_compute_ctx();
@ -321,7 +333,7 @@ public:
ggml_tensor* to_backend(ggml_tensor* tensor);
void cache(const std::string name, ggml_tensor* tensor);
void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr);
ggml_tensor* get_cache_tensor_by_name(const std::string& name) {
return cache_.get(name);
@ -333,10 +345,20 @@ public:
bool no_return = false,
const std::function<bool()>& read_outputs = {});
ggml_status last_compute_status() const { return last_compute_status_; }
void set_flash_attention_enabled(bool enabled) {
flash_attn_enabled = enabled;
}
void set_sage_attention_enabled(bool enabled) {
if (sage_attn_enabled != enabled) {
free_cache_ctx_and_buffer();
graph_cut_plan_cache_.graph_cut_plans.clear();
sage_attn_enabled = enabled;
}
}
void set_scale_overrides(float linear_scale, float attn_scale) {
this->linear_scale = linear_scale;
this->attn_scale = attn_scale;

30
otherarch/sdcpp/src/core/runner_cache.cpp
View file

@ -26,10 +26,13 @@ namespace sd {
std::unique_ptr<CachedTensor> CachedTensor::copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source) {
ggml_tensor* source,
ggml_status& status) {
status = GGML_STATUS_FAILED;
if (ggml_graph_cut::tensor_buffer(source) == nullptr) {
return nullptr;
}
status = GGML_STATUS_ALLOC_FAILED;
auto entry = std::make_unique<CachedTensor>();
entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true});
if (entry->context == nullptr) {
@ -50,6 +53,7 @@ namespace sd {
} else {
ggml_backend_tensor_copy(source, entry->tensor);
}
status = GGML_STATUS_SUCCESS;
return entry;
}
@ -106,9 +110,9 @@ namespace sd {
return pending > SIZE_MAX - committed ? SIZE_MAX : committed + pending;
}
bool RunnerCache::capture(ggml_cgraph* graph) {
ggml_status RunnerCache::capture(ggml_cgraph* graph) {
if (outputs_.empty()) {
return true;
return GGML_STATUS_SUCCESS;
}
const auto tensors = cache_graph_tensors(graph);
for (const auto& output : outputs_) {
@ -116,14 +120,15 @@ namespace sd {
continue;
}
GGML_ASSERT(ggml_is_contiguous(output.second));
auto entry = CachedTensor::copy(backend_, output.first, output.second);
ggml_status status;
auto entry = CachedTensor::copy(backend_, output.first, output.second, status);
if (entry == nullptr) {
return false;
return status;
}
pending_[output.first] = std::move(entry);
}
ggml_backend_synchronize(backend_);
return true;
return GGML_STATUS_SUCCESS;
}
void RunnerCache::graph_end(bool success) {
@ -180,9 +185,9 @@ namespace sd {
}
}
bool GraphCutTensorCache::capture(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment,
const char* log_desc) {
ggml_status GraphCutTensorCache::capture(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment,
const char* log_desc) {
size_t copied_bytes = 0;
size_t copied_count = 0;
for (int index : segment.output_node_indices) {
@ -191,10 +196,11 @@ namespace sd {
!segment.future_cut_names.count(output->name)) {
continue;
}
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output));
ggml_status status;
auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status);
if (entry == nullptr) {
LOG_ERROR("%s failed to capture graph cut tensor: %s", log_desc, output->name);
return false;
return status;
}
const size_t size = ggml_backend_buffer_get_size(entry->buffer);
copied_bytes = size > SIZE_MAX - copied_bytes ? SIZE_MAX : copied_bytes + size;
@ -206,6 +212,6 @@ namespace sd {
LOG_DEBUG("%s graph cut cache added %6.2f MB (%zu tensors)",
log_desc, copied_bytes / (1024.f * 1024.f), copied_count);
}
return true;
return GGML_STATUS_SUCCESS;
}
}

8
otherarch/sdcpp/src/core/runner_cache.h
View file

@ -20,7 +20,8 @@ namespace sd {
~CachedTensor();
static std::unique_ptr<CachedTensor> copy(ggml_backend_t backend,
const std::string& name,
ggml_tensor* source);
ggml_tensor* source,
ggml_status& status);
};
using CachedTensors = std::map<std::string, std::unique_ptr<CachedTensor>>;
@ -41,7 +42,8 @@ namespace sd {
const std::map<std::string, ggml_tensor*>& outputs() const { return outputs_; }
size_t pending_bytes(ggml_cgraph* graph) const;
size_t resident_bytes(ggml_backend_dev_t device) const;
bool capture(ggml_cgraph* graph);
bool empty() const { return committed_.empty(); }
ggml_status capture(ggml_cgraph* graph);
void graph_end(bool success);
void clear();
};
@ -57,7 +59,7 @@ namespace sd {
size_t resident_bytes(ggml_backend_dev_t device) const;
size_t estimate_output_bytes(ggml_cgraph* graph,
const ggml_graph_cut::Segment& segment) const;
bool capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc);
void prune(const std::unordered_set<std::string>& keep_names);
void clear() { tensors_.clear(); }
};

1
otherarch/sdcpp/src/detailer.cpp
View file

@ -970,6 +970,7 @@ bool adetail_image(adetailer_ctx_t* context,
generation.pm_params = {};
generation.pulid_params = {};
generation.hires.enabled = false;
generation.image_preprocess = {};
if (params.steps > 0) {
generation.sample_params.sample_steps = params.steps;
generation.sample_params.custom_sigmas = nullptr;

19
otherarch/sdcpp/src/model.h
View file

@ -60,6 +60,7 @@ enum SDVersion {
VERSION_KREA2,
VERSION_MAGE_FLOW,
VERSION_SENSENOVA_U1_5,
VERSION_LLADA_IMAGE,
VERSION_ESRGAN,
VERSION_COUNT,
};
@ -173,6 +174,13 @@ static inline bool sd_version_is_z_image(SDVersion version) {
return false;
}
static inline bool sd_version_is_llada_image(SDVersion version) {
if (version == VERSION_LLADA_IMAGE) {
return true;
}
return false;
}
static inline bool sd_version_is_boogu_image(SDVersion version) {
if (version == VERSION_BOOGU_IMAGE) {
return true;
@ -244,6 +252,14 @@ static inline bool sd_version_is_sensenova_u1(SDVersion version) {
return version == VERSION_SENSENOVA_U1_5;
}
static inline bool sd_version_supports_video_generation(SDVersion version) {
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version);
}
static inline bool sd_version_supports_image_generation(SDVersion version) {
return !sd_version_supports_video_generation(version);
}
static inline bool sd_version_uses_flux_vae(SDVersion version) {
if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
return true;
@ -252,7 +268,7 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) {
}
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) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) {
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version)) {
return true;
}
return false;
@ -293,6 +309,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
version == VERSION_HIDREAM_O1 ||
sd_version_is_anima(version) ||
sd_version_is_z_image(version) ||
sd_version_is_llada_image(version) ||
sd_version_is_boogu_image(version) ||
sd_version_is_ernie_image(version) ||
sd_version_is_lens(version) ||

21
otherarch/sdcpp/src/model/adapter/lora.hpp
View file

@ -163,6 +163,27 @@ struct LoraModel : public GGMLRunner {
lora_tensors = std::move(new_lora_tensors);
}
std::unordered_map<std::string, ggml_tensor*> new_lora_tensors;
for (const auto& [old_name, tensor] : lora_tensors) {
std::string new_name = old_name;
if (starts_with(old_name, "lora.model.diffusion_model.transformer_blocks.")) {
// Qwen Image 2.1 stores the gate before the projection in fused MLP weights.
for (const auto& suffix : {std::string(".img_mlp.gate_layer.weight."), std::string(".img_mlp.proj.weight.")}) {
size_t pos = old_name.find(suffix);
if (pos == std::string::npos) {
continue;
}
std::string fused_name = old_name.substr(5, pos - 5) + ".img_mlp.gate_up.weight";
if (model_tensor_names.find(fused_name) != model_tensor_names.end()) {
new_name = "lora." + fused_name + (suffix == ".img_mlp.proj.weight." ? ".1." : ".") + old_name.substr(pos + suffix.size());
}
break;
}
}
new_lora_tensors[new_name] = tensor;
}
lora_tensors = std::move(new_lora_tensors);
}
ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {

15
otherarch/sdcpp/src/model/common/ggml_block.hpp
View file

@ -839,21 +839,30 @@ class RMSNorm : public UnaryBlock {
protected:
int64_t hidden_size;
float eps;
bool elementwise_affine;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
this->prefix = prefix;
this->prefix = prefix;
if (!elementwise_affine) {
return;
}
enum ggml_type wtype = GGML_TYPE_F32;
params["weight"] = ggml_new_tensor_1d(ctx, wtype, hidden_size);
}
public:
RMSNorm(int64_t hidden_size,
float eps = 1e-06f)
float eps = 1e-06f,
bool elementwise_affine = true)
: hidden_size(hidden_size),
eps(eps) {}
eps(eps),
elementwise_affine(elementwise_affine) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
if (!elementwise_affine) {
return ggml_rms_norm(ctx->ggml_ctx, x, eps);
}
ggml_tensor* w = params["weight"];
if (ctx->weight_adapter) {
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");

685
otherarch/sdcpp/src/model/common/rope.hpp
View file

@ -5,6 +5,7 @@
#include <cassert>
#include <cmath>
#include <set>
#include <utility>
#include <vector>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
@ -16,6 +17,45 @@ namespace Rope {
ErnieImage,
};
struct SpatialRegion {
size_t begin;
size_t count;
float height_period;
float width_period;
int height_axis = 1;
int width_axis = 2;
};
struct PositionLayout {
// Token ranges are relative to one batch item.
std::vector<SpatialRegion> images;
size_t token_count = 0;
void append_tokens(size_t count) {
token_count += count;
}
void append_image(int height, int width, int frames = 1, float height_step = 1.f, float width_step = 1.f) {
size_t count = static_cast<size_t>(height) * width * frames;
images.push_back({token_count, count, height * height_step, width * width_step});
append_tokens(count);
}
};
struct Frequency {
size_t axis;
float omega;
};
struct Embedding {
std::vector<float> values;
std::vector<std::vector<float>> ids;
PositionLayout positions;
std::vector<Frequency> frequencies;
EmbedNDLayout layout = EmbedNDLayout::Matrix;
int batch_size = 1;
};
enum class RefIndexMode {
FIXED,
INCREASE,
@ -56,40 +96,25 @@ namespace Rope {
return flat_vec;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
int dim,
float theta,
const std::vector<int>& axis_wrap_dims = {}) {
__STATIC_INLINE__ std::vector<float> rope_frequencies(int dim, float theta) {
assert(dim % 2 == 0);
int half_dim = dim / 2;
int half_dim = dim / 2;
std::vector<float> scale = linspace(0.f, (dim * 1.f - 2) / dim, half_dim);
std::vector<float> omega(half_dim);
for (int i = 0; i < half_dim; ++i) {
omega[i] = 1.0f / ::powf(1.f * theta, scale[i]);
}
return omega;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
const std::vector<float>& omega) {
int half_dim = static_cast<int>(omega.size());
size_t pos_size = pos.size();
std::vector<std::vector<float>> out(pos_size, std::vector<float>(half_dim));
for (size_t i = 0; i < pos_size; ++i) {
for (size_t j = 0; j < half_dim; ++j) {
float angle = pos[i] * omega[j];
if (!axis_wrap_dims.empty()) {
size_t wrap_size = axis_wrap_dims.size();
// mod batch size since we only store this for one item in the batch
size_t wrap_idx = wrap_size > 0 ? (i % wrap_size) : 0;
int wrap_dim = axis_wrap_dims[wrap_idx];
if (wrap_dim > 0) {
constexpr float TWO_PI = 6.28318530717958647692f;
float cycles = omega[j] * wrap_dim / TWO_PI;
// closest periodic harmonic, necessary to ensure things neatly tile
// without this round, things don't tile at the boundaries and you end up
// with the model knowing what is "center"
float rounded = std::round(cycles);
angle = pos[i] * TWO_PI * rounded / wrap_dim;
}
}
out[i][j] = angle;
}
@ -108,6 +133,12 @@ namespace Rope {
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> rope(const std::vector<float>& pos,
int dim,
float theta) {
return rope(pos, rope_frequencies(dim, theta));
}
// Generate IDs for image patches and text
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_txt_ids(int bs, int context_len, int axes_dim_num, std::set<int> arange_dims) {
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
@ -136,12 +167,16 @@ namespace Rope {
int patch_size,
int bs,
int axes_dim_num,
int index = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false) {
int index = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false,
PositionLayout* layout = nullptr) {
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
if (layout) {
layout->append_image(h_len, w_len);
}
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
int h_start = h_offset;
@ -192,8 +227,8 @@ namespace Rope {
int bs,
const std::vector<float>& axis_thetas,
const std::vector<int>& axes_dim,
const std::vector<std::vector<int>>& wrap_dims = {},
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
EmbedNDLayout layout = EmbedNDLayout::Matrix,
std::vector<Frequency>* frequencies = nullptr) {
std::vector<std::vector<float>> trans_ids = transpose(ids);
size_t pos_len = ids.size() / bs;
size_t num_axes = axes_dim.size();
@ -205,19 +240,25 @@ namespace Rope {
for (int d : axes_dim)
emb_dim += d / 2;
if (frequencies) {
frequencies->clear();
frequencies->reserve(emb_dim);
}
std::vector<std::vector<float>> emb(bs * pos_len, std::vector<float>(emb_dim * 2 * 2, 0.0));
size_t offset = 0;
for (size_t i = 0; i < num_axes; ++i) {
std::vector<int> axis_wrap_dims;
if (!wrap_dims.empty() && i < (int)wrap_dims.size()) {
axis_wrap_dims = wrap_dims[i];
}
float axis_theta = 10000.0f;
if (!axis_thetas.empty()) {
axis_theta = axis_thetas[std::min(i, axis_thetas.size() - 1)];
}
auto omega = rope_frequencies(axes_dim[i], axis_theta);
if (frequencies) {
for (float frequency : omega) {
frequencies->push_back({i, frequency});
}
}
std::vector<std::vector<float>> rope_emb =
rope(trans_ids[i], axes_dim[i], axis_theta, axis_wrap_dims); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
rope(trans_ids[i], omega); // [bs*pos_len, axes_dim[i]/2 * 2 * 2]
for (int b = 0; b < bs; ++b) {
for (int j = 0; j < pos_len; ++j) {
for (int k = 0; k < rope_emb[0].size(); ++k) {
@ -253,10 +294,10 @@ namespace Rope {
int bs,
float theta,
const std::vector<int>& axes_dim,
const std::vector<std::vector<int>>& wrap_dims = {},
EmbedNDLayout layout = EmbedNDLayout::Matrix) {
EmbedNDLayout layout = EmbedNDLayout::Matrix,
std::vector<Frequency>* frequencies = nullptr) {
std::vector<float> axis_thetas(axes_dim.size(), theta);
return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
return embed_nd(ids, bs, axis_thetas, axes_dim, layout, frequencies);
}
__STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
@ -264,7 +305,7 @@ namespace Rope {
float theta,
int head_dim,
const std::vector<int>& mrope_section,
const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
std::vector<Frequency>* frequencies = nullptr) {
GGML_ASSERT(bs > 0);
GGML_ASSERT(head_dim % 2 == 0);
GGML_ASSERT(mrope_section.size() >= 3);
@ -273,20 +314,26 @@ namespace Rope {
size_t pos_len = ids.size() / bs;
int half_dim = head_dim / 2;
auto omega = rope_frequencies(head_dim, theta);
if (frequencies) {
frequencies->clear();
for (float frequency : omega) {
frequencies->push_back({0, frequency});
}
}
std::vector<std::vector<std::vector<float>>> axis_embs;
axis_embs.reserve(3);
for (int axis = 0; axis < 3; ++axis) {
std::vector<int> axis_wrap;
if (axis < static_cast<int>(axis_wrap_dims.size())) {
axis_wrap = axis_wrap_dims[axis];
}
axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
axis_embs.push_back(rope(trans_ids[axis], omega));
}
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) {
if (frequencies) {
(*frequencies)[freq_idx].axis = axis;
}
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];
@ -298,13 +345,13 @@ namespace Rope {
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) {
__STATIC_INLINE__ Embedding 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;
@ -318,6 +365,10 @@ namespace Rope {
w_ntk = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
}
Embedding result;
result.positions.append_image(height, width, 1,
height > 1 ? scale / (height - 1) : 1.f,
width > 1 ? scale / (width - 1) : 1.f);
std::vector<float> x_pos;
std::vector<float> y_pos;
x_pos.reserve(static_cast<size_t>(height) * width);
@ -326,13 +377,20 @@ namespace Rope {
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);
result.ids.push_back({0.f, y, x});
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);
auto x_freq = rope_frequencies(dim_axis, theta * w_ntk);
auto y_freq = rope_frequencies(dim_axis, theta * h_ntk);
auto x_emb = rope(x_pos, x_freq);
auto y_emb = rope(y_pos, y_freq);
for (int i = 0; i < axis_half_dim; ++i) {
result.frequencies.push_back({2, x_freq[i]});
result.frequencies.push_back({1, y_freq[i]});
}
std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
for (int pos = 0; pos < height * width; ++pos) {
@ -348,7 +406,8 @@ namespace Rope {
}
}
}
return out;
result.values = std::move(out);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
@ -359,7 +418,8 @@ namespace Rope {
RefIndexMode ref_index_mode,
float ref_index_scale,
bool scale_rope,
int base_offset = 0) {
int base_offset = 0,
PositionLayout* layout = nullptr) {
std::vector<std::vector<float>> ids;
int curr_h_offset = 0;
int curr_w_offset = 0;
@ -386,7 +446,8 @@ namespace Rope {
static_cast<int>(index * ref_index_scale),
h_offset + base_offset,
w_offset + base_offset,
scale_rope);
scale_rope,
layout);
ids = concat_ids(ids, ref_ids, bs);
if (ref_index_mode == RefIndexMode::INCREASE) {
@ -409,88 +470,53 @@ namespace Rope {
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
float ref_index_scale,
bool is_longcat) {
bool is_longcat,
PositionLayout* layout = nullptr) {
if (layout) {
layout->append_tokens(context_len);
}
int x_index = is_longcat ? 1 : 0;
auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
int offset = is_longcat ? context_len : 0;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset, false, layout);
auto ids = concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset, layout);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
}
// Generate flux positional embeddings
__STATIC_INLINE__ std::vector<float> gen_flux_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
float ref_index_scale,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim,
bool is_longcat) {
std::vector<std::vector<float>> ids = gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
context_len,
txt_arange_dims,
ref_latents,
ref_index_mode,
ref_index_scale,
is_longcat);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len; // text first
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
cursor += img_tokens;
// reference latents
for (ggml_tensor* ref : ref_latents) {
if (ref == nullptr) {
continue;
}
int ref_h = static_cast<int>(ref->ne[1]);
int ref_w = static_cast<int>(ref->ne[0]);
int ref_h_l = (ref_h + (patch_size / 2)) / patch_size;
int ref_w_l = (ref_w + (patch_size / 2)) / patch_size;
size_t ref_tokens = static_cast<size_t>(ref_h_l) * static_cast<size_t>(ref_w_l);
for (size_t token_i = 0; token_i < ref_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = ref_h_l;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = ref_w_l;
}
}
cursor += ref_tokens;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
__STATIC_INLINE__ Embedding gen_flux_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
std::set<int> txt_arange_dims,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
float ref_index_scale,
int theta,
const std::vector<int>& axes_dim,
bool is_longcat) {
Embedding result;
result.batch_size = bs;
result.ids = gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
context_len,
txt_arange_dims,
ref_latents,
ref_index_mode,
ref_index_scale,
is_longcat, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
@ -500,14 +526,18 @@ namespace Rope {
int ph,
int pw,
int bs,
int t_offset = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false) {
int t_offset = 0,
int h_offset = 0,
int w_offset = 0,
bool scale_rope = false,
PositionLayout* layout = nullptr) {
int t_len = (t + (pt / 2)) / pt;
int h_len = (h + (ph / 2)) / ph;
int w_len = (w + (pw / 2)) / pw;
if (layout) {
layout->append_image(h_len, w_len, t_len);
}
std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
if (scale_rope) {
@ -573,7 +603,11 @@ namespace Rope {
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) {
RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
if (layout) {
layout->append_tokens(context_len);
}
int h_len = (h + (patch_size / 2)) / patch_size;
int w_len = (w + (patch_size / 2)) / patch_size;
int txt_id_start = std::max(h_len, w_len) / 2;
@ -585,90 +619,49 @@ namespace Rope {
}
}
int axes_dim_num = 3;
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
auto img_ids = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true, layout);
auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
if (ref_latents.size() > 0) {
int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true, 0, layout);
ids = concat_ids(ids, refs_ids, bs);
}
return ids;
}
// Generate qwen_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
// This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
const size_t total_tokens = ids.size();
// Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
size_t cursor = context_len; // ignore text tokens
const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
cursor += img_tokens;
// For each reference image, store wrap sizes as well
for (ggml_tensor* ref : ref_latents) {
if (ref == nullptr) {
continue;
}
int ref_h = static_cast<int>(ref->ne[1]);
int ref_w = static_cast<int>(ref->ne[0]);
int ref_pad_h = (patch_size - (ref_h % patch_size)) % patch_size;
int ref_pad_w = (patch_size - (ref_w % patch_size)) % patch_size;
int ref_h_len = (ref_h + ref_pad_h) / patch_size;
int ref_w_len = (ref_w + ref_pad_w) / patch_size;
size_t ref_n_tokens = static_cast<size_t>(ref_h_len) * static_cast<size_t>(ref_w_len);
for (size_t token_i = 0; token_i < ref_n_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = ref_h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = ref_w_len;
}
}
cursor += ref_n_tokens;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
__STATIC_INLINE__ Embedding gen_qwen_image_pe(int t,
int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
int w,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
__STATIC_INLINE__ Embedding gen_mage_flow_pe(int h,
int w,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
const int axes_dim_num = static_cast<int>(axes_dim.size());
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) {
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
std::vector<float>(axes_dim_num, 0.f));
result.positions.append_image(image_h, image_w);
int h_start = -(image_h - image_h / 2);
int w_start = -(image_w - image_w / 2);
for (int b = 0; b < bs; ++b) {
@ -692,15 +685,18 @@ namespace Rope {
static_cast<int>(i + 1));
ids = concat_ids(ids, ref_ids, bs);
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
result.ids = std::move(ids);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
int w,
int bs,
int context_len,
bool scale_rope = true) {
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
bool scale_rope = true,
PositionLayout* layout = nullptr) {
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope, layout);
int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
@ -711,44 +707,37 @@ namespace Rope {
}
}
if (layout) {
layout->append_tokens(context_len);
}
return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
}
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
int w,
int bs,
int context_len,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][token_i] = h;
}
if (circular_w) {
wrap_dims[2][token_i] = w;
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
__STATIC_INLINE__ Embedding gen_lens_pe(int h,
int w,
int bs,
int context_len,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.ids = gen_lens_ids(h, w, bs, context_len, true, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
int w,
int patch_size,
int bs,
int context_len) {
int context_len,
PositionLayout* layout = nullptr) {
int h_len = h / patch_size;
int w_len = w / patch_size;
if (layout) {
layout->append_image(h_len, w_len);
}
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(3, 0.0f));
std::vector<float> h_ids = linspace<float>(0.f, static_cast<float>(h_len - 1), h_len);
std::vector<float> w_ids = linspace<float>(0.f, static_cast<float>(w_len - 1), w_len);
@ -774,39 +763,25 @@ namespace Rope {
}
}
if (layout) {
layout->append_tokens(context_len);
}
return concat_ids(img_ids_repeated, txt_ids, bs);
}
__STATIC_INLINE__ std::vector<float> gen_ernie_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int h_len = h / patch_size;
int w_len = w / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][token_i] = w_len;
}
}
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
__STATIC_INLINE__ Embedding gen_ernie_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.layout = EmbedNDLayout::ErnieImage;
result.ids = gen_ernie_image_ids(h, w, patch_size, bs, context_len, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
// Generate wan positional embeddings
@ -905,7 +880,8 @@ namespace Rope {
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode) {
RefIndexMode ref_index_mode,
PositionLayout* layout = nullptr) {
SD_UNUSED(ref_index_mode);
int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
@ -913,11 +889,17 @@ namespace Rope {
txt_ids[i][0] = (i % padded_context_len) + 1.f;
}
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3;
int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index);
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (layout) {
layout->append_tokens(img_pad_len);
}
if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs);
@ -929,43 +911,160 @@ namespace Rope {
return ids;
}
// Generate z_image positional embeddings
__STATIC_INLINE__ std::vector<float> gen_z_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
int theta,
bool circular_h,
bool circular_w,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
std::vector<std::vector<int>> wrap_dims;
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
int pad_h = (patch_size - (h % patch_size)) % patch_size;
int pad_w = (patch_size - (w % patch_size)) % patch_size;
int h_len = (h + pad_h) / patch_size;
int w_len = (w + pad_w) / patch_size;
if (h_len > 0 && w_len > 0) {
size_t pos_len = ids.size() / bs;
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
size_t cursor = context_len + bound_mod(context_len, seq_multi_of); // skip text (and its padding)
size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
if (circular_h) {
wrap_dims[1][cursor + token_i] = h_len;
}
if (circular_w) {
wrap_dims[2][cursor + token_i] = w_len;
}
}
// LLaDA-Image shares Lumina2/z_image's axes layout, but assigns position (0,0,0) to the
// padding slots of the caption stream instead of continuing the caption ramp through them.
__STATIC_INLINE__ std::vector<std::vector<float>> gen_llada_image_ids(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of,
PositionLayout* layout = nullptr) {
int context_pad_len = bound_mod(context_len, seq_multi_of);
int padded_context_len = context_len + context_pad_len;
auto txt_ids = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
for (int i = 0; i < bs * padded_context_len; i++) {
int pos = i % padded_context_len;
if (pos < context_len) {
txt_ids[i][0] = pos + 1.f;
}
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
if (layout) {
layout->append_tokens(padded_context_len);
}
int axes_dim_num = 3;
int index = padded_context_len + 1;
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, index, 0, 0, false, layout);
int img_pad_len = bound_mod(static_cast<int>(img_ids.size() / bs), seq_multi_of);
if (layout) {
layout->append_tokens(img_pad_len);
}
if (img_pad_len > 0) {
std::vector<std::vector<float>> img_pad_ids(bs * img_pad_len, std::vector<float>(3, 0.f));
img_ids = concat_ids(img_ids, img_pad_ids, bs);
}
return concat_ids(txt_ids, img_ids, bs);
}
// LLaDA-Image editing packs two caption copies (clean and noisy), the source and target
// latents anchored at their own caption's end position, and the SigVQ stream after both.
// Padding slots keep position (0,0,0), as in the text-only layout.
__STATIC_INLINE__ std::vector<std::vector<float>> gen_llada_image_edit_ids(int h,
int w,
int patch_size,
int context_len,
int sigvq_len,
int seq_multi_of,
PositionLayout* layout = nullptr) {
const int context_pad = bound_mod(context_len, seq_multi_of);
const int padded_context = context_len + context_pad;
const int h_len = (h + (patch_size / 2)) / patch_size;
const int w_len = (w + (patch_size / 2)) / patch_size;
const int image_len = h_len * w_len;
const int image_pad = bound_mod(image_len, seq_multi_of);
const int padded_image = image_len + image_pad;
const int sigvq_pad = bound_mod(sigvq_len, seq_multi_of);
std::vector<std::vector<float>> cap_ids;
std::vector<int> cap_end_positions;
int cursor = 1;
for (int copy = 0; copy < 2; ++copy) {
for (int i = 0; i < padded_context; ++i) {
std::vector<float> id(3, 0.f);
if (i < context_len) {
id[0] = static_cast<float>(cursor + i);
}
cap_ids.push_back(id);
}
cursor += context_len;
cap_end_positions.push_back(cursor);
cursor += 2;
}
if (layout) {
layout->append_tokens(cap_ids.size());
}
std::vector<std::vector<float>> img_ids;
for (int copy = 0; copy < 2; ++copy) {
auto ids = gen_flux_img_ids(h, w, patch_size, 1, 3, cap_end_positions[copy], 0, 0, false, layout);
img_ids.insert(img_ids.end(), ids.begin(), ids.end());
img_ids.insert(img_ids.end(), image_pad, std::vector<float>(3, 0.f));
if (layout) {
layout->append_tokens(image_pad);
}
}
const int sigvq_start = static_cast<int>(cap_ids.size() + img_ids.size()) + 1;
std::vector<std::vector<float>> sigvq_ids;
for (int i = 0; i < sigvq_len + sigvq_pad; ++i) {
std::vector<float> id(3, 0.f);
if (i < sigvq_len) {
id[0] = static_cast<float>(sigvq_start + i);
}
sigvq_ids.push_back(id);
}
std::vector<std::vector<float>> ids;
ids.reserve(cap_ids.size() + img_ids.size() + sigvq_ids.size());
ids.insert(ids.end(), cap_ids.begin(), cap_ids.end());
ids.insert(ids.end(), img_ids.begin(), img_ids.end());
ids.insert(ids.end(), sigvq_ids.begin(), sigvq_ids.end());
if (layout) {
layout->append_tokens(sigvq_ids.size());
}
SD_UNUSED(padded_image);
return ids;
}
__STATIC_INLINE__ Embedding gen_llada_image_edit_pe(int h,
int w,
int patch_size,
int context_len,
int sigvq_len,
int seq_multi_of,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = 1;
result.ids = gen_llada_image_edit_ids(h, w, patch_size, context_len, sigvq_len, seq_multi_of, &result.positions);
result.values = embed_nd(result.ids, 1, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ Embedding gen_llada_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.ids = gen_llada_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
// Generate z_image positional embeddings
__STATIC_INLINE__ Embedding gen_z_image_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
int seq_multi_of,
const std::vector<ggml_tensor*>& ref_latents,
RefIndexMode ref_index_mode,
int theta,
const std::vector<int>& axes_dim) {
Embedding result;
result.batch_size = bs;
result.ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode, &result.positions);
result.values = embed_nd(result.ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
__STATIC_INLINE__ ggml_tensor* apply_rope(ggml_context* ctx,

65
otherarch/sdcpp/src/model/common/rope_circular.hpp Normal file
View file

@ -0,0 +1,65 @@
#ifndef __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
#define __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__
#include "model/common/rope.hpp"
namespace Rope {
__STATIC_INLINE__ void apply_circular(Embedding& embedding, bool circular_x, bool circular_y) {
if (!circular_x && !circular_y) {
return;
}
GGML_ASSERT(embedding.batch_size > 0);
GGML_ASSERT(embedding.ids.size() % embedding.batch_size == 0);
size_t pos_len = embedding.ids.size() / embedding.batch_size;
size_t half_dim = embedding.frequencies.size();
GGML_ASSERT(embedding.positions.token_count == pos_len);
GGML_ASSERT(embedding.values.size() == embedding.ids.size() * half_dim * 4);
constexpr float TWO_PI = 6.28318530717958647692f;
for (const auto& region : embedding.positions.images) {
GGML_ASSERT(region.begin <= pos_len && region.count <= pos_len - region.begin);
for (size_t j = 0; j < half_dim; ++j) {
const auto& frequency = embedding.frequencies[j];
float period = 0.f;
if (circular_y && frequency.axis == static_cast<size_t>(region.height_axis)) {
period = region.height_period;
} else if (circular_x && frequency.axis == static_cast<size_t>(region.width_axis)) {
period = region.width_period;
}
if (period <= 0) {
continue;
}
// Quantize to periodic harmonics while preserving the original coordinate offsets.
float rounded = std::round(frequency.omega * period / TWO_PI);
for (int b = 0; b < embedding.batch_size; ++b) {
size_t begin = b * pos_len + region.begin;
for (size_t i = begin; i < begin + region.count; ++i) {
GGML_ASSERT(frequency.axis < embedding.ids[i].size());
float angle = embedding.ids[i][frequency.axis] * TWO_PI * rounded / period;
float cos_val = std::cos(angle);
float sin_val = std::sin(angle);
if (embedding.layout == EmbedNDLayout::ErnieImage) {
size_t cos_offset = (i * half_dim + j) * 2;
size_t sin_offset = embedding.ids.size() * half_dim * 2 + cos_offset;
embedding.values[cos_offset] = cos_val;
embedding.values[cos_offset + 1] = cos_val;
embedding.values[sin_offset] = sin_val;
embedding.values[sin_offset + 1] = sin_val;
} else {
size_t offset = (i * half_dim + j) * 4;
embedding.values[offset] = cos_val;
embedding.values[offset + 1] = -sin_val;
embedding.values[offset + 2] = sin_val;
embedding.values[offset + 3] = cos_val;
}
}
}
}
}
}
} // namespace Rope
#endif // __SD_MODEL_COMMON_ROPE_CIRCULAR_HPP__

67
otherarch/sdcpp/src/model/diffusion/anima.hpp
View file

@ -603,34 +603,37 @@ namespace Anima {
return std::pow(extrapolation_ratio, static_cast<float>(axis_dim) / static_cast<float>(axis_dim - 2));
}
static std::vector<float> gen_anima_image_pe_vec(int bs,
int h,
int w,
int patch_size,
int theta,
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
0,
{},
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
false);
static Rope::Embedding gen_anima_image_pe_vec(int bs,
int h,
int w,
int patch_size,
int theta,
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
Rope::Embedding result;
result.batch_size = bs;
result.ids = Rope::gen_flux_ids(h,
w,
patch_size,
bs,
static_cast<int>(axes_dim.size()),
0,
{},
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
false, &result.positions);
std::vector<float> axis_thetas = {
static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
static_cast<float>(theta) * calc_ntk_factor(h_extrapolation_ratio, axes_dim[1]),
static_cast<float>(theta) * calc_ntk_factor(w_extrapolation_ratio, axes_dim[2]),
};
return Rope::embed_nd(ids, bs, axis_thetas, axes_dim);
result.values = Rope::embed_nd(result.ids, bs, axis_thetas, axes_dim, result.layout, &result.frequencies);
return result;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@ -657,16 +660,16 @@ namespace Anima {
int64_t h_pad = x->ne[1] + pad_h;
int64_t w_pad = x->ne[0] + pad_w;
image_pe_vec = gen_anima_image_pe_vec(1,
static_cast<int>(h_pad),
static_cast<int>(w_pad),
static_cast<int>(config.patch_size),
config.theta,
config.axes_dim,
4.0f,
4.0f,
1.0f,
ref_latents);
image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1,
static_cast<int>(h_pad),
static_cast<int>(w_pad),
static_cast<int>(config.patch_size),
config.theta,
config.axes_dim,
4.0f,
4.0f,
1.0f,
ref_latents));
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());

42
otherarch/sdcpp/src/model/diffusion/boogu.hpp
View file

@ -720,15 +720,18 @@ namespace Boogu {
}
}
__STATIC_INLINE__ std::vector<float> gen_boogu_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids;
__STATIC_INLINE__ Rope::Embedding gen_boogu_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
auto& ids = result.ids;
ids.reserve(static_cast<size_t>(bs) * context_len);
for (int b = 0; b < bs; b++) {
for (int i = 0; i < context_len; i++) {
@ -741,15 +744,18 @@ namespace Boogu {
for (ggml_tensor* ref : ref_latents) {
int ref_h_tokens = patched_token_count(ref->ne[1], patch_size);
int ref_w_tokens = patched_token_count(ref->ne[0], patch_size);
result.positions.append_image(ref_h_tokens, ref_w_tokens);
append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens);
pe_shift += std::max(ref_h_tokens, ref_w_tokens);
}
int h_tokens = patched_token_count(h, patch_size);
int w_tokens = patched_token_count(w, patch_size);
result.positions.append_image(h_tokens, w_tokens);
append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens);
return Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
result.values = Rope::embed_nd(ids, bs, static_cast<float>(theta), axes_dim, result.layout, &result.frequencies);
return result;
}
struct BooguImageRunner : public DiffusionModelRunner {
@ -793,14 +799,14 @@ namespace Boogu {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = gen_boogu_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ref_latents,
config.theta,
config.axes_dim);
pe_vec = finish_rope_pe(gen_boogu_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ref_latents,
config.theta,
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());

5
otherarch/sdcpp/src/model/diffusion/control.hpp
View file

@ -376,7 +376,7 @@ struct ControlNet : public GGMLRunner {
hint = make_input(hint_tensor);
}
auto runner_ctx = get_context();
auto runner_ctx = get_context(gf);
auto outs = control_net.forward(&runner_ctx,
x,
@ -389,8 +389,7 @@ struct ControlNet : public GGMLRunner {
if (guided_hint_input == nullptr && !outs.empty()) {
guided_hint_output_ggml = outs[0];
ggml_set_output(guided_hint_output_ggml);
cache(guided_hint_cache_name(), guided_hint_output_ggml);
ggml_build_forward_expand(gf, guided_hint_output_ggml);
runner_ctx.persist_cache_tensor(guided_hint_cache_name(), guided_hint_output_ggml);
}
control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0);

16
otherarch/sdcpp/src/model/diffusion/ernie_image.hpp
View file

@ -415,15 +415,13 @@ namespace ErnieImage {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
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());

26
otherarch/sdcpp/src/model/diffusion/flux.hpp
View file

@ -1548,20 +1548,18 @@ namespace Flux {
} else if (version == VERSION_OVIS_IMAGE) {
txt_arange_dims = {1, 2};
}
pe_vec = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
ref_index_mode,
config.ref_index_scale,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim,
sd_version_is_longcat(version));
pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
txt_arange_dims,
ref_latents,
ref_index_mode,
config.ref_index_scale,
config.theta,
config.axes_dim,
sd_version_is_longcat(version)));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

71
otherarch/sdcpp/src/model/diffusion/ideogram4.hpp
View file

@ -149,18 +149,21 @@ namespace Ideogram4 {
return std::make_shared<Linear>(in_features, out_features, bias);
}
__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,
bool circular_x = false,
bool circular_y = false) {
__STATIC_INLINE__ Rope::Embedding 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));
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
result.positions.append_image(grid_h, grid_w);
result.ids.assign(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
std::vector<float>(3, 0.f));
auto& ids = result.ids;
for (int i = 0; i < context_len; ++i) {
ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
@ -175,29 +178,13 @@ namespace Ideogram4 {
}
}
std::vector<std::vector<int>> axis_wrap_dims(3);
if (circular_y || circular_x) {
size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
axis_wrap_dims[1].assign(total_len, 0);
axis_wrap_dims[2].assign(total_len, 0);
if (circular_y) {
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
axis_wrap_dims[1][idx] = grid_h;
}
}
if (circular_x) {
for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
axis_wrap_dims[2][idx] = grid_w;
}
}
}
return Rope::embed_interleaved_mrope(ids,
bs,
static_cast<float>(rope_theta),
head_dim,
mrope_section,
axis_wrap_dims);
result.values = Rope::embed_interleaved_mrope(ids,
bs,
static_cast<float>(rope_theta),
head_dim,
mrope_section,
&result.frequencies);
return result;
}
class Ideogram4Attention : public GGMLBlock {
@ -509,15 +496,13 @@ namespace Ideogram4 {
int64_t head_dim = config.emb_dim / config.num_heads;
auto runner_ctx = get_context();
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,
runner_ctx.circular_x_enabled,
runner_ctx.circular_y_enabled);
pe_vec = finish_rope_pe(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());

47
otherarch/sdcpp/src/model/diffusion/krea2.hpp
View file

@ -689,23 +689,28 @@ namespace Krea2 {
}
};
__STATIC_INLINE__ std::vector<float> gen_krea2_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
float theta,
const std::vector<int>& axes_dim,
const std::vector<ggml_tensor*>& ref_latents,
Rope::RefIndexMode ref_index_mode) {
__STATIC_INLINE__ Rope::Embedding gen_krea2_pe(int h,
int w,
int patch_size,
int bs,
int context_len,
float theta,
const std::vector<int>& axes_dim,
const std::vector<ggml_tensor*>& ref_latents,
Rope::RefIndexMode ref_index_mode) {
Rope::Embedding result;
result.batch_size = bs;
result.positions.append_tokens(context_len);
auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false, &result.positions);
auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
if (ref_latents.size() > 0) {
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0, &result.positions);
ids = Rope::concat_ids(ids, refs_ids, bs);
}
return Rope::embed_nd(ids, bs, theta, axes_dim);
result.ids = std::move(ids);
result.values = Rope::embed_nd(result.ids, bs, theta, axes_dim, result.layout, &result.frequencies);
return result;
}
struct Krea2Runner : public DiffusionModelRunner {
@ -749,15 +754,15 @@ namespace Krea2 {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode);
pe_vec = finish_rope_pe(gen_krea2_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
config.axes_dim,
ref_latents,
ref_image_params.ref_index_mode));
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());

14
otherarch/sdcpp/src/model/diffusion/lens.hpp
View file

@ -384,14 +384,12 @@ namespace Lens {
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
config.theta,
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());

525
otherarch/sdcpp/src/model/diffusion/llada_image.hpp Normal file
View file

@ -0,0 +1,525 @@
#ifndef __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
#define __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__
#include <algorithm>
#include <cinttypes>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
#include "core/util.h"
#include "model/common/ggml_block.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/z_image.hpp"
#include "model_loader.h"
// Ref: https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
//
// The denoiser is Lumina2/z_image's NextDiT with identical hyperparameters, so the blocks are
// reused from ZImage. Two things differ: every norm here is non-parametric (the checkpoint
// carries no norm weights at all), and latents arrive already patchified from the Flux2 VAE,
// so patch_size is 1 over 128 channels.
namespace LLaDAImage {
constexpr int LLADA_IMAGE_GRAPH_SIZE = 20480;
struct LLaDAImageConfig {
int patch_size = 1;
int64_t hidden_size = 3840;
int64_t in_channels = 128;
int64_t out_channels = 128;
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;
int64_t semantic_feat_dim = 4096;
int theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int64_t axes_dim_sum = 128;
static int64_t count_blocks(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
const std::string& block_prefix) {
int64_t count = 0;
for (const auto& [name, _] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
size_t pos = name.find(block_prefix);
if (pos == std::string::npos) {
continue;
}
auto items = split_string(name.substr(pos), '.');
if (items.size() > 1) {
count = std::max<int64_t>(count, atoi(items[1].c_str()) + 1);
}
}
return count;
}
static LLaDAImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
LLaDAImageConfig config;
int64_t detected_q_dim = 0;
int64_t detected_kv_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, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) {
config.semantic_feat_dim = tensor_storage.ne[0];
} else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) {
detected_q_dim = tensor_storage.ne[1];
} else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) {
detected_kv_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;
}
}
int64_t detected_layers = count_blocks(tensor_storage_map, prefix, "layers.");
int64_t detected_refiner = std::max(count_blocks(tensor_storage_map, prefix, "noise_refiner."),
count_blocks(tensor_storage_map, prefix, "context_refiner."));
if (detected_layers > 0) {
config.num_layers = detected_layers;
}
if (detected_refiner > 0) {
config.num_refiner_layers = detected_refiner;
}
if (detected_q_dim > 0) {
config.num_heads = detected_q_dim / config.head_dim;
}
if (detected_kv_dim > 0) {
config.num_kv_heads = detected_kv_dim / config.head_dim;
} else if (detected_q_dim > 0) {
config.num_kv_heads = config.num_heads;
}
LOG_VERBOSE("llada_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64 ", cap_feat_dim = %" PRId64 ", semantic_feat_dim = %" PRId64,
config.num_layers,
config.num_refiner_layers,
config.hidden_size,
config.num_heads,
config.num_kv_heads,
config.in_channels,
config.out_channels,
config.cap_feat_dim,
config.semantic_feat_dim);
return config;
}
};
class LLaDAImageModel : public GGMLBlock {
protected:
LLaDAImageConfig config;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
params["sigvq_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
}
std::shared_ptr<ZImage::JointTransformerBlock> make_block(bool modulation) {
return std::make_shared<ZImage::JointTransformerBlock>(0,
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,
modulation,
false,
true);
}
public:
LLaDAImageModel() = default;
LLaDAImageModel(LLaDAImageConfig 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, ZImage::ADALN_EMBED_DIM);
blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps, false);
blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
blocks["semantic_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
blocks["semantic_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
blocks["sigvq_embedder.0"] = std::make_shared<RMSNorm>(config.semantic_feat_dim, config.norm_eps, false);
blocks["sigvq_embedder.1"] = std::make_shared<Linear>(config.semantic_feat_dim, config.hidden_size);
for (int i = 0; i < config.num_refiner_layers; i++) {
blocks["noise_refiner." + std::to_string(i)] = make_block(true);
blocks["context_refiner." + std::to_string(i)] = make_block(false);
blocks["sigvq_refiner." + std::to_string(i)] = make_block(false);
}
for (int i = 0; i < config.num_layers; i++) {
blocks["layers." + std::to_string(i)] = make_block(true);
}
blocks["final_layer"] = std::make_shared<ZImage::FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
}
ggml_tensor* forward_core(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
auto final_layer = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
auto txt_pad_token = params["cap_pad_token"];
auto img_pad_token = params["x_pad_token"];
int64_t N = x->ne[2];
int64_t n_img_token = x->ne[1];
int64_t n_txt_token = context->ne[1];
// sdcpp's flow denoiser already hands over sigma * 1000, which is the range the
// reference reaches via its own t_scale, so no further scaling here.
auto t_emb = t_embedder->forward(ctx, timestep);
auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context)); // [N, n_txt_token, hidden_size]
auto img = x_embedder->forward(ctx, x); // [N, n_img_token, hidden_size]
sd::ggml_graph_cut::mark_graph_cut(txt, "llada_image.prelude", "txt");
sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.prelude", "img");
sd::ggml_graph_cut::mark_graph_cut(t_emb, "llada_image.prelude", "t_emb");
int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), ZImage::SEQ_MULTI_OF);
if (n_txt_pad_token > 0) {
auto txt_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, txt_pad_token, txt_pad_token->ne[0], n_txt_pad_token, N, 1);
txt = ggml_concat(ctx->ggml_ctx, txt, txt_pad_tokens, 1);
}
int64_t n_img_pad_token = Rope::bound_mod(static_cast<int>(n_img_token), ZImage::SEQ_MULTI_OF);
if (n_img_pad_token > 0) {
auto img_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_img_pad_token, N, 1);
img = ggml_concat(ctx->ggml_ctx, img, img_pad_tokens, 1);
}
GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]);
auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::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, "llada_image.context_refiner." + std::to_string(i), "txt");
}
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
img = block->forward(ctx, img, img_pe, nullptr, t_emb);
sd::ggml_graph_cut::mark_graph_cut(img, "llada_image.noise_refiner." + std::to_string(i), "img");
}
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.prelude", "txt_img");
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
sd::ggml_graph_cut::mark_graph_cut(txt_img, "llada_image.layers." + std::to_string(i), "txt_img");
}
txt_img = final_layer->forward(ctx, txt_img, t_emb);
return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token);
}
ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) {
int64_t n_pad = Rope::bound_mod(static_cast<int>(x->ne[1]), ZImage::SEQ_MULTI_OF);
if (n_pad == 0) {
return x;
}
auto pads = ggml_repeat_4d(ctx->ggml_ctx, pad_token, pad_token->ne[0], n_pad, x->ne[2], 1);
return ggml_concat(ctx->ggml_ctx, x, pads, 1);
}
// Editing runs one joint sequence carrying two timesteps: the caption and source latent
// are clean (t = 0) while the second caption copy and the target latent are noisy. adaLN
// is a linear map of the timestep embedding, so feeding a per-token embedding selects the
// right modulation exactly, without duplicating the modulation projections.
ggml_tensor* forward_editing(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* semantic,
ggml_tensor* source_latent,
ggml_tensor* pe) {
ggml_context* gctx = ctx->ggml_ctx;
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"]);
auto cap_embedder_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["cap_embedder.0"]);
auto cap_embedder_1 = std::dynamic_pointer_cast<Linear>(blocks["cap_embedder.1"]);
auto sigvq_embed_0 = std::dynamic_pointer_cast<RMSNorm>(blocks["sigvq_embedder.0"]);
auto sigvq_embed_1 = std::dynamic_pointer_cast<Linear>(blocks["sigvq_embedder.1"]);
auto final_layer = std::dynamic_pointer_cast<ZImage::FinalLayer>(blocks["final_layer"]);
auto t_noisy = t_embedder->forward(ctx, timestep);
auto t_clean = t_embedder->forward(ctx, ggml_scale(gctx, timestep, 0.f));
auto per_token = [&](ggml_tensor* emb, int64_t n) {
return ggml_repeat_4d(gctx, emb, emb->ne[0], n, 1, 1);
};
auto cap = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));
cap = pad_stream(ctx, cap, params["cap_pad_token"]);
int64_t cap_len = cap->ne[1];
cap = ggml_concat(gctx, cap, cap, 1);
auto src = pad_stream(ctx, x_embedder->forward(ctx, source_latent), params["x_pad_token"]);
auto tgt_embed = x_embedder->forward(ctx, x);
int64_t n_img_token = tgt_embed->ne[1];
auto tgt = pad_stream(ctx, tgt_embed, params["x_pad_token"]);
int64_t img_len = tgt->ne[1];
auto img = ggml_concat(gctx, src, tgt, 1);
ggml_tensor* sig = nullptr;
int64_t sig_len = 0;
if (semantic != nullptr) {
sig = sigvq_embed_1->forward(ctx, sigvq_embed_0->forward(ctx, semantic));
sig = pad_stream(ctx, sig, params["sigvq_pad_token"]);
sig_len = sig->ne[1];
}
GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]);
auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2);
auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2);
auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
cap = block->forward(ctx, cap, cap_pe, nullptr, nullptr);
}
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
img = block->forward(ctx, img, img_pe, nullptr, img_adaln);
}
if (sig != nullptr) {
auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]);
for (int i = 0; i < config.num_refiner_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["sigvq_refiner." + std::to_string(i)]);
sig = block->forward(ctx, sig, sig_pe, nullptr, nullptr);
}
}
auto seq = ggml_concat(gctx, cap, img, 1);
auto cap_adaln = ggml_concat(gctx, per_token(t_clean, cap_len), per_token(t_noisy, cap_len), 1);
auto seq_adaln = ggml_concat(gctx, cap_adaln, img_adaln, 1);
if (sig != nullptr) {
seq = ggml_concat(gctx, seq, sig, 1);
seq_adaln = ggml_concat(gctx, seq_adaln, per_token(t_clean, sig_len), 1);
}
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<ZImage::JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
seq = block->forward(ctx, seq, pe, nullptr, seq_adaln);
sd::ggml_graph_cut::mark_graph_cut(seq, "llada_image.layers." + std::to_string(i), "seq");
}
seq = final_layer->forward(ctx, seq, seq_adaln);
// Only the target latent is denoised; the source half of the image stream is context.
// The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of
// the latent grid that unpatchify reconstructs.
int64_t target_start = cap_len * 2 + img_len;
return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
// x: [N, C, H, W]
// timestep: [N,]
// context: [N, L, cap_feat_dim]
// pe: [L, d_head/2, 2, 2]
// return: [N, C, H, W]
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int patch_size = config.patch_size;
auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
auto out = forward_core(ctx, img, timestep, context, pe);
out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false);
// The reference pipeline negates the model output before the scheduler step.
return ggml_ext_scale(ctx->ggml_ctx, out, -1.f);
}
};
struct LLaDAImageRunner : public DiffusionModelRunner {
public:
LLaDAImageConfig config;
LLaDAImageModel llada_image;
std::vector<float> pe_vec;
LLaDAImageRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(LLaDAImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
llada_image = LLaDAImageModel(config);
llada_image.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
llada_image.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor) {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x->ne[3] == 1);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
ZImage::SEQ_MULTI_OF,
config.theta,
config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, cap_feat_dim]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim());
}
ggml_cgraph* build_edit_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& semantic_tensor,
const sd::Tensor<float>& source_tensor) {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_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* semantic = make_optional_input(semantic_tensor);
ggml_tensor* source = make_input(source_tensor);
GGML_ASSERT(x->ne[3] == 1);
pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(context->ne[1]),
semantic != nullptr ? static_cast<int>(semantic->ne[1]) : 0,
ZImage::SEQ_MULTI_OF,
config.theta,
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();
int64_t W = x->ne[0];
int64_t H = x->ne[1];
auto target = DiT::pad_and_patchify(&runner_ctx, x, config.patch_size, config.patch_size, false);
auto src = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false);
auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe);
out = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false);
out = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
const auto* extra = std::get_if<LLaDAImageDiffusionExtra>(&diffusion_params.extra);
bool has_semantic = extra != nullptr && extra->semantic != nullptr && !extra->semantic->empty();
bool has_ref_latent = diffusion_params.ref_latents != nullptr && !diffusion_params.ref_latents->empty();
if (has_semantic && !has_ref_latent) {
LOG_WARN("llada_image: SigVQ features without a reference latent are not supported; falling back to text to image");
}
if (has_ref_latent) {
const auto& source = diffusion_params.ref_latents->front();
if (source.shape() != diffusion_params.x->shape()) {
LOG_ERROR("llada_image: reference latent must match the target shape; use resize_vae_to_target=1");
return {};
}
auto get_graph = [&]() -> ggml_cgraph* {
return build_edit_graph(*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
tensor_or_empty(extra != nullptr ? extra->semantic : nullptr),
source);
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false),
diffusion_params.x->dim());
}
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context));
}
};
} // namespace LLaDAImage
#endif // __SD_MODEL_DIFFUSION_LLADA_IMAGE_HPP__

14
otherarch/sdcpp/src/model/diffusion/mage_flow.hpp
View file

@ -110,13 +110,13 @@ namespace MageFlow {
}
int batch_size = static_cast<int>(x->ne[3]);
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
config.theta,
config.axes_dim);
pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
config.theta,
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());

9
otherarch/sdcpp/src/model/diffusion/minimax_h3.hpp
View file

@ -264,6 +264,9 @@ namespace MiniMaxH3 {
for (int64_t i = 0; i < num_layers; ++i) {
auto block = std::dynamic_pointer_cast<TokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
x = block->forward(ctx, x);
sd::ggml_graph_cut::mark_graph_cut(x,
"minimax_h3.token_refiner.blocks." + std::to_string(i),
"hidden_states");
}
return std::dynamic_pointer_cast<RMSNorm>(blocks["final_norm"])->forward(ctx, x);
}
@ -527,7 +530,11 @@ namespace MiniMaxH3 {
GGML_ASSERT(context->ne[0] == config.text_dim);
auto condition_proj = std::dynamic_pointer_cast<Linear>(blocks["condition_proj"]);
auto token_refiner = std::dynamic_pointer_cast<TokenRefiner>(blocks["token_refiner"]);
return token_refiner->forward(ctx, condition_proj->forward(ctx, context));
auto projected = condition_proj->forward(ctx, context);
sd::ggml_graph_cut::mark_graph_cut(projected,
"minimax_h3.condition_proj",
"hidden_states");
return token_refiner->forward(ctx, projected);
}
ggml_tensor* time_embedding(GGMLRunnerContext* ctx,

21
otherarch/sdcpp/src/model/diffusion/minit2i.hpp
View file

@ -154,18 +154,26 @@ namespace MiniT2I {
return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), head_dim, 10000.f));
}
inline std::vector<float> make_vision_rope(int side, int head_dim) {
inline Rope::Embedding make_vision_rope(int side, int head_dim) {
GGML_ASSERT(head_dim % 4 == 0);
int dim = head_dim / 2;
int quarter = dim / 2;
int length = side * side;
Rope::Embedding result;
result.positions.append_image(side, side);
std::vector<float> out(static_cast<size_t>(length) * (head_dim / 2) * 4);
std::vector<float> freqs(quarter);
for (int i = 0; i < quarter; ++i) {
freqs[i] = 1.0f / std::pow(10000.0f, static_cast<float>(2 * i) / static_cast<float>(dim));
}
for (int axis : {1, 2}) {
for (float frequency : freqs) {
result.frequencies.push_back({static_cast<size_t>(axis), frequency});
}
}
for (int y = 0; y < side; ++y) {
for (int x = 0; x < side; ++x) {
result.ids.push_back({0.f, static_cast<float>(y), static_cast<float>(x)});
int pos = y * side + x;
size_t base = static_cast<size_t>(pos) * (head_dim / 2) * 4;
for (int i = 0; i < quarter; ++i) {
@ -182,7 +190,8 @@ namespace MiniT2I {
}
}
}
return out;
result.values = std::move(out);
return result;
}
struct SwiGLUMlp : public GGMLBlock {
@ -475,6 +484,8 @@ namespace MiniT2I {
int64_t cached_txt_len = -1;
int64_t cached_hidden_size = -1;
int64_t cached_head_dim = -1;
bool cached_circular_x = false;
bool cached_circular_y = false;
MiniT2IRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@ -521,6 +532,8 @@ namespace MiniT2I {
cached_txt_len == txt_len &&
cached_hidden_size == config.hidden_size &&
cached_head_dim == config.head_dim &&
cached_circular_x == circular_x_enabled &&
cached_circular_y == circular_y_enabled &&
cached_pos_embed != nullptr &&
cached_txt_pe != nullptr &&
cached_joint_pe != nullptr) {
@ -531,7 +544,7 @@ namespace MiniT2I {
auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast<int>(img_side), static_cast<int>(config.hidden_size));
auto txt_pe_vec = make_text_rope(static_cast<int>(txt_len), static_cast<int>(config.head_dim));
auto img_pe_vec = make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim));
auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast<int>(img_side), static_cast<int>(config.head_dim)));
auto joint_pe_vec = txt_pe_vec;
joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end());
@ -561,6 +574,8 @@ namespace MiniT2I {
cached_txt_len = txt_len;
cached_hidden_size = config.hidden_size;
cached_head_dim = config.head_dim;
cached_circular_x = circular_x_enabled;
cached_circular_y = circular_y_enabled;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,

20
otherarch/sdcpp/src/model/diffusion/model.hpp
View file

@ -7,7 +7,7 @@
#include "core/ggml_runner.h"
#include "core/tensor_ggml.hpp"
#include "model/common/rope.hpp"
#include "model/common/rope_circular.hpp"
#include "model_manager.h"
enum class RefImageResizeMode {
@ -39,6 +39,9 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
{"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
// pass_to_vlm routes the reference image to the conditioner, which is where LLaDA-Image's
// SigVQ encoder lives; it does its own half-resolution resize.
{"llada_image", {true, true, Rope::RefIndexMode::FIXED, true, true, -1, RefImageResizeMode::NONE, -1, -1, true}},
{"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
};
@ -68,6 +71,8 @@ struct AnimaDiffusionExtra {
struct QwenImage21DiffusionExtra {
const sd::Tensor<int32_t>* image_slots = nullptr;
// Nonzero IDs identify immutable prefix inputs within one sampling run.
uint64_t prefix_id = 0;
};
struct WanDiffusionExtra {
@ -131,6 +136,11 @@ struct HunyuanVideoDiffusionExtra {
const sd::Tensor<float>* timestep_r = nullptr;
};
struct LLaDAImageDiffusionExtra {
// SigVQ semantic features of the reference image; present only in editing mode.
const sd::Tensor<float>* semantic = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
@ -143,7 +153,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
MiniMaxH3DiffusionExtra,
MiniT2IDiffusionExtra,
SenseNovaU1DiffusionExtra,
HunyuanVideoDiffusionExtra>;
HunyuanVideoDiffusionExtra,
LLaDAImageDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;
@ -173,6 +184,11 @@ struct DiffusionModelRunner : public GGMLRunner {
protected:
std::string prefix;
std::vector<float> finish_rope_pe(Rope::Embedding embedding) {
Rope::apply_circular(embedding, circular_x_enabled, circular_y_enabled);
return std::move(embedding.values);
}
public:
DiffusionModelRunner(ggml_backend_t backend,
const std::string& prefix,

42
otherarch/sdcpp/src/model/diffusion/pid.hpp
View file

@ -135,13 +135,13 @@ namespace Pid {
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) {
inline Rope::Embedding 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);
}
@ -867,13 +867,13 @@ namespace Pid {
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));
pos_img_vec = finish_rope_pe(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,
@ -904,13 +904,13 @@ namespace Pid {
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));
pixel_pos_comp_vec = finish_rope_pe(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,

22
otherarch/sdcpp/src/model/diffusion/qwen_image.hpp
View file

@ -635,18 +635,16 @@ namespace Qwen {
ref_index_mode = Rope::RefIndexMode::DECREASE;
}
pe_vec = Rope::gen_qwen_image_pe(time_len,
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len,
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
ref_index_mode,
config.theta,
config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

356
otherarch/sdcpp/src/model/diffusion/qwen_image_2_1.hpp
View file

@ -68,6 +68,7 @@ namespace Qwen {
std::vector<QwenImage21Segment> segments;
std::vector<std::vector<float>> positions;
int64_t prefix_length = 0;
Rope::PositionLayout rope_layout;
static QwenImage21Layout build(int64_t text_length,
const sd::Tensor<int32_t>& image_slots,
@ -82,6 +83,7 @@ namespace Qwen {
auto [height, width] = image_shapes[index];
int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + height * width, context_start, index});
layout.rope_layout.append_image(static_cast<int>(height), static_cast<int>(width));
for (int64_t h = 0; h < height; ++h) {
for (int64_t w = 0; w < width; ++w) {
layout.positions.push_back({static_cast<float>(position),
@ -106,6 +108,7 @@ namespace Qwen {
} else {
int64_t start = static_cast<int64_t>(layout.positions.size());
layout.segments.push_back({start, start + i - begin, begin, -1});
layout.rope_layout.append_tokens(i - begin);
for (int64_t j = begin; j < i; ++j, ++position) {
float p = static_cast<float>(position);
layout.positions.push_back({p, p, p});
@ -121,6 +124,20 @@ namespace Qwen {
}
};
struct QwenImage21PrefixCache {
enum class Mode {
NONE,
STORE,
REUSE
};
Mode mode = Mode::NONE;
std::string name;
std::string cut_group;
int64_t prefix_length = 0;
ggml_type type = GGML_TYPE_F32;
bool* flash_attn_used = nullptr;
};
class QwenImage21ZeroCenterRMSNorm : public RMSNorm {
public:
using RMSNorm::RMSNorm;
@ -160,27 +177,70 @@ namespace Qwen {
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<QwenImage21Segment>& segments, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
int64_t heads = x->ne[0] / dim_head;
auto project = [&](const char* name) {
auto h = std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
return ggml_reshape_4d(ctx->ggml_ctx, h, dim_head, heads, x->ne[1], x->ne[2]);
};
auto q = project("to_q");
auto k = project("to_k");
auto v = project("to_v");
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
auto q = project("to_q");
auto k = project("to_k");
auto v = project("to_v");
q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"])->forward(ctx, k);
q = Rope::apply_rope(ctx->ggml_ctx, q, pe);
k = Rope::apply_rope(ctx->ggml_ctx, k, pe);
if (cache.mode == QwenImage21PrefixCache::Mode::STORE) {
// Preserve query-first attention evaluation while writing each layer's
// prefix before its full-sequence K/V can accumulate across layers.
ctx->expand_graph(q);
auto persist = [&](ggml_tensor* tensor, int axis, const char* name) {
auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length);
// Pack the contiguous data into wider rows so quantization blocks
// can exceed head_dim without padding or changing element order.
part = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length);
auto copy = ggml_cast(ctx->ggml_ctx, part, cache.type);
// Keep the copy in this layer's segment so graph cuts do not
// retain or recompute the full-sequence K/V in the final segment.
sd::ggml_graph_cut::mark_graph_cut(copy, cache.cut_group, name);
ctx->persist_cache_tensor(cache.name + "." + name, copy);
};
persist(k, 1, "k");
persist(v, 2, "v");
}
auto attend = [&](ggml_tensor* aq, ggml_tensor* ak, ggml_tensor* av, ggml_tensor* mask) {
bool used_flash_attn = false;
auto out = ggml_ext_attention_ext(ctx, aq, ak, av, heads, mask, true, ctx->flash_attn_enabled, 1.f, &used_flash_attn);
if (cache.flash_attn_used != nullptr) {
*cache.flash_attn_used &= used_flash_attn;
}
return out;
};
ggml_tensor* result = nullptr;
for (size_t i = 0; i < segments.size(); ++i) {
const auto& segment = segments[i];
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
auto out = ggml_ext_attention_ext(ctx, sq, sk, sv, heads, masks[i], true, ctx->flash_attn_enabled);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
auto prefix_k = ctx->load_cache_tensor(cache.name + ".k");
auto prefix_v = ctx->load_cache_tensor(cache.name + ".v");
GGML_ASSERT(prefix_k != nullptr && prefix_v != nullptr);
if (prefix_k->type != k->type) {
prefix_k = ggml_cast(ctx->ggml_ctx, prefix_k, k->type);
}
if (prefix_v->type != v->type) {
prefix_v = ggml_cast(ctx->ggml_ctx, prefix_v, v->type);
}
prefix_k = ggml_reshape_4d(ctx->ggml_ctx, prefix_k, dim_head, cache.prefix_length, heads, k->ne[3]);
prefix_v = ggml_reshape_4d(ctx->ggml_ctx, prefix_v, dim_head, heads, cache.prefix_length, v->ne[3]);
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 1);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
result = attend(q, k, v, nullptr);
} else {
for (size_t i = 0; i < segments.size(); ++i) {
const auto& segment = segments[i];
auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end);
auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end);
auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end);
auto out = attend(sq, sk, sv, masks[i]);
result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1);
}
}
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
@ -193,16 +253,16 @@ namespace Qwen {
class QwenImage21TransformerBlock : public GGMLBlock {
public:
QwenImage21TransformerBlock(const QwenImage21Config& config) {
blocks["img_norm1"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
blocks["img_norm2"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
blocks["attn"] = std::make_shared<QwenImage21Attention>(config);
blocks["img_norm1"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
blocks["img_norm2"] = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
blocks["attn"] = std::make_shared<QwenImage21Attention>(config);
if (config.fused_mlp) {
blocks["img_mlp.gate_up"] = std::make_shared<Linear>(config.hidden_size, 2 * config.intermediate_size, false);
} else {
blocks["img_mlp.proj"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, false);
blocks["img_mlp.gate_layer"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, false);
}
blocks["img_mlp.out"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, false);
blocks["img_mlp.out"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, false);
}
static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) {
@ -219,13 +279,14 @@ namespace Qwen {
return ggml_concat(ctx, prefix, target, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) {
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
h = modulate(ctx->ggml_ctx, h, modulation[0], layout.prefix_length);
h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks);
x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], layout.prefix_length, true));
h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
h = modulate(ctx->ggml_ctx, h, modulation[2], layout.prefix_length);
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector<ggml_tensor*>& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length;
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm1"])->forward(ctx, x);
h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length);
h = std::dynamic_pointer_cast<QwenImage21Attention>(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache);
x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true));
h = std::dynamic_pointer_cast<LayerNorm>(blocks["img_norm2"])->forward(ctx, x);
h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length);
ggml_tensor* gate;
auto fused = blocks.find("img_mlp.gate_up");
if (fused != blocks.end()) {
@ -237,9 +298,9 @@ namespace Qwen {
gate = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.gate_layer"])->forward(ctx, h);
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.proj"])->forward(ctx, h);
}
h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], layout.prefix_length, true));
h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate));
h = std::dynamic_pointer_cast<Linear>(blocks["img_mlp.out"])->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true));
}
};
@ -261,7 +322,7 @@ namespace Qwen {
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector<ggml_tensor*>& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector<ggml_tensor*>& masks, const QwenImage21PrefixCache& cache) {
auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0);
// Runtime flow timesteps already use the [0, 1000] scale.
time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f);
@ -269,27 +330,37 @@ namespace Qwen {
time = ggml_silu(ctx->ggml_ctx, time);
auto modulation = std::dynamic_pointer_cast<Linear>(blocks["modulation.1"])->forward(ctx, time);
auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0);
auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
ggml_tensor* joint = nullptr;
for (const auto& segment : layout.segments) {
ggml_tensor* h;
if (segment.image_index < 0) {
h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
segment.context_start + segment.end - segment.start);
} else {
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) {
joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1));
} else {
auto text = std::dynamic_pointer_cast<QwenImage21TextProjection>(blocks["txt_in"])->forward(ctx, context);
for (const auto& segment : layout.segments) {
ggml_tensor* h;
if (segment.image_index < 0) {
h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start,
segment.context_start + segment.end - segment.start);
} else {
auto image = segment.image_index == static_cast<int>(refs.size()) ? x : refs[segment.image_index];
h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1));
}
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
}
joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1);
}
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.prelude", "joint");
for (int i = 0; i < config.num_layers; ++i) {
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
joint = block->forward(ctx, joint, mod, pe, layout, masks);
sd::ggml_graph_cut::mark_graph_cut(joint, "qwen_image_2_1.transformer_blocks." + std::to_string(i), "joint");
const std::string layer = "transformer_blocks." + std::to_string(i);
auto layer_cache = cache;
layer_cache.name = cache.name + "." + std::to_string(i);
layer_cache.cut_group = "qwen_image_2_1." + layer;
auto block = std::dynamic_pointer_cast<QwenImage21TransformerBlock>(blocks[layer]);
joint = block->forward(ctx, joint, mod, pe, layout, masks, layer_cache);
sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint");
}
if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) {
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
}
joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]);
auto scale = std::dynamic_pointer_cast<Linear>(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]);
joint = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out.norm"])->forward(ctx, joint);
joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f));
@ -303,11 +374,61 @@ namespace Qwen {
QwenImage21Model model;
std::vector<float> pe_data;
std::vector<sd::Tensor<float>> mask_data;
ggml_type prefix_cache_type = GGML_TYPE_COUNT;
bool prefix_cache_enabled = true;
bool prefix_cache_disabled = false;
bool prefix_cache_auto_f32 = false;
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
static bool supports_prefix_cache_type(ggml_type type) {
if (type == GGML_TYPE_F32) {
return true;
}
const auto* traits = ggml_get_type_traits(type);
if (traits->from_float_ref == nullptr || traits->to_float == nullptr) {
return false;
}
auto cpu = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
if (cpu == nullptr) {
return false;
}
auto ctx = std::unique_ptr<ggml_context, decltype(&ggml_free)>(
ggml_init({3 * ggml_tensor_overhead(), nullptr, true}), ggml_free);
if (ctx == nullptr) {
return false;
}
// Some reference quantizers have no runtime copy support. Query the
// device through the registry so dynamically loaded CPU backends work.
auto source = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, ggml_blck_size(type));
auto encoded = ggml_cast(ctx.get(), source, type);
auto decoded = ggml_cast(ctx.get(), encoded, GGML_TYPE_F32);
return ggml_backend_dev_supports_op(cpu, encoded) && ggml_backend_dev_supports_op(cpu, decoded);
}
QwenImage21Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr, const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(QwenImage21Config::detect_from_weights(weights, prefix)),
model(config) {
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
if (key == "qwen_image_2_1_prefix_cache" && !parse_strict_bool(value, prefix_cache_enabled)) {
LOG_WARN("ignoring invalid Qwen Image 2.1 model arg '%s=%s'", key.c_str(), value.c_str());
} else if (key == "qwen_image_2_1_prefix_cache_type") {
if (value == "auto") {
prefix_cache_type = GGML_TYPE_COUNT;
continue;
}
const auto type = sd_type_to_ggml_type(str_to_sd_type(value.c_str()));
if (type == GGML_TYPE_COUNT) {
LOG_WARN("ignoring unknown Qwen Image 2.1 cache type '%s'", value.c_str());
} else if (!supports_prefix_cache_type(type)) {
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': runtime conversion to and from F32 is unavailable", value.c_str());
} else if (config.hidden_size % ggml_blck_size(type) != 0) {
LOG_WARN("ignoring Qwen Image 2.1 cache type '%s': block size %" PRId64 " does not divide hidden size %" PRId64,
value.c_str(), ggml_blck_size(type), config.hidden_size);
} else {
prefix_cache_type = type;
}
}
}
model.init(params_ctx, weights, prefix);
}
@ -317,6 +438,20 @@ namespace Qwen {
model.get_param_tensors(tensors, prefix);
}
bool has_prefix_cache(const QwenImage21PrefixCache& cache) {
for (int i = 0; i < config.num_layers; ++i) {
const auto name = cache.name + "." + std::to_string(i);
auto k = get_cache_tensor_by_name(name + ".k");
auto v = get_cache_tensor_by_name(name + ".v");
if (k == nullptr || v == nullptr || k->type != cache.type || v->type != cache.type ||
k->ne[0] != config.hidden_size || k->ne[1] != cache.prefix_length || k->ne[2] != 1 || k->ne[3] != 1 ||
v->ne[0] != config.hidden_size || v->ne[1] != cache.prefix_length || v->ne[2] != 1 || v->ne[3] != 1) {
return false;
}
}
return true;
}
sd::Tensor<float> compute(int n_threads, const DiffusionParams& inputs) override {
const auto& x = tensor_or_empty(inputs.x);
const auto& context = tensor_or_empty(inputs.context);
@ -345,38 +480,113 @@ namespace Qwen {
LOG_ERROR("%s", error.what());
return {};
}
pe_data = Rope::embed_nd(layout.positions, 1, 10000.f, config.axes_dim);
mask_data.clear();
for (const auto& segment : layout.segments) {
sd::Tensor<float> mask;
if (segment.image_index < 0) {
mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
for (int64_t q = segment.start; q < segment.end; ++q) {
for (int64_t k = q + 1; k < segment.end; ++k) {
mask[k + segment.end * (q - segment.start)] = -INFINITY;
}
if (!runner_started()) {
prefix_cache_disabled = false;
prefix_cache_auto_f32 = false;
}
QwenImage21PrefixCache cache;
if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) {
cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) +
".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled);
cache.prefix_length = layout.prefix_length;
if (prefix_cache_type != GGML_TYPE_COUNT) {
cache.type = prefix_cache_type;
} else if (!prefix_cache_auto_f32 && flash_attn_enabled && !sage_attn_enabled &&
(attn_scale <= 0.f || attn_scale == 1.f)) {
cache.type = GGML_TYPE_F16;
}
cache.mode = has_prefix_cache(cache) ? QwenImage21PrefixCache::Mode::REUSE : QwenImage21PrefixCache::Mode::STORE;
}
bool flash_attn_used = true;
auto run = [&](const QwenImage21PrefixCache& active_cache) {
flash_attn_used = true;
auto checked_cache = active_cache;
if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) {
checked_cache.flash_attn_used = &flash_attn_used;
}
const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE;
const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0);
Rope::Embedding embedding;
embedding.ids.assign(first_position, layout.positions.end());
const size_t offset = cached ? static_cast<size_t>(layout.prefix_length) : 0;
embedding.positions.token_count = embedding.ids.size();
for (auto region : layout.rope_layout.images) {
if (region.begin >= offset) {
region.begin -= offset;
embedding.positions.images.push_back(region);
}
}
mask_data.push_back(std::move(mask));
}
auto build = [&]() {
auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, layout.positions.size());
set_backend_tensor_data(pe, pe_data.data());
std::vector<ggml_tensor*> masks, ref_inputs;
for (const auto& mask : mask_data) {
masks.push_back(mask.empty() ? nullptr : make_input(mask));
embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies);
pe_data = finish_rope_pe(std::move(embedding));
mask_data.clear();
if (!cached) {
for (const auto& segment : layout.segments) {
sd::Tensor<float> mask;
if (segment.image_index < 0) {
mask = sd::Tensor<float>::zeros({segment.end, segment.end - segment.start});
for (int64_t q = segment.start; q < segment.end; ++q) {
for (int64_t k = q + 1; k < segment.end; ++k) {
mask[k + segment.end * (q - segment.start)] = -INFINITY;
}
}
}
mask_data.push_back(std::move(mask));
}
}
for (const auto& ref : refs) {
ref_inputs.push_back(make_input(ref));
}
auto ctx = get_context();
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), make_input(context),
ref_inputs, pe, layout, masks);
ggml_build_forward_expand(graph, out);
return graph;
auto build = [&]() {
auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2,
layout.positions.size() - (cached ? layout.prefix_length : 0));
set_backend_tensor_data(pe, pe_data.data());
std::vector<ggml_tensor*> masks, ref_inputs;
for (const auto& mask : mask_data) {
masks.push_back(mask.empty() ? nullptr : make_input(mask));
}
if (!cached) {
for (const auto& ref : refs) {
ref_inputs.push_back(make_input(ref));
}
}
auto ctx = get_context(graph);
auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context),
ref_inputs, pe, layout, masks, checked_cache);
if (!flash_attn_used) {
return static_cast<ggml_cgraph*>(nullptr);
}
ggml_build_forward_expand(graph, out);
return graph;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
};
return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), x.dim());
auto result = run(cache);
if (result.empty() && !flash_attn_used) {
// Casting an F16 cache back to F32 cannot recover its original values.
// Recompute the prefix before executing a graph that falls back from FA.
free_cache_ctx_and_buffer();
prefix_cache_auto_f32 = true;
cache.type = GGML_TYPE_F32;
cache.mode = QwenImage21PrefixCache::Mode::STORE;
LOG_DEBUG("Qwen Image 2.1: Flash Attention unavailable; using F32 prefix caching for this sampling run");
result = run(cache);
}
if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED &&
(cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) {
// The failed graph has ended before persistent inputs are released.
free_cache_ctx_and_buffer();
prefix_cache_disabled = true;
LOG_WARN("Qwen Image 2.1: insufficient memory for prefix caching; retrying without it for this sampling run");
return run(QwenImage21PrefixCache{});
}
if (!result.empty() && cache.mode == QwenImage21PrefixCache::Mode::STORE) {
if (!has_prefix_cache(cache)) {
free_cache_ctx_and_buffer();
prefix_cache_disabled = true;
LOG_WARN("Qwen Image 2.1: incomplete prefix cache; disabling it for this sampling run");
} else {
LOG_DEBUG("Qwen Image 2.1: cached prefix %" PRIu64 " (%" PRId64 " tokens, %s)", extra->prefix_id, layout.prefix_length, ggml_type_name(cache.type));
}
}
return result;
}
};
}

15
otherarch/sdcpp/src/model/diffusion/sensenova_u1.h
View file

@ -442,16 +442,9 @@ namespace SenseNovaU1 {
k = ggml_concat(ctx->ggml_ctx, prefix_k, k, 2);
v = ggml_concat(ctx->ggml_ctx, prefix_v, v, 2);
} else {
// Keep dedicated graph outputs alive until the runner copies them
// into its persistent cache buffer after graph execution.
auto cache_k = ggml_dup_tensor(ctx->ggml_ctx, k);
cache_k = ggml_cpy(ctx->ggml_ctx, k, cache_k);
ggml_set_output(cache_k);
auto cache_v = ggml_dup_tensor(ctx->ggml_ctx, v);
cache_v = ggml_cpy(ctx->ggml_ctx, v, cache_v);
ggml_set_output(cache_v);
ctx->persist_cache_tensor(layer_cache + ".k", cache_k);
ctx->persist_cache_tensor(layer_cache + ".v", cache_v);
ctx->expand_graph(q);
ctx->persist_cache_tensor(layer_cache + ".k", k);
ctx->persist_cache_tensor(layer_cache + ".v", v);
}
q = ggml_cont(ctx->ggml_ctx,
@ -687,7 +680,7 @@ namespace SenseNovaU1 {
ggml_set_name(attention_mask, "snu15.prefix.attention_mask");
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
auto runner_ctx = get_context();
auto runner_ctx = get_context(graph);
auto text_model = model.text_model();
auto hidden = text_model->embed(&runner_ctx, ids);
hidden = text_model->forward(&runner_ctx,

104
otherarch/sdcpp/src/model/diffusion/z_image.hpp
View file

@ -131,16 +131,30 @@ namespace ZImage {
int64_t num_heads;
int64_t num_kv_heads;
bool qk_norm;
bool split_qkv;
public:
JointAttention(int64_t hidden_size, int64_t head_dim, int64_t num_heads, int64_t num_kv_heads, bool qk_norm)
: head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) {
blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
float scale = 1.f;
blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
JointAttention(int64_t hidden_size,
int64_t head_dim,
int64_t num_heads,
int64_t num_kv_heads,
bool qk_norm,
bool norm_elementwise_affine = true,
bool split_qkv = false)
: head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm), split_qkv(split_qkv) {
float scale = 1.f;
if (split_qkv) {
blocks["to_q"] = std::make_shared<Linear>(hidden_size, num_heads * head_dim, false);
blocks["to_k"] = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
blocks["to_v"] = std::make_shared<Linear>(hidden_size, num_kv_heads * head_dim, false);
blocks["to_out.0"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
} else {
blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
}
if (qk_norm) {
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-06f, norm_elementwise_affine);
}
}
@ -151,8 +165,35 @@ namespace ZImage {
// x: [N, n_token, hidden_size]
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks[split_qkv ? "to_out.0" : "out"]);
if (split_qkv) {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
if (sd_backend_is(ctx->backend, "ROCm")) {
out_proj->set_scale(1.f / 16.f);
out_proj->set_force_prec_f32(true);
q_proj->set_force_prec_f32(true);
k_proj->set_force_prec_f32(true);
v_proj->set_force_prec_f32(true);
}
auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N);
auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
auto v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N);
if (qk_norm) {
q = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"])->forward(ctx, q);
k = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"])->forward(ctx, k);
}
auto out = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f);
return out_proj->forward(ctx, out);
}
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
if (sd_backend_is(ctx->backend, "ROCm")) {
out_proj->set_scale(1.f / 16.f);
@ -252,9 +293,12 @@ namespace ZImage {
ggml_tensor* x,
ggml_tensor* scale) {
// x: [N, L, C]
// scale: [N, C]
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
// scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing
// feeds a per-token timestep embedding so each segment carries its own modulation).
if (scale->ne[1] != x->ne[1]) {
scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C]
}
x = ggml_add(ctx, x, ggml_mul(ctx, x, scale));
return x;
}
@ -272,14 +316,16 @@ namespace ZImage {
float ffn_dim_multiplier,
float norm_eps,
bool qk_norm,
bool modulation = true)
bool modulation = true,
bool norm_elementwise_affine = true,
bool split_qkv = false)
: modulation(modulation) {
blocks["attention"] = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm);
blocks["attention"] = std::make_shared<JointAttention>(hidden_size, head_dim, num_heads, num_kv_heads, qk_norm, norm_elementwise_affine, split_qkv);
blocks["feed_forward"] = std::make_shared<FeedForward>(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps);
blocks["attention_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
blocks["ffn_norm1"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
blocks["attention_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
blocks["ffn_norm2"] = std::make_shared<RMSNorm>(hidden_size, norm_eps, norm_elementwise_affine);
if (modulation) {
blocks["adaLN_modulation.0"] = std::make_shared<Linear>(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size);
}
@ -596,18 +642,16 @@ namespace ZImage {
ref_latents.push_back(make_input(ref_latent_tensor));
}
pe_vec = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
SEQ_MULTI_OF,
ref_latents,
ref_index_mode,
config.theta,
circular_y_enabled,
circular_x_enabled,
config.axes_dim);
pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
config.patch_size,
static_cast<int>(x->ne[3]),
static_cast<int>(context->ne[1]),
SEQ_MULTI_OF,
ref_latents,
ref_index_mode,
config.theta,
config.axes_dim));
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_VERBOSE("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);

604
otherarch/sdcpp/src/model/te/llada_image_te.hpp Normal file
View file

@ -0,0 +1,604 @@
#ifndef __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
#define __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__
#include <algorithm>
#include <array>
#include <cmath>
#include "core/ggml_extend.h"
#include "core/ggml_runner.h"
#include "model/common/ggml_block.hpp"
#include "model_loader.h"
// The conditioning components LLaDA-Image puts around its LLaDA2-MoE backbone.
// Ref: LLaDAImageQueryFormerModel / LLaDAImageTextProjectionModel in
// https://github.com/inclusionAI/LLaDA-Image/blob/main/src/models/transformer_llada_image.py
//
// QueryFormer turns the LLaDA token embeddings into 256 learned queries that the pipeline
// appends to the backbone input; TextProjection maps the backbone hidden states to the
// denoiser's caption dimension. Neither uses RoPE, and every norm is parameter-free.
// Both MLPs use the tanh GELU approximation, so ggml_gelu (not ggml_gelu_erf).
//
// SigVQ is the editing-only image encoder: a 40-layer ViT whose output is quantized against a
// 16384-entry codebook, with the resulting ids embedded and projected into the semantic features
// the denoiser consumes. Its MLP uses the exact erf GELU, unlike the two above.
namespace LLaDAImageTE {
constexpr int LLADA_IMAGE_TE_GRAPH_SIZE = 16384;
struct QueryFormerConfig {
int64_t num_queries = 256;
int64_t hidden_size = 2048;
int64_t num_layers = 1;
int64_t num_heads = 16;
int64_t intermediate_size = 8192;
float norm_eps = 1e-6f;
};
struct TextProjectionConfig {
int64_t hidden_size = 2048;
int64_t intermediate_size = 8960;
int64_t num_layers = 6;
int64_t num_heads = 32;
int64_t projection_dim = 2560;
float norm_eps = 1e-6f;
};
// Cross-attention with a single fused in_proj over q (from the queries) and k/v (from the
// token embeddings). The checkpoint stores in_proj as one [3*hidden, hidden] parameter.
struct QueryAttention : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t num_heads;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
enum ggml_type wtype = get_type(prefix + "in_proj_weight", tensor_storage_map, GGML_TYPE_F32);
params["in_proj_weight"] = ggml_new_tensor_2d(ctx, wtype, hidden_size, hidden_size * 3);
params["in_proj_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size * 3);
}
public:
QueryAttention(int64_t hidden_size, int64_t num_heads)
: hidden_size(hidden_size), num_heads(num_heads) {
blocks["out_proj"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* query,
ggml_tensor* context,
ggml_tensor* mask = nullptr) {
// query: [N, num_queries, hidden_size], context: [N, n_token, hidden_size]
ggml_context* gctx = ctx->ggml_ctx;
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
auto w = params["in_proj_weight"];
auto b = params["in_proj_bias"];
auto slice_w = [&](int64_t index) {
return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size);
};
auto slice_b = [&](int64_t index) {
return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size);
};
auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0));
auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1));
auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2));
auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, num_queries, hidden_size]
return out_proj->forward(ctx, x);
}
};
struct QueryFormerBlock : public GGMLBlock {
protected:
QueryFormerConfig config;
public:
QueryFormerBlock(const QueryFormerConfig& config)
: config(config) {
blocks["norm_q"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["norm_k"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["cross_attn"] = std::make_shared<QueryAttention>(config.hidden_size, config.num_heads);
blocks["norm1"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* query,
ggml_tensor* context,
ggml_tensor* mask = nullptr) {
auto norm_q = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_k"]);
auto cross_attn = std::dynamic_pointer_cast<QueryAttention>(blocks["cross_attn"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
// The reference overwrites query_embeds with its normalized value before the
// residual add, so both residuals here are on normalized activations.
query = norm_q->forward(ctx, query);
auto ctx_n = norm_k->forward(ctx, context);
query = ggml_add(ctx->ggml_ctx, query, cross_attn->forward(ctx, query, ctx_n, mask));
query = norm1->forward(ctx, query);
auto h = fc1->forward(ctx, query);
h = ggml_gelu(ctx->ggml_ctx, h);
h = fc2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, query, h);
}
};
struct QueryFormerModel : public GGMLBlock {
protected:
QueryFormerConfig config;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["meta_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, config.num_queries);
}
public:
QueryFormerModel() = default;
QueryFormerModel(const QueryFormerConfig& config)
: config(config) {
for (int i = 0; i < config.num_layers; i++) {
blocks["query_blocks." + std::to_string(i)] = std::make_shared<QueryFormerBlock>(config);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* inputs_embeds,
ggml_tensor* mask = nullptr) {
// inputs_embeds: [N, n_token, hidden_size] -> [N, num_queries, hidden_size]
auto query = params["meta_queries"];
query = ggml_reshape_3d(ctx->ggml_ctx, query, config.hidden_size, config.num_queries, 1);
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<QueryFormerBlock>(blocks["query_blocks." + std::to_string(i)]);
query = block->forward(ctx, query, inputs_embeds, mask);
}
return query;
}
};
struct TextProjectionAttention : public GGMLBlock {
protected:
int64_t num_heads;
int64_t head_dim;
public:
TextProjectionAttention(const TextProjectionConfig& config)
: num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) {
blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["out_proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, config.norm_eps, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
ggml_context* gctx = ctx->ggml_ctx;
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["out_proj"]);
auto q_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
auto k_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
auto q = q_proj->forward(ctx, x);
auto k = k_proj->forward(ctx, x);
auto v = v_proj->forward(ctx, x);
q = ggml_reshape_4d(gctx, q, head_dim, num_heads, n_token, N);
k = ggml_reshape_4d(gctx, k, head_dim, num_heads, n_token, N);
q = q_norm->forward(ctx, q);
k = k_norm->forward(ctx, k);
q = ggml_reshape_3d(gctx, q, head_dim * num_heads, n_token, N);
k = ggml_reshape_3d(gctx, k, head_dim * num_heads, n_token, N);
auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
return out_proj->forward(ctx, out);
}
};
struct TextProjectionBlock : public GGMLBlock {
public:
TextProjectionBlock(const TextProjectionConfig& config) {
blocks["self_attn"] = std::make_shared<TextProjectionAttention>(config);
blocks["layer_norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
blocks["layer_norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps, false);
blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto self_attn = std::dynamic_pointer_cast<TextProjectionAttention>(blocks["self_attn"]);
auto layer_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm1"]);
auto layer_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["layer_norm2"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
x = ggml_add(ctx->ggml_ctx, x, self_attn->forward(ctx, layer_norm1->forward(ctx, x)));
auto h = fc1->forward(ctx, layer_norm2->forward(ctx, x));
h = ggml_gelu(ctx->ggml_ctx, h);
h = fc2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct TextProjectionModel : public GGMLBlock {
protected:
TextProjectionConfig config;
public:
TextProjectionModel() = default;
TextProjectionModel(const TextProjectionConfig& config)
: config(config) {
for (int i = 0; i < config.num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::make_shared<TextProjectionBlock>(config);
}
blocks["projector"] = std::make_shared<Linear>(config.hidden_size, config.projection_dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size] -> [N, n_token, projection_dim]
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<TextProjectionBlock>(blocks["layers." + std::to_string(i)]);
x = block->forward(ctx, x);
}
auto projector = std::dynamic_pointer_cast<Linear>(blocks["projector"]);
return projector->forward(ctx, x);
}
};
struct SigVQConfig {
int64_t image_size = 2048;
int patch_size = 16;
int64_t in_channels = 3;
int64_t hidden_size = 1536;
int64_t intermediate_size = 6144;
int64_t num_layers = 40;
int64_t num_heads = 16;
int64_t codebook_size = 16384;
int64_t codebook_embed_dim = 2048;
int64_t semantic_embed_dim = 4096;
float norm_eps = 1e-6f;
};
struct SigVQAttention : public GGMLBlock {
protected:
int64_t num_heads;
int64_t head_dim;
public:
SigVQAttention(const SigVQConfig& config)
: num_heads(config.num_heads), head_dim(config.hidden_size / config.num_heads) {
blocks["qkv"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size * 3, true);
blocks["proj"] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
ggml_context* gctx = ctx->ggml_ctx;
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
int64_t hidden_size = num_heads * head_dim;
auto qkv = qkv_proj->forward(ctx, x);
auto q = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size);
auto k = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2);
auto v = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3);
auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads);
return out_proj->forward(ctx, out);
}
};
struct SigVQBlock : public GGMLBlock {
public:
SigVQBlock(const SigVQConfig& config) {
blocks["norm1"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
blocks["norm2"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps);
blocks["attn"] = std::make_shared<SigVQAttention>(config);
blocks["mlp.fc1"] = std::make_shared<Linear>(config.hidden_size, config.intermediate_size, true);
blocks["mlp.fc2"] = std::make_shared<Linear>(config.intermediate_size, config.hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto attn = std::dynamic_pointer_cast<SigVQAttention>(blocks["attn"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.fc2"]);
x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, norm1->forward(ctx, x)));
auto h = fc1->forward(ctx, norm2->forward(ctx, x));
h = ggml_gelu_erf(ctx->ggml_ctx, h);
h = fc2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct SigVQModel : public GGMLBlock {
protected:
SigVQConfig config;
public:
SigVQModel() = default;
SigVQModel(const SigVQConfig& config)
: config(config) {
blocks["visual.patch_embed.proj"] = std::make_shared<Conv2d>(config.in_channels,
config.hidden_size,
std::make_pair(config.patch_size, config.patch_size),
std::make_pair(config.patch_size, config.patch_size));
for (int i = 0; i < config.num_layers; i++) {
blocks["visual.blocks." + std::to_string(i)] = std::make_shared<SigVQBlock>(config);
}
blocks["vqmodel.quant_conv"] = std::make_shared<Conv2d>(config.hidden_size,
config.codebook_embed_dim,
std::make_pair(1, 1));
blocks["prior_projector.net.0.proj"] = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
blocks["prior_projector.net.2"] = std::make_shared<Linear>(config.semantic_embed_dim, config.semantic_embed_dim, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["visual.embeddings.position_embedding.weight"] =
ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, (config.image_size / config.patch_size) * (config.image_size / config.patch_size));
params["vqmodel.quantize.embedding.weight"] =
ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.codebook_embed_dim, config.codebook_size);
params["prior_token_embedding.weight"] =
ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.semantic_embed_dim, config.codebook_size);
}
// Bilinear-resamples the square position-embedding grid onto the image's patch grid.
// The reference uses grid_sample(align_corners=False, padding_mode="border"); the source
// coordinate for output index j is therefore (j + 0.5) * side / out - 0.5, clamped.
ggml_tensor* resample_pos_embed(GGMLRunnerContext* ctx,
ggml_tensor* pos_idx,
ggml_tensor* pos_weight) {
auto pos_embed = params["visual.embeddings.position_embedding.weight"];
auto gathered = ggml_get_rows(ctx->ggml_ctx, pos_embed, pos_idx);
return ggml_mul(ctx->ggml_ctx, gathered, pos_weight);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
const std::vector<ggml_tensor*>& pos_idx,
const std::vector<ggml_tensor*>& pos_weight) {
// pixel_values: [N, in_channels, H, W] -> [N, grid_h * grid_w, semantic_embed_dim]
ggml_context* gctx = ctx->ggml_ctx;
auto patch_embed = std::dynamic_pointer_cast<Conv2d>(blocks["visual.patch_embed.proj"]);
auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["vqmodel.quant_conv"]);
auto proj_0 = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.0.proj"]);
auto proj_2 = std::dynamic_pointer_cast<Linear>(blocks["prior_projector.net.2"]);
auto x = patch_embed->forward(ctx, pixel_values); // [N, hidden_size, grid_h, grid_w]
int64_t grid_w = x->ne[0];
int64_t grid_h = x->ne[1];
int64_t n_token = grid_h * grid_w;
int64_t N = x->ne[3];
x = ggml_reshape_3d(gctx, x, n_token, config.hidden_size, N);
x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, hidden_size]
ggml_tensor* pos = nullptr;
for (size_t i = 0; i < pos_idx.size(); i++) {
auto corner = resample_pos_embed(ctx, pos_idx[i], pos_weight[i]);
pos = pos == nullptr ? corner : ggml_add(gctx, pos, corner);
}
x = ggml_add(gctx, x, ggml_reshape_3d(gctx, pos, config.hidden_size, n_token, N));
for (int i = 0; i < config.num_layers; i++) {
auto block = std::dynamic_pointer_cast<SigVQBlock>(blocks["visual.blocks." + std::to_string(i)]);
x = block->forward(ctx, x);
}
// quant_conv is 1x1, so run it as a per-token projection rather than reshaping to 2-D.
x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, hidden_size, n_token]
x = ggml_reshape_4d(gctx, x, n_token, 1, config.hidden_size, N);
x = quant_conv->forward(ctx, x); // [N, codebook_embed_dim, 1, n_token]
x = ggml_reshape_3d(gctx, x, n_token, config.codebook_embed_dim, N);
x = ggml_cont(gctx, ggml_permute(gctx, x, 1, 0, 2, 3)); // [N, n_token, codebook_embed_dim]
// Both sides are L2-normalized, so the nearest codebook entry by euclidean distance
// is the one with the largest dot product.
auto codebook = ggml_l2_norm(gctx, params["vqmodel.quantize.embedding.weight"], 1e-12f);
auto normed = ggml_l2_norm(gctx, x, 1e-12f);
auto logits = ggml_mul_mat(gctx, codebook, normed); // [N, n_token, codebook_size]
auto token_ids = ggml_argmax(gctx, ggml_reshape_2d(gctx, logits, config.codebook_size, n_token * N));
auto semantic = ggml_get_rows(gctx, params["prior_token_embedding.weight"], token_ids);
semantic = ggml_reshape_3d(gctx, semantic, config.semantic_embed_dim, n_token, N);
auto h = proj_0->forward(ctx, semantic);
h = ggml_silu(gctx, h);
return proj_2->forward(ctx, h);
}
};
struct QueryFormerRunner : public GGMLRunner {
public:
QueryFormerConfig config;
QueryFormerModel query_former;
QueryFormerRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
query_former = QueryFormerModel(config);
query_former.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image_queryformer";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
query_former.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& inputs_embeds) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
ggml_tensor* x = make_input(inputs_embeds);
auto runner_ctx = get_context();
ggml_tensor* out = query_former.forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
inputs_embeds.dim());
}
};
struct TextProjectionRunner : public GGMLRunner {
public:
TextProjectionConfig config;
TextProjectionModel text_projection;
TextProjectionRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
text_projection = TextProjectionModel(config);
text_projection.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image_text_projection";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
text_projection.get_param_tensors(tensors, prefix);
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& hidden_states) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
ggml_tensor* x = make_input(hidden_states);
auto runner_ctx = get_context();
ggml_tensor* out = text_projection.forward(&runner_ctx, x);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
hidden_states.dim());
}
};
struct SigVQRunner : public GGMLRunner {
public:
SigVQConfig config;
SigVQModel sigvq;
std::array<std::vector<int32_t>, 4> pos_idx_data;
std::array<std::vector<float>, 4> pos_weight_data;
SigVQRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {
sigvq = SigVQModel(config);
sigvq.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "llada_image_sigvq";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
sigvq.get_param_tensors(tensors, prefix);
}
// Precomputes the four bilinear taps that resample the square position-embedding grid
// onto a grid_h x grid_w patch grid, matching grid_sample(align_corners=False,
// padding_mode="border").
void build_pos_embed_taps(int64_t grid_h, int64_t grid_w) {
const int64_t side = config.image_size / config.patch_size;
for (auto& v : pos_idx_data) {
v.clear();
}
for (auto& v : pos_weight_data) {
v.clear();
}
auto clamp_index = [side](int64_t v) {
return static_cast<int32_t>(std::min<int64_t>(std::max<int64_t>(v, 0), side - 1));
};
for (int64_t i = 0; i < grid_h; ++i) {
double src_h = (static_cast<double>(i) + 0.5) * side / static_cast<double>(grid_h) - 0.5;
int64_t h_floor = static_cast<int64_t>(std::floor(src_h));
double dh = src_h - static_cast<double>(h_floor);
for (int64_t j = 0; j < grid_w; ++j) {
double src_w = (static_cast<double>(j) + 0.5) * side / static_cast<double>(grid_w) - 0.5;
int64_t w_floor = static_cast<int64_t>(std::floor(src_w));
double dw = src_w - static_cast<double>(w_floor);
int32_t h0 = clamp_index(h_floor);
int32_t h1 = clamp_index(h_floor + 1);
int32_t w0 = clamp_index(w_floor);
int32_t w1 = clamp_index(w_floor + 1);
pos_idx_data[0].push_back(h0 * static_cast<int32_t>(side) + w0);
pos_idx_data[1].push_back(h0 * static_cast<int32_t>(side) + w1);
pos_idx_data[2].push_back(h1 * static_cast<int32_t>(side) + w0);
pos_idx_data[3].push_back(h1 * static_cast<int32_t>(side) + w1);
pos_weight_data[0].push_back(static_cast<float>((1.0 - dh) * (1.0 - dw)));
pos_weight_data[1].push_back(static_cast<float>((1.0 - dh) * dw));
pos_weight_data[2].push_back(static_cast<float>(dh * (1.0 - dw)));
pos_weight_data[3].push_back(static_cast<float>(dh * dw));
}
}
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& pixel_values) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLADA_IMAGE_TE_GRAPH_SIZE);
ggml_tensor* x = make_input(pixel_values);
int64_t grid_h = x->ne[1] / config.patch_size;
int64_t grid_w = x->ne[0] / config.patch_size;
build_pos_embed_taps(grid_h, grid_w);
std::vector<ggml_tensor*> pos_idx;
std::vector<ggml_tensor*> pos_weight;
for (int i = 0; i < 4; i++) {
auto idx = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast<int64_t>(pos_idx_data[i].size()));
set_backend_tensor_data(idx, pos_idx_data[i].data());
auto w = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast<int64_t>(pos_weight_data[i].size()));
set_backend_tensor_data(w, pos_weight_data[i].data());
pos_idx.push_back(idx);
pos_weight.push_back(w);
}
auto runner_ctx = get_context();
ggml_tensor* out = sigvq.forward(&runner_ctx, x, pos_idx, pos_weight);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3);
}
};
} // namespace LLaDAImageTE
#endif // __SD_MODEL_TE_LLADA_IMAGE_TE_HPP__

369
otherarch/sdcpp/src/model/te/llm.hpp
View file

@ -49,6 +49,7 @@ namespace LLM {
GEMMA2_2B,
GEMMA4_12B,
GPT_OSS_20B,
LLADA2_MOE,
ARCH_COUNT,
};
@ -62,6 +63,7 @@ namespace LLM {
"gemma2_2b",
"gemma4_12b",
"gpt_oss_20b",
"llada2_moe",
};
enum class MLPActivation {
@ -125,6 +127,17 @@ namespace LLM {
std::vector<int> sliding_attention;
int64_t num_experts = 0;
int64_t num_experts_per_tok = 0;
bool qkv_fused = false;
bool bidirectional = false;
float partial_rotary = 1.f;
// DeepSeek-V3-style grouped-sigmoid MoE routing (LLaDA2)
int64_t moe_intermediate_size = 0;
int64_t num_shared_experts = 0;
int64_t first_k_dense_replace = 0;
int64_t n_group = 0;
int64_t topk_group = 0;
float routed_scaling_factor = 1.f;
LLMVisionConfig vision;
bool have_vision_weight = false;
bool llama_cpp_style = false;
@ -212,6 +225,31 @@ namespace LLM {
config.intermediate_size = 9216;
config.num_layers = 26;
config.vocab_size = 256000;
} else if (arch == LLMArch::LLADA2_MOE) {
config.head_dim = 128;
config.num_heads = 16;
config.num_kv_heads = 4;
config.qkv_bias = false;
config.attention_out_bias = false;
config.qk_norm = true;
config.rms_norm_eps = 1e-6f;
config.hidden_size = 2048;
config.intermediate_size = 5120;
config.num_layers = 20;
config.vocab_size = 173568;
config.max_position_embeddings = 16384;
config.rope_thetas = {600000.f};
config.qkv_fused = true;
config.bidirectional = true;
config.partial_rotary = 0.5f;
config.num_experts = 256;
config.num_experts_per_tok = 8;
config.moe_intermediate_size = 512;
config.num_shared_experts = 1;
config.first_k_dense_replace = 1;
config.n_group = 8;
config.topk_group = 4;
config.routed_scaling_factor = 2.5f;
} else if (arch == LLMArch::GPT_OSS_20B) {
config.head_dim = 64;
config.num_heads = 64;
@ -419,6 +457,195 @@ namespace LLM {
}
};
// LLaDA2's MoE differs from GPT-OSS's in three ways that all change the result:
// routing scores are sigmoid (not softmax over the selected logits), expert selection is
// group-limited and uses a bias term that the returned weights do NOT include, and the
// experts carry no biases. Ref: LLaDA2MoeGate / LLaDA2MoeSparseMoeBlock in
// modeling_llada2uni_moe.py.
struct LLaDA2MoEMLP : public GGMLBlock {
protected:
int64_t hidden_size;
int64_t moe_intermediate_size;
int64_t num_experts;
int64_t num_experts_per_tok;
int64_t n_group;
int64_t topk_group;
float routed_scaling_factor;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
auto supported_type = [](ggml_type wtype, int64_t in_features) {
if (in_features % ggml_blck_size(wtype) != 0) {
return GGML_TYPE_F32;
}
return wtype;
};
// The reference runs the router in fp32; keep the weight in fp32 so the sigmoid
// scores and the group sums match.
params["gate.weight"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts);
params["gate.expert_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
ggml_type up_type = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
ggml_type down_type = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), moe_intermediate_size);
// HF ships the stacked experts as 3-D nn.Parameters, while the ComfyUI GGUF repack
// flattens the expert axis into ne[1]. Declare whichever the file holds - the two are
// bit-identical, and forward() reshapes to 3-D for ggml_mul_mat_id either way.
auto declare_experts = [&](const std::string& name, ggml_type type, int64_t in_dim, int64_t out_dim) {
auto storage = tensor_storage_map.find(prefix + name);
if (storage != tensor_storage_map.end() && storage->second.n_dims == 2) {
GGML_ASSERT(storage->second.nelements() == in_dim * out_dim * num_experts);
params[name] = ggml_new_tensor_2d(ctx, type, in_dim, out_dim * num_experts);
} else {
params[name] = ggml_new_tensor_3d(ctx, type, in_dim, out_dim, num_experts);
}
};
declare_experts("experts.gate_proj.weight", gate_type, hidden_size, moe_intermediate_size);
declare_experts("experts.up_proj.weight", up_type, hidden_size, moe_intermediate_size);
declare_experts("experts.down_proj.weight", down_type, moe_intermediate_size, hidden_size);
}
public:
LLaDA2MoEMLP(const LLMConfig& config)
: hidden_size(config.hidden_size),
moe_intermediate_size(config.moe_intermediate_size),
num_experts(config.num_experts),
num_experts_per_tok(config.num_experts_per_tok),
n_group(config.n_group),
topk_group(config.topk_group),
routed_scaling_factor(config.routed_scaling_factor) {
if (config.num_shared_experts > 0) {
blocks["shared_experts"] = std::make_shared<MLP>(config.hidden_size,
config.moe_intermediate_size * config.num_shared_experts,
false,
config.mlp_activation);
}
}
// Reproduces group_limited_topk(): keep the topk_group groups with the highest
// "sum of the two best scores in the group", then take the global top-k among them.
ggml_tensor* group_limited_mask(GGMLRunnerContext* ctx,
ggml_tensor* routing_scores,
int64_t n_token_total) {
ggml_context* gctx = ctx->ggml_ctx;
const int64_t per_group = num_experts / n_group;
// [experts_per_group, n_group * tokens] so top-2 runs per (group, token) row.
auto grouped = ggml_reshape_2d(gctx, routing_scores, per_group, n_group * n_token_total);
auto best2_idx = ggml_argsort_top_k(gctx, grouped, 2); // [2, n_group * tokens]
auto grouped_val = ggml_reshape_3d(gctx, grouped, 1, per_group, n_group * n_token_total);
auto best2 = ggml_get_rows(gctx, grouped_val, best2_idx); // [1, 2, n_group * tokens]
best2 = ggml_reshape_2d(gctx, best2, 2, n_group * n_token_total);
auto group_score = ggml_reshape_2d(gctx, ggml_sum_rows(gctx, best2), n_group, n_token_total); // [n_group, tokens]
// Threshold = the topk_group-th largest group score, taken from the sorted top-k.
auto top_groups = ggml_argsort_top_k(gctx, group_score, (int)topk_group); // [topk_group, tokens]
auto group_val = ggml_reshape_3d(gctx, group_score, 1, n_group, n_token_total);
auto top_scores = ggml_get_rows(gctx, group_val, top_groups); // [1, topk_group, tokens]
top_scores = ggml_reshape_2d(gctx, top_scores, topk_group, n_token_total);
auto threshold = ggml_view_2d(gctx,
top_scores,
1,
n_token_total,
top_scores->nb[1],
(topk_group - 1) * top_scores->nb[0]); // [1, tokens]
threshold = ggml_cont(gctx, threshold);
// keep = 1 - step(threshold - score). step(0) == 0, so the group sitting exactly on
// the threshold is kept without needing an epsilon.
auto diff = ggml_sub(gctx, ggml_repeat(gctx, threshold, group_score), group_score);
auto keep = ggml_scale_bias(gctx, ggml_step(gctx, diff), -1.f, 1.f); // [n_group, tokens]
// 0 for kept groups, a large negative for dropped ones, broadcast over the group.
auto additive = ggml_scale_bias(gctx, keep, 1e30f, -1e30f);
additive = ggml_reshape_3d(gctx, additive, 1, n_group, n_token_total);
auto expanded = ggml_repeat_4d(gctx, additive, per_group, n_group, n_token_total, 1);
return ggml_reshape_2d(gctx, expanded, num_experts, n_token_total);
}
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
const std::string& weight_name,
ggml_tensor* x,
ggml_tensor* selected_experts) {
ggml_tensor* w = params[weight_name];
if (w->ne[2] != num_experts) {
// Flattened layout: split the expert axis back out. ne[0] is untouched, so this
// stays valid for quantized types.
w = ggml_reshape_3d(ctx->ggml_ctx, w, w->ne[0], w->ne[1] / num_experts, num_experts);
}
return ggml_mul_mat_id(ctx->ggml_ctx, w, x, selected_experts);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, n_token, hidden_size]
GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0);
GGML_ASSERT(n_group > 0 && topk_group > 0 && num_experts % n_group == 0);
ggml_context* gctx = ctx->ggml_ctx;
const int64_t n_token = x->ne[1];
const int64_t N = x->ne[2];
const int64_t n_token_total = n_token * N;
auto identity = x;
auto logits = ggml_mul_mat(gctx, params["gate.weight"], x);
logits = ggml_reshape_2d(gctx, logits, num_experts, n_token_total);
auto scores = ggml_sigmoid(gctx, logits); // [num_experts, tokens]
// The bias steers selection only; the combine weights come from the unbiased scores.
auto routing = ggml_add(gctx, scores, params["gate.expert_bias"]);
routing = ggml_add(gctx, routing, group_limited_mask(ctx, routing, n_token_total));
auto selected_experts = ggml_argsort_top_k(gctx, routing, (int)num_experts_per_tok); // [top_k, tokens]
auto score_rows = ggml_reshape_3d(gctx, scores, 1, num_experts, n_token_total);
auto weights = ggml_get_rows(gctx, score_rows, selected_experts); // [1, top_k, tokens]
weights = ggml_reshape_2d(gctx, weights, num_experts_per_tok, n_token_total);
if (num_experts_per_tok > 1) {
auto denom = ggml_scale_bias(gctx, ggml_sum_rows(gctx, weights), 1.f, 1e-20f); // [1, tokens]
weights = ggml_div(gctx, weights, ggml_repeat(gctx, denom, weights));
}
weights = ggml_scale(gctx, weights, routed_scaling_factor);
weights = ggml_reshape_3d(gctx, weights, 1, num_experts_per_tok, n_token_total);
auto xf = ggml_reshape_3d(gctx, x, hidden_size, 1, n_token_total);
auto gate = expert_linear(ctx, "experts.gate_proj.weight", xf, selected_experts);
auto up = expert_linear(ctx, "experts.up_proj.weight", xf, selected_experts);
auto activated = ggml_swiglu_split(gctx, gate, up);
auto experts = expert_linear(ctx, "experts.down_proj.weight", activated, selected_experts);
experts = ggml_mul(gctx, experts, weights);
ggml_tensor* out = nullptr;
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
auto expert_out = ggml_view_2d(gctx,
experts,
hidden_size,
n_token_total,
experts->nb[2],
i * experts->nb[1]);
out = out == nullptr ? expert_out : ggml_add(gctx, out, expert_out);
}
if (num_experts_per_tok == 1) {
out = ggml_cont(gctx, out);
}
out = ggml_reshape_3d(gctx, out, hidden_size, n_token, N);
auto shared_it = blocks.find("shared_experts");
if (shared_it != blocks.end()) {
auto shared_experts = std::dynamic_pointer_cast<MLP>(shared_it->second);
out = ggml_add(gctx, out, shared_experts->forward(ctx, identity));
}
return out;
}
};
struct GPTOSSMLP : public GGMLBlock {
protected:
int64_t hidden_size;
@ -605,21 +832,31 @@ namespace LLM {
}
txt_token_end = image_embeds[i].first;
auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
if (input_embed == nullptr) {
input_embed = txt_embed;
} else {
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
// An embed can sit flush against the previous one or at the very start/end of the
// sequence, leaving no text tokens to splice around it.
if (txt_token_end > txt_token_start) {
auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
if (input_embed == nullptr) {
input_embed = txt_embed;
} else {
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
}
}
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
if (input_embed == nullptr) {
input_embed = image_embeds[i].second;
} else {
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
}
}
txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
txt_token_end = raw_x->ne[1];
auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
if (txt_token_end > txt_token_start) {
auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
}
GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
return input_embed;
}
@ -1122,6 +1359,7 @@ namespace LLM {
bool k_eq_v;
bool v_norm;
bool unscaled_attention;
bool qkv_fused;
float rms_norm_eps;
int rope_pairs;
@ -1147,12 +1385,20 @@ namespace LLM {
k_eq_v(global_layer && config.global_k_eq_v),
v_norm(config.v_norm),
unscaled_attention(config.unscaled_attention),
qkv_fused(config.qkv_fused),
rms_norm_eps(config.rms_norm_eps),
rope_pairs(0) {
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);
if (!k_eq_v) {
blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
if (qkv_fused) {
// The checkpoint ships q, k and v as one tensor and the loader cannot split a
// source tensor, so keep it fused and slice it in forward().
GGML_ASSERT(!k_eq_v);
blocks["query_key_value"] = std::make_shared<Linear>(config.hidden_size, (num_heads + num_kv_heads * 2) * head_dim, config.qkv_bias);
} else {
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);
if (!k_eq_v) {
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) {
@ -1161,7 +1407,7 @@ namespace LLM {
}
// Proportional RoPE rotates only the leading `rope_pairs` dimension pairs of the head;
// the rest are left unrotated through freq_factors (see rope_freq_factors()).
float partial = global_layer ? config.global_partial_rotary : 1.f;
float partial = global_layer && config.global_partial_rotary != 1.f ? config.global_partial_rotary : config.partial_rotary;
rope_pairs = static_cast<int>(partial * head_dim / 2.f);
}
@ -1186,14 +1432,28 @@ namespace LLM {
// x: [N, n_token, hidden_size]
int64_t n_token = x->ne[1];
int64_t N = x->ne[2];
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o_proj"]);
auto q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim]
auto k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
auto v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
ggml_tensor* q = nullptr;
ggml_tensor* k = nullptr;
ggml_tensor* v = nullptr;
if (qkv_fused) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["query_key_value"]);
auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
int64_t q_len = num_heads * head_dim;
int64_t k_len = num_kv_heads * head_dim;
q = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len);
k = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len);
v = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2);
} else {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["q_proj"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["k_proj"]);
auto v_proj = k_eq_v ? nullptr : std::dynamic_pointer_cast<Linear>(blocks["v_proj"]);
q = q_proj->forward(ctx, x); // [N, n_token, num_heads*head_dim]
k = k_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
v = k_eq_v ? k : v_proj->forward(ctx, x); // [N, n_token, num_kv_heads*head_dim]
}
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, num_heads, head_dim]
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N); // [N, n_token, num_kv_heads, head_dim]
@ -1336,6 +1596,38 @@ namespace LLM {
1.f,
32.f,
1.f);
} else if (arch == LLMArch::LLADA2_MOE) {
// LLaDA2 slices the head (query[..., :rotary_dim]) instead of zero-padding
// inv_freq like gemma does, so rotate_half pairs i with i + rotary_dim/2 and the
// frequencies use rotary_dim as the exponent denominator. Passing n_dims =
// rotary_dim reproduces both; freq_factors would give the wrong pairing.
int rotary_dim = rope_pairs * 2;
q = ggml_rope_ext(ctx->ggml_ctx,
q,
input_pos,
nullptr,
rotary_dim,
GGML_ROPE_TYPE_NEOX,
static_cast<int>(max_position_embeddings),
rope_thetas[0],
1.f,
0.f,
1.f,
32.f,
1.f);
k = ggml_rope_ext(ctx->ggml_ctx,
k,
input_pos,
nullptr,
rotary_dim,
GGML_ROPE_TYPE_NEOX,
static_cast<int>(max_position_embeddings),
rope_thetas[0],
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);
@ -1432,6 +1724,8 @@ namespace LLM {
blocks["self_attn"] = std::make_shared<Attention>(config, sliding_attention == 0);
if (config.arch == LLMArch::GPT_OSS_20B) {
blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
} else if (config.arch == LLMArch::LLADA2_MOE && layer_index >= config.first_k_dense_replace) {
blocks["mlp"] = std::make_shared<LLaDA2MoEMLP>(config);
} else {
blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
config.intermediate_size,
@ -1485,6 +1779,10 @@ namespace LLM {
if (arch == LLMArch::GPT_OSS_20B) {
auto mlp = std::dynamic_pointer_cast<GPTOSSMLP>(blocks["mlp"]);
x = mlp->forward(ctx, x);
} else if (auto moe_mlp = std::dynamic_pointer_cast<LLaDA2MoEMLP>(blocks["mlp"])) {
// LLaDA2 is dense for the first first_k_dense_replace layers and MoE afterwards,
// so the block type varies per layer rather than per arch.
x = moe_mlp->forward(ctx, x);
} else {
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
x = mlp->forward(ctx, x);
@ -1650,6 +1948,11 @@ namespace LLM {
return x;
}
ggml_tensor* embed(GGMLRunnerContext* ctx, ggml_tensor* input_ids) {
auto model = std::dynamic_pointer_cast<TextModel>(blocks["model"]);
return model->embed(ctx, input_ids);
}
std::shared_ptr<VisionModel> vision_model() {
GGML_ASSERT(enable_vision);
return std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
@ -1990,7 +2293,8 @@ namespace LLM {
config.arch == LLMArch::GEMMA3_12B ||
config.arch == LLMArch::GEMMA4_12B ||
config.arch == LLMArch::GEMMA2_2B ||
config.arch == LLMArch::GPT_OSS_20B) {
config.arch == LLMArch::GPT_OSS_20B ||
config.arch == LLMArch::LLADA2_MOE) {
input_pos_vec.resize(n_tokens);
for (int i = 0; i < n_tokens; ++i) {
input_pos_vec[i] = i;
@ -2042,8 +2346,9 @@ namespace LLM {
attention_mask_vec.resize(n_tokens * n_tokens);
for (int i0 = 0; i0 < n_tokens; i0++) {
for (int i1 = 0; i1 < n_tokens; i1++) {
// Diffusion LLMs attend in both directions; only causal LMs get the triangle.
float value = 0.f;
if (i0 > i1) {
if (!config.bidirectional && i0 > i1) {
value = -INFINITY;
}
attention_mask_vec[i1 * n_tokens + i0] = value;
@ -2115,6 +2420,22 @@ namespace LLM {
input_ids.dim() + 1);
}
// LLaDA-Image's QueryFormer consumes the raw token embeddings before the backbone runs,
// so it needs the embedding lookup on its own.
sd::Tensor<float> compute_input_embeds(const int n_threads,
const sd::Tensor<int32_t>& input_ids) {
auto get_graph = [&]() -> ggml_cgraph* {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* ids = make_input(input_ids);
auto runner_ctx = get_context();
ggml_tensor* out = model.embed(&runner_ctx, ids);
ggml_build_forward_expand(gf, out);
return gf;
};
return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true),
input_ids.dim() + 1);
}
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 / config.vision.patch_size;
@ -2370,11 +2691,13 @@ namespace LLM {
pad_id = 199999;
} else if (arch == LLMArch::GEMMA2_2B) {
pad_id = 0;
} else if (arch == LLMArch::LLADA2_MOE) {
pad_id = 156892;
}
tokenizer = tokenizers.create(TokenizerConfig::MAIN, model.config.vocab_size, pad_id);
if (!tokenizer) {
if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B) {
throw std::runtime_error("GPT-OSS and Gemma 2 require an external tokenizer.json in the main tokenizer slot");
if (arch == LLMArch::GPT_OSS_20B || arch == LLMArch::GEMMA2_2B || arch == LLMArch::LLADA2_MOE) {
throw std::runtime_error("GPT-OSS, Gemma 2 and LLaDA2 require an external tokenizer.json in the main tokenizer slot");
}
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
tokenizer = std::make_shared<MistralTokenizer>();

5
otherarch/sdcpp/src/model/vae/ltx_vae.hpp
View file

@ -1300,7 +1300,7 @@ struct LTXVideoVAE : public VAE {
feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
}
auto runner_ctx = get_context();
auto runner_ctx = get_context(gf);
int feat_count = 0;
ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
z,
@ -1313,8 +1313,7 @@ struct LTXVideoVAE : public VAE {
for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
if (feat_cache != nullptr) {
cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
ggml_build_forward_expand(gf, feat_cache);
runner_ctx.persist_cache_tensor(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
}
}

49
otherarch/sdcpp/src/model/vae/wan_vae.hpp
View file

@ -24,6 +24,7 @@ namespace WAN {
std::tuple<int, int, int> padding;
std::tuple<int, int, int> dilation;
bool bias;
float scale = 1.f;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
auto weight = tensor_storage_map.find(prefix + "weight");
@ -60,6 +61,10 @@ namespace WAN {
dilation(std::move(dilation)),
bias(bias) {}
void set_scale(float scale_value) {
scale = scale_value;
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* cache_x = nullptr) {
// x: [N*IC, ID, IH, IW]
// result: x: [N*OC, ID, IH, IW]
@ -82,11 +87,25 @@ namespace WAN {
}
x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
if (w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == in_channels) {
// One frame through a one-frame-deep kernel is a 2D conv; backends without
// im2col_3d (Metal) otherwise fall back to a much slower direct conv_3d.
if (!ggml_is_contiguous(x)) {
x = ggml_cont(ctx->ggml_ctx, x);
}
ggml_tensor* x2 = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], in_channels, 1);
ggml_tensor* w2 = ggml_reshape_4d(ctx->ggml_ctx, w, w->ne[0], w->ne[1], in_channels, out_channels);
x2 = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b,
std::get<2>(stride), std::get<1>(stride), 0, 0,
std::get<2>(dilation), std::get<1>(dilation),
ctx->conv2d_direct_enabled, false, false, scale);
return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels);
}
return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
0, 0, 0,
std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
false, ctx->conv3d_direct_enabled);
false, ctx->conv3d_direct_enabled, scale);
}
};
@ -1070,11 +1089,11 @@ namespace WAN {
}
if (version == VERSION_QWEN_IMAGE_2_1) {
wan2_2 = true;
dec_dim = 144;
z_dim = 64;
input_channels = 4;
dim_mult = {1, 2, 4, 8, 8};
wan2_2 = true;
dec_dim = 144;
z_dim = 64;
input_channels = 4;
dim_mult = {1, 2, 4, 8, 8};
}
if (is_2D) {
@ -1103,6 +1122,19 @@ namespace WAN {
} else {
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
}
if (version == VERSION_QWEN_IMAGE_2_1) {
// Keep large VAE activations within the FP16 convolution range.
const float conv_scale = 1.f / 128.f;
std::vector<GGMLBlock*> all_blocks;
get_all_blocks(all_blocks);
for (auto block : all_blocks) {
if (auto conv = dynamic_cast<Conv2d*>(block)) {
conv->set_scale(conv_scale);
} else if (auto conv = dynamic_cast<CausalConv3d*>(block)) {
conv->set_scale(conv_scale);
}
}
}
}
static ggml_tensor* patchify(ggml_context* ctx,
@ -1429,15 +1461,14 @@ namespace WAN {
ggml_tensor* z = make_input(z_tensor);
auto runner_ctx = get_context();
auto runner_ctx = get_context(gf);
ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx);
for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) {
ggml_tensor* feat_cache = ae._feat_map[feat_idx];
if (feat_cache != nullptr) {
cache("feat_idx:" + std::to_string(feat_idx), feat_cache);
ggml_build_forward_expand(gf, feat_cache);
runner_ctx.persist_cache_tensor("feat_idx:" + std::to_string(feat_idx), feat_cache);
}
}

78
otherarch/sdcpp/src/model_loader.cpp
View file

@ -339,7 +339,13 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path,
}
for (const std::string& shard_path : shard_paths) {
if (!parse_file(shard_path, prefix)) {
FileStamp stamp;
if (!read_file_stamp(shard_path, stamp)) {
return false;
}
parsed_dependencies_.push_back(stamp);
LOG_INFO("load %s using safetensors format", shard_path.c_str());
if (!init_from_safetensors_file(shard_path, prefix)) {
return false;
}
}
@ -541,6 +547,9 @@ SDVersion ModelLoader::get_sd_version() const {
if (tensor_storage.name.find("model.diffusion_model.double_blocks.0.img_mlp.gate_proj.weight") != std::string::npos) {
return VERSION_OVIS_IMAGE;
}
if (tensor_storage.name.find("model.diffusion_model.sigvq_embedder.1.weight") != std::string::npos) {
return VERSION_LLADA_IMAGE;
}
if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) {
return VERSION_Z_IMAGE;
}
@ -898,7 +907,8 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors,
bool writable_mmap) {
bool writable_mmap,
ggml_backend_dev_t device) {
std::set<std::string> names;
for (const auto& entry : tensors) {
names.insert(entry.first);
@ -920,6 +930,39 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
if (!fdata.mmbuffer)
continue;
// Wrapped on first use: a device buffer makes the whole file resident on that device.
std::shared_ptr<struct ggml_backend_buffer> file_buffer = device == nullptr ? fdata.mmbuffer : nullptr;
bool file_unmappable = false;
auto buffer_for_file = [&]() -> ggml_backend_buffer_t {
if (file_buffer || file_unmappable) {
return file_buffer.get();
}
auto cached = fdata.device_mmbuffers.find(device);
if (cached != fdata.device_mmbuffers.end()) {
file_buffer = cached->second;
return file_buffer.get();
}
size_t max_tensor_size = 0;
for (const auto& ts : fdata.tensors) {
max_tensor_size = std::max(max_tensor_size, static_cast<size_t>(ts.nbytes()));
}
ggml_backend_buffer_t buf = sd_backend_dev_buffer_from_host_ptr(device,
fdata.mmapped->writable_data(),
fdata.mmapped->size(),
max_tensor_size);
if (buf == nullptr) {
LOG_WARN("mmap: %s cannot map '%s', loading it instead",
ggml_backend_dev_name(device), fdata.path.c_str());
file_unmappable = true;
return nullptr;
}
LOG_INFO("mmap: mapped '%s' for %s", fdata.path.c_str(), ggml_backend_dev_name(device));
file_buffer = std::shared_ptr<struct ggml_backend_buffer>(buf, ggml_backend_buffer_free);
fdata.device_mmbuffers[device] = file_buffer;
return file_buffer.get();
};
const std::vector<TensorStorage>& file_tensors = fdata.tensors;
size_t file_mapped_bytes = 0;
@ -969,10 +1012,13 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
continue;
}
ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
dst_tensor->buffer = buf_mmap;
dst_tensor->data = mmap_data + tensor_offset;
ggml_backend_buffer_t buf_mmap = buffer_for_file();
if (buf_mmap == nullptr) {
break;
}
uint8_t* mmap_data = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
dst_tensor->buffer = buf_mmap;
dst_tensor->data = mmap_data + tensor_offset;
file_mapped_bytes += tensor_size;
file_mapped_tensors++;
@ -981,7 +1027,7 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
if (file_mapped_bytes > 0) {
mapped_tensors += file_mapped_tensors;
mapped_bytes += file_mapped_bytes;
result.push_back({fdata.mmapped, fdata.mmbuffer});
result.push_back({fdata.mmapped, file_buffer});
}
}
@ -997,6 +1043,16 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
return result;
}
std::vector<ggml_backend_buffer_t> ModelLoader::get_device_mmap_buffers() const {
std::vector<ggml_backend_buffer_t> buffers;
for (const auto& fdata : file_data) {
for (const auto& entry : fdata.device_mmbuffers) {
buffers.push_back(entry.second.get());
}
}
return buffers;
}
bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool enable_mmap,
const std::set<std::string>* target_tensor_names,
@ -1142,6 +1198,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
continue;
}
if (dst_tensor->buffer != nullptr &&
std::any_of(fdata.device_mmbuffers.begin(), fdata.device_mmbuffers.end(),
[&](const auto& entry) { return entry.second.get() == dst_tensor->buffer; })) {
continue;
}
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
@ -1599,6 +1660,9 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
// Pass, do not convert. For Unet
} else if (contains(name, "embedding")) {
// Pass, do not convert embedding
} else if (ends_with(name, "_pad_token")) {
// Pass, do not convert. LLaDA-Image stores its pad tokens far outside the f16
// range, so any format with an f16 scale or payload turns them into inf.
} else {
return true;
}

6
otherarch/sdcpp/src/model_loader.h
View file

@ -20,6 +20,8 @@ struct ModelFileData {
std::vector<TensorStorage> tensors;
std::shared_ptr<MmapWrapper> mmapped;
std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
// mmapped wrapped by devices that can use host memory in place (buffer_from_host_ptr)
std::map<ggml_backend_dev_t, std::shared_ptr<struct ggml_backend_buffer>> device_mmbuffers;
bool is_zip;
};
@ -120,7 +122,9 @@ public:
void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
std::set<std::string> ignore_tensors = {},
bool writable = true);
bool writable = true,
ggml_backend_dev_t device = nullptr);
std::vector<ggml_backend_buffer_t> get_device_mmap_buffers() const;
bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
bool use_mmap = false,
const std::set<std::string>* target_tensor_names = nullptr,

2
otherarch/sdcpp/src/model_loader_files.cpp
View file

@ -147,7 +147,7 @@ bool ModelLoader::add_file_impl(const std::string& path, const std::string& pref
}
if (tensor.index_in_zip < 0) {
const auto& stamp = physical_files[tensor.file_index];
if (tensor.offset > stamp.size || static_cast<uint64_t>(tensor.nbytes_to_read()) > stamp.size - tensor.offset) { //kcpp int8 fp8
if (tensor.offset > stamp.size || static_cast<uint64_t>(tensor.nbytes_to_read()) > stamp.size - tensor.offset) { // kcpp int8 fp8
LOG_ERROR("tensor '%s' extends beyond its model file", tensor.name.c_str());
return false;
}

98
otherarch/sdcpp/src/model_manager.cpp
View file

@ -780,38 +780,52 @@ bool ModelManager::validate_tensor(const TensorState& state) const {
bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
std::vector<ParamsStorageBlock*>& created_storage_blocks) {
std::map<std::string, ggml_tensor*> mmap_candidates;
std::map<std::string, TensorState*> mmap_states;
// A GPU that computes on mmapped params in place cannot address a CPU buffer, and nothing
// stages them for it, so they are mapped through a buffer of that GPU's device.
struct MmapGroup {
std::map<std::string, ggml_tensor*> candidates;
std::map<std::string, TensorState*> states;
};
std::map<ggml_backend_dev_t, MmapGroup> groups;
for (TensorState* state : states) {
if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr ||
state->tensor->data != nullptr || state->tensor->view_src != nullptr) {
continue;
}
mmap_candidates[state->name] = state->tensor;
mmap_states[state->name] = state;
}
if (mmap_candidates.empty()) {
return true;
}
auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_);
if (mmap_store.empty()) {
return true;
}
auto block = std::make_unique<ParamsStorageBlock>();
block->mmap_tensor_stores = std::move(mmap_store);
ParamsStorageBlock* raw = block.get();
for (const auto& pair : mmap_states) {
TensorState* state = pair.second;
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
block->states.push_back(state);
ggml_backend_dev_t device = nullptr;
if (!sd_backend_is_cpu(state->compute_backend) && !sd_backend_is_cpu(state->params_backend)) {
device = ggml_backend_get_device(state->compute_backend);
}
MmapGroup& group = groups[device];
group.candidates[state->name] = state->tensor;
group.states[state->name] = state;
}
if (!block->states.empty()) {
params_storage_blocks_.push_back(std::move(block));
created_storage_blocks.push_back(raw);
for (auto& [device, group] : groups) {
// Device buffers wrap read-only mappings only; params that LoRAs are merged into in place
// are loaded instead.
if (device != nullptr && writable_mmap_) {
continue;
}
auto mmap_store = model_loader_.mmap_tensors(group.candidates, {}, writable_mmap_, device);
if (mmap_store.empty()) {
continue;
}
auto block = std::make_unique<ParamsStorageBlock>();
block->mmap_tensor_stores = std::move(mmap_store);
ParamsStorageBlock* raw = block.get();
for (const auto& pair : group.states) {
TensorState* state = pair.second;
if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
block->states.push_back(state);
}
}
if (!block->states.empty()) {
params_storage_blocks_.push_back(std::move(block));
created_storage_blocks.push_back(raw);
}
}
return true;
}
@ -1353,15 +1367,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
}
size_t total_size = 0;
auto add_buffer = [&](ggml_backend_buffer_t buffer) {
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer)) {
std::unordered_set<ggml_backend_buffer_t> seen;
auto add_buffer = [&](ggml_backend_buffer_t buffer) {
if (buffer == nullptr || ggml_backend_buffer_is_host(buffer) || !seen.insert(buffer).second) {
return;
}
ggml_backend_buffer_type_t buffer_type = ggml_backend_buffer_get_type(buffer);
auto split_devices = split_buffer_devices_.find(buffer_type);
const bool on_device = split_devices == split_buffer_devices_.end()
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device
: std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) {
return ggml_backend_get_device(entry.first) == compute_device;
});
if (!on_device) {
@ -1371,9 +1386,16 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe
total_size = buffer_size > SIZE_MAX - total_size ? SIZE_MAX : total_size + buffer_size;
};
// The loader may retain device mappings after their parameter blocks are released.
for (ggml_backend_buffer_t buffer : model_loader_.get_device_mmap_buffers()) {
add_buffer(buffer);
}
for (const auto& block : params_storage_blocks_) {
if (block != nullptr) {
add_buffer(block->buffer);
for (const auto& store : block->mmap_tensor_stores) {
add_buffer(store.mmbuffer.get());
}
}
}
for (const auto& block : compute_staging_blocks_) {
@ -1613,7 +1635,8 @@ void ModelManager::remove_runtime_owner(uintptr_t owner_id) {
ModelManager::CapacityCheck ModelManager::check_capacity(
const DeviceMemoryRequest& request,
const std::vector<TensorState*>& states) const {
const std::vector<TensorState*>& states,
bool log_details) const {
CapacityCheck result;
if (request.compute_backend == nullptr || sd_backend_is_cpu(request.compute_backend)) {
return result;
@ -1631,16 +1654,23 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
}
size_t free_bytes = 0, total_bytes = 0;
ggml_backend_dev_memory(device, &free_bytes, &total_bytes);
const size_t weights_resident = compute_backend_resident_bytes(backend);
const size_t other_runtime = other_runtime_resident_bytes(request.owner_id, backend);
const size_t resident = add(weights_resident, add(other_runtime, request.runtime_resident_bytes));
if (log_details) {
LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB",
ggml_backend_name(backend),
free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0),
weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0),
request.runtime_resident_bytes / (1024.0 * 1024.0));
}
if (free_bytes == 0 && total_bytes == 0) {
return SIZE_MAX;
}
// Vulkan's heap budget subtraction can underflow when usage exceeds the budget.
if (total_bytes > 0 && free_bytes > total_bytes) {
if (total_bytes > 0 && free_bytes > total_bytes && sd_backend_is(backend, "Vulkan")) {
return size_t{0};
}
const size_t resident = add(compute_backend_resident_bytes(backend),
add(other_runtime_resident_bytes(request.owner_id, backend),
request.runtime_resident_bytes));
if (total_bytes > 0) {
free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0);
}
@ -1786,7 +1816,7 @@ bool ModelManager::ensure_compute_backend_capacity(
}
}
const auto capacity = check_capacity(request, required_states);
const auto capacity = check_capacity(request, required_states, true);
const std::string available_device = capacity.available_device_bytes == SIZE_MAX
? "unknown"
: sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0));

3
otherarch/sdcpp/src/model_manager.h
View file

@ -157,7 +157,8 @@ private:
}
};
CapacityCheck check_capacity(const DeviceMemoryRequest& request,
const std::vector<TensorState*>& states) const;
const std::vector<TensorState*>& states,
bool log_details = false) const;
ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const;

59
otherarch/sdcpp/src/name_conversion.cpp
View file

@ -103,6 +103,8 @@ std::string convert_open_clip_to_hf_clip_name(std::string name) {
return name;
}
std::string convert_llada2_moe_te_name(std::string name);
std::string convert_cond_stage_model_name(std::string name, std::string prefix) {
static const std::vector<std::pair<std::string, std::string>> clip_name_map{
{"transformer.text_projection.weight", "transformer.text_model.text_projection"},
@ -177,6 +179,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
replace_with_name_map(name, llm_vision_name_map);
} else {
replace_with_name_map(name, llm_name_map);
name = convert_llada2_moe_te_name(name);
}
} else {
name = convert_open_clip_to_hf_clip_name(name);
@ -749,6 +752,52 @@ std::string convert_hunyuan_video_to_original_flux(std::string name) {
return name;
}
// LLaDA-Image's LLaDA2-MoE text encoder. Both published layouts use these names; the ComfyUI
// GGUF repack differs only by appending ".weight" to the bare 3-D expert parameters.
// Called with the "text_encoders." prefix already stripped, so the name still carries "llm.".
std::string convert_llada2_moe_te_name(std::string name) {
static const std::vector<std::pair<std::string, std::string>> name_map = {
{"model.language_model.word_embeddings.", "model.embed_tokens."},
{"model.language_model.norm.", "model.norm."},
{"model.language_model.lm_head.", "lm_head."},
{"model.language_model.layers.", "model.layers."},
{"attention.query_key_value.", "self_attn.query_key_value."},
{"attention.dense.", "self_attn.o_proj."},
{"attention.query_layernorm.", "self_attn.q_norm."},
{"attention.key_layernorm.", "self_attn.k_norm."},
};
replace_with_name_map(name, name_map);
// The HF checkpoint stores the stacked experts as bare nn.Parameters with no ".weight".
static const std::vector<std::string> bare_expert_params = {
"mlp.experts.gate_proj",
"mlp.experts.up_proj",
"mlp.experts.down_proj",
};
for (const auto& suffix : bare_expert_params) {
if (ends_with(name, suffix)) {
name += ".weight";
break;
}
}
return name;
}
// The attention projections keep their diffusers names (JointAttention's split_qkv mode), so
// only the patch-size-keyed dicts need flattening. Latents arrive already patchified from the
// Flux2 VAE, so the only patch key is 1-1.
std::string convert_diffusers_dit_to_original_llada_image(std::string name) {
static const std::vector<std::pair<std::string, std::string>> prefix_map = {
{"all_x_embedder.1-1.", "x_embedder."},
{"all_final_layer.1-1.", "final_layer."},
};
replace_with_prefix_map(name, prefix_map);
return name;
}
std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
int num_layers = 30;
int num_refiner_layers = 2;
@ -896,6 +945,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
name = convert_hunyuan_video_to_original_flux(name);
} else if (sd_version_is_z_image(version)) {
name = convert_diffusers_dit_to_original_lumina2(name);
} else if (sd_version_is_llada_image(version)) {
name = convert_diffusers_dit_to_original_llada_image(name);
} else if (sd_version_is_anima(version)) {
name = convert_other_dit_to_original_anima(name);
} else if (sd_version_is_krea2(version)) {
@ -1298,9 +1349,8 @@ bool is_first_stage_model_name(const std::string& name) {
}
static std::string convert_esrgan_tensor_name(std::string name) {
static std::unordered_map<std::string, std::string> esrgan_name_map;
if (esrgan_name_map.empty()) {
static const auto esrgan_name_map = [] {
std::unordered_map<std::string, std::string> esrgan_name_map;
esrgan_name_map["model.0."] = "conv_first.";
constexpr int max_num_blocks = 64;
@ -1326,7 +1376,8 @@ static std::string convert_esrgan_tensor_name(std::string name) {
esrgan_name_map["model.7."] = "conv_last.";
esrgan_name_map["model.8."] = "conv_hr.";
esrgan_name_map["model.10."] = "conv_last.";
}
return esrgan_name_map;
}();
replace_with_prefix_map(name, esrgan_name_map);
return name;

221
otherarch/sdcpp/src/pipeline/diffusion_engine.cpp
View file

@ -7,6 +7,7 @@
#include <list>
#include <mutex>
#include <set>
#include <tuple>
#include <type_traits>
#include <unordered_set>
#include <utility>
@ -29,6 +30,7 @@
#include "stable-diffusion.h"
#include "conditioning/conditioner.hpp"
#include "conditioning/conditioning_cache.h"
#include "core/backend_fit.h"
#include "extensions/generation_extension.h"
#include "model/adapter/ip_adapter.hpp"
@ -101,6 +103,7 @@ const char* model_version_to_str[] = {
"Krea2",
"Mage Flow",
"SenseNova U1.5",
"LLaDA-Image",
"ESRGAN",
};
@ -134,6 +137,7 @@ static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
StableDiffusionGGML::StableDiffusionGGML()
: rng(std::make_shared<PhiloxRNG>()),
conditioning_cache_(std::make_unique<ConditioningCache>()),
denoiser(std::make_shared<CompVisDenoiser>()) {}
StableDiffusionGGML::~StableDiffusionGGML() = default;
@ -203,6 +207,8 @@ void StableDiffusionGGML::end_runners() {
}
bool StableDiffusionGGML::reset_runners(const RunnerGroups& groups) {
conditioning_cache_->clear();
conditioning_loras_.clear();
end_runners();
clear_lora_adapters();
runtime_lora_models.clear();
@ -883,7 +889,7 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
{
to_replace = "taesd_f2.embd";
}
else if(sd_version_uses_wan_vae(tempver))
else if(sd_version_uses_wan_vae(tempver) && tempver != VERSION_QWEN_IMAGE_2_1) // qwen 2.1 + tae crashing as of master-917
{
to_replace = "taesd_w21.embd";
}
@ -1105,6 +1111,47 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf
return true;
}
bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) {
if (!diffusion_model) {
return false;
}
if (enabled) {
#ifndef SD_USE_UPSTREAM_GGML
auto* ctx = ggml_init({4 * ggml_tensor_overhead(), nullptr, true});
if (ctx == nullptr) {
return false;
}
auto* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
auto* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1);
auto* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1);
auto* op = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO);
bool supported = true;
for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) {
if (!ggml_backend_supports_op(backend, op)) {
LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels",
ggml_backend_name(backend));
supported = false;
}
}
ggml_free(ctx);
if (!supported) {
return false;
}
#else
LOG_ERROR("SageAttention requires -DSD_USE_UPSTREAM_GGML=OFF and a CUDA backend");
return false;
#endif
}
diffusion_model->set_sage_attention_enabled(enabled);
if (high_noise_diffusion_model) {
high_noise_diffusion_model->set_sage_attention_enabled(enabled);
}
if (enabled) {
LOG_INFO("Using SageAttention in the diffusion model; CUDA selects the supported kernel, unsupported layers use flash/default attention");
}
return true;
}
bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
#ifdef SD_USE_UPSTREAM_GGML
LOG_WARN(
@ -1121,6 +1168,11 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
return false;
}
}
if (sd_ctx_params->conditioning_cache_size < 0) {
LOG_ERROR("conditioning_cache_size must be non-negative");
return false;
}
conditioning_cache_->set_capacity(static_cast<size_t>(sd_ctx_params->conditioning_cache_size));
auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
n_threads = sd_ctx_params->n_threads;
tensor_executor = std::make_unique<sd::ParallelExecutor>(n_threads > 0 ? n_threads : sd_get_num_physical_cores());
@ -1386,6 +1438,9 @@ bool StableDiffusionGGML::validate_and_load_runners() {
high_noise_diffusion_model->set_flash_attention_enabled(true);
}
}
if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) {
return false;
}
LOG_VERBOSE("validating model metadata");
std::set<std::string> ignore_tensors;
@ -1549,6 +1604,7 @@ bool StableDiffusionGGML::build_denoiser() {
sd_version_is_anima(version) ||
sd_version_is_ernie_image(version) ||
sd_version_is_z_image(version) ||
sd_version_is_llada_image(version) ||
sd_version_is_boogu_image(version) ||
sd_version_is_pid(version) ||
sd_version_is_ideogram4(version)) {
@ -1569,6 +1625,8 @@ bool StableDiffusionGGML::build_denoiser() {
default_flow_shift = 3.16f;
} else if (sd_version_is_mage_flow(version)) {
default_flow_shift = 6.f;
} else if (sd_version_is_llada_image(version)) {
default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid
} else {
default_flow_shift = 3.f;
}
@ -1986,8 +2044,19 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
int64_t t0 = ggml_time_ms();
end_runners();
clear_lora_adapters();
if (!model_manager->prepare_lora_sources(all_loras))
if (!model_manager->prepare_lora_sources(all_loras)) {
conditioning_cache_->clear();
return false;
}
if (!std::equal(all_loras.begin(), all_loras.end(),
conditioning_loras_.begin(), conditioning_loras_.end(),
[](const ModelManager::LoraSpec& a, const ModelManager::LoraSpec& b) {
return a.file_id == b.file_id && a.file_revision == b.file_revision &&
a.multiplier == b.multiplier && a.is_high_noise == b.is_high_noise &&
a.tensor_name_prefix_filter == b.tensor_name_prefix_filter;
})) {
conditioning_cache_->clear();
}
runtime_lora_models.erase(std::remove_if(runtime_lora_models.begin(), runtime_lora_models.end(), [&](const RuntimeLora& entry) {
return std::none_of(all_loras.begin(), all_loras.end(), [&](const ModelManager::LoraSpec& spec) {
return entry.matches(spec);
@ -1997,6 +2066,7 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
const bool success = apply_lora_immediately ? apply_loras_immediately(all_loras)
: apply_loras_at_runtime(all_loras);
if (!success) {
conditioning_cache_->clear();
clear_lora_adapters();
runtime_lora_models.clear();
return false;
@ -2006,9 +2076,14 @@ bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_coun
if (!all_loras.empty()) {
LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
}
conditioning_loras_ = std::move(all_loras);
return true;
}
SDCondition StableDiffusionGGML::get_learned_condition(const ConditionerParams& params) {
return conditioning_cache_->get(*cond_stage_model, n_threads, params);
}
void StableDiffusionGGML::reset_generation_extensions() {
for (auto& extension : generation_extensions) {
extension->reset_runtime_condition();
@ -2193,6 +2268,8 @@ void StableDiffusionGGML::preview_image(int step,
int patch_sz = 1;
const float(*latent_rgb_proj)[3] = nullptr;
float* latent_rgb_bias = nullptr;
const float* latent_alpha_proj = nullptr;
float latent_alpha_bias = 1.f;
if (channels == 128) {
if (sd_version_uses_flux2_vae(version)) {
@ -2206,6 +2283,16 @@ void StableDiffusionGGML::preview_image(int step,
LOG_WARN("No latent to RGB projection known for this model");
return;
}
} else if (channels == 64) {
if (version == VERSION_QWEN_IMAGE_2_1) {
latent_rgb_proj = qwen21_latent_rgb_proj;
latent_rgb_bias = qwen21_latent_rgb_bias;
latent_alpha_proj = qwen21_latent_alpha_proj;
latent_alpha_bias = qwen21_latent_alpha_bias;
} else {
LOG_WARN("No latent to RGB projection known for this model");
return;
}
} else if (channels == 48) {
if (sd_version_is_wan(version)) {
latent_rgb_proj = wan_22_latent_rgb_proj;
@ -2256,13 +2343,14 @@ void StableDiffusionGGML::preview_image(int step,
uint32_t img_width = static_cast<uint32_t>(_latents.shape()[0]) * patch_sz;
uint32_t img_height = static_cast<uint32_t>(_latents.shape()[1]) * patch_sz;
uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * 3 * sizeof(uint8_t));
uint32_t img_channels = latent_alpha_proj != nullptr ? 4 : 3;
uint8_t* data = (uint8_t*)malloc(frames * img_width * img_height * img_channels * sizeof(uint8_t));
GGML_ASSERT(data != nullptr);
preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz);
preview_latent_video(data, _latents, latent_rgb_proj, latent_rgb_bias, patch_sz, latent_alpha_proj, latent_alpha_bias);
sd_image_t* images = (sd_image_t*)malloc(frames * sizeof(sd_image_t));
GGML_ASSERT(images != nullptr);
for (uint32_t i = 0; i < frames; i++) {
images[i] = {img_width, img_height, 3, data + i * img_width * img_height * 3};
images[i] = {img_width, img_height, img_channels, data + i * img_width * img_height * img_channels};
}
step_callback(step, frames, images, is_noisy, step_callback_data);
free(data);
@ -2437,6 +2525,15 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
};
RunnerEndOnExit sample_diffusion_runner_end{work_diffusion_model.get()};
// These inputs are immutable for this sampling run. Extensions may replace or
// modify them per step, so those paths need an explicit stability contract first.
const bool cache_qwen_prefix = version == VERSION_QWEN_IMAGE_2_1 &&
std::none_of(generation_extensions.begin(), generation_extensions.end(),
[](const auto& extension) { return extension->is_enabled(); });
using QwenPrefixInputs = std::tuple<const sd::Tensor<float>*, const sd::Tensor<int32_t>*,
const std::vector<sd::Tensor<float>>*>;
std::vector<QwenPrefixInputs> qwen_prefix_inputs;
RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr};
const bool apply_denoise_mask = !denoise_mask.empty() &&
@ -2613,6 +2710,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
timesteps_tensor,
cond,
&controls);
bool uncond_controls_ready = false;
static const std::vector<sd::Tensor<float>> empty_ref_latents;
bool uncond_without_ref_latents = !img_uncond.empty() &&
@ -2668,6 +2766,9 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask,
condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds};
} else if (sd_version_is_llada_image(version)) {
diffusion_params.extra = LLaDAImageDiffusionExtra{
condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0]};
} else if (sd_version_is_minimax_h3(version)) {
diffusion_params.extra = MiniMaxH3DiffusionExtra{
condition.c_token_types.empty() ? nullptr : &condition.c_token_types,
@ -2699,10 +2800,33 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
return std::move(cached_output);
}
// A re-enabled condition can miss the cache even when the positive pass was reused.
if (!uncond_controls_ready && !uncond.empty() &&
(&condition == &uncond || &condition == &img_uncond)) {
compute_sample_controls(control_image,
noised_input,
timesteps_tensor,
uncond,
&controls);
uncond_controls_ready = true;
}
for (const auto& extension : generation_extensions) {
extension->before_diffusion(diffusion_params, step);
}
if (cache_qwen_prefix) {
auto* extra = std::get_if<QwenImage21DiffusionExtra>(&diffusion_params.extra);
if (extra != nullptr) {
auto key = std::make_tuple(diffusion_params.context, extra->image_slots,
diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr);
auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key);
extra->prefix_id = static_cast<uint64_t>(entry - qwen_prefix_inputs.begin()) + 1;
if (entry == qwen_prefix_inputs.end()) {
qwen_prefix_inputs.push_back(key);
}
}
}
auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params);
if (output_opt.empty()) {
LOG_ERROR("diffusion model compute failed");
@ -2726,41 +2850,69 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
}
}
float effective_guidance_scale = guidance_schedule.empty()
? cfg_scale
: guidance_schedule[guidance_schedule.size() - 1 - step];
float image_guidance_scale = img_cfg_scale;
constexpr float kEpsilon = 1e-5f;
bool skip_uncond = false;
if (!uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
if (!img_uncond.empty()) {
skip_uncond = std::abs(image_guidance_scale - effective_guidance_scale) < kEpsilon;
} else {
skip_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
}
}
bool skip_img_uncond = false;
if (!img_uncond.empty() && !needs_uncond_denoised && !use_apg_guidance) {
if (!uncond.empty()) {
skip_img_uncond = std::abs(image_guidance_scale - 1.0f) < kEpsilon;
} else {
skip_img_uncond = std::abs(effective_guidance_scale - 1.0f) < kEpsilon;
}
}
cond_out = run_condition(*positive_condition, c_concat_override);
if (cond_out.empty()) {
return {};
}
if (!uncond.empty()) {
if (!step_cache.is_step_skipped()) {
compute_sample_controls(control_image,
noised_input,
timesteps_tensor,
uncond,
&controls);
}
const std::vector<int>* uncond_skip_layers = nullptr;
if (is_skiplayer_step && slg_uncond) {
LOG_VERBOSE("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,
nullptr,
true);
if (uncond_out.empty()) {
return {};
if (!skip_uncond) {
const std::vector<int>* uncond_skip_layers = nullptr;
if (is_skiplayer_step && slg_uncond) {
LOG_VERBOSE("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,
nullptr,
true);
if (uncond_out.empty()) {
return {};
}
} else {
step_cache.invalidate_condition(&uncond);
}
}
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,
true);
if (img_uncond_out.empty()) {
return {};
if (!skip_img_uncond) {
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,
true);
if (img_uncond_out.empty()) {
return {};
}
} else {
step_cache.invalidate_condition(&img_uncond);
}
}
sd::guidance::GuidanceInput guidance_input;
@ -2770,7 +2922,7 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
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 = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]);
sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {}, effective_guidance_scale);
if (guided.pred.empty()) {
return {};
}
@ -2984,7 +3136,8 @@ sd::Tensor<float> StableDiffusionGGML::decode_first_stage(const sd::Tensor<float
auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
const bool prefer_temporal_tiling = decode_video && first_stage_model->can_temporal_tile_decode();
while (decoded.empty() &&
sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling,
first_stage_model->last_compute_status())) {
decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
}
return decoded;
@ -3047,6 +3200,8 @@ std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version)
return "mage_flow";
} else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
return "z_image_omni";
} else if (sd_version_is_llada_image(version)) {
return "llada_image";
} else if (sd_version_is_krea2(version)) {
// have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit)
// and "krea2_ostris_edit" (for krea2 ostris edit)

6
otherarch/sdcpp/src/pipeline/diffusion_engine.h
View file

@ -26,6 +26,7 @@ class RNG;
struct Denoiser;
struct LoraModel;
struct ConditionerParams;
class ConditioningCache;
struct SDCondition;
struct RefImageParams;
namespace Wav2Vec2 {
@ -184,6 +185,8 @@ public:
std::recursive_mutex execution_mutex;
std::unique_ptr<ModelConfig> config_;
RunnerState runner_state_;
std::unique_ptr<ConditioningCache> conditioning_cache_;
std::vector<ModelManager::LoraSpec> conditioning_loras_;
bool executing_ = false;
std::shared_ptr<Denoiser> denoiser;
@ -318,6 +321,7 @@ public:
bool init_model_loader(ModelLoader& model_loader, ModelConfig& configuration);
bool init(const sd_ctx_params_t* sd_ctx_params);
bool set_sage_attention_enabled(bool enabled);
bool uses_tae() const;
@ -366,6 +370,8 @@ public:
bool apply_loras(const sd_lora_t* loras, uint32_t lora_count);
SDCondition get_learned_condition(const ConditionerParams& params);
void reset_generation_extensions();
void prepare_generation_extensions(const sd_pm_params_t& pm_params,

32
otherarch/sdcpp/src/pipeline/image.cpp
View file

@ -10,6 +10,7 @@
#include "model/vae/vae.hpp"
#include "request.h"
#include "runtime/denoiser.hpp"
#include "runtime/image_preprocess.h"
#include "upscaler.h"
namespace sd::pipeline {
@ -440,8 +441,7 @@ namespace sd::pipeline {
sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength);
int64_t prepare_start_ms = ggml_time_ms();
condition_params.zero_out_masked = false;
auto cond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
auto cond = sd->get_learned_condition(condition_params);
if (cond.empty()) {
LOG_ERROR("failed to encode prompt");
return std::nullopt;
@ -475,8 +475,11 @@ namespace sd::pipeline {
}
condition_params.text = request->negative_prompt;
condition_params.zero_out_masked = zero_out_masked;
uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
if (sd_version_is_llada_image(sd->version)) {
// LLaDA-Image CFG keeps the source latent but drops its SigVQ features.
condition_params.ref_images = nullptr;
}
uncond = sd->get_learned_condition(condition_params);
if (uncond.empty()) {
LOG_ERROR("failed to encode negative prompt");
return std::nullopt;
@ -504,8 +507,7 @@ namespace sd::pipeline {
if (use_ref_latent_img_cfg) {
condition_params.ref_images = &empty_ref_images;
}
img_uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
img_uncond = sd->get_learned_condition(condition_params);
if (img_uncond.empty()) {
LOG_ERROR("failed to encode image guidance prompt");
return std::nullopt;
@ -785,15 +787,9 @@ namespace sd::pipeline {
return false;
}
// MiniMax-H3 is video-only. Its denoiser always splits the packed latent into a video and an
// audio half, and only generate_video ever computes the audio length, so reaching this
// function with an H3 checkpoint is guaranteed to die on
// GGML_ASSERT(!audio_input_cache.empty()) with a core dump, after the several minutes it
// takes to load the weights, and with nothing in the output pointing at the missing --mode.
// (The AnimateDiff path below routes vid_gen back through here, but that is SD1.5 plus a
// motion module, never H3.)
if (sd_version_is_minimax_h3(sd->version)) {
LOG_ERROR("MiniMax-H3 is a video model and cannot be run in img_gen mode; use --mode vid_gen");
if (!sd_version_supports_image_generation(sd->version)) {
LOG_ERROR("%s cannot be run with generate_image(); use generate_video() or --mode vid_gen in the CLI",
model_version_to_str[sd->version]);
return false;
}
@ -802,6 +798,12 @@ namespace sd::pipeline {
int64_t t0 = ggml_time_ms();
sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params;
GenerationRequest request(sd, sd_img_gen_params);
sd::ImagePreprocessor preprocessing(sd_img_gen_params->image_preprocess.rules);
sd_img_gen_params_t processed_params = *sd_img_gen_params;
if (!preprocessing.prepare_inputs(processed_params, request.width, request.height))
return false;
sd_img_gen_params = &processed_params;
request.pm_params = processed_params.pm_params;
LOG_INFO("generate_image %dx%d", request.width, request.height);
sd->rng->manual_seed(request.seed);

17
otherarch/sdcpp/src/pipeline/model_builders.cpp
View file

@ -20,6 +20,7 @@
#include "model/diffusion/krea2.hpp"
#include "model/diffusion/lens.hpp"
#include "model/diffusion/lingbot_video.hpp"
#include "model/diffusion/llada_image.hpp"
#include "model/diffusion/ltxv.hpp"
#include "model/diffusion/mage_flow.hpp"
#include "model/diffusion/minimax_h3.hpp"
@ -290,7 +291,8 @@ namespace sd::model_builders {
result.diffusion = std::make_shared<Qwen::QwenImage21Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
weight_manager,
sd_ctx_params->model_args);
} else {
result.diffusion = std::make_shared<Qwen::QwenImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
@ -370,6 +372,19 @@ namespace sd::model_builders {
"model.diffusion_model",
version,
weight_manager);
} else if (sd_version_is_llada_image(version)) {
result.conditioner = std::make_shared<LLaDAImageEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,
"text_encoders.llm",
"queryformer",
"text_projection",
"sigvq",
weight_manager,
tokenizers);
result.diffusion = std::make_shared<LLaDAImage::LLaDAImageRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
tensor_storage_map,
"model.diffusion_model",
weight_manager);
} else if (sd_version_is_boogu_image(version)) {
result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
tensor_storage_map,

2
otherarch/sdcpp/src/pipeline/request.cpp
View file

@ -73,6 +73,8 @@ namespace sd::pipeline {
return LTX2_SCHEDULER;
} else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) {
return LOGIT_NORMAL_SCHEDULER;
} else if (sd != nullptr && sd_version_is_llada_image(sd->version)) {
return LLADA_IMAGE_SCHEDULER;
}
return DISCRETE_SCHEDULER;
}

29
otherarch/sdcpp/src/pipeline/video.cpp
View file

@ -15,6 +15,7 @@
#include "model/vae/vae.hpp"
#include "request.h"
#include "runtime/denoiser.hpp"
#include "runtime/image_preprocess.h"
// kcpp
#include "model/vae/ltx_audio_vae.hpp"
@ -555,11 +556,15 @@ namespace sd::pipeline {
} runner_guard{sd->audio_vae_model.get()};
if (sd_vid_gen_params->init_image.data) {
start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height);
start_image = ensure_image_tensor_channels(
sd_image_to_tensor(sd_vid_gen_params->init_image, request->width, request->height),
sd->get_image_channels());
}
if (sd_vid_gen_params->end_image.data) {
end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height);
end_image = ensure_image_tensor_channels(
sd_image_to_tensor(sd_vid_gen_params->end_image, request->width, request->height),
sd->get_image_channels());
}
if (sd_version_is_minimax_h3(sd->version)) {
@ -1289,8 +1294,7 @@ namespace sd::pipeline {
}
int64_t prepare_start_ms = ggml_time_ms();
embeds.cond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
embeds.cond = sd->get_learned_condition(condition_params);
if (embeds.cond.empty()) {
LOG_ERROR("failed to encode video prompt");
return std::nullopt;
@ -1315,8 +1319,7 @@ namespace sd::pipeline {
}
if (request.use_uncond) {
condition_params.text = request.negative_prompt;
embeds.uncond = sd->cond_stage_model->get_learned_condition(sd->n_threads,
condition_params);
embeds.uncond = sd->get_learned_condition(condition_params);
if (embeds.uncond.empty()) {
LOG_ERROR("failed to encode negative video prompt");
return std::nullopt;
@ -1548,7 +1551,9 @@ namespace sd::pipeline {
sd::Tensor<float> video_mask = make_ltxav_video_denoise_mask(video_latent, 1.f);
if (sd_vid_gen_params->init_image.data != nullptr) {
sd::Tensor<float> start_image = sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height);
sd::Tensor<float> start_image = ensure_image_tensor_channels(
sd_image_to_tensor(sd_vid_gen_params->init_image, image_width, image_height),
sd->get_image_channels());
if (!apply_ltxav_condition_image_by_latent_index(sd,
start_image,
&video_latent,
@ -1561,7 +1566,9 @@ namespace sd::pipeline {
}
if (sd_vid_gen_params->end_image.data != nullptr) {
sd::Tensor<float> end_image = sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height);
sd::Tensor<float> end_image = ensure_image_tensor_channels(
sd_image_to_tensor(sd_vid_gen_params->end_image, image_width, image_height),
sd->get_image_channels());
sd::Tensor<float> end_image_latent = encode_ltxav_condition_image(sd, end_image, "end");
if (end_image_latent.empty()) {
return false;
@ -1654,6 +1661,7 @@ namespace sd::pipeline {
img_gen_params.qwen_image_layers = 0;
img_gen_params.circular_x = sd_vid_gen_params->circular_x;
img_gen_params.circular_y = sd_vid_gen_params->circular_y;
img_gen_params.image_preprocess = sd_vid_gen_params->image_preprocess;
sd->animatediff_num_frames = n_frames;
bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out);
@ -1684,6 +1692,11 @@ namespace sd::pipeline {
sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params;
sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y);
GenerationRequest request(sd, sd_vid_gen_params);
sd::ImagePreprocessor preprocessing(sd_vid_gen_params->image_preprocess.rules);
sd_vid_gen_params_t processed_params = *sd_vid_gen_params;
if (!preprocessing.prepare_inputs(processed_params, request.width, request.height))
return false;
sd_vid_gen_params = &processed_params;
if (fps_out != nullptr) {
*fps_out = request.fps;
}

89
otherarch/sdcpp/src/runtime/denoiser.hpp
View file

@ -728,12 +728,20 @@ inline float flux_time_shift(float mu, float sigma, float t) {
// https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289
struct FluxScheduler : SigmaScheduler {
int image_seq_len = 0;
float base_shift = 0.5f;
float max_shift = 1.15f;
int image_seq_len = 0;
int base_image_seq_len = 256;
int max_image_seq_len = 4096;
float base_shift = 0.5f;
float max_shift = 1.15f;
float shift_terminal = 0.0f;
explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr)
FluxScheduler(int image_seq_len, SDVersion version, const char* extra_sample_args = nullptr)
: image_seq_len(image_seq_len) {
if (version == VERSION_QWEN_IMAGE_2_1) {
max_image_seq_len = 8192;
max_shift = 0.9f;
shift_terminal = 0.02f;
}
parse_extra_sample_args(extra_sample_args);
}
@ -752,10 +760,8 @@ struct FluxScheduler : SigmaScheduler {
}
float compute_mu() const {
constexpr float base_shift_anchor = 256.0f;
constexpr float max_shift_anchor = 4096.0f;
float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
float b = base_shift - m * base_shift_anchor;
float m = (max_shift - base_shift) / static_cast<float>(max_image_seq_len - base_image_seq_len);
float b = base_shift - m * static_cast<float>(base_image_seq_len);
return static_cast<float>(image_seq_len) * m + b;
}
@ -764,7 +770,7 @@ struct FluxScheduler : SigmaScheduler {
sigmas.reserve(n + 1);
float mu = compute_mu();
LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu);
LOG_VERBOSE("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f, shift_terminal=%.3f", image_seq_len, n, mu, shift_terminal);
if (n == 0) {
sigmas.push_back(1.0f);
@ -780,12 +786,70 @@ struct FluxScheduler : SigmaScheduler {
}
}
if (shift_terminal > 0.0f && n > 1) {
// The terminal shift applies to the last model evaluation, not the final zero sigma.
float scale_factor = (1.0f - sigmas[n - 1]) / (1.0f - shift_terminal);
if (std::isfinite(scale_factor) && scale_factor > 0.0f) {
for (uint32_t i = 0; i < n; ++i) {
sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor;
}
}
}
sigmas[n] = 0.0f;
return sigmas;
}
};
// https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244
// LLaDA-Image does not use a shift-based flow schedule. The reference pipeline builds a
// Kumaraswamy-shaped grid over t = linspace(0.001, 1, n + 1)[:-1]:
// schedule = (1 - (1 - t^1.17)^0.8)^1.1
// sigma = 1 - schedule
// Its scheduler config can also set use_uniform_sigmas, which replaces the whole curve with a
// plain linspace(1, 0, n + 1)[:-1] pre-shift grid.
struct LLaDAImageScheduler : SigmaScheduler {
bool uniform_sigmas = false;
explicit LLaDAImageScheduler(const char* extra_sample_args = nullptr) {
parse_extra_sample_args(extra_sample_args);
}
void parse_extra_sample_args(const char* extra_sample_args) {
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "llada_image scheduler arg")) {
if (key == "uniform") {
if (!parse_strict_bool(value, uniform_sigmas)) {
LOG_WARN("ignoring invalid llada_image scheduler arg '%s=%s'", key.c_str(), value.c_str());
}
}
}
}
std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
std::vector<float> sigmas;
sigmas.reserve(n + 1);
if (n == 0) {
sigmas.push_back(1.0f);
return sigmas;
}
for (uint32_t i = 0; i < n; ++i) {
float progress = static_cast<float>(i) / static_cast<float>(n);
if (uniform_sigmas) {
sigmas.push_back(1.0f - progress);
} else {
float t = 0.001f + progress * (1.0f - 0.001f);
float schedule = powf(1.0f - powf(1.0f - powf(t, 1.17f), 0.8f), 1.1f);
sigmas.push_back(1.0f - schedule);
}
}
sigmas.push_back(0.0f);
return sigmas;
}
};
struct Flux2Scheduler : SigmaScheduler {
int image_seq_len = 0;
@ -1123,9 +1187,14 @@ struct Denoiser {
scheduler = std::make_shared<Flux2Scheduler>(image_seq_len);
break;
}
case LLADA_IMAGE_SCHEDULER: {
LOG_INFO("get_sigmas with LLaDA-Image scheduler");
scheduler = std::make_shared<LLaDAImageScheduler>(extra_sample_args);
break;
}
case FLUX_SCHEDULER: {
LOG_INFO("get_sigmas with Flux scheduler");
scheduler = std::make_shared<FluxScheduler>(image_seq_len, extra_sample_args);
scheduler = std::make_shared<FluxScheduler>(image_seq_len, version, extra_sample_args);
break;
}
default:

447
otherarch/sdcpp/src/runtime/image_preprocess.cpp Normal file
View file

@ -0,0 +1,447 @@
#include "image_preprocess.h"
#include "core/util.h"
#include <climits>
#include <set>
namespace sd {
static constexpr std::pair<const char*, ImageTarget> image_targets[] = {
{"init", ImageTarget::Init},
{"end", ImageTarget::End},
{"mask", ImageTarget::Mask},
{"control", ImageTarget::Control},
{"ref", ImageTarget::Ref},
{"ip-adapter", ImageTarget::IPAdapter},
{"id", ImageTarget::ID},
{"control-frame", ImageTarget::ControlFrame},
};
static constexpr std::pair<const char*, ImageResizeMode> image_resize_modes[] = {
{"auto", ImageResizeMode::Auto},
{"none", ImageResizeMode::None},
{"stretch", ImageResizeMode::Stretch},
{"crop", ImageResizeMode::Crop},
{"crop-resize", ImageResizeMode::CropResize},
{"fit-pad", ImageResizeMode::FitPad},
};
template <typename T, size_t N>
static bool parse_enum(const std::string& text, const std::pair<const char*, T> (&names)[N], T& value) {
for (const auto& entry : names) {
if (text == entry.first) {
value = entry.second;
return true;
}
}
return false;
}
template <typename T, size_t N>
static const char* enum_name(T value, const std::pair<const char*, T> (&names)[N]) {
for (const auto& entry : names) {
if (value == entry.second)
return entry.first;
}
return "unknown";
}
template <typename T>
static bool one_of(T value, std::initializer_list<T> choices) {
return std::find(choices.begin(), choices.end(), value) != choices.end();
}
static bool one_of(const std::string& value, std::initializer_list<const char*> choices) {
for (const char* choice : choices) {
if (value == choice)
return true;
}
return false;
}
static ImageResizeMode resolve_mode(const std::map<std::string, std::string>& options, ImageResizeMode default_mode) {
auto it = options.find("mode");
ImageResizeMode mode = ImageResizeMode::Auto;
if (it != options.end())
parse_enum(it->second, image_resize_modes, mode);
if (mode != ImageResizeMode::Auto)
return mode;
return options.count("width") && default_mode == ImageResizeMode::None ? ImageResizeMode::Stretch : default_mode;
}
bool ImagePreprocessor::fail(const std::string& message) const {
LOG_ERROR("image preprocessing: %s", message.c_str());
valid_ = false;
return false;
}
ImagePreprocessor::ImagePreprocessor(const char* text) {
if (text == nullptr || trim(text).empty())
return;
for (const auto& part : split_string(text, ';')) {
ImagePreprocessRule rule;
std::set<std::string> keys;
if (trim(part).empty()) {
fail("empty rule");
return;
}
for (const auto& entry : split_string(part, ',')) {
size_t equal = entry.find('=');
if (equal == std::string::npos) {
fail("expected key=value: " + entry);
return;
}
std::string key = trim(entry.substr(0, equal));
std::string value = trim(entry.substr(equal + 1));
bool ok = !value.empty() && keys.insert(key).second;
int number = 0;
if (key == "target") {
ok &= parse_enum(value, image_targets, rule.target);
} else if (key == "index") {
ok &= parse_strict_int(value, rule.index) && rule.index >= 0;
} else {
if (key == "mode") {
ImageResizeMode mode;
ok &= parse_enum(value, image_resize_modes, mode);
} else if (key == "filter") {
ok &= one_of(value, {"auto", "nearest", "nearest-exact", "bilinear", "bicubic", "lanczos"});
} else if (key == "antialias") {
ok &= one_of(value, {"auto", "true", "false"});
} else if (key == "canny") {
ok &= one_of(value, {"true", "false"});
} else if (key == "anchor") {
ok &= one_of(value, {"center", "top", "bottom", "left", "right"});
} else if (key == "width" || key == "height") {
ok &= parse_strict_int(value, number) && number > 0;
} else if (key == "pad_color") {
ok &= value.size() == 7 || value.size() == 9;
ok &= !value.empty() && value[0] == '#';
for (size_t i = 1; i < value.size(); ++i)
ok &= std::isxdigit(static_cast<unsigned char>(value[i])) != 0;
} else {
ok = false;
}
rule.options[key] = value;
}
if (!ok) {
fail("invalid or duplicate option: " + entry);
return;
}
}
if (!keys.count("target") || rule.options.empty() ||
rule.options.count("width") != rule.options.count("height") ||
(rule.index >= 0 && !one_of(rule.target, {ImageTarget::Ref, ImageTarget::ID, ImageTarget::ControlFrame}))) {
fail("invalid target, index, or incomplete dimensions: " + part);
return;
}
rules_.push_back(std::move(rule));
}
for (const auto& rule : rules_) {
const auto options = resolve_options(rule.target, std::max(0, rule.index));
if (options.count("antialias") && options.at("antialias") == "true" && options.count("filter") &&
one_of(options.at("filter"), {"nearest", "nearest-exact"})) {
fail("antialias requires bilinear, bicubic, or lanczos");
return;
}
}
}
std::map<std::string, std::string> ImagePreprocessor::resolve_options(ImageTarget target, int index) const {
std::map<std::string, std::string> options;
for (int specificity = 0; specificity < 2; ++specificity) {
for (const auto& rule : rules_) {
if (rule.target == target &&
rule.index == (specificity == 0 ? -1 : index)) {
for (const auto& entry : rule.options)
options[entry.first] = entry.second;
}
}
}
return options;
}
bool ImagePreprocessor::validate_inputs(const sd_img_gen_params_t& params) const {
const std::map<ImageTarget, int> counts = {
{ImageTarget::Init, params.init_image.data != nullptr},
{ImageTarget::Mask, params.mask_image.data != nullptr},
{ImageTarget::Control, params.control_image.data != nullptr},
{ImageTarget::IPAdapter, params.ip_adapter_image.data != nullptr},
{ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0},
{ImageTarget::ID, params.pm_params.id_images != nullptr ? params.pm_params.id_images_count : 0},
};
for (const auto& rule : rules_) {
auto it = counts.find(rule.target);
int count = it == counts.end() ? 0 : it->second;
if (count <= 0 || rule.index >= count) {
return fail(std::string("rule targets an unavailable image: ") + enum_name(rule.target, image_targets));
}
}
return valid_;
}
bool ImagePreprocessor::validate_inputs(const sd_vid_gen_params_t& params) const {
const std::map<ImageTarget, int> counts = {
{ImageTarget::Init, params.init_image.data != nullptr},
{ImageTarget::End, params.end_image.data != nullptr},
{ImageTarget::Ref, params.ref_images != nullptr ? params.ref_images_count : 0},
{ImageTarget::ControlFrame, params.control_frames != nullptr ? params.control_frames_size : 0},
};
for (const auto& rule : rules_) {
auto it = counts.find(rule.target);
int count = it == counts.end() ? 0 : it->second;
if (count <= 0 || rule.index >= count)
return fail(std::string("rule targets an unavailable video input: ") + enum_name(rule.target, image_targets));
}
return valid_;
}
static int anchor_offset(int remaining, const std::string& anchor, bool horizontal) {
if (anchor == (horizontal ? "left" : "top"))
return 0;
if (anchor == (horizontal ? "right" : "bottom"))
return remaining;
return remaining / 2;
}
Tensor<float> ImagePreprocessor::apply_transform(const Tensor<float>& image, const std::map<std::string, std::string>& options, ImageTransform p, const std::string& label, ops::InterpolateMode default_filter) const {
auto value = [&](const char* key, const char* fallback) {
auto it = options.find(key);
return it == options.end() ? std::string(fallback) : it->second;
};
std::string filter = value("filter", "auto");
ops::InterpolateMode mode = default_filter;
if (filter == "nearest")
mode = ops::InterpolateMode::Nearest;
if (filter == "nearest-exact")
mode = ops::InterpolateMode::NearestExact;
if (filter == "bilinear")
mode = ops::InterpolateMode::Bilinear;
if (filter == "bicubic")
mode = ops::InterpolateMode::Bicubic;
if (filter == "lanczos")
mode = ops::InterpolateMode::Lanczos;
bool filtered = ops::is_2d_filter_interpolate_mode(mode);
bool antialias = value("antialias", "auto") == "true" ||
(value("antialias", "auto") == "auto" && filtered &&
(p.resize_width < p.crop_width || p.resize_height < p.crop_height));
if (antialias && !filtered) {
fail(label + ": antialias requires bilinear, bicubic, or lanczos");
return {};
}
auto cropped = ops::slice(ops::slice(image, 0, p.x, p.x + p.crop_width), 1, p.y, p.y + p.crop_height);
int channels = static_cast<int>(image.shape()[2]);
bool resize = p.resize_width != p.crop_width || p.resize_height != p.crop_height;
if (resize && channels == 4 && filtered) {
for (int64_t i = 0, pixels = cropped.shape()[0] * cropped.shape()[1]; i < pixels; ++i) {
for (int c = 0; c < 3; ++c)
cropped[i + c * pixels] *= cropped[i + 3 * pixels];
}
}
auto resized = ops::interpolate(cropped, {p.resize_width, p.resize_height, channels, 1}, mode, false, antialias);
if (resize && channels == 4 && filtered) {
for (int64_t i = 0, pixels = resized.shape()[0] * resized.shape()[1]; i < pixels; ++i) {
float alpha = std::clamp(resized[i + 3 * pixels], 0.f, 1.f);
for (int c = 0; c < 3; ++c)
resized[i + c * pixels] = alpha > 1e-6f ? resized[i + c * pixels] / alpha : 0.f;
}
}
resized = ops::clamp(resized, 0.f, 1.f);
Tensor<float> output({p.width, p.height, channels, 1});
std::string color = value("pad_color", "#000000ff");
if (color.size() == 7)
color += "ff";
uint8_t rgba[4];
for (int c = 0; c < 4; ++c)
rgba[c] = static_cast<uint8_t>(std::strtoul(color.substr(1 + c * 2, 2).c_str(), nullptr, 16));
for (int c = 0; c < channels; ++c) {
float fill = rgba[channels == 1 ? 0 : c] / 255.f;
for (int y = 0; y < p.height; ++y) {
for (int x = 0; x < p.width; ++x) {
output.index(x, y, c, 0) = x >= p.pad_x && x < p.pad_x + p.resize_width && y >= p.pad_y && y < p.pad_y + p.resize_height
? resized.index(x - p.pad_x, y - p.pad_y, c, 0)
: fill;
}
}
}
LOG_INFO("preprocess %s: %dx%d crop=(%d,%d,%d,%d) resize=%dx%d pad=(%d,%d) output=%dx%d filter=%s(%d) antialias=%s",
label.c_str(), p.source_width, p.source_height, p.x, p.y, p.crop_width, p.crop_height,
p.resize_width, p.resize_height, p.pad_x, p.pad_y, p.width, p.height, filter.c_str(), static_cast<int>(mode), BOOL_STR(antialias));
return output;
}
Tensor<float> ImagePreprocessor::apply_geometry(const Tensor<float>& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode, ops::InterpolateMode default_filter, ImageTransform* plan_out) const {
if (!valid_ || image.empty())
return {};
const std::string label = std::string(enum_name(target, image_targets)) + "[" + std::to_string(index) + "]";
auto options = resolve_options(target, index);
if (image.dim() != 4 || image.shape()[3] != 1 || image.shape()[2] < 1 || image.shape()[2] > 4) {
fail(label + ": expected one image with 1 to 4 channels");
return {};
}
ImageTransform p;
p.source_width = p.crop_width = static_cast<int>(image.shape()[0]);
p.source_height = p.crop_height = static_cast<int>(image.shape()[1]);
int target_width = width > 0 ? width : p.source_width;
int target_height = height > 0 ? height : p.source_height;
if (options.count("width")) {
parse_strict_int(options.at("width"), target_width);
parse_strict_int(options.at("height"), target_height);
}
ImageResizeMode mode = resolve_mode(options, default_mode);
std::string anchor = options.count("anchor") ? options.at("anchor") : "center";
p.width = p.resize_width = target_width;
p.height = p.resize_height = target_height;
if (mode == ImageResizeMode::None) {
if (options.count("width") && (target_width != p.source_width || target_height != p.source_height)) {
fail(label + ": mode=none conflicts with requested dimensions");
return {};
}
p.width = p.resize_width = p.source_width;
p.height = p.resize_height = p.source_height;
} else if (mode == ImageResizeMode::Crop || mode == ImageResizeMode::CropResize) {
if (mode == ImageResizeMode::Crop) {
p.crop_width = target_width;
p.crop_height = target_height;
} else if (int64_t(p.source_width) * target_height > int64_t(p.source_height) * target_width) {
p.crop_width = std::max(1, static_cast<int>(int64_t(p.source_height) * target_width / target_height));
} else {
p.crop_height = std::max(1, static_cast<int>(int64_t(p.source_width) * target_height / target_width));
}
if (p.crop_width > p.source_width || p.crop_height > p.source_height) {
fail(label + ": crop exceeds source dimensions");
return {};
}
p.x = anchor_offset(p.source_width - p.crop_width, anchor, true);
p.y = anchor_offset(p.source_height - p.crop_height, anchor, false);
} else if (mode == ImageResizeMode::FitPad) {
double scale = std::min(double(target_width) / p.source_width, double(target_height) / p.source_height);
p.resize_width = std::max(1, std::min(target_width, static_cast<int>(std::round(p.source_width * scale))));
p.resize_height = std::max(1, std::min(target_height, static_cast<int>(std::round(p.source_height * scale))));
p.pad_x = anchor_offset(target_width - p.resize_width, anchor, true);
p.pad_y = anchor_offset(target_height - p.resize_height, anchor, false);
}
if (p.width <= 0 || p.height <= 0) {
fail(label + ": invalid output dimensions");
return {};
}
uint64_t max_pixels = std::min<uint64_t>(INT64_MAX, SIZE_MAX / sizeof(float)) / static_cast<uint64_t>(image.shape()[2]);
if (uint64_t(p.width) * p.height > max_pixels || uint64_t(p.resize_width) * p.resize_height > max_pixels) {
fail(label + ": image allocation size overflows");
return {};
}
if (plan_out != nullptr)
*plan_out = p;
return apply_transform(image, options, p, label, default_filter);
}
Tensor<float> ImagePreprocessor::preprocess_input(sd_image_t image, ImageTarget target, int index, int width, int height) {
if (image.data == nullptr || image.width == 0 || image.height == 0 || image.width > INT_MAX || image.height > INT_MAX || image.channel < 1 || image.channel > 4) {
fail(std::string(enum_name(target, image_targets)) + ": invalid input image");
return {};
}
auto tensor = sd_image_to_tensor(image);
if (target == ImageTarget::Mask && has_init_transform_) {
auto options = resolve_options(target, index);
if (image.width != init_transform_.source_width || image.height != init_transform_.source_height) {
fail("mask and init source dimensions must match");
return {};
}
bool geometry_override = options.count("width") || options.count("anchor") ||
(options.count("mode") && options.at("mode") != "auto");
if (geometry_override) {
ImageTransform p;
auto init_options = resolve_options(ImageTarget::Init, 0);
ImageResizeMode default_mode = resolve_mode(init_options, ImageResizeMode::CropResize);
auto result = apply_geometry(tensor, target, index, init_transform_.width, init_transform_.height, default_mode, ops::InterpolateMode::NearestExact, &p);
if (result.empty())
return {};
const auto& q = init_transform_;
if (p.x != q.x || p.y != q.y || p.crop_width != q.crop_width || p.crop_height != q.crop_height ||
p.resize_width != q.resize_width || p.resize_height != q.resize_height || p.pad_x != q.pad_x || p.pad_y != q.pad_y || p.width != q.width || p.height != q.height) {
fail("mask geometry conflicts with init; configure geometry on init and filter on mask");
return {};
}
return result;
}
return apply_transform(tensor, options, init_transform_, "mask[0]", ops::InterpolateMode::NearestExact);
}
auto result = apply_geometry(tensor, target, index, width, height, width > 0 ? ImageResizeMode::CropResize : ImageResizeMode::None,
target == ImageTarget::Mask ? ops::InterpolateMode::NearestExact : ops::InterpolateMode::Nearest,
target == ImageTarget::Init ? &init_transform_ : nullptr);
if (target == ImageTarget::Init)
has_init_transform_ = !result.empty();
return result;
}
ImagePreprocessor::~ImagePreprocessor() {
for (const auto& image : owned_images_)
std::free(image.data);
}
bool ImagePreprocessor::prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height) {
if (image.data == nullptr)
return true;
auto options = resolve_options(target, index);
bool canny = options.count("canny") && options.at("canny") == "true";
auto tensor = preprocess_input(image, target, index, width, height);
if (tensor.empty())
return false;
auto output = tensor_to_sd_image(tensor);
if (output.data == nullptr)
return fail("could not allocate input preprocessing buffer");
owned_images_.push_back(output);
if (canny && !preprocess_canny(output, 0.08f, 0.08f, 0.8f, 1.f, false))
return fail("Canny preprocessing failed");
image = output;
return true;
}
bool ImagePreprocessor::prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector<sd_image_t>& storage, int width, int height) {
if (count < 0 || (count > 0 && images == nullptr))
return fail(std::string("invalid image array: ") + enum_name(target, image_targets));
if (count == 0)
return true;
storage.assign(images, images + count);
for (int i = 0; i < count; ++i) {
if (storage[i].data == nullptr)
return fail(std::string("empty image in array: ") + enum_name(target, image_targets));
if (!prepare_image(storage[i], target, i, width, height))
return false;
}
images = storage.data();
return true;
}
bool ImagePreprocessor::prepare_inputs(sd_img_gen_params_t& params, int width, int height) {
if (prepared_)
return fail("inputs have already been prepared");
prepared_ = true;
if (!valid_ || !validate_inputs(params))
return false;
if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) ||
!prepare_image(params.mask_image, ImageTarget::Mask, 0, width, height) ||
!prepare_image(params.control_image, ImageTarget::Control, 0, width, height) ||
!prepare_image(params.ip_adapter_image, ImageTarget::IPAdapter, 0, -1, -1) ||
!prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) ||
!prepare_array(params.pm_params.id_images, params.pm_params.id_images_count, ImageTarget::ID, id_images_))
return false;
params.image_preprocess = {};
return true;
}
bool ImagePreprocessor::prepare_inputs(sd_vid_gen_params_t& params, int width, int height) {
if (prepared_)
return fail("inputs have already been prepared");
prepared_ = true;
if (!valid_ || !validate_inputs(params))
return false;
if (!prepare_image(params.init_image, ImageTarget::Init, 0, width, height) ||
!prepare_image(params.end_image, ImageTarget::End, 0, width, height) ||
!prepare_array(params.ref_images, params.ref_images_count, ImageTarget::Ref, ref_images_) ||
!prepare_array(params.control_frames, params.control_frames_size, ImageTarget::ControlFrame, control_frames_, width, height))
return false;
params.image_preprocess = {};
return true;
}
} // namespace sd

88
otherarch/sdcpp/src/runtime/image_preprocess.h Normal file
View file

@ -0,0 +1,88 @@
#ifndef __SD_RUNTIME_IMAGE_PREPROCESS_H__
#define __SD_RUNTIME_IMAGE_PREPROCESS_H__
#include <map>
#include <string>
#include <vector>
#include "core/tensor.hpp"
#include "stable-diffusion.h"
namespace sd {
enum class ImageTarget {
Init,
End,
Mask,
Control,
Ref,
IPAdapter,
ID,
ControlFrame,
};
enum class ImageResizeMode {
Auto,
None,
Stretch,
Crop,
CropResize,
FitPad,
};
struct ImageTransform {
int source_width = 0;
int source_height = 0;
int x = 0;
int y = 0;
int crop_width = 0;
int crop_height = 0;
int resize_width = 0;
int resize_height = 0;
int width = 0;
int height = 0;
int pad_x = 0;
int pad_y = 0;
};
struct ImagePreprocessRule {
ImageTarget target = ImageTarget::Init;
int index = -1;
std::map<std::string, std::string> options;
};
class ImagePreprocessor {
std::vector<ImagePreprocessRule> rules_;
mutable bool valid_ = true;
ImageTransform init_transform_;
bool has_init_transform_ = false;
bool prepared_ = false;
std::vector<sd_image_t> owned_images_;
std::vector<sd_image_t> ref_images_;
std::vector<sd_image_t> id_images_;
std::vector<sd_image_t> control_frames_;
bool fail(const std::string& message) const;
std::map<std::string, std::string> resolve_options(ImageTarget target, int index) const;
Tensor<float> apply_transform(const Tensor<float>& image, const std::map<std::string, std::string>& options, ImageTransform plan, const std::string& label, ops::InterpolateMode default_filter) const;
bool prepare_image(sd_image_t& image, ImageTarget target, int index, int width, int height);
bool prepare_array(sd_image_t*& images, int count, ImageTarget target, std::vector<sd_image_t>& storage, int width = -1, int height = -1);
public:
explicit ImagePreprocessor(const char* rules = nullptr);
~ImagePreprocessor();
ImagePreprocessor(const ImagePreprocessor&) = delete;
ImagePreprocessor& operator=(const ImagePreprocessor&) = delete;
bool prepare_inputs(sd_img_gen_params_t& params, int width, int height);
bool prepare_inputs(sd_vid_gen_params_t& params, int width, int height);
bool is_valid() const { return valid_; }
bool validate_inputs(const sd_img_gen_params_t& params) const;
bool validate_inputs(const sd_vid_gen_params_t& params) const;
Tensor<float> apply_geometry(const Tensor<float>& image, ImageTarget target, int index, int width, int height, ImageResizeMode default_mode = ImageResizeMode::Stretch, ops::InterpolateMode default_filter = ops::InterpolateMode::Nearest, ImageTransform* plan_out = nullptr) const;
Tensor<float> preprocess_input(sd_image_t image, ImageTarget target, int index = 0, int width = -1, int height = -1);
};
} // namespace sd
#endif // __SD_RUNTIME_IMAGE_PREPROCESS_H__

124
otherarch/sdcpp/src/runtime/latent-preview.h
View file

@ -4,6 +4,86 @@
#include "core/tensor.hpp"
#include "ggml.h"
// RGB is projected to [-1, 1]; alpha is projected directly to [0, 1].
const float qwen21_latent_rgb_proj[64][3] = {
{0.00860495522f, 0.01219501462f, -0.00321337196f},
{0.01889233090f, 0.01246581216f, 0.01074959482f},
{0.1255941446f, 0.1176879344f, -0.0332212352f},
{0.0418238528f, 0.1043427754f, 0.0121666316f},
{0.02025338f, 0.01453670296f, -0.000224336044f},
{-0.01896720702f, -0.0206099030f, -0.0322728584f},
{0.00438984796f, -0.01374969766f, 0.02849196f},
{-0.0374495856f, -0.0286777126f, -0.0693192810f},
{0.01511914734f, 0.0242979386f, 0.0553878870f},
{-0.1138629518f, -0.020391466f, 0.001550520522f},
{-0.0233650696f, -0.0417292018f, -0.0362361182f},
{-0.0351603342f, -0.0243595924f, -0.00216261038f},
{0.01093355288f, -0.0373466924f, 0.00241315350f},
{0.01778704744f, -0.00401984678f, -0.0343259192f},
{0.0486059334f, 0.0253144f, 0.0672564966f},
{0.0309463558f, 0.0277963166f, 0.0520869622f},
{0.0374485008f, 0.0551753676f, 0.0225853902f},
{-0.0090809962f, -0.004756176f, 0.00636443612f},
{-0.0270455652f, -0.0384966954f, -0.00905908082f},
{-0.00553493756f, 0.01484553684f, -0.0211502468f},
{0.01319502562f, 0.00948005666f, 0.0483789212f},
{-0.00931847104f, -0.00276452734f, -0.01011985302f},
{0.0180478258f, 0.01614954356f, -0.0209424690f},
{-0.0214530434f, -0.00272961176f, 0.0217887476f},
{-0.0636772304f, -0.0208893548f, 0.0479167742f},
{-0.0250321236f, -0.0286715676f, 0.0530110146f},
{-0.01853078078f, 0.01647272818f, -0.00207747588f},
{0.0023101082f, 0.01228800748f, 0.01303505006f},
{-0.01243671408f, -0.0258638728f, -0.0379116264f},
{0.00598934710f, 0.00642563550f, -0.01234514304f},
{-0.0296733996f, -0.0234698050f, 0.00060018212f},
{-0.0322019498f, -0.0529200462f, -0.00344987414f},
{-0.00205026458f, -0.00846599446f, 0.00455971038f},
{-0.01082227064f, 0.0315661948f, -0.0677753362f},
{0.0645553474f, 0.1109666998f, 0.0674744864f},
{0.01036801108f, -0.00484841210f, -0.001529168474f},
{0.01264353566f, 0.01548126338f, -0.00966374324f},
{-0.0223892408f, -0.00871751526f, -0.000306421670f},
{0.0271322742f, 0.03496524f, -0.0089692858f},
{0.0512178672f, 0.0173080034f, 0.00804227746f},
{0.01210987192f, 0.00758025926f, -0.00281712586f},
{0.1897278390f, 0.1210261828f, 0.062603892f},
{0.0208058822f, 0.00547548182f, 0.01262955638f},
{0.00813332858f, 0.01015930914f, 0.01301771290f},
{-0.000927236014f, -0.00152540594f, -0.00599213302f},
{0.01663314616f, -0.00582789626f, 0.0163958132f},
{-0.0252546342f, -0.0604193732f, -0.1606919922f},
{-0.091722686f, -0.0409201224f, -0.0959576198f},
{0.0282963112f, -0.01387223872f, -0.01648814464f},
{0.0552316818f, 0.0967547788f, 0.0413586632f},
{0.00922849292f, 0.00451467542f, -0.0529172378f},
{0.0558600768f, 0.0122988308f, -0.01445942422f},
{0.000210660902f, -0.01295958782f, -0.01804761764f},
{0.0358136250f, -0.0472505970f, -0.1156405142f},
{-0.0506390696f, -0.0471914842f, 0.0349791468f},
{-0.0480143168f, 0.00628389868f, -0.0545163826f},
{0.0315499582f, 0.0564846606f, -0.0430850488f},
{-0.0362330316f, -0.01267788554f, 0.0061024772f},
{0.0038627542f, 0.00911055916f, -0.00758526008f},
{-0.0447103298f, -0.00835411408f, 0.01545872328f},
{-0.015006738f, 0.00270612302f, -0.00784361356f},
{-0.0221755048f, -0.0513344748f, -0.0475317424f},
{-0.01036656294f, -0.00422146068f, -0.0213499052f},
{0.01788952706f, 0.01191944190f, 0.0397205238f},
};
float qwen21_latent_rgb_bias[3] = {-0.043293118f, -0.02695978f, -0.11986706f};
const float qwen21_latent_alpha_proj[64] = {
-0.0416241114f, -0.00678954612f, -0.0169095515f, -0.0230551401f, 0.0100882595f, 0.00655586802f, 0.0401166874f, -0.0055510216f,
0.0224234441f, -0.0389640963f, -0.0114492163f, -0.00721128977f, -0.0029064082f, 0.0150300547f, -0.00321615308f, -0.0498856338f,
-0.0215251401f, 0.0240220482f, 0.0117338008f, -0.0460420624f, 0.0387872889f, 0.0131517207f, 0.0147100836f, 0.0266985286f,
0.0153097324f, -0.0418119757f, 0.0421013917f, 0.0401724499f, 0.00972515915f, 0.011718495f, 0.0117622291f, 0.0136505134f,
-0.0350017363f, -0.0100692606f, -0.0131582529f, -0.00660639315f, 0.00253611396f, -0.0195736368f, -0.04240184f, 0.0321299262f,
0.0106089414f, -0.0179845306f, -0.00806212708f, 0.0135889057f, -0.0157393098f, -0.0267791344f, 0.0109068534f, 0.0283931966f,
-0.0435370078f, 0.00187883536f, -0.0108995378f, -0.0450757676f, -0.0699481501f, 0.0123562106f, -0.0222592249f, 0.0216155907f,
0.0563755424f, -0.0073379912f, 0.0160012921f, 0.0411637742f, 0.0189607258f, -0.024025029f, -0.0161487905f, -0.016913203f};
const float qwen21_latent_alpha_bias = 0.871268134f;
const float minimax_latent_rgb_proj[24][3] = {
{0.19819857f, 0.11584999f, 0.07929777f},
{-0.16047224f, -0.10601170f, -0.15996324f},
@ -324,7 +404,7 @@ const float sd_latent_rgb_proj[4][3] = {
{-0.178022f, -0.200862f, -0.678514f}};
float sd_latent_rgb_bias[3] = {-0.017478f, -0.055834f, -0.105825f};
void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size, const float* latent_alpha_proj = nullptr, float latent_alpha_bias = 1.f) {
size_t buffer_head = 0;
uint32_t latent_width = static_cast<uint32_t>(latents->ne[0]);
@ -338,7 +418,8 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
uint32_t rgb_width = latent_width * patch_size;
uint32_t rgb_height = latent_height * patch_size;
uint32_t unpatched_dim = dim / (patch_size * patch_size);
uint32_t unpatched_dim = dim / (patch_size * patch_size);
const uint32_t output_channels = latent_alpha_proj != nullptr ? 4 : 3;
for (uint32_t k = 0; k < frames; k++) {
for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
@ -356,13 +437,16 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
// should be incremented by 1 for each pixel
size_t pixel_id = k * rgb_width * rgb_height + rgb_y * rgb_width + rgb_x;
float r = 0, g = 0, b = 0;
float r = 0, g = 0, b = 0, a = 0;
if (latent_rgb_proj != nullptr) {
for (uint32_t d = 0; d < unpatched_dim; d++) {
float value = *(float*)((char*)latents->data + latent_id + (d * patch_size * patch_size + channel_offset) * latents->nb[ggml_n_dims(latents) - 1]);
r += value * latent_rgb_proj[d][0];
g += value * latent_rgb_proj[d][1];
b += value * latent_rgb_proj[d][2];
if (latent_alpha_proj != nullptr) {
a += value * latent_alpha_proj[d];
}
}
} else {
// interpret first 3 channels as RGB
@ -386,9 +470,13 @@ void preview_latent_video(uint8_t* buffer, ggml_tensor* latents, const float (*l
g = g >= 0 ? g <= 1 ? g : 1 : 0;
b = b >= 0 ? b <= 1 ? b : 1 : 0;
buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255);
buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255);
buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255);
buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255);
buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 255);
buffer[pixel_id * output_channels + 2] = (uint8_t)(b * 255);
if (latent_alpha_proj != nullptr) {
a = std::min(1.0f, std::max(0.0f, a + latent_alpha_bias));
buffer[pixel_id * output_channels + 3] = (uint8_t)(a * 255);
}
}
}
}
@ -398,16 +486,17 @@ static inline bool preview_latent_tensor_is_video(const sd::Tensor<float>& laten
return latents.dim() == 5;
}
void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size) {
void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, const float (*latent_rgb_proj)[3], const float latent_rgb_bias[3], int patch_size, const float* latent_alpha_proj = nullptr, float latent_alpha_bias = 1.f) {
uint32_t latent_width = static_cast<uint32_t>(latents.shape()[0]);
uint32_t latent_height = static_cast<uint32_t>(latents.shape()[1]);
bool is_video = preview_latent_tensor_is_video(latents);
uint32_t frames = is_video ? static_cast<uint32_t>(latents.shape()[2]) : 1;
uint32_t dim = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
uint32_t rgb_width = latent_width * patch_size;
uint32_t rgb_height = latent_height * patch_size;
uint32_t unpatched_dim = dim / (patch_size * patch_size);
uint32_t rgb_width = latent_width * patch_size;
uint32_t rgb_height = latent_height * patch_size;
uint32_t unpatched_dim = dim / (patch_size * patch_size);
const uint32_t output_channels = latent_alpha_proj != nullptr ? 4 : 3;
for (uint32_t k = 0; k < frames; k++) {
for (uint32_t rgb_x = 0; rgb_x < rgb_width; rgb_x++) {
@ -427,7 +516,7 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, con
: latents.values()[latent_x + latent_width * (latent_y + latent_height * latent_channel)];
};
float r = 0.f, g = 0.f, b = 0.f;
float r = 0.f, g = 0.f, b = 0.f, a = 0.f;
if (latent_rgb_proj != nullptr) {
for (uint32_t d = 0; d < unpatched_dim; d++) {
uint32_t latent_channel = d * patch_size * patch_size + channel_offset;
@ -435,6 +524,9 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, con
r += value * latent_rgb_proj[d][0];
g += value * latent_rgb_proj[d][1];
b += value * latent_rgb_proj[d][2];
if (latent_alpha_proj != nullptr) {
a += value * latent_alpha_proj[d];
}
}
} else {
r = latent_value(0);
@ -450,9 +542,13 @@ void preview_latent_video(uint8_t* buffer, const sd::Tensor<float>& latents, con
g = std::min(1.0f, std::max(0.0f, g * .5f + .5f));
b = std::min(1.0f, std::max(0.0f, b * .5f + .5f));
buffer[pixel_id * 3 + 0] = (uint8_t)(r * 255);
buffer[pixel_id * 3 + 1] = (uint8_t)(g * 255);
buffer[pixel_id * 3 + 2] = (uint8_t)(b * 255);
buffer[pixel_id * output_channels + 0] = (uint8_t)(r * 255);
buffer[pixel_id * output_channels + 1] = (uint8_t)(g * 255);
buffer[pixel_id * output_channels + 2] = (uint8_t)(b * 255);
if (latent_alpha_proj != nullptr) {
a = std::min(1.0f, std::max(0.0f, a + latent_alpha_bias));
buffer[pixel_id * output_channels + 3] = (uint8_t)(a * 255);
}
}
}
}

25
otherarch/sdcpp/src/runtime/preprocessing.hpp
View file

@ -165,16 +165,18 @@ static inline sd::Tensor<float> convolve_tensor(const sd::Tensor<float>& input,
return output;
}
static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& rgb_img) {
GGML_ASSERT(rgb_img.dim() == 4);
GGML_ASSERT(rgb_img.shape()[2] >= 3);
sd::Tensor<float> grayscale({rgb_img.shape()[0], rgb_img.shape()[1], 1, rgb_img.shape()[3]});
for (int64_t iy = 0; iy < rgb_img.shape()[1]; ++iy) {
for (int64_t ix = 0; ix < rgb_img.shape()[0]; ++ix) {
float r = preprocessing_get_4d(rgb_img, ix, iy, 0, 0);
float g = preprocessing_get_4d(rgb_img, ix, iy, 1, 0);
float b = preprocessing_get_4d(rgb_img, ix, iy, 2, 0);
float gray = 0.2989f * r + 0.5870f * g + 0.1140f * b;
static inline sd::Tensor<float> grayscale_tensor(const sd::Tensor<float>& image) {
GGML_ASSERT(image.dim() == 4);
GGML_ASSERT(image.shape()[2] >= 1);
sd::Tensor<float> grayscale({image.shape()[0], image.shape()[1], 1, image.shape()[3]});
for (int64_t iy = 0; iy < image.shape()[1]; ++iy) {
for (int64_t ix = 0; ix < image.shape()[0]; ++ix) {
float gray = preprocessing_get_4d(image, ix, iy, 0, 0);
if (image.shape()[2] >= 3) {
float g = preprocessing_get_4d(image, ix, iy, 1, 0);
float b = preprocessing_get_4d(image, ix, iy, 2, 0);
gray = 0.2989f * gray + 0.5870f * g + 0.1140f * b;
}
preprocessing_set_4d(grayscale, gray, ix, iy, 0, 0);
}
}
@ -317,11 +319,12 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
image_gray = non_max_supression(G, theta);
threshold_hystersis(&image_gray, high_threshold, low_threshold, weak, strong);
const uint32_t color_channels = img.channel == 2 || img.channel == 4 ? img.channel - 1 : img.channel;
for (uint32_t iy = 0; iy < img.height; ++iy) {
for (uint32_t ix = 0; ix < img.width; ++ix) {
float gray = preprocessing_get_4d(image_gray, ix, iy, 0, 0);
gray = inverse ? 1.0f - gray : gray;
for (uint32_t c = 0; c < img.channel; ++c) {
for (uint32_t c = 0; c < color_channels; ++c) {
preprocessing_set_4d(image, gray, ix, iy, c, 0);
}
}

20
otherarch/sdcpp/src/runtime/sample-cache.cpp
View file

@ -275,6 +275,26 @@ namespace sd_sample {
}
}
void SampleStepCacheDispatcher::invalidate_condition(const void* condition) {
if (condition == nullptr) {
return;
}
switch (runtime.mode) {
case SampleCacheMode::EASYCACHE:
runtime.easycache.cache_diffs.erase(condition);
break;
case SampleCacheMode::UCACHE:
runtime.ucache.cache_diffs.erase(condition);
break;
case SampleCacheMode::CACHEDIT:
runtime.cachedit.cache_diffs.erase(condition);
break;
case SampleCacheMode::NONE:
break;
}
}
bool SampleStepCacheDispatcher::is_step_skipped() const {
switch (runtime.mode) {
case SampleCacheMode::EASYCACHE:

1
otherarch/sdcpp/src/runtime/sample-cache.h
View file

@ -46,6 +46,7 @@ namespace sd_sample {
bool before_condition(const void* condition, const sd::Tensor<float>& input, sd::Tensor<float>* output);
void after_condition(const void* condition, const sd::Tensor<float>& input, const sd::Tensor<float>& output);
void invalidate_condition(const void* condition);
bool is_step_skipped() const;
};

41
otherarch/sdcpp/src/stable-diffusion.cpp
View file

@ -137,6 +137,7 @@ const char* scheduler_to_str[] = {
"flux2",
"flux",
"beta",
"llada_image",
};
static_assert(SCHEDULER_COUNT == sizeof(scheduler_to_str) / sizeof(scheduler_to_str[0]),
@ -325,6 +326,7 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) {
void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
*sd_ctx_params = {};
sd_ctx_params->n_threads = sd_get_num_physical_cores();
sd_ctx_params->conditioning_cache_size = 4;
sd_ctx_params->wtype = SD_TYPE_COUNT;
sd_ctx_params->rng_type = CUDA_RNG;
sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT;
@ -336,6 +338,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
sd_ctx_params->eager_load = false;
sd_ctx_params->enable_mmap = false;
sd_ctx_params->diffusion_flash_attn = false;
sd_ctx_params->sage_attn = false;
sd_ctx_params->linear_scale = 0.f;
sd_ctx_params->attn_scale = 0.f;
sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO;
@ -376,6 +379,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"pulid_weights_path: %s\n"
"tensor_type_rules: %s\n"
"n_threads: %d\n"
"conditioning_cache_size: %d\n"
"wtype: %s\n"
"rng_type: %s\n"
"sampler_rng_type: %s\n"
@ -391,6 +395,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
"auto_fit: %s\n"
"flash_attn: %s\n"
"diffusion_flash_attn: %s\n"
"sage_attn: %s\n"
"linear_scale: %g\n"
"attn_scale: %g\n"
"vae_format: %s\n",
@ -415,6 +420,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
SAFE_STR(sd_ctx_params->pulid_weights_path),
SAFE_STR(sd_ctx_params->tensor_type_rules),
sd_ctx_params->n_threads,
sd_ctx_params->conditioning_cache_size,
sd_type_name(sd_ctx_params->wtype),
sd_rng_type_name(sd_ctx_params->rng_type),
sd_rng_type_name(sd_ctx_params->sampler_rng_type),
@ -430,6 +436,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
BOOL_STR(sd_ctx_params->auto_fit),
BOOL_STR(sd_ctx_params->flash_attn),
BOOL_STR(sd_ctx_params->diffusion_flash_attn),
BOOL_STR(sd_ctx_params->sage_attn),
sd_ctx_params->linear_scale,
sd_ctx_params->attn_scale,
sd_vae_format_name(sd_ctx_params->vae_format));
@ -627,14 +634,6 @@ struct sd_ctx_t {
StableDiffusionGGML* sd = nullptr;
};
static bool sd_version_supports_video_generation(SDVersion version) {
return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version);
}
static bool sd_version_supports_image_generation(SDVersion version) {
return !sd_version_supports_video_generation(version);
}
sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) {
sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
if (sd_ctx == nullptr) {
@ -756,28 +755,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
int* num_frames_out,
sd_audio_t** audio_out,
int* fps_out) {
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
if (fps_out != nullptr) {
*fps_out = 0;
}
return false;
}
if (frames_out != nullptr) {
if (frames_out != nullptr)
*frames_out = nullptr;
}
if (audio_out != nullptr) {
if (audio_out != nullptr)
*audio_out = nullptr;
}
if (num_frames_out != nullptr) {
if (num_frames_out != nullptr)
*num_frames_out = 0;
if (fps_out != nullptr)
*fps_out = 0;
if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) {
return false;
}
StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd);
if (!execution.ready) {
if (fps_out != nullptr) {
*fps_out = 0;
}
return false;
}
@ -827,7 +818,7 @@ namespace kcpp_sd {
model_info res = {};
SDVersion loadedsdver = (SDVersion)get_loaded_sd_version(ctx);
res.is_wan = (loadedsdver == SDVersion::VERSION_WAN2 || loadedsdver == SDVersion::VERSION_WAN2_2_I2V || loadedsdver == SDVersion::VERSION_WAN2_2_TI2V);
res.is_qwenimg = (loadedsdver == SDVersion::VERSION_QWEN_IMAGE);
res.is_qwenimg = sd_version_is_qwen_image(loadedsdver);
res.is_chroma = loaded_model_is_chroma(ctx);
res.is_kontext = (loadedsdver==SDVersion::VERSION_FLUX && !res.is_chroma);
res.is_flux2 = (loadedsdver == SDVersion::VERSION_FLUX2 || loadedsdver == SDVersion::VERSION_FLUX2_KLEIN);
@ -840,7 +831,7 @@ namespace kcpp_sd {
res.is_ltx = sd_version_is_ltxav(loadedsdver);
res.is_minimaxh3 = sd_version_is_minimax_h3(loadedsdver);
res.is_boogu = sd_version_is_boogu_image(loadedsdver);
res.supports_ref_image = sd_version_supports_ref_latent_img_cfg(loadedsdver) || sd_version_is_pid(loadedsdver);
res.supports_ref_image = sd_version_supports_ref_latent_img_cfg(loadedsdver) || sd_version_is_pid(loadedsdver) || res.is_qwenimg;
res.vae_scale_factor = ctx->sd->get_vae_scale_factor();
res.spatial_multiple = get_spatial_multiple(ctx);
return res;

48
otherarch/sdcpp/src/upscaler.cpp
View file

@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
sd::ParallelScope tensor_scope(&tensor_executor);
if (input_tensor.empty() || input_tensor.dim() != 4 ||
(input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) {
LOG_ERROR("esrgan expects a 4D RGB or RGBA image tensor");
return {};
}
const bool has_alpha = input_tensor.shape()[2] == 4;
sd::Tensor<float> rgb;
if (has_alpha) {
rgb = sd::ops::slice(input_tensor, 2, 0, 3);
}
const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
sd::Tensor<float> upscaled;
const int scale = esrgan_upscaler->config.scale;
if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
upscaled = esrgan_upscaler->compute(n_threads, input_tensor);
upscaled = esrgan_upscaler->compute(n_threads, model_input);
} else {
auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> {
auto output_tile = esrgan_upscaler->compute(n_threads, input_tile);
@ -125,7 +137,7 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
return output_tile;
};
upscaled = process_tiles_2d(input_tensor,
upscaled = process_tiles_2d(model_input,
static_cast<int>(input_tensor.shape()[0] * scale),
static_cast<int>(input_tensor.shape()[1] * scale),
scale,
@ -141,6 +153,14 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
LOG_ERROR("esrgan compute failed");
return {};
}
if (has_alpha) {
auto alpha = sd::ops::slice(input_tensor, 2, 3, 4);
auto alpha_shape = alpha.shape();
alpha_shape[0] = upscaled.shape()[0];
alpha_shape[1] = upscaled.shape()[1];
alpha = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear);
upscaled = sd::ops::concat(upscaled, alpha, 2);
}
return upscaled;
}
@ -242,6 +262,30 @@ int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) {
return upscaler_ctx->upscaler->esrgan_upscaler->config.scale;
}
int get_upscaler_model_scale(const char* model_path) {
if (model_path == nullptr || model_path[0] == '\0') {
return 0;
}
try {
ModelLoader loader;
if (!loader.init_from_file_and_convert_name(model_path, "", VERSION_ESRGAN)) {
return 0;
}
const auto& tensors = loader.get_tensor_storage_map();
auto first = tensors.find("conv_first.weight");
auto last = tensors.find("conv_last.weight");
if (first == tensors.end() || last == tensors.end() ||
tensors.count("body.0.rdb1.conv1.weight") == 0 ||
first->second.n_dims != 4 || last->second.n_dims != 4 ||
first->second.ne[2] != 3 || last->second.ne[3] != 3) {
return 0;
}
return ESRGANConfig::detect_from_weights(tensors).scale;
} catch (const std::exception&) {
return 0;
}
}
void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) {
if (upscaler_ctx->upscaler != nullptr) {
delete upscaler_ctx->upscaler;