From 4a326a25482ddd2479d0fdf895f6ab4d060f0243 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Wed, 8 Mar 2023 22:07:44 +0000 Subject: [PATCH] Added UploadImage+Mask nodes --- .gitignore | 1 + nodes.py | 12 ++++- server.py | 32 ++++++++++-- web/extensions/core/uploadImage.js | 10 ++++ web/scripts/app.js | 2 +- web/scripts/widgets.js | 80 ++++++++++++++++++++++++++++++ 6 files changed, 132 insertions(+), 5 deletions(-) create mode 100644 web/extensions/core/uploadImage.js diff --git a/.gitignore b/.gitignore index 6da32900..5c050404 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ models/checkpoints models/vae models/embeddings models/loras +uploads/ \ No newline at end of file diff --git a/nodes.py b/nodes.py index bb53cf05..49d76fae 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)), )}, } @@ -830,7 +832,10 @@ class LoadImage: with open(image_path, 'rb') as f: m.update(f.read()) return m.digest().hex() - + +class UploadImage(LoadImage): + input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "uploads") + class LoadImageMask: input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") @classmethod @@ -865,6 +870,9 @@ class LoadImageMask: with open(image_path, 'rb') as f: m.update(f.read()) return m.digest().hex() + +class UploadImageMask(LoadImageMask): + input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "uploads") class ImageScale: upscale_methods = ["nearest-exact", "bilinear", "area"] @@ -917,7 +925,9 @@ NODE_CLASS_MAPPINGS = { "LatentUpscale": LatentUpscale, "SaveImage": SaveImage, "LoadImage": LoadImage, + "UploadImage": UploadImage, "LoadImageMask": LoadImageMask, + "UploadImageMask": UploadImageMask, "ImageScale": ImageScale, "ImageInvert": ImageInvert, "ConditioningCombine": ConditioningCombine, diff --git a/server.py b/server.py index 84b0941f..3dce72e5 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,10 +70,37 @@ 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__)), "uploads") + + 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) + + with open(os.path.join(upload_dir, filename), "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 != "uploads": + 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.join(output_dir, file) diff --git a/web/extensions/core/uploadImage.js b/web/extensions/core/uploadImage.js new file mode 100644 index 00000000..aca40b3e --- /dev/null +++ b/web/extensions/core/uploadImage.js @@ -0,0 +1,10 @@ +import { app } from "/scripts/app.js"; + +app.registerExtension({ + name: "Comfy.UploadImage", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "UploadImage" || nodeData.name === "UploadImageMask") { + 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..76933f9d 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=uploads`; + } + + // 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 }; + }, };