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.00101
COV-CD (higher is better)	56 %
1-NNA-CD (50 % is ideal)	57.5 % (generated 15 %, real 100 %)
Median distance to the true surface	0.0053 (real shapes: 0.0004)

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. MMD, COV and 1-NNA are the standard point-cloud generation metrics of PointFlow, computed between 200 generated shapes and 200 real shapes drawn separately (scripts/evaluate_generation.py). The surface distance is the signed distance of each point to the nearest bone cylinder, computed from the skeleton.

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.0090	0.0051
Points more than 30 % of the radius outside the surface	8.8 %	2.2 %
MMD-CD	0.00150	0.00098
COV-CD	50 %	56 %
1-NNA-CD	99.6 %	53.6 %
of which generated shapes correctly classified	99 %	7 %
of which real shapes correctly classified	100 %	100 %

A 1-NNA near 50 % does not mean the generated shapes are indistinguishable here. It averages two opposite errors. A real shape always has another real shape as its nearest neighbour (100 %), because the more diffuse generated shapes stay farther from it than its real neighbours. A generated shape, in turn, almost always has a real shape as its nearest neighbour (7 %): it sits closer to the real shapes than to other generated ones. The gap that remains is visible in the split, which evaluate_generation.py prints.

Number of ODE steps

Each Heun step costs two network evaluations. With this checkpoint, Heun 10 keeps almost all the quality for half the compute of the default (MMD 0.00108, COV 57 %, surface distance 0.0059), and beats Euler 20 at the same cost (MMD 0.00110, COV 35 %). Below 5 Heun steps quality collapses (Heun 3: MMD 0.00176, COV 40 %).

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). A nearest-neighbour classifier still recognises every real shape (1-NNA on the real side: 100 %).
• 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.