Merge pull request #321 from one-some/ui2-img-prompter

Img gen manual prompt

aiserver.py

@@ -9162,7 +9162,8 @@ def UI_2_save_revision(data):
 #==================================================================#
 @socketio.on("generate_image")
 @logger.catch
-def UI_2_generate_image(data):
+def UI_2_generate_image_from_story(data):
+    # Independant of generate_story_image as summarization is rather time consuming
     koboldai_vars.generating_image = True
     eventlet.sleep(0)
     
@@ -9213,36 +9214,51 @@ def UI_2_generate_image(data):
         keys = [summarize(text, max_length=max_length)]
         logger.debug("Text from summarizer: {}".format(keys[0]))
     
-    
+    generate_story_image(", ".join(keys), art_guide=art_guide)
 
-   
+@socketio.on("generate_image_from_prompt")
+@logger.catch
+def UI_2_generate_image_from_prompt(prompt: str):
+    eventlet.sleep(0)
+    generate_story_image(prompt)
+
+def generate_story_image(prompt: str, art_guide: str = "") -> None:
+    # This function is a wrapper around generate_image() that integrates the
+    # result with the story (read: puts it in the corner of the screen).
 
-    #If we don't have a GPU, use horde if we're allowed to
     start_time = time.time()
-    # Check if stable-diffusion-webui API option selected and use that if found.
-    if koboldai_vars.img_gen_priority == 4:
-        b64_data = text2img_api(", ".join(keys), art_guide = art_guide)
-    elif ((not koboldai_vars.hascuda or not os.path.exists("models/stable-diffusion-v1-4")) and koboldai_vars.img_gen_priority != 0) or  koboldai_vars.img_gen_priority == 3:
-        b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
-    else:
-        if ((not koboldai_vars.hascuda or not os.path.exists("models/stable-diffusion-v1-4")) and koboldai_vars.img_gen_priority != 0) or  koboldai_vars.img_gen_priority == 3:
-            b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
-        else:
-            import psutil
-            #We aren't being forced to use horde, so now let's figure out if we should use local
-            if torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_reserved(0) >= 6000000000:
-                #He have enough vram, just do it locally
-                b64_data = text2img_local(", ".join(keys), art_guide = art_guide)
-            elif torch.cuda.get_device_properties(0).total_memory > 6000000000 and koboldai_vars.img_gen_priority <= 1:
-                #We could do it locally by swapping the model out
-                print("Could do local or online")
-                b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
-            elif koboldai_vars.img_gen_priority != 0:
-                b64_data = text2img_horde(", ".join(keys), art_guide = art_guide)
+    koboldai_vars.generating_image = True
+
+    b64_data = generate_image(prompt, art_guide=art_guide)
+
     logger.debug("Time to Generate Image {}".format(time.time()-start_time))
+
     koboldai_vars.picture = b64_data
-    koboldai_vars.picture_prompt = ", ".join(keys)
+    koboldai_vars.picture_prompt = prompt
     koboldai_vars.generating_image = False
+
+def generate_image(prompt: str, art_guide: str = "") -> Optional[str]:
+    if koboldai_vars.img_gen_priority == 4:
+        # Check if stable-diffusion-webui API option selected and use that if found.
+        return text2img_api(prompt, art_guide=art_guide)
+    elif ((not koboldai_vars.hascuda or not os.path.exists("models/stable-diffusion-v1-4")) and koboldai_vars.img_gen_priority != 0) or  koboldai_vars.img_gen_priority == 3:
+        # If we don't have a GPU, use horde if we're allowed to
+        return text2img_horde(prompt, art_guide=art_guide)
+
+    memory = torch.cuda.get_device_properties(0).total_memory
+
+    # We aren't being forced to use horde, so now let's figure out if we should use local
+    if memory - torch.cuda.memory_reserved(0) >= 6000000000:
+        # We have enough vram, just do it locally
+        return text2img_local(prompt, art_guide=art_guide)
+    elif memory > 6000000000 and koboldai_vars.img_gen_priority <= 1:
+        # We could do it locally by swapping the model out
+        print("Could do local or online")
+        return text2img_horde(prompt, art_guide=art_guide)
+    elif koboldai_vars.img_gen_priority != 0:
+        return text2img_horde(prompt, art_guide=art_guide)
+
+    raise RuntimeError("Unable to decide image generation backend. Please report this.")
     
 
 @logger.catch
@@ -9335,9 +9351,7 @@ def text2img_horde(prompt,
         logger.error(submit_req.text)
 
 @logger.catch
-def text2img_api(prompt,
-             art_guide = "",
-             filename = "story_art.png"):
+def text2img_api(prompt, art_guide=""):
     logger.debug("Generating Image using Local SD-WebUI API")
     koboldai_vars.generating_image = True
     #The following list are valid properties with their defaults, to add/modify in final_imgen_params. Will refactor configuring values into UI element in future.

static/koboldai.js

@@ -79,6 +79,11 @@ var finder_actions = [
 	{name: "Download Story", icon: "file_download", type: "action", func: function() { document.getElementById('download_iframe').src = 'json'; }},
 	{name: "Import Story", icon: "file_download", desc: "Import a prompt from aetherroom.club, formerly prompts.aidg.club", type: "action", func: openClubImport },
 
+	// Imggen
+	{name: "Download Generated Image", icon: "file_download", type: "action", func: imgGenDownload},
+	{name: "View Generated Image", icon: "image", type: "action", func: imgGenView},
+	{name: "Clear Generated Image", icon: "image_not_supported", type: "action", func: imgGenClear},
+
 	// Locations
 	{name: "Setting Presets", icon: "open_in_new", type: "location", func: function() { highlightEl(".var_sync_model_selected_preset") }},
 	{name: "Memory", icon: "open_in_new", type: "location", func: function() { highlightEl("#memory") }},
@@ -5087,8 +5092,14 @@ function sendScratchpadPrompt(prompt) {
 	socket.emit("scratchpad_prompt", prompt);
 }
 
+function finderSendImgPrompt(prompt) {
+	closePopups();
+	$el("#image-loading").classList.remove("hidden");
+	socket.emit("generate_image_from_prompt", prompt);
+}
+
 function updateSearchListings() {
-	if (finder_mode === "scratchpad") return;
+	if (["scratchpad", "imgPrompt"].includes(finder_mode)) return;
 	if (this.value === finder_last_input) return;
 	finder_last_input = this.value;
 	finder_selection_index = -1;
@@ -5125,8 +5136,13 @@ function updateFinderMode(mode) {
 	const finderInput = document.querySelector("#finder-input");
 	const finderScratchpad = document.querySelector("#finder-scratchpad");
 
-	finderIcon.innerText = {ui: "search", wi: "auto_stories", scratchpad: "speaker_notes"}[mode];
-	finderInput.placeholder = {ui: "Search for something...", wi: "Search for a World Info entry...", scratchpad: "Prompt the AI..."}[mode];
+	finderIcon.innerText = {ui: "search", wi: "auto_stories", scratchpad: "speaker_notes", "imgPrompt": "image"}[mode];
+	finderInput.placeholder = {
+		ui: "Search for something...",
+		wi: "Search for a World Info entry...",
+		scratchpad: "Prompt the AI...",
+		imgPrompt: "Generate an image..."
+	}[mode];
 	finderScratchpad.classList.add("hidden");
 
 	finder_mode = mode;
@@ -5134,7 +5150,7 @@ function updateFinderMode(mode) {
 
 function cycleFinderMode() {
 	// Initiated by clicking on icon
-	updateFinderMode({ui: "wi", wi: "scratchpad", scratchpad: "ui"}[finder_mode]);
+	updateFinderMode({ui: "wi", wi: "scratchpad", scratchpad: "imgPrompt", imgPrompt: "ui"}[finder_mode]);
 }
 
 function open_finder() {
@@ -5401,7 +5417,7 @@ process_cookies();
 		let delta = 0;
 		const actions = document.getElementsByClassName("finder-result");
 
-		let newMode = {">": "wi", "#": "ui", "!": "scratchpad"}[event.key];
+		let newMode = {">": "wi", "#": "ui", "!": "scratchpad", "?": "imgPrompt"}[event.key];
 		if (newMode && !finderInput.value) {
 			event.preventDefault();
 			updateFinderMode(newMode);
@@ -5412,10 +5428,15 @@ process_cookies();
 			if (finder_mode === "scratchpad") {
 				sendScratchpadPrompt(finderInput.value);
 				return;
+			} else if (finder_mode === "imgPrompt") {
+				finderSendImgPrompt(finderInput.value);
+				return;
 			} else if (finder_mode === "ui") {
 				let index = finder_selection_index >= 0 ? finder_selection_index : 0;
 				if (!actions[index]) return;
 				actions[index].click();
+			} else {
+				return;
 			}
 		} else if (event.key === "ArrowUp") {
 			delta = -1;

templates/settings flyout.html

@@ -112,7 +112,7 @@
 		</div>
 		<span id="debug-dump" class="cursor" onclick="openPopup('debug-file-prompt');">Download debug dump</span>
 		<div id="Images">
-			<button id="generate-image-button" class="settings_button var_sync_alt_system_generating_image" onclick="socket.emit('generate_image', {})">Generate Image</button>
+			<button id="generate-image-button" class="settings_button var_sync_alt_system_generating_image">Generate Image</button>
 			<div id="action image">
 				<div id="image-loading" class="hidden">
 					<div class="spinner material-icons-outlined">autorenew</div>