File size: 1,722 Bytes
59630ba | 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 47 48 49 50 51 52 53 54 55 56 57 58 59 | #!/usr/bin/env bash
set -o pipefail
# CONFIG SETTING
NUM_GPUS=${NUM_GPUS:-1}
MASTER_PORT=${MASTER_PORT:-29110}
OFFLOAD=${OFFLOAD:-0} # 1 or true = mutually-exclusive GPU residency for generator/VAE
EXP_NAME=decmem
CONFIG_PATH=${CONFIG_PATH:-configs/${EXP_NAME}.yaml}
CHECKPOINT_PATH=checkpoints/decmem.pt
DATA_PATH=example/example.csv
NUM_CONDITION_FRAMES=${NUM_CONDITION_FRAMES:-56}
NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
GUIDANCE_SCALE=${GUIDANCE_SCALE:-5}
SEED=${SEED:-42}
START_INDEX=${START_INDEX:-0}
N=${N:-2}
T_LAT=${T_LAT:-160} # optional; empty = use config value
T_TAG=${T_LAT:+_T${T_LAT}}
OUTPUT_DIR=${OUTPUT_DIR:-test/res/${EXP_NAME}_${T_TAG}_con_${NUM_CONDITION_FRAMES}_inf_${NUM_INFERENCE_STEPS}_gs_${GUIDANCE_SCALE}}
mkdir -p "$OUTPUT_DIR"
RANK0_LOG="${OUTPUT_DIR}/rank0.log"
SCRIPT_ARGS=(
--config_path "$CONFIG_PATH"
--checkpoint_path "$CHECKPOINT_PATH"
--data_path "$DATA_PATH"
--output_dir "$OUTPUT_DIR"
--num_condition_frames "$NUM_CONDITION_FRAMES"
--guidance_scale "$GUIDANCE_SCALE"
--num_inference_steps "$NUM_INFERENCE_STEPS"
--seed "$SEED")
SCRIPT_ARGS+=(--start_index "$START_INDEX")
if [[ -n "$N" ]]; then
SCRIPT_ARGS+=(--num_entries "$N")
fi
if [[ -n "$T_LAT" ]]; then
SCRIPT_ARGS+=(--num_latent_frames "$T_LAT")
fi
if [[ "$NUM_GPUS" -gt 1 ]]; then
echo "Multi-GPU inference: NUM_GPUS=${NUM_GPUS} MASTER_PORT=${MASTER_PORT}"
torchrun \
--standalone \
--nproc_per_node="$NUM_GPUS" \
--master_port="$MASTER_PORT" \
--local-ranks-filter=0 \
infer.py \
"${SCRIPT_ARGS[@]}" \
2>&1 | tee "$RANK0_LOG"
else
torchrun infer.py "${SCRIPT_ARGS[@]}" 2>&1 | tee "$RANK0_LOG"
fi |