diff --git "a/1.3B/code/model/distribution_matching.py" "b/1.3B/code/model/distribution_matching.py" new file mode 100644--- /dev/null +++ "b/1.3B/code/model/distribution_matching.py" @@ -0,0 +1,6004 @@ +"""distribution matching: match generated videos in frozen encoder feature spaces. + +Read the active path in this order: + collect_rollout_features -> hybrid MMD -> prepare_latent_gradients + -> replay_generator_gradients. + +The encoders and VAE are frozen, but their input derivatives connect the +feature-space objective to the generator. Modules/attributes used in saved +state dictionaries keep their original names for checkpoint compatibility. +See docs/paper-to-code.md for the paper mapping and estimator conventions. +""" + +from __future__ import annotations + +import copy +from contextlib import contextmanager +import hashlib +import json +import os +import random +import time +from typing import Mapping, Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist + +from losses.multilayer_video_nystrom_mmd import ( + MultiLayerVideoNystromMMD, +) +from losses.wan_kv_mmd import EMARBFMMD +from model.generation import FewStepGenerationModel +from utils.model_identity import representation_runtime_identity +from video_fd.encoders.vjepa2_dense_metadata import ( + canonical_vjepa2_dense_feature_metadata, + uniform_dense_token_indices, +) +from video_fd.encoders.vjepa2_temporal_sampling import ( + stateless_temporal_window_start, +) +from video_fd.replay.replay_video_encoder_vjp import replay_decoder_vjp + + +def _canonicalize_shared_asset_paths(value): + """Normalize the host-local and cluster-wide mount names recursively.""" + local_root = "/data/shared/video-generation-assets" + logical_root = "/assets/video-generation-assets" + if isinstance(value, str): + if value == local_root or value.startswith(local_root + "/"): + return logical_root + value[len(local_root):] + return value + if isinstance(value, Mapping): + return { + key: _canonicalize_shared_asset_paths(item) + for key, item in value.items() + } + if isinstance(value, tuple): + return tuple(_canonicalize_shared_asset_paths(item) for item in value) + if isinstance(value, list): + return [_canonicalize_shared_asset_paths(item) for item in value] + return value + + +def _accumulate_fp32_pixel_cotangent( + accumulated: Optional[torch.Tensor], + value: torch.Tensor, + expected_shape: Optional[tuple[int, ...]], + *, + branch_name: str, +) -> tuple[torch.Tensor, tuple[int, ...], torch.Tensor]: + """Accumulate one branch immediately and retain only its scalar norm.""" + if not isinstance(value, torch.Tensor) or not value.is_floating_point(): + raise TypeError(f"{branch_name} pixel cotangent must be floating") + detached = value.detach().float() + shape = tuple(detached.shape) + if expected_shape is not None and shape != expected_shape: + raise ValueError( + "triple encoder pixel cotangent shapes differ: " + f"{branch_name}={shape} != {expected_shape}") + if not bool(torch.isfinite(detached).all().item()): + raise FloatingPointError( + f"{branch_name} pixel cotangent is non-finite") + branch_norm = detached.norm().detach() + if accumulated is None: + accumulated = detached + else: + accumulated.add_(detached) + return accumulated, shape, branch_norm + + +class DistributionMatchingModel(FewStepGenerationModel): + """Match generated and teacher latents in a configured dense feature space. + + The generator remains the only trainable model. In ``post_block_latent`` + mode ``real_score`` is reused as the frozen Wan encoder so the trainer can + shard it normally with FSDP. In ``vjepa2_dense_latent`` mode ``real_score`` + is not constructed; a frozen, replicated FP32 V-JEPA2 prefix is used + instead. + """ + + uses_critic = False + uses_vae_for_loss = False + uses_vae_for_visualization = True + requires_distribution_state_on_resume = True + + _DENSE_VIDEO_FEATURE_TYPES = frozenset({ + "vjepa2_dense_latent", + "dinov3_dense_latent", + "videomae_v2_dense_latent", + "dual_vjepa2_videomae_v2_dense_latent", + "dual_dinov3_vjepa2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent", + }) + + @staticmethod + def _dense_feature_prefix_for(feature_type: str) -> Optional[str]: + return { + "vjepa2_dense_latent": "vjepa2", + "dinov3_dense_latent": "dinov3", + "videomae_v2_dense_latent": "videomae_v2", + "dual_vjepa2_videomae_v2_dense_latent": "dual", + "dual_dinov3_vjepa2_dense_latent": "dinov3_vjepa2", + "triple_dinov3_vjepa2_videomae_v2_dense_latent": "triple", + "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent": "triple", + }.get(str(feature_type)) + + def _dense_config_value(self, args, suffix: str, default): + """Read a backbone-specific public config key. + + The established replay implementation keeps its historical internal + ``vjepa2_*`` attributes for compatibility. Public VideoMAE-v2 configs, + metadata, logs, and checkpoint keys remain independently namespaced. + """ + prefix = self.dense_feature_prefix + return getattr(args, f"{prefix}_{suffix}", default) + + @property + def dense_state_prefix(self) -> str: + """Public artifact namespace for the active dense-video backbone.""" + prefix = getattr(self, "dense_feature_prefix", None) + if prefix is None: + prefix = self._dense_feature_prefix_for( + getattr(self, "feature_type", "")) + return prefix or "vjepa2" + + @property + def dense_nystrom_version_key(self) -> str: + return f"{self.dense_state_prefix}_dense_nystrom_mmd_version" + + @property + def dense_ema_version_key(self) -> str: + return f"{self.dense_state_prefix}_dense_mmd_version" + + @property + def dense_ema_payload_key(self) -> str: + return f"{self.dense_state_prefix}_dense_mmd" + + @property + def dense_feature_display_name(self) -> str: + return ( + "VideoMAE-v2" + if self.dense_state_prefix == "videomae_v2" + else ( + "DINOv3 + V-JEPA2 + VideoMAE-v2" + if self.dense_state_prefix == "triple" + else ( + "DINOv3 + V-JEPA2" + if self.dense_state_prefix == "dinov3_vjepa2" + else ( + "V-JEPA2 + VideoMAE-v2" + if self.dense_state_prefix == "dual" + else ("DINOv3" if self.dense_state_prefix == "dinov3" else "V-JEPA2") + ) + ) + ) + ) + + @property + def dense_feature_metric_prefix(self) -> str: + return f"{self.dense_state_prefix}_dense" + + @property + def dense_sampled_tokens_per_layer(self) -> int: + return int(self.vjepa2_dense_num_tokens) * int( + getattr(self, "dense_feature_segment_count", 1)) + + @property + def dense_raw_tokens_total(self) -> int: + return int(self.vjepa2_dense_raw_tokens) * int( + getattr(self, "dense_feature_segment_count", 1)) + + def _namespace_dense_metrics(self, values: Mapping[str, object]) -> dict: + """Translate legacy internal metric names at the public boundary.""" + if self.dense_state_prefix == "vjepa2": + return dict(values) + if self.dense_state_prefix in { + "dual", "triple", "dinov3", "dinov3_vjepa2" + }: + return dict(values) + return { + ( + key.replace("vjepa2_", "videomae_v2_", 1) + if isinstance(key, str) and key.startswith("vjepa2_") + else key + ): value + for key, value in values.items() + } + + @staticmethod + def _all_ranks_true(local_value: bool, device) -> bool: + flag = torch.tensor( + int(bool(local_value)), device=device, dtype=torch.int32) + if ( + dist.is_available() + and dist.is_initialized() + and dist.get_world_size() > 1 + ): + dist.all_reduce(flag, op=dist.ReduceOp.MIN) + return bool(flag.item()) + + def _synchronize_dinov3_local_failure( + self, + error: Optional[BaseException], + device, + *, + operation: str, + ) -> None: + """Make a rank-local DINOv3 failure an all-rank branch decision. + + DINOv3 startup performs expensive local imports, file inspection and + metadata validation before it hashes shared artifacts. A rank that + raises before a peer enters one of those hash broadcasts would strand + the peer. Every such local-only phase terminates through this helper. + Single-process callers retain the original exception type and message. + """ + if self._all_ranks_true(error is None, device): + return + is_distributed = ( + dist.is_available() + and dist.is_initialized() + and dist.get_world_size() > 1 + ) + if error is not None and not is_distributed: + raise error + detail = ( + f": {type(error).__name__}: {error}" + if error is not None + else " on another rank" + ) + raise RuntimeError(operation + detail) from error + @staticmethod + def _configured_current_teacher_global_samples(args) -> int: + """Resolve the exact live-teacher geometry for this process group. + + Legacy GradCache updates use every forward sample for the MMD VJP, so + their teacher count remains ``total_batch_size``. Partial replay + explicitly configures a per-rank gradient suffix; only that suffix has + a matched teacher and therefore the exact teacher kernel has + ``gradient_local * world_size`` rows. + """ + configured_total = int(getattr(args, "total_batch_size", 0)) + gradient_local = getattr( + args, "fresh_video_generator_grad_local_samples", None) + if gradient_local is None: + return configured_total + if type(gradient_local) is not int or gradient_local <= 0: + raise ValueError( + "fresh_video_generator_grad_local_samples must be a " + "positive integer") + forward_local = getattr(args, "fresh_video_local_samples", None) + if type(forward_local) is not int or forward_local <= 0: + raise ValueError( + "fresh_video_local_samples must be a positive integer when " + "partial GradCache replay is configured") + if gradient_local > forward_local: + raise ValueError( + "fresh_video_generator_grad_local_samples must not exceed " + "fresh_video_local_samples") + + if dist.is_available() and dist.is_initialized(): + world_size = int(dist.get_world_size()) + else: + if configured_total <= 0 or configured_total % forward_local != 0: + raise ValueError( + "total_batch_size must be divisible by " + "fresh_video_local_samples before distributed " + "initialization") + world_size = configured_total // forward_local + expected_forward = forward_local * world_size + if configured_total != expected_forward: + raise ValueError( + "total_batch_size does not match the partial GradCache " + "forward geometry: " + f"{configured_total} != {forward_local} * {world_size}") + return gradient_local * world_size + + def __init__(self, args, device): + p0_value = os.environ.get("SF_FRESH_VIDEO_P0_REPLAY") + if p0_value is None: + self._fresh_video_p0_requested = True + else: + normalized = p0_value.strip().lower() + if normalized not in {"0", "1", "false", "true", "no", "yes"}: + raise ValueError( + "SF_FRESH_VIDEO_P0_REPLAY must be one of " + "0/1/false/true/no/yes") + self._fresh_video_p0_requested = normalized in { + "1", "true", "yes"} + # Final eligibility is known only after the estimator is initialized. + self.fresh_video_prepared_replay_enabled = False + self.fresh_video_pixel_cache_device = os.environ.get( + "SF_FRESH_VIDEO_PIXEL_CACHE_DEVICE", "cuda" + ).strip().lower() + if self.fresh_video_pixel_cache_device not in {"cpu", "cuda"}: + raise ValueError( + "SF_FRESH_VIDEO_PIXEL_CACHE_DEVICE must be cpu or cuda") + requested_feature_type = str( + getattr(args, "wan_feature_type", "post_block_latent")) + supported_feature_types = { + "post_block_latent", + "vjepa2_dense_latent", + "dinov3_dense_latent", + "videomae_v2_dense_latent", + "dual_vjepa2_videomae_v2_dense_latent", + "dual_dinov3_vjepa2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent", + } + if requested_feature_type not in supported_feature_types: + raise ValueError( + "wan_feature_type must be one of " + f"{sorted(supported_feature_types)}, got " + f"{requested_feature_type!r}") + self.dense_feature_prefix = self._dense_feature_prefix_for( + requested_feature_type) + is_dense_video_feature = bool( + requested_feature_type in self._DENSE_VIDEO_FEATURE_TYPES) + self.mmd_estimator = str( + getattr(args, "wan_mmd_estimator", "ema_rff")) + supported_estimators = {"ema_rff", "nystrom_exact_repulsion"} + if self.mmd_estimator not in supported_estimators: + raise ValueError( + "wan_mmd_estimator must be one of " + f"{sorted(supported_estimators)}, got " + f"{self.mmd_estimator!r}") + if ( + self.mmd_estimator == "nystrom_exact_repulsion" + and not is_dense_video_feature + ): + raise ValueError( + "nystrom_exact_repulsion currently requires " + "a dense video feature backend") + self.is_vjepa2_dense = ( + requested_feature_type == "vjepa2_dense_latent") + self.is_dinov3_dense = ( + requested_feature_type == "dinov3_dense_latent") + self.is_videomae_v2_dense = ( + requested_feature_type == "videomae_v2_dense_latent") + self.is_dinov3_vjepa2_dense = ( + requested_feature_type == "dual_dinov3_vjepa2_dense_latent") + self.triple_t5_joint_enabled = ( + requested_feature_type + == "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent") + self.is_triple_dense = requested_feature_type in { + "triple_dinov3_vjepa2_videomae_v2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent", + } + self.is_dense_video_feature = is_dense_video_feature + if self.is_dinov3_dense and ( + self.mmd_estimator != "nystrom_exact_repulsion" + ): + raise ValueError( + "wan_feature_type=dinov3_dense_latent currently requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if self.is_dense_video_feature: + # BaseModel otherwise constructs full real/fake Wan score models. + # Neither participates in this generator-only replay objective. + args.uses_score_models = False + args.uses_real_score = False + args.uses_fake_score = False + super().__init__(args, device) + self.feature_type = requested_feature_type + self.dense_encoder_name = self.dense_state_prefix + self.dense_log_prefix = f"{self.dense_encoder_name}_dense" + if self.triple_t5_joint_enabled: + self._initialize_triple_t5_joint_configuration(args) + + if self.feature_type == "post_block_latent": + self._initialize_wan_features(args) + elif self.is_vjepa2_dense: + self._initialize_vjepa2_features(args, device) + elif self.is_dinov3_dense: + self._initialize_dinov3_features(args, device) + elif self.feature_type == "videomae_v2_dense_latent": + self._initialize_videomae_v2_features(args, device) + elif self.is_dinov3_vjepa2_dense: + self._initialize_dual_dinov3_vjepa2_features(args, device) + else: + self._initialize_multi_encoder_features(args, device) + + self.latent_loss_weight = float( + getattr(args, "wan_kv_loss_weight", 1.0)) + self.latent_self_normalize = bool( + getattr(args, "wan_kv_self_normalize", True)) + self.latent_norm_eps = float( + getattr(args, "wan_kv_norm_eps", 1e-4)) + if ( + getattr(self, "vjepa2_multilayer_enabled", False) + and self.mmd_estimator != "nystrom_exact_repulsion" + ): + raise ValueError( + "multi-layer V-JEPA2 MMD requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if ( + getattr(self, "vjepa2_multilayer_enabled", False) + and not self.latent_self_normalize + ): + raise ValueError( + "multi-layer V-JEPA2 MMD requires " + "wan_kv_self_normalize=true so layer losses are normalized " + "independently before summation") + if self.is_vjepa2_dense: + self.vjepa2_token_sampling = str(getattr( + args, "vjepa2_token_sampling", "uniform_fixed")) + if self.vjepa2_token_sampling != "uniform_fixed": + raise ValueError( + "this release supports vjepa2_token_sampling=" + "uniform_fixed") + else: + self.vjepa2_token_sampling = None + self.uses_fresh_video_gradcache = ( + self.mmd_estimator == "nystrom_exact_repulsion") + self.nystrom_attraction_ema_enabled = bool( + self._dense_config_value( + args, "nystrom_attraction_ema_enabled", False)) + self.nystrom_attraction_ema_beta = float( + self._dense_config_value( + args, "nystrom_attraction_ema_beta", 0.99)) + self.nystrom_attraction_weight = float( + self._dense_config_value( + args, "nystrom_attraction_weight", 1.0)) + self.current_teacher_attraction_enabled = bool( + self._dense_config_value( + args, "nystrom_current_teacher_attraction_enabled", False)) + self.current_teacher_attraction_weight = float( + self._dense_config_value( + args, "nystrom_current_teacher_attraction_weight", 1.0)) + self.expected_current_teacher_samples = int(getattr( + args, "expected_current_teacher_samples", 0)) + self.current_teacher_source = str(getattr( + args, "current_teacher_source", "auto")) + if self.current_teacher_source not in { + "auto", "none", "lmdb_latent", "openvid_video", + "matched_manifest_video", "matched_dual_feature_bank", + "matched_triple_feature_bank", + "matched_dinov3_vjepa2_feature_bank", + }: + raise ValueError( + "unsupported current_teacher_source: " + f"{self.current_teacher_source!r}") + self.nystrom_repulsion_fifo_size = int( + self._dense_config_value( + args, "nystrom_repulsion_fifo_size", 0)) + if ( + not np.isfinite(self.nystrom_attraction_weight) + or self.nystrom_attraction_weight < 0.0 + ): + raise ValueError( + f"{self.dense_state_prefix}_nystrom_attraction_weight must " + "be finite and non-negative") + if ( + not np.isfinite(self.current_teacher_attraction_weight) + or self.current_teacher_attraction_weight < 0.0 + ): + raise ValueError( + f"{self.dense_state_prefix}_nystrom_current_teacher_" + "attraction_weight must be finite and non-negative") + if ( + self.current_teacher_attraction_enabled + and self.nystrom_attraction_weight + + self.current_teacher_attraction_weight <= 0.0 + ): + raise ValueError( + "fixed/current teacher attraction weights must have a " + "positive sum") + if self.nystrom_repulsion_fifo_size < 0: + raise ValueError( + "dense Nyström repulsion FIFO size must be non-negative") + self.nystrom_repulsion_fifo_enabled = bool( + self.nystrom_repulsion_fifo_size > 0) + if ( + not np.isfinite(self.nystrom_attraction_ema_beta) + or not 0.0 <= self.nystrom_attraction_ema_beta < 1.0 + ): + raise ValueError( + "dense Nyström attraction EMA beta must be in [0, 1)") + if ( + self.nystrom_attraction_ema_enabled + and not self.uses_fresh_video_gradcache + ): + raise ValueError( + "vjepa2_nystrom_attraction_ema_enabled requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if ( + self.current_teacher_attraction_enabled + and not self.uses_fresh_video_gradcache + ): + raise ValueError( + "current teacher attraction requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if ( + self.current_teacher_attraction_enabled + and getattr(self, "vjepa2_multilayer_enabled", False) + ): + raise ValueError( + "current teacher attraction currently supports one dense " + "feature layer") + if self.current_teacher_attraction_enabled: + configured_teacher_total = ( + self._configured_current_teacher_global_samples(args) + ) + if self.expected_current_teacher_samples <= 0: + raise ValueError( + "expected_current_teacher_samples must be explicitly " + "positive when current teacher attraction is enabled") + if configured_teacher_total != self.expected_current_teacher_samples: + raise ValueError( + "the configured GradCache gradient geometry must match " + "expected_current_teacher_samples for exact current " + "teacher MMD: " + f"{configured_teacher_total} != " + f"{self.expected_current_teacher_samples}") + if ( + self.nystrom_repulsion_fifo_enabled + and not self.uses_fresh_video_gradcache + ): + raise ValueError( + "vjepa2_nystrom_repulsion_fifo_size requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if ( + self.nystrom_repulsion_fifo_enabled + and getattr(self, "vjepa2_multilayer_enabled", False) + ): + raise ValueError( + "repulsion FIFO currently supports one fixed V-JEPA2 layer") + if ( + self.nystrom_repulsion_fifo_enabled + and self.nystrom_attraction_ema_enabled + ): + raise ValueError( + "repulsion FIFO training requires attraction EMA to be disabled") + self.fresh_video_prepared_replay_enabled = bool( + self._fresh_video_p0_requested + and self.uses_fresh_video_gradcache + ) + if self.uses_fresh_video_gradcache: + if not self.causal_generator: + raise NotImplementedError( + "nystrom_exact_repulsion currently requires the causal " + "generator so every stochastic rollout input can be " + "captured explicitly") + from losses.video_nystrom_mmd import ( + VideoNystromMMD, + load_reference_pack, + ) + kernel_chunk_size = int( + getattr( + args, + "dinov3_nystrom_kernel_chunk_size", + 32, + ) + if self.is_dinov3_dense + else self._dense_config_value( + args, "nystrom_kernel_chunk_size", 32) + ) + if self.is_dinov3_dense: + reference_load_error = None + reference_path = None + reference_pack = None + try: + configured_reference = str(getattr( + args, + "dinov3_nystrom_reference_path", + "", + )).strip() + reference_path = ( + os.path.realpath(os.path.expanduser( + configured_reference)) + if configured_reference else "" + ) + if ( + not reference_path + or not os.path.isfile(reference_path) + ): + raise FileNotFoundError( + "DINOv3 dense Nyström requires a precomputed fixed " + "reference pack under " + "dinov3_nystrom_reference_path; " + f"got {reference_path!r}") + reference_pack = load_reference_pack( + reference_path, map_location="cpu") + except BaseException as exc: + reference_load_error = exc + self._synchronize_dinov3_local_failure( + reference_load_error, + device, + operation="DINOv3 reference-pack loading failed", + ) + self._validate_dinov3_reference_runtime_identity( + args, reference_pack["metadata"]) + objective_error = None + try: + reference_section = reference_pack["metadata"].get( + "reference", {}) + reference_source = ( + reference_section.get("source", {}) + if isinstance(reference_section, dict) + else {} + ) + feature_bank_hash = reference_source.get( + "feature_bank_fp16_sha256") + if ( + not isinstance(feature_bank_hash, str) + or len(feature_bank_hash) != 64 + or any( + character not in "0123456789abcdef" + for character in feature_bank_hash + ) + ): + raise ValueError( + "DINOv3 Nyström reference is missing a valid " + "frozen teacher feature_bank_fp16_sha256") + objective = VideoNystromMMD( + reference_pack, + expected_metadata={ + "feature": self._dinov3_feature_metadata(), + "nystrom": { + "num_landmarks": int(getattr( + args, "dinov3_nystrom_landmarks", 256)), + "seed": int(getattr( + args, "dinov3_nystrom_seed", 0)), + }, + }, + kernel_chunk_size=kernel_chunk_size, + attraction_ema_enabled=False, + repulsion_fifo_size=0, + ) + reference_count = int( + objective.reference_metadata[ + "nystrom"]["reference_count"]) + if int(reference_source.get( + "selected_samples", -1 + )) != reference_count: + raise ValueError( + "DINOv3 Nyström reference source count does not " + "match the frozen kernel reference count") + self.dinov3_nystrom_reference_path = reference_path + self.dinov3_nystrom_reference_hash = str( + objective.reference_hash) + self.dinov3_teacher_feature_bank_hash = feature_bank_hash + self.dinov3_nystrom_reference_count = reference_count + self.latent_mmd = objective.to(device=device) + except BaseException as exc: + objective_error = exc + self._synchronize_dinov3_local_failure( + objective_error, + device, + operation="DINOv3 reference objective setup failed", + ) + else: + def load_fixed_reference(layer_index, configured_path): + reference_path = ( + os.path.realpath(os.path.expanduser( + str(configured_path).strip())) + if str(configured_path).strip() else "" + ) + if ( + not reference_path + or not os.path.isfile(reference_path) + ): + raise FileNotFoundError( + "nystrom_exact_repulsion requires a precomputed " + "reference pack; got " + f"{reference_path!r} for layer {layer_index}") + reference_pack = load_reference_pack( + reference_path, map_location="cpu") + reference_section = reference_pack["metadata"].get( + "reference", {}) + reference_source = ( + reference_section.get("source", {}) + if isinstance(reference_section, dict) + else {} + ) + feature_bank_hash = reference_source.get( + "feature_bank_fp16_sha256") + if ( + not isinstance(feature_bank_hash, str) + or len(feature_bank_hash) != 64 + or any( + character not in "0123456789abcdef" + for character in feature_bank_hash + ) + ): + raise ValueError( + "Nyström reference is missing a valid frozen " + "teacher feature_bank_fp16_sha256 for layer " + f"{layer_index}") + objective = VideoNystromMMD( + reference_pack, + expected_metadata={ + "feature": self._vjepa2_feature_metadata( + layer_index=layer_index), + "nystrom": { + "num_landmarks": int( + self._dense_config_value( + args, "nystrom_landmarks", 256)), + "seed": int( + self._dense_config_value( + args, "nystrom_seed", 0)), + }, + }, + kernel_chunk_size=kernel_chunk_size, + attraction_ema_enabled=( + self.nystrom_attraction_ema_enabled), + attraction_ema_beta=( + self.nystrom_attraction_ema_beta), + attraction_weight=( + self.nystrom_attraction_weight), + current_teacher_attraction_enabled=( + self.current_teacher_attraction_enabled), + current_teacher_attraction_weight=( + self.current_teacher_attraction_weight), + repulsion_fifo_size=( + self.nystrom_repulsion_fifo_size), + ) + reference_count = int( + objective.reference_metadata[ + "nystrom"]["reference_count"]) + if int(reference_source.get( + "selected_samples", -1 + )) != reference_count: + raise ValueError( + "Nyström reference source count does not match " + "the frozen kernel reference count for layer " + f"{layer_index}") + source_identity = { + key: value + for key, value in reference_source.items() + if key != "feature_bank_fp16_sha256" + } + return ( + reference_path, + objective, + feature_bank_hash, + reference_count, + source_identity, + ) + + if getattr(self, "vjepa2_multilayer_enabled", False): + configured_paths = getattr( + args, "vjepa2_nystrom_reference_paths", None) + if ( + configured_paths is None + or isinstance(configured_paths, (str, bytes)) + ): + raise ValueError( + "multi-layer V-JEPA2 MMD requires " + "vjepa2_nystrom_reference_paths in layer order") + configured_paths = tuple(configured_paths) + if len(configured_paths) != len( + self.vjepa2_dense_layers + ): + raise ValueError( + "vjepa2_nystrom_reference_paths must match " + "vjepa2_dense_layers one-for-one") + loaded = [ + load_fixed_reference(layer, path) + for layer, path in zip( + self.vjepa2_dense_layers, + configured_paths, + ) + ] + objectives = { + layer: item[1] + for layer, item in zip( + self.vjepa2_dense_layers, loaded) + } + self.latent_mmd = MultiLayerVideoNystromMMD( + objectives, + self_normalize=self.latent_self_normalize, + norm_eps=self.latent_norm_eps, + ).to(device=device) + self.vjepa2_nystrom_reference_paths = tuple( + item[0] for item in loaded) + self.vjepa2_nystrom_reference_hashes = tuple( + objective.reference_hash + for objective in objectives.values()) + self.vjepa2_teacher_feature_bank_hashes = tuple( + item[2] for item in loaded) + reference_counts = tuple( + item[3] for item in loaded) + if len(set(reference_counts)) != 1: + raise ValueError( + "multi-layer Nyström references must contain " + "the same teacher sample count") + source_identities = tuple( + item[4] for item in loaded) + if any( + identity != source_identities[0] + for identity in source_identities[1:] + ): + raise ValueError( + "multi-layer Nyström references must use the " + "same teacher LMDB rows and prompt index") + self.vjepa2_nystrom_reference_count = ( + reference_counts[0]) + # Singular aliases keep generic diagnostics compatible; + # checkpoint identity uses the complete tuples below. + self.vjepa2_nystrom_reference_path = ( + self.vjepa2_nystrom_reference_paths[-1]) + self.vjepa2_nystrom_reference_hash = ( + self.vjepa2_nystrom_reference_hashes[-1]) + self.vjepa2_teacher_feature_bank_hash = ( + self.vjepa2_teacher_feature_bank_hashes[-1]) + else: + configured_reference_path = self._dense_config_value( + args, "nystrom_reference_path", "") + ( + reference_path, + objective, + feature_bank_hash, + reference_count, + _, + ) = load_fixed_reference( + self.vjepa2_dense_layer, + configured_reference_path, + ) + self.vjepa2_nystrom_reference_path = reference_path + self.latent_mmd = objective.to(device=device) + self.vjepa2_teacher_feature_bank_hash = ( + feature_bank_hash) + self.vjepa2_nystrom_reference_hash = str( + objective.reference_hash) + self.vjepa2_nystrom_reference_count = reference_count + if self.is_dinov3_dense and ( + not dist.is_available() + or not dist.is_initialized() + or dist.get_rank() == 0 + ): + print( + "[dinov3-nystrom] loaded frozen reference " + f"n={self.dinov3_nystrom_reference_count}, " + f"m={self.latent_mmd.landmarks.shape[0]}, " + f"sigma={float(self.latent_mmd.sigma):.6g}, " + f"hash={self.dinov3_nystrom_reference_hash}, " + "teacher_features=" + f"{self.dinov3_teacher_feature_bank_hash}", + flush=True, + ) + elif ( + not dist.is_available() + or not dist.is_initialized() + or dist.get_rank() == 0 + ): + if getattr(self, "vjepa2_multilayer_enabled", False): + reference_summary = ", ".join( + ( + f"layer={layer}:" + f"sigma={float(objective.sigma):.6g}:" + f"hash={objective.reference_hash}" + ) + for layer, objective in zip( + self.vjepa2_dense_layers, + self.latent_mmd.objectives.values(), + ) + ) + print( + "[vjepa2-nystrom] loaded frozen multi-layer " + f"references n={self.vjepa2_nystrom_reference_count}, " + f"m={self.latent_mmd.landmarks.shape[0]}, " + f"normalized={self.latent_self_normalize}, " + "attraction_ema=" + f"{self.nystrom_attraction_ema_enabled}, " + "attraction_weight=" + f"{self.nystrom_attraction_weight:g}, " + "repulsion_fifo=" + f"{self.nystrom_repulsion_fifo_size}, " + f"{reference_summary}", + flush=True, + ) + else: + print( + f"[{self.dense_encoder_name}-nystrom] loaded frozen reference " + f"n={self.vjepa2_nystrom_reference_count}, " + f"m={self.latent_mmd.landmarks.shape[0]}, " + f"sigma={float(self.latent_mmd.sigma):.6g}, " + f"hash={self.vjepa2_nystrom_reference_hash}, " + "attraction_ema=" + f"{self.nystrom_attraction_ema_enabled}, " + "repulsion_fifo=" + f"{self.nystrom_repulsion_fifo_size}, " + "teacher_features=" + f"{self.vjepa2_teacher_feature_bank_hash}", + flush=True, + ) + else: + self.latent_mmd = EMARBFMMD( + beta=float(getattr(args, "wan_kv_ema_beta", 0.99)), + num_rff_features=int( + getattr(args, "wan_kv_rff_features", 256)), + bandwidth=getattr(args, "wan_kv_rbf_bandwidth", None), + bandwidth_multipliers=getattr( + args, "wan_kv_rbf_multipliers", [0.5, 1.0, 2.0]), + seed=int(getattr(args, "wan_kv_rff_seed", 0)), + row_normalize=bool( + getattr(args, "wan_kv_row_normalize", False)), + sample_mode=( + "video_matrix" + if self.is_dense_video_feature + else "rows" + ), + ).to(device) + if self.is_triple_dense: + self._build_multi_encoder_objective(args, device) + elif self.is_dinov3_vjepa2_dense: + self._replace_with_dual_dinov3_vjepa2_nystrom_objective( + args, device) + # Dense video backends always decode videos. The Wan path only needs VAE + # loss decoding for longer causal rollouts. + self.uses_vae_for_loss = ( + self.is_dense_video_feature + or bool(self.causal_generator and self.num_training_frames > 21) + ) + @staticmethod + def _canonical_sha256(value) -> str: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + @staticmethod + def _canonical_identity_value(value): + """Normalize an artifact identity through strict canonical JSON. + + Offline identities are loaded from JSON, where sequence shapes are + lists. Live inspectors naturally return tuples for those same shapes. + Comparing the two Python containers directly creates a false mismatch + even though their canonical artifact identity is identical. + """ + try: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ValueError( + "artifact identity must contain only finite " + "JSON-serializable values" + ) from exc + return json.loads(payload) + + @staticmethod + def _file_sha256(path: str) -> str: + digest = hashlib.sha256() + with open(path, "rb") as handle: + while True: + chunk = handle.read(8 * 1024 * 1024) + if not chunk: + break + digest.update(chunk) + return digest.hexdigest() + + @classmethod + def _distributed_file_sha256(cls, path: str) -> str: + """Hash a shared file once on rank zero and broadcast the result.""" + if ( + not dist.is_available() + or not dist.is_initialized() + or dist.get_world_size() == 1 + ): + return cls._file_sha256(path) + payload = [None] + if dist.get_rank() == 0: + try: + payload[0] = { + "ok": True, + "sha256": cls._file_sha256(path), + } + except BaseException as exc: + payload[0] = { + "ok": False, + "error": f"{type(exc).__name__}: {exc}", + } + dist.broadcast_object_list(payload, src=0) + result = payload[0] + if not isinstance(result, dict) or not result.get("ok"): + error = ( + result.get("error", "unknown error") + if isinstance(result, dict) + else "no rank-zero result" + ) + raise RuntimeError(f"rank-zero file hashing failed: {error}") + return str(result["sha256"]) + + @classmethod + def _runtime_file_identity(cls, path, *, file_hasher=None) -> dict: + """Rebuild the offline tool's strict path/stat/content identity.""" + resolved = os.path.realpath(os.path.expanduser(os.fspath(path))) + if not os.path.isfile(resolved): + raise FileNotFoundError( + f"runtime identity target is not a file: {resolved}") + before = os.stat(resolved) + hasher = file_hasher or cls._distributed_file_sha256 + digest = str(hasher(resolved)) + after = os.stat(resolved) + if ( + int(before.st_size) != int(after.st_size) + or int(before.st_mtime_ns) != int(after.st_mtime_ns) + ): + raise RuntimeError( + f"runtime identity target changed while hashing: {resolved}") + if ( + len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise ValueError( + f"runtime identity returned an invalid SHA256 for {resolved}") + return { + "path": resolved, + "size_bytes": int(after.st_size), + "mtime_ns": int(after.st_mtime_ns), + "sha256": digest, + } + + @classmethod + def _distributed_runtime_file_identity(cls, path) -> dict: + """Build one shared-file identity on rank zero and publish it safely. + + Unlike calling ``_runtime_file_identity`` with a distributed hasher, + nonzero ranks do not perform a fallible ``stat`` before entering the + broadcast. Rank zero owns the complete stat/hash/stat transaction and + broadcasts either its validated identity or a serialized failure. + """ + if ( + not dist.is_available() + or not dist.is_initialized() + or dist.get_world_size() == 1 + ): + return cls._runtime_file_identity( + path, file_hasher=cls._file_sha256) + payload = [None] + if dist.get_rank() == 0: + try: + payload[0] = { + "ok": True, + "identity": cls._runtime_file_identity( + path, file_hasher=cls._file_sha256), + } + except BaseException as exc: + payload[0] = { + "ok": False, + "error": f"{type(exc).__name__}: {exc}", + } + dist.broadcast_object_list(payload, src=0) + result = payload[0] + if not isinstance(result, dict) or not result.get("ok"): + error = ( + result.get("error", "unknown error") + if isinstance(result, dict) + else "no rank-zero result" + ) + raise RuntimeError( + f"rank-zero runtime file identity failed: {error}") + identity = result.get("identity") + if not isinstance(identity, dict): + raise RuntimeError( + "rank-zero runtime file identity returned invalid data") + return dict(identity) + + def _validate_dinov3_reference_runtime_identity( + self, + args, + reference_metadata: Mapping, + *, + inspect_lmdb=None, + model_path_resolver=None, + wrapper_source_path=None, + file_hasher=None, + ) -> dict: + """Bind a DINOv3 reference to the active CausVid LMDB and Wan VAE. + + All fallible rank-local inspection completes and reaches consensus + before any rank enters a shared-file hash broadcast. The large VAE + files are then identified entirely on rank zero. Injection seams keep + the artifact checks unit-testable without loading a real LMDB or VAE. + """ + from pathlib import Path + + collective_device = getattr(self, "device", None) + if collective_device is None: + collective_device = torch.device("cpu") + preflight = None + preflight_error = None + try: + if not isinstance(reference_metadata, Mapping): + raise TypeError("DINOv3 reference metadata must be a mapping") + reference = reference_metadata.get("reference") + if not isinstance(reference, Mapping): + raise ValueError( + "DINOv3 reference metadata is missing reference identity") + source = reference.get("source") + saved_vae = reference.get("vae") + if not isinstance(source, Mapping) or not isinstance( + saved_vae, Mapping + ): + raise ValueError( + "DINOv3 reference must contain source and VAE identities") + source = dict(source) + saved_vae = dict(saved_vae) + + teacher_path = os.path.realpath(os.path.expanduser(str(getattr( + args, "teacher_latent_data_path", "")))) + if not teacher_path or not os.path.isdir(teacher_path): + raise FileNotFoundError( + "DINOv3 reference validation requires the configured " + f"teacher LMDB directory; got {teacher_path!r}") + active_inspector = inspect_lmdb + if active_inspector is None: + from tools.precompute_dinov3_dense_nystrom_reference import ( + _inspect_lmdb, + ) + + active_inspector = _inspect_lmdb + live_lmdb = dict(active_inspector(Path(teacher_path))) + live_lmdb_path = os.path.realpath(str( + live_lmdb.get("path", ""))) + live_latents_shape = [ + int(value) + for value in live_lmdb.get("latents_shape", ()) + ] + live_prompts_shape = [ + int(value) + for value in live_lmdb.get("prompts_shape", ()) + ] + if ( + live_lmdb_path != teacher_path + or len(live_latents_shape) != 6 + or live_prompts_shape != [live_latents_shape[0]] + ): + raise ValueError( + "active teacher LMDB has an invalid CausVid array identity") + teacher_count = int(getattr( + args, "teacher_latent_count", live_latents_shape[0])) + data_mdb_path = os.path.join(teacher_path, "data.mdb") + if not os.path.isfile(data_mdb_path): + raise FileNotFoundError( + f"teacher LMDB data file is missing: {data_mdb_path}") + data_stat = os.stat(data_mdb_path) + expected_source_keys = { + "kind", + "lmdb_path", + "latents_shape", + "prompts_shape", + "selected_samples", + "selected_index_start", + "selected_index_stop_exclusive", + "trajectory_index", + "lmdb_storage_dtype", + "data_mdb_size_bytes", + "data_mdb_mtime_ns", + "prompt_index_sha256", + "feature_bank_fp16_sha256", + "identity_sha256", + } + source_without_hash = dict(source) + saved_source_hash = source_without_hash.pop( + "identity_sha256", None) + source_hash_valid = ( + isinstance(saved_source_hash, str) + and saved_source_hash + == self._canonical_sha256(source_without_hash) + ) + source_sha_fields_valid = all( + isinstance(source.get(key), str) + and len(source[key]) == 64 + and all( + character in "0123456789abcdef" + for character in source[key] + ) + for key in ( + "prompt_index_sha256", + "feature_bank_fp16_sha256", + ) + ) + if ( + set(source) != expected_source_keys + or source.get("kind") + != "causvid_ode_teacher_lmdb_final_latent" + or os.path.realpath(str(source.get("lmdb_path", ""))) + != teacher_path + or source.get("latents_shape") != live_latents_shape + or source.get("prompts_shape") != live_prompts_shape + or int(source.get("selected_samples", -1)) + != live_latents_shape[0] + or live_latents_shape[0] != 1537 + or teacher_count != live_latents_shape[0] + or int(source.get("selected_index_start", -1)) != 0 + or int(source.get("selected_index_stop_exclusive", -1)) + != live_latents_shape[0] + or int(source.get("trajectory_index", -2)) != -1 + or source.get("lmdb_storage_dtype") != "float16" + or int(source.get("data_mdb_size_bytes", -1)) + != int(data_stat.st_size) + or int(source.get("data_mdb_mtime_ns", -1)) + != int(data_stat.st_mtime_ns) + or int(live_lmdb.get("data_mdb_size_bytes", -1)) + != int(data_stat.st_size) + or int(live_lmdb.get("data_mdb_mtime_ns", -1)) + != int(data_stat.st_mtime_ns) + or not source_sha_fields_valid + or not source_hash_valid + ): + raise ValueError( + "DINOv3 reference teacher LMDB identity does not match " + "the active CausVid dataset") + + active_model_path_resolver = model_path_resolver + if active_model_path_resolver is None: + from utils.paths import model_path + + active_model_path_resolver = model_path + vae_model_name = "Wan2.1-T2V-1.3B" + weights_path = ( + Path(active_model_path_resolver(vae_model_name)) + / "Wan2.1_VAE.pth" + ) + active_wrapper_source_path = wrapper_source_path + if active_wrapper_source_path is None: + active_wrapper_source_path = ( + Path(__file__).resolve().parents[1] + / "utils" / "wan_wrapper.py" + ) + # Prove every rank can resolve/stat both identities before rank zero + # enters either hash transaction. + for identity_path in (weights_path, active_wrapper_source_path): + resolved = os.path.realpath(os.path.expanduser( + os.fspath(identity_path))) + if not os.path.isfile(resolved): + raise FileNotFoundError( + f"runtime identity target is not a file: {resolved}") + os.stat(resolved) + + vae_mean = getattr(self.vae, "mean", None) + vae_std = getattr(self.vae, "std", None) + if not isinstance(vae_mean, torch.Tensor) or not isinstance( + vae_std, torch.Tensor + ): + raise TypeError( + "active Wan VAE wrapper is missing tensor mean/std " + "constants") + runtime_vae_base = { + "wrapper": ( + f"{type(self.vae).__module__}." + f"{type(self.vae).__qualname__}"), + "model_name": vae_model_name, + "latent_channels": int(vae_mean.numel()), + "latent_mean": [ + round(float(value), 4) for value in + vae_mean.detach().float().cpu().reshape(-1).tolist() + ], + "latent_std": [ + round(float(value), 4) for value in + vae_std.detach().float().cpu().reshape(-1).tolist() + ], + "mean_std_contract": "WanVAEWrapper.inline_constants.v1", + "compute_dtype": str(self.dtype).replace("torch.", ""), + "decode_use_cache": False, + "decode_output_range": "[-1,1]", + } + preflight = { + "saved_source_hash": saved_source_hash, + "saved_vae": saved_vae, + "weights_path": weights_path, + "wrapper_source_path": active_wrapper_source_path, + "runtime_vae_base": runtime_vae_base, + } + except BaseException as exc: + preflight_error = exc + self._synchronize_dinov3_local_failure( + preflight_error, + collective_device, + operation="DINOv3 runtime identity preflight failed", + ) + + identity_error = None + weights_identity = None + wrapper_identity = None + try: + if file_hasher is None: + weights_identity = self._distributed_runtime_file_identity( + preflight["weights_path"]) + wrapper_identity = self._distributed_runtime_file_identity( + preflight["wrapper_source_path"]) + else: + weights_identity = self._runtime_file_identity( + preflight["weights_path"], file_hasher=file_hasher) + wrapper_identity = self._runtime_file_identity( + preflight["wrapper_source_path"], + file_hasher=file_hasher, + ) + except BaseException as exc: + identity_error = exc + self._synchronize_dinov3_local_failure( + identity_error, + collective_device, + operation="DINOv3 shared-file identity failed", + ) + + receipt = None + comparison_error = None + try: + runtime_vae = dict(preflight["runtime_vae_base"]) + runtime_vae["wrapper_source"] = wrapper_identity + runtime_vae["weights"] = weights_identity + runtime_vae["identity_sha256"] = self._canonical_sha256( + runtime_vae) + if preflight["saved_vae"] != runtime_vae: + raise ValueError( + "DINOv3 reference Wan VAE identity does not match the " + "active decoder") + receipt = { + "source_identity_sha256": preflight["saved_source_hash"], + "vae_identity_sha256": runtime_vae["identity_sha256"], + } + except BaseException as exc: + comparison_error = exc + self._synchronize_dinov3_local_failure( + comparison_error, + collective_device, + operation="DINOv3 runtime identity comparison failed", + ) + return receipt + + @classmethod + def _resolve_vjepa2_model_configuration( + cls, + args, + ) -> Tuple[str, str, str]: + """Resolve the load source independently from reference metadata. + + vjepa2_model_id identifies the model that transformers actually opens. + vjepa2_reference_model_id is the canonical identity stored in an + already-created feature/reference pack; it can legitimately be a path + from the machine that produced that pack. Keeping the two separate + lets a byte-identical local mirror satisfy strict reference metadata + without asking transformers to open the producer's path. + + When configured, vjepa2_model_sha256 authenticates the local + model.safetensors before the adapter is constructed. + """ + configured_load_id = str(getattr( + args, + "vjepa2_model_id", + "assets/models/vjepa2-vitl-fpc64-256", + )).strip() + if not configured_load_id: + raise ValueError("vjepa2_model_id must be a non-empty string") + expanded_load_id = os.path.expanduser(configured_load_id) + model_load_path = ( + os.path.realpath(expanded_load_id) + if os.path.exists(expanded_load_id) + else configured_load_id + ) + + configured_reference_id = getattr( + args, "vjepa2_reference_model_id", None) + reference_model_id = ( + "" if configured_reference_id is None + else str(configured_reference_id).strip() + ) + if not reference_model_id: + # Backward compatibility: old configs used one field for both + # loading and canonical reference metadata. + reference_model_id = model_load_path + + expected_sha256 = str(getattr( + args, "vjepa2_model_sha256", "")).strip().lower() + if expected_sha256: + if ( + len(expected_sha256) != 64 + or any( + character not in "0123456789abcdef" + for character in expected_sha256 + ) + ): + raise ValueError( + "vjepa2_model_sha256 must be a 64-character SHA-256") + if not os.path.isdir(model_load_path): + raise ValueError( + "vjepa2_model_sha256 requires vjepa2_model_id to be a " + "local model directory") + weight_path = os.path.join( + model_load_path, "model.safetensors") + if not os.path.isfile(weight_path): + raise FileNotFoundError( + "vjepa2_model_sha256 requires the local weight file " + f"{weight_path}") + actual_sha256 = cls._distributed_file_sha256(weight_path) + if actual_sha256 != expected_sha256: + raise RuntimeError( + "V-JEPA2 model.safetensors SHA-256 mismatch: " + f"expected {expected_sha256}, got {actual_sha256}") + + return model_load_path, reference_model_id, expected_sha256 + + @staticmethod + def _distributed_representation_identity(module) -> dict: + """Hash one replicated encoder once and share its strict identity.""" + if ( + not dist.is_available() + or not dist.is_initialized() + or dist.get_world_size() == 1 + ): + return representation_runtime_identity(module) + payload = [None] + if dist.get_rank() == 0: + try: + payload[0] = { + "ok": True, + "identity": representation_runtime_identity(module), + } + except BaseException as exc: + payload[0] = { + "ok": False, + "error": f"{type(exc).__name__}: {exc}", + } + dist.broadcast_object_list(payload, src=0) + result = payload[0] + if not isinstance(result, dict) or not result.get("ok"): + error = ( + result.get("error", "unknown error") + if isinstance(result, dict) + else "no rank-zero result" + ) + raise RuntimeError( + f"rank-zero representation hashing failed: {error}") + return dict(result["identity"]) + + @staticmethod + def _config_list(value) -> list: + if isinstance(value, str): + return [value] + return list(value) + + def _dense_config(self, args, suffix: str, default): + """Read an encoder-scoped dense option without changing old YAMLs.""" + return getattr( + args, f"{self._dense_encoder_kind()}_{suffix}", default) + + def _dense_key(self, suffix: str) -> str: + prefix = getattr( + self, + "dense_log_prefix", + f"{self._dense_encoder_kind()}_dense", + ) + return f"{prefix}_{suffix}" + + def _dense_encoder_kind(self) -> str: + return getattr( + self, + "dense_encoder_name", + ( + "motionformer" + if getattr(self, "feature_type", None) + == "motionformer_dense_latent" + else "vjepa2" + ), + ) + + def _dense_encoder_display_name(self) -> str: + return ( + "MotionFormer" + if self._dense_encoder_kind() == "motionformer" + else "V-JEPA2" + ) + + def _dense_attr(self, suffix: str): + generic_name = f"dense_{suffix}" + if hasattr(self, generic_name): + return getattr(self, generic_name) + return getattr(self, f"vjepa2_dense_{suffix}") + + def _dense_ema_version_key(self) -> str: + return self._dense_key("mmd_version") + + def _dense_ema_state_key(self) -> str: + return self._dense_key("mmd") + + def _dense_nystrom_version_key(self) -> str: + return self._dense_key("nystrom_mmd_version") + + def _initialize_wan_features(self, args) -> None: + if self.real_score is None: + raise ValueError( + "post_block_latent requires uses_real_score: true") + self.latent_layer = int(getattr(args, "wan_latent_layer", 15)) + num_encoder_blocks = len(self.real_score.model.blocks) + if self.latent_layer < 0 or self.latent_layer >= num_encoder_blocks: + raise ValueError( + f"wan_latent_layer={self.latent_layer} outside frozen Wan's " + f"0..{num_encoder_blocks - 1} range") + if bool(getattr(args, "wan_latent_prune_unused_blocks", True)): + # The feature path returns after the requested block, so later + # frozen blocks can never be used. Removing them before FSDP avoids + # sharding/loading another ~650M parameters for layers 16..29. + keep = self.latent_layer + 1 + self.real_score.model.blocks = self.real_score.model.blocks[:keep] + self.latent_num_tokens = int( + getattr(args, "wan_latent_num_tokens", 128)) + if self.latent_num_tokens <= 1: + raise ValueError("wan_latent_num_tokens must be > 1") + self.latent_timestep = float( + getattr(args, "wan_latent_timestep", 0.0)) + self.latent_patch_size = tuple( + int(v) for v in self.real_score.model.patch_size) + self.latent_feature_dim = int(self.real_score.model.dim) + # We need input gradients through the frozen encoder, and checkpointing is + # the main activation-memory control for deeper prefixes. + if bool(getattr(args, "wan_latent_gradient_checkpointing", True)): + self.real_score.enable_gradient_checkpointing() + + def _initialize_vjepa2_features(self, args, device) -> None: + # Lazy import keeps the transformers/V-JEPA2 dependency completely out + # of the unchanged Wan feature path. + from video_fd.encoders.vjepa2_adapter import VJEPA2Adapter + + ( + self.vjepa2_model_load_path, + self.vjepa2_reference_model_id, + self.vjepa2_model_sha256, + ) = self._resolve_vjepa2_model_configuration(args) + # Retain the historical runtime attribute for reference-bank identity + # checks and downstream code that expects the actual load source. + self.vjepa2_model_id = self.vjepa2_model_load_path + configured_dense_layers = getattr( + args, "vjepa2_dense_layers", None) + if configured_dense_layers is None: + dense_layers = ( + int(getattr(args, "vjepa2_dense_layer", 15)), + ) + else: + if isinstance(configured_dense_layers, (str, bytes)): + raise TypeError( + "vjepa2_dense_layers must be a sequence of integers") + dense_layers = tuple( + int(layer) for layer in configured_dense_layers) + if not dense_layers: + raise ValueError( + "vjepa2_dense_layers must not be empty") + if len(set(dense_layers)) != len(dense_layers): + raise ValueError( + "vjepa2_dense_layers must contain unique indices") + if tuple(sorted(dense_layers)) != dense_layers: + raise ValueError( + "vjepa2_dense_layers must be sorted") + self.vjepa2_dense_layers = dense_layers + self.vjepa2_multilayer_enabled = len(dense_layers) > 1 + # The maximum layer controls prefix pruning. Single-layer callers keep + # the historical scalar attribute unchanged. + self.vjepa2_dense_layer = dense_layers[-1] + self.vjepa2_dense_num_frames = int( + getattr(args, "vjepa2_dense_num_frames", 64)) + self.vjepa2_dense_frame_sampling = str(getattr( + args, "vjepa2_dense_frame_sampling", "center_contiguous")) + self.vjepa2_dense_source_num_frames = int(getattr( + args, + "vjepa2_dense_source_num_frames", + self.vjepa2_dense_num_frames, + )) + configured_student_frame_sampling = getattr( + args, "vjepa2_student_frame_sampling", None) + if self.vjepa2_dense_frame_sampling == "last_contiguous": + if ( + configured_student_frame_sampling is not None + and str(configured_student_frame_sampling) != "last_contiguous" + ): + raise ValueError( + "vjepa2_dense_frame_sampling=last_contiguous conflicts " + "with vjepa2_student_frame_sampling") + # The dense contract identifies the frozen center64 reference. + # Treat this alias as a student-only last64 request. + self.vjepa2_dense_frame_sampling = "center_contiguous" + self.vjepa2_student_frame_sampling = "last_contiguous" + else: + self.vjepa2_student_frame_sampling = str( + configured_student_frame_sampling + if configured_student_frame_sampling is not None + else self.vjepa2_dense_frame_sampling + ) + default_student_source_frames = ( + 81 + if self.vjepa2_student_frame_sampling == "last_contiguous" + else self.vjepa2_dense_source_num_frames + ) + self.vjepa2_student_source_num_frames = int(getattr( + args, "vjepa2_student_source_num_frames", + default_student_source_frames, + )) + self.vjepa2_dense_window_seed = str(getattr( + args, "vjepa2_dense_window_seed", "sf-vjepa2-window-v1")) + self.vjepa2_dense_num_tokens = int( + getattr(args, "vjepa2_dense_num_tokens", 128)) + self.vjepa2_dense_resize_short_side = int( + getattr(args, "vjepa2_dense_resize_short_side", 292)) + self.vjepa2_dense_crop_size = int( + getattr(args, "vjepa2_dense_crop_size", 256)) + self.vjepa2_dense_gradient_checkpointing = bool( + getattr(args, "vjepa2_dense_gradient_checkpointing", True)) + + if any(layer < 0 for layer in self.vjepa2_dense_layers): + raise ValueError( + "V-JEPA2 dense layer indices must be non-negative") + if self.vjepa2_dense_num_frames <= 0: + raise ValueError("vjepa2_dense_num_frames must be positive") + if self.vjepa2_dense_frame_sampling not in { + "center_contiguous", + "random_contiguous_per_video", + }: + raise ValueError( + "vjepa2_dense_frame_sampling must be center_contiguous or " + "random_contiguous_per_video") + if self.vjepa2_student_frame_sampling not in { + "center_contiguous", + "random_contiguous_per_video", + "last_contiguous", + }: + raise ValueError( + "vjepa2_student_frame_sampling must be center_contiguous, " + "random_contiguous_per_video, or last_contiguous") + if ( + self.vjepa2_student_frame_sampling + != self.vjepa2_dense_frame_sampling + and self.vjepa2_student_frame_sampling != "last_contiguous" + ): + raise ValueError( + "only last_contiguous may intentionally differ from the " + "frozen-reference frame sampling") + if ( + self.vjepa2_student_frame_sampling == "last_contiguous" + and self.vjepa2_dense_frame_sampling != "center_contiguous" + ): + raise ValueError("last_contiguous requires a center64 reference") + if ( + self.vjepa2_student_source_num_frames + < self.vjepa2_dense_num_frames + ): + raise ValueError( + "vjepa2_student_source_num_frames must be at least " + "vjepa2_dense_num_frames") + if self.vjepa2_student_frame_sampling == "last_contiguous" and ( + self.vjepa2_dense_num_frames != 64 + or self.vjepa2_student_source_num_frames != 81 + ): + raise ValueError( + "last_contiguous is the strict 81-frame last64 contract and " + "requires vjepa2_dense_num_frames=64 plus " + "vjepa2_student_source_num_frames=81") + if self.vjepa2_dense_source_num_frames < self.vjepa2_dense_num_frames: + raise ValueError( + "vjepa2_dense_source_num_frames must be at least " + "vjepa2_dense_num_frames") + if not self.vjepa2_dense_window_seed: + raise ValueError( + "vjepa2_dense_window_seed must be a non-empty string") + if ( + self.vjepa2_student_frame_sampling + == "random_contiguous_per_video" + and self.mmd_estimator != "nystrom_exact_repulsion" + ): + raise ValueError( + "random V-JEPA2 temporal windows require " + "wan_mmd_estimator=nystrom_exact_repulsion") + if self.vjepa2_dense_num_tokens <= 1: + raise ValueError("vjepa2_dense_num_tokens must be > 1") + if ( + self.vjepa2_dense_crop_size <= 0 + or self.vjepa2_dense_resize_short_side + < self.vjepa2_dense_crop_size + ): + raise ValueError( + "V-JEPA2 resize_short_side must be >= positive crop_size") + + adapter_config = { + "name": "vjepa2_dense", + "model_id": self.vjepa2_model_load_path, + "input_frames": self.vjepa2_dense_num_frames, + "dense_num_frames": self.vjepa2_dense_num_frames, + "input_size": self.vjepa2_dense_crop_size, + "resize_short_side": self.vjepa2_dense_resize_short_side, + "crop_mode": "center", + "dense_only": True, + "dense_layer_index": self.vjepa2_dense_layer, + } + self.vjepa2_encoder = VJEPA2Adapter(adapter_config) + self.vjepa2_encoder.to(device=device, dtype=torch.float32) + self.vjepa2_encoder.eval().requires_grad_(False) + + actual_geometry = ( + int(self.vjepa2_encoder.dense_layer_index), + int(self.vjepa2_encoder.dense_num_frames), + int(self.vjepa2_encoder.resize_short_side), + int(self.vjepa2_encoder.input_size), + int(self.vjepa2_encoder.retained_dense_blocks), + ) + expected_geometry = ( + self.vjepa2_dense_layer, + self.vjepa2_dense_num_frames, + self.vjepa2_dense_resize_short_side, + self.vjepa2_dense_crop_size, + self.vjepa2_dense_layer + 1, + ) + if actual_geometry != expected_geometry: + raise RuntimeError( + "loaded V-JEPA2 adapter geometry does not match the requested " + "dense-prefix configuration: " + f"actual={actual_geometry}, expected={expected_geometry}") + self.vjepa2_dense_raw_tokens = int( + self.vjepa2_encoder.dense_raw_tokens) + self.vjepa2_dense_feature_dim = int( + self.vjepa2_encoder.dense_feature_dim) + if self.vjepa2_dense_num_tokens > self.vjepa2_dense_raw_tokens: + raise ValueError( + "vjepa2_dense_num_tokens cannot exceed the full dense grid: " + f"{self.vjepa2_dense_num_tokens} > " + f"{self.vjepa2_dense_raw_tokens}") + self._bind_dense_runtime_aliases() + + def _initialize_motionformer_features(self, args, device) -> None: + """Build the optional frozen MotionFormer dense-prefix extractor.""" + from video_fd.encoders.motionformer_adapter import ( + MotionFormerAdapter, + SSV2_MOTIONFORMER_224_16X4_SHA256, + ) + + configured_path = str(getattr( + args, "motionformer_checkpoint_path", "")).strip() + checkpoint_path = os.path.realpath( + os.path.expanduser(configured_path)) if configured_path else "" + if not checkpoint_path or not os.path.isfile(checkpoint_path): + raise FileNotFoundError( + "wan_feature_type=motionformer_dense_latent requires " + "motionformer_checkpoint_path; got " + f"{checkpoint_path!r}") + checkpoint_sha256 = str(getattr( + args, "motionformer_checkpoint_sha256", "")).strip().lower() + if ( + len(checkpoint_sha256) != 64 + or any(c not in "0123456789abcdef" for c in checkpoint_sha256) + ): + raise ValueError( + "motionformer_checkpoint_sha256 must be a lowercase SHA-256") + if checkpoint_sha256 != SSV2_MOTIONFORMER_224_16X4_SHA256: + raise ValueError( + "motionformer_checkpoint_sha256 must identify the official " + "SSv2 Base 16x4 checkpoint") + + self.motionformer_checkpoint_path = checkpoint_path + self.motionformer_checkpoint_sha256 = checkpoint_sha256 + self.motionformer_dense_layer = int(getattr( + args, "motionformer_dense_layer", 11)) + self.motionformer_dense_num_frames = int(getattr( + args, "motionformer_dense_num_frames", 16)) + self.motionformer_dense_num_tokens = int(getattr( + args, "motionformer_dense_num_tokens", 128)) + self.motionformer_dense_resize_short_side = int(getattr( + args, "motionformer_dense_resize_short_side", 224)) + self.motionformer_dense_crop_size = int(getattr( + args, "motionformer_dense_crop_size", 224)) + self.motionformer_dense_gradient_checkpointing = bool(getattr( + args, "motionformer_dense_gradient_checkpointing", True)) + if self.motionformer_dense_layer != 11: + raise ValueError( + "the supported SSv2 MotionFormer feature contract fixes " + "motionformer_dense_layer=11") + if self.motionformer_dense_num_frames != 16: + raise ValueError( + "the supported SSv2 MotionFormer feature contract fixes " + "motionformer_dense_num_frames=16") + if ( + self.motionformer_dense_resize_short_side != 224 + or self.motionformer_dense_crop_size != 224 + ): + raise ValueError( + "the supported SSv2 MotionFormer preprocessing contract fixes " + "resize_short_side=crop_size=224") + if self.motionformer_dense_num_tokens <= 1: + raise ValueError("motionformer_dense_num_tokens must be > 1") + + adapter_config = { + "name": "motionformer_dense", + "checkpoint_path": checkpoint_path, + "checkpoint_sha256": checkpoint_sha256, + "input_frames": self.motionformer_dense_num_frames, + "dense_num_frames": self.motionformer_dense_num_frames, + "input_size": self.motionformer_dense_crop_size, + "resize_short_side": self.motionformer_dense_resize_short_side, + "crop_mode": "center", + "dense_only": True, + "dense_layer_index": self.motionformer_dense_layer, + "use_original_trajectory_attention": True, + } + self.motionformer_encoder = MotionFormerAdapter(adapter_config) + self.motionformer_encoder.to(device=device, dtype=torch.float32) + self.motionformer_encoder.eval().requires_grad_(False) + actual_geometry = ( + int(self.motionformer_encoder.dense_layer_index), + int(self.motionformer_encoder.dense_num_frames), + int(self.motionformer_encoder.input_size), + int(self.motionformer_encoder.retained_dense_blocks), + ) + expected_geometry = ( + self.motionformer_dense_layer, + self.motionformer_dense_num_frames, + self.motionformer_dense_crop_size, + self.motionformer_dense_layer + 1, + ) + if actual_geometry != expected_geometry: + raise RuntimeError( + "loaded MotionFormer adapter geometry does not match the " + "requested dense-prefix configuration: " + f"actual={actual_geometry}, expected={expected_geometry}") + self.motionformer_dense_raw_tokens = int( + self.motionformer_encoder.dense_raw_tokens) + self.motionformer_dense_feature_dim = int( + self.motionformer_encoder.dense_feature_dim) + if self.motionformer_dense_raw_tokens != 1568: + raise RuntimeError( + "SSv2 MotionFormer must expose 1568 non-CLS dense tokens, got " + f"{self.motionformer_dense_raw_tokens}") + if self.motionformer_dense_feature_dim != 768: + raise RuntimeError( + "SSv2 MotionFormer must expose 768-dimensional features, got " + f"{self.motionformer_dense_feature_dim}") + if ( + self.motionformer_dense_num_tokens + > self.motionformer_dense_raw_tokens + ): + raise ValueError( + "motionformer_dense_num_tokens cannot exceed the full dense " + f"grid: {self.motionformer_dense_num_tokens} > " + f"{self.motionformer_dense_raw_tokens}") + + # The replay/MMD implementation below deliberately consumes one common + # dense-encoder contract. These private aliases keep the already + # validated V-JEPA2 implementation and its legacy tests byte-compatible; + # public generic aliases are what the trainer uses. + self.vjepa2_model_id = checkpoint_path + self.vjepa2_dense_layer = self.motionformer_dense_layer + self.vjepa2_dense_layers = (self.motionformer_dense_layer,) + self.vjepa2_multilayer_enabled = False + self.vjepa2_dense_num_frames = self.motionformer_dense_num_frames + # MotionFormer receives the complete decoded video and performs its + # official 16-segment midpoint sampling inside the adapter. Mark this + # as non-random so the V-JEPA2-only contiguous-window machinery stays + # dormant. + self.vjepa2_dense_frame_sampling = "center_contiguous" + self.vjepa2_dense_source_num_frames = ( + self.motionformer_dense_num_frames) + self.vjepa2_dense_window_seed = "" + self.vjepa2_dense_num_tokens = self.motionformer_dense_num_tokens + self.vjepa2_dense_resize_short_side = ( + self.motionformer_dense_resize_short_side) + self.vjepa2_dense_crop_size = self.motionformer_dense_crop_size + self.vjepa2_dense_gradient_checkpointing = ( + self.motionformer_dense_gradient_checkpointing) + self.vjepa2_encoder = self.motionformer_encoder + self.vjepa2_dense_raw_tokens = self.motionformer_dense_raw_tokens + self.vjepa2_dense_feature_dim = self.motionformer_dense_feature_dim + self._bind_dense_runtime_aliases() + + def _bind_dense_runtime_aliases(self) -> None: + """Publish the encoder-neutral runtime interface used by DistributionMatchingTrainer.""" + self.dense_encoder = self.vjepa2_encoder + self.dense_layer = self.vjepa2_dense_layer + self.dense_num_frames = self.vjepa2_dense_num_frames + self.dense_num_tokens = self.vjepa2_dense_num_tokens + self.dense_resize_short_side = self.vjepa2_dense_resize_short_side + self.dense_crop_size = self.vjepa2_dense_crop_size + self.dense_gradient_checkpointing = ( + self.vjepa2_dense_gradient_checkpointing) + self.dense_raw_tokens = self.vjepa2_dense_raw_tokens + self.dense_feature_dim = self.vjepa2_dense_feature_dim + + def _initialize_dual_dinov3_vjepa2_features( + self, + args, + device, + ) -> None: + """Build only the frozen DINOv3 and V-JEPA2 prefixes.""" + self.fresh_video_total_samples = int(getattr( + args, "total_batch_size", 0)) + if self.fresh_video_total_samples <= 0: + raise ValueError( + "DINOv3+V-JEPA2 MMD requires positive total_batch_size") + self._initialize_vjepa2_features(args, device) + expected_vjepa = { + "layers": (15,), + "frames": 64, + "source_frames": 81, + "student_sampling": "last_contiguous", + "tokens": 1024, + "feature_dim": 1024, + } + actual_vjepa = { + "layers": tuple(self.vjepa2_dense_layers), + "frames": int(self.vjepa2_dense_num_frames), + "source_frames": int(self.vjepa2_student_source_num_frames), + "student_sampling": str(self.vjepa2_student_frame_sampling), + "tokens": int(self.vjepa2_dense_num_tokens), + "feature_dim": int(self.vjepa2_dense_feature_dim), + } + if actual_vjepa != expected_vjepa: + raise ValueError( + "dual DINOv3+V-JEPA2 contract fixes V-JEPA2 layer15, " + "frames17..80, K=1024 and D=1024: " + f"actual={actual_vjepa}, expected={expected_vjepa}") + self._initialize_dinov3_features(args, device) + actual_dinov3 = { + "geometry": self.dinov3_dense_geometry, + "source_frames": self.dinov3_dense_source_num_frames, + "sample_shape": ( + self.dinov3_dense_num_tokens, + self.dinov3_dense_feature_dim, + ), + } + expected_dinov3 = { + "geometry": "temporal_mean5_full256", + "source_frames": 81, + "sample_shape": (1024, 1024), + } + if actual_dinov3 != expected_dinov3: + raise ValueError( + "dual DINOv3+V-JEPA2 contract fixes DINOv3 to 20 frames, " + "mean5, K=1024 and D=1024: " + f"actual={actual_dinov3}, expected={expected_dinov3}") + self.dual_branch_registry = { + "dinov3": { + "encoder_attribute": "dinov3_encoder", + "source_frame_indices": list( + self.dinov3_dense_frame_indices), + "sample_shape": (1024, 1024), + "tokens_per_temporal_group": 256, + "temporal_groups": 4, + "model_id": self.dinov3_model_id, + "source_commit": self.dinov3_source_commit, + "checkpoint_sha256": getattr( + self.dinov3_encoder, "checkpoint_sha256", ""), + }, + "vjepa2": { + "encoder_attribute": "vjepa2_encoder", + "layer": 15, + "source_frame_range": [17, 80], + "sample_shape": (1024, 1024), + "tokens_per_segment": 1024, + "segments": 1, + "model_id": self.vjepa2_reference_model_id, + "model_sha256": self.vjepa2_model_sha256, + }, + } + + def _replace_with_dual_dinov3_vjepa2_nystrom_objective( + self, args, device + ) -> None: + """Bind two Wan-video references and sum normalized branch MMDs.""" + from losses.multi_encoder_mmd import ( + MultiEncoderMMD, + ) + from losses.video_nystrom_mmd import ( + VideoNystromMMD, + load_reference_pack, + ) + + if not self.uses_fresh_video_gradcache: + raise ValueError( + "DINOv3+V-JEPA2 MMD requires nystrom_exact_repulsion") + if not self.latent_self_normalize: + raise ValueError( + "DINOv3+V-JEPA2 MMD requires self normalization") + if ( + self.nystrom_attraction_ema_enabled + or self.nystrom_repulsion_fifo_enabled + ): + raise ValueError( + "DINOv3+V-JEPA2 MMD requires EMA and FIFO disabled") + if not self.current_teacher_attraction_enabled: + raise ValueError( + "DINOv3+V-JEPA2 MMD requires matched random teachers") + if ( + self.nystrom_attraction_weight != 2.0 + or self.current_teacher_attraction_weight != 1.0 + ): + raise ValueError( + "DINOv3+V-JEPA2 fixed/random teacher weights must be 2:1") + if self.expected_current_teacher_samples != 128: + raise ValueError( + "DINOv3+V-JEPA2 objective requires 128 matched teachers") + + def strict_sha256(value, label): + normalized = str(value).strip().lower() + if ( + len(normalized) != 64 + or any(c not in "0123456789abcdef" for c in normalized) + ): + raise ValueError(f"{label} must be a lowercase SHA-256") + return normalized + + def checked_reference_path(path_value, sha_value, label): + path = os.path.realpath(os.path.expanduser(str(path_value))) + if not os.path.isfile(path): + raise FileNotFoundError( + f"DINOv3+V-JEPA2 {label} reference missing: {path!r}") + expected_sha = strict_sha256(sha_value, f"{label} reference SHA-256") + actual_sha = self._distributed_file_sha256(path) + if actual_sha != expected_sha: + raise ValueError( + f"{label} reference SHA-256 mismatch: " + f"{actual_sha} != {expected_sha}") + return path, expected_sha, load_reference_pack( + path, map_location="cpu") + + # Each encoder owns its reference summary; the wrapper combines normalized losses. + vjepa_objective = self.latent_mmd + if not isinstance(vjepa_objective, VideoNystromMMD): + raise TypeError( + "dual bootstrap V-JEPA2 objective must be VideoNystromMMD") + vjepa_path = self.vjepa2_nystrom_reference_path + vjepa_file_sha = strict_sha256( + self._dense_config_value( + args, "nystrom_reference_sha256", ""), + "V-JEPA2 reference SHA-256", + ) + if self._distributed_file_sha256(vjepa_path) != vjepa_file_sha: + raise ValueError( + "dual V-JEPA2 reference file SHA-256 mismatch") + + dinov3_path, dinov3_file_sha, dinov3_pack = checked_reference_path( + getattr(args, "dinov3_nystrom_reference_path", ""), + getattr(args, "dinov3_nystrom_reference_sha256", ""), + "DINOv3", + ) + dinov3_objective = VideoNystromMMD( + dinov3_pack, + expected_metadata={ + "feature": self._dinov3_feature_metadata(), + "nystrom": { + "num_landmarks": int(getattr( + args, "dinov3_nystrom_landmarks", 256)), + "seed": int(getattr(args, "dinov3_nystrom_seed", 0)), + }, + }, + kernel_chunk_size=int(getattr( + args, "dinov3_nystrom_kernel_chunk_size", 32)), + attraction_weight=2.0, + current_teacher_attraction_enabled=True, + current_teacher_attraction_weight=1.0, + attraction_ema_enabled=False, + repulsion_fifo_size=0, + ).to(device=device) + objectives = { + "dinov3": dinov3_objective, + "vjepa2": vjepa_objective, + } + expected_shapes = { + "dinov3": (1024, 1024), + "vjepa2": (1024, 1024), + } + actual_shapes = { + name: tuple(objective.sample_shape) + for name, objective in objectives.items() + } + if actual_shapes != expected_shapes: + raise ValueError( + "dual reference sample-shape mismatch: " + f"actual={actual_shapes}, expected={expected_shapes}") + + configured_bank_hash = strict_sha256( + getattr(args, "dinov3_vjepa2_teacher_bank_content_hash", ""), + "dinov3_vjepa2_teacher_bank_content_hash", + ) + expected_manifest_hash = strict_sha256( + getattr(args, "matched_manifest_sha256", ""), + "matched_manifest_sha256", + ) + for name, objective in objectives.items(): + source = objective.reference_metadata.get( + "reference", {}).get("source", {}) + if ( + source.get("dinov3_vjepa2_teacher_bank_content_hash") + != configured_bank_hash + or source.get("source_manifest_sha256") + != expected_manifest_hash + or int(source.get("selected_samples", -1)) != 5000 + or source.get("branch_name") != name + ): + raise ValueError( + f"{name} reference is not bound to the configured " + "frozen 5000-Wan-video DINOv3+V-JEPA2 bank") + # Paper Eq. (3) motivates encoder weighting; this release also uses detached-scale normalization. + branch_weights = { + "dinov3": float(getattr( + args, "dinov3_vjepa2_dinov3_weight", 10.0)), + "vjepa2": float(getattr( + args, "dinov3_vjepa2_vjepa2_weight", 1.0)), + } + if branch_weights != {"dinov3": 10.0, "vjepa2": 1.0}: + raise ValueError( + "DINOv3+V-JEPA2 representation weights must be exactly 10:1") + self.latent_mmd = MultiEncoderMMD( + objectives, + weights=branch_weights, + normalization_eps=self.latent_norm_eps, + ).to(device=device) + self.dinov3_nystrom_reference_path = dinov3_path + self.dinov3_nystrom_reference_hash = dinov3_objective.reference_hash + self.dinov3_vjepa2_teacher_bank_content_hash = configured_bank_hash + self.dual_teacher_bank_content_hash = configured_bank_hash + reference_files = { + "dinov3": dinov3_file_sha, + "vjepa2": vjepa_file_sha, + } + reference_paths = { + "dinov3": dinov3_path, + "vjepa2": vjepa_path, + } + for name, objective in objectives.items(): + self.dual_branch_registry[name].update({ + "reference_path": reference_paths[name], + "reference_file_sha256": reference_files[name], + "reference_hash": objective.reference_hash, + "nystrom_landmarks": int(objective.landmarks.shape[0]), + }) + + def _initialize_triple_t5_joint_configuration(self, args) -> None: + """Freeze the UMT5 row-concatenation protocol before encoders load.""" + if not bool(getattr(args, "joint_t5_enabled", True)): + raise ValueError( + "the triple T5 joint feature type requires joint_t5_enabled") + configured = getattr(args, "joint_t5_num_tokens_by_branch", None) + if not isinstance(configured, Mapping): + raise TypeError( + "joint_t5_num_tokens_by_branch must be a branch mapping") + branch_tokens = { + str(name): int(value) for name, value in configured.items() + } + expected_tokens = { + "dinov3": 85, + "vjepa2": 85, + "videomae_v2": 107, + } + if branch_tokens != expected_tokens: + raise ValueError( + "triple joint T5 token counts must be exactly " + f"{expected_tokens}, got {branch_tokens}") + sampling = str(getattr( + args, + "joint_t5_token_sampling", + "uniform_valid_linspace_round_zero_pad", + )) + if sampling != "uniform_valid_linspace_round_zero_pad": + raise ValueError( + "joint_t5_token_sampling must be " + "uniform_valid_linspace_round_zero_pad") + layout = str(getattr( + args, + "joint_t5_layout", + "reshape_4096_to_4x1024_then_row_concat", + )) + if layout != "reshape_4096_to_4x1024_then_row_concat": + raise ValueError( + "joint_t5_layout must be " + "reshape_4096_to_4x1024_then_row_concat") + + def strict_sha256(value, label): + normalized = str(value).strip().lower() + if ( + len(normalized) != 64 + or any(c not in "0123456789abcdef" for c in normalized) + ): + raise ValueError(f"{label} must be a lowercase SHA-256") + return normalized + + self.joint_t5_num_tokens_by_branch = branch_tokens + self.joint_t5_token_sampling = sampling + self.joint_t5_layout = layout + self.joint_t5_raw_tokens = 512 + self.joint_t5_feature_dim = 4096 + self.joint_t5_model_id = str(getattr(args, "joint_t5_model_id")) + self.joint_t5_tokenizer_id = str( + getattr(args, "joint_t5_tokenizer_id")) + self.joint_t5_reference_model_id = str(getattr( + args, "joint_t5_reference_model_id", self.joint_t5_model_id)) + self.joint_t5_reference_tokenizer_id = str(getattr( + args, + "joint_t5_reference_tokenizer_id", + self.joint_t5_tokenizer_id, + )) + self.joint_t5_weights_sha256 = strict_sha256( + getattr(args, "joint_t5_weights_sha256"), + "joint_t5_weights_sha256", + ) + self.joint_t5_tokenizer_sha256 = strict_sha256( + getattr(args, "joint_t5_tokenizer_sha256"), + "joint_t5_tokenizer_sha256", + ) + + def _initialize_multi_encoder_features( + self, + args, + device, + ) -> None: + """Build the frozen feature encoders and their video/token geometry. + + The active encoder list distinguishes the three-encoder 1.3B release + from the two-encoder 14B release; saved attribute names stay compatible.""" + self.fresh_video_total_samples = int(getattr( + args, "total_batch_size", 0)) + if self.fresh_video_total_samples <= 0: + raise ValueError( + "triple video MMD requires positive total_batch_size") + self._initialize_vjepa2_features(args, device) + vjepa_attribute_names = ( + "vjepa2_model_load_path", + "vjepa2_reference_model_id", + "vjepa2_model_sha256", + "vjepa2_model_id", + "vjepa2_dense_layers", + "vjepa2_multilayer_enabled", + "vjepa2_dense_layer", + "vjepa2_dense_num_frames", + "vjepa2_dense_frame_sampling", + "vjepa2_dense_source_num_frames", + "vjepa2_student_frame_sampling", + "vjepa2_student_source_num_frames", + "vjepa2_dense_window_seed", + "vjepa2_dense_num_tokens", + "vjepa2_dense_resize_short_side", + "vjepa2_dense_crop_size", + "vjepa2_dense_gradient_checkpointing", + "vjepa2_dense_raw_tokens", + "vjepa2_dense_feature_dim", + ) + vjepa_state = { + name: getattr(self, name) for name in vjepa_attribute_names + } + vjepa_encoder = self.vjepa2_encoder + expected_vjepa = { + "layers": (15,), + "frames": 64, + "source_frames": 81, + "student_sampling": "last_contiguous", + "tokens": 1024, + "feature_dim": 1024, + } + actual_vjepa = { + "layers": tuple(self.vjepa2_dense_layers), + "frames": int(self.vjepa2_dense_num_frames), + "source_frames": int(self.vjepa2_student_source_num_frames), + "student_sampling": str(self.vjepa2_student_frame_sampling), + "tokens": int(self.vjepa2_dense_num_tokens), + "feature_dim": int(self.vjepa2_dense_feature_dim), + } + if actual_vjepa != expected_vjepa: + raise ValueError( + "triple V-JEPA2 contract is fixed to layer15, frames17..80, " + "K=1024 and D=1024: " + f"actual={actual_vjepa}, expected={expected_vjepa}") + + self._initialize_videomae_v2_features(args, device) + expected_videomae = { + "layers": (15,), + "segments": 5, + "skip": 1, + "frames_per_segment": 16, + "source_frames": 81, + "tokens_per_segment": 256, + "feature_dim": 1024, + } + actual_videomae = { + "layers": tuple(self.videomae_v2_dense_layers), + "segments": int(self.videomae_v2_encoder_count), + "skip": int(self.videomae_v2_skip_leading_frames), + "frames_per_segment": int(self.videomae_v2_segment_num_frames), + "source_frames": int(self.videomae_v2_source_num_frames), + "tokens_per_segment": int(self.videomae_v2_dense_num_tokens), + "feature_dim": int(self.videomae_v2_encoder.dense_feature_dim), + } + if actual_videomae != expected_videomae: + raise ValueError( + "triple VideoMAE-v2 contract is fixed to layer15, frames1..80 " + "as 5x16, K=256 per segment and D=1024: " + f"actual={actual_videomae}, expected={expected_videomae}") + + for name, value in vjepa_state.items(): + setattr(self, name, value) + self.vjepa2_encoder = vjepa_encoder + self.dense_feature_segment_count = 1 + self._bind_dense_runtime_aliases() + self._initialize_dinov3_features(args, device) + actual_dinov3 = { + "geometry": self.dinov3_dense_geometry, + "source_frames": self.dinov3_dense_source_num_frames, + "sample_shape": ( + self.dinov3_dense_num_tokens, + self.dinov3_dense_feature_dim, + ), + } + expected_dinov3 = { + "geometry": "temporal_mean5_full256", + "source_frames": 81, + "sample_shape": (1024, 1024), + } + if actual_dinov3 != expected_dinov3: + raise ValueError( + "triple DINOv3 contract is fixed to 20 frames, mean-pool " + "consecutive groups of 5, K=1024 and D=1024: " + f"actual={actual_dinov3}, expected={expected_dinov3}") + self.dual_branch_registry = { + "dinov3": { + "encoder_attribute": "dinov3_encoder", + "source_frame_indices": list( + self.dinov3_dense_frame_indices), + "sample_shape": (1024, 1024), + "tokens_per_temporal_group": 256, + "temporal_groups": 4, + "model_id": self.dinov3_model_id, + "source_commit": self.dinov3_source_commit, + "weight_sha256": getattr( + self.dinov3_encoder, "weight_sha256", ""), + }, + "vjepa2": { + "encoder_attribute": "vjepa2_encoder", + "layer": 15, + "source_frame_range": [17, 80], + "sample_shape": (1024, 1024), + "tokens_per_segment": 1024, + "segments": 1, + "model_id": self.vjepa2_reference_model_id, + "model_sha256": self.vjepa2_model_sha256, + }, + "videomae_v2": { + "encoder_attribute": "videomae_v2_encoder", + "layer": 15, + "source_frame_range": [1, 80], + "sample_shape": (1280, 1024), + "tokens_per_segment": 256, + "segments": 5, + "model_id": self.videomae_v2_reference_model_id, + "model_revision": str(getattr( + args, "videomae_v2_model_revision", "")), + "download_manifest_content_hash": ( + self.videomae_v2_download_manifest_content_hash), + }, + } + for name, branch in self.dual_branch_registry.items(): + raw_shape = tuple(int(value) for value in branch["sample_shape"]) + branch["raw_sample_shape"] = raw_shape + if self.triple_t5_joint_enabled: + text_tokens = self.joint_t5_num_tokens_by_branch[name] + text_rows = text_tokens * ( + self.joint_t5_feature_dim // raw_shape[1]) + branch.update({ + "text_tokens": text_tokens, + "text_rows": text_rows, + "sample_shape": ( + raw_shape[0] + text_rows, raw_shape[1]), + }) + + def _build_multi_encoder_objective(self, args, device) -> None: + """Bind each encoder to its fixed reference summary and combine its loss. + + The hybrid reference uses both Nyström statistics and a sampled reference + minibatch. Encoder weights follow the release's detached-scale normalization.""" + from losses.multi_encoder_mmd import ( + MultiEncoderMMD, + ) + from losses.video_nystrom_mmd import ( + VideoNystromMMD, + load_reference_pack, + ) + from video_fd.encoders.videomae_v2_dense_metadata import ( + canonical_videomae_v2_dense_feature_metadata, + uniform_segment_token_indices, + ) + + if not self.uses_fresh_video_gradcache: + raise ValueError( + "triple video MMD requires nystrom_exact_repulsion") + if not self.latent_self_normalize: + raise ValueError( + "triple video MMD requires wan_kv_self_normalize=true") + if ( + self.nystrom_attraction_ema_enabled + or self.nystrom_repulsion_fifo_enabled + ): + raise ValueError( + "triple video MMD requires EMA and FIFO disabled") + if not self.current_teacher_attraction_enabled: + raise ValueError( + "triple video MMD requires exact matched teacher MMD") + if ( + self.nystrom_attraction_weight != 2.0 + or self.current_teacher_attraction_weight != 1.0 + ): + raise ValueError( + "triple fixed/current teacher weights must be exactly 2:1") + if self.expected_current_teacher_samples != 128: + raise ValueError( + "triple objective requires exactly 128 matched teachers") + + def strict_sha256(value, label): + normalized = str(value).strip().lower() + if ( + len(normalized) != 64 + or any(c not in "0123456789abcdef" for c in normalized) + ): + raise ValueError(f"{label} must be a lowercase SHA-256") + return normalized + + def checked_reference_path(path_value, sha_value, label): + path = os.path.realpath(os.path.expanduser(str(path_value))) + if not os.path.isfile(path): + raise FileNotFoundError( + f"triple {label} reference pack is missing: {path!r}") + expected_sha = strict_sha256(sha_value, f"{label} reference SHA-256") + actual_sha = self._distributed_file_sha256(path) + if actual_sha != expected_sha: + raise ValueError( + f"triple {label} reference file SHA-256 mismatch: " + f"{actual_sha} != {expected_sha}") + return path, expected_sha, load_reference_pack( + path, map_location="cpu") + + # Each encoder owns its reference summary; the wrapper combines normalized losses. + vjepa_objective = self.latent_mmd + if not isinstance(vjepa_objective, VideoNystromMMD): + raise TypeError( + "triple bootstrap V-JEPA2 objective must be VideoNystromMMD") + vjepa_path = self.vjepa2_nystrom_reference_path + configured_vjepa_reference_sha = strict_sha256( + self._dense_config_value( + args, "nystrom_reference_sha256", ""), + "V-JEPA2 reference SHA-256", + ) + if self._distributed_file_sha256( + vjepa_path + ) != configured_vjepa_reference_sha: + raise ValueError( + "triple V-JEPA2 reference file SHA-256 mismatch") + + dinov3_path, dinov3_file_sha, dinov3_pack = checked_reference_path( + getattr(args, "dinov3_nystrom_reference_path", ""), + getattr(args, "dinov3_nystrom_reference_sha256", ""), + "DINOv3", + ) + dinov3_objective = VideoNystromMMD( + dinov3_pack, + expected_metadata={ + "feature": self._triple_t5_joint_feature_metadata( + "dinov3", self._dinov3_feature_metadata() + ), + "nystrom": { + "num_landmarks": int(getattr( + args, "dinov3_nystrom_landmarks", 256)), + "seed": int(getattr(args, "dinov3_nystrom_seed", 0)), + }, + }, + kernel_chunk_size=int(getattr( + args, "dinov3_nystrom_kernel_chunk_size", 32)), + attraction_weight=2.0, + current_teacher_attraction_enabled=True, + current_teacher_attraction_weight=1.0, + attraction_ema_enabled=False, + repulsion_fifo_size=0, + ).to(device=device) + + videomae_path, videomae_file_sha, videomae_pack = checked_reference_path( + getattr(args, "videomae_v2_nystrom_reference_path", ""), + getattr(args, "videomae_v2_nystrom_reference_sha256", ""), + "VideoMAE-v2", + ) + video_tokens = uniform_segment_token_indices( + self.videomae_v2_encoder.dense_raw_tokens, + self.videomae_v2_dense_num_tokens, + ) + video_feature_metadata = canonical_videomae_v2_dense_feature_metadata( + model_id=self.videomae_v2_reference_model_id, + layer_index=15, + num_tokens=256, + token_indices=video_tokens, + num_segments=5, + frames_per_segment=16, + raw_tokens_per_segment=( + self.videomae_v2_encoder.dense_raw_tokens), + feature_dim=self.videomae_v2_encoder.dense_feature_dim, + resize_short_side=224, + crop_size=224, + skip_first_frame=True, + ) + videomae_objective = VideoNystromMMD( + videomae_pack, + expected_metadata={ + "feature": self._triple_t5_joint_feature_metadata( + "videomae_v2", video_feature_metadata + ), + "nystrom": { + "num_landmarks": int(getattr( + args, "videomae_v2_nystrom_landmarks", 256)), + "seed": int(getattr( + args, "videomae_v2_nystrom_seed", 0)), + }, + }, + kernel_chunk_size=int(getattr( + args, "videomae_v2_nystrom_kernel_chunk_size", 32)), + attraction_weight=2.0, + current_teacher_attraction_enabled=True, + current_teacher_attraction_weight=1.0, + attraction_ema_enabled=False, + repulsion_fifo_size=0, + ).to(device=device) + + objectives = { + "dinov3": dinov3_objective, + "vjepa2": vjepa_objective, + "videomae_v2": videomae_objective, + } + expected_shapes = { + name: tuple(branch["sample_shape"]) + for name, branch in self.dual_branch_registry.items() + } + actual_shapes = { + name: tuple(objective.sample_shape) + for name, objective in objectives.items() + } + if actual_shapes != expected_shapes: + raise ValueError( + "triple reference sample-shape mismatch: " + f"actual={actual_shapes}, expected={expected_shapes}") + + configured_bank_hash = strict_sha256( + getattr(args, "triple_teacher_bank_content_hash", ""), + "triple_teacher_bank_content_hash", + ) + expected_manifest_hash = strict_sha256( + getattr(args, "matched_manifest_sha256", ""), + "matched_manifest_sha256", + ) + sources = { + name: objective.reference_metadata.get( + "reference", {}).get("source", {}) + for name, objective in objectives.items() + } + for name, source in sources.items(): + if ( + source.get("triple_teacher_bank_content_hash") + != configured_bank_hash + or source.get("source_manifest_sha256") + != expected_manifest_hash + or int(source.get("selected_samples", -1)) != int(getattr( + args, "matched_manifest_samples", 5000)) + or source.get("branch_name") != name + ): + raise ValueError( + f"{name} reference is not bound to the configured frozen " + "triple teacher bank") + + # Paper Eq. (3) motivates encoder weighting; this release also uses detached-scale normalization. + branch_weights = { + "dinov3": float(getattr(args, "triple_dinov3_weight", 15.0)), + "vjepa2": float(getattr(args, "triple_vjepa2_weight", 1.0)), + "videomae_v2": float(getattr( + args, "triple_videomae_v2_weight", 2.0)), + } + aggregation_divisor = float(getattr( + args, "triple_aggregation_divisor", 10.0)) + aggregation = (branch_weights, aggregation_divisor) + legacy_aggregation = ({ + "dinov3": 15.0, + "vjepa2": 1.0, + "videomae_v2": 2.0, + }, 10.0) + requested_aggregation = ({ + "dinov3": 10.0, + "vjepa2": 1.0, + "videomae_v2": 1.0, + }, 3.0) + if aggregation not in (legacy_aggregation, requested_aggregation): + raise ValueError( + "triple representation aggregation must be exactly either " + "(15*DINO + V-JEPA2 + 2*VideoMAE)/10 or " + "(10*DINO + V-JEPA2 + VideoMAE)/3") + self.latent_mmd = MultiEncoderMMD( + objectives, + weights=branch_weights, + normalization_eps=self.latent_norm_eps, + aggregation_divisor=aggregation_divisor, + ).to(device=device) + + self.dinov3_nystrom_reference_path = dinov3_path + self.dinov3_nystrom_reference_hash = dinov3_objective.reference_hash + self.videomae_v2_nystrom_reference_path = videomae_path + self.videomae_v2_nystrom_reference_hash = ( + videomae_objective.reference_hash) + self.triple_teacher_bank_content_hash = configured_bank_hash + # Compatibility aliases are internal only; public metadata is triple. + self.dual_teacher_bank_content_hash = configured_bank_hash + reference_files = { + "dinov3": dinov3_file_sha, + "vjepa2": configured_vjepa_reference_sha, + "videomae_v2": videomae_file_sha, + } + reference_paths = { + "dinov3": dinov3_path, + "vjepa2": vjepa_path, + "videomae_v2": videomae_path, + } + for name, objective in objectives.items(): + self.dual_branch_registry[name].update({ + "reference_path": reference_paths[name], + "reference_file_sha256": reference_files[name], + "reference_hash": objective.reference_hash, + "nystrom_landmarks": int(objective.landmarks.shape[0]), + }) + + def _initialize_videomae_v2_features(self, args, device) -> None: + """Build the fixed five-replica VideoMAE-v2 Large feature map. + + Public state is VideoMAE-namespaced. A small set of historical + ``vjepa2_*`` aliases is retained only inside this class and its trainer + so the already validated GradCache/VAE replay implementation can be + reused without changing the distribution-loss formula. + """ + from video_fd.encoders.segmented_videomae_v2 import ( + SegmentedVideoMAEv2Ensemble, + ) + + configured_model_id = str(getattr( + args, + "videomae_v2_model_id", + "assets/models/VideoMAEv2-Large", + )) + expanded_model_id = os.path.expanduser(configured_model_id) + self.videomae_v2_model_id = ( + os.path.realpath(expanded_model_id) + if os.path.exists(expanded_model_id) + else configured_model_id + ) + self.videomae_v2_reference_model_id = str(getattr( + args, "videomae_v2_reference_model_id", + "VideoMAEv2-Large")) + self.videomae_v2_local_files_only = bool(getattr( + args, "videomae_v2_local_files_only", True)) + if ( + self.videomae_v2_local_files_only + and not os.path.isdir(self.videomae_v2_model_id) + ): + raise FileNotFoundError( + "VideoMAE-v2 training requires the pinned local checkpoint " + "directory; prepare the local VideoMAEv2-Large files first: " + f"{self.videomae_v2_model_id}") + self.videomae_v2_model_revision = str(getattr( + args, "videomae_v2_model_revision", "")) + self.videomae_v2_download_manifest_content_hash = str(getattr( + args, "videomae_v2_download_manifest_content_hash", "")) + if self.dual_video_mmd_enabled: + expected_revision = ( + "9981a9c8f77118c421e5228e1b219468a4b0238d") + if self.videomae_v2_model_revision != expected_revision: + raise ValueError( + "VideoMAE-v2 revision must be pinned to " + f"{expected_revision}") + download_manifest_path = os.path.join( + self.videomae_v2_model_id, "download_manifest.json") + if not os.path.isfile(download_manifest_path): + raise FileNotFoundError( + "VideoMAE-v2 download_manifest.json is required") + with open(download_manifest_path, "r", encoding="utf-8") as handle: + download_manifest = json.load(handle) + saved_content_hash = download_manifest.get("content_hash") + manifest_payload = dict(download_manifest) + manifest_payload.pop("content_hash", None) + if ( + download_manifest.get("revision") != expected_revision + or saved_content_hash + != self._canonical_sha256(manifest_payload) + or self.videomae_v2_download_manifest_content_hash + != saved_content_hash + ): + raise ValueError( + "VideoMAE-v2 download manifest identity mismatch") + for file_identity in download_manifest.get("files", []): + relative = str(file_identity.get("path", "")) + path = os.path.realpath(os.path.join( + self.videomae_v2_model_id, relative)) + if ( + not relative + or not path.startswith(self.videomae_v2_model_id + os.sep) + or not os.path.isfile(path) + or os.path.getsize(path) != int( + file_identity.get("size_bytes", -1)) + or self._distributed_file_sha256(path) + != file_identity.get("sha256") + ): + raise ValueError( + "VideoMAE-v2 downloaded file identity mismatch: " + f"{relative}") + + configured_dense_layers = getattr( + args, "videomae_v2_dense_layers", None) + if configured_dense_layers is None: + dense_layers = ( + int(getattr(args, "videomae_v2_dense_layer", 15)), + ) + else: + if isinstance(configured_dense_layers, (str, bytes)): + raise TypeError( + "videomae_v2_dense_layers must be a sequence of integers") + dense_layers = tuple( + int(layer) for layer in configured_dense_layers) + if not dense_layers: + raise ValueError( + "videomae_v2_dense_layers must not be empty") + if len(set(dense_layers)) != len(dense_layers): + raise ValueError( + "videomae_v2_dense_layers must contain unique indices") + if tuple(sorted(dense_layers)) != dense_layers: + raise ValueError( + "videomae_v2_dense_layers must be sorted") + if any(layer < 0 or layer >= 24 for layer in dense_layers): + raise ValueError( + "VideoMAE-v2 Large dense layers use zero-based indices in " + "[0,24)") + + self.videomae_v2_dense_layers = dense_layers + self.videomae_v2_multilayer_enabled = len(dense_layers) > 1 + self.videomae_v2_dense_layer = dense_layers[-1] + self.videomae_v2_encoder_count = int(getattr( + args, "videomae_v2_encoder_count", 5)) + self.videomae_v2_skip_leading_frames = int(getattr( + args, "videomae_v2_skip_leading_frames", 1)) + self.videomae_v2_segment_num_frames = int(getattr( + args, "videomae_v2_segment_num_frames", 16)) + derived_source_num_frames = ( + self.videomae_v2_skip_leading_frames + + self.videomae_v2_encoder_count + * self.videomae_v2_segment_num_frames + ) + self.videomae_v2_source_num_frames = int(getattr( + args, + "videomae_v2_source_num_frames", + derived_source_num_frames, + )) + self.videomae_v2_dense_num_tokens = int(getattr( + args, "videomae_v2_dense_num_tokens", 128)) + self.videomae_v2_dense_resize_short_side = int(getattr( + args, "videomae_v2_dense_resize_short_side", 224)) + self.videomae_v2_dense_crop_size = int(getattr( + args, "videomae_v2_dense_crop_size", 224)) + self.videomae_v2_dense_gradient_checkpointing = bool(getattr( + args, "videomae_v2_dense_gradient_checkpointing", True)) + + fixed_geometry = ( + self.videomae_v2_encoder_count, + self.videomae_v2_skip_leading_frames, + self.videomae_v2_segment_num_frames, + self.videomae_v2_source_num_frames, + self.videomae_v2_dense_resize_short_side, + self.videomae_v2_dense_crop_size, + ) + expected_geometry = (5, 1, 16, 81, 224, 224) + if fixed_geometry != expected_geometry: + raise ValueError( + "VideoMAE-v2 whole-video geometry is fixed to five physical " + "16-frame encoders over decoded frames [1:81]: " + f"actual={fixed_geometry}, expected={expected_geometry}") + if self.videomae_v2_dense_num_tokens <= 1: + raise ValueError( + "videomae_v2_dense_num_tokens must be > 1 per segment") + + adapter_config = { + "name": "videomae_v2_five_segment_dense", + "model_id": self.videomae_v2_model_id, + "encoder_count": self.videomae_v2_encoder_count, + "skip_leading_frames": self.videomae_v2_skip_leading_frames, + "segment_num_frames": self.videomae_v2_segment_num_frames, + "source_num_frames": self.videomae_v2_source_num_frames, + "input_frames": self.videomae_v2_segment_num_frames, + "dense_num_frames": self.videomae_v2_segment_num_frames, + "input_size": self.videomae_v2_dense_crop_size, + "resize_short_side": ( + self.videomae_v2_dense_resize_short_side), + "crop_mode": "center", + "dense_only": True, + # ``15`` means the 16th block and captures its post-residual + # dense state, before pooling/fc_norm. + "dense_layer_index": self.videomae_v2_dense_layer, + "local_files_only": self.videomae_v2_local_files_only, + } + if self.videomae_v2_model_revision: + adapter_config["revision"] = self.videomae_v2_model_revision + self.videomae_v2_encoder = SegmentedVideoMAEv2Ensemble( + adapter_config) + self.videomae_v2_encoder.to(device=device, dtype=torch.float32) + self.videomae_v2_encoder.eval().requires_grad_(False) + + actual_adapter_geometry = ( + int(self.videomae_v2_encoder.dense_layer_index), + int(self.videomae_v2_encoder.dense_num_frames), + int(self.videomae_v2_encoder.resize_short_side), + int(self.videomae_v2_encoder.input_size), + int(self.videomae_v2_encoder.retained_dense_blocks), + int(self.videomae_v2_encoder.dense_raw_tokens), + int(self.videomae_v2_encoder.dense_feature_dim), + len(self.videomae_v2_encoder.encoders), + ) + expected_adapter_geometry = ( + self.videomae_v2_dense_layer, + 16, + 224, + 224, + self.videomae_v2_dense_layer + 1, + 1568, + 1024, + 5, + ) + if actual_adapter_geometry != expected_adapter_geometry: + raise RuntimeError( + "loaded VideoMAE-v2 Large ensemble does not match the fixed " + "dense-prefix geometry: " + f"actual={actual_adapter_geometry}, " + f"expected={expected_adapter_geometry}") + configured_depths = { + int(getattr(encoder, "_configured_depth", -1)) + for encoder in self.videomae_v2_encoder.encoders + } + if configured_depths != {24}: + raise RuntimeError( + "all five encoders must be VideoMAE-v2 Large with 24 blocks; " + f"got depths={sorted(configured_depths)}") + + # Historical internal aliases. Public metadata, config, logs, and + # checkpoints are translated back to ``videomae_v2_*``. + self.dense_feature_segment_count = 5 + self.vjepa2_model_id = self.videomae_v2_model_id + self.vjepa2_dense_layers = self.videomae_v2_dense_layers + self.vjepa2_multilayer_enabled = ( + self.videomae_v2_multilayer_enabled) + self.vjepa2_dense_layer = self.videomae_v2_dense_layer + self.vjepa2_dense_num_frames = 80 + self.vjepa2_dense_source_num_frames = 81 + self.vjepa2_dense_frame_sampling = "center_contiguous" + self.vjepa2_dense_window_seed = "fixed-five-segment-videomae-v2" + self.vjepa2_dense_num_tokens = ( + self.videomae_v2_dense_num_tokens) + self.vjepa2_dense_resize_short_side = ( + self.videomae_v2_dense_resize_short_side) + self.vjepa2_dense_crop_size = self.videomae_v2_dense_crop_size + self.vjepa2_dense_gradient_checkpointing = ( + self.videomae_v2_dense_gradient_checkpointing) + self.vjepa2_encoder = self.videomae_v2_encoder + self.vjepa2_dense_raw_tokens = int( + self.videomae_v2_encoder.dense_raw_tokens) + self.vjepa2_dense_feature_dim = int( + self.videomae_v2_encoder.dense_feature_dim) + if self.vjepa2_dense_num_tokens > self.vjepa2_dense_raw_tokens: + raise ValueError( + "videomae_v2_dense_num_tokens cannot exceed the native " + "per-segment token grid: " + f"{self.vjepa2_dense_num_tokens} > " + f"{self.vjepa2_dense_raw_tokens}") + + def _initialize_dinov3_features(self, args, device) -> None: + """Load the independent frozen DINOv3 dense-image backend. + + The compatibility geometry consumes sixteen fixed frames and retains a + row-major 8x8 spatial patch subset from each frame. The optional + ``temporal_mean5_full256`` geometry consumes twenty fixed frames, + averages every five consecutive frame-feature grids, and retains the + complete 16x16 patch grid for each of the four temporal groups. Both + geometries produce one aligned [1024, 1024] matrix per video; their + adapter metadata keeps references and checkpoints mutually exclusive. + """ + compatibility_frames = ( + 0, 5, 11, 16, 21, 27, 32, 37, + 43, 48, 53, 59, 64, 69, 75, 80, + ) + temporal_pool_frames = ( + 0, 4, 8, 13, 17, 21, 25, 29, 34, 38, + 42, 46, 51, 55, 59, 63, 67, 72, 76, 80, + ) + compatibility_spatial_axis = (0, 2, 4, 6, 9, 11, 13, 15) + full_spatial_axis = tuple(range(16)) + initialization_error = None + try: + geometry = str(getattr( + args, "dinov3_dense_geometry", "f16x64")).strip() + if geometry == "f16x64": + from video_fd.encoders.dinov3_adapter import ( + DINOv3DenseAdapter as Adapter, + ) + + canonical_frames = compatibility_frames + canonical_spatial_axis = compatibility_spatial_axis + elif geometry == "temporal_mean5_full256": + from video_fd.encoders.dinov3_temporal_pool_adapter import ( + DINOv3TemporalPoolAdapter as Adapter, + ) + + canonical_frames = temporal_pool_frames + canonical_spatial_axis = full_spatial_axis + else: + raise ValueError( + "dinov3_dense_geometry must be 'f16x64' or " + "'temporal_mean5_full256', got " + f"{geometry!r}") + + configured_frames = tuple(int(value) for value in getattr( + args, "dinov3_dense_frame_indices", canonical_frames)) + configured_spatial_axis = tuple(int(value) for value in getattr( + args, "dinov3_dense_spatial_axis_indices", + canonical_spatial_axis)) + adapter_config = { + "repo_path": str(getattr(args, "dinov3_repo_path", "")), + "weight_path": str(getattr(args, "dinov3_weight_path", "")), + "source_commit": str(getattr( + args, "dinov3_source_commit", "")), + "source_num_frames": int(getattr( + args, "dinov3_dense_source_num_frames", 81)), + "frame_indices": configured_frames, + "input_size": int(getattr( + args, "dinov3_dense_input_size", 256)), + "patch_size": int(getattr( + args, "dinov3_dense_patch_size", 16)), + "spatial_axis_indices": configured_spatial_axis, + "forward_chunk_frames": int(getattr( + args, "dinov3_dense_forward_chunk_frames", 4)), + } + if geometry == "temporal_mean5_full256": + adapter_config.update({ + "geometry": "temporal_mean5_full256", + "temporal_groups": 4, + "temporal_pool_frames": 5, + }) + configured_checkpoint_sha256 = str(getattr( + args, "dinov3_checkpoint_sha256", "")).strip() + if configured_checkpoint_sha256: + adapter_config["checkpoint_sha256"] = ( + configured_checkpoint_sha256) + + self.dinov3_encoder = Adapter(adapter_config) + self.dinov3_encoder.to(device=device, dtype=torch.float32) + self.dinov3_encoder.eval().requires_grad_(False) + + self.dinov3_dense_geometry = geometry + self.dinov3_model_id = str(self.dinov3_encoder.model_id) + self.dinov3_repo_path = str(self.dinov3_encoder.repo_path) + self.dinov3_weight_path = str(self.dinov3_encoder.weight_path) + self.dinov3_source_commit = str( + self.dinov3_encoder.source_commit) + self.dinov3_dense_source_num_frames = int( + self.dinov3_encoder.source_num_frames) + self.dinov3_dense_frame_indices = tuple( + int(value) for value in self.dinov3_encoder.frame_indices) + self.dinov3_dense_num_frames = int( + self.dinov3_encoder.dense_num_frames) + self.dinov3_dense_input_size = int( + self.dinov3_encoder.input_size) + self.dinov3_dense_patch_size = int( + self.dinov3_encoder.patch_size) + self.dinov3_dense_spatial_axis_indices = tuple( + int(value) for value in + self.dinov3_encoder.spatial_axis_indices) + self.dinov3_dense_forward_chunk_frames = int( + self.dinov3_encoder.forward_chunk_frames) + self.dinov3_dense_raw_tokens = int( + self.dinov3_encoder.dense_raw_tokens) + self.dinov3_dense_num_tokens = int( + self.dinov3_encoder.dense_num_tokens) + self.dinov3_dense_feature_dim = int( + self.dinov3_encoder.dense_feature_dim) + self.dinov3_dense_gradient_checkpointing = bool(getattr( + args, "dinov3_dense_gradient_checkpointing", False)) + + if geometry == "temporal_mean5_full256": + self.dinov3_dense_temporal_groups = int( + self.dinov3_encoder.temporal_groups) + self.dinov3_dense_temporal_pool_size = int( + self.dinov3_encoder.temporal_pool_frames) + self.dinov3_dense_tokens_per_group = int( + self.dinov3_encoder.per_frame_tokens) + + actual_contract = { + "model_id": self.dinov3_model_id, + "source_num_frames": self.dinov3_dense_source_num_frames, + "frame_indices": self.dinov3_dense_frame_indices, + "num_frames": self.dinov3_dense_num_frames, + "input_size": self.dinov3_dense_input_size, + "patch_size": self.dinov3_dense_patch_size, + "spatial_axis_indices": ( + self.dinov3_dense_spatial_axis_indices), + "forward_chunk_frames": ( + self.dinov3_dense_forward_chunk_frames), + "raw_tokens": self.dinov3_dense_raw_tokens, + "num_tokens": self.dinov3_dense_num_tokens, + "feature_dim": self.dinov3_dense_feature_dim, + } + if geometry == "temporal_mean5_full256": + actual_contract.update({ + "geometry": str(self.dinov3_encoder.geometry), + "temporal_groups": ( + self.dinov3_dense_temporal_groups), + "temporal_pool_size": ( + self.dinov3_dense_temporal_pool_size), + "tokens_per_group": ( + self.dinov3_dense_tokens_per_group), + }) + expected_contract = { + "model_id": "dinov3_vitl16", + "source_num_frames": 81, + "frame_indices": canonical_frames, + "num_frames": 20, + "input_size": 256, + "patch_size": 16, + "spatial_axis_indices": canonical_spatial_axis, + "forward_chunk_frames": 4, + "raw_tokens": 5120, + "num_tokens": 1024, + "feature_dim": 1024, + "geometry": "temporal_mean5_full256", + "temporal_groups": 4, + "temporal_pool_size": 5, + "tokens_per_group": 256, + } + else: + expected_contract = { + "model_id": "dinov3_vitl16", + "source_num_frames": 81, + "frame_indices": canonical_frames, + "num_frames": 16, + "input_size": 256, + "patch_size": 16, + "spatial_axis_indices": canonical_spatial_axis, + "forward_chunk_frames": 4, + "raw_tokens": 4096, + "num_tokens": 1024, + "feature_dim": 1024, + } + if actual_contract != expected_contract: + raise ValueError( + "loaded DINOv3 adapter does not satisfy the public dense " + f"feature contract: actual={actual_contract}, " + f"expected={expected_contract}") + except BaseException as exc: + initialization_error = exc + self._synchronize_dinov3_local_failure( + initialization_error, + device, + operation="DINOv3 adapter initialization failed", + ) + + def _dinov3_feature_metadata(self) -> dict: + metadata = self.dinov3_encoder.feature_metadata() + if not isinstance(metadata, Mapping): + raise TypeError("DINOv3 feature_metadata() must return a mapping") + metadata = dict(metadata) + required_identity = { + "feature_type": "dinov3_dense_latent", + "model_id": self.dinov3_model_id, + "raw_tokens": self.dinov3_dense_raw_tokens, + "num_tokens": self.dinov3_dense_num_tokens, + "feature_dim": self.dinov3_dense_feature_dim, + "matrix_shape": [ + self.dinov3_dense_num_tokens, + self.dinov3_dense_feature_dim, + ], + } + mismatched = { + key: {"saved": metadata.get(key), "runtime": value} + for key, value in required_identity.items() + if metadata.get(key) != value + } + if mismatched: + raise ValueError( + "DINOv3 adapter metadata violates the runtime feature " + f"contract: {mismatched}") + # Normalize equivalent local and shared asset paths for reference metadata. + local_root = "/data/shared/video-generation-assets" + logical_root = "/assets/video-generation-assets" + aliased = False + for key in ("repo_path", "weight_path"): + value = metadata.get(key) + if isinstance(value, str) and value.startswith(local_root + "/"): + metadata[key] = logical_root + value[len(local_root):] + aliased = True + if aliased: + metadata.pop("feature_contract_sha256", None) + payload = json.dumps( + metadata, sort_keys=True, separators=(",", ":"), + ensure_ascii=True, allow_nan=False).encode("utf-8") + metadata["feature_contract_sha256"] = hashlib.sha256(payload).hexdigest() + return metadata + + def _triple_t5_joint_feature_metadata( + self, + branch_name: str, + video_feature_metadata: Mapping[str, object], + ) -> dict: + """Wrap one raw video protocol in the frozen branch-specific T5 layout.""" + if not bool(getattr( + self, "triple_t5_joint_enabled", False + )): + return dict(video_feature_metadata) + from video_fd.encoders.triple_t5_joint import ( + canonical_triple_t5_joint_feature_metadata, + ) + + if branch_name == "dinov3": + branch_encoder_identity = { + "model_id": self.dinov3_model_id, + "weight_sha256": getattr( + self.dinov3_encoder, "weight_sha256", ""), + "source_commit": self.dinov3_source_commit, + } + elif branch_name == "vjepa2": + branch_encoder_identity = { + "model_id": self.vjepa2_reference_model_id, + "model_sha256": self.vjepa2_model_sha256, + } + elif branch_name == "videomae_v2": + branch_encoder_identity = { + "model_id": self.videomae_v2_reference_model_id, + "revision": self.videomae_v2_model_revision, + "download_manifest_content_hash": ( + self.videomae_v2_download_manifest_content_hash), + } + else: + raise ValueError(f"unsupported triple joint branch {branch_name!r}") + return canonical_triple_t5_joint_feature_metadata( + branch_name=branch_name, + video_feature_metadata=video_feature_metadata, + branch_encoder_identity=branch_encoder_identity, + t5_model_id=self.joint_t5_reference_model_id, + tokenizer_id=self.joint_t5_reference_tokenizer_id, + t5_weights_sha256=self.joint_t5_weights_sha256, + tokenizer_sha256=self.joint_t5_tokenizer_sha256, + ) + + def _dinov3_descriptor_frame_indices( + self, + replay_descriptor: Mapping, + *, + device: torch.device, + ) -> torch.Tensor: + saved = replay_descriptor.get("dinov3_frame_indices") + if not isinstance(saved, torch.Tensor): + raise TypeError("DINOv3 replay frame indices must be a tensor") + saved = saved.to(device=device, dtype=torch.long) + expected = torch.tensor( + self.dinov3_dense_frame_indices, + device=device, + dtype=torch.long, + ) + if saved.ndim != 1 or not torch.equal(saved, expected): + raise ValueError( + "DINOv3 replay frame indices changed from the fixed public " + "sampling contract") + return saved + + def _extract_dinov3_dense(self, video01: torch.Tensor) -> torch.Tensor: + if ( + video01.ndim != 5 + or int(video01.shape[1]) + != self.dinov3_dense_source_num_frames + or int(video01.shape[2]) != 3 + ): + raise ValueError( + "DINOv3 dense extraction requires full decoded videos with " + f"shape [B,{self.dinov3_dense_source_num_frames},3,H,W], " + f"got {tuple(video01.shape)}") + extract = ( + self.dinov3_encoder.forward_dense_no_grad + if not torch.is_grad_enabled() + else self.dinov3_encoder.forward_dense + ) + features = extract( + video01, + gradient_checkpointing=( + self.dinov3_dense_gradient_checkpointing), + ) + expected_shape = ( + int(video01.shape[0]), + self.dinov3_dense_num_tokens, + self.dinov3_dense_feature_dim, + ) + if features.ndim != 3 or tuple(features.shape) != expected_shape: + raise ValueError( + "frozen DINOv3 returned an invalid dense feature matrix: " + f"got {tuple(features.shape)}, expected {expected_shape}") + return features + + def _replay_dinov3_dense_to_pixels( + self, + video01: torch.Tensor, + *, + expected_features: torch.Tensor, + feature_cotangent: torch.Tensor, + replay_atol: float, + replay_rtol: float, + ) -> Tuple[torch.Tensor, torch.Tensor]: + replay_features, video_cotangent = ( + self.dinov3_encoder.replay_dense_vjp( + video01.detach().float(), + feature_cotangent.detach().float(), + expected_features=expected_features.detach().float(), + gradient_checkpointing=( + self.dinov3_dense_gradient_checkpointing), + replay_atol=float(replay_atol), + replay_rtol=float(replay_rtol), + ) + ) + expected_feature_shape = tuple(expected_features.shape) + if tuple(replay_features.shape) != expected_feature_shape: + raise RuntimeError( + "DINOv3 feature replay changed matrix shape: " + f"{tuple(replay_features.shape)} != " + f"{expected_feature_shape}") + if tuple(video_cotangent.shape) != tuple(video01.shape): + raise RuntimeError( + "DINOv3 replay returned an invalid full-video pixel " + f"cotangent: {tuple(video_cotangent.shape)} != " + f"{tuple(video01.shape)}") + if not bool(torch.isfinite(video_cotangent).all().item()): + raise FloatingPointError( + "DINOv3 replay returned a non-finite pixel cotangent") + replay_max_abs = ( + replay_features.detach().float() + - expected_features.detach().float() + ).abs().max() + return video_cotangent.detach().float(), replay_max_abs.detach() + + def _dense_metric_prefix(self) -> str: + return ( + "dinov3_dense" + if self._uses_dinov3_dense_backend() + else "vjepa2_dense" + ) + + def _uses_dinov3_dense_backend(self) -> bool: + """Support legacy/test instances created without the new init flags.""" + return bool(getattr( + self, + "is_dinov3_dense", + getattr(self, "feature_type", None) == "dinov3_dense_latent", + )) + + def _token_indices(self, latent: torch.Tensor) -> torch.Tensor: + patch = self.latent_patch_size + tokens = ( + (latent.shape[1] // int(patch[0])) + * (latent.shape[3] // int(patch[1])) + * (latent.shape[4] // int(patch[2])) + ) + n = min(tokens, self.latent_num_tokens) + return torch.linspace( + 0, tokens - 1, steps=n, device=latent.device).round().long() + + def _extract(self, latent, conditional_dict, token_indices): + timestep = latent.new_full( + (latent.shape[0], latent.shape[1]), + self.latent_timestep, + dtype=torch.float32, + ) + features = self.real_score( + noisy_image_or_video=latent, + conditional_dict=conditional_dict, + timestep=timestep, + return_latent_features=True, + latent_feature_layer=self.latent_layer, + latent_token_indices=token_indices, + ) + expected_shape = ( + latent.shape[0], token_indices.numel(), self.latent_feature_dim) + if features.ndim != 3 or tuple(features.shape) != expected_shape: + raise ValueError( + "frozen Wan returned an invalid post-block latent shape: " + f"got {tuple(features.shape)}, expected {expected_shape}") + return features + + def _vjepa2_token_indices(self, device: torch.device) -> torch.Tensor: + return uniform_dense_token_indices( + self.vjepa2_dense_raw_tokens, + self.vjepa2_dense_num_tokens, + device=device, + ) + + def _vjepa2_feature_metadata( + self, + token_indices: Optional[torch.Tensor] = None, + *, + layer_index: Optional[int] = None, + ) -> dict: + resolved_layer = ( + ( + self.motionformer_dense_layer + if self._dense_encoder_kind() == "motionformer" + else self.vjepa2_dense_layer + ) + if layer_index is None + else int(layer_index) + ) + if self._dense_encoder_kind() == "motionformer": + if resolved_layer != self.motionformer_dense_layer: + raise ValueError( + "requested feature metadata for an inactive MotionFormer " + f"layer: {resolved_layer}") + from video_fd.encoders.motionformer_dense_metadata import ( + canonical_motionformer_dense_feature_metadata, + ) + + return canonical_motionformer_dense_feature_metadata( + checkpoint_path=self.motionformer_checkpoint_path, + checkpoint_sha256=self.motionformer_checkpoint_sha256, + layer_index=self.motionformer_dense_layer, + num_frames=self.motionformer_dense_num_frames, + raw_tokens=self.motionformer_dense_raw_tokens, + num_tokens=self.motionformer_dense_num_tokens, + feature_dim=self.motionformer_dense_feature_dim, + token_indices=token_indices, + ) + if resolved_layer not in getattr( + self, "vjepa2_dense_layers", (self.vjepa2_dense_layer,) + ): + raise ValueError( + "requested feature metadata for an inactive V-JEPA2 layer: " + f"{resolved_layer}") + if ( + token_indices is None + and self.dense_state_prefix in { + "dual", "triple", "dinov3_vjepa2" + } + ): + # Multi-encoder Wan banks serialize the exact uniform subset so + # the reference and replay contracts attest row order explicitly. + token_indices = uniform_dense_token_indices( + self.vjepa2_dense_raw_tokens, + self.vjepa2_dense_num_tokens, + ) + metadata_kwargs = { + "model_id": getattr( + self, "vjepa2_reference_model_id", self.vjepa2_model_id), + "layer_index": resolved_layer, + "num_frames": self.vjepa2_dense_num_frames, + "resize_short_side": self.vjepa2_dense_resize_short_side, + "crop_size": self.vjepa2_dense_crop_size, + "raw_tokens": self.vjepa2_dense_raw_tokens, + "num_tokens": self.vjepa2_dense_num_tokens, + "feature_dim": self.vjepa2_dense_feature_dim, + "token_indices": token_indices, + } + frame_sampling = ( + "last_contiguous" + if self.dense_state_prefix in { + "dual", "triple", "dinov3_vjepa2" + } + else getattr( + self, + "vjepa2_dense_frame_sampling", + "center_contiguous", + ) + ) + if frame_sampling == "last_contiguous": + metadata_kwargs.update({ + "frame_sampling": frame_sampling, + "source_num_frames": self.vjepa2_dense_source_num_frames, + }) + elif frame_sampling == "random_contiguous_per_video": + metadata_kwargs.update({ + "frame_sampling": frame_sampling, + "source_num_frames": self.vjepa2_dense_source_num_frames, + "window_seed": self.vjepa2_dense_window_seed, + }) + video_metadata = canonical_vjepa2_dense_feature_metadata( + **metadata_kwargs) + return self._triple_t5_joint_feature_metadata( + "vjepa2", video_metadata + ) if bool(getattr( + self, "triple_t5_joint_enabled", False)) else video_metadata + + def _dense_feature_metadata( + self, + token_indices: Optional[torch.Tensor] = None, + *, + layer_index: Optional[int] = None, + ) -> dict: + """Encoder-neutral alias for strict dense feature identity.""" + return self._vjepa2_feature_metadata( + token_indices, layer_index=layer_index) + + def _vjepa2_temporal_window_indices( + self, + sample_keys, + *, + batch_size: int, + device: torch.device, + ) -> Optional[torch.Tensor]: + sampling = getattr( + self, + "vjepa2_student_frame_sampling", + getattr(self, "vjepa2_dense_frame_sampling", "center_contiguous"), + ) + if sampling == "center_contiguous": + if sample_keys is not None: + raise ValueError( + "window sample keys are invalid for center sampling") + return None + if sampling == "last_contiguous": + if sample_keys is not None: + raise ValueError( + "window sample keys are invalid for last sampling") + start = ( + int(self.vjepa2_student_source_num_frames) + - int(self.vjepa2_dense_num_frames) + ) + offsets = torch.arange( + self.vjepa2_dense_num_frames, + device=device, + dtype=torch.long, + ) + return (offsets + start).unsqueeze(0).expand( + int(batch_size), -1).clone() + if sample_keys is None: + raise ValueError( + "random V-JEPA2 windows require one sample key per video") + if isinstance(sample_keys, torch.Tensor): + if sample_keys.ndim != 1: + raise ValueError("window sample keys must have shape [B]") + keys = sample_keys.detach().cpu().tolist() + else: + keys = list(sample_keys) + if len(keys) != int(batch_size): + raise ValueError( + "window sample key count does not match the video batch: " + f"{len(keys)} != {int(batch_size)}") + starts = [ + stateless_temporal_window_start( + self.vjepa2_dense_window_seed, + "student", + int(key), + source_num_frames=getattr( + self, "vjepa2_student_source_num_frames", + self.vjepa2_dense_source_num_frames), + window_num_frames=self.vjepa2_dense_num_frames, + ) + for key in keys + ] + start_tensor = torch.tensor( + starts, device=device, dtype=torch.long) + offsets = torch.arange( + self.vjepa2_dense_num_frames, + device=device, + dtype=torch.long, + ) + return start_tensor[:, None] + offsets[None, :] + + def _vjepa2_descriptor_window_indices( + self, + replay_descriptor: dict, + *, + device: torch.device, + ) -> Optional[torch.Tensor]: + if getattr( + self, + "vjepa2_student_frame_sampling", + getattr(self, "vjepa2_dense_frame_sampling", "center_contiguous"), + ) == "center_contiguous": + return None + return replay_descriptor["vjepa2_frame_indices"].to( + device=device, dtype=torch.long) + + def _extract_vjepa2_dense( + self, + video01: torch.Tensor, + token_indices: torch.Tensor, + *, + gradient_checkpointing: bool, + temporal_window_indices: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self._dense_encoder_kind() == "motionformer": + if temporal_window_indices is not None: + raise ValueError( + "MotionFormer performs fixed segment-midpoint sampling " + "and does not accept V-JEPA2 temporal windows") + features = self.motionformer_encoder.forward_dense( + video01, + layer_index=self.motionformer_dense_layer, + token_indices=token_indices, + gradient_checkpointing=gradient_checkpointing, + ) + expected_shape = ( + video01.shape[0], + token_indices.numel(), + self.motionformer_dense_feature_dim, + ) + if features.ndim != 3 or tuple(features.shape) != expected_shape: + raise ValueError( + "frozen MotionFormer returned an invalid intermediate " + f"dense shape: got {tuple(features.shape)}, expected " + f"{expected_shape}") + return features + + dense_layers = getattr( + self, "vjepa2_dense_layers", (self.vjepa2_dense_layer,)) + if len(dense_layers) == 1: + features_by_layer = { + dense_layers[0]: self.vjepa2_encoder.forward_dense( + video01, + layer_index=dense_layers[0], + token_indices=token_indices, + gradient_checkpointing=gradient_checkpointing, + temporal_window_indices=temporal_window_indices, + ) + } + else: + features_by_layer = self.vjepa2_encoder.forward_dense_layers( + video01, + layer_indices=dense_layers, + token_indices=token_indices, + gradient_checkpointing=gradient_checkpointing, + temporal_window_indices=temporal_window_indices, + ) + expected_shape = ( + video01.shape[0], + token_indices.numel(), + self.vjepa2_dense_feature_dim, + ) + ordered_features = [] + for layer in dense_layers: + features = features_by_layer[layer] + if ( + features.ndim != 3 + or tuple(features.shape) != expected_shape + ): + raise ValueError( + "frozen V-JEPA2 returned an invalid intermediate dense " + f"shape at layer {layer}: got {tuple(features.shape)}, " + f"expected {expected_shape}") + ordered_features.append(features) + # The GradCache transports one tensor. The loss module splits this + # layer-major concatenation and evaluates independent kernels. + return ( + ordered_features[0] + if len(ordered_features) == 1 + else torch.cat(ordered_features, dim=1) + ) + + def _replay_dense_encoder_vjp( + self, + video01: torch.Tensor, + token_indices: torch.Tensor, + expected_features: torch.Tensor, + feature_cotangent: torch.Tensor, + *, + gradient_checkpointing: bool, + temporal_window_indices: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Replay frozen V-JEPA2 and return its features and pixel VJP.""" + if expected_features.shape != feature_cotangent.shape: + raise ValueError( + "dense encoder replay feature/cotangent shapes differ: " + f"{tuple(expected_features.shape)} != " + f"{tuple(feature_cotangent.shape)}") + student_video_leaf = video01.detach().float().requires_grad_(True) + replay_kwargs = { + "gradient_checkpointing": gradient_checkpointing, + } + if temporal_window_indices is not None: + replay_kwargs["temporal_window_indices"] = temporal_window_indices + replay_features = self._extract_vjepa2_dense( + student_video_leaf, token_indices, **replay_kwargs).float() + encoder_surrogate = (replay_features * feature_cotangent).sum() + (video_cotangent,) = torch.autograd.grad( + encoder_surrogate, student_video_leaf, + retain_graph=False, create_graph=False) + replay_features = replay_features.detach().float() + video_cotangent = video_cotangent.detach().float() + del encoder_surrogate, student_video_leaf + return replay_features, video_cotangent + + def _decode_video01(self, latent: torch.Tensor) -> torch.Tensor: + pixels = self.vae.decode_to_pixel(latent.detach()) + if pixels.ndim != 5 or pixels.shape[2] != 3: + raise ValueError( + "Wan VAE returned an invalid video shape: expected " + f"[B,T,3,H,W], got {tuple(pixels.shape)}") + if ( + self.is_triple_dense + or getattr(self, "is_dinov3_vjepa2_dense", False) + or self._uses_dinov3_dense_backend() + ): + if int(pixels.shape[1]) != 81: + raise ValueError( + "DINOv3/triple dense features require exactly 81 " + f"decoded Wan frames; got {int(pixels.shape[1])}") + elif self.dense_state_prefix == "videomae_v2": + if int(pixels.shape[1]) != 81: + raise ValueError( + "segmented VideoMAE-v2 requires exactly 81 decoded " + f"frames; got {int(pixels.shape[1])}") + else: + if pixels.shape[1] < self.vjepa2_dense_num_frames: + raise ValueError( + "decoded video is too short for the dense-video clip") + student_sampling = getattr( + self, "vjepa2_student_frame_sampling", + getattr(self, "vjepa2_dense_frame_sampling", "center_contiguous")) + if ( + student_sampling in {"random_contiguous_per_video", "last_contiguous"} + and int(pixels.shape[1]) != self.vjepa2_student_source_num_frames + ): + raise ValueError( + f"{student_sampling} V-JEPA2 sampling requires " + f"{self.vjepa2_student_source_num_frames} frames") + return ((pixels + 1.0) * 0.5).clamp(0.0, 1.0) + + def _capture_generation_rng(self) -> dict: + return { + "cpu": torch.get_rng_state().clone(), + "cuda": torch.cuda.get_rng_state(self.device).clone(), + "numpy": np.random.get_state(), + "python": random.getstate(), + } + + @contextmanager + def _replay_generation_rng(self, rng_state: dict): + if set(rng_state) != {"cpu", "cuda", "numpy", "python"}: + raise ValueError( + "invalid fresh-video replay RNG descriptor") + # fork_rng restores the ambient stream afterwards. Consequently the + # second differentiable pass does not consume new training randomness. + ambient_numpy = np.random.get_state() + ambient_python = random.getstate() + try: + with torch.random.fork_rng(devices=[self.device], enabled=True): + torch.set_rng_state(rng_state["cpu"]) + torch.cuda.set_rng_state(rng_state["cuda"], self.device) + np.random.set_state(rng_state["numpy"]) + random.setstate(rng_state["python"]) + yield + finally: + np.random.set_state(ambient_numpy) + random.setstate(ambient_python) + + def _generate_student_latent( + self, + *, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + initial_latent: Optional[torch.Tensor] = None, + generation_descriptor: Optional[dict] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if self.causal_generator: + pred, gradient_mask, _, _ = self._run_generator( + image_or_video_shape=image_or_video_shape, + conditional_dict=dict(conditional_dict), + initial_latent=initial_latent, + generation_descriptor=generation_descriptor, + ) + return pred, gradient_mask + if initial_latent is not None: + raise NotImplementedError( + "fresh V-JEPA2 GradCache currently supports T2V only") + pred = self._run_few_step_generator( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + ) + return pred, None + + def _validate_current_teacher_feature_tensor( + self, + features: torch.Tensor, + *, + student_latent: torch.Tensor, + token_indices: torch.Tensor, + source: str, + ) -> torch.Tensor: + """Return one frozen FP32 teacher matrix with exact student geometry.""" + if not isinstance(features, torch.Tensor): + raise TypeError(f"{source} must be a torch.Tensor") + if not features.is_floating_point(): + raise TypeError(f"{source} must be floating point") + expected_shape = ( + int(student_latent.shape[0]), + int(token_indices.numel()), + int(self._dense_attr("feature_dim")), + ) + if features.ndim != 3 or tuple(features.shape) != expected_shape: + raise ValueError( + f"{source} must have exact [B,N,D] shape " + f"{expected_shape}, got {tuple(features.shape)}") + if features.device != student_latent.device: + raise ValueError( + f"{source} and student latent must share a device") + frozen = features.detach() + if not bool(torch.isfinite(frozen).all().item()): + raise FloatingPointError(f"{source} contains non-finite values") + return frozen.float() + + def _current_teacher_features_for_fresh_video( + self, + *, + student_latent: torch.Tensor, + token_indices: torch.Tensor, + extract_kwargs: Mapping[str, object], + current_teacher_features: Optional[torch.Tensor], + current_teacher_video01: Optional[torch.Tensor], + current_teacher_latent: Optional[torch.Tensor], + ) -> Optional[torch.Tensor]: + """Resolve the mutually exclusive matched-teacher input contract.""" + candidates = { + "current_teacher_features": current_teacher_features, + "current_teacher_video01": current_teacher_video01, + "current_teacher_latent": current_teacher_latent, + } + supplied = [ + name for name, value in candidates.items() if value is not None + ] + if not getattr(self, "current_teacher_attraction_enabled", False): + if supplied: + raise ValueError( + f"{', '.join(supplied)} provided while exact " + "current-teacher attraction is disabled") + return None + if len(supplied) != 1: + raise ValueError( + "exactly one of current_teacher_features, " + "current_teacher_video01, or current_teacher_latent must be " + f"provided; got {supplied}") + + source = supplied[0] + del candidates + if source == "current_teacher_features": + return self._validate_current_teacher_feature_tensor( + current_teacher_features, + student_latent=student_latent, + token_indices=token_indices, + source=source, + ) + + if source == "current_teacher_video01": + teacher_video01 = current_teacher_video01 + if not isinstance(teacher_video01, torch.Tensor): + raise TypeError( + "current_teacher_video01 must be a torch.Tensor") + if not teacher_video01.is_floating_point(): + raise TypeError( + "current_teacher_video01 must be floating point") + if ( + teacher_video01.ndim != 5 + or int(teacher_video01.shape[0]) + != int(student_latent.shape[0]) + or int(teacher_video01.shape[2]) != 3 + or int(teacher_video01.shape[1]) + < int(self._dense_attr("num_frames")) + or int(teacher_video01.shape[3]) <= 0 + or int(teacher_video01.shape[4]) <= 0 + ): + raise ValueError( + "current_teacher_video01 must have shape " + f"[B,T>={int(self._dense_attr('num_frames'))},3,H,W] " + "with the student batch size; got " + f"{tuple(teacher_video01.shape)}") + student_sampling = getattr( + self, + "vjepa2_student_frame_sampling", + getattr( + self, "vjepa2_dense_frame_sampling", "center_contiguous"), + ) + if ( + student_sampling == "last_contiguous" + and int(teacher_video01.shape[1]) not in { + int(self.vjepa2_dense_num_frames), + int(self.vjepa2_student_source_num_frames), + } + ): + raise ValueError( + "last_contiguous teacher video must be either the exact " + "64-frame contract clip or its complete 81-frame source") + if teacher_video01.device != student_latent.device: + raise ValueError( + "current_teacher_video01 and student latent must share " + "a device") + frozen_video01 = teacher_video01.detach() + if not bool(torch.isfinite(frozen_video01).all().item()): + raise FloatingPointError( + "current_teacher_video01 contains non-finite values") + if ( + float(frozen_video01.min().item()) < 0.0 + or float(frozen_video01.max().item()) > 1.0 + ): + raise ValueError( + "current_teacher_video01 must be normalized to [0,1]") + extracted = self._extract_vjepa2_dense( + frozen_video01, + token_indices, + **extract_kwargs, + ) + del current_teacher_video01, frozen_video01, teacher_video01 + else: + teacher_latent = current_teacher_latent + if not isinstance(teacher_latent, torch.Tensor): + raise TypeError( + "current_teacher_latent must be a torch.Tensor") + if not teacher_latent.is_floating_point(): + raise TypeError( + "current_teacher_latent must be floating point") + if tuple(teacher_latent.shape) != tuple(student_latent.shape): + raise ValueError( + "current teacher/student latent shapes must match: " + f"{tuple(teacher_latent.shape)} != " + f"{tuple(student_latent.shape)}") + if teacher_latent.device != student_latent.device: + raise ValueError( + "current teacher/student latents must share a device") + frozen_latent = teacher_latent.detach() + if not bool(torch.isfinite(frozen_latent).all().item()): + raise FloatingPointError( + "current teacher latent contains non-finite values") + teacher_video01 = self._decode_video01(frozen_latent) + extracted = self._extract_vjepa2_dense( + teacher_video01, + token_indices, + **extract_kwargs, + ) + del current_teacher_latent, teacher_latent, frozen_latent, teacher_video01 + + return self._validate_current_teacher_feature_tensor( + extracted, + student_latent=student_latent, + token_indices=token_indices, + source=f"V-JEPA2 output for {source}", + ) + + @property + def dual_video_mmd_enabled(self) -> bool: + """Compatibility name for heterogeneous multi-encoder MMD.""" + return getattr(self, "feature_type", None) in { + "dual_vjepa2_videomae_v2_dense_latent", + "dual_dinov3_vjepa2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_dense_latent", + "triple_dinov3_vjepa2_videomae_v2_t5_joint_dense_latent", + } + + def _dual_token_indices( + self, + device: torch.device, + ) -> dict[str, torch.Tensor]: + indices = { + "vjepa2": uniform_dense_token_indices( + self.vjepa2_dense_raw_tokens, 1024, device=device), + } + if "dinov3" in self.dual_branch_registry: + indices["dinov3"] = torch.tensor( + self.dinov3_dense_frame_indices, + device=device, + dtype=torch.long, + ) + if "videomae_v2" in self.dual_branch_registry: + from video_fd.encoders.videomae_v2_dense_metadata import ( + uniform_segment_token_indices, + ) + + indices["videomae_v2"] = uniform_segment_token_indices( + self.videomae_v2_encoder.dense_raw_tokens, + 256, + device=device, + ) + return indices + + def _validate_dual_feature_mapping( + self, + features: Mapping[str, torch.Tensor], + *, + batch_size: int, + device: torch.device, + label: str, + shape_field: str = "sample_shape", + ) -> dict[str, torch.Tensor]: + expected_keys = tuple(self.dual_branch_registry) + if not isinstance(features, Mapping) or set(features) != set(expected_keys): + raise ValueError( + f"{label} must contain exactly branches {expected_keys}") + result = {} + for name in expected_keys: + tensor = features[name] + branch = self.dual_branch_registry[name] + if shape_field not in branch: + raise ValueError( + f"branch {name!r} has no {shape_field!r} geometry") + expected_shape = ( + batch_size, + *branch[shape_field], + ) + if ( + not isinstance(tensor, torch.Tensor) + or not tensor.is_floating_point() + or tuple(tensor.shape) != expected_shape + or tensor.device != device + ): + raise ValueError( + f"{label}[{name!r}] must be floating {expected_shape} " + f"on {device}; got {getattr(tensor, 'shape', None)}") + frozen = tensor.detach().float() + if not bool(torch.isfinite(frozen).all().item()): + raise FloatingPointError(f"{label}[{name!r}] is non-finite") + result[name] = frozen + return result + + def _extract_multi_encoder_features( + self, + video01: torch.Tensor, + token_indices: Mapping[str, torch.Tensor], + *, + gradient_checkpointing: bool, + ) -> dict[str, torch.Tensor]: + if int(video01.shape[1]) != 81: + raise ValueError( + "multi-encoder video feature extraction requires 81 frames") + temporal_indices = torch.arange( + 17, 81, device=video01.device, dtype=torch.long + ).unsqueeze(0).expand(int(video01.shape[0]), -1) + features = {} + if "dinov3" in self.dual_branch_registry: + features["dinov3"] = self._extract_dinov3_dense(video01) + features["vjepa2"] = self._extract_vjepa2_dense( + video01, + token_indices["vjepa2"], + gradient_checkpointing=gradient_checkpointing, + temporal_window_indices=temporal_indices, + ) + if "videomae_v2" in self.dual_branch_registry: + features["videomae_v2"] = ( + self.videomae_v2_encoder.forward_dense( + video01, + layer_index=15, + token_indices=token_indices["videomae_v2"], + gradient_checkpointing=gradient_checkpointing, + ) + ) + return self._validate_dual_feature_mapping( + features, + batch_size=int(video01.shape[0]), + device=video01.device, + label="multi-encoder output", + shape_field=( + "raw_sample_shape" + if self.triple_t5_joint_enabled else "sample_shape" + ), + ) + + def _fuse_triple_t5_joint_feature_mapping( + self, + video_features: Mapping[str, torch.Tensor], + conditional_dict: Mapping[str, torch.Tensor], + t5_attention_mask: Optional[torch.Tensor], + *, + label: str, + ) -> dict[str, torch.Tensor]: + """Append the same aligned prompt representation to every branch.""" + batch_size = int(next(iter(video_features.values())).shape[0]) + device = next(iter(video_features.values())).device + raw = self._validate_dual_feature_mapping( + video_features, + batch_size=batch_size, + device=device, + label=label, + shape_field=( + "raw_sample_shape" + if self.triple_t5_joint_enabled else "sample_shape" + ), + ) + if not self.triple_t5_joint_enabled: + if t5_attention_mask is not None: + raise ValueError( + "T5 attention masks require the joint feature type") + return raw + if not isinstance(conditional_dict, Mapping): + raise TypeError("joint conditional_dict must be a mapping") + prompt_embeds = conditional_dict.get("prompt_embeds") + if not isinstance(prompt_embeds, torch.Tensor): + raise TypeError("joint features require tensor prompt_embeds") + if not isinstance(t5_attention_mask, torch.Tensor): + raise TypeError("joint features require a T5 attention mask") + if ( + tuple(prompt_embeds.shape[:2]) + != tuple(t5_attention_mask.shape) + or int(prompt_embeds.shape[0]) != batch_size + or prompt_embeds.device != device + or t5_attention_mask.device != device + ): + raise ValueError( + "joint prompt embeddings/mask must align with the video batch") + from video_fd.encoders.triple_t5_joint import ( + fuse_video_t5_features, + ) + + joint = { + name: fuse_video_t5_features( + raw[name], + prompt_embeds, + t5_attention_mask, + branch_name=name, + ).detach().float() + for name in self.dual_branch_registry + } + return self._validate_dual_feature_mapping( + joint, + batch_size=batch_size, + device=device, + label=label.replace("raw", "joint"), + ) + + def _triple_t5_joint_replay_components( + self, + expected: Mapping[str, torch.Tensor], + cotangent: Mapping[str, torch.Tensor], + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor], dict]: + """Keep complete-joint diagnostics but replay only video coordinates.""" + if not self.triple_t5_joint_enabled: + return dict(expected), dict(cotangent), {} + from video_fd.encoders.triple_t5_joint import ( + split_video_t5_joint_features, + ) + + expected_video = {} + video_cotangent = {} + diagnostics = {} + for name in self.dual_branch_registry: + expected_video[name], expected_text = ( + split_video_t5_joint_features( + expected[name], branch_name=name)) + video_cotangent[name], text_cotangent = ( + split_video_t5_joint_features( + cotangent[name], branch_name=name)) + diagnostics.update({ + f"{name}_t5_joint_feature_cotangent_norm": + cotangent[name].norm().detach(), + f"{name}_t5_text_feature_cotangent_norm": + text_cotangent.norm().detach(), + f"{name}_t5_text_feature_norm": + expected_text.norm().detach(), + }) + return expected_video, video_cotangent, diagnostics + + @torch.no_grad() + def _collect_multi_encoder_rollout_features( + self, + *, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + t5_attention_mask: Optional[torch.Tensor], + initial_latent: Optional[torch.Tensor], + current_teacher_features: Optional[Mapping[str, torch.Tensor]], + current_teacher_video01: Optional[torch.Tensor], + current_teacher_latent: Optional[torch.Tensor], + require_current_teacher: bool, + ) -> Tuple[Mapping[str, torch.Tensor], dict]: + if current_teacher_video01 is not None or current_teacher_latent is not None: + raise ValueError( + "dual training reads frozen teacher features from the bank; " + "online teacher video/latent decoding is forbidden") + if require_current_teacher != (current_teacher_features is not None): + raise ValueError( + "selected dual suffix rows require exactly one bank feature mapping") + if not self.fresh_video_prepared_replay_enabled: + raise ValueError( + "dual video MMD requires prepared replay so both encoder VJPs " + "share one decoded pixel tensor and one VAE VJP") + + generation_descriptor = self._sample_generator_input_descriptor( + image_or_video_shape, initial_latent=initial_latent) + saved_generation_noise = generation_descriptor["noise"].detach().clone() + rng_state = self._capture_generation_rng() + student_latent, _ = self._generate_student_latent( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + initial_latent=initial_latent, + generation_descriptor=generation_descriptor, + ) + token_indices = self._dual_token_indices(student_latent.device) + teacher = None + if require_current_teacher: + teacher = self._validate_dual_feature_mapping( + current_teacher_features, + batch_size=int(student_latent.shape[0]), + device=student_latent.device, + label="dual teacher bank features", + shape_field=( + "raw_sample_shape" + if self.triple_t5_joint_enabled else "sample_shape" + ), + ) + student_video01 = self._decode_video01(student_latent) + if int(student_video01.shape[1]) != 81: + raise ValueError("dual student decode must produce exactly 81 frames") + raw_features = self._extract_multi_encoder_features( + student_video01, + token_indices, + gradient_checkpointing=False, + ) + features = self._fuse_triple_t5_joint_feature_mapping( + raw_features, + conditional_dict, + t5_attention_mask, + label="raw multi-encoder student features", + ) + if teacher is not None: + teacher = self._fuse_triple_t5_joint_feature_mapping( + teacher, + conditional_dict, + t5_attention_mask, + label="raw multi-encoder teacher features", + ) + descriptor = { + "rng_state": rng_state, + "latent_shape": tuple(int(value) for value in student_latent.shape), + "generation_descriptor": { + "version": int(generation_descriptor["version"]), + "num_generated_blocks": int( + generation_descriptor["num_generated_blocks"]), + "noise": saved_generation_noise, + }, + "student_latent": student_latent.detach().clone(), + } + if teacher is not None: + descriptor["current_teacher_features"] = teacher + cached_video01 = student_video01.detach() + if self.fresh_video_pixel_cache_device == "cpu": + cached_video01 = cached_video01.to( + device="cpu", dtype=torch.float32, copy=True) + descriptor["cached_student_video01"] = cached_video01 + return features, descriptor + + def _prepare_multi_encoder_latent_gradients( + self, + *, + replay_descriptor: dict, + expected_features: Mapping[str, torch.Tensor], + feature_cotangent: Mapping[str, torch.Tensor], + replay_atol: float, + replay_rtol: float, + ) -> dict: + required = { + "rng_state", + "latent_shape", + "generation_descriptor", + "student_latent", + "cached_student_video01", + } + if not isinstance(replay_descriptor, dict) or set(replay_descriptor) != required: + raise ValueError("invalid cached triple-video replay descriptor") + student_latent = replay_descriptor["student_latent"] + cached = replay_descriptor.pop("cached_student_video01") + student_video01 = cached.to( + device=student_latent.device, + dtype=torch.float32, + non_blocking=bool(cached.is_pinned()), + ) + batch_size = int(student_latent.shape[0]) + expected = self._validate_dual_feature_mapping( + expected_features, + batch_size=batch_size, + device=student_latent.device, + label="cached triple student features", + ) + cotangent = self._validate_dual_feature_mapping( + feature_cotangent, + batch_size=batch_size, + device=student_latent.device, + label="triple feature cotangent", + ) + full_feature_cotangent = cotangent + expected, cotangent, joint_replay_log = ( + self._triple_t5_joint_replay_components(expected, cotangent)) + if self.triple_t5_joint_enabled: + expected = self._validate_dual_feature_mapping( + expected, + batch_size=batch_size, + device=student_latent.device, + label="cached triple raw-video student features", + shape_field="raw_sample_shape", + ) + cotangent = self._validate_dual_feature_mapping( + cotangent, + batch_size=batch_size, + device=student_latent.device, + label="triple raw-video feature cotangent", + shape_field="raw_sample_shape", + ) + tokens = self._dual_token_indices(student_latent.device) + temporal_indices = torch.arange( + 17, 81, device=student_latent.device, dtype=torch.long + ).unsqueeze(0).expand(batch_size, -1) + + combined_pixel = None + pixel_shape = None + pixel_norms = {} + replay_maxima = {} + if "dinov3" in self.dual_branch_registry: + with torch.enable_grad(): + pixel_dino, dino_max_abs = self._replay_dinov3_dense_to_pixels( + student_video01, + expected_features=expected["dinov3"], + feature_cotangent=cotangent["dinov3"], + replay_atol=replay_atol, + replay_rtol=replay_rtol, + ) + combined_pixel, pixel_shape, pixel_norms["dinov3"] = ( + _accumulate_fp32_pixel_cotangent( + combined_pixel, + pixel_dino, + pixel_shape, + branch_name="dinov3", + ) + ) + replay_maxima["dinov3"] = dino_max_abs + del pixel_dino + + with torch.enable_grad(): + replay_vj, pixel_vj = self._replay_dense_encoder_vjp( + student_video01, + tokens["vjepa2"], + expected["vjepa2"], + cotangent["vjepa2"], + gradient_checkpointing=( + self.vjepa2_dense_gradient_checkpointing), + temporal_window_indices=temporal_indices, + ) + if not torch.allclose( + replay_vj, expected["vjepa2"], + atol=float(replay_atol), rtol=float(replay_rtol), + ): + raise RuntimeError("triple V-JEPA2 replay changed cached features") + combined_pixel, pixel_shape, pixel_norms["vjepa2"] = ( + _accumulate_fp32_pixel_cotangent( + combined_pixel, + pixel_vj, + pixel_shape, + branch_name="vjepa2", + ) + ) + replay_maxima["vjepa2"] = ( + replay_vj - expected["vjepa2"]).abs().max().detach() + del replay_vj, pixel_vj + + if "videomae_v2" in self.dual_branch_registry: + with torch.enable_grad(): + replay_vm, pixel_vm = self.videomae_v2_encoder.replay_dense_vjp( + student_video01.detach().float(), + layer_indices=(15,), + token_indices=tokens["videomae_v2"], + expected_features=expected["videomae_v2"], + feature_cotangent=cotangent["videomae_v2"], + gradient_checkpointing=( + self.videomae_v2_dense_gradient_checkpointing), + ) + if not torch.allclose( + replay_vm, expected["videomae_v2"], + atol=float(replay_atol), rtol=float(replay_rtol), + ): + raise RuntimeError( + "VideoMAE-v2 replay changed cached features") + combined_pixel, pixel_shape, pixel_norms["videomae_v2"] = ( + _accumulate_fp32_pixel_cotangent( + combined_pixel, + pixel_vm, + pixel_shape, + branch_name="videomae_v2", + ) + ) + replay_maxima["videomae_v2"] = ( + replay_vm - expected["videomae_v2"]).abs().max().detach() + del replay_vm, pixel_vm + del student_video01 + + if set(pixel_norms) != set(self.dual_branch_registry): + raise RuntimeError( + "triple encoder replay did not visit every branch") + if not bool(torch.isfinite(combined_pixel).all().item()): + raise FloatingPointError( + "combined triple pixel cotangent is non-finite") + latent_cotangent = replay_decoder_vjp( + self.vae, + student_latent.detach(), + combined_pixel, + independent_first_frame=self.independent_first_frame, + ).detach().float() + if latent_cotangent.shape != student_latent.shape: + raise ValueError( + "triple Wan VAE VJP returned an invalid latent shape") + if not bool(torch.isfinite(latent_cotangent).all().item()): + raise FloatingPointError( + "triple Wan VAE VJP returned non-finite values") + + prefix = self.dense_log_prefix + feature_norm = torch.sqrt(torch.stack([ + value.norm().float().square() + for value in full_feature_cotangent.values() + ]).sum()) + replay_max_abs = torch.stack([ + value.float() for value in replay_maxima.values() + ]).max() + replay_log = { + f"{prefix}_feature_grad_norm": feature_norm.detach(), + f"{prefix}_video_grad_norm": combined_pixel.norm().detach(), + f"{prefix}_latent_grad_norm": latent_cotangent.norm().detach(), + f"{prefix}_replay_max_abs": replay_max_abs.detach(), + "combined_latent_grad_norm": latent_cotangent.norm().detach(), + "fresh_video_prepared_vae_vjp": latent_cotangent.new_ones(()), + f"{prefix}_encoder_vjp_count": latent_cotangent.new_tensor( + len(pixel_norms), dtype=torch.long), + f"{prefix}_vae_vjp_count": latent_cotangent.new_tensor( + 1, dtype=torch.long), + } + replay_log.update(joint_replay_log) + if self.triple_t5_joint_enabled: + replay_log["triple_t5_text_cotangent_discarded"] = ( + latent_cotangent.new_ones(())) + for name in self.dual_branch_registry: + replay_log[f"{name}_dense_feature_grad_norm"] = ( + full_feature_cotangent[name].norm().detach()) + if self.triple_t5_joint_enabled: + replay_log[f"{name}_dense_video_feature_grad_norm"] = ( + cotangent[name].norm().detach()) + replay_log[f"{name}_dense_video_grad_norm"] = ( + pixel_norms[name]) + replay_log[f"{name}_dense_replay_max_abs"] = ( + replay_maxima[name].detach()) + return { + "latent_cotangent": latent_cotangent, + "replay_log": replay_log, + } + + @torch.no_grad() + def collect_rollout_features( + self, + *, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + t5_attention_mask: Optional[torch.Tensor] = None, + initial_latent: Optional[torch.Tensor] = None, + current_teacher_features: Optional[torch.Tensor] = None, + current_teacher_video01: Optional[torch.Tensor] = None, + current_teacher_latent: Optional[torch.Tensor] = None, + require_current_teacher: bool = True, + vjepa2_window_sample_keys=None, + ) -> Tuple[torch.Tensor, dict]: + """First GradCache pass: generate and encode one fresh microbatch.""" + if not self.uses_fresh_video_gradcache: + raise RuntimeError( + "collect_rollout_features requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if self.dual_video_mmd_enabled: + return self._collect_multi_encoder_rollout_features( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + t5_attention_mask=t5_attention_mask, + initial_latent=initial_latent, + current_teacher_features=current_teacher_features, + current_teacher_video01=current_teacher_video01, + current_teacher_latent=current_teacher_latent, + require_current_teacher=require_current_teacher, + ) + if type(require_current_teacher) is not bool: + raise TypeError("require_current_teacher must be a bool") + supplied_teacher_inputs = [ + name for name, value in { + "current_teacher_features": current_teacher_features, + "current_teacher_video01": current_teacher_video01, + "current_teacher_latent": current_teacher_latent, + }.items() if value is not None + ] + if not require_current_teacher and supplied_teacher_inputs: + raise ValueError( + "require_current_teacher=false rejects all teacher inputs; " + f"got {supplied_teacher_inputs}") + first_pass_error = None + try: + generation_descriptor = self._sample_generator_input_descriptor( + image_or_video_shape, + initial_latent=initial_latent, + ) + saved_generation_noise = ( + generation_descriptor["noise"].detach().clone()) + # Initial Gaussian noise and the block count are explicit. The + # saved state begins *after* sampling them and drives any + # stochastic denoise/renoise choices inside the rollout pipeline. + rng_state = self._capture_generation_rng() + student_latent, _ = self._generate_student_latent( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + initial_latent=initial_latent, + generation_descriptor=generation_descriptor, + ) + matched_teacher_features = None + teacher_extract_seconds = None + student_video01 = self._decode_video01(student_latent) + token_indices = None + temporal_window_indices = None + if self._uses_dinov3_dense_backend(): + if vjepa2_window_sample_keys is not None: + raise ValueError( + "DINOv3 fixed-frame sampling rejects V-JEPA2 keys") + features = self._extract_dinov3_dense( + student_video01).detach().float() + if require_current_teacher: + if ( + current_teacher_video01 is not None + or current_teacher_latent is not None + or not isinstance(current_teacher_features, torch.Tensor) + ): + raise ValueError( + "DINOv3 current teacher must be detached features") + matched_teacher_features = ( + current_teacher_features.detach().float()) + if tuple(matched_teacher_features.shape) != tuple(features.shape): + raise ValueError( + "DINOv3 teacher/student feature shapes differ") + else: + token_indices = self._vjepa2_token_indices( + student_latent.device) + temporal_window_indices = ( + self._vjepa2_temporal_window_indices( + vjepa2_window_sample_keys, + batch_size=int(student_latent.shape[0]), + device=student_latent.device, + ) + ) + extract_kwargs = {"gradient_checkpointing": False} + if temporal_window_indices is not None: + extract_kwargs["temporal_window_indices"] = ( + temporal_window_indices) + # A live-video cache miss supplies an already selected 64-frame + # clip. A latent or complete 81-frame video uses the student's + # explicit temporal indices. + teacher_extract_kwargs = dict(extract_kwargs) + if ( + current_teacher_video01 is not None + and temporal_window_indices is not None + and int(current_teacher_video01.shape[1]) + == int(self.vjepa2_dense_num_frames) + ): + teacher_extract_kwargs.pop("temporal_window_indices", None) + teacher_requires_extraction = bool( + require_current_teacher + and ( + current_teacher_video01 is not None + or current_teacher_latent is not None + ) + ) + if teacher_requires_extraction and student_latent.is_cuda: + torch.cuda.synchronize(student_latent.device) + teacher_extract_started = time.perf_counter() + matched_teacher_features = ( + self._current_teacher_features_for_fresh_video( + student_latent=student_latent, + token_indices=token_indices, + extract_kwargs=teacher_extract_kwargs, + current_teacher_features=current_teacher_features, + current_teacher_video01=current_teacher_video01, + current_teacher_latent=current_teacher_latent, + ) + if require_current_teacher + else None + ) + if teacher_requires_extraction and student_latent.is_cuda: + torch.cuda.synchronize(student_latent.device) + teacher_extract_seconds = ( + time.perf_counter() - teacher_extract_started + if teacher_requires_extraction else None + ) + # The trainer transfers ownership with dict.pop(). Drop the three + # input aliases now so an uncached live video can be reclaimed + # before student VAE decode/V-JEPA2 extraction begins. + current_teacher_features = None + current_teacher_video01 = None + current_teacher_latent = None + + student_video01 = self._decode_video01(student_latent) + features = self._extract_vjepa2_dense( + student_video01, + token_indices, + **extract_kwargs, + ).detach().float() + if ( + matched_teacher_features is not None + and tuple(matched_teacher_features.shape) + != tuple(features.shape) + ): + raise ValueError( + "matched teacher/student V-JEPA2 feature shapes must " + "match exactly: " + f"{tuple(matched_teacher_features.shape)} != " + f"{tuple(features.shape)}") + except BaseException as exc: + first_pass_error = exc + if not self._all_ranks_true( + first_pass_error is None, self.device + ): + message = ( + f": {type(first_pass_error).__name__}: {first_pass_error}" + if first_pass_error is not None + else " on another rank" + ) + raise RuntimeError( + "fresh video generation/VAE/dense-encoder first pass failed" + + message + ) from first_pass_error + local_features_finite = bool(torch.isfinite(features).all().item()) + if not self._all_ranks_true( + local_features_finite, features.device + ): + raise FloatingPointError( + "non-finite dense feature in fresh GradCache first pass " + "on at least one rank") + if matched_teacher_features is not None: + local_teacher_features_finite = bool( + torch.isfinite(matched_teacher_features).all().item()) + if not self._all_ranks_true( + local_teacher_features_finite, + matched_teacher_features.device, + ): + raise FloatingPointError( + "non-finite matched teacher feature in fresh GradCache " + "first pass on at least one rank") + descriptor = { + "rng_state": rng_state, + "latent_shape": tuple(int(v) for v in student_latent.shape), + "generation_descriptor": { + "version": int(generation_descriptor["version"]), + "num_generated_blocks": int( + generation_descriptor["num_generated_blocks"]), + "noise": saved_generation_noise, + }, + "student_latent": student_latent.detach().clone(), + } + if matched_teacher_features is not None: + # The trainer moves this detached value into the update-local + # GradCache immediately; teacher latents/pixels are never replayed. + descriptor["current_teacher_features"] = ( + matched_teacher_features) + if teacher_extract_seconds is not None: + descriptor["current_teacher_extract_seconds"] = float( + teacher_extract_seconds) + if temporal_window_indices is not None: + descriptor["vjepa2_frame_indices"] = ( + temporal_window_indices.detach().clone()) + if self._uses_dinov3_dense_backend(): + descriptor["dinov3_frame_indices"] = torch.tensor( + self.dinov3_dense_frame_indices, + device=student_latent.device, + dtype=torch.long, + ) + if self.fresh_video_prepared_replay_enabled: + cache_error = None + try: + # Keep the exact first-pass FP32 pixels until update-wide + # feature cotangents are available. They are consumed before + # any differentiable generator replay, so this storage never + # overlaps the generator graph peak. + cached_video01 = student_video01.detach() + if self.fresh_video_pixel_cache_device == "cpu": + cached_video01 = cached_video01.to( + device="cpu", dtype=torch.float32, copy=True) + descriptor["cached_student_video01"] = cached_video01 + except BaseException as exc: + cache_error = exc + if not self._all_ranks_true( + cache_error is None, student_latent.device + ): + message = ( + f": {type(cache_error).__name__}: {cache_error}" + if cache_error is not None + else " on another rank" + ) + raise RuntimeError( + "fresh-video pixel caching failed" + message + ) from cache_error + return features, descriptor + + def prepare_latent_gradients( + self, + *, + replay_descriptor: dict, + expected_features: torch.Tensor, + feature_cotangent: torch.Tensor, + replay_atol: float = 1e-5, + replay_rtol: float = 1e-4, + ) -> dict: + """Replay frozen encoders and prepare a detached latent cotangent.""" + if self.dual_video_mmd_enabled: + return self._prepare_multi_encoder_latent_gradients( + replay_descriptor=replay_descriptor, + expected_features=expected_features, + feature_cotangent=feature_cotangent, + replay_atol=replay_atol, + replay_rtol=replay_rtol, + ) + prepared_enabled = self.fresh_video_prepared_replay_enabled + if not prepared_enabled or not self.uses_fresh_video_gradcache: + raise RuntimeError( + "prepared replay is unavailable for this fresh-video path") + prepare_error = None + prepared = None + try: + required_descriptor_keys = { + "rng_state", + "latent_shape", + "generation_descriptor", + "student_latent", + "cached_student_video01", + } + if getattr( + self, "vjepa2_student_frame_sampling", + getattr(self, "vjepa2_dense_frame_sampling", "center_contiguous"), + ) != "center_contiguous": + required_descriptor_keys.add("vjepa2_frame_indices") + if self._uses_dinov3_dense_backend(): + required_descriptor_keys.add("dinov3_frame_indices") + if ( + not isinstance(replay_descriptor, dict) + or set(replay_descriptor) != required_descriptor_keys + ): + raise ValueError( + "invalid cached fresh-video replay descriptor") + + cached_student_video01 = replay_descriptor.pop( + "cached_student_video01") + student_latent = replay_descriptor["student_latent"] + if ( + not isinstance(student_latent, torch.Tensor) + or tuple(student_latent.shape) + != tuple(replay_descriptor["latent_shape"]) + ): + raise ValueError( + "cached first-pass latent has an invalid shape") + if not isinstance(cached_student_video01, torch.Tensor): + raise TypeError("cached first-pass pixels must be a tensor") + student_video01 = cached_student_video01.to( + device=student_latent.device, + dtype=torch.float32, + non_blocking=bool(cached_student_video01.is_pinned()), + ) + del cached_student_video01 + if ( + student_video01.ndim != 5 + or student_video01.shape[0] != student_latent.shape[0] + or student_video01.shape[2] != 3 + ): + raise ValueError( + "cached first-pass pixels have an invalid shape") + if not bool(torch.isfinite(student_video01).all().item()): + raise FloatingPointError( + "cached first-pass pixels are non-finite") + + temporal_window_indices = None + token_indices = None + if self._uses_dinov3_dense_backend(): + self._dinov3_descriptor_frame_indices( + replay_descriptor, device=student_latent.device) + else: + temporal_window_indices = ( + self._vjepa2_descriptor_window_indices( + replay_descriptor, device=student_latent.device) + ) + token_indices = self._vjepa2_token_indices( + student_latent.device) + + expected_features = expected_features.to( + device=student_latent.device, dtype=torch.float32) + feature_cotangent = feature_cotangent.to( + device=student_latent.device, dtype=torch.float32) + if feature_cotangent.shape != expected_features.shape: + raise ValueError( + "feature cotangent shape does not match cached features: " + f"{tuple(feature_cotangent.shape)} != " + f"{tuple(expected_features.shape)}") + if not bool(torch.isfinite(feature_cotangent).all().item()): + raise FloatingPointError( + "non-finite fresh dense feature cotangent") + layer_feature_grad_norms = {} + if getattr(self, "vjepa2_multilayer_enabled", False): + expected_transport_tokens = ( + len(self.vjepa2_dense_layers) + * self.vjepa2_dense_num_tokens + ) + if int(feature_cotangent.shape[1]) != expected_transport_tokens: + raise ValueError( + "multi-layer feature cotangent has an invalid " + "transport token count") + layer_feature_grad_norms = { + f"vjepa2_dense_l{layer}_feature_grad_norm": ( + part.norm().detach()) + for layer, part in zip( + self.vjepa2_dense_layers, + feature_cotangent.split( + self.vjepa2_dense_num_tokens, dim=1), + ) + } + + if self._uses_dinov3_dense_backend(): + video_cotangent, replay_max_abs = ( + self._replay_dinov3_dense_to_pixels( + student_video01, + expected_features=expected_features, + feature_cotangent=feature_cotangent, + replay_atol=float(replay_atol), + replay_rtol=float(replay_rtol), + ) + ) + else: + with torch.enable_grad(): + student_video_leaf = ( + student_video01.detach().requires_grad_(True)) + replay_extract_kwargs = { + "gradient_checkpointing": ( + self.vjepa2_dense_gradient_checkpointing), + } + if temporal_window_indices is not None: + replay_extract_kwargs["temporal_window_indices"] = ( + temporal_window_indices) + replay_features = self._extract_vjepa2_dense( + student_video_leaf, + token_indices, + **replay_extract_kwargs, + ).float() + if replay_features.shape != expected_features.shape: + raise RuntimeError( + "fresh feature replay changed matrix shape: " + f"{tuple(replay_features.shape)} != " + f"{tuple(expected_features.shape)}") + replay_max_abs = ( + replay_features.detach() + - expected_features).abs().max() + if not bool(torch.allclose( + replay_features.detach(), + expected_features, + atol=float(replay_atol), + rtol=float(replay_rtol), + )): + raise RuntimeError( + "cached pixels changed V-JEPA2 features " + f"(max_abs={float(replay_max_abs):.6g})") + encoder_surrogate = ( + replay_features * feature_cotangent).sum() + (video_cotangent,) = torch.autograd.grad( + encoder_surrogate, + student_video_leaf, + retain_graph=False, + create_graph=False, + ) + del encoder_surrogate, student_video_leaf, replay_features + video_cotangent = video_cotangent.detach().float() + if not bool(torch.isfinite(video_cotangent).all().item()): + raise FloatingPointError( + "non-finite fresh dense-encoder pixel cotangent") + video_grad_norm = video_cotangent.norm().detach() + del student_video01 + + latent_cotangent = replay_decoder_vjp( + self.vae, + student_latent.detach(), + video_cotangent, + independent_first_frame=self.independent_first_frame, + ).detach().float() + del video_cotangent + if latent_cotangent.shape != student_latent.shape: + raise ValueError( + "Wan VAE VJP returned an invalid latent shape") + if not bool(torch.isfinite(latent_cotangent).all().item()): + raise FloatingPointError( + "Wan VAE VJP returned a non-finite cotangent") + + replay_log = { + self._dense_key("feature_grad_norm"): ( + feature_cotangent.norm().detach()), + self._dense_key("video_grad_norm"): video_grad_norm, + self._dense_key("latent_grad_norm"): ( + latent_cotangent.norm().detach()), + "combined_latent_grad_norm": ( + latent_cotangent.norm().detach()), + self._dense_key("replay_max_abs"): replay_max_abs.detach(), + "fresh_video_prepared_vae_vjp": ( + latent_cotangent.new_ones(())), + } + replay_log.update(layer_feature_grad_norms) + prepared = { + "latent_cotangent": latent_cotangent.detach(), + "replay_log": replay_log, + } + except BaseException as exc: + prepare_error = exc + prepare_ok = self._all_ranks_true( + prepare_error is None, self.device) + if not prepare_ok: + all_errors_are_nonfinite = self._all_ranks_true( + prepare_error is None + or isinstance(prepare_error, FloatingPointError), + self.device, + ) + message = ( + f": {type(prepare_error).__name__}: {prepare_error}" + if prepare_error is not None + else " on another rank" + ) + if all_errors_are_nonfinite: + raise FloatingPointError( + "fresh-video prepared replay produced a non-finite value" + + message + ) from prepare_error + raise RuntimeError( + "fresh-video replay preparation failed" + message + ) from prepare_error + return prepared + + + def replay_generator_gradients( + self, + *, + replay_descriptor: dict, + expected_features: torch.Tensor, + feature_cotangent: torch.Tensor, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + initial_latent: Optional[torch.Tensor] = None, + prepared_replay: Optional[dict] = None, + replay_atol: float = 1e-5, + replay_rtol: float = 1e-4, + ) -> Tuple[torch.Tensor, dict]: + """Second GradCache pass: replay one video microbatch and encoder/VAE VJP.""" + if not self.uses_fresh_video_gradcache: + raise RuntimeError( + "replay_generator_gradients requires " + "wan_mmd_estimator=nystrom_exact_repulsion") + if self.dual_video_mmd_enabled and prepared_replay is None: + raise RuntimeError("dual video MMD requires prepared replay") + required_descriptor_keys = { + "rng_state", + "latent_shape", + "generation_descriptor", + "student_latent", + } + # Prepared multi-encoder replay has already consumed encoder-specific + # frame fields. Generator replay retains only the common identity keys. + if not self.dual_video_mmd_enabled: + if getattr( + self, "vjepa2_student_frame_sampling", + getattr(self, "vjepa2_dense_frame_sampling", "center_contiguous"), + ) != "center_contiguous": + required_descriptor_keys.add("vjepa2_frame_indices") + if self._uses_dinov3_dense_backend(): + required_descriptor_keys.add("dinov3_frame_indices") + if ( + not isinstance(replay_descriptor, dict) + or set(replay_descriptor) != required_descriptor_keys + ): + raise ValueError("invalid fresh-video replay descriptor") + generator_replay_error = None + try: + with self._replay_generation_rng(replay_descriptor["rng_state"]): + student_latent, gradient_mask = self._generate_student_latent( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + initial_latent=initial_latent, + generation_descriptor=replay_descriptor[ + "generation_descriptor"], + ) + except BaseException as exc: + generator_replay_error = exc + if not self._all_ranks_true( + generator_replay_error is None, self.device + ): + message = ( + f": {type(generator_replay_error).__name__}: " + f"{generator_replay_error}" + if generator_replay_error is not None + else " on another rank" + ) + raise RuntimeError( + "fresh generator replay failed" + message + ) from generator_replay_error + expected_latent_shape = tuple(replay_descriptor["latent_shape"]) + local_latent_shape_ok = ( + tuple(student_latent.shape) == expected_latent_shape) + if not self._all_ranks_true( + local_latent_shape_ok, student_latent.device + ): + raise RuntimeError( + "fresh generator replay changed latent shape on at least one " + f"rank; local={tuple(student_latent.shape)}, " + f"expected={expected_latent_shape}") + expected_student_latent = replay_descriptor["student_latent"].to( + device=student_latent.device, dtype=student_latent.dtype) + latent_replay_max_abs = ( + student_latent.detach() - expected_student_latent + ).abs().max() + local_latent_replay_matches = bool(torch.allclose( + student_latent.detach(), + expected_student_latent, + atol=float(replay_atol), + rtol=float(replay_rtol), + )) + if not self._all_ranks_true( + local_latent_replay_matches, student_latent.device + ): + raise RuntimeError( + "fresh generator replay changed the generated latent on at " + "least one rank " + f"(local max_abs={float(latent_replay_max_abs):.6g})") + student_latent = self._apply_gradient_mask( + student_latent, gradient_mask) + + if prepared_replay is not None: + # The detached decoder VJP was evaluated at the exact first-pass + # latent. Unlike the legacy value-side replay, an allclose-only + # match is insufficient because J_D(z) changes with z. + local_prepared_latent_exact = bool(torch.equal( + student_latent.detach(), expected_student_latent)) + if not self._all_ranks_true( + local_prepared_latent_exact, student_latent.device + ): + raise RuntimeError( + "prepared fresh-video replay requires bitwise-identical generator " + "latents on every rank " + f"(local max_abs={float(latent_replay_max_abs):.6g})") + + prepared_error = None + latent_cotangent = None + replay_log = None + try: + prepared_enabled = self.fresh_video_prepared_replay_enabled + if not prepared_enabled: + raise RuntimeError( + "prepared replay was supplied while P0 is disabled") + if ( + not isinstance(prepared_replay, dict) + or set(prepared_replay) + != {"latent_cotangent", "replay_log"} + or not isinstance(prepared_replay.get("replay_log"), dict) + ): + raise ValueError( + "invalid prepared fresh-video replay payload") + latent_cotangent = prepared_replay["latent_cotangent"] + if ( + not isinstance(latent_cotangent, torch.Tensor) + or latent_cotangent.requires_grad + or latent_cotangent.grad_fn is not None + or latent_cotangent.shape != student_latent.shape + ): + raise ValueError( + "prepared fresh-video latent cotangent has invalid structure") + latent_cotangent = latent_cotangent.to( + device=student_latent.device, dtype=torch.float32) + if not bool(torch.isfinite( + latent_cotangent).all().item() + ): + raise FloatingPointError( + "prepared fresh-video latent cotangent is non-finite") + replay_log = dict(prepared_replay["replay_log"]) + except BaseException as exc: + prepared_error = exc + prepared_ok = self._all_ranks_true( + prepared_error is None, student_latent.device) + if not prepared_ok: + all_errors_are_nonfinite = self._all_ranks_true( + prepared_error is None + or isinstance(prepared_error, FloatingPointError), + student_latent.device, + ) + message = ( + f": {type(prepared_error).__name__}: {prepared_error}" + if prepared_error is not None + else " on another rank" + ) + if all_errors_are_nonfinite: + raise FloatingPointError( + "prepared fresh-video replay produced a non-finite value" + + message + ) from prepared_error + raise RuntimeError( + "prepared fresh-video replay validation failed" + message + ) from prepared_error + replay_surrogate = ( + student_latent.float() * latent_cotangent).sum() + loss = replay_surrogate - replay_surrogate.detach() + dense_prefix = self._dense_metric_prefix() + replay_log.update({ + self._dense_key("latent_replay_max_abs"): ( + latent_replay_max_abs.detach()), + "fresh_video_prepared_replay": ( + latent_cotangent.new_ones(())), + }) + return loss, replay_log + + token_indices = None + temporal_window_indices = None + if not self._uses_dinov3_dense_backend(): + token_indices = self._vjepa2_token_indices( + student_latent.device) + temporal_window_indices = ( + self._vjepa2_descriptor_window_indices( + replay_descriptor, device=student_latent.device) + ) + + # Recreate the value-side VAE/V-JEPA2 path and prove that the replayed + # generator sample is exactly the sample scored in the first pass. + with torch.no_grad(): + student_video01 = self._decode_video01(student_latent) + expected_features = expected_features.to( + device=student_latent.device, dtype=torch.float32) + + feature_cotangent = feature_cotangent.to( + device=student_latent.device, dtype=torch.float32) + local_cotangent_shape_ok = ( + feature_cotangent.shape == expected_features.shape) + if not self._all_ranks_true( + local_cotangent_shape_ok, student_latent.device + ): + raise ValueError( + "feature cotangent shape does not match the cached video " + "matrix on at least one rank: " + f"{tuple(feature_cotangent.shape)} != " + f"{tuple(expected_features.shape)}") + local_cotangent_finite = bool( + torch.isfinite(feature_cotangent).all().item()) + if not self._all_ranks_true( + local_cotangent_finite, student_latent.device + ): + raise FloatingPointError( + "non-finite fresh dense feature cotangent on at least one " + "rank") + + # Frozen encoder VJP to decoded pixels. No rank enters a subsequent + # collective until every locally fallible encoder/autograd operation + # has either completed or been captured, preventing a one-rank replay + # exception from stranding peers in an all-reduce. + encoder_replay_error = None + try: + if self._uses_dinov3_dense_backend(): + self._dinov3_descriptor_frame_indices( + replay_descriptor, device=student_latent.device) + video_cotangent, replay_max_abs = ( + self._replay_dinov3_dense_to_pixels( + student_video01, + expected_features=expected_features, + feature_cotangent=feature_cotangent, + replay_atol=float(replay_atol), + replay_rtol=float(replay_rtol), + ) + ) + else: + with torch.enable_grad(): + student_video_leaf = ( + student_video01.detach().float().requires_grad_(True)) + replay_extract_kwargs = { + "gradient_checkpointing": ( + self.vjepa2_dense_gradient_checkpointing), + } + if temporal_window_indices is not None: + replay_extract_kwargs["temporal_window_indices"] = ( + temporal_window_indices) + replay_features = self._extract_vjepa2_dense( + student_video_leaf, + token_indices, + **replay_extract_kwargs, + ).float() + if replay_features.shape != expected_features.shape: + raise RuntimeError( + "fresh feature replay changed matrix shape: " + f"{tuple(replay_features.shape)} != " + f"{tuple(expected_features.shape)}") + replay_max_abs = ( + replay_features.detach() + - expected_features).abs().max() + if not bool(torch.allclose( + replay_features.detach(), + expected_features, + atol=float(replay_atol), + rtol=float(replay_rtol), + )): + raise RuntimeError( + "fresh generator replay is not deterministic in " + "V-JEPA2 space " + f"(max_abs={float(replay_max_abs):.6g})") + encoder_surrogate = ( + replay_features * feature_cotangent).sum() + (video_cotangent,) = torch.autograd.grad( + encoder_surrogate, + student_video_leaf, + retain_graph=False, + create_graph=False, + ) + del encoder_surrogate, student_video_leaf, replay_features + except BaseException as exc: + encoder_replay_error = exc + if not self._all_ranks_true( + encoder_replay_error is None, student_latent.device + ): + message = ( + f": {type(encoder_replay_error).__name__}: " + f"{encoder_replay_error}" + if encoder_replay_error is not None + else " on another rank" + ) + raise RuntimeError( + "frozen dense encoder replay failed" + message + ) from encoder_replay_error + video_cotangent = video_cotangent.detach().float() + local_video_finite = bool( + torch.isfinite(video_cotangent).all().item()) + if not self._all_ranks_true( + local_video_finite, video_cotangent.device + ): + raise FloatingPointError( + "non-finite fresh dense-encoder pixel cotangent on at least one " + "rank") + mmd_video_grad_norm = video_cotangent.norm().detach() + # ``encoder_surrogate`` is the autograd root and therefore keeps the + # leaf (and its full decoded-video storage) alive even after + # ``autograd.grad`` has consumed the graph. Drop the already-consumed + # root together with the leaf so ``student_video01 = None`` below can + # actually release that storage before the frame VAE VJP. No tensor + # operation or derivative changes: all three values are dead here. + # Frozen causal Wan VAE VJP to the generated latent, with the same + # rank-consensus exception boundary as the encoder replay above. + vae_replay_error = None + try: + mmd_latent_cotangent = replay_decoder_vjp( + self.vae, + student_latent.detach(), + video_cotangent, + independent_first_frame=self.independent_first_frame, + ).detach().float() + if mmd_latent_cotangent.shape != student_latent.shape: + raise ValueError( + "Wan VAE decoder VJP returned an invalid latent shape: " + f"{tuple(mmd_latent_cotangent.shape)} != " + f"{tuple(student_latent.shape)}") + except BaseException as exc: + vae_replay_error = exc + if not self._all_ranks_true( + vae_replay_error is None, student_latent.device + ): + message = ( + f": {type(vae_replay_error).__name__}: {vae_replay_error}" + if vae_replay_error is not None + else " on another rank" + ) + raise RuntimeError( + "frozen Wan VAE replay failed" + message + ) from vae_replay_error + local_latent_finite = bool( + torch.isfinite(mmd_latent_cotangent).all().item()) + if not self._all_ranks_true( + local_latent_finite, mmd_latent_cotangent.device + ): + raise FloatingPointError( + "non-finite fresh dense-feature latent cotangent on at least one " + "rank") + del video_cotangent + + del student_video01 + latent_cotangent = mmd_latent_cotangent + + # The global scalar value is logged from FreshVideoGradCache. This + # zero-valued surrogate carries only its exact replayed derivative. + replay_surrogate = ( + student_latent.float() * latent_cotangent).sum() + loss = replay_surrogate - replay_surrogate.detach() + dense_prefix = self._dense_metric_prefix() + return loss, { + f"{dense_prefix}_feature_grad_norm": ( + feature_cotangent.norm().detach()), + self._dense_key("video_grad_norm"): ( + mmd_video_grad_norm), + self._dense_key("latent_grad_norm"): ( + mmd_latent_cotangent.norm().detach()), + "combined_latent_grad_norm": ( + latent_cotangent.norm().detach()), + self._dense_key("replay_max_abs"): replay_max_abs.detach(), + self._dense_key("latent_replay_max_abs"): ( + latent_replay_max_abs.detach()), + } + + def compute_latent_mmd_loss( + self, + student_latent: torch.Tensor, + teacher_latent: torch.Tensor, + conditional_dict: dict, + gradient_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, dict]: + if self.uses_fresh_video_gradcache: + raise RuntimeError( + "nystrom_exact_repulsion must be evaluated through the " + "trainer's full fresh-video GradCache batch") + student_latent = self._apply_gradient_mask(student_latent, gradient_mask) + teacher_latent = teacher_latent.to( + device=student_latent.device, dtype=student_latent.dtype) + if teacher_latent.shape != student_latent.shape: + raise ValueError( + f"teacher latent {tuple(teacher_latent.shape)} != " + f"student latent {tuple(student_latent.shape)}") + + if self.feature_type == "post_block_latent": + return self._compute_wan_latent_mmd_loss( + student_latent, teacher_latent, conditional_dict) + return self._compute_vjepa2_dense_mmd_loss( + student_latent, teacher_latent) + + def _compute_wan_latent_mmd_loss( + self, + student_latent: torch.Tensor, + teacher_latent: torch.Tensor, + conditional_dict: dict, + ) -> Tuple[torch.Tensor, dict]: + indices = self._token_indices(student_latent) + # Teacher first: its no-grad temporaries are released before the student + # checkpoint graph is constructed. + with torch.no_grad(): + teacher_features = self._extract( + teacher_latent, conditional_dict, indices) + student_features = self._extract( + student_latent, conditional_dict, indices) + mmd, mmd_meta = self.latent_mmd( + student_features, teacher_features, stage=True) + normalized_mmd = ( + mmd / (mmd.detach() + self.latent_norm_eps) + if self.latent_self_normalize else mmd) + loss = self.latent_loss_weight * normalized_mmd + return loss, { + "dmdtrain_gradient_norm": loss.new_zeros(()), + "wan_latent_mmd": mmd.detach(), + "wan_latent_loss_normalized": normalized_mmd.detach(), + "wan_latent_bandwidth": mmd_meta["bandwidth"], + "wan_latent_history_updates": mmd_meta["history_updates"], + "wan_latent_current_coefficient": mmd_meta["current_coefficient"], + "wan_latent_tokens": loss.new_tensor(student_features.shape[1]), + "wan_latent_feature_dim": loss.new_tensor( + student_features.shape[2]), + "wan_latent_prefix_blocks": loss.new_tensor( + self.latent_layer + 1), + } + + def _compute_vjepa2_dense_mmd_loss( + self, + student_latent: torch.Tensor, + teacher_latent: torch.Tensor, + ) -> Tuple[torch.Tensor, dict]: + token_indices = self._vjepa2_token_indices(student_latent.device) + + # Stage A: decode and extract both sides without retaining either the + # causal VAE or the V-JEPA2 graph. Teacher is processed first and its + # large decoded video is released before the student decode. + with torch.no_grad(): + teacher_video01 = self._decode_video01(teacher_latent) + teacher_features = self._extract_vjepa2_dense( + teacher_video01, + token_indices, + gradient_checkpointing=False, + ).detach().float() + del teacher_video01 + + student_video01 = self._decode_video01(student_latent) + student_features_detached = self._extract_vjepa2_dense( + student_video01, + token_indices, + gradient_checkpointing=False, + ).detach().float() + + # Stage A.2: calculate the MMD and its cotangent only on a small dense + # feature leaf. EMARBFMMD detaches historical kernel means and teacher + # features, so only this latest student leaf receives a gradient. + with torch.enable_grad(): + student_features_leaf = ( + student_features_detached.detach().requires_grad_(True)) + mmd, mmd_meta = self.latent_mmd( + student_features_leaf, teacher_features, stage=True) + normalized_mmd = ( + mmd / (mmd.detach() + self.latent_norm_eps) + if self.latent_self_normalize else mmd) + objective = self.latent_loss_weight * normalized_mmd + (feature_cotangent,) = torch.autograd.grad( + objective, + student_features_leaf, + retain_graph=False, + create_graph=False, + ) + feature_cotangent = feature_cotangent.detach().float() + if not torch.isfinite(feature_cotangent).all(): + raise FloatingPointError( + "non-finite V-JEPA2 dense-feature MMD cotangent") + + # Stage B.1: replay only the frozen FP32 V-JEPA2 prefix. The adapter + # repeats the exact deterministic center-frame/resize/crop/normalize + # pipeline used above; non-reentrant block checkpointing bounds memory. + with torch.enable_grad(): + student_video_leaf = ( + student_video01.detach().float().requires_grad_(True)) + replay_features = self._extract_vjepa2_dense( + student_video_leaf, + token_indices, + gradient_checkpointing=( + self.vjepa2_dense_gradient_checkpointing), + ).float() + encoder_surrogate = ( + replay_features * feature_cotangent).sum() + (video_cotangent,) = torch.autograd.grad( + encoder_surrogate, + student_video_leaf, + retain_graph=False, + create_graph=False, + ) + video_cotangent = video_cotangent.detach().float() + if not torch.isfinite(video_cotangent).all(): + raise FloatingPointError( + "non-finite V-JEPA2 replay pixel cotangent") + del encoder_surrogate, student_video_leaf, replay_features, student_video01 + + # Stage B.2: replay the frozen stateful Wan VAE decoder one latent frame + # at a time. Its temporal cache is carried by value and detached between + # frames, matching the repository's validated bounded-memory replay. + latent_cotangent = replay_decoder_vjp( + self.vae, + student_latent.detach(), + video_cotangent, + independent_first_frame=self.independent_first_frame, + ).detach().float() + if latent_cotangent.shape != student_latent.shape: + raise ValueError( + "Wan VAE decoder VJP returned an invalid latent shape: " + f"got {tuple(latent_cotangent.shape)}, expected " + f"{tuple(student_latent.shape)}") + if not torch.isfinite(latent_cotangent).all(): + raise FloatingPointError( + "non-finite V-JEPA2 replay latent cotangent") + + # Preserve the true objective as the scalar value reported by the + # trainer while replacing its derivative with the replayed latent VJP. + replay_surrogate = ( + student_latent.float() * latent_cotangent).sum() + loss = replay_surrogate + ( + objective.detach() - replay_surrogate.detach()) + return loss, { + "dmdtrain_gradient_norm": latent_cotangent.norm().detach(), + self._dense_key("mmd"): mmd.detach(), + self._dense_key("loss_normalized"): normalized_mmd.detach(), + self._dense_key("bandwidth"): mmd_meta["bandwidth"], + self._dense_key("history_updates"): mmd_meta["history_updates"], + self._dense_key("current_coefficient"): ( + mmd_meta["current_coefficient"]), + "vjepa2_dense_raw_tokens": loss.new_tensor( + self.vjepa2_dense_raw_tokens), + "vjepa2_dense_sampled_tokens": loss.new_tensor( + student_features_detached.shape[1]), + self._dense_key("feature_dim"): loss.new_tensor( + student_features_detached.shape[2]), + self._dense_key("local_video_samples"): mmd_meta["local_samples"], + self._dense_key("global_video_samples"): mmd_meta["global_samples"], + self._dense_key("num_frames"): loss.new_tensor( + self.vjepa2_dense_num_frames), + self._dense_key("layer"): loss.new_tensor( + self.vjepa2_dense_layer), + self._dense_key("prefix_blocks"): loss.new_tensor( + self.vjepa2_dense_layer + 1), + self._dense_key("feature_grad_norm"): ( + feature_cotangent.norm().detach()), + self._dense_key("video_grad_norm"): ( + video_cotangent.norm().detach()), + self._dense_key("latent_grad_norm"): ( + latent_cotangent.norm().detach()), + } + + def distribution_state_pending(self) -> bool: + latent_mmd = getattr(self, "latent_mmd", None) + return bool( + getattr(self, "uses_fresh_video_gradcache", False) + and ( + getattr(latent_mmd, "attraction_ema_pending", False) + or getattr(latent_mmd, "repulsion_fifo_pending", False) + ) + ) + + def prepare_distribution_state(self): + """Validate staged Nyström state before optimizer mutation.""" + latent_mmd = getattr(self, "latent_mmd", None) + if ( + getattr(self, "uses_fresh_video_gradcache", False) + and ( + getattr(latent_mmd, "attraction_ema_enabled", False) + or getattr(latent_mmd, "repulsion_fifo_enabled", False) + ) + ): + latent_mmd.prepare_commit() + + def commit_distribution_state(self) -> None: + if not self.uses_fresh_video_gradcache: + self.latent_mmd.commit() + return + latent_mmd = getattr(self, "latent_mmd", None) + if ( + getattr(latent_mmd, "attraction_ema_enabled", False) + or getattr(latent_mmd, "repulsion_fifo_enabled", False) + ): + latent_mmd.commit() + + def discard_distribution_state(self) -> None: + if not self.uses_fresh_video_gradcache: + self.latent_mmd.discard() + return + latent_mmd = getattr(self, "latent_mmd", None) + if ( + getattr(latent_mmd, "attraction_ema_enabled", False) + or getattr(latent_mmd, "repulsion_fifo_enabled", False) + ): + latent_mmd.discard() + + def generator_loss( + self, + image_or_video_shape, + conditional_dict: dict, + unconditional_dict: dict, + clean_latent: torch.Tensor, + initial_latent: torch.Tensor = None, + ): + if self.uses_fresh_video_gradcache: + raise RuntimeError( + "nystrom_exact_repulsion generator_loss requires " + "FreshVideoGradCache orchestration in DistributionMatchingTrainer") + if clean_latent is None: + raise ValueError( + "wan_kv_mmd needs aligned teacher latents; configure " + "teacher_latent_data_path") + if self.causal_generator: + pred, gradient_mask, _, _ = self._run_generator( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + initial_latent=initial_latent, + ) + else: + pred = self._run_few_step_generator( + image_or_video_shape=image_or_video_shape, + conditional_dict=conditional_dict, + unconditional_dict=unconditional_dict, + ) + gradient_mask = None + return self.compute_latent_mmd_loss( + pred, clean_latent, conditional_dict, gradient_mask=gradient_mask) + + def _fd_metadata(self) -> dict: + if self.feature_type == "dinov3_dense_latent": + return self._dinov3_feature_metadata() + if self.feature_type in self._DENSE_VIDEO_FEATURE_TYPES: + if self.dual_video_mmd_enabled: + branches = {} + for name, registry in self.dual_branch_registry.items(): + branches[name] = { + "reference_path": registry["reference_path"], + "reference_hash": registry["reference_hash"], + "sample_shape": list(registry["sample_shape"]), + } + if "vjepa2" in branches: + branches["vjepa2"]["feature"] = ( + self._vjepa2_feature_metadata()) + if "dinov3" in branches: + branches["dinov3"]["feature"] = ( + self._triple_t5_joint_feature_metadata( + "dinov3", self._dinov3_feature_metadata())) + if "videomae_v2" in branches: + from video_fd.encoders.videomae_v2_dense_metadata import ( + canonical_videomae_v2_dense_feature_metadata, + uniform_segment_token_indices, + ) + video_indices = uniform_segment_token_indices( + self.videomae_v2_encoder.dense_raw_tokens, 256) + branches["videomae_v2"]["feature"] = ( + self._triple_t5_joint_feature_metadata( + "videomae_v2", + canonical_videomae_v2_dense_feature_metadata( + model_id=self.videomae_v2_reference_model_id, + layer_index=15, + num_tokens=256, + token_indices=video_indices, + num_segments=5, + frames_per_segment=16, + raw_tokens_per_segment=( + self.videomae_v2_encoder.dense_raw_tokens), + feature_dim=( + self.videomae_v2_encoder.dense_feature_dim), + resize_short_side=224, + crop_size=224, + skip_first_frame=True, + ))) + branch_weights = { + name: float(self.latent_mmd.branch_weights[name]) + for name in sorted(branches) + } + teacher_bank_content_hash = ( + self.triple_teacher_bank_content_hash + if self.is_triple_dense + else self.dual_teacher_bank_content_hash + ) + aggregation_divisor = float( + self.latent_mmd.aggregation_divisor) + combination = { + "operation": "sum", + "branch_weights": branch_weights, + "divide_by_branch_count": ( + aggregation_divisor == float(len(branches))), + "normalization": ( + "raw_div_stop_gradient_raw_plus_epsilon"), + "normalization_epsilon": self.latent_norm_eps, + } + if aggregation_divisor != 1.0: + combination["aggregation_divisor"] = aggregation_divisor + return { + "feature_type": self.feature_type, + "branches": branches, + "teacher_bank_content_hash": teacher_bank_content_hash, + "teacher_mixture": { + "fixed_reference_weight": 2.0, + "current_teacher_weight": 1.0, + "normalization": "divide_by_weight_sum", + }, + "combination": combination, + "partial_gradient": { + "forward_global_samples": int( + self.fresh_video_total_samples), + "gradient_global_samples": 128, + "expected_teacher_global_samples": ( + self.expected_current_teacher_samples), + "full_student_population": True, + "replay_cotangent_compensation": 2.0, + "denominators": { + "kss": 65536, + "fixed_attraction": 256, + "teacher_cross": 32768, + "teacher_self": 16384, + }, + }, + } + if getattr(self, "vjepa2_multilayer_enabled", False): + return { + "feature_type": "vjepa2_dense_multilayer_latent", + "features": [ + self._vjepa2_feature_metadata( + layer_index=layer) + for layer in self.vjepa2_dense_layers + ], + "combination": { + "operation": "sum", + "per_layer_objective": ( + "nystrom_exact_repulsion"), + "normalization": ( + "mmd_over_abs_detached_mmd_plus_epsilon" + if self.latent_self_normalize + else "none" + ), + "normalization_epsilon": self.latent_norm_eps, + "layer_weights": [ + 1.0 for _ in self.vjepa2_dense_layers + ], + }, + } + metadata = self._vjepa2_feature_metadata() + reference_sampling = getattr( + self, "vjepa2_dense_frame_sampling", "center_contiguous") + student_sampling = getattr( + self, + "vjepa2_student_frame_sampling", + reference_sampling, + ) + if student_sampling != reference_sampling: + student_start = ( + int(self.vjepa2_student_source_num_frames) + - int(self.vjepa2_dense_num_frames) + ) + metadata = dict(metadata) + metadata["student_temporal_sampling"] = { + "frame_sampling": student_sampling, + "source_num_frames": int( + self.vjepa2_student_source_num_frames), + "selected_frame_indices": [ + student_start, + student_start + int(self.vjepa2_dense_num_frames) - 1, + ], + "reference_frame_sampling": ( + reference_sampling), + } + if getattr(self, "nystrom_attraction_weight", 1.0) != 1.0: + metadata = dict(metadata) + metadata["nystrom_attraction_weight"] = ( + self.nystrom_attraction_weight) + if getattr( + self, "current_teacher_attraction_enabled", False + ): + metadata = dict(metadata) + metadata["nystrom_current_teacher_attraction"] = { + "enabled": True, + "estimator": "exact_biased_gaussian_mmd", + "fixed_weight": self.nystrom_attraction_weight, + "current_weight": ( + self.current_teacher_attraction_weight), + "normalization": "divide_by_group_weight_sum", + "expected_global_samples": ( + self.expected_current_teacher_samples), + "ema_gradient_mode": ( + "straight_through_full_current_mix" + if self.nystrom_attraction_ema_enabled + else "direct_full_current_mix"), + } + if self.current_teacher_source == "matched_manifest_video": + metadata[ + "nystrom_current_teacher_attraction" + ]["source"] = "matched_manifest_video" + # Preserve the tuple used by legacy v2 checkpoints. The new + # canonical reference metadata stores JSON-like lists instead. + if getattr(self, "mmd_estimator", "ema_rff") == "ema_rff": + metadata = dict(metadata) + metadata.pop("feature_metadata_version", None) + metadata.pop("vae_id", None) + metadata.pop("token_indices_sha256", None) + metadata.pop("kernel", None) + metadata["matrix_shape"] = tuple(metadata["matrix_shape"]) + return metadata + return { + "feature_type": self.feature_type, + "layer_index": self.latent_layer, + "prefix_blocks": self.latent_layer + 1, + "num_tokens": self.latent_num_tokens, + "feature_dim": self.latent_feature_dim, + "timestep": self.latent_timestep, + "patch_size": self.latent_patch_size, + } + + def _nystrom_attraction_ema_active(self) -> bool: + return bool( + getattr(self, "uses_fresh_video_gradcache", False) + and getattr( + getattr(self, "latent_mmd", None), + "attraction_ema_enabled", + False, + ) + ) + + def _nystrom_repulsion_fifo_active(self) -> bool: + return bool( + getattr(self, "uses_fresh_video_gradcache", False) + and getattr( + getattr(self, "latent_mmd", None), + "repulsion_fifo_enabled", + False, + ) + ) + + def _nystrom_fd_state_version(self) -> int: + if getattr(self, "triple_t5_joint_enabled", False): + base_version = 11 + elif getattr(self, "is_triple_dense", False): + base_version = 10 + elif self.dual_video_mmd_enabled: + base_version = 9 + elif getattr(self, "vjepa2_multilayer_enabled", False): + base_version = 4 + else: + base_version = 1 + if self._nystrom_repulsion_fifo_active(): + if base_version != 1 or self._nystrom_attraction_ema_active(): + raise ValueError( + "repulsion FIFO requires the single-layer fixed-reference " + "Nyström state without attraction EMA") + return 8 + if not self._nystrom_attraction_ema_active(): + return base_version + enabled_versions = {1: 5, 4: 7} + if base_version not in enabled_versions: + raise ValueError( + "attraction EMA is incompatible with this Nyström state " + f"variant v{base_version}") + return enabled_versions[base_version] + + def _with_nystrom_attraction_ema_state(self, state: dict) -> dict: + if self._nystrom_attraction_ema_active(): + state["attraction_ema"] = ( + self.latent_mmd.attraction_ema_state_dict()) + return state + + def _with_nystrom_distribution_state(self, state: dict) -> dict: + state = self._with_nystrom_attraction_ema_state(state) + if self._nystrom_repulsion_fifo_active(): + state["repulsion_fifo"] = ( + self.latent_mmd.repulsion_fifo_state_dict()) + return state + + def fd_state_dict(self) -> dict: + nystrom_version_key = self._dense_nystrom_version_key() + if getattr(self, "uses_fresh_video_gradcache", False): + if self.distribution_state_pending(): + raise RuntimeError( + "refusing to checkpoint pending distribution state") + if self._uses_dinov3_dense_backend(): + return { + "dinov3_dense_nystrom_mmd_version": 1, + "metadata": self._fd_metadata(), + "reference_path": self.dinov3_nystrom_reference_path, + "reference_hash": self.dinov3_nystrom_reference_hash, + "teacher_feature_bank_hash": ( + self.dinov3_teacher_feature_bank_hash), + "estimator": self.mmd_estimator, + } + if getattr(self, "vjepa2_multilayer_enabled", False): + return self._with_nystrom_distribution_state({ + "vjepa2_dense_nystrom_mmd_version": ( + self._nystrom_fd_state_version()), + "metadata": self._fd_metadata(), + "reference_paths": list( + self.vjepa2_nystrom_reference_paths), + "reference_hashes": list( + self.vjepa2_nystrom_reference_hashes), + "teacher_feature_bank_hashes": list( + self.vjepa2_teacher_feature_bank_hashes), + "estimator": self.mmd_estimator, + }) + if self.dual_video_mmd_enabled: + return { + self.dense_nystrom_version_key: ( + self._nystrom_fd_state_version()), + "metadata": self._fd_metadata(), + "references": { + name: { + "path": branch["reference_path"], + "hash": branch["reference_hash"], + "landmarks": branch["nystrom_landmarks"], + } + for name, branch in self.dual_branch_registry.items() + }, + "teacher_bank_content_hash": ( + self.triple_teacher_bank_content_hash + if self.is_triple_dense + else self.dual_teacher_bank_content_hash), + "estimator": self.mmd_estimator, + } + return self._with_nystrom_distribution_state({ + "vjepa2_dense_nystrom_mmd_version": ( + self._nystrom_fd_state_version()), + "metadata": self._fd_metadata(), + "reference_path": self.vjepa2_nystrom_reference_path, + "reference_hash": self.vjepa2_nystrom_reference_hash, + "teacher_feature_bank_hash": ( + self.vjepa2_teacher_feature_bank_hash), + "estimator": self.mmd_estimator, + }) + if self.feature_type == "vjepa2_dense_latent": + return { + "vjepa2_dense_mmd_version": 2, + "metadata": self._fd_metadata(), + "vjepa2_dense_mmd": self.latent_mmd.state_dict(), + } + return { + "wan_kv_mmd_version": 3, + "metadata": self._fd_metadata(), + "wan_latent_mmd": self.latent_mmd.state_dict(), + } + + def load_fd_state_dict(self, state_dict: dict) -> None: + if self._uses_dinov3_dense_backend(): + if not isinstance(state_dict, dict): + raise TypeError("DINOv3 distribution state must be a mapping") + if any(key in state_dict for key in ( + "wan_kv_mmd_version", + "vjepa2_dense_mmd_version", + "vjepa2_dense_nystrom_mmd_version", + )): + raise ValueError( + "Wan/V-JEPA2 distribution state cannot be loaded into " + "the DINOv3 dense-latent feature space") + version = state_dict.get( + "dinov3_dense_nystrom_mmd_version") + if version != 1: + raise ValueError( + "resume checkpoint has no compatible DINOv3 dense " + f"Nyström state v1 (got {version!r})") + expected = self.fd_state_dict() + if set(state_dict) != set(expected): + raise ValueError( + "DINOv3 Nyström checkpoint/reference keys mismatch: " + f"saved={sorted(state_dict)}, " + f"current={sorted(expected)}") + if state_dict != expected: + raise ValueError( + "DINOv3 Nyström checkpoint/reference identity mismatch: " + f"saved={state_dict}, current={expected}") + return + + dense_version_key = self._dense_ema_version_key() + dense_state_key = self._dense_ema_state_key() + nystrom_version_key = self._dense_nystrom_version_key() + foreign_dense_version_keys = { + "vjepa2_dense_mmd_version", + "motionformer_dense_mmd_version", + "vjepa2_dense_nystrom_mmd_version", + "motionformer_dense_nystrom_mmd_version", + "dual_dense_mmd_version", + "dual_dense_nystrom_mmd_version", + "triple_dense_nystrom_mmd_version", + } - {dense_version_key, nystrom_version_key} + if getattr(self, "uses_fresh_video_gradcache", False): + present_foreign = sorted( + foreign_dense_version_keys.intersection(state_dict)) + if present_foreign: + raise ValueError( + "checkpoint belongs to another dense feature namespace: " + f"{present_foreign}") + if "dinov3_dense_nystrom_mmd_version" in state_dict: + raise ValueError( + "DINOv3 distribution state cannot be loaded into the " + "V-JEPA2 dense-latent feature space") + if ( + "wan_kv_mmd_version" in state_dict + or dense_version_key in state_dict + ): + raise ValueError( + "EMA/RFF distribution history cannot be loaded into the " + "fresh Nyström exact-repulsion estimator") + version = state_dict.get(nystrom_version_key) + expected_version = self._nystrom_fd_state_version() + if version != expected_version: + raise ValueError( + "resume checkpoint has no compatible dense feature " + f"Nyström state v{expected_version} for " + f"{nystrom_version_key} (got {version!r})") + validated_attraction_ema = None + if self._nystrom_attraction_ema_active(): + if "attraction_ema" not in state_dict: + raise ValueError( + "resume checkpoint is missing attraction EMA state") + validated_attraction_ema = ( + self.latent_mmd.validate_attraction_ema_state_dict( + state_dict["attraction_ema"]) + ) + validated_repulsion_fifo = None + if self._nystrom_repulsion_fifo_active(): + if "repulsion_fifo" not in state_dict: + raise ValueError( + "resume checkpoint is missing repulsion FIFO state") + validated_repulsion_fifo = ( + self.latent_mmd.validate_repulsion_fifo_state_dict( + state_dict["repulsion_fifo"]) + ) + expected = self.fd_state_dict() + if set(state_dict) != set(expected): + raise ValueError( + "V-JEPA2 Nyström checkpoint/reference keys mismatch: " + f"saved={sorted(state_dict)}, " + f"current={sorted(expected)}") + identity_keys = set(expected) + identity_keys.discard("attraction_ema") + identity_keys.discard("repulsion_fifo") + saved_identity = { + key: state_dict.get(key) for key in identity_keys} + expected_identity = { + key: expected.get(key) for key in identity_keys} + # Runtime live-teacher batch size is not frozen-reference identity. + saved_identity = copy.deepcopy(saved_identity) + expected_identity = copy.deepcopy(expected_identity) + for identity in (saved_identity, expected_identity): + current_teacher = identity.get("metadata", {}).get( + "nystrom_current_teacher_attraction", {}) + current_teacher.pop("expected_global_samples", None) + saved_identity = _canonicalize_shared_asset_paths(saved_identity) + expected_identity = _canonicalize_shared_asset_paths( + expected_identity) + if saved_identity != expected_identity: + raise ValueError( + "V-JEPA2 Nyström checkpoint/reference mismatch: " + f"saved={saved_identity}, current={expected_identity}") + if validated_attraction_ema is not None: + self.latent_mmd.load_attraction_ema_state_dict( + validated_attraction_ema) + if validated_repulsion_fifo is not None: + self.latent_mmd.load_repulsion_fifo_state_dict( + validated_repulsion_fifo) + return + + if self.feature_type == "vjepa2_dense_latent": + if "dinov3_dense_nystrom_mmd_version" in state_dict: + raise ValueError( + "DINOv3 distribution state cannot be loaded into the " + "V-JEPA2 dense-latent feature space") + if "wan_kv_mmd_version" in state_dict: + raise ValueError( + "Wan post-block latent MMD history cannot be loaded into " + "the V-JEPA2 dense-latent feature space") + version = state_dict.get("vjepa2_dense_mmd_version") + if version != 2: + detail = ( + " V-JEPA2 dense MMD v1 treated individual tokens as " + "distribution samples and is intentionally incompatible." + if version == 1 else "" + ) + raise ValueError( + "resume checkpoint has no compatible V-JEPA2 dense " + f"latent MMD v2 state (got version {version!r})." + f"{detail}") + if "vjepa2_dense_mmd" not in state_dict: + raise ValueError( + "resume checkpoint is missing fd.vjepa2_dense_mmd history") + expected = self._fd_metadata() + saved = state_dict.get("metadata") + if saved != expected: + raise ValueError( + "V-JEPA2 dense latent checkpoint/config mismatch: " + f"saved metadata={saved}, current={expected}") + self.latent_mmd.load_state_dict( + state_dict["vjepa2_dense_mmd"], strict=True) + return + + if "dinov3_dense_nystrom_mmd_version" in state_dict: + raise ValueError( + "DINOv3 distribution state cannot be loaded into the Wan " + "post-block latent feature space") + version = state_dict.get("wan_kv_mmd_version") + if version != 3: + detail = ( + " Wan KV MMD v2 history is intentionally incompatible." + if version == 2 else "") + raise ValueError( + "resume checkpoint has no compatible Wan post-block latent " + f"MMD v3 state.{detail}") + if "wan_latent_mmd" not in state_dict: + raise ValueError( + "resume checkpoint is missing fd.wan_latent_mmd history") + expected = self._fd_metadata() + saved = state_dict.get("metadata") + if saved != expected: + raise ValueError( + "Wan post-block latent checkpoint/config mismatch: " + f"saved metadata={saved}, current={expected}. " + "KV EMA checkpoints are intentionally incompatible.") + self.latent_mmd.load_state_dict( + state_dict["wan_latent_mmd"], strict=True) + + def can_load_initialization_fd_state(self, state_dict: dict) -> bool: + """Whether an initialization-only checkpoint carries our estimator.""" + if not isinstance(state_dict, dict): + return False + if self._uses_dinov3_dense_backend(): + return state_dict.get( + "dinov3_dense_nystrom_mmd_version") == 1 + if getattr(self, "uses_fresh_video_gradcache", False): + return ( + state_dict.get(self._dense_nystrom_version_key()) + == self._nystrom_fd_state_version() + ) + if self.feature_type == "vjepa2_dense_latent": + return state_dict.get("vjepa2_dense_mmd_version") == 2 + return state_dict.get("wan_kv_mmd_version") == 3 + + # Historical entry points remain aliases for external callers. + fresh_video_gradcache_prepare_replay = prepare_latent_gradients + fresh_video_gradcache_replay = replay_generator_gradients + fresh_video_gradcache_forward = collect_rollout_features + + _initialize_triple_dinov3_vjepa2_videomae_v2_features = _initialize_multi_encoder_features + _replace_with_triple_nystrom_objective = _build_multi_encoder_objective + _extract_dual_video_features = _extract_multi_encoder_features + _fresh_video_dual_gradcache_forward = _collect_multi_encoder_rollout_features + _dual_gradcache_prepare_replay = _prepare_multi_encoder_latent_gradients + +# Old import/class spelling remains valid for existing callers. +WanKVMMD = DistributionMatchingModel