| 1255 | |
| 1256 | # This way is quicker when batch grows up |
| 1257 | def _batch_encode_vae(pixel_values): |
| 1258 | pixel_values = rearrange(pixel_values, "b f c h w -> b c f h w") |
| 1259 | bs = args.vae_mini_batch |
| 1260 | new_pixel_values = [] |
| 1261 | for i in range(0, pixel_values.shape[0], bs): |
| 1262 | pixel_values_bs = pixel_values[i: i + bs] |
| 1263 | pixel_values_bs = vae.encode(pixel_values_bs)[0] |
| 1264 | pixel_values_bs = pixel_values_bs.sample() |
| 1265 | new_pixel_values.append(pixel_values_bs) |
| 1266 | del pixel_values_bs |
| 1267 | if args.low_vram: |
| 1268 | torch.cuda.empty_cache() |
| 1269 | result = torch.cat(new_pixel_values, dim=0) |
| 1270 | del new_pixel_values |
| 1271 | return result |
| 1272 | |
| 1273 | if vae_stream_1 is not None: |
| 1274 | vae_stream_1.wait_stream(torch.cuda.current_stream()) |