(t, vae, scale=True)
| 366 | |
| 367 | |
| 368 | def tensor_to_vae_latent(t, vae, scale=True): |
| 369 | t = t.to(vae.dtype) |
| 370 | if len(t.shape) == 5: |
| 371 | video_length = t.shape[1] |
| 372 | |
| 373 | t = rearrange(t, "b f c h w -> (b f) c h w") |
| 374 | latents = vae.encode(t).latent_dist.sample() |
| 375 | latents = rearrange(latents, "(b f) c h w -> b f c h w", f=video_length) |
| 376 | elif len(t.shape) == 4: |
| 377 | latents = vae.encode(t).latent_dist.sample() |
| 378 | if scale: |
| 379 | latents = latents * vae.config.scaling_factor |
| 380 | return latents |
| 381 | |
| 382 | |
| 383 | def parse_args(): |