From 0403aa4107250e99f623c71a33c07fce2029a13d Mon Sep 17 00:00:00 2001
From: DoctorDiffusion <139909432+DoctorDiffusion@users.noreply.github.com>
Date: Thu, 31 Oct 2024 13:07:50 -0400
Subject: [PATCH] Added Video Merge node

---
 __init__.py    | 22 ++++++++++++++++++++--
 video_merge.py | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)
 create mode 100644 video_merge.py

diff --git a/__init__.py b/__init__.py
index d90d0f1..3ac7b76 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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"]
\ No newline at end of file
+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
+}
\ No newline at end of file
diff --git a/video_merge.py b/video_merge.py
new file mode 100644
index 0000000..b1854ac
--- /dev/null
+++ b/video_merge.py
@@ -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"
+}
\ No newline at end of file