# Author: Rui Heng Yang from typing import Any, Optional, Sequence, Union import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from fastwam.utils.logging_config import get_logger from .action_dit import ActionDiT from .helpers.loader import load_wan22_ti2v_5b_components from .mot import MoT from .schedulers.scheduler_continuous import WanContinuousFlowMatchScheduler logger = get_logger(__name__) def _infer_new_fused_kv_projection_mode(mot_state: dict[str, torch.Tensor]) -> Optional[str]: """Infer the new_fused_kv projection mode from its unique state keys.""" detected_modes = [] if "k_video_pos_projection" in mot_state: detected_modes.append("simple+PE") elif "simple_kv_fusing_layer" in mot_state: detected_modes.append("simple") if any( key in mot_state for key in ( "per_head_kv_fusing_layer", "k_head_channel_projection", "k_head_channel_bias", ) ): detected_modes.append("per_head_channel") if any( key in mot_state for key in ( "kv_fusing_layer", "k_channel_projection", "v_channel_projection", "k_channel_bias", "v_channel_bias", ) ): detected_modes.append("full") has_head_fused_kv_channel_projection = any( key in mot_state for key in ( "head_fused_kv_k_channel_projection", "head_fused_kv_v_channel_projection", ) ) has_head_fused_kv_layer_mixing = "head_fused_kv_layer_mixing" in mot_state has_head_fused_kv_sin2d = any( key.startswith("head_fused_kv_sin2d_pe_mlps.") for key in mot_state ) if has_head_fused_kv_sin2d: detected_modes.append("HeadFusedKV+Sin2DPE") elif has_head_fused_kv_channel_projection or has_head_fused_kv_layer_mixing: if has_head_fused_kv_channel_projection: detected_modes.append("HeadFusedKV") else: detected_modes.append("simple_head_fused") if any(key.startswith("mlp_mixer_fused_kv_blocks.") for key in mot_state): detected_modes.append("MLPMixerFusedKV") if len(detected_modes) > 1: raise RuntimeError( "Checkpoint contains parameters from multiple new_fused_kv projection " f"modes: {detected_modes}. Refusing to infer an ambiguous architecture." ) return detected_modes[0] if detected_modes else None def _new_fused_kv_projection_signature_matches( metadata_mode: str, inferred_mode: str, ) -> bool: if metadata_mode == inferred_mode: return True shared_signature_modes = {"simple_head_fused", "simple_head_softmax"} return {metadata_mode, inferred_mode} <= shared_signature_modes class FastWAM(torch.nn.Module): """MoT world model with video/action experts.""" def __init__( self, video_expert, action_expert: ActionDiT, mot: MoT, vae, text_encoder=None, tokenizer=None, text_dim: Optional[int] = None, proprio_dim: Optional[int] = None, device: str = "cpu", torch_dtype: torch.dtype = torch.float32, video_train_shift: float = 5.0, video_infer_shift: float = 5.0, video_num_train_timesteps: int = 1000, action_train_shift: float = 5.0, action_infer_shift: float = 5.0, action_num_train_timesteps: int = 1000, loss_lambda_video: float = 1.0, loss_lambda_action: float = 1.0, freeze_video_backbone: bool = False, ): super().__init__() self.video_expert = video_expert self.action_expert = action_expert self.mot = mot # Keep trainer compatibility: optimizer and freeze logic use `model.dit`. self.dit = self.mot self.freeze_video_backbone = freeze_video_backbone self.vae = vae self.text_encoder = text_encoder self.tokenizer = tokenizer if text_dim is None: if self.text_encoder is None: raise ValueError("`text_dim` is required when `text_encoder` is not loaded.") text_dim = int(self.text_encoder.dim) self.text_dim = int(text_dim) self.proprio_dim = None if proprio_dim is None else int(proprio_dim) if self.proprio_dim is not None: self.proprio_encoder = nn.Linear(self.proprio_dim, self.text_dim).to(torch_dtype) else: self.proprio_encoder = None self.train_video_scheduler = WanContinuousFlowMatchScheduler( num_train_timesteps=video_num_train_timesteps, shift=video_train_shift, ) self.infer_video_scheduler = WanContinuousFlowMatchScheduler( num_train_timesteps=video_num_train_timesteps, shift=video_infer_shift, ) self.train_action_scheduler = WanContinuousFlowMatchScheduler( num_train_timesteps=action_num_train_timesteps, shift=action_train_shift, ) self.infer_action_scheduler = WanContinuousFlowMatchScheduler( num_train_timesteps=action_num_train_timesteps, shift=action_infer_shift, ) # Optional aliases for consistency with Wan22Core naming. self.train_scheduler = self.train_video_scheduler self.infer_scheduler = self.infer_video_scheduler self.device = torch.device(device) self.torch_dtype = torch_dtype self.loss_lambda_video = float(loss_lambda_video) self.loss_lambda_action = float(loss_lambda_action) self.to(self.device) @classmethod def from_wan22_pretrained( cls, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, model_id: str = "Wan-AI/Wan2.2-TI2V-5B", tokenizer_model_id: str = "Wan-AI/Wan2.1-T2V-1.3B", tokenizer_max_len: int = 512, load_text_encoder: bool = True, proprio_dim: Optional[int] = None, redirect_common_files: bool = True, video_dit_config: dict[str, Any] | None = None, action_dit_config: dict[str, Any] | None = None, action_dit_pretrained_path: str | None = None, skip_dit_load_from_pretrain: bool = False, freeze_video_backbone: bool = False, mot_checkpoint_mixed_attn: bool = True, video_train_shift: float = 5.0, video_infer_shift: float = 5.0, video_num_train_timesteps: int = 1000, action_train_shift: float = 5.0, action_infer_shift: float = 5.0, action_num_train_timesteps: int = 1000, loss_lambda_video: float = 1.0, loss_lambda_action: float = 1.0, decoupled: bool = False, kv_source_mapping: list[int] | None = None, kv_source_mode: str = "final_only", kv_fusion: "nn.Module | None" = None, ): if video_dit_config is None: raise ValueError("`video_dit_config` is required for FastWAM.from_wan22_pretrained().") if "text_dim" not in video_dit_config: raise ValueError("`video_dit_config['text_dim']` is required for FastWAM.") components = load_wan22_ti2v_5b_components( device=device, torch_dtype=torch_dtype, model_id=model_id, tokenizer_model_id=tokenizer_model_id, tokenizer_max_len=tokenizer_max_len, redirect_common_files=redirect_common_files, dit_config=video_dit_config, skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, load_text_encoder=load_text_encoder, ) video_expert = components.dit layer_init_mapping = kv_source_mapping if decoupled and kv_source_mode in {"fused_kv", "new_fused_kv"}: from .mot_decoupled import compute_kv_source_mapping action_config_for_init = action_dit_config or {} action_num_layers = int(action_config_for_init.get("num_layers", 5)) if kv_source_mapping is not None and len(kv_source_mapping) == action_num_layers: layer_init_mapping = kv_source_mapping else: layer_init_mapping = compute_kv_source_mapping( mode="uniform_end", video_num_layers=len(video_expert.blocks), action_num_layers=action_num_layers, ) action_expert = ActionDiT.from_pretrained( action_dit_config=action_dit_config, action_dit_pretrained_path=action_dit_pretrained_path, skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, device=device, torch_dtype=torch_dtype, layer_init_mapping=layer_init_mapping if decoupled else None, ) if int(action_expert.num_heads) != int(video_expert.num_heads): raise ValueError("ActionDiT `num_heads` must match video expert for MoT mixed attention.") if int(action_expert.attn_head_dim) != int(video_expert.attn_head_dim): raise ValueError("ActionDiT `attn_head_dim` must match video expert for MoT mixed attention.") if decoupled: # Decoupled MoT: asymmetric layer counts (e.g., video=30, action=5). # Skip the equal-layer-count assertion since experts have different depths. from .mot_decoupled import MoTDecoupled mot = MoTDecoupled( mixtures={"video": video_expert, "action": action_expert}, video_num_layers=len(video_expert.blocks), action_num_layers=len(action_expert.blocks), num_heads=int(video_expert.num_heads), attn_head_dim=int(video_expert.attn_head_dim), mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, kv_source_mapping=kv_source_mapping, kv_source_mode=kv_source_mode, kv_fusion=kv_fusion, ) else: # Standard MoT: both experts must have the same number of layers. if int(len(action_expert.blocks)) != int(len(video_expert.blocks)): raise ValueError("ActionDiT `num_layers` must match video expert.") mot = MoT( mixtures={"video": video_expert, "action": action_expert}, mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, ) model = cls( video_expert=video_expert, action_expert=action_expert, mot=mot, vae=components.vae, text_encoder=components.text_encoder, tokenizer=components.tokenizer, text_dim=int(video_dit_config["text_dim"]), proprio_dim=proprio_dim, device=device, torch_dtype=torch_dtype, video_train_shift=video_train_shift, video_infer_shift=video_infer_shift, video_num_train_timesteps=video_num_train_timesteps, action_train_shift=action_train_shift, action_infer_shift=action_infer_shift, action_num_train_timesteps=action_num_train_timesteps, loss_lambda_video=loss_lambda_video, loss_lambda_action=loss_lambda_action, freeze_video_backbone=freeze_video_backbone, ) model.model_paths = { "video_dit": components.dit_path, "vae": components.vae_path, "text_encoder": components.text_encoder_path, "tokenizer": components.tokenizer_path, "action_dit_backbone": ( "SKIPPED_PRETRAIN" if skip_dit_load_from_pretrain else action_dit_pretrained_path ), } return model def to(self, *args, **kwargs): super().to(*args, **kwargs) self.mot.to(*args, **kwargs) if self.text_encoder is not None: self.text_encoder.to(*args, **kwargs) self.vae.to(*args, **kwargs) return self @staticmethod def _check_resize_height_width(height, width, num_frames): if height % 16 != 0: height = (height + 15) // 16 * 16 if width % 16 != 0: width = (width + 15) // 16 * 16 if num_frames % 4 != 1: num_frames = (num_frames + 3) // 4 * 4 + 1 return height, width, num_frames @torch.no_grad() def encode_prompt(self, prompt: Union[str, Sequence[str]]): if self.text_encoder is None or self.tokenizer is None: raise ValueError( "Prompt encoding requires loaded text encoder/tokenizer. " "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." ) ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) ids = ids.to(self.device) mask = mask.to(self.device, dtype=torch.bool) prompt_emb = self.text_encoder(ids, mask) # FIXME: original implementation's zero padding is visible in cross-attn. seq_lens = mask.gt(0).sum(dim=1).long() for i, v in enumerate(seq_lens): prompt_emb[i, v:] = 0 mask = torch.ones_like(mask) return prompt_emb.to(device=self.device), mask def _append_proprio_to_context( self, context: torch.Tensor, context_mask: torch.Tensor, proprio: Optional[torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: if self.proprio_encoder is None or proprio is None: return context, context_mask if proprio.ndim != 2: raise ValueError(f"`proprio` must be 2D [B, D], got shape {tuple(proprio.shape)}") if self.proprio_dim is None or proprio.shape[1] != self.proprio_dim: raise ValueError( f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}" ) proprio_token = self.proprio_encoder( proprio.to(device=self.device, dtype=context.dtype).unsqueeze(1) ).to(dtype=context.dtype) # [B, 1, D] proprio_mask = torch.ones((context_mask.shape[0], 1), dtype=torch.bool, device=context_mask.device) return ( torch.cat([context, proprio_token], dim=1), torch.cat([context_mask, proprio_mask], dim=1), ) @torch.no_grad() def _encode_video_latents(self, video_tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): z = self.vae.encode( video_tensor, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride, ) return z @torch.no_grad() def _encode_input_image_latents_tensor(self, input_image: torch.Tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): if input_image.ndim == 3: input_image = input_image.unsqueeze(0) if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: raise ValueError( f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" ) image = input_image.to(device=self.device)[0].unsqueeze(1) z = self.vae.encode([image], device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) if isinstance(z, list): z = z[0].unsqueeze(0) return z def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): video_tensor = self.vae.decode(latents, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) video_tensor = video_tensor.squeeze(0).detach().float().clamp(-1, 1) video_tensor = ((video_tensor + 1.0) * 127.5).to(torch.uint8).cpu() frames = [] for t in range(video_tensor.shape[1]): frame = video_tensor[:, t].permute(1, 2, 0).numpy() frames.append(Image.fromarray(frame)) return frames def build_inputs(self, sample, tiled: bool = False): video = sample.get("video") video_latents = sample.get("video_latents") if (video is None) == (video_latents is None): raise ValueError( "FastWAM training requires exactly one of `sample['video']` or " "`sample['video_latents']`." ) if "context" not in sample or "context_mask" not in sample: raise ValueError( "FastWAM training requires `sample['context']` and `sample['context_mask']`." ) context = sample["context"] context_mask = sample["context_mask"] proprio = sample.get("proprio", None) if video is not None: if video.ndim != 5: raise ValueError(f"`sample['video']` must be 5D [B, 3, T, H, W], got shape {tuple(video.shape)}") if video.shape[1] != 3: raise ValueError(f"`sample['video']` channel dimension must be 3, got shape {tuple(video.shape)}") batch_size, _, num_frames, height, width = video.shape else: if video_latents.ndim != 5: raise ValueError( "`sample['video_latents']` must be 5D [B, C, T, H, W], " f"got shape {tuple(video_latents.shape)}" ) source_video_shape = sample.get("source_video_shape") if source_video_shape is None: raise ValueError( "`sample['source_video_shape']` is required with cached video latents." ) source_video_shape = torch.as_tensor(source_video_shape) if source_video_shape.ndim == 1: source_video_shape = source_video_shape.unsqueeze(0) if source_video_shape.ndim != 2 or source_video_shape.shape[1] != 4: raise ValueError( "`sample['source_video_shape']` must be [B,4] containing [C,T,H,W], " f"got {tuple(source_video_shape.shape)}" ) batch_size = int(video_latents.shape[0]) if source_video_shape.shape[0] != batch_size: raise ValueError( "Cached latent/source shape batch mismatch: " f"{batch_size} vs {source_video_shape.shape[0]}" ) if not torch.equal(source_video_shape, source_video_shape[:1].expand_as(source_video_shape)): raise ValueError("All cached samples in a batch must have the same source video shape.") channels, num_frames, height, width = ( int(value) for value in source_video_shape[0].tolist() ) if channels != 3: raise ValueError(f"Cached source video channel count must be 3, got {channels}") if height % 16 != 0 or width % 16 != 0: raise ValueError( f"Video spatial dims must be multiples of 16, got H={height}, W={width}" ) if num_frames % 4 != 1: raise ValueError(f"Video T must satisfy T % 4 == 1, got T={num_frames}") if num_frames <= 1: raise ValueError(f"Video T must be > 1 for action-conditioned training, got T={num_frames}") if "action" not in sample: raise ValueError("`sample['action']` is required for FastWAM training.") action = sample["action"] if action.ndim != 3: raise ValueError(f"`sample['action']` must be 3D [B, T, a_dim], got shape {tuple(action.shape)}") action_horizon = int(action.shape[1]) if action_horizon % (num_frames - 1) != 0: raise ValueError( f"`sample['action']` temporal dimension must be divisible by video transitions ({num_frames - 1}), got {action_horizon}" ) action_is_pad = sample.get("action_is_pad", None) if action_is_pad is not None: if action_is_pad.ndim != 2: raise ValueError( f"`sample['action_is_pad']` must be 2D [B, T], got shape {tuple(action_is_pad.shape)}" ) if action_is_pad.shape[0] != batch_size or action_is_pad.shape[1] != action_horizon: raise ValueError( "`sample['action_is_pad']` shape mismatch: " f"got {tuple(action_is_pad.shape)} vs expected ({batch_size}, {action_horizon})" ) image_is_pad = sample.get("image_is_pad", None) if image_is_pad is not None: if image_is_pad.ndim != 2: raise ValueError( f"`sample['image_is_pad']` must be 2D [B, T], got shape {tuple(image_is_pad.shape)}" ) if image_is_pad.shape[0] != batch_size or image_is_pad.shape[1] != num_frames: raise ValueError( "`sample['image_is_pad']` shape mismatch: " f"got {tuple(image_is_pad.shape)} vs expected ({batch_size}, {num_frames})" ) if video is not None: input_video = video.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) input_latents = self._encode_video_latents(input_video, tiled=tiled) else: temporal_factor = int(self.vae.temporal_downsample_factor) spatial_factor = int(self.vae.upsampling_factor) expected_shape = ( batch_size, int(self.vae.model.z_dim), (num_frames - 1) // temporal_factor + 1, height // spatial_factor, width // spatial_factor, ) if tuple(video_latents.shape) != expected_shape: raise ValueError( "Cached video latent shape mismatch: " f"got {tuple(video_latents.shape)} vs expected {expected_shape}" ) input_latents = video_latents.to( device=self.device, dtype=self.torch_dtype, non_blocking=True ) first_frame_latents = None fuse_flag = False if getattr(self.video_expert, "fuse_vae_embedding_in_latents", False): first_frame_latents = input_latents[:, :, 0:1] fuse_flag = True if context.ndim != 3 or context_mask.ndim != 2: raise ValueError( f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" ) context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) if self.proprio_encoder is not None: if proprio is None: raise ValueError("`sample['proprio']` is required when `proprio_dim` is enabled.") if proprio.ndim != 3: raise ValueError(f"`sample['proprio']` must be 3D [B, T, d], got shape {tuple(proprio.shape)}") if proprio.shape[2] != self.proprio_dim: raise ValueError( f"`sample['proprio']` last dim must be {self.proprio_dim}, got {proprio.shape[2]}" ) proprio = proprio[:, 0, :] # [B, D] context, context_mask = self._append_proprio_to_context( context=context, context_mask=context_mask, proprio=proprio.to(device=self.device, dtype=self.torch_dtype), ) action = action.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) if action_is_pad is not None: action_is_pad = action_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) if image_is_pad is not None: image_is_pad = image_is_pad.to(device=self.device, dtype=torch.bool, non_blocking=True) return { "context": context, "context_mask": context_mask, "input_latents": input_latents, "first_frame_latents": first_frame_latents, "fuse_vae_embedding_in_latents": fuse_flag, "action": action, "action_is_pad": action_is_pad, "image_is_pad": image_is_pad, } @torch.no_grad() def _build_mot_attention_mask( self, video_seq_len: int, action_seq_len: int, video_tokens_per_frame: int, device: torch.device, ) -> torch.Tensor: total_seq_len = video_seq_len + action_seq_len mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device) # video -> video mask[:video_seq_len, :video_seq_len] = self.video_expert.build_video_to_video_mask( video_seq_len=video_seq_len, video_tokens_per_frame=video_tokens_per_frame, device=device, ) # action -> action mask[video_seq_len:, video_seq_len:] = True # action -> first-frame video only first_frame_tokens = min(video_tokens_per_frame, video_seq_len) mask[video_seq_len:, :first_frame_tokens] = True return mask def _compute_video_loss_per_sample( self, pred_video: torch.Tensor, target_video: torch.Tensor, image_is_pad: Optional[torch.Tensor], include_initial_video_step: bool, ) -> torch.Tensor: video_loss_token = F.mse_loss(pred_video.float(), target_video.float(), reduction="none").mean(dim=(1, 3, 4)) if image_is_pad is None: return video_loss_token.mean(dim=1) temporal_factor = int(self.vae.temporal_downsample_factor) if temporal_factor <= 0: raise ValueError(f"`vae.temporal_downsample_factor` must be positive, got {temporal_factor}.") if image_is_pad.shape[1] < 1: raise ValueError("`image_is_pad` must contain at least one frame.") if (image_is_pad.shape[1] - 1) % temporal_factor != 0: raise ValueError( "Cannot align `image_is_pad` with video latent steps: " f"num_frames={image_is_pad.shape[1]}, temporal_downsample_factor={temporal_factor}." ) tail_is_pad = image_is_pad[:, 1:] latent_tail_is_pad = tail_is_pad.view(image_is_pad.shape[0], -1, temporal_factor).all(dim=2) if include_initial_video_step: video_is_pad = torch.cat([image_is_pad[:, :1], latent_tail_is_pad], dim=1) else: video_is_pad = latent_tail_is_pad if video_is_pad.shape[1] != video_loss_token.shape[1]: raise ValueError( "Video-loss mask shape mismatch: " f"mask steps={video_is_pad.shape[1]}, loss steps={video_loss_token.shape[1]}." ) valid = (~video_is_pad).to(device=video_loss_token.device, dtype=video_loss_token.dtype) valid_sum = valid.sum(dim=1).clamp(min=1.0) return (video_loss_token * valid).sum(dim=1) / valid_sum def training_loss(self, sample, tiled: bool = False): inputs = self.build_inputs(sample, tiled=tiled) input_latents = inputs["input_latents"] batch_size = input_latents.shape[0] context = inputs["context"] context_mask = inputs["context_mask"] action = inputs["action"] action_is_pad = inputs["action_is_pad"] image_is_pad = inputs["image_is_pad"] noise_video = torch.randn_like(input_latents) timestep_video = self.train_video_scheduler.sample_training_t( batch_size=batch_size, device=self.device, dtype=input_latents.dtype, ) latents = self.train_video_scheduler.add_noise(input_latents, noise_video, timestep_video) target_video = self.train_video_scheduler.training_target(input_latents, noise_video, timestep_video) if inputs["first_frame_latents"] is not None: latents[:, :, 0:1] = inputs["first_frame_latents"] noise_action = torch.randn_like(action) timestep_action = self.train_action_scheduler.sample_training_t( batch_size=batch_size, device=self.device, dtype=action.dtype, ) noisy_action = self.train_action_scheduler.add_noise(action, noise_action, timestep_action) target_action = self.train_action_scheduler.training_target(action, noise_action, timestep_action) video_pre = self.video_expert.pre_dit( x=latents, timestep=timestep_video, context=context, context_mask=context_mask, action=action, fuse_vae_embedding_in_latents=inputs["fuse_vae_embedding_in_latents"], ) action_pre = self.action_expert.pre_dit( action_tokens=noisy_action, timestep=timestep_action, context=context, context_mask=context_mask, ) video_tokens = video_pre["tokens"] action_tokens = action_pre["tokens"] attention_mask = self._build_mot_attention_mask( video_seq_len=video_tokens.shape[1], action_seq_len=action_tokens.shape[1], video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), device=video_tokens.device, ) tokens_out = self.mot( embeds_all={ "video": video_tokens, "action": action_tokens, }, attention_mask=attention_mask, freqs_all={ "video": video_pre["freqs"], "action": action_pre["freqs"], }, context_all={ "video": { "context": video_pre["context"], "mask": video_pre["context_mask"], }, "action": { "context": action_pre["context"], "mask": action_pre["context_mask"], } if action_pre["context"] is not None else None, }, t_mod_all={ "video": video_pre["t_mod"], "action": action_pre["t_mod"], }, ) pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) include_initial_video_step = inputs["first_frame_latents"] is None if inputs["first_frame_latents"] is not None: pred_video = pred_video[:, :, 1:] target_video = target_video[:, :, 1:] loss_video_per_sample = self._compute_video_loss_per_sample( pred_video=pred_video, target_video=target_video, image_is_pad=image_is_pad, include_initial_video_step=include_initial_video_step, ) video_weight = self.train_video_scheduler.training_weight(timestep_video).to( loss_video_per_sample.device, dtype=loss_video_per_sample.dtype ) loss_video = (loss_video_per_sample * video_weight).mean() action_loss_token = F.mse_loss(pred_action.float(), target_action.float(), reduction="none").mean(dim=2) # [B, T] if action_is_pad is not None: valid = (~action_is_pad).to(device=action_loss_token.device, dtype=action_loss_token.dtype) valid_sum = valid.sum(dim=1).clamp(min=1.0) action_loss_per_sample = (action_loss_token * valid).sum(dim=1) / valid_sum else: action_loss_per_sample = action_loss_token.mean(dim=1) action_weight = self.train_action_scheduler.training_weight(timestep_action).to( action_loss_per_sample.device, dtype=action_loss_per_sample.dtype ) loss_action = (action_loss_per_sample * action_weight).mean() loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action loss_dict = { "loss_video": self.loss_lambda_video * float(loss_video.detach().item()), "loss_action": self.loss_lambda_action * float(loss_action.detach().item()), } return loss_total, loss_dict @torch.no_grad() def _predict_joint_noise( self, latents_video: torch.Tensor, latents_action: torch.Tensor, timestep_video: torch.Tensor, timestep_action: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, fuse_vae_embedding_in_latents: bool, gt_action: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: video_pre = self.video_expert.pre_dit( x=latents_video, timestep=timestep_video, context=context, context_mask=context_mask, action=gt_action, fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, ) action_pre = self.action_expert.pre_dit( action_tokens=latents_action, timestep=timestep_action, context=context, context_mask=context_mask, ) attention_mask = self._build_mot_attention_mask( video_seq_len=video_pre["tokens"].shape[1], action_seq_len=action_pre["tokens"].shape[1], video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), device=video_pre["tokens"].device, ) tokens_out = self.mot( embeds_all={ "video": video_pre["tokens"], "action": action_pre["tokens"], }, attention_mask=attention_mask, freqs_all={ "video": video_pre["freqs"], "action": action_pre["freqs"], }, context_all={ "video": { "context": video_pre["context"], "mask": video_pre["context_mask"], }, "action": { "context": action_pre["context"], "mask": action_pre["context_mask"], } if action_pre["context"] is not None else None, }, t_mod_all={ "video": video_pre["t_mod"], "action": action_pre["t_mod"], }, ) pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) return pred_video, pred_action @torch.no_grad() def _predict_action_noise( self, first_frame_latents: torch.Tensor, latents_action: torch.Tensor, timestep_action: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, fuse_vae_embedding_in_latents: bool, ) -> torch.Tensor: timestep_video = torch.zeros_like(timestep_action, dtype=first_frame_latents.dtype, device=self.device) video_pre = self.video_expert.pre_dit( x=first_frame_latents, timestep=timestep_video, context=context, context_mask=context_mask, action=None, fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, ) action_pre = self.action_expert.pre_dit( action_tokens=latents_action, timestep=timestep_action, context=context, context_mask=context_mask, ) attention_mask = self._build_mot_attention_mask( video_seq_len=video_pre["tokens"].shape[1], action_seq_len=action_pre["tokens"].shape[1], video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), device=video_pre["tokens"].device, ) tokens_out = self.mot( embeds_all={ "video": video_pre["tokens"], "action": action_pre["tokens"], }, attention_mask=attention_mask, freqs_all={ "video": video_pre["freqs"], "action": action_pre["freqs"], }, context_all={ "video": { "context": video_pre["context"], "mask": video_pre["context_mask"], }, "action": { "context": action_pre["context"], "mask": action_pre["context_mask"], } if action_pre["context"] is not None else None, }, t_mod_all={ "video": video_pre["t_mod"], "action": action_pre["t_mod"], }, ) pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre) return pred_action @torch.no_grad() def _predict_action_noise_with_cache( self, latents_action: torch.Tensor, timestep_action: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, video_kv_cache: list[dict[str, torch.Tensor]], attention_mask: torch.Tensor, video_seq_len: int, ) -> torch.Tensor: action_pre = self.action_expert.pre_dit( action_tokens=latents_action, timestep=timestep_action, context=context, context_mask=context_mask, ) action_tokens = self.mot.forward_action_with_video_cache( action_tokens=action_pre["tokens"], action_freqs=action_pre["freqs"], action_t_mod=action_pre["t_mod"], action_context_payload={ "context": action_pre["context"], "mask": action_pre["context_mask"], } if action_pre["context"] is not None else None, video_kv_cache=video_kv_cache, attention_mask=attention_mask, video_seq_len=video_seq_len, ) return self.action_expert.post_dit(action_tokens, action_pre) @torch.no_grad() def infer_joint( self, prompt: Optional[str], input_image: torch.Tensor, num_video_frames: int, action_horizon: int, action: Optional[torch.Tensor] = None, # NOTE: this is gt action for conditioning videos, not for action expert proprio: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, negative_prompt: Optional[str] = None, text_cfg_scale: float = 1.0, num_inference_steps: int = 20, sigma_shift: Optional[float] = None, seed: Optional[int] = None, rand_device: str = "cpu", tiled: bool = False, test_action_with_infer_action: bool = True, ) -> dict[str, Any]: self.eval() if test_action_with_infer_action: if seed is None: raise ValueError("`test_action_with_infer_action=True` requires non-null `seed`.") action_only_out = self.infer_action( prompt=prompt, input_image=input_image.clone(), action_horizon=action_horizon, context=context.clone() if context is not None else None, context_mask=context_mask.clone() if context_mask is not None else None, num_inference_steps=num_inference_steps, sigma_shift=sigma_shift, seed=seed, rand_device=rand_device, tiled=tiled, proprio=proprio.clone() if proprio is not None else None, )["action"] if input_image.ndim == 3: input_image = input_image.unsqueeze(0) if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: raise ValueError( f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" ) _, _, height, width = input_image.shape checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) if (checked_h, checked_w) != (height, width): raise ValueError( f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" ) if checked_t != num_video_frames: raise ValueError( f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}" ) if action is not None: if action.ndim == 2: action = action.unsqueeze(0) if action.ndim != 3 or action.shape[0] != 1 or action.shape[1] != action_horizon: # NOTE: This enforces action condition to have the same shape as action horizon to predict, which may be unnecessary raise ValueError( f"`action` must have shape [1, T, a_dim] or [T, a_dim], got {tuple(action.shape)} with action_horizon={action_horizon}" ) action = action.to(device=self.device, dtype=self.torch_dtype) if proprio is not None: if self.proprio_dim is None: raise ValueError("`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled.") if proprio.ndim == 1: proprio = proprio.unsqueeze(0) elif proprio.ndim == 2 and proprio.shape[0] == 1: pass else: raise ValueError(f"`proprio` must be [D] or [1,D], got shape {tuple(proprio.shape)}") if proprio.shape[1] != self.proprio_dim: raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") proprio = proprio.to(device=self.device, dtype=self.torch_dtype) latent_t = (num_video_frames - 1) // self.vae.temporal_downsample_factor + 1 latent_h = height // self.vae.upsampling_factor latent_w = width // self.vae.upsampling_factor video_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) action_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) latents_video = torch.randn( (1, self.vae.model.z_dim, latent_t, latent_h, latent_w), generator=video_generator, device=rand_device, dtype=torch.float32, ).to(device=self.device, dtype=self.torch_dtype) latents_action = torch.randn( (1, action_horizon, self.action_expert.action_dim), generator=action_generator, device=rand_device, dtype=torch.float32, ).to(device=self.device, dtype=self.torch_dtype) input_image = input_image.to(device=self.device, dtype=self.torch_dtype) first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) latents_video[:, :, 0:1] = first_frame_latents.clone() fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) use_prompt = prompt is not None use_context = context is not None or context_mask is not None if use_prompt and use_context: raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") if not use_prompt and not use_context: raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") if use_prompt: context, context_mask = self.encode_prompt(prompt) else: if context is None or context_mask is None: raise ValueError("`context` and `context_mask` must be both provided together.") if context.ndim == 2: context = context.unsqueeze(0) if context_mask.ndim == 1: context_mask = context_mask.unsqueeze(0) if context.ndim != 3 or context_mask.ndim != 2: raise ValueError( f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" ) context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) if proprio is not None: context, context_mask = self._append_proprio_to_context( context=context, context_mask=context_mask, proprio=proprio, ) infer_timesteps_video, infer_deltas_video = self.infer_video_scheduler.build_inference_schedule( num_inference_steps=num_inference_steps, device=self.device, dtype=latents_video.dtype, shift_override=sigma_shift, ) infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( num_inference_steps=num_inference_steps, device=self.device, dtype=latents_action.dtype, shift_override=sigma_shift, ) for step_t_video, step_delta_video, step_t_action, step_delta_action in zip( infer_timesteps_video, infer_deltas_video, infer_timesteps_action, infer_deltas_action, ): timestep_video = step_t_video.unsqueeze(0).to(dtype=latents_video.dtype, device=self.device) timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) pred_video_posi, pred_action_posi = self._predict_joint_noise( latents_video=latents_video, latents_action=latents_action, timestep_video=timestep_video, timestep_action=timestep_action, context=context, context_mask=context_mask, fuse_vae_embedding_in_latents=fuse_flag, gt_action=action, ) pred_video = pred_video_posi pred_action = pred_action_posi latents_video = self.infer_video_scheduler.step(pred_video, step_delta_video, latents_video) latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) latents_video[:, :, 0:1] = first_frame_latents.clone() action_out = latents_action[0].detach().to(device="cpu", dtype=torch.float32) if test_action_with_infer_action: if not torch.allclose(action_out, action_only_out, atol=1e-2, rtol=1e-2): max_abs_diff = (action_out - action_only_out).abs().max().item() logger.warning( f"Action from infer_joint and infer_action differ with max abs diff {max_abs_diff:.6f}. " ) return { "video": self._decode_latents(latents_video, tiled=tiled), "action": action_out, } @torch.no_grad() def infer_action( self, prompt: Optional[str], input_image: torch.Tensor, action_horizon: int, proprio: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, negative_prompt: Optional[str] = None, text_cfg_scale: float = 1.0, num_inference_steps: int = 20, sigma_shift: Optional[float] = None, seed: Optional[int] = None, rand_device: str = "cpu", tiled: bool = False, ) -> dict[str, Any]: self.eval() if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal": raise ValueError( "`infer_action` requires `video_attention_mask_mode='first_frame_causal'`." ) if input_image.ndim == 3: input_image = input_image.unsqueeze(0) if input_image.ndim != 4 or input_image.shape[0] != 1 or input_image.shape[1] != 3: raise ValueError( f"`input_image` must have shape [1,3,H,W] or [3,H,W], got {tuple(input_image.shape)}" ) _, _, height, width = input_image.shape if height % 16 != 0 or width % 16 != 0: raise ValueError( f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" ) if proprio is not None: if self.proprio_dim is None: raise ValueError("`proprio` was provided but `proprio_dim=None` so `proprio_encoder` is disabled.") if proprio.ndim == 1: proprio = proprio.unsqueeze(0) elif proprio.ndim == 2 and proprio.shape[0] == 1: pass else: raise ValueError(f"`proprio` must be [D] or [1,D], got shape {tuple(proprio.shape)}") if proprio.shape[1] != self.proprio_dim: raise ValueError(f"`proprio` last dim must be {self.proprio_dim}, got {proprio.shape[1]}") proprio = proprio.to(device=self.device, dtype=self.torch_dtype) import time as _time generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) latents_action = torch.randn( (1, action_horizon, self.action_expert.action_dim), generator=generator, device=rand_device, dtype=torch.float32, ).to(device=self.device, dtype=self.torch_dtype) def _sync_ms(t0: float) -> float: torch.cuda.synchronize() return (_time.perf_counter() - t0) * 1000.0 torch.cuda.synchronize() t0 = _time.perf_counter() input_image = input_image.to(device=self.device, dtype=self.torch_dtype) first_frame_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) vae_ms = _sync_ms(t0) fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) use_prompt = prompt is not None use_context = context is not None or context_mask is not None if use_prompt and use_context: raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") if not use_prompt and not use_context: raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") torch.cuda.synchronize() t0 = _time.perf_counter() if use_prompt: context, context_mask = self.encode_prompt(prompt) else: if context is None or context_mask is None: raise ValueError("`context` and `context_mask` must be both provided together.") if context.ndim == 2: context = context.unsqueeze(0) if context_mask.ndim == 1: context_mask = context_mask.unsqueeze(0) if context.ndim != 3 or context_mask.ndim != 2: raise ValueError( f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" ) context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) text_enc_ms = _sync_ms(t0) torch.cuda.synchronize() t0 = _time.perf_counter() if proprio is not None: context, context_mask = self._append_proprio_to_context( context=context, context_mask=context_mask, proprio=proprio, ) proprio_ms = _sync_ms(t0) torch.cuda.synchronize() t0 = _time.perf_counter() timestep_video = torch.zeros( (first_frame_latents.shape[0],), dtype=first_frame_latents.dtype, device=self.device, ) video_pre = self.video_expert.pre_dit( x=first_frame_latents, timestep=timestep_video, context=context, context_mask=context_mask, action=None, fuse_vae_embedding_in_latents=fuse_flag, ) video_seq_len = int(video_pre["tokens"].shape[1]) attention_mask = self._build_mot_attention_mask( video_seq_len=video_seq_len, action_seq_len=latents_action.shape[1], video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), device=video_pre["tokens"].device, ) video_kv_cache = self.mot.prefill_video_cache( video_tokens=video_pre["tokens"], video_freqs=video_pre["freqs"], video_t_mod=video_pre["t_mod"], video_context_payload={ "context": video_pre["context"], "mask": video_pre["context_mask"], }, video_attention_mask=attention_mask[:video_seq_len, :video_seq_len], ) video_prefill_ms = _sync_ms(t0) torch.cuda.synchronize() t0 = _time.perf_counter() infer_timesteps_action, infer_deltas_action = self.infer_action_scheduler.build_inference_schedule( num_inference_steps=num_inference_steps, device=self.device, dtype=latents_action.dtype, shift_override=sigma_shift, ) for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action): timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) pred_action_posi = self._predict_action_noise_with_cache( latents_action=latents_action, timestep_action=timestep_action, context=context, context_mask=context_mask, video_kv_cache=video_kv_cache, attention_mask=attention_mask, video_seq_len=video_seq_len, ) pred_action = pred_action_posi latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action) denoise_ms = _sync_ms(t0) return { "action": latents_action[0].detach().to(device="cpu", dtype=torch.float32), "timing_ms": { "text_enc_ms": text_enc_ms, "proprio_ms": proprio_ms, "vae_ms": vae_ms, "video_prefill_ms": video_prefill_ms, "denoise_ms": denoise_ms, "num_denoise_steps": num_inference_steps, }, } @torch.no_grad() def infer( self, prompt: Optional[str], input_image: torch.Tensor, num_frames: int, action: Optional[torch.Tensor] = None, action_horizon: Optional[int] = None, proprio: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, negative_prompt: Optional[str] = None, text_cfg_scale: float = 5.0, action_cfg_scale: float = 1.0, num_inference_steps: int = 20, sigma_shift: Optional[float] = None, seed: Optional[int] = None, rand_device: str = "cpu", tiled: bool = False, ): return self.infer_joint( prompt=prompt, input_image=input_image, num_video_frames=num_frames, action_horizon=action_horizon, action=action, proprio=proprio, context=context, context_mask=context_mask, negative_prompt=negative_prompt, text_cfg_scale=text_cfg_scale, num_inference_steps=num_inference_steps, sigma_shift=sigma_shift, seed=seed, rand_device=rand_device, tiled=tiled, ) def save_checkpoint(self, path, optimizer=None, step=None): """Save model weights and optional optimizer state to a PyTorch checkpoint. For decoupled MoT variants the KV routing metadata (``kv_source_mode`` and ``kv_source_mapping``) is persisted alongside the weights so ``load_checkpoint`` can reject a checkpoint whose routing disagrees with the current model. Same-length mappings produce identically shaped fused_kv tensors, so shape checks alone cannot catch a semantically wrong routing (mot_decoupled.py:284). """ payload = { "mot": self.mot.state_dict(), "step": step, "torch_dtype": str(self.torch_dtype), } kv_source_mode = getattr(self.mot, "kv_source_mode", None) if kv_source_mode is not None: payload["kv_source_mode"] = kv_source_mode kv_source_mapping = getattr(self.mot, "kv_source_mapping", None) if kv_source_mapping is not None: payload["kv_source_mapping"] = list(kv_source_mapping) # Persist the MoT class name as the single authoritative identity # discriminator. Sibling MoT classes (e.g. the base MoTDecoupled vs the # fixed-RoPE MoTDecoupledActionAlignedVideoRoPE subclass) share identical # state_dict keys/shapes, so a cross-class load is shape-compatible and # would otherwise pass silently despite differing KV-cache RoPE semantics. # A legible class-name string lets load_checkpoint reject that case. payload["mot_class"] = type(self.mot).__name__ # Persist the new_fused_kv RoPE mode alongside mot_class. The two rope # modes ("aligned_3d" vs "original_3d") live on the SAME MoT class # (MoTDecoupledActionAlignedVideoRoPE) and produce byte-identical # state_dicts, so mot_class alone cannot discriminate them: an # original_3d checkpoint would otherwise load and eval silently as the # default aligned_3d, applying the WRONG RoPE to fused video K. A # legible mode string lets load_checkpoint reject that mismatch. Models # without the attribute (base MoTDecoupled, non-decoupled) write None. payload["new_fused_kv_rope_mode"] = getattr( self.mot, "new_fused_kv_rope_mode", None ) payload["aligned_3d_action_spatial_anchor_layout"] = getattr( self.mot, "aligned_3d_action_spatial_anchor_layout", None ) if kv_source_mode == "new_fused_kv": payload["new_fused_kv_projection_mode"] = getattr( self.mot, "new_fused_kv_projection_mode", "full" ) payload["new_fused_kv_simple_head_softmax_fuse_mode"] = getattr( self.mot, "new_fused_kv_simple_head_softmax_fuse_mode", "all" ) payload["new_fused_kv_head_fused_kv_fuse_mode"] = getattr( self.mot, "new_fused_kv_head_fused_kv_fuse_mode", "all" ) # EEF-relative RoPE geometry identity. A CONJUNCTION, not just the # calibration digest: a dataset re-exported at a different raw # resolution changes every anchor while the digest stays identical # (plan Section 20.1). eef_identity = getattr(self.mot, "eef_geometry_identity", None) if eef_identity is not None: payload["eef_geometry_identity"] = dict(eef_identity) if self.proprio_encoder is not None: payload["proprio_encoder"] = self.proprio_encoder.state_dict() if optimizer is not None: payload["optimizer"] = optimizer.state_dict() torch.save(payload, path) def load_checkpoint(self, path, optimizer=None): """Load a FastWAM checkpoint. Fused-MLP checkpoints must include ``kv_fusion.*`` weights whenever the current model has ``mot.kv_fusion`` enabled. This prevents accidental eval/resume with randomly initialized fusion layers. Fused-KV checkpoints must include the five top-level fused_kv mixing parameters (``kv_fusing_layer``, ``k_channel_projection``, ``v_channel_projection``, ``k_channel_bias``, ``v_channel_bias``) whenever the current model is in ``kv_source_mode="fused_kv"``. Those parameters live directly on the MoT (not under a ``kv_fusion`` module, which stays ``None`` for fused_kv), so the ``kv_fusion``-based guard above cannot catch them. Mirroring the fused_mlp guard, a fused_kv checkpoint missing any of the five is rejected rather than silently loaded at random init (B1-class failure). """ payload = torch.load(path, map_location="cpu", weights_only=True) # Detect fused_kv from the authoritative mode attribute on the MoT # (mot_decoupled.py:250). In fused_kv mode the five mixing tensors are # real nn.Parameters and hence appear as top-level state_dict keys; in # every other mode they are None non-persistent buffers (absent from the # state_dict), so this guard only arms for genuine fused_kv models. current_kv_source_mode = getattr(self.mot, "kv_source_mode", None) current_kv_source_mapping = getattr(self.mot, "kv_source_mapping", None) if current_kv_source_mode is not None: # Routing-metadata guard: fused_kv/selected-KV tensor shapes depend # only on the mapping LENGTH, so a same-length but semantically # different routing would load silently if we relied on shapes. # Checkpoints written before this metadata existed lack the keys: # warn (cannot verify) instead of raising, for backward compat. if "kv_source_mode" in payload: ckpt_mode = payload["kv_source_mode"] if ckpt_mode != current_kv_source_mode: raise RuntimeError( "Checkpoint kv_source_mode " f"{ckpt_mode!r} does not match the current model's " f"kv_source_mode {current_kv_source_mode!r}. Refusing to " f"load a checkpoint with mismatched KV routing. Checkpoint: {path}" ) else: logger.warning( "Checkpoint has no `kv_source_mode` metadata (legacy format); " "cannot verify KV routing against the current model " "(kv_source_mode=%r). Checkpoint: %s", current_kv_source_mode, path, ) if current_kv_source_mapping is not None: if "kv_source_mapping" in payload: ckpt_mapping = list(payload["kv_source_mapping"]) if ckpt_mapping != list(current_kv_source_mapping): raise RuntimeError( "Checkpoint kv_source_mapping " f"{ckpt_mapping} does not match the current model's " f"kv_source_mapping {list(current_kv_source_mapping)}. " "Same-length mappings load shape-compatibly but route KV from " f"the WRONG video layers. Checkpoint: {path}" ) elif "kv_source_mode" in payload: # Partial metadata: mode present but mapping absent. New # save_checkpoint always writes both together, so this is a # tampered/hand-built payload; warn (routing unverifiable) # rather than raise. The mode-absent case is already covered # by the legacy warning above. logger.warning( "Checkpoint has `kv_source_mode` but no `kv_source_mapping` " "metadata (partial/hand-built payload); cannot verify KV " "routing mapping against the current model (%s). Checkpoint: %s", list(current_kv_source_mapping), path, ) is_fused_kv = current_kv_source_mode in {"fused_kv", "new_fused_kv"} current_new_fused_kv_projection_mode = getattr( self.mot, "new_fused_kv_projection_mode", "full" ) if ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode == "simple" ): fused_kv_param_names = ["simple_kv_fusing_layer"] elif ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode in {"simple+PE", "simple+PE-postnorm"} ): fused_kv_param_names = ["simple_kv_fusing_layer", "k_video_pos_projection"] elif ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode == "per_head_channel" ): fused_kv_param_names = [ "per_head_kv_fusing_layer", "k_head_channel_projection", "k_head_channel_bias", ] elif ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode in {"simple_head_fused", "simple_head_softmax"} ): fused_kv_param_names = ["head_fused_kv_layer_mixing"] elif ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode == "HeadFusedKV" ): fused_kv_param_names = [ "head_fused_kv_k_channel_projection", "head_fused_kv_v_channel_projection", "head_fused_kv_layer_mixing", ] elif ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode == "HeadFusedKV+Sin2DPE" ): fused_kv_param_names = [ "head_fused_kv_k_channel_projection", "head_fused_kv_v_channel_projection", "head_fused_kv_layer_mixing", ] + [ key for key in self.mot.state_dict() if key.startswith("head_fused_kv_sin2d_pe_mlps.") ] elif ( current_kv_source_mode == "new_fused_kv" and current_new_fused_kv_projection_mode == "MLPMixerFusedKV" ): fused_kv_param_names = [ key for key in self.mot.state_dict() if key.startswith("mlp_mixer_fused_kv_blocks.") ] else: fused_kv_param_names = [ "kv_fusing_layer", "k_channel_projection", "v_channel_projection", "k_channel_bias", "v_channel_bias", ] k_fused_norm = getattr(self.mot, "k_fused_norm", None) if k_fused_norm is not None: fused_kv_param_names.extend( f"k_fused_norm.{i}.weight" for i in range(len(k_fused_norm)) ) if "mot" in payload: mot_state = payload["mot"] # MoT identity guard (mirrors the kv_source_mode guard above). Sibling # MoT classes share identical state_dict keys/shapes, so a cross-class # load is shape-compatible and load_state_dict would NOT catch it. The # semantics differ at inference: the fixed-RoPE MoT caches RAW K and # re-applies RoPE per action step, while the base MoT caches already # RoPE'd K -- loading one class's weights into the other silently # produces WRONG inference. Backward-compat matrix: field present + # mismatch -> raise; present + match -> silent load; absent (legacy # checkpoint written before this metadata) -> warn (cannot verify) and # load, preserving prior behavior. current_mot_class = type(self.mot).__name__ if "mot_class" in payload: ckpt_mot_class = payload["mot_class"] if ckpt_mot_class != current_mot_class: raise RuntimeError( "Checkpoint mot_class " f"{ckpt_mot_class!r} does not match the current model's " f"mot_class {current_mot_class!r}. These MoT classes share " "identical state_dict shapes but have DIFFERENT KV-cache " "RoPE semantics (fixed-RoPE MoT caches raw K and re-applies " "RoPE per step; base MoT caches RoPE'd K), so a cross-class " "load loads shape-compatibly yet produces WRONG inference. " f"Refusing to load. Checkpoint: {path}" ) else: logger.warning( "Checkpoint has no `mot_class` metadata (legacy format); " "cannot verify the MoT class against the current model " "(mot_class=%r). Checkpoint: %s", current_mot_class, path, ) # new_fused_kv RoPE-mode identity guard (mirrors the mot_class guard # above). The two rope modes ("aligned_3d" vs "original_3d") share the # SAME MoT class and produce byte-identical state_dicts, so mot_class # cannot discriminate them and load_state_dict sees no shape symptom. # They differ only at inference: original_3d keeps the video 3D RoPE on # fused video K and the action 1D RoPE on action Q/K, while aligned_3d # re-positions video K into the action-attention 3D basis -- loading one # mode's weights under the other silently produces WRONG inference. # Backward-compat matrix mirrors mot_class: field present + mismatch -> # raise ONLY when the current model actually runs new_fused_kv (for # every other kv_source_mode the attribute is INERT -- the subclass # sets it on all fixed-RoPE instances, so e.g. a fused_kv checkpoint # with a CLI-overridden rope-mode value must not be refused); present + # match -> silent load; None==None (base MoTDecoupled / non-decoupled # models, which never carry the attribute) -> silent pass. # Absent key AND the current model is genuinely new_fused_kv (a legacy # pre-metadata new_fused_kv checkpoint) -> warn (cannot verify) and load. # Cross-kv_source_mode and cross-class loads never reach this guard: # the kv_source_mode and mot_class guards above raise first. current_rope_mode = getattr(self.mot, "new_fused_kv_rope_mode", None) if "new_fused_kv_rope_mode" in payload: ckpt_rope_mode = payload["new_fused_kv_rope_mode"] if ( ckpt_rope_mode != current_rope_mode and current_kv_source_mode == "new_fused_kv" ): raise RuntimeError( "Checkpoint new_fused_kv_rope_mode " f"{ckpt_rope_mode!r} does not match the current model's " f"new_fused_kv_rope_mode {current_rope_mode!r}. Both modes " "share the same MoT class and identical state_dict shapes " "but apply DIFFERENT RoPE to the fused video K / action " "Q/K, so a cross-mode load loads shape-compatibly yet " f"produces WRONG inference. Refusing to load. Checkpoint: {path}" ) elif current_kv_source_mode == "new_fused_kv": # Legacy new_fused_kv checkpoint written before this metadata # existed: cannot verify which rope mode it was trained with. Warn # (not raise) for backward compat. Base/non-new_fused_kv models pass # silently: their None==None comparison never reaches this branch. logger.warning( "Checkpoint has no `new_fused_kv_rope_mode` metadata (legacy " "format); cannot verify the new_fused_kv RoPE mode against the " "current model (new_fused_kv_rope_mode=%r). Checkpoint: %s", current_rope_mode, path, ) # EEF-relative geometry identity. Checked as a conjunction: any one # field changing invalidates the checkpoint, because every one of # them moves the anchors (plan Section 20.1). This is the only # geometry check that exists -- anchors are computed at load time, # so there is no sidecar or dataset fingerprint behind it. current_eef_identity = getattr(self.mot, "eef_geometry_identity", None) if current_eef_identity is not None or "eef_geometry_identity" in payload: ckpt_eef_identity = payload.get("eef_geometry_identity") if ckpt_eef_identity is None: raise RuntimeError( "Current model uses EEF-relative camera RoPE " f"(new_fused_kv_rope_mode={current_rope_mode!r}) but the " "checkpoint carries no `eef_geometry_identity`. It cannot " "be verified to have been trained against this geometry. " f"Refusing to load. Checkpoint: {path}" ) if current_eef_identity is None: raise RuntimeError( "Checkpoint carries `eef_geometry_identity` but the current " "model has none; it was not configured for EEF-relative " f"camera RoPE. Refusing to load. Checkpoint: {path}" ) differing = { field: (ckpt_eef_identity.get(field), current_eef_identity.get(field)) for field in current_eef_identity if ckpt_eef_identity.get(field) != current_eef_identity.get(field) } if differing: detail = "; ".join( f"{field}: checkpoint={ckpt!r} current={cur!r}" for field, (ckpt, cur) in sorted(differing.items()) ) raise RuntimeError( "Checkpoint EEF geometry identity does not match the " f"current model ({detail}). Every field here moves the " "anchors, so the loaded weights were trained against a " "different spatial origin and would produce WRONG " f"inference. Refusing to load. Checkpoint: {path}" ) current_anchor_layout = getattr( self.mot, "aligned_3d_action_spatial_anchor_layout", None ) if ( current_kv_source_mode == "new_fused_kv" and current_rope_mode in {"aligned_3d", "aligned_3dp"} ): if "aligned_3d_action_spatial_anchor_layout" in payload: checkpoint_anchor_layout = payload[ "aligned_3d_action_spatial_anchor_layout" ] if checkpoint_anchor_layout != current_anchor_layout: raise RuntimeError( "Checkpoint aligned_3d_action_spatial_anchor_layout " f"{checkpoint_anchor_layout!r} does not match the " "current model's " "aligned_3d_action_spatial_anchor_layout " f"{current_anchor_layout!r}. The layouts share " "identical state_dict shapes but assign different " "camera anchors to action-attention heads. Refusing " f"to load. Checkpoint: {path}" ) else: logger.warning( "Checkpoint has no " "`aligned_3d_action_spatial_anchor_layout` metadata " "(legacy format); cannot verify the %s action " "anchor layout against the current model " "(aligned_3d_action_spatial_anchor_layout=%r). " "Checkpoint: %s", current_rope_mode, current_anchor_layout, path, ) if current_kv_source_mode == "new_fused_kv": inferred_projection_mode = _infer_new_fused_kv_projection_mode(mot_state) checkpoint_projection_mode = payload.get( "new_fused_kv_projection_mode", inferred_projection_mode ) if checkpoint_projection_mode is None: # Legacy / unstamped checkpoints: trust the model config the # caller already built (CLI/Hydra). Safer than hard-failing # eval on otherwise loadable weights. logger.warning( "Checkpoint has no `new_fused_kv_projection_mode` metadata " "and its projection mode cannot be inferred from the MoT " "state; trusting the current model's " "new_fused_kv_projection_mode=%r. Checkpoint: %s", current_new_fused_kv_projection_mode, path, ) checkpoint_projection_mode = current_new_fused_kv_projection_mode if ( "new_fused_kv_projection_mode" in payload and inferred_projection_mode is not None and not _new_fused_kv_projection_signature_matches( checkpoint_projection_mode, inferred_projection_mode, ) ): raise RuntimeError( "Checkpoint new_fused_kv projection metadata disagrees with its " f"parameter signature: metadata={checkpoint_projection_mode!r}, " f"inferred={inferred_projection_mode!r}. Checkpoint: {path}" ) if checkpoint_projection_mode != current_new_fused_kv_projection_mode: raise RuntimeError( "Checkpoint new_fused_kv_projection_mode " f"{checkpoint_projection_mode!r} does not match the current " "model's new_fused_kv_projection_mode " f"{current_new_fused_kv_projection_mode!r}. Refusing to load " f"a checkpoint with a mismatched projection architecture. Checkpoint: {path}" ) if "new_fused_kv_projection_mode" not in payload: logger.warning( "Checkpoint has no `new_fused_kv_projection_mode` metadata " "(legacy format); inferred %r from its MoT parameter signature. " "Checkpoint: %s", inferred_projection_mode, path, ) if current_new_fused_kv_projection_mode == "simple_head_softmax": current_softmax_fuse_mode = getattr( self.mot, "new_fused_kv_simple_head_softmax_fuse_mode", "all", ) if "new_fused_kv_simple_head_softmax_fuse_mode" in payload: checkpoint_softmax_fuse_mode = payload[ "new_fused_kv_simple_head_softmax_fuse_mode" ] if checkpoint_softmax_fuse_mode == "uni_end": checkpoint_softmax_fuse_mode = "uniform_end" if checkpoint_softmax_fuse_mode != current_softmax_fuse_mode: raise RuntimeError( "Checkpoint new_fused_kv_simple_head_softmax_fuse_mode " f"{checkpoint_softmax_fuse_mode!r} does not match " "the current model's " "new_fused_kv_simple_head_softmax_fuse_mode " f"{current_softmax_fuse_mode!r}. The modes share " "the same state_dict shapes but route video layers " "to action DiT layers differently. Refusing to load. " f"Checkpoint: {path}" ) elif current_softmax_fuse_mode != "all": raise RuntimeError( "Checkpoint has no " "`new_fused_kv_simple_head_softmax_fuse_mode` metadata " "(legacy format, treated as 'all') but the current " "model requests " f"{current_softmax_fuse_mode!r}. Refusing to load a " "shape-compatible checkpoint with unverifiable " f"simple_head_softmax routing. Checkpoint: {path}" ) if current_new_fused_kv_projection_mode in { "HeadFusedKV", "HeadFusedKV+Sin2DPE", }: current_head_fused_fuse_mode = getattr( self.mot, "new_fused_kv_head_fused_kv_fuse_mode", "all", ) if "new_fused_kv_head_fused_kv_fuse_mode" in payload: checkpoint_head_fused_fuse_mode = payload[ "new_fused_kv_head_fused_kv_fuse_mode" ] if checkpoint_head_fused_fuse_mode != current_head_fused_fuse_mode: raise RuntimeError( "Checkpoint new_fused_kv_head_fused_kv_fuse_mode " f"{checkpoint_head_fused_fuse_mode!r} does not " "match the current model's " "new_fused_kv_head_fused_kv_fuse_mode " f"{current_head_fused_fuse_mode!r}. The modes " "share the same state_dict shapes but route " "video layers to action DiT layers differently. " f"Refusing to load. Checkpoint: {path}" ) elif current_head_fused_fuse_mode != "all": raise RuntimeError( "Checkpoint has no " "`new_fused_kv_head_fused_kv_fuse_mode` metadata " "(legacy format, treated as 'all') but the current " "model requests " f"{current_head_fused_fuse_mode!r}. Refusing to load " "a shape-compatible checkpoint with unverifiable " f"HeadFusedKV routing. Checkpoint: {path}" ) if is_fused_kv: # Pre-load guard: the fused_kv mixing params must be present # in the incoming state before we attempt a (strict=False) load, # otherwise strict=False would silently leave them at random init. missing_fused_kv_keys = [ key for key in fused_kv_param_names if key not in mot_state ] if missing_fused_kv_keys: raise RuntimeError( "Checkpoint is missing fused_kv mixing parameters for a " "fused_kv model. Refusing to leave the KV-mixing tensors " f"at random init. Missing: {', '.join(missing_fused_kv_keys)}. " f"Checkpoint: {path}" ) if getattr(self.mot, "kv_fusion", None) is not None: expected_fusion_keys = [ key for key in self.mot.state_dict().keys() if key.startswith("kv_fusion.") ] missing_fusion_keys = [ key for key in expected_fusion_keys if key not in mot_state ] if missing_fusion_keys: preview = ", ".join(missing_fusion_keys[:5]) if len(missing_fusion_keys) > 5: preview += f", ... ({len(missing_fusion_keys)} total)" raise RuntimeError( "Checkpoint is missing kv_fusion.* weights for a fused_mlp model. " f"Refusing to leave fusion layers at random init. Missing: {preview}. " f"Checkpoint: {path}" ) incompatible = self.mot.load_state_dict(mot_state, strict=False) if getattr(self.mot, "kv_fusion", None) is not None: missing_fusion_keys = [ key for key in incompatible.missing_keys if key.startswith("kv_fusion.") ] if missing_fusion_keys: raise RuntimeError( "Checkpoint load reported missing kv_fusion.* weights for a fused_mlp " f"model: {missing_fusion_keys}. Checkpoint: {path}" ) if is_fused_kv: # Post-load defense: even after the pre-load key check, confirm # load_state_dict did not report any fused_kv mixing tensors as # missing (e.g. a partial/renamed state). Any missing # entry means the tensor stayed at random init. missing_fused_kv_after = [ key for key in incompatible.missing_keys if key in fused_kv_param_names ] if missing_fused_kv_after: raise RuntimeError( "Checkpoint load reported missing fused_kv mixing " f"parameters for a fused_kv model: {missing_fused_kv_after}. " f"Checkpoint: {path}" ) if incompatible.missing_keys: # General missing-keys guard: strict=False tolerates ANY absent # key, so a full `mot` checkpoint missing non-fused weights # (e.g. a truncated or foreign-variant state dict) would leave # those tensors silently at initialization (B1-class failure). # The fused-specific guards above fire first with more # actionable messages; this catches everything else. Unexpected # EXTRA keys stay tolerated (forward compat with newer saves). preview = ", ".join(incompatible.missing_keys[:5]) if len(incompatible.missing_keys) > 5: preview += f", ... ({len(incompatible.missing_keys)} total)" raise RuntimeError( "Checkpoint `mot` state is missing weights present in the " f"current model; those tensors would silently stay at " f"initialization. Missing: {preview}. Checkpoint: {path}" ) elif "dit" in payload: if getattr(self.mot, "kv_fusion", None) is not None: raise RuntimeError( "Cannot load legacy `dit` checkpoint into a fused_mlp model because " "`dit` checkpoints do not contain kv_fusion.* weights. Load a full " f"`mot` checkpoint instead. Checkpoint: {path}" ) # Mirror the fused_mlp guard above for fused_kv models. A legacy `dit` # checkpoint only carries the video expert's weights, so it can never # contain the five top-level fused_kv mixing params (kv_fusing_layer / # k/v_channel_projection / k/v_channel_bias); loading it would leave # them at random init (a B1-class silent corruption). is_fused_kv was # computed above from the authoritative mot.kv_source_mode attribute. if is_fused_kv: raise RuntimeError( "Cannot load legacy `dit` checkpoint into a fused_kv model because " "`dit` checkpoints do not contain the fused_kv mixing parameters " f"({', '.join(fused_kv_param_names)}). Load a full `mot` checkpoint " f"instead. Checkpoint: {path}" ) logger.warning("Loading legacy `dit` checkpoint into video expert only.") self.video_expert.load_state_dict(payload["dit"], strict=False) else: raise ValueError(f"Checkpoint missing both `mot` and `dit` keys: {path}") if self.proprio_encoder is not None: if "proprio_encoder" in payload: self.proprio_encoder.load_state_dict(payload["proprio_encoder"], strict=True) else: logger.warning("Checkpoint has no `proprio_encoder` weights; keeping current `proprio_encoder` params.") elif "proprio_encoder" in payload: logger.warning("Checkpoint contains `proprio_encoder` weights but current model has `proprio_dim=None`; ignoring.") if optimizer is not None and "optimizer" in payload: optimizer.load_state_dict(payload["optimizer"]) return payload def forward(self, *args, **kwargs): return self.training_loss(*args, **kwargs)