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)