Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """HDMap-only diffusion loading, sampling, and output for native sensor grids.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import time | |
| from collections import Counter | |
| from copy import deepcopy | |
| from datetime import timedelta | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from tqdm import tqdm | |
| from utils import inference as common | |
| from utils.lidar_rangemap import unnormalize_and_reduce_channels | |
| # Dataset-ID order is a checkpoint contract, not inferred from grid dimensions. | |
| SENSORS = { | |
| 0: ("pandar", 128, 1024, 10.0), | |
| 1: ("pandaset", 64, 1024, 10.0), | |
| 2: ("nuscenes", 32, 1024, 20.0), | |
| } | |
| def resolve_checkpoint(path: Path) -> Path: | |
| if path.is_file(): | |
| return path.resolve() | |
| folder = path / "checkpoints" if (path / "checkpoints").is_dir() else path | |
| marker = folder / "latest_checkpoint.txt" | |
| stem = marker.read_text().strip().removesuffix(".pt") | |
| if not stem.startswith("iter_") or not stem[5:].isdigit(): | |
| raise ValueError(f"Invalid checkpoint marker in {marker}: {stem!r}") | |
| source = folder / f"{stem}_ema_model.pt" | |
| if not source.is_file() or source.stat().st_size == 0: | |
| raise FileNotFoundError(f"No complete EMA checkpoint: {source}") | |
| # Match the research launcher: pin the EMA before training retention unlinks it. | |
| snapshot_dir = folder.parent / "eval_snapshots" | |
| snapshot_dir.mkdir(parents=True, exist_ok=True) | |
| snapshot = snapshot_dir / source.name | |
| try: | |
| os.link(source, snapshot) | |
| except FileExistsError: | |
| pass | |
| return snapshot.resolve() | |
| def load_model(checkpoint: Path, tokenizer_checkpoint: Path, latent_stats: Path): | |
| from diffusion.training.config.base.lidarverse_vae import get_lidarverse_tokenizer | |
| from diffusion.training.config.base.model import LidarModelConfig | |
| from diffusion.training.config.text2world.registry import ( | |
| MININET_2B_Config, VideoConditionerFpsSizePaddingConfig, | |
| ) | |
| from diffusion.training.models.hdmaptolidar_model_rf import HdmapToLidarRectifiedFlowModel | |
| from utils.lazy_config import LazyDict | |
| net = deepcopy(MININET_2B_Config) | |
| net.update( | |
| in_channels=128, out_channels=64, patch_spatial=(1, 1), | |
| rope_scale_positions_by_patch=True, num_dataset_embeddings=3, | |
| timestep_scale=0.001, rope_h_extrapolation_ratio=3.0, | |
| rope_w_extrapolation_ratio=3.0, rope_t_extrapolation_ratio=1.0, | |
| rope_enable_fps_modulation=False, use_wan_fp32_strategy=True, | |
| ) | |
| net.sac_config.mode = "none" | |
| vae = get_lidarverse_tokenizer() | |
| vae.update(checkpoint_path=str(tokenizer_checkpoint), mean_std_fp=str(latent_stats), | |
| strict_checkpoint=True, use_ema=False) | |
| # No training experiment registration, dataloaders, FSDP, or extra EMA model. | |
| config = LidarModelConfig( | |
| net=net, vae=vae, conditioner=deepcopy(VideoConditionerFpsSizePaddingConfig), | |
| ema=LazyDict(dict(enabled=False)), latent_shape=[64, 21, 16, 128], | |
| hdmap_channels=3, additional_concat_ch=64, encode_hdmap_with_vae=True, | |
| hdmap_dropout_prob=0.1, sigma_data=1.0, shift=5, | |
| shift_clamp_percentiles=(0.5, 99.9), rf_prediction_type="velocity", | |
| loss_reduce="mean", input_video_has_prepend=True, | |
| ) | |
| model = HdmapToLidarRectifiedFlowModel(config) | |
| model.model.to(dtype=torch.bfloat16) | |
| # Local/trusted PyTorch checkpoints only. Load strictly: never adapt a head. | |
| payload = torch.load(checkpoint, map_location="cpu", weights_only=False) | |
| state = payload.get("model", payload) | |
| if isinstance(payload.get("ema"), dict): | |
| state = payload["ema"] | |
| model.model.load_state_dict(state, strict=True) | |
| del state, payload | |
| model.cuda().eval() | |
| model.vae.float() | |
| print(f"Loaded all diffusion weights strictly: {checkpoint}", flush=True) | |
| return model | |
| def load_hdmap_sample(path: Path, num_frames: int) -> dict: | |
| with np.load(path, allow_pickle=False) as payload: | |
| sample = {key: payload[key] for key in payload.files} | |
| for key in ("hdmap", "ray_directions", "sensor_dataset_id", "fps"): | |
| if key not in sample: | |
| raise ValueError(f"{path}: missing {key}; tokenizer-only NPZs have no HDMap condition") | |
| dataset_id = int(sample["sensor_dataset_id"]) | |
| if dataset_id not in SENSORS: | |
| raise ValueError(f"{path}: unknown sensor_dataset_id {dataset_id}") | |
| sensor, height, width, fps = SENSORS[dataset_id] | |
| hdmap = np.asarray(sample["hdmap"], dtype=np.float32) | |
| if hdmap.shape != (num_frames, height, width): | |
| raise ValueError(f"{path}: expected {num_frames} real {height}x{width} HDMap frames; got {hdmap.shape}") | |
| if bool(sample.get("has_prepend", False)): | |
| raise ValueError(f"{path}: store only real frames, not the duplicated VAE prefix") | |
| if str(sample.get("dataset", sensor)) != sensor or float(sample["fps"]) != fps: | |
| raise ValueError(f"{path}: sensor name/FPS does not match dataset ID {dataset_id}") | |
| if float(sample.get("min_range_m", 2.5)) != 2.5 or float(sample.get("max_range_m", 100.0)) != 100.0: | |
| raise ValueError(f"{path}: this checkpoint uses 2.5--100 m") | |
| if not np.isfinite(hdmap).all() or hdmap.min() < -1 or hdmap.max() > 1: | |
| raise ValueError(f"{path}: hdmap must be normalized to [-1,1], invalid cells=-1") | |
| rays = np.asarray(sample["ray_directions"], dtype=np.float32) | |
| if rays.shape not in ((height, width, 3), (num_frames, height, width, 3)): | |
| raise ValueError(f"{path}: rays do not match the native HDMap grid: {rays.shape}") | |
| if not np.isfinite(rays).all() or not np.allclose(np.linalg.norm(rays, axis=-1), 1, atol=1e-4): | |
| raise ValueError(f"{path}: ray_directions must be finite unit vectors") | |
| sample.update(hdmap=hdmap, ray_directions=rays, dataset=sensor, | |
| visualization_yaw_deg=common.display_yaw(path, sensor)) | |
| return sample | |
| def _metric_ranges(video): | |
| # Same validity cutoff (including near/far buffers) as the working launcher. | |
| ranges, valid = unnormalize_and_reduce_channels( | |
| video, max_range=100.0, min_range=2.5, min_value=-1.0, | |
| inverse_depth=False, process_output="mono", | |
| ) | |
| return ranges[0].float().cpu().numpy(), valid[0].cpu().numpy() | |
| def _preview_frames(sample, generated, valid, output): | |
| hdmap, hd_valid = _metric_ranges(torch.from_numpy(sample["hdmap"])[None, None]) | |
| rays = np.broadcast_to(sample["ray_directions"], hdmap.shape + (3,)) | |
| for index in range(len(generated)): | |
| panels = [] | |
| clouds = [("hdmap", "HDMap", hdmap, hd_valid), | |
| ("generation", "Generated", generated, valid)] | |
| for name, title, values, mask in clouds: | |
| points = common.points_from_ranges(values[index], mask[index], rays[index]) | |
| common.write_ply(output / "ply" / f"{name}_{index:03d}.ply", points) | |
| panels.append(common.label_image( | |
| common.render_points(points, yaw_deg=sample["visualization_yaw_deg"]), | |
| f"{sample['dataset']} | {title}", | |
| )) | |
| yield np.concatenate(panels, axis=1) | |
| def _generate_sequence(model, sample, args, seed, console, label): | |
| dataset_id = int(sample["sensor_dataset_id"]) | |
| torch.cuda.synchronize() | |
| phase_started = time.perf_counter() | |
| # The diffusion path sees HDMap + null text, never GT. | |
| video = torch.from_numpy(sample["hdmap"])[None, None].cuda() | |
| video = torch.cat([video[:, :, :1], video], dim=2) | |
| _, _, _, height, width = video.shape | |
| batch = { | |
| # The shared preparation API encodes a target to obtain its latent shape. | |
| # Supply an empty target, never GT. Its latent is not a sampling condition. | |
| "video": torch.full_like(video, -1), | |
| "lidar_hdmap": video, | |
| "is_preprocessed": True, | |
| "sensor_dataset_id": torch.tensor([dataset_id], device="cuda", dtype=torch.long), | |
| "fps": torch.tensor([float(sample["fps"])], device="cuda"), | |
| "num_frames": torch.tensor([args.sample_n_frames], device="cuda"), | |
| "image_size": torch.tensor([[height, width]], device="cuda"), | |
| "padding_mask": torch.zeros(1, 1, height, width, device="cuda"), | |
| "t5_text_embeddings": torch.zeros(1, 512, 1024, device="cuda"), | |
| "t5_text_mask": torch.zeros(1, 512, device="cuda"), | |
| } | |
| batch["t5_text_mask"][:, 0] = 1 | |
| # The VAE samples a Gaussian posterior even in eval mode. Seed preparation | |
| # too, retaining the reference target-then-HDMap encode/RNG order. | |
| torch.manual_seed(seed) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| _, shape_latent, _ = model.get_data_and_condition(batch) | |
| torch.cuda.synchronize() | |
| encode_seconds = time.perf_counter() - phase_started | |
| phase_started = time.perf_counter() | |
| # Both supported samplers make one velocity query per step. Ordinary | |
| # CFG makes two DiT forwards; CFG=0/1 makes one. The hook observes only | |
| # completion and MUST return None so it cannot replace network outputs. | |
| forwards_per_step = 1 if args.guidance in (0.0, 1.0) else 2 | |
| forwards = 0 | |
| with tqdm(total=args.num_steps, desc=f"{label} sampling", unit="step", | |
| file=console, dynamic_ncols=True, leave=False, | |
| disable=not console.isatty()) as progress: | |
| def on_forward(_module, _inputs, _output): | |
| nonlocal forwards | |
| forwards += 1 | |
| if forwards % forwards_per_step == 0: | |
| progress.update(1) | |
| hook = model.net.register_forward_hook(on_forward) | |
| try: | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| latent = model.generate_samples_from_batch( | |
| batch, state_shape=shape_latent.shape[1:], n_sample=1, | |
| guidance=args.guidance, seed=seed, num_steps=args.num_steps, | |
| sampler=args.sampler, sde_eta=args.sde_eta, | |
| ) | |
| torch.cuda.synchronize() | |
| finally: | |
| hook.remove() | |
| sampling_seconds = time.perf_counter() - phase_started | |
| phase_started = time.perf_counter() | |
| decoded = model.decode(latent.float()) | |
| if decoded.shape != video.shape or not torch.isfinite(decoded).all(): | |
| raise RuntimeError(f"Invalid decoded range video: {tuple(decoded.shape)}") | |
| generated, valid = _metric_ranges(decoded[:, :, 1:]) | |
| torch.cuda.synchronize() | |
| decode_seconds = time.perf_counter() - phase_started | |
| return generated, valid, { | |
| "encode_seconds": encode_seconds, | |
| "sampling_seconds": sampling_seconds, | |
| "decode_seconds": decode_seconds, | |
| } | |
| def run(args, console, say, started): | |
| sample_root, output_root = args.sample_path.resolve(), args.output_folder | |
| # Count selected clips without decompressing their range videos. Preserve | |
| # the original file order and per-sensor seed/index convention. | |
| samples, selected_counts = [], Counter() | |
| for path in common.sample_paths(sample_root): | |
| with np.load(path, allow_pickle=False) as payload: | |
| dataset_id = int(payload["sensor_dataset_id"]) | |
| if dataset_id not in SENSORS: | |
| raise ValueError(f"{path}: unknown sensor_dataset_id {dataset_id}") | |
| if selected_counts[dataset_id] < args.num_samples: | |
| samples.append(path) | |
| selected_counts[dataset_id] += 1 | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError("A CUDA GPU is required; use the diffusion-capable container") | |
| checkpoint = resolve_checkpoint(args.checkpoint.resolve()) | |
| model = load_model(checkpoint, args.tokenizer_checkpoint.resolve(), args.latent_stats.resolve()) | |
| torch.cuda.synchronize() | |
| startup_seconds = time.perf_counter() - started | |
| say(f"Ready in {startup_seconds:.1f}s | {checkpoint.name} | {len(samples)} clips, {args.sample_n_frames} frames each") | |
| metadata = {key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()} | |
| metadata.update(resolved_checkpoint=str(checkpoint), shift=5, dataset_ids={k: v[0] for k, v in SENSORS.items()}) | |
| (output_root / "run.json").write_text(json.dumps(metadata, indent=2) + "\n") | |
| counts = {key: 0 for key in SENSORS} | |
| sensor_seconds = Counter() | |
| records = [] | |
| timing_path = output_root / "timings.json" | |
| for clip_index, path in enumerate(samples, 1): | |
| clip_started = time.perf_counter() | |
| sample = load_hdmap_sample(path, args.sample_n_frames) | |
| dataset_id = int(sample["sensor_dataset_id"]) | |
| seed = args.seed + int(sample.get("sample_index", counts[dataset_id])) | |
| counts[dataset_id] += 1 | |
| label = f"[{clip_index:02d}/{len(samples):02d}] {sample['dataset']}/{path.stem}" | |
| say(f"{label} | encoding, {sample['hdmap'].shape[1]}x{sample['hdmap'].shape[2]}, seed {seed}") | |
| generated, valid, generation_timing = _generate_sequence( | |
| model, sample, args, seed, console, label | |
| ) | |
| encode_seconds = generation_timing["encode_seconds"] | |
| sampling_seconds = generation_timing["sampling_seconds"] | |
| decode_seconds = generation_timing["decode_seconds"] | |
| phase_started = time.perf_counter() | |
| output = common.sample_output(path, sample_root, output_root) | |
| output.mkdir(parents=True, exist_ok=True) | |
| saved = {key: sample[key] for key in ( | |
| "ray_directions", "fps", "sensor_dataset_id", "dataset", "clip_name", | |
| "lidar_to_world", "pose_sensor_start_end", "frame_indices", | |
| ) if key in sample} | |
| np.savez_compressed(output / "generation.npz", range_m=generated, valid_mask=valid, | |
| seed=np.int64(seed), **saved) | |
| save_seconds = time.perf_counter() - phase_started | |
| phase_started = time.perf_counter() | |
| if args.vis_pcd: | |
| with tqdm(_preview_frames(sample, generated, valid, output), total=len(generated), | |
| desc=f"{label} rendering", unit="frame", file=console, | |
| dynamic_ncols=True, leave=False, disable=not console.isatty()) as preview: | |
| common.write_video(output / "generation.mp4", preview, float(sample["fps"])) | |
| render_seconds = time.perf_counter() - phase_started if args.vis_pcd else 0.0 | |
| total_seconds = time.perf_counter() - clip_started | |
| sensor_seconds[dataset_id] += total_seconds | |
| remaining = selected_counts[dataset_id] - counts[dataset_id] | |
| eta = sensor_seconds[dataset_id] / counts[dataset_id] * remaining | |
| generation_seconds = encode_seconds + sampling_seconds + decode_seconds | |
| records.append(dict(sample=str(path), dataset=sample["dataset"], seed=seed, | |
| encode_seconds=encode_seconds, sampling_seconds=sampling_seconds, | |
| decode_seconds=decode_seconds, save_seconds=save_seconds, | |
| render_seconds=render_seconds, total_seconds=total_seconds, | |
| generation_fps=args.sample_n_frames / generation_seconds)) | |
| report = dict(startup_seconds=startup_seconds, elapsed_seconds=time.perf_counter() - started, | |
| completed_clips=len(records), selected_clips=len(samples), samples=records) | |
| timing_path.with_suffix(".json.tmp").write_text(json.dumps(report, indent=2) + "\n") | |
| timing_path.with_suffix(".json.tmp").replace(timing_path) | |
| say(f" Done {total_seconds:.1f}s | encode {encode_seconds:.1f}s | sample {sampling_seconds:.1f}s" | |
| f" | decode {decode_seconds:.1f}s | save/render {save_seconds + render_seconds:.1f}s" | |
| f" | {sample['dataset']} remaining ~{timedelta(seconds=round(eta))}") | |
| model._store_hdmap_latent(None) | |
| say(f"Finished {len(records)} clips in {timedelta(seconds=round(time.perf_counter() - started))} (including loading and rendering).") | |
| say(f"Outputs: {output_root}\nTimings: {timing_path}") | |