From 1da4ea589851ed9b171c9d06c49b505d20f2963f Mon Sep 17 00:00:00 2001
From: "Lt.Dr.Data" <lt.dr.data@gmail.com>
Date: Tue, 21 Mar 2023 12:44:31 +0900
Subject: [PATCH] Add support for sending output image to loaded image.
 Simplify workflow of img2img with new feature.

* ADDED: 'recv img' button to LoadImage, LoadImageMask
* ADDED: 'send to img' button to SaveImage
---
 nodes.py               |  9 +++++++--
 server.py              | 13 +++++++++++++
 web/scripts/api.js     |  5 +++++
 web/scripts/widgets.js | 44 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 1658a4c88..f5ec4e780 100644
--- a/nodes.py
+++ b/nodes.py
@@ -724,7 +724,9 @@ class SaveImage:
     def INPUT_TYPES(s):
         return {"required": 
                     {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
+                     "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+                     "send to img": ("IMAGESEND", )
+                     },
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                 }
 
@@ -806,7 +808,9 @@ class LoadImage:
         if not os.path.exists(s.input_dir):
             os.makedirs(s.input_dir)
         return {"required":
-                    {"image": (sorted(os.listdir(s.input_dir)), )},
+                    {"image": (sorted(os.listdir(s.input_dir)), ), 
+                    "recv img": (["disable", "enable"], ) 
+                    },
                 }
 
     CATEGORY = "image"
@@ -840,6 +844,7 @@ class LoadImageMask:
     def INPUT_TYPES(s):
         return {"required":
                     {"image": (sorted(os.listdir(s.input_dir)), ),
+                    "recv img": (["disable", "enable"], ),
                     "channel": (["alpha", "red", "green", "blue"], ),}
                 }
 
diff --git a/server.py b/server.py
index 6615a39e4..b1ec1fcf5 100644
--- a/server.py
+++ b/server.py
@@ -7,6 +7,8 @@ import execution
 import uuid
 import json
 import glob
+from shutil import copyfile
+
 try:
     import aiohttp
     from aiohttp import web
@@ -84,6 +86,17 @@ class PromptServer():
             files = glob.glob(os.path.join(self.web_root, 'extensions/**/*.js'), recursive=True)
             return web.json_response(list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files)))
 
+        @routes.get("/image/output_to_input/{name}")
+        async def copy_output_to_input_image(request):
+            name  = request.match_info["name"]
+            
+            src_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output", name)
+            dest_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input", name)
+
+            copyfile(src_dir,dest_dir)
+
+            return web.Response(status=200)
+            
         @routes.post("/upload/image")
         async def upload_image(request):
             upload_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
diff --git a/web/scripts/api.js b/web/scripts/api.js
index b90b1c656..1645272fb 100644
--- a/web/scripts/api.js
+++ b/web/scripts/api.js
@@ -124,6 +124,11 @@ class ComfyApi extends EventTarget {
 		return await resp.json();
 	}
 
+
+	async sendOutputToInputImage(name) {
+		await fetch(`/image/output_to_input/${name}`, { cache: "no-store" });
+	}
+	
 	/**
 	 *
 	 * @param {number} number The index at which to queue the prompt, passing -1 will insert the prompt at the front of the queue
diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js
index e1f637637..dca088d3d 100644
--- a/web/scripts/widgets.js
+++ b/web/scripts/widgets.js
@@ -1,3 +1,6 @@
+import { api } from "./api.js";
+import { app } from "../../scripts/app.js";
+
 function getNumberDefaults(inputData, defaultStep) {
 	let defaultVal = inputData[1]["default"];
 	let { min, max, step } = inputData[1];
@@ -126,6 +129,47 @@ export const ComfyWidgets = {
 			return { widget: node.addWidget("text", inputName, defaultVal, () => {}, {}) };
 		}
 	},
+	IMAGESEND(node, inputName) {
+		function showImage(node,uploadWidget,name) {
+			// Position the image somewhere sensible
+			if (!node.imageOffset) {
+				node.imageOffset = uploadWidget.last_y ? uploadWidget.last_y + 25 : 75;
+			}
+
+			const img = new Image();
+			img.onload = () => {
+				node.imgs = [img];
+				app.graph.setDirtyCanvas(true);
+			};
+			img.src = `/view?filename=${name}&type=input`;
+		}
+
+		async function callback() {
+			const imageWidget = node.widgets.find((w) => w.name === "image");
+
+			const copied = false;
+
+			for(let i in app.graph._nodes) {
+				var n = app.graph._nodes[i];
+				if(n.type == "LoadImage" || n.type == "LoadImageMask") {
+					const recvWidget = n.widgets.find((w) => w.name === "recv img");
+
+					if(recvWidget.value == "enable") {
+						// copy current node image to 'recv img' enabled node
+						
+						if(!copied) {
+							api.sendOutputToInputImage(imageWidget.value);
+						}
+
+						const thatImageWidget = n.widgets.find((w) => w.value === "image");
+						await showImage(n,thatImageWidget,imageWidget.value);
+					}
+				}
+			}
+		}
+
+		return { widget: node.addWidget("button", inputName, "", () => { callback(); }, {}) };
+	},
 	IMAGEUPLOAD(node, inputName, inputData, app) {
 		const imageWidget = node.widgets.find((w) => w.name === "image");
 		let uploadWidget;