From a330d2069fd7b7e6803cc4266eef52981346aba0 Mon Sep 17 00:00:00 2001 From: ebolam Date: Fri, 16 Sep 2022 20:02:15 -0400 Subject: [PATCH] Possible relay fix :) --- aiserver.py | 28 +++++++++++++++------------- koboldai_settings.py | 1 + 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/aiserver.py b/aiserver.py index 4b1b73ce..f330d18d 100644 --- a/aiserver.py +++ b/aiserver.py @@ -472,11 +472,13 @@ def emit(*args, **kwargs): #replacement for tpool.execute to maintain request contexts def replacement_tpool_execute(function, *args, **kwargs): + temp = {} + socketio.start_background_task(replacement_tpool_execute_2, function, temp, *args, **kwargs).join() + print(temp) + return temp[1] - -def replacement_tpool_execute_2(function, queue, *args, **kwargs): - return_data = function(*args, **kwargs) - queue.put(return_data) +def replacement_tpool_execute_2(function, temp, *args, **kwargs): + temp[1] = function(*args, **kwargs) # marshmallow/apispec setup from apispec import APISpec @@ -3024,8 +3026,8 @@ def load_lua_scripts(): try: koboldai_vars.lua_koboldbridge.obliterate_multiverse() - tpool.execute(koboldai_vars.lua_koboldbridge.load_corescript, koboldai_vars.corescript) - koboldai_vars.has_genmod = tpool.execute(koboldai_vars.lua_koboldbridge.load_userscripts, filenames, modulenames, descriptions) + replacement_tpool_execute(koboldai_vars.lua_koboldbridge.load_corescript, koboldai_vars.corescript) + koboldai_vars.has_genmod = replacement_tpool_execute(koboldai_vars.lua_koboldbridge.load_userscripts, filenames, modulenames, descriptions) koboldai_vars.lua_running = True except lupa.LuaError as e: try: @@ -3523,7 +3525,7 @@ def execute_inmod(): koboldai_vars.lua_edited = set() koboldai_vars.lua_deleted = set() try: - tpool.execute(koboldai_vars.lua_koboldbridge.execute_inmod) + replacement_tpool_execute(koboldai_vars.lua_koboldbridge.execute_inmod) except lupa.LuaError as e: koboldai_vars.lua_koboldbridge.obliterate_multiverse() koboldai_vars.lua_running = False @@ -3542,7 +3544,7 @@ def execute_outmod(): setgamesaved(False) emit('from_server', {'cmd': 'hidemsg', 'data': ''}, broadcast=True, room="UI_1") try: - tpool.execute(koboldai_vars.lua_koboldbridge.execute_outmod) + replacement_tpool_execute(koboldai_vars.lua_koboldbridge.execute_outmod) except lupa.LuaError as e: koboldai_vars.lua_koboldbridge.obliterate_multiverse() koboldai_vars.lua_running = False @@ -4409,7 +4411,7 @@ def apiactionsubmit_generate(txt, minimum, maximum): torch.cuda.empty_cache() # Submit input text to generator - _genout, already_generated = tpool.execute(_generate, txt, minimum, maximum, set()) + _genout, already_generated = replacement_tpool_execute(_generate, txt, minimum, maximum, set()) genout = [applyoutputformatting(utils.decodenewlines(tokenizer.decode(tokens[-already_generated:]))) for tokens in _genout] @@ -4434,7 +4436,7 @@ def apiactionsubmit_tpumtjgenerate(txt, minimum, maximum): # Submit input text to generator soft_tokens = tpumtjgetsofttokens() - genout = tpool.execute( + genout = replacement_tpool_execute( tpu_mtj_backend.infer_static, np.uint32(txt), gen_len = maximum-minimum+1, @@ -4927,7 +4929,7 @@ def generate(txt, minimum, maximum, found_entries=None): # Submit input text to generator try: - genout, already_generated = tpool.execute(_generate, txt, minimum, maximum, found_entries) + genout, already_generated = replacement_tpool_execute(_generate, txt, minimum, maximum, found_entries) except Exception as e: if(issubclass(type(e), lupa.LuaError)): koboldai_vars.lua_koboldbridge.obliterate_multiverse() @@ -5352,7 +5354,7 @@ def tpumtjgenerate(txt, minimum, maximum, found_entries=None): past = np.empty((koboldai_vars.numseqs, 0), dtype=np.uint32) while(True): - genout, n_generated, regeneration_required, halt = tpool.execute( + genout, n_generated, regeneration_required, halt = replacement_tpool_execute( tpu_mtj_backend.infer_dynamic, context, gen_len = maximum-minimum+1, @@ -5394,7 +5396,7 @@ def tpumtjgenerate(txt, minimum, maximum, found_entries=None): ) else: - genout = tpool.execute( + genout = replacement_tpool_execute( tpu_mtj_backend.infer_static, np.uint32(txt), gen_len = maximum-minimum+1, diff --git a/koboldai_settings.py b/koboldai_settings.py index a98544eb..183b3083 100644 --- a/koboldai_settings.py +++ b/koboldai_settings.py @@ -34,6 +34,7 @@ def process_variable_changes(socketio, classname, name, value, old_value, debug_ 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(): + print("Had to use relay :(") data = ["var_changed", {"classname": classname, "name": name, "old_value": clean_var_for_emit(old_value), "value": clean_var_for_emit(value)}, {"include_self":True, "broadcast":True, "room":"UI_2"}] if queue is not None: queue.put(data)