Merge pull request #210 from one-some/UI2

Download check & fixes

aiserver.py

@@ -737,6 +737,11 @@ def get_config_filename(model_name = None):
     else:
         logger.warning(f"Empty configfile name sent back. Defaulting to ReadOnly")
         return(f"settings/ReadOnly.settings")
+
+def is_model_downloaded(model_name: str) -> bool:
+    model_stub = model_name.replace("/", "_")
+    return os.path.isdir(os.path.join("models", model_stub))
+
 #==================================================================#
 # Function to get model selection at startup
 #==================================================================#
@@ -755,10 +760,26 @@ def sendModelSelection(menu="mainmenu", folder="./models"):
         else:
             showdelete=False
         emit('from_server', {'cmd': 'show_model_menu', 'data': menu_list, 'menu': menu, 'breadcrumbs': breadcrumbs, "showdelete": showdelete}, broadcast=True, room="UI_1")
-        emit('show_model_menu', {'data': menu_list_ui_2, 'menu': menu, 'breadcrumbs': breadcrumbs, "showdelete": showdelete}, broadcast=False)
+
+        p_menu = [{
+            "label": m[0],
+            "name": m[1],
+            "size": m[2],
+            "isMenu": m[3],
+            "isDownloaded": True,
+        } for m in menu_list_ui_2]
+        emit('show_model_menu', {'data': p_menu, 'menu': menu, 'breadcrumbs': breadcrumbs, "showdelete": showdelete}, broadcast=False)
     else:
         emit('from_server', {'cmd': 'show_model_menu', 'data': model_menu[menu], 'menu': menu, 'breadcrumbs': [], "showdelete": False}, broadcast=True, room="UI_1")
-        emit('show_model_menu', {'data': model_menu[menu], 'menu': menu, 'breadcrumbs': [], "showdelete": False}, broadcast=False)
+
+        p_menu = [{
+            "label": m[0],
+            "name": m[1],
+            "size": m[2],
+            "isMenu": m[3],
+            "isDownloaded": is_model_downloaded(m[1]) if not m[3] else False,
+        } for m in model_menu[menu]]
+        emit('show_model_menu', {'data': p_menu, 'menu': menu, 'breadcrumbs': [], "showdelete": False}, broadcast=False)
 
 def get_folder_path_info(base):
     if base == 'This PC':
@@ -1545,7 +1566,7 @@ def get_layer_count(model, directory=""):
             from transformers import AutoConfig
             if(os.path.isdir(model.replace('/', '_'))):
                 model_config = AutoConfig.from_pretrained(model.replace('/', '_'), revision=koboldai_vars.revision, cache_dir="cache")
-            elif(os.path.isdir("models/{}".format(model.replace('/', '_')))):
+            elif(is_model_downloaded(model)):
                 model_config = AutoConfig.from_pretrained("models/{}".format(model.replace('/', '_')), revision=koboldai_vars.revision, cache_dir="cache")
             elif(os.path.isdir(directory)):
                 model_config = AutoConfig.from_pretrained(directory, revision=koboldai_vars.revision, cache_dir="cache")

static/koboldai.css

@@ -1595,8 +1595,8 @@ body {
 	border-radius: var(--radius_item_popup);
 	padding: 2px;
 	display: grid;
-	grid-template-areas: "folder_icon delete_icon edit_icon rename_icon file gpu_size warning_icon";
-	grid-template-columns: 30px 0px 0px 0px auto 50px 30px;
+	grid-template-areas: "folder_icon delete_icon edit_icon rename_icon file gpu_size warning_icon downloaded_icon";
+	grid-template-columns: 30px 0px 0px 0px auto 50px 30px 30px;
 	
 }
 

static/koboldai.js

@@ -1213,6 +1213,8 @@ function oai_engines(data) {
 }
 
 function getModelParameterCount(modelName) {
+	if (!modelName) return null;
+
 	// The "T" and "K" may be a little optimistic...
 	let paramsString = modelName.toUpperCase().match(/[\d.]+[TBMK]/)
 	if (!paramsString) return null;
@@ -1275,37 +1277,40 @@ function show_model_menu(data) {
 		var folder_icon = document.createElement("span");
 		folder_icon.classList.add("material-icons-outlined");
 		folder_icon.classList.add("cursor");
-		if ((item[3]) || (item[0] == 'Load a model from its directory') || (item[0] == 'Load an old GPT-2 model (eg CloverEdition)')) {
-			folder_icon.textContent = "folder";
-		} else {
-			folder_icon.textContent = "psychology";
-		}
+
+		let isModel = !(
+			item.isMenu ||
+			item.label === "Load a model from its directory" ||
+			item.label === "Load an old GPT-2 model (eg CloverEdition)"
+		);
+
+		folder_icon.textContent = isModel ? "psychology" : "folder";
 		list_item.append(folder_icon);
 		
 		
 		//create the actual item
 		var popup_item = document.createElement("span");
 		popup_item.classList.add("model");
-		popup_item.setAttribute("display_name", item[0]);
-		popup_item.id = item[1];
+		popup_item.setAttribute("display_name", item.label);
+		popup_item.id = item.name;
 		
 		popup_item.setAttribute("Menu", data.menu)
 		//name text
 		var text = document.createElement("span");
 		text.style="grid-area: item;";
-		text.textContent = item[0];
+		text.textContent = item.label;
 		popup_item.append(text);
 		//model size text
 		var text = document.createElement("span");
-		text.textContent = item[2];
+		text.textContent = item.size;
 		text.style="grid-area: gpu_size;padding: 2px;";
 		popup_item.append(text);
 
 		(function() {
 			// Anon function to avoid unreasonable indentation
-			if (folder_icon.innerText !== "psychology") return;
+			if (!isModel) return;
 
-			let parameterCount = getModelParameterCount(item[0]);
+			let parameterCount = getModelParameterCount(item.label);
 			if (!parameterCount) return;
 
 			let warningText = "";
@@ -1316,11 +1321,25 @@ function show_model_menu(data) {
 
 			if (!warningText) return;
 			$e("span", list_item, {
-				classes: ["material-icons-outlined"],
+				classes: ["material-icons-outlined", "model-size-warning"],
 				innerText: "warning",
 				"style.grid-area": "warning_icon",
 				tooltip: warningText
 			});
+
+		})();
+
+		(function() {
+			// Anon function to avoid unreasonable indentation
+			if (!item.isDownloaded) return;
+			if (!isModel) return;
+
+			$e("span", list_item, {
+				classes: ["material-icons-outlined", "model-download-notification"],
+				innerText: "download_done",
+				"style.grid-area": "downloaded_icon",
+				tooltip: "This model is already downloaded."
+			});
 		})();
 		
 		popup_item.onclick = function () {
@@ -2652,9 +2671,13 @@ function calc_token_usage(
 function Change_Theme(theme) {
 	var css = document.getElementById("CSSTheme");
     css.setAttribute("href", "/themes/"+theme+".css");
-	setTimeout(() => {
+
+	// We must wait for the style to load before we read it
+	css.onload = function() {
+		recolorTokens();
 		create_theming_elements();
-	}, "1000")
+	}
+
 	setCookie("theme", theme);
 	select = document.getElementById("selected_theme");
 	for (element of select.childNodes) {
@@ -2664,7 +2687,6 @@ function Change_Theme(theme) {
 			element.selected = false;
 		}
 	}
-	recolorTokens();
 }
 
 function palette_color(item) {