Merge branch 'uploadimages' of https://github.com/pythongosssss/ComfyUI

nodes.py

@@ -807,6 +807,8 @@ class LoadImage:
     input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
     @classmethod
     def INPUT_TYPES(s):
+        if not os.path.exists(s.input_dir):
+            os.makedirs(s.input_dir)
         return {"required":
                     {"image": (sorted(os.listdir(s.input_dir)), )},
                 }

server.py

@@ -6,7 +6,6 @@ import execution
 import uuid
 import json
 import glob
-
 try:
     import aiohttp
     from aiohttp import web
@@ -27,7 +26,7 @@ class PromptServer():
         self.loop = loop
         self.messages = asyncio.Queue()
         self.number = 0
-        self.app = web.Application()
+        self.app = web.Application(client_max_size=20971520)
         self.sockets = dict()
         self.web_root = os.path.join(os.path.dirname(
             os.path.realpath(__file__)), "web")
@@ -71,12 +70,47 @@ 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.post("/upload/image")
+        async def upload_image(request):
+            upload_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input")
+
+            if not os.path.exists(upload_dir):
+                os.makedirs(upload_dir)
+            
+            post = await request.post()
+            image = post.get("image")
+
+            if image and image.file:
+                filename = image.filename
+                if not filename:
+                    return web.Response(status=400)
+
+                split = os.path.splitext(filename)
+                i = 1
+                while os.path.exists(os.path.join(upload_dir, filename)):
+                    filename = f"{split[0]} ({i}){split[1]}"
+                    i += 1
+
+                filepath = os.path.join(upload_dir, filename)
+
+                with open(filepath, "wb") as f:
+                    f.write(image.file.read())
+                
+                return web.json_response({"name" : filename})
+            else:
+                return web.Response(status=400)
+
+
         @routes.get("/view/{file}")
         async def view_image(request):
             if "file" in request.match_info:
-                output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")
+                type = request.rel_url.query.get("type", "output")
+                if type != "output" and type != "input":
+                    return web.Response(status=400)
+
+                output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type)
                 file = request.match_info["file"]
-                file = os.path.splitext(os.path.basename(file))[0] + ".png"
+                file = os.path.basename(file)
                 file = os.path.join(output_dir, file)
                 if os.path.isfile(file):
                     return web.FileResponse(file)

web/extensions/core/uploadImage.js

@@ -0,0 +1,12 @@
+import { app } from "/scripts/app.js";
+
+// Adds an upload button to the nodes
+
+app.registerExtension({
+	name: "Comfy.UploadImage",
+	async beforeRegisterNodeDef(nodeType, nodeData, app) {
+		if (nodeData.name === "LoadImage" || nodeData.name === "LoadImageMask") {
+			nodeData.input.required.upload = ["IMAGEUPLOAD"];
+		}
+	},
+});

web/scripts/app.js

@@ -142,7 +142,7 @@ class ComfyApp {
 					if (numImages === 1 && !imageIndex) {
 						this.imageIndex = imageIndex = 0;
 					}
-					let shiftY = this.type === "SaveImage" ? 55 : 0;
+					let shiftY = this.type === "SaveImage" ? 55 : this.imageOffset || 0;
 					let dw = this.size[0];
 					let dh = this.size[1];
 					dh -= shiftY;

web/scripts/widgets.js

@@ -126,4 +126,84 @@ export const ComfyWidgets = {
 			return { widget: node.addWidget("text", inputName, defaultVal, () => {}, {}) };
 		}
 	},
+	IMAGEUPLOAD(node, inputName, inputData, app) {
+		const imageWidget = node.widgets.find((w) => w.name === "image");
+		let uploadWidget;
+
+		function showImage(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/${name}?type=input`;
+		}
+
+		// Add our own callback to the combo widget to render an image when it changes
+		const cb = node.callback;
+		imageWidget.callback = function () {
+			showImage(imageWidget.value);
+			if (cb) {
+				return cb.apply(this, arguments);
+			}
+		};
+
+		// On load if we have a value then render the image
+		// The value isnt set immediately so we need to wait a moment
+		// No change callbacks seem to be fired on initial setting of the value
+		requestAnimationFrame(() => {
+			if (imageWidget.value) {
+				showImage(imageWidget.value);
+			}
+		});
+
+		const fileInput = document.createElement("input");
+		Object.assign(fileInput, {
+			type: "file",
+			accept: "image/jpeg,image/png",
+			style: "display: none",
+			onchange: async () => {
+				if (fileInput.files.length) {
+					try {
+						// Wrap file in formdata so it includes filename
+						const body = new FormData();
+						body.append("image", fileInput.files[0]);
+						const resp = await fetch("/upload/image", {
+							method: "POST",
+							body,
+						});
+
+						if (resp.status === 200) {
+							const data = await resp.json();
+							showImage(data.name);
+
+							// Add the file as an option and update the widget value
+							if (!imageWidget.options.values.includes(data.name)) {
+								imageWidget.options.values.push(data.name);
+							}
+							imageWidget.value = data.name;
+						} else {
+							alert(resp.status + " - " + resp.statusText);
+						}
+					} catch (error) {
+						alert(error);
+					}
+				}
+			},
+		});
+		document.body.append(fileInput);
+
+		// Create the button widget for selecting the files
+		uploadWidget = node.addWidget("button", "choose file to upload", "image", () => {
+			fileInput.click();
+		});
+		uploadWidget.serialize = false;
+
+		return { widget: uploadWidget };
+	},
 };