diff --git a/nodes.py b/nodes.py index bb53cf05..002d022d 100644 --- a/nodes.py +++ b/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)), )}, } diff --git a/server.py b/server.py index 84b0941f..fbab4e0c 100644 --- a/server.py +++ b/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) diff --git a/web/extensions/core/uploadImage.js b/web/extensions/core/uploadImage.js new file mode 100644 index 00000000..45fabb78 --- /dev/null +++ b/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"]; + } + }, +}); diff --git a/web/scripts/app.js b/web/scripts/app.js index 1cf81b8f..e70e1c15 100644 --- a/web/scripts/app.js +++ b/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; diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js index f2e094ae..55bdd8f1 100644 --- a/web/scripts/widgets.js +++ b/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 }; + }, };