Skip to content

Commit

Permalink
added --remote not allowing navigation outside of the model folder fo…
Browse files Browse the repository at this point in the history
…r custom models.

added a delete custom models option (will not delete models outside of the models directory, nor will it delete non-model directories)
  • Loading branch information
ebolam committed Jun 14, 2022
1 parent 780548f commit 462206f
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 8 deletions.
20 changes: 18 additions & 2 deletions aiserver.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,7 @@ def write(self, bar):
from flask_socketio import SocketIO, emit
app = Flask(__name__, root_path=os.getcwd())
app.config['SECRET KEY'] = 'secret!'
app.config['TEMPLATES_AUTO_RELOAD'] = True
socketio = SocketIO(app, async_method="eventlet")
print("{0}OK!{1}".format(colors.GREEN, colors.END))

Expand All @@ -364,6 +365,8 @@ def sendModelSelection(menu="mainmenu", folder="./models"):
#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'):
(paths, breadcrumbs) = get_folder_path_info(folder)
if args.remote:
breadcrumbs = []
menu_list = [[folder, menu, "", False] for folder in paths]
menu_list.append(["Return to Main Menu", "mainmenu", "", True])
emit('from_server', {'cmd': 'show_model_menu', 'data': menu_list, 'menu': menu, 'breadcrumbs': breadcrumbs}, broadcast=True)
Expand Down Expand Up @@ -2023,7 +2026,10 @@ def tpumtjgenerate_settings_callback() -> dict:
@app.route('/')
@app.route('/index')
def index():
return render_template('index.html', hide_ai_menu=args.noaimenu)
if 'new_ui' in request.args:
return render_template('index_new.html', hide_ai_menu=args.noaimenu)
else:
return render_template('index.html', hide_ai_menu=args.noaimenu)
@app.route('/favicon.ico')
def favicon():
return send_from_directory(app.root_path,
Expand Down Expand Up @@ -3090,7 +3096,17 @@ def get_message(msg):
get_model_info(msg['data'], directory=msg['path'])
else:
get_model_info(vars.model)

elif(msg['cmd'] == 'delete_model'):
if "{}/models".format(os.getcwd()) in msg['data'] or "{}\\models".format(os.getcwd()) in msg['data']:
if check_if_dir_is_model(msg['data']):
print("It's a model, now we really will kill it")
import shutil
shutil.rmtree(msg['data'])
sendModelSelection(menu=msg['menu'])
else:
print("Not a model, don't delete")
else:
print("Ah ah ah, you didn't say the magic word: The selected directory is not in the KoboldAI Models directory, not doing anything.")
elif(msg['cmd'] == 'OAI_Key_Update'):
get_oai_models(msg['key'])
elif(msg['cmd'] == 'loadselect'):
Expand Down
2 changes: 1 addition & 1 deletion pytest.ini
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
[pytest]
addopts = --ignore=miniconda3 --html=unit_test_report.html --self-contained-html -v
addopts = --ignore=miniconda3 --ignore=runtime --html=unit_test_report.html --self-contained-html -v
19 changes: 14 additions & 5 deletions static/application.js
Original file line number Diff line number Diff line change
Expand Up @@ -1017,6 +1017,14 @@ function buildLoadModelList(ar, menu, breadcrumbs) {
$("#loadmodellistbreadcrumbs").append("<hr size='1'>")
}
for(i=0; i<ar.length; i++) {
if (Array.isArray(ar[i][0])) {
full_path = ar[i][0][0];
folder = ar[i][0][1];
} else {
full_path = "";
folder = ar[i][0];
}

var html
html = "<div class=\"flex\">\
<div class=\"loadlistpadding\"></div>"
Expand All @@ -1027,13 +1035,14 @@ function buildLoadModelList(ar, menu, breadcrumbs) {
//this is a model
html = html + "<div class=\"loadlistpadding\"></div>"
}
if (Array.isArray(ar[i][0])) {
full_path = ar[i][0][0];
folder = ar[i][0][1];

//now let's do the delete icon if applicable
if (['NeoCustom', 'GPT2Custom'].includes(menu) && !ar[i][3]) {
html = html + "<span class=\"loadlisticon loadmodellisticon-folder oi oi-x allowed\" aria-hidden=\"true\" onclick='if(confirm(\"This will delete the selected folder with all contents. Are you sure?\")) { socket.send({\"cmd\": \"delete_model\", \"data\": \""+full_path.replaceAll("\\", "\\\\")+"\", \"menu\": \""+menu+"\"});}'></span>"
} else {
full_path = "";
folder = ar[i][0];
html = html + "<div class=\"loadlistpadding\"></div>"
}

html = html + "<div class=\"loadlistpadding\"></div>\
<div class=\"loadlistitem\" id=\"loadmodel"+i+"\" name=\""+ar[i][1]+"\" pretty_name=\""+full_path+"\">\
<div>"+folder+"</div>\
Expand Down
51 changes: 51 additions & 0 deletions static/custom.css
Original file line number Diff line number Diff line change
Expand Up @@ -1474,4 +1474,55 @@ body.connected .popupfooter, .popupfooter.always-available {

.model_layers:focus {
color: #cdf;
}

.menu_icon {
position: fixed;
top:10px;
left: 5px;
z-index:100;
display: inline-block;
cursor: pointer;
}

.SideMenu {
height: 100%;
width: 0;
position: fixed;
z-index: 1;
top: 0;
left: 0;
background-color: #111;
overflow-x: hidden;
transition: 0.5s;
padding-top: 60px;
}

.SideMenu.open {
width: 450px;
}

@media only screen and (max-width: 768px) {
.SideMenu.open {
width: 100%;
}
}


.menubar1, .menubar2, .menubar3 {
width: 21px;
height: 3px;
background-color: #999;
margin: 3px 0;
transition: 0.4s;
}

.change .menubar1 {
transform: translate(0px, 6px) rotate(-45deg);
}

.change .menubar2 {opacity: 0;}

.change .menubar3 {
transform: translate(0px, -6px) rotate(45deg);
}

0 comments on commit 462206f

Please sign in to comment.