From 51f112168157cb527f8008acaefe396bdf25f061 Mon Sep 17 00:00:00 2001 From: Llama <34464159+pi6am@users.noreply.github.com> Date: Thu, 24 Nov 2022 11:40:57 -0800 Subject: [PATCH 1/9] Several fixes to calc_ai_text and lua_compute_context The method lua_compute_context no longer ignores its arguments. This restores support for overriding the submitted text and allowed WI entries and folders to lua_compute_context. Add optional allowed_wi_entries and allowed_wi_folders parameters to calc_ai_text. These allow restricting which WI is added to the context. The default is no restriction. Fix an issue in which streaming tokens were returned by calc_ai_text. This fixes an issue in which lua_compute_context would receive streamed tokens (only if token streaming was turned on). Add docstrings to calc_ai_text and to_sentences. Improve the regular expression used to split actions by sentence. The new regular expression allows quotation marks after punctuation to be attached to the sentence. Improve the interaction between AN and WI depth and actions split by sentence. The author's note was added after the sentence that pushed the action count up to or greater than the depth. This included "empty" actions from undoing actions, as well as continuation actions. As a result, the author's note text was often inserted at the end of the context, which often prevented the model from generating coherent continuation text. Consider sentence count in addition to action count for AN and WI depth. Ignore completely empty actions when determining AN and WI depth. Insert AN text before the sentence that crosses the AN depth, instead of after the sentence. This means that at least one sentence from the action alwasy appears at the end of the context, which gives the AI something to continue from. A few extremely minor optimizations from reducing redundant work. Pre-compile the sentence splitting regular expression. Don't join action sentences multiple times just to compute their length. Fix a few typos and remove some commented out code. --- aiserver.py | 21 ++------- koboldai_settings.py | 101 ++++++++++++++++++++++++++++++++----------- 2 files changed, 80 insertions(+), 42 deletions(-) diff --git a/aiserver.py b/aiserver.py index 7221a3da..3d99a9a5 100644 --- a/aiserver.py +++ b/aiserver.py @@ -3477,21 +3477,9 @@ def lua_compute_context(submission, entries, folders, kwargs): while(folders[i] is not None): allowed_folders.add(int(folders[i])) i += 1 - #winfo, mem, anotetxt, _ = calcsubmitbudgetheader( - # submission, - # allowed_entries=allowed_entries, - # allowed_folders=allowed_folders, - # force_use_txt=True, - # scan_story=kwargs["scan_story"] if kwargs["scan_story"] != None else True, - #) - txt, _, _, found_entries = koboldai_vars.calc_ai_text() - #txt, _, _ = calcsubmitbudget( - # len(actions), - # winfo, - # mem, - # anotetxt, - # actions, - #) + txt, _, _, found_entries = koboldai_vars.calc_ai_text(submitted_text=submission, + allowed_wi_entries=allowed_entries, + allowed_wi_folders=allowed_folders) return utils.decodenewlines(tokenizer.decode(txt)) #==================================================================# @@ -4725,7 +4713,7 @@ def actionsubmit(data, actionmode=0, force_submit=False, force_prompt_gen=False, if(len(koboldai_vars.prompt.strip()) == 0): koboldai_vars.prompt = data else: - koboldai_vars.actions.append(data) + koboldai_vars.actions.append_submission(data) update_story_chunk('last') send_debug() @@ -6079,7 +6067,6 @@ def generate(txt, minimum, maximum, found_entries=None): if(len(genout) == 1): genresult(genout[0]["generated_text"]) - #koboldai_vars.actions.append(applyoutputformatting(genout[0]["generated_text"])) else: koboldai_vars.actions.append_options([applyoutputformatting(x["generated_text"]) for x in genout]) genout = [{"generated_text": x['text']} for x in koboldai_vars.actions.get_current_options()] diff --git a/koboldai_settings.py b/koboldai_settings.py index e057698e..1bc91da7 100644 --- a/koboldai_settings.py +++ b/koboldai_settings.py @@ -202,7 +202,16 @@ class koboldai_vars(object): # TODO: This might be ineffecient, should we cache some of this? return [[token, self.tokenizer.decode(token)] for token in encoded] - def calc_ai_text(self, submitted_text="", return_text=False, send_context=True): + def calc_ai_text(self, submitted_text=None, return_text=False, send_context=True, allowed_wi_entries=None, allowed_wi_folders=None): + """Compute the context that would be submitted to the AI. + + submitted_text: Optional override to the player-submitted text. + return_text: If True, return the context as a string. Otherwise, return a tuple consisting of: + (tokens, used_tokens, used_tokens+self.genamt, set(used_world_info)) + send_context: If True, the context is prepared for submission to the AI (by marking used world info and setting self.context + allowed_wi_entries: If not None, only world info entries with uids in the given set are allowed to be used. + allowed_wi_folders: If not None, only world info folders with uids in the given set are allowed to be used. + """ #start_time = time.time() if self.tokenizer is None: if return_text: @@ -227,6 +236,17 @@ class koboldai_vars(object): if send_context: self.worldinfo_v2.reset_used_in_game() + def allow_wi(wi): + """Return True if adding this world info entry is allowed.""" + wi_uid = wi['uid'] + if wi_uid in used_world_info: + return False + if allowed_wi_entries is not None and wi_uid not in allowed_wi_entries: + return False + if allowed_wi_folders is not None and wi['folder'] not in allowed_wi_folders: + return False + return True + ######################################### Add memory ######################################################## memory_text = self.memory if memory_text != "": @@ -256,7 +276,7 @@ class koboldai_vars(object): ######################################### Constant World Info ######################################################## #Add constant world info entries to memory for wi in self.worldinfo_v2: - if wi['constant']: + if wi['constant'] and allow_wi(wi): wi_length = len(self.tokenizer.encode(wi['content'])) if used_tokens + wi_length <= token_budget: used_tokens+=wi_length @@ -287,9 +307,8 @@ class koboldai_vars(object): if self.is_chat_v2(): action_text_split[-1][0] = action_text_split[-1][0].strip() + "\n" - ######################################### Prompt ######################################################## - #Add prompt lenght/text if we're set to always use prompt + #Add prompt length/text if we're set to always use prompt if self.useprompt: prompt_length = 0 prompt_data = [] @@ -310,7 +329,7 @@ class koboldai_vars(object): self.prompt_in_ai = True #Find World Info entries in prompt for wi in self.worldinfo_v2: - if wi['uid'] not in used_world_info: + if allow_wi(wi): #Check to see if we have the keys/secondary keys in the text so far match = False for key in wi['key']: @@ -359,19 +378,24 @@ class koboldai_vars(object): actions_seen = [] #Used to track how many actions we've seen so we can insert author's note in the appropriate place as well as WI depth stop inserted_author_note = False + sentences_seen = 0 for i in range(len(action_text_split)-1, -1, -1): if action_text_split[i][3]: #We've hit an item we've already included. Stop break + sentences_seen += 1 + + #Add our author's note if we've hit andepth + if not inserted_author_note and len(actions_seen) >= self.andepth and sentences_seen > self.andepth and self.authornote != "": + game_context.insert(0, {"type": "authors_note", "text": authors_note_text, "tokens": authors_note_data, "attention_multiplier": self.an_attn_bias}) + inserted_author_note = True + + # Add to actions_seen after potentially inserting the author note, since we want to insert the author note + # after the sentence that pushes our action count above the author note depth. for action in action_text_split[i][1]: if action not in actions_seen: actions_seen.append(action) - #Add our author's note if we've hit andepth - if not inserted_author_note and len(actions_seen) >= self.andepth and self.authornote != "": - game_context.insert(0, {"type": "authors_note", "text": authors_note_text, "tokens": authors_note_data, "attention_multiplier": self.an_attn_bias}) - inserted_author_note = True - action_data = [[x, self.tokenizer.decode(x)] for x in self.tokenizer.encode(action_text_split[i][0])] length = len(action_data) @@ -391,13 +415,11 @@ class koboldai_vars(object): "tokens": action_data, "action_ids": action_text_split[i][1] }) - #wi_search = re.sub("[^A-Za-z\ 0-9\'\"]", "", action_text_split[i][0]) wi_search = action_text_split[i][0] - #Now we need to check for used world info entries for wi in self.worldinfo_v2: - if wi['uid'] not in used_world_info: + if allow_wi(wi): #Check to see if we have the keys/secondary keys in the text so far match = False for key in wi['key']: @@ -411,7 +433,7 @@ class koboldai_vars(object): match=True break if method == 1: - if len(actions_seen) > self.widepth: + if len(actions_seen) > self.widepth and sentences_seen > self.widepth: match = False if match: wi_length = len(self.tokenizer.encode(wi['content'])) @@ -1213,6 +1235,9 @@ class KoboldStoryRegister(object): self.actions = {} #keys = "Selected Text", "Wi_highlighted_text", "Options", "Selected Text Length", "Probabilities". #Options being a list of dict with keys of "text", "Pinned", "Previous Selection", "Edited", "Probabilities" self.action_count = -1 + # The id of the last submission action, or 0 if the last append was not a submission + self.submission_id = 0 + self.sentence_re = re.compile(r"[^.!?]*[.!?]+\"?\s*", re.S) self.story_settings = story_settings self.tts_model = None self.make_audio_thread = None @@ -1375,7 +1400,15 @@ class KoboldStoryRegister(object): self.set_game_saved() self.gen_all_audio() + def append_submission(self, text, action_id_offset=0, recalc=True): + self._append(text, action_id_offset=action_id_offset, recalc=recalc) + self.submission_id = self.action_count + action_id_offset + def append(self, text, action_id_offset=0, recalc=True): + self._append(text, action_id_offset=action_id_offset, recalc=recalc) + self.submission_id = 0 + + def _append(self, text, action_id_offset=0, recalc=True): if self.koboldai_vars.remove_double_space: while " " in text: text = text.replace(" ", " ") @@ -1667,7 +1700,11 @@ class KoboldStoryRegister(object): self.actions[action_id]["Options"][option_number]['Probabilities'].append(probabilities) process_variable_changes(self.socketio, "story", 'actions', {"id": action_id, 'action': self.actions[action_id]}, None) - def to_sentences(self, submitted_text=""): + def to_sentences(self, submitted_text=None): + """Return a list of the actions split into sentences. + submitted_text: Optional additional text to append to the actions, the text just submitted by the player. + returns: List of [sentence text, actions used, token length, included in AI context] + """ #start_time = time.time() #we're going to split our actions by sentence for better context. We'll add in which actions the sentence covers. Prompt will be added at a -1 ID actions = {i: self.actions[i]['Selected Text'] for i in self.actions} @@ -1675,16 +1712,29 @@ class KoboldStoryRegister(object): actions[-1] = "" else: actions[-1] = self.story_settings.prompt - if submitted_text != "": - actions[self.action_count+1] = submitted_text - action_text = self.__str__() - action_text = "{}{}{}".format("" if self.story_settings is None else self.story_settings.prompt, action_text, submitted_text) + # During generation, the action with id action_count+1 holds streamed tokens, and action_count holds the last + # submitted text. Outside of generation, action_count+1 is missing/empty. Strip streaming tokens and + # replace submitted if specified. + if self.action_count+1 in self.actions: + actions[self.action_count+1] = "" + if submitted_text: + if self.submission_id: + # Text has been submitted + actions[self.submission_id] = submitted_text + elif self.action_count == -1: + # The only submission is the prompt + actions[-1] = submitted_text + else: + # Add submitted_text to the end + actions[self.action_count + 1] = submitted_text + action_text = "".join(txt for _, txt in sorted(actions.items())) ###########action_text_split = [sentence, actions used in sentence, token length, included in AI context]################ - action_text_split = [[x, [], 0, False] for x in re.findall(".*?[.!?]\s+", action_text, re.S)] + action_text_split = [[x, [], 0, False] for x in self.sentence_re.findall(action_text)] #The above line can trim out the last sentence if it's incomplete. Let's check for that and add it back in - if len("".join([x[0] for x in action_text_split])) < len(action_text): - action_text_split.append([action_text[len("".join([x[0] for x in action_text_split])):], [], 0, False]) - #The last action shouldn't have the extra space from the sentence splitting, so let's remove it + text_len = sum(len(x[0]) for x in action_text_split) + if text_len < len(action_text): + action_text_split.append([action_text[text_len:], [], 0, False]) + #If we don't have any actions at this point, just return. if len(action_text_split) == 0: return [] @@ -1700,7 +1750,9 @@ class KoboldStoryRegister(object): advance_sentence = True if Action_Position[0] not in action_text_split[Sentence_Position[0]][1]: #Since this action is in the sentence, add it to the list if it's not already there - action_text_split[Sentence_Position[0]][1].append(Action_Position[0]) + #Only add actions with non-empty text + if len(actions[Action_Position[0]]) > 0: + action_text_split[Sentence_Position[0]][1].append(Action_Position[0]) #Fix the text length leftovers first since they interact with each other if not advance_action: Action_Position[1] -= Sentence_Position[1] @@ -1718,7 +1770,6 @@ class KoboldStoryRegister(object): break Sentence_Position[1] = len(action_text_split[Sentence_Position[0]][0]) #OK, action_text_split now contains a list of [sentence including trailing space if needed, [action IDs that sentence includes]] - #logger.debug("to_sentences: {}s".format(time.time()-start_time)) return action_text_split def gen_audio(self, action_id=None, overwrite=True): From 44385730f683c3b087d502c186a558a12f0ad9cf Mon Sep 17 00:00:00 2001 From: somebody Date: Thu, 24 Nov 2022 23:17:35 -0600 Subject: [PATCH 2/9] Allow logging to server log from show_error_notification --- aiserver.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/aiserver.py b/aiserver.py index 7221a3da..d3f0c139 100644 --- a/aiserver.py +++ b/aiserver.py @@ -768,7 +768,10 @@ api_v1 = KoboldAPISpec( tags=tags, ) -def show_error_notification(title: str, text: str) -> None: +def show_error_notification(title: str, text: str, do_log: bool = False) -> None: + if do_log: + logger.error(f"{title}: {text}") + if has_request_context(): socketio.emit("show_error_notification", {"title": title, "text": text}, broadcast=True, room="UI_2") else: From 3f0511d7bdd69c995e72643698a2f67de4a88d0f Mon Sep 17 00:00:00 2001 From: somebody Date: Thu, 24 Nov 2022 23:18:24 -0600 Subject: [PATCH 3/9] Don't add commas to join prompts when one is empty --- aiserver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiserver.py b/aiserver.py index d3f0c139..301df9cd 100644 --- a/aiserver.py +++ b/aiserver.py @@ -9372,7 +9372,7 @@ def text2img_api(prompt, #"override_settings": {}, #"sampler_index": "Euler" final_imgen_params = { - "prompt": "{}, {}".format(prompt, art_guide), + "prompt": ", ".join(filter(bool, [prompt, art_guide])), "n_iter": 1, "width": 512, "height": 512, From a48d31c746313b74a4b9739fb20b482ce8b93fc3 Mon Sep 17 00:00:00 2001 From: somebody Date: Thu, 24 Nov 2022 23:19:07 -0600 Subject: [PATCH 4/9] Changes on sd webui api stuff - Strip url of trailing slashes, this caused issues when your api url ended with a slash (such as the default! at least i think) - Notify the user if they didn't call sd webui with the --api argument, which is required to use the api - Don't try to write to non-existant art directory, just send b64 image like the diffusers function does it. In the future we probably want to save these, but it would be in standardized manner most likely tied to story files somehow --- aiserver.py | 48 ++++++++++++++++++++++++------------------------ 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/aiserver.py b/aiserver.py index 301df9cd..6fa2307e 100644 --- a/aiserver.py +++ b/aiserver.py @@ -9381,33 +9381,33 @@ def text2img_api(prompt, "negative_prompt": "{}".format(koboldai_vars.img_gen_negative_prompt), "sampler_index": "Euler a" } - apiaddress = '{}/sdapi/v1/txt2img'.format(koboldai_vars.img_gen_api_url) + apiaddress = '{}/sdapi/v1/txt2img'.format(koboldai_vars.img_gen_api_url.rstrip("/")) payload_json = json.dumps(final_imgen_params) logger.debug(final_imgen_params) - #print("payload_json contains " + payload_json) - submit_req = requests.post(url=f'{apiaddress}', data=payload_json).json() - if submit_req: - results = submit_req - for i in results['images']: - final_src_img = Image.open(BytesIO(base64.b64decode(i.split(",",1)[0]))) - buffer = BytesIO() - final_src_img.save(buffer, format="Webp", quality=95) - b64img = base64.b64encode(buffer.getvalue()).decode("utf8") - base64_bytes = b64img.encode('utf-8') - img_bytes = base64.b64decode(base64_bytes) - img = Image.open(BytesIO(img_bytes)) - dt_string = datetime.datetime.now().strftime("%H%M%S%d%m%Y") - final_filename = "stories/art/{}_{}".format(dt_string,filename) - pnginfo = PngImagePlugin.PngInfo() - prompttext = results.get('info').split("\",")[0].split("\"")[3] - pnginfo.add_text("parameters","prompttext") - img.save(final_filename, pnginfo=pnginfo) - logger.debug("Saved Image") - koboldai_vars.generating_image = False - return(b64img) - else: + submit_req = requests.post(url=apiaddress, data=payload_json) + + if submit_req.status_code == 404: + show_error_notification( + "SD Web API Failure", + f"The SD Web UI was not called with --api. Unable to connect.", + do_log=True + ) + return None + elif not submit_req.ok: + show_error_notification("SD Web API Failure", f"HTTP Code {submit_req.status_code} -- See console for details") + logger.error(f"SD Web API Failure: HTTP Code {submit_req.status_code}, Body:\n{submit_req.text}") koboldai_vars.generating_image = False - logger.error(submit_req.text) + return None + + results = submit_req.json() + + try: + base64_image = results["images"][0] + except (IndexError, KeyError): + show_error_notification("SD Web API Failure", "SD Web API returned no images", do_log=True) + return None + + return base64_image #@logger.catch def get_items_locations_from_text(text): From ec20b3193143d8c5c47b9b78fde9a4e33eea5560 Mon Sep 17 00:00:00 2001 From: somebody Date: Thu, 24 Nov 2022 23:34:48 -0600 Subject: [PATCH 5/9] Remove reset switch for sd web api Automatically recover from errors --- aiserver.py | 7 +++++-- gensettings.py | 17 ----------------- koboldai_settings.py | 3 ++- 3 files changed, 7 insertions(+), 20 deletions(-) diff --git a/aiserver.py b/aiserver.py index 6fa2307e..b3e0710b 100644 --- a/aiserver.py +++ b/aiserver.py @@ -9384,7 +9384,11 @@ def text2img_api(prompt, apiaddress = '{}/sdapi/v1/txt2img'.format(koboldai_vars.img_gen_api_url.rstrip("/")) payload_json = json.dumps(final_imgen_params) logger.debug(final_imgen_params) - submit_req = requests.post(url=apiaddress, data=payload_json) + + try: + submit_req = requests.post(url=apiaddress, data=payload_json) + finally: + koboldai_vars.generating_image = False if submit_req.status_code == 404: show_error_notification( @@ -9396,7 +9400,6 @@ def text2img_api(prompt, elif not submit_req.ok: show_error_notification("SD Web API Failure", f"HTTP Code {submit_req.status_code} -- See console for details") logger.error(f"SD Web API Failure: HTTP Code {submit_req.status_code}, Body:\n{submit_req.text}") - koboldai_vars.generating_image = False return None results = submit_req.json() diff --git a/gensettings.py b/gensettings.py index 7ca496c7..6af24f21 100644 --- a/gensettings.py +++ b/gensettings.py @@ -670,23 +670,6 @@ gensettingstf = [ "ui_level": 2 }, { - "UI_V2_Only": True, - "uitype": "toggle", - "unit": "bool", - "label": "Reset Image Generating", - "id": "generating_image", - "min": 0, - "max": 1, - "step": 1, - "default": 0, - "tooltip": "Use to reset image gen flag in case of error.", - "menu_path": "Interface", - "sub_path": "Images", - "classname": "system", - "name": "generating_image", - "ui_level": 2 - }, - { "UI_V2_Only": True, "uitype": "toggle", "unit": "bool", diff --git a/koboldai_settings.py b/koboldai_settings.py index e057698e..a28fbcba 100644 --- a/koboldai_settings.py +++ b/koboldai_settings.py @@ -1060,7 +1060,8 @@ class system_settings(settings): no_save_variables = ['socketio', 'lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'sp', 'sp_length', '_horde_pid', 'horde_share', 'aibusy', 'serverstarted', 'inference_config', 'image_pipeline', 'summarizer', - 'summary_tokenizer', 'use_colab_tpu', 'noai', 'disable_set_aibusy', 'cloudflare_link', 'tts_model'] + 'summary_tokenizer', 'use_colab_tpu', 'noai', 'disable_set_aibusy', 'cloudflare_link', 'tts_model', + 'generating_image'] settings_name = "system" def __init__(self, socketio, koboldai_var): self.socketio = socketio From 9ca741b09a9abc92f5f3264d4ed4471407f8faf6 Mon Sep 17 00:00:00 2001 From: somebody Date: Thu, 24 Nov 2022 23:36:04 -0600 Subject: [PATCH 6/9] Add loading thing for image gen --- static/koboldai.css | 33 ++++++++++++++++++++++++++++++++- static/koboldai.js | 14 +++++++++++--- templates/settings flyout.html | 8 ++++++-- 3 files changed, 49 insertions(+), 6 deletions(-) diff --git a/static/koboldai.css b/static/koboldai.css index be6fdf74..a534b688 100644 --- a/static/koboldai.css +++ b/static/koboldai.css @@ -2988,8 +2988,39 @@ input[type='range'] { pointer-events:none; } +#action\ image { + width: 100%; + aspect-ratio: 1/1; + position: relative; +} + +#image-loading { + display: flex; + align-items: center; + justify-content: center; + position: absolute; + top: 0px; + left: 0px; + width: 100%; + height: 100% !important; + aspect-ratio: 1/1; + background-color: rgba(0, 0, 0, 0.6); + + /* Defer pointer events to the image in case a user wants to quickly + save / view an image after asking for a new image */ + pointer-events: none; +} + +.spinner { + animation: spin 1s linear infinite; + opacity: 0.7; +} + +@keyframes spin { 100% { transform: rotate(360deg) } } + + .action_image { - width: var(--flyout_menu_width); + width: 100%; } /*Tooltip based on attribute diff --git a/static/koboldai.js b/static/koboldai.js index c2805a15..4a8b42d8 100644 --- a/static/koboldai.js +++ b/static/koboldai.js @@ -790,9 +790,12 @@ function var_changed(data) { //Special Case for story picture } else if (data.classname == "story" && data.name == "picture") { image_area = document.getElementById("action image"); - while (image_area.firstChild) { - image_area.removeChild(image_area.firstChild); - } + + let maybeImage = image_area.getElementsByClassName("action_image")[0]; + if (maybeImage) maybeImage.remove(); + + $el("#image-loading").classList.add("hidden"); + if (data.value != "") { var image = new Image(); image.src = 'data:image/png;base64,'+data.value; @@ -5991,6 +5994,11 @@ $el("#aidgpromptnum").addEventListener("keydown", function(event) { event.preventDefault(); }); +$el("#generate-image-button").addEventListener("click", function() { + $el("#image-loading").classList.remove("hidden"); + socket.emit("generate_image", {}); +}); + /* -- Shiny New Chat -- */ function addMessage(author, content, actionId, afterMsgEl=null, time=null) { if (!time) time = Number(new Date()); diff --git a/templates/settings flyout.html b/templates/settings flyout.html index 3ddf3c9a..8dcfcf52 100644 --- a/templates/settings flyout.html +++ b/templates/settings flyout.html @@ -112,8 +112,12 @@ Download debug dump
- -
+ +
+ +