VBVR Latent Cache (832×832 × 33f, Wan2.2-TI2V-5B VAE + UMT5-XXL)

Pre-encoded latent cache for the Video-Reason/VBVR-Dataset geometric / logical reasoning video corpus, prepared for Equilibrium Matching (EqM) post-training of Wan-AI/Wan2.2-TI2V-5B-Diffusers on AWS Trainium2.

This is a working cache, not a primary dataset. It exists to skip the ~5 s/sample VAE+T5 encode cost during training. The original videos + prompts live in the upstream VBVR-Dataset repo.

Source → Tensor pipeline

For each VBVR sample (one directory inside an upstream tar shard), we transform two raw inputs into two encoded tensors. Nothing else from the sample dir is used (we ignore first_frame.png, final_frame.png, metadata.json).

─── Per-sample inputs (from upstream VBVR-Dataset) ────────────────────────
  ground_truth.mp4 : 1024×1024 RGB video, 36–96 frames @ 16 fps
  prompt.txt       : English instruction describing the task

─── Step 1: Decode mp4 ─────────────────────────────────────────────────────
  imageio.v3.imiter(mp4_bytes)
    → np.ndarray (T_src, 1024, 1024, 3) uint8           # all frames

─── Step 2: Uniform temporal subsample ─────────────────────────────────────
  idx = np.linspace(0, T_src - 1, 33).round()           # full-clip span
  arr = arr[idx]
    → np.ndarray (33, 1024, 1024, 3) uint8              # 33 frames spanning full clip

─── Step 3: Normalize + permute ────────────────────────────────────────────
  vid = torch.from_numpy(arr).float() / 127.5 - 1.0     # → [-1, 1]
  vid = vid.permute(0, 3, 1, 2)                         # (T, C, H, W)
    → torch.Tensor (33, 3, 1024, 1024) float32

─── Step 4: Bilinear spatial resize ────────────────────────────────────────
  vid = F.interpolate(vid, size=(832, 832),
                      mode="bilinear", align_corners=False)
    → torch.Tensor (33, 3, 832, 832) float32            # square downsample

─── Step 5: Add batch dim, channels-before-time ────────────────────────────
  vid = vid.permute(1, 0, 2, 3).unsqueeze(0).contiguous()
    → torch.Tensor (1, 3, 33, 832, 832) float32         # Wan VAE input shape

─── Step 6: VAE encode → normalize → cast bf16 ─────────────────────────────
  x = vid.to(neuron, bf16)                              # CPU → Neuron device, bf16
  posterior = vae.encode(x).latent_dist                 # AutoencoderKLWan
  z = posterior.mode()                                  # deterministic (NOT .sample())
  z = z.to(fp32).cpu()
  z = (z - latents_mean) / latents_std                  # per-channel from vae.config
  x0 = z.squeeze(0).bf16().contiguous()
    → torch.Tensor x0 = (48, 9, 52, 52) bf16            # latent: C=48, T_lat=9, H/W_lat=52
                                                        #   T_lat = 1 + (33-1)/4
                                                        #   H_lat = 832 / 16 (Wan VAE 16× spatial)

─── Step 7: Tokenize prompt ────────────────────────────────────────────────
  tok = tokenizer([prompt], padding="max_length",
                  max_length=512, truncation=True,
                  add_special_tokens=True)
    → ids (1, 512) int64, mask (1, 512) int64

─── Step 8: UMT5-XXL encode → cast bf16 ────────────────────────────────────
  out = text_encoder(input_ids=ids, attention_mask=mask).last_hidden_state
  out = out.to(fp32).cpu().to(bf16).contiguous()
  text_embeds = out[0]                                  # drop batch dim
    → torch.Tensor text_embeds = (512, 4096) bf16

─── Step 9: Save dict via torch.save ───────────────────────────────────────
  torch.save({"x0": ..., "text_embeds": ..., "prompt": ..., ...},
             out_path)
    → file: <task_family>/<sample_id>.pt   ~6.24 MiB

What's in each .pt

import torch
sample = torch.load(
    "G-11_handle_object_reappearance_data-generator/handle_object_reappearance_00000023.pt",
    weights_only=False,
)
sample.keys()
# dict_keys(['x0', 'text_embeds', 'prompt', 'source_shape', 'latent_shape'])

Per-sample disk size: ~6.24 MiB (latent ≈ 2.23 MiB + text embeds ≈ 4.0 MiB

• small dict overhead).

Models / configs used

What's deliberately NOT done

• No augmentation. No random crop, flip, color jitter, etc. The video → latent is a deterministic function of the source mp4.
• No noise added. EqM-specific noise scheduling happens at training time, not at cache time.
• No frame extrapolation / interpolation. Just uniform-spaced subsample.
• No aspect-ratio preservation. Source is 1024×1024 (1:1); output is 832×832 (1:1). Bilinear resize is identity-aspect.
• No per-task-family weighting. Each task family contributes proportional to its source tar size.

Layout

<repo root>/
├── README.md                                            ← this file
├── G-11_handle_object_reappearance_data-generator/      ← one dir per task family
│   ├── handle_object_reappearance_00000023.pt
│   ├── handle_object_reappearance_00000156.pt
│   └── ... (up to ~5000 per task family at the 50% target)
├── G-12_grid_obtaining_award_data-generator/
├── ...
├── O-87_fluid_diffusion_reasoning_data-generator/
└── O-8_shape_rotation_data-generator/

100 task families total (matches the 100 tar shards in upstream VBVR-Dataset). Filenames are the original VBVR sample directory IDs from inside each tar (no renaming, deterministic mapping).

Status snapshot at upload time

Upload is incremental — re-running hf upload-large-folder only sends new files (content-hash dedup). If preprocessing later resumes, this repo will be re-synced without overwriting existing files.

Estimated full sizes

Caveat for HF Hub users: full snapshots can be very large. For training on a subset, prefer:

hf download Central-Cat/vbvr-latent-cache-832x832x33f \
    --repo-type dataset \
    --include "G-11_*"          # one task family
    --local-dir /your/local/dir

Reproducing the encode pipeline

The encoder script lives at projects/Video_Energy_Model/training/precompute_latents.py in the Neuron_Knowledge_Base repo (branch Trainium-48xlarge).

Single-sample dry-run (prints shapes, doesn't write):

source /home/ubuntu/workspace/native_venv/bin/activate
python -m projects.Video_Energy_Model.training.precompute_latents \
    --device neuron --dtype bf16 \
    --height 832 --width 832 --num-frames 33 \
    --shard-filter G-11 --max-samples-per-tar 1 --dry-run

Real run on Trainium2 (16 worker, full coverage):

projects/Video_Energy_Model/tools/setup_dual_disk_cache.sh
projects/Video_Energy_Model/tools/preprocess_50pct.sh

License

Inherits the upstream VBVR-Dataset license (Apache-2.0). The encode scripts in the linked repo are MIT.