| 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) |
|
|
