MissawB/creature-flow-point-dit

Creature Flow: skeleton-conditioned Point-DiT

A flow matching model that generates a 2048-point surface of a quadruped from its kinematic skeleton alone (35 joints, SMAL-like topology). Because it only reads the skeleton, it can also grow bodies for skeletons it was never trained on, such as a chimera grafted from a lion body and a horse neck.

Try it: creature-flow-viewer plays the whole flow from Gaussian noise (t = 0) to the finished specimen (t = 1).

2048-point clouds generated in 20 Heun steps from the skeleton alone (red). The last column is a chimera: lion body, horse neck and head.

Model

Architecture	Point-DiT, 6 blocks, width 256, 8 heads, AdaLN-Zero time modulation
Conditioning	Transformer skeleton encoder, then point-to-joint cross-attention with a multi-scale distance bias
Parameters	10.4 M
Objective	Conditional optimal-transport flow matching, $x_t = (1-t),x_0 + t,x_1$, target $x_1 - x_0$
Sampling	Heun ODE solver, 20 steps
Weights	Exponential moving average (decay 0.995) after 90 epochs
Data scale	Points and joints are multiplied by 4 inside the flow (data_scale), so the data matches the unit-variance noise
Training data	procedural-quadrupeds, 500 shapes
Training cost	About 48 min on an RTX 4060 Laptop GPU

Each attention head $h$ subtracts $\lambda_h \lVert \hat{x}_{1,i} - J_k \rVert^2$ from the point-to-joint score, with $\lambda_h = 2 \cdot 2^h$, so the 8 heads range from near-global (2) to very local (256) attention. The distance uses an estimate of the final position $\hat{x}_1 = x_t + (1-t),\hat{v}$ rather than $x_t$, which is pure noise early in the trajectory.

Results

Evaluated on 40 skeletons never seen in training (8 per family). Each generated shape is assigned to the family whose 8 reference shapes are closest on average in Chamfer distance.

	This checkpoint
Family recognised	40/40 (100 %)
Chamfer distance to its family	0.0011
MMD-CD (lower is better)	0.00103
COV-CD (higher is better)	50 %
1-NNA-CD (50 % is ideal)	53.8 %

For scale, two real shapes of the same family are 0.0007 apart on average, and an untrained model scores about 1.9. Three training seeds of the same recipe all recognised 40/40, with a Chamfer distance of 0.0010 to 0.0011 and a 1-NNA of 50.0 to 53.8 %. MMD, COV and 1-NNA are the standard point-cloud generation metrics of PointFlow, computed between 40 generated shapes and 40 real shapes drawn separately (scripts/evaluate_generation.py).

In the ablation, the multi-scale distance bias accounts for most of the gain: a single fixed bias reached only 40 %.

The last step came from the data, not the model. The model only reads the skeleton, and bovids differed from horses mainly by body thickness, which no joint encoded: about half of the generated bovids were taken for horses (90 % overall). Generator version 4 makes that difference visible in the skeleton. Legs stand wider under a thicker body, and bovids carry the neck at 20° instead of 49°. With the same model and training, bovids went from 13/24 to 24/24 over three seeds, with no loss on the other families.

The surfaces were then too thick. Normalised shapes have a standard deviation of about 0.25, while the starting noise has a standard deviation of 1. Multiplying the coordinates by 4 inside the flow (data_scale: 4, in the spirit of EDM's $\sigma_{data}$) improves every measure over three seeds each:

Mean over 3 seeds	data_scale: 1	data_scale: 4
Median distance to the true surface	0.0129	0.0092
Points more than 30 % of the radius outside the surface	15.9 %	8.8 %
MMD-CD	0.00224	0.00154
COV-CD	48 %	51 %
1-NNA-CD	98.4 %	96.3 %

A scale of 8 (one run) was worse than 4.

The remaining blur came from undertraining. More ODE steps changed nothing, and Fourier input features were slightly worse, but shapes kept improving with every training stage while the loss barely moved (1.30 to 1.29). Training 90 epochs instead of 30 fixed most of it, over three runs each:

Mean over 3 seeds (data_scale: 4)	30 epochs	90 epochs
Median distance to the true surface	0.0092	0.0053
Points more than 30 % of the radius outside the surface	8.8 %	2.2 %
MMD-CD	0.00154	0.00100
COV-CD	51 %	50 %
1-NNA-CD	96.3 %	51.7 %

At 90 epochs, a nearest-neighbour classifier can barely tell generated shapes from real ones.

Use

git clone https://huggingface.co/MissawB/creature-flow-point-dit
cd creature-flow-point-dit
pip install -r requirements.txt safetensors

import torch
from safetensors.torch import load_file
from data.smal_quadrupeds import SMALQuadrupedGenerator
from models.point_dit import PointDiT
from models.flow_matching import OptimalTransportFlowMatching

model = PointDiT(point_dim=3, hidden_dim=256, num_layers=6, num_heads=8,
                 num_joints=35, spatial_distance_scale=2.0 / 4**2, data_scale=4.0)
model.load_state_dict(load_file("model.safetensors"))
model.eval()

joints = SMALQuadrupedGenerator(seed=0).generate_instance("felidae")["joints"]
fm = OptimalTransportFlowMatching(model)
_, points = fm.sample(torch.from_numpy(joints).unsqueeze(0), num_points=2048,
                      num_steps=20, solver="heun")
print(points.shape)  # (1, 2048, 3)

To build a chimera, graft a donor subtree onto a host skeleton with pipeline.hybrid_splicer.KinematicHybridSplicer().splice(host, donor, junction_joint="neck_base") and sample from the result.

The repository also holds the training (scripts/train.py), evaluation (scripts/evaluate_families.py, scripts/evaluate_generation.py) and viewer export (scripts/sample_and_export.py) scripts. Those scripts read .pt checkpoints produced by training. This repo ships only the EMA weights as model.safetensors.

Limitations

• Synthetic data. Shapes are cylinders sampled around a parametric SMAL-like skeleton, not real animal scans or the SMAL model itself. Within-family variation is small (about ±3 % on proportions).
• An easy test. The five families are separated by construction (proportions, stance, head carriage) and vary little inside a family, so 100 % shows that the model follows the skeleton, not that it generalises to real animals.
• Surface precision. Generated surfaces are still more diffuse than real ones (median distance 0.0053 against 0.0004), and shapes about 1.5 times as far from their family as two real shapes are from each other (0.0011 against 0.0007).
• Saturated evaluation. With 40 generated against 40 real shapes, a 1-NNA near 50 % can no longer measure further progress.
• Grafting aligns the donor subtree by translation only, with no scaling or rotation.

Collection

Source code, experiments and full results: GitHub.

Part of the Skeleton-conditioned creature flow collection with the viewer Space and the training dataset.