Spaces:
Running on Zero
Running on Zero
File size: 1,801 Bytes
f110cb1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | """
Functional API for VAE encoding of reference images.
"""
import torch
from einops import rearrange, repeat
from uranus.modules.vae import WanVideoVAE
def encode_reference_images(
vae: WanVideoVAE,
ref_images: list[torch.Tensor], # N_REF tensors of shape [3, H, W] in [-1, 1]
device: torch.device,
tiled: bool = True,
tile_size: tuple[int, int] = (30, 52),
tile_stride: tuple[int, int] = (15, 26),
) -> torch.Tensor:
"""
VAE-encode reference images into latent space.
Args:
vae: WanVideoVAE instance
ref_images: List of reference images, each [3, H, W] in [-1, 1]
device: Target device
tiled: Whether to use tiled VAE encoding
tile_size: Tile size in latent space
tile_stride: Tile stride in latent space
Returns:
Reference latents of shape [1, N_CAM, z_dim, N_REF, H_lat, W_lat]
where N_CAM is 1 (replicated per camera by the caller if needed).
"""
if not isinstance(ref_images, list):
ref_images = [ref_images]
n_ref = len(ref_images)
# Stack: [N_REF, 3, H, W] -> [B=1, 3, N_REF, H, W]
ref_tensor = torch.stack(ref_images, dim=0).unsqueeze(0)
ref_tensor = rearrange(ref_tensor, "B N_REF C H W -> B C N_REF H W")
# VAE encode expects [B, C, T, H, W] with T=1 per image
ref_tensor = rearrange(ref_tensor, "B C N_REF H W -> (B N_REF) C 1 H W")
ref_latents = vae.encode(ref_tensor, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride)
# Reshape: [B*N_REF, z_dim, 1, H_lat, W_lat] -> [1, 1, z_dim, N_REF, H_lat, W_lat]
ref_latents = rearrange(ref_latents, "(B N_REF) C 1 H W -> B N_REF C H W", N_REF=n_ref)
ref_latents = rearrange(ref_latents, "B N_REF C H W -> B C N_REF H W").unsqueeze(1)
return ref_latents
|