From 1fe99f33b2eb4e3a9a2c4bfd7c66b67e00a9f7c2 Mon Sep 17 00:00:00 2001
From: shadowcz007 <chizhiwei007@163.com>
Date: Tue, 26 Mar 2024 16:48:45 +0800
Subject: [PATCH] add VAEEncodeForInpaint_Frames

---
 __init__.py                   |  9 +++--
 nodes/Video.py                | 75 ++++++++++++++++++++++++++++++++++-
 web/javascript/videoupload.js |  2 +-
 3 files changed, 79 insertions(+), 7 deletions(-)

diff --git a/__init__.py b/__init__.py
index 9f04fe3..023a604 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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的节点功能
diff --git a/nodes/Video.py b/nodes/Video.py
index 51b187b..87f7ccc 100644
--- a/nodes/Video.py
+++ b/nodes/Video.py
@@ -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, )
\ No newline at end of file
diff --git a/web/javascript/videoupload.js b/web/javascript/videoupload.js
index 5d9bfe1..2265b46 100644
--- a/web/javascript/videoupload.js
+++ b/web/javascript/videoupload.js
@@ -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) {