From a40ccd22379e899377b77dcd1632e0fbcd4cbeda Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 18 Nov 2024 15:28:37 +0530 Subject: [PATCH] updates --- training/dataset.py | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/training/dataset.py b/training/dataset.py index 1d250dc..411fdfc 100644 --- a/training/dataset.py +++ b/training/dataset.py @@ -22,8 +22,8 @@ decord.bridge.set_bridge("torch") logger = get_logger(__name__) HEIGHT_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536] -WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 848, 960, 1024, 1280, 1536] -FRAME_BUCKETS = [16, 24, 32, 48, 64, 80, 84] +WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 960, 1024, 1280, 1536] +FRAME_BUCKETS = [16, 24, 32, 48, 64, 80] class VideoDataset(Dataset): @@ -144,17 +144,16 @@ class VideoDataset(Dataset): } else: image, video, _ = self._preprocess_video(self.video_paths[index]) - if video is not None: - return { - "prompt": self.id_token + self.prompts[index], - "image": image, - "video": video, - "video_metadata": { - "num_frames": video.shape[0], - "height": video.shape[2], - "width": video.shape[3], - }, - } + return { + "prompt": self.id_token + self.prompts[index], + "image": image, + "video": video, + "video_metadata": { + "num_frames": video.shape[0], + "height": video.shape[2], + "width": video.shape[3], + }, + } def _load_dataset_from_local_path(self) -> Tuple[List[str], List[str]]: if not self.data_root.exists(): @@ -276,13 +275,9 @@ class VideoDatasetWithResizing(VideoDataset): else: video_reader = decord.VideoReader(uri=path.as_posix()) video_num_frames = len(video_reader) - nearest_frame_bucket = min( self.frame_buckets, key=lambda x: abs(x - min(video_num_frames, self.max_num_frames)) ) - if video_num_frames < nearest_frame_bucket: - # TODO: we could handle this by padding zero frames or duplicating the existing frames? - return None, None, None frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket)) frames = video_reader.get_batch(frame_indices)