(_)
| 349 | |
| 350 | |
| 351 | def worker_init_fn(_): |
| 352 | worker_info = torch.utils.data.get_worker_info() |
| 353 | |
| 354 | dataset = worker_info.dataset |
| 355 | worker_id = worker_info.id |
| 356 | |
| 357 | if isinstance(dataset, Txt2ImgIterableBaseDataset): |
| 358 | split_size = dataset.num_records // worker_info.num_workers |
| 359 | # reset num_records to the true number to retain reliable length information |
| 360 | dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size] |
| 361 | current_id = np.random.choice(len(np.random.get_state()[1]), 1) |
| 362 | return np.random.seed(np.random.get_state()[1][current_id] + worker_id) |
| 363 | else: |
| 364 | return np.random.seed(np.random.get_state()[1][0] + worker_id) |
| 365 | |
| 366 | |
| 367 | class ConcatDataset(Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected