From 961be2f2b80b3b4c9c95eb823e8ed32908aabb82 Mon Sep 17 00:00:00 2001 From: bruefire Date: Fri, 8 Dec 2023 05:17:51 +0900 Subject: [PATCH] implemented sketch logics. --- AnimateDiffInpaint.py | 90 ++++++++++++++++++++++++++++------ js/MaskEditorForVframes.js | 99 ++++++++++++++++++-------------------- js/main.js | 42 ++++++++++------ 3 files changed, 147 insertions(+), 84 deletions(-) diff --git a/AnimateDiffInpaint.py b/AnimateDiffInpaint.py index 61f1758..ebec8ce 100644 --- a/AnimateDiffInpaint.py +++ b/AnimateDiffInpaint.py @@ -29,29 +29,33 @@ class LoadImagesFromDirInpaint: comfy_root_dir = os.path.join(comfy_custom_dir, "..", "..") pathParts = images_id.split(":") - vframeId = pathParts[0] - filenames = pathParts[1:] - imageDir = os.path.join(comfy_root_dir, "input", "extVideoFrame" + vframeId) - maskDir = os.path.join(comfy_root_dir, "input", "extVideoFrame-mask" + vframeId) + vframeIds = pathParts[len(pathParts) - 1].split("/") + filenames = pathParts[:-1] + + images = [] + masks = [] + if len(vframeIds) == 1: + imageDir = os.path.join(comfy_root_dir, "input", "extVideoFrame" + vframeIds[0]) + + images = self.load_from_dir(imageDir, image_load_cap, start_index) + masks = [image.clone() for image in images] + for mask in masks: + mask.fill_(0) + else: + imageDir = os.path.join(comfy_root_dir, "input", "extVideoFrame" + vframeIds[0]) + maskDir = os.path.join(comfy_root_dir, "input", "extVideoFrame" + vframeIds[1]) + sketchDir = os.path.join(comfy_root_dir, "input", "extVideoFrame" + vframeIds[2]) + + images = self.load_from_dir_sketch(imageDir, sketchDir, image_load_cap, start_index) + masks = self.load_from_dir(maskDir, image_load_cap, start_index) - images = self.load_from_dir(imageDir, image_load_cap, start_index) - masks = self.load_from_dir(maskDir, image_load_cap, start_index) - print(len(masks)) return (torch.cat(images, dim=0), torch.cat(masks, dim=0), len(images)) def load_from_dir(self, imageDir, image_load_cap, start_index): - # confirm paths - if not os.path.isdir(imageDir): - raise FileNotFoundError(f"Directory '{imageDir} cannot be found.'") - img_files = os.listdir(imageDir) - if len(img_files) == 0: - raise FileNotFoundError(f"No files in directory '{imageDir}'.") - # start at start_index - img_files = sorted(img_files) - img_files = [os.path.join(imageDir, x) for x in img_files] + img_files = self.getImagePaths(imageDir) img_files = img_files[start_index:] images = [] @@ -79,3 +83,57 @@ class LoadImagesFromDirInpaint: raise FileNotFoundError(f"No images could be loaded from directory '{imageDir}'.") return images + + + def load_from_dir_sketch(self, imageDir, sketchDir, image_load_cap, start_index): + + img_files = self.getImagePaths(imageDir) + img_files = img_files[start_index:] + skt_files = self.getImagePaths(sketchDir) + skt_files = skt_files[start_index:] + + images = [] + + limit_images = False + if image_load_cap > 0: + limit_images = True + image_count = 0 + + for idx, image_path in enumerate(img_files): + if os.path.isdir(image_path): + continue + if limit_images and image_count >= image_load_cap: + break + i = Image.open(image_path) + i = ImageOps.exif_transpose(i).convert("RGBA") + s = Image.open(skt_files[idx]) + s = ImageOps.exif_transpose(s).convert("RGBA") + s = s.resize(i.size) + image = Image.alpha_composite(i, s) + image.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + + images.append(image) + image_count += 1 + + if len(images) == 0: + raise FileNotFoundError(f"No images could be loaded from directory '{imageDir}'.") + + return images + + + def getImagePaths(self, imageDir): + # confirm paths + if not os.path.isdir(imageDir): + raise FileNotFoundError(f"Directory '{imageDir} cannot be found.'") + img_files = os.listdir(imageDir) + if len(img_files) == 0: + raise FileNotFoundError(f"No files in directory '{imageDir}'.") + + # start at start_index + img_files = sorted(img_files) + img_files = [os.path.join(imageDir, x) for x in img_files] + + return img_files + diff --git a/js/MaskEditorForVframes.js b/js/MaskEditorForVframes.js index 328aed3..6977e96 100644 --- a/js/MaskEditorForVframes.js +++ b/js/MaskEditorForVframes.js @@ -17,39 +17,6 @@ function dataURLToBlob(dataURL) { return new Blob([arrayBuffer], { type: contentType }); } -function loadedImageToBlob(image) { - const canvas = document.createElement('canvas'); - - canvas.width = image.width; - canvas.height = image.height; - - const ctx = canvas.getContext('2d'); - - ctx.drawImage(image, 0, 0); - - const dataURL = canvas.toDataURL('image/png', 1); - const blob = dataURLToBlob(dataURL); - - return blob; -} - -async function uploadMask(filepath, formData) { - await api.fetchApi('/upload/mask', { - method: 'POST', - body: formData - }).then(response => {}).catch(error => { - console.error('Error:', error); - }); - - ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']] = new Image(); - ComfyApp.clipspace.imgs[ComfyApp.clipspace['selectedIndex']].src = api.apiURL("/view?" + new URLSearchParams(filepath).toString() + app.getPreviewFormatParam()); - - if(ComfyApp.clipspace.images) - ComfyApp.clipspace.images[ComfyApp.clipspace['selectedIndex']] = filepath; - - ClipspaceDialog.invalidatePreview(); -} - function prepareRGB(image, backupCanvas, backupCtx) { // paste mask data into alpha channel backupCtx.drawImage(image, 0, 0, backupCanvas.width, backupCanvas.height); @@ -87,7 +54,8 @@ class MaskEditorDialog extends ComfyDialog { set path_data(value) { const parts = value?.split(":"); if (parts) { - const id = parts.shift(); + const dirIds = parts.pop(); + const id = dirIds.split("/")[0]; this.#vframeId = id; this.#paths = parts.map(file_name => { return { @@ -220,7 +188,7 @@ class MaskEditorDialog extends ComfyDialog { document.body.appendChild(brush); - var colorPicker = this.createLeftColorPicker("sketch", + var colorPicker = this.colorPicker = this.createLeftColorPicker("sketch", () => { }); var modeButton = this.createLeftButton("sketch", @@ -228,6 +196,7 @@ class MaskEditorDialog extends ComfyDialog { if (!this.is_sketch) { ev.target.innerText = "inpaint"; colorPicker.style.display = "inline"; + maskCanvas.style.opacity = "1.0"; this.storeActiveToBack(); const bSketchCanvas = this.backSketchCanvases[this.#selectedIndex]; @@ -236,6 +205,7 @@ class MaskEditorDialog extends ComfyDialog { } else { ev.target.innerText = "sketch"; colorPicker.style.display = "none"; + maskCanvas.style.opacity = "0.66"; this.storeActiveToBack(); const bMaskCanvas = this.backMaskCanvases[this.#selectedIndex]; @@ -416,8 +386,15 @@ class MaskEditorDialog extends ComfyDialog { maskCanvas.height = drawHeight; maskCanvas.style.top = imgCanvas.offsetTop + "px"; maskCanvas.style.left = imgCanvas.offsetLeft + "px"; - // backupCtx.drawImage(maskCanvas, 0, 0, maskCanvas.width, maskCanvas.height, 0, 0, backupCanvas.width, backupCanvas.height); - // maskCtx.drawImage(backupCanvas, 0, 0, backupCanvas.width, backupCanvas.height, 0, 0, maskCanvas.width, maskCanvas.height); + + const resizeBackCanvas = backCanvas => { + if (backCanvas.width != maskCanvas.width || backCanvas.height != maskCanvas.height) { + backCanvas.width = maskCanvas.width; + backCanvas.height = maskCanvas.height; + } + }; + this.backMaskCanvases.forEach(resizeBackCanvas); + this.backSketchCanvases.forEach(resizeBackCanvas); maskCtx.drawImage(this.getBackCanvasForCurrentMode(this.#selectedIndex), 0, 0, maskCanvas.width, maskCanvas.height); }); @@ -551,7 +528,7 @@ class MaskEditorDialog extends ComfyDialog { if(diff > 20 && !this.drawing_mode) requestAnimationFrame(() => { self.maskCtx.beginPath(); - self.maskCtx.fillStyle = "rgb(255, 255, 255)"; + self.maskCtx.fillStyle = this.getActiveBrushColor(); self.maskCtx.globalCompositeOperation = "source-over"; self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false); self.maskCtx.fill(); @@ -561,7 +538,7 @@ class MaskEditorDialog extends ComfyDialog { else requestAnimationFrame(() => { self.maskCtx.beginPath(); - self.maskCtx.fillStyle = "rgb(255, 255, 255)"; + self.maskCtx.fillStyle = this.getActiveBrushColor(); self.maskCtx.globalCompositeOperation = "source-over"; var dx = x - self.lastx; @@ -652,7 +629,7 @@ class MaskEditorDialog extends ComfyDialog { self.maskCtx.beginPath(); if (event.button == 0) { - self.maskCtx.fillStyle = "rgb(255, 255, 255)"; + self.maskCtx.fillStyle = this.getActiveBrushColor(); self.maskCtx.globalCompositeOperation = "source-over"; } else { self.maskCtx.globalCompositeOperation = "destination-out"; @@ -666,20 +643,22 @@ class MaskEditorDialog extends ComfyDialog { } getBackCanvasForCurrentMode(index) { - if (this.is_sketch) { - return this.backSketchCanvases[index]; - } else { - return this.backMaskCanvases[index]; - } + return this.is_sketch + ? this.backSketchCanvases[index] + : this.backMaskCanvases[index]; } storeActiveToBack() { const backCanvas = this.getBackCanvasForCurrentMode(this.#selectedIndex); - backCanvas.width = this.maskCanvas.width; - backCanvas.height = this.maskCanvas.height; backCanvas.getContext('2d').drawImage(this.maskCanvas, 0, 0, backCanvas.width, backCanvas.height); } + getActiveBrushColor() { + return this.is_sketch + ? this.colorPicker.value + : "rgb(255, 255, 255)"; + } + async save() { const backupCtx = this.backupCanvas.getContext('2d'); @@ -688,8 +667,7 @@ class MaskEditorDialog extends ComfyDialog { 0, 0, this.maskCanvas.width, this.maskCanvas.height, 0, 0, this.backupCanvas.width, this.backupCanvas.height); - - this.backMaskCanvases.forEach(async (maskCanvas, i) => { + const uploadImages = async (maskCanvas, i, idSitr) => { const body = new FormData(); const dataURL = maskCanvas.toDataURL(); const blob = dataURLToBlob(dataURL); @@ -697,19 +675,36 @@ class MaskEditorDialog extends ComfyDialog { const filename = this.#paths[i].filename.substr(0, extPos) + ".png"; body.append("image", blob, filename); - body.append("subfolder", "extVideoFrame-mask" + String(this.#vframeId)); + body.append("subfolder", "extVideoFrame" + idSitr); const resp = await api.fetchApi("/upload/image", { method: "POST", body, }); - }); - + }; this.saveButton.innerText = "Saving..."; + + const maskDirId = Date.now(); + const sketchDirId = maskDirId + 1; + for (let i = 0; i < this.backMaskCanvases.length; i++) + await uploadImages(this.backMaskCanvases[i], i, String(maskDirId)); + for (let i = 0; i < this.backSketchCanvases.length; i++) + await uploadImages(this.backSketchCanvases[i], i, String(sketchDirId)); + + this.#updatePathDataHnadler(maskDirId, sketchDirId); this.saveButton.disabled = true; this.close(); } + + + #_updatePathDataHnadler = null; + get #updatePathDataHnadler() { + return this.#_updatePathDataHnadler ?? ((maskDirId, sketchDirId) => {}); + } + set updatePathDataHnadler(value) { + this.#_updatePathDataHnadler = value; + } } diff --git a/js/main.js b/js/main.js index 64e04b0..1f2c7b6 100644 --- a/js/main.js +++ b/js/main.js @@ -9,20 +9,27 @@ app.registerExtension({ name: "ComfyUI.animatediff.inpaint", async init(app) - { - app.extVideoFrameIdCount = 0; - + { ComfyApp.ext_open_maskeditor_for_vframes = - function (pathData) { + function (idWidget) { + const dlg = MaskEditorForVframes.getInstance(); - dlg.path_data = pathData; + dlg.path_data = idWidget.value; + dlg.updatePathDataHnadler = (maskDirId, sketchDirId) => { + let imageData = idWidget.value.split(":"); + let dirIdStr = imageData.pop(); + let dirIds = dirIdStr.split("/"); + + dirIds = [dirIds[0], maskDirId, sketchDirId]; + dirIdStr = dirIds.join("/"); + imageData.push(dirIdStr); + + idWidget.value = imageData.join(":"); + }; if(!dlg.isOpened()) { dlg.show(); } }; - - // const context_predicate = () => ComfyApp.clipspace && ComfyApp.clipspace.imgs && ComfyApp.clipspace.imgs.length > 0 - // ClipspaceDialog.registerButton("MaskEditor", context_predicate, ComfyApp.open_maskeditor); }, async setup() { @@ -46,7 +53,7 @@ app.registerExtension({ // ComfyApp.copyToClipspace(this); // ComfyApp.clipspace_return_node = this; - ComfyApp.ext_open_maskeditor_for_vframes(this.widgets_values[0]); + ComfyApp.ext_open_maskeditor_for_vframes(this.widgets[0]); } }); } @@ -85,15 +92,17 @@ app.registerExtension({ // Add our own callback to the combo widget to render an image when it changes idWidget.callback = function () { - let imageData = idWidget.value.split(":", 3); - showImage(imageData[0] + "/" + imageData[1]); + let imageData = idWidget.value.split(":"); + let dirIds = imageData[imageData.length - 1].split("/"); + showImage(dirIds[0] + "/" + imageData[0]); }; // On load if we have a value then render the image requestAnimationFrame(() => { if (idWidget.value) { - let imageData = idWidget.value.split(":", 3); - showImage(imageData[0] + "/" + imageData[1]); + let imageData = idWidget.value.split(":"); + let dirIds = imageData[imageData.length - 1].split("/"); + showImage(dirIds[0] + "/" + imageData[0]); } }); @@ -135,14 +144,15 @@ app.registerExtension({ webkitdirectory: true, onchange: async () => { if (fileInput.files.length) { - const vframeId = app.extVideoFrameIdCount++; - let vframeData = String(vframeId); + const vframeId = Date.now(); + let vframeData = ""; for (let i = 0; i < fileInput.files.length; i++) { - vframeData += ":" + fileInput.files[i].name + vframeData += fileInput.files[i].name + ":"; await uploadFile( fileInput.files[i], i === 0, vframeId); } + vframeData += String(vframeId); idWidget.value = vframeData; } },