Fix for model loading on web ui and removing AI menu when using remote/colab methods

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(vars.model == "GPT2Custom"):
-        if(os.path.isdir(directory)):
+            model_config = open(vars.custmodpth + "/config.json", "r")
-            try:
+        # Get the model_type from the config or assume a model type if it isn't present
-                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
         else:
-            try:
+            from transformers import AutoConfig
-                model_config = AutoConfig.from_pretrained(model, cache_dir="cache/")
+            if vars.custmodpth == "":
-            except ValueError as e:
+                model_config = AutoConfig.from_pretrained(vars.model, revision=vars.revision, cache_dir="cache")
-                model_config = None
+            elif(os.path.isdir(vars.custmodpth.replace('/', '_'))):
-        if model_config is None:
+                model_config = AutoConfig.from_pretrained(vars.custmodpth.replace('/', '_'), revision=vars.revision, cache_dir="cache")
-            return None
+            elif(os.path.isdir("models/{}".format(vars.custmodpth.replace('/', '_')))):
-        try:
+                model_config = AutoConfig.from_pretrained("models/{}".format(vars.custmodpth.replace('/', '_')), revision=vars.revision, cache_dir="cache")
-            layers = model_config.num_layers if hasattr(model_config, "num_layers") else model_config.n_layer
+            else:
-        except:
+                model_config = AutoConfig.from_pretrained(vars.custmodpth, revision=vars.revision, cache_dir="cache")
-            layers = None
+        
-            pass
+        
-        return layers
+        
+        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()

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">