INNER CODE UNIT · Python
__init__
TIGER-AI-Lab/ImagenHub · src/imagen_hub/pipelines/custom_diffusion/custom_diffusion_src/stable-diffusion/main.py:139
def __init__(self, dataset):
self.data = dataset
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
return self.data[idx]
def worker_init_fn(_):
worker_info = torch.utils.data.get_worker_info()
dataset = worker_info.dataset
worker_id = worker_info.id
if isinstance(dataset, Txt2ImgIterableBaseDataset):
split_size = dataset.num_records // worker_info.num_workers