File size: 1,445 Bytes
ec0a9aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | #!/usr/bin/env bash
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "${REPO_ROOT}"
if [[ -f "${REPO_ROOT}/models/DreamDojo/.venv/bin/activate" ]]; then
source "${REPO_ROOT}/models/DreamDojo/.venv/bin/activate"
fi
PROCESSED_ROOT="${PROCESSED_ROOT:-datasets/processed/ctrl_world_robo_multiview}"
SVD_MODEL_PATH="${SVD_MODEL_PATH:-checkpoints/stabilityai/stable-video-diffusion-img2vid}"
CLIP_MODEL_PATH="${CLIP_MODEL_PATH:-checkpoints/openai/clip-vit-base-patch32}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROCESSED_ROOT}/checkpoints}"
GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
NNODES="${NNODES:-1}"
NODE_RANK="${NODE_RANK:-0}"
MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
MASTER_PORT="${MASTER_PORT:-29501}"
TRAIN_ARGS=(
-m finetuning.train_ctrl_world
--processed-root "${PROCESSED_ROOT}"
--manifest-path "${PROCESSED_ROOT}/manifest.json"
--svd-model-path "${SVD_MODEL_PATH}"
--clip-model-path "${CLIP_MODEL_PATH}"
--output-dir "${OUTPUT_DIR}"
--action-key eef_sim_pose_action
--state-key eef_sim_pose_state
--action-dim 12
--n-views 4
--fps 5
)
if [[ "${NNODES}" -gt 1 || "${GPUS_PER_NODE}" -gt 1 ]]; then
accelerate launch \
--num_processes "$((NNODES * GPUS_PER_NODE))" \
--num_machines "${NNODES}" \
--machine_rank "${NODE_RANK}" \
--main_process_ip "${MASTER_ADDR}" \
--main_process_port "${MASTER_PORT}" \
"${TRAIN_ARGS[@]}" "$@"
else
python "${TRAIN_ARGS[@]}" "$@"
fi
|