Spaces:
Running on Zero
Running on Zero
File size: 4,533 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """
Functional API for VAE decoding of video latents to pixel frames.
"""
import inspect
import torch
from einops import rearrange
from uranus.modules.vae import VideoVAE38_, WanVideoVAE, count_conv3d, unpatchify
def decode_latents(
vae: WanVideoVAE,
latents: torch.Tensor, # [B, N_CAM, z_dim, F_lat, H_lat, W_lat]
device: torch.device,
tiled: bool = True,
tile_size: tuple[int, int] = (34, 34),
tile_stride: tuple[int, int] = (18, 16),
) -> torch.Tensor:
"""
VAE-decode latent frames to video frames.
VAE temporal expansion:
- 1st latent frame → 1 video frame
- subsequent latent frames → 4 video frames each
Args:
vae: WanVideoVAE instance
latents: Video latents [B, N_CAM, z_dim, F_lat, H_lat, W_lat]
device: Target device
tiled: Whether to use tiled VAE decoding
tile_size: Tile size in latent space
tile_stride: Tile stride in latent space
Returns:
Video frames [B, N_CAM, 3, F_video, H_pix, W_pix] in [-1, 1]
"""
batch, n_cam, _, _, _, _ = latents.shape
latents_flat = rearrange(latents, "B N_CAM C F H W -> (B N_CAM) C F H W")
videos = vae.decode(
latents_flat,
device=device,
tiled=tiled,
tile_size=tile_size,
tile_stride=tile_stride,
)
videos = rearrange(
videos, "(B N_CAM) C F H W -> B N_CAM C F H W", B=batch, N_CAM=n_cam
)
return videos.clamp_(-1, 1)
def _apply_vae_scale(z: torch.Tensor, vae: WanVideoVAE) -> torch.Tensor:
scale = vae.scale
if isinstance(scale[0], torch.Tensor):
scale = [s.to(dtype=z.dtype, device=z.device) for s in scale]
return z / scale[1].view(1, vae.z_dim, 1, 1, 1) + scale[0].view(
1, vae.z_dim, 1, 1, 1
)
scale = scale.to(dtype=z.dtype, device=z.device)
return z / scale[1] + scale[0]
def _decode_single_chunk(
vae: WanVideoVAE,
latents: torch.Tensor,
dec_feat_map,
decoded_latent_frames: int,
) -> tuple[torch.Tensor, list, list]:
model = vae.model
z = _apply_vae_scale(latents, vae)
x = model.conv2(z)
# Flash-VAED decoder exposes its own cache-slot count (its causal convs
# are depthwise-separable); fall back to the legacy count otherwise.
count_fn = getattr(model.decoder, "count_conv3d", None) or count_conv3d
if dec_feat_map is None:
dec_feat_map = [None] * count_fn(model.decoder)
decoder_signature = inspect.signature(model.decoder.forward)
supports_first_chunk = "first_chunk" in decoder_signature.parameters
outputs = []
for frame_idx in range(x.shape[2]):
feat_idx = [0]
decoder_kwargs = {
"feat_cache": dec_feat_map,
"feat_idx": feat_idx,
}
if supports_first_chunk:
decoder_kwargs["first_chunk"] = (
decoded_latent_frames == 0 and frame_idx == 0
)
out, dec_feat_map, feat_idx = model.decoder(
x[:, :, frame_idx : frame_idx + 1], **decoder_kwargs
)
outputs.append(out)
video = torch.cat(outputs, dim=2) if len(outputs) > 1 else outputs[0]
if isinstance(model, VideoVAE38_):
video = unpatchify(video, patch_size=2)
return video, dec_feat_map, feat_idx
def decode_latents_chunk(
vae: WanVideoVAE,
latents: torch.Tensor,
device: torch.device,
dec_feat_map=None,
dec_feat_idx: list | None = None,
decoded_latent_frames: int = 0,
tiled: bool = False,
tile_size: tuple[int, int] = (34, 34),
tile_stride: tuple[int, int] = (18, 16),
) -> tuple[torch.Tensor, list, list, int]:
"""
Incrementally decode a latent chunk while reusing decoder feature cache.
This mirrors `VideoVAE_.decode()` / `VideoVAE38_.decode()` but externalizes
decoder cache so streaming callers can carry it in their own state.
"""
del tiled, tile_size, tile_stride
batch, n_cam, _, f_lat, _, _ = latents.shape
latents_flat = rearrange(latents, "B N_CAM C F H W -> (B N_CAM) C F H W").to(
device=device
)
video, dec_feat_map, dec_feat_idx = _decode_single_chunk(
vae=vae,
latents=latents_flat,
dec_feat_map=dec_feat_map,
decoded_latent_frames=decoded_latent_frames,
)
videos = rearrange(
video, "(B N_CAM) C F H W -> B N_CAM C F H W", B=batch, N_CAM=n_cam
)
return (
videos.clamp_(-1, 1),
dec_feat_map,
dec_feat_idx,
decoded_latent_frames + f_lat,
)
|