Merge pull request #241 from one-some/screenshot

Screenshot tool

aiserver.py

@@ -8526,7 +8526,12 @@ def UI_2_download_story():
 def UI_2_Set_Selected_Text(data):
     if not koboldai_vars.quiet:
         print("Updating Selected Text: {}".format(data))
-    koboldai_vars.actions[int(data['id'])] = data['text']
+    action_id = int(data["id"])
+
+    if not koboldai_vars.actions.actions[action_id].get("Original Text"):
+        koboldai_vars.actions.actions[action_id]["Original Text"] = data["text"]
+
+    koboldai_vars.actions[action_id] = data['text']
 
 #==================================================================#
 # Event triggered when Option is Selected
@@ -9093,6 +9098,41 @@ def UI_2_set_commentator_image(commentator_id):
         file.write(data)
     return ":)"
 
+@app.route("/image_db.json", methods=["GET"])
+@logger.catch
+def UI_2_get_image_db():
+    try:
+        return send_file(os.path.join(koboldai_vars.save_paths.generated_images, "db.json"))
+    except FileNotFoundError:
+        return jsonify([])
+
+@app.route("/action_composition.json", methods=["GET"])
+@logger.catch
+def UI_2_get_action_composition():
+    try:
+        actions = request.args.get("actions").split(",")
+        if not actions:
+            raise ValueError()
+    except (ValueError, AttributeError):
+        return "No actions", 400
+
+    try:
+        actions = [int(action) for action in actions]
+    except TypeError:
+        return "Not all actions int", 400
+
+    ret = []
+    for action_id in actions:
+        try:
+            ret.append(koboldai_vars.actions.get_action_composition(action_id))
+        except KeyError:
+            ret.append([])
+    return jsonify(ret)
+
+@app.route("/generated_images/<path:path>")
+def UI_2_send_generated_images(path):
+    return send_from_directory(koboldai_vars.save_paths.generated_images, path)
+
 @socketio.on("scratchpad_prompt")
 @logger.catch
 def UI_2_scratchpad_prompt(data):

koboldai_settings.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 from dataclasses import dataclass
+import difflib
 import importlib
 import os, re, time, threading, json, pickle, base64, copy, tqdm, datetime, sys
 import shutil
@@ -1166,6 +1167,12 @@ class user_settings(settings):
         self.cluster_requested_models = [] # The models which we allow to generate during cluster mode
         self.wigen_use_own_wi = False
         self.wigen_amount = 80
+        self.screenshot_show_attribution = True
+        self.screenshot_show_story_title = True
+        self.screenshot_show_author_name = True
+        self.screenshot_author_name = "Anonymous"
+        self.screenshot_show_model_name = True
+        self.screenshot_use_boring_colors = False
         
         
     def __setattr__(self, name, value):
@@ -1504,6 +1511,9 @@ class KoboldStoryRegister(object):
             if "Time" not in json_data["actions"][item]:
                 json_data["actions"][item]["Time"] = int(time.time())
 
+            if "Original Text" not in json_data["actions"][item]:
+                json_data["actions"][item]["Original Text"] = json_data["actions"][item]["Selected Text"]
+
             temp[int(item)] = json_data['actions'][item]
             if int(item) >= self.action_count-100: #sending last 100 items to UI
                 data_to_send.append({"id": item, 'action':  temp[int(item)]})
@@ -1539,6 +1549,9 @@ class KoboldStoryRegister(object):
         self.action_count+=1
         action_id = self.action_count + action_id_offset
         if action_id in self.actions:
+            if not self.actions[action_id].get("Original Text"):
+                self.actions[action_id]["Original Text"] = text
+
             if self.actions[action_id]["Selected Text"] != text:
                 self.actions[action_id]["Selected Text"] = text
                 self.actions[action_id]["Time"] = self.actions[action_id].get("Time", int(time.time()))
@@ -1567,6 +1580,8 @@ class KoboldStoryRegister(object):
                 "Options": [],
                 "Probabilities": [],
                 "Time": int(time.time()),
+                "Original Text": text,
+                "Origin": "user" if submission else "ai"
             }
 
         if submission:
@@ -2095,6 +2110,44 @@ class KoboldStoryRegister(object):
             return filename, prompt
         return None, None
     
+    def get_action_composition(self, action_id: int) -> List[dict]:
+        """
+        Returns a list of chunks that comprise an action in dictionaries
+        formatted as follows:
+            type: string identifying chunk type ("ai", "user", "edit", or "prompt")
+            content: the actual content of the chunk
+        """
+        # Prompt doesn't need standard edit data
+        if action_id == -1:
+            if self.koboldai_vars.prompt:
+                return [{"type": "prompt", "content": self.koboldai_vars.prompt}]
+            return []
+
+        current_text = self.actions[action_id]["Selected Text"]
+        action_original_type = self.actions[action_id].get("Origin", "ai")
+        original = self.actions[action_id]["Original Text"]
+
+        matching_blocks = difflib.SequenceMatcher(
+            None,
+            self.actions[action_id]["Original Text"],
+            current_text
+        ).get_matching_blocks()
+
+        chunks = []
+        base = 0
+        for chunk_match in matching_blocks:
+            inserted = current_text[base:chunk_match.b]
+            content = current_text[chunk_match.b:chunk_match.b + chunk_match.size]
+
+            base = chunk_match.b + chunk_match.size
+
+            if inserted:
+                chunks.append({"type": "edit", "content": inserted})
+            if content:
+                chunks.append({"type": action_original_type, "content": content})
+
+        return chunks
+
     def __setattr__(self, name, value):
         new_variable = name not in self.__dict__
         old_value = getattr(self, name, None)

static/koboldai.css

@@ -1220,6 +1220,182 @@ td.server_vars {
 	margin-right: 5px;
 }
 
+#screenshot-wizard {
+	width: 50%;
+	background-color: var(--background);
+}
+
+#screenshot-wizard > .popup-header {
+	padding-top: 5px;
+	padding-left: 5px;
+	margin-bottom: 6px;
+}
+
+#screenshot-wizard > .help_text {
+	margin-top: 0.7em;
+}
+
+#screenshot-wizard > .popup-header > h1 {
+	margin-top: 0px;
+	margin-bottom: 0px;
+}
+
+#screenshot-target-container {
+	position: relative;
+	margin: 20px;
+	margin-top: 10px;
+	max-height: 25vh;
+	overflow-y: auto;
+}
+
+#screenshot-target {
+	min-height: 130px;
+	background-color: var(--gamescreen_background);
+	padding: 20px;
+	/* width: 1080px; */
+	display: flex;
+	overflow: hidden;
+}
+
+#screenshot-target > #screenshot-left {
+	display: flex;
+	flex-direction: column;
+	justify-content: space-between;
+	flex-grow: 1;
+}
+
+#screenshot-target > #screenshot-images {
+	border-left: 1px solid var(--background);
+	flex-shrink: 0;
+	max-width: 25%;
+	padding-left: 20px;
+	margin-left: 20px;
+	position: relative;
+	overflow: hidden;
+}
+
+#screenshot-target > #screenshot-images > img {
+	display: block;
+	width: 100%; /* the buggy line */
+	margin-bottom: 12px;
+}
+
+#screenshot-target .action-text {
+	font-size: 1.1em;
+	font-weight: bold;
+}
+
+#robot-attribution {
+	margin-top: 30px;
+}
+
+#screenshot-target:not(.boring-colors) .human-text,
+#screenshot-target:not(.boring-colors) .edit-text,
+#screenshot-target:not(.boring-colors) .prompt-text,
+#robot-attribution #human-attribution {
+	color: var(--text_edit);
+}
+
+#screenshot-target .ai-text,
+#robot-attribution #model-name {
+	color: var(--text);
+}
+
+#robot-attribution #model-name {
+	font-family: monospace;
+}
+
+#robot-attribution #story-attribution {
+	font-style: italic;
+}
+
+#robot-attribution #human-attribution,
+#robot-attribution #model-name,
+#robot-attribution #story-attribution {
+	font-weight: bold;
+}
+
+#robot-attribution .boring-part-of-madlibs {
+	opacity: 0.9;
+}
+
+#screenshot-options {
+	display: flex;
+	flex-direction: column;
+	margin-left: 8px;
+	margin-bottom: 8px;
+}
+
+#screenshot-wizard .header {
+	font-weight: bold;
+	opacity: 0.6;
+}
+
+#screenshot-wizard .option {
+	margin-bottom: 3px;
+}
+
+#screenshot-options .indent {
+	margin-left: 15px;
+}
+
+#screenshot-options .warning {
+	color: rgb(185, 70, 70);
+	font-style: italic;
+}
+
+/* Screenshot Image Picker */
+#screenshot-image-picker {
+	display: grid;
+	grid-template-columns: auto auto auto auto auto;
+	max-height: 200px;
+	overflow-y: auto;
+	overflow-x: hidden;
+	/* TODO: We REALLY need to figure out a better system for colors and themeing */
+	background-color: var(--flyout_background_pinned);
+}
+
+#screenshot-image-picker > .img-container {
+	display: inline-block;
+	position: relative;
+	cursor: pointer;
+	border: 2px solid var(----enabled_button_background_color);
+}
+
+#screenshot-image-picker > .img-container > img {
+	height: 100%;
+	width: 100%;
+	object-fit: cover;
+	overflow: hidden;
+	aspect-ratio: 1/1;
+}
+
+#screenshot-image-picker > .img-container > input {
+	position: absolute;
+	bottom: 5px;
+	right: 5px;
+	z-index: 2;
+}
+
+#screenshot-image-picker > .img-container > input:checked ~ img {
+	opacity: 0.6;
+}
+
+#sw-download {
+	margin-top: 12px;
+	margin: 3px;
+	padding: 12px;
+	display: flex;
+	justify-content: center;
+	align-items: center;
+	background-color: var(--button_background);
+	font-size: 1.2em;
+	cursor: pointer;
+}
+
+#screenshot-text-container {
+	line-height: 2;
+}
 
 /* ---------------------------- OVERALL PAGE CONFIG ------------------------------*/
 body {

static/koboldai.js

@@ -117,6 +117,8 @@ const context_menu_actions = {
 		{label: "Add to World Info Entry", icon: "auto_stories", enabledOn: "SELECTION", click: push_selection_to_world_info},
 		{label: "Add as Bias", icon: "insights", enabledOn: "SELECTION", click: push_selection_to_phrase_bias},
 		{label: "Retry from here", icon: "refresh", enabledOn: "CARET", click: retry_from_here},
+		null,
+		{label: "Take Screenshot", icon: "screenshot_monitor", enabledOn: "SELECTION", click: screenshot_selection},
 		// Not implemented! See view_selection_probabiltiies
 		// null,
 		// {label: "View Token Probabilities", icon: "assessment", enabledOn: "SELECTION", click: view_selection_probabilities},
@@ -3179,7 +3181,7 @@ function privacy_mode(enabled) {
 		document.getElementById('SideMenu').classList.remove("superblur");
 		document.getElementById('main-grid').classList.remove("superblur");
 		document.getElementById('rightSideMenu').classList.remove("superblur");
-		closePopups();
+		if (!$el("#privacy_mode").classList.contains("hidden")) closePopups();
 		document.getElementById('privacy_password').value = "";
 	}
 }
@@ -4651,7 +4653,7 @@ async function downloadDebugFile(redact=true) {
 
 		// actions - "Selected Text", Options, Probabilities
 		for (const key of Object.keys(varsData.story_settings.actions.actions)) {
-			for (const redactKey of ["Selected Text", "Options", "Probabilities"]) {
+			for (const redactKey of ["Selected Text", "Options", "Probabilities", "Original Text"]) {
 				varsData.story_settings.actions.actions[key][redactKey] = getRedactedValue(varsData.story_settings.actions.actions[key][redactKey]);
 			}
 		}
@@ -6937,4 +6939,231 @@ $el(".gametext").addEventListener("keydown", function(event) {
 	// contentEditable="plaintext-only" we're just gonna have to roll with it
 	document.execCommand("insertLineBreak");
 	event.preventDefault();
-});
+});
+
+/* Screenshot */
+const screenshotTarget = $el("#screenshot-target");
+const screenshotImagePicker = $el("#screenshot-image-picker");
+const screenshotImageContainer = $el("#screenshot-images");
+const robotAttribution = $el("#robot-attribution");
+const screenshotTextContainer = $el("#screenshot-text-container");
+
+sync_hooks.push({
+	class: "story",
+	name: "story_name",
+	func: function(title) {
+		$el("#story-attribution").innerText = title;
+	}
+});
+
+sync_hooks.push({
+	class: "model",
+	name: "model",
+	func: function(modelName) {
+		$el("#model-name").innerText = modelName
+	}
+})
+
+sync_hooks.push({
+	class: "user",
+	name: "screenshot_author_name",
+	func: function(name) {
+		$el("#human-attribution").innerText = name;
+	}
+});
+
+sync_hooks.push({
+	class: "user",
+	name: "screenshot_show_attribution",
+	func: function(show) {
+		robotAttribution.classList.toggle("hidden", !show);
+		$el("#screenshot-options-attribution").classList.toggle("disabled", !show);
+		if (show) robotAttribution.scrollIntoView();
+	}
+});
+
+sync_hooks.push({
+	class: "user",
+	name: "screenshot_show_story_title",
+	func: function(show) {
+		$el("#story-title-vis").classList.toggle("hidden", !show);
+		robotAttribution.scrollIntoView();
+	}
+});
+
+sync_hooks.push({
+	class: "user",
+	name: "screenshot_show_author_name",
+	func: function(show) {
+		$el("#author-name-vis").classList.toggle("hidden", !show);
+		$el("#screenshot-options-author-name").classList.toggle("disabled", !show);
+		robotAttribution.scrollIntoView();
+	}
+});
+
+sync_hooks.push({
+	class: "user",
+	name: "screenshot_show_model_name",
+	func: function(show) {
+		$el("#model-name-vis").classList.toggle("hidden", !show);
+		robotAttribution.scrollIntoView();
+	}
+});
+
+sync_hooks.push({
+	class: "user",
+	name: "screenshot_use_boring_colors",
+	func: function(boring) {
+		screenshotTarget.classList.toggle("boring-colors", boring);
+	}
+});
+
+async function showScreenshotWizard(actionComposition, startDebt, endDebt) {
+	// startDebt is the amount we need to shave off the front, and endDebt the
+	// same for the end
+	
+	screenshotTextContainer.innerHTML = "";
+	let charCount = startDebt;
+	let i = 0;
+	for (const action of actionComposition) {
+		for (const chunk of action) {
+			// Account for debt
+			if (startDebt > 0) {
+				if (chunk.content.length <= startDebt) {
+					startDebt -= chunk.content.length;
+					continue;
+				} else {
+					// Slice up chunk
+					chunk.content = chunk.content.slice(startDebt);
+					startDebt = 0;
+				}
+			}
+
+			if (charCount > endDebt) {
+				break;
+			} else if (charCount + chunk.content.length > endDebt) {
+				let charsLeft = endDebt - charCount
+				chunk.content = chunk.content.slice(0, charsLeft).trimEnd();
+				endDebt = -1;
+			}
+
+
+			if (i == 0) chunk.content = chunk.content.trimStart();
+			i++;
+
+			charCount += chunk.content.length;
+
+			let actionClass = {
+				ai: "ai-text",
+				user: "human-text",
+				edit: "edit-text",
+				prompt: "prompt-text",
+			}[chunk.type];
+
+			$e("span", screenshotTextContainer, {
+				innerText: chunk.content,
+				classes: ["action-text", actionClass]
+			});
+		}
+	}
+
+	let imageData = await (await fetch("/image_db.json")).json();
+	screenshotImagePicker.innerHTML = "";
+
+	for (const image of imageData) {
+		if (!image) continue;
+
+		const imgContainer = $e("div", screenshotImagePicker, {classes: ["img-container"]});
+		const checkbox = $e("input", imgContainer, {type: "checkbox"});
+		const imageEl = $e("img", imgContainer, {
+			src: `/generated_images/${image.fileName}`,
+			draggable: false,
+			tooltip: image.displayPrompt
+		});
+
+		imgContainer.addEventListener("click", function(event) {
+			// TODO: Preventdefault if too many images selected and checked is false
+			checkbox.click();
+		});
+
+		checkbox.addEventListener("click", function(event) {
+			event.stopPropagation();
+			screenshotWizardUpdateShownImages();
+		});
+	}
+	openPopup("screenshot-wizard");
+}
+
+function screenshotWizardUpdateShownImages() {
+	screenshotImageContainer.innerHTML = "";
+
+	for (const imgCont of screenshotImagePicker.children) {
+		const checked = imgCont.querySelector("input").checked;
+		if (!checked) continue;
+		const src = imgCont.querySelector("img").src;
+		$e("img", screenshotImageContainer, {src: src});
+	}
+}
+
+async function downloadScreenshot() {
+	// TODO: Upscale (eg transform with given ratio like 1.42 to make image
+	// bigger via screenshotTarget cloning)
+	const canvas = await html2canvas(screenshotTarget, {
+		width: screenshotTarget.clientWidth,
+		height: screenshotTarget.clientHeight - 1
+	});
+
+	canvas.style.display = "none";
+	document.body.appendChild(canvas);
+	$e("a", null, {download: "screenshot.png", href: canvas.toDataURL("image/png")}).click();
+	canvas.remove();
+}
+$el("#sw-download").addEventListener("click", downloadScreenshot);
+
+// Other side of screenshot-options hack
+for (const el of document.getElementsByClassName("screenshot-setting")) {
+	// yeah this really sucks but bootstrap toggle only works with this
+	el.setAttribute("onchange", "sync_to_server(this);")
+}
+
+async function screenshot_selection(summonEvent) {
+	// Adapted from https://stackoverflow.com/a/4220888
+	let selection = window.getSelection();
+	let range = selection.getRangeAt(0);
+	let commonAncestorContainer = range.commonAncestorContainer;
+
+	if (commonAncestorContainer.nodeName === "#text") commonAncestorContainer = commonAncestorContainer.parentNode;
+
+	let rangeParentChildren = commonAncestorContainer.childNodes;
+	// Array of STRING actions ids
+	let selectedActionIds = [];
+
+	for (let el of rangeParentChildren) {
+		if (!selection.containsNode(el, true)) continue;
+		// When selecting a portion of a singular action, el can be a text
+		// node rather than an action span
+		if (el.nodeName === "#text") el = el.parentNode.closest("[chunk]");
+		let actionId = el.getAttribute("chunk");
+
+		if (!actionId) continue;
+		if (selectedActionIds.includes(actionId)) continue;
+
+		selectedActionIds.push(actionId);
+	}
+
+	let actionComposition = await (await fetch(`/action_composition.json?actions=${selectedActionIds.join(",")}`)).json();
+
+	let totalText = "";
+
+	for (const action of actionComposition) {
+		for (const chunk of action) totalText += chunk.content;
+	}
+
+	let selectionContent = selection.toString();
+	let startDebt = totalText.indexOf(selectionContent);
+	// lastIndexOf??
+	// endDebt is distance from the end of selection.
+	let endDebt = totalText.indexOf(selectionContent) + selectionContent.length;
+
+	await showScreenshotWizard(actionComposition, startDebt=startDebt, endDebt=endDebt, totalText);
+}

templates/popups.html

@@ -320,6 +320,86 @@
 
 	<!-- Big Image -->
 	<img id="big-image"></img>
+
+	<!-- Screenshot Wizard -->
+	<div id="screenshot-wizard" class="popup-window">
+		<script src="/static/html2canvas.min.js"></script>
+		<div class="popup-header">
+			<h1>Screenshot</h1>
+		</div>
+		<span class="help_text">
+			This tool is used to create screenshots of story segments that can easily be shared. The options below will help you customize the screenshot to your liking.
+		</span>
+		<div id="screenshot-target-container">
+			<div id="screenshot-target" class="noselect">
+				<div id="screenshot-left">
+					<div id="screenshot-text-container"></div>
+					<span id="robot-attribution">
+						<span class="boring-part-of-madlibs">A snippit</span>
+
+						<span id="story-title-vis">
+							<span class="boring-part-of-madlibs">of</span>
+							<span id="story-attribution">The Great John Koblad</span>,
+						</span>
+
+						<span class="boring-part-of-madlibs">written</span>
+
+						<span id="author-name-vis">
+							<span class="boring-part-of-madlibs">by</span>
+							<span id="human-attribution">Anonymous</span>
+						</span>
+
+						<span class="boring-part-of-madlibs">in KoboldAI</span>
+
+						<span id="model-name-vis">
+							<span class="boring-part-of-madlibs">with</span>
+							<span id="model-name">facebook/OPT_175B</span>
+						</span>
+					</span>
+				</div>
+				<div id="screenshot-images"></div>
+			</div>
+		</div>
+
+		<div id="screenshot-options">
+			<!--
+				FIXME: screenshot-setting is a lazy hack until var_sync can converted to an attribute for scanning and such
+			-->
+			<span class="header">Options</span>
+			<div class="option">
+				Show Attribution
+				<input id="sw-attribution" class="screenshot-setting var_sync_user_screenshot_show_attribution" type="checkbox" data-size="mini" data-onstyle="success" data-toggle="toggle">
+			</div>
+			<div id="screenshot-options-attribution" class="indent">
+				<div class="option">
+					Show Story Title
+					<input class="screenshot-setting var_sync_user_screenshot_show_story_title" type="checkbox" data-size="mini" data-onstyle="success" data-toggle="toggle">
+				</div>
+				<div class="option">
+					Show Author Name
+					<input class="screenshot-setting var_sync_user_screenshot_show_author_name" type="checkbox" data-size="mini" data-onstyle="success" data-toggle="toggle">
+				</div>
+				<div class="indent">
+					<div id="screenshot-options-author-name" class="option">
+						Author Name
+						<input class="screenshot-setting var_sync_user_screenshot_author_name">
+					</div>
+				</div>
+				<div class="option">
+					Show Model Name
+					<input class="screenshot-setting var_sync_user_screenshot_show_model_name" type="checkbox" data-size="mini" data-onstyle="success" data-toggle="toggle">
+				</div>
+				<div class="option">
+					Boring Colors
+					<input class="screenshot-setting var_sync_user_screenshot_use_boring_colors" type="checkbox" data-size="mini" data-onstyle="success" data-toggle="toggle">
+				</div>
+			</div>
+		</div>
+		<div class="header">Generated Images</div>
+		<div id="screenshot-image-picker" class="noselect"></div>
+
+		<div id="sw-download">Download Screenshot</div>
+	</div>
 </div>
 
 <div id="notification-container"></div>