from __future__ import annotations from dataclasses import dataclass import difflib import importlib import os, re, time, threading, json, pickle, base64, copy, tqdm, datetime, sys import shutil from typing import List, Union from io import BytesIO from flask import has_request_context, session from flask_socketio import join_room, leave_room from collections import OrderedDict import multiprocessing from logger import logger import torch import numpy as np import random import inspect serverstarted = False queue = None multi_story = False if importlib.util.find_spec("tortoise") is not None: from tortoise import api from tortoise.utils.audio import load_voices def clean_var_for_emit(value): if isinstance(value, KoboldStoryRegister) or isinstance(value, KoboldWorldInfo): return value.to_json() elif isinstance(value, set): return list(value) elif isinstance(value, datetime.datetime): return str(value) else: return value def process_variable_changes(socketio, classname, name, value, old_value, debug_message=None): global multi_story if serverstarted and name != "serverstarted": transmit_time = str(datetime.datetime.now()) if debug_message is not None: print("{} {}: {} changed from {} to {}".format(debug_message, classname, name, old_value, value)) if value != old_value: #Get which room we'll send the messages to if multi_story: if classname != 'story': room = 'UI_2' else: if has_request_context(): room = 'default' if 'story' not in session else session['story'] else: logger.error("We tried to access the story register outside of an http context. Will not work in multi-story mode") return else: room = "UI_2" #logger.debug("sending data to room (multi_story={},classname={}): {}".format(multi_story, classname, room)) #Special Case for KoboldStoryRegister if isinstance(value, KoboldStoryRegister): #To speed up loading time we will only transmit the last 100 actions to the UI, then rely on scrolling triggers to load more as needed if not has_request_context(): if queue is not None: #logger.debug("Had to use queue") queue.put(["var_changed", {"classname": "actions", "name": "Action Count", "old_value": None, "value":value.action_count, "transmit_time": transmit_time}, {"broadcast":True, "room":room}]) data_to_send = [] for i in list(value.actions)[-100:]: data_to_send.append({"id": i, "action": value.actions[i]}) queue.put(["var_changed", {"classname": "story", "name": "actions", "old_value": None, "value":data_to_send, "transmit_time": transmit_time}, {"broadcast":True, "room":room}]) else: if socketio is not None: socketio.emit("var_changed", {"classname": "actions", "name": "Action Count", "old_value": None, "value":value.action_count, "transmit_time": transmit_time}, broadcast=True, room=room) data_to_send = [] for i in list(value.actions)[-100:]: data_to_send.append({"id": i, "action": value.actions[i]}) if socketio is not None: socketio.emit("var_changed", {"classname": "story", "name": "actions", "old_value": None, "value": data_to_send, "transmit_time": transmit_time}, broadcast=True, room=room) elif isinstance(value, KoboldWorldInfo): value.send_to_ui() else: #If we got a variable change from a thread other than what the app is run it, eventlet seems to block and no further messages are sent. Instead, we'll rely the message to the app and have the main thread send it if not has_request_context(): data = ["var_changed", {"classname": classname, "name": name, "old_value": clean_var_for_emit(old_value), "value": clean_var_for_emit(value), "transmit_time": transmit_time}, {"include_self":True, "broadcast":True, "room":room}] if queue is not None: #logger.debug("Had to use queue") queue.put(data) else: if socketio is not None: socketio.emit("var_changed", {"classname": classname, "name": name, "old_value": clean_var_for_emit(old_value), "value": clean_var_for_emit(value), "transmit_time": transmit_time}, include_self=True, broadcast=True, room=room) class koboldai_vars(object): def __init__(self, socketio): self._model_settings = model_settings(socketio, self) self._user_settings = user_settings(socketio) self._system_settings = system_settings(socketio, self) self._story_settings = {'default': story_settings(socketio, self)} self._undefined_settings = undefined_settings() self._socketio = socketio self.tokenizer = None def get_story_name(self): global multi_story if multi_story: if has_request_context(): story = 'default' if 'story' not in session else session['story'] else: logger.error("We tried to access the story register outside of an http context. Will not work in multi-story mode") assert multi_story and has_request_context(), "Tried to access story data outside context in multi_story mode" else: story = "default" return story def to_json(self, classname): if classname == 'story_settings': return self._story_settings[self.get_story_name()].to_json() # data = {} # for story in self._story_settings: # data[story] = json.loads(self._story_settings[story].to_json()) # return json.dumps(data) return self.__dict__["_{}".format(classname)].to_json() def load_story(self, story_name, json_data): #Story name here is intended for multiple users on multiple stories. Now always uses default #If we can figure out a way to get flask sessions into/through the lua bridge we could re-enable global multi_story original_story_name = story_name if not multi_story: story_name = 'default' #Leave the old room and join the new one logger.debug("Leaving room {}".format(session['story'])) leave_room(session['story']) logger.debug("Joining room {}".format(story_name)) join_room(story_name) session['story'] = story_name logger.debug("Sending story reset") self._story_settings[story_name]._socketio.emit("reset_story", {}, broadcast=True, room=story_name) if story_name in self._story_settings: self._story_settings[story_name].no_save = True self._story_settings[story_name].worldinfo_v2.reset() self._story_settings[story_name].from_json(json_data) logger.debug("Calcing AI text after load story") ignore = self.calc_ai_text() self._story_settings[story_name].no_save = False else: #self._story_settings[story_name].no_save = True self.create_story(story_name, json_data=json_data) logger.debug("Calcing AI text after create story") ignore = self.calc_ai_text() #self._story_settings[story_name].no_save = False self._system_settings.story_loads[original_story_name] = datetime.datetime.now().strftime("%m/%d/%Y, %H:%M:%S") with open("settings/system_settings.v2_settings", "w") as settings_file: settings_file.write(self._system_settings.to_json()) def save_story(self): logger.debug("Saving story from koboldai_vars.save_story()") self._story_settings[self.get_story_name()].save_story() def download_story(self): return self._story_settings[self.get_story_name()].to_json() def save_revision(self): self._story_settings[self.get_story_name()].save_revision() def create_story(self, story_name, json_data=None): #Story name here is intended for multiple users on multiple stories. Now always uses default #If we can figure out a way to get flask sessions into/through the lua bridge we could re-enable #story_name = 'default' if not multi_story: story_name = 'default' self._story_settings[story_name] = story_settings(self._socketio, self) #self._story_settings[story_name].reset() if json_data is not None: self.load_story(story_name, json_data) else: #Leave the old room and join the new one logger.debug("Leaving room {}".format(session['story'])) leave_room(session['story']) logger.debug("Joining room {}".format(story_name)) join_room(story_name) session['story'] = story_name logger.debug("Sending story reset") self._story_settings[story_name]._socketio.emit("reset_story", {}, broadcast=True, room=story_name) self._story_settings[story_name].send_to_ui() session['story'] = story_name def story_list(self): return [x for x in self._story_settings] def send_to_ui(self): self._model_settings.send_to_ui() self._user_settings.send_to_ui() self._system_settings.send_to_ui() self._story_settings[self.get_story_name()].send_to_ui() def reset_model(self): self._model_settings.reset_for_model_load() def is_chat_v2(self): return self.chat_style > 0 and self.chatmode def get_token_representation(self, text: Union[str, list, None]) -> list: if not self.tokenizer or not text: return [] if isinstance(text, str): encoded = self.tokenizer.encode(text) else: encoded = text # 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=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: return "" return [], 0, 0+self.genamt, [] if self.alt_gen: method = 2 else: method = 1 #Context and Game Context are lists of chunks of text that will go to the AI. Game Context will be appended after context when we're done context = [] game_context = [] token_budget = self.max_length - self.genamt used_world_info = [] used_tokens = 0 if self.sp_length is None else self.sp_length + len(self.tokenizer._koboldai_header) if self.sp_length > 0: context.append({"type": "soft_prompt", "text": f"<{self.sp_length} tokens of Soft Prompt.>", "tokens": [[-1, ""]] * self.sp_length}) 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 wi["type"] == "commentator" and not (allowed_wi_entries and wi_uid in allowed_wi_entries): 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 Genres # if self.genres: # Erebus, Nerys, Janeway, Picard (probably) genre_template = "[Genre: %s]" model_name = self.model.lower() if "skein" in model_name or "adventure" in model_name: genre_template = "[Themes: %s]" elif "shinen" in model_name: genre_template = "[Theme: %s]" genre_text = genre_template % (", ".join(self.genres)) genre_tokens = self.tokenizer.encode(genre_text) genre_data = [[x, self.tokenizer.decode(x)] for x in genre_tokens] context.append({ "type": "genre", "text": genre_text, "tokens": genre_data, }) used_tokens += len(genre_tokens) ######################################### Add memory ######################################################## max_memory_length = int(token_budget * self.max_memory_fraction) memory_text = self.memory if memory_text != "": if memory_text[-1] not in [" ", '\n']: memory_text += " " memory_tokens = self.tokenizer.encode(memory_text) else: memory_tokens = [] if len(memory_tokens) > max_memory_length: memory_tokens = memory_tokens[:max_memory_length] memory_length = max_memory_length memory_data = [[x, self.tokenizer.decode(x)] for x in memory_tokens] #We have so much memory that we've run out of budget without going over the max_memory_length. Just stop if len(memory_tokens) > token_budget: return [], 0, 0+self.genamt, [] #Actually Add Memory if len(memory_tokens) != 0: context.append({"type": "memory", "text": "".join([x[1] for x in memory_data]), "tokens": memory_data, "attention_multiplier": self.memory_attn_bias}) used_tokens += len(memory_tokens) ######################################### Constant World Info ######################################################## #Add constant world info entries to memory for wi in self.worldinfo_v2: 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 used_world_info.append(wi['uid']) if send_context: self.worldinfo_v2.set_world_info_used(wi['uid']) wi_text = wi['content']+" " if wi['content'] != "" and wi['content'][-1] not in [" ", "\n"] else wi['content'] wi_tokens = self.tokenizer.encode(wi_text) context.append({ "type": "world_info", "text": wi_text, "uid": wi['uid'], "tokens": [[x, self.tokenizer.decode(x)] for x in wi_tokens], }) used_tokens += len(wi_tokens) ######################################### Get Action Text by Sentence ######################################################## action_text_split = self.actions.to_sentences(submitted_text=submitted_text) # Always add newlines on chat v2 if action_text_split and self.is_chat_v2(): action_text_split[-1][0] = action_text_split[-1][0].strip() + "\n" ######################################### Prompt ######################################################## #Add prompt length/text if we're set to always use prompt if self.useprompt: prompt_length = 0 prompt_data = [] for item in reversed(action_text_split): if -1 in item[1]: tokenized_data = [[x, self.tokenizer.decode(x)] for x in self.tokenizer.encode(item[0])] item[2] = len(tokenized_data) if prompt_length + item[2] <= self.max_prompt_length: prompt_length += item[2] item[3] = True prompt_data = tokenized_data + prompt_data prompt_text = self.tokenizer.decode([x[0] for x in prompt_data]) #wi_search = re.sub("[^A-Za-z\ 0-9\'\"]", "", prompt_text) wi_search = prompt_text if prompt_length + used_tokens <= token_budget: used_tokens += prompt_length #We'll add the prompt text AFTER we go through the game text as the world info needs to come first if we're in method 1 rather than method 2 self.prompt_in_ai = True #Find World Info entries in prompt for wi in self.worldinfo_v2: 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']: if key in wi_search: match = True break if wi['selective'] and match: match = False for key in wi['keysecondary']: if key in wi_search: match=True break if match: wi_length = len(self.tokenizer.encode(wi['content'])) if used_tokens+wi_length <= token_budget: used_tokens+=wi_length used_world_info.append(wi['uid']) wi_text = wi['content']+" " if wi['content'] != "" and wi['content'][-1] not in [" ", "\n"] else wi['content'] wi_tokens = self.tokenizer.encode(wi_text) context.append({ "type": "world_info", "text": wi_text, "uid": wi['uid'], "tokens": [[x, self.tokenizer.decode(x)] for x in wi_tokens], }) used_tokens += len(wi_tokens) if send_context: self.worldinfo_v2.set_world_info_used(wi['uid']) else: self.prompt_in_ai = False ######################################### Setup Author's Note Data ######################################################## authors_note_text = self.authornotetemplate.replace("<|>", self.authornote) if len(authors_note_text) > 0 and authors_note_text[-1] not in [" ", "\n"]: authors_note_text += " " authors_note_data = [[x, self.tokenizer.decode(x)] for x in self.tokenizer.encode(authors_note_text)] if used_tokens + len(authors_note_data) <= token_budget: used_tokens += len(authors_note_data) ######################################### Actions ######################################################## #Start going through the actions backwards, adding it to the text if it fits and look for world info entries used_all_tokens = False 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) action_data = [[x, self.tokenizer.decode(x)] for x in self.tokenizer.encode(action_text_split[i][0])] length = len(action_data) if length+used_tokens <= token_budget: used_tokens += length action_text_split[i][3] = True action_type = "action" if action_text_split[i][1] == [self.actions.action_count+1]: action_type = "submit" elif -1 in action_text_split[i][1]: action_type = "prompt" game_context.insert(0, { "type": action_type, "text": action_text_split[i][0], "tokens": action_data, "action_ids": action_text_split[i][1] }) wi_search = action_text_split[i][0] #Now we need to check for used world info entries for wi in self.worldinfo_v2: 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']: if key in wi_search: match = True break if wi['selective'] and match: match = False for key in wi['keysecondary']: if key in wi_search: match=True break if method == 1: if len(actions_seen) > self.widepth and sentences_seen > self.widepth: match = False if match: wi_length = len(self.tokenizer.encode(wi['content'])) if used_tokens+wi_length <= token_budget: used_tokens+=wi_length used_world_info.append(wi['uid']) wi_text = wi['content']+" " if wi['content'] != "" and wi['content'][-1] not in [" ", "\n"] else wi['content'] wi_tokens = self.tokenizer.encode(wi_text) if method == 1: context.append({ "type": "world_info", "text": wi_text, "uid": wi['uid'], "tokens": [[x, self.tokenizer.decode(x)] for x in wi_tokens], }) else: #for method 2 we add the game text before the current action game_context.insert(0, { "type": "world_info", "text": wi_text, "uid": wi['uid'], "tokens": [[x, self.tokenizer.decode(x)] for x in wi_tokens], }) used_tokens += len(wi_tokens) if send_context: self.worldinfo_v2.set_world_info_used(wi['uid']) else: used_all_tokens = True break ######################################### Verify Author's Note Data in AI Text ######################################################## #if we don't have enough actions to get to author's note depth then we just add it right before the game text if not inserted_author_note and self.authornote != "": game_context.insert(0, {"type": "authors_note", "text": authors_note_text, "tokens": authors_note_data, "attention_multiplier": self.an_attn_bias}) ######################################### Add our prompt data ######################################################## if self.useprompt and len(prompt_data) != 0: context.append({"type": "prompt", "text": prompt_text, "tokens": prompt_data}) context += game_context if len(context) == 0: tokens = [] else: tokens = [] for item in context: if item['type'] != 'soft_prompt': tokens.extend([x[0] for x in item['tokens']]) if send_context: self.context = context #logger.debug("Calc_AI_text: {}s".format(time.time()-start_time)) logger.debug("Token Budget: {}. Used Tokens: {}".format(token_budget, used_tokens)) if return_text: return "".join([x['text'] for x in context]) return tokens, used_tokens, used_tokens+self.genamt, set(used_world_info) def is_model_torch(self) -> bool: if self.use_colab_tpu: return False if self.model in ["Colab", "API", "CLUSTER", "ReadOnly", "OAI"]: return False return True def assign_world_info_to_actions(self, *args, **kwargs): self._story_settings[self.get_story_name()].assign_world_info_to_actions(*args, **kwargs) def reset_for_model_load(self): self._model_settings.reset_for_model_load() def __setattr__(self, name, value): if name[0] == "_" or name == "tokenizer": super().__setattr__(name, value) if name[0] != "_": #Send it to the corrent _setting class if name in self._model_settings.__dict__: setattr(self._model_settings, name, value) elif name in self._user_settings.__dict__: setattr(self._user_settings, name, value) elif name in self._system_settings.__dict__: setattr(self._system_settings, name, value) elif name in self._story_settings[self.get_story_name()].__dict__: setattr(self._story_settings[self.get_story_name()], name, value) else: setattr(self._undefined_settings, name, value) def __getattr__(self, name): if name in self.__dict__: return getattr(self, name) elif name in self._model_settings.__dict__ or hasattr(self._model_settings, name): return getattr(self._model_settings, name) elif name in self._user_settings.__dict__ or hasattr(self._user_settings, name): return getattr(self._user_settings, name) elif name in self._system_settings.__dict__ or hasattr(self._system_settings, name): return getattr(self._system_settings, name) elif name in self._story_settings[self.get_story_name()].__dict__ or hasattr(self._story_settings[self.get_story_name()], name): return getattr(self._story_settings[self.get_story_name()], name) else: return getattr(self._undefined_settings, name) class settings(object): def to_json(self): json_data = {'file_version': 2} for (name, value) in vars(self).items(): if name not in self.no_save_variables and name[0] != "_": json_data[name] = value def to_base64(data): if isinstance(data, KoboldStoryRegister): return data.to_json() elif isinstance(data, KoboldWorldInfo): return data.to_json() elif isinstance(data, datetime.datetime): return str(data) output = BytesIO() pickle.dump(data, output) output.seek(0) return "base64:{}".format(base64.encodebytes(output.read()).decode()) return json.dumps(json_data, default=to_base64, indent="\t") def from_json(self, data): if isinstance(data, str): json_data = json.loads(data) else: json_data = data #since loading will trigger the autosave, we need to disable it if 'no_save' in self.__dict__: setattr(self, 'no_save', True) for key, value in json_data.items(): start_time = time.time() if key in self.__dict__ and key not in self.no_save_variables: if key == 'sampler_order': if(len(value) < 7): value = [6] + value if key == 'autosave': autosave = value if isinstance(value, str): if value[:7] == 'base64:': value = pickle.loads(base64.b64decode(value[7:])) #Need to fix the data type of value to match the module if type(getattr(self, key)) == int: setattr(self, key, int(value)) elif type(getattr(self, key)) == float: setattr(self, key, float(value)) elif type(getattr(self, key)) == bool: setattr(self, key, bool(value)) elif type(getattr(self, key)) == str: setattr(self, key, str(value)) elif isinstance(getattr(self, key), KoboldStoryRegister): getattr(self, key).load_json(value) elif isinstance(getattr(self, key), KoboldWorldInfo): getattr(self, key).load_json(value) else: setattr(self, key, value) logger.debug("Loading {} took {}s".format(key, time.time()- start_time)) #check from prompt_wi_highlighted_text since that wasn't always in the V2 save format if 'prompt' in json_data and 'prompt_wi_highlighted_text' not in json_data: self.prompt_wi_highlighted_text[0]['text'] = self.prompt self.assign_world_info_to_actions(action_id=-1) process_variable_changes(self._socketio, "story", 'prompt_wi_highlighted_text', self.prompt_wi_highlighted_text, None) if 'no_save' in self.__dict__: setattr(self, 'no_save', False) def send_to_ui(self): for (name, value) in vars(self).items(): if name not in self.local_only_variables and name[0] != "_": try: process_variable_changes(self._socketio, self.__class__.__name__.replace("_settings", ""), name, value, None) except: print("{} is of type {} and I can't transmit".format(name, type(value))) raise class model_settings(settings): local_only_variables = ['badwordsids', 'apikey', 'default_preset'] no_save_variables = ['modelconfig', 'custmodpth', 'generated_tkns', 'loaded_layers', 'total_layers', 'total_download_chunks', 'downloaded_chunks', 'presets', 'default_preset', 'welcome', 'welcome_default', 'simple_randomness', 'simple_creativity', 'simple_repitition', 'badwordsids', 'uid_presets', 'model', 'model_type', 'lazy_load', 'fp32_model', 'modeldim', 'horde_wait_time', 'horde_queue_position', 'horde_queue_size', 'newlinemode', 'tqdm_progress', 'tqdm_rem_time', '_tqdm'] settings_name = "model" default_settings = {"rep_pen" : 1.1, "rep_pen_slope": 0.7, "rep_pen_range": 1024, "temp": 0.5, "top_p": 0.9, "top_k": 0, "top_a": 0.0, "tfs": 1.0, "typical": 1.0, "sampler_order": [6,0,1,2,3,4,5]} def __init__(self, socketio, koboldai_vars): self._socketio = socketio self.reset_for_model_load() self.model = "" # Model ID string chosen at startup self.model_type = "" # Model Type (Automatically taken from the model config) self.modelconfig = {} # Raw contents of the model's config.json, or empty dictionary if none found self.custmodpth = "" # Filesystem location of custom model to run self.generated_tkns = 0 # If using a backend that supports Lua generation modifiers, how many tokens have already been generated, otherwise 0 self.loaded_layers = 0 # Used in UI 2 to show model loading progress self.total_layers = 0 # Same as above self.total_download_chunks = 0 # tracks how much of the model has downloaded for the UI 2 self.downloaded_chunks = 0 #as above self._tqdm = tqdm.tqdm(total=self.genamt, file=self.ignore_tqdm()) # tqdm agent for generating tokens. This will allow us to calculate the remaining time self.tqdm_progress = 0 # TQDP progress self.tqdm_rem_time = 0 # tqdm calculated reemaining time self.configname = None self.online_model = '' self.welcome_default = """
Please load a model from the left.
If you encounter any issues, please click the Download debug dump link in the Home tab on the left flyout and attach the downloaded file to your error report on Github, Reddit, or Discord. A redacted version (without story text) is available.
""" # Custom Welcome Text self.welcome = self.welcome_default self._koboldai_vars = koboldai_vars self.alt_multi_gen = False def reset_for_model_load(self): self.simple_randomness = 0 #Set first as this affects other outputs self.simple_creativity = 0 #Set first as this affects other outputs self.simple_repitition = 0 #Set first as this affects other outputs self.max_length = 1024 # Maximum number of tokens to submit per action self.ikmax = 3000 # Maximum number of characters to submit to InferKit self.genamt = 80 # Amount of text for each action to generate self.ikgen = 200 # Number of characters for InferKit to generate self.rep_pen = 1.1 # Default generator repetition_penalty self.rep_pen_slope = 0.7 # Default generator repetition penalty slope self.rep_pen_range = 1024 # Default generator repetition penalty range self.temp = 0.5 # Default generator temperature self.top_p = 0.9 # Default generator top_p self.top_k = 0 # Default generator top_k self.top_a = 0.0 # Default generator top-a self.tfs = 1.0 # Default generator tfs (tail-free sampling) self.typical = 1.0 # Default generator typical sampling threshold self.numseqs = 1 # Number of sequences to ask the generator to create self.generated_tkns = 0 # If using a backend that supports Lua generation modifiers, how many tokens have already been generated, otherwise 0 self.badwordsids = [] self.fp32_model = False # Whether or not the most recently loaded HF model was in fp32 format self.modeldim = -1 # Embedding dimension of your model (e.g. it's 4096 for GPT-J-6B and 2560 for GPT-Neo-2.7B) self.sampler_order = [6, 0, 1, 2, 3, 4, 5] self.newlinemode = "n" self.lazy_load = True # Whether or not to use torch_lazy_loader.py for transformers models in order to reduce CPU memory usage self.presets = [] # Holder for presets self.selected_preset = "" self.uid_presets = [] self.default_preset = {} self.horde_wait_time = 0 self.horde_queue_position = 0 self.horde_queue_size = 0 self.use_alt_rep_pen = False #dummy class to eat the tqdm output class ignore_tqdm(object): def write(self, bar): pass def __setattr__(self, name, value): new_variable = name not in self.__dict__ old_value = getattr(self, name, None) super().__setattr__(name, value) #Put variable change actions here if name in ['simple_randomness', 'simple_creativity', 'simple_repitition'] and not new_variable: #We need to disable all of the samplers self.top_k = 0 self.top_a = 0.0 self.tfs = 1.0 self.typical = 1.0 self.rep_pen_range = 1024 self.rep_pen_slope = 0.7 #Now we setup our other settings if self.simple_randomness < 0: self.temp = default_rand_range[1]+(default_rand_range[1]-default_rand_range[0])/100*self.simple_randomness else: self.temp = default_rand_range[1]+(default_rand_range[2]-default_rand_range[1])/100*self.simple_randomness self.top_p = (default_creativity_range[1]-default_creativity_range[0])/200*self.simple_creativity+sum(default_creativity_range)/2 self.rep_pen = (default_rep_range[1]-default_rep_range[0])/200*self.simple_repitition+sum(default_rep_range)/2 if not new_variable and (name == 'max_length' or name == 'genamt'): ignore = self._koboldai_vars.calc_ai_text() #set preset values if name == 'selected_preset' and value != "": logger.info("Changing preset to {}".format(value)) if value in self.uid_presets: for default_key, default_value in self.default_settings.items(): setattr(self, default_key, default_value) for preset_key, preset_value in self.uid_presets[value].items(): if preset_key in self.__dict__: if type(getattr(self, preset_key)) == int: preset_value = int(preset_value) elif type(getattr(self, preset_key)) == float: preset_value = float(preset_value) elif type(getattr(self, preset_key)) == bool: preset_value = bool(preset_value) elif type(getattr(self, preset_key)) == str: preset_value = str(preset_value) if preset_key == "sampler_order": if 6 not in preset_value: preset_value.insert(0, 6) setattr(self, preset_key, preset_value) #Setup TQDP for token generation elif name == "generated_tkns" and '_tqdm' in self.__dict__: if value == 0: self._tqdm.reset(total=self.genamt * (self.numseqs if self.alt_multi_gen else 1) ) self.tqdm_progress = 0 else: self._tqdm.update(value-self._tqdm.n) self.tqdm_progress = int(float(self.generated_tkns)/float(self.genamt * (self.numseqs if self.alt_multi_gen else 1))*100) if self._tqdm.format_dict['rate'] is not None: self.tqdm_rem_time = str(datetime.timedelta(seconds=int(float((self.genamt * (self.numseqs if self.alt_multi_gen else 1))-self.generated_tkns)/self._tqdm.format_dict['rate']))) #Setup TQDP for model loading elif name == "loaded_layers" and '_tqdm' in self.__dict__: if value == 0: self._tqdm.reset(total=self.total_layers) self.tqdm_progress = 0 else: self._tqdm.update(1) self.tqdm_progress = int(float(self.loaded_layers)/float(self.total_layers)*100) if self._tqdm.format_dict['rate'] is not None: self.tqdm_rem_time = str(datetime.timedelta(seconds=int(float(self.total_layers-self.loaded_layers)/self._tqdm.format_dict['rate']))) #Setup TQDP for model downloading elif name == "total_download_chunks" and '_tqdm' in self.__dict__: self._tqdm.reset(total=value) self.tqdm_progress = 0 elif name == "downloaded_chunks" and '_tqdm' in self.__dict__: if value == 0: self._tqdm.reset(total=self.total_download_chunks) self.tqdm_progress = 0 else: self._tqdm.update(value-old_value) if self.total_download_chunks is not None: if self.total_download_chunks==0: self.tqdm_progress = 0 elif float(self.downloaded_chunks) > float(self.total_download_chunks): self.tqdm_progress = 100 else: self.tqdm_progress = round(float(self.downloaded_chunks)/float(self.total_download_chunks)*100, 1) else: self.tqdm_progress = 0 if self._tqdm.format_dict['rate'] is not None: self.tqdm_rem_time = str(datetime.timedelta(seconds=int(float(self.total_download_chunks-self.downloaded_chunks)/self._tqdm.format_dict['rate']))) if name not in self.local_only_variables and name[0] != "_" and not new_variable: process_variable_changes(self._socketio, self.__class__.__name__.replace("_settings", ""), name, value, old_value) class story_settings(settings): local_only_variables = ['tokenizer', 'no_save', 'revisions', 'prompt', 'save_paths'] no_save_variables = ['tokenizer', 'context', 'no_save', 'prompt_in_ai', 'authornote_length', 'prompt_length', 'memory_length', 'save_paths'] settings_name = "story" def __init__(self, socketio, koboldai_vars, tokenizer=None): self._socketio = socketio self.tokenizer = tokenizer self._koboldai_vars = koboldai_vars self.privacy_mode = False self.privacy_password = "" self.story_name = "New Game" # Title of the story self.lastact = "" # The last action received from the user self.submission = "" # Same as above, but after applying input formatting self.lastctx = "" # The last context submitted to the generator self.gamestarted = False # Whether the game has started (disables UI elements) self.gamesaved = True # Whether or not current game is saved self.autosave = False # Whether or not to automatically save after each action self.prompt = "" # Prompt self.prompt_wi_highlighted_text = [{"text": self.prompt, "WI matches": None, "WI Text": ""}] self.memory = "" # Text submitted to memory field self.auto_memory = "" self.authornote = "" # Text submitted to Author's Note field self.authornotetemplate = "[Author's note: <|>]" # Author's note template self.setauthornotetemplate = self.authornotetemplate # Saved author's note template in settings self.andepth = 3 # How far back in history to append author's note self.actions = KoboldStoryRegister(socketio, self, koboldai_vars, tokenizer=tokenizer) # Actions submitted by user and AI self.actions_metadata = {} # List of dictonaries, one dictonary for every action that contains information about the action like alternative options. # Contains at least the same number of items as actions. Back action will remove an item from actions, but not actions_metadata # Dictonary keys are: # Selected Text: (text the user had selected. None when this is a newly generated action) # Alternative Generated Text: {Text, Pinned, Previous Selection, Edited} # self.worldinfo_v2 = KoboldWorldInfo(socketio, self, koboldai_vars, tokenizer=self.tokenizer) self.worldinfo = [] # List of World Info key/value objects self.worldinfo_i = [] # List of World Info key/value objects sans uninitialized entries self.worldinfo_u = {} # Dictionary of World Info UID - key/value pairs self.wifolders_d = {} # Dictionary of World Info folder UID-info pairs self.wifolders_l = [] # List of World Info folder UIDs self.wifolders_u = {} # Dictionary of pairs of folder UID - list of WI UID self.lua_edited = set() # Set of chunk numbers that were edited from a Lua generation modifier self.lua_deleted = set() # Set of chunk numbers that were deleted from a Lua generation modifier self.deletewi = None # Temporary storage for UID to delete self.mode = "play" # Whether the interface is in play, memory, or edit mode self.editln = 0 # Which line was last selected in Edit Mode self.genseqs = [] # Temporary storage for generated sequences self.recentback = False # Whether Back button was recently used without Submitting or Retrying after self.recentrng = None # If a new random game was recently generated without Submitting after, this is the topic used (as a string), otherwise this is None self.recentrngm = None # If a new random game was recently generated without Submitting after, this is the memory used (as a string), otherwise this is None self.useprompt = False # Whether to send the full prompt with every submit action self.chatmode = False self.chatname = "You" self.botname = "Bot" self.adventure = False self.actionmode = 0 self.storymode = 0 self.dynamicscan = False self.recentedit = False self.notes = "" #Notes for the story. Does nothing but save self.biases = {} # should look like {"phrase": [score, completion_threshold]} self.story_id = int.from_bytes(os.urandom(16), 'little', signed=True) # this is a unique ID for the story. We'll use this to ensure when we save that we're saving the same story self.memory_length = 0 self.prompt_length = 0 self.authornote_length = 0 self.max_memory_fraction = 0.5 # Tokens from memory are allowed to use this much of the token budget self.max_prompt_length = 512 self.max_authornote_length = 512 self.prompt_in_ai = False self.context = [] self.last_story_load = None self.revisions = [] self.picture = "" #base64 of the image shown for the story self.picture_prompt = "" #Prompt used to create picture self.substitutions = [ {"target": "--", "substitution": "–", "enabled": False}, {"target": "---", "substitution": "—", "enabled": False}, {"target": "...", "substitution": "…", "enabled": False}, # {"target": "(c)", "substitution": "©", "enabled": False}, # {"target": "(r)", "substitution": "®", "enabled": False}, # {"target": "(tm)", "substitution": "™", "enabled": False}, ] self.gen_audio = False self.prompt_picture_filename = "" self.prompt_picture_prompt = "" # It's important to only use "=" syntax on this to ensure syncing; no # .append() or the like self.genres = [] # bias experiment self.memory_attn_bias = 1 self.an_attn_bias = 1 self.chat_style = 0 # In percent!!! self.commentary_chance = 0 self.commentary_enabled = False self.save_paths = SavePaths(os.path.join("stories", self.story_name or "Untitled")) ################### must be at bottom ######################### self.no_save = False #Temporary disable save (doesn't save with the file) def save_story(self) -> None: if self.no_save: return if not any([self.prompt, self.memory, self.authornote, len(self.actions), len(self.worldinfo_v2)]): return logger.info("Saving") save_name = self.story_name or "Untitled" # Disambiguate stories by adding (n) if needed disambiguator = 0 self.save_paths.base = os.path.join("stories", save_name) while os.path.exists(self.save_paths.base): try: # If the stories share a story id, overwrite the existing one. with open(self.save_paths.story, "r", encoding="utf-8") as file: j = json.load(file) if self.story_id == j["story_id"]: break except FileNotFoundError: logger.error(f"Malformed save file: Missing story.json in {self.save_paths.base}. Populating it with new data.") break disambiguator += 1 self.save_paths.base = os.path.join("stories", save_name + (f" ({disambiguator})" if disambiguator else "")) # Setup the directory structure. for path in self.save_paths.required_paths: try: os.mkdir(path) except FileExistsError: pass # Convert v2 if applicable v2_path = os.path.join("stories", f"{self.story_name}_v2.json") if os.path.exists(v2_path): logger.info("Migrating v2 save") with open(v2_path, "r", encoding="utf-8") as file: v2j = json.load(file) if v2j["story_id"] == self.story_id: shutil.move(v2_path, os.path.join(self.save_paths.base, ".v2_old.json")) else: logger.warning(f"Story mismatch in v2 migration. Existing file had story id {v2j['story_id']} but we have {self.story_id}") with open(self.save_paths.story, "w", encoding="utf-8") as file: file.write(self.to_json()) self.gamesaved = True def update_story_path_structure(self, path: str) -> None: # Upon loading a file, makes directories that are required for certain # functionality. sp = SavePaths(path) for path in sp.required_paths: try: os.mkdir(path) except FileExistsError: pass def save_revision(self): game = json.loads(self.to_json()) del game['revisions'] self.revisions.append(game) self.gamesaved = False def reset(self): self.no_save = True self._socketio.emit("reset_story", {}, broadcast=True, room="UI_2") self.__init__(self._socketio, self._koboldai_vars, tokenizer=self.tokenizer) self.no_save = False def sync_worldinfo_v1_to_v2(self): new_world_info = KoboldWorldInfo(None, self, self._koboldai_vars, tokenizer=self.tokenizer) for wi in self.worldinfo: if wi['init'] == True: new_world_info.add_item([x.strip() for x in wi["key"].split(",")][0], wi["key"], wi.get("keysecondary", ""), "root" if wi["folder"] is None else self.wifolders_d[str(wi['folder'])]['name'], wi.get("constant", False), wi["content"], wi.get("comment", ""), v1_uid=wi['uid'], sync=False) new_world_info._socketio = self._socketio self.worldinfo_v2 = new_world_info def assign_world_info_to_actions(self, action_id=None, wuid=None, no_transmit=False): logger.debug("Calcing WI Assignment for action_id: {} wuid: {}".format(action_id, wuid)) if action_id != -1 and (action_id is None or action_id not in self.actions.actions): actions_to_check = self.actions.actions elif action_id == -1: actions_to_check = {} else: actions_to_check = {action_id: self.actions.actions[action_id]} if wuid is None or wuid not in self.worldinfo_v2.world_info: wi_to_check = self.worldinfo_v2.world_info else: wi_to_check = {wuid: self.worldinfo_v2.world_info[wuid]} for action_id, action in actions_to_check.items(): for uid, wi in wi_to_check.items(): for key in sorted(wi['key'], key=len, reverse=True): if key in action['Selected Text']: self.actions.add_wi_to_action(action_id, key, wi['content'], uid, no_transmit=no_transmit) break #Do prompt if no action_id was sent if action_id is None or action_id == -1: for uid, wi in wi_to_check.items(): for key in sorted(wi['key'], key=len, reverse=True): if key in self.prompt: if wi['keysecondary'] != []: for key2 in wi['keysecondary']: if key2 in self.prompt: self.actions.add_wi_to_action(-1, key, wi['content'], uid, no_transmit=no_transmit) break else: self.actions.add_wi_to_action(-1, key, wi['content'], uid, no_transmit=no_transmit) break def __setattr__(self, name, value): new_variable = name not in self.__dict__ old_value = getattr(self, name, None) super().__setattr__(name, value) #Put variable change actions here if name not in self.local_only_variables and name[0] != "_" and not new_variable and old_value != value: process_variable_changes(self._socketio, self.__class__.__name__.replace("_settings", ""), name, value, old_value) #We want to automatically set gamesaved to false if something happens to the actions list (pins, redos, generations, text, etc) #To do that we need to give the actions list a copy of this data so it can set the gamesaved variable as needed #autosave action if name == "gamesaved" and value == False and self.autosave: logger.debug("Saving story from gamesaved change and on autosave") self.save_story() if not new_variable and old_value != value: #Change game save state if name in ['story_name', 'prompt', 'memory', 'authornote', 'authornotetemplate', 'andepth', 'chatname', 'actionmode', 'dynamicscan', 'notes', 'biases']: self.gamesaved = False if name == "gen_audio" and value: self.actions.gen_all_audio() elif name == 'useprompt': ignore = self._koboldai_vars.calc_ai_text() elif name == 'actions': self.actions.story_settings = self logger.debug("Calcing AI text after setting actions") ignore = self._koboldai_vars.calc_ai_text() elif name == 'story_name': #reset the story id if we change the name self.story_id = int.from_bytes(os.urandom(16), 'little', signed=True) # Story name influences save base self.save_paths.base = os.path.join("stories", self.story_name or "Untitled") #Recalc AI Text elif name == 'authornote': ignore = self._koboldai_vars.calc_ai_text() elif name == 'authornotetemplate': ignore = self._koboldai_vars.calc_ai_text() elif name == 'memory': ignore = self._koboldai_vars.calc_ai_text() elif name == "genres": self._koboldai_vars.calc_ai_text() elif name == 'prompt': self.prompt_wi_highlighted_text = [{"text": self.prompt, "WI matches": None, "WI Text": ""}] self.assign_world_info_to_actions(action_id=-1, wuid=None) process_variable_changes(self._socketio, self.__class__.__name__.replace("_settings", ""), 'prompt_wi_highlighted_text', self.prompt_wi_highlighted_text, None) ignore = self._koboldai_vars.calc_ai_text() self.actions.gen_audio(action_id=-1) #Because we have seperate variables for action types, this syncs them elif name == 'storymode': if value == 0: self.adventure = False self.chatmode = False self.actionmode = 0 elif value == 1: self.adventure = True self.chatmode = False elif value == 2: self.adventure = False self.chatmode = True self.actionmode = 0 elif name == 'adventure' and value == True: self.chatmode = False self.storymode = 1 elif name == 'adventure' and value == False and self.chatmode == False: self.storymode = 0 elif name == 'chatmode' and value == True: self.adventure = False self.storymode = 2 self.actionmode = 0 elif name == 'chatmode' and value == False and self.adventure == False: self.storymode = 0 self.actionmode = 0 class user_settings(settings): local_only_variables = ['importjs'] no_save_variables = ['importnum', 'importjs', 'loadselect', 'spselect', 'svowname', 'saveow', 'laststory', 'sid', "revision"] settings_name = "user" def __init__(self, socketio): self._socketio = socketio self.wirmvwhtsp = False # Whether to remove leading whitespace from WI entries self.widepth = 3 # How many historical actions to scan for WI hits self.formatoptns = {'frmttriminc': True, 'frmtrmblln': False, 'frmtrmspch': False, 'frmtadsnsp': True, 'singleline': False} # Container for state of formatting options self.frmttriminc = True self.frmtrmblln = False self.frmtrmspch = False self.frmtadsnsp = True self.singleline = False self.remove_double_space = True self.importnum = -1 # Selection on import popup list self.importjs = {} # Temporary storage for import data self.loadselect = "" # Temporary storage for story filename to load self.spselect = "" # Temporary storage for soft prompt filename to load self.svowname = "" # Filename that was flagged for overwrite confirm self.saveow = False # Whether or not overwrite confirm has been displayed self.laststory = None # Filename (without extension) of most recent story JSON file we loaded self.sid = "" # session id for the socketio client (request.sid) self.username = "Default User" # Displayed Username self.nopromptgen = False self.rngpersist = False self.nogenmod = False self.debug = False # If set to true, will send debug information to the client for display self.output_streaming = True self.show_probs = False # Whether or not to show token probabilities self.beep_on_complete = False self.img_gen_priority = 1 self.show_budget = False self.ui_level = 2 self.img_gen_api_url = "http://127.0.0.1:7860" self.img_gen_art_guide = "masterpiece, digital painting, <|>, dramatic lighting, highly detailed, trending" self.img_gen_negative_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name" self.img_gen_api_username = "" self.img_gen_api_password = "" self.img_gen_steps = 30 self.img_gen_cfg_scale = 7.0 self.cluster_requested_models = [] # The models which we allow to generate during cluster mode self.wigen_use_own_wi = False self.wigen_amount = 80 self.screenshot_show_story_title = True self.screenshot_show_author_name = True self.screenshot_author_name = "Anonymous" self.screenshot_use_boring_colors = False self.oaiurl = "" # OpenAI API URL self.revision = None self.oaiengines = "https://api.openai.com/v1/engines" self.url = "https://api.inferkit.com/v1/models/standard/generate" # InferKit API URL self.colaburl = "" # Ngrok url for Google Colab mode self.apikey = "" # API key to use for InferKit API calls self.oaiapikey = "" # API key to use for OpenAI API calls self.horde_api_key = "0000000000" self.horde_worker_name = "My Awesome Instance" self.horde_url = "https://horde.koboldai.net" def __setattr__(self, name, value): new_variable = name not in self.__dict__ old_value = getattr(self, name, None) super().__setattr__(name, value) #Put variable change actions here if name not in self.local_only_variables and name[0] != "_" and not new_variable: process_variable_changes(self._socketio, self.__class__.__name__.replace("_settings", ""), name, value, old_value) class undefined_settings(settings): def __init__(self): pass def __setattr__(self, name, value): super().__setattr__(name, value) logger.error("{} just set {} to {} in koboldai_vars. That variable isn't defined!".format(inspect.stack()[1].function, name, value)) class system_settings(settings): local_only_variables = ['lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 'lua_koboldcore', 'regex_sl', 'acregex_ai', 'acregex_ui', 'comregex_ai', 'comregex_ui', 'sp', '_horde_pid', 'inference_config', 'image_pipeline', 'summarizer', 'summary_tokenizer', 'tts_model'] no_save_variables = ['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', 'generating_image', 'bit_8_available', 'host', 'hascuda', 'usegpu'] settings_name = "system" def __init__(self, socketio, koboldai_var): self._socketio = socketio self.noai = False # Runs the script without starting up the transformers pipeline self.aibusy = False # Stops submissions while the AI is working self.status_message = "" self.serverstarted = False # Whether or not the Flask server has started self.lua_state = None # Lua state of the Lua scripting system self.lua_koboldbridge = None # `koboldbridge` from bridge.lua self.lua_kobold = None # `kobold` from` bridge.lua self.lua_koboldcore = None # `koboldcore` from bridge.lua self.lua_logname = ... # Name of previous userscript that logged to terminal self.lua_running = False # Whether or not Lua is running (i.e. wasn't stopped due to an error) self.abort = False # Whether or not generation was aborted by clicking on the submit button during generation self.compiling = False # If using a TPU Colab, this will be set to True when the TPU backend starts compiling and then set to False again self.checking = False # Whether or not we are actively checking to see if TPU backend is compiling or not self.sp_changed = False # This gets set to True whenever a userscript changes the soft prompt so that check_for_sp_change() can alert the browser that the soft prompt has changed self.spfilename = "" # Filename of soft prompt to load, or an empty string if not using a soft prompt self.userscripts = [] # List of userscripts to load self.last_userscripts = [] # List of previous userscript filenames from the previous time userscripts were send via usstatitems self.corescript = "default.lua" # Filename of corescript to load self.gpu_device = 0 # Which PyTorch device to use when using pure GPU generation self.savedir = os.getcwd()+"\\stories" self.hascuda = False # Whether torch has detected CUDA on the system self.usegpu = False # Whether to launch pipeline with GPU support self.splist = [] self.spselect = "" # Temporary storage for soft prompt filename to load self.spmeta = None # Metadata of current soft prompt, or None if not using a soft prompt self.spname = "Not in Use" # Name of the soft prompt self.sp = None # Current soft prompt tensor (as a NumPy array) self.sp_length = 0 # Length of current soft prompt in tokens, or 0 if not using a soft prompt self.has_genmod = False # Whether or not at least one loaded Lua userscript has a generation modifier self.breakmodel = False # For GPU users, whether to use both system RAM and VRAM to conserve VRAM while offering speedup compared to CPU-only self.bmsupported = False # Whether the breakmodel option is supported (GPT-Neo/GPT-J/XGLM/OPT only, currently) self.nobreakmodel = False # Something specifically requested Breakmodel to be disabled (For example a models config) self.smandelete = False # Whether stories can be deleted from inside the browser self.smanrename = False # Whether stories can be renamed from inside the browser self.allowsp = False # Whether we are allowed to use soft prompts (by default enabled if we're using GPT-2, GPT-Neo or GPT-J) self.regex_sl = re.compile(r'\n*(?<=.) *\n(.|\n)*') # Pattern for limiting the output to a single line self.acregex_ai = re.compile(r'\n* *>(.|\n)*') # Pattern for matching adventure actions from the AI so we can remove them self.acregex_ui = re.compile(r'^ *(>.*)$', re.MULTILINE) # Pattern for matching actions in the HTML-escaped story so we can apply colouring, etc (make sure to encase part to format in parentheses) self.comregex_ai = re.compile(r'(?:\n<\|(?:.|\n)*?\|>(?=\n|$))|(?:<\|(?:.|\n)*?\|>\n?)') # Pattern for matching comments to remove them before sending them to the AI self.comregex_ui = re.compile(r'(<\|(?:.|\n)*?\|>)') # Pattern for matching comments in the editor self.host = False self.flaskwebgui = False self.quiet = False # If set will suppress any story text from being printed to the console (will only be seen on the client web page) self.use_colab_tpu = os.environ.get("COLAB_TPU_ADDR", "") != "" or os.environ.get("TPU_NAME", "") != "" # Whether or not we're in a Colab TPU instance or Kaggle TPU instance and are going to use the TPU rather than the CPU self.aria2_port = 6799 #Specify the port on which aria2's RPC interface will be open if aria2 is installed (defaults to 6799) self.standalone = False self.api_tokenizer_id = None self.disable_set_aibusy = False self.disable_input_formatting = False self.disable_output_formatting = False self.full_determinism = False # Whether or not full determinism is enabled self.seed_specified = False # Whether or not the current RNG seed was specified by the user (in their settings file) self.seed = None # The current RNG seed (as an int), or None if unknown self.alt_gen = False # Use the calc_ai_text method for generating text to go to the AI self.theme_list = [".".join(f.split(".")[:-1]) for f in os.listdir("./themes") if os.path.isfile(os.path.join("./themes", f))] self.cloudflare_link = "" self.story_loads = {} #dict of when each story was last loaded self.standalone = False self.disable_set_aibusy = False self.disable_input_formatting = False self.disable_output_formatting = False self.api_tokenizer_id = None self.port = 5000 try: import google.colab self.on_colab = True except: self.on_colab = False pass print("Colab Check: {}".format(self.on_colab)) self.horde_share = False self._horde_pid = None self.generating_image = False #The current status of image generation self.image_pipeline = None self.summarizer = None self.summary_tokenizer = None self.keep_img_gen_in_memory = False self.cookies = {} #cookies for colab since colab's URL changes, cookies are lost self.experimental_features = False #check if bitsandbytes is installed self.bit_8_available = False if importlib.util.find_spec("bitsandbytes") is not None and sys.platform.startswith('linux'): #We can install bitsandbytes, but it doesn't work on windows, so limit it here if torch.cuda.is_available(): for device in range(torch.cuda.device_count()): if torch.cuda.get_device_properties(device).major > 7: self.bit_8_available = True break elif torch.cuda.get_device_properties(device).major == 7 and torch.cuda.get_device_properties(device).minor >= 2: self.bit_8_available = True break self.seen_messages = [] @dataclass class _inference_config: do_streaming: bool = False do_dynamic_wi: bool = False # Genamt stopping is mostly tied to Dynamic WI stop_at_genamt: bool = False do_core: bool = True self.inference_config = _inference_config() self._koboldai_var = koboldai_var def __setattr__(self, name, value): new_variable = name not in self.__dict__ old_value = getattr(self, name, None) super().__setattr__(name, value) #Put variable change actions here if name == 'serverstarted': global serverstarted serverstarted = value if name not in self.local_only_variables and name[0] != "_" and not new_variable: process_variable_changes(self._socketio, self.__class__.__name__.replace("_settings", ""), name, value, old_value) #if name == "aibusy" and value == False and self.abort == True: # koboldai_vars.abort = False #for original UI if name == 'sp_changed': self._socketio.emit('from_server', {'cmd': 'spstatitems', 'data': {self.spfilename: self.spmeta} if self.allowsp and len(self.spfilename) else {}}, namespace=None, broadcast=True, room="UI_1") super().__setattr__("sp_changed", False) if name == 'keep_img_gen_in_memory' and value == False: self.image_pipeline = None if name == 'alt_gen' and old_value != value: logger.debug("Calcing AI Text from setting alt_gen") self._koboldai_var.calc_ai_text() if name == 'horde_share': if self.on_colab == False: if os.path.exists("./KoboldAI-Horde-Bridge"): if value == True: if self._horde_pid is None: logger.info("Starting Horde bridge") bridge = importlib.import_module("KoboldAI-Horde-Bridge.bridge") self._horde_pid = bridge.kai_bridge() try: bridge_cd = importlib.import_module("KoboldAI-Horde-Bridge.clientData") cluster_url = bridge_cd.cluster_url kai_name = bridge_cd.kai_name if kai_name == "My Awesome Instance": kai_name = f"KoboldAI UI Instance #{random.randint(-100000000, 100000000)}" api_key = bridge_cd.api_key priority_usernames = bridge_cd.priority_usernames except: cluster_url = "https://horde.koboldai.net" kai_name = self._koboldai_var.horde_worker_name if kai_name == "My Awesome Instance": kai_name = f"KoboldAI UI Instance #{random.randint(-100000000, 100000000)}" api_key = self._koboldai_var.horde_api_key priority_usernames = [] # Always use the local URL & port kai_url = f'http://127.0.0.1:{self.port}' logger.info(f"Name: {kai_name} on {kai_url}") threading.Thread(target=self._horde_pid.bridge, args=(1, api_key, kai_name, kai_url, cluster_url, priority_usernames)).run() else: if self._horde_pid is not None: logger.info("Killing Horde bridge") self._horde_pid.stop() self._horde_pid = None class KoboldStoryRegister(object): def __init__(self, socketio, story_settings, koboldai_vars, tokenizer=None, sequence=[]): self._socketio = socketio self._koboldai_vars = koboldai_vars #### DO NOT DIRECTLY EDIT THE ACTIONS DICT. IT WILL NOT TRANSMIT TO CLIENT. USE FUCTIONS BELOW TO DO SO ### #### doing actions[x] = game text is OK 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.tortoise = None self.make_audio_thread = None self.make_audio_queue = multiprocessing.Queue() self.make_audio_thread_slow = None self.make_audio_queue_slow = multiprocessing.Queue() self.probability_buffer = None for item in sequence: self.append(item) def reset(self, sequence=[]): self.__init__(self._socketio, self.story_settings, self._koboldai_vars, sequence=sequence) def add_wi_to_action(self, action_id, key, content, uid, no_transmit=False): old = self.story_settings.prompt_wi_highlighted_text.copy() if action_id == -1 else self.actions[action_id].copy() force_changed = False #First check to see if we have the wi_highlighted_text variable if action_id != -1: if 'wi_highlighted_text' not in self.actions[action_id]: self.actions[action_id]['wi_highlighted_text'] = [{"text": self.actions[action_id]['Selected Text'], "WI matches": None, "WI Text": ""}] action = self.actions[action_id]['wi_highlighted_text'] else: action = self.story_settings.prompt_wi_highlighted_text i=0 while i < len(action): if action[i]['WI matches'] is None and key in action[i]['text'] and action[i]['text'] != "": chunk = action[i] for j in range(len(chunk['text'])-1, -1, -1): if key in chunk['text'][j:]: #OK, we found the start of our key. We'll split our text into before key, key, after key pre_text = chunk['text'][:j] if j != 0 else None post_text = chunk['text'][j+len(key):] if j != len(chunk['text']) else None text = chunk['text'][j:j+len(key)] action[i]['text'] = text action[i]['WI matches'] = uid action[i]['WI Text'] = content adder = 1 if pre_text is not None: action.insert(i, {"text": pre_text, "WI matches": None, "WI Text": ""}) adder += 1 if post_text is not None: action.insert(i+adder, {"text": post_text, "WI matches": None, "WI Text": ""}) break; elif action[i]['WI matches'] == uid: action[i]['WI Text'] = content force_changed = True i+=1 else: i+=1 if action_id != -1: if old != self.actions[action_id] or force_changed: if not no_transmit: process_variable_changes(self._socketio, "story", 'actions', {"id": action_id, 'action': self.actions[action_id]}, old) else: if old != self.story_settings.prompt_wi_highlighted_text or force_changed: if not no_transmit: process_variable_changes(self._socketio, "story", 'prompt_wi_highlighted_text', self.story_settings.prompt_wi_highlighted_text, old) def __str__(self): if len(self.actions) > 0: return "".join([x['Selected Text'] for ignore, x in sorted(self.actions.items())]) else: return "" def __repr__(self): return self.__str__() def __iter__(self): self.itter = -1 return self def __next__(self): self.itter += 1 if self.itter <= self.action_count: return self.itter else: raise StopIteration def __getitem__(self, i): if isinstance(i, slice): temp = [self.actions[x]["Selected Text"] for x in list(self.actions)[i]] return temp else: if i < 0: return self.actions[self.action_count+i+1]["Selected Text"] return self.actions[i]["Selected Text"] def __setitem__(self, i, text): if self._koboldai_vars.remove_double_space: while " " in text: text = text.replace(" ", " ") if i > 0 and text != "" and i-1 in self.actions and self.actions[i-1]['Selected Text'] != "": if self.actions[i-1]['Selected Text'][-1] == " " and text[0] == " ": text = text[1:] if i in self.actions: old = self.actions[i] old_text = self.actions[i]["Selected Text"] if self.actions[i]["Selected Text"] != text: self.actions[i]["Selected Text"] = text if 'wi_highlighted_text' in self.actions[i]: del self.actions[i]['wi_highlighted_text'] if self._koboldai_vars.tokenizer is not None: tokens = self._koboldai_vars.tokenizer.encode(text) if 'Probabilities' in self.actions[i]: for token_num in range(len(self.actions[i]["Probabilities"])): for token_option in range(len(self.actions[i]["Probabilities"][token_num])): if token_num < len(tokens): self.actions[i]["Probabilities"][token_num][token_option]["Used"] = tokens[token_num] == self.actions[i]["Probabilities"][token_num][token_option]["tokenId"] else: self.actions[i]["Probabilities"][token_num][token_option]["Used"] = False if "Options" in self.actions[i]: for option in self.actions[i]["Options"][:]: if option["text"] == text: self.actions[i]["Options"].remove(option) if old_text != "": self.actions[i]["Options"].append({"text": old_text, "Pinned": False, "Previous Selection": False, "Edited": True}) else: old_text = None old_length = None old = None self.actions[i] = {"Selected Text": text, "Probabilities": [], "Options": [], "Time": int(time.time())} self.story_settings.assign_world_info_to_actions(action_id=i, no_transmit=True) process_variable_changes(self._socketio, "story", 'actions', {"id": i, 'action': self.actions[i]}, old) logger.debug("Calcing AI Text from Action __setitem__") ignore = self._koboldai_vars.calc_ai_text() self.set_game_saved() self.gen_audio(i) def __len__(self): return self.action_count+1 if self.action_count >=0 else 0 def __reversed__(self): return reversed(range(self.action_count+1)) def values(self): return [self.actions[k]["Selected Text"] for k in self.actions] def options(self, ui_version=2): if ui_version == 1: return [{"Selected Text": self.actions[k]["Selected Text"], "Alternative Text": self.actions[k]["Options"]} for k in self.actions] else: return [self.actions[k]["Options"] for k in self.actions] def to_json(self): return {"action_count": self.action_count, "actions": self.actions} def load_json(self, json_data): if type(json_data) == str: import json json_data = json.loads(json_data) self.action_count = int(json_data['action_count']) #JSON forces keys to be strings, so let's fix that temp = {} data_to_send = [] for item in json_data['actions']: if "Time" not in json_data["actions"][item]: json_data["actions"][item]["Time"] = int(time.time()) if "Original Text" not in json_data["actions"][item]: json_data["actions"][item]["Original Text"] = json_data["actions"][item]["Selected Text"] temp[int(item)] = json_data['actions'][item] if int(item) >= self.action_count-100: #sending last 100 items to UI data_to_send.append({"id": item, 'action': temp[int(item)]}) process_variable_changes(self._socketio, "story", 'actions', data_to_send, None) self.actions = temp #Check if our json has our new world info highlighting data if len(self.actions) > 0: if 'wi_highlighted_text' not in self.actions[0]: self.story_settings.assign_world_info_to_actions(no_transmit=True) logger.debug("Calcing AI Text from Action load from json") ignore = self._koboldai_vars.calc_ai_text() self.set_game_saved() self.gen_all_audio() def append(self, text, action_id_offset=0, recalc=True, submission=False): """Create a new action and append it to the table of actions. text: The text of the action. action_id_offset: An optional offset added to action_count when assiging an action_id. recalc: If True, recalculate the context and generate audio. submission: If True, mark this action as a user submission. """ if self._koboldai_vars.remove_double_space: while " " in text: text = text.replace(" ", " ") if action_id_offset > 0: if self.actions[action_id_offset-1]['Selected Text'][-1] == " " and text[0] == " ": text = text[1:] self.clear_unused_options(clear_probabilities=False) self.action_count+=1 action_id = self.action_count + action_id_offset if action_id in self.actions: if not self.actions[action_id].get("Original Text"): self.actions[action_id]["Original Text"] = text if self.actions[action_id]["Selected Text"] != text: self.actions[action_id]["Selected Text"] = text self.actions[action_id]["Time"] = self.actions[action_id].get("Time", int(time.time())) if 'buffer' in self.actions[action_id]: if self._koboldai_vars.tokenizer is not None: tokens = self._koboldai_vars.tokenizer.encode(text) for token_num in range(len(self.actions[action_id]["Probabilities"])): for token_option in range(len(self.actions[action_id]["Probabilities"][token_num])): if token_num < len(tokens): self.actions[action_id]["Probabilities"][token_num][token_option]["Used"] = tokens[token_num] == self.actions[action_id]["Probabilities"][token_num][token_option]["tokenId"] else: self.actions[action_id]["Probabilities"][token_num][token_option]["Used"] = False selected_text_length = 0 self.actions[action_id]["Selected Text Length"] = selected_text_length for item in self.actions[action_id]["Options"]: if item['text'] == text: old_options = self.actions[action_id]["Options"] del item else: selected_text_length = 0 self.actions[action_id] = { "Selected Text": text, "Selected Text Length": selected_text_length, "Options": [], "Probabilities": [], "Time": int(time.time()), "Original Text": text, "Origin": "user" if submission else "ai" } if submission: self.submission_id = action_id else: self.submission_id = 0 if self.story_settings is not None: self.story_settings.assign_world_info_to_actions(action_id=action_id, no_transmit=True) process_variable_changes(self._socketio, "story", 'actions', {"id": action_id, 'action': self.actions[action_id]}, None) self.set_game_saved() if recalc: logger.debug("Calcing AI Text from Action Append") ignore = self._koboldai_vars.calc_ai_text() self.gen_audio(action_id) def append_options(self, option_list): if self.action_count+1 in self.actions: #First let's check if we did streaming, as we can just replace those items with these old_options = copy.deepcopy(self.actions[self.action_count+1]["Options"]) i=-1 for option in option_list: i+=1 found = False for item in self.actions[self.action_count+1]["Options"]: if 'stream_id' in item and item['stream_id'] == i: item['text'] = option del item['stream_id'] found = True if 'Probabilities' in item: if self._koboldai_vars.tokenizer is not None: tokens = self._koboldai_vars.tokenizer.encode(option) for token_num in range(len(item["Probabilities"])): for token_option in range(len(item["Probabilities"][token_num])): if token_num < len(tokens): item["Probabilities"][token_num][token_option]["Used"] = tokens[token_num] == item["Probabilities"][token_num][token_option]["tokenId"] else: item["Probabilities"][token_num][token_option]["Used"] = False break elif item['text'] == option: found = True if 'stream_id' in item: del item['stream_id'] found = True if 'Probabilities' in item: if self._koboldai_vars.tokenizer is not None: tokens = self._koboldai_vars.tokenizer.encode(option) for token_num in range(len(item["Probabilities"])): for token_option in range(len(item["Probabilities"][token_num])): if token_num < len(tokens): item["Probabilities"][token_num][token_option]["Used"] = tokens[token_num] == item["Probabilities"][token_num][token_option]["tokenId"] else: item["Probabilities"][token_num][token_option]["Used"] = False break if not found: self.actions[self.action_count+1]['Options'].append({"text": option, "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": []}) else: old_options = None self.actions[self.action_count+1] = { "Selected Text": "", "Selected Text Length": 0, "Options": [{"text": x, "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": []} for x in option_list], "Time": int(time.time()), } process_variable_changes(self._socketio, "story", 'actions', {"id": self.action_count+1, 'action': self.actions[self.action_count+1]}, None) self.set_game_saved() def set_options(self, option_list, action_id): if action_id not in self.actions: old_options = None self.actions[action_id] = { "Selected Text": "", "Options": option_list, "Time": int(time.time()), } else: old_options = self.actions[action_id]["Options"] self.actions[action_id]["Options"] = [] for item in option_list: for old_item in old_options: if item['text'] == old_item['text']: #We already have this option, so we need to save the probabilities if 'Probabilities' in old_item: item['Probabilities'] = old_item['Probabilities'] self.actions[action_id]["Options"].append(item) process_variable_changes(self._socketio, "story", 'actions', {"id": action_id, 'action': self.actions[action_id]}, None) def clear_unused_options(self, pointer=None, clear_probabilities=True): new_options = [] old_options = None if pointer is None: pointer = self.action_count+1 if pointer in self.actions: old_options = copy.deepcopy(self.actions[pointer]["Options"]) self.actions[pointer]["Options"] = [x for x in self.actions[pointer]["Options"] if x["Pinned"] or x["Previous Selection"] or x["Edited"]] new_options = self.actions[pointer]["Options"] if clear_probabilities: self.actions[pointer]['Probabilities'] = [] process_variable_changes(self._socketio, "story", 'actions', {"id": pointer, 'action': self.actions[pointer]}, None) self.set_game_saved() def clear_all_options(self): for i in self.actions: self.actions[i]["Options"] = [] def set_pin(self, action_step, option_number): if action_step in self.actions: if option_number < len(self.actions[action_step]['Options']): old_options = copy.deepcopy(self.actions[action_step]["Options"]) self.actions[action_step]['Options'][option_number]['Pinned'] = True process_variable_changes(self._socketio, "story", 'actions', {"id": action_step, 'action': self.actions[action_step]}, None) self.set_game_saved() def unset_pin(self, action_step, option_number): if action_step in self.actions: old_options = copy.deepcopy(self.actions[action_step]["Options"]) if option_number < len(self.actions[action_step]['Options']): self.actions[action_step]['Options'][option_number]['Pinned'] = False process_variable_changes(self._socketio, "story", 'actions', {"id": action_step, 'action': self.actions[action_step]}, None) self.set_game_saved() def toggle_pin(self, action_step, option_number): if action_step in self.actions: old_options = copy.deepcopy(self.actions[action_step]["Options"]) if option_number < len(self.actions[action_step]['Options']): self.actions[action_step]['Options'][option_number]['Pinned'] = not self.actions[action_step]['Options'][option_number]['Pinned'] process_variable_changes(self._socketio, "story", 'actions', {"id": action_step, 'action': self.actions[action_step]}, None) self.set_game_saved() def use_option(self, option_number, action_step=None): if action_step is None: action_step = self.action_count+1 if action_step in self.actions: old_options = copy.deepcopy(self.actions[action_step]["Options"]) old_text = self.actions[action_step]["Selected Text"] old_length = self.actions[action_step]["Selected Text Length"] if option_number < len(self.actions[action_step]['Options']): self.actions[action_step]["Selected Text"] = self.actions[action_step]['Options'][option_number]['text'] if 'Probabilities' in self.actions[action_step]['Options'][option_number]: self.actions[action_step]["Probabilities"] = self.actions[action_step]['Options'][option_number]['Probabilities'] if self._koboldai_vars.tokenizer is not None: self.actions[action_step]['Selected Text Length'] = len(self._koboldai_vars.tokenizer.encode(self.actions[action_step]['Options'][option_number]['text'])) else: self.actions[action_step]['Selected Text Length'] = 0 del self.actions[action_step]['Options'][option_number] #If this is the current spot in the story, advance if action_step-1 == self.action_count: self.action_count+=1 self._socketio.emit("var_changed", {"classname": "actions", "name": "Action Count", "old_value": None, "value":self.action_count, "transmit_time": str(datetime.datetime.now())}, broadcast=True, room="UI_2") self.story_settings.assign_world_info_to_actions(action_id=action_step, no_transmit=True) self.clear_unused_options(pointer=action_step) #process_variable_changes(self._socketio, "story", 'actions', {"id": action_step, 'action': self.actions[action_step]}, None) self.set_game_saved() logger.debug("Calcing AI Text from Action Use Option") ignore = self._koboldai_vars.calc_ai_text() self.gen_audio(action_step) def delete_option(self, option_number, action_step=None): if action_step is None: action_step = self.action_count+1 if action_step in self.actions: if option_number < len(self.actions[action_step]['Options']): del self.actions[action_step]['Options'][option_number] process_variable_changes(self._socketio, "story", 'actions', {"id": action_step, 'action': self.actions[action_step]}, None) self.set_game_saved() def delete_action(self, action_id, keep=True): if action_id in self.actions: old_options = copy.deepcopy(self.actions[action_id]["Options"]) old_text = self.actions[action_id]["Selected Text"] old_length = self.actions[action_id]["Selected Text Length"] if keep: self.actions[action_id]["Options"].append({"text": self.actions[action_id]["Selected Text"], "Pinned": False, "Previous Selection": True, "Edited": False}) self.actions[action_id]["Selected Text"] = "" if "wi_highlighted_text" in self.actions[action_id]: del self.actions[action_id]["wi_highlighted_text"] self.actions[action_id]['Selected Text Length'] = 0 self.action_count -= 1 process_variable_changes(self._socketio, "story", 'actions', {"id": action_id, 'action': self.actions[action_id]}, None) self.set_game_saved() logger.debug("Calcing AI Text from Action Delete") ignore = self._koboldai_vars.calc_ai_text() def pop(self, keep=True): if self.action_count >= 0: text = self.actions[self.action_count]['Selected Text'] self.delete_action(self.action_count, keep=keep) logger.debug("Calcing AI Text from Action Pop") return text else: return None def get_last_key(self): return self.action_count def get_current_options(self): if self.action_count+1 in self.actions: return self.actions[self.action_count+1]["Options"] else: return [] def get_current_options_no_edits(self, ui=2): if ui==2: if self.action_count+1 in self.actions: return [x for x in self.actions[self.action_count+1]["Options"] if x["Edited"] == False and x['Previous Selection'] == False] else: return [] else: if self.action_count+1 in self.actions: return [[x['text'], "pinned" if x['Pinned'] else 'normal'] for x in self.actions[self.action_count+1]["Options"] if x["Edited"] == False and x['Previous Selection'] == False] else: return [] def get_redo_options(self): if self.action_count+1 in self.actions: return [x for x in self.actions[self.action_count+1]['Options'] if x['Pinned'] or x['Previous Selection']] else: return [] def set_game_saved(self): if self.story_settings is not None: self.story_settings.gamesaved = False def stream_tokens(self, text_list): if len(text_list) > 1 or (self._koboldai_vars.alt_multi_gen and self._koboldai_vars.numseqs > 1): if self._koboldai_vars.alt_multi_gen: #since alt_multi_gen is really just several single gens the text list is always 1 deep, so we need some #other way to figure out wich spot in our options list we're on. We'll figure it out by seeing how many #tokens we generated vs how many each option should take stream_offset = int((self._koboldai_vars.generated_tkns-1) / self._koboldai_vars.genamt) else: stream_offset = 0 option_offset = 0 if self.action_count+1 in self.actions: for x in range(len(self.actions[self.action_count+1]['Options'])): option = self.actions[self.action_count+1]['Options'][x] if option['Pinned'] or option["Previous Selection"] or option["Edited"]: option_offset = x+1 if self.action_count+1 in self.actions: for i in range(len(text_list)): if i+stream_offset+option_offset < len(self.actions[self.action_count+1]['Options']): self.actions[self.action_count+1]['Options'][i+stream_offset+option_offset]['text'] = "{}{}".format(self.actions[self.action_count+1]['Options'][i+stream_offset+option_offset]['text'], text_list[i]) else: logger.info(self.actions[self.action_count+1]['Options']) self.actions[self.action_count+1]['Options'].append({"text": text_list[i], "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [], "stream_id": i+stream_offset}) else: self.actions[self.action_count+1] = {"Selected Text": "", "Selected Text Length": 0, "Options": [], "Time": int(time.time())} for i in range(len(text_list)): self.actions[self.action_count+1]['Options'].append({"text": text_list[i], "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [], "stream_id": i+stream_offset}) #We need to see if this is the last token being streamed. If so due to the rely it will come in AFTER the actual trimmed final text overwriting it in the UI if self._koboldai_vars.tokenizer is not None: if len(self._koboldai_vars.tokenizer.encode(self.actions[self.action_count+1]["Options"][0]['text'])) != self._koboldai_vars.genamt: #process_variable_changes(self._socketio, "actions", "Options", {"id": self.action_count+1, "options": self.actions[self.action_count+1]["Options"]}, {"id": self.action_count+1, "options": None}) process_variable_changes(self._socketio, "story", 'actions', {"id": self.action_count+1, 'action': self.actions[self.action_count+1]}, None) else: #We're streaming single options so our output is our selected #First we need to see if this is actually the prompt. If so we'll just not do streaming: if self.story_settings.prompt != "": if self.action_count+1 in self.actions: if self._koboldai_vars.tokenizer is not None: selected_text_length = len(self._koboldai_vars.tokenizer.encode(self.actions[self.action_count+1]['Selected Text'])) else: selected_text_length = 0 self.actions[self.action_count+1]['Selected Text'] = "{}{}".format(self.actions[self.action_count+1]['Selected Text'], text_list[0]) self.actions[self.action_count+1]['Selected Text Length'] = selected_text_length else: if self._koboldai_vars.tokenizer is not None: selected_text_length = len(self._koboldai_vars.tokenizer.encode(text_list[0])) else: selected_text_length = 0 self.actions[self.action_count+1] = {"Selected Text": text_list[0], "Selected Text Length": selected_text_length, "Options": [], "Time": int(time.time())} if self._koboldai_vars.tokenizer is not None: if len(self._koboldai_vars.tokenizer.encode(self.actions[self.action_count+1]['Selected Text'])) != self._koboldai_vars.genamt: #ui1 if queue is not None: queue.put(["from_server", {"cmd": "streamtoken", "data": [{ "decoded": text_list[0], "probabilities": self.probability_buffer }]}, {"broadcast":True, "room":"UI_1"}]) #process_variable_changes(self._socketio, "actions", "Options", {"id": self.action_count+1, "options": self.actions[self.action_count+1]["Options"]}, {"id": self.action_count+1, "options": None}) process_variable_changes(self._socketio, "story", 'actions', {"id": self.action_count+1, 'action': self.actions[self.action_count+1]}, None) def set_probabilities(self, probabilities, action_id=None): self.probability_buffer = probabilities if action_id is None: action_id = self.action_count+1 if action_id in self.actions: if 'Probabilities' not in self.actions[action_id]: self.actions[action_id]['Probabilities'] = [] self.actions[action_id]['Probabilities'].append(probabilities) else: self.actions[action_id] = { "Selected Text": "", "Selected Text Length": 0, "Options": [], "Probabilities": [probabilities], "Time": int(time.time()), } process_variable_changes(self._socketio, "story", 'actions', {"id": action_id, 'action': self.actions[action_id]}, None) def set_option_probabilities(self, probabilities, option_number, action_id=None): if action_id is None: action_id = self.action_count+1 if action_id in self.actions: old_options = self.actions[action_id]["Options"] if option_number < len(self.actions[action_id]["Options"]): if "Probabilities" not in self.actions[action_id]["Options"][option_number]: self.actions[action_id]["Options"][option_number]["Probabilities"] = [] 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) else: self.actions[action_id]['Options'].append({"temp_prob": True, "text": "", "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [probabilities], "stream_id": option_number}) else: self.actions[action_id] = { "Selected Text": "", "Selected Text Length": 0, "Options": [{"temp_prob": True, "text": "", "Pinned": False, "Previous Selection": False, "Edited": False, "Probabilities": [probabilities], "stream_id": option_number}], "Time": int(time.time()), } def to_sentences(self, submitted_text=None, max_action_id=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 if max_action_id is None: actions = {i: self.actions[i]['Selected Text'] for i in self.actions} else: actions = {i: self.actions[i]['Selected Text'] for i in self.actions if i <= max_action_id} if self.story_settings is None: actions[-1] = "" else: actions[-1] = self.story_settings.prompt # 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 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 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 [] Action_Position = [-1, len(actions[-1])] #First element is the action item, second is how much text is left Sentence_Position = [0, len(action_text_split[0][0])] while True: advance_action = False advance_sentence = False if Action_Position[1] <= Sentence_Position[1]: #We have enough text in the sentence to completely cover the action. Advance it to the next action advance_action = True if Sentence_Position[1] <= Action_Position[1]: 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 #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] if not advance_sentence: Sentence_Position[1] -= Action_Position[1] if advance_action: Action_Position[0] += 1 if Action_Position[0] > max(actions): break Action_Position[1] = len(actions[Action_Position[0]]) if advance_sentence: Sentence_Position[0] += 1 if Sentence_Position[0] >= len(action_text_split): 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]] return action_text_split def gen_audio(self, action_id=None, overwrite=True): if self.story_settings.gen_audio and self._koboldai_vars.experimental_features: if action_id is None: action_id = self.action_count if self.tts_model is None: language = 'en' model_id = 'v3_en' self.tts_model, example_text = torch.hub.load(repo_or_dir='snakers4/silero-models', model='silero_tts', language=language, speaker=model_id) #self.tts_model.to(torch.device(0)) # gpu or cpu self.tts_model.to(torch.device("cpu")) # gpu or cpu filename = os.path.join(self._koboldai_vars.save_paths.generated_audio, f"{action_id}.ogg") filename_slow = os.path.join(self._koboldai_vars.save_paths.generated_audio, f"{action_id}_slow.ogg") if overwrite or not os.path.exists(filename): if action_id == -1: self.make_audio_queue.put((self._koboldai_vars.prompt, filename)) else: self.make_audio_queue.put((self.actions[action_id]['Selected Text'], filename)) if self.make_audio_thread_slow is None or not self.make_audio_thread_slow.is_alive(): self.make_audio_thread_slow = threading.Thread(target=self.create_wave_slow, args=(self.make_audio_queue_slow, )) self.make_audio_thread_slow.start() if overwrite or not os.path.exists(filename_slow): if action_id == -1: self.make_audio_queue_slow.put((self._koboldai_vars.prompt, filename_slow)) else: self.make_audio_queue_slow.put((self.actions[action_id]['Selected Text'], filename_slow)) if self.make_audio_thread_slow is None or not self.make_audio_thread_slow.is_alive(): self.make_audio_thread_slow = threading.Thread(target=self.create_wave_slow, args=(self.make_audio_queue_slow, )) self.make_audio_thread_slow.start() def create_wave(self, make_audio_queue): import pydub sample_rate = 24000 speaker = 'en_5' while not make_audio_queue.empty(): (text, filename) = make_audio_queue.get() logger.info("Creating audio for {}".format(os.path.basename(filename))) if text.strip() == "": shutil.copy("data/empty_audio.ogg", filename) else: if len(text) > 2000: text = self.sentence_re.findall(text) else: text = [text] output = None for process_text in text: audio = self.tts_model.apply_tts(text=process_text, speaker=speaker, sample_rate=sample_rate) #audio_path=filename) channels = 2 if (audio.ndim == 2 and audio.shape[1] == 2) else 1 if output is None: output = pydub.AudioSegment(np.int16(audio * 2 ** 15).tobytes(), frame_rate=sample_rate, sample_width=2, channels=channels) else: output = output + pydub.AudioSegment(np.int16(audio * 2 ** 15).tobytes(), frame_rate=sample_rate, sample_width=2, channels=channels) output.export(filename, format="ogg", bitrate="16k") def create_wave_slow(self, make_audio_queue_slow): import pydub sample_rate = 24000 speaker = 'train_daws' if self.tortoise is None and importlib.util.find_spec("tortoise") is not None: self.tortoise=api.TextToSpeech() if importlib.util.find_spec("tortoise") is not None: voice_samples, conditioning_latents = load_voices([speaker]) while not make_audio_queue_slow.empty(): start_time = time.time() (text, filename) = make_audio_queue_slow.get() text_length = len(text) logger.info("Creating audio for {}".format(os.path.basename(filename))) if text.strip() == "": shutil.copy("data/empty_audio.ogg", filename) else: if len(text) > 20000: text = self.sentence_re.findall(text) else: text = [text] output = None for process_text in text: audio = self.tortoise.tts_with_preset(process_text, preset='ultra_fast', voice_samples=voice_samples, conditioning_latents=conditioning_latents).numpy() channels = 2 if (audio.ndim == 2 and audio.shape[1] == 2) else 1 if output is None: output = pydub.AudioSegment(np.int16(audio * 2 ** 15).tobytes(), frame_rate=sample_rate, sample_width=2, channels=channels) else: output = output + pydub.AudioSegment(np.int16(audio * 2 ** 15).tobytes(), frame_rate=sample_rate, sample_width=2, channels=channels) output.export(filename, format="ogg", bitrate="16k") logger.info("Slow audio took {} for {} characters".format(time.time()-start_time, text_length)) def gen_all_audio(self, overwrite=False): if self.story_settings.gen_audio and self._koboldai_vars.experimental_features: for i in reversed([-1]+list(self.actions.keys())): self.gen_audio(i, overwrite=False) #else: # print("{} and {}".format(self.story_settings.gen_audio, self._koboldai_vars.experimental_features)) def set_picture(self, action_id, filename, prompt): if action_id == -1: self.story_settings.prompt_picture_filename = filename self.story_settings.prompt_picture_prompt = prompt elif action_id in self.actions: self.actions[action_id]['picture_filename'] = filename self.actions[action_id]['picture_prompt'] = prompt def get_picture(self, action_id): if action_id == -1: if self.story_settings.prompt_picture_filename == "": return None, None filename = os.path.join(self._koboldai_vars.save_paths.generated_images, self.story_settings.prompt_picture_filename) prompt = self.story_settings.prompt_picture_prompt elif action_id in self.actions and 'picture_filename' in self.actions[action_id]: filename = os.path.join(self._koboldai_vars.save_paths.generated_images, self.actions[action_id]['picture_filename']) prompt = self.actions[action_id]['picture_prompt'] else: #Let's find the last picture if there is one found = False for i in reversed(range(-1, action_id)): if i in self.actions and 'picture_filename' in self.actions[i]: filename = os.path.join(self._koboldai_vars.save_paths.generated_images, self.actions[i]['picture_filename']) prompt = self.actions[i]['picture_prompt'] found = True break elif i == -1: if self.story_settings.prompt_picture_filename == "": return None, None filename = os.path.join(self._koboldai_vars.save_paths.generated_images, self.story_settings.prompt_picture_filename) prompt = self.story_settings.prompt_picture_prompt found = True if not found: return None, None if os.path.exists(filename): return filename, prompt return None, None def get_action_composition(self, action_id: int) -> List[dict]: """ Returns a list of chunks that comprise an action in dictionaries formatted as follows: type: string identifying chunk type ("ai", "user", "edit", or "prompt") content: the actual content of the chunk """ # Prompt doesn't need standard edit data if action_id == -1: if self._koboldai_vars.prompt: return [{"type": "prompt", "content": self._koboldai_vars.prompt}] return [] current_text = self.actions[action_id]["Selected Text"] action_original_type = self.actions[action_id].get("Origin", "ai") original = self.actions[action_id]["Original Text"] matching_blocks = difflib.SequenceMatcher( None, self.actions[action_id]["Original Text"], current_text ).get_matching_blocks() chunks = [] base = 0 for chunk_match in matching_blocks: inserted = current_text[base:chunk_match.b] content = current_text[chunk_match.b:chunk_match.b + chunk_match.size] base = chunk_match.b + chunk_match.size if inserted: chunks.append({"type": "edit", "content": inserted}) if content: chunks.append({"type": action_original_type, "content": content}) return chunks def __setattr__(self, name, value): new_variable = name not in self.__dict__ old_value = getattr(self, name, None) super().__setattr__(name, value) if name == 'action_count' and not new_variable: process_variable_changes(self._socketio, "actions", "Action Count", value, old_value) class KoboldWorldInfo(object): def __init__(self, socketio, story_settings, koboldai_vars, tokenizer=None): self._socketio = socketio self._koboldai_vars = koboldai_vars self.world_info = {} self.world_info_folder = OrderedDict() self.world_info_folder['root'] = [] self.story_settings = story_settings def reset(self): self.__init__(self._socketio, self.story_settings, self._koboldai_vars) def __iter__(self): self.itter = -1 return self def __next__(self): self.itter += 1 if self.itter < len(self.world_info): try: return self.world_info[list(self.world_info)[self.itter]].copy() except: print(self.itter) print(list(self.world_info)) raise else: raise StopIteration def __getitem__(self, i): return self.self.world_info[i].copy() def __len__(self): return len(self.world_info) def add_folder(self, folder): if folder in self.world_info_folder: i=0 while "{} {}".format(folder, i) in self.world_info_folder: i+=1 folder = "{} {}".format(folder, i) self.world_info_folder[folder] = [] self.story_settings.gamesaved = False self.sync_world_info_to_old_format() if self._socketio is not None: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") def delete_folder(self, folder): keys = [key for key in self.world_info] for key in keys: if self.world_info[key]['folder'] == folder: self.delete(key) if folder in self.world_info_folder: del self.world_info_folder[folder] if self._socketio is not None: self._socketio.emit("delete_world_info_folder", folder, broadcast=True, room="UI_2") logger.debug("Calcing AI Text from WI Folder Delete") ignore = self._koboldai_vars.calc_ai_text() def add_item_to_folder(self, uid, folder, before=None): if uid in self.world_info: #fiirst we need to remove the item from whatever folder it's in for temp in self.world_info_folder: if uid in self.world_info_folder[temp]: self.world_info_folder[temp].remove(uid) #Now we add it to the folder list if folder not in self.world_info_folder: self.world_info_folder[folder] = [] if before is None: self.world_info_folder[folder].append(uid) else: for i in range(len(self.world_info_folder[folder])): if self.world_info_folder[folder][i] == before: self.world_info_folder[folder].insert(i, uid) break #Finally, adjust the folder tag in the element self.world_info[uid]['folder'] = folder self.story_settings.gamesaved = False self.sync_world_info_to_old_format() if self._socketio is not None: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") def add_item(self, title, key, keysecondary, folder, constant, manual_text, comment, wi_type="wi", use_wpp=False, wpp={'name': "", 'type': "", 'format': "W++", 'attributes': {}}, v1_uid=None, recalc=True, sync=True, send_to_ui=True, object_type=None): if len(self.world_info) == 0: uid = 0 else: uid = max(self.world_info)+1 if use_wpp: if wpp['format'] == "W++": content = '[{}("{}")\n{{\n'.format(wpp['type'], wpp['name']) for attribute in wpp['attributes']: content = "{}{}({})\n".format(content, attribute, " + ".join(['"{}"'.format(x) for x in wpp['attributes'][attribute]])) content = "{}}}]".format(content) else: content = '[ {}: "{}";'.format(wpp['type'], wpp['name']) for attribute in wpp['attributes']: content = "{} {}: {};".format(content, attribute, ", ".join(['"{}"'.format(x) for x in wpp['attributes'][attribute]])) content = "{} ]".format(content[:-1]) else: content = manual_text if self._koboldai_vars.tokenizer is not None: token_length = len(self._koboldai_vars.tokenizer.encode(content)) else: token_length = 0 if folder is None: folder = "root" if isinstance(key, str): key = [x.strip() for x in key.split(",")] if key == [""]: key = [] if isinstance(keysecondary, str): keysecondary = [x.strip() for x in keysecondary.split(",")] if keysecondary == [""]: keysecondary = [] try: self.world_info[uid] = {"uid": uid, "title": title, "key": key, "keysecondary": keysecondary, "folder": folder, "constant": constant, 'manual_text': manual_text, "content": content, "comment": comment, "type": wi_type, "token_length": token_length, "selective": len(keysecondary) > 0, "used_in_game": constant, 'wpp': wpp, 'use_wpp': use_wpp, 'v1_uid': v1_uid, "object_type": object_type, } except: print("Error:") print(key) print(title) raise if folder not in self.world_info_folder: self.world_info_folder[folder] = [] self.world_info_folder[folder].append(uid) self.story_settings.gamesaved = False if sync: self.sync_world_info_to_old_format() self.story_settings.assign_world_info_to_actions(wuid=uid) if self._socketio is not None and send_to_ui: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") self._socketio.emit("world_info_entry", self.world_info[uid], broadcast=True, room="UI_2") if recalc: logger.debug("Calcing AI Text from WI Add") ignore = self._koboldai_vars.calc_ai_text() return uid def edit_item( self, uid, title, key, keysecondary, folder, constant, manual_text, comment, wi_type, use_wpp=False, before=None, wpp={'name': "", 'type': "", 'format': "W++", 'attributes': {}}, object_type=None, ): logger.debug("Editing World Info {}: {}".format(uid, title)) old_folder = self.world_info[uid]['folder'] #move the world info entry if the folder changed or if there is a new order requested if old_folder != folder or before is not None: self.add_item_to_folder(uid, folder, before=before) if use_wpp: if wpp['format'] == "W++": content = '[{}("{}")\n{{\n'.format(wpp['type'], wpp['name']) for attribute in wpp['attributes']: content = "{}{}({})\n".format(content, attribute, " + ".join(['"{}"'.format(x) for x in wpp['attributes'][attribute]])) content = "{}}}]".format(content) else: content = '[ {}: "{}";'.format(wpp['type'], wpp['name']) for attribute in wpp['attributes']: content = "{} {}: {};".format(content, attribute, ", ".join(['"{}"'.format(x) for x in wpp['attributes'][attribute]])) content = "{} ]".format(content[:-1]) else: content = manual_text if self._koboldai_vars.tokenizer is not None: token_length = len(self._koboldai_vars.tokenizer.encode(content)) else: token_length = 0 if folder is None: folder = "root" self.world_info[uid] = {"uid": uid, "title": title, "key": key, "keysecondary": keysecondary, "folder": folder, "constant": constant, 'manual_text': manual_text, "content": content, "comment": comment, "type": wi_type, "token_length": token_length, "selective": len(keysecondary) > 0, "used_in_game": constant, 'wpp': wpp, 'use_wpp': use_wpp, "object_type": object_type, } self.story_settings.gamesaved = False self.sync_world_info_to_old_format() self.story_settings.assign_world_info_to_actions(wuid=uid) logger.debug("Calcing AI Text from WI Edit") ignore = self._koboldai_vars.calc_ai_text() if self._socketio is not None: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") self._socketio.emit("world_info_entry", self.world_info[uid], broadcast=True, room="UI_2") def delete(self, uid): del self.world_info[uid] try: os.remove(os.path.join(self._koboldai_vars.save_paths.wi_images, str(uid))) except FileNotFoundError: pass for folder in self.world_info_folder: if uid in self.world_info_folder[folder]: self.world_info_folder[folder].remove(uid) self.story_settings.gamesaved = False self.sync_world_info_to_old_format() if self._socketio is not None: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") self._socketio.emit("delete_world_info_entry", uid, broadcast=True, room="UI_2") logger.debug("Calcing AI Text from WI Delete") ignore = self._koboldai_vars.calc_ai_text() def rename_folder(self, old_folder, folder): self.story_settings.gamesaved = False if folder in self.world_info_folder: i=0 while "{} {}".format(folder, i) in self.world_info_folder: i+=1 folder = "{} {}".format(folder, i) self.world_info_folder[folder] = self.world_info_folder[old_folder] #The folder dict is ordered and we want the new key to be in the same location as the last key. To do that we need to find all of the keys after old_folder after_old = False folder_order = [x for x in self.world_info_folder] for check_folder in folder_order: if check_folder == old_folder: after_old = True elif after_old and check_folder != folder: self.world_info_folder.move_to_end(check_folder) del self.world_info_folder[old_folder] # Need to change the folder properties on each affected world info item for uid in self.world_info: if self.world_info[uid]['folder'] == old_folder: self.world_info[uid]['folder'] = folder self.story_settings.gamesaved = False self.sync_world_info_to_old_format() if self._socketio is not None: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") def reorder(self, uid, before): self.add_item_to_folder(uid, self.world_info[before]['folder'], before=before) self.sync_world_info_to_old_format() def send_to_ui(self): if self._socketio is not None: self._socketio.emit("world_info_folder", {x: self.world_info_folder[x] for x in self.world_info_folder}, broadcast=True, room="UI_2") logger.debug("Sending all world info from send_to_ui") self._socketio.emit("world_info_entry", [self.world_info[uid] for uid in self.world_info], broadcast=True, room="UI_2") def to_json(self, folder=None): if folder is None: return { "folders": {x: self.world_info_folder[x] for x in self.world_info_folder}, "entries": self.world_info, } else: return { "folders": {x: self.world_info_folder[x] for x in self.world_info_folder if x == folder}, "entries": {x: self.world_info[x] for x in self.world_info if self.world_info[x]['folder'] == folder}, } def upgrade_entry(self, wi_entry: dict) -> dict: # If we do not have a type, or it is incorrect, set to WI. if wi_entry.get("type") not in ["constant", "chatcharacter", "wi", "commentator"]: wi_entry["type"] = "wi" if wi_entry["type"] in ["commentator", "constant"]: wi_entry["constant"] = True return wi_entry def load_json(self, data, folder=None): # Legacy WI images (stored in json) if "images" in data: for uid, image_b64 in data["images"].items(): image_b64 = image_b64.split(",")[-1] image_path = os.path.join( self._koboldai_vars.save_paths.wi_images, str(uid) ) with open(image_path, "wb") as file: file.write(base64.b64decode(image_b64)) data["entries"] = {k: self.upgrade_entry(v) for k,v in data["entries"].items()} if folder is None: self.world_info_folder = data['folders'] #Add the item start_time = time.time() for uid, item in data['entries'].items(): self.add_item(item['title'] if 'title' in item else item['key'][0], item['key'] if 'key' in item else [], item['keysecondary'] if 'keysecondary' in item else [], folder if folder is not None else item['folder'] if 'folder' in item else 'root', item['constant'] if 'constant' in item else False, item['manual_text'] if 'manual_text' in item else item['content'], item['comment'] if 'comment' in item else '', wi_type=item["type"], use_wpp=item['use_wpp'] if 'use_wpp' in item else False, wpp=item['wpp'] if 'wpp' in item else {'name': "", 'type': "", 'format': "W++", 'attributes': {}}, object_type=item.get("object_type"), recalc=False, sync=False) if folder is None: #self.world_info = {int(x): data['entries'][x] for x in data['entries']} self.world_info_folder = data['folders'] logger.debug("Load World Info took {}s".format(time.time()-start_time)) try: start_time = time.time() self.sync_world_info_to_old_format() logger.debug("Syncing WI2 to WI1 took {}s".format(time.time()-start_time)) except: print(self.world_info) print(data) raise self.send_to_ui() def sync_world_info_to_old_format(self): #Since the old UI uses world info entries for folders, we need to make some up folder_entries = {} i=-1 for folder in self.world_info_folder: folder_entries[folder] = i i-=1 #self.worldinfo_i = [] # List of World Info key/value objects sans uninitialized entries self.story_settings.worldinfo_i = [{ "comment": self.world_info[x]['comment'], "constant": self.world_info[x]['constant'], "content": self.world_info[x]['content'], "folder": folder_entries[self.world_info[x]['folder']], "init": True, "key": ",".join(self.world_info[x]['key']), "keysecondary": ",".join(self.world_info[x]['keysecondary']), "num": x, "selective": len(self.world_info[x]['keysecondary'])>0, "uid": self.world_info[x]['uid'] if 'v1_uid' not in self.world_info[x] or self.world_info[x]['v1_uid'] is None else self.world_info[x]['v1_uid'] } for x in self.world_info] #self.worldinfo = [] # List of World Info key/value objects self.story_settings.worldinfo = [x for x in self.story_settings.worldinfo_i] #We have to have an uninitialized blank entry for every folder or the old method craps out for folder in folder_entries: self.story_settings.worldinfo.append({ "comment": "", "constant": False, "content": "", "folder": folder_entries[folder], "init": False, "key": "", "keysecondary": "", "num": (0 if len(self.world_info) == 0 else max(self.world_info))+(folder_entries[folder]*-1), "selective": False, "uid": folder_entries[folder] }) #self.wifolders_d = {} # Dictionary of World Info folder UID-info pairs self.story_settings.wifolders_d = {folder_entries[x]: {'collapsed': False, 'name': x} for x in folder_entries} #self.worldinfo_u = {} # Dictionary of World Info UID - key/value pairs self.story_settings.worldinfo_u = {x['uid']: x for x in self.story_settings.worldinfo} #self.wifolders_l = [] # List of World Info folder UIDs self.story_settings.wifolders_l = [folder_entries[x] for x in folder_entries] #self.wifolders_u = {} # Dictionary of pairs of folder UID - list of WI UID self.story_settings.wifolders_u = {folder_entries[x]: [y for y in self.story_settings.worldinfo if y['folder'] == x] for x in folder_entries} def reset_used_in_game(self): for key in self.world_info: if self.world_info[key]["used_in_game"] != self.world_info[key]["constant"]: self.world_info[key]["used_in_game"] = self.world_info[key]["constant"] if self._socketio is not None: self._socketio.emit("world_info_entry_used_in_game", {"uid": key, "used_in_game": False}, broadcast=True, room="UI_2") def set_world_info_used(self, uid): if uid in self.world_info: self.world_info[uid]["used_in_game"] = True else: logger.warning("Something tried to set world info UID {} to in game, but it doesn't exist".format(uid)) if self._socketio is not None: self._socketio.emit("world_info_entry_used_in_game", {"uid": uid, "used_in_game": True}, broadcast=True, room="UI_2") def get_used_wi(self): return [x['content'] for x in self.world_info if x['used_in_game']] def to_wi_fewshot_format(self, excluding_uid: int) -> List[str]: """ Returns a list of strings representing applicable (has title, text, and object type) World Info entries. Intended for feeding into the fewshot generator. """ the_collection = [] for entry in self.world_info.values(): if entry["uid"] == excluding_uid: continue if not ( entry["title"] and entry["manual_text"] and entry["object_type"] ): continue processed_desc = entry["manual_text"].replace("\n", " ") while " " in processed_desc: processed_desc = processed_desc.replace(" ", " ") processed_desc = processed_desc.strip() the_collection.append( f"Title: {entry['title']}\n" \ f"Type: {entry['object_type']}\n" f"Description: {processed_desc}" ) return the_collection def get_commentators(self) -> List[dict]: ret = [] for entry in self.world_info.values(): if entry["type"] != "commentator": continue ret.append(entry) return ret @dataclass class SavePaths: base: str @property def required_paths(self) -> List[str]: return [ self.base, self.generated_audio, self.generated_images, self.wi_images ] @property def story(self) -> str: return os.path.join(self.base, "story.json") @property def generated_audio(self) -> str: return os.path.join(self.base, "generated_audio") @property def generated_images(self) -> str: return os.path.join(self.base, "generated_images") @property def wi_images(self) -> str: return os.path.join(self.base, "wi_images") default_rand_range = [0.44, 1, 2] default_creativity_range = [0.5, 1] default_rep_range = [1.0, 1.3] default_preset = { "preset": "Default", "description": "Known working settings.", "Match": "Recommended", "Preset Category": "Official", "temp": 0.5, "genamt": 80, "top_k": 0, "top_p": 0.9, "top_a": 0.0, "typical": 1.0, "tfs": 1.0, "rep_pen": 1.1, "rep_pen_range": 1024, "rep_pen_slope": 0.7, "sampler_order": [ 6, 0, 1, 2, 3, 4, 5 ] } badwordsids_default = [[6880], [50256], [42496], [4613], [17414], [22039], [16410], [27], [29], [38430], [37922], [15913], [24618], [28725], [58], [47175], [36937], [26700], [12878], [16471], [37981], [5218], [29795], [13412], [45160], [3693], [49778], [4211], [20598], [36475], [33409], [44167], [32406], [29847], [29342], [42669], [685], [25787], [7359], [3784], [5320], [33994], [33490], [34516], [43734], [17635], [24293], [9959], [23785], [21737], [28401], [18161], [26358], [32509], [1279], [38155], [18189], [26894], [6927], [14610], [23834], [11037], [14631], [26933], [46904], [22330], [25915], [47934], [38214], [1875], [14692], [41832], [13163], [25970], [29565], [44926], [19841], [37250], [49029], [9609], [44438], [16791], [17816], [30109], [41888], [47527], [42924], [23984], [49074], [33717], [31161], [49082], [30138], [31175], [12240], [14804], [7131], [26076], [33250], [3556], [38381], [36338], [32756], [46581], [17912], [49146]] # Tokenized array of badwords used to prevent AI artifacting badwordsids_neox = [[0], [1], [44162], [9502], [12520], [31841], [36320], [49824], [34417], [6038], [34494], [24815], [26635], [24345], [3455], [28905], [44270], [17278], [32666], [46880], [7086], [43189], [37322], [17778], [20879], [49821], [3138], [14490], [4681], [21391], [26786], [43134], [9336], [683], [48074], [41256], [19181], [29650], [28532], [36487], [45114], [46275], [16445], [15104], [11337], [1168], [5647], [29], [27482], [44965], [43782], [31011], [42944], [47389], [6334], [17548], [38329], [32044], [35487], [2239], [34761], [7444], [1084], [12399], [18990], [17636], [39083], [1184], [35830], [28365], [16731], [43467], [47744], [1138], [16079], [40116], [45564], [18297], [42368], [5456], [18022], [42696], [34476], [23505], [23741], [39334], [37944], [45382], [38709], [33440], [26077], [43600], [34418], [36033], [6660], [48167], [48471], [15775], [19884], [41533], [1008], [31053], [36692], [46576], [20095], [20629], [31759], [46410], [41000], [13488], [30952], [39258], [16160], [27655], [22367], [42767], [43736], [49694], [13811], [12004], [46768], [6257], [37471], [5264], [44153], [33805], [20977], [21083], [25416], [14277], [31096], [42041], [18331], [33376], [22372], [46294], [28379], [38475], [1656], [5204], [27075], [50001], [16616], [11396], [7748], [48744], [35402], [28120], [41512], [4207], [43144], [14767], [15640], [16595], [41305], [44479], [38958], [18474], [22734], [30522], [46267], [60], [13976], [31830], [48701], [39822], [9014], [21966], [31422], [28052], [34607], [2479], [3851], [32214], [44082], [45507], [3001], [34368], [34758], [13380], [38363], [4299], [46802], [30996], [12630], [49236], [7082], [8795], [5218], [44740], [9686], [9983], [45301], [27114], [40125], [1570], [26997], [544], [5290], [49193], [23781], [14193], [40000], [2947], [43781], [9102], [48064], [42274], [18772], [49384], [9884], [45635], [43521], [31258], [32056], [47686], [21760], [13143], [10148], [26119], [44308], [31379], [36399], [23983], [46694], [36134], [8562], [12977], [35117], [28591], [49021], [47093], [28653], [29013], [46468], [8605], [7254], [25896], [5032], [8168], [36893], [38270], [20499], [27501], [34419], [29547], [28571], [36586], [20871], [30537], [26842], [21375], [31148], [27618], [33094], [3291], [31789], [28391], [870], [9793], [41361], [47916], [27468], [43856], [8850], [35237], [15707], [47552], [2730], [41449], [45488], [3073], [49806], [21938], [24430], [22747], [20924], [46145], [20481], [20197], [8239], [28231], [17987], [42804], [47269], [29972], [49884], [21382], [46295], [36676], [34616], [3921], [26991], [27720], [46265], [654], [9855], [40354], [5291], [34904], [44342], [2470], [14598], [880], [19282], [2498], [24237], [21431], [16369], [8994], [44524], [45662], [13663], [37077], [1447], [37786], [30863], [42854], [1019], [20322], [4398], [12159], [44072], [48664], [31547], [18736], [9259], [31], [16354], [21810], [4357], [37982], [5064], [2033], [32871], [47446], [62], [22158], [37387], [8743], [47007], [17981], [11049], [4622], [37916], [36786], [35138], [29925], [14157], [18095], [27829], [1181], [22226], [5709], [4725], [30189], [37014], [1254], [11380], [42989], [696], [24576], [39487], [30119], [1092], [8088], [2194], [9899], [14412], [21828], [3725], [13544], [5180], [44679], [34398], [3891], [28739], [14219], [37594], [49550], [11326], [6904], [17266], [5749], [10174], [23405], [9955], [38271], [41018], [13011], [48392], [36784], [24254], [21687], [23734], [5413], [41447], [45472], [10122], [17555], [15830], [47384], [12084], [31350], [47940], [11661], [27988], [45443], [905], [49651], [16614], [34993], [6781], [30803], [35869], [8001], [41604], [28118], [46462], [46762], [16262], [17281], [5774], [10943], [5013], [18257], [6750], [4713], [3951], [11899], [38791], [16943], [37596], [9318], [18413], [40473], [13208], [16375]]