| | import os |
| | import tempfile |
| | import unittest |
| |
|
| | from finetrainers.data import ( |
| | InMemoryDistributedDataPreprocessor, |
| | PrecomputedDistributedDataPreprocessor, |
| | VideoCaptionFilePairDataset, |
| | initialize_preprocessor, |
| | wrap_iterable_dataset_for_preprocessing, |
| | ) |
| | from finetrainers.data.precomputation import PRECOMPUTED_DATA_DIR |
| | from finetrainers.utils import find_files |
| |
|
| | from .utils import create_dummy_directory_structure |
| |
|
| |
|
| | class PreprocessorFastTests(unittest.TestCase): |
| | def setUp(self): |
| | self.rank = 0 |
| | self.world_size = 1 |
| | self.num_items = 3 |
| | self.processor_fn = { |
| | "latent": self._latent_processor_fn, |
| | "condition": self._condition_processor_fn, |
| | } |
| | self.save_dir = tempfile.TemporaryDirectory() |
| |
|
| | directory_structure = [ |
| | "0.mp4", |
| | "1.mp4", |
| | "2.mp4", |
| | "0.txt", |
| | "1.txt", |
| | "2.txt", |
| | ] |
| | create_dummy_directory_structure( |
| | directory_structure, self.save_dir, self.num_items, "a cat ruling the world", "mp4" |
| | ) |
| |
|
| | dataset = VideoCaptionFilePairDataset(self.save_dir.name, infinite=True) |
| | dataset = wrap_iterable_dataset_for_preprocessing( |
| | dataset, |
| | dataset_type="video", |
| | config={ |
| | "video_resolution_buckets": [[2, 32, 32]], |
| | "reshape_mode": "bicubic", |
| | }, |
| | ) |
| | self.dataset = dataset |
| |
|
| | def tearDown(self): |
| | self.save_dir.cleanup() |
| |
|
| | @staticmethod |
| | def _latent_processor_fn(**data): |
| | video = data["video"] |
| | video = video[:, :, :16, :16] |
| | data["video"] = video |
| | return data |
| |
|
| | @staticmethod |
| | def _condition_processor_fn(**data): |
| | caption = data["caption"] |
| | caption = caption + " surrounded by mystical aura" |
| | data["caption"] = caption |
| | return data |
| |
|
| | def test_initialize_preprocessor(self): |
| | preprocessor = initialize_preprocessor( |
| | self.rank, |
| | self.world_size, |
| | self.num_items, |
| | self.processor_fn, |
| | self.save_dir.name, |
| | enable_precomputation=False, |
| | ) |
| | self.assertIsInstance(preprocessor, InMemoryDistributedDataPreprocessor) |
| |
|
| | preprocessor = initialize_preprocessor( |
| | self.rank, |
| | self.world_size, |
| | self.num_items, |
| | self.processor_fn, |
| | self.save_dir.name, |
| | enable_precomputation=True, |
| | ) |
| | self.assertIsInstance(preprocessor, PrecomputedDistributedDataPreprocessor) |
| |
|
| | def test_in_memory_preprocessor_consume(self): |
| | data_iterator = iter(self.dataset) |
| | preprocessor = initialize_preprocessor( |
| | self.rank, |
| | self.world_size, |
| | self.num_items, |
| | self.processor_fn, |
| | self.save_dir.name, |
| | enable_precomputation=False, |
| | ) |
| |
|
| | condition_iterator = preprocessor.consume( |
| | "condition", components={}, data_iterator=data_iterator, cache_samples=True |
| | ) |
| | latent_iterator = preprocessor.consume( |
| | "latent", components={}, data_iterator=data_iterator, use_cached_samples=True, drop_samples=True |
| | ) |
| |
|
| | self.assertFalse(preprocessor.requires_data) |
| | for _ in range(self.num_items): |
| | condition_item = next(condition_iterator) |
| | latent_item = next(latent_iterator) |
| | self.assertIn("caption", condition_item) |
| | self.assertIn("video", latent_item) |
| | self.assertEqual(condition_item["caption"], "a cat ruling the world surrounded by mystical aura") |
| | self.assertEqual(latent_item["video"].shape[-2:], (16, 16)) |
| | self.assertTrue(preprocessor.requires_data) |
| |
|
| | def test_in_memory_preprocessor_consume_once(self): |
| | data_iterator = iter(self.dataset) |
| | preprocessor = initialize_preprocessor( |
| | self.rank, |
| | self.world_size, |
| | self.num_items, |
| | self.processor_fn, |
| | self.save_dir.name, |
| | enable_precomputation=False, |
| | ) |
| |
|
| | condition_iterator = preprocessor.consume_once( |
| | "condition", components={}, data_iterator=data_iterator, cache_samples=True |
| | ) |
| | latent_iterator = preprocessor.consume_once( |
| | "latent", components={}, data_iterator=data_iterator, use_cached_samples=True, drop_samples=True |
| | ) |
| |
|
| | self.assertFalse(preprocessor.requires_data) |
| | for _ in range(self.num_items): |
| | condition_item = next(condition_iterator) |
| | latent_item = next(latent_iterator) |
| | self.assertIn("caption", condition_item) |
| | self.assertIn("video", latent_item) |
| | self.assertEqual(condition_item["caption"], "a cat ruling the world surrounded by mystical aura") |
| | self.assertEqual(latent_item["video"].shape[-2:], (16, 16)) |
| | self.assertFalse(preprocessor.requires_data) |
| |
|
| | def test_precomputed_preprocessor_consume(self): |
| | data_iterator = iter(self.dataset) |
| | preprocessor = initialize_preprocessor( |
| | self.rank, |
| | self.world_size, |
| | self.num_items, |
| | self.processor_fn, |
| | self.save_dir.name, |
| | enable_precomputation=True, |
| | ) |
| |
|
| | condition_iterator = preprocessor.consume( |
| | "condition", components={}, data_iterator=data_iterator, cache_samples=True |
| | ) |
| | latent_iterator = preprocessor.consume( |
| | "latent", components={}, data_iterator=data_iterator, use_cached_samples=True, drop_samples=True |
| | ) |
| |
|
| | precomputed_data_dir = os.path.join(self.save_dir.name, PRECOMPUTED_DATA_DIR) |
| | condition_file_list = find_files(precomputed_data_dir, "condition-*") |
| | latent_file_list = find_files(precomputed_data_dir, "latent-*") |
| | self.assertEqual(len(condition_file_list), 3) |
| | self.assertEqual(len(latent_file_list), 3) |
| |
|
| | self.assertFalse(preprocessor.requires_data) |
| | for _ in range(self.num_items): |
| | condition_item = next(condition_iterator) |
| | latent_item = next(latent_iterator) |
| | self.assertIn("caption", condition_item) |
| | self.assertIn("video", latent_item) |
| | self.assertEqual(condition_item["caption"], "a cat ruling the world surrounded by mystical aura") |
| | self.assertEqual(latent_item["video"].shape[-2:], (16, 16)) |
| | self.assertTrue(preprocessor.requires_data) |
| |
|
| | def test_precomputed_preprocessor_consume_once(self): |
| | data_iterator = iter(self.dataset) |
| | preprocessor = initialize_preprocessor( |
| | self.rank, |
| | self.world_size, |
| | self.num_items, |
| | self.processor_fn, |
| | self.save_dir.name, |
| | enable_precomputation=True, |
| | ) |
| |
|
| | condition_iterator = preprocessor.consume_once( |
| | "condition", components={}, data_iterator=data_iterator, cache_samples=True |
| | ) |
| | latent_iterator = preprocessor.consume_once( |
| | "latent", components={}, data_iterator=data_iterator, use_cached_samples=True, drop_samples=True |
| | ) |
| |
|
| | precomputed_data_dir = os.path.join(self.save_dir.name, PRECOMPUTED_DATA_DIR) |
| | condition_file_list = find_files(precomputed_data_dir, "condition-*") |
| | latent_file_list = find_files(precomputed_data_dir, "latent-*") |
| | self.assertEqual(len(condition_file_list), 3) |
| | self.assertEqual(len(latent_file_list), 3) |
| |
|
| | self.assertFalse(preprocessor.requires_data) |
| | for _ in range(self.num_items): |
| | condition_item = next(condition_iterator) |
| | latent_item = next(latent_iterator) |
| | self.assertIn("caption", condition_item) |
| | self.assertIn("video", latent_item) |
| | self.assertEqual(condition_item["caption"], "a cat ruling the world surrounded by mystical aura") |
| | self.assertEqual(latent_item["video"].shape[-2:], (16, 16)) |
| | self.assertFalse(preprocessor.requires_data) |
| |
|
