#!/usr/bin/env bash
set -euo pipefail

UMM_ROOT="${UMM_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
VIDEO_GEN_ROOT="${VIDEO_GEN_ROOT:-/rczhang/rczhang/zhangrch/video_gen}"
HF_CACHE_ROOT="${HF_CACHE_ROOT:-${VIDEO_GEN_ROOT}/omnigen/cache/huggingface/hub}"

unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY
mkdir -p \
  "${UMM_ROOT}/checkpoints/stage1" \
  "${UMM_ROOT}/models/base/Qwen3-VL-4B-Instruct" \
  "${UMM_ROOT}/models/teachers/VGGT-1B" \
  "${UMM_ROOT}/models/teachers/siglip2-large-patch16-256" \
  "${UMM_ROOT}/models/teachers/dinov2-large-processor" \
  "${UMM_ROOT}/cache/torch/hub/checkpoints" \
  "${UMM_ROOT}/cache/torch/hub/facebookresearch_dinov2_main" \
  "${UMM_ROOT}/data/covt"

rsync -a "${VIDEO_GEN_ROOT}/gen_checkpoint/stage1_algin_v2/embeddings_connector/checkpoint-32000.bin" \
  "${UMM_ROOT}/checkpoints/stage1/checkpoint-32000.bin"
rsync -a "${VIDEO_GEN_ROOT}/pretrined_models/Qwen/Qwen3-VL-4B-Instruct/" \
  "${UMM_ROOT}/models/base/Qwen3-VL-4B-Instruct/"
rsync -a "${VIDEO_GEN_ROOT}/pretrined_models/covt_moe/VGGT-1B/model.pt" \
  "${UMM_ROOT}/models/teachers/VGGT-1B/model.pt"
rsync -a "${VIDEO_GEN_ROOT}/pretrined_models/covt_moe/sam_vit_h_4b8939.pth" \
  "${UMM_ROOT}/models/teachers/sam_vit_h_4b8939.pth"
rsync -a "${VIDEO_GEN_ROOT}/pretrined_models/covt_moe/table5_baseline.pth" \
  "${UMM_ROOT}/models/teachers/table5_baseline.pth"

rsync -aL \
  "${HF_CACHE_ROOT}/models--google--siglip2-large-patch16-256/snapshots/787800c8990e6f058423089178e718139608408c/" \
  "${UMM_ROOT}/models/teachers/siglip2-large-patch16-256/"
rsync -aL \
  "${HF_CACHE_ROOT}/models--facebook--dinov2-large/snapshots/47b73eefe95e8d44ec3623f8890bd894b6ea2d6c/" \
  "${UMM_ROOT}/models/teachers/dinov2-large-processor/"
rsync -a "/root/.cache/torch/hub/checkpoints/dinov2_vitl14_pretrain.pth" \
  "${UMM_ROOT}/cache/torch/hub/checkpoints/dinov2_vitl14_pretrain.pth"
rsync -a --exclude='.git' --exclude='__pycache__' "/root/.cache/torch/hub/facebookresearch_dinov2_main/" \
  "${UMM_ROOT}/cache/torch/hub/facebookresearch_dinov2_main/"

"/rczhang/rczhang/conda_packs/unvideo/bin/python" \
  "${UMM_ROOT}/scripts/make_portable_data.py" \
  "${VIDEO_GEN_ROOT}/gen_dataset/covt/data.json" \
  "${VIDEO_GEN_ROOT}/gen_dataset/covt/images" \
  "${UMM_ROOT}/data/covt/data.json"

echo "Local model, checkpoint, teacher and metadata assets collected"