#!/usr/bin/expect

set timeout -1
set host "156.225.27.52"
set port "5204"
set user "root"
set password $env(DEPLOY_SSH_PASSWORD)

puts "--> Connecting to finish installation..."

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"

puts "--> Installing PyTorch (retrying)..."
# 使用 pip 的 --timeout 参数防止超时
send "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 --timeout 1000\r"
expect "#"

puts "--> Installing other dependencies..."
send "pip install transformers diffusers accelerate sentencepiece protobuf huggingface_hub modelscope\r"
expect "#"

puts "--> Verifying installation..."
send "python -c 'import torch; print(\"CUDA available:\", torch.cuda.is_available())'\r"
expect "#"

send "exit\r"
expect eof

puts "Deployment fix completed."