"""Sampling from a trained Rectified Point Flow.""" import logging from pathlib import Path import os import warnings import time import numpy as np import hydra import lightning as L import torch from omegaconf import DictConfig from rectified_point_flow.utils import load_checkpoint_for_module, download_rap_checkpoint, print_eval_table from rectified_point_flow.visualizer import VisualizationCallback logger = logging.getLogger("Sample") warnings.filterwarnings("ignore", module="lightning") # ignore warning from lightning' connectors warnings.filterwarnings("ignore", category=FutureWarning) # Suppress Lightning's device info (GPU/TPU available) and other verbose output logging.getLogger("lightning.pytorch").setLevel(logging.WARNING) logging.getLogger("lightning_fabric").setLevel(logging.WARNING) # Optimize for performance torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.backends.cudnn.benchmark = True DEFAULT_CKPT_PATH_HF = "rap_model.ckpt" def get_time(): """ :return: get timing statistics """ cuda_available = torch.cuda.is_available() if cuda_available: torch.cuda.synchronize() return time.time() def setup(cfg: DictConfig): """Setup evaluation components.""" ckpt_path = cfg.get("ckpt_path", None) if ckpt_path is None: ckpt_path = download_rap_checkpoint(DEFAULT_CKPT_PATH_HF, './weights') elif not os.path.exists(ckpt_path): logger.error(f"Checkpoint not found: {ckpt_path}") logger.error("Please provide a valid checkpoint in the config or via ckpt_path='...' argument") exit(1) seed = cfg.get("seed", None) if seed is not None: L.seed_everything(seed, workers=True, verbose=False) logger.debug(f"Seed set to {seed} for sampling") datamodule: L.LightningDataModule = hydra.utils.instantiate(cfg.data) model = hydra.utils.instantiate(cfg.model) load_checkpoint_for_module(model, ckpt_path) model.eval() vis_config = cfg.get("visualizer", {}) callbacks = [] if vis_config: vis_callback: VisualizationCallback = hydra.utils.instantiate(vis_config) callbacks.append(vis_callback) trainer: L.Trainer = hydra.utils.instantiate( cfg.trainer, callbacks=callbacks, enable_checkpointing=False, logger=False, ) return model, datamodule, trainer @hydra.main(version_base="1.3", config_path="./config", config_name="RAP_base") def main(cfg: DictConfig): """Entry point for evaluating the model on validation set. Visualization Options: - To disable all visualization saving: visualizer.renderer=none - To limit visualization to first N samples per batch: visualizer.max_samples_per_batch=N - To combine both: visualizer.renderer=none visualizer.max_samples_per_batch=10 Examples: - Evaluation only (no visualizations): python sample.py visualizer.renderer=none - First 5 samples only: python sample.py visualizer.max_samples_per_batch=5 - First 10 samples with no saving: python sample.py visualizer.renderer=none visualizer.max_samples_per_batch=10 """ model, datamodule, trainer = setup(cfg) # Add timing callback to track inference times class TimingCallback(L.Callback): def __init__(self): self.inference_times = [] self.generation_times = [] self.batch_sizes = [] # Track batch sizes for per-sample calculations self.generation_batch_sizes = [] # Track batch sizes for each generation time def _get_batch_size(self, batch): """Get the batch size from the batch data.""" if isinstance(batch, dict): # For dynamic batching, check if we have cu_seqlens_batch if "cu_seqlens_batch" in batch: cu_seqlens_batch = batch["cu_seqlens_batch"] batch_size = len(cu_seqlens_batch) - 1 # cu_seqlens_batch has shape (B + 1,) return batch_size # For fixed batching, use points_per_part shape elif "points_per_part" in batch: return batch["points_per_part"].shape[0] else: # Fallback: try to get batch size from any tensor in the batch for key, value in batch.items(): if isinstance(value, torch.Tensor) and value.ndim >= 1: return value.shape[0] return 1 # Default fallback def on_test_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx=0): # Record start time for this batch self.batch_start_time = get_time() def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0): # Record end time and calculate inference time batch_end_time = get_time() inference_time = batch_end_time - self.batch_start_time self.inference_times.append(inference_time) # Get dataset name from dataloader index dataset_names = getattr(trainer.datamodule, 'dataset_names', []) dataset_name = dataset_names[dataloader_idx] if dataloader_idx < len(dataset_names) else f"dataset_{dataloader_idx}" # Get batch size for per-sample timing calculations batch_size = self._get_batch_size(batch) self.batch_sizes.append(batch_size) # Extract generation times and overlap ratios if outputs: # Handle generation times if 'generation_times' in outputs: gen_times = outputs['generation_times'] self.generation_times.extend(gen_times) # Track batch size for each generation time self.generation_batch_sizes.extend([batch_size] * len(gen_times)) # Calculate per-sample generation times per_sample_gen_times = [t / batch_size for t in gen_times] # Print generation times for this batch if len(gen_times) > 1: avg_gen_time = np.mean(gen_times) avg_per_sample_gen_time = np.mean(per_sample_gen_times) logger.info(f"Test Dataloader {dataloader_idx} ({dataset_name}) - Batch {batch_idx}: Generation times = {[f'{t:.4f}s' for t in gen_times]}") logger.info(f"Test Dataloader {dataloader_idx} ({dataset_name}) - Batch {batch_idx}: Per-sample generation times = {[f'{t:.4f}s' for t in per_sample_gen_times]}") logger.info(f"Test Dataloader {dataloader_idx} ({dataset_name}) - Batch {batch_idx}: Average generation time = {avg_gen_time:.4f}s") logger.info(f"Test Dataloader {dataloader_idx} ({dataset_name}) - Batch {batch_idx}: Average per-sample generation time = {avg_per_sample_gen_time:.4f}s") else: logger.info(f"Test Dataloader {dataloader_idx} ({dataset_name}) - Batch {batch_idx}: Generation time = {gen_times[0]:.4f}s") logger.info(f"Test Dataloader {dataloader_idx} ({dataset_name}) - Batch {batch_idx}: Per-sample generation time = {per_sample_gen_times[0]:.4f}s") # Calculate per-sample inference time (but don't log by default) per_sample_inference_time = inference_time / batch_size def on_test_end(self, trainer, pl_module): if self.inference_times and self.batch_sizes: # Calculate batch-level statistics avg_time = np.mean(self.inference_times) std_time = np.std(self.inference_times) total_time = sum(self.inference_times) # Calculate per-sample statistics per_sample_times = [inf_time / batch_size for inf_time, batch_size in zip(self.inference_times, self.batch_sizes)] avg_per_sample_time = np.mean(per_sample_times) std_per_sample_time = np.std(per_sample_times) total_samples = sum(self.batch_sizes) # Inference time summary removed by default # logger.info(f"=== INFERENCE TIME (include IO, visualization, etc.) SUMMARY ===") # logger.info(f"Average inference time per batch: {avg_time:.4f}s ± {std_time:.4f}s") # logger.info(f"Average inference time per sample: {avg_per_sample_time:.4f}s ± {std_per_sample_time:.4f}s") # logger.info(f"Total inference time: {total_time:.4f}s") # logger.info(f"Total samples processed: {total_samples}") if self.generation_times and self.generation_batch_sizes: # Calculate batch-level statistics avg_gen_time = np.mean(self.generation_times) std_gen_time = np.std(self.generation_times) total_gen_time = sum(self.generation_times) # Calculate per-sample statistics using the correct batch sizes for each generation per_sample_gen_times = [gen_time / batch_size for gen_time, batch_size in zip(self.generation_times, self.generation_batch_sizes)] avg_per_sample_gen_time = np.mean(per_sample_gen_times) std_per_sample_gen_time = np.std(per_sample_gen_times) logger.info(f"=== GENERATION TIME SUMMARY ===") logger.info(f"Average generation time per batch: {avg_gen_time:.4f}s ± {std_gen_time:.4f}s") logger.info(f"Average generation time per sample: {avg_per_sample_gen_time:.4f}s ± {std_per_sample_gen_time:.4f}s") logger.info(f"Total generation time: {total_gen_time:.4f}s") # Add timing callback to trainer timing_callback = TimingCallback() trainer.callbacks.append(timing_callback) eval_results = trainer.test( model=model, datamodule=datamodule, verbose=False, ) # Extract sample counts from model and convert to list in dataset order sample_counts = [] part_count_ranges = [] for dataset_name in datamodule.dataset_names: count = model.last_sample_counts.get(dataset_name, 0) sample_counts.append(count) part_range = model.last_part_count_ranges.get(dataset_name, (0, 0)) part_count_ranges.append(part_range) print_eval_table(eval_results, datamodule.dataset_names, sample_counts=sample_counts, part_count_ranges=part_count_ranges) logger.info("Visualizations saved to:" + str(Path(cfg.get('log_dir')) / "visualizations")) logger.info("Evaluation results saved to:" + str(Path(cfg.get('log_dir')) / "results")) if __name__ == "__main__": main()