stable-diffusion-webui/ldm/data/base.py
hlky 010b27ce9a
Repo merge (#712)
* repo-merge

* cutdown size

* Create setup.py

* webui.cmd

* ldm

* Update environment.yaml

* Update environment.yaml
2022-09-06 23:50:14 +01:00

23 lines
693 B
Python

from abc import abstractmethod
from torch.utils.data import Dataset, ConcatDataset, ChainDataset, IterableDataset
class Txt2ImgIterableBaseDataset(IterableDataset):
'''
Define an interface to make the IterableDatasets for text2img data chainable
'''
def __init__(self, num_records=0, valid_ids=None, size=256):
super().__init__()
self.num_records = num_records
self.valid_ids = valid_ids
self.sample_ids = valid_ids
self.size = size
print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
def __len__(self):
return self.num_records
@abstractmethod
def __iter__(self):
pass