storytold/FineTrainers-Conditioning
2
0
Fork 0
mirror of https://github.com/storytold/FineTrainers-Conditioning.git synced 2026-10-18 00:33:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

updates

Browse Source
This commit is contained in:
sayakpaul
2024-11-17 10:18:20 +05:30
parent a9adf22386
commit 5ba510e061
3 changed files with 88 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/dataset.py
+30 -24
View File
@@ -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, 960, 1024, 1280, 1536]
FRAME_BUCKETS = [16, 24, 32, 48, 64, 80]
WIDTH_BUCKETS = [256, 320, 384, 480, 512, 576, 720, 768, 848, 960, 1024, 1280, 1536]
FRAME_BUCKETS = [16, 24, 32, 48, 64, 80, 84]
class VideoDataset(Dataset):
@@ -144,17 +144,17 @@ class VideoDataset(Dataset):
}
else:
image, video, _ = self._preprocess_video(self.video_paths[index])
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],
},
}
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],
},
}
def _load_dataset_from_local_path(self) -> Tuple[List[str], List[str]]:
if not self.data_root.exists():
@@ -276,12 +276,15 @@ 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)
frames = frames[:nearest_frame_bucket].float()
frames = frames.permute(0, 3, 1, 2).contiguous()
@@ -344,6 +347,8 @@ class VideoDatasetWithResizeAndRectangleCrop(VideoDataset):
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:
return None, None, None
frame_indices = list(range(0, video_num_frames, video_num_frames // nearest_frame_bucket))
@@ -404,16 +409,17 @@ class BucketSampler(Sampler):
def __iter__(self):
for index, data in enumerate(self.data_source):
video_metadata = data["video_metadata"]
f, h, w = video_metadata["num_frames"], video_metadata["height"], video_metadata["width"]
if data is not None:
video_metadata = data["video_metadata"]
f, h, w = video_metadata["num_frames"], video_metadata["height"], video_metadata["width"]
self.buckets[(f, h, w)].append(data)
if len(self.buckets[(f, h, w)]) == self.batch_size:
if self.shuffle:
random.shuffle(self.buckets[(f, h, w)])
yield self.buckets[(f, h, w)]
del self.buckets[(f, h, w)]
self.buckets[(f, h, w)] = []
self.buckets[(f, h, w)].append(data)
if len(self.buckets[(f, h, w)]) == self.batch_size:
if self.shuffle:
random.shuffle(self.buckets[(f, h, w)])
yield self.buckets[(f, h, w)]
del self.buckets[(f, h, w)]
self.buckets[(f, h, w)] = []
if self.drop_last:
return