fix: SD3 conditioning crash when clip_l text encoder is missing (#1638)

format-code.ps1

@@ -0,0 +1,54 @@
+$patterns = @(
+    "src/*.cpp"
+    "src/*.h"
+    "src/*.hpp"
+    "src/conditioning/*.cpp"
+    "src/conditioning/*.h"
+    "src/conditioning/*.hpp"
+    "src/core/*.cpp"
+    "src/core/*.h"
+    "src/core/*.hpp"
+    "src/extensions/*.cpp"
+    "src/extensions/*.h"
+    "src/extensions/*.hpp"
+    "src/runtime/*.cpp"
+    "src/runtime/*.h"
+    "src/runtime/*.hpp"
+    "src/model/*/*.cpp"
+    "src/model/*/*.h"
+    "src/model/*/*.hpp"
+    "src/tokenizers/*.h"
+    "src/tokenizers/*.cpp"
+    "src/tokenizers/vocab/*.h"
+    "src/tokenizers/vocab/*.cpp"
+    "src/model_io/*.h"
+    "src/model_io/*.cpp"
+    "examples/cli/*.cpp"
+    "examples/cli/*.h"
+    "examples/server/*.cpp"
+    "examples/common/*.hpp"
+    "examples/common/*.h"
+    "examples/common/*.cpp"
+)
+
+$root = (Get-Location).Path
+
+foreach ($pattern in $patterns) {
+    $files = Get-ChildItem -Path $pattern -File -ErrorAction SilentlyContinue | Sort-Object FullName
+
+    foreach ($file in $files) {
+        $relativePath = $file.FullName.Substring($root.Length).TrimStart('\', '/') -replace '\\', '/'
+
+        if ($relativePath -like "vocab*") {
+            continue
+        }
+
+        Write-Host "formatting '$relativePath'"
+
+        # if ($relativePath -ne "stable-diffusion.h") {
+        #     clang-tidy -fix -p build_linux/ "$relativePath"
+        # }
+
+        & clang-format -style=file -i $relativePath
+    }
+}

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

src/conditioning/conditioner.hpp

@@ -853,7 +853,7 @@ struct SD3CLIPEmbedder : public Conditioner {
             } else {
                 chunk_hidden_states_l = sd::Tensor<float>::zeros({768, static_cast<int64_t>(chunk_len), 1});
                 if (chunk_idx == 0) {
-                    pooled = sd::Tensor<float>::zeros({768, 1});
+                    pooled_l = sd::Tensor<float>::zeros({768, 1});
                 }
             }
 

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);
         }
 

src/model/common/rope.hpp

@@ -253,7 +253,8 @@ namespace Rope {
                                                                  int bs,
                                                                  float theta,
                                                                  int head_dim,
-                                                                 const std::vector<int>& mrope_section) {
+                                                                 const std::vector<int>& mrope_section,
+                                                                 const std::vector<std::vector<int>>& 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<std::vector<std::vector<float>>> axis_embs;
         axis_embs.reserve(3);
         for (int axis = 0; axis < 3; ++axis) {
-            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta));
+            std::vector<int> axis_wrap;
+            if (axis < static_cast<int>(axis_wrap_dims.size())) {
+                axis_wrap = axis_wrap_dims[axis];
+            }
+            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
         }
 
         std::vector<std::vector<float>> emb = axis_embs[0];

src/model/diffusion/ideogram4.hpp

@@ -151,7 +151,9 @@ namespace Ideogram4 {
                                                           int context_len,
                                                           int head_dim,
                                                           int rope_theta,
-                                                          const std::vector<int>& mrope_section) {
+                                                          const std::vector<int>& mrope_section,
+                                                          bool circular_x = false,
+                                                          bool circular_y = false) {
         GGML_ASSERT(bs == 1);
         std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
                                             std::vector<float>(3, 0.f));
@@ -169,7 +171,29 @@ namespace Ideogram4 {
             }
         }
 
-        return Rope::embed_interleaved_mrope(ids, bs, static_cast<float>(rope_theta), head_dim, mrope_section);
+        std::vector<std::vector<int>> axis_wrap_dims(3);
+        if (circular_y || circular_x) {
+            size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
+            axis_wrap_dims[1].assign(total_len, 0);
+            axis_wrap_dims[2].assign(total_len, 0);
+            if (circular_y) {
+                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
+                    axis_wrap_dims[1][idx] = grid_h;
+                }
+            }
+            if (circular_x) {
+                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
+                    axis_wrap_dims[2][idx] = grid_w;
+                }
+            }
+        }
+
+        return Rope::embed_interleaved_mrope(ids,
+                                            bs,
+                                            static_cast<float>(rope_theta),
+                                            head_dim,
+                                            mrope_section,
+                                            axis_wrap_dims);
     }
 
     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<int>(grid_h),
                                        static_cast<int>(grid_w),
                                        static_cast<int>(x->ne[3]),
                                        static_cast<int>(context_len),
                                        static_cast<int>(head_dim),
                                        static_cast<int>(config.rope_theta),
-                                       config.mrope_section);
+                                       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;

src/stable-diffusion.cpp

@@ -575,15 +575,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<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),
                                                                         params_backend_for(SDBackendModule::TE),
                                                                         tensor_storage_map,
@@ -5570,3 +5561,18 @@ 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);
+}