Added Video Merge node

This commit is contained in:
DoctorDiffusion
2024-10-31 13:07:50 -04:00
committed by GitHub
parent 5315a0c47b
commit 0403aa4107
2 changed files with 60 additions and 2 deletions
+20 -2
View File
@@ -1,3 +1,21 @@
from .final_frame_selector import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
from .final_frame_selector import (
FinalFrameSelector,
NODE_CLASS_MAPPINGS as FINAL_FRAME_SELECTOR_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as FINAL_FRAME_SELECTOR_DISPLAY_MAPPINGS
)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
from .video_merge import (
VideoMerge,
NODE_CLASS_MAPPINGS as VIDEO_MERGE_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS as VIDEO_MERGE_DISPLAY_MAPPINGS
)
NODE_CLASS_MAPPINGS = {
**FINAL_FRAME_SELECTOR_MAPPINGS,
**VIDEO_MERGE_MAPPINGS
}
NODE_DISPLAY_NAME_MAPPINGS = {
**FINAL_FRAME_SELECTOR_DISPLAY_MAPPINGS,
**VIDEO_MERGE_DISPLAY_MAPPINGS
}
+40
View File
@@ -0,0 +1,40 @@
import torch
class VideoMerge:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Video_A": ("IMAGE",),
"Video_B": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "merge_videos"
CATEGORY = "Doctor-Tools"
def merge_videos(self, Video_A, Video_B):
# Ensure both inputs are 4D tensors
if Video_A.dim() == 3:
Video_A = Video_A.unsqueeze(0)
if Video_B.dim() == 3:
Video_B = Video_B.unsqueeze(0)
# Concatenate the images along the first dimension (batch dimension)
Video = torch.cat((Video_A, Video_B), dim=0)
return (Video,)
# A dictionary that contains all nodes you want to export with their names
NODE_CLASS_MAPPINGS = {
"VideoMerge": VideoMerge
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"VideoMerge": "Video Merge"
}