''' Swati G, Feb 2026. Autoregressive inference handler for Any4D model. ''' import copy import torch from imaginaire.utils import log, misc from custom.any4d.a4d_logistics import unpack_entries_from_streams class AutoregressiveHandler: """ Handles autoregressive inference for Any4D model. Generates long video sequences by: 1. Precomputing full action trajectory upfront using backtracking 2. Predicting frames segment-by-segment 3. Using last N predicted frames as context for next segment 4. Accumulating non-overlapping predictions into final output Per-segment explanation (if context = conditioning = 4 in latent space): - Segment 0: Input 11 frames (4 cond + 7 pred), with input_mask set accordingly, output has 11 frames, out of which last 7 are actual predictions --> total 11 frames - Segment 1: Use last 4 frame from predictions of segment=0 as context (i.e. overlap), generate 7 new -> so current_size = (4+7)+7 → total 18 frames in val_dict - Segment 2: Again use last 4 frames from predicts of segment-1 as context, generate 7 new -> so current_size = (4+7)+7+7 → total 25 frames in val_dict and so on... """ def __init__(self, config): """ :param config: Model configuration (Any4DConfig) """ self.config = config def precompute_full_action_trajectory( self, original_actions: torch.Tensor, num_segments: int, cond_frames_raw: int, total_frames_raw: int, extrapolation_strategy: str = "backtrack", ) -> torch.Tensor: """ Precompute the full action trajectory for all segments. Supports two extrapolation strategies: - "backtrack": Reverses the arm's trajectory by computing negative velocities, making the arm retrace its exact path backwards (undo the forward motion). - "extrapolate": Continues forward motion by applying recent velocities in the same direction, extrapolating xyz positions for both arms. :param original_actions: Original action sequence from data_batch (1, 41, 20) :param num_segments: Total number of segments to generate :param cond_frames_raw: Number of context frames in raw space (e.g., 13) :param total_frames_raw: Total frames in raw space per segment (e.g., 41) :param extrapolation_strategy: Strategy to use ("backtrack" or "extrapolate") :return: Full precomputed action trajectory (1, T_total, 20) where T_total = 41 + (num_segments - 1) * (41 - 13) """ if num_segments <= 1: return original_actions.clone() # Start with the original actions full_trajectory = original_actions.clone() # (1, 41, 20) num_new_frames_per_segment = total_frames_raw - cond_frames_raw # 28 frames log.debug(f"[Precompute Actions] Computing full trajectory for {num_segments} segments using '{extrapolation_strategy}' strategy") log.debug(f" Original actions: {original_actions.shape}") log.debug(f" New frames per segment: {num_new_frames_per_segment}") # For each additional segment, extrapolate using selected strategy for seg_idx in range(1, num_segments): T_current = full_trajectory.shape[1] # Use the most recent motion history to compute velocities lookback_frames = min(num_new_frames_per_segment + 1, T_current) if lookback_frames > 1: # Extract the most recent motion history history_start = T_current - lookback_frames motion_history = full_trajectory[:, history_start:T_current, :] # (1, lookback, 20) # Compute velocities from the forward motion: v[i] = pos[i+1] - pos[i] forward_velocities = motion_history[:, 1:, :] - motion_history[:, :-1, :] if extrapolation_strategy == "backtrack": # Backtracking: apply negative velocities in reverse order current_pos = full_trajectory[:, -1:, :].clone() # Last position (1, 1, 20) extrapolated_actions = [] # Apply velocities in reverse: most recent velocity first for i in range(min(num_new_frames_per_segment, forward_velocities.shape[1])): vel_idx = forward_velocities.shape[1] - 1 - i velocity = forward_velocities[:, vel_idx:vel_idx+1, :] current_pos = current_pos - velocity extrapolated_actions.append(current_pos.clone()) # If we need more frames than we have velocities, hold the last position while len(extrapolated_actions) < num_new_frames_per_segment: extrapolated_actions.append(extrapolated_actions[-1].clone()) new_actions = torch.cat(extrapolated_actions, dim=1) # (1, 28, 20) log.debug(f" Segment {seg_idx}: Backtracked {num_new_frames_per_segment} frames " f"using {lookback_frames} frames of history (frames {history_start}-{T_current-1})") elif extrapolation_strategy == "extrapolate": # Simple extrapolation: continue forward motion by applying recent velocities # Compute the average velocity from recent history (use last few velocities) num_vel_frames = min(5, forward_velocities.shape[1]) # Use last 5 velocities recent_velocities = forward_velocities[:, -num_vel_frames:, :] # (1, N, 20) avg_velocity = recent_velocities.mean(dim=1, keepdim=True) # (1, 1, 20) # Start from the last position and extrapolate forward current_pos = full_trajectory[:, -1:, :].clone() # Last position (1, 1, 20) extrapolated_actions = [] for i in range(num_new_frames_per_segment): # Apply the average velocity to continue forward motion current_pos = current_pos + avg_velocity extrapolated_actions.append(current_pos.clone()) new_actions = torch.cat(extrapolated_actions, dim=1) # (1, 28, 20) log.debug(f" Segment {seg_idx}: Extrapolated {num_new_frames_per_segment} frames forward " f"using average of last {num_vel_frames} velocities") else: raise ValueError(f"Unknown extrapolation strategy: {extrapolation_strategy}. " f"Must be 'backtrack' or 'extrapolate'.") else: # Fallback: repeat last action new_actions = full_trajectory[:, -1:, :].expand(-1, num_new_frames_per_segment, -1).clone() log.debug(f" Segment {seg_idx}: Not enough history, repeating last action") # Append to full trajectory full_trajectory = torch.cat([full_trajectory, new_actions], dim=1) log.debug(f" Final trajectory shape: {full_trajectory.shape}") return full_trajectory def get_segment_actions( self, full_action_trajectory: torch.Tensor, segment_idx: int, cond_frames_raw: int, total_frames_raw: int, ) -> torch.Tensor: """ Extract the action window for a specific segment from the precomputed trajectory. :param full_action_trajectory: Precomputed full action trajectory (1, T_total, 20) :param segment_idx: Current segment index (0-based) :param cond_frames_raw: Number of context frames in raw space (e.g., 13) :param total_frames_raw: Total frames in raw space per segment (e.g., 41) :return: Action window for this segment (1, 41, 20) """ num_new_frames = total_frames_raw - cond_frames_raw # 28 if segment_idx == 0: # First segment: use first total_frames_raw start_idx = 0 else: # Subsequent segments: start from where context begins # Segment 1: context starts at 41 - 13 = 28, window is [28:69] # Segment 2: context starts at 69 - 13 = 56, window is [56:97] start_idx = total_frames_raw + (segment_idx - 1) * num_new_frames - cond_frames_raw end_idx = start_idx + total_frames_raw segment_actions = full_action_trajectory[:, start_idx:end_idx, :] return segment_actions def prepare_next_segment( self, current_data_batch: dict, segment_val_dict: dict, cond_frames_latent: int, cond_frames_raw: int, total_frames_raw: int, total_frames_latent: int, full_action_trajectory: torch.Tensor, segment_idx: int, ) -> dict: """ Extract last cond_frames from predictions and prepare next segment data_batch. :param current_data_batch: Current segment's data batch :param segment_val_dict: Output from pipe.generate() for current segment :param cond_frames_latent: Number of context frames in latent space (e.g., 4) :param cond_frames_raw: Number of context frames in raw/pixel space (e.g., 13) :param total_frames_raw: Total number of frames in raw/pixel space (e.g., 41) :param total_frames_latent: Total number of frames in latent space (e.g., 11) :param full_action_trajectory: Precomputed full action trajectory :param segment_idx: Current segment index (0 for first segment) :return: Modified data_batch for next segment """ # Deepcopy to avoid modifying original batch next_segment_batch = copy.deepcopy(current_data_batch) # Unpack predictions from streams to entries y0_pred_entries = unpack_entries_from_streams( segment_val_dict['y0_pred_streams'], self.config, detach=True, ignore_masks=True ) # For each RGB view, extract last cond_frames_latent predictions as new context for view_idx in range(self.config.num_views): view_key = f'rgb{view_idx}' if view_key not in y0_pred_entries: continue # Get predicted latent frames: (1, C, T, H, W) where T=11 pred_latent = y0_pred_entries[view_key] # (1, 16, 11, 48, 64) # Extract last cond_frames_latent frames as context for next segment context_latent = pred_latent[:, :, -cond_frames_latent:, :, :] # (1, 16, 4, 48, 64) # Create remaining frames as zeros (no ground truth for future) B, C, _, H, W = context_latent.shape num_remaining = total_frames_latent - cond_frames_latent # 7 frames remaining_latent = torch.zeros(B, C, num_remaining, H, W, dtype=context_latent.dtype, device=context_latent.device) # Concatenate: context from prediction + zeros for output region new_latent = torch.cat([context_latent, remaining_latent], dim=2) next_segment_batch['a4d_latent'][view_key] = new_latent # Update masks: context frames are INPUT (=1), rest are OUTPUT (=1) input_mask = torch.zeros_like(next_segment_batch['a4d_latent'][f'{view_key}_input_mask']) input_mask[:, :, :cond_frames_latent, :, :] = 1.0 next_segment_batch['a4d_latent'][f'{view_key}_input_mask'] = input_mask output_mask = torch.zeros_like(next_segment_batch['a4d_latent'][f'{view_key}_output_mask']) output_mask[:, :, cond_frames_latent:, :, :] = 1.0 next_segment_batch['a4d_latent'][f'{view_key}_output_mask'] = output_mask next_segment_batch['a4d_latent'][f'{view_key}_supervise_mask'] = output_mask.clone() # Handle actions: extract from precomputed trajectory (next segment = segment_idx + 1) if 'action' in next_segment_batch['a4d_latent'] and full_action_trajectory is not None: next_segment_idx = segment_idx + 1 segment_actions = self.get_segment_actions( full_action_trajectory, next_segment_idx, cond_frames_raw, total_frames_raw, ) next_segment_batch['a4d_latent']['action'] = segment_actions # Action masks: For world model, ALL action frames should be conditioning B, T, _ = segment_actions.shape action_input_mask = torch.ones(B, T, 1, dtype=segment_actions.dtype, device=segment_actions.device) next_segment_batch['a4d_latent']['action_input_mask'] = action_input_mask action_output_mask = torch.zeros(B, T, 1, dtype=segment_actions.dtype, device=segment_actions.device) next_segment_batch['a4d_latent']['action_output_mask'] = action_output_mask action_supervise_mask = torch.zeros(B, T, 1, dtype=segment_actions.dtype, device=segment_actions.device) next_segment_batch['a4d_latent']['action_supervise_mask'] = action_supervise_mask log.debug(f"Segment {next_segment_idx}: Extracted actions of shape {segment_actions.shape} " f"from precomputed trajectory") return next_segment_batch def combine_segments( self, predictions_so_far: list, cond_frames_latent: int, cond_frames_raw: int, full_action_trajectory: torch.Tensor, ) -> dict: """ Combine accumulated segment predictions into a single temporally extended val_dict. Handles overlap by keeping only non-overlapping frames from each segment. :param predictions_so_far: List of segment_val_dict for all the segments :param cond_frames_latent: Number of overlapping context frames (e.g., 4) :param cond_frames_raw: Number of overlapping context frames in raw space (e.g., 13) :param full_action_trajectory: Precomputed full action trajectory :return: Combined val_dict with extended temporal dimension """ # Initialize combined val_dict combined_val_dict = {} # Combine y0_pred_streams combined_streams = {} # Get stream keys from first segment first_segment_val_dict = predictions_so_far[0] pred_stream_keys = list(first_segment_val_dict['y0_pred_streams'].keys()) # ['v0', 'v1'] for stream_key in pred_stream_keys: accumulated_frames = [] for seg_idx, segment_val_dict in enumerate(predictions_so_far): y0_pred_stream_tensor = segment_val_dict['y0_pred_streams'][stream_key] # (1, C, T, H, W) if seg_idx == 0: # First segment: keep all frames accumulated_frames.append(y0_pred_stream_tensor) else: # Subsequent segments: skip first cond_frames_latent (overlapping context) new_frames = y0_pred_stream_tensor[:, :, cond_frames_latent:, :, :] accumulated_frames.append(new_frames) # Concatenate along temporal dimension combined_frame_tensor = torch.cat(accumulated_frames, dim=2) # (1, C, T_total, H, W) combined_streams[stream_key] = combined_frame_tensor combined_val_dict['y0_pred_streams'] = combined_streams # Use precomputed full action trajectory (slice to match video frames) # Video: 11 + (N-1) * 7 latent frames, Actions: 41 + (N-1) * 28 raw frames # Combine x0_streams and xt_streams, handling overlap first_x0_streams = first_segment_val_dict.get('x0_streams', {}) x0_stream_keys = [k for k in first_x0_streams.keys() if k in first_x0_streams] for stream_key in x0_stream_keys: accumulated_frames_x0 = [] accumulated_frames_xt = [] for seg_idx, segment_val_dict in enumerate(predictions_so_far): x0_tensor = segment_val_dict['x0_streams'][stream_key] xt_tensor = segment_val_dict['xt_streams'][stream_key] if seg_idx == 0: # First segment: keep all frames accumulated_frames_x0.append(x0_tensor) accumulated_frames_xt.append(xt_tensor) else: # Subsequent segments: skip first cond_frames (overlapping context) if stream_key == 'adaln': # adaln (action) has shape (1, T=41, C) where T is in raw/pixel space new_frames_x0 = x0_tensor[:, cond_frames_raw:, :] new_frames_xt = xt_tensor[:, cond_frames_raw:, :] accumulated_frames_x0.append(new_frames_x0) accumulated_frames_xt.append(new_frames_xt) else: # Regular streams have shape (1, C, T, H, W) where T is in latent space new_frames_x0 = x0_tensor[:, :, cond_frames_latent:, :, :] new_frames_xt = xt_tensor[:, :, cond_frames_latent:, :, :] accumulated_frames_x0.append(new_frames_x0) accumulated_frames_xt.append(new_frames_xt) # Concatenate along temporal dimension if stream_key == 'adaln': # adaln: temporal dimension is 1, shape (1, T, C) combined_x0_tensor = torch.cat(accumulated_frames_x0, dim=1) # (1, T_total, C) combined_xt_tensor = torch.cat(accumulated_frames_xt, dim=1) # (1, T_total, C) # Verify action portion of adaln stream matches accumulated_actions action_dim = full_action_trajectory.shape[-1] # typically 20 assert torch.allclose(combined_x0_tensor[:, :, :action_dim].float(), full_action_trajectory.float()), \ f"adaln x0_stream action mismatch with accumulated_actions: " \ f"x0={combined_x0_tensor.shape}, actions={full_action_trajectory.shape}" assert torch.allclose(combined_xt_tensor[:, :, :action_dim].float(), full_action_trajectory.float()), \ f"adaln xt_stream action mismatch with accumulated_actions: " \ f"xt={combined_xt_tensor.shape}, actions={full_action_trajectory.shape}" else: # Regular streams: temporal dimension is 2, shape (1, C, T, H, W) combined_x0_tensor = torch.cat(accumulated_frames_x0, dim=2) # (1, C, T_total, H, W) combined_xt_tensor = torch.cat(accumulated_frames_xt, dim=2) # (1, C, T_total, H, W) combined_val_dict.setdefault('x0_streams', {})[stream_key] = combined_x0_tensor combined_val_dict.setdefault('xt_streams', {})[stream_key] = combined_xt_tensor # Create masks structure masks = { 'is_mask': {}, 'input': {}, 'output': {}, 'supervise': {} } for stream_key, stream_tensor in combined_streams.items(): B, C, T, H, W = stream_tensor.shape device = stream_tensor.device dtype = stream_tensor.dtype masks['is_mask'][stream_key] = torch.zeros((B, C, T, H, W), device=device, dtype=dtype) masks['input'][stream_key] = torch.zeros((B, C, T, H, W), device=device, dtype=dtype) masks['output'][stream_key] = torch.ones((B, C, T, H, W), device=device, dtype=dtype) masks['supervise'][stream_key] = torch.ones((B, C, T, H, W), device=device, dtype=dtype) combined_val_dict['masks'] = masks # Use sampling_params from first segment combined_val_dict['sampling_params'] = first_segment_val_dict['sampling_params'] return combined_val_dict def run_autoregressive_inference( self, data_batch: dict, pipe, config, directives: dict, seed: int = 0, ) -> tuple[dict, dict]: """ Run full autoregressive inference loop. :param data_batch: Input data batch (already preprocessed) :param pipe: Pipeline object with generate() method :param config: Model configuration :param directives: Inference directives including num_segments :param seed: Random seed for generation :return: (combined val_dict, updated data_batch with extended masks) """ predictions_so_far = [] num_segments = directives.get('num_segments', 2) # Default 2 segments # Extract metadata cond_frames_raw = data_batch['meta']['cond_frames_raw'] cond_frames_latent = data_batch['meta']['cond_frames_latent'] total_frames_raw = data_batch['meta']['total_frames_raw'] total_frames_latent = data_batch['meta']['total_frames_latent'] log.debug(f'[Autoregressive] Running {num_segments} segments with ' f'cond_frames_latent={cond_frames_latent}') # STEP 1: Precompute full action trajectory upfront full_action_trajectory = None if 'action' in data_batch['a4d_latent']: original_actions = data_batch['a4d_latent']['action'].clone() extrapolation_strategy = directives.get('extrapolation_strategy', 'backtrack') full_action_trajectory = self.precompute_full_action_trajectory( original_actions, num_segments, cond_frames_raw, total_frames_raw, extrapolation_strategy=extrapolation_strategy, ) # Initialize with original batch for first segment current_segment_batch = copy.deepcopy(data_batch) # STEP 2: Run segment-by-segment inference for segment_idx in range(num_segments): log.debug(f'Segment {segment_idx + 1}/{num_segments}...') # For segment 0, actions are already in data_batch # For segments 1+, actions were set in prepare_next_segment if segment_idx > 0 and full_action_trajectory is not None: # Verify segment actions match precomputed trajectory expected_actions = self.get_segment_actions( full_action_trajectory, segment_idx, cond_frames_raw, total_frames_raw ) actual_actions = current_segment_batch['a4d_latent']['action'] assert torch.allclose(expected_actions.float(), actual_actions.float()), \ f"Segment {segment_idx} action mismatch!" # Generate for current segment with misc.timer('pipe.generate()'): segment_val_dict = pipe.generate( current_segment_batch, seed=seed, num_sampling_steps=config.val_num_steps, guidance=config.val_cfg_scale, sigma_max=config.val_sigma_max, sigma_min=config.val_sigma_min, cfg_zero_entries=config.val_cfg_zero_entries, is_negative_prompt=False, phase='val', cond_aug_sigma=config.val_cond_aug_sigma, ) # Store prediction for later combination predictions_so_far.append(segment_val_dict) # Prepare next segment batch (if not last segment) if segment_idx < num_segments - 1: current_segment_batch = self.prepare_next_segment( current_segment_batch, segment_val_dict, cond_frames_latent, cond_frames_raw, total_frames_raw, total_frames_latent, full_action_trajectory, segment_idx ) # STEP 3: Combine all segments into final val_dict val_dict = self.combine_segments( predictions_so_far, cond_frames_latent, cond_frames_raw, full_action_trajectory, ) log.debug(f'[Autoregressive] Combined {num_segments} segments into val_dict') return val_dict, data_batch def update_masks_for_visualization( self, data_batch: dict, val_dict: dict, ) -> dict: """ Update data_batch masks to match extended temporal dimension for visualization. :param data_batch: Original data batch :param val_dict: Combined val_dict with extended predictions :return: Updated data_batch with extended masks """ # Update RGB masks for each view for view_idx in range(self.config.num_views): view_key = f'rgb{view_idx}' stream_key = f'v{view_idx}' if stream_key in val_dict['y0_pred_streams']: # Get the extended temporal dimension from combined streams T_extended = val_dict['y0_pred_streams'][stream_key].shape[2] # Get original mask properties orig_input_mask = data_batch['a4d_latent'][f'{view_key}_input_mask'] B, C_mask, T_orig, H, W = orig_input_mask.shape # Create extended masks data_batch['a4d_latent'][f'{view_key}_input_mask'] = torch.zeros( (B, C_mask, T_extended, H, W), device=orig_input_mask.device, dtype=orig_input_mask.dtype ) data_batch['a4d_latent'][f'{view_key}_input_mask'][:, :, :T_orig, :, :] = orig_input_mask data_batch['a4d_latent'][f'{view_key}_output_mask'] = torch.ones( (B, C_mask, T_extended, H, W), device=orig_input_mask.device, dtype=orig_input_mask.dtype ) data_batch['a4d_latent'][f'{view_key}_output_mask'][:, :, :T_orig, :, :] = 1.0 - orig_input_mask data_batch['a4d_latent'][f'{view_key}_supervise_mask'] = torch.ones( (B, C_mask, T_extended, H, W), device=orig_input_mask.device, dtype=orig_input_mask.dtype ) data_batch['a4d_latent'][f'{view_key}_supervise_mask'][:, :, :T_orig, :, :] = 1.0 - orig_input_mask # Update action masks if present if 'action' in data_batch['a4d_latent']: orig_action_mask = data_batch['a4d_latent']['action_input_mask'] B, T_orig, C_mask = orig_action_mask.shape total_action_frames = val_dict['x0_streams']['adaln'].shape[1] data_batch['a4d_latent']['action_input_mask'] = torch.ones( (B, total_action_frames, C_mask), device=orig_action_mask.device, dtype=orig_action_mask.dtype ) data_batch['a4d_latent']['action_output_mask'] = torch.zeros( (B, total_action_frames, C_mask), device=orig_action_mask.device, dtype=orig_action_mask.dtype ) data_batch['a4d_latent']['action_supervise_mask'] = torch.zeros( (B, total_action_frames, C_mask), device=orig_action_mask.device, dtype=orig_action_mask.dtype ) return data_batch