wehub-resource-sync
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from utils.dataset import MultiVideoConcatDataset
|
|
|
|
|
|
class I2VDatasetFrameAccountingTest(unittest.TestCase):
|
|
def _dataset(self, root, latent_frames):
|
|
(root / "video" / "sample").mkdir(parents=True)
|
|
(root / "caption" / "sample").mkdir(parents=True)
|
|
total_raw_frames = 1 + (latent_frames - 1) * 4
|
|
return MultiVideoConcatDataset(
|
|
data_dir=str(root),
|
|
video_size=(704, 1280),
|
|
total_frames=total_raw_frames,
|
|
independent_first_frame=True,
|
|
num_frame_per_block=8,
|
|
temporal_compression_ratio=4,
|
|
)
|
|
|
|
def test_96_frame_i2v_uses_regular_8_frame_blocks(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
dataset = self._dataset(Path(tmp), latent_frames=96)
|
|
|
|
self.assertEqual(dataset.first_chunk_latent_frames, 8)
|
|
self.assertEqual(dataset.first_chunk_frames, 29)
|
|
self.assertEqual(dataset.total_segments, 12)
|
|
|
|
def test_33_frame_i2v_keeps_legacy_one_plus_block_layout(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
dataset = self._dataset(Path(tmp), latent_frames=33)
|
|
|
|
self.assertEqual(dataset.first_chunk_latent_frames, 9)
|
|
self.assertEqual(dataset.first_chunk_frames, 33)
|
|
self.assertEqual(dataset.total_segments, 4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|