1212import math
1313import os
1414import torch
15- from .._util import VideoContainer , VideoCodec , VideoComponents
15+ from .._util import VideoContainer , VideoCodec , VideoComponents , normalize_crop_rect
16+ import comfy .utils
1617import logging
1718
1819
@@ -196,19 +197,25 @@ def webm_streams_compatible(streams) -> bool:
196197 return True
197198
198199
200+ def _rotation_quadrant (frame : av .VideoFrame ) -> int :
201+ return int (round (frame .rotation // 90 )) % 4 if frame .rotation else 0
202+
203+
199204class VideoFromFile (VideoInput ):
200205 """
201206 Class representing video input from a file.
202207 """
203208
204- def __init__ (self , file : str | io .BytesIO , * , start_time : float = 0 , duration : float = 0 ):
209+ def __init__ (self , file : str | io .BytesIO , * , start_time : float = 0 , duration : float = 0 ,
210+ crop : tuple [int , int , int , int ] | None = None ):
205211 """
206212 Initialize the VideoFromFile object based off of either a path on disk or a BytesIO object
207213 containing the file contents.
208214 """
209215 self .__file = file
210216 self .__start_time = start_time
211217 self .__duration = duration
218+ self .__crop = crop
212219
213220 def get_stream_source (self ) -> str | io .BytesIO :
214221 """
@@ -238,7 +245,31 @@ def get_dimensions(self) -> tuple[int, int]:
238245 for stream in container .streams :
239246 if stream .type == 'video' :
240247 assert isinstance (stream , av .VideoStream )
241- return stream .width , stream .height
248+ if self .__crop is None :
249+ return stream .width , stream .height
250+
251+ display_width , display_height = self ._get_display_dimensions ()
252+ rect = normalize_crop_rect (* self .__crop , display_width , display_height )
253+ if rect is not None :
254+ return rect [2 ], rect [3 ]
255+ return display_width , display_height
256+ raise ValueError (f"No video stream found in file '{ self .__file } '" )
257+
258+ def _get_display_dimensions (self ) -> tuple [int , int ]:
259+ if isinstance (self .__file , io .BytesIO ):
260+ self .__file .seek (0 )
261+ with av .open (self .__file , mode = 'r' ) as container :
262+ for stream in container .streams :
263+ if stream .type == 'video' :
264+ assert isinstance (stream , av .VideoStream )
265+ width , height = stream .width , stream .height
266+ try :
267+ frame = next (container .decode (stream ), None )
268+ except av .error .FFmpegError :
269+ frame = None
270+ if frame is not None and _rotation_quadrant (frame ) % 2 :
271+ width , height = height , width
272+ return width , height
242273 raise ValueError (f"No video stream found in file '{ self .__file } '" )
243274
244275 def get_bit_depth (self ) -> int :
@@ -415,6 +446,8 @@ def get_components_internal(self, container: InputContainer) -> VideoComponents:
415446 streams = [video_stream ]
416447 has_first_audio_frame = False
417448 checked_alpha = False
449+ crop_rect = None
450+ crop_resolved = False
418451
419452 # Default to False so we decode until EOF if duration is 0
420453 video_done = False
@@ -485,9 +518,16 @@ def get_components_internal(self, container: InputContainer) -> VideoComponents:
485518 img = np .ascontiguousarray (align_graph [2 ].pull ().to_ndarray (format = image_format )[:frame .height , :frame .width ])
486519 else :
487520 img = frame .to_ndarray (format = image_format )
488- if frame .rotation != 0 :
489- k = int (round (frame .rotation // 90 ))
490- img = np .rot90 (img , k = k , axes = (0 , 1 )).copy ()
521+ rotation_quadrant = _rotation_quadrant (frame )
522+ if rotation_quadrant :
523+ img = np .rot90 (img , k = rotation_quadrant , axes = (0 , 1 )).copy ()
524+ if self .__crop is not None :
525+ if not crop_resolved :
526+ crop_rect = normalize_crop_rect (* self .__crop , img .shape [1 ], img .shape [0 ])
527+ crop_resolved = True
528+ if crop_rect is not None :
529+ cx , cy , cw , ch = crop_rect
530+ img = np .ascontiguousarray (img [cy :cy + ch , cx :cx + cw ])
491531 if alphas is None :
492532 frames .append (torch .from_numpy (img ))
493533 else :
@@ -586,6 +626,8 @@ def save_to(
586626 reuse_streams = False
587627 if self .__start_time or self .__duration :
588628 reuse_streams = False
629+ if self .__crop is not None :
630+ reuse_streams = False
589631
590632 if not reuse_streams :
591633 if bit_depth is None :
@@ -673,6 +715,16 @@ def _save_transcoded(
673715 if duration :
674716 duration_cap = math .ceil (duration * sample_rate )
675717
718+ if duration :
719+ window_seconds = duration
720+ else :
721+ try :
722+ window_seconds = max (self ._get_raw_duration () - start_time , 0.0 )
723+ except ValueError :
724+ window_seconds = 0.0
725+ progress_total = max (1 , int (round (window_seconds * float (rate ))))
726+ pbar = comfy .utils .ProgressBar (progress_total )
727+
676728 streams = [video_stream ] if audio_stream is None else [video_stream , audio_stream ]
677729 pts_step = max (1 , int (round ((1 / rate ) / video_stream .time_base )))
678730 video_done = False
@@ -685,6 +737,8 @@ def _save_transcoded(
685737 source_size = None
686738 rotation_k = 0
687739 rotation_filter = None
740+ crop_rect = None
741+ crop_filter = None
688742 audio_started = False
689743 samples_written = 0
690744 pending_audio = []
@@ -758,13 +812,27 @@ def drain_audio(final=False):
758812 if end_pts is not None and frame .pts is not None :
759813 frame_duration = min (frame_duration , end_pts - frame .pts )
760814 if output is None :
761- rotation_k = int ( round ( frame . rotation // 90 )) % 4 if frame . rotation else 0
815+ rotation_k = _rotation_quadrant ( frame )
762816 if rotation_k % 2 :
763817 out_width , out_height = frame .height , frame .width
764818 else :
765819 out_width , out_height = frame .width , frame .height
820+ if self .__crop is not None :
821+ crop_rect = normalize_crop_rect (* self .__crop , out_width , out_height )
822+ if crop_rect is not None :
823+ out_width , out_height = crop_rect [2 ], crop_rect [3 ]
824+ if (out_width % 2 or out_height % 2 ) and crop_rect is None :
825+ even_width = out_width - out_width % 2
826+ even_height = out_height - out_height % 2
827+ if even_width > 0 and even_height > 0 :
828+ crop_rect = (0 , 0 , even_width , even_height )
829+ out_width , out_height = even_width , even_height
766830 if out_width % 2 or out_height % 2 :
767831 raise ValueError (f"{ output_codec .value .upper ()} output requires even dimensions, got { out_width } x{ out_height } " )
832+ if any (component .is_alpha for component in frame .format .components ):
833+ logging .warning (
834+ "Transcoded video output does not support alpha; the alpha channel will be discarded."
835+ )
768836 source_size = (frame .width , frame .height )
769837 output = av .open (path , ** open_kwargs )
770838 # Add metadata before writing any streams
@@ -810,6 +878,19 @@ def drain_audio(final=False):
810878 rotation_filter = (g_src , g_sink )
811879 rotation_filter [0 ].push (frame )
812880 frame = rotation_filter [1 ].pull ()
881+ if crop_rect is not None :
882+ if crop_filter is None :
883+ g = av .filter .Graph ()
884+ g_src = g .add_buffer (width = frame .width , height = frame .height ,
885+ format = frame .format .name , time_base = video_stream .time_base )
886+ g_crop = g .add ("crop" , f"{ crop_rect [2 ]} :{ crop_rect [3 ]} :{ crop_rect [0 ]} :{ crop_rect [1 ]} " )
887+ g_sink = g .add ("buffersink" )
888+ g_src .link_to (g_crop )
889+ g_crop .link_to (g_sink )
890+ g .configure ()
891+ crop_filter = (g_src , g_sink )
892+ crop_filter [0 ].push (frame )
893+ frame = crop_filter [1 ].pull ()
813894 if frame .color_range == ColorRange .JPEG and not preserve_source_color :
814895 # compress full-range sources (yuvj/MJPEG) to limited range
815896 frame = frame .reformat (format = pix_fmt , src_color_range = "JPEG" , dst_color_range = "MPEG" )
@@ -852,6 +933,7 @@ def drain_audio(final=False):
852933 out_packet .duration = video_frame_durations .pop (out_packet .pts , 0 )
853934 output .mux (out_packet )
854935 drain_audio ()
936+ pbar .update (1 )
855937
856938 elif packet .stream == audio_stream and not audio_done :
857939 for resampled in itertools .chain .from_iterable (map (resampler .resample , packet .decode ())):
@@ -921,11 +1003,42 @@ def as_trimmed(
9211003 self .get_stream_source (),
9221004 start_time = start_time + self .__start_time ,
9231005 duration = duration ,
1006+ crop = self .__crop ,
9241007 )
925- if trimmed .get_duration () < duration and strict_duration :
1008+ if strict_duration and duration and trimmed .get_duration () < duration :
9261009 return None
9271010 return trimmed
9281011
1012+ def as_cropped (
1013+ self , x : int = 0 , y : int = 0 , width : int = 0 , height : int = 0
1014+ ) -> VideoInput :
1015+ if int (width ) <= 0 or int (height ) <= 0 :
1016+ return self
1017+
1018+ display_width , display_height = self ._get_display_dimensions ()
1019+ outer = (
1020+ normalize_crop_rect (* self .__crop , display_width , display_height )
1021+ if self .__crop is not None
1022+ else None
1023+ )
1024+ if outer is None :
1025+ rect = normalize_crop_rect (x , y , width , height , display_width , display_height )
1026+ else :
1027+ inner = normalize_crop_rect (x , y , width , height , outer [2 ], outer [3 ])
1028+ rect = (
1029+ (outer [0 ] + inner [0 ], outer [1 ] + inner [1 ], inner [2 ], inner [3 ])
1030+ if inner is not None
1031+ else None
1032+ )
1033+ if rect is None :
1034+ return self
1035+ return VideoFromFile (
1036+ self .get_stream_source (),
1037+ start_time = self .__start_time ,
1038+ duration = self .__duration ,
1039+ crop = rect ,
1040+ )
1041+
9291042
9301043class VideoFromComponents (VideoInput ):
9311044 """
@@ -942,6 +1055,8 @@ def get_components(self) -> VideoComponents:
9421055 images = self .__components .images ,
9431056 audio = self .__components .audio ,
9441057 frame_rate = self .__components .frame_rate ,
1058+ metadata = self .__components .metadata ,
1059+ alpha = self .__components .alpha ,
9451060 )
9461061
9471062 def get_bit_depth (self ) -> int :
0 commit comments