| """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"} |
| |
| 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: |
| |
| |
| 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]) |
| |
| |
| 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) |
| |
| |
| 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" |
| ) |
| |
| |
| 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: |
| |
| |
| 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)): |
| |
| |
| |
| 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) |
| |
| |
| if bool(getattr(args, "wan_latent_gradient_checkpointing", True)): |
| self.real_score.enable_gradient_checkpointing() |
|
|
| def _initialize_vjepa2_features(self, args, device) -> None: |
| |
| |
| 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) |
| |
| |
| 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 |
| |
| |
| 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") |
| |
| |
| 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}") |
|
|
| |
| |
| |
| |
| 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 |
| |
| |
| |
| |
| 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") |
|
|
| |
| 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") |
| |
| 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") |
|
|
| |
| 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") |
|
|
| |
| 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 |
| |
| 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, |
| |
| |
| "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)}") |
|
|
| |
| |
| 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}") |
| |
| 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" |
| } |
| ): |
| |
| |
| 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) |
| |
| |
| 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") |
| |
| |
| 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()) |
| |
| |
| |
| 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) |
| |
| |
| |
| 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 |
| ) |
| |
| |
| |
| 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: |
| |
| |
| 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: |
| |
| |
| |
| |
| 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", |
| } |
| |
| |
| 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: |
| |
| |
| |
| 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) |
| ) |
|
|
| |
| |
| 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") |
|
|
| |
| |
| |
| |
| 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() |
| |
| |
| |
| |
| |
| |
| |
| |
| 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 |
|
|
| |
| |
| 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) |
| |
| |
| 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) |
|
|
| |
| |
| |
| 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() |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| 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" |
| |
| |
| 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} |
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| WanKVMMD = DistributionMatchingModel |
|
|
