| 373 | raise NotImplementedError("latent data not supported") |
| 374 | |
| 375 | def get_cap_generator(): |
| 376 | while True: |
| 377 | for data in tqdm( |
| 378 | loader, |
| 379 | disable=not accelerator.is_main_process, |
| 380 | initial=0, |
| 381 | desc="generate_captions", |
| 382 | ): |
| 383 | if has_text(args): |
| 384 | _cap_feats, _cap = ( |
| 385 | data["caption_feature"].to(device), |
| 386 | data["caption"], |
| 387 | ) |
| 388 | B, N, T, C = _cap_feats.shape # each image has N captions |
| 389 | _p = random.randint(0, N - 1) |
| 390 | yield _cap_feats[:, _p], [_cap[i][_p] for i in range(len(_cap))] |
| 391 | else: |
| 392 | raise NotImplementedError("current dataset doesnt have captions") |
| 393 | |
| 394 | train_dg = get_data_generator() |
| 395 | real_img_dg = get_real_img_generator() |