File size: 2,449 Bytes
50ee618 | 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 | """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."
)
|