| 1606 | with torch.no_grad(): |
| 1607 | # This way is quicker when batch grows up |
| 1608 | def _batch_encode_vae(pixel_values): |
| 1609 | pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") |
| 1610 | bs = args.vae_mini_batch |
| 1611 | new_pixel_values = [] |
| 1612 | for i in range(0, pixel_values.shape[0], bs): |
| 1613 | pixel_values_bs = pixel_values[i : i + bs] |
| 1614 | pixel_values_bs = vae.encode(pixel_values_bs)[0] |
| 1615 | pixel_values_bs = pixel_values_bs.sample() |
| 1616 | new_pixel_values.append(pixel_values_bs) |
| 1617 | return torch.cat(new_pixel_values, dim = 0) |
| 1618 | if vae_stream_1 is not None: |
| 1619 | vae_stream_1.wait_stream(torch.cuda.current_stream()) |
| 1620 | with torch.cuda.stream(vae_stream_1): |