From e3f47a799b23b54ab23a8d85afc5ec51d2d97a7c Mon Sep 17 00:00:00 2001 From: twu Date: Thu, 21 Aug 2025 09:13:45 +0000 Subject: [PATCH] make it more efficient to locate where to sample the frame --- diffsynth/trainers/utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/diffsynth/trainers/utils.py b/diffsynth/trainers/utils.py index b55b258..c4b5d92 100644 --- a/diffsynth/trainers/utils.py +++ b/diffsynth/trainers/utils.py @@ -283,11 +283,16 @@ class VideoDataset(torch.utils.data.Dataset): # make a ((start,end),frameid) struct start_end_idx_map = [((sum(delays[:i]), sum(delays[:i+1])), i) for i in range(len(delays))] _frames = [] + # according gemini-code-assist, make it more efficient to locate + # where to sample the frame + last_match = 0 for i in range(sum(delays) // minimal_interval): current_time = minimal_interval * i - for ((start, end), frame_idx) in start_end_idx_map: + for idx, ((start, end), frame_idx) in enumerate(start_end_idx_map[last_match:]): if start <= current_time < end: _frames.append(frames[frame_idx]) + last_match = idx + last_match + break frames = _frames return frames