From bfb47cbcd8fbe3712a283624ddf71427ca98a63a Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Sat, 14 Jun 2025 10:10:34 +0800
Subject: [PATCH] Revert "revert padding change for sd chroma"

This reverts commit 7de88802f90e3940638de8f6d292261cfe56f1fa.
---
 otherarch/sdcpp/conditioner.hpp | 30 ++++++++++++++++++++++++++++++
 otherarch/sdcpp/flux.hpp        | 27 ---------------------------
 2 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/otherarch/sdcpp/conditioner.hpp b/otherarch/sdcpp/conditioner.hpp
index 8f67e56b1..f48f4f493 100644
--- a/otherarch/sdcpp/conditioner.hpp
+++ b/otherarch/sdcpp/conditioner.hpp
@@ -1288,6 +1288,21 @@ struct PixArtCLIPEmbedder : public Conditioner {
         return {t5_tokens, t5_weights, t5_mask};
     }
 
+    void modify_mask_to_attend_padding(struct ggml_tensor* mask, int max_seq_length, int num_extra_padding = 8) {
+        float* mask_data = (float*)mask->data;
+        int num_pad      = 0;
+        for (int64_t i = 0; i < max_seq_length; i++) {
+            if (num_pad >= num_extra_padding) {
+                break;
+            }
+            if (std::isinf(mask_data[i])) {
+                mask_data[i] = 0;
+                ++num_pad;
+            }
+        }
+        // LOG_DEBUG("PAD: %d", num_pad);
+    }
+
     SDCondition get_learned_condition_common(ggml_context* work_ctx,
                                              int n_threads,
                                              std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> token_and_weights,
@@ -1374,6 +1389,21 @@ struct PixArtCLIPEmbedder : public Conditioner {
             hidden_states = ggml_new_tensor_2d(work_ctx, GGML_TYPE_F32, 4096, 256);
             ggml_set_f32(hidden_states, 0.f);
         }
+
+        int mask_pad                            = 1;
+        const char* SD_CHROMA_MASK_PAD_OVERRIDE = getenv("SD_CHROMA_MASK_PAD_OVERRIDE");
+        if (SD_CHROMA_MASK_PAD_OVERRIDE != nullptr) {
+            std::string mask_pad_str = SD_CHROMA_MASK_PAD_OVERRIDE;
+            try {
+                mask_pad = std::stoi(mask_pad_str);
+            } catch (const std::invalid_argument&) {
+                LOG_WARN("SD_CHROMA_MASK_PAD_OVERRIDE environment variable is not a valid integer (%s). Falling back to default (%d)", SD_CHROMA_MASK_PAD_OVERRIDE, mask_pad);
+            } catch (const std::out_of_range&) {
+                LOG_WARN("SD_CHROMA_MASK_PAD_OVERRIDE environment variable value is out of range for `int` type (%s). Falling back to default (%d)", SD_CHROMA_MASK_PAD_OVERRIDE, mask_pad);
+            }
+        }
+        modify_mask_to_attend_padding(t5_attn_mask, ggml_nelements(t5_attn_mask), mask_pad);
+
         return SDCondition(hidden_states, t5_attn_mask, NULL);
     }
 
diff --git a/otherarch/sdcpp/flux.hpp b/otherarch/sdcpp/flux.hpp
index 1df9ca11f..12f99c51c 100644
--- a/otherarch/sdcpp/flux.hpp
+++ b/otherarch/sdcpp/flux.hpp
@@ -709,20 +709,6 @@ namespace Flux {
             return ids;
         }
 
-        void chroma_modify_mask_to_attend_padding(struct ggml_tensor* mask, int max_seq_length, int num_extra_padding = 8) {
-            float* mask_data = (float*)mask->data;
-            int num_pad      = 0;
-            for (int64_t i = 0; i < max_seq_length; i++) {
-                if (num_pad >= num_extra_padding) {
-                    break;
-                }
-                if (std::isinf(mask_data[i])) {
-                    mask_data[i] = 0;
-                    ++num_pad;
-                }
-            }
-            // LOG_DEBUG("PAD: %d", num_pad);
-        }
 
         // Generate positional embeddings
         std::vector<float> gen_pe(int h, int w, int patch_size, int bs, int context_len, int theta, const std::vector<int>& axes_dim) {
@@ -1098,19 +1084,6 @@ namespace Flux {
                     guidance = ggml_set_f32(guidance, 0);
                 }
 
-                int mask_pad                            = 1;
-                const char* SD_CHROMA_MASK_PAD_OVERRIDE = getenv("SD_CHROMA_MASK_PAD_OVERRIDE");
-                if (SD_CHROMA_MASK_PAD_OVERRIDE != nullptr) {
-                    std::string mask_pad_str = SD_CHROMA_MASK_PAD_OVERRIDE;
-                    try {
-                        mask_pad = std::stoi(mask_pad_str);
-                    } catch (const std::invalid_argument&) {
-                        LOG_WARN("SD_CHROMA_MASK_PAD_OVERRIDE environment variable is not a valid integer (%s). Falling back to default (%d)", SD_CHROMA_MASK_PAD_OVERRIDE, mask_pad);
-                    } catch (const std::out_of_range&) {
-                        LOG_WARN("SD_CHROMA_MASK_PAD_OVERRIDE environment variable value is out of range for `int` type (%s). Falling back to default (%d)", SD_CHROMA_MASK_PAD_OVERRIDE, mask_pad);
-                    }
-                }
-                flux.chroma_modify_mask_to_attend_padding(y, ggml_nelements(y), mask_pad);
 
                 const char* SD_CHROMA_USE_DIT_MASK = getenv("SD_CHROMA_USE_DIT_MASK");
                 if (SD_CHROMA_USE_DIT_MASK != nullptr) {