scripts/setup_wan_train_env.sh

#!/usr/bin/env bash
# Creates the `wan-train` conda env used to LoRA-finetune Wan2.1-VACE-1.3B via
# DiffSynth-Studio (third_party/diffsynth, pinned commit in
# third_party/PINNED_COMMITS.txt) -- scripts/train_wan_vace_lora.sh runs the
# actual training loop (src/fpgm/training/train.py) inside this env.
#
# Idempotent: `conda create` is skipped if the env already exists, and every
# pip install below is safe to re-run. Never touches the `fpgm` or `wan-vace`
# conda envs.
#
# Why a THIRD env, not reusing `wan-vace` (inference, scripts/setup_wan_env.sh):
# training pulls a materially different, heavier dependency stack --
# `accelerate`, `peft` (LoRA injection), `modelscope`, `datasets`, `pandas` --
# that `wan-vace` has no reason to carry, and this repo's convention is one
# env per heavy external model/tool (fpgm, mujoco, sam3d-objects, trellis2,
# wan-vace are all already separate). Reusing wan-vace's torch build
# (2.6.0+cu124) anyway removes a variable: it's the exact build already
# proven on this host's H200s (see scripts/setup_wan_env.sh's own note), and
# DiffSynth's own pyproject.toml only requires torch>=2.0.0 -- nothing here
# needs a newer CUDA toolkit.
#
# Why conda (not uv): matches every other GPU-model env on this host except
# cosmos-transfer2.5's uv .venv, whose reason (transformer_engine's
# CUDA-version-matched NVRTC) doesn't apply to Wan2.1/DiffSynth.
#
# flash_attn is installed best-effort, NOT required (unlike wan-vace's setup
# script, where it is load-bearing). Read from DiffSynth's own code
# (diffsynth/core/attention/attention.py::initialize_attention_priority): if
# flash_attn isn't importable, DiffSynth genuinely falls back to
# torch.nn.functional.scaled_dot_product_attention (ATTENTION_IMPLEMENTATION
# = "torch"), not a hard assert like third_party/wan2.1's native
# generate.py -- see src/fpgm/training/train.py's module docstring for the
# exact line reference. Installed anyway when possible because SDPA is
# slower, and pinned to the same wheel already proven for wan-vace/trellis2
# (identical torch/CUDA/python build here).
#
# Usage:
#   scripts/setup_wan_train_env.sh

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_NAME="wan-train"
CONDA_ROOT="${CONDA_ROOT:-/home/quang/miniconda3}"
ENV_PY="${CONDA_ROOT}/envs/${ENV_NAME}/bin/python"
DIFFSYNTH_DIR="${REPO_ROOT}/third_party/diffsynth"
PINNED_SHA="6e2b14bc73ff317229b2a28487fe09250bbf463f"
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/7] sanity checks"
if [ ! -d "${DIFFSYNTH_DIR}/.git" ]; then
    echo "ERROR: DiffSynth-Studio not cloned at ${DIFFSYNTH_DIR}." >&2
    echo "       Run:" >&2
    echo "         git clone https://github.com/modelscope/DiffSynth-Studio.git ${DIFFSYNTH_DIR}" >&2
    echo "         (cd ${DIFFSYNTH_DIR} && git checkout ${PINNED_SHA})" >&2
    exit 1
fi
checked_out="$(cd "${DIFFSYNTH_DIR}" && git rev-parse HEAD)"
if [ "${checked_out}" != "${PINNED_SHA}" ]; then
    echo "WARNING: ${DIFFSYNTH_DIR} is at ${checked_out}, not the pinned ${PINNED_SHA}." >&2
    echo "         Proceeding anyway, but the VACE-branch verification and tensor-format" >&2
    echo "         contract recorded in src/fpgm/training/train.py's docstring were read" >&2
    echo "         against the pin -- a moved checkout can silently invalidate both." >&2
fi

echo "==> [1/7] 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/7] torch==2.6.0+cu124 / torchvision==0.21.0+cu124 (same build as wan-vace)"
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/7] numpy<2 (pinned first, same reasoning as wan-vace's setup script:"
echo "    installed before transformers/modelscope/datasets so nothing pulls numpy>=2"
echo "    that a later step would then have to downgrade). NOTE: step 6's"
echo "    opencv-python-headless has its own numpy>=2 floor and overrides this"
echo "    pin anyway -- confirmed by direct test that diffsynth/fpgm.training"
echo "    both import and run fine under numpy 2.x here, unlike the SAM3.1"
echo "    constraint that forces numpy<2 in the fpgm env. Kept as the first"
echo "    install anyway so a fresh run's dependency resolution starts from"
echo "    the same base as wan-vace's proven env."
PY -m pip install "numpy<2"

echo "==> [4/7] DiffSynth-Studio itself, --no-deps (its declared deps installed"
echo "    explicitly next, so its unpinned torch/torchvision entries can never"
echo "    silently upgrade the pinned build from step 2)"
PY -m pip install -e "${DIFFSYNTH_DIR}" --no-deps

echo "==> [5/7] DiffSynth's declared dependencies (diffsynth/pyproject.toml's"
echo "    [project.dependencies], minus torch/torchvision/numpy already pinned above)"
PY -m pip install \
    transformers imageio "imageio[ffmpeg]" safetensors einops sentencepiece \
    protobuf modelscope ftfy pandas accelerate peft datasets

echo "==> [6/7] opencv (fpgm.training.bundle_io's only non-numpy dependency)"
echo "    NOT 'pip install -e REPO_ROOT': fpgm's own pyproject.toml requires"
echo "    python>=3.12 (this env is 3.10, matching wan-vace/DiffSynth's proven"
echo "    build), so an editable install of the whole fpgm package would fail"
echo "    the requires-python gate outright. src/fpgm/training/{dataset,manifest,"
echo "    augment,types,bundle_io}.py + the handful of fpgm.geometry/fpgm.types/"
echo "    fpgm.utils.logging modules they import need only numpy (already"
echo "    installed above) and opencv beyond the stdlib -- installed directly,"
echo "    and REPO_ROOT/src is put on PYTHONPATH by"
echo "    scripts/train_wan_vace_lora.sh at run time (same mechanism"
echo "    conftest.py uses for tests, no install required either way)."
PY -m pip install "opencv-python-headless>=4.9.0.80"

echo "==> [7/7] flash-attn (best-effort, NOT load-bearing -- see header)"
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}" || \
        echo "    WARNING: flash_attn wheel install failed -- DiffSynth will fall back to" \
             "torch SDPA (slower, not incorrect). Not treated as fatal, see header." >&2
fi

echo "==> smoke test: import diffsynth's VACE model + our training dataset package; check CUDA"
LD_LIBRARY_PATH="${REPO_ROOT}/.nvshim:${LD_LIBRARY_PATH:-}" PYTHONPATH="${REPO_ROOT}/src:${PYTHONPATH:-}" PY - <<PYEOF
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 diffsynth
from diffsynth.models.wan_video_vace import VaceWanModel
from diffsynth.pipelines.wan_video import WanVideoPipeline
print("diffsynth", diffsynth.__version__ if hasattr(diffsynth, "__version__") else "(no __version__)")
print("VaceWanModel, WanVideoPipeline import OK")

from diffsynth.core.attention.attention import (
    FLASH_ATTN_2_AVAILABLE, FLASH_ATTN_3_AVAILABLE, initialize_attention_priority,
)
print("flash_attn 2 available:", FLASH_ATTN_2_AVAILABLE, "| flash_attn 3 available:", FLASH_ATTN_3_AVAILABLE)
print("attention implementation that will be used:", initialize_attention_priority())

import fpgm.training.dataset  # noqa: F401  -- must import with no GPU touched
import fpgm.training.train  # noqa: F401    -- lazy-imports diffsynth internally, must still import cleanly
print("fpgm.training package imports OK")
PYEOF

echo "==> verify wan-vace and fpgm envs' torch are untouched"
for other_env in wan-vace fpgm; do
    OTHER_PY="${CONDA_ROOT}/envs/${other_env}/bin/python"
    if [ -x "${OTHER_PY}" ]; then
        OTHER_TORCH="$("${OTHER_PY}" -c 'import torch; print(torch.__version__)' 2>/dev/null || echo '(import failed)')"
        echo "    ${other_env} torch: ${OTHER_TORCH}"
    else
        echo "    WARNING: ${other_env} env python not found at ${OTHER_PY}, skipping check." >&2
    fi
done

echo "==> done. Run training with scripts/train_wan_vace_lora.sh"