RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

When doing auto frame count. Ensure the time is not squeezed or expanded to fit tempooral spacing. trime the few extra frames. Also fixed frame counts of buckets.

This commit is contained in:
Jaret Burkett 2026-08-09 12:31:02 -06:00
parent 682b27c6ee
commit 72623ed3d6
3 changed files with 95 additions and 34 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/config_modules.py
View File

@ -1063,6 +1063,11 @@ class DatasetConfig:
# this wont work with bucketing for now until I can handle this before bucketing.
self.auto_frame_count: bool = kwargs.get('auto_frame_count', False)
# old behavior shrank the video to fit the temporal spacing of the model. Which fits the whole video, but
# can lead to fast motion/chipmunking. This will prevent the video from shrinking to fit, and instead, trim
# the tail of the video. Usually only a few frames.
self.trim_auto_frame_count_tail: bool = kwargs.get('trim_auto_frame_count_tail', True)
# debug the frame count and frame selection. You dont need this. It is for debugging.
self.debug: bool = kwargs.get('debug', False)

42
toolkit/data_transfer_object/data_loader.py
View File

@ -97,6 +97,7 @@ class FileItemDTO(
raise Exception("Error: Could not get file signature for {self.path}")
use_db_entry = False
db_entry = None
if file_key in size_database:
db_entry = size_database[file_key]
if (
@ -105,6 +106,8 @@ class FileItemDTO(
and db_entry[2] == file_signature
):
use_db_entry = True
video_total_frames = None
video_fps = None
if self.is_audio_model:
# get the length of the audio file in ms
with av.open(self.path) as c:
@ -114,24 +117,31 @@ class FileItemDTO(
s = c.streams.audio[0]
w = int(float(s.duration * s.time_base) * 1_000)
h = 1
elif use_db_entry:
w, h, _ = size_database[file_key]
elif self.is_video:
# Open the video file
video = cv2.VideoCapture(self.path)
# video entries also carry (total_frames, fps); older 3-item entries
# get re-read and upgraded here
if use_db_entry and len(db_entry) >= 5:
w, h, _, video_total_frames, video_fps = db_entry[:5]
else:
# Open the video file
video = cv2.VideoCapture(self.path)
# Check if video opened successfully
if not video.isOpened():
raise Exception(f"Error: Could not open video file {self.path}")
# Check if video opened successfully
if not video.isOpened():
raise Exception(f"Error: Could not open video file {self.path}")
# Get width and height
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
w, h = width, height
# Get width and height
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
w, h = width, height
video_total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
video_fps = video.get(cv2.CAP_PROP_FPS)
# Release the video capture object immediately
video.release()
size_database[file_key] = (width, height, file_signature)
# Release the video capture object immediately
video.release()
size_database[file_key] = (width, height, file_signature, video_total_frames, video_fps)
elif use_db_entry:
w, h, _ = db_entry[:3]
else:
if self.dataset_config.fast_image_size:
# original method is significantly faster, but some images are read sideways. Not sure why. Do slow method by default.
@ -150,6 +160,10 @@ class FileItemDTO(
size_database[file_key] = (w, h, file_signature)
self.width: int = w
self.height: int = h
if self.is_video and self.dataset_config.auto_frame_count:
# compute the real frame count now (same math as load time) so buckets
# are keyed on the frame count this video will actually train at
self.num_frames = self.get_auto_frame_count(video_total_frames, video_fps)
self.dataloader_transforms = kwargs.get("dataloader_transforms", None)
super().__init__(*args, **kwargs)

82
toolkit/dataloader_mixins.py
View File

@ -470,6 +470,33 @@ class AudioProcessingDTOMixin:
class ImageProcessingDTOMixin:
def get_auto_frame_count(self: 'FileItemDTO', total_frames: int, video_fps: float) -> int:
# frame count this video will train at with auto_frame_count. Also called at
# FileItemDTO init so bucket keys carry the real frame count, so it must give
# the same answer at bucketing time and at load time.
# 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)
if getattr(self, 'frame_count_snapper', None) is not None:
# model-specific valid-frame-count grid (e.g. minimax_h3's 17n+5)
desired_num_frames = self.frame_count_snapper(desired_num_frames)
else:
# make sure it is divisible by temporal_compression
if self.dataset_config.trim_auto_frame_count_tail:
# snap to the largest valid count that fits inside the video (after the
# key frame +1 below) so trim mode never overshoots the source, which
# would freeze the last frame and pad the audio tail with silence
desired_num_frames = max(0, desired_num_frames - 1) // self.temporal_compression * self.temporal_compression
else:
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
return desired_num_frames
def load_and_process_video(
self: 'FileItemDTO',
transform: Union[None, transforms.Compose],
@ -511,26 +538,18 @@ 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
self.num_frames = self.get_auto_frame_count(total_frames, video_fps)
desired_num_frames = int(vid_length_seconds * self.dataset_config.fps)
if getattr(self, 'frame_count_snapper', None) is not None:
# model-specific valid-frame-count grid (e.g. minimax_h3's 17n+5)
desired_num_frames = self.frame_count_snapper(desired_num_frames)
else:
# 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.num_frames:
if self.dataset_config.auto_frame_count and self.dataset_config.trim_auto_frame_count_tail:
# preserve real time: pull frames at the dataset fps from the start of the
# video and trim the tail that didn't fit the snapped frame count, instead of
# shrinking the whole video to fit (which speeds up motion / chipmunks audio).
# Critical for audio models (e.g. minimax_h3) where audio must stay in sync.
fps_ratio = video_fps / self.dataset_config.fps if video_fps and video_fps > 0 else 1.0
frames_to_extract = [min(round(i * fps_ratio), max_frame_index) for i in range(self.num_frames)]
elif 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.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)]
@ -733,10 +752,20 @@ class ImageProcessingDTOMixin:
gain = target_peak / (peak + eps)
waveform = waveform * gain
trim_tail_audio = (
self.dataset_config.auto_frame_count
and self.dataset_config.trim_auto_frame_count_tail
)
# Slice to the selected clip region (when we have a meaningful time range)
if source_duration > 0.0:
start_sample = int(round(clip_start_time * sample_rate))
end_sample = int(round(clip_end_time * sample_rate))
if trim_tail_audio and target_duration > 0.0:
# time must stay 1:1 with the video — cut exactly the
# training duration so no stretch is needed below
end_sample = start_sample + round(target_duration * sample_rate)
else:
end_sample = round(clip_end_time * sample_rate)
start_sample = max(0, min(start_sample, waveform.shape[-1]))
end_sample = max(0, min(end_sample, waveform.shape[-1]))
if end_sample > start_sample:
@ -749,10 +778,19 @@ class ImageProcessingDTOMixin:
waveform = None
if waveform is not None and waveform.numel() > 0:
target_samples = int(round(target_duration * sample_rate))
target_samples = round(target_duration * sample_rate)
if target_samples > 0 and waveform.shape[-1] != target_samples:
# Time-stretch/shrink to match the video clip duration implied by dataset FPS.
if self.dataset_config.audio_preserve_pitch:
if trim_tail_audio:
# never stretch/contract in trim mode. The waveform can only be
# short here (audio/video ended a hair before the target) —
# pad the tail with silence, or cut any rounding overshoot
pad = target_samples - waveform.shape[-1]
if pad > 0:
waveform = F.pad(waveform, (0, pad))
else:
waveform = waveform[..., :target_samples]
elif self.dataset_config.audio_preserve_pitch:
waveform = time_stretch_preserve_pitch(waveform, sample_rate, target_samples) # waveform is [C, L]
else:
# Use linear interpolation over the time axis.
@ -1732,6 +1770,10 @@ class LatentCachingFileItemDTOMixin:
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
if self.dataset_config.trim_auto_frame_count_tail:
# changes frame selection; only added when on so caches made before
# this option existed stay valid when it is off
item["trim_auto_frame_count_tail"] = True
is_video = True
elif self.is_video and self.dataset_config.num_frames > 1:
item["num_frames"] = self.dataset_config.num_frames