(_)
| 326 | |
| 327 | |
| 328 | def worker_init_fn(_): |
| 329 | worker_info = torch.utils.data.get_worker_info() |
| 330 | |
| 331 | dataset = worker_info.dataset |
| 332 | worker_id = worker_info.id |
| 333 | |
| 334 | if isinstance(dataset, Txt2ImgIterableBaseDataset): |
| 335 | split_size = dataset.num_records // worker_info.num_workers |
| 336 | # reset num_records to the true number to retain reliable length information |
| 337 | dataset.sample_ids = dataset.valid_ids[worker_id * split_size:(worker_id + 1) * split_size] |
| 338 | current_id = np.random.choice(len(np.random.get_state()[1]), 1) |
| 339 | return np.random.seed(np.random.get_state()[1][current_id] + worker_id) |
| 340 | else: |
| 341 | return np.random.seed(np.random.get_state()[1][0] + worker_id) |
| 342 | |
| 343 | |
| 344 | class ConcatDataset(Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected