augment/vace/README.md

# VACE augmentation — 8 objects × 32 episodes, 3 cameras

Swaps a target mesh into RoboCasa episodes with the VACE video-diffusion model, producing
the same folder layout as the MimicGen augmentation run so the two baselines can be
compared file-for-file.

**Same-category first, random for the remainder.** For `donut_5`, episodes that already
contain a donut (`donut_8`, `donut_9`, …) are used first — those are the swaps VACE was
designed for. Only after that pool runs out does the rest of the 32 come from randomly
drawn episodes of any category (apple → donut_5, milk → donut_5, …). The split is recorded
per assignment as `fill: same_category | random`.

How far the same-category pool goes, after `make_gt_masks.py` filled in the 58 missing
same-category masks (episodes with a complete set of GT masks, excluding the target mesh):

| target | same-category pool | same-category used | random fill | reaches 32 same-category? |
|---|---|---|---|---|
| MeasuringCup009 | 57 | **32** | 0 | **yes** |
| steak_8 | 18 | 18 | 14 | no |
| SoapDispenser010 | 5 | 5 | 27 | no |
| teapot_6 | 5 | 5 | 27 | no |
| teapot_7 | 5 | 5 | 27 | no |
| donut_5 | 3 | 3 | 29 | no |
| Jar023 | 2 | 2 | 30 | no |
| SyrupBottle005 | 2 | 2 | 30 | no |
| **total** | | **72** | **184** | |

Only MeasuringCup009 can be filled from its own category. This is a property of the
dataset, not of the masks: across all 2,517 episodes there are only 5 teapot episodes,
3 donut, 2 jar, 2 syrup_bottle. Generating the other 1,830 missing masks would not move
these numbers — only 58 of the 1,888 were same-category for these eight targets (55
measuring_cup, 3 steak), and those are already done.

## Run

```bash
./run_aug32.sh <work_dir> <out_dir> <gpu> [num_wan_workers]

# one GPU
./run_aug32.sh /scratch/vace_aug32 \
    /home/nvidia/minha/pose6daug/baseline/mimicgen/output_aug_256 2

# four GPUs: split-wan writes four manifests, one worker each
./run_aug32.sh /scratch/vace_aug32 /path/out 2 4     # then move the workers onto 4 GPUs by hand
```

Every stage is idempotent — re-running after an interruption resumes rather than redoing.

| stage | what it does | cost |
|---|---|---|
| 1 `build` | pick sources, same-category first | seconds |
| 2 `prep` | render the target mesh's reference views, cut frames, per-view manifests | ~4 s / assignment |
| 3 `split` | chop into 81-frame WAN jobs | seconds |
| 4 `wan` | the diffusion pass | **everything** |
| 5 `stitch` | glue each camera's chunks back together | seconds |
| 6 `collect` | copy into the output layout | seconds |

**Runtime.** 256 assignments × 3 cameras ≈ **3,084 WAN jobs**. At the measured 51.4 s/job
that is **~44 h on one GPU**, ~11 h on four. Levers, in order of how much they buy and how
little they cost:

- `--offload_model False` is already passed by the script. The upstream default is `True`
  whenever `world_size == 1`, which moves the DiT to CPU and back around every job: 81.8 s
  vs 51.4 s measured. Never drop this flag.
- `--sample_steps 25` (upstream default 50) → 32.1 s/job, ~27 h on one GPU. Quality
  trade-off; look at the plates before committing.
- Dropping the two exo cameras cuts the job count to a third. At 256×256 the target-object
  mask is ~0.2 % of an exo frame and VACE tends to erase the object instead of replacing it;
  only `robot0_eye_in_hand` (~9 %) reliably gets a real swap. If the exo views are not
  needed, edit `CAMS` in `../vace_object_insert_batch.py`.

## Output

```
<out_dir>/<task>-ep<N>_<object the episode had>_<object it becomes>/
    robot0_agentview_left/robot0_agentview_left.mp4
    robot0_agentview_right/robot0_agentview_right.mp4
    robot0_eye_in_hand/robot0_eye_in_hand.mp4
    .done
```

e.g. `c2cab-ep000151_donut_donut_5` (same-category) and
`snk2c-ep000004_egg_teapot_7` (random fill). Task tags are `c2cab`, `c2sto`, `drw2c`,
`snk2c`, `tst2c`; the hyphen keeps the name to three underscore-separated parts, matching
MimicGen's `ep098_apple_teapot_6`. `.done` carries the real task, episode, source mesh,
target, language instruction, and which fill it came from.

## Files

| file | |
|---|---|
| `build_aug32_assignments.py` | source selection — the only new logic |
| `collect_outputs.py` | VACE plates → the output layout |
| `run_aug32.sh` | all six stages |
| `SETUP.md` | **start here on a new machine** — the setup path, in order |
| `preflight.sh` | checks environments, weights, datasets, GT-mask coverage, GPU |
| `make_gt_masks.py` | re-renders the missing GT masks from the simulator |
| `vace_wan_inference_batch.patch` | the upstream diff, for porting onto a fresh clone |
| `vendor/render_mesh_views_side.py` | mesh reference-view renderer, vendored from the original machine |

Stages 2–5 are `../vace_object_insert_batch.py` and `../VACE/vace/vace_wan_inference.py`
unmodified; this directory adds selection and collection around them.

## Prerequisites

Run `./preflight.sh` — it checks every item below and prints GT-mask coverage and GPU
state. Nothing here needs to be installed or cloned; it all already exists on this box and
is referenced by absolute path, the same convention the other baselines use.

### Environments (do NOT create a venv)

| what | python |
|---|---|
| batch driver (`vace_object_insert_batch.py`, the builder, the collector) | `/lp-dev/jonghoon/robocasa_calib/envs/robocasa/bin/python` (3.10, mujoco + robocasa) |
| WAN inference | `/lp-dev/jonghoon/VACE/venv/bin/python` (3.12, torch 2.5.1+cu124, flash_attn 2.8.3) |
| mesh reference-view rendering (called by `prep-worker`) | `/lp-dev/jonghoon/isaac-gr00t/venv/bin/python` — override with `VACE_RENDER_PYTHON`; it only needs `trimesh`, `pyrender`, `numpy`, `pillow` (measured: 4.12.1 / 0.1.45 / 1.26.4 / 12.1.1), so any env with those will do |

`prep-worker` renders the target mesh's reference views by importing `render_views` from
`render_mesh_views_side.py`, which lived outside this tree on the original machine. A copy
ships in **`vendor/render_mesh_views_side.py`**; point at it with

```bash
export VACE_RENDER_SCRIPT_DIR=$PWD/aug32/vendor
export VACE_RENDER_PYTHON=/path/to/a/python/with/trimesh+pyrender
```

### VACE code and weights (already cloned — do NOT clone again)

`../VACE/` is a fork of `https://github.com/ali-vilab/VACE.git` carrying the batch flags
(`--batch_manifest / --skip_existing / --batch_continue_on_error`) this pipeline depends on;
a fresh upstream clone will not have them. Weights come in through
`../VACE/models -> /lp-dev/jonghoon/VACE-home-models` (31 GB: `Wan2.1-VACE-1.3B`,
`VACE-Annotators`). ~27 GB of GPU memory is needed with `--offload_model False`.

### Datasets

| path | what |
|---|---|
| `/home/nvidia/jonghoon/robocasa_full/pickplace_target_human/` | source episodes — `videos/chunk-000/observation.images.<camera>/episode_NNNNNN.mp4` and `extras/episode_NNNNNN/ep_meta.json` (symlink to `/lp-dev/...`) |
| `/home/nvidia/jonghoon/robocasa_full/aug_materials/robocasa/<task>/ep<N>/` | `mask_gt_target_object_<camera>.mp4` — the GT masks |
| `/lp-dev/jonghoon/robocasa_calib/repos/robocasa/robocasa/models/assets/objects/{objaverse,lightwheel}/` | target meshes |

GT-mask coverage, measured: `PickPlaceCounterToCabinet` 502, `CounterToStove` 50,
`DrawerToCounter` 43, `SinkToCounter` 34, `ToasterToCounter` 0 — 629 usable episodes, and
the builder draws only from those.

## Porting to another machine

Do **not** clone `ali-vilab/VACE` fresh — upstream has no `--batch_manifest`, which every
stage past `split-wan` depends on.

**Copy the whole `baseline/vace/` directory**, not just `aug32/`. Nothing outside it is
needed except the 31 GB of WAN weights. To rebuild from a clone instead:
`git clone https://github.com/ali-vilab/VACE.git && git checkout 48eb44f`, then
`git apply vace_wan_inference_batch.patch`, and copy `build_vace_wan_batch_manifest.py`,
`vace_inference_full.py` and this `aug32/` directory alongside it.

### Unpack the data

```bash
cd baseline/vace && mkdir -p data && cd data
tar -xf  ../datasets_pack/source_episodes.tar       # 2.7 GB -> pickplace_target_human/
tar -xzf ../datasets_pack/aug_materials_ext.tar.gz  # 121 MB -> aug_materials_ext/
tar -xzf ../datasets_pack/target_meshes.tar.gz      #  14 MB -> objaverse/, lightwheel/
cd ..
```

### Point everything at it

Every path is an environment variable with the original machine's value as the default, so
this is the whole configuration:

| variable | default | set it to |
|---|---|---|
| `VACE_ROBOCASA_ROOT` | original machine | `$PWD/data/pickplace_target_human` |
| `VACE_AUG_MATERIALS` | original machine | `$PWD/data/aug_materials_ext` |
| `VACE_ASSET_ROOTS` | original machine | `$PWD/data/lightwheel:$PWD/data/objaverse` |
| `VACE_PY_ROBOCASA` | original machine | a python with `robocasa`, `mujoco`, `trimesh`, `pyrender` |
| `VACE_PY_WAN` | original machine | a python with torch 2.5.1+cu124 and `flash_attn` |
| `VACE_RENDER_SCRIPT_DIR` | **`aug32/vendor/`** | — already inside the tree |
| `VACE_RENDER_PYTHON` | **`VACE_PY_ROBOCASA`** | — only if that env lacks trimesh/pyrender |

```bash
export VACE_ROBOCASA_ROOT=$PWD/data/pickplace_target_human
export VACE_AUG_MATERIALS=$PWD/data/aug_materials_ext
export VACE_ASSET_ROOTS=$PWD/data/lightwheel:$PWD/data/objaverse
export VACE_PY_ROBOCASA=/path/to/robocasa/python
export VACE_PY_WAN=/path/to/wan/python
./aug32/preflight.sh          # says MISS for anything still wrong
```

The mesh reference-view renderer (`render_mesh_views_side.py`, which `prep-worker` imports)
is vendored at `aug32/vendor/` and defaults to the driver's own interpreter, so it needs no
separate environment — `trimesh` and `pyrender` are its only requirements beyond numpy/PIL.

### Still needed from elsewhere

The 31 GB of WAN weights (`Wan2.1-VACE-1.3B`, `VACE-Annotators`) — link them in at
`VACE/models`. And, only if you want to re-run `make_gt_masks.py` there, the meshes of
objects other than the eight targets, since that rebuilds each episode's own scene.

### Gotchas

- An episode without GT masks makes `prep-worker` raise `FileNotFoundError`. The builder
  filters on this, so it cannot happen through `run_aug32.sh` — but it will if you hand-edit
  `assignments.jsonl`.
- The WAN stage must run with `vace/VACE` as the working directory (imports and `--ckpt_dir`
  are relative to it). `run_aug32.sh` does the `cd`.
- `--save_file` in `vace_wan_inference.py` is written relative to the *current* directory,
  not `--save_dir`. The batch path does not use it; a one-off call should pass an absolute
  path.