github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Allow images to work with auto frame count, and include images in video datasets if they exist.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-08 20:26:28 -06:00
parent 6d8afa5684
commit 8c1a4082fd
3 changed files with 27 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
toolkit/data_loader.py
View File

@@ -431,8 +431,9 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
# only look for audio files
extensions = audio_extensions
elif self.is_video:
# only look for videos
extensions = video_extensions
# look for videos and images. Video models can train on both;
# images are bucketed separately as single-frame items
extensions = video_extensions + image_extensions
file_list = [os.path.join(root, file) for root, _, files in os.walk(self.dataset_path) for file in files if file.lower().endswith(tuple(extensions)) and not file.startswith('.')]
else:
# assume json
@@ -555,7 +556,12 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
json.dump(self.size_database, f)
if self.is_video:
print_acc(f" - Found {len(self.file_list)} videos")
num_videos = len([x for x in self.file_list if x.is_video])
num_images = len(self.file_list) - num_videos
if num_images > 0:
print_acc(f" - Found {num_videos} videos and {num_images} images")
else:
print_acc(f" - Found {num_videos} videos")
assert len(self.file_list) > 0, f"no videos found in {self.dataset_path}"
else:
print_acc(f" - Found {len(self.file_list)} images")