diff --git a/wan_va/modules/model.py b/wan_va/modules/model.py index 25b45e57..675833d9 100644 --- a/wan_va/modules/model.py +++ b/wan_va/modules/model.py @@ -361,6 +361,12 @@ def init_kv_cache(self, cache_name, total_tolen, num_head, head_dim, torch.zeros((total_tolen, ), dtype=torch.bool, device=device), "is_pred": torch.zeros((total_tolen, ), dtype=torch.bool, device=device), + "frame_id": + torch.full((total_tolen, ), -1, device=device, dtype=torch.long), + "current_frame_id": + -1, + "window_size": + 999999, # default: no window limit } def allocate_slots(self, cache_name, key_size): @@ -406,8 +412,49 @@ def update_cache(self, cache_name, key, value, is_pred): cache['mask'][slots] = True cache['id'][slots] = new_id cache['is_pred'][slots] = is_pred + cache['frame_id'][slots] = cache.get('current_frame_id', -1) return slots + def set_frame_context(self, cache_name, frame_id, window_size): + """Set current frame_id and window_size for subsequent cache updates and attention filtering.""" + if self.attn_caches is None or cache_name not in self.attn_caches: + return + cache = self.attn_caches[cache_name] + if cache is not None: + cache['current_frame_id'] = frame_id + cache['window_size'] = window_size + + @staticmethod + def _select_valid_cache_slots(cache, update_cache): + """Filter cache slots to match training-time attention visibility. + + During training, a noisy query may attend to: + - past clean tokens (real observations / real actions from earlier frames) + - same-frame predicted tokens (within the same chunk) + - strictly-past predicted tokens + and only within the attention window. The original streaming-inference + code instead surfaced every allocated slot, which let noisy queries + attend to same-frame clean tokens and to predicted tokens from future + frames -- a train/inference distribution mismatch. + """ + valid = cache['mask'].nonzero(as_tuple=False).squeeze(-1) + if update_cache not in (0, 1) or valid.numel() == 0: + return valid + + is_pred = cache['is_pred'][valid] + frame_ids = cache['frame_id'][valid] + current_frame_id = cache.get('current_frame_id', -1) + window_size = cache.get('window_size', 999999) + + is_clean = ~is_pred + is_past = (frame_ids >= 0) & (frame_ids < current_frame_id) + is_same_frame = frame_ids == current_frame_id + within_window = (frame_ids - current_frame_id).abs() <= window_size + # Training mask: noisy query -> clean key must be past; noisy query -> noisy + # key must be same-frame or past; all subject to the attention window. + allowed = (is_clean & is_past) | (is_pred & (is_same_frame | is_past)) + return valid[allowed & within_window] + def restore_cache(self, cache_name, slots): self.attn_caches[cache_name]['mask'][slots] = False @@ -447,8 +494,10 @@ def apply_rotary_emb(x, freqs): is_pred=(update_cache == 1)) key_pool = self.attn_caches[cache_name]['k'] value_pool = self.attn_caches[cache_name]['v'] - mask = self.attn_caches[cache_name]['mask'] - valid = mask.nonzero(as_tuple=False).squeeze(-1) + valid = self._select_valid_cache_slots( + self.attn_caches[cache_name], update_cache + ) + key = key_pool[:, valid] value = value_pool[:, valid] @@ -658,6 +707,11 @@ def clear_pred_cache(self, cache_name): for block in self.blocks: block.attn1.clear_pred_cache(cache_name) + def set_frame_context(self, cache_name, frame_id, window_size): + """Set frame context for inference attention masking across all layers.""" + for block in self.blocks: + block.attn1.set_frame_context(cache_name, frame_id, window_size) + def create_empty_cache(self, cache_name, attn_window, latent_token_per_chunk, action_token_per_chunk, device, dtype, batch_size): diff --git a/wan_va/wan_va_server.py b/wan_va/wan_va_server.py index 4c95ace4..bd6e52a2 100644 --- a/wan_va/wan_va_server.py +++ b/wan_va/wan_va_server.py @@ -250,7 +250,37 @@ def postprocess_action(self, action): raise NotImplementedError action = action.squeeze(0).detach().cpu().numpy() return action[self.job_config.used_action_channel_ids] - + + def _make_zero_action_condition(self): + """Build the cold-start action conditioning token. + + At the very first inference step (``frame_st_id == 0``) there is no + real action history to feed the model, so we construct an all-zero + action in *physical* space. Crucially, the model operates in + quantile-normalized space, where physical zero is *not* the tensor + ``0`` -- it is ``(0 - q01) / (q99 - q01) * 2 - 1`` per channel. + Directly feeding ``torch.zeros`` therefore decoded to the per-channel + statistical midpoint ``(q01 + q99) / 2`` rather than to physical zero, + seeding the diffusion loop with a spurious non-zero action. That + erroneous conditioning entered both the action denoising loop and the + KV cache, corrupting every subsequent chunk. + """ + action_cond = torch.zeros( + [1, self.job_config.action_dim, 1, self.action_per_frame, 1], + device=self.device, + dtype=torch.float32, + ) + if self.action_norm_method == 'quantiles': + q01 = self.actions_q01.to(self.device).view(1, -1, 1, 1, 1) + q99 = self.actions_q99.to(self.device).view(1, -1, 1, 1, 1) + action_cond = (action_cond - q01) / (q99 - q01 + 1e-6) * 2. - 1. + else: + raise NotImplementedError + + action_mask = self.action_mask.to(self.device) + action_cond[:, ~action_mask] = 0 + return action_cond.to(self.dtype) + def _repeat_input_for_cfg(self, input_dict): if self.use_cfg: input_dict['noisy_latents'] = input_dict['noisy_latents'].repeat(2, 1, 1, 1, 1) @@ -486,6 +516,11 @@ def _infer(self, obs, frame_st_id=0): torch.no_grad(), ): # 1. Video Generation Loop + # Set frame context so noisy latent queries cannot attend to + # predicted tokens of future frames in the KV cache. + latent_frame_id = frame_st_id + self.transformer.set_frame_context(self.cache_name, latent_frame_id, + self.job_config.attn_window) for i, t in enumerate(tqdm(timesteps)): last_step = i == len(timesteps) - 1 latent_cond = init_latent[:, :, 0:1].to( @@ -521,15 +556,16 @@ def _infer(self, obs, frame_st_id=0): latents[:, :, 0:1] = latent_cond if frame_st_id == 0 else latents[:, :, 0:1] + # 2. Action Generation Loop + # Action tokens live at the *next* frame id relative to the video + # tokens they accompany; advance the frame context accordingly. + action_frame_id = frame_st_id + 1 + self.transformer.set_frame_context(self.cache_name, action_frame_id, + self.job_config.attn_window) for i, t in enumerate(tqdm(action_timesteps)): last_step = i == len(action_timesteps) - 1 - action_cond = torch.zeros( - [ - 1, self.job_config.action_dim, 1, - self.action_per_frame, 1 - ], - device=self.device, - dtype=self.dtype) if frame_st_id == 0 else None + action_cond = self._make_zero_action_condition( + ) if frame_st_id == 0 else None input_dict = self._prepare_latent_input( None, @@ -591,11 +627,17 @@ def _compute_kv_cache(self, obs): with ( torch.no_grad(), ): + # Set latent frame context (clean observations, update_cache=2) + self.transformer.set_frame_context(self.cache_name, self.frame_st_id, + self.job_config.attn_window) self.transformer(self._repeat_input_for_cfg(input_dict['latent_res_lst']), update_cache=2, cache_name=self.cache_name, action_mode=False) + # Set action frame context (clean observations, update_cache=2) + self.transformer.set_frame_context(self.cache_name, self.frame_st_id + 1, + self.job_config.attn_window) self.transformer(self._repeat_input_for_cfg(input_dict['action_res_lst']), update_cache=2, cache_name=self.cache_name,