Blame for 14B/code/tools/extract_triple_matched_teacher_bank.py

#!/usr/bin/env python3
"""Decode each frozen 81-frame MP4 once and write a resumable triple bank."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import tempfile
import time

import numpy as np
import torch
import torch.distributed as dist
from safetensors import safe_open
from safetensors.torch import save_file

from utils.triple_video_feature_bank import (
    BANK_KIND,
    BANK_SCHEMA_VERSION,
    canonical_sha256,
    file_sha256,
)
from utils.matched_video_manifest import (
    load_matched_video_manifest,
    sha256_file,
)
from video_fd.encoders.dinov3_temporal_pool_adapter import (
    DINOv3TemporalPoolAdapter,
)
from video_fd.encoders.segmented_videomae_v2 import (
    SegmentedVideoMAEv2Ensemble,
)
from video_fd.encoders.videomae_v2_dense_metadata import (
    canonical_videomae_v2_dense_feature_metadata,
    uniform_segment_token_indices,
)
from video_fd.encoders.vjepa2_adapter import VJEPA2Adapter
from video_fd.encoders.vjepa2_dense_metadata import (
    canonical_vjepa2_dense_feature_metadata,
    uniform_dense_token_indices,
)


VIDEOMAE_REVISION = "9981a9c8f77118c421e5228e1b219468a4b0238d"


def atomic_json(path: Path, value) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    descriptor, name = tempfile.mkstemp(
        dir=path.parent, prefix=f".{path.name}.", suffix=".tmp")
    try:
        with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
            json.dump(value, handle, sort_keys=True, indent=2, ensure_ascii=False)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(name, path)
    finally:
        Path(name).unlink(missing_ok=True)


def tensor_sha256(tensor: torch.Tensor) -> str:
    value = tensor.detach().cpu().contiguous()
    return hashlib.sha256(value.numpy().tobytes()).hexdigest()


def decode_81(path: Path, *, size: int, mtime_ns: int) -> torch.Tensor:
    stat = path.stat()
    if stat.st_size != size or stat.st_mtime_ns != mtime_ns:
        raise ValueError(f"frozen video stat changed: {path}")
    import av

    frames = []
    with av.open(str(path), mode="r") as container:
        if len(container.streams.video) != 1:
            raise ValueError(f"expected one video stream: {path}")
        stream = container.streams.video[0]
        for frame in container.decode(stream):
            array = frame.to_ndarray(format="rgb24")
            if array.dtype != np.uint8 or array.ndim != 3 or array.shape[2] != 3:
                raise ValueError(f"invalid RGB frame in {path}")
            frames.append(torch.from_numpy(
                np.ascontiguousarray(array)).permute(2, 0, 1))
    if len(frames) != 81:
        raise ValueError(f"{path.name} has {len(frames)} frames, expected 81")
    shape = tuple(frames[0].shape)
    if any(tuple(frame.shape) != shape for frame in frames):
        raise ValueError(f"resolution changes within {path}")
    return torch.stack(frames).float().div_(255.0).unsqueeze(0)


def load_download_manifest(model_dir: Path) -> dict:
    path = model_dir / "download_manifest.json"
    with path.open("r", encoding="utf-8") as handle:
        manifest = json.load(handle)
    payload = dict(manifest)
    saved = payload.pop("content_hash")
    if (
        manifest.get("revision") != VIDEOMAE_REVISION
        or canonical_sha256(payload) != saved
    ):
        raise ValueError("invalid pinned VideoMAE download manifest")
    for item in manifest["files"]:
        path = (model_dir / item["path"]).resolve()
        if (
            model_dir not in path.parents
            or path.stat().st_size != int(item["size_bytes"])
            or file_sha256(path) != item["sha256"]
        ):
            raise ValueError(f"VideoMAE file identity mismatch: {path}")
    return manifest


def build_encoders(args, device):
    dino = DINOv3TemporalPoolAdapter({
        "repo_path": str(args.dinov3_repo),
        "weight_path": str(args.dinov3_weights),
        "source_commit": str(args.dinov3_source_commit),
        "checkpoint_sha256": str(args.dinov3_checkpoint_sha256),
        "source_num_frames": 81,
        "frame_indices": (
            0, 4, 8, 13, 17, 21, 25, 29, 34, 38,
            42, 46, 51, 55, 59, 63, 67, 72, 76, 80),
        "input_size": 256,
        "patch_size": 16,
        "spatial_axis_indices": tuple(range(16)),
        "forward_chunk_frames": 4,
        "geometry": "temporal_mean5_full256",
        "temporal_groups": 4,
        "temporal_pool_frames": 5,
    }).to(device=device, dtype=torch.float32).eval().requires_grad_(False)
    vj = VJEPA2Adapter({
        "name": "vjepa2_dense",
        "model_id": str(args.vjepa_model),
        "input_frames": 64,
        "dense_num_frames": 64,
        "input_size": 256,
        "resize_short_side": 292,
        "crop_mode": "center",
        "dense_only": True,
        "dense_layer_index": 15,
    }).to(device=device, dtype=torch.float32).eval().requires_grad_(False)
    vm = SegmentedVideoMAEv2Ensemble({
        "name": "videomae_v2_five_segment_dense",
        "model_id": str(args.videomae_model),
        "revision": VIDEOMAE_REVISION,
        "local_files_only": True,
        "encoder_count": 5,
        "skip_leading_frames": 1,
        "segment_num_frames": 16,
        "source_num_frames": 81,
        "input_frames": 16,
        "dense_num_frames": 16,
        "input_size": 224,
        "resize_short_side": 224,
        "crop_mode": "center",
        "dense_only": True,
        "dense_layer_index": 15,
    }).to(device=device, dtype=torch.float32).eval().requires_grad_(False)
    if (dino.dense_feature_dim, vj.dense_feature_dim, vm.dense_feature_dim) != (1024, 1024, 1024):
        raise RuntimeError("encoder feature dimensions differ from contract")
    return dino, vj, vm


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--manifest-sha256", required=True)
    parser.add_argument("--source-csv-sha256", required=True)
    parser.add_argument("--video-root", required=True)
    parser.add_argument("--dinov3-repo", required=True)
    parser.add_argument("--dinov3-weights", required=True)
    parser.add_argument("--dinov3-source-commit", required=True)
    parser.add_argument("--dinov3-checkpoint-sha256", required=True)
    parser.add_argument("--vjepa-model", required=True)
    parser.add_argument("--vjepa-reference-model-id", required=True)
    parser.add_argument("--vjepa-model-sha256", required=True)
    parser.add_argument("--videomae-model", required=True)
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--shard-rows", type=int, default=32)
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--expected-count", type=int, default=8000)
    args = parser.parse_args()
    if args.shard_rows != 32:
        raise ValueError("triple teacher bank requires shard-rows=32")
    manifest_path = Path(args.manifest).expanduser().resolve()
    if sha256_file(manifest_path) != args.manifest_sha256:
        raise ValueError("source manifest SHA-256 mismatch")
    entries, source_metadata = load_matched_video_manifest(
        manifest_path,
        expected_source_csv_sha256=args.source_csv_sha256,
        expected_count=args.expected_count,
    )
    video_root = Path(args.video_root).expanduser().resolve()
    output = Path(args.output_dir).expanduser().resolve()
    shard_dir = output / "shards"
    shard_dir.mkdir(parents=True, exist_ok=True)
    progress_path = output / "progress.json"
    final_path = output / "manifest.json"
    vm_download = load_download_manifest(
        Path(args.videomae_model).expanduser().resolve())
    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    distributed = world_size > 1
    if distributed:
        local_rank = int(os.environ["LOCAL_RANK"])
        torch.cuda.set_device(local_rank)
        dist.init_process_group(backend="nccl")
        rank = int(dist.get_rank())
        world_size = int(dist.get_world_size())
        device = torch.device("cuda", local_rank)
    else:
        rank = 0
        device = torch.device(args.device)
    dino, vj, vm = build_encoders(args, device)
    vj_indices = uniform_dense_token_indices(
        vj.dense_raw_tokens, 1024, device=device)
    vm_indices = uniform_segment_token_indices(
        vm.dense_raw_tokens, 256, device=device)
    vj_contract = canonical_vjepa2_dense_feature_metadata(
        model_id=args.vjepa_reference_model_id,
        layer_index=15,
        num_frames=64,
        resize_short_side=292,
        crop_size=256,
        raw_tokens=vj.dense_raw_tokens,
        num_tokens=1024,
        feature_dim=1024,
        token_indices=vj_indices.cpu(),
        frame_sampling="last_contiguous",
        source_num_frames=81,
    )
    vm_contract = canonical_videomae_v2_dense_feature_metadata(
        model_id=str(Path(args.videomae_model).expanduser().resolve()),
        layer_index=15,
        num_tokens=256,
        token_indices=vm_indices.cpu(),
    )
    fixed = {
        "schema_version": BANK_SCHEMA_VERSION,
        "kind": BANK_KIND,
        "sample_count": len(entries),
        "source_manifest": {
            "path": str(manifest_path),
            "sha256": args.manifest_sha256,
            "source_csv_sha256": args.source_csv_sha256,
            "sample_identity_sha256": source_metadata.get(
                "sample_identity_sha256"),
        },
        "branches": {
            "dinov3": {
                "sample_shape": [1024, 1024],
                "dtype": "float16",
                "feature_contract": dino.feature_metadata(),
                "model_identity": {
                    "model_id": "dinov3_vitl16",
                    "repo_path": str(Path(args.dinov3_repo).resolve()),
                    "weight_path": str(Path(args.dinov3_weights).resolve()),
                    "source_commit": args.dinov3_source_commit,
                    "checkpoint_sha256": args.dinov3_checkpoint_sha256,
                },
            },
            "vjepa2": {
                "sample_shape": [1024, 1024],
                "dtype": "float16",
                "feature_contract": vj_contract,
                "model_identity": {
                    "model_id": args.vjepa_reference_model_id,
                    "load_path": str(Path(args.vjepa_model).resolve()),
                    "sha256": args.vjepa_model_sha256,
                },
            },
            "videomae_v2": {
                "sample_shape": [1280, 1024],
                "dtype": "float16",
                "feature_contract": vm_contract,
                "model_identity": {
                    "model_id": "VideoMAEv2-Large",
                    "revision": VIDEOMAE_REVISION,
                    "download_manifest_content_hash": vm_download[
                        "content_hash"],
                },
            },
        },
        "samples": [{
            "index": entry.manifest_index,
            "sample_id": entry.sample_id,
            "prompt": entry.caption,
        } for entry in entries],
    }
    progress_identity = canonical_sha256(fixed)
    receipt_dir = output / "receipts"
    receipt_dir.mkdir(parents=True, exist_ok=True)
    shard_count = (len(entries) + args.shard_rows - 1) // args.shard_rows
    started = time.time()
    for shard_index in range(rank, shard_count, world_size):
        start = shard_index * args.shard_rows
        end = min(start + args.shard_rows, len(entries))
        relative = f"shards/shard-{shard_index:05d}.safetensors"
        path = output / relative
        receipt_path = receipt_dir / f"shard-{shard_index:05d}.json"
        if receipt_path.is_file():
            receipt = json.loads(receipt_path.read_text(encoding="utf-8"))
            if (
                receipt.get("progress_identity") != progress_identity
                or int(receipt.get("index", -1)) != shard_index
                or int(receipt.get("sample_start", -1)) != start
                or int(receipt.get("sample_end", -1)) != end
                or not path.is_file()
                or file_sha256(path) != receipt.get("sha256")
            ):
                raise ValueError(
                    f"resumable triple shard identity mismatch: {path}")
            print(json.dumps({
                "rank": rank,
                "resumed_shard": shard_index,
                "completed": end,
            }), flush=True)
            continue

        branch_rows = {
            "dinov3": [],
            "vjepa2": [],
            "videomae_v2": [],
        }
        for entry in entries[start:end]:
            video = decode_81(
                video_root / entry.video_name,
                size=entry.video_size_bytes,
                mtime_ns=entry.video_mtime_ns,
            ).to(device=device, dtype=torch.float32)
            temporal = torch.arange(
                17, 81, device=device, dtype=torch.long).unsqueeze(0)
            with torch.inference_mode():
                dino_feature = dino.forward_dense_no_grad(
                    video, gradient_checkpointing=False)
                vj_feature = vj.forward_dense(
                    video,
                    layer_index=15,
                    token_indices=vj_indices,
                    gradient_checkpointing=False,
                    temporal_window_indices=temporal,
                )
                vm_feature = vm.forward_dense(
                    video,
                    layer_index=15,
                    token_indices=vm_indices,
                    gradient_checkpointing=False,
                )
            observed = {
                "dinov3": tuple(dino_feature.shape),
                "vjepa2": tuple(vj_feature.shape),
                "videomae_v2": tuple(vm_feature.shape),
            }
            expected = {
                "dinov3": (1, 1024, 1024),
                "vjepa2": (1, 1024, 1024),
                "videomae_v2": (1, 1280, 1024),
            }
            if observed != expected:
                raise RuntimeError(
                    f"triple bank row shape mismatch: {observed}")
            branch_rows["dinov3"].append(
                dino_feature.squeeze(0).half().cpu().contiguous())
            branch_rows["vjepa2"].append(
                vj_feature.squeeze(0).half().cpu().contiguous())
            branch_rows["videomae_v2"].append(
                vm_feature.squeeze(0).half().cpu().contiguous())
            del video, dino_feature, vj_feature, vm_feature

        tensors = {
            name: torch.stack(rows, dim=0)
            for name, rows in branch_rows.items()
        }
        ids = [entry.sample_id for entry in entries[start:end]]
        shard_payload = {
            "index": shard_index,
            "sample_start": start,
            "sample_end": end,
            "sample_ids_sha256": canonical_sha256(ids),
            "tensor_sha256": {
                name: tensor_sha256(value)
                for name, value in tensors.items()
            },
            "branches": fixed["branches"],
            "source_manifest_sha256": args.manifest_sha256,
            "progress_identity": progress_identity,
        }
        shard_content_hash = canonical_sha256(shard_payload)
        temporary = path.with_name(f".{path.name}.rank{rank}.tmp")
        save_file(tensors, str(temporary), metadata={
            "content_hash": shard_content_hash,
            "sample_start": str(start),
            "sample_end": str(end),
            "source_manifest_sha256": args.manifest_sha256,
            "progress_identity": progress_identity,
        })
        os.replace(temporary, path)
        shard = {
            "index": shard_index,
            "path": relative,
            "sample_start": start,
            "sample_end": end,
            "sample_ids_sha256": canonical_sha256(ids),
            "sha256": file_sha256(path),
            "content_hash": shard_content_hash,
        }
        atomic_json(receipt_path, {
            **shard,
            "progress_identity": progress_identity,
            "rank": rank,
        })
        print(json.dumps({
            "rank": rank,
            "completed": end,
            "total": len(entries),
            "elapsed_seconds": time.time() - started,
            "shard": relative,
        }), flush=True)

    if distributed:
        dist.barrier()
    if rank == 0:
        shards = []
        for shard_index in range(shard_count):
            receipt_path = receipt_dir / f"shard-{shard_index:05d}.json"
            if not receipt_path.is_file():
                raise FileNotFoundError(
                    f"missing completed shard receipt: {receipt_path}")
            receipt = json.loads(receipt_path.read_text(encoding="utf-8"))
            if receipt.pop("progress_identity", None) != progress_identity:
                raise ValueError("triple shard progress identity mismatch")
            receipt.pop("rank", None)
            path = output / receipt["path"]
            if file_sha256(path) != receipt["sha256"]:
                raise ValueError(f"triple shard SHA-256 mismatch: {path}")
            shards.append(receipt)

        branch_hashers = {
            name: hashlib.sha256() for name in fixed["branches"]
        }
        for shard in shards:
            with safe_open(
                str(output / shard["path"]),
                framework="pt",
                device="cpu",
            ) as handle:
                for name in branch_hashers:
                    value = handle.get_tensor(name).contiguous()
                    branch_hashers[name].update(value.numpy().tobytes())
        for name, hasher in branch_hashers.items():
            fixed["branches"][name]["feature_bank_fp16_sha256"] = (
                hasher.hexdigest())
        final_payload = {**fixed, "shards": shards}
        final_payload["content_hash"] = canonical_sha256(final_payload)
        atomic_json(final_path, final_payload)
        print(json.dumps({
            "manifest": str(final_path),
            "manifest_sha256": file_sha256(final_path),
            "content_hash": final_payload["content_hash"],
            "samples": len(entries),
            "world_size": world_size,
        }, sort_keys=True), flush=True)
    if distributed:
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    main()