(video_tokenizer_path, latent_actions_path, dynamics_path, device, use_actions=True)
| 8 | from einops import repeat |
| 9 | |
| 10 | def load_models(video_tokenizer_path, latent_actions_path, dynamics_path, device, use_actions=True): |
| 11 | # Load tokenizer and dynamics, and Latent Actions if using actions |
| 12 | video_tokenizer, _vt_ckpt = load_videotokenizer_from_checkpoint(video_tokenizer_path, device) |
| 13 | video_tokenizer.eval() |
| 14 | latent_action_model = None |
| 15 | if use_actions: |
| 16 | latent_action_model, _latent_action_ckpt = load_latent_actions_from_checkpoint(latent_actions_path, device) |
| 17 | latent_action_model.eval() |
| 18 | dynamics_model, _dyn_ckpt = load_dynamics_from_checkpoint(dynamics_path, device) |
| 19 | dynamics_model.eval() |
| 20 | return video_tokenizer, latent_action_model, dynamics_model |
| 21 | |
| 22 | |
| 23 | def visualize_inference(predicted_frames, ground_truth_frames, inferred_actions, fps, use_actions=True): |
no test coverage detected