From 2074018edb92da3db24d880f5f5b032c7329665b Mon Sep 17 00:00:00 2001
From: Jacob Segal <jacob.e.segal@gmail.com>
Date: Sun, 27 Apr 2025 21:06:09 -0700
Subject: [PATCH] Add optimization to avoid parsing entire video

---
 comfy/comfy_types/input_types.py | 26 ++++++++++++++++++++++++++
 comfy_extras/nodes_video.py      |  4 +---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/comfy/comfy_types/input_types.py b/comfy/comfy_types/input_types.py
index b5c30299c..31c7e83a0 100644
--- a/comfy/comfy_types/input_types.py
+++ b/comfy/comfy_types/input_types.py
@@ -101,6 +101,18 @@ class VideoInput(ABC):
         """
         pass
 
+    # Provide a default implementation, but subclasses can provide optimized versions
+    # if possible.
+    def get_dimensions(self) -> tuple[int, int]:
+        """
+        Returns the dimensions of the video input.
+
+        Returns:
+            Tuple of (width, height)
+        """
+        components = self.get_components()
+        return components.images.shape[2], components.images.shape[1]
+
 class VideoFromFile(VideoInput):
     """
     Class representing video input from a file.
@@ -113,6 +125,20 @@ class VideoFromFile(VideoInput):
         """
         self.file = file
 
+    def get_dimensions(self) -> tuple[int, int]:
+        """
+        Returns the dimensions of the video input.
+
+        Returns:
+            Tuple of (width, height)
+        """
+        with av.open(self.file, mode='r') as container:
+            for stream in container.streams:
+                if stream.type == 'video':
+                    assert isinstance(stream, av.VideoStream)
+                    return stream.width, stream.height
+        raise ValueError(f"No video stream found in file '{self.file}'")
+
     def get_components_internal(self, container: InputContainer) -> VideoComponents:
         # Get video frames
         frames = []
diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py
index 028fadf36..8694bdadb 100644
--- a/comfy_extras/nodes_video.py
+++ b/comfy_extras/nodes_video.py
@@ -107,9 +107,7 @@ class SaveVideo(ComfyNodeABC):
 
     def save_video(self, video: VideoInput, filename_prefix, format, codec, prompt=None, extra_pnginfo=None):
         filename_prefix += self.prefix_append
-        components = video.get_components()
-        width = components.images.shape[2]
-        height = components.images.shape[1]
+        width, height = video.get_dimensions()
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
             filename_prefix,
             self.output_dir,