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

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

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
Show Stats 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(
[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(
[
x._cached_first_frame_latent.unsqueeze(0)
if x._cached_first_frame_latent is not None
else torch.zeros_like(
self.file_items[0]._cached_first_frame_latent
).unsqueeze(0)
else torch.zeros_like(base_first_frame_latent).unsqueeze(0)
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(
[
x._cached_audio_latent.unsqueeze(0)
if x._cached_audio_latent is not None
else torch.zeros_like(
self.file_items[0]._cached_audio_latent
).unsqueeze(0)
else torch.zeros_like(base_audio_latent).unsqueeze(0)
for x in self.file_items
]
)