github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Recover from issue when a video model first fram may not have been cached properly

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-01 12:00:33 -06:00
parent 5baa495585
commit 6b95282097
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
toolkit/data_transfer_object/data_loader.py
View File

@@ -232,24 +232,32 @@ class DataLoaderBatchDTO:
if any( if any(
[x._cached_first_frame_latent is not None for x in self.file_items] [x._cached_first_frame_latent is not None for x in self.file_items]
): ):
# find one to use as a base; item 0 may not have one
base_first_frame_latent = None
for x in self.file_items:
if x._cached_first_frame_latent is not None:
base_first_frame_latent = x._cached_first_frame_latent
break
self.first_frame_latents = torch.cat( self.first_frame_latents = torch.cat(
[ [
x._cached_first_frame_latent.unsqueeze(0) x._cached_first_frame_latent.unsqueeze(0)
if x._cached_first_frame_latent is not None if x._cached_first_frame_latent is not None
else torch.zeros_like( else torch.zeros_like(base_first_frame_latent).unsqueeze(0)
self.file_items[0]._cached_first_frame_latent
).unsqueeze(0)
for x in self.file_items for x in self.file_items
] ]
) )
if any([x._cached_audio_latent is not None for x in self.file_items]): if any([x._cached_audio_latent is not None for x in self.file_items]):
# find one to use as a base; item 0 may not have one
base_audio_latent = None
for x in self.file_items:
if x._cached_audio_latent is not None:
base_audio_latent = x._cached_audio_latent
break
self.audio_latents = torch.cat( self.audio_latents = torch.cat(
[ [
x._cached_audio_latent.unsqueeze(0) x._cached_audio_latent.unsqueeze(0)
if x._cached_audio_latent is not None if x._cached_audio_latent is not None
else torch.zeros_like( else torch.zeros_like(base_audio_latent).unsqueeze(0)
self.file_items[0]._cached_audio_latent
).unsqueeze(0)
for x in self.file_items for x in self.file_items
] ]
) )