Merge pull request #241 from one-some/screenshot

Screenshot tool
This commit is contained in:
henk717
2023-01-04 23:49:24 +01:00
committed by GitHub
6 changed files with 602 additions and 4 deletions

View File

@@ -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):

View File

@@ -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)

20
static/html2canvas.min.js vendored Normal file

File diff suppressed because one or more lines are too long

View File

@@ -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 {

View File

@@ -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);
}

View File

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