diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index d2b8a59b..20a1f28f 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -1063,6 +1063,11 @@ class DatasetConfig: # this wont work with bucketing for now until I can handle this before bucketing. self.auto_frame_count: bool = kwargs.get('auto_frame_count', False) + # old behavior shrank the video to fit the temporal spacing of the model. Which fits the whole video, but + # can lead to fast motion/chipmunking. This will prevent the video from shrinking to fit, and instead, trim + # the tail of the video. Usually only a few frames. + self.trim_auto_frame_count_tail: bool = kwargs.get('trim_auto_frame_count_tail', True) + # debug the frame count and frame selection. You dont need this. It is for debugging. self.debug: bool = kwargs.get('debug', False) diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 4ed4608b..0f104522 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -97,6 +97,7 @@ class FileItemDTO( raise Exception("Error: Could not get file signature for {self.path}") use_db_entry = False + db_entry = None if file_key in size_database: db_entry = size_database[file_key] if ( @@ -105,6 +106,8 @@ class FileItemDTO( and db_entry[2] == file_signature ): use_db_entry = True + video_total_frames = None + video_fps = None if self.is_audio_model: # get the length of the audio file in ms with av.open(self.path) as c: @@ -114,24 +117,31 @@ class FileItemDTO( s = c.streams.audio[0] w = int(float(s.duration * s.time_base) * 1_000) h = 1 - elif use_db_entry: - w, h, _ = size_database[file_key] elif self.is_video: - # Open the video file - video = cv2.VideoCapture(self.path) + # video entries also carry (total_frames, fps); older 3-item entries + # get re-read and upgraded here + if use_db_entry and len(db_entry) >= 5: + w, h, _, video_total_frames, video_fps = db_entry[:5] + else: + # Open the video file + video = cv2.VideoCapture(self.path) - # Check if video opened successfully - if not video.isOpened(): - raise Exception(f"Error: Could not open video file {self.path}") + # Check if video opened successfully + if not video.isOpened(): + raise Exception(f"Error: Could not open video file {self.path}") - # Get width and height - width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) - height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) - w, h = width, height + # Get width and height + width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) + w, h = width, height + video_total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + video_fps = video.get(cv2.CAP_PROP_FPS) - # Release the video capture object immediately - video.release() - size_database[file_key] = (width, height, file_signature) + # Release the video capture object immediately + video.release() + size_database[file_key] = (width, height, file_signature, video_total_frames, video_fps) + elif use_db_entry: + w, h, _ = db_entry[:3] else: if self.dataset_config.fast_image_size: # original method is significantly faster, but some images are read sideways. Not sure why. Do slow method by default. @@ -150,6 +160,10 @@ class FileItemDTO( size_database[file_key] = (w, h, file_signature) self.width: int = w self.height: int = h + if self.is_video and self.dataset_config.auto_frame_count: + # compute the real frame count now (same math as load time) so buckets + # are keyed on the frame count this video will actually train at + self.num_frames = self.get_auto_frame_count(video_total_frames, video_fps) self.dataloader_transforms = kwargs.get("dataloader_transforms", None) super().__init__(*args, **kwargs) diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 08fe466a..398be98a 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -470,6 +470,33 @@ class AudioProcessingDTOMixin: class ImageProcessingDTOMixin: + def get_auto_frame_count(self: 'FileItemDTO', total_frames: int, video_fps: float) -> int: + # frame count this video will train at with auto_frame_count. Also called at + # FileItemDTO init so bucket keys carry the real frame count, so it must give + # the same answer at bucketing time and at load time. + # allow for any length video here but make sure it is temporally compressable. + vid_length_seconds = total_frames / video_fps + + desired_num_frames = int(vid_length_seconds * self.dataset_config.fps) + + if getattr(self, 'frame_count_snapper', None) is not None: + # model-specific valid-frame-count grid (e.g. minimax_h3's 17n+5) + desired_num_frames = self.frame_count_snapper(desired_num_frames) + else: + # make sure it is divisible by temporal_compression + if self.dataset_config.trim_auto_frame_count_tail: + # snap to the largest valid count that fits inside the video (after the + # key frame +1 below) so trim mode never overshoots the source, which + # would freeze the last frame and pad the audio tail with silence + desired_num_frames = max(0, desired_num_frames - 1) // self.temporal_compression * self.temporal_compression + else: + desired_num_frames = desired_num_frames // self.temporal_compression * self.temporal_compression + + # TODO, all models currently add a key frame, but future models may not, update here if this changes. + desired_num_frames += 1 # add one for the key frame that is always added + + return desired_num_frames + def load_and_process_video( self: 'FileItemDTO', transform: Union[None, transforms.Compose], @@ -511,26 +538,18 @@ class ImageProcessingDTOMixin: frames_to_extract = [] if self.dataset_config.auto_frame_count: - # allow for any length video here but make sure it is temporally compressable. - vid_length_seconds = total_frames / video_fps + self.num_frames = self.get_auto_frame_count(total_frames, video_fps) - desired_num_frames = int(vid_length_seconds * self.dataset_config.fps) - if getattr(self, 'frame_count_snapper', None) is not None: - # model-specific valid-frame-count grid (e.g. minimax_h3's 17n+5) - desired_num_frames = self.frame_count_snapper(desired_num_frames) - else: - # make sure it is divisible by temporal_compression - desired_num_frames = desired_num_frames // self.temporal_compression * self.temporal_compression - - # TODO, all models currently add a key frame, but future models may not, update here if this changes. - desired_num_frames += 1 # add one for the key frame that is always added - - self.num_frames = desired_num_frames - - # Always stretch/shrink to the requested number of frames if needed - if self.dataset_config.shrink_video_to_frames or total_frames < self.num_frames: + if self.dataset_config.auto_frame_count and self.dataset_config.trim_auto_frame_count_tail: + # preserve real time: pull frames at the dataset fps from the start of the + # video and trim the tail that didn't fit the snapped frame count, instead of + # shrinking the whole video to fit (which speeds up motion / chipmunks audio). + # Critical for audio models (e.g. minimax_h3) where audio must stay in sync. + fps_ratio = video_fps / self.dataset_config.fps if video_fps and video_fps > 0 else 1.0 + frames_to_extract = [min(round(i * fps_ratio), max_frame_index) for i in range(self.num_frames)] + elif self.dataset_config.shrink_video_to_frames or total_frames < self.num_frames: # Distribute frames evenly across the entire video interval = max_frame_index / (self.num_frames - 1) if self.num_frames > 1 else 0 frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.num_frames)] @@ -733,10 +752,20 @@ class ImageProcessingDTOMixin: gain = target_peak / (peak + eps) waveform = waveform * gain + trim_tail_audio = ( + self.dataset_config.auto_frame_count + and self.dataset_config.trim_auto_frame_count_tail + ) + # Slice to the selected clip region (when we have a meaningful time range) if source_duration > 0.0: start_sample = int(round(clip_start_time * sample_rate)) - end_sample = int(round(clip_end_time * sample_rate)) + if trim_tail_audio and target_duration > 0.0: + # time must stay 1:1 with the video — cut exactly the + # training duration so no stretch is needed below + end_sample = start_sample + round(target_duration * sample_rate) + else: + end_sample = round(clip_end_time * sample_rate) start_sample = max(0, min(start_sample, waveform.shape[-1])) end_sample = max(0, min(end_sample, waveform.shape[-1])) if end_sample > start_sample: @@ -749,10 +778,19 @@ class ImageProcessingDTOMixin: waveform = None if waveform is not None and waveform.numel() > 0: - target_samples = int(round(target_duration * sample_rate)) + target_samples = round(target_duration * sample_rate) if target_samples > 0 and waveform.shape[-1] != target_samples: # Time-stretch/shrink to match the video clip duration implied by dataset FPS. - if self.dataset_config.audio_preserve_pitch: + if trim_tail_audio: + # never stretch/contract in trim mode. The waveform can only be + # short here (audio/video ended a hair before the target) — + # pad the tail with silence, or cut any rounding overshoot + pad = target_samples - waveform.shape[-1] + if pad > 0: + waveform = F.pad(waveform, (0, pad)) + else: + waveform = waveform[..., :target_samples] + elif self.dataset_config.audio_preserve_pitch: waveform = time_stretch_preserve_pitch(waveform, sample_rate, target_samples) # waveform is [C, L] else: # Use linear interpolation over the time axis. @@ -1732,6 +1770,10 @@ class LatentCachingFileItemDTOMixin: if self.is_video and self.dataset_config.auto_frame_count: # don't store num frames here as it is calculated dynamically item["auto_frame_count"] = True + if self.dataset_config.trim_auto_frame_count_tail: + # changes frame selection; only added when on so caches made before + # this option existed stay valid when it is off + item["trim_auto_frame_count_tail"] = True is_video = True elif self.is_video and self.dataset_config.num_frames > 1: item["num_frames"] = self.dataset_config.num_frames