Spaces:
Sleeping
Sleeping
| """Offline prebuild of the whole-clip VAE latent cache. | |
| Training encodes the first latents of a window fresh and slices the tail from this cache. | |
| Usage (single or multi-GPU sharding, resumable): | |
| PYTHONPATH=. python scripts/tools/precache_vae_latents.py \ | |
| --config configs/stage1_pretrain_bidir.yaml [--rank 0 --world 2] [--device cuda:0] [--limit N] | |
| Pixel processing mirrors the dataset: uniform_sample on the 24fps grid -> ToTensor -> | |
| bicubic(antialias) resize (H,W) → [-1,1] → StreamingVAEEncoder(chunk 33)。 | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| sys.path.insert(0, os.getcwd()) | |
| from alaya.config.loader import load_config | |
| from alaya.data.vae_latent_cache import entry_paths, save_entry | |
| from alaya.model.loader import load_vae | |
| from fastvideo.ltx2_streaming_vae import StreamingVAEEncoder | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--config", required=True) | |
| ap.add_argument("--rank", type=int, default=int(os.environ.get("RANK", 0))) | |
| ap.add_argument("--world", type=int, default=int(os.environ.get("WORLD_SIZE", 1))) | |
| ap.add_argument("--device", default=None) | |
| ap.add_argument("--limit", type=int, default=None, help="only process the first N videos (smoke test)") | |
| ap.add_argument("--batch_chunk", type=int, default=33) | |
| args = ap.parse_args() | |
| cfg = load_config(args.config) | |
| cache_dir = cfg.runtime.vae_latent_cache_dir | |
| assert cache_dir, "runtime.vae_latent_cache_dir must be set in the config" | |
| dev = torch.device(args.device or f"cuda:{args.rank % max(1, torch.cuda.device_count())}") | |
| dtype = torch.bfloat16 | |
| H, W = int(cfg.sample.height), int(cfg.sample.width) | |
| fps = float(cfg.sample.fps) | |
| stride = int(cfg.sample.temporal_stride) | |
| # Enumerate unique videos from the training dataset so source/video_id match training | |
| from alaya.data.dataloader import build_train_dataloader | |
| from alaya.utils.distributed import DistributedState | |
| dist_state = DistributedState(rank=0, local_rank=0, world_size=1, device=torch.device("cpu")) | |
| loader = build_train_dataloader(cfg, dist_state) | |
| ds = loader.dataset | |
| subsets = [ds] if hasattr(ds, "samples") else list(getattr(ds, "datasets", [])) | |
| videos: dict[tuple[str, str], str] = {} | |
| for sub in subsets: | |
| for video_path, _cap, _pose, source, video_id in sub.samples: | |
| videos.setdefault((str(source), str(video_id)), str(video_path)) | |
| items = sorted(videos.items()) | |
| if args.limit: | |
| items = items[: args.limit] | |
| missing = [ | |
| (k, v) for k, v in items | |
| if not os.path.exists(entry_paths(cache_dir, k[0], k[1], H, W, fps)[0]) | |
| ] | |
| shard = missing[args.rank :: max(1, args.world)] | |
| print( | |
| f"[VAECache] unique videos {len(items)}, already cached {len(items) - len(missing)}, " | |
| f"this rank ({args.rank}/{args.world}) will encode {len(shard)} (dir={cache_dir})", | |
| flush=True, | |
| ) | |
| if not shard: | |
| return | |
| vae_enc_raw, _dec = load_vae(cfg.paths.vae, device=dev, dtype=dtype) | |
| enc = StreamingVAEEncoder(vae_enc_raw, device=dev, dtype=dtype) | |
| import decord # noqa: E402 - imported after torch to avoid a thread conflict | |
| t0 = time.time() | |
| done = 0 | |
| total = len(shard) | |
| report_every = max(1, total // 100) | |
| for (source, video_id), video_path in shard: | |
| try: | |
| vr = decord.VideoReader(video_path, num_threads=2) | |
| video_fps = float(vr.get_avg_fps()) # raw value, matching the dataset fps ratio (no rounding) | |
| ratio = video_fps / fps | |
| total_src = len(vr) | |
| # Largest 24fps grid length n with round((n-1)*ratio) inside the clip and (n-1)%stride==0 | |
| n = int((total_src - 1) / ratio) + 1 | |
| n = ((n - 1) // stride) * stride + 1 | |
| if n < stride + 1: | |
| done += 1 | |
| continue | |
| ids = [min(int(j * ratio), total_src - 1) for j in range(n)] | |
| frames = vr.get_batch(ids).asnumpy() # [T,H0,W0,C] uint8 | |
| px = torch.from_numpy(frames).float().div_(255.0).permute(0, 3, 1, 2) # [T,C,H0,W0] | |
| px = torch.nn.functional.interpolate( | |
| px, size=(H, W), mode="bicubic", align_corners=False, antialias=True | |
| ) | |
| px = px.sub_(0.5).div_(0.5).permute(1, 0, 2, 3).to(dev, dtype) # [C,T,H,W] | |
| with torch.no_grad(): | |
| latent = enc.encode(px.unsqueeze(0), chunk_size=args.batch_chunk, verbose=False)[0] | |
| save_entry(cache_dir, source, video_id, H, W, fps, latent, ratio) | |
| del px, latent, frames | |
| except Exception as e: # one failed video must not abort the shard | |
| print(f"[VAECache] SKIP {source}/{video_id}: {type(e).__name__} {e}", flush=True) | |
| done += 1 | |
| if done % report_every == 0 or done == total: | |
| el = time.time() - t0 | |
| ips = done / max(el, 1e-6) | |
| eta = (total - done) / max(ips, 1e-6) | |
| bar_n = int(30 * done / total) | |
| print( | |
| f"[VAECache] rank{args.rank} |{'█' * bar_n}{'.' * (30 - bar_n)}| " | |
| f"{done}/{total} ({100 * done / total:.1f}%) {ips:.2f} vid/s ETA {eta / 3600:.2f}h", | |
| flush=True, | |
| ) | |
| print(f"[VAECache] rank{args.rank} done in {(time.time() - t0) / 60:.1f}min", flush=True) | |
| if __name__ == "__main__": | |
| main() | |