scripts/setup_wan_env.sh

#!/usr/bin/env bash
# Creates the `wan-vace` conda env used by scripts/run_wan_vace.sh to run
# Wan2.1-VACE-1.3B inference (third_party/wan2.1, pinned commit in
# third_party/PINNED_COMMITS.txt).
#
# Idempotent: `conda create` is skipped if the env already exists, and every
# pip install below is safe to re-run. Never touches the `fpgm` conda env,
# whose torch/numpy stack is load-bearing for SAM3/TAPNext/pyrender.
#
# Why conda (not uv, unlike third_party/cosmos-transfer2.5/.venv): every
# other GPU-model env on this host (fpgm, mujoco, sam3d-objects, trellis2)
# is conda, and Wan2.1 has no analogue of cosmos-transfer2.5's
# transformer_engine problem (a CUDA-toolkit-version-matched NVRTC dlopen
# that uv's isolated venv made tractable to pin). Wan2.1 compiles nothing
# itself; the one native extension it benefits from (flash-attn) installs
# from a prebuilt wheel, no compiler needed either way.
#
# Why torch==2.6.0+cu124 specifically: it's the exact build already proven
# to work on this host's H200s in the `trellis2` conda env, including with a
# prebuilt flash-attn wheel (see scripts/setup_trellis_env.sh). Reusing it
# removes a variable -- Wan2.1 itself only requires torch>=2.4.0 and has no
# compiled extensions that would pin the CUDA toolkit version tighter.
#
# Usage:
#   scripts/setup_wan_env.sh

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_NAME="wan-vace"
CONDA_ROOT="${CONDA_ROOT:-/home/quang/miniconda3}"
ENV_PY="${CONDA_ROOT}/envs/${ENV_NAME}/bin/python"
WAN_DIR="${REPO_ROOT}/third_party/wan2.1"
PINNED_SHA="9737cba9c1c3c4d04b33fcad41c111989865d315"
FLASH_ATTN_WHEEL_URL="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"

echo "==> [0/6] sanity checks"
if [ ! -d "${WAN_DIR}/.git" ]; then
    echo "ERROR: Wan2.1 not cloned at ${WAN_DIR}." >&2
    echo "       Run:" >&2
    echo "         git clone https://github.com/Wan-Video/Wan2.1 ${WAN_DIR}" >&2
    echo "         (cd ${WAN_DIR} && git checkout ${PINNED_SHA})" >&2
    exit 1
fi
checked_out="$(cd "${WAN_DIR}" && git rev-parse HEAD)"
if [ "${checked_out}" != "${PINNED_SHA}" ]; then
    echo "WARNING: ${WAN_DIR} is at ${checked_out}, not the pinned ${PINNED_SHA}." >&2
    echo "         Proceeding anyway, but the flash_attn-not-optional finding" >&2
    echo "         and other workarounds below were verified against the pin." >&2
fi

echo "==> [1/6] conda env '${ENV_NAME}' (python 3.10)"
if [ -x "${ENV_PY}" ]; then
    echo "    already exists at ${ENV_PY}, reusing"
else
    source "${CONDA_ROOT}/etc/profile.d/conda.sh"
    conda create -y -n "${ENV_NAME}" python=3.10
fi

PY() { "${ENV_PY}" "$@"; }

echo "==> [2/6] torch==2.6.0+cu124 / torchvision==0.21.0+cu124"
PY -c "import torch" 2>/dev/null && echo "    torch already installed, skipping" || \
    PY -m pip install torch==2.6.0 torchvision==0.21.0 \
        --index-url https://download.pytorch.org/whl/cu124

echo "==> [3/6] python dependencies"
# numpy<2 pinned per Wan2.1's own requirements.txt. decord and einops are
# both imported by wan/ but missing from that requirements.txt (see
# run_wan_vace.sh header for exactly which files fail without each).
# dashscope is imported unconditionally by generate.py at module load time
# even though this launcher never uses --use_prompt_extend. flash_attn and
# gradio are the two requirements.txt entries deliberately skipped here --
# see the run_wan_vace.sh header for why gradio is safe to skip; flash_attn
# is handled separately below because it is NOT safe to skip (see step 4).
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

echo "==> [4/6] flash-attn (prebuilt wheel -- NOT optional, see run_wan_vace.sh header)"
# wan/modules/model.py calls the raw flash_attention() function directly,
# not the attention() dispatcher that has an SDPA fallback -- so every
# vace-1.3B run crashes with AssertionError at wan/modules/attention.py:112
# without this, even though the SDPA fallback code exists elsewhere in the
# same file. Confirmed on this host by hitting that exact assertion first.
if PY -c "import flash_attn" 2>/dev/null; then
    echo "    flash_attn already importable, skipping"
else
    PY -m pip install "${FLASH_ATTN_WHEEL_URL}"
fi

echo "==> [5/6] smoke test: import wan, decord, flash_attn; check CUDA"
# .nvshim is needed here only because torch's own NVML probe at import time
# would otherwise print a (harmless, for this torch/CUDA build) warning --
# see run_wan_vace.sh header for the full story on what .nvshim is and is
# not load-bearing for on this host.
LD_LIBRARY_PATH="${REPO_ROOT}/.nvshim:${LD_LIBRARY_PATH:-}" PY - <<PYEOF
import sys
sys.path.insert(0, "${WAN_DIR}")
import torch
print("torch", torch.__version__, "| cuda available:", torch.cuda.is_available())
assert torch.__version__.startswith("2.6.0"), f"unexpected torch version {torch.__version__}"
import wan
from wan.configs import WAN_CONFIGS, SUPPORTED_SIZES
print("vace-1.3B supported sizes:", SUPPORTED_SIZES["vace-1.3B"])
print("vace-1.3B sample_fps:", WAN_CONFIGS["vace-1.3B"].sample_fps)
import decord, flash_attn
print("decord", decord.__version__, "| flash_attn", flash_attn.__version__)
PYEOF

echo "==> [6/6] verify fpgm env's torch is untouched"
FPGM_PY="${CONDA_ROOT}/envs/fpgm/bin/python"
if [ -x "${FPGM_PY}" ]; then
    FPGM_TORCH="$("${FPGM_PY}" -c 'import torch; print(torch.__version__)')"
    echo "    fpgm torch: ${FPGM_TORCH}"
else
    echo "    WARNING: fpgm env python not found at ${FPGM_PY}, skipping check." >&2
fi

echo "==> done. Run inference with scripts/run_wan_vace.sh"