File size: 1,920 Bytes
c335050 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | #!/bin/bash
# Run all memory baseline inferences sequentially.
# Requires: CKPT_DIR pointing to a directory with per-row checkpoint folders
# (same layout as HF repo: context_k1/epoch-0.safetensors, etc.)
#
# Usage:
# CKPT_DIR=./ckpts bash inference/memory_baselines_basic/run_infer_all.sh
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
: "${CKPT_DIR:?ERROR: set CKPT_DIR to your checkpoints directory}"
declare -A CKPT_MAP=(
[run_infer_no_memory.sh]="no_memory_extra_two_chunk"
[run_infer_framepack_weight.sh]="framepack_weight_only"
[run_infer_framepack_len_r2.sh]="framepack_lencompress_r2"
[run_infer_framepack_len_r4.sh]="framepack_lencompress_r4"
[run_infer_framepack_hybrid_r2.sh]="framepack_hybrid_r2_weight_two_chunk"
[run_infer_framepack_hybrid_r4.sh]="framepack_hybrid_r4_weight_two_chunk"
[run_infer_spatial_mem.sh]="spatial_mem"
[run_infer_spatial_concat_text.sh]="spatial_concat_text_two_chunk"
[run_infer_spatial_inject_none.sh]="spatial_inject_none_two_chunk"
[run_infer_spatial_cross_attn_readout.sh]="spatial_cross_attn_readout_two_chunk"
[run_infer_videossm_hybrid.sh]="videossm_hybrid"
[run_infer_block_wise_ssm.sh]="block_wise_ssm_two_chunk"
)
for script in "${SCRIPT_DIR}"/run_infer_*.sh; do
name="$(basename "${script}")"
[[ "${name}" == "run_infer_all.sh" ]] && continue
ckpt_folder="${CKPT_MAP[${name}]:-}"
if [ -z "${ckpt_folder}" ]; then
echo "[run_infer_all] SKIP: no checkpoint mapping for ${name}"
continue
fi
ckpt_file="${CKPT_DIR}/${ckpt_folder}/epoch-0.safetensors"
if [ ! -f "${ckpt_file}" ]; then
echo "[run_infer_all] SKIP: ${ckpt_file} not found"
continue
fi
echo "========================================"
echo "[run_infer_all] Running ${name} with ${ckpt_file}"
echo "========================================"
CKPT="${ckpt_file}" bash "${script}"
done
echo "[run_infer_all] Done."
|