Add support for audo frame count so datasets can have varrying length videos. Varous ltx 2.3 VAE optimizations such as removing tiling articacts, and doing frame split encoding to reduce vram on encoding/decoding.
This commit is contained in:
Jaret Burkett committed
1 parent
4ad14d211a
commit
7f3309b291
10 files changed
+132
-36
No files matched your search
@@ -507,11 +507,26 @@ class ImageProcessingDTOMixin:
|
||||
|
||||
frames_to_extract = []
|
||||
|
||||
if self.dataset_config.auto_frame_count:
|
||||
# allow for any length video here but make sure it is temporally compressable.
|
||||
vid_length_seconds = total_frames / video_fps
|
||||
|
||||
desired_num_frames = int(vid_length_seconds * self.dataset_config.fps)
|
||||
|
||||
# make sure it is divisible by temporal_compression
|
||||
desired_num_frames = desired_num_frames // self.temporal_compression * self.temporal_compression
|
||||
|
||||
# TODO, all models currently add a key frame, but future models may not, update here if this changes.
|
||||
desired_num_frames += 1 # add one for the key frame that is always added
|
||||
|
||||
self.num_frames = desired_num_frames
|
||||
|
||||
|
||||
# Always stretch/shrink to the requested number of frames if needed
|
||||
if self.dataset_config.shrink_video_to_frames or total_frames < self.dataset_config.num_frames:
|
||||
if self.dataset_config.shrink_video_to_frames or total_frames < self.num_frames:
|
||||
# Distribute frames evenly across the entire video
|
||||
interval = max_frame_index / (self.dataset_config.num_frames - 1) if self.dataset_config.num_frames > 1 else 0
|
||||
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.dataset_config.num_frames)]
|
||||
interval = max_frame_index / (self.num_frames - 1) if self.num_frames > 1 else 0
|
||||
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.num_frames)]
|
||||
else:
|
||||
# Calculate frame interval based on FPS ratio
|
||||
fps_ratio = video_fps / self.dataset_config.fps
|
||||
@@ -520,17 +535,17 @@ class ImageProcessingDTOMixin:
|
||||
# Calculate max consecutive frames we can extract at desired FPS
|
||||
max_consecutive_frames = (total_frames // frame_interval)
|
||||
|
||||
if max_consecutive_frames < self.dataset_config.num_frames:
|
||||
if max_consecutive_frames < self.num_frames:
|
||||
# Not enough frames at desired FPS, so stretch instead
|
||||
interval = max_frame_index / (self.dataset_config.num_frames - 1) if self.dataset_config.num_frames > 1 else 0
|
||||
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.dataset_config.num_frames)]
|
||||
interval = max_frame_index / (self.num_frames - 1) if self.num_frames > 1 else 0
|
||||
frames_to_extract = [min(int(round(i * interval)), max_frame_index) for i in range(self.num_frames)]
|
||||
else:
|
||||
# Calculate max start frame to ensure we can get all num_frames
|
||||
max_start_frame = max_frame_index - ((self.dataset_config.num_frames - 1) * frame_interval)
|
||||
max_start_frame = max_frame_index - ((self.num_frames - 1) * frame_interval)
|
||||
start_frame = random.randint(0, max(0, max_start_frame))
|
||||
|
||||
# Generate list of frames to extract
|
||||
frames_to_extract = [start_frame + (i * frame_interval) for i in range(self.dataset_config.num_frames)]
|
||||
frames_to_extract = [start_frame + (i * frame_interval) for i in range(self.num_frames)]
|
||||
|
||||
# Final safety check - ensure no frame exceeds max valid index
|
||||
frames_to_extract = [min(frame_idx, max_frame_index) for frame_idx in frames_to_extract]
|
||||
@@ -642,7 +657,7 @@ class ImageProcessingDTOMixin:
|
||||
# Target duration is how this sampled/stretched clip is interpreted for training
|
||||
# (i.e. num_frames at the configured dataset FPS).
|
||||
if hasattr(self.dataset_config, "fps") and self.dataset_config.fps and self.dataset_config.fps > 0:
|
||||
target_duration = float(self.dataset_config.num_frames) / float(self.dataset_config.fps)
|
||||
target_duration = float(self.num_frames) / float(self.dataset_config.fps)
|
||||
else:
|
||||
target_duration = source_duration
|
||||
|
||||
@@ -757,7 +772,7 @@ class ImageProcessingDTOMixin:
|
||||
if self.has_unconditional:
|
||||
self.load_unconditional_image()
|
||||
return
|
||||
if self.dataset_config.num_frames > 1:
|
||||
if self.dataset_config.num_frames > 1 or self.dataset_config.auto_frame_count:
|
||||
self.load_and_process_video(transform, only_load_latents)
|
||||
return
|
||||
try:
|
||||
@@ -1751,16 +1766,25 @@ class LatentCachingFileItemDTOMixin:
|
||||
("latent_space_version", self.latent_space_version),
|
||||
("latent_version", self.latent_version),
|
||||
])
|
||||
is_video = False
|
||||
# when adding items, do it after so we dont change old latents
|
||||
if self.flip_x:
|
||||
item["flip_x"] = True
|
||||
if self.flip_y:
|
||||
item["flip_y"] = True
|
||||
if self.dataset_config.num_frames > 1:
|
||||
if 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:
|
||||
item["num_frames"] = self.dataset_config.num_frames
|
||||
if self.dataset_config.do_i2v:
|
||||
is_video = True
|
||||
if is_video and self.dataset_config.fps != 24:
|
||||
# only add fps if it deviates from the default
|
||||
item["fps"] = self.dataset_config.fps
|
||||
if is_video and self.dataset_config.do_i2v:
|
||||
item["do_i2v"] = True
|
||||
if self.dataset_config.do_audio:
|
||||
if is_video and self.dataset_config.do_audio:
|
||||
item["do_audio"] = True
|
||||
if self.dataset_config.audio_normalize:
|
||||
item["audio_normalize"] = True
|
||||
@@ -1815,6 +1839,8 @@ class LatentCachingFileItemDTOMixin:
|
||||
self._cached_first_frame_latent = state_dict['first_frame_latent']
|
||||
if 'audio_latent' in state_dict:
|
||||
self._cached_audio_latent = state_dict['audio_latent']
|
||||
if 'num_frames' in state_dict:
|
||||
self.num_frames = int(state_dict['num_frames'].item())
|
||||
return self._encoded_latent
|
||||
|
||||
|
||||
@@ -1866,6 +1892,7 @@ class LatentCachingMixin:
|
||||
state_dict = OrderedDict()
|
||||
first_frame_latent = None
|
||||
audio_latent = None
|
||||
frames = None
|
||||
# add batch dimension
|
||||
try:
|
||||
imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
|
||||
@@ -1877,7 +1904,8 @@ class LatentCachingMixin:
|
||||
print_acc(f"Error: {str(e)}")
|
||||
raise e
|
||||
# do first frame
|
||||
if self.dataset_config.num_frames > 1 and self.dataset_config.do_i2v:
|
||||
is_video = self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1
|
||||
if is_video and self.dataset_config.do_i2v:
|
||||
frames = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
|
||||
if len(frames.shape) == 4:
|
||||
first_frames = frames
|
||||
@@ -1895,6 +1923,9 @@ class LatentCachingMixin:
|
||||
if to_disk:
|
||||
state_dict['audio_latent'] = audio_latent.clone().detach().cpu()
|
||||
|
||||
if is_video:
|
||||
state_dict['num_frames'] = torch.tensor(file_item.num_frames, dtype=torch.int32)
|
||||
|
||||
# save_latent
|
||||
if to_disk:
|
||||
# metadata
|
||||
@@ -1912,7 +1943,11 @@ class LatentCachingMixin:
|
||||
|
||||
del imgs
|
||||
del latent
|
||||
del frames
|
||||
del file_item.tensor
|
||||
del state_dict
|
||||
del first_frame_latent
|
||||
del audio_latent
|
||||
file_item.cleanup()
|
||||
|
||||
file_item.is_latent_cached = True
|
||||
|
||||
Reference in new issue
Block a user