From db2fbbb17a261f85b9a3016ed6a3496aa33e03a4 Mon Sep 17 00:00:00 2001
From: ebolam <ebolam@gmail.com>
Date: Mon, 24 Oct 2022 19:16:03 -0400
Subject: [PATCH] Colab Debug

---
 aiserver.py        | 5 ++++-
 tpu_mtj_backend.py | 1 +
 utils.py           | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index bfeb35b4..41dff108 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -519,6 +519,7 @@ logger.add(UI_2_log_history, serialize=True, colorize=True, enqueue=True, level=
 koboldai_vars = koboldai_settings.koboldai_vars(socketio)
 
 utils.koboldai_vars = koboldai_vars
+utils.socketio = socketio
 
 old_socketio_on = socketio.on
 def new_socketio_on(*a, **k):
@@ -3853,7 +3854,7 @@ def execute_outmod():
 #==================================================================#
 @socketio.on('connect')
 def do_connect():
-    logger.info("Client connected!")
+    logger.info("Client connected! UI_{}".format(request.args.get('ui')))
     if request.args.get("rely") == "true":
         return
     join_room("UI_{}".format(request.args.get('ui')))
@@ -9595,6 +9596,7 @@ def _generate_text(body: GenerationInputSchema):
         }}), mimetype="application/json", status=503))
     if koboldai_vars.use_colab_tpu:
         import tpu_mtj_backend
+        tpu_mtj_backend.socketio = socketio
     if hasattr(body, "sampler_seed"):
         # If a seed was specified, we need to save the global RNG state so we
         # can restore it later
@@ -11950,6 +11952,7 @@ def put_config_sampler_seed(body: SamplerSeedSettingSchema):
     """
     if koboldai_vars.use_colab_tpu:
         import tpu_mtj_backend
+        tpu_mtj_backend.socketio = socketio
         tpu_mtj_backend.set_rng_seed(body.value)
     else:
         import torch
diff --git a/tpu_mtj_backend.py b/tpu_mtj_backend.py
index 492102a8..b1cddca0 100644
--- a/tpu_mtj_backend.py
+++ b/tpu_mtj_backend.py
@@ -118,6 +118,7 @@ def show_spinner():
     bar = progressbar.ProgressBar(max_value=progressbar.UnknownLength, widgets=[progressbar.Timer(), '  ', progressbar.BouncingBar(left='[', right=']', marker='█')])
     i = 0
     while True:
+        print("Sending to client")
         if i % 2 == 0:
             socketio.emit('from_server', {'cmd': 'model_load_status', 'data': "Connecting to TPU..." }, broadcast=True, room="UI_1")
         else:
diff --git a/utils.py b/utils.py
index d07fd5be..b8baa786 100644
--- a/utils.py
+++ b/utils.py
@@ -190,7 +190,7 @@ def _download_with_aria2(aria2_config: str, total_length: int, directory: str =
                 try:
                     print('\r' + bar, end='')
                     try:
-                        emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", "&nbsp;")}, broadcast=True, room="UI_1")
+                        socketio.emit('from_server', {'cmd': 'model_load_status', 'data': bar.replace(" ", "&nbsp;")}, broadcast=True, room="UI_1")
                     except:
                         pass
                     eventlet.sleep(seconds=0)