Spaces:
Running
Running
| import pathlib | |
| import tempfile | |
| import unittest | |
| import torch | |
| from PIL import Image | |
| from finetrainers.data import ( | |
| ImageCaptionFilePairDataset, | |
| ImageFileCaptionFileListDataset, | |
| ImageFolderDataset, | |
| ValidationDataset, | |
| VideoCaptionFilePairDataset, | |
| VideoFileCaptionFileListDataset, | |
| VideoFolderDataset, | |
| VideoWebDataset, | |
| initialize_dataset, | |
| ) | |
| from finetrainers.utils import find_files | |
| from .utils import create_dummy_directory_structure | |
| class DatasetTesterMixin: | |
| num_data_files = None | |
| directory_structure = None | |
| caption = "A cat ruling the world" | |
| metadata_extension = None | |
| def setUp(self): | |
| if self.num_data_files is None: | |
| raise ValueError("num_data_files is not defined") | |
| if self.directory_structure is None: | |
| raise ValueError("dataset_structure is not defined") | |
| self.tmpdir = tempfile.TemporaryDirectory() | |
| create_dummy_directory_structure( | |
| self.directory_structure, self.tmpdir, self.num_data_files, self.caption, self.metadata_extension | |
| ) | |
| def tearDown(self): | |
| self.tmpdir.cleanup() | |
| class ImageDatasetTesterMixin(DatasetTesterMixin): | |
| metadata_extension = "jpg" | |
| class VideoDatasetTesterMixin(DatasetTesterMixin): | |
| metadata_extension = "mp4" | |
| class ImageCaptionFilePairDatasetFastTests(ImageDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "0.jpg", | |
| "1.jpg", | |
| "2.jpg", | |
| "0.txt", | |
| "1.txt", | |
| "2.txt", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = ImageCaptionFilePairDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(self.num_data_files): | |
| item = next(iterator) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["image"])) | |
| self.assertEqual(item["image"].shape, (3, 64, 64)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) | |
| self.assertIsInstance(dataset, ImageCaptionFilePairDataset) | |
| class ImageFileCaptionFileListDatasetFastTests(ImageDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "prompts.txt", | |
| "images.txt", | |
| "images/", | |
| "images/0.jpg", | |
| "images/1.jpg", | |
| "images/2.jpg", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = ImageFileCaptionFileListDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for i in range(3): | |
| item = next(iterator) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["image"])) | |
| self.assertEqual(item["image"].shape, (3, 64, 64)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) | |
| self.assertIsInstance(dataset, ImageFileCaptionFileListDataset) | |
| class ImageFolderDatasetFastTests___CSV(ImageDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "metadata.csv", | |
| "0.jpg", | |
| "1.jpg", | |
| "2.jpg", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = ImageFolderDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(3): | |
| item = next(iterator) | |
| self.assertIn("caption", item) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["image"])) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) | |
| self.assertIsInstance(dataset, ImageFolderDataset) | |
| class ImageFolderDatasetFastTests___JSONL(ImageDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "metadata.jsonl", | |
| "0.jpg", | |
| "1.jpg", | |
| "2.jpg", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = ImageFolderDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(3): | |
| item = next(iterator) | |
| self.assertIn("caption", item) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["image"])) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "image", infinite=False) | |
| self.assertIsInstance(dataset, ImageFolderDataset) | |
| class VideoCaptionFilePairDatasetFastTests(VideoDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "0.mp4", | |
| "1.mp4", | |
| "2.mp4", | |
| "0.txt", | |
| "1.txt", | |
| "2.txt", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = VideoCaptionFilePairDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(self.num_data_files): | |
| item = next(iterator) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["video"])) | |
| self.assertEqual(len(item["video"]), 4) | |
| self.assertEqual(item["video"][0].shape, (3, 64, 64)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) | |
| self.assertIsInstance(dataset, VideoCaptionFilePairDataset) | |
| class VideoFileCaptionFileListDatasetFastTests(VideoDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "prompts.txt", | |
| "videos.txt", | |
| "videos/", | |
| "videos/0.mp4", | |
| "videos/1.mp4", | |
| "videos/2.mp4", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = VideoFileCaptionFileListDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(3): | |
| item = next(iterator) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["video"])) | |
| self.assertEqual(len(item["video"]), 4) | |
| self.assertEqual(item["video"][0].shape, (3, 64, 64)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) | |
| self.assertIsInstance(dataset, VideoFileCaptionFileListDataset) | |
| class VideoFolderDatasetFastTests___CSV(VideoDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "metadata.csv", | |
| "0.mp4", | |
| "1.mp4", | |
| "2.mp4", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = VideoFolderDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(3): | |
| item = next(iterator) | |
| self.assertIn("caption", item) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["video"])) | |
| self.assertEqual(len(item["video"]), 4) | |
| self.assertEqual(item["video"][0].shape, (3, 64, 64)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) | |
| self.assertIsInstance(dataset, VideoFolderDataset) | |
| class VideoFolderDatasetFastTests___JSONL(VideoDatasetTesterMixin, unittest.TestCase): | |
| num_data_files = 3 | |
| directory_structure = [ | |
| "metadata.jsonl", | |
| "0.mp4", | |
| "1.mp4", | |
| "2.mp4", | |
| ] | |
| def setUp(self): | |
| super().setUp() | |
| self.dataset = VideoFolderDataset(self.tmpdir.name, infinite=False) | |
| def test_getitem(self): | |
| iterator = iter(self.dataset) | |
| for _ in range(3): | |
| item = next(iterator) | |
| self.assertIn("caption", item) | |
| self.assertEqual(item["caption"], self.caption) | |
| self.assertTrue(torch.is_tensor(item["video"])) | |
| self.assertEqual(len(item["video"]), 4) | |
| self.assertEqual(item["video"][0].shape, (3, 64, 64)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset(self.tmpdir.name, "video", infinite=False) | |
| self.assertIsInstance(dataset, VideoFolderDataset) | |
| class ImageWebDatasetFastTests(unittest.TestCase): | |
| # TODO(aryan): setup a dummy dataset | |
| pass | |
| class VideoWebDatasetFastTests(unittest.TestCase): | |
| def setUp(self): | |
| self.num_data_files = 15 | |
| self.dataset = VideoWebDataset("finetrainers/dummy-squish-wds", infinite=False) | |
| def test_getitem(self): | |
| for index, item in enumerate(self.dataset): | |
| if index > 2: | |
| break | |
| self.assertIn("caption", item) | |
| self.assertIn("video", item) | |
| self.assertTrue(torch.is_tensor(item["video"])) | |
| self.assertEqual(len(item["video"]), 121) | |
| self.assertEqual(item["video"][0].shape, (3, 720, 1280)) | |
| def test_initialize_dataset(self): | |
| dataset = initialize_dataset("finetrainers/dummy-squish-wds", "video", infinite=False) | |
| self.assertIsInstance(dataset, VideoWebDataset) | |
| class DatasetUtilsFastTests(unittest.TestCase): | |
| def test_find_files_depth_0(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| file1 = tempfile.NamedTemporaryFile(dir=tmpdir, suffix=".txt", delete=False) | |
| file2 = tempfile.NamedTemporaryFile(dir=tmpdir, suffix=".txt", delete=False) | |
| file3 = tempfile.NamedTemporaryFile(dir=tmpdir, suffix=".txt", delete=False) | |
| files = find_files(tmpdir, "*.txt") | |
| self.assertEqual(len(files), 3) | |
| self.assertIn(file1.name, files) | |
| self.assertIn(file2.name, files) | |
| self.assertIn(file3.name, files) | |
| def test_find_files_depth_n(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| dir1 = tempfile.TemporaryDirectory(dir=tmpdir) | |
| dir2 = tempfile.TemporaryDirectory(dir=dir1.name) | |
| file1 = tempfile.NamedTemporaryFile(dir=dir1.name, suffix=".txt", delete=False) | |
| file2 = tempfile.NamedTemporaryFile(dir=dir2.name, suffix=".txt", delete=False) | |
| files = find_files(tmpdir, "*.txt", depth=1) | |
| self.assertEqual(len(files), 1) | |
| self.assertIn(file1.name, files) | |
| self.assertNotIn(file2.name, files) | |
| files = find_files(tmpdir, "*.txt", depth=2) | |
| self.assertEqual(len(files), 2) | |
| self.assertIn(file1.name, files) | |
| self.assertIn(file2.name, files) | |
| self.assertNotIn(dir1.name, files) | |
| self.assertNotIn(dir2.name, files) | |
| class ValidationDatasetFastTests(unittest.TestCase): | |
| def setUp(self): | |
| num_data_files = 3 | |
| self.tmpdir = tempfile.TemporaryDirectory() | |
| metadata_filename = pathlib.Path(self.tmpdir.name) / "metadata.csv" | |
| with open(metadata_filename, "w") as f: | |
| f.write("caption,image_path,video_path\n") | |
| for i in range(num_data_files): | |
| Image.new("RGB", (64, 64)).save((pathlib.Path(self.tmpdir.name) / f"{i}.jpg").as_posix()) | |
| f.write(f"test caption,{self.tmpdir.name}/{i}.jpg,\n") | |
| self.dataset = ValidationDataset(metadata_filename.as_posix()) | |
| def tearDown(self): | |
| self.tmpdir.cleanup() | |
| def test_getitem(self): | |
| for i, data in enumerate(self.dataset): | |
| self.assertEqual(data["image_path"], f"{self.tmpdir.name}/{i}.jpg") | |
| self.assertIsInstance(data["image"], Image.Image) | |
| self.assertEqual(data["image"].size, (64, 64)) | |