File size: 4,642 Bytes
af8c481 | 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 | """Encode every frame of one T-Rex episode with the frozen Wan2.2 VAE.
Replicates the baseline observation path exactly: each view is resized to
160x320 (bilinear, antialias), tiled into the DreamZero 2x2 grid
[head | right ; left | black] (320x640), resized back to 160x320 (bilinear,
as ``encode_video`` does) and encoded as a single-frame video into a
48x10x20 latent. Output: ``latents`` [N,48,10,20] float16.
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
import numpy as np
import torch
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as TF
from fact.perception.wan_video_vae import WanVideoVAE38
VIEW_DIRS = (
"observation.images.head_left",
"observation.images.left_wrist",
"observation.images.right_wrist",
)
def load_vae(path: str, device: torch.device) -> WanVideoVAE38:
vae = WanVideoVAE38(z_dim=48, dim=160, vae_pretrained_path=path)
state = torch.load(path, map_location="cpu")
if isinstance(state, dict):
for key in ("state_dict", "model", "module"):
if isinstance(state.get(key), dict):
state = state[key]
break
state = {k.removeprefix("module."): v for k, v in state.items() if torch.is_tensor(v)}
current = vae.model.state_dict()
compatible = {k: v for k, v in state.items() if k in current and current[k].shape == v.shape}
coverage = sum(v.numel() for v in compatible.values()) / sum(v.numel() for v in current.values())
if coverage < 0.99:
raise ValueError(f"VAE coverage {coverage:.3f} from {path}")
vae.model.load_state_dict(compatible, strict=False)
vae = vae.to(device).eval()
vae.mean = vae.mean.to(device)
vae.std = vae.std.to(device)
vae.scale = [vae.mean, vae.std.reciprocal()]
return vae
def read_all_frames(path: Path) -> np.ndarray:
import decord
reader = decord.VideoReader(str(path), num_threads=2)
return reader.get_batch(list(range(len(reader)))).asnumpy().astype(np.uint8)
def composite(views: list[np.ndarray], height: int, width: int) -> torch.Tensor:
"""3 x [N,H,W,3] uint8 -> [N,3,H,W] float in [-1,1] at the model size."""
resized = []
for frames in views:
t = torch.from_numpy(frames).permute(0, 3, 1, 2) # N,3,H,W uint8
t = TF.resize(t, [height, width], interpolation=InterpolationMode.BILINEAR, antialias=True)
resized.append(t)
head, left, right = resized
n = head.shape[0]
grid = torch.zeros(n, 3, 2 * height, 2 * width, dtype=torch.uint8)
grid[:, :, :height, :width] = head
grid[:, :, height:, :width] = left
grid[:, :, :height, width:] = right
grid = grid.float() / 255.0 * 2 - 1
grid = torch.nn.functional.interpolate(grid, size=(height, width), mode="bilinear", align_corners=False)
return grid
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--dataset-root", required=True)
parser.add_argument("--episode", type=int, default=0)
parser.add_argument("--vae", required=True)
parser.add_argument("--out", required=True)
parser.add_argument("--height", type=int, default=160)
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--batch", type=int, default=16)
args = parser.parse_args()
device = torch.device("cuda")
root = Path(args.dataset_root)
views = []
for view in VIEW_DIRS:
path = root / "videos" / "chunk-000" / view / f"episode_{args.episode:06d}.mp4"
frames = read_all_frames(path)
print(f"{view}: {frames.shape}", flush=True)
views.append(frames)
n = min(v.shape[0] for v in views)
views = [v[:n] for v in views]
images = composite(views, args.height, args.width) # N,3,H,W
vae = load_vae(args.vae, device)
latents = []
start = time.time()
with torch.no_grad():
for i in range(0, n, args.batch):
chunk = images[i : i + args.batch].to(device)
video = chunk.unsqueeze(2) # B,3,1,H,W (single-frame videos)
z = vae.encode(video, tiled=False) # B,48,1,h,w
latents.append(z[:, :, 0].float().cpu())
if (i // args.batch) % 10 == 0:
print(f" {i}/{n} ({time.time() - start:.1f}s)", flush=True)
latents = torch.cat(latents).numpy().astype(np.float16)
print("latents", latents.shape, "mean", float(latents.mean()), "std", float(latents.std()))
np.savez_compressed(args.out, latents=latents, episode=args.episode, dataset_root=str(root))
print("wrote", args.out)
if __name__ == "__main__":
main()
|