Allow images to work with auto frame count, and include images in video datasets if they exist.
This commit is contained in:
Jaret Burkett
parent
6d8afa5684
commit
8c1a4082fd
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue