Skip to content

Repository files navigation

Relightable Avatar Pipeline

A self-contained, 5-stage pipeline that trains Generative Relightable Avatars and generates novel relit videos with independent pose/camera/lighting control:

Stage Script What it does Produces
1 01_train_microfacet.py Fits a per-subject microfacet (GGX) material -- albedo, roughness, environment-light scale -- against multi-view lightstage photos, with a differentiable renderer. Material parameter files
2 02_train_relightnet.py Trains a UNet ("RelightNet") to refine stage 1's physically-based render into the final relit UV texture. DDP across all visible GPUs. RelightNet checkpoints
3 03_dump_vace_training_data.py Renders matched ground-truth / conditioning / lighting video triplets used to fine-tune the video diffusion model. Training video clips
4 04_train_wan_vace.py Fine-tunes Wan2.1-VACE-1.3B (LoRA on the DiT + full fine-tune of the VACE conditioning branch and a lighting cross-attention module) on stage 3's clips, with an "error recycling" replay-buffer regularizer. DDP across all visible GPUs. Wan-VACE checkpoints
5 05_generate_video.py Test-time, autoregressive generation of an arbitrarily long relit video, given independent pose / camera / lighting indices. A relit .mp4

Run the stages in order once per subject. Stages 1-2 are fit once; stage 3 can be re-run any time to add more training clips; stage 4 is the expensive fine-tune; stage 5 is repeatable test-time generation (different pose/camera/lighting per run, no retraining needed).

Everything under common/ and the numbered scripts is self-contained project code, including a vendored copy of the DiffSynth-Studio subset needed for the Wan-VACE video model (common/diffsynth_vendor/, it has a different LICENSE) -- no pip install diffsynth is required.

This document covers: setting up the conda environment, MatAnyone2, what every path argument means, and how to run all 5 stages end to end. See project page for data.


1. Requirements

  • Linux, NVIDIA GPU(s) with a driver new enough for CUDA 11.8+. RTX 20-series (Turing) or newer is required for the OptiX-based shadow-ray extension used by stages 1, 3, and 5.
  • conda (or mamba) and git.
  • One conda environment (Python 3.13) is used for all 5 stages -- it needs both the rendering stack (PyTorch3D, torchoptix) for stages 1-3 and 5, and the video-model stack for stages 4-5, since stage 5 runs the render and the video model together in the same process. See section 2.
  • Enough disk space for: the Wan2.1-VACE-1.3B base checkpoint (a few GB, auto-downloaded on first use), MatAnyone2's pretrained weights (~140MB), and this pipeline's own outputs (checkpoints and dumped training videos are the bulk of it -- easily 100s of GB for a full run).
  • Multi-GPU is optional but strongly recommended for stages 2 and 4 (both auto-detect and use every visible GPU via torch.cuda.device_count() + DDP; single-GPU works too, just slower).

2. Environment setup

2.1 Clone this repository

Clone the repository, and:

cd GRA

Nothing else needs to be cloned alongside it for the pipeline code itself -- common/diffsynth_vendor/ already contains everything needed from DiffSynth-Studio. You only need to separately clone MatAnyone2 (see 2.3).

2.2 The conda environment (all 5 stages)

conda create -n relighting python=3.13 -y
conda activate relighting

# PyTorch (CUDA 11.8 build)
pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu118

# PyTorch3D -- version must match your torch/CUDA/Python combo; see
# https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md for the current recommended
# install command for your exact combination.
pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"

# torchoptix -- custom CUDA/OptiX shadow-ray visibility extension used by the microfacet renderer
# (common/microfacet.py). Requires an OptiX-capable NVIDIA GPU + driver.
pip install torchoptix

# Wan-VACE / DiffSynth stack (no `diffsynth` package itself -- it's vendored)
pip install peft transformers safetensors huggingface_hub modelscope accelerate \
            ftfy regex typing_extensions pandas einops

# Everything else
pip install opencv-python trimesh scipy numpy Pillow imageio imageio-ffmpeg tensorboard tqdm

# Optional, faster attention backends for stage 4/5 -- used automatically if importable, silently
# skipped otherwise. Only install if your GPU/CUDA setup supports them.
pip install flash-attn        # or: sageattention / xformers

Package versions this pipeline was developed and tested against, for reference: torch==2.7.1+cu118, torchvision==0.22.1+cu118, torchaudio==2.7.1+cu118, pytorch3d==0.7.9, torchoptix==0.0.1, trimesh==4.10.1, opencv-python==4.12.0.88, scipy==1.16.3, numpy==1.26.4, Pillow==9.4.0, imageio==2.35.1, imageio-ffmpeg==0.6.0, tensorboard==2.20.0, peft==0.18.0, transformers==4.57.3, safetensors==0.7.0, huggingface_hub==0.36.0, modelscope==1.32.0, accelerate==1.12.0, ftfy==6.3.1, einops==0.8.1, pandas==2.3.3.

2.3 MatAnyone2 (stage 3's video matting)

Stage 3 mats ground-truth clips using MatAnyone2 (video matting given a first-frame mask). It's imported in-process from within stage 3, not called as a separate program, so its dependencies need to be installed into this same environment:

git clone https://github.com/pq-yang/MatAnyone2 /path/to/MatAnyone2
cd /path/to/MatAnyone2
conda activate relighting
pip install -e .

# Download the pretrained checkpoint into pretrained_models/ (see MatAnyone2's own README for
# the current release URL/version)
mkdir -p pretrained_models
# place matanyone2.pth here -- pretrained_models/matanyone2.pth

Then point stage 3 at it: --matanyone_repo_path /path/to/MatAnyone2 --matanyone_ckpt_path /path/to/MatAnyone2/pretrained_models/matanyone2.pth.

2.4 Wan2.1-VACE-1.3B base weights

Stages 4 and 5 need the base Wan2.1-VACE-1.3B checkpoint (DiT + T5 text encoder + VAE). This is downloaded automatically on first use via ModelScope -- no manual download or account needed. Control where it's cached with these environment variables (set before running stage 4/5, or export in your shell profile):

export DIFFSYNTH_MODEL_BASE_PATH=/path/to/model/cache/   # where weights are stored
export MODELSCOPE_CACHE=/path/to/modelscope/cache/
export MODELSCOPE_HOME=/path/to/modelscope/cache/

To download from Hugging Face instead of ModelScope, set DIFFSYNTH_DOWNLOAD_SOURCE=huggingface.


3. Path reference: what each argument means

Every stage script takes its paths as plain CLI flags -- nothing is hardcoded. The convenience script prepare_subject.py (see section 5) can fill all of these in automatically if you adapt its hardcoded path templates to your own data layout first; otherwise, pass every flag below explicitly per stage.

Inputs (capture data -- read-only, never written to by this pipeline)

Flag Meaning
--mesh_path The subject's template mesh (UV-mapped).
--camera_calibration Camera calibration file for the training capture rig.
--reference_pointcloud One reference point cloud, composed of mesh vertices, used to establish the scene's global scale (can be any pointcloud from pcs_dir, see below).
--pcs_dir Directory of per-frame tracked mesh vertices (one file per frame).
--geometry_dir Directory of per-frame UV-space geometry maps, {frame}.npy, float32, channels: [0:3] = position of frame t, [3:6] = normal of frame t-2, [6:9] = normal of frame t-1, [9:12] = normal of frame t. The loader spatially downsamples 2x on load (every other pixel), to match the pipeline's 512x512 UV working resolution -- render/store at 1024x1024 so the effective resolution comes out to 512x512.
--envmap_dir Directory of per-frame environment-light samples used for physically-based shading (stages 1, 2, 3): each {frame}.npy is the envmap downsampled to a 16x32x3 lat-long RGB grid, progressively with Gaussian blur at every stage, concatenated with a 64-dim 3D sinusoidal positional encoding per grid cell (16x32x67 total), flattened to (512, 67) and stored as float32.
--envmap_hdr_dir Directory of per-frame full-resolution .hdr environment panoramas (not downsampled/encoded like --envmap_dir), used as the video model's lighting conditioning signal (stage 3, 5) and to render stage 3's cosmetic envmap-preview video.
--sample_envmap_hdr One .hdr file used to initialize stage 3's envmap projector.
--ground_truth_dir / --input_dir Directory of ground-truth multi-view photos (per camera, per frame).
--mask_dir_stage1 Foreground masks of the relit frames -- used by stage 1/2 as a per-pixel training loss weight (every frame is read).
--mask_dir The same foreground masks -- used by stage 3 only to seed MatAnyone2's video matting (only each sampled clip's first frame is read). Two separate flags since they're read by different stages, but you can point both at the same directory.
--interp_frames_dir Optional motion-compensated/interpolated frames used to drive MatAnyone2's matting (the resulting mask is still applied to the raw photos in --ground_truth_dir); omit to mat directly from the raw frames.
--active_cameras Comma-separated camera IDs to train on (e.g. excluding held-out evaluation cameras).
--min_frame / --max_frame Inclusive frame-number range to train on.

Test-time variants used only by stage 5 (a held-out capture, different from the training one above): --camera_calibration, --pcs_dir, --geometry_dir, --envmap_dir, --envmap_hdr_dir -- same meanings, just pointed at the test-time capture.

This pipeline's own outputs (created automatically; point these anywhere with free disk space)

Flag Meaning
--output_dir (stage 1, in example_config.json) Where stage 1 writes its material parameters (model_params/) and PNG previews.
--microfacet_params_dir Stage 1's {output_dir}/model_params -- stage 2/3/5's input.
--checkpoint_dir (stage 2) Where RelightNet checkpoints (checkpoint_{iter}.pt) are saved.
--relightnet_checkpoint (stage 3, 5) A specific RelightNet checkpoint file; if omitted, auto-resolves to the highest-iteration checkpoint in the relevant checkpoint directory.
--gt_output_dir / --gt_video_path Stage 3's masked ground-truth video clips; stage 4's dataset input.
--vace_output_dir / --vace_video_path Stage 3's VACE-conditioning render video clips (stage 1+2 composite); stage 4's dataset input.
--envmap_output_dir / --envmap_video_path Stage 3's cosmetic lighting-conditioning video clips; stage 4's dataset input.
--vace_reference_image_path Not produced by this pipeline -- a single fixed reference photo (e.g. a well-lit frontal frame of the subject) used by VACE conditioning in stages 4 and 5. Place one manually before running stage 4.
--checkpoint_path (stage 4) Where Wan-VACE checkpoints (checkpoint_{iter}.pt) are saved.
--vace_checkpoint (stage 5) A specific Wan-VACE checkpoint file; if omitted, auto-resolves to the highest-iteration checkpoint in --checkpoint_path.
--output_path (stage 4) Stage 4's own output directory (alongside --checkpoint_path; used for any non-checkpoint artifacts).
--tensorboard_dir / --log_dir Shared TensorBoard log directory -- stages 1, 2, and 4 all write here (tags prefixed stage1/, stage2/, stage4/), so one tensorboard --logdir <dir> shows all three together.
--output_dir (stage 5) Where the final generated .mp4s are written.

4. Running the pipeline

Every stage script also accepts --subj <N> plus an auto-filled default path set if you've adapted prepare_subject.py to your data (section 5) -- the commands below show the fully explicit form so they work regardless of that. conda activate relighting once; every stage below runs in that same environment.

Stage 1 -- microfacet material fitting

python 01_train_microfacet.py --config my_config.json

Takes a single JSON config rather than flat CLI flags (see example_config.json for the schema and every field's meaning -- paths, active cameras/frame range, and training/model/learning-rate hyperparameters). Runs single-GPU, self-terminates at training.max_steps (default 20000).

Stage 2 -- RelightNet training

python 02_train_relightnet.py \
    --mesh_path ... --camera_calibration ... --reference_pointcloud ... \
    --microfacet_params_dir ... --pcs_dir ... --geometry_dir ... --envmap_dir ... \
    --ground_truth_dir ... --mask_dir_stage1 ... \
    --active_cameras 0,1,2,3,4,5,7,8,9,10 --min_frame 2 --max_frame 14200 \
    --checkpoint_dir ...

DDP across every visible GPU (torch.cuda.device_count()). Has no built-in stopping point -- it runs --num_epochs (default 10) full passes over the dataset, which for a large per-frame, per-camera dataset can far exceed any reasonable single job's time limit. Saves a checkpoint every --checkpoint_every iterations (default 10000); to resume, pass --resume_checkpoint_path <checkpoint file> (use common.subject_config.find_latest_checkpoint to locate the newest one, or just pick a filename). Just keep resubmitting/resuming until TensorBoard shows the loss has converged -- there's no need to run every epoch.

Stage 3 -- dump Wan-VACE training data

python 03_dump_vace_training_data.py \
    --input_dir ... --mask_dir ... \
    --mesh_path ... --camera_calibration ... --reference_pointcloud ... --microfacet_params_dir ... \
    --relightnet_checkpoint ... \
    --pcs_dir ... --geometry_dir ... --envmap_dir ... --envmap_hdr_dir ... \
    --gt_output_dir ... --vace_output_dir ... --envmap_output_dir ... \
    --num_sequences 60000 --num_frames 17

Single GPU. Samples --num_sequences random (camera, start-frame) clips and renders the matching GT/VACE/envmap video triplet for each. --relightnet_checkpoint can be omitted to auto-resolve the latest checkpoint. This can be sharded across multiple parallel jobs by giving each one a disjoint --start_seq_id range and distinct --seed. Can be re-run at any time to add more training clips (stage 4's dataset just globs whatever's in --gt_video_path/--vace_video_path).

Stage 4 -- Wan-VACE fine-tuning

export DIFFSYNTH_MODEL_BASE_PATH=... MODELSCOPE_CACHE=... MODELSCOPE_HOME=...
python 04_train_wan_vace.py \
    --gt_video_path ... --vace_video_path ... --envmap_video_path ... \
    --vace_reference_image_path ... \
    --checkpoint_path ... --output_path ... --log_dir ... \
    --use_error_recycling --error_buffer_k 500 --buffer_warmup_iter 50 \
    --buffer_replacement_strategy l2_batch --motion_error_sample_from_all_grids \
    --noise_prob 0.01 --motion_error_prob 0.9 --latent_prob 0.9 --clean_prob 0.2 \
    --clean_buffer_update_prob 0.1 --learning_rate 1e-4

DDP across every visible GPU. See common/vace_config.py's VaceTrainingConfig class for every tunable field and its default (LoRA rank/targets, lighting-encoder type, error-recycling probabilities, checkpoint cadence, etc.) -- every field is also settable as --<field_name> on the CLI. Like stage 2, has no built-in stopping point (--num_epochs, default 5, over the full dumped-clip dataset) -- checkpoints every --checkpoint_steps iterations (default 500); to resume, pass --resume --resume_checkpoint_path <checkpoint file>.

--use_error_recycling enables the replay-buffer regularizer described in the module docstring of common/vace_model.py; omit it for plain fine-tuning.

Stage 5 -- test-time generation

export DIFFSYNTH_MODEL_BASE_PATH=... MODELSCOPE_CACHE=... MODELSCOPE_HOME=...
python 05_generate_video.py \
    --prompt "photorealistic human portrait" \
    --mesh_path ... --camera_calibration ... --reference_pointcloud ... \
    --microfacet_params_dir ... --relightnet_checkpoint ... --vace_checkpoint ... \
    --vace_reference_image_path ... \
    --pcs_dir ... --geometry_dir ... --envmap_dir ... --envmap_hdr_dir ... \
    --camera_number 18 --pose_start_frame 7 --total_frames 209 \
    --output_dir ... --output_name my_video

--relightnet_checkpoint / --vace_checkpoint can be omitted to auto-resolve the latest checkpoint in the relevant directory. Generation is chunked ("multi-segment", --segment_size frames per chunk, default 17) with a 5-frame overlap carried between chunks for temporal consistency, so --total_frames can be arbitrarily long. Pose, camera, and lighting are independently controllable:

  • Pose: --pose_start_frame + --total_frames (stepping by --frame_stride, default 2) select which pose frames to drive the mesh with, from --pcs_dir / --geometry_dir.
  • Camera: --camera_number (looked up in --camera_calibration).
  • Lighting: --envmap_start_idx + --envmap_stride select which lighting sample(s) to use, advanced independently of the pose -- hold it fixed, animate it, or desync it from the pose.

Produces two videos: {output_name}_generated.mp4 (the final relit result) and {output_name}_vace.mp4 (the VACE-conditioning input, for comparison/debugging).


5. prepare_subject.py (optional convenience)

If you're running many subjects with a consistent directory layout, it's worth adapting prepare_subject.py once so every stage script can be run with zero path arguments:

  1. If your data follows the same {data_root}/subj{N}/... / {scratch_root}/... layout the subject_paths() templates encode, just point --data_root/--scratch_root at your own roots. Otherwise, open prepare_subject.py and edit subject_paths() directly to match your own layout, keyed by subject number.
  2. python prepare_subject.py --subj 1 --data_root /your/data --scratch_root /your/scratch generates subjects/subj1/{microfacet_config.json, pipeline_paths.json} and marks subject 1 as active (subjects/active_subject.txt).
  3. Every stage script then needs no path flags at all: python 01_train_microfacet.py, python 02_train_relightnet.py, etc. -- an explicit --subj N, or any individual path flag, still overrides the prepared defaults.
  4. All of this pipeline's own outputs default to outputs/subj{N}/ under the repo root (override with prepare_subject.py --output_root). outputs/ and subjects/ are gitignored -- this is per-machine, generated state, not something to commit.

slurm/*.slurm are example SLURM launcher scripts: PIPELINE_ROOT is resolved relative to the script itself, conda/cache/model-cache paths default to repo-relative locations, and they all activate the single relighting environment from section 2.2. The one thing you must still edit per script is #SBATCH -p <partition> -- GPU partition names are cluster-specific and SLURM doesn't let a directive read an environment variable. Otherwise, treat them as a reference for resource sizing (GPUs/CPUs/time per stage) and NCCL environment setup for your own scheduler.


About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages