| """Load checkpoint and run autoregressive pose generation.""" |
|
|
| from __future__ import annotations |
|
|
| import os |
| from typing import Tuple |
|
|
| import numpy as np |
| import torch |
|
|
| from audio_features import CONTEXT_LEN |
| from model import Music2PoseTransformer |
|
|
| DEFAULT_CHECKPOINT = "pytorch_model.pt" |
|
|
|
|
| def load_checkpoint( |
| ckpt_path: str = DEFAULT_CHECKPOINT, |
| device: torch.device | None = None, |
| ) -> Tuple[Music2PoseTransformer, dict]: |
| """Load model weights and normalisation stats from a checkpoint file.""" |
| if device is None: |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| ckpt = torch.load(ckpt_path, map_location=device, weights_only=False) |
| model = Music2PoseTransformer( |
| audio_dim=ckpt["audio_dim"], |
| pose_dim=ckpt["pose_dim"], |
| **ckpt["config"], |
| ).to(device) |
| model.load_state_dict(ckpt["model_state"]) |
| model.eval() |
| return model, ckpt |
|
|
|
|
| def generate_poses( |
| model: Music2PoseTransformer, |
| audio_feat: np.ndarray, |
| x_mean: np.ndarray, |
| x_std: np.ndarray, |
| y_mean: np.ndarray, |
| y_std: np.ndarray, |
| device: torch.device, |
| ) -> np.ndarray: |
| """ |
| Autoregressive inference. |
| |
| Returns |
| ------- |
| poses_xyz : (T, 33, 3) normalised MediaPipe-style joint positions. |
| """ |
| T = audio_feat.shape[0] |
| af_norm = (audio_feat - x_mean) / x_std |
| af_t = torch.tensor(af_norm, dtype=torch.float32, device=device) |
| pose_dim = model.pose_dim |
|
|
| audio_ctx = torch.zeros(1, CONTEXT_LEN, af_t.shape[-1], device=device) |
| pose_ctx = torch.zeros(1, CONTEXT_LEN, pose_dim, device=device) |
| preds = [] |
|
|
| with torch.no_grad(): |
| for t in range(T): |
| audio_ctx = torch.roll(audio_ctx, -1, dims=1) |
| audio_ctx[0, -1] = af_t[t] |
| next_pose = model.step(audio_ctx, pose_ctx)[0] |
| preds.append(next_pose.cpu().numpy()) |
| pose_ctx = torch.roll(pose_ctx, -1, dims=1) |
| pose_ctx[0, -1] = next_pose |
|
|
| pred_norm = np.stack(preds) |
| pred_raw = pred_norm * y_std + y_mean |
| return pred_raw.reshape(T, 33, 9)[:, :, :3] |
|
|
|
|
| def resolve_checkpoint(path: str | None = None) -> str: |
| """Return checkpoint path, defaulting to the bundled weights.""" |
| if path: |
| return path |
| if os.path.exists(DEFAULT_CHECKPOINT): |
| return DEFAULT_CHECKPOINT |
| raise FileNotFoundError( |
| f"No checkpoint found. Download weights or pass --checkpoint explicitly." |
| ) |
|
|
