Spaces:
Running on Zero
Running on Zero
File size: 7,617 Bytes
dcc4594 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | 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()
|