video_types: split up saver to streamable state

Split this up into start -> chunk -> finish so it can be saved piece by piece.

comfy_api/latest/_input_impl/video_types.py

@@ -386,6 +386,7 @@ class VideoFromComponents(VideoInput):
 
     def __init__(self, components: VideoComponents):
         self.__components = components
+        self._frame_counter = 0
 
     def get_components(self) -> VideoComponents:
         return VideoComponents(
@@ -394,14 +395,13 @@ class VideoFromComponents(VideoInput):
             frame_rate=self.__components.frame_rate,
         )
 
-    def save_to(
+    def save_start(
         self,
         path: str,
         format: VideoContainer = VideoContainer.AUTO,
         codec: VideoCodec = VideoCodec.AUTO,
         metadata: Optional[dict] = None,
     ):
-        """Save the video to a file path or BytesIO buffer."""
         if format != VideoContainer.AUTO and format != VideoContainer.MP4:
             raise ValueError("Only MP4 format is supported for now")
         if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
@@ -413,7 +413,12 @@ class VideoFromComponents(VideoInput):
             # BytesIO has no file extension, so av.open can't infer the format.
             # Default to mp4 since that's the only supported format anyway.
             extra_kwargs["format"] = "mp4"
-        with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
+
+        width, height = self.get_dimensions()
+
+        output = av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs)
+        if True:
+
             # Add metadata before writing any streams
             if metadata is not None:
                 for key, value in metadata.items():
@@ -422,8 +427,8 @@ class VideoFromComponents(VideoInput):
             frame_rate = Fraction(round(self.__components.frame_rate * 1000), 1000)
             # Create a video stream
             video_stream = output.add_stream('h264', rate=frame_rate)
-            video_stream.width = self.__components.images.shape[2]
+            video_stream.width = width
-            video_stream.height = self.__components.images.shape[1]
+            video_stream.height = height
             video_stream.pix_fmt = 'yuv420p'
 
             # Create an audio stream
@@ -432,23 +437,33 @@ class VideoFromComponents(VideoInput):
             if self.__components.audio:
                 audio_sample_rate = int(self.__components.audio['sample_rate'])
                 waveform = self.__components.audio['waveform']
-                waveform = waveform[0, :, :math.ceil((audio_sample_rate / frame_rate) * self.__components.images.shape[0])]
+                waveform = waveform[0]
                 layout = {1: 'mono', 2: 'stereo', 6: '5.1'}.get(waveform.shape[0], 'stereo')
                 audio_stream = output.add_stream('aac', rate=audio_sample_rate, layout=layout)
 
+        self._frame_counter = 0
+        return output, video_stream, audio_stream, audio_sample_rate, frame_rate
+
+    def save_add(self, output, video_stream, images) -> None:
             # Encode video
-            for i, frame in enumerate(self.__components.images):
+            for frame in images:
                 img = (frame * 255).clamp(0, 255).byte().cpu().numpy() # shape: (H, W, 3)
                 frame = av.VideoFrame.from_ndarray(img, format='rgb24')
                 frame = frame.reformat(format='yuv420p')  # Convert to YUV420P as required by h264
                 packet = video_stream.encode(frame)
                 output.mux(packet)
+                self._frame_counter += 1
 
+    def save_finalize(self, output, video_stream, audio_stream, audio_sample_rate, frame_rate) -> None:
             # Flush video
             packet = video_stream.encode(None)
             output.mux(packet)
 
             if audio_stream and self.__components.audio:
+                waveform = self.__components.audio['waveform']
+                waveform = waveform[0]
+                layout = {1: 'mono', 2: 'stereo', 6: '5.1'}.get(waveform.shape[0], 'stereo')
+                waveform = waveform[:, :math.ceil((audio_sample_rate / frame_rate) * self._frame_counter)]
                 frame = av.AudioFrame.from_ndarray(waveform.float().cpu().contiguous().numpy(), format='fltp', layout=layout)
                 frame.sample_rate = audio_sample_rate
                 frame.pts = 0
@@ -457,6 +472,29 @@ class VideoFromComponents(VideoInput):
                 # Flush encoder
                 output.mux(audio_stream.encode(None))
 
+            output.close()
+
+    def save_to(
+        self,
+        path: str,
+        format: VideoContainer = VideoContainer.AUTO,
+        codec: VideoCodec = VideoCodec.AUTO,
+        metadata: Optional[dict] = None,
+    ):
+        """Save the video to a file path or BytesIO buffer."""
+        output, video_stream, audio_stream, audio_sample_rate, frame_rate = self.save_start(
+            path,
+            format=format,
+            codec=codec,
+            metadata=metadata,
+        )
+        try:
+            self.save_add(output, video_stream, self.__components.images)
+            self.save_finalize(output, video_stream, audio_stream, audio_sample_rate, frame_rate)
+        except Exception:
+            output.close()
+            raise
+
     def as_trimmed(
         self,
         start_time: float | None = None,