#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"

GPU=0
CONDA_ENV="/mnt/s3files/s3-us-west2-default/zoubin/cz/envs/self_forcing"
CONFIG_PATH="configs/self_forcing_dmd.yaml"
CHECKPOINT_PATH="checkpoints/self_forcing_dmd.pt"
OUTPUT_FOLDER="videos/self_forcing_dmd_single"
DATA_PATH="$OUTPUT_FOLDER/prompt.txt"
NUM_OUTPUT_FRAMES=21
NUM_SAMPLES=1
SEED=0

PROMPT="A stylish woman strolls down a bustling Tokyo street, neon reflections shimmering on wet pavement, cinematic lighting, detailed motion, medium shot."

mkdir -p "$OUTPUT_FOLDER"
printf '%s\n' "$PROMPT" > "$DATA_PATH"

for path in "$CONFIG_PATH" "$CHECKPOINT_PATH" "$DATA_PATH"; do
  if [ ! -e "$path" ]; then
    echo "Missing required path: $path" >&2
    exit 1
  fi
done

CUDA_VISIBLE_DEVICES="$GPU" \
LD_LIBRARY_PATH="$CONDA_ENV/lib/python3.10/site-packages/nvidia/cublas/lib:$CONDA_ENV/lib/python3.10/site-packages/nvidia/cuda_runtime/lib:$CONDA_ENV/lib:${LD_LIBRARY_PATH:-}" \
python inference.py \
  --config_path "$CONFIG_PATH" \
  --output_folder "$OUTPUT_FOLDER" \
  --checkpoint_path "$CHECKPOINT_PATH" \
  --data_path "$DATA_PATH" \
  --num_output_frames "$NUM_OUTPUT_FRAMES" \
  --num_samples "$NUM_SAMPLES" \
  --seed "$SEED" \
  --use_ema \
  --save_with_index