| 29 | |
| 30 | class DataConfig: |
| 31 | def __init__( |
| 32 | self, |
| 33 | grouped_datasets, |
| 34 | text_cond_dropout_prob=0.1, |
| 35 | vit_cond_dropout_prob=0.4, |
| 36 | dino_cond_dropout_prob=0.4, |
| 37 | vit_patch_size=14, |
| 38 | dino_patch_size=14, |
| 39 | vit_max_num_patch_per_side=70, |
| 40 | dino_max_num_patch_per_side=37, |
| 41 | ): |
| 42 | self.grouped_datasets = grouped_datasets |
| 43 | self.text_cond_dropout_prob = text_cond_dropout_prob |
| 44 | self.vit_cond_dropout_prob = vit_cond_dropout_prob |
| 45 | self.vit_patch_size = vit_patch_size |
| 46 | self.dino_cond_dropout_prob = dino_cond_dropout_prob |
| 47 | self.dino_patch_size = dino_patch_size |
| 48 | self.vit_max_num_patch_per_side = vit_max_num_patch_per_side |
| 49 | self.dino_max_num_patch_per_side = dino_max_num_patch_per_side |
| 50 | |
| 51 | |
| 52 | class PackedDataset(torch.utils.data.IterableDataset): |