From b33042bc189c0b390c432ecca0de8b8425bbe8c2 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Mon, 15 Jun 2026 13:33:00 -0300 Subject: [PATCH] sd: sync with master-690-3a54597 (#2270) * sd: sync with master-690-3a54597 * sd: drop Chroma workaround for Flash Attention Fixed in master-686-1b702a5 . * sd: skip setting cpu flags (soon to be removed) --- otherarch/sdcpp/include/stable-diffusion.h | 4 +++ otherarch/sdcpp/sdtype_adapter.cpp | 24 ++----------- .../sdcpp/src/conditioning/conditioner.hpp | 2 +- otherarch/sdcpp/src/core/ggml_extend.hpp | 16 ++++++--- otherarch/sdcpp/src/model/common/rope.hpp | 9 +++-- .../sdcpp/src/model/diffusion/ideogram4.hpp | 34 ++++++++++++++++--- otherarch/sdcpp/src/stable-diffusion.cpp | 24 +++++++------ 7 files changed, 70 insertions(+), 43 deletions(-) diff --git a/otherarch/sdcpp/include/stable-diffusion.h b/otherarch/sdcpp/include/stable-diffusion.h index 17596f849..2175f895a 100644 --- a/otherarch/sdcpp/include/stable-diffusion.h +++ b/otherarch/sdcpp/include/stable-diffusion.h @@ -491,6 +491,10 @@ SD_API bool preprocess_canny(sd_image_t image, SD_API const char* sd_commit(void); SD_API const char* sd_version(void); +// for C API, caller needs to call free_sd_images to free the memory after use +// This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT. +SD_API void free_sd_images(sd_image_t* result_images, int num_images); + #ifdef __cplusplus } #endif diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp index f7a5722e3..c37c9fcb8 100644 --- a/otherarch/sdcpp/sdtype_adapter.cpp +++ b/otherarch/sdcpp/sdtype_adapter.cpp @@ -124,8 +124,6 @@ struct SDParams { bool diffusion_conv_direct = false; bool vae_conv_direct = false; - bool chroma_use_dit_mask = true; - LoraMap lora_map; bool lora_dynamic = false; @@ -481,17 +479,12 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { params.diffusion_flash_attn = sd_params->diffusion_flash_attn; params.diffusion_conv_direct = sd_params->diffusion_conv_direct; params.vae_conv_direct = sd_params->vae_conv_direct; - params.chroma_use_dit_mask = sd_params->chroma_use_dit_mask; + params.chroma_use_dit_mask = true; params.max_vram = inputs.max_vram; params.stream_layers = inputs.stream_layers; params.enable_mmap = inputs.use_mmap; - // the _cpu flags are only used if the backend string is empty, but - // we always set both for consistency - params.offload_params_to_cpu = inputs.offload_cpu; params.params_backend = inputs.offload_cpu ? "CPU" : ""; - params.keep_vae_on_cpu = (inputs.kcpp_vae_device <= -2); backends += get_device_override(inputs.kcpp_vae_device, "VAE"); - params.keep_clip_on_cpu = (inputs.kcpp_clip_device <= -2); backends += get_device_override(inputs.kcpp_clip_device, "CLIP"); if (backends.rfind(",", 0) == 0) { backends = "auto" + backends; @@ -505,11 +498,6 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { // also switches flash attn for the vae and conditioner params.flash_attn = params.diffusion_flash_attn; - if (params.chroma_use_dit_mask && params.diffusion_flash_attn) { - // note we don't know yet if it's a Chroma model - params.chroma_use_dit_mask = false; - } - if(inputs.debugmode==1) { char* buf = sd_ctx_params_to_str(¶ms); @@ -529,14 +517,6 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { auto info = get_model_info(sd_ctx); - if (!sd_is_quiet) { - if (info.is_chroma && sd_params->diffusion_flash_attn && sd_params->chroma_use_dit_mask) - { - printf("Chroma: flash attention is on, disabling DiT mask (this will lower image quality)\n"); - // disabled before loading - } - } - if (info.is_wan || info.is_ltx) { printf("\nSetting to Video Generation Mode!\n"); @@ -559,7 +539,7 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) { if (upscaler_filename!="") { const int upscale_tile_size = 128; upscaler_ctx = new_upscaler_ctx(upscaler_filename.c_str(), - params.offload_params_to_cpu, + inputs.offload_cpu, params.diffusion_conv_direct, params.n_threads, upscale_tile_size, diff --git a/otherarch/sdcpp/src/conditioning/conditioner.hpp b/otherarch/sdcpp/src/conditioning/conditioner.hpp index 0cb3172b9..5e74af073 100644 --- a/otherarch/sdcpp/src/conditioning/conditioner.hpp +++ b/otherarch/sdcpp/src/conditioning/conditioner.hpp @@ -853,7 +853,7 @@ struct SD3CLIPEmbedder : public Conditioner { } else { chunk_hidden_states_l = sd::Tensor::zeros({768, static_cast(chunk_len), 1}); if (chunk_idx == 0) { - pooled = sd::Tensor::zeros({768, 1}); + pooled_l = sd::Tensor::zeros({768, 1}); } } diff --git a/otherarch/sdcpp/src/core/ggml_extend.hpp b/otherarch/sdcpp/src/core/ggml_extend.hpp index 893ed1400..c85728529 100644 --- a/otherarch/sdcpp/src/core/ggml_extend.hpp +++ b/otherarch/sdcpp/src/core/ggml_extend.hpp @@ -1346,10 +1346,18 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); if (mask_in != nullptr) { - mask_in = ggml_transpose(ctx, mask_in); - } - - if (mask_in != nullptr) { + // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped + // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and, + // unlike the manual-attention path, does not broadcast the query dimension. + // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over + // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the + // kernel indexes it correctly. (A bare ggml_transpose here produced a + // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank + // output for masked flash attention.) + if (mask_in->ne[1] != L_q) { + mask_in = ggml_repeat(ctx, mask_in, + ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3])); + } mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16); } diff --git a/otherarch/sdcpp/src/model/common/rope.hpp b/otherarch/sdcpp/src/model/common/rope.hpp index 3ab2551d0..c0077de33 100644 --- a/otherarch/sdcpp/src/model/common/rope.hpp +++ b/otherarch/sdcpp/src/model/common/rope.hpp @@ -253,7 +253,8 @@ namespace Rope { int bs, float theta, int head_dim, - const std::vector& mrope_section) { + const std::vector& mrope_section, + const std::vector>& axis_wrap_dims = {}) { GGML_ASSERT(bs > 0); GGML_ASSERT(head_dim % 2 == 0); GGML_ASSERT(mrope_section.size() >= 3); @@ -265,7 +266,11 @@ namespace Rope { std::vector>> axis_embs; axis_embs.reserve(3); for (int axis = 0; axis < 3; ++axis) { - axis_embs.push_back(rope(trans_ids[axis], head_dim, theta)); + std::vector axis_wrap; + if (axis < static_cast(axis_wrap_dims.size())) { + axis_wrap = axis_wrap_dims[axis]; + } + axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap)); } std::vector> emb = axis_embs[0]; diff --git a/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp b/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp index 5149ff72d..2f53c787c 100644 --- a/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp +++ b/otherarch/sdcpp/src/model/diffusion/ideogram4.hpp @@ -151,7 +151,9 @@ namespace Ideogram4 { int context_len, int head_dim, int rope_theta, - const std::vector& mrope_section) { + const std::vector& mrope_section, + bool circular_x = false, + bool circular_y = false) { GGML_ASSERT(bs == 1); std::vector> ids(static_cast(bs) * (context_len + grid_h * grid_w), std::vector(3, 0.f)); @@ -169,7 +171,29 @@ namespace Ideogram4 { } } - return Rope::embed_interleaved_mrope(ids, bs, static_cast(rope_theta), head_dim, mrope_section); + std::vector> axis_wrap_dims(3); + if (circular_y || circular_x) { + size_t total_len = static_cast(bs) * (context_len + grid_h * grid_w); + axis_wrap_dims[1].assign(total_len, 0); + axis_wrap_dims[2].assign(total_len, 0); + if (circular_y) { + for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { + axis_wrap_dims[1][idx] = grid_h; + } + } + if (circular_x) { + for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { + axis_wrap_dims[2][idx] = grid_w; + } + } + } + + return Rope::embed_interleaved_mrope(ids, + bs, + static_cast(rope_theta), + head_dim, + mrope_section, + axis_wrap_dims); } class Ideogram4Attention : public GGMLBlock { @@ -480,13 +504,16 @@ namespace Ideogram4 { int64_t pos_len = context_len + grid_h * grid_w; int64_t head_dim = config.emb_dim / config.num_heads; + auto runner_ctx = get_context(); pe_vec = gen_ideogram4_pe(static_cast(grid_h), static_cast(grid_w), static_cast(x->ne[3]), static_cast(context_len), static_cast(head_dim), static_cast(config.rope_theta), - config.mrope_section); + config.mrope_section, + runner_ctx.circular_x_enabled, + runner_ctx.circular_y_enabled); auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); @@ -497,7 +524,6 @@ namespace Ideogram4 { auto indicator = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_I32, pos_len, x->ne[3]); set_backend_tensor_data(indicator, image_indicator_vec.data()); - auto runner_ctx = get_context(); ggml_tensor* out = active_model.forward(&runner_ctx, x, timesteps, context, pe, indicator); ggml_build_forward_expand(gf, out); return gf; diff --git a/otherarch/sdcpp/src/stable-diffusion.cpp b/otherarch/sdcpp/src/stable-diffusion.cpp index c59831c41..678787c08 100644 --- a/otherarch/sdcpp/src/stable-diffusion.cpp +++ b/otherarch/sdcpp/src/stable-diffusion.cpp @@ -813,15 +813,6 @@ public: } } if (is_chroma) { - if ((sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) && sd_ctx_params->chroma_use_dit_mask) { - LOG_WARN( - "!!!It looks like you are using Chroma with flash attention. " - "This is currently unsupported. " - "If you find that the generated images are broken, " - "try either disabling flash attention or specifying " - "--chroma-disable-dit-mask as a workaround."); - } - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), params_backend_for(SDBackendModule::TE), tensor_storage_map, @@ -5848,6 +5839,20 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, return true; } +SD_API void free_sd_images(sd_image_t* result_images, int num_images) { + if (result_images == nullptr) { + return; + } + + for (int i = 0; i < num_images; ++i) { + if (result_images[i].data != nullptr) { + free(result_images[i].data); + result_images[i].data = nullptr; + } + } + + free(result_images); +} #include "kcpp_sd_extensions.h" @@ -5909,4 +5914,3 @@ namespace kcpp_sd { } } -