Add scripts for extracting EasyCom session features: run_extract_easycom_direct.sh for full session extraction and run_smoke_easycom_direct.sh for smoke testing on a single GPU. Both scripts utilize streaming_turn_model_directly.py for processing and include necessary configurations and logging. This enhances the EasyCom pipeline for better session management and feature extraction.

2026-08-17 22:16:27 +08:00
parent de29e75734
commit 8145be4108
2 changed files with 200 additions and 0 deletions
@@ -0,0 +1,134 @@
+#!/usr/bin/env bash
+# Extract open-stream Thinker states for all 12 EasyCom sessions via
+# streaming_turn_model_directly.py (original clips, no prepared 2 FPS MP4).
+#
+# Recommended for full sessions: --memory-seconds 20 (bounded KV).
+# Do not default to --memory-seconds 0 on multi-minute sessions.
+set -euo pipefail
+
+PROJECT_ROOT="${PROJECT_ROOT:-/mnt/58462F4C462F29E8/Code/icassp2027-experiment}"
+MODEL="${MODEL:-/mnt/58462F4C462F29E8/Models/Qwen2.5-Omni-3B}"
+VIDEO_ROOT="${VIDEO_ROOT:-/mnt/58462F4C462F29E8/Datasets/easycom/Main/Video_Compressed}"
+TARGET_ROOT="${TARGET_ROOT:-/mnt/58462F4C462F29E8/Code/EgoSpeak/dataset/EasyCom/target_perframe}"
+FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}"
+LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}"
+
+MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
+NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}"
+VIDEO_FPS="${VIDEO_FPS:-2}"
+MODE="${MODE:-open-stream-replay}"
+DTYPE="${DTYPE:-bfloat16}"
+# Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time.
+GPUS="${GPUS:-0}"
+# Optional: only run selected session ids, e.g. SESSIONS="1 2 3"
+SESSIONS="${SESSIONS:-1 2 3 4 5 6 7 8 9 10 11 12}"
+
+# StreamConfig requires memory / native_block to be an integer (e.g. 20/2=10).
+# Values like 15 yield 7.5 and crash before loading the model.
+python - "${MEMORY_SECONDS}" "${NATIVE_BLOCK_SECONDS}" <<'PY'
+import sys
+memory, block = float(sys.argv[1]), float(sys.argv[2])
+ratio = memory / block
+if abs(ratio - round(ratio)) > 1e-6:
+    raise SystemExit(
+        f"MEMORY_SECONDS={memory} is invalid: memory/native_block={ratio} "
+        f"(native_block={block}). Use a multiple of {block:g}, e.g. 16, 18, 20."
+    )
+print(f"memory_seconds={memory:g} native_block={block:g} blocks={int(round(ratio))}")
+PY
+
+mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
+cd "${PROJECT_ROOT}"
+
+run_one() {
+  local gpu="$1"
+  local sid="$2"
+  local session="Session_${sid}"
+  local session_dir="${VIDEO_ROOT}/${session}"
+  local target="${TARGET_ROOT}/${session}.npy"
+  local output="${FEATURE_ROOT}/${session}.pt"
+  local log="${LOG_ROOT}/${session}.log"
+
+  if [[ ! -d "${session_dir}" ]]; then
+    echo "Missing session dir: ${session_dir}" >&2
+    return 1
+  fi
+  if [[ ! -s "${target}" ]]; then
+    echo "Missing target: ${target}" >&2
+    return 1
+  fi
+  if [[ -s "${output}" ]]; then
+    echo "Skipping existing output: ${output}"
+    return 0
+  fi
+
+  {
+    echo "$(date): GPU ${gpu} starting ${session}"
+    echo "model=${MODEL}"
+    echo "session_dir=${session_dir}"
+    echo "target=${target}"
+    echo "output=${output}"
+    echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}"
+
+    if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
+      python streaming_turn_model_directly.py \
+        --model-id "${MODEL}" \
+        --easycom-session-dir "${session_dir}" \
+        --target "${target}" \
+        --use-audio-in-video \
+        --mode "${MODE}" \
+        --memory-seconds "${MEMORY_SECONDS}" \
+        --native-block-seconds "${NATIVE_BLOCK_SECONDS}" \
+        --video-fps "${VIDEO_FPS}" \
+        --dtype "${DTYPE}" \
+        --flash-attention-2 \
+        --device cuda:0 \
+        --output "${output}"
+    then
+      echo "$(date): GPU ${gpu} FAILED ${session}"
+      exit 1
+    fi
+
+    echo "$(date): GPU ${gpu} completed ${session}"
+  } >"${log}" 2>&1
+}
+
+# Round-robin sessions across GPUs.
+mapfile -t GPU_ARR < <(echo "${GPUS}" | tr ' ' '\n' | sed '/^$/d')
+mapfile -t SID_ARR < <(echo "${SESSIONS}" | tr ' ' '\n' | sed '/^$/d')
+
+declare -A WORKERS=()
+for gpu in "${GPU_ARR[@]}"; do
+  WORKERS["${gpu}"]=""
+done
+
+i=0
+for sid in "${SID_ARR[@]}"; do
+  gpu="${GPU_ARR[$((i % ${#GPU_ARR[@]}))]}"
+  WORKERS["${gpu}"]+="${sid} "
+  i=$((i + 1))
+done
+
+pids=()
+status=0
+for gpu in "${GPU_ARR[@]}"; do
+  sids="${WORKERS[${gpu}]}"
+  [[ -z "${sids// }" ]] && continue
+  (
+    for sid in ${sids}; do
+      run_one "${gpu}" "${sid}" || exit 1
+    done
+  ) &
+  pids+=("$!")
+done
+
+for pid in "${pids[@]}"; do
+  wait "${pid}" || status=1
+done
+
+if [[ "${status}" -ne 0 ]]; then
+  echo "At least one EasyCom direct extraction failed. Check ${LOG_ROOT}." >&2
+  exit 1
+fi
+
+echo "ALL EasyCom direct extractions completed -> ${FEATURE_ROOT}"