#!/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"
set flux_dir "$base_dir/checkpoints/FLUX_Model"

puts "--> Starting DreamOmni2 Inference Test (Offline Mode)..."

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 using LOCAL models..."
# 指定所有本地模型路径，防止在线下载
send "python inference_gen.py \\
    --vlm_path '$model_dir/vlm-model' \\
    --gen_lora_path '$model_dir/gen_lora' \\
    --base_model_path '$flux_dir' \\
    --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!" }
    "Saved" { puts "\n[SUCCESS] Image saved successfully!" }
    "#"
}

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

send "exit\r"
expect eof

puts "Test run finished."