#!/bin/bash
set -e

MODEL_PATH="tarun7r/vibevoice-hindi-1.5B"
LORA_PATH="/home/ubuntu/modi_processed/lora_output/lora"
OUTPUT_DIR="/home/ubuntu/modi_processed/test_outputs"
INPUT_DIR="/home/ubuntu/modi_processed/test_inputs"

export PYTHONPATH=/home/ubuntu/vibevoice-community:$PYTHONPATH
export CUDA_VISIBLE_DEVICES=0

mkdir -p "$OUTPUT_DIR"

for txt_file in "$INPUT_DIR"/test_*.txt; do
    fname=$(basename "$txt_file" .txt)
    echo ""
    echo "=============================================="
    echo "  TESTING: $fname"
    echo "=============================================="
    echo ""
    python3 /home/ubuntu/vibevoice-community/demo/inference_from_file.py \
        --model_path "$MODEL_PATH" \
        --txt_path "$txt_file" \
        --output_dir "$OUTPUT_DIR" \
        --checkpoint_path "$LORA_PATH" \
        --disable_prefill \
        --device cuda \
        --seed 42
    echo ""
    echo "  Done: $fname"
    echo ""
done

echo "=============================================="
echo "  ALL TESTS COMPLETE"
echo "=============================================="

for wav in "$OUTPUT_DIR"/*.wav; do
    duration=$(python3 -c "
import wave
with wave.open('$wav', 'r') as w:
    frames = w.getnframes()
    rate = w.getframerate()
    print(f'{frames/rate:.2f}s')
" 2>/dev/null || echo "N/A")
    echo "  $(basename $wav): $duration"
done