From e53d21d74836ffa861e77a3a3b2936b142b5cb88 Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Sat, 2 Mar 2024 12:05:59 +0800
Subject: [PATCH] sanitize SD prompt to avoid segfault

---
 klite.embd                         | 10 +++++++---
 otherarch/sdcpp/sdtype_adapter.cpp | 23 +++++++++++++++++++++--
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/klite.embd b/klite.embd
index fe86cf6b5..6ea7709fe 100644
--- a/klite.embd
+++ b/klite.embd
@@ -8866,7 +8866,7 @@ Current version: 118
 	function do_auto_gen_image(truncated_context)
 	{
 		var tclen = truncated_context.length;
-		var sentence = truncated_context.substring(tclen - 300, tclen);
+		var sentence = truncated_context.substring(tclen - 350, tclen);
 		sentence = start_trim_to_sentence(sentence);
 		sentence = end_trim_to_sentence(sentence,true);
 		if (sentence.length > 0) {
@@ -8895,7 +8895,7 @@ Current version: 118
 			let userinput = getInputBoxValue();
 			if(userinput.trim()!="")
 			{
-				var sentence = userinput.trim().substring(0, 300);
+				var sentence = userinput.trim().substring(0, 350);
 				do_manual_gen_image(sentence);
 			}
 		},false);
@@ -10221,6 +10221,10 @@ Current version: 118
 		}
 
 		let negprompt = localsettings.image_negprompt?(" ### "+localsettings.image_negprompt):" ### ugly, deformed, poorly, censor, blurry, lowres, malformed, watermark, duplicated, grainy, distorted, signature";
+		if(localsettings.image_negprompt=="none")
+		{
+			negprompt = "";
+		}
 
 		let genimg_payload = {
 			"prompt": (sentence + negprompt),
@@ -14360,7 +14364,7 @@ Current version: 118
 			<div class="aidgpopuplistheader anotelabel">Style tags to use for generating images:<br>(E.g. Sketch, Realistic, Anime, 3D Render, Drawing)<br></div>
 			<input class="form-control" type="text" placeholder="Default Style" value="" id="imagestyleinput">
 			<div class="aidgpopuplistheader anotelabel">Negative Prompt<br></div>
-			<input class="form-control" type="text" placeholder="Default Negative Prompt" value="" id="negpromptinput">
+			<input class="form-control" type="text" placeholder="Default Negative Prompt. Put &quot;none&quot; to skip" value="" id="negpromptinput">
 
 			<div class="inlinelabel">
 				<div class="justifyleft" style="padding:4px">Number of Steps: </div>
diff --git a/otherarch/sdcpp/sdtype_adapter.cpp b/otherarch/sdcpp/sdtype_adapter.cpp
index 3341c53be..e770b64ba 100644
--- a/otherarch/sdcpp/sdtype_adapter.cpp
+++ b/otherarch/sdcpp/sdtype_adapter.cpp
@@ -8,6 +8,7 @@
 
 #include <inttypes.h>
 #include <cinttypes>
+#include <algorithm>
 
 #include "model_adapter.h"
 
@@ -225,6 +226,20 @@ bool sdtype_load_model(const sd_load_model_inputs inputs) {
 
 }
 
+std::string clean_input_prompt(const std::string& input) {
+    std::string result;
+    result.reserve(input.size());
+    for (char ch : input) {
+        // Check if the character is an ASCII or extended ASCII character
+        if (static_cast<unsigned char>(ch) <= 0x7F || (ch >= 0xC2 && ch <= 0xF4)) {
+            result.push_back(ch);
+        }
+    }
+    //limit to max 800 chars
+    result = result.substr(0, 800);
+    return result;
+}
+
 sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
 {
     sd_generation_outputs output;
@@ -240,8 +255,12 @@ sd_generation_outputs sdtype_generate(const sd_generation_inputs inputs)
     sd_image_t * results;
     sd_image_t* control_image = NULL;
 
-    sd_params->prompt = inputs.prompt;
-    sd_params->negative_prompt = inputs.negative_prompt;
+    //sanitize prompts, remove quotes and limit lengths
+    std::string cleanprompt = clean_input_prompt(inputs.prompt);
+    std::string cleannegprompt = clean_input_prompt(inputs.negative_prompt);
+
+    sd_params->prompt = cleanprompt;
+    sd_params->negative_prompt = cleannegprompt;
     sd_params->cfg_scale = inputs.cfg_scale;
     sd_params->sample_steps = inputs.sample_steps;
     sd_params->seed = inputs.seed;