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

# Self-Forcing DMD 单卡文本生成视频推理脚本。
# 可通过同名环境变量覆盖下面的默认值，例如：
#   GPU=1 PROMPT="A cat running on the grass." \
#     OUTPUT_FOLDER=videos/cat bash scripts/run_inference_dmd.sh
# 也可以直接使用已有的 prompt 文件：
#   DATA_PATH=prompts/MovieGenVideoBench_extended.txt \
#     bash scripts/run_inference_dmd.sh

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

GPU="${GPU:-0}"
PYTHON_BIN="${PYTHON_BIN:-python}"
CONFIG_PATH="${CONFIG_PATH:-configs/self_forcing_dmd.yaml}"
CHECKPOINT_PATH="${CHECKPOINT_PATH:-checkpoints/self_forcing_dmd.pt}"
OUTPUT_FOLDER="${OUTPUT_FOLDER:-videos/self_forcing_dmd}"
NUM_OUTPUT_FRAMES="${NUM_OUTPUT_FRAMES:-21}"
NUM_SAMPLES="${NUM_SAMPLES:-1}"
SEED="${SEED:-0}"

PROMPT="${PROMPT:-A cinematic shot of a stylish woman walking down a neon-lit Tokyo street at night, reflections shimmering on the wet pavement, natural body motion, detailed background, smooth camera movement.}"
DATA_PATH="${DATA_PATH:-}"

if [[ -z "$DATA_PATH" ]]; then
  DATA_PATH="$OUTPUT_FOLDER/prompts.txt"
  mkdir -p "$OUTPUT_FOLDER"
  printf '%s\n' "$PROMPT" > "$DATA_PATH"
fi

for required_path in "$CONFIG_PATH" "$CHECKPOINT_PATH" "$DATA_PATH"; do
  if [[ ! -f "$required_path" ]]; then
    echo "错误：找不到所需文件：$required_path" >&2
    exit 1
  fi
done

if ! command -v "$PYTHON_BIN" >/dev/null 2>&1; then
  echo "错误：找不到 Python 命令：$PYTHON_BIN" >&2
  echo "请先激活 Self-Forcing 环境，或通过 PYTHON_BIN 指定 Python。" >&2
  exit 1
fi

mkdir -p "$OUTPUT_FOLDER"

echo "DMD 推理配置："
echo "  GPU:        $GPU"
echo "  Config:     $CONFIG_PATH"
echo "  Checkpoint: $CHECKPOINT_PATH"
echo "  Prompts:    $DATA_PATH"
echo "  Output:     $OUTPUT_FOLDER"

CUDA_VISIBLE_DEVICES="$GPU" "$PYTHON_BIN" inference.py \
  --config_path "$CONFIG_PATH" \
  --checkpoint_path "$CHECKPOINT_PATH" \
  --data_path "$DATA_PATH" \
  --output_folder "$OUTPUT_FOLDER" \
  --num_output_frames "$NUM_OUTPUT_FRAMES" \
  --num_samples "$NUM_SAMPLES" \
  --seed "$SEED" \
  --use_ema \
  --save_with_index