#!/usr/bin/env python3 """Build the fixed DINOv3 ViT-L/16 whole-video Nyström reference. The teacher and objective are unchanged from the paired V-JEPA2 experiment: the final latent from every row of the 1,537-video CausVid ODE LMDB is decoded with the frozen Wan VAE, then represented by a frozen DINOv3 image encoder. The DINOv3 feature map is fixed to 16 uniformly distributed source frames and an 8x8 subset of the 16x16 patch grid per frame. Frame-major concatenation therefore produces one aligned ``[16 * 64, 1024]`` matrix per video. Every torchrun rank atomically writes one indexed FP16 feature shard. Rank 0 validates and merges all shards in LMDB-index order, constructs the existing video-matrix Nyström objective (m=256, seed=0), atomically publishes the pack, and reloads the *final path* for validation. Temporary shards are removed only after that final validation succeeds; failed runs retain them for diagnosis. """ from __future__ import annotations import argparse from datetime import timedelta import gc import hashlib import json import math import os from pathlib import Path import shutil import subprocess import sys import tempfile import time from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Tuple import torch import torch.distributed as dist REPO_ROOT = Path(__file__).resolve().parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) TOOL_FORMAT_VERSION = 1 MIB = 1024 ** 2 EXPECTED_TEACHER_COUNT = 1537 EXPECTED_SOURCE_COMMIT = "6876159a11b4df116f30f667f8c9888617df0751" SOURCE_NUM_FRAMES = 81 FRAME_INDICES = (0, 5, 11, 16, 21, 27, 32, 37, 43, 48, 53, 59, 64, 69, 75, 80) INPUT_SIZE = 256 PATCH_SIZE = 16 PATCH_GRID_SIZE = 16 SPATIAL_AXIS_INDICES = (0, 2, 4, 6, 9, 11, 13, 15) SPATIAL_TOKEN_INDICES = tuple( row * PATCH_GRID_SIZE + column for row in SPATIAL_AXIS_INDICES for column in SPATIAL_AXIS_INDICES ) PER_FRAME_TOKENS = 64 TOTAL_TOKENS = len(FRAME_INDICES) * PER_FRAME_TOKENS FEATURE_DIM = 1024 FORWARD_CHUNK_FRAMES = 4 NUM_LANDMARKS = 256 NYSTROM_SEED = 0 MEDIAN_MAX_SAMPLES = EXPECTED_TEACHER_COUNT KMEANS_ITERATIONS = 20 NYSTROM_RIDGE = 1.0e-6 KERNEL_CHUNK_SIZE = 32 WAN_VAE_MODEL_NAME = "Wan2.1-T2V-1.3B" WAN_VAE_MEAN = ( -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921, ) WAN_VAE_STD = ( 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160, ) def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description=( "Build the fixed full-1537 DINOv3 ViT-L/16 video-matrix " "Nyström teacher reference." ) ) parser.add_argument("--teacher-lmdb", type=Path, required=True) parser.add_argument("--repo-path", type=Path, required=True) parser.add_argument("--weight-path", type=Path, required=True) parser.add_argument( "--source-commit", type=str, default=EXPECTED_SOURCE_COMMIT) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--batch-size", type=int, default=1) parser.add_argument( "--expected-teacher-count", type=int, default=EXPECTED_TEACHER_COUNT) parser.add_argument( "--forward-chunk-frames", type=int, default=FORWARD_CHUNK_FRAMES) parser.add_argument("--num-landmarks", type=int, default=NUM_LANDMARKS) parser.add_argument("--seed", type=int, default=NYSTROM_SEED) parser.add_argument( "--median-max-samples", type=int, default=MEDIAN_MAX_SAMPLES) parser.add_argument("--bandwidth", type=float, default=None) parser.add_argument( "--kmeans-iterations", type=int, default=KMEANS_ITERATIONS) parser.add_argument("--ridge", type=float, default=NYSTROM_RIDGE) parser.add_argument("--chunk-size", type=int, default=KERNEL_CHUNK_SIZE) parser.add_argument("--tmp-dir", type=Path, default=None) parser.add_argument("--overwrite", action="store_true") parser.add_argument("--log-every", type=int, default=25) parser.add_argument("--distributed-timeout-minutes", type=int, default=360) parser.add_argument("--dry-run", action="store_true") return parser def _validate_args(args: argparse.Namespace) -> None: positive = { "batch_size": args.batch_size, "expected_teacher_count": args.expected_teacher_count, "forward_chunk_frames": args.forward_chunk_frames, "num_landmarks": args.num_landmarks, "median_max_samples": args.median_max_samples, "kmeans_iterations": args.kmeans_iterations, "chunk_size": args.chunk_size, "distributed_timeout_minutes": args.distributed_timeout_minutes, } invalid = {name: value for name, value in positive.items() if int(value) <= 0} if invalid: raise ValueError(f"arguments must be positive: {invalid}") if int(args.expected_teacher_count) != EXPECTED_TEACHER_COUNT: raise ValueError( "the canonical reference requires exactly " f"{EXPECTED_TEACHER_COUNT} teachers") if int(args.forward_chunk_frames) != FORWARD_CHUNK_FRAMES: raise ValueError( f"--forward-chunk-frames must be {FORWARD_CHUNK_FRAMES}") if int(args.num_landmarks) != NUM_LANDMARKS: raise ValueError(f"--num-landmarks must be {NUM_LANDMARKS}") if int(args.seed) != NYSTROM_SEED: raise ValueError(f"--seed must be {NYSTROM_SEED}") if int(args.median_max_samples) != MEDIAN_MAX_SAMPLES: raise ValueError( f"--median-max-samples must be {MEDIAN_MAX_SAMPLES}") if args.bandwidth is not None: raise ValueError( "the canonical reference requires median-estimated bandwidth") if int(args.kmeans_iterations) != KMEANS_ITERATIONS: raise ValueError( f"--kmeans-iterations must be {KMEANS_ITERATIONS}") if not math.isclose( float(args.ridge), NYSTROM_RIDGE, rel_tol=0.0, abs_tol=0.0 ): raise ValueError(f"--ridge must be {NYSTROM_RIDGE}") if int(args.chunk_size) != KERNEL_CHUNK_SIZE: raise ValueError(f"--chunk-size must be {KERNEL_CHUNK_SIZE}") if str(args.source_commit) != EXPECTED_SOURCE_COMMIT: raise ValueError( "--source-commit must pin the approved DINOv3 source commit " f"{EXPECTED_SOURCE_COMMIT}") if float(args.ridge) <= 0.0 or not math.isfinite(float(args.ridge)): raise ValueError("--ridge must be finite and positive") if args.bandwidth is not None and ( float(args.bandwidth) <= 0.0 or not math.isfinite(float(args.bandwidth)) ): raise ValueError("--bandwidth must be finite and positive") if int(args.log_every) < 0: raise ValueError("--log-every must be non-negative") def _git_head(repo_path: Path) -> str: try: result = subprocess.run( ["git", "-C", str(repo_path), "rev-parse", "HEAD"], check=True, capture_output=True, text=True, ) except (OSError, subprocess.CalledProcessError) as exc: raise RuntimeError( f"failed to inspect DINOv3 source checkout {repo_path}") from exc return result.stdout.strip() def _git_source_status(repo_path: Path) -> str: try: result = subprocess.run( [ "git", "-C", str(repo_path), "status", "--porcelain", "--untracked-files=all", ], check=True, capture_output=True, text=True, ) except (OSError, subprocess.CalledProcessError) as exc: raise RuntimeError( f"failed to inspect DINOv3 source worktree {repo_path}") from exc return result.stdout def _inspect_assets(args: argparse.Namespace) -> Dict[str, Any]: repo_path = args.repo_path.expanduser().resolve() weight_path = args.weight_path.expanduser().resolve() if not repo_path.is_dir(): raise FileNotFoundError(f"DINOv3 source directory not found: {repo_path}") if not weight_path.is_file(): raise FileNotFoundError(f"DINOv3 weight file not found: {weight_path}") source_commit = _git_head(repo_path) if source_commit != str(args.source_commit): raise ValueError( "DINOv3 source checkout commit mismatch: " f"{source_commit} != {args.source_commit}") source_status = _git_source_status(repo_path) if source_status.strip(): raise ValueError( "DINOv3 source checkout has worktree changes; refusing " "to build a reference from unpinned source") stat = weight_path.stat() return { "repo_path": str(repo_path), "weight_path": str(weight_path), "source_commit": source_commit, "weight_size_bytes": int(stat.st_size), "weight_mtime_ns": int(stat.st_mtime_ns), } def _inspect_lmdb(path: Path) -> Dict[str, Any]: try: import lmdb except ImportError as exc: raise ImportError( "python-lmdb is required; use the same Python environment as training" ) from exc from utils.lmdb import get_array_shape_from_lmdb resolved = path.expanduser().resolve() if not resolved.is_dir(): raise FileNotFoundError(f"teacher LMDB directory not found: {resolved}") env = lmdb.open( str(resolved), readonly=True, lock=False, readahead=False, meminit=False) try: latents_shape = tuple( int(value) for value in get_array_shape_from_lmdb(env, "latents")) prompts_shape = tuple( int(value) for value in get_array_shape_from_lmdb(env, "prompts")) finally: env.close() if len(latents_shape) != 6: raise ValueError( "teacher LMDB latents must be [N,trajectory,T,C,H,W], got " f"{latents_shape}") if prompts_shape != (latents_shape[0],): raise ValueError( f"teacher prompts/latents count mismatch: {prompts_shape} vs " f"{latents_shape}") data_path = resolved / "data.mdb" stat = data_path.stat() return { "path": str(resolved), "latents_shape": latents_shape, "prompts_shape": prompts_shape, "data_mdb_size_bytes": int(stat.st_size), "data_mdb_mtime_ns": int(stat.st_mtime_ns), } def _rank_indices(count: int, rank: int, world_size: int) -> range: return range(int(rank), int(count), int(world_size)) def _required_temp_bytes(count: int) -> int: feature_bytes = int(count) * TOTAL_TOKENS * FEATURE_DIM * 2 return int(math.ceil(feature_bytes * 1.15)) + 64 * MIB def _required_output_bytes() -> int: landmark_bytes = NUM_LANDMARKS * TOTAL_TOKENS * FEATURE_DIM * 4 return int(math.ceil(landmark_bytes * 1.25)) + 32 * MIB def _choose_temp_parent( requested: Optional[Path], required_bytes: int ) -> Tuple[Path, int]: candidates: Sequence[Path] = ( (requested.expanduser(),) if requested is not None else (Path("/dev/shm"), Path(tempfile.gettempdir())) ) errors = [] for candidate in candidates: try: candidate = candidate.resolve() if not candidate.is_dir(): errors.append(f"{candidate}: not a directory") continue if not os.access(candidate, os.W_OK | os.X_OK): errors.append(f"{candidate}: not writable") continue free = int(shutil.disk_usage(candidate).free) if free < required_bytes: errors.append( f"{candidate}: {free / MIB:.1f} MiB free, " f"{required_bytes / MIB:.1f} MiB required") continue return candidate, free except OSError as exc: errors.append(f"{candidate}: {exc}") raise RuntimeError( "no temporary shard directory has sufficient capacity: " + "; ".join(errors)) def _check_output( output: Path, *, required_bytes: int, overwrite: bool ) -> Dict[str, Any]: output = output.expanduser().resolve() if output.exists() and not overwrite: raise FileExistsError( f"output already exists: {output}; pass --overwrite to replace it") existing = output.parent while not existing.exists() and existing != existing.parent: existing = existing.parent if not existing.is_dir(): raise FileNotFoundError( f"no existing output parent for capacity check: {output.parent}") free = int(shutil.disk_usage(existing).free) if free < required_bytes: raise RuntimeError( f"output filesystem has only {free / MIB:.1f} MiB free; " f"{required_bytes / MIB:.1f} MiB required") return { "path": str(output), "capacity_parent": str(existing), "free_bytes": free, "required_bytes": required_bytes, } def _plan(args: argparse.Namespace, world_size: int) -> Dict[str, Any]: lmdb = _inspect_lmdb(args.teacher_lmdb) count = int(lmdb["latents_shape"][0]) if count != EXPECTED_TEACHER_COUNT: raise ValueError( "canonical DINOv3 reference requires the complete 1,537-row LMDB; " f"found {count} rows") if int(args.num_landmarks) > count: raise ValueError("number of landmarks exceeds teacher population") assets = _inspect_assets(args) temp_required = _required_temp_bytes(count) temp_parent, temp_free = _choose_temp_parent(args.tmp_dir, temp_required) output = _check_output( args.output, required_bytes=_required_output_bytes(), overwrite=bool(args.overwrite), ) return { "tool_format_version": TOOL_FORMAT_VERSION, "teacher_lmdb": lmdb, "assets": assets, "selected_samples": count, "world_size": int(world_size), "samples_per_rank": [ len(_rank_indices(count, rank, world_size)) for rank in range(world_size) ], "feature_geometry": { "source_num_frames": SOURCE_NUM_FRAMES, "frame_indices": list(FRAME_INDICES), "input_size": INPUT_SIZE, "patch_size": PATCH_SIZE, "patch_grid_size": PATCH_GRID_SIZE, "spatial_axis_indices": list(SPATIAL_AXIS_INDICES), "spatial_token_indices": list(SPATIAL_TOKEN_INDICES), "raw_tokens_per_frame": PATCH_GRID_SIZE ** 2, "per_frame_tokens": PER_FRAME_TOKENS, "num_tokens": TOTAL_TOKENS, "feature_dim": FEATURE_DIM, "sample_shape": [TOTAL_TOKENS, FEATURE_DIM], "concat_order": "frame_major_then_spatial_row_major", }, "num_landmarks": NUM_LANDMARKS, "temporary": { "parent": str(temp_parent), "required_bytes": temp_required, "free_bytes": temp_free, }, "output": output, } def _setup_distributed(timeout_minutes: int) -> Tuple[torch.device, int, int]: world_size = int(os.environ.get("WORLD_SIZE", "1")) rank = int(os.environ.get("RANK", "0")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) if not torch.cuda.is_available(): raise RuntimeError( "reference extraction requires CUDA; use --dry-run on CPU") torch.cuda.set_device(local_rank) if world_size > 1: dist.init_process_group( backend="nccl", timeout=timedelta(minutes=int(timeout_minutes)), ) rank = dist.get_rank() world_size = dist.get_world_size() return torch.device("cuda", local_rank), rank, world_size def _barrier() -> None: if dist.is_available() and dist.is_initialized(): dist.barrier() def _broadcast_rank0_object(value: Any, rank: int) -> Any: if not (dist.is_available() and dist.is_initialized()): return value payload = [value if rank == 0 else None] dist.broadcast_object_list(payload, src=0) return payload[0] def _atomic_torch_save(value: Any, path: Path) -> None: temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}") try: torch.save(value, temporary) os.replace(temporary, path) finally: if temporary.exists(): temporary.unlink() def _canonical_sha256(value: Dict[str, Any]) -> str: payload = json.dumps( value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, allow_nan=False, ).encode("utf-8") return hashlib.sha256(payload).hexdigest() def _file_identity(path: Path) -> Dict[str, Any]: resolved = path.expanduser().resolve(strict=True) if not resolved.is_file(): raise ValueError(f"identity target is not a regular file: {resolved}") before = resolved.stat() digest = hashlib.sha256() with resolved.open("rb") as handle: while True: block = handle.read(8 << 20) if not block: break digest.update(block) after = resolved.stat() if ( int(before.st_size) != int(after.st_size) or int(before.st_mtime_ns) != int(after.st_mtime_ns) ): raise RuntimeError(f"identity target changed while hashing: {resolved}") return { "path": str(resolved), "size_bytes": int(after.st_size), "mtime_ns": int(after.st_mtime_ns), "sha256": digest.hexdigest(), } def _vae_identity() -> Dict[str, Any]: """Resolve the exact Wan decoder used before DINOv3 extraction.""" from utils.paths import model_path weights = model_path(WAN_VAE_MODEL_NAME) / "Wan2.1_VAE.pth" identity = { "wrapper": "utils.wan_wrapper.WanVAEWrapper", "wrapper_source": _file_identity(REPO_ROOT / "utils" / "wan_wrapper.py"), "model_name": WAN_VAE_MODEL_NAME, "weights": _file_identity(weights), "latent_channels": 16, "latent_mean": list(WAN_VAE_MEAN), "latent_std": list(WAN_VAE_STD), "mean_std_contract": "WanVAEWrapper.inline_constants.v1", "compute_dtype": "bfloat16", "decode_use_cache": False, "decode_output_range": "[-1,1]", } identity["identity_sha256"] = _canonical_sha256(identity) return identity def _validate_reference_current_assets( pack: Mapping[str, Any], *, repo_path: Path, weight_path: Path, teacher_lmdb: Path, ) -> Dict[str, Any]: """Require a validated reference to match every current production asset. Callers must first run ``load_reference_pack`` (which verifies schema and content hash). This second gate prevents a valid but stale pack from being accepted after the DINO checkpoint, Wan VAE, or teacher LMDB changes. """ if not isinstance(pack, Mapping): raise TypeError("reference pack must be a mapping") metadata = pack.get("metadata") if not isinstance(metadata, Mapping): raise ValueError("reference pack metadata is missing") feature = metadata.get("feature") reference = metadata.get("reference") if not isinstance(feature, Mapping) or not isinstance(reference, Mapping): raise ValueError("reference feature/source metadata is missing") source = reference.get("source") saved_vae = reference.get("vae") if not isinstance(source, Mapping) or not isinstance(saved_vae, Mapping): raise ValueError("reference source or Wan VAE identity is missing") from video_fd.encoders.dinov3_dense_metadata import file_sha256 resolved_repo = repo_path.expanduser().resolve(strict=True) resolved_weight = weight_path.expanduser().resolve(strict=True) current_weight_sha256 = file_sha256(resolved_weight) expected_feature_identity = { "repo_path": str(resolved_repo), "weight_path": str(resolved_weight), "weight_sha256": current_weight_sha256, } feature_mismatches = { key: {"saved": feature.get(key), "current": value} for key, value in expected_feature_identity.items() if feature.get(key) != value } if feature_mismatches: raise ValueError( "reference DINOv3 asset identity does not match current assets: " f"{feature_mismatches}") current_vae = _vae_identity() if dict(saved_vae) != current_vae: raise ValueError( "reference Wan VAE identity does not match current assets") current_lmdb = _inspect_lmdb(teacher_lmdb) expected_source_identity = { "lmdb_path": current_lmdb["path"], "latents_shape": list(current_lmdb["latents_shape"]), "prompts_shape": list(current_lmdb["prompts_shape"]), "selected_samples": int(current_lmdb["latents_shape"][0]), "data_mdb_size_bytes": int(current_lmdb["data_mdb_size_bytes"]), "data_mdb_mtime_ns": int(current_lmdb["data_mdb_mtime_ns"]), } source_mismatches = { key: {"saved": source.get(key), "current": value} for key, value in expected_source_identity.items() if source.get(key) != value } if source_mismatches: raise ValueError( "reference teacher LMDB identity does not match current assets: " f"{source_mismatches}") return { "dinov3_weight_sha256": current_weight_sha256, "vae_identity_sha256": current_vae["identity_sha256"], "teacher_lmdb": current_lmdb, } def _prompt_entry_digest(index: int, prompt: str) -> bytes: digest = hashlib.sha256() digest.update(int(index).to_bytes(8, byteorder="little", signed=False)) digest.update(b"\0") digest.update(str(prompt).encode("utf-8")) return digest.digest() def _load_dense_encoder(args: argparse.Namespace, device: torch.device): """The sole runtime boundary with the separately-owned DINOv3 adapter.""" from video_fd.encoders.dinov3_adapter import DINOv3DenseAdapter encoder = DINOv3DenseAdapter({ "name": "dinov3_dense", "repo_path": str(args.repo_path.expanduser().resolve()), "weight_path": str(args.weight_path.expanduser().resolve()), "source_commit": str(args.source_commit), "source_num_frames": SOURCE_NUM_FRAMES, "frame_indices": list(FRAME_INDICES), "input_size": INPUT_SIZE, "patch_size": PATCH_SIZE, "patch_grid_size": PATCH_GRID_SIZE, "spatial_axis_indices": list(SPATIAL_AXIS_INDICES), "forward_chunk_frames": int(args.forward_chunk_frames), "compute_dtype": "float32", }) return encoder.to(device=device, dtype=torch.float32).eval().requires_grad_(False) @torch.no_grad() def _extract_rank_shard( args: argparse.Namespace, plan: Dict[str, Any], *, device: torch.device, rank: int, world_size: int, shard_path: Path, vae_identity: Dict[str, Any], ) -> None: from utils.dataset import ODERegressionLMDBDataset from utils.wan_wrapper import WanVAEWrapper selected_count = int(plan["selected_samples"]) assigned = list(_rank_indices(selected_count, rank, world_size)) dataset = ODERegressionLMDBDataset( str(args.teacher_lmdb.expanduser().resolve()), max_pair=selected_count) if len(dataset) != selected_count: raise RuntimeError( f"LMDB length changed after planning: {len(dataset)} != " f"{selected_count}") vae = ( WanVAEWrapper() .to(device=device, dtype=torch.bfloat16) .eval() .requires_grad_(False) ) if not torch.equal(vae.mean, torch.tensor(WAN_VAE_MEAN)): raise RuntimeError("WanVAEWrapper latent mean changed from the reference contract") if not torch.equal(vae.std, torch.tensor(WAN_VAE_STD)): raise RuntimeError("WanVAEWrapper latent std changed from the reference contract") encoder = _load_dense_encoder(args, device) feature_metadata = encoder.feature_metadata() if not isinstance(feature_metadata, dict): raise TypeError("DINOv3 adapter feature_metadata() must return a dict") spatial_indices = torch.tensor( SPATIAL_TOKEN_INDICES, device=device, dtype=torch.long) feature_chunks = [] index_chunks = [] prompt_digests = [] expected_tail = tuple(plan["teacher_lmdb"]["latents_shape"][1:]) started = time.monotonic() for offset in range(0, len(assigned), int(args.batch_size)): batch_indices = assigned[offset:offset + int(args.batch_size)] items = [dataset[index] for index in batch_indices] trajectories = [item["ode_latent"] for item in items] for index, trajectory in zip(batch_indices, trajectories): if tuple(trajectory.shape) != expected_tail: raise ValueError( f"LMDB row {index} shape {tuple(trajectory.shape)} != " f"{expected_tail}") if not bool(torch.isfinite(trajectory).all().item()): raise FloatingPointError( f"LMDB row {index} contains non-finite latents") latent = torch.stack( [trajectory[-1] for trajectory in trajectories], dim=0 ).to(device=device, dtype=torch.bfloat16) pixels = vae.decode_to_pixel(latent.detach()) video01 = ((pixels + 1.0) * 0.5).clamp(0.0, 1.0) if ( video01.ndim != 5 or int(video01.shape[1]) != SOURCE_NUM_FRAMES or int(video01.shape[2]) != 3 ): raise ValueError( "Wan VAE must decode [B,81,3,H,W], got " f"{tuple(video01.shape)}") features = encoder.forward_dense( video01, token_indices=spatial_indices, gradient_checkpointing=False, ).float() expected_shape = (len(batch_indices), TOTAL_TOKENS, FEATURE_DIM) if tuple(features.shape) != expected_shape: raise RuntimeError( f"DINOv3 features {tuple(features.shape)} != {expected_shape}") if not bool(torch.isfinite(features).all().item()): raise FloatingPointError( f"non-finite DINOv3 features for indices {batch_indices}") features_fp16 = features.to(dtype=torch.float16, device="cpu") if not bool(torch.isfinite(features_fp16).all().item()): raise FloatingPointError( f"FP16 conversion overflowed for indices {batch_indices}") feature_chunks.append(features_fp16) index_chunks.append(torch.tensor(batch_indices, dtype=torch.long)) prompt_digests.extend( _prompt_entry_digest(index, item["prompts"]) for index, item in zip(batch_indices, items) ) del latent, pixels, video01, features, features_fp16, items, trajectories completed = min(offset + len(batch_indices), len(assigned)) if int(args.log_every) > 0 and ( completed == len(assigned) or completed % int(args.log_every) == 0 ): print( f"[extract rank {rank}] {completed}/{len(assigned)} " f"({time.monotonic() - started:.1f}s)", flush=True, ) if feature_chunks: shard_features = torch.cat(feature_chunks, dim=0).contiguous() shard_indices = torch.cat(index_chunks, dim=0).contiguous() digest_tensor = torch.tensor( [list(value) for value in prompt_digests], dtype=torch.uint8) else: shard_features = torch.empty( 0, TOTAL_TOKENS, FEATURE_DIM, dtype=torch.float16) shard_indices = torch.empty(0, dtype=torch.long) digest_tensor = torch.empty(0, 32, dtype=torch.uint8) payload = { "shard_format_version": TOOL_FORMAT_VERSION, "rank": int(rank), "world_size": int(world_size), "indices": shard_indices, "features": shard_features, "prompt_entry_sha256": digest_tensor, "sample_shape": [TOTAL_TOKENS, FEATURE_DIM], "feature_dtype": "float16", "feature_metadata": feature_metadata, "vae_identity": vae_identity, } _atomic_torch_save(payload, shard_path) print( f"[extract rank {rank}] wrote {shard_path} " f"{tuple(shard_features.shape)}", flush=True, ) del encoder, vae, dataset, payload, feature_chunks, shard_features gc.collect() torch.cuda.empty_cache() def _load_and_merge_shards( shard_dir: Path, *, world_size: int, selected_count: int, sample_shape: Sequence[int], ) -> Tuple[torch.Tensor, str, str, Dict[str, Any], Dict[str, Any]]: feature_parts = [] index_parts = [] digest_parts = [] feature_metadata = None vae_identity = None expected_tail = tuple(int(value) for value in sample_shape) for rank in range(int(world_size)): path = shard_dir / f"rank_{rank:05d}.pt" if not path.is_file(): raise FileNotFoundError(f"missing rank shard: {path}") payload = torch.load(path, map_location="cpu", weights_only=False) if int(payload.get("shard_format_version", -1)) != TOOL_FORMAT_VERSION: raise ValueError(f"incompatible shard format in {path}") if int(payload.get("rank", -1)) != rank: raise ValueError(f"rank identity mismatch in {path}") if int(payload.get("world_size", -1)) != int(world_size): raise ValueError(f"world-size mismatch in {path}") features = payload.get("features") indices = payload.get("indices") digests = payload.get("prompt_entry_sha256") shard_feature_metadata = payload.get("feature_metadata") shard_vae_identity = payload.get("vae_identity") if ( not isinstance(features, torch.Tensor) or features.dtype != torch.float16 or features.ndim != 3 or tuple(features.shape[1:]) != expected_tail or not bool(torch.isfinite(features).all().item()) ): raise ValueError(f"invalid FP16 feature shard {path}") if ( not isinstance(indices, torch.Tensor) or indices.dtype != torch.long or tuple(indices.shape) != (features.shape[0],) ): raise ValueError(f"invalid indices in {path}") if ( not isinstance(digests, torch.Tensor) or digests.dtype != torch.uint8 or tuple(digests.shape) != (features.shape[0], 32) ): raise ValueError(f"invalid prompt hashes in {path}") if not isinstance(shard_feature_metadata, dict): raise ValueError(f"missing DINOv3 feature metadata in {path}") if feature_metadata is None: feature_metadata = shard_feature_metadata elif shard_feature_metadata != feature_metadata: raise ValueError( f"DINOv3 feature metadata differs across rank shards: {path}") if not isinstance(shard_vae_identity, dict): raise ValueError(f"missing Wan VAE identity in {path}") if vae_identity is None: vae_identity = shard_vae_identity elif shard_vae_identity != vae_identity: raise ValueError( f"Wan VAE identity differs across rank shards: {path}") feature_parts.append(features) index_parts.append(indices) digest_parts.append(digests) features = torch.cat(feature_parts, dim=0) indices = torch.cat(index_parts, dim=0) digests = torch.cat(digest_parts, dim=0) if int(features.shape[0]) != int(selected_count): raise ValueError( f"rank shards contain {features.shape[0]} samples, expected " f"{selected_count}") order = torch.argsort(indices, stable=True) indices = indices.index_select(0, order) if not torch.equal(indices, torch.arange(selected_count, dtype=torch.long)): raise ValueError( "rank shards do not cover every selected LMDB index exactly once") features = features.index_select(0, order).contiguous() digests = digests.index_select(0, order).contiguous() prompt_hash = hashlib.sha256(digests.numpy().tobytes()).hexdigest() feature_hash = hashlib.sha256( features.view(torch.uint8).numpy().tobytes()).hexdigest() if feature_metadata is None: raise ValueError("rank shards contained no DINOv3 feature metadata") if vae_identity is None: raise ValueError("rank shards contained no Wan VAE identity") return features, prompt_hash, feature_hash, feature_metadata, vae_identity def _validate_feature_metadata( args: argparse.Namespace, feature_metadata: Dict[str, Any], ) -> Dict[str, Any]: """Rebuild the canonical schema and reject even consistent bad shards.""" from video_fd.encoders.dinov3_dense_metadata import ( canonical_dinov3_dense_feature_metadata, file_sha256, ) repo_path = args.repo_path.expanduser().resolve() if _git_head(repo_path) != str(args.source_commit): raise RuntimeError("DINOv3 source commit changed during extraction") if _git_source_status(repo_path).strip(): raise RuntimeError("DINOv3 source worktree changed during extraction") current_weight_sha256 = file_sha256(args.weight_path) expected = canonical_dinov3_dense_feature_metadata( repo_path=str(repo_path), weight_path=str(args.weight_path.expanduser().resolve()), source_commit=str(args.source_commit), weight_sha256=current_weight_sha256, source_num_frames=SOURCE_NUM_FRAMES, frame_indices=list(FRAME_INDICES), input_size=INPUT_SIZE, patch_size=PATCH_SIZE, patch_grid_size=PATCH_GRID_SIZE, spatial_axis_indices=list(SPATIAL_AXIS_INDICES), spatial_token_indices=list(SPATIAL_TOKEN_INDICES), per_frame_tokens=PER_FRAME_TOKENS, num_tokens=TOTAL_TOKENS, forward_chunk_frames=FORWARD_CHUNK_FRAMES, feature_dim=FEATURE_DIM, ) if feature_metadata != expected: raise ValueError( "rank-shard DINOv3 metadata does not equal the canonical schema") return expected def _reference_metadata( args: argparse.Namespace, plan: Dict[str, Any], *, feature_metadata: Dict[str, Any], vae_identity: Dict[str, Any], prompt_index_hash: str, feature_bank_hash: str, ) -> Dict[str, Any]: feature_metadata = _validate_feature_metadata(args, feature_metadata) if vae_identity != _vae_identity(): raise RuntimeError( "Wan VAE asset or wrapper changed during teacher extraction") data_mdb = Path(plan["teacher_lmdb"]["path"]) / "data.mdb" data_stat = data_mdb.stat() if ( int(data_stat.st_size) != int(plan["teacher_lmdb"]["data_mdb_size_bytes"]) or int(data_stat.st_mtime_ns) != int(plan["teacher_lmdb"]["data_mdb_mtime_ns"]) ): raise RuntimeError("teacher LMDB changed during DINOv3 extraction") source = { "kind": "causvid_ode_teacher_lmdb_final_latent", "lmdb_path": plan["teacher_lmdb"]["path"], "latents_shape": list(plan["teacher_lmdb"]["latents_shape"]), "prompts_shape": list(plan["teacher_lmdb"]["prompts_shape"]), "selected_samples": EXPECTED_TEACHER_COUNT, "selected_index_start": 0, "selected_index_stop_exclusive": EXPECTED_TEACHER_COUNT, "trajectory_index": -1, "lmdb_storage_dtype": "float16", "data_mdb_size_bytes": int( plan["teacher_lmdb"]["data_mdb_size_bytes"]), "data_mdb_mtime_ns": int( plan["teacher_lmdb"]["data_mdb_mtime_ns"]), "prompt_index_sha256": prompt_index_hash, "feature_bank_fp16_sha256": feature_bank_hash, } source["identity_sha256"] = _canonical_sha256(source) return { "feature": feature_metadata, "reference": { "source": source, "vae": vae_identity, "extraction": { "tool": "tools/precompute_dinov3_dense_nystrom_reference.py", "tool_format_version": TOOL_FORMAT_VERSION, "world_size": int(plan["world_size"]), "vae_compute_dtype": "bfloat16", "vae_decode_use_cache": False, "decoded_pixel_transform": "((pixels+1)*0.5).clamp(0,1)", "dinov3_compute_dtype": "float32", "forward_chunk_frames": FORWARD_CHUNK_FRAMES, "rank_shard_dtype": "float16", }, "build_request": { "num_landmarks": NUM_LANDMARKS, "seed": NYSTROM_SEED, "bandwidth": ( None if args.bandwidth is None else float(args.bandwidth)), "median_max_samples": int(args.median_max_samples), "kmeans_iterations": int(args.kmeans_iterations), "ridge": float(args.ridge), "chunk_size": int(args.chunk_size), }, }, } def _build_and_save_reference( args: argparse.Namespace, plan: Dict[str, Any], shard_dir: Path, *, device: torch.device, ) -> Dict[str, Any]: from losses.video_nystrom_mmd import ( build_video_nystrom_reference, load_reference_pack, save_reference_pack, validate_reference_pack, ) features, prompt_hash, feature_hash, feature_metadata, vae_identity = ( _load_and_merge_shards( shard_dir, world_size=int(plan["world_size"]), selected_count=EXPECTED_TEACHER_COUNT, sample_shape=(TOTAL_TOKENS, FEATURE_DIM), ) ) metadata = _reference_metadata( args, plan, feature_metadata=feature_metadata, vae_identity=vae_identity, prompt_index_hash=prompt_hash, feature_bank_hash=feature_hash, ) print( f"[merge] building m={NUM_LANDMARKS} reference from " f"{tuple(features.shape)} FP16 matrices on {device}", flush=True, ) reference = features.to(device=device, non_blocking=False) pack = build_video_nystrom_reference( reference, num_landmarks=NUM_LANDMARKS, metadata=metadata, bandwidth=args.bandwidth, median_max_samples=int(args.median_max_samples), kmeans_iterations=int(args.kmeans_iterations), ridge=float(args.ridge), chunk_size=int(args.chunk_size), seed=NYSTROM_SEED, ) validate_reference_pack(pack) del reference torch.cuda.empty_cache() output = args.output.expanduser().resolve() output.parent.mkdir(parents=True, exist_ok=True) temporary = output.with_name(f".{output.name}.tmp-{os.getpid()}") try: save_reference_pack(pack, temporary) staged = load_reference_pack(temporary, map_location="cpu") validate_reference_pack(staged) if output.exists() and not bool(args.overwrite): raise FileExistsError( f"output appeared during construction: {output}") os.replace(temporary, output) finally: if temporary.exists(): temporary.unlink() # This final-path reload is the gate for temporary-shard deletion. published = load_reference_pack(output, map_location="cpu") validate_reference_pack(published) _validate_published_pack_identity(pack, published) return { "output": str(output), "samples": int(features.shape[0]), "sample_shape": list(features.shape[1:]), "num_landmarks": int(pack["landmarks"].shape[0]), "sigma": float(torch.as_tensor(pack["sigma"]).item()), "content_hash": str(pack["content_hash"]), "prompt_index_sha256": prompt_hash, "feature_bank_fp16_sha256": feature_hash, "final_path_validated": True, } def _validate_published_pack_identity( pack: Mapping[str, Any], published: Mapping[str, Any], ) -> None: """Require the final-path reload to preserve the built pack identity.""" if published.get("content_hash") != pack.get("content_hash"): raise ValueError("published reference content hash changed after rename") # The builder canonicalizes user metadata and adds its reserved ``nystrom`` # block. Compare against that exact pack payload, not the pre-build input. if published.get("metadata") != pack.get("metadata"): raise ValueError("published reference metadata changed after rename") def _remove_validated_shards( shard_dir: Path, build_summary: Dict[str, Any], ) -> None: """Delete shards only with the explicit final-path validation receipt.""" if build_summary.get("final_path_validated") is not True: raise RuntimeError( "refusing to delete shards before final reference validation") output = Path(str(build_summary.get("output", ""))).expanduser().resolve() if not output.is_file(): raise FileNotFoundError( f"refusing to delete shards; published reference is missing: {output}") shutil.rmtree(shard_dir) def _run(args: argparse.Namespace) -> None: _validate_args(args) env_world_size = int(os.environ.get("WORLD_SIZE", "1")) env_rank = int(os.environ.get("RANK", "0")) plan = _plan(args, env_world_size) if bool(args.dry_run): if env_rank == 0: print(json.dumps(plan, indent=2, sort_keys=True), flush=True) return device, rank, world_size = _setup_distributed( int(args.distributed_timeout_minutes)) if world_size != int(plan["world_size"]): raise RuntimeError( f"world size changed after planning: {world_size} != " f"{plan['world_size']}") shard_dir = None success = False validated_summary = None try: vae_result = None if rank == 0: try: vae_result = {"identity": _vae_identity(), "error": None} except Exception as exc: vae_result = { "identity": None, "error": f"{type(exc).__name__}: {exc}", } vae_result = _broadcast_rank0_object(vae_result, rank) if vae_result["error"] is not None: raise RuntimeError( "rank 0 could not resolve the Wan VAE identity: " f"{vae_result['error']}") vae_identity = vae_result["identity"] temp_result = None if rank == 0: try: created = Path(tempfile.mkdtemp( prefix="dinov3_dense_nystrom_", dir=Path(plan["temporary"]["parent"]), )) temp_result = {"path": str(created), "error": None} except Exception as exc: temp_result = { "path": None, "error": f"{type(exc).__name__}: {exc}", } temp_result = _broadcast_rank0_object(temp_result, rank) if temp_result["error"] is not None: raise RuntimeError( "rank 0 could not create shard directory: " f"{temp_result['error']}") shard_dir = Path(temp_result["path"]) _extract_rank_shard( args, plan, device=device, rank=rank, world_size=world_size, shard_path=shard_dir / f"rank_{rank:05d}.pt", vae_identity=vae_identity, ) _barrier() summary = None build_error = None if rank == 0: try: summary = _build_and_save_reference( args, plan, shard_dir, device=device) except Exception as exc: import traceback traceback.print_exc() build_error = f"{type(exc).__name__}: {exc}" outcome = _broadcast_rank0_object( {"summary": summary, "error": build_error}, rank) if outcome["error"] is not None: raise RuntimeError( f"rank 0 reference construction failed: {outcome['error']}") _barrier() validated_summary = outcome["summary"] success = True if rank == 0: print(json.dumps( {"status": "ok", **outcome["summary"]}, indent=2, sort_keys=True, ), flush=True) finally: if rank == 0 and shard_dir is not None: if success: _remove_validated_shards(shard_dir, validated_summary) print( f"[cleanup] removed validated temporary shards {shard_dir}", flush=True, ) else: print( f"[cleanup] preserving failed-run shards at {shard_dir}", flush=True, ) if dist.is_available() and dist.is_initialized(): dist.destroy_process_group() def main(argv: Optional[Iterable[str]] = None) -> None: _run(build_parser().parse_args(argv)) if __name__ == "__main__": main()