diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 03a35c9a7..00f6777cb 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -928,6 +928,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", @@ -1701,6 +1705,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); @@ -2143,6 +2148,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) { @@ -2277,6 +2287,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(ref_image_views.size()); @@ -2410,6 +2421,7 @@ std::string SDGenerationParams::to_string() const { << " prompt: \"" << prompt << "\",\n" << " negative_prompt: \"" << negative_prompt << "\",\n" << " clip_skip: " << clip_skip << ",\n" + << " text_ctx: " << text_ctx << ",\n" << " width: " << width << ",\n" << " height: " << height << ",\n" << " batch_count: " << batch_count << ",\n" @@ -2558,6 +2570,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["auto_resize_ref_image"] = gen_params.auto_resize_ref_image; diff --git a/examples/common/common.h b/examples/common/common.h index 86fcc1627..aab8c6282 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -191,6 +191,7 @@ struct SDGenerationParams { std::string prompt; std::string negative_prompt; 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; diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index c60e9da65..3870f2bf6 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/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}, diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 674c9d63a..9077d8c64 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -349,6 +349,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; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index c74d73634..5dfb33249 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -2775,6 +2775,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->width = 512; sd_img_gen_params->height = 512; @@ -2800,6 +2801,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" @@ -2817,6 +2819,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), @@ -3065,6 +3068,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; @@ -3100,6 +3104,7 @@ struct GenerationRequest { seed = sd_img_gen_params->seed; batch_count = sd_img_gen_params->batch_count; 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; @@ -3730,6 +3735,32 @@ struct ConditionerRunnerDoneOnExit { } }; +static void truncate_text_context_tensor(sd::Tensor* 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; @@ -4095,6 +4126,10 @@ static std::optional 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);