from atomgit_hub import snapshot_download
snapshot_download("Ascend-SACT/tacotron2",  local_dir = './download')

# 1. 克隆代码库
cd /tmp
git clone https://github.com/biubug6/Pytorch_Retinaface.git
cd Pytorch_Retinaface

# 2. 安装必要依赖
pip install opencv-python-headless huggingface_hub
pip install torch==2.1.0 torchvision==0.16.0 torch_npu==2.1.0

# 1. 全局替换 .cuda() 为 .to("npu")
find . -name "*.py" | xargs sed -i 's/\.cuda()/.to("npu")/g'

# 2. 修复 train.py 中的 import 顺序 (避免 SyntaxError)
# 先删除第一行可能错误的插入，再在 correct 位置插入 torch_npu
sed -i '1d' train.py
sed -i '/from __future__ import print_function/a import torch_npu' train.py

# 3. 禁用 Cudnn Benchmark (NPU 不支持)
sed -i 's/cudnn\.benchmark = True/cudnn.benchmark = False/g' train.py

# 4. 开启 NPU 内存优化配置
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

python -c "
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
from huggingface_hub import hf_hub_download

print('⬇️ 正在下载 ResNet50 权重...')
try:
    os.makedirs('weights', exist_ok=True)
    hf_hub_download(repo_id='akhaliq/RetinaFace-R50', filename='RetinaFace_R50.pth', local_dir='./weights', local_dir_use_symlinks=False)
    # 重命名为代码预期的文件名
    if os.path.exists('weights/RetinaFace_R50.pth'):
        os.rename('weights/RetinaFace_R50.pth', 'weights/Resnet50_Final.pth')
    print('✅ 权重下载并重命名成功！')
except Exception as e:
    print(f'❌ 下载失败: {e}')
"

cat <<EOF > test_npu.py
import torch
import torch_npu
import numpy as np
import time
from models.retinaface import RetinaFace
from data import cfg_re50

def load_model(model, pretrained_path, device):
    pretrained_dict = torch.load(pretrained_path, map_location=device)
    if "state_dict" in pretrained_dict.keys():
        pretrained_dict = {k.replace('module.', ''): v for k, v in pretrained_dict['state_dict'].items()}
    else:
        pretrained_dict = {k.replace('module.', ''): v for k, v in pretrained_dict.items()}
    model.load_state_dict(pretrained_dict, strict=False)
    return model

if __name__ == '__main__':
    torch.set_grad_enabled(False)
    device = torch.device("npu:0")
    print(f"🚀 Device: {device}")

    # Init
    net = RetinaFace(cfg=cfg_re50, phase='test')
    net = load_model(net, './weights/Resnet50_Final.pth', device)
    net = net.to(device)
    net.eval()

    # Data
    img = torch.randn(1, 3, 640, 640).to(device) # Dummy image

    # Warmup
    print("🔥 Warming up...")
    _ = net(img)

    # Inference
    print("⚡ Running Inference...")
    t0 = time.time()
    _ = net(img)
    torch.npu.synchronize()
    t1 = time.time()
    
    print(f"⏱️ Cost: {(t1 - t0) * 1000:.2f} ms")
    print("🎉 Inference Validation Passed!")
EOF

# 运行验证
python test_npu.py

# 1. 下载数据脚本
python -c "
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
from huggingface_hub import hf_hub_download

repo_id = 'CUHK-CSE/wider_face'
files = ['data/WIDER_train.zip', 'data/WIDER_val.zip']

print('🚀 正在下载数据集 (约 3GB)...')
for f in files:
    try:
        hf_hub_download(repo_id=repo_id, filename=f, repo_type='dataset', local_dir='.', local_dir_use_symlinks=False)
        print(f'✅ {f} 下载完成')
    except Exception as e:
        print(f'❌ {f} 失败: {e}')
"

# 2. 解压并整理目录
mkdir -p data/widerface/train
mkdir -p data/widerface/val
unzip -q data/WIDER_train.zip -d data/widerface/train/
mv data/widerface/train/WIDER_train/images data/widerface/train/images

# 1. 下载标签文件
python -c "
import urllib.request
url = 'https://mirror.ghproxy.com/https://raw.githubusercontent.com/wujixs/RetinaFace-Pytorch/master/data/widerface/train/label.txt'
try:
    urllib.request.urlretrieve(url, 'data/widerface/train/label.txt')
    print('✅ Label 下载成功')
except Exception as e:
    print(f'❌ Label 下载失败: {e}')
"
# 2. 裁剪数据集 (仅保留 50 张图)
python -c "
path = 'data/widerface/train/label.txt'
with open(path, 'r') as f:
    lines = f.readlines()
new_lines = []
count = 0
for line in lines:
    if line.startswith('#'):
        count += 1
        if count > 50: break
    new_lines.append(line)
with open(path, 'w') as f:
    f.writelines(new_lines)
print('✂️ 数据集已裁剪为 50 张图片')
"

# 1. 修改 Batch Size 为 8
sed -i "s/'batch_size': [0-9]\+/'batch_size': 8/g" data/config.py

# 2. 启动训练 (50 Epoch)
# --save_folder: 权重保存路径
python train.py --network resnet50 --num_workers 4 --save_folder ./weights/