mlnha/vace-aug

vace-aug — VACE object-swap augmentation for RoboCasa

Code and packed inputs for swapping a target mesh into RoboCasa pick-and-place episodes with the VACE video-diffusion model (8 target objects x 32 episodes x 3 cameras), built as a baseline to compare against MimicGen-based augmentation.

aug32/                        the augmentation pipeline — start at aug32/README.md
VACE/                         ali-vilab/VACE fork carrying the batch flags this needs
vace_object_insert_batch.py   stock batch driver (build / prep-worker / split-wan / stitch)
build_assignments_*.py        earlier selection variants (gt-mask-only, cross-category)
e2e_videos/                   a worked end-to-end example, three cameras
datasets_pack/                everything the pipeline reads, trimmed from 140 GB to 2.9 GB

New machine? Read [`aug32/SETUP.md`](aug32/SETUP.md) first — it is the setup path in order: unpack, link weights, build the two environments, export five variables, preflight, smoke test.

Unpack

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

export VACE_ROBOCASA_ROOT=$PWD/data/pickplace_target_human
export VACE_AUG_MATERIALS=$PWD/data/aug_materials_ext

Not included: the 31 GB of WAN weights (Wan2.1-VACE-1.3B, VACE-Annotators) — link them in at VACE/models — and the meshes of objects other than the eight targets, needed only to re-render GT masks.

The rest of this page is aug32/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 → donut5, milk → donut5, …). 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):

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 syrupbottle. 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 measuringcup, 3 steak), and those are already done.

Run

./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.

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

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)

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

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

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

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:

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.