Spaces:
Running on Zero
Running on Zero
| import torch | |
| from datasets.data_utils import load_data_and_data_loaders | |
| import matplotlib.pyplot as plt | |
| import time | |
| import os | |
| import random | |
| import glob | |
| import re | |
| from utils.utils import load_videotokenizer_from_checkpoint, load_latent_actions_from_checkpoint, load_dynamics_from_checkpoint, find_latest_checkpoint | |
| from utils.config import InferenceConfig, load_config | |
| from utils.inference_utils import load_models, visualize_inference, sample_random_action, get_action_latent | |
| from einops import repeat | |
| from typing import Optional | |
| def main(): | |
| # load inference config | |
| args: InferenceConfig = load_config(InferenceConfig, default_config_path=os.path.join(os.getcwd(), 'configs', 'inference.yaml')) | |
| # enable tf32 if requested | |
| if args.tf32: | |
| torch.backends.cuda.matmul.allow_tf32 = True | |
| torch.backends.cudnn.allow_tf32 = True | |
| # whether any setting requires using action tokens | |
| use_latent_actions = (args.use_actions or args.use_gt_actions or args.use_interactive_mode) | |
| # check if any path is missing | |
| def missing(path: Optional[str]) -> bool: | |
| return (path is None) or (not os.path.exists(path)) | |
| # resolve latest checkpoints if requested or any path missing | |
| base_dir = os.getcwd() | |
| if args.use_latest_checkpoints or missing(args.video_tokenizer_path): | |
| vt_ckpt = find_latest_checkpoint(base_dir, "video_tokenizer") | |
| args.video_tokenizer_path = vt_ckpt | |
| if (args.use_latest_checkpoints or missing(args.latent_actions_path)) and use_latent_actions: | |
| lam_ckpt = find_latest_checkpoint(base_dir, "latent_actions") | |
| args.latent_actions_path = lam_ckpt | |
| if args.use_latest_checkpoints or missing(args.dynamics_path): | |
| dyn_ckpt = find_latest_checkpoint(base_dir, "dynamics") | |
| args.dynamics_path = dyn_ckpt | |
| # confirm which ckpts are being used | |
| print(f"Using video_tokenizer checkpoint: {args.video_tokenizer_path}") | |
| if use_latent_actions: | |
| print(f"Using latent_actions checkpoint: {args.latent_actions_path}") | |
| print(f"Using dynamics checkpoint: {args.dynamics_path}") | |
| # validate required paths | |
| if missing(args.video_tokenizer_path): | |
| raise FileNotFoundError("video_tokenizer_path is not set or not a file. Set it in configs/inference.yaml or enable use_latest_checkpoints with available runs.") | |
| if use_latent_actions and missing(args.latent_actions_path): | |
| raise FileNotFoundError("latent_actions_path is not set or not a file while actions are requested. Set it in configs/inference.yaml or enable use_latest_checkpoints.") | |
| if missing(args.dynamics_path): | |
| raise FileNotFoundError("dynamics_path is not set or not a file. Set it in configs/inference.yaml or enable use_latest_checkpoints.") | |
| # load models, optionally compile | |
| video_tokenizer, latent_action_model, dynamics_model = load_models(args.video_tokenizer_path, args.latent_actions_path, args.dynamics_path, args.device, use_actions=use_latent_actions) | |
| if args.compile: | |
| video_tokenizer = torch.compile(video_tokenizer, mode="reduce-overhead", fullgraph=False, dynamic=True) | |
| if use_latent_actions: | |
| latent_action_model = torch.compile(latent_action_model, mode="reduce-overhead", fullgraph=False, dynamic=True) | |
| dynamics_model = torch.compile(dynamics_model, mode="reduce-overhead", fullgraph=False, dynamic=True) | |
| print("Compiled all models for inference.") | |
| # determine how many ground-truth frames we need in each batch: context + generation steps + prediction horizon | |
| frames_to_load = args.context_window + args.generation_steps * args.prediction_horizon | |
| # dataloader | |
| if hasattr(args, 'preload_ratio') and args.preload_ratio is not None: | |
| data_overrides = {'preload_ratio': args.preload_ratio} | |
| else: | |
| data_overrides = {} | |
| _, _, data_loader, _, _ = load_data_and_data_loaders( | |
| dataset=args.dataset, batch_size=1, num_frames=frames_to_load, **data_overrides) | |
| # sample random batch | |
| random_idx = random.randint(0, len(data_loader.dataset) - 1) | |
| og_ground_truth_frames = data_loader.dataset[random_idx][0] # full sequence | |
| og_ground_truth_frames = og_ground_truth_frames.unsqueeze(0).to(args.device) # [1, frames_to_load, C, H, W] | |
| ground_truth_frames = og_ground_truth_frames[:, :frames_to_load, :, :, :] # [1, frames_to_load, C, H, W] | |
| # start with initial context (first context_window GT frames) | |
| context_frames = ground_truth_frames[:, :args.context_window, :, :, :] | |
| generated_frames = context_frames.clone() | |
| # initialize actions | |
| n_actions = None | |
| inferred_actions = [] | |
| if use_latent_actions: | |
| n_actions = latent_action_model.quantizer.codebook_size | |
| # ensure we don’t exceed available GT frames if in teacher-forced mode | |
| max_possible_steps = ground_truth_frames.shape[1] - args.context_window | |
| if args.teacher_forced and args.generation_steps > max_possible_steps: | |
| print(f"[WARN] Requested {args.generation_steps} generation steps but only {max_possible_steps} are possible with teacher-forced context. Clamping.") | |
| effective_steps = args.generation_steps if not args.teacher_forced else min(args.generation_steps, max_possible_steps) | |
| for i in range(effective_steps): | |
| print(f"Inferring frame {i+1}/{effective_steps}") | |
| # select context depending on teacher-forced flag | |
| if args.teacher_forced: | |
| context_start = i # shift window along ground truth | |
| context_frames = ground_truth_frames[:, context_start:context_start+args.context_window, :, :, :] # [1, context_window, C, H, W] | |
| else: | |
| # autoregressive: last context_window frames from generated sequence | |
| context_frames = generated_frames[:, -args.context_window:, :, :, :] # [1, context_window, C, H, W] | |
| # encode context frames each iteration | |
| video_indices = video_tokenizer.tokenize(context_frames) | |
| video_latents = video_tokenizer.quantizer.get_latents_from_indices(video_indices) | |
| sampled_action_index, action_latent = get_action_latent(args, inferred_actions, n_actions, context_frames, latent_action_model, i) | |
| # dynamics forward inference needs idx -> latents fun | |
| def idx_to_latents(idx): | |
| return video_tokenizer.quantizer.get_latents_from_indices(idx, dim=-1) | |
| # autocast for inference if amp enabled (bfloat16 on CUDA by default) | |
| autocast_dtype = torch.bfloat16 if args.amp else None | |
| with torch.amp.autocast('cuda', enabled=args.amp, dtype=autocast_dtype): | |
| next_video_latents = dynamics_model.forward_inference( | |
| context_latents=video_latents, | |
| prediction_horizon=args.prediction_horizon, | |
| num_steps=10, | |
| index_to_latents_fn=idx_to_latents, | |
| conditioning=action_latent, | |
| temperature=args.temperature, | |
| ) | |
| # decode next video tokens to frames | |
| next_frames = video_tokenizer.detokenize(next_video_latents) # [1, T, C, H, W] | |
| generated_frames = torch.cat([generated_frames, next_frames[:, -args.prediction_horizon:, :, :]], dim=1) | |
| # TODO: if using interactive mode, visualize next_frames[:, -1] (recently inferred frame) every time, probably with matplotlib is easiest | |
| # point is for user to be able to interact with it in real time | |
| # visualize inference | |
| visualize_inference(generated_frames, ground_truth_frames, inferred_actions, args.fps, use_actions=use_latent_actions) | |
| if __name__ == "__main__": | |
| main() | |