README.md

---
title: PhysFormer
sdk: gradio
app_file: app.py
python_version: 3.10.13
---

<h1 align="center">PhysFormer: Learning to Simulate Mechanics in World Space</h1>

<p align="center">
  Yiming Chen, Yushi Lan, Andrea Vedaldi<br>
  Visual Geometry Group, University of Oxford
</p>

<p align="center">
  <a href="TODO_ARXIV_URL"><img src="https://img.shields.io/badge/arXiv-Paper-b31b1b" alt="arXiv"></a>
  <a href="TODO_PROJECT_PAGE_URL"><img src="https://img.shields.io/badge/Project-Page-orange" alt="Project Page"></a>
  <a href="TODO_GITHUB_URL"><img src="https://img.shields.io/badge/GitHub-Repo-blue" alt="GitHub"></a>
  <a href="TODO_HUGGINGFACE_URL"><img src="https://img.shields.io/badge/HuggingFace-Model-green?logo=huggingface" alt="Models"></a>
</p>

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

```bash
# 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
```

```bash
# 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:

```text
checkpoints/checkpoint-best.pt
```

Packaged demo data is local to this repository:

```text
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:

```bash
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:

```yaml
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
bash scripts/run_ood_example.sh
```

This uses the OOD folders:

```text
ood_examples/2obj_cow_horse
ood_examples/3obj_teapot_fish_bunny
```

Default OOD materials are:

```text
elastic: horse fish bunny
rigid:   cow teapot
```

Run in-distribution example inference:

```bash
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:

```bash
# 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:

```bash
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:

```bash
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:

```bash
# 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:

```bash
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.