New worker fixes

aiserver.py

@@ -61,6 +61,12 @@ import gc
 import traceback
 
 import lupa
+# Hack to make the new Horde worker understand its imports...
+try:
+    sys.path.append(os.path.abspath("AI-Horde-Worker"))
+    print(os.path.abspath("AI-Horde-Worker/"))
+except:
+    pass
 
 # KoboldAI
 import fileops

environments/huggingface.yml

@@ -54,3 +54,5 @@ dependencies:
     - einops
     - peft==0.3.0
     - scipy
+    - windows-curses; sys_platform == 'win32'
+    - pynvml

environments/rocm.yml

@@ -44,3 +44,5 @@ dependencies:
     - git+https://github.com/0cc4m/hf_bleeding_edge/
     - einops
     - peft==0.3.0
+    - windows-curses; sys_platform == 'win32'
+    - pynvml

logger.py

@@ -17,6 +17,8 @@ class Colors:
 STDOUT_LEVELS = ["GENERATION", "PROMPT"]
 INIT_LEVELS = ["INIT", "INIT_OK", "INIT_WARN", "INIT_ERR"]
 MESSAGE_LEVELS = ["MESSAGE"]
+STATS_LEVELS = ["STATS"]
+
 # By default we're at error level or higher
 verbosity = 20
 quiet = 0
@@ -54,6 +56,16 @@ def is_msg_log(record):
         return(False)
     return(True)
 
+def is_stats_log(record):
+    if record["level"].name not in STATS_LEVELS:
+        return False
+    return True
+
+def is_not_stats_log(record):
+    if record["level"].name in STATS_LEVELS:
+        return False
+    return True
+    
 def is_stderr_log(record):
     if record["level"].name in STDOUT_LEVELS + INIT_LEVELS + MESSAGE_LEVELS:
         return(False)
@@ -91,6 +103,7 @@ logger.level("INIT_ERR", no=31, color="<red>")
 # Messages contain important information without which this application might not be able to be used
 # As such, they have the highest priority
 logger.level("MESSAGE", no=61, color="<green>")
+logger.level("STATS", no=19, color="<blue>")
 
 logger.__class__.generation = partialmethod(logger.__class__.log, "GENERATION")
 logger.__class__.prompt = partialmethod(logger.__class__.log, "PROMPT")

requirements.txt

@@ -42,3 +42,5 @@ git+https://github.com/0cc4m/hf_bleeding_edge/
 einops
 peft==0.3.0
 scipy
+windows-curses; sys_platform == 'win32'
+pynvml