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