tarn59/UniMate-Weights

UniMate uniml3d_60frames_graph_adaln — independently trained checkpoint

A from-scratch training run of the main UniMate model ("UniMate: One Unified Model to Animate Diverse Skeletons", Mou et al., SIGGRAPH Asia 2026, arXiv 2609.05415) on a single GPU, using the authors' released code and the shipped configs/uniml3d_60frames_graph_adaln.json config unchanged.

This is not the authors' checkpoint. It was trained separately, by a third party, on one workstation GPU. The paper's example launch uses 8 GPUs, which at the same per-GPU batch size would see 8× more samples over the same 120k steps. Expect this model to differ from the official release and from any other independent run: different hardware, precision, sampling order and random seed all change the final weights. Numbers reported in the paper do not apply to this checkpoint unless you evaluate it yourself.

Which file to use

File	Use it for	Size
model_ema.safetensors	Inference. The exponential-moving-average weights, which the UniMate sampler loads by default and which the authors use for all results.	283 MB
model.safetensors	The raw (non-EMA) weights at step 120,000. Slightly noisier than the EMA copy; mainly useful for fine-tuning or comparison.	283 MB
checkpoints/checkpoint_step_120000.pt	Resuming or continuing training with --resume. Contains model weights, EMA shadow weights, AdamW state, LR-scheduler state and the step counter.	1.13 GB
config.json	The resolved run config, including the auto-computed max_joints=61 / max_depth=19. UniMate's inference scripts read this to rebuild the exact architecture.
dataset_stats.npy	Per-channel normalization statistics computed from the training data. Required at inference.
train_config_as_launched.json	The config file exactly as passed to accelerate launch, for reference.
logs/events.out.tfevents.*	TensorBoard scalars for the whole run (loss terms, grad norm, LR, EMA decay every 50 steps).
training_curves.png	The plot below.

Running inference with the UniMate code

UniMate's sampler expects an experiment directory laid out the way training writes it, and this repository mirrors that layout. Clone or download it, then point --exp_dir at it:

git clone https://github.com/Friedrich-M/UniMate && cd UniMate
# set up the environment per the UniMate README, and build dataset/features/ with stage 4 of
# its data pipeline (the sampler reads the target skeleton's T-pose and topology from there)
hf download tarn59/UniMate-Weights --local-dir outputs/uniml3d_60frames_graph_adaln

python -m unimate.inference.sample \
    --exp_dir outputs/uniml3d_60frames_graph_adaln \
    --test_cases_json test_cases.json \
    --num_repetitions 3

The sampler picks the latest file in checkpoints/ (here the single step-120000 checkpoint) and swaps in its EMA weights automatically, so nothing else needs to be passed.

Loading the safetensors directly

The .pt checkpoint is the format UniMate's scripts read; the safetensors files are provided as a smaller, framework-neutral copy of the weights. To load one into the model class:

import torch
from safetensors.torch import load_model
from unimate.configs.schema import MainConfig
from unimate.models.factory import create_model

cfg = MainConfig.from_json("outputs/uniml3d_60frames_graph_adaln/config.json")
model = create_model(cfg.dataset, cfg.model)
load_model(model, "outputs/uniml3d_60frames_graph_adaln/model_ema.safetensors", strict=True)
model.eval()

Use safetensors.torch.load_model, not load_file + load_state_dict: the spectral-RoPE encoder is one module shared by all ten transformer blocks, so state_dict() lists it under eleven names. The safetensors files store it once, under rope_j.*, and load_model re-ties the aliases.

Model

74.1 M parameters. A factored spatio-temporal diffusion transformer: 10 blocks, 512-wide latent, 8 heads, SwiGLU MLP (2048 nominal), RMSNorm with QK-norm, adaLN-Zero conditioning. Spatial attention runs within each frame across joints with Graphormer-style graph-distance and edge-type biases and a sign-invariant spectral RoPE over 8 Laplacian eigenvectors; temporal attention runs within each joint across frames with 1-D RoPE. The caption (google/flan-t5-base, mean-pooled) and the pooled T-pose enter through the adaLN vector. Trained with flow matching (linear interpolant, velocity prediction) under a masked L2 loss plus a geodesic rotation loss (weight 0.5), with 10 % caption dropout for classifier-free guidance. See the UniMate repository for the full description.

Training run

Config	configs/uniml3d_60frames_graph_adaln.json, unmodified
Data	UniML3D export (Linzhan/UniML3D) → UniMate stage-4 features. Truebones ZOO motions rebuilt locally from a licensed copy.
Training clips	10,204 from 5,756 object types: Objaverse 7,308 (5,698 rigs), Mixamo 2,169 (1 rig, 22-joint core), Truebones 727 (57 species)
Filtering	rigs outside 5–60 joints dropped (17 Truebones species, 781 Objaverse rigs); 34 Truebones clips skipped at load for a representation-recovery mismatch; KingCobra-Walk absent (its FBX needs the authors' one-off mesh-graft repair)
Eval split	none (test_split_ratio = 0)
Hardware	1× NVIDIA RTX PRO 6000 Blackwell Workstation Edition (96 GB), driver 580.173
Software	PyTorch 2.7.1+cu128 (the repo pins 2.5.1+cu124, which has no sm_120 kernels), Accelerate 1.14.0, Transformers 5.16.1
Precision	fp32 (--mixed_precision no), single process, no gradient accumulation
Batch / steps	16 × 120,000 optimizer steps = 1.92 M samples ≈ 188 passes over the training set
Optimizer	AdamW, lr 1e-4, betas (0.9, 0.99), weight decay 1e-5, grad-clip 1.0
Schedule	cosine with 3 % warm-up (3,600 steps) to 5 % of peak
EMA	decay 0.9999 with warm-up (0.98 at step 0, reaches ≈0.9999 near 90k)
Sampler / augmentation	balanced per-object-type sampler (α = 0.5); joint addition, leaf removal, chain pooling, bone-length perturbation
Seed	10
Throughput	1.16 steps/s (≈0.86 s/step)
Wall-clock	28 h 47 min (2026-09-17 19:45 → 2026-09-19 00:33 local)
Peak GPU memory	≈61 GB

Loss

Window (steps)	total loss	velocity L2	geodesic	pre-clip grad norm (mean)
0 – 10k	0.515	0.440	0.149	4.6
10k – 20k	0.242	0.184	0.115	0.69
20k – 40k	0.212	0.160	0.105	0.47
40k – 60k	0.185 †	0.137 †	0.098	0.40 †
60k – 80k	0.172 †	0.126 †	0.091	0.35 †
80k – 100k	0.160	0.116	0.088	0.38
100k – 120k	0.163 †	0.120 †	0.084	0.58 †
final 2k steps, median	0.141

† excluding single-step spikes. Four isolated spikes occurred during the run (steps 43,600, 78,100, 106,900 and 112,900; the largest reached loss 91 with a pre-clip gradient norm of 588). Each recovered to baseline by the next logged step. They match the README's description of defective Objaverse rigs and clips; gradient clipping to norm 1.0 bounds their effect on the weights, and the EMA averages them out. The trainer's non-finite-loss guard never fired.

Sample visualizations (EMA weights, CFG scale 3.0) were rendered at every 10k-step checkpoint and show coherent poses for humanoid, quadruped, insect and rigged-object skeletons. No quantitative evaluation has been run on this checkpoint.

Example samples

Four of the visualizations the trainer rendered at the final checkpoint (step 120,000), sampled from the EMA weights with classifier-free guidance 3.0 and the dopri5 ODE solver. Each is a stick-figure render of the generated motion via forward kinematics, with the prompt in the title. The rigs are training rigs; the prompts are their training captions.

File	Rig	Prompt
examples/mixamo_jumps_while_running_arm_extended.mp4	Mixamo humanoid, 22 joints	"A person jumps while running with one arm extended as if holding something."
examples/truebones_crocodile_bounces_then_collapses.mp4	Truebones Crocodile	"An animal bounces and then collapses."
examples/objaverse_669a73d2_sways_hips_and_legs.mp4	Objaverse rig 669a73d257bbceb71e6d3194_fbx, biped	"An object sways its hips and legs in place."
examples/objaverse_eeb9bae9_backflip.mp4	Objaverse rig eeb9bae905d047699933b8b67a213bdf	"An object performs a backflip."

These are the trainer's own debug renders, not a curated or cherry-picked set: they are the first four of the six rigs the visualization step drew at random for that checkpoint.

License and data provenance

The weights, configs, logs and example renders in this repository are released under the MIT License, matching the license of the UniMate code they were produced with. The training data keeps its sources' terms: Mixamo assets under Adobe's Mixamo terms, Objaverse-XL assets under each object's own license, and the Truebones ZOO animal motions under Truebones' commercial terms, which permit royalty-free use including research but forbid redistributing the motion files. No motion data is included in this repository; only weights derived from it, the per-channel normalization statistics, and the run config. Review the source licenses before using this model commercially.

Citation

If you use this checkpoint, please cite the original work:

@article{mou2026unimate,
  title   = {UniMate: One Unified Model to Animate Diverse Skeletons},
  author  = {Mou, Linzhan and Lei, Jiahui and Dou, Zhiyang and Cai, Chenyue and Song, Chaoyue and Finkelstein, Adam and Rusinkiewicz, Szymon},
  journal = {arXiv preprint arXiv:2609.05415},
  year    = {2026}
}