From 78378d8cfbde9c733ed10e4d52460f8616e4984d Mon Sep 17 00:00:00 2001 From: bruefire Date: Wed, 6 Dec 2023 02:00:13 +0900 Subject: [PATCH] initial commit --- AnimateDiffInpaint.py | 81 +++++ __init__.py | 17 + js/MaskEditorForVframes.js | 652 +++++++++++++++++++++++++++++++++++++ js/main.js | 186 +++++++++++ 4 files changed, 936 insertions(+) create mode 100644 AnimateDiffInpaint.py create mode 100644 __init__.py create mode 100644 js/MaskEditorForVframes.js create mode 100644 js/main.js diff --git a/AnimateDiffInpaint.py b/AnimateDiffInpaint.py new file mode 100644 index 0000000..61f1758 --- /dev/null +++ b/AnimateDiffInpaint.py @@ -0,0 +1,81 @@ +import os +import torch +import numpy as np +from PIL import Image, ImageOps + + +class LoadImagesFromDirInpaint: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images_id": ("STRING", {"frames_upload": True}), + }, + "optional": { + "image_load_cap": ("INT", {"default": 0, "min": 0, "step": 1}), + "start_index": ("INT", {"default": 0, "min": 0, "step": 1}), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE", "INT") + RETURN_NAMES = ("images", "mask_images", "image_count",) + FUNCTION = "load_images" + + CATEGORY = "image" + + + def load_images(self, images_id: str, image_load_cap: int = 0, start_index: int = 0): + comfy_custom_dir = os.path.dirname(os.path.abspath(__file__)) + 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) + + 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 = img_files[start_index:] + + images = [] + + limit_images = False + if image_load_cap > 0: + limit_images = True + image_count = 0 + + for image_path in 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) + image = i.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 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..b7a0bc3 --- /dev/null +++ b/__init__.py @@ -0,0 +1,17 @@ +import os, shutil +import folder_paths +from .AnimateDiffInpaint import LoadImagesFromDirInpaint + + +NODE_CLASS_MAPPINGS = { "VFrame Loader With Mask Editor" : LoadImagesFromDirInpaint } +NODE_DISPLAY_NAME_MAPPINGS = { "VFrame Loader With Mask Editor" : "AnimateDiffInpaint" } +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] + + +module_js_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "js") +application_root_directory = os.path.dirname(folder_paths.__file__) +application_web_extensions_directory = os.path.join(application_root_directory, "web", "extensions", "AnimateDiffInpaint") + +if os.path.exists(application_web_extensions_directory): + shutil.rmtree(application_web_extensions_directory) +shutil.copytree(module_js_directory, application_web_extensions_directory, dirs_exist_ok=True) diff --git a/js/MaskEditorForVframes.js b/js/MaskEditorForVframes.js new file mode 100644 index 0000000..90cc1f4 --- /dev/null +++ b/js/MaskEditorForVframes.js @@ -0,0 +1,652 @@ +import { app } from "../../scripts/app.js"; +import { ComfyDialog, $el } from "../../scripts/ui.js"; +import { ComfyApp } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js" +//import { ClipspaceDialog } from "./clipspace.js"; + +// Helper function to convert a data URL to a Blob object +function dataURLToBlob(dataURL) { + const parts = dataURL.split(';base64,'); + const contentType = parts[0].split(':')[1]; + const byteString = atob(parts[1]); + const arrayBuffer = new ArrayBuffer(byteString.length); + const uint8Array = new Uint8Array(arrayBuffer); + for (let i = 0; i < byteString.length; i++) { + uint8Array[i] = byteString.charCodeAt(i); + } + 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); + const backupData = backupCtx.getImageData(0, 0, backupCanvas.width, backupCanvas.height); + + // refine mask image + for (let i = 0; i < backupData.data.length; i += 4) { + if(backupData.data[i+3] == 255) + backupData.data[i+3] = 0; + else + backupData.data[i+3] = 255; + + backupData.data[i] = 0; + backupData.data[i+1] = 0; + backupData.data[i+2] = 0; + } + + backupCtx.globalCompositeOperation = 'source-over'; + backupCtx.putImageData(backupData, 0, 0); +} + +class MaskEditorDialog extends ComfyDialog { + static instance = null; + + static getInstance() { + if(!MaskEditorDialog.instance) { + MaskEditorDialog.instance = new MaskEditorDialog(app); + } + + return MaskEditorDialog.instance; + } + + is_layout_created = false; + + set path_data(value) { + const parts = value?.split(":"); + if (parts) { + const id = parts.shift(); + this.#vframeId = id; + this.#paths = parts.map(file_name => { + return { + filename: file_name, + subfolder: "extVideoFrame" + id, + type: "input" + } + }); + } + } + #paths = null; + #vframeId = -1; + + #selectedIndex = 0; + + constructor() { + super(); + this.element = $el("div.comfy-modal", { parent: document.body }, + [ $el("div.comfy-modal-content", + [...this.createButtons()]), + ]); + } + + createButtons() { + return []; + } + + createButton(name, callback) { + var button = document.createElement("button"); + button.innerText = name; + button.addEventListener("click", callback); + return button; + } + + createLeftButton(name, callback) { + var button = this.createButton(name, callback); + button.style.cssFloat = "left"; + button.style.marginRight = "4px"; + return button; + } + + createRightButton(name, callback) { + var button = this.createButton(name, callback); + button.style.cssFloat = "right"; + button.style.marginLeft = "4px"; + return button; + } + + createLeftSlider(self, name, callback) { + const divElement = document.createElement('div'); + divElement.id = "maskeditor-slider"; + divElement.style.cssFloat = "left"; + divElement.style.fontFamily = "sans-serif"; + divElement.style.marginRight = "4px"; + divElement.style.color = "var(--input-text)"; + divElement.style.backgroundColor = "var(--comfy-input-bg)"; + divElement.style.borderRadius = "8px"; + divElement.style.borderColor = "var(--border-color)"; + divElement.style.borderStyle = "solid"; + divElement.style.fontSize = "15px"; + divElement.style.height = "21px"; + divElement.style.padding = "1px 6px"; + divElement.style.display = "flex"; + divElement.style.position = "relative"; + divElement.style.top = "2px"; + self.brush_slider_input = document.createElement('input'); + self.brush_slider_input.setAttribute('type', 'range'); + self.brush_slider_input.setAttribute('min', '1'); + self.brush_slider_input.setAttribute('max', '100'); + self.brush_slider_input.setAttribute('value', '10'); + const labelElement = document.createElement("label"); + labelElement.textContent = name; + + divElement.appendChild(labelElement); + divElement.appendChild(self.brush_slider_input); + + self.brush_slider_input.addEventListener("change", callback); + + return divElement; + } + + setlayout(imgCanvas, maskCanvas) { + const self = this; + + // If it is specified as relative, using it only as a hidden placeholder for padding is recommended + // to prevent anomalies where it exceeds a certain size and goes outside of the window. + var placeholder = document.createElement("div"); + placeholder.style.position = "relative"; + placeholder.style.height = "50px"; + + var bottom_panel = document.createElement("div"); + bottom_panel.style.position = "absolute"; + bottom_panel.style.bottom = "0px"; + bottom_panel.style.left = "20px"; + bottom_panel.style.right = "20px"; + bottom_panel.style.height = "50px"; + + var brush = document.createElement("div"); + brush.id = "brush"; + brush.style.backgroundColor = "transparent"; + brush.style.outline = "1px dashed black"; + brush.style.boxShadow = "0 0 0 1px white"; + brush.style.borderRadius = "50%"; + brush.style.MozBorderRadius = "50%"; + brush.style.WebkitBorderRadius = "50%"; + brush.style.position = "absolute"; + brush.style.zIndex = 8889; + brush.style.pointerEvents = "none"; + this.brush = brush; + this.element.appendChild(imgCanvas); + this.element.appendChild(maskCanvas); + this.element.appendChild(placeholder); // must below z-index than bottom_panel to avoid covering button + this.element.appendChild(bottom_panel); + document.body.appendChild(brush); + + var brush_size_slider = this.createLeftSlider(self, "Thickness", (event) => { + self.brush_size = event.target.value; + self.updateBrushPreview(self, null, null); + }); + var clearButton = this.createLeftButton("Clear", + () => { + self.maskCtx.clearRect(0, 0, self.maskCanvas.width, self.maskCanvas.height); + self.backupCtx.clearRect(0, 0, self.backupCanvas.width, self.backupCanvas.height); + }); + var prevButton = this.createLeftButton("Prev", + () => { + if (this.#selectedIndex > 0) { + const backMask = this.backMaskCanvases[this.#selectedIndex]; + backMask.width = this.maskCanvas.width; + backMask.height = this.maskCanvas.height; + backMask.getContext('2d').drawImage(this.maskCanvas, 0, 0, backMask.width, backMask.height); + + const params = new URLSearchParams(this.#paths[--this.#selectedIndex]); + this.image.src = new URL(api.apiURL("/view?" + params.toString()), window.location.href); + } + }); + var nextButton = this.createLeftButton("Next", + () => { + if (this.#selectedIndex < this.#paths.length - 1) { + const backMask = this.backMaskCanvases[this.#selectedIndex]; + backMask.width = this.maskCanvas.width; + backMask.height = this.maskCanvas.height; + backMask.getContext('2d').drawImage(this.maskCanvas, 0, 0, backMask.width, backMask.height); + + const params = new URLSearchParams(this.#paths[++this.#selectedIndex]); + this.image.src = new URL(api.apiURL("/view?" + params.toString()), window.location.href); + } + }); + var cancelButton = this.createRightButton("Cancel", () => { + document.removeEventListener("mouseup", MaskEditorDialog.handleMouseUp); + document.removeEventListener("keydown", MaskEditorDialog.handleKeyDown); + self.close(); + }); + + this.saveButton = this.createRightButton("Save", () => { + document.removeEventListener("mouseup", MaskEditorDialog.handleMouseUp); + document.removeEventListener("keydown", MaskEditorDialog.handleKeyDown); + self.save(); + }); + + this.element.appendChild(imgCanvas); + this.element.appendChild(maskCanvas); + this.element.appendChild(placeholder); // must below z-index than bottom_panel to avoid covering button + this.element.appendChild(bottom_panel); + + bottom_panel.appendChild(clearButton); + bottom_panel.appendChild(this.saveButton); + bottom_panel.appendChild(cancelButton); + bottom_panel.appendChild(brush_size_slider); + bottom_panel.appendChild(prevButton); + bottom_panel.appendChild(nextButton); + + imgCanvas.style.position = "relative"; + imgCanvas.style.top = "200"; + imgCanvas.style.left = "0"; + + maskCanvas.style.position = "absolute"; + maskCanvas.style.opacity = "0.66"; + } + + show() { + if(!this.is_layout_created) { + // layout + const imgCanvas = document.createElement('canvas'); + const maskCanvas = document.createElement('canvas'); + const backupCanvas = document.createElement('canvas'); + const backMaskCanvases = [...new Array(this.#paths.length)].map(_ => document.createElement("canvas")); + + imgCanvas.id = "imageCanvas"; + maskCanvas.id = "maskCanvas"; + backupCanvas.id = "backupCanvas"; + + this.setlayout(imgCanvas, maskCanvas); + + // prepare content + this.imgCanvas = imgCanvas; + this.maskCanvas = maskCanvas; + this.backupCanvas = backupCanvas; + this.backMaskCanvases = backMaskCanvases; + this.maskCtx = maskCanvas.getContext('2d'); + this.backupCtx = backupCanvas.getContext('2d'); + + this.setEventHandler(maskCanvas); + + this.is_layout_created = true; + + // replacement of onClose hook since close is not real close + const self = this; + const observer = new MutationObserver(function(mutations) { + mutations.forEach(function(mutation) { + if (mutation.type === 'attributes' && mutation.attributeName === 'style') { + if(self.last_display_style && self.last_display_style != 'none' && self.element.style.display == 'none') { + ComfyApp.onClipspaceEditorClosed(); + } + + self.last_display_style = self.element.style.display; + } + }); + }); + + const config = { attributes: true }; + observer.observe(this.element, config); + } + + this.setImages(this.imgCanvas, this.backupCanvas); + + if(ComfyApp.clipspace_return_node) { + this.saveButton.innerText = "Save to node"; + } + else { + this.saveButton.innerText = "Save"; + } + this.saveButton.disabled = false; + + this.element.style.display = "block"; + this.element.style.zIndex = 8888; // NOTE: alert dialog must be high priority. + } + + isOpened() { + return this.element.style.display == "block"; + } + + setImages(imgCanvas, backupCanvas) { + const imgCtx = imgCanvas.getContext('2d'); + const backupCtx = backupCanvas.getContext('2d'); + const maskCtx = this.maskCtx; + const maskCanvas = this.maskCanvas; + + backupCtx.clearRect(0,0,this.backupCanvas.width,this.backupCanvas.height); + imgCtx.clearRect(0,0,this.imgCanvas.width,this.imgCanvas.height); + maskCtx.clearRect(0,0,this.maskCanvas.width,this.maskCanvas.height); + + // image load + const orig_image = new Image(); + window.addEventListener("resize", () => { + // repositioning + imgCanvas.width = window.innerWidth - 250; + imgCanvas.height = window.innerHeight - 200; + + // redraw image + let drawWidth = orig_image.width; + let drawHeight = orig_image.height; + if (orig_image.width > imgCanvas.width) { + drawWidth = imgCanvas.width; + drawHeight = (drawWidth / orig_image.width) * orig_image.height; + } + + if (drawHeight > imgCanvas.height) { + drawHeight = imgCanvas.height; + drawWidth = (drawHeight / orig_image.height) * orig_image.width; + } + + imgCtx.drawImage(orig_image, 0, 0, drawWidth, drawHeight); + + // update mask + maskCanvas.width = drawWidth; + 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); + maskCtx.drawImage(this.backMaskCanvases[this.#selectedIndex], 0, 0, maskCanvas.width, maskCanvas.height); + }); + + const touched_image = new Image(); + + touched_image.onload = function() { + backupCanvas.width = touched_image.width; + backupCanvas.height = touched_image.height; + + prepareRGB(touched_image, backupCanvas, backupCtx); + }; + + const alpha_url = new URL(api.apiURL("/view?" + new URLSearchParams(this.#paths[this.#selectedIndex]).toString()), window.location.href); + alpha_url.searchParams.delete('channel'); + alpha_url.searchParams.delete('preview'); + alpha_url.searchParams.set('channel', 'a'); + touched_image.src = alpha_url; + + // original image load + orig_image.onload = function() { + window.dispatchEvent(new Event('resize')); + }; + + const rgb_url = new URL(api.apiURL("/view?" + new URLSearchParams(this.#paths[this.#selectedIndex]).toString()), window.location.href); + rgb_url.searchParams.delete('channel'); + rgb_url.searchParams.set('channel', 'rgb'); + orig_image.src = rgb_url; + this.image = orig_image; + } + + setEventHandler(maskCanvas) { + maskCanvas.addEventListener("contextmenu", (event) => { + event.preventDefault(); + }); + + const self = this; + maskCanvas.addEventListener('wheel', (event) => this.handleWheelEvent(self,event)); + maskCanvas.addEventListener('pointerdown', (event) => this.handlePointerDown(self,event)); + document.addEventListener('pointerup', MaskEditorDialog.handlePointerUp); + maskCanvas.addEventListener('pointermove', (event) => this.draw_move(self,event)); + maskCanvas.addEventListener('touchmove', (event) => this.draw_move(self,event)); + maskCanvas.addEventListener('pointerover', (event) => { this.brush.style.display = "block"; }); + maskCanvas.addEventListener('pointerleave', (event) => { this.brush.style.display = "none"; }); + document.addEventListener('keydown', MaskEditorDialog.handleKeyDown); + } + + brush_size = 10; + drawing_mode = false; + lastx = -1; + lasty = -1; + lasttime = 0; + + static handleKeyDown(event) { + const self = MaskEditorDialog.instance; + if (event.key === ']') { + self.brush_size = Math.min(self.brush_size+2, 100); + } else if (event.key === '[') { + self.brush_size = Math.max(self.brush_size-2, 1); + } else if(event.key === 'Enter') { + self.save(); + } + + self.updateBrushPreview(self); + } + + static handlePointerUp(event) { + event.preventDefault(); + MaskEditorDialog.instance.drawing_mode = false; + } + + updateBrushPreview(self) { + const brush = self.brush; + + var centerX = self.cursorX; + var centerY = self.cursorY; + + brush.style.width = self.brush_size * 2 + "px"; + brush.style.height = self.brush_size * 2 + "px"; + brush.style.left = (centerX - self.brush_size) + "px"; + brush.style.top = (centerY - self.brush_size) + "px"; + } + + handleWheelEvent(self, event) { + if(event.deltaY < 0) + self.brush_size = Math.min(self.brush_size+2, 100); + else + self.brush_size = Math.max(self.brush_size-2, 1); + + self.brush_slider_input.value = self.brush_size; + + self.updateBrushPreview(self); + } + + draw_move(self, event) { + event.preventDefault(); + + this.cursorX = event.pageX; + this.cursorY = event.pageY; + + self.updateBrushPreview(self); + + if (window.TouchEvent && event instanceof TouchEvent || event.buttons == 1) { + var diff = performance.now() - self.lasttime; + + const maskRect = self.maskCanvas.getBoundingClientRect(); + + var x = event.offsetX; + var y = event.offsetY + + if(event.offsetX == null) { + x = event.targetTouches[0].clientX - maskRect.left; + } + + if(event.offsetY == null) { + y = event.targetTouches[0].clientY - maskRect.top; + } + + var brush_size = this.brush_size; + if(event instanceof PointerEvent && event.pointerType == 'pen') { + brush_size *= event.pressure; + this.last_pressure = event.pressure; + } + else if(window.TouchEvent && event instanceof TouchEvent && diff < 20){ + // The firing interval of PointerEvents in Pen is unreliable, so it is supplemented by TouchEvents. + brush_size *= this.last_pressure; + } + else { + brush_size = this.brush_size; + } + + if(diff > 20 && !this.drawing_mode) + requestAnimationFrame(() => { + self.maskCtx.beginPath(); + self.maskCtx.fillStyle = "rgb(255, 255, 255)"; + self.maskCtx.globalCompositeOperation = "source-over"; + self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false); + self.maskCtx.fill(); + self.lastx = x; + self.lasty = y; + }); + else + requestAnimationFrame(() => { + self.maskCtx.beginPath(); + self.maskCtx.fillStyle = "rgb(255, 255, 255)"; + self.maskCtx.globalCompositeOperation = "source-over"; + + var dx = x - self.lastx; + var dy = y - self.lasty; + + var distance = Math.sqrt(dx * dx + dy * dy); + var directionX = dx / distance; + var directionY = dy / distance; + + for (var i = 0; i < distance; i+=5) { + var px = self.lastx + (directionX * i); + var py = self.lasty + (directionY * i); + self.maskCtx.arc(px, py, brush_size, 0, Math.PI * 2, false); + self.maskCtx.fill(); + } + self.lastx = x; + self.lasty = y; + }); + + self.lasttime = performance.now(); + } + else if(event.buttons == 2 || event.buttons == 5 || event.buttons == 32) { + const maskRect = self.maskCanvas.getBoundingClientRect(); + const x = event.offsetX || event.targetTouches[0].clientX - maskRect.left; + const y = event.offsetY || event.targetTouches[0].clientY - maskRect.top; + + var brush_size = this.brush_size; + if(event instanceof PointerEvent && event.pointerType == 'pen') { + brush_size *= event.pressure; + this.last_pressure = event.pressure; + } + else if(window.TouchEvent && event instanceof TouchEvent && diff < 20){ + brush_size *= this.last_pressure; + } + else { + brush_size = this.brush_size; + } + + if(diff > 20 && !drawing_mode) // cannot tracking drawing_mode for touch event + requestAnimationFrame(() => { + self.maskCtx.beginPath(); + self.maskCtx.globalCompositeOperation = "destination-out"; + self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false); + self.maskCtx.fill(); + self.lastx = x; + self.lasty = y; + }); + else + requestAnimationFrame(() => { + self.maskCtx.beginPath(); + self.maskCtx.globalCompositeOperation = "destination-out"; + + var dx = x - self.lastx; + var dy = y - self.lasty; + + var distance = Math.sqrt(dx * dx + dy * dy); + var directionX = dx / distance; + var directionY = dy / distance; + + for (var i = 0; i < distance; i+=5) { + var px = self.lastx + (directionX * i); + var py = self.lasty + (directionY * i); + self.maskCtx.arc(px, py, brush_size, 0, Math.PI * 2, false); + self.maskCtx.fill(); + } + self.lastx = x; + self.lasty = y; + }); + + self.lasttime = performance.now(); + } + } + + handlePointerDown(self, event) { + var brush_size = this.brush_size; + if(event instanceof PointerEvent && event.pointerType == 'pen') { + brush_size *= event.pressure; + this.last_pressure = event.pressure; + } + + if ([0, 2, 5].includes(event.button)) { + self.drawing_mode = true; + + event.preventDefault(); + const maskRect = self.maskCanvas.getBoundingClientRect(); + const x = event.offsetX || event.targetTouches[0].clientX - maskRect.left; + const y = event.offsetY || event.targetTouches[0].clientY - maskRect.top; + + self.maskCtx.beginPath(); + if (event.button == 0) { + self.maskCtx.fillStyle = "rgb(255, 255, 255)"; + self.maskCtx.globalCompositeOperation = "source-over"; + } else { + self.maskCtx.globalCompositeOperation = "destination-out"; + } + self.maskCtx.arc(x, y, brush_size, 0, Math.PI * 2, false); + self.maskCtx.fill(); + self.lastx = x; + self.lasty = y; + self.lasttime = performance.now(); + } + } + + async save() { + const backupCtx = this.backupCanvas.getContext('2d'); + + backupCtx.clearRect(0,0,this.backupCanvas.width,this.backupCanvas.height); + backupCtx.drawImage(this.maskCanvas, + 0, 0, this.maskCanvas.width, this.maskCanvas.height, + 0, 0, this.backupCanvas.width, this.backupCanvas.height); + + + this.backMaskCanvases.forEach(async (maskCanvas, i) => { + const body = new FormData(); + const dataURL = maskCanvas.toDataURL(); + const blob = dataURLToBlob(dataURL); + const extPos = this.#paths[i].filename.lastIndexOf("."); + const filename = this.#paths[i].filename.substr(0, extPos) + ".png"; + + body.append("image", blob, filename); + body.append("subfolder", "extVideoFrame-mask" + String(this.#vframeId)); + const resp = await api.fetchApi("/upload/image", { + method: "POST", + body, + }); + }); + + + this.saveButton.innerText = "Saving..."; + this.saveButton.disabled = true; + + this.close(); + } +} + + +export default MaskEditorDialog; diff --git a/js/main.js b/js/main.js new file mode 100644 index 0000000..64e04b0 --- /dev/null +++ b/js/main.js @@ -0,0 +1,186 @@ + +import { app } from "../../../scripts/app.js"; +import { ComfyApp } from "../../../scripts/app.js"; +import { api } from "../../../scripts/api.js"; +import MaskEditorForVframes from "./MaskEditorForVframes.js"; + +app.registerExtension({ + + name: "ComfyUI.animatediff.inpaint", + + async init(app) + { + app.extVideoFrameIdCount = 0; + + ComfyApp.ext_open_maskeditor_for_vframes = + function (pathData) { + const dlg = MaskEditorForVframes.getInstance(); + dlg.path_data = pathData; + 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() + { + }, + async beforeRegisterNodeDef(nodeType, nodeData, app) + { + if (nodeData.name !== "VFrame Loader With Mask Editor") + return; + + // menu initializing + const originalProto = nodeType.prototype.getExtraMenuOptions; + nodeType.prototype.getExtraMenuOptions = function (_, options){ + // do original processing + originalProto.bind(this)(_, options); + + // add mask menu + if (this.extVframes) { + options.push({ + content: "Open in MaskEditor", + callback: (obj) => { + // ComfyApp.copyToClipspace(this); + // ComfyApp.clipspace_return_node = this; + + ComfyApp.ext_open_maskeditor_for_vframes(this.widgets_values[0]); + } + }); + } + } + + // head image drawing + // ... + + // widget setting + // Adds an upload button to the nodes + if (nodeData?.input?.required?.images_id?.[1]?.frames_upload === true) { + nodeData.input.required.upload = ["VFRAMESUPLOAD"]; + } + }, + async getCustomWidgets() + { + return { + VFRAMESUPLOAD(node, inputName, inputData, app) { + const idWidget = node.widgets.find((w) => w.name === "images_id"); + + function showImage(name) { + const img = new Image(); + img.onload = () => { + node.extVframes = [img]; + app.graph.setDirtyCanvas(true); + }; + let folder_separator = name.lastIndexOf("/"); + let subfolder = ""; + if (folder_separator > -1) { + subfolder = name.substring(0, folder_separator); + name = name.substring(folder_separator + 1); + } + img.src = api.apiURL(`/view?filename=${encodeURIComponent(name)}&type=input&subfolder=${subfolder}${app.getPreviewFormatParam()}`); + node.setSizeForImage?.(); + } + + // 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]); + }; + + // 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]); + } + }); + + async function uploadFile(file, updateNode, vframeId) { + if (!file.type.startsWith("image/")) + return; + + try { + // Wrap file in formdata so it includes filename + const body = new FormData(); + body.append("image", file, file.name); + body.append("subfolder", "extVideoFrame" + String(vframeId)); + const resp = await api.fetchApi("/upload/image", { + method: "POST", + body, + }); + + if (resp.status === 200) { + const data = await resp.json(); + // update the widget value + let path = data.name; + if (data.subfolder) path = data.subfolder + "/" + path; + + if (updateNode) { + showImage(path); + } + } else { + alert(resp.status + " - " + resp.statusText); + } + } catch (error) { + alert(error); + } + } + + const fileInput = document.createElement("input"); + Object.assign(fileInput, { + type: "file", + style: "display: none", + webkitdirectory: true, + onchange: async () => { + if (fileInput.files.length) { + const vframeId = app.extVideoFrameIdCount++; + let vframeData = String(vframeId); + + for (let i = 0; i < fileInput.files.length; i++) { + vframeData += ":" + fileInput.files[i].name + await uploadFile( + fileInput.files[i], i === 0, vframeId); + } + idWidget.value = vframeData; + } + }, + }); + document.body.append(fileInput); + + // Create the button widget for selecting the files + let uploadWidget = node.addWidget("button", "choose image directory to upload", "image", () => { + fileInput.click(); + }); + uploadWidget.serialize = false; + + // Add handler to check if an image is being dragged over our node + node.onDragOver = function (e) { + if (e.dataTransfer && e.dataTransfer.items) { + const image = [...e.dataTransfer.items].find((f) => f.kind === "file"); + return !!image; + } + + return false; + }; + + // On drop upload files + node.onDragDrop = function (e) { + console.log("onDragDrop called"); + let handled = false; + for (const file of e.dataTransfer.files) { + if (file.type.startsWith("image/")) { + uploadFile(file, !handled); // Dont await these, any order is fine, only update on first one + handled = true; + } + } + + return handled; + }; + + return { widget: uploadWidget }; + } + } + } +}) \ No newline at end of file