diff --git a/aiserver.py b/aiserver.py index 248f049d..3f111f07 100644 --- a/aiserver.py +++ b/aiserver.py @@ -26,7 +26,6 @@ import gensettings from utils import debounce import utils import structures -import breakmodel #==================================================================# # Variables & Storage @@ -111,6 +110,8 @@ class vars: useprompt = True # Whether to send the full prompt with every submit action breakmodel = False # For GPU users, whether to use both system RAM and VRAM to conserve VRAM while offering speedup compared to CPU-only bmsupported = False # Whether the breakmodel option is supported (GPT-Neo/GPT-J only, currently) + smandelete = False # Whether stories can be deleted from inside the browser + smanrename = False # Whether stories can be renamed from inside the browser acregex_ai = re.compile(r'\n* *>(.|\n)*') # Pattern for matching adventure actions from the AI so we can remove them acregex_ui = re.compile(r'^ *(>.*)$', re.MULTILINE) # Pattern for matching actions in the HTML-escaped story so we can apply colouring, etc (make sure to encase part to format in parentheses) actionmode = 1 @@ -173,12 +174,17 @@ parser.add_argument("--path", help="Specify the Path for local models (For model parser.add_argument("--cpu", action='store_true', help="By default unattended launches are on the GPU use this option to force CPU usage.") parser.add_argument("--breakmodel", action='store_true', help="For models that support GPU-CPU hybrid generation, use this feature instead of GPU or CPU generation") parser.add_argument("--breakmodel_layers", type=int, help="Specify the number of layers to commit to system RAM if --breakmodel is used") +parser.add_argument("--override_delete", action='store_true', help="Deleting stories from inside the browser is disabled if you are using --remote and enabled otherwise. Using this option will instead allow deleting stories if using --remote and prevent deleting stories otherwise.") +parser.add_argument("--override_rename", action='store_true', help="Renaming stories from inside the browser is disabled if you are using --remote and enabled otherwise. Using this option will instead allow renaming stories if using --remote and prevent renaming stories otherwise.") args = parser.parse_args() vars.model = args.model; if args.remote: vars.remote = True; +vars.smandelete = vars.remote == args.override_delete +vars.smanrename = vars.remote == args.override_rename + # Select a model to run if args.model: print("Welcome to KoboldAI!\nYou have selected the following Model:", vars.model) @@ -363,7 +369,7 @@ log.setLevel(logging.ERROR) # Start flask & SocketIO print("{0}Initializing Flask... {1}".format(colors.PURPLE, colors.END), end="") -from flask import Flask, render_template +from flask import Flask, render_template, Response, request from flask_socketio import SocketIO, emit app = Flask(__name__) app.config['SECRET KEY'] = 'secret!' @@ -385,6 +391,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]): if(vars.usegpu): generator = pipeline('text-generation', model=model, tokenizer=tokenizer, device=0) elif(vars.breakmodel): # Use both RAM and VRAM (breakmodel) + import breakmodel n_layers = model.config.num_layers breakmodel.total_blocks = n_layers model.half().to('cpu') @@ -435,6 +442,7 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]): if(vars.usegpu): generator = pipeline('text-generation', model=vars.model, device=0) elif(vars.breakmodel): # Use both RAM and VRAM (breakmodel) + import breakmodel model = AutoModel.from_pretrained(vars.model) n_layers = model.config.num_layers breakmodel.total_blocks = n_layers @@ -494,9 +502,17 @@ def index(): return render_template('index.html') @app.route('/download') def download(): - # Leave Edit/Memory mode before continuing - exitModes() - + save_format = request.args.get("format", "json").strip().lower() + + if(save_format == "plaintext"): + txt = vars.prompt + "".join(vars.actions.values()) + save = Response(txt) + filename = path.basename(vars.savedir) + if filename[-5:] == ".json": + filename = filename[:-5] + save.headers.set('Content-Disposition', 'attachment', filename='%s.txt' % filename) + return(save) + # Build json to write js = {} js["gamestarted"] = vars.gamestarted @@ -516,8 +532,12 @@ def download(): "selective": wi["selective"], "constant": wi["constant"] }) - save = flask.Response(json.dumps(js, indent=3)) - save.headers.set('Content-Disposition', 'attachment', filename='%s.json' % path.basename(vars.savedir)) + + save = Response(json.dumps(js, indent=3)) + filename = path.basename(vars.savedir) + if filename[-5:] == ".json": + filename = filename[:-5] + save.headers.set('Content-Disposition', 'attachment', filename='%s.json' % filename) return(save) #============================ METHODS =============================# @@ -528,7 +548,7 @@ def download(): @socketio.on('connect') def do_connect(): print("{0}Client connected!{1}".format(colors.GREEN, colors.END)) - emit('from_server', {'cmd': 'connected'}) + emit('from_server', {'cmd': 'connected', 'smandelete': vars.smandelete, 'smanrename': vars.smanrename}) if(vars.remote): emit('from_server', {'cmd': 'runs_remotely'}) @@ -597,11 +617,11 @@ def get_message(msg): deleterequest() elif(msg['cmd'] == 'memory'): togglememorymode() - elif(msg['cmd'] == 'savetofile'): + elif(not vars.remote and msg['cmd'] == 'savetofile'): savetofile() - elif(msg['cmd'] == 'loadfromfile'): + elif(not vars.remote and msg['cmd'] == 'loadfromfile'): loadfromfile() - elif(msg['cmd'] == 'import'): + elif(not vars.remote and msg['cmd'] == 'import'): importRequest() elif(msg['cmd'] == 'newgame'): newGameRequest() @@ -714,7 +734,11 @@ def get_message(msg): elif(msg['cmd'] == 'loadselect'): vars.loadselect = msg["data"] elif(msg['cmd'] == 'loadrequest'): - loadRequest(getcwd()+"/stories/"+vars.loadselect+".json") + loadRequest(fileops.storypath(vars.loadselect)) + elif(msg['cmd'] == 'deletestory'): + deletesave(msg['data']) + elif(msg['cmd'] == 'renamestory'): + renamesave(msg['data'], msg['newname']) elif(msg['cmd'] == 'clearoverwrite'): vars.svowname = "" vars.saveow = False @@ -738,7 +762,7 @@ def get_message(msg): vars.adventure = msg['data'] settingschanged() refresh_settings() - elif(msg['cmd'] == 'importwi'): + elif(not vars.remote and msg['cmd'] == 'importwi'): wiimportrequest() #==================================================================# @@ -1848,16 +1872,62 @@ def saveas(name): name = utils.cleanfilename(name) if(not fileops.saveexists(name) or (vars.saveow and vars.svowname == name)): # All clear to save - saveRequest(getcwd()+"/stories/"+name+".json") - emit('from_server', {'cmd': 'hidesaveas', 'data': ''}) + e = saveRequest(fileops.storypath(name)) vars.saveow = False vars.svowname = "" + if(e is None): + emit('from_server', {'cmd': 'hidesaveas', 'data': ''}) + else: + print("{0}{1}{2}".format(colors.RED, str(e), colors.END)) + emit('from_server', {'cmd': 'popuperror', 'data': str(e)}) else: # File exists, prompt for overwrite vars.saveow = True vars.svowname = name emit('from_server', {'cmd': 'askforoverwrite', 'data': ''}) +#==================================================================# +# Launch in-browser story-delete prompt +#==================================================================# +def deletesave(name): + name = utils.cleanfilename(name) + e = fileops.deletesave(name) + if(e is None): + if(vars.smandelete): + emit('from_server', {'cmd': 'hidepopupdelete', 'data': ''}) + getloadlist() + else: + emit('from_server', {'cmd': 'popuperror', 'data': "The server denied your request to delete this story"}) + else: + print("{0}{1}{2}".format(colors.RED, str(e), colors.END)) + emit('from_server', {'cmd': 'popuperror', 'data': str(e)}) + +#==================================================================# +# Launch in-browser story-rename prompt +#==================================================================# +def renamesave(name, newname): + # Check if filename exists already + name = utils.cleanfilename(name) + newname = utils.cleanfilename(newname) + if(not fileops.saveexists(newname) or name == newname or (vars.saveow and vars.svowname == newname)): + e = fileops.renamesave(name, newname) + vars.saveow = False + vars.svowname = "" + if(e is None): + if(vars.smanrename): + emit('from_server', {'cmd': 'hidepopuprename', 'data': ''}) + getloadlist() + else: + emit('from_server', {'cmd': 'popuperror', 'data': "The server denied your request to rename this story"}) + else: + print("{0}{1}{2}".format(colors.RED, str(e), colors.END)) + emit('from_server', {'cmd': 'popuperror', 'data': str(e)}) + else: + # File exists, prompt for overwrite + vars.saveow = True + vars.svowname = newname + emit('from_server', {'cmd': 'askforoverwrite', 'data': ''}) + #==================================================================# # Save the currently running story #==================================================================# @@ -1906,33 +1976,30 @@ def saveRequest(savpath): "constant": wi["constant"] }) - ln = len(vars.actions) - - if(ln > 0): - chunks = collections.deque() - i = 0 - for key in reversed(vars.actions): - chunk = vars.actions[key] - chunks.appendleft(chunk) - i += 1 - - if(ln > 0): - txt = vars.prompt + "".join(chunks) - elif(ln == 0): - txt = vars.prompt + txt = vars.prompt + "".join(vars.actions.values()) # Write it - file = open(savpath, "w") + try: + file = open(savpath, "w") + except Exception as e: + return e try: file.write(json.dumps(js, indent=3)) - finally: + except Exception as e: file.close() - - file = open(txtpath, "w") + return e + file.close() + + try: + file = open(txtpath, "w") + except Exception as e: + return e try: file.write(txt) - finally: + except Exception as e: file.close() + return e + file.close() print("{0}Story saved to {1}!{2}".format(colors.GREEN, path.basename(savpath), colors.END)) diff --git a/fileops.py b/fileops.py index 7ae96836..7fcc44cb 100644 --- a/fileops.py +++ b/fileops.py @@ -1,6 +1,7 @@ import tkinter as tk from tkinter import filedialog from os import getcwd, listdir, path +import os import json #==================================================================# @@ -54,6 +55,12 @@ def getdirpath(dir, title): else: return None +#==================================================================# +# Returns the path (as a string) to the given story by its name +#==================================================================# +def storypath(name): + return path.join(path.dirname(path.realpath(__file__)), "stories", name + ".json") + #==================================================================# # Returns an array of dicts containing story files in /stories #==================================================================# @@ -83,4 +90,22 @@ def getstoryfiles(): # Returns True if json file exists with requested save name #==================================================================# def saveexists(name): - return path.exists(path.dirname(path.realpath(__file__))+"/stories/"+name+".json") + return path.exists(storypath(name)) + +#==================================================================# +# Delete save file by name; returns None if successful, or the exception if not +#==================================================================# +def deletesave(name): + try: + os.remove(storypath(name)) + except Exception as e: + return e + +#==================================================================# +# Rename save file; returns None if successful, or the exception if not +#==================================================================# +def renamesave(name, new_name): + try: + os.replace(storypath(name), storypath(new_name)) + except Exception as e: + return e diff --git a/static/application.js b/static/application.js index b9b6a527..3a0a457f 100644 --- a/static/application.js +++ b/static/application.js @@ -49,7 +49,6 @@ var saveasinput; var topic; var saveas_accept; var saveas_close; -var saveasoverwrite; var loadpopup; var loadcontent; var load_accept; @@ -70,6 +69,8 @@ var connected = false; var newly_loaded = true; var current_editing_chunk = null; var chunk_conflict = false; +var sman_allow_delete = false; +var sman_allow_rename = false; // Key states var shift_down = false; @@ -286,7 +287,7 @@ function addWiLine(ob) { disableWiSelective(ob.num); }); $("#constant-key-"+ob.num).on("click", function () { - element = $("#constant-key-"+ob.num); + var element = $("#constant-key-"+ob.num); if(element.hasClass("constant-key-icon-enabled")) { socket.send({'cmd': 'wiconstantoff', 'data': ob.num}); element.removeClass("constant-key-icon-enabled") @@ -542,7 +543,7 @@ function hideSaveAsPopup() { saveaspopup.removeClass("flex"); saveaspopup.addClass("hidden"); saveasinput.val(""); - hide([saveasoverwrite]); + hide([$(".saveasoverwrite"), $(".popuperror")]); } function sendSaveAsRequest() { @@ -566,20 +567,67 @@ function buildLoadList(ar) { showLoadPopup(); var i; for(i=0; i\ -
"+ar[i].name+"
\ -
"+ar[i].actions+"
\ + loadcontent.append("
\ +
\ + \ +
\ + \ +
\ +
\ +
"+ar[i].name+"
\ +
"+ar[i].actions+"
\ +
\
"); $("#load"+i).on("click", function () { enableButtons([load_accept]); socket.send({'cmd': 'loadselect', 'data': $(this).attr("name")}); highlightLoadLine($(this)); }); + + $("#loaddelete"+i).off("click").on("click", (function (name) { + return function () { + if(!sman_allow_delete) { + return; + } + $("#loadcontainerdelete-storyname").text(name); + $("#btn_dsaccept").off("click").on("click", (function (name) { + return function () { + hide([$(".saveasoverwrite"), $(".popuperror")]); + socket.send({'cmd': 'deletestory', 'data': name}); + } + })(name)); + $("#loadcontainerdelete").removeClass("hidden").addClass("flex"); + } + })(ar[i].name)); + + $("#loadrename"+i).off("click").on("click", (function (name) { + return function () { + if(!sman_allow_rename) { + return; + } + $("#newsavename").val("") + $("#loadcontainerrename-storyname").text(name); + var submit = (function (name) { + return function () { + hide([$(".saveasoverwrite"), $(".popuperror")]); + socket.send({'cmd': 'renamestory', 'data': name, 'newname': $("#newsavename").val()}); + } + })(name); + $("#btn_rensaccept").off("click").on("click", submit); + $("#newsavename").off("keydown").on("keydown", function (ev) { + if (ev.which == 13 && $(this).val() != "") { + submit(); + } + }); + $("#loadcontainerrename").removeClass("hidden").addClass("flex"); + $("#newsavename").val(name).select(); + } + })(ar[i].name)); } } function highlightLoadLine(ref) { - $("#loadlistcontent > div").removeClass("popuplistselected"); + $("#loadlistcontent > div > div.popuplistselected").removeClass("popuplistselected"); ref.addClass("popuplistselected"); } @@ -687,26 +735,26 @@ function chunkOnKeyDown(event) { switch(event.keyCode) { case 37: // left case 39: // right - old_range = getSelection().getRangeAt(0); - old_range_start = old_range.startOffset; - old_range_end = old_range.endOffset; - old_range_ancestor = old_range.commonAncestorContainer; - old_range_start_container = old_range.startContainer; - old_range_end_container = old_range.endContainer; + var old_range = getSelection().getRangeAt(0); + var old_range_start = old_range.startOffset; + var old_range_end = old_range.endOffset; + var old_range_ancestor = old_range.commonAncestorContainer; + var old_range_start_container = old_range.startContainer; + var old_range_end_container = old_range.endContainer; setTimeout(function () { // Wait a few milliseconds and check if the caret has moved - new_selection = getSelection(); - new_range = new_selection.getRangeAt(0); + var new_selection = getSelection(); + var new_range = new_selection.getRangeAt(0); if(old_range_start != new_range.startOffset || old_range_end != new_range.endOffset || old_range_ancestor != new_range.commonAncestorContainer || old_range_start_container != new_range.startContainer || old_range_end_container != new_range.endContainer) { return; } // If it hasn't moved, we're at the beginning or end of a chunk // and the caret must be moved to a different chunk - chunk = document.activeElement; + var chunk = document.activeElement; switch(event.keyCode) { case 37: // left if((chunk = chunk.previousSibling) && chunk.tagName == "CHUNK") { - range = document.createRange(); + var range = document.createRange(); range.selectNodeContents(chunk); range.collapse(false); new_selection.removeAllRanges(); @@ -723,7 +771,7 @@ function chunkOnKeyDown(event) { return; case 8: // backspace - old_length = document.activeElement.innerText.length; + var old_length = document.activeElement.innerText.length; setTimeout(function () { // Wait a few milliseconds and compare the chunk's length if(old_length != document.activeElement.innerText.length) { @@ -731,8 +779,8 @@ function chunkOnKeyDown(event) { } // If it's the same, we're at the beginning of a chunk if((chunk = document.activeElement.previousSibling) && chunk.tagName == "CHUNK") { - range = document.createRange(); - selection = getSelection(); + var range = document.createRange(); + var selection = getSelection(); range.selectNodeContents(chunk); range.collapse(false); selection.removeAllRanges(); @@ -771,7 +819,7 @@ function submitEditedChunk(event) { return; } - chunk = current_editing_chunk; + var chunk = current_editing_chunk; current_editing_chunk = null; // Submit the edited chunk if it's not empty, otherwise delete it @@ -795,6 +843,8 @@ $(document).ready(function(){ button_save = $('#btn_save'); button_saveas = $('#btn_saveas'); button_savetofile = $('#btn_savetofile'); + button_download = $('#btn_download'); + button_downloadtxt= $('#btn_downloadtxt'); button_load = $('#btn_load'); button_loadfrfile = $('#btn_loadfromfile'); button_import = $("#btn_import"); @@ -833,7 +883,6 @@ $(document).ready(function(){ topic = $("#topic"); saveas_accept = $("#btn_saveasaccept"); saveas_close = $("#btn_saveasclose"); - saveasoverwrite = $("#saveasoverwrite"); loadpopup = $("#loadcontainer"); loadcontent = $("#loadlistcontent"); load_accept = $("#btn_loadaccept"); @@ -853,6 +902,8 @@ $(document).ready(function(){ socket.on('from_server', function(msg) { if(msg.cmd == "connected") { // Connected to Server Actions + sman_allow_delete = msg.hasOwnProperty("smandelete") && msg.smandelete; + sman_allow_rename = msg.hasOwnProperty("smanrename") && msg.smanrename; connected = true; connect_status.html("Connected to KoboldAI Process!"); connect_status.removeClass("color_orange"); @@ -870,7 +921,7 @@ $(document).ready(function(){ } }); } else if(msg.cmd == "updatescreen") { - _gamestarted = gamestarted; + var _gamestarted = gamestarted; gamestarted = msg.gamestarted; if(_gamestarted != gamestarted) { action_mode = 0; @@ -1048,6 +1099,14 @@ $(document).ready(function(){ } else if(msg.cmd == "popupshow") { // Show/Hide Popup popupShow(msg.data); + } else if(msg.cmd == "hidepopupdelete") { + // Hide the dialog box that asks you to confirm deletion of a story + $("#loadcontainerdelete").removeClass("flex").addClass("hidden"); + hide([$(".saveasoverwrite"), $(".popuperror")]); + } else if(msg.cmd == "hidepopuprename") { + // Hide the story renaming dialog box + $("#loadcontainerrename").removeClass("flex").addClass("hidden"); + hide([$(".saveasoverwrite"), $(".popuperror")]); } else if(msg.cmd == "addimportline") { // Add import popup entry addImportLine(msg.data); @@ -1081,7 +1140,11 @@ $(document).ready(function(){ buildLoadList(msg.data); } else if(msg.cmd == "askforoverwrite") { // Show overwrite warning - show([saveasoverwrite]); + show([$(".saveasoverwrite")]); + } else if(msg.cmd == "popuperror") { + // Show error in the current dialog box + $(".popuperror").text(msg.data); + show([$(".popuperror")]); } else if(msg.cmd == "genseqs") { // Parse generator sequences to UI parsegenseqs(msg.data); @@ -1211,6 +1274,14 @@ $(document).ready(function(){ saveas_accept.on("click", function(ev) { sendSaveAsRequest(); }); + + button_download.on("click", function(ev) { + window.open("/download", "_blank"); + }); + + button_downloadtxt.on("click", function(ev) { + window.open("/download?format=plaintext", "_blank"); + }); button_load.on("click", function(ev) { socket.send({'cmd': 'loadlistrequest', 'data': ''}); @@ -1238,6 +1309,25 @@ $(document).ready(function(){ ns_close.on("click", function(ev) { hideNewStoryPopup(); }); + + $("#btn_dsclose").on("click", function () { + $("#loadcontainerdelete").removeClass("flex").addClass("hidden"); + hide([$(".saveasoverwrite"), $(".popuperror")]); + }); + + $("#newsavename").on("input", function (ev) { + if($(this).val() == "") { + disableButtons([$("#btn_rensaccept")]); + } else { + enableButtons([$("#btn_rensaccept")]); + } + hide([$(".saveasoverwrite"), $(".popuperror")]); + }); + + $("#btn_rensclose").on("click", function () { + $("#loadcontainerrename").removeClass("flex").addClass("hidden"); + hide([$(".saveasoverwrite"), $(".popuperror")]); + }); button_rndgame.on("click", function(ev) { showRandomStoryPopup(); @@ -1262,7 +1352,7 @@ $(document).ready(function(){ } else { enableButtons([saveas_accept]); } - hide([saveasoverwrite]); + hide([$(".saveasoverwrite"), $(".popuperror")]); }); // Bind Enter button to submit diff --git a/static/custom.css b/static/custom.css index 19349702..d9363914 100644 --- a/static/custom.css +++ b/static/custom.css @@ -271,12 +271,6 @@ chunk, chunk * { margin-top: 200px; } -#saveasoverwrite { - color: #ff9900; - font-weight: bold; - text-align: center; -} - #loadpopup { width: 500px; background-color: #262626; @@ -291,6 +285,18 @@ chunk, chunk * { } } +#loadpopupdelete { + width: 350px; + background-color: #262626; + margin-top: 200px; +} + +#loadpopuprename { + width: 350px; + background-color: #262626; + margin-top: 200px; +} + #loadlistcontent { height: 325px; overflow-y: scroll; @@ -319,6 +325,12 @@ chunk, chunk * { text-align: center; } +.dialogheader { + padding: 10px 40px 10px 40px; + color: #737373; + text-align: center; +} + .anotelabel { font-size: 10pt; color: #ffffff; @@ -556,16 +568,16 @@ chunk, chunk * { } .loadlistheader { - padding-left: 10px; - display: grid; - grid-template-columns: 80% 20%; + padding-left: 68px; + padding-right: 20px; + display: flex; color: #737373; } .loadlistitem { padding: 5px 10px 5px 10px; - display: grid; - grid-template-columns: 80% 20%; + display: flex; + flex-grow: 1; color: #ffffff; -moz-transition: background-color 0.25s ease-in; @@ -579,6 +591,30 @@ chunk, chunk * { background-color: #688f1f; } +.loadlistpadding { + padding-right: 10px; +} + +.loadlisticon { + color: #333 +} + +.loadlisticon.allowed { + color: #ddd +} + +.loadlisticon.allowed:hover { + cursor: pointer; +} + +.loadlisticon-delete.allowed:hover { + color: #ef2929 +} + +.loadlisticon-rename.allowed:hover { + color: #fce94f +} + .navbar .navbar-nav .nav-link:hover { border-radius: 5px; background-color: #98bcdb; @@ -678,6 +714,11 @@ chunk, chunk * { font-size: 12pt; } +.popuperror { + color: #ef2929; + text-align: center; +} + .popupfooter { width: 100%; padding: 10px; @@ -692,6 +733,12 @@ chunk, chunk * { margin-right: 10px; } +.saveasoverwrite { + color: #ff9900; + font-weight: bold; + text-align: center; +} + .seqselheader { color: #737373; } diff --git a/templates/index.html b/templates/index.html index 085d9245..e9bf94e4 100644 --- a/templates/index.html +++ b/templates/index.html @@ -41,6 +41,8 @@ Save Save As Save To File... + Download Story as JSON + Download Story as Plaintext
  • @@ -193,7 +195,10 @@
    - + +