Merge pull request #5 from henk717/united

Merge united.

aiserver.py

@@ -452,6 +452,7 @@ def emit(*args, **kwargs):
         return _emit(*args, **kwargs)
     except AttributeError:
         return socketio.emit(*args, **kwargs)
+utils.emit = emit
 
 # marshmallow/apispec setup
 from apispec import APISpec
@@ -879,7 +880,7 @@ def device_config(config):
                     print(f"{colors.RED}Please enter an integer between -1 and {n_layers}.{colors.END}")
 
     logger.init_ok("Final device configuration:", status="Info")
-    device_list(n_layers)
+    device_list(n_layers, primary=breakmodel.primary_device)
 
     # If all layers are on the same device, use the old GPU generation mode
     while(len(breakmodel.gpu_blocks) and breakmodel.gpu_blocks[-1] == 0):
@@ -1360,6 +1361,8 @@ def general_startup(override_args=None):
         args = parser.parse_args(shlex.split(os.environ["KOBOLDAI_ARGS"]))
     else:
         args = parser.parse_args()
+    
+    utils.args = args
 
     set_logger_verbosity(args.verbosity)
     quiesce_logger(args.quiesce)
@@ -1796,7 +1799,9 @@ def patch_transformers():
         if not args.no_aria2:
             utils.aria2_hook(pretrained_model_name_or_path, **kwargs)
         return old_from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
-    PreTrainedModel.from_pretrained = new_from_pretrained
+    if(not hasattr(PreTrainedModel, "_kai_patched")):
+        PreTrainedModel.from_pretrained = new_from_pretrained
+        PreTrainedModel._kai_patched = True
     if(hasattr(modeling_utils, "get_checkpoint_shard_files")):
         old_get_checkpoint_shard_files = modeling_utils.get_checkpoint_shard_files
         def new_get_checkpoint_shard_files(pretrained_model_name_or_path, index_filename, *args, **kwargs):
@@ -2662,7 +2667,9 @@ def load_model(use_gpu=True, gpu_layers=None, disk_layers=None, initial_load=Fal
             if not args.no_aria2:
                 utils.aria2_hook(pretrained_model_name_or_path, **kwargs)
             return old_from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
-        PreTrainedModel.from_pretrained = new_from_pretrained
+        if(not hasattr(PreTrainedModel, "_kai_patched")):
+            PreTrainedModel.from_pretrained = new_from_pretrained
+            PreTrainedModel._kai_patched = True
         if(hasattr(modeling_utils, "get_checkpoint_shard_files")):
             old_get_checkpoint_shard_files = modeling_utils.get_checkpoint_shard_files
             def new_get_checkpoint_shard_files(pretrained_model_name_or_path, index_filename, *args, **kwargs):

colab/GPU.ipynb

@@ -70,10 +70,18 @@
         "Model = \"Nerys 2.7B\" #@param [\"Nerys 2.7B\", \"AID 2.7B\", \"Erebus 2.7B\", \"Janeway 2.7B\", \"Picard 2.7B\", \"Horni LN 2.7B\", \"Horni 2.7B\", \"Shinen 2.7B\", \"OPT 2.7B\", \"Fairseq Dense 2.7B\", \"Neo 2.7B\"] {allow-input: true}\n",
         "Version = \"Official\" #@param [\"Official\", \"United\"] {allow-input: true}\n",
         "Provider = \"Localtunnel\" #@param [\"Localtunnel\", \"Cloudflare\"]\n",
+		"use_google_drive = True #@param {type:\"boolean\"}\n",
         "\n",
         "!nvidia-smi\n",
         "from google.colab import drive\n",
-        "drive.mount('/content/drive/')\n",
+        "if use_google_drive:\n",
+		"  drive.mount('/content/drive/')\n",
+		"else:\n",
+		"  import os\n",
+		"  if not os.path.exists(\"/content/drive\"):\n",
+		"    os.mkdir(\"/content/drive\")\n",
+		"  if not os.path.exists(\"/content/drive/MyDrive/\"):\n",
+		"    os.mkdir(\"/content/drive/MyDrive/\")\n",
         "\n",
         "if Model == \"Nerys 2.7B\":\n",
         "  Model = \"KoboldAI/fairseq-dense-2.7B-Nerys\"\n",

colab/TPU.ipynb

@@ -69,6 +69,7 @@
         "Model = \"Nerys 13B V2\" #@param [\"Nerys 13B V2\", \"Erebus 13B\", \"Janeway 13B\", \"Shinen 13B\", \"Skein 20B\", \"Erebus 20B\", \"Skein 6B\", \"Janeway 6B\", \"Adventure 6B\", \"Shinen 6B\", \"Lit V2 6B\", \"Lit 6B\", \"NeoX 20B\", \"OPT 13B\", \"Fairseq Dense 13B\", \"GPT-J-6B\"] {allow-input: true}\n",
         "Version = \"Official\" #@param [\"Official\", \"United\"] {allow-input: true}\n",
         "Provider = \"Localtunnel\" #@param [\"Localtunnel\", \"Cloudflare\"]\n",
+		"use_google_drive = True #@param {type:\"boolean\"}\n",
         "\n",
         "import os\n",
         "try:\n",
@@ -79,7 +80,14 @@
         "    raise RuntimeError(\"⚠️You can not run this notebook without the TPU accelerator, go to Runtime->Sessions, terminate your session and then try again.⚠️\")\n",
         "print('Now we will need your Google Drive to store settings and saves, you must login with the same account you used for Colab.')\n",
         "from google.colab import drive\n",
-        "drive.mount('/content/drive/')\n",
+        "if use_google_drive:\n",
+		"  drive.mount('/content/drive/')\n",
+		"else:\n",
+		"  import os\n",
+		"  if not os.path.exists(\"/content/drive\"):\n",
+		"    os.mkdir(\"/content/drive\")\n",
+		"  if not os.path.exists(\"/content/drive/MyDrive/\"):\n",
+		"    os.mkdir(\"/content/drive/MyDrive/\")\n",
         "\n",
         "if Model == \"Janeway 13B\":\n",
         "  Model = \"KoboldAI/fairseq-dense-13B-Janeway\"\n",

commandline.bat

@@ -1,5 +1,7 @@
 @echo off
 cd /D %~dp0
+SET CONDA_SHLVL=
+
 TITLE CMD for KoboldAI Runtime
 SET /P M=<loader.settings
 IF %M%==1 GOTO drivemap

environments/huggingface.yml

@@ -20,9 +20,12 @@ dependencies:
   - marshmallow>=3.13
   - apispec-webframeworks
   - loguru
+  - termcolor
   - pip:
     - flask-cloudflared
     - flask-ngrok
     - lupa==1.10
     - transformers>=4.20.1
+    - huggingface_hub>=0.10.1
     - accelerate
+    - git+https://github.com/VE-FORBRYDERNE/mkultra

environments/rocm.yml

@@ -17,6 +17,7 @@ dependencies:
   - marshmallow>=3.13
   - apispec-webframeworks
   - loguru
+  - termcolor
   - pip:
     - --extra-index-url https://download.pytorch.org/whl/rocm5.1.1
     - torch
@@ -25,4 +26,6 @@ dependencies:
     - flask-ngrok
     - lupa==1.10
     - transformers>=4.20.1
+    - huggingface_hub>=0.10.1
     - accelerate
+    - git+https://github.com/VE-FORBRYDERNE/mkultra

fileops.py

@@ -86,7 +86,7 @@ def uspath(filename):
 def getstoryfiles():
     list = []
     for file in listdir("stories"):
-        if file.endswith(".json"):
+        if file.endswith(".json") and not file.endswith(".v2.json"):
             ob = {}
             ob["name"] = file.replace(".json", "")
             f = open("stories/"+file, "r")

install_requirements.bat

@@ -8,6 +8,7 @@ echo.
 
 Reg add "HKLM\SYSTEM\CurrentControlSet\Control\FileSystem" /v "LongPathsEnabled" /t REG_DWORD /d "1" /f 2>nul
 cd /D %~dp0
+SET CONDA_SHLVL=
 
 if exist miniconda3\ (
   echo Delete existing installation?

play.bat

@@ -1,5 +1,7 @@
 @echo off
 cd /D %~dp0
+SET CONDA_SHLVL=
+
 TITLE KoboldAI - Server
 SET /P M=<loader.settings
 IF %M%==1 GOTO drivemap

requirements.txt

@@ -1,4 +1,5 @@
 transformers>=4.20.1
+huggingface_hub>=0.10.1
 Flask
 Flask-SocketIO
 requests
@@ -16,3 +17,5 @@ flask-session
 marshmallow>=3.13
 apispec-webframeworks
 loguru
+termcolor
+git+https://github.com/VE-FORBRYDERNE/mkultra

requirements_mtj.txt

@@ -6,6 +6,7 @@ dm-haiku == 0.0.5
 jax == 0.2.21
 jaxlib >= 0.1.69, <= 0.3.7
 transformers >= 4.20.1
+huggingface_hub >= 0.10.1
 progressbar2
 git+https://github.com/VE-FORBRYDERNE/mesh-transformer-jax@ck
 flask

update-koboldai.bat

@@ -1,5 +1,7 @@
 @echo off
 cd /d %~dp0
+SET CONDA_SHLVL=
+
 TITLE KoboldAI - Updater
 SET /P M=<loader.settings
 IF %M%==1 GOTO drivemap

utils.py

@@ -27,6 +27,7 @@ except ImportError:
     HAS_ACCELERATE = False
 
 vars = None
+args = None
 num_shards: Optional[int] = None
 current_shard = 0
 from_pretrained_model_name = ""
@@ -40,6 +41,8 @@ named_buffers: Optional[List[tuple]] = None
 
 default_sampler_order = [6, 0, 1, 2, 3, 4, 5]
 
+emit = None
+
 #==================================================================#
 # Decorator to prevent a function's actions from being run until
 # at least x seconds have passed without the function being called
@@ -198,6 +201,7 @@ def _download_with_aria2(aria2_config: str, total_length: int, directory: str =
             pass
     
     import transformers
+    aria2_port = 6799 if vars is None else vars.aria2_port
     lengths = {}
     s = requests.Session()
     s.mount("http://", requests.adapters.HTTPAdapter(max_retries=requests.adapters.Retry(total=120, backoff_factor=1)))
@@ -208,9 +212,9 @@ def _download_with_aria2(aria2_config: str, total_length: int, directory: str =
         with tempfile.NamedTemporaryFile("w+b", delete=False) as f:
             f.write(aria2_config)
             f.flush()
-            p = subprocess.Popen(["aria2c", "-x", "10", "-s", "10", "-j", "10", "--enable-rpc=true", f"--rpc-secret={secret}", "--rpc-listen-port", str(vars.aria2_port), "--disable-ipv6", "--file-allocation=trunc", "--allow-overwrite", "--auto-file-renaming=false", "-d", directory, "-i", f.name, "-U", transformers.file_utils.http_user_agent(user_agent)] + (["-c"] if not force_download else []) + ([f"--header='Authorization: Bearer {use_auth_token}'"] if use_auth_token else []), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            p = subprocess.Popen(["aria2c", "-x", "10", "-s", "10", "-j", "10", "--enable-rpc=true", f"--rpc-secret={secret}", "--rpc-listen-port", str(aria2_port), "--disable-ipv6", "--file-allocation=trunc", "--allow-overwrite", "--auto-file-renaming=false", "-d", directory, "-i", f.name, "-U", transformers.file_utils.http_user_agent(user_agent)] + (["-c"] if not force_download else []) + ([f"--header='Authorization: Bearer {use_auth_token}'"] if use_auth_token else []), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
             while p.poll() is None:
-                r = s.post(f"http://localhost:{vars.aria2_port}/jsonrpc", json={"jsonrpc": "2.0", "id": "kai", "method": "aria2.tellActive", "params": [f"token:{secret}"]}).json()["result"]
+                r = s.post(f"http://localhost:{aria2_port}/jsonrpc", json={"jsonrpc": "2.0", "id": "kai", "method": "aria2.tellActive", "params": [f"token:{secret}"]}).json()["result"]
                 if not r:
                     s.close()
                     if bar is not None:
@@ -602,4 +606,4 @@ def get_missing_module_names(model: PreTrainedModel, names: List[str]) -> List[s
             else:
                 recurse(c[1], head=name + ".")
     recurse(model)
-    return missing_names
+    return missing_names