nodes/bruefire-ComfyUI-SeqImageLoader
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

implemented sketch logics.

Browse Source
This commit is contained in:
bruefire
2023-12-08 05:17:51 +09:00
parent 49abcb8070
commit 961be2f2b8
3 changed files with 147 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files
AnimateDiffInpaint.py
+74 -16
View File
@@ -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
js/MaskEditorForVframes.js
+47 -52
View File
@@ -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;
}
}
js/main.js
+26 -16
View File
@@ -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;
}
},