Merge pull request #310 from one-some/ui2-automatic1111-fixes

SD WebUI API fixes + Image gen loading indicator

aiserver.py

@@ -768,7 +768,10 @@ api_v1 = KoboldAPISpec(
     tags=tags,
 )
 
-def show_error_notification(title: str, text: str) -> None:
+def show_error_notification(title: str, text: str, do_log: bool = False) -> None:
+    if do_log:
+        logger.error(f"{title}: {text}")
+
     if has_request_context():
         socketio.emit("show_error_notification", {"title": title, "text": text}, broadcast=True, room="UI_2")
     else:
@@ -9369,42 +9372,45 @@ def text2img_api(prompt,
       #"override_settings": {},
       #"sampler_index": "Euler"
     final_imgen_params = {
-        "prompt": "{}, {}".format(prompt, art_guide),
+        "prompt": ", ".join(filter(bool, [prompt, art_guide])),
         "n_iter": 1,
         "width": 512,
         "height": 512,
         "steps": koboldai_vars.img_gen_steps,
         "cfg_scale": koboldai_vars.img_gen_cfg_scale,
-        "negative_prompt": "{}".format(koboldai_vars.img_gen_negative_prompt),
+        "negative_prompt": koboldai_vars.img_gen_negative_prompt,
         "sampler_index": "Euler a"
     }
-    apiaddress = '{}/sdapi/v1/txt2img'.format(koboldai_vars.img_gen_api_url)
+    apiaddress = '{}/sdapi/v1/txt2img'.format(koboldai_vars.img_gen_api_url.rstrip("/"))
     payload_json = json.dumps(final_imgen_params)
     logger.debug(final_imgen_params)
-    #print("payload_json contains " + payload_json)
-    submit_req = requests.post(url=f'{apiaddress}', data=payload_json).json()
-    if submit_req:
-        results = submit_req
-        for i in results['images']:
-            final_src_img = Image.open(BytesIO(base64.b64decode(i.split(",",1)[0])))
-            buffer = BytesIO()
-            final_src_img.save(buffer, format="Webp", quality=95)
-            b64img = base64.b64encode(buffer.getvalue()).decode("utf8")
-            base64_bytes = b64img.encode('utf-8')
-            img_bytes = base64.b64decode(base64_bytes)
-            img = Image.open(BytesIO(img_bytes))
-            dt_string = datetime.datetime.now().strftime("%H%M%S%d%m%Y")
-            final_filename = "stories/art/{}_{}".format(dt_string,filename)
-            pnginfo = PngImagePlugin.PngInfo()
-            prompttext = results.get('info').split("\",")[0].split("\"")[3]
-            pnginfo.add_text("parameters","prompttext")
-            img.save(final_filename, pnginfo=pnginfo)
-            logger.debug("Saved Image")
+
+    try:
+        submit_req = requests.post(url=apiaddress, data=payload_json)
+    finally:
         koboldai_vars.generating_image = False
-            return(b64img)
-    else:
-        koboldai_vars.generating_image = False
-        logger.error(submit_req.text)
+
+    if submit_req.status_code == 404:
+        show_error_notification(
+            "SD Web API Failure",
+            f"The SD Web UI was not called with --api. Unable to connect.",
+            do_log=True
+        )
+        return None
+    elif not submit_req.ok:
+        show_error_notification("SD Web API Failure", f"HTTP Code {submit_req.status_code} -- See console for details")
+        logger.error(f"SD Web API Failure: HTTP Code {submit_req.status_code}, Body:\n{submit_req.text}")
+        return None
+
+    results = submit_req.json()
+
+    try:
+        base64_image = results["images"][0]
+    except (IndexError, KeyError):
+        show_error_notification("SD Web API Failure", "SD Web API returned no images", do_log=True)
+        return None
+
+    return base64_image
 
 #@logger.catch
 def get_items_locations_from_text(text):

gensettings.py

@@ -670,23 +670,6 @@ gensettingstf = [
     "ui_level": 2
  	},
     {
-    "UI_V2_Only": True,
- 	"uitype": "toggle",
- 	"unit": "bool",
- 	"label": "Reset Image Generating",
- 	"id": "generating_image",
- 	"min": 0,
- 	"max": 1,
- 	"step": 1,
- 	"default": 0,
-	"tooltip": "Use to reset image gen flag in case of error.",
-    "menu_path": "Interface",
-    "sub_path": "Images",
-    "classname": "system",
-    "name": "generating_image",
-    "ui_level": 2
- 	},
-    {
     "UI_V2_Only": True,
  	"uitype": "toggle",
  	"unit": "bool",

koboldai_settings.py

@@ -1060,7 +1060,8 @@ class system_settings(settings):
     no_save_variables = ['socketio', 'lua_state', 'lua_logname', 'lua_koboldbridge', 'lua_kobold', 
                          'lua_koboldcore', 'sp', 'sp_length', '_horde_pid', 'horde_share', 'aibusy', 
                          'serverstarted', 'inference_config', 'image_pipeline', 'summarizer', 
-                         'summary_tokenizer', 'use_colab_tpu', 'noai', 'disable_set_aibusy', 'cloudflare_link', 'tts_model']
+                         'summary_tokenizer', 'use_colab_tpu', 'noai', 'disable_set_aibusy', 'cloudflare_link', 'tts_model',
+                         'generating_image']
     settings_name = "system"
     def __init__(self, socketio, koboldai_var):
         self.socketio = socketio

static/koboldai.css

@@ -2988,8 +2988,39 @@ input[type='range'] {
 	pointer-events:none;
 }
 
+#action\ image {
+	width: 100%;
+	aspect-ratio: 1/1;
+	position: relative;
+}
+
+#image-loading {
+	display: flex;
+	align-items: center;
+	justify-content: center;
+	position: absolute;
+	top: 0px;
+	left: 0px;
+	width: 100%;
+	height: 100% !important;
+	aspect-ratio: 1/1;
+	background-color: rgba(0, 0, 0, 0.6);
+
+	/* Defer pointer events to the image in case a user wants to quickly
+	save / view an image after asking for a new image */
+	pointer-events: none;
+}
+
+.spinner {
+	animation: spin 1s linear infinite;
+	opacity: 0.7;
+}
+
+@keyframes spin { 100% { transform: rotate(360deg) } }
+
+
 .action_image {
-	width: var(--flyout_menu_width);
+	width: 100%;
 }
 
 /*Tooltip based on attribute

static/koboldai.js

@@ -790,9 +790,12 @@ function var_changed(data) {
 	//Special Case for story picture
 	} else if (data.classname == "story" && data.name == "picture") {
 		image_area = document.getElementById("action image");
-		while (image_area.firstChild) { 
-			image_area.removeChild(image_area.firstChild);
-		}
+
+		let maybeImage = image_area.getElementsByClassName("action_image")[0];
+		if (maybeImage) maybeImage.remove();
+
+		$el("#image-loading").classList.add("hidden");
+
 		if (data.value != "") {
 			var image = new Image();
 			image.src = 'data:image/png;base64,'+data.value;
@@ -800,9 +803,7 @@ function var_changed(data) {
 			image_area.appendChild(image);
 		}
 	}  else if (data.classname == "story" && data.name == "picture_prompt") {
-		if (document.getElementById("action image").firstChild) {
-			document.getElementById("action image").firstChild.setAttribute("title", data.value);
-		}
+		if (data.value) document.getElementById("action image").setAttribute("tooltip", data.value);
 	//special case for welcome text since we want to allow HTML
 	} else if (data.classname == 'model' && data.name == 'welcome') {
 		document.getElementById('welcome_text').innerHTML = data.value;
@@ -5991,6 +5992,11 @@ $el("#aidgpromptnum").addEventListener("keydown", function(event) {
 	event.preventDefault();
 });
 
+$el("#generate-image-button").addEventListener("click", function() {
+	$el("#image-loading").classList.remove("hidden");
+	socket.emit("generate_image", {});
+});
+
 /* -- Shiny New Chat -- */
 function addMessage(author, content, actionId, afterMsgEl=null, time=null) {
 	if (!time) time = Number(new Date());

templates/settings flyout.html

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