#!/usr/bin/env bash
# Phase 1: generate the Extended-251 videos for every strategy, one base at a time,
# sharded by prompt across the GPUs in $GPUS (default 0-3).  Generation runs in the base self_forcing env --
# the same interpreter and package set the speed study used -- so the eval videos
# are produced by exactly the model stack that was benchmarked.
set -u
C=/local/zoubin/cz/projects/comparison
PY=/local/zoubin/cz/envs/self_forcing/bin/python
LOG="$1"; shift
NSHARDS=${NSHARDS:-4}
GPUS=(${GPUS:-0 1 2 3})

export SSL_CERT_FILE=/etc/pki/tls/certs/ca-bundle.crt
export TORCH_HOME=/local/zoubin/cz/.cache/torch
export HF_HOME=/local/zoubin/cz/.cache/huggingface
export HF_HUB_CACHE=/local/zoubin/cz/.cache/huggingface/hub
export TRANSFORMERS_CACHE=/local/zoubin/cz/.cache/huggingface/hub
export XDG_CACHE_HOME=/local/zoubin/cz/.cache

cd "$C"
for BASE in ${BASES:-self_forcing causal_forcing}; do
  echo "=== generation: $BASE ($(date +%H:%M:%S)) ==="
  pids=()
  for i in $(seq 0 $((NSHARDS-1))); do
    CUDA_VISIBLE_DEVICES=${GPUS[$i]} nohup $PY eval/generate_eval.py \
      --base "$BASE" --shard $i --num-shards $NSHARDS --out-root eval_out \
      > "$LOG/gen_${BASE}_$i.log" 2>&1 &
    pids+=($!)
  done
  for p in "${pids[@]}"; do wait "$p"; done
  echo "=== generation done: $BASE ($(date +%H:%M:%S)) ==="
done
echo "ALL GENERATION DONE"