add VAEEncodeForInpaint_Frames

This commit is contained in:
shadowcz007
2024-03-26 16:48:45 +08:00
parent a2ece25ac0
commit 1fe99f33b2
3 changed files with 79 additions and 7 deletions
+5 -4
View File
@@ -603,7 +603,7 @@ from .nodes.Mask import MaskListReplace,MaskListMerge,OutlineMask,FeatheredMask
from .nodes.Style import ApplyVisualStylePrompting,StyleAlignedReferenceSampler,StyleAlignedBatchAlign,StyleAlignedSampleReferenceLatents
from .nodes.Video import LoadVideoAndSegment,ImageListReplace
from .nodes.Video import LoadVideoAndSegment,ImageListReplace,VAEEncodeForInpaint_Frames
# 要导出的所有节点及其名称的字典
@@ -683,8 +683,8 @@ NODE_CLASS_MAPPINGS = {
"LoadVideoAndSegment_":LoadVideoAndSegment,
"ListSplit_":ListSplit,
"MaskListReplace_":MaskListReplace,
"ImageListReplace_":ImageListReplace
# "LaMaInpainting":LaMaInpainting
"ImageListReplace_":ImageListReplace,
"VAEEncodeForInpaint_Frames":VAEEncodeForInpaint_Frames
# "GamePal":GamePal
}
@@ -733,7 +733,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"GridDisplayAndSave":"Grid Display And Save",
"GridInput":"Grid Input",
"GridOutput":"Grid Output",
"GetImageSize_":"Get Image Size"
"GetImageSize_":"Get Image Size",
"VAEEncodeForInpaint_Frames":"VAE Encode For Inpaint Frames"
}
# web ui的节点功能
+73 -2
View File
@@ -4,7 +4,7 @@ import json
import subprocess
import shutil
import re
import time
import time,math
import numpy as np
from typing import List
import torch
@@ -256,12 +256,15 @@ class LoadVideoAndSegment:
def load_video(self, video,video_segment_frames,transition_frames ):
frame_load_cap=0
skip_first_frames=0
video_path = folder_paths.get_annotated_filepath(video)
# check if video is a gif - will need to use cv fallback to read frames
# use cv fallback if ffmpeg not installed or gif
if ffmpeg_path is None:
return self.load_video_cv_fallback(video, frame_load_cap, skip_first_frames)
# otherwise, continue with ffmpeg
video_path = folder_paths.get_annotated_filepath(video)
args_dummy = [ffmpeg_path, "-i", video_path, "-f", "null", "-"]
try:
with subprocess.Popen(args_dummy, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) as proc:
@@ -335,3 +338,71 @@ class LoadVideoAndSegment:
return True
class VAEEncodeForInpaint_Frames:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"vae": ("VAE", ),
"images": ("IMAGE", ),
"masks": ("MASK", ),
"grow_mask_by": ("INT", {"default": 6, "min": 0, "max": 64, "step": 1}),
}}
FUNCTION = "encode"
RETURN_TYPES = ("LATENT",)
RETURN_NAMES = ("LATENT",)
CATEGORY = "♾️Mixlab/Video"
OUTPUT_NODE = True
INPUT_IS_LIST = True
OUTPUT_IS_LIST = (True,)
def encode(self, vae, images, masks, grow_mask_by=[6]):
vae=vae[0]
grow_mask_by=grow_mask_by[0]
result=[]
for i in range(len(images)):
pixels=images[i]
mask=masks[i]
x = (pixels.shape[1] // 8) * 8
y = (pixels.shape[2] // 8) * 8
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
pixels = pixels.clone()
if pixels.shape[1] != x or pixels.shape[2] != y:
x_offset = (pixels.shape[1] % 8) // 2
y_offset = (pixels.shape[2] % 8) // 2
pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
#grow mask by a few pixels to keep things seamless in latent space
if grow_mask_by == 0:
mask_erosion = mask
else:
kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
padding = math.ceil((grow_mask_by - 1) / 2)
mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)
m = (1.0 - mask.round()).squeeze(1)
for i in range(3):
pixels[:,:,:,i] -= 0.5
pixels[:,:,:,i] *= m
pixels[:,:,:,i] += 0.5
t = vae.encode(pixels)
result.append({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())})
return (result, )
+1 -1
View File
@@ -180,7 +180,7 @@ function videoUpload (node, inputName, inputData, app) {
const fileInput = document.createElement('input')
Object.assign(fileInput, {
type: 'file',
accept: 'video/webm,video/mp4,video/mkv,image/gif',
accept: 'video/*,.mkv,video/webm,video/mp4,video/x-matroska,image/gif',
style: 'display: none',
onchange: async () => {
if (fileInput.files.length) {