| 325 | yield data["image"].to(device), None |
| 326 | |
| 327 | def get_real_img_generator(): # [0,255] |
| 328 | while True: |
| 329 | for data in tqdm( |
| 330 | loader, |
| 331 | disable=not accelerator.is_main_process, |
| 332 | initial=0, |
| 333 | desc="generate_real_img", |
| 334 | ): |
| 335 | if has_text(args): |
| 336 | yield out2img(data["image"]).to(device) |
| 337 | elif "facehq" in str(args.data.name): |
| 338 | yield out2img(data["image"]).to(device) |
| 339 | elif "ucf101" in str(args.data.name): |
| 340 | _video = data["frame_feature256"].to(device) |
| 341 | if args.is_latent: |
| 342 | with torch.no_grad(): |
| 343 | _video = torch.stack( |
| 344 | [ |
| 345 | vae.decode(_video[:, _v]).sample |
| 346 | for _v in range(_video.shape[1]) |
| 347 | ], |
| 348 | dim=1, |
| 349 | ) |
| 350 | yield out2img(_video) |
| 351 | elif "celebav" in str(args.data.name): |
| 352 | _start = random.randint( |
| 353 | 0, |
| 354 | data["frame_feature256"].shape[1] |
| 355 | - args.model.params.video_frames |
| 356 | - 1, |
| 357 | ) |
| 358 | _video = data["frame_feature256"][ |
| 359 | :, _start : _start + args.model.params.video_frames |
| 360 | ].to(device) |
| 361 | if args.is_latent: |
| 362 | with torch.no_grad(): |
| 363 | _video = torch.stack( |
| 364 | [ |
| 365 | vae.decode(_video[:, _v]).sample |
| 366 | for _v in range(_video.shape[1]) |
| 367 | ], |
| 368 | dim=1, |
| 369 | ) |
| 370 | yield out2img(_video) |
| 371 | |
| 372 | else: |
| 373 | raise NotImplementedError("latent data not supported") |
| 374 | |
| 375 | def get_cap_generator(): |
| 376 | while True: |