File size: 2,345 Bytes
3c3a938 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | #!/bin/bash
#
# Batch SV4D generation across 8 GPUs.
# Each GPU gets a chunk of videos and processes them sequentially.
#
# Usage (run from project root):
# bash scripts/sampling/batch_sv4d.sh ./cogvideo_batch ./outputs_4d
# bash scripts/sampling/batch_sv4d.sh ./cogvideo_batch ./outputs_4d 8 50 # num_gpus, num_steps
INPUT_DIR="${1:?Usage: bash scripts/sampling/batch_sv4d.sh <input_dir> <output_dir> [num_gpus] [num_steps]}"
OUTPUT_DIR="${2:?Usage: bash scripts/sampling/batch_sv4d.sh <input_dir> <output_dir> [num_gpus] [num_steps]}"
NUM_GPUS="${3:-8}"
NUM_STEPS="${4:-50}"
MODEL_PATH="checkpoints/sv4d2.safetensors"
SCRIPT="scripts/sampling/simple_video_sample_4d2.py"
mapfile -t ALL_VIDEOS < <(find "$INPUT_DIR" -maxdepth 1 -name "*.mp4" | sort)
TOTAL=${#ALL_VIDEOS[@]}
if [ "$TOTAL" -eq 0 ]; then
echo "No .mp4 files found in $INPUT_DIR"
exit 1
fi
echo "Found $TOTAL videos, distributing across $NUM_GPUS GPUs"
mkdir -p "$OUTPUT_DIR"
process_gpu() {
local gpu_id=$1
shift
local videos=("$@")
local count=${#videos[@]}
for i in "${!videos[@]}"; do
local vpath="${videos[$i]}"
local vname=$(basename "${vpath%.*}")
local out_dir="${OUTPUT_DIR}/${vname}"
if [ -f "${out_dir}/_done" ]; then
echo "[GPU $gpu_id] ($((i+1))/$count) Skip done: $vname"
continue
fi
echo "[GPU $gpu_id] ($((i+1))/$count) Processing: $vname"
CUDA_VISIBLE_DEVICES=$gpu_id python $SCRIPT \
"--input_path=$vpath" \
"--model_path=$MODEL_PATH" \
"--output_folder=$out_dir" \
"--num_steps=$NUM_STEPS"
if [ $? -eq 0 ]; then
touch "${out_dir}/_done"
echo "[GPU $gpu_id] ($((i+1))/$count) Done: $vname"
else
echo "[GPU $gpu_id] ($((i+1))/$count) FAILED: $vname"
fi
done
}
# Distribute videos round-robin across GPUs
for ((g=0; g<NUM_GPUS; g++)); do
GPU_VIDEOS=()
for ((j=g; j<TOTAL; j+=NUM_GPUS)); do
GPU_VIDEOS+=("${ALL_VIDEOS[$j]}")
done
if [ ${#GPU_VIDEOS[@]} -gt 0 ]; then
process_gpu $g "${GPU_VIDEOS[@]}" &
echo "Launched GPU $g with ${#GPU_VIDEOS[@]} videos"
fi
done
echo "All $NUM_GPUS GPU workers launched. Waiting..."
wait
echo "All done! Results in $OUTPUT_DIR/<video_name>/sv4d2/"
|