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