diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index b0737601..11dc879b 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -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") diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index 9f8e2eed..4ed4608b 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -30,6 +30,10 @@ if TYPE_CHECKING: printed_messages = [] +# keep in sync with video_extensions in toolkit/data_loader.py (importing it +# here would be circular) +video_extensions = ['.mp4', '.avi', '.mov', '.webm', '.mkv', '.wmv', '.m4v', '.flv'] + def print_once(msg): global printed_messages @@ -55,10 +59,13 @@ class FileItemDTO( def __init__(self, *args, **kwargs): self.path = kwargs.get("path", "") self.dataset_config: "DatasetConfig" = kwargs.get("dataset_config", None) - self.is_video = self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count + # a video dataset can contain both videos and images. Images are + # treated as single-frame items and bucketed separately from videos + dataset_is_video = self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count + self.is_video = dataset_is_video and os.path.splitext(self.path)[1].lower() in video_extensions self.is_audio_model = kwargs.get("is_audio_model", False) self.sample_rate = kwargs.get("sample_rate", 48000) - self.num_frames = self.dataset_config.num_frames + self.num_frames = self.dataset_config.num_frames if self.is_video else 1 self.temporal_compression = kwargs.get("temporal_compression", 8) # module-level function (picklable) for models whose valid frame # counts are not temporal_compression * n + 1; None = default math diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 53ebdf05..08fe466a 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -296,6 +296,9 @@ class BucketsMixin: # check if bucket exists, if not, create it bucket_key = f'{file_item.crop_width}x{file_item.crop_height}' + if self.is_video: + # images (1 frame) and videos must not mix in a batch + bucket_key += f'x{file_item.num_frames}f' if bucket_key not in self.buckets: self.buckets[bucket_key] = Bucket(file_item.crop_width, file_item.crop_height) self.buckets[bucket_key].file_list_idx.append(idx) @@ -836,7 +839,7 @@ class ImageProcessingDTOMixin: if self.is_audio_model: self.load_and_process_audio() return - if self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count: + if self.is_video: self.load_and_process_video(transform, only_load_latents) return try: @@ -1726,11 +1729,11 @@ class LatentCachingFileItemDTOMixin: item["flip_x"] = True if self.flip_y: item["flip_y"] = True - if self.dataset_config.auto_frame_count: + 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 is_video = True - elif self.dataset_config.num_frames > 1: + elif self.is_video and self.dataset_config.num_frames > 1: item["num_frames"] = self.dataset_config.num_frames is_video = True if is_video and self.dataset_config.fps != 24: @@ -1975,7 +1978,7 @@ class LatentCachingMixin: print_acc(f"Error: {str(e)}") raise e # do first frame - is_video = self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1 + is_video = file_item.is_video if is_video and self.dataset_config.do_i2v: frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype) if len(frames.shape) == 4: @@ -2053,7 +2056,7 @@ class TextEmbeddingFileItemDTOMixin: elif ( getattr(self, "encode_first_frame_in_text_embeddings", False) and self.dataset_config.do_i2v - and (self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1) + and self.is_video ): item["first_frame_in_te"] = True return item @@ -2144,7 +2147,7 @@ class TextEmbeddingCachingMixin: elif ( getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) and self.dataset_config.do_i2v - and (self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1) + and file_item.is_video ): # video item: encode the clip's FIRST FRAME into the text embeddings # as a vision reference, matching sampling (where the ctrl image goes