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