diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index 1b4993aa7..4cde7763a 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/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.height = self.__components.images.shape[1] + video_stream.width = width + 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,