cedpsam/latentdiffusion

1from abc import abstractmethod2from torch.utils.data import Dataset, ConcatDataset, ChainDataset, IterableDataset3 4 5class Txt2ImgIterableBaseDataset(IterableDataset):6    '''7    Define an interface to make the IterableDatasets for text2img data chainable8    '''9    def __init__(self, num_records=0, valid_ids=None, size=256):10        super().__init__()11        self.num_records = num_records12        self.valid_ids = valid_ids13        self.sample_ids = valid_ids14        self.size = size15 16        print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')17 18    def __len__(self):19        return self.num_records20 21    @abstractmethod22    def __iter__(self):23        pass