From d9815d4b1f374d0f1947316c55f9d609e6b01b3f Mon Sep 17 00:00:00 2001
From: Henk <henk@henk.tech>
Date: Mon, 21 Aug 2023 17:52:21 +0200
Subject: [PATCH] New worker fixes

---
 aiserver.py                  |  6 ++++++
 environments/huggingface.yml |  2 ++
 environments/rocm.yml        |  2 ++
 logger.py                    | 13 +++++++++++++
 requirements.txt             |  2 ++
 5 files changed, 25 insertions(+)

diff --git a/aiserver.py b/aiserver.py
index 6c0456b2..24fb9146 100644
--- a/aiserver.py
+++ b/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
diff --git a/environments/huggingface.yml b/environments/huggingface.yml
index 004c7ecc..0ceb43ec 100644
--- a/environments/huggingface.yml
+++ b/environments/huggingface.yml
@@ -54,3 +54,5 @@ dependencies:
     - einops
     - peft==0.3.0
     - scipy
+    - windows-curses; sys_platform == 'win32'
+    - pynvml
diff --git a/environments/rocm.yml b/environments/rocm.yml
index 9538a615..9bf2813e 100644
--- a/environments/rocm.yml
+++ b/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
\ No newline at end of file
diff --git a/logger.py b/logger.py
index c0b8b8b0..6d37f780 100644
--- a/logger.py
+++ b/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")
diff --git a/requirements.txt b/requirements.txt
index 8dc7f9a2..dff40042 100644
--- a/requirements.txt
+++ b/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