PhysiFormer / README.md
yimingc9's picture
Use ZeroGPU-compatible PyTorch
7a7e899
|
Raw
History Blame Contribute Delete
8.99 kB

A newer version of the Gradio SDK is available: 6.22.0

Upgrade
metadata
title: PhysFormer
sdk: gradio
app_file: app.py
python_version: 3.10.13

PhysFormer: Learning to Simulate Mechanics in World Space

Yiming Chen, Yushi Lan, Andrea Vedaldi
Visual Geometry Group, University of Oxford

arXiv Project Page GitHub Models

PhysFormer is a unified diffusion transformer that generates 4D multi-object mesh dynamics directly in world coordinates for both rigid and elastic materials. Rather than predicting future frames in pixel space or rolling out next-step system states autoregressively, PhysFormer models motion as full-trajectory coordinate diffusion: given initial per-vertex positions, velocities, and material conditions, it denoises entire future vertex trajectories in one process, with mesh topology imposed at inference. This design enables physically plausible object-object and object-environment interactions without hard-coded constraints, simulator priors, or learned shape latents. Its DiT-style backbone uses factorized temporal, spatial, and object-level attention to capture coherent structure across time, vertices, and objects. Trained on over 100k collision-rich, single-material simulated trajectories, PhysFormer generalizes to unseen real-world geometries, larger object counts, and mixed-material scenes.

Installation

# Clone the repo
git clone TODO
cd physformer

# Create conda environment
conda create -n physformer python=3.10 -y
conda activate physformer

# Install requirements. The Hugging Face Space uses the ZeroGPU-supported PyTorch version pinned here.
pip install -r requirements.txt
# Optional: FlashAttention
pip install flash-attn==2.8.3 --no-build-isolation

FlashAttention is recommended for better speed/memory efficiency, but not required. Without it, the code falls back to PyTorch native SDPA.

The demo expects a checkpoint at:

checkpoints/checkpoint-best.pt

Packaged demo data is local to this repository:

ood_examples/         OOD inference inputs
indistri_examples/    in-distribution visualization inputs
eval_assets/          compact inputs used directly by eval_publish_losses.py

Run all commands below from the repository root:

cd physformer

Inference

The example scripts write predicted rollout samples into each input sample directory as sample_00/, sample_01/, etc. Existing outputs are kept unless OVERWRITE_FLAG=--overwrite is set.

Hugging Face ZeroGPU Demo

This repository includes a minimal Gradio Space app in app.py. It downloads checkpoint-best.pt from yslan/physformer if needed, runs one small inference rollout, and returns the rendered inference.mp4. Evaluation precomputation is not used by this demo.

Create a Gradio Space and select ZeroGPU in the Space hardware settings. Use this README configuration block if this repository is pushed directly as a Space:

sdk: gradio
app_file: app.py
python_version: 3.10.13

If the checkpoint repo is private or gated, add HF_TOKEN as a Space secret. The app also accepts:

  • PHYSFORMER_CKPT_REPO_ID: checkpoint repo, default yslan/physformer.
  • PHYSFORMER_CKPT_FILENAME: checkpoint file in that repo, default checkpoint-best.pt.
  • PHYSFORMER_AMP: inference precision, default bf16.

Run OOD inference:

bash scripts/run_ood_example.sh

This uses the OOD folders:

ood_examples/2obj_cow_horse
ood_examples/3obj_teapot_fish_bunny

Default OOD materials are:

elastic: horse fish bunny
rigid:   cow teapot

Run in-distribution example inference:

bash scripts/run_indistri_example.sh

This runs indistri_examples/rigid as rigid and indistri_examples/soft as elastic. It writes inference-only renders as inference.mp4 and ground-truth-only renders as GT.mp4. OOD example runs write only inference renders.

Common user controls:

# Fast smoke test: one rollout sample, one input sample, fewer denoising steps.
GENERATIONS=1 MAX_SAMPLES=1 SAMPLING_STEPS=5 bash scripts/run_ood_example.sh

# Change OOD object material assignment.
OOD_ELASTIC="horse bunny" OOD_RIGID="cow teapot fish" bash scripts/run_ood_example.sh

# Skip MP4 rendering.
RENDER_FLAG="" bash scripts/run_indistri_example.sh

# Replace existing sample_* outputs.
OVERWRITE_FLAG=--overwrite bash scripts/run_indistri_example.sh

GENERATIONS is the number of independent rollout samples generated per input sample. It maps to the launcher argument --generations.

Direct launcher form, if you do not want to use the scripts:

python run_official_demo_inference.py \
  --include ood \
  --generations 3 \
  --elastic horse --elastic fish --elastic bunny \
  --rigid cow --rigid teapot \
  --save-mp4

Useful direct launcher flags:

  • --generations: number of rollout samples per input sample.
  • --max-samples: limit how many input samples are run; 0 means all.
  • --num-sampling-steps: override denoising steps; useful for quick tests.
  • --elastic PATTERN, --rigid PATTERN: set OOD object materials by object-name substring.
  • --save-mp4: render prediction MP4s as inference.mp4; omit it to save only vertices.npz.
  • --save-gt-mp4: render ground-truth MP4s as GT.mp4.
  • --overwrite / --no-overwrite: replace or preserve existing outputs.
  • --dry-run: print selected samples and command without running the model.
  • --attention-debug: report the first PyTorch attention backend used.

Evaluation

The repository packages compact evaluation inputs in eval_assets/. The evaluator checks for eval_split.json, eval_precomp/, and eval_data/ before loading the model. If they are already present, it prints that preparation is skipped. Run the evaluator to generate rollout samples and compute losses:

python eval_publish_losses.py \
  --ckpt checkpoints/checkpoint-best.pt \
  --out_json reports/publication_losses.json \
  --out_tsv reports/publication_losses.tsv \
  -k 4

The evaluator reports:

  • mse: masked MSE on raw vertex positions.
  • rigidity: per-object Kabsch residual over frames 1..--rigidity_last_frame.
  • conservation_of_momentum: normalized system momentum drift over frames 1..--momentum_last_frame.

Common evaluation controls:

# Fast evaluation smoke test.
python eval_publish_losses.py \
  --ckpt checkpoints/checkpoint-best.pt \
  --out_json reports/smoke_losses.json \
  --out_tsv reports/smoke_losses.tsv \
  -k 1 \
  --limit 1 \
  --num_sampling_steps 5

Useful eval flags:

  • -k, --num_generations: number of rollout samples evaluated per input sample.
  • --limit: evaluate only the first N split entries.
  • --num_sampling_steps, --cfg_scale: sampling controls.
  • --split_file, --split_name: choose the evaluation split.
  • --precomp_root, --data_root: choose prepared evaluation assets.
  • --sample_root: raw ground-truth sample root used only if prepared eval assets are missing.
  • --prepare_sample_names: comma-separated raw sample folders to prepare from --sample_root.
  • --prepare_overwrite: regenerate prepared eval assets even if they already exist.
  • --device, --amp: runtime device and precision.
  • --out_json, --out_tsv: report paths.

prepare_publish_eval_inputs.py is the standalone version of the same preparation step. It is only needed to regenerate eval_assets/ from raw ground-truth sample folders. Those raw samples are not packaged by default. If you have them separately, either pass --sample_root to eval_publish_losses.py or run:

python prepare_publish_eval_inputs.py --sample_root /path/to/raw_eval_samples

Runtime Dependencies

Use the project inference environment with PyTorch installed. The copied code expects:

  • Python 3.10+
  • PyTorch 2.5.1 installed from the CUDA 12.4 wheel index for practical runtime
  • NumPy
  • tqdm
  • matplotlib
  • imageio and imageio-ffmpeg only if saving GIF/MP4 renders

requirements.txt includes a ZeroGPU-supported PyTorch version so Hugging Face Spaces can build directly from this repository. The copied model uses PyTorch scaled-dot-product attention and requires a CUDA fast-attention backend; it does not directly import the external flash-attn package.

The packaged src/official_demo_inference/configs/vertex_counts_multiobj_all.json replaces the checkpoint's original training-machine absolute vertex-count path. A legacy copy is also kept under src/mesh_primitives/ so direct calls into the copied PhysFormer script still have a local fallback.