From 01e3ba05cfb27d27f087176f6bfc0d0e48a133f4 Mon Sep 17 00:00:00 2001
From: somebody <onesome01@protonmail.com>
Date: Sat, 26 Nov 2022 22:17:24 -0600
Subject: [PATCH] WI and Generated Image link up

---
 aiserver.py              |  6 +++++
 static/koboldai.js       | 48 ++++++++++++++++++++++++++++++++++++++--
 templates/templates.html |  8 +++++--
 3 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 4c1e11de..7222943e 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -9406,6 +9406,12 @@ def text2img_api(prompt,
 
     return base64_image
 
+@socketio.on("clear_generated_image")
+@logger.catch
+def UI2_clear_generated_image(data):
+    koboldai_vars.picture = ""
+    koboldai_vars.picture_prompt = ""
+
 #@logger.catch
 def get_items_locations_from_text(text):
     # load model and tokenizer
diff --git a/static/koboldai.js b/static/koboldai.js
index a95ff6dc..948b7c79 100644
--- a/static/koboldai.js
+++ b/static/koboldai.js
@@ -113,9 +113,13 @@ const context_menu_actions = {
 	],
 	"generated-image": [
 		{label: "View", icon: "search", enabledOn: "ALWAYS", click: imgGenView},
-		{label: "Download Image", icon: "download", enabledOn: "ALWAYS", click: imgGenView},
-		{label: "Use as World Info image", icon: "auto_stories", enabledOn: "ALWAYS", click: imgGenView},
+		{label: "Download", icon: "download", enabledOn: "ALWAYS", click: imgGenDownload},
+		{label: "Clear", icon: "clear", enabledOn: "ALWAYS", click: imgGenClear},
 	],
+	"wi-img-upload-button": [
+		{label: "Upload Image", icon: "file_upload", enabledOn: "ALWAYS", click: wiImageReplace},
+		{label: "Use Generated Image", icon: "image", enabledOn: "GENERATED-IMAGE", click: wiImageUseGeneratedImage},
+	]
 };
 
 // CTRL-[X]
@@ -5475,6 +5479,9 @@ process_cookies();
 		// The caret is placed
 		if (get_caret_position(target) !== null) $(".context-menu-item[enabled-on=CARET]").removeClass("disabled");
 
+		// The generated image is present
+		if ($el(".action_image")) $(".context-menu-item[enabled-on=GENERATED-IMAGE]").removeClass("disabled");
+
 		$(".context-menu-item[enabled-on=ALWAYS]").removeClass("disabled");
 
 		contextMenu.classList.remove("hidden");
@@ -5794,6 +5801,7 @@ function initalizeTooltips() {
 		let text = el.getAttribute("tooltip");
 
 		el.addEventListener("mouseenter", function(event) {
+			if (!el.hasAttribute("tooltip")) return;
 			tooltip.innerText = text;
 			let specialClass = "tooltip-standard";
 
@@ -6287,6 +6295,8 @@ function wiImageView(summonEvent) {
 }
 
 function wiImageReplace(summonEvent) {
+	// This is also used for the "Upload" context menu action on the placeholder.
+
 	// NOTE: WI image context menu stuff is pretty reliant on the current
 	// element structure, be sure to update this code if that's changed.
 	summonEvent.target.parentElement.click();
@@ -6302,9 +6312,43 @@ async function wiImageClear(summonEvent) {
 	});
 }
 
+async function wiImageUseGeneratedImage(summonEvent) {
+	// summonEvent is placeholder icon
+	const generatedImage = $el(".action_image");
+	if (!generatedImage) return;
+
+	let uid = parseInt(summonEvent.target.closest(".world_info_card").getAttribute("uid"));
+	summonEvent.target.classList.add("hidden");
+
+	let image = summonEvent.target.parentElement.getElementsByTagName("img")[0];
+	image.src = generatedImage.src;
+
+	let r = await fetch(`/set_wi_image/${uid}`, {
+		method: "POST",
+		body: generatedImage.src
+	});
+}
+
 function imgGenView() {
 	const image = $el(".action_image");
 	if (!image) return;
 	$el("#big-image").src = image.src;
 	openPopup("big-image");
+}
+
+function imgGenDownload() {
+	const image = $el(".action_image");
+	if (!image) return;
+	const a = $e("a", null, {href: image.src, download: "generated.png"});
+	a.click();
+}
+
+function imgGenClear() {
+	const image = $el(".action_image");
+	if (!image) return;
+	image.remove();
+
+	const container = $el("#action\\ image");
+	container.removeAttribute("tooltip");
+	socket.emit("clear_generated_image", {});
 }
\ No newline at end of file
diff --git a/templates/templates.html b/templates/templates.html
index bd137833..08d8c313 100644
--- a/templates/templates.html
+++ b/templates/templates.html
@@ -2,8 +2,12 @@
 <div draggable="true" class="world_info_card" id="world_info_">
 	<div class="world_info_title_area">
 		<div>
-			<div class="world_info_image_container" tooltip="Upload a picture for this World Info entry">
-				<span class="placeholder material-icons-outlined">add</span>
+			<div class="world_info_image_container">
+				<span
+					class="placeholder material-icons-outlined"
+					context-menu="wi-img-upload-button"
+					tooltip="Upload a picture for this World Info entry"
+				>add</span>
 				<img class="world_info_image">
 			</div>
 			<h4 draggable="true" ondragstart="event.preventDefault();event.stopPropagation();" id="world_info_title_"