| #!/usr/bin/env bash |
| |
| |
| set -Eeuo pipefail |
| cd "$(dirname "$0")/.." |
| source scripts/mv_env.sh |
| BACKBONE=${BACKBONE:-qwen21} |
| RESOLUTION=${RESOLUTION:-2048} |
| LEARNING_RATE=${LEARNING_RATE:-5e-5} |
| MAX_TRAIN_STEPS=${MAX_TRAIN_STEPS:-1000000} |
| CHECKPOINTING_STEPS=${CHECKPOINTING_STEPS:-2000} |
| TRAIN_JSON=${TRAIN_JSON:-data/train_demo.json} |
| EVAL_JSON=${EVAL_JSON:-data/eval_demo.json} |
| TEXT_COND=${TEXT_COND:-checkpoints/${BACKBONE}/text_cond.safetensors} |
| DECODER_CKPT=${DECODER_CKPT:-} |
| PROJECT_DIR=${PROJECT_DIR:-outputs/train_${BACKBONE}_albedo_res=${RESOLUTION}_lr=${LEARNING_RATE}} |
| SPARSE_ATTENTION=${SPARSE_ATTENTION:-true} |
| WORKERS=${WORKERS:-8} |
|
|
| for f in "${TEXT_COND}" "${TRAIN_JSON}" "${EVAL_JSON}"; do |
| [ -f "$f" ] || { echo "missing required file: $f (text cache: run scripts/precompute_text_mv.sh)" >&2; exit 1; } |
| done |
| EXTRA=() |
| [ -n "${DECODER_CKPT}" ] && EXTRA+=(--decoder_ckpt "${DECODER_CKPT}") |
| [ -n "${BUCKET_METADATA_JSON:-}" ] && EXTRA+=(--bucket_metadata_json "${BUCKET_METADATA_JSON}") |
| [ -n "${DATA_ROOT:-}" ] && EXTRA+=(--data_root "${DATA_ROOT}") |
| if [ -z "${RESUME_FROM_CHECKPOINT+x}" ]; then |
| LATEST=$(ls -d "${PROJECT_DIR}"/checkpoint-* 2>/dev/null | sed 's/.*checkpoint-//' | grep -E '^[0-9]+$' | sort -n | tail -1 || true) |
| [ -n "${LATEST}" ] && RESUME_FROM_CHECKPOINT="${PROJECT_DIR}/checkpoint-${LATEST}" |
| fi |
| [ -n "${RESUME_FROM_CHECKPOINT:-}" ] && EXTRA+=(--resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}") |
|
|
| echo "backbone=${BACKBONE} resolution=${RESOLUTION} lr=${LEARNING_RATE} steps=${MAX_TRAIN_STEPS} out=${PROJECT_DIR} resume=${RESUME_FROM_CHECKPOINT:-<none>}" |
| exec "${MV_PYTHON}" train_mv.py \ |
| --backbone "${BACKBONE}" \ |
| --text_cond_path "${TEXT_COND}" \ |
| --train_data_json "${TRAIN_JSON}" \ |
| --eval_data_json "${EVAL_JSON}" \ |
| --resolution "${RESOLUTION}" \ |
| --learning_rate "${LEARNING_RATE}" \ |
| --max_train_steps "${MAX_TRAIN_STEPS}" \ |
| --checkpointing_steps "${CHECKPOINTING_STEPS}" \ |
| --num_workers "${WORKERS}" \ |
| --project_dir "${PROJECT_DIR}" \ |
| --batch_size 1 --lora_rank 64 \ |
| --sparse_attention "${SPARSE_ATTENTION}" --sparse_attention_topk 0.20 --sparse_attention_blkq 128 --sparse_attention_blkk 64 \ |
| --inference_num_steps 25 --inference_drop_background_tokens true --inference_remove_background true \ |
| "${EXTRA[@]}" "$@" |
|
|
