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

updates

Browse Source
This commit is contained in:
sayakpaul
2024-11-18 15:32:04 +05:30
parent a8dd94ee1c
commit ac83c78a5a
1 changed files with 11 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/dataset.py
+11 -12
View File
@@ -144,7 +144,7 @@ class VideoDataset(Dataset):
}
else:
image, video, _ = self._preprocess_video(self.video_paths[index])
return {
"prompt": self.id_token + self.prompts[index],
"image": image,
@@ -281,7 +281,7 @@ class VideoDatasetWithResizing(VideoDataset):
)
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()
@@ -404,17 +404,16 @@ class BucketSampler(Sampler):
def __iter__(self):
for index, data in enumerate(self.data_source):
if data is not None:
video_metadata = data["video_metadata"]
f, h, w = video_metadata["num_frames"], video_metadata["height"], video_metadata["width"]
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