feat: add text context CLI parameter

examples/common/common.cpp

@@ -1045,6 +1045,10 @@ ArgOptions SDGenerationParams::get_options() {
          "ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer (default: -1). "
          "<= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x",
          &clip_skip},
+        {"",
+         "--text-ctx",
+         "maximum text/cross-attention context length; 0 disables truncation (default: 0)",
+         &text_ctx},
         {"-b",
          "--batch-count",
          "batch count",
@@ -1890,6 +1894,7 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("scm_mask", scm_mask);
 
     load_if_exists("clip_skip", clip_skip);
+    load_if_exists("text_ctx", text_ctx);
     load_if_exists("width", width);
     load_if_exists("height", height);
     load_if_exists("batch_count", batch_count);
@@ -2343,6 +2348,11 @@ bool SDGenerationParams::validate(SDMode mode) {
         return false;
     }
 
+    if (text_ctx < 0) {
+        LOG_ERROR("error: text_ctx must be non-negative");
+        return false;
+    }
+
     if (!cache_mode.empty()) {
         if (cache_mode == "easycache" || cache_mode == "ucache") {
             if (cache_params.reuse_threshold < 0.0f) {
@@ -2504,6 +2514,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.prompt              = prompt.c_str();
     params.negative_prompt     = negative_prompt.c_str();
     params.clip_skip           = clip_skip;
+    params.text_ctx            = text_ctx;
     params.init_image          = init_image.get();
     params.ref_images          = ref_image_views.empty() ? nullptr : ref_image_views.data();
     params.ref_images_count    = static_cast<int>(ref_image_views.size());
@@ -2648,6 +2659,7 @@ std::string SDGenerationParams::to_string() const {
         << "  ad_negative_prompt: \"" << ad_negative_prompt << "\",\n"
         << "  extra_ad_args: \"" << extra_ad_args << "\",\n"
         << "  clip_skip: " << clip_skip << ",\n"
+        << "  text_ctx: " << text_ctx << ",\n"
         << "  width: " << width << ",\n"
         << "  height: " << height << ",\n"
         << "  batch_count: " << batch_count << ",\n"
@@ -2810,6 +2822,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
     root["models"] = std::move(models);
 
     root["clip_skip"]             = gen_params.clip_skip;
+    root["text_ctx"]              = gen_params.text_ctx;
     root["strength"]              = gen_params.strength;
     root["control_strength"]      = gen_params.control_strength;
     root["ip_adapter_strength"]   = gen_params.ip_adapter_strength;

examples/common/common.h

@@ -194,6 +194,7 @@ struct SDGenerationParams {
     std::string ad_negative_prompt;
     std::string extra_ad_args;
     int clip_skip              = -1;  // <= 0 represents unspecified
+    int text_ctx               = 0;   // <= 0 disables text context truncation
     int width                  = -1;
     int height                 = -1;
     int batch_count            = 1;

examples/server/routes_sdcpp.cpp

@@ -121,6 +121,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"prompt", defaults.prompt},
         {"negative_prompt", defaults.negative_prompt},
         {"clip_skip", defaults.clip_skip},
+        {"text_ctx", defaults.text_ctx},
         {"width", defaults.width > 0 ? defaults.width : 512},
         {"height", defaults.height > 0 ? defaults.height : 512},
         {"strength", defaults.strength},

include/stable-diffusion.h

@@ -363,6 +363,7 @@ typedef struct {
     const char* prompt;
     const char* negative_prompt;
     int clip_skip;
+    int text_ctx;
     sd_image_t init_image;
     sd_image_t* ref_images;
     int ref_images_count;

src/stable-diffusion.cpp

@@ -3604,6 +3604,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     *sd_img_gen_params = {};
     sd_sample_params_init(&sd_img_gen_params->sample_params);
     sd_img_gen_params->clip_skip           = -1;
+    sd_img_gen_params->text_ctx            = 0;
     sd_img_gen_params->ref_images_count    = 0;
     sd_img_gen_params->ref_image_args      = "";
     sd_img_gen_params->width               = 512;
@@ -3635,6 +3636,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "prompt: %s\n"
              "negative_prompt: %s\n"
              "clip_skip: %d\n"
+             "text_ctx: %d\n"
              "width: %d\n"
              "height: %d\n"
              "sample_params: %s\n"
@@ -3654,6 +3656,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              SAFE_STR(sd_img_gen_params->prompt),
              SAFE_STR(sd_img_gen_params->negative_prompt),
              sd_img_gen_params->clip_skip,
+             sd_img_gen_params->text_ctx,
              sd_img_gen_params->width,
              sd_img_gen_params->height,
              SAFE_STR(sample_params_str),
@@ -3947,6 +3950,7 @@ struct GenerationRequest {
     int width                                = -1;
     int height                               = -1;
     int clip_skip                            = -1;
+    int text_ctx                             = 0;
     int vae_scale_factor                     = -1;
     int diffusion_model_down_factor          = -1;
     int64_t seed                             = -1;
@@ -3983,6 +3987,7 @@ struct GenerationRequest {
         batch_count                 = sd_img_gen_params->batch_count;
         qwen_image_layers           = std::max(0, sd_img_gen_params->qwen_image_layers);
         clip_skip                   = sd_img_gen_params->clip_skip;
+        text_ctx                    = sd_img_gen_params->text_ctx;
         shifted_timestep            = sd_img_gen_params->sample_params.shifted_timestep;
         strength                    = sd_img_gen_params->strength;
         control_strength            = sd_img_gen_params->control_strength;
@@ -4612,6 +4617,32 @@ struct ConditionerRunnerDoneOnExit {
     }
 };
 
+static void truncate_text_context_tensor(sd::Tensor<float>* tensor, int text_ctx, const char* name) {
+    if (tensor == nullptr || text_ctx <= 0 || tensor->empty() || tensor->dim() < 2) {
+        return;
+    }
+
+    const int64_t old_ctx = tensor->shape()[1];
+    if (old_ctx <= text_ctx) {
+        return;
+    }
+
+    *tensor = sd::ops::slice(*tensor, 1, 0, text_ctx);
+    LOG_INFO("truncated %s text context from %" PRId64 " to %d", name, old_ctx, text_ctx);
+}
+
+static void truncate_text_context_condition(SDCondition* condition, int text_ctx, const char* name) {
+    if (condition == nullptr || text_ctx <= 0) {
+        return;
+    }
+
+    truncate_text_context_tensor(&condition->c_crossattn, text_ctx, name);
+    for (size_t i = 0; i < condition->extra_c_crossattns.size(); ++i) {
+        std::string extra_name = std::string(name) + ".extra_c_crossattns[" + std::to_string(i) + "]";
+        truncate_text_context_tensor(&condition->extra_c_crossattns[i], text_ctx, extra_name.c_str());
+    }
+}
+
 struct CircularAxesState {
     bool circular_x = false;
     bool circular_y = false;
@@ -5114,6 +5145,10 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
         }
     }
 
+    truncate_text_context_condition(&cond, request->text_ctx, "cond");
+    truncate_text_context_condition(&uncond, request->text_ctx, "uncond");
+    truncate_text_context_condition(&img_uncond, request->text_ctx, "img_uncond");
+
     int64_t t1 = ggml_time_ms();
     LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000);