# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Causal P1B4 sampling with clean K/V commits and streaming spatial memory.""" import torch from tqdm import tqdm from diffusion.inference.hdmap import _metric_ranges from diffusion.inference.pai import SPATIAL_CACHE, make_batch, mode_spec from diffusion.training.datasets.lidar_dataset.forecast_warp import TorchSpatialPointMap from utils import misc def _spatial_map(model, batch): rays, elev, azim, row_off = model._stage2_spatial_map_geometry(batch["video"]) policy = {key: getattr(model.config, f"blockwise_spatial_map_{key}") for key in SPATIAL_CACHE} policy["voxel_size_m"] = policy.pop("voxel_m") return TorchSpatialPointMap( ray_dirs=rays, row_elevations_rad=elev, column_azimuths_rad=azim, row_azimuth_offsets_rad=row_off, device=rays.device, min_range=2.5, max_range=100, min_value=-1, inverse_depth=False, **policy, ) def _commit(model, latent, condition): model._current_sfdmd_store_self_attn_kv = True try: velocity = model._rf_network_forward( latent.to(**model.tensor_kwargs), torch.zeros((1, latent.shape[2]), device=latent.device), condition, ) if not torch.isfinite(velocity).all(): raise FloatingPointError("Clean K/V commit produced non-finite values") finally: model._current_sfdmd_store_self_attn_kv = False def generate(model, sample, mode, args, seed, console): spec = mode_spec(mode, args.stage) length = spec["input_frames"] + spec["output_frames"] # Complete P1+4N blocks using the extra supplied pose, then crop for display. padded_length = 1 + ((length - 1 + 3) // 4) * 4 batch = make_batch(sample, mode, args.stage, length=padded_length) video = batch["video"] known = torch.arange(padded_length, device=video.device) < spec["input_frames"] torch.manual_seed(seed) model._current_stage2_data_batch = batch model._current_stage2_geometry_source_mask = known model._current_sensor_dataset_ids = batch["sensor_dataset_id"] model._current_stage2_geometry_mode = "frozen_object" if mode == "resim" else "rolling_static" model.factorized_static_only_geometry = False model._hdmap_force_zero_latent = spec["hdmap"] == "null" model._current_sfdmd_self_attn_kv_cache = model.new_teacher_forcing_kv_cache() model._current_sfdmd_self_attn_kv_max_frames = args.kv_cache_frames model._current_sfdmd_self_attn_kv_sink_frames = 1 model._current_sfdmd_store_self_attn_kv = False model._current_sfdmd_detach_self_attn_kv = True model._current_sfdmd_return_self_attn_kv = False blocks = [(0, 1)] + [(start, start + 4) for start in range(1, padded_length, 4)] spatial_map = _spatial_map(model, batch) if mode != "resim" else None frozen = None if mode == "resim": # The training builder reads points only from the observed prefix. # Target actor poses place the observed canonical object clouds. frozen, _, _ = model._build_torch_spatial_map_history(video, batch["lidar_to_world"]) elif mode == "forecast": # Prefill all supplied observations, then append only newly generated sweeps. for start, end in blocks: end = min(end, spec["input_frames"]) if start >= end: break spatial_map.append_token_block( video[0, :, start:end], batch["lidar_to_world"][0, start:end], start, input_has_prepend=False, token_pose_sensor_start_end=batch["pose_sensor_start_end"][0, start:end], ) observed = torch.zeros((1, 64, padded_length, 16, 128), device=video.device) observed_cache = None for t in range(spec["input_frames"]): latent, observed_cache = model.encode_stream( video[:, :, t:t + 1].float(), temporal_kv_cache=observed_cache, max_cache_frames=9, ) if not torch.isfinite(latent).all(): raise FloatingPointError(f"Non-finite observed latent at frame {t}") observed[:, :, t:t + 1] = latent.float() del observed_cache condition_cache = decoder_cache = None decoded_blocks = [] try: for block_index, (start, end) in enumerate(tqdm( blocks, desc=f"{mode} causal sampling", unit="block", file=console, leave=False, disable=not console.isatty(), )): poses = batch["lidar_to_world"][0, start:end] motion = batch["pose_sensor_start_end"][0, start:end] if spatial_map is not None: geometry, _ = spatial_map.render(poses, target_pose_sensor_start_end=motion) geometry = geometry[None] else: geometry = frozen[:, :, start:end] active = dict(batch) active.update( video=torch.full_like(video[:, :, start:end], -1), lidar_hdmap=batch["lidar_hdmap"][:, :, start:end], lidar_blockwise_warped_history=geometry, num_frames=batch["num_frames"].new_tensor([end - start]), ) model._current_sfdmd_temporal_positions = torch.arange(start, end, device=video.device) model._current_stage2_known_frame_mask = known[start:end] with torch.autocast("cuda", dtype=torch.bfloat16): condition, condition_cache = model.get_online_condition_stream( active, temporal_kv_cache=condition_cache, max_cache_frames=9, ) block_known = known[start:end] known_latents = observed[:, :, start:end] if bool(block_known.all()): sampled = known_latents else: noise = misc.arch_invariant_rand( (1, 64, end - start, 16, 128), torch.float32, "cuda", seed + block_index, ) sampled = model._generate_sde( latents=noise, condition=condition, uncondition=None, guidance=1.0, use_cfg=False, B=1, num_steps=args.num_steps, shift=5, seed=seed + block_index, sde_eta=args.sde_eta, known_latents=known_latents if bool(block_known.any()) else None, known_mask=block_known if bool(block_known.any()) else None, known_latents_mode="clean_mask", ) if not torch.isfinite(sampled).all(): raise FloatingPointError(f"Non-finite latents in block {block_index}") if bool(block_known.any()) and not torch.equal( sampled[:, :, block_known].float(), known_latents[:, :, block_known].float(), ): raise RuntimeError(f"Clean-prefix replacement failed in block {block_index}") # Commit once after denoising, synchronously, before the next block. _commit(model, sampled, condition) decoded, decoder_cache = model.decode_stream( sampled.float(), temporal_kv_cache=decoder_cache, max_cache_frames=30, ) if decoded.shape != video[:, :, start:end].shape or not torch.isfinite(decoded).all(): raise RuntimeError(f"Invalid decoded LiDAR in block {block_index}") decoded = torch.where(block_known[None, None, :, None, None], video[:, :, start:end], decoded) decoded_blocks.append(decoded) new_start = max(start, spec["input_frames"]) if spatial_map is not None and new_start < end: # No current-block or future GT points can enter its own condition. offset = new_start - start spatial_map.append_token_block( decoded[0, :, offset:], poses[offset:], new_start, input_has_prepend=False, token_pose_sensor_start_end=motion[offset:], ) finally: model._current_sfdmd_temporal_positions = None model._current_sfdmd_self_attn_kv_cache = None model._current_sfdmd_store_self_attn_kv = False model._current_stage2_data_batch = None return _metric_ranges(torch.cat(decoded_blocks, dim=2)[:, :, :length])