scripts/run_wan_vace.sh

#!/usr/bin/env bash
# Run Wan2.1-VACE-1.3B inference (generate.py --task vace-1.3B) on this host.
#
# Upstream: https://github.com/Wan-Video/Wan2.1, pinned at
#   9737cba9c1c3c4d04b33fcad41c111989865d315
# cloned into third_party/wan2.1 (see third_party/PINNED_COMMITS.txt).
#
# Environment: a dedicated conda env `wan-vace` (python 3.10), separate from
# `fpgm`, created with:
#   conda create -y -n wan-vace python=3.10
#   PY=~/miniconda3/envs/wan-vace/bin/python
#   $PY -m pip install torch==2.6.0 torchvision==0.21.0 \
#       --index-url https://download.pytorch.org/whl/cu124
#   $PY -m pip install "numpy<2" "opencv-python-headless>=4.9.0.80" \
#       "diffusers>=0.31.0" "transformers>=4.49.0" "tokenizers>=0.20.3" \
#       "accelerate>=1.1.1" tqdm imageio easydict ftfy imageio-ffmpeg \
#       decord einops huggingface_hub dashscope
#   $PY -m pip install "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.3/flash_attn-2.7.3%2Bcu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"
# `decord` and `einops` are both imported by wan/ but missing from Wan2.1's
# own requirements.txt -- without them, `import wan` fails at
# wan/modules/vae.py (einops) and any --src_video/--src_mask run fails inside
# VaceVideoProcessor (decord). `dashscope` IS needed too, despite being a
# prompt-extension-only dependency: generate.py imports it unconditionally
# at module load time (`from wan.utils.prompt_extend import
# DashScopePromptExpander, ...` -> `import dashscope`), so it must be
# installed even though this launcher never calls --use_prompt_extend.
#
# flash_attn is NOT optional, despite appearances. wan/modules/attention.py
# does define an `attention()` dispatcher that falls back cleanly to
# torch.nn.functional.scaled_dot_product_attention when flash_attn is
# missing -- but at this pinned commit, wan/modules/model.py imports and
# calls the raw `flash_attention` function directly (see its `from
# .attention import flash_attention` and every self/cross-attn call site),
# never the `attention()` wrapper. That raw function unconditionally
# `assert FLASH_ATTN_2_AVAILABLE`s, so every vace-1.3B run crashes ~90s in,
# after both models are already loaded, if flash_attn isn't importable. The
# SDPA fallback in attention.py is dead code from this entrypoint's
# perspective. Confirmed by hitting the AssertionError in
# wan/modules/attention.py:112 through generate.py -> wan/vace.py -> the
# VaceWanModel's forward -> WanAttentionBlock.self_attn -> flash_attention,
# then fixing it with the prebuilt wheel above (same
# cu12/torch2.6/cxx11abiFALSE/cp310 build already proven working for the
# `trellis2` env on this host -- installs from a wheel, no source build).
# `gradio` from requirements.txt is skipped: it is only the web demo, which
# this launcher doesn't use.
# torch==2.6.0+cu124 mirrors the repo's other proven conda env (`trellis2`);
# nothing here needs a newer CUDA toolkit since VACE-1.3B has no custom CUDA
# extensions to compile.
#
# Chosen conda (not uv) because every other GPU-model env on this host
# (fpgm, mujoco, sam3d-objects, trellis2) is conda; cosmos-transfer2.5's uv
# .venv is the outlier, driven by its need for transformer_engine's own
# CUDA-version-matched NVRTC, which Wan2.1 does not have.
#
# Talks to the rest of the repo over subprocess + files only (CLI flags in,
# an mp4 out) -- the fpgm env's torch/numpy stack is load-bearing for
# SAM3/TAPNext/pyrender and must never be touched by this script.
#
# Why the environment workarounds below exist:
#
#   .nvshim FIRST on LD_LIBRARY_PATH -- this host's NVIDIA kernel module is
#   570.172.08 but /usr/lib's libnvidia-ml.so.1 points at 580.173.02, so
#   nvmlInit_v2() fails ("Failed to initialize NVML: Driver/library version
#   mismatch", confirmed by running plain `nvidia-smi` without the shim).
#   scripts/nvidia_lib_shim.sh builds $REPO_ROOT/.nvshim from the matching
#   570.172.08 libraries that are still on disk, no root needed.
#
#   IMPORTANT correction to the Cosmos-Transfer2.5 launcher's version of this
#   comment: for THIS stack (torch==2.6.0+cu124, no transformer_engine),
#   skipping .nvshim does NOT reproduce the CUDACachingAllocator.cpp:983
#   hard assert. Tested directly, including with
#   PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True set: torch just emits
#   `UserWarning: Can't initialize NVML` once at import and every CUDA op
#   after that -- conv2d, large allocations, and a full real 50-step
#   vace-1.3B run -- completed normally. The hard assert is apparently
#   specific to transformer_engine's own direct NVML usage in the Cosmos
#   stack, not a general torch/CUDA-allocator property. .nvshim is kept here
#   anyway because it's cheap and because `nvidia-smi` ITSELF hard-fails
#   (exit 18) without it, which breaks this script's own GPU auto-pick /
#   wait-for-headroom logic below -- that dependency, not a torch crash, is
#   the real reason it's required for this launcher.
#
#   PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True -- this is a busy
#   4xH200 host shared with other users; free memory is fragmented, and a
#   large contiguous reservation is exactly what fails without this.
#
#   GPU auto-pick + wait-for-headroom -- measured VACE-1.3B peak VRAM at
#   832x480/81 frames on this host: see the number this script prints and
#   the report that shipped with it. Typical free memory per card here is
#   only 4-25 GB out of 143 GB (heavily contended). Launching into
#   insufficient memory doesn't just fail this run, it raises memory
#   pressure on whatever else is resident on that GPU -- so this polls for
#   real headroom instead of gambling (same pattern as
#   scripts/wait_and_run_cosmos.sh).
#
# Usage:
#   scripts/run_wan_vace.sh <save_file.mp4> [extra generate.py args...]
#
# Example (control video + reference image + prompt, VACE's native
# 480P bucket):
#   scripts/run_wan_vace.sh outputs/wan_vace/demo.mp4 \
#       --src_video /abs/path/depth_control.mp4 \
#       --src_ref_images /abs/path/plate.png \
#       --prompt "..."
#
# generate.py's own --task/--size/--ckpt_dir defaults are set below and can
# be overridden by passing them again in the extra-args tail (argparse takes
# the last occurrence).
#
# Environment overrides:
#   WAN_GPU          CUDA device index. Default: auto-pick the card with the
#                     most free memory at launch time.
#   WAN_MIN_FREE_GB   Minimum free VRAM required before launching (default 12).
#   WAN_WAIT_POLL_S   Poll interval in seconds while waiting for headroom
#                     (default 60).
#   WAN_WAIT_MAX_S    Give up waiting after this many seconds (default 3600).
#   WAN_CKPT_DIR      Wan2.1-VACE-1.3B checkpoint dir (default: the snapshot
#                     already downloaded under ~/.cache/huggingface/hub).
#   WAN_TASK          generate.py --task (default vace-1.3B).
#   WAN_SIZE          generate.py --size (default 832*480; VACE-1.3B only
#                     supports 480*832 / 832*480 -- see SUPPORTED_SIZES in
#                     third_party/wan2.1/wan/configs/__init__.py).

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
WAN_DIR="$REPO_ROOT/third_party/wan2.1"
ENV_PY="/home/quang/miniconda3/envs/wan-vace/bin/python"
NVSHIM="$REPO_ROOT/.nvshim"

WAN_CKPT_DIR="${WAN_CKPT_DIR:-/home/quang/.cache/huggingface/hub/models--Wan-AI--Wan2.1-VACE-1.3B/snapshots/574e6a744642ce3bee319afc31496b88bde8aac4}"
WAN_TASK="${WAN_TASK:-vace-1.3B}"
WAN_SIZE="${WAN_SIZE:-832*480}"
WAN_MIN_FREE_GB="${WAN_MIN_FREE_GB:-12}"
WAN_WAIT_POLL_S="${WAN_WAIT_POLL_S:-60}"
WAN_WAIT_MAX_S="${WAN_WAIT_MAX_S:-3600}"

if [[ $# -lt 1 ]]; then
    echo "usage: $0 <save_file.mp4> [extra generate.py args...]" >&2
    exit 2
fi

# Resolved to an absolute path *before* the cd below, exactly like
# run_cosmos_transfer.sh -- a relative --save_file would otherwise land
# under third_party/wan2.1 instead of where the caller meant.
SAVE_FILE_ARG="$1"; shift
mkdir -p "$(dirname "$SAVE_FILE_ARG")"
SAVE_FILE="$(readlink -f "$SAVE_FILE_ARG")"

# Any remaining path-shaped args (--src_video, --src_mask, --src_ref_images
# take comma-separated paths) must also be resolved before the cd. We only
# rewrite the args we recognize; anything else passes through untouched.
EXTRA_ARGS=()
prev=""
for arg in "$@"; do
    case "$prev" in
        --src_video|--src_mask)
            EXTRA_ARGS+=("$(readlink -f "$arg")")
            prev=""
            continue
            ;;
        --src_ref_images)
            resolved=""
            IFS=',' read -ra _refs <<<"$arg"
            for r in "${_refs[@]}"; do
                resolved+="$(readlink -f "$r"),"
            done
            EXTRA_ARGS+=("${resolved%,}")
            prev=""
            continue
            ;;
    esac
    EXTRA_ARGS+=("$arg")
    case "$arg" in
        --src_video|--src_mask|--src_ref_images) prev="$arg" ;;
        *) prev="" ;;
    esac
done

[[ -x "$ENV_PY" ]] || { echo "no wan-vace conda env at $ENV_PY -- see the header of this script for setup commands" >&2; exit 1; }
[[ -d "$WAN_DIR/wan" ]] || { echo "Wan2.1 not cloned at $WAN_DIR -- git clone https://github.com/Wan-Video/Wan2.1 $WAN_DIR && (cd $WAN_DIR && git checkout 9737cba9c1c3c4d04b33fcad41c111989865d315)" >&2; exit 1; }
[[ -f "$WAN_CKPT_DIR/config.json" ]] || { echo "checkpoint dir not found or incomplete: $WAN_CKPT_DIR" >&2; exit 1; }

export LD_LIBRARY_PATH="$NVSHIM:${LD_LIBRARY_PATH:-}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"

pick_gpu() {
    # Prints "index free_mib" for the card with the most free memory, or
    # "-1 0" if nvidia-smi (even shimmed) can't be read.
    LD_LIBRARY_PATH="$NVSHIM" nvidia-smi --query-gpu=index,memory.free \
        --format=csv,noheader,nounits 2>/dev/null \
        | awk -F', *' '{print $1, $2}' | sort -k2 -rn | head -1
}

if [[ -n "${WAN_GPU:-}" ]]; then
    GPU="$WAN_GPU"
    echo "WAN_GPU set explicitly: using GPU $GPU (no headroom check)"
else
    need_mib=$(( WAN_MIN_FREE_GB * 1024 ))
    waited=0
    while true; do
        read -r idx free_mib < <(pick_gpu)
        if [[ -z "${idx:-}" || "$idx" == "-1" ]]; then
            echo "WARNING: could not read GPU memory via nvidia-smi (even with .nvshim)" >&2
            GPU=0
            break
        fi
        if (( free_mib >= need_mib )); then
            echo "$(date '+%H:%M:%S') GPU $idx has $((free_mib/1024)) GB free (need ${WAN_MIN_FREE_GB} GB) -- using it"
            GPU="$idx"
            break
        fi
        echo "$(date '+%H:%M:%S') best card is GPU $idx with $((free_mib/1024)) GB free, need ${WAN_MIN_FREE_GB} GB -- waiting ${WAN_WAIT_POLL_S}s (${waited}/${WAN_WAIT_MAX_S}s waited)"
        if (( waited >= WAN_WAIT_MAX_S )); then
            echo "ERROR: no GPU reached ${WAN_MIN_FREE_GB} GB free within ${WAN_WAIT_MAX_S}s -- giving up rather than launching into a card that will likely OOM" >&2
            exit 1
        fi
        sleep "$WAN_WAIT_POLL_S"
        waited=$(( waited + WAN_WAIT_POLL_S ))
    done
fi
export CUDA_VISIBLE_DEVICES="$GPU"

cd "$WAN_DIR"

echo "task:     $WAN_TASK"
echo "size:     $WAN_SIZE"
echo "ckpt_dir: $WAN_CKPT_DIR"
echo "save to:  $SAVE_FILE"
echo "GPU:      $CUDA_VISIBLE_DEVICES"

exec "$ENV_PY" generate.py \
    --task "$WAN_TASK" \
    --size "$WAN_SIZE" \
    --ckpt_dir "$WAN_CKPT_DIR" \
    --save_file "$SAVE_FILE" \
    "${EXTRA_ARGS[@]}"