From dd07b10b73728c1fb476a6dc1f12a82dd26fb8a8 Mon Sep 17 00:00:00 2001
From: ebolam <ebolam@gmail.com>
Date: Mon, 6 Jun 2022 13:57:19 -0400
Subject: [PATCH] Fix for model loading on web ui and removing AI menu when
 using remote/colab methods

---
 aiserver.py          | 48 +++++++++++++++++++++-----------------------
 templates/index.html |  2 ++
 2 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 54805596..c0bd9fd9 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -348,7 +348,9 @@ print("{0}OK!{1}".format(colors.GREEN, colors.END))
 def sendModelSelection(menu="mainmenu"):
     #If we send one of the manual load options, send back the list of model directories, otherwise send the menu
     if menu in ('NeoCustom', 'GPT2Custom'):
-        emit('from_server', {'cmd': 'show_model_menu', 'data': [[folder, menu, "", False] for folder in next(os.walk('./models'))[1]], 'menu': 'custom'}, broadcast=True)
+        menu_list = [[folder, menu, "", False] for folder in next(os.walk('./models'))[1]]
+        menu_list.append(["Read Only (No AI)", "ReadOnly", "", True])
+        emit('from_server', {'cmd': 'show_model_menu', 'data': menu_list, 'menu': 'custom'}, broadcast=True)
     else:
         emit('from_server', {'cmd': 'show_model_menu', 'data': model_menu[menu], 'menu': menu}, broadcast=True)
 
@@ -932,30 +934,26 @@ def general_startup():
 #==================================================================#  
 def get_layer_count(model, directory=""):
     if(model not in ["InferKit", "Colab", "OAI", "GooseAI" , "ReadOnly", "TPUMeshTransformerGPTJ"]):
-        from transformers import AutoConfig
-        if(os.path.isdir(directory)):
-            try:
-                model_config = AutoConfig.from_pretrained(directory, cache_dir="cache/")
-            except ValueError as e:
-                model_config = None
-        elif(os.path.isdir("models/{}".format(model.replace('/', '_')))):
-            try:
-                model_config = AutoConfig.from_pretrained("models/{}".format(model.replace("/", "_")), cache_dir="cache/")
-            except ValueError as e:
-                model_config = None
+        if(vars.model == "GPT2Custom"):
+            model_config = open(vars.custmodpth + "/config.json", "r")
+        # Get the model_type from the config or assume a model type if it isn't present
         else:
-            try:
-                model_config = AutoConfig.from_pretrained(model, cache_dir="cache/")
-            except ValueError as e:
-                model_config = None
-        if model_config is None:
-            return None
-        try:
-            layers = model_config.num_layers if hasattr(model_config, "num_layers") else model_config.n_layer
-        except:
-            layers = None
-            pass
-        return layers
+            from transformers import AutoConfig
+            if vars.custmodpth == "":
+                model_config = AutoConfig.from_pretrained(vars.model, revision=vars.revision, cache_dir="cache")
+            elif(os.path.isdir(vars.custmodpth.replace('/', '_'))):
+                model_config = AutoConfig.from_pretrained(vars.custmodpth.replace('/', '_'), revision=vars.revision, cache_dir="cache")
+            elif(os.path.isdir("models/{}".format(vars.custmodpth.replace('/', '_')))):
+                model_config = AutoConfig.from_pretrained("models/{}".format(vars.custmodpth.replace('/', '_')), revision=vars.revision, cache_dir="cache")
+            else:
+                model_config = AutoConfig.from_pretrained(vars.custmodpth, revision=vars.revision, cache_dir="cache")
+        
+        
+        
+        return utils.num_layers(model_config)
+    else:
+        return None
+
 
 def get_oai_models(dummy=True):
     if vars.oaiapikey != "":
@@ -1891,7 +1889,7 @@ def load_model(use_gpu=True, key='', gpu_layers=None, initial_load=False):
 @app.route('/')
 @app.route('/index')
 def index():
-    return render_template('index.html')
+    return render_template('index.html', hide_ai_menu=args.remote)
 @app.route('/download')
 def download():
     save_format = request.args.get("format", "json").strip().lower()
diff --git a/templates/index.html b/templates/index.html
index eca0c815..933decf6 100644
--- a/templates/index.html
+++ b/templates/index.html
@@ -33,6 +33,7 @@
 						</button>
 						<div class="collapse navbar-collapse" id="navbarNavDropdown">
 							<ul class="nav navbar-nav">
+								{% if not hide_ai_menu %}
 								<li class="nav-item dropdown">
 									<a class="nav-link dropdown-toggle" href="#" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">AI</a>
 									<div class="dropdown-menu">
@@ -40,6 +41,7 @@
 										<a class="dropdown-item" href="#" id="btn_showmodel">Model Info</a>
 									</div>
 								</li>
+								{% endif %}
 								<li class="nav-item dropdown">
 									<a class="nav-link dropdown-toggle" href="#" data-toggle="dropdown" aria-haspopup="true" aria-expanded="false">New Game</a>
 									<div class="dropdown-menu">