JLNTV_IPD/CogVideo
4
0
Fork 0
mirror of https://github.com/THUDM/CogVideo.git synced 2025-04-06 03:57:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #119 from cly2625/dev_cly2

Browse Source 
fix out-of-index bugs
This commit is contained in:
zR 2024-08-18 02:29:12 +08:00 committed by GitHub
commit 2825d9b707
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
sat/data_video.py
View File

@ -378,8 +378,9 @@ class SFTDataset(Dataset):
num_frames = max_num_frames
start = int(skip_frms_num)
end = int(start + num_frames / fps * actual_fps)
indices = np.arange(start, end, (end - start) / num_frames).astype(int)
temp_frms = vr.get_batch(np.arange(start, end))
end_safty = min(int(start + num_frames / fps * actual_fps), int(ori_vlen))
indices = np.arange(start, end, (end - start) // num_frames).astype(int)
temp_frms = vr.get_batch(np.arange(start, end_safty))
assert temp_frms is not None
tensor_frms = torch.from_numpy(temp_frms) if type(temp_frms) is not torch.Tensor else temp_frms
tensor_frms = tensor_frms[torch.tensor((indices - start).tolist())]
@ -388,7 +389,7 @@ class SFTDataset(Dataset):
num_frames = max_num_frames
start = int(skip_frms_num)
end = int(ori_vlen - skip_frms_num)
indices = np.arange(start, end, (end - start) / num_frames).astype(int)
indices = np.arange(start, end, (end - start) // num_frames).astype(int)
temp_frms = vr.get_batch(np.arange(start, end))
assert temp_frms is not None
tensor_frms = (