From 84765f59674fce3616854595e4529f366968726e Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Sat, 27 Dec 2025 05:30:52 -0300 Subject: [PATCH] sd: sync to master-447-ccb6b0a (#1898) * sd: sync to master-438-298b110 * sd: sync to master-440-3e81246 * sd: sync to master-444-a0adcfb * sd: sync to master-447-ccb6b0a --- otherarch/sdcpp/cache_dit.hpp | 975 +++++++++++++++++++++++++++ otherarch/sdcpp/common.hpp | 2 +- otherarch/sdcpp/common/common.hpp | 298 +++++--- otherarch/sdcpp/denoiser.hpp | 6 +- otherarch/sdcpp/diffusion_model.hpp | 34 +- otherarch/sdcpp/flux.hpp | 41 +- otherarch/sdcpp/ggml_extend.hpp | 119 +++- otherarch/sdcpp/lora.hpp | 14 +- otherarch/sdcpp/main.cpp | 89 ++- otherarch/sdcpp/mmdit.hpp | 2 +- otherarch/sdcpp/model.cpp | 11 +- otherarch/sdcpp/name_conversion.cpp | 18 +- otherarch/sdcpp/qwen_image.hpp | 132 +++- otherarch/sdcpp/rope.hpp | 190 +++++- otherarch/sdcpp/stable-diffusion.cpp | 473 ++++++++++--- otherarch/sdcpp/stable-diffusion.h | 38 +- otherarch/sdcpp/ucache.hpp | 404 +++++++++++ otherarch/sdcpp/wan.hpp | 15 +- otherarch/sdcpp/z_image.hpp | 14 +- 19 files changed, 2568 insertions(+), 307 deletions(-) create mode 100644 otherarch/sdcpp/cache_dit.hpp create mode 100644 otherarch/sdcpp/ucache.hpp diff --git a/otherarch/sdcpp/cache_dit.hpp b/otherarch/sdcpp/cache_dit.hpp new file mode 100644 index 000000000..f5a1f180b --- /dev/null +++ b/otherarch/sdcpp/cache_dit.hpp @@ -0,0 +1,975 @@ +#ifndef __CACHE_DIT_HPP__ +#define __CACHE_DIT_HPP__ + +#include +#include +#include +#include +#include +#include + +#include "ggml_extend.hpp" + +struct DBCacheConfig { + bool enabled = false; + int Fn_compute_blocks = 8; + int Bn_compute_blocks = 0; + float residual_diff_threshold = 0.08f; + int max_warmup_steps = 8; + int max_cached_steps = -1; + int max_continuous_cached_steps = -1; + float max_accumulated_residual_diff = -1.0f; + std::vector steps_computation_mask; + bool scm_policy_dynamic = true; +}; + +struct TaylorSeerConfig { + bool enabled = false; + int n_derivatives = 1; + int max_warmup_steps = 2; + int skip_interval_steps = 1; +}; + +struct CacheDitConfig { + DBCacheConfig dbcache; + TaylorSeerConfig taylorseer; + int double_Fn_blocks = -1; + int double_Bn_blocks = -1; + int single_Fn_blocks = -1; + int single_Bn_blocks = -1; +}; + +struct TaylorSeerState { + int n_derivatives = 1; + int current_step = -1; + int last_computed_step = -1; + std::vector> dY_prev; + std::vector> dY_current; + + void init(int n_deriv, size_t hidden_size) { + n_derivatives = n_deriv; + int order = n_derivatives + 1; + dY_prev.resize(order); + dY_current.resize(order); + for (int i = 0; i < order; i++) { + dY_prev[i].clear(); + dY_current[i].clear(); + } + current_step = -1; + last_computed_step = -1; + } + + void reset() { + for (auto& v : dY_prev) + v.clear(); + for (auto& v : dY_current) + v.clear(); + current_step = -1; + last_computed_step = -1; + } + + bool can_approximate() const { + return last_computed_step >= n_derivatives && !dY_prev.empty() && !dY_prev[0].empty(); + } + + void update_derivatives(const float* Y, size_t size, int step) { + int order = n_derivatives + 1; + dY_prev = dY_current; + dY_current[0].resize(size); + for (size_t i = 0; i < size; i++) { + dY_current[0][i] = Y[i]; + } + + int window = step - last_computed_step; + if (window <= 0) + window = 1; + + for (int d = 0; d < n_derivatives; d++) { + if (!dY_prev[d].empty() && dY_prev[d].size() == size) { + dY_current[d + 1].resize(size); + for (size_t i = 0; i < size; i++) { + dY_current[d + 1][i] = (dY_current[d][i] - dY_prev[d][i]) / static_cast(window); + } + } else { + dY_current[d + 1].clear(); + } + } + + current_step = step; + last_computed_step = step; + } + + void approximate(float* output, size_t size, int target_step) const { + if (!can_approximate() || dY_prev[0].size() != size) { + return; + } + + int elapsed = target_step - last_computed_step; + if (elapsed <= 0) + elapsed = 1; + + std::fill(output, output + size, 0.0f); + float factorial = 1.0f; + int order = static_cast(dY_prev.size()); + + for (int o = 0; o < order; o++) { + if (dY_prev[o].empty() || dY_prev[o].size() != size) + continue; + if (o > 0) + factorial *= static_cast(o); + float coeff = std::pow(static_cast(elapsed), o) / factorial; + for (size_t i = 0; i < size; i++) { + output[i] += coeff * dY_prev[o][i]; + } + } + } +}; + +struct BlockCacheEntry { + std::vector residual_img; + std::vector residual_txt; + std::vector residual; + std::vector prev_img; + std::vector prev_txt; + std::vector prev_output; + bool has_prev = false; +}; + +struct CacheDitState { + CacheDitConfig config; + bool initialized = false; + + int total_double_blocks = 0; + int total_single_blocks = 0; + size_t hidden_size = 0; + + int current_step = -1; + int total_steps = 0; + int warmup_remaining = 0; + std::vector cached_steps; + int continuous_cached_steps = 0; + float accumulated_residual_diff = 0.0f; + + std::vector double_block_cache; + std::vector single_block_cache; + + std::vector Fn_residual_img; + std::vector Fn_residual_txt; + std::vector prev_Fn_residual_img; + std::vector prev_Fn_residual_txt; + bool has_prev_Fn_residual = false; + + std::vector Bn_buffer_img; + std::vector Bn_buffer_txt; + std::vector Bn_buffer; + bool has_Bn_buffer = false; + + TaylorSeerState taylor_state; + + bool can_cache_this_step = false; + bool is_caching_this_step = false; + + int total_blocks_computed = 0; + int total_blocks_cached = 0; + + void init(const CacheDitConfig& cfg, int num_double_blocks, int num_single_blocks, size_t h_size) { + config = cfg; + total_double_blocks = num_double_blocks; + total_single_blocks = num_single_blocks; + hidden_size = h_size; + + initialized = cfg.dbcache.enabled || cfg.taylorseer.enabled; + + if (!initialized) + return; + + warmup_remaining = cfg.dbcache.max_warmup_steps; + double_block_cache.resize(total_double_blocks); + single_block_cache.resize(total_single_blocks); + + if (cfg.taylorseer.enabled) { + taylor_state.init(cfg.taylorseer.n_derivatives, h_size); + } + + reset_runtime(); + } + + void reset_runtime() { + current_step = -1; + total_steps = 0; + warmup_remaining = config.dbcache.max_warmup_steps; + cached_steps.clear(); + continuous_cached_steps = 0; + accumulated_residual_diff = 0.0f; + + for (auto& entry : double_block_cache) { + entry.residual_img.clear(); + entry.residual_txt.clear(); + entry.prev_img.clear(); + entry.prev_txt.clear(); + entry.has_prev = false; + } + + for (auto& entry : single_block_cache) { + entry.residual.clear(); + entry.prev_output.clear(); + entry.has_prev = false; + } + + Fn_residual_img.clear(); + Fn_residual_txt.clear(); + prev_Fn_residual_img.clear(); + prev_Fn_residual_txt.clear(); + has_prev_Fn_residual = false; + + Bn_buffer_img.clear(); + Bn_buffer_txt.clear(); + Bn_buffer.clear(); + has_Bn_buffer = false; + + taylor_state.reset(); + + can_cache_this_step = false; + is_caching_this_step = false; + + total_blocks_computed = 0; + total_blocks_cached = 0; + } + + bool enabled() const { + return initialized && (config.dbcache.enabled || config.taylorseer.enabled); + } + + void begin_step(int step_index, float sigma = 0.0f) { + if (!enabled()) + return; + if (step_index == current_step) + return; + + current_step = step_index; + total_steps++; + + bool in_warmup = warmup_remaining > 0; + if (in_warmup) { + warmup_remaining--; + } + + bool scm_allows_cache = true; + if (!config.dbcache.steps_computation_mask.empty()) { + if (step_index < static_cast(config.dbcache.steps_computation_mask.size())) { + scm_allows_cache = (config.dbcache.steps_computation_mask[step_index] == 0); + if (!config.dbcache.scm_policy_dynamic && scm_allows_cache) { + can_cache_this_step = true; + is_caching_this_step = false; + return; + } + } + } + + bool max_cached_ok = (config.dbcache.max_cached_steps < 0) || + (static_cast(cached_steps.size()) < config.dbcache.max_cached_steps); + + bool max_cont_ok = (config.dbcache.max_continuous_cached_steps < 0) || + (continuous_cached_steps < config.dbcache.max_continuous_cached_steps); + + bool accum_ok = (config.dbcache.max_accumulated_residual_diff < 0.0f) || + (accumulated_residual_diff < config.dbcache.max_accumulated_residual_diff); + + can_cache_this_step = !in_warmup && scm_allows_cache && max_cached_ok && max_cont_ok && accum_ok && has_prev_Fn_residual; + is_caching_this_step = false; + } + + void end_step(bool was_cached) { + if (was_cached) { + cached_steps.push_back(current_step); + continuous_cached_steps++; + } else { + continuous_cached_steps = 0; + } + } + + static float calculate_residual_diff(const float* prev, const float* curr, size_t size) { + if (size == 0) + return 0.0f; + + float sum_diff = 0.0f; + float sum_abs = 0.0f; + + for (size_t i = 0; i < size; i++) { + sum_diff += std::fabs(prev[i] - curr[i]); + sum_abs += std::fabs(prev[i]); + } + + return sum_diff / (sum_abs + 1e-6f); + } + + static float calculate_residual_diff(const std::vector& prev, const std::vector& curr) { + if (prev.size() != curr.size() || prev.empty()) + return 1.0f; + return calculate_residual_diff(prev.data(), curr.data(), prev.size()); + } + + int get_double_Fn_blocks() const { + return (config.double_Fn_blocks >= 0) ? config.double_Fn_blocks : config.dbcache.Fn_compute_blocks; + } + + int get_double_Bn_blocks() const { + return (config.double_Bn_blocks >= 0) ? config.double_Bn_blocks : config.dbcache.Bn_compute_blocks; + } + + int get_single_Fn_blocks() const { + return (config.single_Fn_blocks >= 0) ? config.single_Fn_blocks : config.dbcache.Fn_compute_blocks; + } + + int get_single_Bn_blocks() const { + return (config.single_Bn_blocks >= 0) ? config.single_Bn_blocks : config.dbcache.Bn_compute_blocks; + } + + bool is_Fn_double_block(int block_idx) const { + return block_idx < get_double_Fn_blocks(); + } + + bool is_Bn_double_block(int block_idx) const { + int Bn = get_double_Bn_blocks(); + return Bn > 0 && block_idx >= (total_double_blocks - Bn); + } + + bool is_Mn_double_block(int block_idx) const { + return !is_Fn_double_block(block_idx) && !is_Bn_double_block(block_idx); + } + + bool is_Fn_single_block(int block_idx) const { + return block_idx < get_single_Fn_blocks(); + } + + bool is_Bn_single_block(int block_idx) const { + int Bn = get_single_Bn_blocks(); + return Bn > 0 && block_idx >= (total_single_blocks - Bn); + } + + bool is_Mn_single_block(int block_idx) const { + return !is_Fn_single_block(block_idx) && !is_Bn_single_block(block_idx); + } + + void store_Fn_residual(const float* img, const float* txt, size_t img_size, size_t txt_size, const float* input_img, const float* input_txt) { + Fn_residual_img.resize(img_size); + Fn_residual_txt.resize(txt_size); + + for (size_t i = 0; i < img_size; i++) { + Fn_residual_img[i] = img[i] - input_img[i]; + } + for (size_t i = 0; i < txt_size; i++) { + Fn_residual_txt[i] = txt[i] - input_txt[i]; + } + } + + bool check_cache_decision() { + if (!can_cache_this_step) { + is_caching_this_step = false; + return false; + } + + if (!has_prev_Fn_residual || prev_Fn_residual_img.empty()) { + is_caching_this_step = false; + return false; + } + + float diff_img = calculate_residual_diff(prev_Fn_residual_img, Fn_residual_img); + float diff_txt = calculate_residual_diff(prev_Fn_residual_txt, Fn_residual_txt); + float diff = (diff_img + diff_txt) / 2.0f; + + if (diff < config.dbcache.residual_diff_threshold) { + is_caching_this_step = true; + accumulated_residual_diff += diff; + return true; + } + + is_caching_this_step = false; + return false; + } + + void update_prev_Fn_residual() { + prev_Fn_residual_img = Fn_residual_img; + prev_Fn_residual_txt = Fn_residual_txt; + has_prev_Fn_residual = !prev_Fn_residual_img.empty(); + } + + void store_double_block_residual(int block_idx, const float* img, const float* txt, size_t img_size, size_t txt_size, const float* prev_img, const float* prev_txt) { + if (block_idx < 0 || block_idx >= static_cast(double_block_cache.size())) + return; + + BlockCacheEntry& entry = double_block_cache[block_idx]; + + entry.residual_img.resize(img_size); + entry.residual_txt.resize(txt_size); + for (size_t i = 0; i < img_size; i++) { + entry.residual_img[i] = img[i] - prev_img[i]; + } + for (size_t i = 0; i < txt_size; i++) { + entry.residual_txt[i] = txt[i] - prev_txt[i]; + } + + entry.prev_img.resize(img_size); + entry.prev_txt.resize(txt_size); + for (size_t i = 0; i < img_size; i++) { + entry.prev_img[i] = img[i]; + } + for (size_t i = 0; i < txt_size; i++) { + entry.prev_txt[i] = txt[i]; + } + entry.has_prev = true; + } + + void apply_double_block_cache(int block_idx, float* img, float* txt, size_t img_size, size_t txt_size) { + if (block_idx < 0 || block_idx >= static_cast(double_block_cache.size())) + return; + + const BlockCacheEntry& entry = double_block_cache[block_idx]; + if (entry.residual_img.size() != img_size || entry.residual_txt.size() != txt_size) + return; + + for (size_t i = 0; i < img_size; i++) { + img[i] += entry.residual_img[i]; + } + for (size_t i = 0; i < txt_size; i++) { + txt[i] += entry.residual_txt[i]; + } + + total_blocks_cached++; + } + + void store_single_block_residual(int block_idx, const float* output, size_t size, const float* input) { + if (block_idx < 0 || block_idx >= static_cast(single_block_cache.size())) + return; + + BlockCacheEntry& entry = single_block_cache[block_idx]; + + entry.residual.resize(size); + for (size_t i = 0; i < size; i++) { + entry.residual[i] = output[i] - input[i]; + } + + entry.prev_output.resize(size); + for (size_t i = 0; i < size; i++) { + entry.prev_output[i] = output[i]; + } + entry.has_prev = true; + } + + void apply_single_block_cache(int block_idx, float* output, size_t size) { + if (block_idx < 0 || block_idx >= static_cast(single_block_cache.size())) + return; + + const BlockCacheEntry& entry = single_block_cache[block_idx]; + if (entry.residual.size() != size) + return; + + for (size_t i = 0; i < size; i++) { + output[i] += entry.residual[i]; + } + + total_blocks_cached++; + } + + void store_Bn_buffer(const float* img, const float* txt, size_t img_size, size_t txt_size, const float* Bn_start_img, const float* Bn_start_txt) { + Bn_buffer_img.resize(img_size); + Bn_buffer_txt.resize(txt_size); + + for (size_t i = 0; i < img_size; i++) { + Bn_buffer_img[i] = img[i] - Bn_start_img[i]; + } + for (size_t i = 0; i < txt_size; i++) { + Bn_buffer_txt[i] = txt[i] - Bn_start_txt[i]; + } + has_Bn_buffer = true; + } + + void apply_Bn_buffer(float* img, float* txt, size_t img_size, size_t txt_size) { + if (!has_Bn_buffer) + return; + if (Bn_buffer_img.size() != img_size || Bn_buffer_txt.size() != txt_size) + return; + + for (size_t i = 0; i < img_size; i++) { + img[i] += Bn_buffer_img[i]; + } + for (size_t i = 0; i < txt_size; i++) { + txt[i] += Bn_buffer_txt[i]; + } + } + + void taylor_update(const float* hidden_state, size_t size) { + if (!config.taylorseer.enabled) + return; + taylor_state.update_derivatives(hidden_state, size, current_step); + } + + bool taylor_can_approximate() const { + return config.taylorseer.enabled && taylor_state.can_approximate(); + } + + void taylor_approximate(float* output, size_t size) { + if (!config.taylorseer.enabled) + return; + taylor_state.approximate(output, size, current_step); + } + + bool should_use_taylor_this_step() const { + if (!config.taylorseer.enabled) + return false; + if (current_step < config.taylorseer.max_warmup_steps) + return false; + + int interval = config.taylorseer.skip_interval_steps; + if (interval <= 0) + interval = 1; + + return (current_step % (interval + 1)) != 0; + } + + void log_metrics() const { + if (!enabled()) + return; + + int total_blocks = total_blocks_computed + total_blocks_cached; + float cache_ratio = (total_blocks > 0) ? (static_cast(total_blocks_cached) / total_blocks * 100.0f) : 0.0f; + + float step_cache_ratio = (total_steps > 0) ? (static_cast(cached_steps.size()) / total_steps * 100.0f) : 0.0f; + + LOG_INFO("CacheDIT: steps_cached=%zu/%d (%.1f%%), blocks_cached=%d/%d (%.1f%%), accum_diff=%.4f", + cached_steps.size(), total_steps, step_cache_ratio, + total_blocks_cached, total_blocks, cache_ratio, + accumulated_residual_diff); + } + + std::string get_summary() const { + char buf[256]; + snprintf(buf, sizeof(buf), + "CacheDIT[thresh=%.2f]: cached %zu/%d steps, %d/%d blocks", + config.dbcache.residual_diff_threshold, + cached_steps.size(), total_steps, + total_blocks_cached, total_blocks_computed + total_blocks_cached); + return std::string(buf); + } +}; + +inline std::vector parse_scm_mask(const std::string& mask_str) { + std::vector mask; + if (mask_str.empty()) + return mask; + + size_t pos = 0; + size_t start = 0; + while ((pos = mask_str.find(',', start)) != std::string::npos) { + std::string token = mask_str.substr(start, pos - start); + mask.push_back(std::stoi(token)); + start = pos + 1; + } + if (start < mask_str.length()) { + mask.push_back(std::stoi(mask_str.substr(start))); + } + + return mask; +} + +inline std::vector generate_scm_mask( + const std::vector& compute_bins, + const std::vector& cache_bins, + int total_steps) { + std::vector mask; + size_t c_idx = 0, cache_idx = 0; + + while (static_cast(mask.size()) < total_steps) { + if (c_idx < compute_bins.size()) { + for (int i = 0; i < compute_bins[c_idx] && static_cast(mask.size()) < total_steps; i++) { + mask.push_back(1); + } + c_idx++; + } + if (cache_idx < cache_bins.size()) { + for (int i = 0; i < cache_bins[cache_idx] && static_cast(mask.size()) < total_steps; i++) { + mask.push_back(0); + } + cache_idx++; + } + if (c_idx >= compute_bins.size() && cache_idx >= cache_bins.size()) + break; + } + + if (!mask.empty()) { + mask.back() = 1; + } + + return mask; +} + +inline std::vector get_scm_preset(const std::string& preset, int total_steps) { + struct Preset { + std::vector compute_bins; + std::vector cache_bins; + }; + + Preset slow = {{8, 3, 3, 2, 1, 1}, {1, 2, 2, 2, 3}}; + Preset medium = {{6, 2, 2, 2, 2, 1}, {1, 3, 3, 3, 3}}; + Preset fast = {{6, 1, 1, 1, 1, 1}, {1, 3, 4, 5, 4}}; + Preset ultra = {{4, 1, 1, 1, 1}, {2, 5, 6, 7}}; + + Preset* p = nullptr; + if (preset == "slow" || preset == "s" || preset == "S") + p = &slow; + else if (preset == "medium" || preset == "m" || preset == "M") + p = &medium; + else if (preset == "fast" || preset == "f" || preset == "F") + p = &fast; + else if (preset == "ultra" || preset == "u" || preset == "U") + p = &ultra; + else + return {}; + + if (total_steps != 28 && total_steps > 0) { + float scale = static_cast(total_steps) / 28.0f; + std::vector scaled_compute, scaled_cache; + + for (int v : p->compute_bins) { + scaled_compute.push_back(std::max(1, static_cast(v * scale + 0.5f))); + } + for (int v : p->cache_bins) { + scaled_cache.push_back(std::max(1, static_cast(v * scale + 0.5f))); + } + + return generate_scm_mask(scaled_compute, scaled_cache, total_steps); + } + + return generate_scm_mask(p->compute_bins, p->cache_bins, total_steps); +} + +inline float get_preset_threshold(const std::string& preset) { + if (preset == "slow" || preset == "s" || preset == "S") + return 0.20f; + if (preset == "medium" || preset == "m" || preset == "M") + return 0.25f; + if (preset == "fast" || preset == "f" || preset == "F") + return 0.30f; + if (preset == "ultra" || preset == "u" || preset == "U") + return 0.34f; + return 0.08f; +} + +inline int get_preset_warmup(const std::string& preset) { + if (preset == "slow" || preset == "s" || preset == "S") + return 8; + if (preset == "medium" || preset == "m" || preset == "M") + return 6; + if (preset == "fast" || preset == "f" || preset == "F") + return 6; + if (preset == "ultra" || preset == "u" || preset == "U") + return 4; + return 8; +} + +inline int get_preset_Fn(const std::string& preset) { + if (preset == "slow" || preset == "s" || preset == "S") + return 8; + if (preset == "medium" || preset == "m" || preset == "M") + return 8; + if (preset == "fast" || preset == "f" || preset == "F") + return 6; + if (preset == "ultra" || preset == "u" || preset == "U") + return 4; + return 8; +} + +inline int get_preset_Bn(const std::string& preset) { + (void)preset; + return 0; +} + +inline void parse_dbcache_options(const std::string& opts, DBCacheConfig& cfg) { + if (opts.empty()) + return; + + int Fn = 8, Bn = 0, warmup = 8, max_cached = -1, max_cont = -1; + float thresh = 0.08f; + + sscanf(opts.c_str(), "%d,%d,%f,%d,%d,%d", + &Fn, &Bn, &thresh, &warmup, &max_cached, &max_cont); + + cfg.Fn_compute_blocks = Fn; + cfg.Bn_compute_blocks = Bn; + cfg.residual_diff_threshold = thresh; + cfg.max_warmup_steps = warmup; + cfg.max_cached_steps = max_cached; + cfg.max_continuous_cached_steps = max_cont; +} + +inline void parse_taylorseer_options(const std::string& opts, TaylorSeerConfig& cfg) { + if (opts.empty()) + return; + + int n_deriv = 1, warmup = 2, interval = 1; + sscanf(opts.c_str(), "%d,%d,%d", &n_deriv, &warmup, &interval); + + cfg.n_derivatives = n_deriv; + cfg.max_warmup_steps = warmup; + cfg.skip_interval_steps = interval; +} + +struct CacheDitConditionState { + DBCacheConfig config; + TaylorSeerConfig taylor_config; + bool initialized = false; + + int current_step_index = -1; + bool step_active = false; + bool skip_current_step = false; + bool initial_step = true; + int warmup_remaining = 0; + std::vector cached_steps; + int continuous_cached_steps = 0; + float accumulated_residual_diff = 0.0f; + int total_steps_skipped = 0; + + const void* anchor_condition = nullptr; + + struct CacheEntry { + std::vector diff; + std::vector prev_input; + std::vector prev_output; + bool has_prev = false; + }; + std::unordered_map cache_diffs; + + TaylorSeerState taylor_state; + + float start_sigma = std::numeric_limits::max(); + float end_sigma = 0.0f; + + void reset_runtime() { + current_step_index = -1; + step_active = false; + skip_current_step = false; + initial_step = true; + warmup_remaining = config.max_warmup_steps; + cached_steps.clear(); + continuous_cached_steps = 0; + accumulated_residual_diff = 0.0f; + total_steps_skipped = 0; + anchor_condition = nullptr; + cache_diffs.clear(); + taylor_state.reset(); + } + + void init(const DBCacheConfig& dbcfg, const TaylorSeerConfig& tcfg) { + config = dbcfg; + taylor_config = tcfg; + initialized = dbcfg.enabled || tcfg.enabled; + reset_runtime(); + + if (taylor_config.enabled) { + taylor_state.init(taylor_config.n_derivatives, 0); + } + } + + void set_sigmas(const std::vector& sigmas) { + if (!initialized || sigmas.size() < 2) + return; + + float start_percent = 0.15f; + float end_percent = 0.95f; + + size_t n_steps = sigmas.size() - 1; + size_t start_step = static_cast(start_percent * n_steps); + size_t end_step = static_cast(end_percent * n_steps); + + if (start_step >= n_steps) + start_step = n_steps - 1; + if (end_step >= n_steps) + end_step = n_steps - 1; + + start_sigma = sigmas[start_step]; + end_sigma = sigmas[end_step]; + + if (start_sigma < end_sigma) { + std::swap(start_sigma, end_sigma); + } + } + + bool enabled() const { + return initialized && (config.enabled || taylor_config.enabled); + } + + void begin_step(int step_index, float sigma) { + if (!enabled()) + return; + if (step_index == current_step_index) + return; + + current_step_index = step_index; + skip_current_step = false; + step_active = false; + + if (sigma > start_sigma) + return; + if (!(sigma > end_sigma)) + return; + + step_active = true; + + if (warmup_remaining > 0) { + warmup_remaining--; + return; + } + + if (!config.steps_computation_mask.empty()) { + if (step_index < static_cast(config.steps_computation_mask.size())) { + if (config.steps_computation_mask[step_index] == 1) { + return; + } + } + } + + if (config.max_cached_steps >= 0 && + static_cast(cached_steps.size()) >= config.max_cached_steps) { + return; + } + + if (config.max_continuous_cached_steps >= 0 && + continuous_cached_steps >= config.max_continuous_cached_steps) { + return; + } + } + + bool step_is_active() const { + return enabled() && step_active; + } + + bool is_step_skipped() const { + return enabled() && step_active && skip_current_step; + } + + bool has_cache(const void* cond) const { + auto it = cache_diffs.find(cond); + return it != cache_diffs.end() && !it->second.diff.empty(); + } + + void update_cache(const void* cond, const float* input, const float* output, size_t size) { + CacheEntry& entry = cache_diffs[cond]; + entry.diff.resize(size); + for (size_t i = 0; i < size; i++) { + entry.diff[i] = output[i] - input[i]; + } + + entry.prev_input.resize(size); + entry.prev_output.resize(size); + for (size_t i = 0; i < size; i++) { + entry.prev_input[i] = input[i]; + entry.prev_output[i] = output[i]; + } + entry.has_prev = true; + } + + void apply_cache(const void* cond, const float* input, float* output, size_t size) { + auto it = cache_diffs.find(cond); + if (it == cache_diffs.end() || it->second.diff.empty()) + return; + if (it->second.diff.size() != size) + return; + + for (size_t i = 0; i < size; i++) { + output[i] = input[i] + it->second.diff[i]; + } + } + + bool before_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output, float sigma, int step_index) { + if (!enabled() || step_index < 0) + return false; + + if (step_index != current_step_index) { + begin_step(step_index, sigma); + } + + if (!step_active) + return false; + + if (initial_step) { + anchor_condition = cond; + initial_step = false; + } + + bool is_anchor = (cond == anchor_condition); + + if (skip_current_step) { + if (has_cache(cond)) { + apply_cache(cond, (float*)input->data, (float*)output->data, + static_cast(ggml_nelements(output))); + return true; + } + return false; + } + + if (!is_anchor) + return false; + + auto it = cache_diffs.find(cond); + if (it == cache_diffs.end() || !it->second.has_prev) + return false; + + size_t ne = static_cast(ggml_nelements(input)); + if (it->second.prev_input.size() != ne) + return false; + + float* input_data = (float*)input->data; + float diff = CacheDitState::calculate_residual_diff( + it->second.prev_input.data(), input_data, ne); + + float effective_threshold = config.residual_diff_threshold; + if (config.Fn_compute_blocks > 0) { + float fn_confidence = 1.0f + 0.02f * (config.Fn_compute_blocks - 8); + fn_confidence = std::max(0.5f, std::min(2.0f, fn_confidence)); + effective_threshold *= fn_confidence; + } + if (config.Bn_compute_blocks > 0) { + float bn_quality = 1.0f - 0.03f * config.Bn_compute_blocks; + bn_quality = std::max(0.5f, std::min(1.0f, bn_quality)); + effective_threshold *= bn_quality; + } + + if (diff < effective_threshold) { + skip_current_step = true; + total_steps_skipped++; + cached_steps.push_back(current_step_index); + continuous_cached_steps++; + accumulated_residual_diff += diff; + apply_cache(cond, input_data, (float*)output->data, ne); + return true; + } + + continuous_cached_steps = 0; + return false; + } + + void after_condition(const void* cond, struct ggml_tensor* input, struct ggml_tensor* output) { + if (!step_is_active()) + return; + + size_t ne = static_cast(ggml_nelements(output)); + update_cache(cond, (float*)input->data, (float*)output->data, ne); + + if (cond == anchor_condition && taylor_config.enabled) { + taylor_state.update_derivatives((float*)output->data, ne, current_step_index); + } + } + + void log_metrics() const { + if (!enabled()) + return; + + LOG_INFO("CacheDIT: steps_skipped=%d/%d (%.1f%%), accum_residual_diff=%.4f", + total_steps_skipped, + current_step_index + 1, + (current_step_index > 0) ? (100.0f * total_steps_skipped / (current_step_index + 1)) : 0.0f, + accumulated_residual_diff); + } +}; + +#endif diff --git a/otherarch/sdcpp/common.hpp b/otherarch/sdcpp/common.hpp index 74b218ab7..b17c11e35 100644 --- a/otherarch/sdcpp/common.hpp +++ b/otherarch/sdcpp/common.hpp @@ -28,7 +28,7 @@ public: if (vae_downsample) { auto conv = std::dynamic_pointer_cast(blocks["conv"]); - x = ggml_pad(ctx->ggml_ctx, x, 1, 1, 0, 0); + x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); x = conv->forward(ctx, x); } else { auto conv = std::dynamic_pointer_cast(blocks["op"]); diff --git a/otherarch/sdcpp/common/common.hpp b/otherarch/sdcpp/common/common.hpp index 6c32de030..001d2e076 100644 --- a/otherarch/sdcpp/common/common.hpp +++ b/otherarch/sdcpp/common/common.hpp @@ -454,10 +454,16 @@ struct SDContextParams { bool diffusion_conv_direct = false; bool vae_conv_direct = false; + bool circular = false; + bool circular_x = false; + bool circular_y = false; + bool chroma_use_dit_mask = true; bool chroma_use_t5_mask = false; int chroma_t5_mask_pad = 1; + bool qwen_image_zero_cond_t = false; + prediction_t prediction = PREDICTION_COUNT; lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO; @@ -610,10 +616,26 @@ struct SDContextParams { "--vae-conv-direct", "use ggml_conv2d_direct in the vae model", true, &vae_conv_direct}, + {"", + "--circular", + "enable circular padding for convolutions", + true, &circular}, + {"", + "--circularx", + "enable circular RoPE wrapping on x-axis (width) only", + true, &circular_x}, + {"", + "--circulary", + "enable circular RoPE wrapping on y-axis (height) only", + true, &circular_y}, {"", "--chroma-disable-dit-mask", "disable dit mask for chroma", false, &chroma_use_dit_mask}, + {"", + "--qwen-image-zero-cond-t", + "enable zero_cond_t for qwen image", + true, &qwen_image_zero_cond_t}, {"", "--chroma-enable-t5-mask", "enable t5 mask for chroma", @@ -873,7 +895,11 @@ struct SDContextParams { << " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" << " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n" + << " circular: " << (circular ? "true" : "false") << ",\n" + << " circular_x: " << (circular_x ? "true" : "false") << ",\n" + << " circular_y: " << (circular_y ? "true" : "false") << ",\n" << " chroma_use_dit_mask: " << (chroma_use_dit_mask ? "true" : "false") << ",\n" + << " qwen_image_zero_cond_t: " << (qwen_image_zero_cond_t ? "true" : "false") << ",\n" << " chroma_use_t5_mask: " << (chroma_use_t5_mask ? "true" : "false") << ",\n" << " chroma_t5_mask_pad: " << chroma_t5_mask_pad << ",\n" << " prediction: " << sd_prediction_name(prediction) << ",\n" @@ -933,10 +959,13 @@ struct SDContextParams { taesd_preview, diffusion_conv_direct, vae_conv_direct, + circular || circular_x, + circular || circular_y, force_sdxl_vae_conv_scale, chroma_use_dit_mask, chroma_use_t5_mask, chroma_t5_mask_pad, + qwen_image_zero_cond_t, flow_shift, }; return sd_ctx_params; @@ -1002,8 +1031,12 @@ struct SDGenerationParams { std::vector custom_sigmas; - std::string easycache_option; - sd_easycache_params_t easycache_params; + std::string cache_mode; + std::string cache_option; + std::string cache_preset; + std::string scm_mask; + bool scm_policy_dynamic = true; + sd_cache_params_t cache_params{}; float moe_boundary = 0.875f; int video_frames = 1; @@ -1365,36 +1398,64 @@ struct SDGenerationParams { return 1; }; - auto on_easycache_arg = [&](int argc, const char** argv, int index) { - const std::string default_values = "0.2,0.15,0.95"; - auto looks_like_value = [](const std::string& token) { - if (token.empty()) { - return false; - } - if (token[0] != '-') { - return true; - } - if (token.size() == 1) { - return false; - } - unsigned char next = static_cast(token[1]); - return std::isdigit(next) || token[1] == '.'; - }; + auto on_cache_mode_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + cache_mode = argv_to_utf8(index, argv); + if (cache_mode != "easycache" && cache_mode != "ucache" && + cache_mode != "dbcache" && cache_mode != "taylorseer" && cache_mode != "cache-dit") { + fprintf(stderr, "error: invalid cache mode '%s', must be 'easycache', 'ucache', 'dbcache', 'taylorseer', or 'cache-dit'\n", cache_mode.c_str()); + return -1; + } + return 1; + }; - std::string option_value; - int consumed = 0; - if (index + 1 < argc) { - std::string next_arg = argv[index + 1]; - if (looks_like_value(next_arg)) { - option_value = argv_to_utf8(index + 1, argv); - consumed = 1; - } + auto on_cache_option_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; } - if (option_value.empty()) { - option_value = default_values; + cache_option = argv_to_utf8(index, argv); + return 1; + }; + + auto on_scm_mask_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; } - easycache_option = option_value; - return consumed; + scm_mask = argv_to_utf8(index, argv); + return 1; + }; + + auto on_scm_policy_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + std::string policy = argv_to_utf8(index, argv); + if (policy == "dynamic") { + scm_policy_dynamic = true; + } else if (policy == "static") { + scm_policy_dynamic = false; + } else { + fprintf(stderr, "error: invalid scm policy '%s', must be 'dynamic' or 'static'\n", policy.c_str()); + return -1; + } + return 1; + }; + + auto on_cache_preset_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + cache_preset = argv_to_utf8(index, argv); + if (cache_preset != "slow" && cache_preset != "s" && cache_preset != "S" && + cache_preset != "medium" && cache_preset != "m" && cache_preset != "M" && + cache_preset != "fast" && cache_preset != "f" && cache_preset != "F" && + cache_preset != "ultra" && cache_preset != "u" && cache_preset != "U") { + fprintf(stderr, "error: invalid cache preset '%s', must be 'slow'/'s', 'medium'/'m', 'fast'/'f', or 'ultra'/'u'\n", cache_preset.c_str()); + return -1; + } + return 1; }; options.manual_options = { @@ -1433,9 +1494,25 @@ struct SDGenerationParams { "reference image for Flux Kontext models (can be used multiple times)", on_ref_image_arg}, {"", - "--easycache", - "enable EasyCache for DiT models with optional \"threshold,start_percent,end_percent\" (default: 0.2,0.15,0.95)", - on_easycache_arg}, + "--cache-mode", + "caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level)", + on_cache_mode_arg}, + {"", + "--cache-option", + "named cache params (key=value format, comma-separated):\n - easycache/ucache: threshold=,start=,end=,decay=,relative=,reset=\n - dbcache/taylorseer/cache-dit: Fn=,Bn=,threshold=,warmup=\n Examples: \"threshold=0.25\" or \"threshold=1.5,reset=0\"", + on_cache_option_arg}, + {"", + "--cache-preset", + "cache-dit preset: 'slow'/'s', 'medium'/'m', 'fast'/'f', 'ultra'/'u'", + on_cache_preset_arg}, + {"", + "--scm-mask", + "SCM steps mask for cache-dit: comma-separated 0/1 (e.g., \"1,1,1,0,0,1,0,0,1,0\") - 1=compute, 0=can cache", + on_scm_mask_arg}, + {"", + "--scm-policy", + "SCM policy: 'dynamic' (default) or 'static'", + on_scm_policy_arg}, }; @@ -1478,7 +1555,10 @@ struct SDGenerationParams { load_if_exists("prompt", prompt); load_if_exists("negative_prompt", negative_prompt); - load_if_exists("easycache_option", easycache_option); + load_if_exists("cache_mode", cache_mode); + load_if_exists("cache_option", cache_option); + load_if_exists("cache_preset", cache_preset); + load_if_exists("scm_mask", scm_mask); load_if_exists("clip_skip", clip_skip); load_if_exists("width", width); @@ -1618,57 +1698,118 @@ struct SDGenerationParams { return false; } - if (!easycache_option.empty()) { - float values[3] = {0.0f, 0.0f, 0.0f}; - std::stringstream ss(easycache_option); + sd_cache_params_init(&cache_params); + + auto parse_named_params = [&](const std::string& opt_str) -> bool { + std::stringstream ss(opt_str); std::string token; - int idx = 0; while (std::getline(ss, token, ',')) { - auto trim = [](std::string& s) { - const char* whitespace = " \t\r\n"; - auto start = s.find_first_not_of(whitespace); - if (start == std::string::npos) { - s.clear(); - return; - } - auto end = s.find_last_not_of(whitespace); - s = s.substr(start, end - start + 1); - }; - trim(token); - if (token.empty()) { - LOG_ERROR("error: invalid easycache option '%s'", easycache_option.c_str()); - return false; - } - if (idx >= 3) { - LOG_ERROR("error: easycache expects exactly 3 comma-separated values (threshold,start,end)\n"); + size_t eq_pos = token.find('='); + if (eq_pos == std::string::npos) { + LOG_ERROR("error: cache option '%s' missing '=' separator", token.c_str()); return false; } + std::string key = token.substr(0, eq_pos); + std::string val = token.substr(eq_pos + 1); try { - values[idx] = std::stof(token); + if (key == "threshold") { + if (cache_mode == "easycache" || cache_mode == "ucache") { + cache_params.reuse_threshold = std::stof(val); + } else { + cache_params.residual_diff_threshold = std::stof(val); + } + } else if (key == "start") { + cache_params.start_percent = std::stof(val); + } else if (key == "end") { + cache_params.end_percent = std::stof(val); + } else if (key == "decay") { + cache_params.error_decay_rate = std::stof(val); + } else if (key == "relative") { + cache_params.use_relative_threshold = (std::stof(val) != 0.0f); + } else if (key == "reset") { + cache_params.reset_error_on_compute = (std::stof(val) != 0.0f); + } else if (key == "Fn" || key == "fn") { + cache_params.Fn_compute_blocks = std::stoi(val); + } else if (key == "Bn" || key == "bn") { + cache_params.Bn_compute_blocks = std::stoi(val); + } else if (key == "warmup") { + cache_params.max_warmup_steps = std::stoi(val); + } else { + LOG_ERROR("error: unknown cache parameter '%s'", key.c_str()); + return false; + } } catch (const std::exception&) { - LOG_ERROR("error: invalid easycache value '%s'", token.c_str()); + LOG_ERROR("error: invalid value '%s' for parameter '%s'", val.c_str(), key.c_str()); return false; } - idx++; } - if (idx != 3) { - LOG_ERROR("error: easycache expects exactly 3 comma-separated values (threshold,start,end)\n"); - return false; + return true; + }; + + if (!cache_mode.empty()) { + if (cache_mode == "easycache") { + cache_params.mode = SD_CACHE_EASYCACHE; + cache_params.reuse_threshold = 0.2f; + cache_params.start_percent = 0.15f; + cache_params.end_percent = 0.95f; + cache_params.error_decay_rate = 1.0f; + cache_params.use_relative_threshold = true; + cache_params.reset_error_on_compute = true; + } else if (cache_mode == "ucache") { + cache_params.mode = SD_CACHE_UCACHE; + cache_params.reuse_threshold = 1.0f; + cache_params.start_percent = 0.15f; + cache_params.end_percent = 0.95f; + cache_params.error_decay_rate = 1.0f; + cache_params.use_relative_threshold = true; + cache_params.reset_error_on_compute = true; + } else if (cache_mode == "dbcache") { + cache_params.mode = SD_CACHE_DBCACHE; + cache_params.Fn_compute_blocks = 8; + cache_params.Bn_compute_blocks = 0; + cache_params.residual_diff_threshold = 0.08f; + cache_params.max_warmup_steps = 8; + } else if (cache_mode == "taylorseer") { + cache_params.mode = SD_CACHE_TAYLORSEER; + cache_params.Fn_compute_blocks = 8; + cache_params.Bn_compute_blocks = 0; + cache_params.residual_diff_threshold = 0.08f; + cache_params.max_warmup_steps = 8; + } else if (cache_mode == "cache-dit") { + cache_params.mode = SD_CACHE_CACHE_DIT; + cache_params.Fn_compute_blocks = 8; + cache_params.Bn_compute_blocks = 0; + cache_params.residual_diff_threshold = 0.08f; + cache_params.max_warmup_steps = 8; } - if (values[0] < 0.0f) { - LOG_ERROR("error: easycache threshold must be non-negative\n"); - return false; + + if (!cache_option.empty()) { + if (!parse_named_params(cache_option)) { + return false; + } } - if (values[1] < 0.0f || values[1] >= 1.0f || values[2] <= 0.0f || values[2] > 1.0f || values[1] >= values[2]) { - LOG_ERROR("error: easycache start/end percents must satisfy 0.0 <= start < end <= 1.0\n"); - return false; + + if (cache_mode == "easycache" || cache_mode == "ucache") { + if (cache_params.reuse_threshold < 0.0f) { + LOG_ERROR("error: cache threshold must be non-negative"); + return false; + } + if (cache_params.start_percent < 0.0f || cache_params.start_percent >= 1.0f || + cache_params.end_percent <= 0.0f || cache_params.end_percent > 1.0f || + cache_params.start_percent >= cache_params.end_percent) { + LOG_ERROR("error: cache start/end percents must satisfy 0.0 <= start < end <= 1.0"); + return false; + } } - easycache_params.enabled = true; - easycache_params.reuse_threshold = values[0]; - easycache_params.start_percent = values[1]; - easycache_params.end_percent = values[2]; - } else { - easycache_params.enabled = false; + } + + if (cache_params.mode == SD_CACHE_DBCACHE || + cache_params.mode == SD_CACHE_TAYLORSEER || + cache_params.mode == SD_CACHE_CACHE_DIT) { + if (!scm_mask.empty()) { + cache_params.scm_mask = scm_mask.c_str(); + } + cache_params.scm_policy_dynamic = scm_policy_dynamic; } sample_params.guidance.slg.layers = skip_layers.data(); @@ -1770,12 +1911,13 @@ struct SDGenerationParams { << " high_noise_skip_layers: " << vec_to_string(high_noise_skip_layers) << ",\n" << " high_noise_sample_params: " << high_noise_sample_params_str << ",\n" << " custom_sigmas: " << vec_to_string(custom_sigmas) << ",\n" - << " easycache_option: \"" << easycache_option << "\",\n" - << " easycache: " - << (easycache_params.enabled ? "enabled" : "disabled") - << " (threshold=" << easycache_params.reuse_threshold - << ", start=" << easycache_params.start_percent - << ", end=" << easycache_params.end_percent << "),\n" + << " cache_mode: \"" << cache_mode << "\",\n" + << " cache_option: \"" << cache_option << "\",\n" + << " cache: " + << (cache_params.mode != SD_CACHE_DISABLED ? "enabled" : "disabled") + << " (threshold=" << cache_params.reuse_threshold + << ", start=" << cache_params.start_percent + << ", end=" << cache_params.end_percent << "),\n" << " moe_boundary: " << moe_boundary << ",\n" << " video_frames: " << video_frames << ",\n" << " fps: " << fps << ",\n" diff --git a/otherarch/sdcpp/denoiser.hpp b/otherarch/sdcpp/denoiser.hpp index fc5230d7b..7a8242e7d 100644 --- a/otherarch/sdcpp/denoiser.hpp +++ b/otherarch/sdcpp/denoiser.hpp @@ -366,18 +366,18 @@ struct KLOptimalScheduler : SigmaScheduler { for (uint32_t i = 0; i < n; ++i) { // t goes from 0.0 to 1.0 - float t = static_cast(i) / static_cast(n-1); + float t = static_cast(i) / static_cast(n - 1); // Interpolate in the angle domain float angle = t * alpha_min + (1.0f - t) * alpha_max; // Convert back to sigma sigmas.push_back(std::tan(angle)); - } + } // Append the final zero to sigma sigmas.push_back(0.0f); - + return sigmas; } }; diff --git a/otherarch/sdcpp/diffusion_model.hpp b/otherarch/sdcpp/diffusion_model.hpp index 8c741fdc4..06cbecc28 100644 --- a/otherarch/sdcpp/diffusion_model.hpp +++ b/otherarch/sdcpp/diffusion_model.hpp @@ -37,8 +37,9 @@ struct DiffusionModel { virtual void get_param_tensors(std::map& tensors) = 0; virtual size_t get_params_buffer_size() = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter){}; - virtual int64_t get_adm_in_channels() = 0; - virtual void set_flash_attn_enabled(bool enabled) = 0; + virtual int64_t get_adm_in_channels() = 0; + virtual void set_flash_attn_enabled(bool enabled) = 0; + virtual void set_circular_axes(bool circular_x, bool circular_y) = 0; }; struct UNetModel : public DiffusionModel { @@ -87,6 +88,10 @@ struct UNetModel : public DiffusionModel { unet.set_flash_attention_enabled(enabled); } + void set_circular_axes(bool circular_x, bool circular_y) override { + unet.set_circular_axes(circular_x, circular_y); + } + bool compute(int n_threads, DiffusionParams diffusion_params, struct ggml_tensor** output = nullptr, @@ -148,6 +153,10 @@ struct MMDiTModel : public DiffusionModel { mmdit.set_flash_attention_enabled(enabled); } + void set_circular_axes(bool circular_x, bool circular_y) override { + mmdit.set_circular_axes(circular_x, circular_y); + } + bool compute(int n_threads, DiffusionParams diffusion_params, struct ggml_tensor** output = nullptr, @@ -210,6 +219,10 @@ struct FluxModel : public DiffusionModel { flux.set_flash_attention_enabled(enabled); } + void set_circular_axes(bool circular_x, bool circular_y) override { + flux.set_circular_axes(circular_x, circular_y); + } + bool compute(int n_threads, DiffusionParams diffusion_params, struct ggml_tensor** output = nullptr, @@ -277,6 +290,10 @@ struct WanModel : public DiffusionModel { wan.set_flash_attention_enabled(enabled); } + void set_circular_axes(bool circular_x, bool circular_y) override { + wan.set_circular_axes(circular_x, circular_y); + } + bool compute(int n_threads, DiffusionParams diffusion_params, struct ggml_tensor** output = nullptr, @@ -303,8 +320,9 @@ struct QwenImageModel : public DiffusionModel { bool offload_params_to_cpu, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "model.diffusion_model", - SDVersion version = VERSION_QWEN_IMAGE) - : prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version) { + SDVersion version = VERSION_QWEN_IMAGE, + bool zero_cond_t = false) + : prefix(prefix), qwen_image(backend, offload_params_to_cpu, tensor_storage_map, prefix, version, zero_cond_t) { } std::string get_desc() override { @@ -343,6 +361,10 @@ struct QwenImageModel : public DiffusionModel { qwen_image.set_flash_attention_enabled(enabled); } + void set_circular_axes(bool circular_x, bool circular_y) override { + qwen_image.set_circular_axes(circular_x, circular_y); + } + bool compute(int n_threads, DiffusionParams diffusion_params, struct ggml_tensor** output = nullptr, @@ -406,6 +428,10 @@ struct ZImageModel : public DiffusionModel { z_image.set_flash_attention_enabled(enabled); } + void set_circular_axes(bool circular_x, bool circular_y) override { + z_image.set_circular_axes(circular_x, circular_y); + } + bool compute(int n_threads, DiffusionParams diffusion_params, struct ggml_tensor** output = nullptr, diff --git a/otherarch/sdcpp/flux.hpp b/otherarch/sdcpp/flux.hpp index 7ce263569..86e2007ad 100644 --- a/otherarch/sdcpp/flux.hpp +++ b/otherarch/sdcpp/flux.hpp @@ -233,14 +233,17 @@ namespace Flux { __STATIC_INLINE__ struct ggml_tensor* modulate(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* shift, - struct ggml_tensor* scale) { + struct ggml_tensor* scale, + bool skip_reshape = false) { // x: [N, L, C] // scale: [N, C] // shift: [N, C] - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] - shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C] - x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); - x = ggml_add(ctx, x, shift); + if (!skip_reshape) { + scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] + shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C] + } + x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + x = ggml_add(ctx, x, shift); return x; } @@ -860,14 +863,14 @@ namespace Flux { } } - struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx, + struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, struct ggml_tensor* x) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; int pad_h = (params.patch_size - H % params.patch_size) % params.patch_size; int pad_w = (params.patch_size - W % params.patch_size) % params.patch_size; - x = ggml_pad(ctx, x, pad_w, pad_h, 0, 0); // [N, C, H + pad_h, W + pad_w] + x = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return x; } @@ -893,11 +896,11 @@ namespace Flux { return x; } - struct ggml_tensor* process_img(struct ggml_context* ctx, + struct ggml_tensor* process_img(GGMLRunnerContext* ctx, struct ggml_tensor* x) { // img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size) x = pad_to_patch_size(ctx, x); - x = patchify(ctx, x); + x = patchify(ctx->ggml_ctx, x); return x; } @@ -1076,7 +1079,7 @@ namespace Flux { int pad_h = (patch_size - H % patch_size) % patch_size; int pad_w = (patch_size - W % patch_size) % patch_size; - auto img = pad_to_patch_size(ctx->ggml_ctx, x); + auto img = pad_to_patch_size(ctx, x); auto orig_img = img; if (params.chroma_radiance_params.use_patch_size_32) { @@ -1150,7 +1153,7 @@ namespace Flux { int pad_h = (patch_size - H % patch_size) % patch_size; int pad_w = (patch_size - W % patch_size) % patch_size; - auto img = process_img(ctx->ggml_ctx, x); + auto img = process_img(ctx, x); uint64_t img_tokens = img->ne[1]; if (params.version == VERSION_FLUX_FILL) { @@ -1158,8 +1161,8 @@ namespace Flux { ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0); ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C); - masked = process_img(ctx->ggml_ctx, masked); - mask = process_img(ctx->ggml_ctx, mask); + masked = process_img(ctx, masked); + mask = process_img(ctx, mask); img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, masked, mask, 0), 0); } else if (params.version == VERSION_FLEX_2) { @@ -1168,21 +1171,21 @@ namespace Flux { ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C); ggml_tensor* control = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * (C + 1)); - masked = process_img(ctx->ggml_ctx, masked); - mask = process_img(ctx->ggml_ctx, mask); - control = process_img(ctx->ggml_ctx, control); + masked = process_img(ctx, masked); + mask = process_img(ctx, mask); + control = process_img(ctx, control); img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, masked, mask, 0), control, 0), 0); } else if (params.version == VERSION_FLUX_CONTROLS) { GGML_ASSERT(c_concat != nullptr); - auto control = process_img(ctx->ggml_ctx, c_concat); + auto control = process_img(ctx, c_concat); img = ggml_concat(ctx->ggml_ctx, img, control, 0); } if (ref_latents.size() > 0) { for (ggml_tensor* ref : ref_latents) { - ref = process_img(ctx->ggml_ctx, ref); + ref = process_img(ctx, ref); img = ggml_concat(ctx->ggml_ctx, img, ref, 1); } } @@ -1472,6 +1475,8 @@ namespace Flux { increase_ref_index, flux_params.ref_index_scale, flux_params.theta, + circular_y_enabled, + circular_x_enabled, flux_params.axes_dim); int pos_len = pe_vec.size() / flux_params.axes_dim_sum / 2; // LOG_DEBUG("pos_len %d", pos_len); diff --git a/otherarch/sdcpp/ggml_extend.hpp b/otherarch/sdcpp/ggml_extend.hpp index 7b355d3ee..f6480fd2f 100644 --- a/otherarch/sdcpp/ggml_extend.hpp +++ b/otherarch/sdcpp/ggml_extend.hpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -993,6 +994,48 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_ext_linear(struct ggml_context* ctx, return x; } +__STATIC_INLINE__ struct ggml_tensor* ggml_ext_pad_ext(struct ggml_context* ctx, + struct ggml_tensor* x, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3, + bool circular_x = false, + bool circular_y = false) { + if (circular_x && circular_y) { + return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + } + + if (circular_x && (lp0 != 0 || rp0 != 0)) { + x = ggml_pad_ext_circular(ctx, x, lp0, rp0, 0, 0, 0, 0, 0, 0); + lp0 = rp0 = 0; + } + if (circular_y && (lp1 != 0 || rp1 != 0)) { + x = ggml_pad_ext_circular(ctx, x, 0, 0, lp1, rp1, 0, 0, 0, 0); + lp1 = rp1 = 0; + } + + if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) { + x = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + } + return x; +} + +__STATIC_INLINE__ struct ggml_tensor* ggml_ext_pad(struct ggml_context* ctx, + struct ggml_tensor* x, + int p0, + int p1, + int p2 = 0, + int p3 = 0, + bool circular_x = false, + bool circular_y = false) { + return ggml_ext_pad_ext(ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); +} + // w: [OC，IC, KH, KW] // x: [N, IC, IH, IW] // b: [OC,] @@ -1001,20 +1044,29 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_ext_conv_2d(struct ggml_context* ctx, struct ggml_tensor* x, struct ggml_tensor* w, struct ggml_tensor* b, - int s0 = 1, - int s1 = 1, - int p0 = 0, - int p1 = 0, - int d0 = 1, - int d1 = 1, - bool direct = false, - float scale = 1.f) { + int s0 = 1, + int s1 = 1, + int p0 = 0, + int p1 = 0, + int d0 = 1, + int d1 = 1, + bool direct = false, + bool circular_x = false, + bool circular_y = false, + float scale = 1.f) { if (scale != 1.f) { x = ggml_scale(ctx, x, scale); } if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) { w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]); } + + if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) { + x = ggml_ext_pad_ext(ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); + p0 = 0; + p1 = 0; + } + if (direct) { x = ggml_conv_2d_direct(ctx, w, x, s0, s1, p0, p1, d0, d1); } else { @@ -1521,14 +1573,16 @@ struct WeightAdapter { float scale = 1.f; } linear; struct { - int s0 = 1; - int s1 = 1; - int p0 = 0; - int p1 = 0; - int d0 = 1; - int d1 = 1; - bool direct = false; - float scale = 1.f; + int s0 = 1; + int s1 = 1; + int p0 = 0; + int p1 = 0; + int d0 = 1; + int d1 = 1; + bool direct = false; + bool circular_x = false; + bool circular_y = false; + float scale = 1.f; } conv2d; }; virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_tensor* weight, const std::string& weight_name) = 0; @@ -1546,6 +1600,8 @@ struct GGMLRunnerContext { ggml_context* ggml_ctx = nullptr; bool flash_attn_enabled = false; bool conv2d_direct_enabled = false; + bool circular_x_enabled = false; + bool circular_y_enabled = false; std::shared_ptr weight_adapter = nullptr; }; @@ -1582,6 +1638,8 @@ protected: bool flash_attn_enabled = false; bool conv2d_direct_enabled = false; + bool circular_x_enabled = false; + bool circular_y_enabled = false; void alloc_params_ctx() { struct ggml_init_params params; @@ -1859,6 +1917,8 @@ public: runner_ctx.backend = runtime_backend; runner_ctx.flash_attn_enabled = flash_attn_enabled; runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled; + runner_ctx.circular_x_enabled = circular_x_enabled; + runner_ctx.circular_y_enabled = circular_y_enabled; runner_ctx.weight_adapter = weight_adapter; return runner_ctx; } @@ -2003,6 +2063,11 @@ public: conv2d_direct_enabled = enabled; } + void set_circular_axes(bool circular_x, bool circular_y) { + circular_x_enabled = circular_x; + circular_y_enabled = circular_y; + } + void set_weight_adapter(const std::shared_ptr& adapter) { weight_adapter = adapter; } @@ -2266,15 +2331,17 @@ public: } if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; - forward_params.conv2d.s0 = stride.second; - forward_params.conv2d.s1 = stride.first; - forward_params.conv2d.p0 = padding.second; - forward_params.conv2d.p1 = padding.first; - forward_params.conv2d.d0 = dilation.second; - forward_params.conv2d.d1 = dilation.first; - forward_params.conv2d.direct = ctx->conv2d_direct_enabled; - forward_params.conv2d.scale = scale; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; + forward_params.conv2d.s0 = stride.second; + forward_params.conv2d.s1 = stride.first; + forward_params.conv2d.p0 = padding.second; + forward_params.conv2d.p1 = padding.first; + forward_params.conv2d.d0 = dilation.second; + forward_params.conv2d.d1 = dilation.first; + forward_params.conv2d.direct = ctx->conv2d_direct_enabled; + forward_params.conv2d.circular_x = ctx->circular_x_enabled; + forward_params.conv2d.circular_y = ctx->circular_y_enabled; + forward_params.conv2d.scale = scale; return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, x, w, b, prefix, forward_params); } return ggml_ext_conv_2d(ctx->ggml_ctx, @@ -2288,6 +2355,8 @@ public: dilation.second, dilation.first, ctx->conv2d_direct_enabled, + ctx->circular_x_enabled, + ctx->circular_y_enabled, scale); } }; diff --git a/otherarch/sdcpp/lora.hpp b/otherarch/sdcpp/lora.hpp index 43cd2fed0..415765068 100644 --- a/otherarch/sdcpp/lora.hpp +++ b/otherarch/sdcpp/lora.hpp @@ -599,6 +599,8 @@ struct LoraModel : public GGMLRunner { forward_params.conv2d.d0, forward_params.conv2d.d1, forward_params.conv2d.direct, + forward_params.conv2d.circular_x, + forward_params.conv2d.circular_y, forward_params.conv2d.scale); if (lora_mid) { lx = ggml_ext_conv_2d(ctx, @@ -612,6 +614,8 @@ struct LoraModel : public GGMLRunner { 1, 1, forward_params.conv2d.direct, + forward_params.conv2d.circular_x, + forward_params.conv2d.circular_y, forward_params.conv2d.scale); } lx = ggml_ext_conv_2d(ctx, @@ -625,6 +629,8 @@ struct LoraModel : public GGMLRunner { 1, 1, forward_params.conv2d.direct, + forward_params.conv2d.circular_x, + forward_params.conv2d.circular_y, forward_params.conv2d.scale); } @@ -716,11 +722,11 @@ struct LoraModel : public GGMLRunner { /* Don't worry if this message shows up twice in the logs per LoRA, * this function is called once to calculate the required buffer size * and then again to actually generate a graph to be used */ - if (applied_lora_tensors_count != total_lora_tensors_count) { - printf("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s", + if (!at_runntime && applied_lora_tensors_count != total_lora_tensors_count) { + LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s", applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str()); } else { - LOG_WARN("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s", + LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s", applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str()); } } @@ -779,6 +785,8 @@ public: forward_params.conv2d.d0, forward_params.conv2d.d1, forward_params.conv2d.direct, + forward_params.conv2d.circular_x, + forward_params.conv2d.circular_y, forward_params.conv2d.scale); } for (auto& lora_model : lora_models) { diff --git a/otherarch/sdcpp/main.cpp b/otherarch/sdcpp/main.cpp index a1b92f77e..77ef1c935 100644 --- a/otherarch/sdcpp/main.cpp +++ b/otherarch/sdcpp/main.cpp @@ -26,12 +26,16 @@ const char* previews_str[] = { "vae", }; +std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)"); + struct SDCliParams { SDMode mode = IMG_GEN; std::string output_path = "output.png"; + int output_begin_idx = -1; bool verbose = false; bool canny_preprocess = false; + bool convert_name = false; preview_t preview_method = PREVIEW_NONE; int preview_interval = 1; @@ -49,7 +53,7 @@ struct SDCliParams { options.string_options = { {"-o", "--output", - "path to write result image to (default: ./output.png)", + "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png)", &output_path}, {"", "--preview-path", @@ -62,6 +66,10 @@ struct SDCliParams { "--preview-interval", "interval in denoising steps between consecutive updates of the image preview file (default is 1, meaning updating at every step)", &preview_interval}, + {"", + "--output-begin-idx", + "starting index for output image sequence, must be non-negative (default 0 if specified %d in output path, 1 otherwise)", + &output_begin_idx}, }; options.bool_options = { @@ -69,6 +77,10 @@ struct SDCliParams { "--canny", "apply canny preprocessor (edge detection)", true, &canny_preprocess}, + {"", + "--convert-name", + "convert tensor name (for convert mode)", + true, &convert_name}, {"-v", "--verbose", "print extra info", @@ -174,6 +186,7 @@ struct SDCliParams { << " verbose: " << (verbose ? "true" : "false") << ",\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" << " preview_path: \"" << preview_path << "\",\n" @@ -338,6 +351,25 @@ void step_callback(int step, int frame_count, sd_image_t* image, bool is_noisy, } } +std::string format_frame_idx(std::string pattern, int frame_idx) { + std::smatch match; + std::string result = pattern; + while (std::regex_search(result, match, format_specifier_regex)) { + std::string specifier = match.str(1); + char buffer[32]; + snprintf(buffer, sizeof(buffer), specifier.c_str(), frame_idx); + result.replace(match.position(1), match.length(1), buffer); + } + + // Then replace all '%%' with '%' + size_t pos = 0; + while ((pos = result.find("%%", pos)) != std::string::npos) { + result.replace(pos, 2, "%"); + pos += 1; + } + return result; +} + int main(int argc, const char* argv[]) { if (argc > 1 && std::string(argv[1]) == "--version") { std::cout << version_string() << "\n"; @@ -387,7 +419,8 @@ int main(int argc, const char* argv[]) { ctx_params.vae_path.c_str(), cli_params.output_path.c_str(), ctx_params.wtype, - ctx_params.tensor_type_rules.c_str()); + ctx_params.tensor_type_rules.c_str(), + cli_params.convert_name); if (!success) { LOG_ERROR("convert '%s'/'%s' to '%s' failed", ctx_params.model_path.c_str(), @@ -610,7 +643,7 @@ int main(int argc, const char* argv[]) { gen_params.pm_style_strength, }, // pm_params ctx_params.vae_tiling_params, - gen_params.easycache_params, + gen_params.cache_params, }; results = generate_image(sd_ctx, &img_gen_params); @@ -635,7 +668,7 @@ int main(int argc, const char* argv[]) { gen_params.seed, gen_params.video_frames, gen_params.vace_strength, - gen_params.easycache_params, + gen_params.cache_params, }; results = generate_video(sd_ctx, &vid_gen_params, &num_results); @@ -712,25 +745,59 @@ int main(int argc, const char* argv[]) { is_jpg = false; } - if (cli_params.mode == VID_GEN && num_results > 1) { - std::string vid_output_path = cli_params.output_path; - if (file_ext_lower == ".png") { - vid_output_path = base_path + ".avi"; + if (std::regex_search(cli_params.output_path, format_specifier_regex)) { + std::string final_output_path = cli_params.output_path; + if (cli_params.output_begin_idx == -1) { + cli_params.output_begin_idx = 0; } - create_mjpg_avi_from_sd_images(vid_output_path.c_str(), results, num_results, gen_params.fps); - LOG_INFO("save result MJPG AVI video to '%s'\n", vid_output_path.c_str()); + // writing image sequence, default to PNG + if (!is_jpg && file_ext_lower != ".png") { + base_path += file_ext; + file_ext = ".png"; + } + final_output_path = base_path + file_ext; + for (int i = 0; i < num_results; i++) { + if (results[i].data == nullptr) { + continue; + } + std::string final_image_path = format_frame_idx(final_output_path, cli_params.output_begin_idx + i); + if (is_jpg) { + int write_ok = stbi_write_jpg(final_image_path.c_str(), results[i].width, results[i].height, results[i].channel, + results[i].data, 90, get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + i).c_str()); + LOG_INFO("save result JPEG image %d to '%s' (%s)", i, final_image_path.c_str(), write_ok == 0 ? "failure" : "success"); + } else { + int write_ok = stbi_write_png(final_image_path.c_str(), results[i].width, results[i].height, results[i].channel, + results[i].data, 0, get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + i).c_str()); + LOG_INFO("save result PNG image %d to '%s' (%s)", i, final_image_path.c_str(), write_ok == 0 ? "failure" : "success"); + } + } + } else if (cli_params.mode == VID_GEN && num_results > 1) { + std::string final_output_path = cli_params.output_path; + if (file_ext_lower != ".avi") { + if (!is_jpg && file_ext_lower != ".png") { + base_path += file_ext; + } + file_ext = ".avi"; + final_output_path = base_path + file_ext; + } + create_mjpg_avi_from_sd_images(final_output_path.c_str(), results, num_results, gen_params.fps); + LOG_INFO("save result MJPG AVI video to '%s'\n", final_output_path.c_str()); } else { // appending ".png" to absent or unknown extension if (!is_jpg && file_ext_lower != ".png") { base_path += file_ext; file_ext = ".png"; } + if (cli_params.output_begin_idx == -1) { + cli_params.output_begin_idx = 1; + } for (int i = 0; i < num_results; i++) { if (results[i].data == nullptr) { continue; } int write_ok; - std::string final_image_path = i > 0 ? base_path + "_" + std::to_string(i + 1) + file_ext : base_path + file_ext; + std::string final_image_path; + final_image_path = i > 0 ? base_path + "_" + std::to_string(cli_params.output_begin_idx + i) + file_ext : base_path + file_ext; if (is_jpg) { write_ok = stbi_write_jpg(final_image_path.c_str(), results[i].width, results[i].height, results[i].channel, results[i].data, 90, get_image_params(cli_params, ctx_params, gen_params, gen_params.seed + i).c_str()); diff --git a/otherarch/sdcpp/mmdit.hpp b/otherarch/sdcpp/mmdit.hpp index 38bdc2e74..eeb74a268 100644 --- a/otherarch/sdcpp/mmdit.hpp +++ b/otherarch/sdcpp/mmdit.hpp @@ -983,4 +983,4 @@ struct MMDiTRunner : public GGMLRunner { } }; -#endif \ No newline at end of file +#endif diff --git a/otherarch/sdcpp/model.cpp b/otherarch/sdcpp/model.cpp index f2edfb9e2..3eb071786 100644 --- a/otherarch/sdcpp/model.cpp +++ b/otherarch/sdcpp/model.cpp @@ -1837,7 +1837,12 @@ int64_t ModelLoader::get_params_mem_size(ggml_backend_t backend, ggml_type type) return mem_size; } -bool convert(const char* input_path, const char* vae_path, const char* output_path, sd_type_t output_type, const char* tensor_type_rules) { +bool convert(const char* input_path, + const char* vae_path, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name) { ModelLoader model_loader; if (!model_loader.init_from_file(input_path)) { @@ -1851,7 +1856,9 @@ bool convert(const char* input_path, const char* vae_path, const char* output_pa return false; } } - model_loader.convert_tensors_name(); + if (convert_name) { + model_loader.convert_tensors_name(); + } bool success = model_loader.save_to_gguf_file(output_path, (ggml_type)output_type, tensor_type_rules); return success; } diff --git a/otherarch/sdcpp/name_conversion.cpp b/otherarch/sdcpp/name_conversion.cpp index 359a0a6a0..7b22e505e 100644 --- a/otherarch/sdcpp/name_conversion.cpp +++ b/otherarch/sdcpp/name_conversion.cpp @@ -960,6 +960,7 @@ bool is_first_stage_model_name(const std::string& name) { std::string convert_tensor_name(std::string name, SDVersion version) { bool is_lora = false; bool is_lycoris_underline = false; + bool is_underline = false; std::vector lora_prefix_vec = { "lora.lora.", "lora.lora_", @@ -967,12 +968,27 @@ std::string convert_tensor_name(std::string name, SDVersion version) { "lora.lycoris.", "lora.", }; + std::vector underline_lora_prefix_vec = { + "unet_", + "te_", + "te1_", + "te2_", + "te3_", + "vae_", + }; for (const auto& prefix : lora_prefix_vec) { if (starts_with(name, prefix)) { is_lora = true; name = name.substr(prefix.size()); if (contains(prefix, "lycoris_")) { is_lycoris_underline = true; + } else { + for (const auto& underline_lora_prefix : underline_lora_prefix_vec) { + if (starts_with(name, underline_lora_prefix)) { + is_underline = true; + break; + } + } } break; } @@ -1034,7 +1050,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { // LOG_DEBUG("name %s %d", name.c_str(), version); - if (sd_version_is_unet(version) || sd_version_is_flux(version) || is_lycoris_underline) { + if (sd_version_is_unet(version) || is_underline || is_lycoris_underline) { name = convert_sep_to_dot(name); } diff --git a/otherarch/sdcpp/qwen_image.hpp b/otherarch/sdcpp/qwen_image.hpp index eeb823d50..9c095a9a0 100644 --- a/otherarch/sdcpp/qwen_image.hpp +++ b/otherarch/sdcpp/qwen_image.hpp @@ -191,11 +191,16 @@ namespace Qwen { }; class QwenImageTransformerBlock : public GGMLBlock { + protected: + bool zero_cond_t; + public: QwenImageTransformerBlock(int64_t dim, int64_t num_attention_heads, int64_t attention_head_dim, - float eps = 1e-6) { + float eps = 1e-6, + bool zero_cond_t = false) + : zero_cond_t(zero_cond_t) { // img_mod.0 is nn.SiLU() blocks["img_mod.1"] = std::shared_ptr(new Linear(dim, 6 * dim, true)); @@ -220,11 +225,37 @@ namespace Qwen { eps)); } + std::vector get_mod_params_vec(ggml_context* ctx, ggml_tensor* mod_params, ggml_tensor* index = nullptr) { + // index: [N, n_img_token] + // mod_params: [N, hidden_size * 12] + if (index == nullptr) { + return ggml_ext_chunk(ctx, mod_params, 6, 0); + } + mod_params = ggml_reshape_1d(ctx, mod_params, ggml_nelements(mod_params)); + auto mod_params_vec = ggml_ext_chunk(ctx, mod_params, 12, 0); + index = ggml_reshape_3d(ctx, index, 1, index->ne[0], index->ne[1]); // [N, n_img_token, 1] + index = ggml_repeat_4d(ctx, index, mod_params_vec[0]->ne[0], index->ne[1], index->ne[2], index->ne[3]); // [N, n_img_token, hidden_size] + std::vector mod_results; + for (int i = 0; i < 6; i++) { + auto mod_0 = mod_params_vec[i]; + auto mod_1 = mod_params_vec[i + 6]; + + // mod_result = torch.where(index == 0, mod_0, mod_1) + // mod_result = (1 - index)*mod_0 + index*mod_1 + mod_0 = ggml_sub(ctx, ggml_repeat(ctx, mod_0, index), ggml_mul(ctx, index, mod_0)); // [N, n_img_token, hidden_size] + mod_1 = ggml_mul(ctx, index, mod_1); // [N, n_img_token, hidden_size] + auto mod_result = ggml_add(ctx, mod_0, mod_1); + mod_results.push_back(mod_result); + } + return mod_results; + } + virtual std::pair forward(GGMLRunnerContext* ctx, struct ggml_tensor* img, struct ggml_tensor* txt, struct ggml_tensor* t_emb, - struct ggml_tensor* pe) { + struct ggml_tensor* pe, + struct ggml_tensor* modulate_index = nullptr) { // img: [N, n_img_token, hidden_size] // txt: [N, n_txt_token, hidden_size] // pe: [n_img_token + n_txt_token, d_head/2, 2, 2] @@ -244,14 +275,18 @@ namespace Qwen { auto img_mod_params = ggml_silu(ctx->ggml_ctx, t_emb); img_mod_params = img_mod_1->forward(ctx, img_mod_params); - auto img_mod_param_vec = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0); + auto img_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, img_mod_params, modulate_index); + + if (zero_cond_t) { + t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0]; + } auto txt_mod_params = ggml_silu(ctx->ggml_ctx, t_emb); txt_mod_params = txt_mod_1->forward(ctx, txt_mod_params); - auto txt_mod_param_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0); + auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params); auto img_normed = img_norm1->forward(ctx, img); - auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1]); + auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr); auto img_gate1 = img_mod_param_vec[2]; auto txt_normed = txt_norm1->forward(ctx, txt); @@ -264,7 +299,7 @@ namespace Qwen { txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_gate1)); auto img_normed2 = img_norm2->forward(ctx, img); - auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4]); + auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr); auto img_gate2 = img_mod_param_vec[5]; auto txt_normed2 = txt_norm2->forward(ctx, txt); @@ -325,6 +360,7 @@ namespace Qwen { float theta = 10000; std::vector axes_dim = {16, 56, 56}; int64_t axes_dim_sum = 128; + bool zero_cond_t = false; }; class QwenImageModel : public GGMLBlock { @@ -346,7 +382,8 @@ namespace Qwen { auto block = std::shared_ptr(new QwenImageTransformerBlock(inner_dim, params.num_attention_heads, params.attention_head_dim, - 1e-6f)); + 1e-6f, + params.zero_cond_t)); blocks["transformer_blocks." + std::to_string(i)] = block; } @@ -354,14 +391,14 @@ namespace Qwen { blocks["proj_out"] = std::shared_ptr(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels)); } - struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx, + struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, struct ggml_tensor* x) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; int pad_h = (params.patch_size - H % params.patch_size) % params.patch_size; int pad_w = (params.patch_size - W % params.patch_size) % params.patch_size; - x = ggml_pad(ctx, x, pad_w, pad_h, 0, 0); // [N, C, H + pad_h, W + pad_w] + x = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return x; } @@ -387,10 +424,10 @@ namespace Qwen { return x; } - struct ggml_tensor* process_img(struct ggml_context* ctx, + struct ggml_tensor* process_img(GGMLRunnerContext* ctx, struct ggml_tensor* x) { x = pad_to_patch_size(ctx, x); - x = patchify(ctx, x); + x = patchify(ctx->ggml_ctx, x); return x; } @@ -421,7 +458,8 @@ namespace Qwen { struct ggml_tensor* x, struct ggml_tensor* timestep, struct ggml_tensor* context, - struct ggml_tensor* pe) { + struct ggml_tensor* pe, + struct ggml_tensor* modulate_index = nullptr) { auto time_text_embed = std::dynamic_pointer_cast(blocks["time_text_embed"]); auto txt_norm = std::dynamic_pointer_cast(blocks["txt_norm"]); auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); @@ -430,18 +468,26 @@ namespace Qwen { auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); auto t_emb = time_text_embed->forward(ctx, timestep); - auto img = img_in->forward(ctx, x); - auto txt = txt_norm->forward(ctx, context); - txt = txt_in->forward(ctx, txt); + if (params.zero_cond_t) { + auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros(ctx->ggml_ctx, timestep->ne[0], timestep->ne[1], timestep->ne[2], timestep->ne[3])); + t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1); + } + auto img = img_in->forward(ctx, x); + auto txt = txt_norm->forward(ctx, context); + txt = txt_in->forward(ctx, txt); for (int i = 0; i < params.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); - auto result = block->forward(ctx, img, txt, t_emb, pe); + auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index); img = result.first; txt = result.second; } + if (params.zero_cond_t) { + t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0]; + } + img = norm_out->forward(ctx, img, t_emb); img = proj_out->forward(ctx, img); @@ -453,7 +499,8 @@ namespace Qwen { struct ggml_tensor* timestep, struct ggml_tensor* context, struct ggml_tensor* pe, - std::vector ref_latents = {}) { + std::vector ref_latents = {}, + struct ggml_tensor* modulate_index = nullptr) { // Forward pass of DiT. // x: [N, C, H, W] // timestep: [N,] @@ -466,12 +513,12 @@ namespace Qwen { int64_t C = x->ne[2]; int64_t N = x->ne[3]; - auto img = process_img(ctx->ggml_ctx, x); + auto img = process_img(ctx, x); uint64_t img_tokens = img->ne[1]; if (ref_latents.size() > 0) { for (ggml_tensor* ref : ref_latents) { - ref = process_img(ctx->ggml_ctx, ref); + ref = process_img(ctx, ref); img = ggml_concat(ctx->ggml_ctx, img, ref, 1); } } @@ -479,7 +526,7 @@ namespace Qwen { int64_t h_len = ((H + (params.patch_size / 2)) / params.patch_size); int64_t w_len = ((W + (params.patch_size / 2)) / params.patch_size); - auto out = forward_orig(ctx, img, timestep, context, pe); // [N, h_len*w_len, ph*pw*C] + auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index); // [N, h_len*w_len, ph*pw*C] if (out->ne[1] > img_tokens) { out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3)); // [num_tokens, N, C * patch_size * patch_size] @@ -502,19 +549,25 @@ namespace Qwen { QwenImageParams qwen_image_params; QwenImageModel qwen_image; std::vector pe_vec; + std::vector modulate_index_vec; SDVersion version; QwenImageRunner(ggml_backend_t backend, bool offload_params_to_cpu, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "", - SDVersion version = VERSION_QWEN_IMAGE) + SDVersion version = VERSION_QWEN_IMAGE, + bool zero_cond_t = false) : GGMLRunner(backend, offload_params_to_cpu) { - qwen_image_params.num_layers = 0; + qwen_image_params.num_layers = 0; + qwen_image_params.zero_cond_t = zero_cond_t; for (auto pair : tensor_storage_map) { std::string tensor_name = pair.first; if (tensor_name.find(prefix) == std::string::npos) continue; + if (tensor_name.find("__index_timestep_zero__") != std::string::npos) { + qwen_image_params.zero_cond_t = true; + } size_t pos = tensor_name.find("transformer_blocks."); if (pos != std::string::npos) { tensor_name = tensor_name.substr(pos); // remove prefix @@ -529,6 +582,9 @@ namespace Qwen { } } LOG_INFO("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers); + if (qwen_image_params.zero_cond_t) { + LOG_INFO("use zero_cond_t"); + } qwen_image = QwenImageModel(qwen_image_params); qwen_image.init(params_ctx, tensor_storage_map, prefix); } @@ -565,6 +621,8 @@ namespace Qwen { ref_latents, increase_ref_index, qwen_image_params.theta, + circular_y_enabled, + circular_x_enabled, qwen_image_params.axes_dim); int pos_len = pe_vec.size() / qwen_image_params.axes_dim_sum / 2; // LOG_DEBUG("pos_len %d", pos_len); @@ -574,6 +632,31 @@ namespace Qwen { // pe->data = nullptr; set_backend_tensor_data(pe, pe_vec.data()); + ggml_tensor* modulate_index = nullptr; + if (qwen_image_params.zero_cond_t) { + modulate_index_vec.clear(); + + int64_t h_len = ((x->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); + int64_t w_len = ((x->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); + int64_t num_img_tokens = h_len * w_len; + + modulate_index_vec.insert(modulate_index_vec.end(), num_img_tokens, 0.f); + int64_t num_ref_img_tokens = 0; + for (ggml_tensor* ref : ref_latents) { + int64_t h_len = ((ref->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); + int64_t w_len = ((ref->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); + + num_ref_img_tokens += h_len * w_len; + } + + if (num_ref_img_tokens > 0) { + modulate_index_vec.insert(modulate_index_vec.end(), num_ref_img_tokens, 1.f); + } + + modulate_index = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, modulate_index_vec.size()); + set_backend_tensor_data(modulate_index, modulate_index_vec.data()); + } + auto runner_ctx = get_context(); struct ggml_tensor* out = qwen_image.forward(&runner_ctx, @@ -581,7 +664,8 @@ namespace Qwen { timesteps, context, pe, - ref_latents); + ref_latents, + modulate_index); ggml_build_forward_expand(gf, out); @@ -684,4 +768,4 @@ namespace Qwen { } // namespace name -#endif // __QWEN_IMAGE_HPP__ \ No newline at end of file +#endif // __QWEN_IMAGE_HPP__ diff --git a/otherarch/sdcpp/rope.hpp b/otherarch/sdcpp/rope.hpp index 4abc51469..4e6136c11 100644 --- a/otherarch/sdcpp/rope.hpp +++ b/otherarch/sdcpp/rope.hpp @@ -1,6 +1,8 @@ #ifndef __ROPE_HPP__ #define __ROPE_HPP__ +#include +#include #include #include "ggml_extend.hpp" @@ -39,7 +41,10 @@ namespace Rope { return flat_vec; } - __STATIC_INLINE__ std::vector> rope(const std::vector& pos, int dim, int theta) { + __STATIC_INLINE__ std::vector> rope(const std::vector& pos, + int dim, + int theta, + const std::vector& axis_wrap_dims = {}) { assert(dim % 2 == 0); int half_dim = dim / 2; @@ -47,14 +52,31 @@ namespace Rope { std::vector omega(half_dim); for (int i = 0; i < half_dim; ++i) { - omega[i] = 1.0 / std::pow(theta, scale[i]); + omega[i] = 1.0f / std::pow(theta, scale[i]); } int pos_size = pos.size(); std::vector> out(pos_size, std::vector(half_dim)); for (int i = 0; i < pos_size; ++i) { for (int j = 0; j < half_dim; ++j) { - out[i][j] = pos[i] * omega[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; } } @@ -89,16 +111,25 @@ namespace Rope { int patch_size, int bs, int axes_dim_num, - int index = 0, - int h_offset = 0, - int w_offset = 0) { + int index = 0, + int h_offset = 0, + int w_offset = 0, + bool scale_rope = false) { int h_len = (h + (patch_size / 2)) / patch_size; int w_len = (w + (patch_size / 2)) / patch_size; std::vector> img_ids(h_len * w_len, std::vector(axes_dim_num, 0.0)); - std::vector row_ids = linspace(h_offset, h_len - 1 + h_offset, h_len); - std::vector col_ids = linspace(w_offset, w_len - 1 + w_offset, w_len); + int h_start = h_offset; + int w_start = w_offset; + + if (scale_rope) { + h_start -= h_len / 2; + w_start -= w_len / 2; + } + + std::vector row_ids = linspace(h_start, h_start + h_len - 1, h_len); + std::vector col_ids = linspace(w_start, w_start + w_len - 1, w_len); for (int i = 0; i < h_len; ++i) { for (int j = 0; j < w_len; ++j) { @@ -137,7 +168,8 @@ namespace Rope { __STATIC_INLINE__ std::vector embed_nd(const std::vector>& ids, int bs, int theta, - const std::vector& axes_dim) { + const std::vector& axes_dim, + const std::vector>& wrap_dims = {}) { std::vector> trans_ids = transpose(ids); size_t pos_len = ids.size() / bs; int num_axes = axes_dim.size(); @@ -152,7 +184,12 @@ namespace Rope { std::vector> emb(bs * pos_len, std::vector(emb_dim * 2 * 2, 0.0)); int offset = 0; for (int i = 0; i < num_axes; ++i) { - std::vector> rope_emb = rope(trans_ids[i], axes_dim[i], theta); // [bs*pos_len, axes_dim[i]/2 * 2 * 2] + std::vector axis_wrap_dims; + if (!wrap_dims.empty() && i < (int)wrap_dims.size()) { + axis_wrap_dims = wrap_dims[i]; + } + std::vector> rope_emb = + rope(trans_ids[i], axes_dim[i], theta, axis_wrap_dims); // [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) { @@ -171,7 +208,8 @@ namespace Rope { int axes_dim_num, const std::vector& ref_latents, bool increase_ref_index, - float ref_index_scale) { + float ref_index_scale, + bool scale_rope) { std::vector> ids; uint64_t curr_h_offset = 0; uint64_t curr_w_offset = 0; @@ -185,6 +223,7 @@ namespace Rope { } else { h_offset = curr_h_offset; } + scale_rope = false; } auto ref_ids = gen_flux_img_ids(ref->ne[1], @@ -194,7 +233,8 @@ namespace Rope { axes_dim_num, static_cast(index * ref_index_scale), h_offset, - w_offset); + w_offset, + scale_rope); ids = concat_ids(ids, ref_ids, bs); if (increase_ref_index) { @@ -222,7 +262,7 @@ namespace Rope { 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, ref_latents, increase_ref_index, ref_index_scale); + auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, ref_index_scale, false); ids = concat_ids(ids, refs_ids, bs); } return ids; @@ -239,6 +279,8 @@ namespace Rope { bool increase_ref_index, float ref_index_scale, int theta, + bool circular_h, + bool circular_w, const std::vector& axes_dim) { std::vector> ids = gen_flux_ids(h, w, @@ -250,7 +292,47 @@ namespace Rope { ref_latents, increase_ref_index, ref_index_scale); - return embed_nd(ids, bs, theta, axes_dim); + std::vector> 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(pos_len, 0)); + size_t cursor = context_len; // text first + const size_t img_tokens = static_cast(h_len) * static_cast(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(ref->ne[1]); + int ref_w = static_cast(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(ref_h_l) * static_cast(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, theta, axes_dim, wrap_dims); } __STATIC_INLINE__ std::vector> gen_qwen_image_ids(int h, @@ -271,10 +353,10 @@ namespace Rope { } } int axes_dim_num = 3; - auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num); + auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true); auto ids = concat_ids(txt_ids_repeated, img_ids, bs); if (ref_latents.size() > 0) { - auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, 1.f); + auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, 1.f, true); ids = concat_ids(ids, refs_ids, bs); } return ids; @@ -289,9 +371,57 @@ namespace Rope { const std::vector& ref_latents, bool increase_ref_index, int theta, + bool circular_h, + bool circular_w, const std::vector& axes_dim) { std::vector> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index); - return embed_nd(ids, bs, theta, axes_dim); + std::vector> 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(total_tokens / bs, 0)); + size_t cursor = context_len; // ignore text tokens + const size_t img_tokens = static_cast(h_len) * static_cast(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(ref->ne[1]); + int ref_w = static_cast(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(ref_h_len) * static_cast(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, theta, axes_dim, wrap_dims); } __STATIC_INLINE__ std::vector> gen_vid_ids(int t, @@ -428,9 +558,33 @@ namespace Rope { const std::vector& ref_latents, bool increase_ref_index, int theta, + bool circular_h, + bool circular_w, const std::vector& axes_dim) { std::vector> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index); - return embed_nd(ids, bs, theta, axes_dim); + std::vector> 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(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(h_len) * static_cast(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; + } + } + } + } + + return embed_nd(ids, bs, theta, axes_dim, wrap_dims); } __STATIC_INLINE__ struct ggml_tensor* apply_rope(struct ggml_context* ctx, diff --git a/otherarch/sdcpp/stable-diffusion.cpp b/otherarch/sdcpp/stable-diffusion.cpp index 9ff498bf5..35cb40596 100644 --- a/otherarch/sdcpp/stable-diffusion.cpp +++ b/otherarch/sdcpp/stable-diffusion.cpp @@ -7,6 +7,7 @@ #include "stable-diffusion.h" #include "util.h" +#include "cache_dit.hpp" #include "conditioner.hpp" #include "control.hpp" #include "denoiser.hpp" @@ -16,6 +17,7 @@ #include "lora.hpp" #include "pmid.hpp" #include "tae.hpp" +#include "ucache.hpp" #include "vae.hpp" #include "latent-preview.h" @@ -178,7 +180,27 @@ public: #endif #ifdef SD_USE_VULKAN LOG_DEBUG("Using Vulkan backend"); - for (int device = 0; device < ggml_backend_vk_get_device_count(); ++device) { + size_t device = 0; + const int device_count = ggml_backend_vk_get_device_count(); + if (device_count) { + const char* SD_VK_DEVICE = getenv("SD_VK_DEVICE"); + if (SD_VK_DEVICE != nullptr) { + std::string sd_vk_device_str = SD_VK_DEVICE; + try { + device = std::stoull(sd_vk_device_str); + } catch (const std::invalid_argument&) { + LOG_WARN("SD_VK_DEVICE environment variable is not a valid integer (%s). Falling back to device 0.", SD_VK_DEVICE); + device = 0; + } catch (const std::out_of_range&) { + LOG_WARN("SD_VK_DEVICE environment variable value is out of range for `unsigned long long` type (%s). Falling back to device 0.", SD_VK_DEVICE); + device = 0; + } + if (device >= device_count) { + LOG_WARN("Cannot find targeted vulkan device (%llu). Falling back to device 0.", device); + device = 0; + } + } + LOG_INFO("Vulkan: Using device %llu", device); backend = ggml_backend_vk_init(device); } if (!backend) { @@ -520,6 +542,10 @@ public: vae_decode_only = false; } + if (sd_ctx_params->circular_x || sd_ctx_params->circular_y) { + LOG_INFO("Using circular padding for convolutions"); + } + bool clip_on_cpu = sd_ctx_params->keep_clip_on_cpu; { @@ -629,7 +655,8 @@ public: offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", - version); + version, + sd_ctx_params->qwen_image_zero_cond_t); } else if (sd_version_is_z_image(version)) { cond_stage_model = std::make_shared(clip_backend, offload_params_to_cpu, @@ -699,8 +726,8 @@ public: vae_backend = backend; } - if (sd_version_is_wan(version) || sd_version_is_qwen_image(version)) { - if (!use_tiny_autoencoder) { + if (!use_tiny_autoencoder || sd_ctx_params->tae_preview_only) { + if (sd_version_is_wan(version) || sd_version_is_qwen_image(version)) { first_stage_model = std::make_shared(vae_backend, offload_params_to_cpu, tensor_storage_map, @@ -709,57 +736,56 @@ public: version); first_stage_model->alloc_params_buffer(); first_stage_model->get_param_tensors(tensors, "first_stage_model"); + } else if (version == VERSION_CHROMA_RADIANCE) { + first_stage_model = std::make_shared(vae_backend, + offload_params_to_cpu); } else { + first_stage_model = std::make_shared(vae_backend, + offload_params_to_cpu, + tensor_storage_map, + "first_stage_model", + vae_decode_only, + false, + version); + if (sd_ctx_params->vae_conv_direct) { + LOG_INFO("Using Conv2d direct in the vae model"); + first_stage_model->set_conv2d_direct_enabled(true); + } + if (version == VERSION_SDXL && + (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale)) { + float vae_conv_2d_scale = 1.f / 32.f; + LOG_WARN( + "No VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, " + "using Conv2D scale %.3f", + vae_conv_2d_scale); + first_stage_model->set_conv2d_scale(vae_conv_2d_scale); + } + first_stage_model->alloc_params_buffer(); + first_stage_model->get_param_tensors(tensors, "first_stage_model"); + } + } + + if (use_tiny_autoencoder) { + if (sd_version_is_wan(version) || sd_version_is_qwen_image(version)) { tae_first_stage = std::make_shared(vae_backend, offload_params_to_cpu, tensor_storage_map, "decoder", vae_decode_only, version); - if (sd_ctx_params->vae_conv_direct) { - LOG_INFO("Using Conv2d direct in the tae model"); - tae_first_stage->set_conv2d_direct_enabled(true); - } + } else { + tae_first_stage = std::make_shared(vae_backend, + offload_params_to_cpu, + tensor_storage_map, + "decoder.layers", + vae_decode_only, + version); } - } else if (version == VERSION_CHROMA_RADIANCE) { - first_stage_model = std::make_shared(vae_backend, - offload_params_to_cpu); - } else if (!use_tiny_autoencoder || sd_ctx_params->tae_preview_only) { - first_stage_model = std::make_shared(vae_backend, - offload_params_to_cpu, - tensor_storage_map, - "first_stage_model", - vae_decode_only, - false, - version); - if (sd_ctx_params->vae_conv_direct) { - LOG_INFO("Using Conv2d direct in the vae model"); - first_stage_model->set_conv2d_direct_enabled(true); - } - if (version == VERSION_SDXL && - (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale)) { - float vae_conv_2d_scale = 1.f / 32.f; - LOG_WARN( - "No VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, " - "using Conv2D scale %.3f", - vae_conv_2d_scale); - first_stage_model->set_conv2d_scale(vae_conv_2d_scale); - } - first_stage_model->alloc_params_buffer(); - first_stage_model->get_param_tensors(tensors, "first_stage_model"); - } else if (use_tiny_autoencoder) { - tae_first_stage = std::make_shared(vae_backend, - offload_params_to_cpu, - tensor_storage_map, - "decoder.layers", - vae_decode_only, - version); if (sd_ctx_params->vae_conv_direct) { LOG_INFO("Using Conv2d direct in the tae model"); tae_first_stage->set_conv2d_direct_enabled(true); } } - // first_stage_model->get_param_tensors(tensors, "first_stage_model."); if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) { ggml_backend_t controlnet_backend = nullptr; @@ -824,6 +850,20 @@ public: } pmid_model->get_param_tensors(tensors, "pmid"); } + + diffusion_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); + if (high_noise_diffusion_model) { + high_noise_diffusion_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); + } + if (control_net) { + control_net->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); + } + if (first_stage_model) { + first_stage_model->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); + } + if (tae_first_stage) { + tae_first_stage->set_circular_axes(sd_ctx_params->circular_x, sd_ctx_params->circular_y); + } } struct ggml_init_params params; @@ -849,6 +889,7 @@ public: } ignore_tensors.insert("model.diffusion_model.__x0__"); ignore_tensors.insert("model.diffusion_model.__32x32__"); + ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__"); if (vae_decode_only) { ignore_tensors.insert("first_stage_model.encoder"); @@ -1626,19 +1667,19 @@ public: const std::vector& sigmas, int start_merge_step, SDCondition id_cond, - std::vector ref_latents = {}, - bool increase_ref_index = false, - ggml_tensor* denoise_mask = nullptr, - ggml_tensor* vace_context = nullptr, - float vace_strength = 1.f, - const sd_easycache_params_t* easycache_params = nullptr) { + std::vector ref_latents = {}, + bool increase_ref_index = false, + ggml_tensor* denoise_mask = nullptr, + ggml_tensor* vace_context = nullptr, + float vace_strength = 1.f, + const sd_cache_params_t* cache_params = nullptr) { if (shifted_timestep > 0 && !sd_version_is_sdxl(version)) { LOG_WARN("timestep shifting is only supported for SDXL models!"); shifted_timestep = 0; } std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); - float cfg_scale = guidance.txt_cfg; + float cfg_scale = guidance.txt_cfg; if (cfg_scale < 1.f) { if (cfg_scale == 0.f) { // Diffusers follow the convention from the original paper @@ -1659,31 +1700,40 @@ public: } EasyCacheState easycache_state; + UCacheState ucache_state; + CacheDitConditionState cachedit_state; bool easycache_enabled = false; - if (easycache_params != nullptr && easycache_params->enabled) { - bool easycache_supported = sd_version_is_dit(version); - if (!easycache_supported) { - LOG_WARN("EasyCache requested but not supported for this model type"); - } else { - EasyCacheConfig easycache_config; - easycache_config.enabled = true; - easycache_config.reuse_threshold = std::max(0.0f, easycache_params->reuse_threshold); - easycache_config.start_percent = easycache_params->start_percent; - easycache_config.end_percent = easycache_params->end_percent; - bool percent_valid = easycache_config.start_percent >= 0.0f && - easycache_config.start_percent < 1.0f && - easycache_config.end_percent > 0.0f && - easycache_config.end_percent <= 1.0f && - easycache_config.start_percent < easycache_config.end_percent; - if (!percent_valid) { - LOG_WARN("EasyCache disabled due to invalid percent range (start=%.3f, end=%.3f)", - easycache_config.start_percent, - easycache_config.end_percent); + bool ucache_enabled = false; + bool cachedit_enabled = false; + + if (cache_params != nullptr && cache_params->mode != SD_CACHE_DISABLED) { + bool percent_valid = true; + if (cache_params->mode == SD_CACHE_EASYCACHE || cache_params->mode == SD_CACHE_UCACHE) { + percent_valid = cache_params->start_percent >= 0.0f && + cache_params->start_percent < 1.0f && + cache_params->end_percent > 0.0f && + cache_params->end_percent <= 1.0f && + cache_params->start_percent < cache_params->end_percent; + } + + if (!percent_valid) { + LOG_WARN("Cache disabled due to invalid percent range (start=%.3f, end=%.3f)", + cache_params->start_percent, + cache_params->end_percent); + } else if (cache_params->mode == SD_CACHE_EASYCACHE) { + bool easycache_supported = sd_version_is_dit(version); + if (!easycache_supported) { + LOG_WARN("EasyCache requested but not supported for this model type"); } else { + EasyCacheConfig easycache_config; + easycache_config.enabled = true; + easycache_config.reuse_threshold = std::max(0.0f, cache_params->reuse_threshold); + easycache_config.start_percent = cache_params->start_percent; + easycache_config.end_percent = cache_params->end_percent; easycache_state.init(easycache_config, denoiser.get()); if (easycache_state.enabled()) { easycache_enabled = true; - LOG_INFO("EasyCache enabled - threshold: %.3f, start_percent: %.2f, end_percent: %.2f", + LOG_INFO("EasyCache enabled - threshold: %.3f, start: %.2f, end: %.2f", easycache_config.reuse_threshold, easycache_config.start_percent, easycache_config.end_percent); @@ -1691,9 +1741,84 @@ public: LOG_WARN("EasyCache requested but could not be initialized for this run"); } } + } else if (cache_params->mode == SD_CACHE_UCACHE) { + bool ucache_supported = sd_version_is_unet(version); + if (!ucache_supported) { + LOG_WARN("UCache requested but not supported for this model type (only UNET models)"); + } else { + UCacheConfig ucache_config; + ucache_config.enabled = true; + ucache_config.reuse_threshold = std::max(0.0f, cache_params->reuse_threshold); + ucache_config.start_percent = cache_params->start_percent; + ucache_config.end_percent = cache_params->end_percent; + ucache_config.error_decay_rate = std::max(0.0f, std::min(1.0f, cache_params->error_decay_rate)); + ucache_config.use_relative_threshold = cache_params->use_relative_threshold; + ucache_config.reset_error_on_compute = cache_params->reset_error_on_compute; + ucache_state.init(ucache_config, denoiser.get()); + if (ucache_state.enabled()) { + ucache_enabled = true; + LOG_INFO("UCache enabled - threshold: %.3f, start: %.2f, end: %.2f, decay: %.2f, relative: %s, reset: %s", + ucache_config.reuse_threshold, + ucache_config.start_percent, + ucache_config.end_percent, + ucache_config.error_decay_rate, + ucache_config.use_relative_threshold ? "true" : "false", + ucache_config.reset_error_on_compute ? "true" : "false"); + } else { + LOG_WARN("UCache requested but could not be initialized for this run"); + } + } + } else if (cache_params->mode == SD_CACHE_DBCACHE || + cache_params->mode == SD_CACHE_TAYLORSEER || + cache_params->mode == SD_CACHE_CACHE_DIT) { + bool cachedit_supported = sd_version_is_dit(version); + if (!cachedit_supported) { + LOG_WARN("CacheDIT requested but not supported for this model type (only DiT models)"); + } else { + DBCacheConfig dbcfg; + dbcfg.enabled = (cache_params->mode == SD_CACHE_DBCACHE || + cache_params->mode == SD_CACHE_CACHE_DIT); + dbcfg.Fn_compute_blocks = cache_params->Fn_compute_blocks; + dbcfg.Bn_compute_blocks = cache_params->Bn_compute_blocks; + dbcfg.residual_diff_threshold = cache_params->residual_diff_threshold; + dbcfg.max_warmup_steps = cache_params->max_warmup_steps; + dbcfg.max_cached_steps = cache_params->max_cached_steps; + dbcfg.max_continuous_cached_steps = cache_params->max_continuous_cached_steps; + if (cache_params->scm_mask != nullptr && strlen(cache_params->scm_mask) > 0) { + dbcfg.steps_computation_mask = parse_scm_mask(cache_params->scm_mask); + } + dbcfg.scm_policy_dynamic = cache_params->scm_policy_dynamic; + + TaylorSeerConfig tcfg; + tcfg.enabled = (cache_params->mode == SD_CACHE_TAYLORSEER || + cache_params->mode == SD_CACHE_CACHE_DIT); + tcfg.n_derivatives = cache_params->taylorseer_n_derivatives; + tcfg.skip_interval_steps = cache_params->taylorseer_skip_interval; + + cachedit_state.init(dbcfg, tcfg); + if (cachedit_state.enabled()) { + cachedit_enabled = true; + LOG_INFO("CacheDIT enabled - mode: %s, Fn: %d, Bn: %d, threshold: %.3f, warmup: %d", + cache_params->mode == SD_CACHE_CACHE_DIT ? "DBCache+TaylorSeer" : (cache_params->mode == SD_CACHE_DBCACHE ? "DBCache" : "TaylorSeer"), + dbcfg.Fn_compute_blocks, + dbcfg.Bn_compute_blocks, + dbcfg.residual_diff_threshold, + dbcfg.max_warmup_steps); + } else { + LOG_WARN("CacheDIT requested but could not be initialized for this run"); + } + } } } + if (ucache_enabled) { + ucache_state.set_sigmas(sigmas); + } + + if (cachedit_enabled) { + cachedit_state.set_sigmas(sigmas); + } + size_t steps = sigmas.size() - 1; struct ggml_tensor* x = ggml_dup_tensor(work_ctx, init_latent); copy_ggml_tensor(x, init_latent); @@ -1797,6 +1922,91 @@ public: return easycache_step_active && easycache_state.is_step_skipped(); }; + const bool ucache_step_active = ucache_enabled && step > 0; + int ucache_step_index = ucache_step_active ? (step - 1) : -1; + if (ucache_step_active) { + ucache_state.begin_step(ucache_step_index, sigma); + } + + auto ucache_before_condition = [&](const SDCondition* condition, struct ggml_tensor* output_tensor) -> bool { + if (!ucache_step_active || condition == nullptr || output_tensor == nullptr) { + return false; + } + return ucache_state.before_condition(condition, + diffusion_params.x, + output_tensor, + sigma, + ucache_step_index); + }; + + auto ucache_after_condition = [&](const SDCondition* condition, struct ggml_tensor* output_tensor) { + if (!ucache_step_active || condition == nullptr || output_tensor == nullptr) { + return; + } + ucache_state.after_condition(condition, + diffusion_params.x, + output_tensor); + }; + + auto ucache_step_is_skipped = [&]() { + return ucache_step_active && ucache_state.is_step_skipped(); + }; + + const bool cachedit_step_active = cachedit_enabled && step > 0; + int cachedit_step_index = cachedit_step_active ? (step - 1) : -1; + if (cachedit_step_active) { + cachedit_state.begin_step(cachedit_step_index, sigma); + } + + auto cachedit_before_condition = [&](const SDCondition* condition, struct ggml_tensor* output_tensor) -> bool { + if (!cachedit_step_active || condition == nullptr || output_tensor == nullptr) { + return false; + } + return cachedit_state.before_condition(condition, + diffusion_params.x, + output_tensor, + sigma, + cachedit_step_index); + }; + + auto cachedit_after_condition = [&](const SDCondition* condition, struct ggml_tensor* output_tensor) { + if (!cachedit_step_active || condition == nullptr || output_tensor == nullptr) { + return; + } + cachedit_state.after_condition(condition, + diffusion_params.x, + output_tensor); + }; + + auto cachedit_step_is_skipped = [&]() { + return cachedit_step_active && cachedit_state.is_step_skipped(); + }; + + auto cache_before_condition = [&](const SDCondition* condition, struct ggml_tensor* output_tensor) -> bool { + if (easycache_step_active) { + return easycache_before_condition(condition, output_tensor); + } else if (ucache_step_active) { + return ucache_before_condition(condition, output_tensor); + } else if (cachedit_step_active) { + return cachedit_before_condition(condition, output_tensor); + } + return false; + }; + + auto cache_after_condition = [&](const SDCondition* condition, struct ggml_tensor* output_tensor) { + if (easycache_step_active) { + easycache_after_condition(condition, output_tensor); + } else if (ucache_step_active) { + ucache_after_condition(condition, output_tensor); + } else if (cachedit_step_active) { + cachedit_after_condition(condition, output_tensor); + } + }; + + auto cache_step_is_skipped = [&]() { + return easycache_step_is_skipped() || ucache_step_is_skipped() || cachedit_step_is_skipped(); + }; + std::vector scaling = denoiser->get_scalings(sigma); GGML_ASSERT(scaling.size() == 3); float c_skip = scaling[0]; @@ -1872,7 +2082,7 @@ public: active_condition = &id_cond; } - bool skip_model = easycache_before_condition(active_condition, *active_output); + bool skip_model = cache_before_condition(active_condition, *active_output); if (!skip_model) { if (!work_diffusion_model->compute(n_threads, diffusion_params, @@ -1880,10 +2090,10 @@ public: LOG_ERROR("diffusion model compute failed"); return nullptr; } - easycache_after_condition(active_condition, *active_output); + cache_after_condition(active_condition, *active_output); } - bool current_step_skipped = easycache_step_is_skipped(); + bool current_step_skipped = cache_step_is_skipped(); float* negative_data = nullptr; if (has_unconditioned) { @@ -1895,12 +2105,12 @@ public: LOG_ERROR("controlnet compute failed"); } } - current_step_skipped = easycache_step_is_skipped(); + current_step_skipped = cache_step_is_skipped(); diffusion_params.controls = controls; diffusion_params.context = uncond.c_crossattn; diffusion_params.c_concat = uncond.c_concat; diffusion_params.y = uncond.c_vector; - bool skip_uncond = easycache_before_condition(&uncond, out_uncond); + bool skip_uncond = cache_before_condition(&uncond, out_uncond); if (!skip_uncond) { if (!work_diffusion_model->compute(n_threads, diffusion_params, @@ -1908,7 +2118,7 @@ public: LOG_ERROR("diffusion model compute failed"); return nullptr; } - easycache_after_condition(&uncond, out_uncond); + cache_after_condition(&uncond, out_uncond); } negative_data = (float*)out_uncond->data; } @@ -1918,7 +2128,7 @@ public: diffusion_params.context = img_cond.c_crossattn; diffusion_params.c_concat = img_cond.c_concat; diffusion_params.y = img_cond.c_vector; - bool skip_img_cond = easycache_before_condition(&img_cond, out_img_cond); + bool skip_img_cond = cache_before_condition(&img_cond, out_img_cond); if (!skip_img_cond) { if (!work_diffusion_model->compute(n_threads, diffusion_params, @@ -1926,7 +2136,7 @@ public: LOG_ERROR("diffusion model compute failed"); return nullptr; } - easycache_after_condition(&img_cond, out_img_cond); + cache_after_condition(&img_cond, out_img_cond); } img_cond_data = (float*)out_img_cond->data; } @@ -1936,7 +2146,7 @@ public: float* skip_layer_data = has_skiplayer ? (float*)out_skip->data : nullptr; if (is_skiplayer_step) { LOG_DEBUG("Skipping layers at step %d\n", step); - if (!easycache_step_is_skipped()) { + if (!cache_step_is_skipped()) { // skip layer (same as conditioned) diffusion_params.context = cond.c_crossattn; diffusion_params.c_concat = cond.c_concat; @@ -2040,6 +2250,48 @@ public: } } + if (ucache_enabled) { + size_t total_steps = sigmas.size() > 0 ? sigmas.size() - 1 : 0; + if (ucache_state.total_steps_skipped > 0 && total_steps > 0) { + if (ucache_state.total_steps_skipped < static_cast(total_steps)) { + double speedup = static_cast(total_steps) / + static_cast(total_steps - ucache_state.total_steps_skipped); + LOG_INFO("UCache skipped %d/%zu steps (%.2fx estimated speedup)", + ucache_state.total_steps_skipped, + total_steps, + speedup); + } else { + LOG_INFO("UCache skipped %d/%zu steps", + ucache_state.total_steps_skipped, + total_steps); + } + } else if (total_steps > 0) { + LOG_INFO("UCache completed without skipping steps"); + } + } + + if (cachedit_enabled) { + size_t total_steps = sigmas.size() > 0 ? sigmas.size() - 1 : 0; + if (cachedit_state.total_steps_skipped > 0 && total_steps > 0) { + if (cachedit_state.total_steps_skipped < static_cast(total_steps)) { + double speedup = static_cast(total_steps) / + static_cast(total_steps - cachedit_state.total_steps_skipped); + LOG_INFO("CacheDIT skipped %d/%zu steps (%.2fx estimated speedup), accum_diff: %.4f", + cachedit_state.total_steps_skipped, + total_steps, + speedup, + cachedit_state.accumulated_residual_diff); + } else { + LOG_INFO("CacheDIT skipped %d/%zu steps, accum_diff: %.4f", + cachedit_state.total_steps_skipped, + total_steps, + cachedit_state.accumulated_residual_diff); + } + } else if (total_steps > 0) { + LOG_INFO("CacheDIT completed without skipping steps"); + } + } + if (inverse_noise_scaling) { x = denoiser->inverse_noise_scaling(sigmas[sigmas.size() - 1], x); } @@ -2655,12 +2907,25 @@ enum lora_apply_mode_t str_to_lora_apply_mode(const char* str) { return LORA_APPLY_MODE_COUNT; } -void sd_easycache_params_init(sd_easycache_params_t* easycache_params) { - *easycache_params = {}; - easycache_params->enabled = false; - easycache_params->reuse_threshold = 0.2f; - easycache_params->start_percent = 0.15f; - easycache_params->end_percent = 0.95f; +void sd_cache_params_init(sd_cache_params_t* cache_params) { + *cache_params = {}; + cache_params->mode = SD_CACHE_DISABLED; + cache_params->reuse_threshold = 1.0f; + cache_params->start_percent = 0.15f; + cache_params->end_percent = 0.95f; + cache_params->error_decay_rate = 1.0f; + cache_params->use_relative_threshold = true; + cache_params->reset_error_on_compute = true; + cache_params->Fn_compute_blocks = 8; + cache_params->Bn_compute_blocks = 0; + cache_params->residual_diff_threshold = 0.08f; + cache_params->max_warmup_steps = 8; + cache_params->max_cached_steps = -1; + cache_params->max_continuous_cached_steps = -1; + cache_params->taylorseer_n_derivatives = 1; + cache_params->taylorseer_skip_interval = 1; + cache_params->scm_mask = nullptr; + cache_params->scm_policy_dynamic = true; } void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { @@ -2678,6 +2943,8 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { sd_ctx_params->keep_control_net_on_cpu = false; sd_ctx_params->keep_vae_on_cpu = false; sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->circular_x = false; + sd_ctx_params->circular_y = false; sd_ctx_params->chroma_use_dit_mask = true; sd_ctx_params->chroma_use_t5_mask = false; sd_ctx_params->chroma_t5_mask_pad = 1; @@ -2717,6 +2984,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "keep_control_net_on_cpu: %s\n" "keep_vae_on_cpu: %s\n" "diffusion_flash_attn: %s\n" + "circular_x: %s\n" + "circular_y: %s\n" "chroma_use_dit_mask: %s\n" "chroma_use_t5_mask: %s\n" "chroma_t5_mask_pad: %d\n", @@ -2746,6 +3015,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { BOOL_STR(sd_ctx_params->keep_control_net_on_cpu), BOOL_STR(sd_ctx_params->keep_vae_on_cpu), BOOL_STR(sd_ctx_params->diffusion_flash_attn), + BOOL_STR(sd_ctx_params->circular_x), + BOOL_STR(sd_ctx_params->circular_y), BOOL_STR(sd_ctx_params->chroma_use_dit_mask), BOOL_STR(sd_ctx_params->chroma_use_t5_mask), sd_ctx_params->chroma_t5_mask_pad); @@ -2819,7 +3090,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { sd_img_gen_params->control_strength = 0.9f; sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f}; sd_img_gen_params->vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f}; - sd_easycache_params_init(&sd_img_gen_params->easycache); + sd_cache_params_init(&sd_img_gen_params->cache); } char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { @@ -2863,12 +3134,18 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { sd_img_gen_params->pm_params.id_images_count, SAFE_STR(sd_img_gen_params->pm_params.id_embed_path), BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled)); + const char* cache_mode_str = "disabled"; + if (sd_img_gen_params->cache.mode == SD_CACHE_EASYCACHE) { + cache_mode_str = "easycache"; + } else if (sd_img_gen_params->cache.mode == SD_CACHE_UCACHE) { + cache_mode_str = "ucache"; + } snprintf(buf + strlen(buf), 4096 - strlen(buf), - "easycache: %s (threshold=%.3f, start=%.2f, end=%.2f)\n", - sd_img_gen_params->easycache.enabled ? "enabled" : "disabled", - sd_img_gen_params->easycache.reuse_threshold, - sd_img_gen_params->easycache.start_percent, - sd_img_gen_params->easycache.end_percent); + "cache: %s (threshold=%.3f, start=%.2f, end=%.2f)\n", + cache_mode_str, + sd_img_gen_params->cache.reuse_threshold, + sd_img_gen_params->cache.start_percent, + sd_img_gen_params->cache.end_percent); free(sample_params_str); return buf; } @@ -2885,7 +3162,7 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) { sd_vid_gen_params->video_frames = 6; sd_vid_gen_params->moe_boundary = 0.875f; sd_vid_gen_params->vace_strength = 1.f; - sd_easycache_params_init(&sd_vid_gen_params->easycache); + sd_cache_params_init(&sd_vid_gen_params->cache); } struct sd_ctx_t { @@ -2964,9 +3241,9 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx, std::vector ref_images, std::vector ref_latents, bool increase_ref_index, - ggml_tensor* concat_latent = nullptr, - ggml_tensor* denoise_mask = nullptr, - const sd_easycache_params_t* easycache_params = nullptr) { + ggml_tensor* concat_latent = nullptr, + ggml_tensor* denoise_mask = nullptr, + const sd_cache_params_t* cache_params = nullptr) { if (seed < 0) { // Generally, when using the provided command line, the seed is always >0. // However, to prevent potential issues if 'stable-diffusion.cpp' is invoked as a library @@ -3255,7 +3532,7 @@ sd_image_t* generate_image_internal(sd_ctx_t* sd_ctx, denoise_mask, nullptr, 1.0f, - easycache_params); + cache_params); int64_t sampling_end = ggml_time_ms(); if (x_0 != nullptr) { // print_ggml_tensor(x_0); @@ -3593,7 +3870,7 @@ sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_g sd_img_gen_params->increase_ref_index, concat_latent, denoise_mask, - &sd_img_gen_params->easycache); + &sd_img_gen_params->cache); size_t t2 = ggml_time_ms(); @@ -3964,7 +4241,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s denoise_mask, vace_context, sd_vid_gen_params->vace_strength, - &sd_vid_gen_params->easycache); + &sd_vid_gen_params->cache); int64_t sampling_end = ggml_time_ms(); LOG_INFO("sampling(high noise) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); @@ -4001,7 +4278,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s denoise_mask, vace_context, sd_vid_gen_params->vace_strength, - &sd_vid_gen_params->easycache); + &sd_vid_gen_params->cache); int64_t sampling_end = ggml_time_ms(); LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); diff --git a/otherarch/sdcpp/stable-diffusion.h b/otherarch/sdcpp/stable-diffusion.h index adb65a1d2..3e9faf854 100644 --- a/otherarch/sdcpp/stable-diffusion.h +++ b/otherarch/sdcpp/stable-diffusion.h @@ -189,10 +189,13 @@ typedef struct { bool tae_preview_only; bool diffusion_conv_direct; bool vae_conv_direct; + bool circular_x; + bool circular_y; bool force_sdxl_vae_conv_scale; bool chroma_use_dit_mask; bool chroma_use_t5_mask; int chroma_t5_mask_pad; + bool qwen_image_zero_cond_t; float flow_shift; } sd_ctx_params_t; @@ -236,12 +239,34 @@ typedef struct { float style_strength; } sd_pm_params_t; // photo maker +enum sd_cache_mode_t { + SD_CACHE_DISABLED = 0, + SD_CACHE_EASYCACHE, + SD_CACHE_UCACHE, + SD_CACHE_DBCACHE, + SD_CACHE_TAYLORSEER, + SD_CACHE_CACHE_DIT, +}; + typedef struct { - bool enabled; + enum sd_cache_mode_t mode; float reuse_threshold; float start_percent; float end_percent; -} sd_easycache_params_t; + float error_decay_rate; + bool use_relative_threshold; + bool reset_error_on_compute; + int Fn_compute_blocks; + int Bn_compute_blocks; + float residual_diff_threshold; + int max_warmup_steps; + int max_cached_steps; + int max_continuous_cached_steps; + int taylorseer_n_derivatives; + int taylorseer_skip_interval; + const char* scm_mask; + bool scm_policy_dynamic; +} sd_cache_params_t; typedef struct { bool is_high_noise; @@ -271,7 +296,7 @@ typedef struct { float control_strength; sd_pm_params_t pm_params; sd_tiling_params_t vae_tiling_params; - sd_easycache_params_t easycache; + sd_cache_params_t cache; } sd_img_gen_params_t; typedef struct { @@ -293,7 +318,7 @@ typedef struct { int64_t seed; int video_frames; float vace_strength; - sd_easycache_params_t easycache; + sd_cache_params_t cache; } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; @@ -323,7 +348,7 @@ SD_API enum preview_t str_to_preview(const char* str); SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode); SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str); -SD_API void sd_easycache_params_init(sd_easycache_params_t* easycache_params); +SD_API void sd_cache_params_init(sd_cache_params_t* cache_params); SD_API void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params); SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params); @@ -363,7 +388,8 @@ SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, enum sd_type_t output_type, - const char* tensor_type_rules); + const char* tensor_type_rules, + bool convert_name); SD_API bool preprocess_canny(sd_image_t image, float high_threshold, diff --git a/otherarch/sdcpp/ucache.hpp b/otherarch/sdcpp/ucache.hpp new file mode 100644 index 000000000..92e94a16d --- /dev/null +++ b/otherarch/sdcpp/ucache.hpp @@ -0,0 +1,404 @@ +#ifndef __UCACHE_HPP__ +#define __UCACHE_HPP__ + +#include +#include +#include +#include + +#include "denoiser.hpp" +#include "ggml_extend.hpp" + +struct UCacheConfig { + bool enabled = false; + float reuse_threshold = 1.0f; + float start_percent = 0.15f; + float end_percent = 0.95f; + float error_decay_rate = 1.0f; + bool use_relative_threshold = true; + bool adaptive_threshold = true; + float early_step_multiplier = 0.5f; + float late_step_multiplier = 1.5f; + bool reset_error_on_compute = true; +}; + +struct UCacheCacheEntry { + std::vector diff; +}; + +struct UCacheState { + UCacheConfig config; + Denoiser* denoiser = nullptr; + float start_sigma = std::numeric_limits::max(); + float end_sigma = 0.0f; + bool initialized = false; + bool initial_step = true; + bool skip_current_step = false; + bool step_active = false; + const SDCondition* anchor_condition = nullptr; + std::unordered_map cache_diffs; + std::vector prev_input; + std::vector prev_output; + float output_prev_norm = 0.0f; + bool has_prev_input = false; + bool has_prev_output = false; + bool has_output_prev_norm = false; + bool has_relative_transformation_rate = false; + float relative_transformation_rate = 0.0f; + float cumulative_change_rate = 0.0f; + float last_input_change = 0.0f; + bool has_last_input_change = false; + int total_steps_skipped = 0; + int current_step_index = -1; + int steps_computed_since_active = 0; + float accumulated_error = 0.0f; + float reference_output_norm = 0.0f; + + struct BlockMetrics { + float sum_transformation_rate = 0.0f; + float sum_output_norm = 0.0f; + int sample_count = 0; + float min_change_rate = std::numeric_limits::max(); + float max_change_rate = 0.0f; + + void reset() { + sum_transformation_rate = 0.0f; + sum_output_norm = 0.0f; + sample_count = 0; + min_change_rate = std::numeric_limits::max(); + max_change_rate = 0.0f; + } + + void record(float change_rate, float output_norm) { + if (std::isfinite(change_rate) && change_rate > 0.0f) { + sum_transformation_rate += change_rate; + sum_output_norm += output_norm; + sample_count++; + if (change_rate < min_change_rate) + min_change_rate = change_rate; + if (change_rate > max_change_rate) + max_change_rate = change_rate; + } + } + + float avg_transformation_rate() const { + return (sample_count > 0) ? (sum_transformation_rate / sample_count) : 0.0f; + } + + float avg_output_norm() const { + return (sample_count > 0) ? (sum_output_norm / sample_count) : 0.0f; + } + }; + BlockMetrics block_metrics; + int total_active_steps = 0; + + void reset_runtime() { + initial_step = true; + skip_current_step = false; + step_active = false; + anchor_condition = nullptr; + cache_diffs.clear(); + prev_input.clear(); + prev_output.clear(); + output_prev_norm = 0.0f; + has_prev_input = false; + has_prev_output = false; + has_output_prev_norm = false; + has_relative_transformation_rate = false; + relative_transformation_rate = 0.0f; + cumulative_change_rate = 0.0f; + last_input_change = 0.0f; + has_last_input_change = false; + total_steps_skipped = 0; + current_step_index = -1; + steps_computed_since_active = 0; + accumulated_error = 0.0f; + reference_output_norm = 0.0f; + block_metrics.reset(); + total_active_steps = 0; + } + + void init(const UCacheConfig& cfg, Denoiser* d) { + config = cfg; + denoiser = d; + initialized = cfg.enabled && d != nullptr; + reset_runtime(); + if (initialized) { + start_sigma = percent_to_sigma(config.start_percent); + end_sigma = percent_to_sigma(config.end_percent); + } + } + + void set_sigmas(const std::vector& sigmas) { + if (!initialized || sigmas.size() < 2) { + return; + } + size_t n_steps = sigmas.size() - 1; + + size_t start_step = static_cast(config.start_percent * n_steps); + size_t end_step = static_cast(config.end_percent * n_steps); + + if (start_step >= n_steps) + start_step = n_steps - 1; + if (end_step >= n_steps) + end_step = n_steps - 1; + + start_sigma = sigmas[start_step]; + end_sigma = sigmas[end_step]; + + if (start_sigma < end_sigma) { + std::swap(start_sigma, end_sigma); + } + } + + bool enabled() const { + return initialized && config.enabled; + } + + float percent_to_sigma(float percent) const { + if (!denoiser) { + return 0.0f; + } + if (percent <= 0.0f) { + return std::numeric_limits::max(); + } + if (percent >= 1.0f) { + return 0.0f; + } + float t = (1.0f - percent) * (TIMESTEPS - 1); + return denoiser->t_to_sigma(t); + } + + void begin_step(int step_index, float sigma) { + if (!enabled()) { + return; + } + if (step_index == current_step_index) { + return; + } + current_step_index = step_index; + skip_current_step = false; + has_last_input_change = false; + step_active = false; + + if (sigma > start_sigma) { + return; + } + if (!(sigma > end_sigma)) { + return; + } + step_active = true; + total_active_steps++; + } + + bool step_is_active() const { + return enabled() && step_active; + } + + bool is_step_skipped() const { + return enabled() && step_active && skip_current_step; + } + + float get_adaptive_threshold(int estimated_total_steps = 0) const { + float base_threshold = config.reuse_threshold; + + if (!config.adaptive_threshold) { + return base_threshold; + } + + int effective_total = estimated_total_steps; + if (effective_total <= 0) { + effective_total = std::max(20, steps_computed_since_active * 2); + } + + float progress = (effective_total > 0) ? (static_cast(steps_computed_since_active) / effective_total) : 0.0f; + + float multiplier = 1.0f; + if (progress < 0.2f) { + multiplier = config.early_step_multiplier; + } else if (progress > 0.8f) { + multiplier = config.late_step_multiplier; + } + + return base_threshold * multiplier; + } + + bool has_cache(const SDCondition* cond) const { + auto it = cache_diffs.find(cond); + return it != cache_diffs.end() && !it->second.diff.empty(); + } + + void update_cache(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) { + UCacheCacheEntry& entry = cache_diffs[cond]; + size_t ne = static_cast(ggml_nelements(output)); + entry.diff.resize(ne); + float* out_data = (float*)output->data; + float* in_data = (float*)input->data; + + for (size_t i = 0; i < ne; ++i) { + entry.diff[i] = out_data[i] - in_data[i]; + } + } + + void apply_cache(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) { + auto it = cache_diffs.find(cond); + if (it == cache_diffs.end() || it->second.diff.empty()) { + return; + } + + copy_ggml_tensor(output, input); + float* out_data = (float*)output->data; + const std::vector& diff = it->second.diff; + for (size_t i = 0; i < diff.size(); ++i) { + out_data[i] += diff[i]; + } + } + + bool before_condition(const SDCondition* cond, + ggml_tensor* input, + ggml_tensor* output, + float sigma, + int step_index) { + if (!enabled() || step_index < 0) { + return false; + } + if (step_index != current_step_index) { + begin_step(step_index, sigma); + } + if (!step_active) { + return false; + } + + if (initial_step) { + anchor_condition = cond; + initial_step = false; + } + + bool is_anchor = (cond == anchor_condition); + + if (skip_current_step) { + if (has_cache(cond)) { + apply_cache(cond, input, output); + return true; + } + return false; + } + + if (!is_anchor) { + return false; + } + + if (!has_prev_input || !has_prev_output || !has_cache(cond)) { + return false; + } + + size_t ne = static_cast(ggml_nelements(input)); + if (prev_input.size() != ne) { + return false; + } + + float* input_data = (float*)input->data; + last_input_change = 0.0f; + for (size_t i = 0; i < ne; ++i) { + last_input_change += std::fabs(input_data[i] - prev_input[i]); + } + if (ne > 0) { + last_input_change /= static_cast(ne); + } + has_last_input_change = true; + + if (has_output_prev_norm && has_relative_transformation_rate && + last_input_change > 0.0f && output_prev_norm > 0.0f) { + float approx_output_change_rate = (relative_transformation_rate * last_input_change) / output_prev_norm; + accumulated_error = accumulated_error * config.error_decay_rate + approx_output_change_rate; + + float effective_threshold = get_adaptive_threshold(); + if (config.use_relative_threshold && reference_output_norm > 0.0f) { + effective_threshold = effective_threshold * reference_output_norm; + } + + if (accumulated_error < effective_threshold) { + skip_current_step = true; + total_steps_skipped++; + apply_cache(cond, input, output); + return true; + } else if (config.reset_error_on_compute) { + accumulated_error = 0.0f; + } + } + + return false; + } + + void after_condition(const SDCondition* cond, ggml_tensor* input, ggml_tensor* output) { + if (!step_is_active()) { + return; + } + + update_cache(cond, input, output); + + if (cond != anchor_condition) { + return; + } + + size_t ne = static_cast(ggml_nelements(input)); + float* in_data = (float*)input->data; + prev_input.resize(ne); + for (size_t i = 0; i < ne; ++i) { + prev_input[i] = in_data[i]; + } + has_prev_input = true; + + float* out_data = (float*)output->data; + float output_change = 0.0f; + if (has_prev_output && prev_output.size() == ne) { + for (size_t i = 0; i < ne; ++i) { + output_change += std::fabs(out_data[i] - prev_output[i]); + } + if (ne > 0) { + output_change /= static_cast(ne); + } + } + + prev_output.resize(ne); + for (size_t i = 0; i < ne; ++i) { + prev_output[i] = out_data[i]; + } + has_prev_output = true; + + float mean_abs = 0.0f; + for (size_t i = 0; i < ne; ++i) { + mean_abs += std::fabs(out_data[i]); + } + output_prev_norm = (ne > 0) ? (mean_abs / static_cast(ne)) : 0.0f; + has_output_prev_norm = output_prev_norm > 0.0f; + + if (reference_output_norm == 0.0f) { + reference_output_norm = output_prev_norm; + } + + if (has_last_input_change && last_input_change > 0.0f && output_change > 0.0f) { + float rate = output_change / last_input_change; + if (std::isfinite(rate)) { + relative_transformation_rate = rate; + has_relative_transformation_rate = true; + block_metrics.record(rate, output_prev_norm); + } + } + + has_last_input_change = false; + } + + void log_block_metrics() const { + if (block_metrics.sample_count > 0) { + LOG_INFO("UCacheBlockMetrics: samples=%d, avg_rate=%.4f, min=%.4f, max=%.4f, avg_norm=%.4f", + block_metrics.sample_count, + block_metrics.avg_transformation_rate(), + block_metrics.min_change_rate, + block_metrics.max_change_rate, + block_metrics.avg_output_norm()); + } + } +}; + +#endif // __UCACHE_HPP__ diff --git a/otherarch/sdcpp/wan.hpp b/otherarch/sdcpp/wan.hpp index 75333bfe1..31ecf33f7 100644 --- a/otherarch/sdcpp/wan.hpp +++ b/otherarch/sdcpp/wan.hpp @@ -75,7 +75,7 @@ namespace WAN { lp2 -= (int)cache_x->ne[2]; } - x = ggml_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0); + x = ggml_ext_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), 0, 0, 0, @@ -206,9 +206,9 @@ namespace WAN { } else if (mode == "upsample3d") { x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST); } else if (mode == "downsample2d") { - x = ggml_pad(ctx->ggml_ctx, x, 1, 1, 0, 0); + x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); } else if (mode == "downsample3d") { - x = ggml_pad(ctx->ggml_ctx, x, 1, 1, 0, 0); + x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); } x = resample_1->forward(ctx, x); x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // (c, t, h, w) @@ -1826,7 +1826,7 @@ namespace WAN { } } - struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx, + struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, struct ggml_tensor* x) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; @@ -1835,8 +1835,7 @@ namespace WAN { int pad_t = (std::get<0>(params.patch_size) - T % std::get<0>(params.patch_size)) % std::get<0>(params.patch_size); int pad_h = (std::get<1>(params.patch_size) - H % std::get<1>(params.patch_size)) % std::get<1>(params.patch_size); int pad_w = (std::get<2>(params.patch_size) - W % std::get<2>(params.patch_size)) % std::get<2>(params.patch_size); - x = ggml_pad(ctx, x, pad_w, pad_h, pad_t, 0); // [N*C, T + pad_t, H + pad_h, W + pad_w] - + ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return x; } @@ -1986,14 +1985,14 @@ namespace WAN { int64_t T = x->ne[2]; int64_t C = x->ne[3]; - x = pad_to_patch_size(ctx->ggml_ctx, x); + x = pad_to_patch_size(ctx, x); int64_t t_len = ((T + (std::get<0>(params.patch_size) / 2)) / std::get<0>(params.patch_size)); int64_t h_len = ((H + (std::get<1>(params.patch_size) / 2)) / std::get<1>(params.patch_size)); int64_t w_len = ((W + (std::get<2>(params.patch_size) / 2)) / std::get<2>(params.patch_size)); if (time_dim_concat != nullptr) { - time_dim_concat = pad_to_patch_size(ctx->ggml_ctx, time_dim_concat); + time_dim_concat = pad_to_patch_size(ctx, time_dim_concat); x = ggml_concat(ctx->ggml_ctx, x, time_dim_concat, 2); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] t_len = ((x->ne[2] + (std::get<0>(params.patch_size) / 2)) / std::get<0>(params.patch_size)); } diff --git a/otherarch/sdcpp/z_image.hpp b/otherarch/sdcpp/z_image.hpp index bc554f177..af8d57e04 100644 --- a/otherarch/sdcpp/z_image.hpp +++ b/otherarch/sdcpp/z_image.hpp @@ -324,14 +324,14 @@ namespace ZImage { blocks["final_layer"] = std::make_shared(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels); } - struct ggml_tensor* pad_to_patch_size(struct ggml_context* ctx, + struct ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, struct ggml_tensor* x) { int64_t W = x->ne[0]; int64_t H = x->ne[1]; int pad_h = (z_image_params.patch_size - H % z_image_params.patch_size) % z_image_params.patch_size; int pad_w = (z_image_params.patch_size - W % z_image_params.patch_size) % z_image_params.patch_size; - x = ggml_pad(ctx, x, pad_w, pad_h, 0, 0); // [N, C, H + pad_h, W + pad_w] + x = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return x; } @@ -357,10 +357,10 @@ namespace ZImage { return x; } - struct ggml_tensor* process_img(struct ggml_context* ctx, + struct ggml_tensor* process_img(GGMLRunnerContext* ctx, struct ggml_tensor* x) { x = pad_to_patch_size(ctx, x); - x = patchify(ctx, x); + x = patchify(ctx->ggml_ctx, x); return x; } @@ -473,12 +473,12 @@ namespace ZImage { int64_t C = x->ne[2]; int64_t N = x->ne[3]; - auto img = process_img(ctx->ggml_ctx, x); + auto img = process_img(ctx, x); uint64_t n_img_token = img->ne[1]; if (ref_latents.size() > 0) { for (ggml_tensor* ref : ref_latents) { - ref = process_img(ctx->ggml_ctx, ref); + ref = process_img(ctx, ref); img = ggml_concat(ctx->ggml_ctx, img, ref, 1); } } @@ -552,6 +552,8 @@ namespace ZImage { ref_latents, increase_ref_index, z_image_params.theta, + circular_y_enabled, + circular_x_enabled, z_image_params.axes_dim); int pos_len = pe_vec.size() / z_image_params.axes_dim_sum / 2; // LOG_DEBUG("pos_len %d", pos_len);