add img2img tool and refactor tool context menus

Signed-off-by: Victor Seiji Hariki <victorseijih@gmail.com>

Former-commit-id: 3a36d0e0d7004142ee8d1ecdec14bb905db4919f

js/index.js

@@ -55,12 +55,31 @@ function sliderChangeHandlerFactory(
 	textBoxId,
 	dataKey,
 	defaultV,
+	save = true,
 	setter = (k, v) => (stableDiffusionData[k] = v),
 	getter = (k) => stableDiffusionData[k]
 ) {
-	const sliderEl = document.getElementById(sliderId);
+	return sliderChangeHandlerFactoryEl(
-	const textBoxEl = document.getElementById(textBoxId);
+		document.getElementById(sliderId),
-	const savedValue = localStorage.getItem(dataKey);
+		document.getElementById(textBoxId),
+		dataKey,
+		defaultV,
+		save,
+		setter,
+		getter
+	);
+}
+
+function sliderChangeHandlerFactoryEl(
+	sliderEl,
+	textBoxEl,
+	dataKey,
+	defaultV,
+	save = true,
+	setter = (k, v) => (stableDiffusionData[k] = v),
+	getter = (k) => stableDiffusionData[k]
+) {
+	const savedValue = save && localStorage.getItem(dataKey);
 
 	if (savedValue) setter(dataKey, savedValue || defaultV);
 
@@ -70,12 +89,12 @@ function sliderChangeHandlerFactory(
 
 		if (value) setter(dataKey, value);
 
-		if (!eventSource || eventSource.id === textBoxId)
+		if (!eventSource || eventSource === textBoxEl)
 			sliderEl.value = getter(dataKey);
 		setter(dataKey, Number(sliderEl.value));
 		textBoxEl.value = getter(dataKey);
 
-		localStorage.setItem(dataKey, getter(dataKey));
+		if (save) localStorage.setItem(dataKey, getter(dataKey));
 	}
 
 	textBoxEl.onchange = changeHandler;
@@ -198,14 +217,8 @@ function dream(
 	tmpImgXYWH.y = y;
 	tmpImgXYWH.w = prompt.width;
 	tmpImgXYWH.h = prompt.height;
-	console.log(
+	console.info(`dreaming "${prompt.prompt}"`);
-		"dreaming to " +
+	console.debug(prompt);
-			host +
-			url +
-			(extra.method || endpoint) +
-			":\r\n" +
-			JSON.stringify(prompt)
-	);
 	postData(prompt, extra).then((data) => {
 		returnedImages = data.images;
 		totalImagesReturned = data.images.length;
@@ -497,6 +510,7 @@ const changeScaleFactor = sliderChangeHandlerFactory(
 	"scaleFactorTxt",
 	"scaleFactor",
 	8,
+	true,
 	(k, v) => (scaleFactor = v),
 	(k) => scaleFactor
 );

js/input.js

@@ -362,7 +362,9 @@ window.onkeydown = (evn) => {
 		}, inputConfig.keyboardHoldTiming),
 	};
 
-	// Process shortcuts
+	// Process shortcuts if input target is not a text field
+	if (evn.target instanceof HTMLInputElement && evn.type === "text") return;
+
 	const callbacks = keyboard.shortcuts[evn.code];
 
 	if (callbacks)

js/shortcuts.js

@@ -14,3 +14,6 @@ keyboard.onShortcut({key: "KeyD"}, () => {
 keyboard.onShortcut({key: "KeyM"}, () => {
 	tools.maskbrush.enable();
 });
+keyboard.onShortcut({key: "KeyI"}, () => {
+	tools.img2img.enable();
+});

js/ui/tool/dream.js

@@ -85,3 +85,94 @@ const dream_erase_callback = (evn, state) => {
 	);
 	commands.runCommand("eraseImage", "Erase Area", bb);
 };
+
+/**
+ * Image to Image
+ */
+const dream_img2img_callback = (evn, state) => {
+	if (evn.target.id === "overlayCanvas" && !blockNewImages) {
+		const bb = getBoundingBox(
+			evn.x,
+			evn.y,
+			basePixelCount * scaleFactor,
+			basePixelCount * scaleFactor,
+			state.snapToGrid && basePixelCount
+		);
+
+		// Do nothing if no image exists
+		if (isCanvasBlank(bb.x, bb.y, bb.w, bb.h, imgCanvas)) return;
+
+		// Build request to the API
+		const request = {};
+		Object.assign(request, stableDiffusionData);
+
+		request.denoising_strength = state.denoisingStrength;
+		request.inpainting_fill = 1; // For img2img use original
+
+		// Load prompt (maybe we should add some events so we don't have to do this)
+		request.prompt = document.getElementById("prompt").value;
+		request.negative_prompt = document.getElementById("negPrompt").value;
+
+		// Don't allow another image until is finished
+		blockNewImages = true;
+
+		// Setup marching ants
+		stopMarching = march(bb);
+
+		// Setup some basic information for SD
+		request.width = bb.w;
+		request.height = bb.h;
+
+		request.firstphase_width = bb.w / 2;
+		request.firstphase_height = bb.h / 2;
+
+		// Use img2img
+
+		// Temporary canvas for init image and mask generation
+		const auxCanvas = document.createElement("canvas");
+		auxCanvas.width = request.width;
+		auxCanvas.height = request.height;
+		const auxCtx = auxCanvas.getContext("2d");
+
+		auxCtx.fillStyle = "#000F";
+
+		// Get init image
+		auxCtx.fillRect(0, 0, bb.w, bb.h);
+		auxCtx.drawImage(imgCanvas, bb.x, bb.y, bb.w, bb.h, 0, 0, bb.w, bb.h);
+		request.init_images = [auxCanvas.toDataURL()];
+
+		// Get mask image
+		auxCtx.fillRect(0, 0, bb.w, bb.h);
+		auxCtx.globalCompositeOperation = "destination-out";
+		auxCtx.drawImage(maskPaintCanvas, bb.x, bb.y, bb.w, bb.h, 0, 0, bb.w, bb.h);
+
+		// Border Mask
+		if (state.useBorderMask) {
+			auxCtx.fillStyle = "#000F";
+			auxCtx.fillRect(0, 0, state.borderMaskSize, bb.h);
+			auxCtx.fillRect(0, 0, bb.w, state.borderMaskSize);
+			auxCtx.fillRect(
+				bb.w - state.borderMaskSize,
+				0,
+				state.borderMaskSize,
+				bb.h
+			);
+			auxCtx.fillRect(
+				0,
+				bb.h - state.borderMaskSize,
+				bb.w,
+				state.borderMaskSize
+			);
+		}
+
+		auxCtx.globalCompositeOperation = "destination-atop";
+		auxCtx.fillStyle = "#FFFF";
+		auxCtx.fillRect(0, 0, bb.w, bb.h);
+		request.mask = auxCanvas.toDataURL();
+
+		request.inpainting_mask_invert = true;
+
+		// Dream
+		dream(bb.x, bb.y, request, {method: "img2img"});
+	}
+};

js/ui/toolbar.js

@@ -115,6 +115,67 @@ const toolbar = {
 	},
 };
 
+/**
+ * Premade inputs for populating the context menus
+ */
+const _toolbar_input = {
+	checkbox: (state, dataKey, text) => {
+		if (state[dataKey] === undefined) state[dataKey] = false;
+
+		const checkbox = document.createElement("input");
+		checkbox.type = "checkbox";
+		checkbox.checked = state[dataKey];
+		checkbox.onchange = () => (state[dataKey] = checkbox.checked);
+
+		const label = document.createElement("label");
+		label.appendChild(checkbox);
+		label.appendChild(new Text(text));
+
+		return {checkbox, label};
+	},
+
+	slider: (state, dataKey, text, min = 0, max = 1, step = 0.1) => {
+		const slider = document.createElement("input");
+		slider.type = "range";
+		slider.max = max;
+		slider.step = step;
+		slider.min = min;
+		slider.value = state[dataKey];
+
+		const textEl = document.createElement("input");
+		textEl.type = "number";
+		textEl.value = state[dataKey];
+
+		console.log(state[dataKey]);
+
+		sliderChangeHandlerFactoryEl(
+			slider,
+			textEl,
+			dataKey,
+			state[dataKey],
+			false,
+			(k, v) => (state[dataKey] = v),
+			(k) => state[dataKey]
+		);
+
+		const label = document.createElement("label");
+		label.appendChild(new Text(text));
+		label.appendChild(textEl);
+		label.appendChild(slider);
+
+		return {
+			slider,
+			text: textEl,
+			label,
+			setValue(v) {
+				slider.value = v;
+				textEl.value = slider.value;
+				return parseInt(slider.value);
+			},
+		};
+	},
+};
+
 /**
  * Dream and img2img tools
  */
@@ -145,7 +206,7 @@ tools.dream = toolbar.registerTool(
 	(state, opt) => {
 		// Draw new cursor immediately
 		ovCtx.clearRect(0, 0, ovCanvas.width, ovCanvas.height);
-		_reticle_draw({...mouse.canvas.pos, target: {id: "overlayCanvas"}});
+		state.mousemovecb({...mouse.canvas.pos, target: {id: "overlayCanvas"}});
 
 		// Start Listeners
 		mouse.listen.canvas.onmousemove.on(state.mousemovecb);
@@ -170,18 +231,11 @@ tools.dream = toolbar.registerTool(
 		populateContextMenu: (menu, state) => {
 			if (!state.ctxmenu) {
 				state.ctxmenu = {};
-				// Snap To Grid Checkbox
+				state.ctxmenu.snapToGridLabel = _toolbar_input.checkbox(
-				const snapToGridCheckbox = document.createElement("input");
+					state,
-				snapToGridCheckbox.type = "checkbox";
+					"snapToGrid",
-				snapToGridCheckbox.checked = state.snapToGrid;
+					"Snap To Grid"
-				snapToGridCheckbox.onchange = () =>
+				).label;
-					(state.snapToGrid = snapToGridCheckbox.checked);
-				state.ctxmenu.snapToGridCheckbox = snapToGridCheckbox;
-
-				const snapToGridLabel = document.createElement("label");
-				snapToGridLabel.appendChild(snapToGridCheckbox);
-				snapToGridLabel.appendChild(new Text("Snap to Grid"));
-				state.ctxmenu.snapToGridLabel = snapToGridLabel;
 			}
 
 			menu.appendChild(state.ctxmenu.snapToGridLabel);
@@ -190,6 +244,126 @@ tools.dream = toolbar.registerTool(
 	}
 );
 
+tools.img2img = toolbar.registerTool(
+	"res/icons/image.svg",
+	"Img2Img",
+	(state, opt) => {
+		// Draw new cursor immediately
+		ovCtx.clearRect(0, 0, ovCanvas.width, ovCanvas.height);
+		state.mousemovecb({...mouse.canvas.pos, target: {id: "overlayCanvas"}});
+
+		// Start Listeners
+		mouse.listen.canvas.onmousemove.on(state.mousemovecb);
+		mouse.listen.canvas.left.onclick.on(state.dreamcb);
+		mouse.listen.canvas.right.onclick.on(state.erasecb);
+	},
+	(state, opt) => {
+		// Clear Listeners
+		mouse.listen.canvas.onmousemove.clear(state.mousemovecb);
+		mouse.listen.canvas.left.onclick.clear(state.dreamcb);
+		mouse.listen.canvas.right.onclick.clear(state.erasecb);
+	},
+	{
+		init: (state) => {
+			state.snapToGrid = true;
+			state.denoisingStrength = 0.7;
+
+			state.useBorderMask = true;
+			state.borderMaskSize = 64;
+
+			state.mousemovecb = (evn) => {
+				_reticle_draw(evn, state.snapToGrid);
+				const bb = getBoundingBox(
+					evn.x,
+					evn.y,
+					basePixelCount * scaleFactor,
+					basePixelCount * scaleFactor,
+					snapToGrid && basePixelCount
+				);
+
+				// For displaying border mask
+				const auxCanvas = document.createElement("canvas");
+				auxCanvas.width = bb.w;
+				auxCanvas.height = bb.h;
+				const auxCtx = auxCanvas.getContext("2d");
+
+				if (state.useBorderMask) {
+					auxCtx.fillStyle = "#FF6A6A50";
+					auxCtx.fillRect(0, 0, state.borderMaskSize, bb.h);
+					auxCtx.fillRect(0, 0, bb.w, state.borderMaskSize);
+					auxCtx.fillRect(
+						bb.w - state.borderMaskSize,
+						0,
+						state.borderMaskSize,
+						bb.h
+					);
+					auxCtx.fillRect(
+						0,
+						bb.h - state.borderMaskSize,
+						bb.w,
+						state.borderMaskSize
+					);
+				}
+
+				const tmp = ovCtx.globalAlpha;
+				ovCtx.globalAlpha = 0.4;
+				ovCtx.drawImage(auxCanvas, bb.x, bb.y);
+				ovCtx.globalAlpha = tmp;
+			};
+			state.dreamcb = (evn) => {
+				dream_img2img_callback(evn, state);
+			};
+			state.erasecb = (evn) => dream_erase_callback(evn, state);
+		},
+		populateContextMenu: (menu, state) => {
+			if (!state.ctxmenu) {
+				state.ctxmenu = {};
+				// Snap To Grid Checkbox
+				state.ctxmenu.snapToGridLabel = _toolbar_input.checkbox(
+					state,
+					"snapToGrid",
+					"Snap To Grid"
+				).label;
+
+				// Denoising Strength Slider
+				state.ctxmenu.denoisingStrengthLabel = _toolbar_input.slider(
+					state,
+					"denoisingStrength",
+					"Denoising Strength",
+					0,
+					1,
+					0.05
+				).label;
+
+				// Use Border Mask Checkbox
+				state.ctxmenu.useBorderMaskLabel = _toolbar_input.checkbox(
+					state,
+					"useBorderMask",
+					"Use Border Mask"
+				).label;
+				// Border Mask Size Slider
+				state.ctxmenu.borderMaskSize = _toolbar_input.slider(
+					state,
+					"borderMaskSize",
+					"Border Mask Size",
+					0,
+					128,
+					1
+				).label;
+			}
+
+			menu.appendChild(state.ctxmenu.snapToGridLabel);
+			menu.appendChild(document.createElement("br"));
+			menu.appendChild(state.ctxmenu.denoisingStrengthLabel);
+			menu.appendChild(document.createElement("br"));
+			menu.appendChild(state.ctxmenu.useBorderMaskLabel);
+			menu.appendChild(document.createElement("br"));
+			menu.appendChild(state.ctxmenu.borderMaskSize);
+		},
+		shortcut: "I",
+	}
+);
+
 /**
  * Mask Editing tools
  */
@@ -246,15 +420,9 @@ tools.maskbrush = toolbar.registerTool(
 
 			state.wheelcb = (evn) => {
 				if (evn.target.id === "overlayCanvas") {
-					state.setBrushSize(
+					state.brushSize = state.setBrushSize(
-						Math.max(
+						state.brushSize -
-							state.config.minBrushSize,
+							Math.floor(state.config.brushScrollSpeed * evn.delta)
-							Math.min(
-								state.config.maxBrushSize,
-								state.brushSize -
-									Math.floor(state.config.brushScrollSpeed * evn.delta)
-							)
-						)
 					);
 					ovCtx.clearRect(0, 0, ovCanvas.width, ovCanvas.height);
 					state.movecb(evn);
@@ -267,28 +435,16 @@ tools.maskbrush = toolbar.registerTool(
 		populateContextMenu: (menu, state) => {
 			if (!state.ctxmenu) {
 				state.ctxmenu = {};
-				// Brush Size slider
+				const brushSizeSlider = _toolbar_input.slider(
-				const brushSizeRange = document.createElement("input");
+					state,
-				brushSizeRange.type = "range";
+					"brushSize",
-				brushSizeRange.value = state.brushSize;
+					"Brush Size",
-				brushSizeRange.max = state.config.maxBrushSize;
+					state.config.minBrushSize,
-				brushSizeRange.step = 8;
+					state.config.maxBrushSize,
-				brushSizeRange.min = state.config.minBrushSize;
+					1
-				brushSizeRange.oninput = () =>
+				);
-					(state.brushSize = parseInt(brushSizeRange.value));
+				state.ctxmenu.brushSizeLabel = brushSizeSlider.label;
-				state.ctxmenu.brushSizeRange = brushSizeRange;
+				state.setBrushSize = brushSizeSlider.setValue;
-				const brushSizeText = document.createElement("input");
-				brushSizeText.type = "number";
-				brushSizeText.value = state.brushSize;
-				brushSizeText.oninput = () =>
-					(state.brushSize = parseInt(brushSizeText.value));
-				state.ctxmenu.brushSizeText = brushSizeText;
-
-				const brushSizeLabel = document.createElement("label");
-				brushSizeLabel.appendChild(new Text("Brush Size"));
-				brushSizeLabel.appendChild(brushSizeText);
-				brushSizeLabel.appendChild(brushSizeRange);
-				state.ctxmenu.brushSizeLabel = brushSizeLabel;
 			}
 
 			menu.appendChild(state.ctxmenu.brushSizeLabel);