File size: 1,290 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 | #!/bin/bash
# Precompute ctx (1 latent/frame) and target (1 latent/4 frames) latents for Context-as-Memory dataset.
#
# No metadata needed: auto-discovers segments from frames/ + captions.txt
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}"
DATASET_BASE="${DATASET_BASE_PATH:-${REPO_ROOT}/data/Context-as-Memory-Dataset}"
OUTPUT_DIR="${DATASET_BASE}/latents"
WAN_BASE_MODEL="${WAN_BASE_MODEL:-${REPO_ROOT}/checkpoints/Wan2.1-T2V-1.3B}"
MODEL_PATHS="${MODEL_PATHS:-[\"${WAN_BASE_MODEL}/diffusion_pytorch_model.safetensors\",\"${WAN_BASE_MODEL}/models_t5_umt5-xxl-enc-bf16.pth\",\"${WAN_BASE_MODEL}/Wan2.1_VAE.pth\"]}"
NUM_PROCESSES="${NUM_PROCESSES:-1}"
CONTEXT_FRAMES="${CONTEXT_FRAMES:-20}"
cd "${REPO_ROOT}"
accelerate launch --num_processes "${NUM_PROCESSES}" scripts/precompute_ctx_target_latents.py \
--dataset_base_path "${DATASET_BASE}" \
--output_dir "${OUTPUT_DIR}" \
--model_paths "${MODEL_PATHS}" \
--height 352 --width 640 \
--context_frames "${CONTEXT_FRAMES}" \
--target_frames_per_latent 4 \
--no_metadata \
--use_overlap_labels \
--overlap_labels_dense \
--skip_existing
echo "Done. Latents saved to ${OUTPUT_DIR}/ctx_latents/ and ${OUTPUT_DIR}/target_latents/"
|