#!/usr/bin/expect

set timeout -1
set host "156.225.27.52"
set port "5204"
set user "root"
set password "Jsjj@8888"

# 定义正确的模型路径
set base_dir "/root/ai-projects/DreamOmni2"
set model_dir "$base_dir/checkpoints/DreamOmni2"

puts "--> Starting DreamOmni2 Inference Test (with corrected paths)..."

spawn ssh -p $port $user@$host
expect {
    "password:" { send "$password\r" }
    "yes/no" { send "yes\r"; exp_continue }
}

expect "#"
send "source ~/miniconda3/bin/activate\r"
send "conda activate dreamomni2\r"
send "cd $base_dir\r"

puts "--> Running inference with explicit model paths..."
# 指定本地模型路径，并使用默认的在线 FLUX 模型（首次运行会下载）
# 如果服务器无法访问 HuggingFace 下载 FLUX，可能需要额外处理
send "python inference_gen.py \\
    --vlm_path '$model_dir/vlm-model' \\
    --gen_lora_path '$model_dir/gen_lora' \\
    --input_img_path 'example_input/gen_tests/img1.jpg' 'example_input/gen_tests/img2.jpg' \\
    --output_path 'test_result_final.png'\r"

expect {
    "Traceback" { puts "\n[ERROR] Traceback detected!" }
    "Error" { puts "\n[ERROR] General error detected!" }
    "Loading pipeline components" { puts "\n[INFO] Downloading/Loading FLUX model..." }
    "#"
}

# 再次检查文件
send "ls -lh test_result_final.png\r"
expect "#"

send "exit\r"
expect eof

puts "Test run finished."