add model card, conditioning, and training-time processing

0613ebe verified 3 months ago

10 kB

#!/usr/bin/env python
"""Extract Wan2.2 VAE video latents + UMT5 text embeddings for a LIBERO LeRobot dataset.
VLA-vs-WAM study, WAM data pipeline. LingBot-VA trains on pre-extracted video latents, not raw
frames (see wan_va/dataset/lerobot_latent_dataset.py). Robbyant released these for LIBERO-Long
only; this script reproduces the *exact* per-camera `.pth` format so we can train the WAM on
libero_spatial / libero_object / libero_goal too.
Output layout (mirrors the released libero-long-lerobot):
  <root>/latents/chunk-000/<cam_key>/episode_{idx:06d}_{start}_{end}.pth
Each .pth is a dict:
  latent (F*H*W, 48) bf16, latent_num_frames/height/width, video_num_frames/height/width,
  text_emb (512, 4096) bf16, text, frame_ids (np.int), start_frame, end_frame, fps, ori_fps
Also writes <root>/empty_emb.pt (UMT5 embedding of "") for classifier-free guidance.
Wan2.2-TI2V VAE: 16x spatial, 4x temporal; latents normalized as (z - latents_mean)/latents_std.
Video frames are truncated to the largest (4k+1) count so temporal compression is exact.
Usage:
  # Verify the pipeline reproduces the released libero-long episode-0 latent, then stop:
  python extract_latents.py --verify_episode0
  # Extract a whole dataset:
  python extract_latents.py \
      --lerobot_root data/libero-spatial-lerobot \
      --wan_dir ckpts/Wan2.2-TI2V-5B-Diffusers \
      --cam_keys observation.images.agentview_rgb observation.images.eye_in_hand_rgb
"""

import argparse
import importlib.machinery
import json
import sys
import types
from pathlib import Path

import numpy as np
import torch


def _stub_flash_attn():
    """Allow importing diffusers without a built flash-attn -- but prefer the real package.
    Bug fixed here (hit on job 446): the original version always stubbed both module names
    unconditionally, even when a real, working flash-attn was installed. That stub only defines
    flash_attn_func, not flash_attn_varlen_func -- so as soon as diffusers' own internal
    `from flash_attn import flash_attn_func, flash_attn_varlen_func` (attention_dispatch.py, hit
    just from `from diffusers import AutoencoderKLWan`, nothing to do with the transformer/
    flash-attn kernel this script never uses) ran, it got the stub instead of the real package
    and crashed with ImportError. Now we only stub a name if it genuinely isn't importable.
    """
    for name in ("flash_attn", "flash_attn_interface"):
        if name in sys.modules:
            continue
        try:
            importlib.import_module(name)
            continue
        except ImportError:
            pass
        mod = types.ModuleType(name)
        mod.flash_attn_func = lambda *a, **k: None
        mod.__spec__ = importlib.machinery.ModuleSpec(name, None)
        mod.__version__ = "0.0.0-stub"
        sys.modules[name] = mod


def largest_4k_plus_1(n: int) -> int:
    """Largest m <= n with m % 4 == 1 (so temporal compression (m-1)/4+1 is exact)."""
    m = n - ((n - 1) % 4)
    return max(m, 1)


def read_video_frames(mp4_path, target_hw):
    """Return uint8 array [T, H, W, 3] resized to target_hw."""
    import av
    import cv2
    container = av.open(str(mp4_path))
    frames = [f.to_ndarray(format="rgb24") for f in container.decode(video=0)]
    container.close()
    frames = np.stack(frames)  # [T,H,W,3]
    H, W = target_hw
    if frames.shape[1] != H or frames.shape[2] != W:
        frames = np.stack([cv2.resize(f, (W, H), interpolation=cv2.INTER_AREA) for f in frames])
    return frames


class Encoders:
    def __init__(self, wan_dir, device="cuda", dtype=torch.bfloat16):
        from diffusers import AutoencoderKLWan
        from transformers import UMT5EncoderModel, T5TokenizerFast
        self.device, self.dtype = device, dtype
        self.vae = AutoencoderKLWan.from_pretrained(wan_dir, subfolder="vae", torch_dtype=torch.float32).to(device).eval()
        self.text_encoder = UMT5EncoderModel.from_pretrained(wan_dir, subfolder="text_encoder", torch_dtype=dtype).to(device).eval()
        self.tokenizer = T5TokenizerFast.from_pretrained(wan_dir, subfolder="tokenizer")
        cfg = self.vae.config
        self.mean = torch.tensor(cfg.latents_mean, dtype=torch.float32, device=device).view(1, -1, 1, 1, 1)
        self.std = torch.tensor(cfg.latents_std, dtype=torch.float32, device=device).view(1, -1, 1, 1, 1)
        self.text_max_len = 512

    @torch.no_grad()
    def encode_video(self, frames_uint8):
        """frames_uint8 [T,H,W,3] -> normalized latent [C,f,h,w] float32, plus dims."""
        T = largest_4k_plus_1(frames_uint8.shape[0])
        frames = frames_uint8[:T]
        # [T,H,W,3] uint8 -> [1,3,T,H,W] in [-1,1]
        x = torch.from_numpy(frames).to(self.device).float().div(127.5).sub(1.0)
        x = x.permute(3, 0, 1, 2).unsqueeze(0)  # [1,3,T,H,W]
        z = self.vae.encode(x).latent_dist.mode()          # [1,48,f,h,w]
        z = (z - self.mean) / self.std                      # Wan normalization
        z = z[0]                                            # [48,f,h,w]
        return z.cpu(), frames_uint8.shape[0], frames.shape[1], frames.shape[2]

    @torch.no_grad()
    def encode_text(self, text):
        tok = self.tokenizer([text], padding="max_length", max_length=self.text_max_len,
                             truncation=True, return_tensors="pt").to(self.device)
        emb = self.text_encoder(input_ids=tok.input_ids, attention_mask=tok.attention_mask).last_hidden_state
        return emb[0].to(self.dtype).cpu()                  # [512,4096]


def make_latent_dict(z_cfhw, video_num_frames, vheight, vwidth, text_emb, text, start, end, fps, ori_fps):
    C, f, h, w = z_cfhw.shape
    latent_flat = z_cfhw.permute(1, 2, 3, 0).reshape(f * h * w, C).contiguous().to(torch.bfloat16)  # (f h w) c
    used = largest_4k_plus_1(video_num_frames)
    return {
        "latent": latent_flat,
        "latent_num_frames": int(f), "latent_height": int(h), "latent_width": int(w),
        "video_num_frames": int(used), "video_height": int(vheight), "video_width": int(vwidth),
        "text_emb": text_emb, "text": text,
        "frame_ids": np.arange(used, dtype=np.int64),
        "start_frame": int(start), "end_frame": int(end),
        "fps": int(fps), "ori_fps": int(ori_fps),
    }


def verify_episode0(args):
    """Extract libero-long episode-0 and compare to the released latent to confirm conventions."""
    _stub_flash_attn()
    root = Path(args.verify_root)
    ref = torch.load(root / "latents/chunk-000/observation.images.agentview_rgb/episode_000000_0_388.pth",
                     weights_only=False)
    enc = Encoders(args.wan_dir, device=args.device)
    frames = read_video_frames(root / "videos/chunk-000/observation.images.agentview_rgb/episode_000000.mp4",
                               (ref["video_height"], ref["video_width"]))
    z, vnf, vh, vw = enc.encode_video(frames)
    print(f"[verify] my latent shape (C,f,h,w)={tuple(z.shape)} ; released flat={tuple(ref['latent'].shape)} "
          f"num_frames={ref['latent_num_frames']} h={ref['latent_height']} w={ref['latent_width']}")
    mine = z.permute(1, 2, 3, 0).reshape(-1, z.shape[0]).float()
    rel = ref["latent"].float()
    if mine.shape == rel.shape:
        err = (mine - rel).abs()
        corr = torch.corrcoef(torch.stack([mine.flatten(), rel.flatten()]))[0, 1]
        print(f"[verify] shape MATCH | mean|abs err|={err.mean():.4f} max={err.max():.4f} corr={corr:.4f}")
        print("[verify] -> PASS (normalization/convention confirmed)" if corr > 0.98
              else "[verify] -> MISMATCH: inspect VAE version / normalization / frame sampling")
    else:
        print("[verify] shape MISMATCH -> check spatial/temporal compression or frame truncation")


def extract_dataset(args):
    _stub_flash_attn()
    root = Path(args.lerobot_root)
    meta = [json.loads(l) for l in (root / "meta/episodes.jsonl").read_text().splitlines()]
    info = json.load(open(root / "meta/info.json"))
    fps = int(info.get("fps", 60))
    enc = Encoders(args.wan_dir, device=args.device)

    # empty text embedding for CFG
    torch.save(enc.encode_text(""), root / "empty_emb.pt")
    print(f"[extract] wrote {root/'empty_emb.pt'}")

    for ep in meta:
        idx = ep["episode_index"]
        chunk = idx // int(info.get("chunks_size", 1000))
        for seg in ep["action_config"]:
            start, end, text = seg["start_frame"], seg["end_frame"], seg["action_text"]
            text_emb = enc.encode_text(text)
            for cam in args.cam_keys:
                mp4 = root / f"videos/chunk-{chunk:03d}/{cam}/episode_{idx:06d}.mp4"
                frames = read_video_frames(mp4, (info["features"][cam]["shape"][0], info["features"][cam]["shape"][1]))
                frames = frames[start:end]
                z, vnf, vh, vw = enc.encode_video(frames)
                d = make_latent_dict(z, frames.shape[0], vh, vw, text_emb, text, start, end, fps, fps)
                out = root / f"latents/chunk-{chunk:03d}/{cam}/episode_{idx:06d}_{start}_{end}.pth"
                out.parent.mkdir(parents=True, exist_ok=True)
                torch.save(d, out)
        print(f"[extract] episode {idx} done")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--wan_dir", default="ckpts/Wan2.2-TI2V-5B-Diffusers")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--verify_episode0", action="store_true",
                    help="Reproduce released libero-long episode-0 latent and compare (no dataset written)")
    ap.add_argument("--verify_root", default="data/libero-long-lerobot")
    ap.add_argument("--lerobot_root", default=None)
    ap.add_argument("--cam_keys", nargs="+",
                    default=["observation.images.agentview_rgb", "observation.images.eye_in_hand_rgb"])
    args = ap.parse_args()
    if args.verify_episode0:
        verify_episode0(args)
    else:
        assert args.lerobot_root, "--lerobot_root required for extraction"
        extract_dataset(args)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python
	"""Extract Wan2.2 VAE video latents + UMT5 text embeddings for a LIBERO LeRobot dataset.

	VLA-vs-WAM study, WAM data pipeline. LingBot-VA trains on pre-extracted video latents, not raw
	frames (see wan_va/dataset/lerobot_latent_dataset.py). Robbyant released these for LIBERO-Long
	only; this script reproduces the exact per-camera `.pth` format so we can train the WAM on
	libero_spatial / libero_object / libero_goal too.

	Output layout (mirrors the released libero-long-lerobot):
	<root>/latents/chunk-000/<cam_key>/episode_{idx:06d}_{start}_{end}.pth
	Each .pth is a dict:
	latent (FHW, 48) bf16, latent_num_frames/height/width, video_num_frames/height/width,
	text_emb (512, 4096) bf16, text, frame_ids (np.int), start_frame, end_frame, fps, ori_fps
	Also writes <root>/empty_emb.pt (UMT5 embedding of "") for classifier-free guidance.

	Wan2.2-TI2V VAE: 16x spatial, 4x temporal; latents normalized as (z - latents_mean)/latents_std.
	Video frames are truncated to the largest (4k+1) count so temporal compression is exact.

	Usage:
	# Verify the pipeline reproduces the released libero-long episode-0 latent, then stop:
	python extract_latents.py --verify_episode0

	# Extract a whole dataset:
	python extract_latents.py \
	--lerobot_root data/libero-spatial-lerobot \
	--wan_dir ckpts/Wan2.2-TI2V-5B-Diffusers \
	--cam_keys observation.images.agentview_rgb observation.images.eye_in_hand_rgb
	"""

	import argparse
	import importlib.machinery
	import json
	import sys
	import types
	from pathlib import Path

	import numpy as np
	import torch


	def _stub_flash_attn():
	"""Allow importing diffusers without a built flash-attn -- but prefer the real package.

	Bug fixed here (hit on job 446): the original version always stubbed both module names
	unconditionally, even when a real, working flash-attn was installed. That stub only defines
	flash_attn_func, not flash_attn_varlen_func -- so as soon as diffusers' own internal
	`from flash_attn import flash_attn_func, flash_attn_varlen_func` (attention_dispatch.py, hit
	just from `from diffusers import AutoencoderKLWan`, nothing to do with the transformer/
	flash-attn kernel this script never uses) ran, it got the stub instead of the real package
	and crashed with ImportError. Now we only stub a name if it genuinely isn't importable.
	"""
	for name in ("flash_attn", "flash_attn_interface"):
	if name in sys.modules:
	continue
	try:
	importlib.import_module(name)
	continue
	except ImportError:
	pass
	mod = types.ModuleType(name)
	mod.flash_attn_func = lambda a, *k: None
	mod.__spec__ = importlib.machinery.ModuleSpec(name, None)
	mod.__version__ = "0.0.0-stub"
	sys.modules[name] = mod


	def largest_4k_plus_1(n: int) -> int:
	"""Largest m <= n with m % 4 == 1 (so temporal compression (m-1)/4+1 is exact)."""
	m = n - ((n - 1) % 4)
	return max(m, 1)


	def read_video_frames(mp4_path, target_hw):
	"""Return uint8 array [T, H, W, 3] resized to target_hw."""
	import av
	import cv2
	container = av.open(str(mp4_path))
	frames = [f.to_ndarray(format="rgb24") for f in container.decode(video=0)]
	container.close()
	frames = np.stack(frames) # [T,H,W,3]
	H, W = target_hw
	if frames.shape[1] != H or frames.shape[2] != W:
	frames = np.stack([cv2.resize(f, (W, H), interpolation=cv2.INTER_AREA) for f in frames])
	return frames


	class Encoders:
	def __init__(self, wan_dir, device="cuda", dtype=torch.bfloat16):
	from diffusers import AutoencoderKLWan
	from transformers import UMT5EncoderModel, T5TokenizerFast
	self.device, self.dtype = device, dtype
	self.vae = AutoencoderKLWan.from_pretrained(wan_dir, subfolder="vae", torch_dtype=torch.float32).to(device).eval()
	self.text_encoder = UMT5EncoderModel.from_pretrained(wan_dir, subfolder="text_encoder", torch_dtype=dtype).to(device).eval()
	self.tokenizer = T5TokenizerFast.from_pretrained(wan_dir, subfolder="tokenizer")
	cfg = self.vae.config
	self.mean = torch.tensor(cfg.latents_mean, dtype=torch.float32, device=device).view(1, -1, 1, 1, 1)
	self.std = torch.tensor(cfg.latents_std, dtype=torch.float32, device=device).view(1, -1, 1, 1, 1)
	self.text_max_len = 512

	@torch.no_grad()
	def encode_video(self, frames_uint8):
	"""frames_uint8 [T,H,W,3] -> normalized latent [C,f,h,w] float32, plus dims."""
	T = largest_4k_plus_1(frames_uint8.shape[0])
	frames = frames_uint8[:T]
	# [T,H,W,3] uint8 -> [1,3,T,H,W] in [-1,1]
	x = torch.from_numpy(frames).to(self.device).float().div(127.5).sub(1.0)
	x = x.permute(3, 0, 1, 2).unsqueeze(0) # [1,3,T,H,W]
	z = self.vae.encode(x).latent_dist.mode() # [1,48,f,h,w]
	z = (z - self.mean) / self.std # Wan normalization
	z = z[0] # [48,f,h,w]
	return z.cpu(), frames_uint8.shape[0], frames.shape[1], frames.shape[2]

	@torch.no_grad()
	def encode_text(self, text):
	tok = self.tokenizer([text], padding="max_length", max_length=self.text_max_len,
	truncation=True, return_tensors="pt").to(self.device)
	emb = self.text_encoder(input_ids=tok.input_ids, attention_mask=tok.attention_mask).last_hidden_state
	return emb[0].to(self.dtype).cpu() # [512,4096]


	def make_latent_dict(z_cfhw, video_num_frames, vheight, vwidth, text_emb, text, start, end, fps, ori_fps):
	C, f, h, w = z_cfhw.shape
	latent_flat = z_cfhw.permute(1, 2, 3, 0).reshape(f * h * w, C).contiguous().to(torch.bfloat16) # (f h w) c
	used = largest_4k_plus_1(video_num_frames)
	return {
	"latent": latent_flat,
	"latent_num_frames": int(f), "latent_height": int(h), "latent_width": int(w),
	"video_num_frames": int(used), "video_height": int(vheight), "video_width": int(vwidth),
	"text_emb": text_emb, "text": text,
	"frame_ids": np.arange(used, dtype=np.int64),
	"start_frame": int(start), "end_frame": int(end),
	"fps": int(fps), "ori_fps": int(ori_fps),
	}


	def verify_episode0(args):
	"""Extract libero-long episode-0 and compare to the released latent to confirm conventions."""
	_stub_flash_attn()
	root = Path(args.verify_root)
	ref = torch.load(root / "latents/chunk-000/observation.images.agentview_rgb/episode_000000_0_388.pth",
	weights_only=False)
	enc = Encoders(args.wan_dir, device=args.device)
	frames = read_video_frames(root / "videos/chunk-000/observation.images.agentview_rgb/episode_000000.mp4",
	(ref["video_height"], ref["video_width"]))
	z, vnf, vh, vw = enc.encode_video(frames)
	print(f"[verify] my latent shape (C,f,h,w)={tuple(z.shape)} ; released flat={tuple(ref['latent'].shape)} "
	f"num_frames={ref['latent_num_frames']} h={ref['latent_height']} w={ref['latent_width']}")
	mine = z.permute(1, 2, 3, 0).reshape(-1, z.shape[0]).float()
	rel = ref["latent"].float()
	if mine.shape == rel.shape:
	err = (mine - rel).abs()
	corr = torch.corrcoef(torch.stack([mine.flatten(), rel.flatten()]))[0, 1]
	print(f"[verify] shape MATCH \| mean\|abs err\|={err.mean():.4f} max={err.max():.4f} corr={corr:.4f}")
	print("[verify] -> PASS (normalization/convention confirmed)" if corr > 0.98
	else "[verify] -> MISMATCH: inspect VAE version / normalization / frame sampling")
	else:
	print("[verify] shape MISMATCH -> check spatial/temporal compression or frame truncation")


	def extract_dataset(args):
	_stub_flash_attn()
	root = Path(args.lerobot_root)
	meta = [json.loads(l) for l in (root / "meta/episodes.jsonl").read_text().splitlines()]
	info = json.load(open(root / "meta/info.json"))
	fps = int(info.get("fps", 60))
	enc = Encoders(args.wan_dir, device=args.device)

	# empty text embedding for CFG
	torch.save(enc.encode_text(""), root / "empty_emb.pt")
	print(f"[extract] wrote {root/'empty_emb.pt'}")

	for ep in meta:
	idx = ep["episode_index"]
	chunk = idx // int(info.get("chunks_size", 1000))
	for seg in ep["action_config"]:
	start, end, text = seg["start_frame"], seg["end_frame"], seg["action_text"]
	text_emb = enc.encode_text(text)
	for cam in args.cam_keys:
	mp4 = root / f"videos/chunk-{chunk:03d}/{cam}/episode_{idx:06d}.mp4"
	frames = read_video_frames(mp4, (info["features"][cam]["shape"][0], info["features"][cam]["shape"][1]))
	frames = frames[start:end]
	z, vnf, vh, vw = enc.encode_video(frames)
	d = make_latent_dict(z, frames.shape[0], vh, vw, text_emb, text, start, end, fps, fps)
	out = root / f"latents/chunk-{chunk:03d}/{cam}/episode_{idx:06d}_{start}_{end}.pth"
	out.parent.mkdir(parents=True, exist_ok=True)
	torch.save(d, out)
	print(f"[extract] episode {idx} done")


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--wan_dir", default="ckpts/Wan2.2-TI2V-5B-Diffusers")
	ap.add_argument("--device", default="cuda")
	ap.add_argument("--verify_episode0", action="store_true",
	help="Reproduce released libero-long episode-0 latent and compare (no dataset written)")
	ap.add_argument("--verify_root", default="data/libero-long-lerobot")
	ap.add_argument("--lerobot_root", default=None)
	ap.add_argument("--cam_keys", nargs="+",
	default=["observation.images.agentview_rgb", "observation.images.eye_in_hand_rgb"])
	args = ap.parse_args()
	if args.verify_episode0:
	verify_episode0(args)
	else:
	assert args.lerobot_root, "--lerobot_root required for extraction"
	extract_dataset(args)


	if __name__ == "__main__":
	main()