初始化项目，由ModelHub XC社区提供模型
Model: EmpathicRobotics/vla-1.7b-qwen3-v2 Source: Original Platform
This commit is contained in:
ModelHub XC
309
pipeline_pose/phase5_adaptive_pchip.py
Normal file
309
pipeline_pose/phase5_adaptive_pchip.py
Normal file
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
Phase 5 — Adaptive PCHIP per-joint tokenizer.
|
||||
|
||||
For each window from Phase 4 (WINDOW_FRAMES frames -- 8 originally, 24 as of
|
||||
2026-07-22, see REPORT.md #38), each of the 17 joints gets an independent
|
||||
PCHIP compression with adaptive control-point count based on per-joint
|
||||
curvature over the WHOLE window (all WINDOW_FRAMES frames considered, not a
|
||||
fixed sub-grid):
|
||||
|
||||
Tier 2 (2 CPs: start + end) — max curvature < tau_low
|
||||
Tier 4 (4 CPs: start + end + top-2 interior) — tau_low <= curvature < tau_high
|
||||
Tier MAX_CPS (MAX_CPS CPs: start+end+top-(MAX_CPS-2) interior, chosen by
|
||||
curvature out of ALL WINDOW_FRAMES candidates) — curvature >= tau_high
|
||||
|
||||
MAX_CPS is fixed at 8 regardless of WINDOW_FRAMES -- widening the window
|
||||
gives the top tier more candidate positions to pick its best 8 from, not more
|
||||
tokens/joint. See MAX_CPS's own docstring below for why this matters.
|
||||
|
||||
Token stream per window (WINDOW_FRAMES=24 example; t values now range over
|
||||
however many of the window's real frame indices got chosen as CPs, e.g.
|
||||
t_0 and t_23 for tier 2, or t_0/t_5/t_14/t_23 for tier 4 if frames 5 and 14
|
||||
had the highest curvature):
|
||||
<fps_30>
|
||||
<pelvis> <pelvis_t_0> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N>
|
||||
<pelvis_t_23> <pelvis_x_N> <pelvis_y_N> <pelvis_z_N> </pelvis>
|
||||
<r_hip> <r_hip_t_0> <r_hip_x_N> ... </r_hip>
|
||||
...
|
||||
|
||||
Quantization: [-2.0 m, +2.0 m] -> [0, 255] (precision ~15.7 mm)
|
||||
Time tokens: real frame index within the window, 0 to WINDOW_FRAMES-1 --
|
||||
requires <{joint}_t_N> tokens up to WINDOW_FRAMES-1 in the
|
||||
tokenizer vocab (only 0-7 existed before 2026-07-22).
|
||||
|
||||
Input: outputs/yolo_cleaned_30fps/{video_id}_cleaned.jsonl
|
||||
Output: outputs/agent_tokens_adaptive/{video_id}_tokens.jsonl
|
||||
Each line: {"video_id", "window_id", "fps", "token_str", "cp_counts"}
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
# ── Constants ─────────────────────────────────────────────────────────────────
|
||||
|
||||
TARGET_FPS = 30
|
||||
WINDOW_FRAMES = 8
|
||||
N_JOINTS = 17
|
||||
COORD_RANGE = 2.0
|
||||
STRIDE = 8
|
||||
# 2026-07-22 (REPORT.md #38): cap on control points per joint, independent of
|
||||
# WINDOW_FRAMES. Before this change WINDOW_FRAMES==MAX_CPS==8 always (the top
|
||||
# tier was literally "use every frame"), so widening the window to 24 frames
|
||||
# would have silently tripled worst-case tokens/joint (24 CPs instead of 8)
|
||||
# with no code change needed to trigger it. Keeping MAX_CPS fixed at 8 means
|
||||
# the top tier now means "pick the best 8 of WINDOW_FRAMES candidates by
|
||||
# curvature" instead of "use all of them" -- same worst-case token cost as
|
||||
# before, but the 8 chosen points can be anywhere in the (now wider) window
|
||||
# instead of forced onto a fixed 8-slot grid. This is the whole point of
|
||||
# "Option 2" (dense pose, no subsampling before curve-fitting) agreed with
|
||||
# the user: Phase 3 keeps every real frame, Phase 5 decides freely which
|
||||
# frames matter most.
|
||||
MAX_CPS = 8
|
||||
|
||||
TAU_LOW = 0.005
|
||||
TAU_HIGH = 0.05
|
||||
|
||||
JOINT_NAMES = [
|
||||
"pelvis", "r_hip", "r_knee", "r_ankle",
|
||||
"l_hip", "l_knee", "l_ankle",
|
||||
"spine", "thorax", "nose", "head_top",
|
||||
"l_shoulder", "l_elbow", "l_wrist",
|
||||
"r_shoulder", "r_elbow", "r_wrist",
|
||||
]
|
||||
|
||||
# ── Quantization ──────────────────────────────────────────────────────────────
|
||||
|
||||
def quantize(v: float) -> int:
|
||||
return int(np.clip(round((v + COORD_RANGE) / (2.0 * COORD_RANGE) * 255), 0, 255))
|
||||
|
||||
|
||||
def dequantize(n: int) -> float:
|
||||
return n / 255.0 * (2.0 * COORD_RANGE) - COORD_RANGE
|
||||
|
||||
|
||||
# ── Per-joint adaptive CP selection ───────────────────────────────────────────
|
||||
|
||||
def joint_curvature(trajectory: np.ndarray) -> float:
|
||||
"""Max curvature (acceleration norm) for a single joint trajectory (8, 3)."""
|
||||
if trajectory.shape[0] < 3:
|
||||
return 0.0
|
||||
vel = np.diff(trajectory, axis=0)
|
||||
acc = np.diff(vel, axis=0)
|
||||
return float(np.max(np.linalg.norm(acc, axis=1)))
|
||||
|
||||
|
||||
def select_cp_indices(trajectory: np.ndarray, tau_low: float, tau_high: float) -> np.ndarray:
|
||||
"""Choose which frame indices become control points for one joint.
|
||||
|
||||
Tier sizes are fixed at 2 / 4 / MAX_CPS regardless of how many frames are
|
||||
in the window (see MAX_CPS's docstring) -- the top tier picks the
|
||||
MAX_CPS-2 highest-curvature *interior* frames out of every candidate in
|
||||
the window, not a fixed grid position. With WINDOW_FRAMES==MAX_CPS==8
|
||||
(the original config) this is exactly equivalent to the old
|
||||
`np.arange(WINDOW_FRAMES)` behavior, since "top 6 of 6 interior
|
||||
candidates" is all of them.
|
||||
"""
|
||||
curv = joint_curvature(trajectory)
|
||||
n_frames = trajectory.shape[0]
|
||||
|
||||
if curv < tau_low:
|
||||
return np.array([0, n_frames - 1])
|
||||
|
||||
n_interior = 2 if curv < tau_high else (MAX_CPS - 2)
|
||||
n_interior = min(n_interior, max(n_frames - 2, 0))
|
||||
|
||||
vel = np.diff(trajectory, axis=0)
|
||||
acc = np.diff(vel, axis=0)
|
||||
acc_norms = np.linalg.norm(acc, axis=1) # (n_frames-2,)
|
||||
|
||||
# acc[i] corresponds to frame i+1 (second derivative offset)
|
||||
interior_curv = np.zeros(n_frames)
|
||||
for i in range(len(acc_norms)):
|
||||
interior_curv[i + 1] = acc_norms[i]
|
||||
|
||||
# Exclude endpoints (already included), pick top-n_interior interior frames
|
||||
interior_curv[0] = -1.0
|
||||
interior_curv[-1] = -1.0
|
||||
top_n = np.argsort(interior_curv)[-n_interior:] if n_interior > 0 else np.array([], dtype=int)
|
||||
|
||||
indices = np.unique(np.sort(np.concatenate(([0], top_n, [n_frames - 1]))))
|
||||
return indices.astype(int)
|
||||
|
||||
|
||||
# ── Token builder ─────────────────────────────────────────────────────────────
|
||||
|
||||
def build_token_str(
|
||||
states: np.ndarray,
|
||||
fps: int = TARGET_FPS,
|
||||
tau_low: float = TAU_LOW,
|
||||
tau_high: float = TAU_HIGH,
|
||||
) -> tuple:
|
||||
"""
|
||||
states : (8, 17, 3) float32, root-centred metric coordinates
|
||||
Returns (token_str, cp_counts_dict)
|
||||
"""
|
||||
parts = [f"<fps_{fps}>"]
|
||||
cp_counts = {}
|
||||
|
||||
for j in range(N_JOINTS):
|
||||
name = JOINT_NAMES[j]
|
||||
trajectory = states[:, j, :] # (8, 3)
|
||||
cp_idx = select_cp_indices(trajectory, tau_low, tau_high)
|
||||
cp_counts[name] = len(cp_idx)
|
||||
|
||||
parts.append(f"<{name}>")
|
||||
for fi in cp_idx:
|
||||
x, y, z = trajectory[fi]
|
||||
parts.append(f"<{name}_t_{fi}>")
|
||||
parts.append(f"<{name}_x_{quantize(x)}>")
|
||||
parts.append(f"<{name}_y_{quantize(y)}>")
|
||||
parts.append(f"<{name}_z_{quantize(z)}>")
|
||||
parts.append(f"</{name}>")
|
||||
|
||||
return " ".join(parts), cp_counts
|
||||
|
||||
|
||||
# ── Per-file processing ──────────────────────────────────────────────────────
|
||||
|
||||
def process_file(
|
||||
input_path: str,
|
||||
output_jsonl: str,
|
||||
video_id: str,
|
||||
stride: int = STRIDE,
|
||||
tau_low: float = TAU_LOW,
|
||||
tau_high: float = TAU_HIGH,
|
||||
) -> int:
|
||||
records = []
|
||||
|
||||
with open(input_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
data = json.loads(line)
|
||||
window_id = int(data["window_id"])
|
||||
|
||||
if window_id % stride != 0:
|
||||
continue
|
||||
|
||||
states = np.array(data["states"], dtype=np.float32)
|
||||
|
||||
if states.shape != (WINDOW_FRAMES, N_JOINTS, 3):
|
||||
continue
|
||||
if np.isnan(states).any():
|
||||
continue
|
||||
|
||||
token_str, cp_counts = build_token_str(states, TARGET_FPS, tau_low, tau_high)
|
||||
|
||||
records.append({
|
||||
"video_id": video_id,
|
||||
"window_id": window_id,
|
||||
"fps": TARGET_FPS,
|
||||
"token_str": token_str,
|
||||
"cp_counts": cp_counts,
|
||||
})
|
||||
|
||||
if not records:
|
||||
return 0
|
||||
|
||||
tmp = output_jsonl + ".tmp"
|
||||
os.makedirs(os.path.dirname(output_jsonl), exist_ok=True)
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
for rec in records:
|
||||
f.write(json.dumps(rec, ensure_ascii=False) + "\n")
|
||||
os.replace(tmp, output_jsonl)
|
||||
|
||||
return len(records)
|
||||
|
||||
|
||||
# ── Entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(
|
||||
description="Phase 5 — Adaptive PCHIP per-joint tokenizer."
|
||||
)
|
||||
p.add_argument("--input-dir", required=True,
|
||||
help="Directory with *_cleaned.jsonl from Phase 4.")
|
||||
p.add_argument("--output-dir", required=True,
|
||||
help="Directory to write *_tokens.jsonl files.")
|
||||
p.add_argument("--stride", type=int, default=STRIDE,
|
||||
help=f"Keep windows where window_id %% stride == 0. Default: {STRIDE}")
|
||||
p.add_argument("--tau-low", type=float, default=TAU_LOW,
|
||||
help=f"Curvature threshold for 2-CP tier. Default: {TAU_LOW}")
|
||||
p.add_argument("--tau-high", type=float, default=TAU_HIGH,
|
||||
help=f"Curvature threshold for 8-CP tier. Default: {TAU_HIGH}")
|
||||
p.add_argument("--file-list", default=None,
|
||||
help="Optional text file listing specific *_cleaned.jsonl paths.")
|
||||
p.add_argument("--window-frames", type=int, default=WINDOW_FRAMES,
|
||||
help=f"Must match Phase 3/4's --window-size. Default: {WINDOW_FRAMES}. "
|
||||
f"2026-07-22: use 24 to match the wider cosmos chunk window -- "
|
||||
f"see REPORT.md #38. MAX_CPS (top-tier control-point count) is "
|
||||
f"NOT tied to this and stays 8 either way.")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
global WINDOW_FRAMES
|
||||
args = parse_args()
|
||||
WINDOW_FRAMES = args.window_frames
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
if args.file_list:
|
||||
with open(args.file_list) as f:
|
||||
all_files = [l.strip() for l in f if l.strip()]
|
||||
else:
|
||||
all_files = sorted(glob.glob(os.path.join(args.input_dir, "*_cleaned.jsonl")))
|
||||
|
||||
task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", 0))
|
||||
num_tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", 1))
|
||||
|
||||
my_files = [f for i, f in enumerate(all_files) if i % num_tasks == task_id]
|
||||
total = len(my_files)
|
||||
|
||||
print(f"\n[Worker {task_id}/{num_tasks}] {total} files to process.")
|
||||
print("=" * 60)
|
||||
|
||||
processed = skipped = empty = 0
|
||||
tier_counts = {2: 0, 4: 0, 8: 0}
|
||||
|
||||
for idx, input_path in enumerate(my_files, start=1):
|
||||
base = os.path.basename(input_path)
|
||||
video_id = base[: -len("_cleaned.jsonl")]
|
||||
|
||||
out_jsonl = os.path.join(args.output_dir, f"{video_id}_tokens.jsonl")
|
||||
|
||||
if os.path.exists(out_jsonl):
|
||||
skipped += 1
|
||||
print(f"[{idx}/{total}] {video_id} — already done", end="\r")
|
||||
continue
|
||||
|
||||
try:
|
||||
n = process_file(
|
||||
input_path, out_jsonl, video_id,
|
||||
stride=args.stride,
|
||||
tau_low=args.tau_low,
|
||||
tau_high=args.tau_high,
|
||||
)
|
||||
if n > 0:
|
||||
processed += 1
|
||||
pct = (processed + skipped + empty) / total * 100
|
||||
print(f"[{idx}/{total}] {pct:.1f}% | {video_id} — {n} windows")
|
||||
else:
|
||||
empty += 1
|
||||
except Exception as e:
|
||||
print(f"[{idx}/{total}] ERROR {video_id} — {e}")
|
||||
for p in (out_jsonl + ".tmp",):
|
||||
if os.path.exists(p):
|
||||
os.remove(p)
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"[Worker {task_id}] done — processed: {processed}, skipped: {skipped}, empty: {empty}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
582
pipeline_pose/snac_finevideo.py
Normal file
582
pipeline_pose/snac_finevideo.py
Normal file
@@ -0,0 +1,582 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SNAC tokenization for FineVideo-VLA activities.
|
||||
|
||||
Reads final_dataset_adaptive JSONL files, extracts audio from .mp4 videos,
|
||||
and tokenizes each activity segment with SNAC_24kHz in "listen" format.
|
||||
|
||||
Listen format — 3 tokens per SNAC base frame (base rate = 12.5 Hz → 37.5 tokens/sec):
|
||||
token_1 = codes[0][i] + 128266 → <snac_128266> .. <snac_132361>
|
||||
token_2 = codes[1][2*i] + 128266 + 4096 → <snac_132362> .. <snac_136457>
|
||||
token_3 = codes[1][2*i+1] + 128266 +16384 → <snac_144650> .. <snac_148745>
|
||||
|
||||
Same Orpheus offset scheme as MixtureVitae-Omni → tokens are directly compatible.
|
||||
Total unique SNAC token strings: 3 × 4096 = 12,288.
|
||||
|
||||
Output: {OUTPUT_DIR}/{video_id}_snac.jsonl
|
||||
One line per activity:
|
||||
{
|
||||
"video_id":"...", "activity_id":"...", "start_sec":1.0, "end_sec":8.9,
|
||||
"has_agent": true,
|
||||
"snac_by_chunk": {
|
||||
"0": ["<snac_130055>", "<snac_133001>", "<snac_145000>", ...], // ~9-10 tokens
|
||||
"1": [...],
|
||||
...
|
||||
}
|
||||
}
|
||||
|
||||
snac_by_chunk keys are chunk_idx (integer as string), aligned to the same 8-frame
|
||||
grid as cosmos/avclm/agent. Phase7 reads chunk_idx → snac tokens directly.
|
||||
|
||||
Two modes:
|
||||
--build-tasks Scan all final_dataset_adaptive files, write snac_task_list.json.
|
||||
Run once on login node (or task 0) before the array job.
|
||||
(default) Load task list, process this SLURM task's slice of videos.
|
||||
|
||||
SLURM usage:
|
||||
SLURM_ARRAY_TASK_ID = task index (0-based)
|
||||
SLURM_ARRAY_TASK_COUNT = total number of tasks in the array
|
||||
|
||||
Local test:
|
||||
python pipeline_pose/snac_finevideo.py --build-tasks
|
||||
python pipeline_pose/snac_finevideo.py # task_id=0, num_tasks=1
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import multiprocessing
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# ── Paths ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
VIDEO_DIR = "/e/data1/datasets/playground/mmlaion/shared/nguyen38/videos_staging"
|
||||
# Still scanned for activity time_range_sec + has_agent -- those don't depend
|
||||
# on window size. chunk_timing (per-window breakdown) is IGNORED (see
|
||||
# _scan_one_rank_file docstring) because this file is window=8-based (the
|
||||
# pre-pivot merge) and Phase 6 hasn't rerun at window=24 yet; n_chunks is
|
||||
# instead recomputed independently from CHUNK_SIZE below, same pattern as
|
||||
# data_prep/omnivideo_100k/snac_omnivideo.py.
|
||||
INPUT_GLOB = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/"
|
||||
"final_dataset_adaptive/final_vla_adaptive_rank_*.jsonl")
|
||||
OUTPUT_DIR = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/snac_tokens_w24")
|
||||
TASK_CACHE = ("/e/data1/datasets/playground/mmlaion/shared/nguyen38/FineVideo-VLA/"
|
||||
"snac_task_list_w24.json")
|
||||
HF_CACHE = "/e/project1/reformo/nguyen38/jupiter_cache/huggingface"
|
||||
SNAC_MODEL = "hubertsiuzdak/snac_24khz"
|
||||
SAMPLE_RATE = 24000
|
||||
TARGET_FPS = 30
|
||||
CHUNK_SIZE = 24 # 2026-07-23 window=24 pivot -- must match step_a_tokenize_video.py's CHUNK_SIZE
|
||||
|
||||
# ── SNAC listen-format offsets (matches MixtureVitae-Omni) ───────────────────
|
||||
|
||||
OFFSET_L0 = 128266 # codes[0] base
|
||||
OFFSET_L1A = 128266 + 4096 # codes[1] even frames → 132362
|
||||
OFFSET_L1B = 128266 + 4 * 4096 # codes[1] odd frames → 144650
|
||||
|
||||
# 2026-07-23: speak-format offsets for codes[2] (fine, 50Hz, 4 sub-positions
|
||||
# per base frame) -- corrected to match the REAL scheme in Huu/Chien's
|
||||
# production snac_gpu.py on Leonardo (pipeline_video/snac_gpu.py), the
|
||||
# Orpheus-standard SNAC packing layout. Sub-codes 0/1 sit between L1A and
|
||||
# L1B (in what was wrongly assumed to be an unused gap); sub-codes 2/3 sit
|
||||
# after L1B. See data_prep/laion_emotional_roleplay/tokenize_snac.py's
|
||||
# OFFSET_L2 docstring for the full correction history. Must match exactly.
|
||||
OFFSET_L2 = [136458, 140554, 148746, 152842]
|
||||
|
||||
# ── Logging ───────────────────────────────────────────────────────────────────
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s %(levelname)s %(message)s",
|
||||
datefmt="%H:%M:%S",
|
||||
stream=sys.stdout,
|
||||
)
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Audio extraction
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _find_ffmpeg() -> str:
|
||||
"""Return path to ffmpeg binary, trying imageio_ffmpeg as fallback."""
|
||||
try:
|
||||
subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
|
||||
return "ffmpeg"
|
||||
except (FileNotFoundError, subprocess.CalledProcessError):
|
||||
pass
|
||||
try:
|
||||
import imageio_ffmpeg
|
||||
return imageio_ffmpeg.get_ffmpeg_exe()
|
||||
except ImportError:
|
||||
pass
|
||||
raise RuntimeError("ffmpeg not found. Install ffmpeg or imageio_ffmpeg.")
|
||||
|
||||
_FFMPEG = None
|
||||
|
||||
def get_ffmpeg() -> str:
|
||||
global _FFMPEG
|
||||
if _FFMPEG is None:
|
||||
_FFMPEG = _find_ffmpeg()
|
||||
return _FFMPEG
|
||||
|
||||
|
||||
def extract_full_audio(video_path: str) -> np.ndarray | None:
|
||||
"""
|
||||
Extract full mono 24 kHz PCM audio from a video file.
|
||||
Pipes raw float32 PCM from ffmpeg directly to a numpy array — no temp files.
|
||||
Returns float32 array or None on any failure.
|
||||
"""
|
||||
cmd = [
|
||||
get_ffmpeg(), "-y",
|
||||
"-i", video_path,
|
||||
"-vn", # strip video
|
||||
"-ac", "1", # mono
|
||||
"-ar", str(SAMPLE_RATE), # 24 kHz
|
||||
"-f", "f32le", # raw float32 PCM
|
||||
"-",
|
||||
]
|
||||
try:
|
||||
result = subprocess.run(cmd, capture_output=True, timeout=300)
|
||||
if result.returncode != 0 or not result.stdout:
|
||||
return None
|
||||
audio = np.frombuffer(result.stdout, dtype=np.float32).copy()
|
||||
return audio if len(audio) > 0 else None
|
||||
except Exception as e:
|
||||
log.debug(f"ffmpeg failed for {video_path}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def slice_audio(
|
||||
audio: np.ndarray,
|
||||
start_sec: float,
|
||||
end_sec: float,
|
||||
sr: int = SAMPLE_RATE,
|
||||
) -> np.ndarray:
|
||||
"""Slice a float32 audio array to [start_sec, end_sec]."""
|
||||
s = max(0, int(start_sec * sr))
|
||||
e = min(len(audio), int(end_sec * sr))
|
||||
return audio[s:e]
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# SNAC tokenization
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def encode_listen(audio: np.ndarray, model, device: str) -> list[str]:
|
||||
"""
|
||||
Encode a float32 audio array with SNAC_24kHz, return listen-format tokens.
|
||||
|
||||
Listen format (3 tokens per base frame):
|
||||
<snac_{codes[0][i] + 128266}>
|
||||
<snac_{codes[1][2i] + 132362}>
|
||||
<snac_{codes[1][2i+1] + 144650}>
|
||||
|
||||
SNAC_24kHz hierarchy:
|
||||
codes[0] — base codebook, 12.5 Hz
|
||||
codes[1] — mid codebook, 25.0 Hz (2× codes[0])
|
||||
codes[2] — fine codebook, 50.0 Hz (4× codes[0], not used in listen)
|
||||
|
||||
Listen format ignores codes[2] (fine detail) to keep token count low
|
||||
(~37.5 tokens/sec vs 87.5 for full speak format). Matches MV-Omni.
|
||||
"""
|
||||
tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device) # (1,1,T)
|
||||
with torch.inference_mode():
|
||||
codes = model.encode(tensor) # list: [codes[0], codes[1], codes[2]]
|
||||
|
||||
c0 = codes[0] # (1, N0)
|
||||
c1 = codes[1] # (1, N1), N1 == 2*N0
|
||||
|
||||
n0 = c0.shape[1]
|
||||
tokens: list[str] = []
|
||||
for i in range(n0):
|
||||
i1a = 2 * i
|
||||
i1b = 2 * i + 1
|
||||
if i1b >= c1.shape[1]:
|
||||
break # boundary guard: shouldn't happen for valid audio
|
||||
tokens.append(f"<snac_{c0[0, i].item() + OFFSET_L0}>")
|
||||
tokens.append(f"<snac_{c1[0, i1a].item() + OFFSET_L1A}>")
|
||||
tokens.append(f"<snac_{c1[0, i1b].item() + OFFSET_L1B}>")
|
||||
return tokens
|
||||
|
||||
|
||||
def encode_speak(audio: np.ndarray, model, device: str) -> list[str]:
|
||||
"""
|
||||
Encode a float32 audio array with SNAC_24kHz, return full speak-format
|
||||
tokens (7 tokens per base frame, +133% vs encode_listen()) -- 2026-07-22
|
||||
(REPORT.md #37), decided after a real audio A/B
|
||||
(tools/snac_l2_experiment.py) showed audibly better reconstruction.
|
||||
|
||||
Speak format (2026-07-23, Leo-matched order -- L2 interleaved between
|
||||
L1a and L1b, not appended after):
|
||||
<snac_{codes[0][i] + OFFSET_L0}>
|
||||
<snac_{codes[1][2i] + OFFSET_L1A}>
|
||||
<snac_{codes[2][4i] + OFFSET_L2[0]}>
|
||||
<snac_{codes[2][4i+1] + OFFSET_L2[1]}>
|
||||
<snac_{codes[1][2i+1] + OFFSET_L1B}>
|
||||
<snac_{codes[2][4i+2] + OFFSET_L2[2]}>
|
||||
<snac_{codes[2][4i+3] + OFFSET_L2[3]}>
|
||||
"""
|
||||
tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device)
|
||||
with torch.inference_mode():
|
||||
codes = model.encode(tensor)
|
||||
|
||||
c0, c1, c2 = codes[0], codes[1], codes[2]
|
||||
n0 = c0.shape[1]
|
||||
tokens: list[str] = []
|
||||
for i in range(n0):
|
||||
i1a, i1b = 2 * i, 2 * i + 1
|
||||
i2 = [4 * i + k for k in range(4)]
|
||||
if i1b >= c1.shape[1] or i2[-1] >= c2.shape[1]:
|
||||
break
|
||||
tokens.append(f"<snac_{c0[0, i].item() + OFFSET_L0}>")
|
||||
tokens.append(f"<snac_{c1[0, i1a].item() + OFFSET_L1A}>")
|
||||
tokens.append(f"<snac_{c2[0, i2[0]].item() + OFFSET_L2[0]}>")
|
||||
tokens.append(f"<snac_{c2[0, i2[1]].item() + OFFSET_L2[1]}>")
|
||||
tokens.append(f"<snac_{c1[0, i1b].item() + OFFSET_L1B}>")
|
||||
tokens.append(f"<snac_{c2[0, i2[2]].item() + OFFSET_L2[2]}>")
|
||||
tokens.append(f"<snac_{c2[0, i2[3]].item() + OFFSET_L2[3]}>")
|
||||
return tokens
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Task list building (pre-processing step)
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _scan_one_rank_file(fpath: str) -> dict:
|
||||
"""
|
||||
Scan one final_dataset_adaptive rank file for activity time boundaries.
|
||||
Returns {video_id: [activity_dict, ...]} for ALL activities with a valid
|
||||
time_range_sec.
|
||||
|
||||
Each activity dict: activity_id, start_sec, end_sec, has_agent.
|
||||
|
||||
2026-07-23: no longer reads chunk_timing (this file is the pre-window=24-
|
||||
pivot merge; its chunk_timing reflects the OLD 8-frame grid). n_chunks is
|
||||
instead recomputed in process_video() from (end_sec-start_sec) and the
|
||||
current CHUNK_SIZE=24 -- same independent-recompute pattern
|
||||
data_prep/omnivideo_100k/snac_omnivideo.py uses, so this script no longer
|
||||
needs to wait on Phase 6 merge to re-run at window=24 before it can align
|
||||
correctly.
|
||||
|
||||
We tokenize ALL activities (not just agent) because:
|
||||
- Non-agent activities have seed2+cosmos → seed2+cosmos+snac trains modality transitions
|
||||
- Agent-only activities = only 14% of total; skipping the rest wastes 86% of this GPU run
|
||||
"""
|
||||
tasks: dict = {}
|
||||
try:
|
||||
with open(fpath, "r", errors="replace") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
rec = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
vid = rec.get("video_id", "")
|
||||
if not vid:
|
||||
continue
|
||||
for scene in rec.get("scenes", []):
|
||||
for act in scene.get("activities", []):
|
||||
tr = act.get("time_range_sec")
|
||||
if not tr or len(tr) < 2:
|
||||
continue
|
||||
has_agent = "<agent>" in act.get("video_tokens", "")
|
||||
tasks.setdefault(vid, []).append({
|
||||
"activity_id": act.get("activity_id", ""),
|
||||
"start_sec": float(tr[0]),
|
||||
"end_sec": float(tr[1]),
|
||||
"has_agent": has_agent,
|
||||
})
|
||||
except Exception as e:
|
||||
log.warning(f"Error scanning {fpath}: {e}")
|
||||
return tasks
|
||||
|
||||
|
||||
def build_task_list(input_glob: str, cache_path: str, workers: int = 8) -> dict:
|
||||
"""
|
||||
Scan all final_dataset_adaptive rank files in parallel to build a task list.
|
||||
Saves result to cache_path as JSON.
|
||||
Returns {video_id: [activity_dicts]}.
|
||||
|
||||
This is I/O-heavy (~657 GB total) — use multiprocessing to parallelize.
|
||||
Estimated wall time: 5–15 min with 8 workers on shared filesystem.
|
||||
"""
|
||||
rank_files = sorted(glob.glob(input_glob))
|
||||
if not rank_files:
|
||||
raise FileNotFoundError(f"No files matched: {input_glob}")
|
||||
log.info(f"Scanning {len(rank_files)} rank files with {workers} workers...")
|
||||
t0 = time.time()
|
||||
|
||||
with multiprocessing.Pool(workers) as pool:
|
||||
partial_results = pool.map(_scan_one_rank_file, rank_files)
|
||||
|
||||
# merge
|
||||
all_tasks: dict = {}
|
||||
for partial in partial_results:
|
||||
for vid, acts in partial.items():
|
||||
all_tasks.setdefault(vid, []).extend(acts)
|
||||
|
||||
# deduplicate activities by activity_id (in case of overlap across ranks)
|
||||
for vid in all_tasks:
|
||||
seen = set()
|
||||
deduped = []
|
||||
for act in all_tasks[vid]:
|
||||
key = act["activity_id"]
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
deduped.append(act)
|
||||
all_tasks[vid] = deduped
|
||||
|
||||
log.info(
|
||||
f"Task list built: {len(all_tasks)} videos, "
|
||||
f"{sum(len(v) for v in all_tasks.values())} activities "
|
||||
f"({time.time()-t0:.0f}s)"
|
||||
)
|
||||
os.makedirs(os.path.dirname(cache_path), exist_ok=True)
|
||||
with open(cache_path, "w") as f:
|
||||
json.dump(all_tasks, f)
|
||||
log.info(f"Saved task list → {cache_path}")
|
||||
return all_tasks
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Chunk alignment
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def split_snac_by_chunks(tokens: list[str], n_chunks: int) -> dict[int, list[str]]:
|
||||
"""
|
||||
Split a flat SNAC listen token list evenly across n_chunks video chunks.
|
||||
|
||||
Why: SNAC rate (12.5 Hz base × 3 tokens = 37.5 tok/s) does not divide evenly
|
||||
by the video chunk rate (30fps / 8 = 3.75 Hz). Encoding the full activity once
|
||||
preserves audio context; then we split by chunk count snapping to 3-token
|
||||
boundaries (one SNAC base frame = 3 listen tokens).
|
||||
|
||||
Per 8-frame chunk at 30fps (0.267s): ~3.33 SNAC base frames → 9–10 listen tokens.
|
||||
"""
|
||||
n_tokens = len(tokens)
|
||||
n_base = n_tokens // 3 # truncate to complete base frames
|
||||
tokens = tokens[:n_base * 3]
|
||||
|
||||
result: dict[int, list[str]] = {}
|
||||
for k in range(n_chunks):
|
||||
start_base = round(k * n_base / n_chunks)
|
||||
end_base = round((k + 1) * n_base / n_chunks)
|
||||
result[k] = tokens[start_base * 3 : end_base * 3]
|
||||
return result
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Per-video processing
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def process_video(
|
||||
video_id: str,
|
||||
activities: list[dict],
|
||||
model,
|
||||
device: str,
|
||||
video_dir: str,
|
||||
output_dir: str,
|
||||
skip_existing: bool,
|
||||
encode_fn=encode_listen,
|
||||
) -> dict:
|
||||
"""
|
||||
Tokenize all activities for one video.
|
||||
|
||||
Steps:
|
||||
1. Check skip — if output file exists and skip_existing, return immediately.
|
||||
2. Extract full audio from .mp4 once (1 ffmpeg call per video).
|
||||
3. For each activity: slice audio by time_range_sec, run SNAC encode once,
|
||||
then split the flat token list across chunks (preserving audio context).
|
||||
4. Write all results to {output_dir}/{video_id}_snac.jsonl.
|
||||
|
||||
Output per activity: snac_by_chunk {chunk_idx → [tokens]}
|
||||
Phase6 merge uses this directly to inject SNAC tokens per 24-frame chunk,
|
||||
aligned with the cosmos/avclm/agent tokens that fire at the same chunk
|
||||
boundaries (n_chunks recomputed from CHUNK_SIZE, not read from a file).
|
||||
|
||||
Returns stats: {ok, skipped_vid, failed_audio, failed_snac, tokens}
|
||||
"""
|
||||
out_path = os.path.join(output_dir, f"{video_id}_snac.jsonl")
|
||||
if skip_existing and os.path.exists(out_path):
|
||||
return {"ok": 0, "skipped_vid": len(activities), "failed_audio": 0,
|
||||
"failed_snac": 0, "tokens": 0}
|
||||
|
||||
video_path = os.path.join(video_dir, f"{video_id}.mp4")
|
||||
if not os.path.exists(video_path):
|
||||
return {"ok": 0, "skipped_vid": 0, "failed_audio": len(activities),
|
||||
"failed_snac": 0, "tokens": 0}
|
||||
|
||||
# Extract full audio once
|
||||
full_audio = extract_full_audio(video_path)
|
||||
if full_audio is None:
|
||||
log.warning(f"No audio: {video_path}")
|
||||
return {"ok": 0, "skipped_vid": 0, "failed_audio": len(activities),
|
||||
"failed_snac": 0, "tokens": 0}
|
||||
|
||||
stats = {"ok": 0, "skipped_vid": 0, "failed_audio": 0, "failed_snac": 0, "tokens": 0}
|
||||
rows = []
|
||||
|
||||
for act in activities:
|
||||
segment = slice_audio(full_audio, act["start_sec"], act["end_sec"])
|
||||
if len(segment) < int(SAMPLE_RATE * 0.1): # skip segments < 100 ms
|
||||
stats["failed_audio"] += 1
|
||||
continue
|
||||
try:
|
||||
flat_tokens = encode_fn(segment, model, device)
|
||||
except Exception as e:
|
||||
log.warning(f"SNAC failed {video_id}/{act['activity_id']}: {e}")
|
||||
stats["failed_snac"] += 1
|
||||
continue
|
||||
if not flat_tokens:
|
||||
stats["failed_snac"] += 1
|
||||
continue
|
||||
|
||||
# Split flat token list into per-chunk dicts, aligned to the same
|
||||
# 24-frame grid as cosmos/avclm/agent. n_chunks recomputed
|
||||
# independently from CHUNK_SIZE (not from a merged file's
|
||||
# chunk_timing -- see _scan_one_rank_file docstring).
|
||||
total_frames = max(1, round((act["end_sec"] - act["start_sec"]) * TARGET_FPS))
|
||||
n_chunks = math.ceil(total_frames / CHUNK_SIZE)
|
||||
by_chunk = split_snac_by_chunks(flat_tokens, n_chunks)
|
||||
snac_by_chunk = {str(k): v for k, v in by_chunk.items()}
|
||||
|
||||
rows.append({
|
||||
"video_id": video_id,
|
||||
"activity_id": act["activity_id"],
|
||||
"start_sec": round(act["start_sec"], 4),
|
||||
"end_sec": round(act["end_sec"], 4),
|
||||
"has_agent": act.get("has_agent", False),
|
||||
"snac_by_chunk": snac_by_chunk,
|
||||
})
|
||||
stats["ok"] += 1
|
||||
stats["tokens"] += len(flat_tokens)
|
||||
|
||||
if rows:
|
||||
with open(out_path, "w") as f:
|
||||
for row in rows:
|
||||
f.write(json.dumps(row) + "\n")
|
||||
|
||||
return stats
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
# Main
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description="SNAC tokenization for FineVideo-VLA")
|
||||
p.add_argument("--build-tasks", action="store_true",
|
||||
help="Scan final_dataset_adaptive and write snac_task_list.json, then exit.")
|
||||
p.add_argument("--input-glob", default=INPUT_GLOB)
|
||||
p.add_argument("--output-dir", default=OUTPUT_DIR)
|
||||
p.add_argument("--video-dir", default=VIDEO_DIR)
|
||||
p.add_argument("--task-cache", default=TASK_CACHE)
|
||||
p.add_argument("--hf-cache", default=HF_CACHE)
|
||||
p.add_argument("--scan-workers", type=int, default=8,
|
||||
help="CPU workers for --build-tasks scan (default 8)")
|
||||
p.add_argument("--no-skip", action="store_true",
|
||||
help="Re-process videos even if output file exists")
|
||||
p.add_argument("--format", choices=["listen", "speak"], default="listen",
|
||||
help="listen = L0+L1 only (current production, 3 tok/base-frame); "
|
||||
"speak = full L0+L1+L2 (2026-07-22, 7 tok/base-frame, +133%% tokens). "
|
||||
"speak requires L2 tokens added to the tokenizer vocab first.")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
skip_existing = not args.no_skip
|
||||
|
||||
# ── Set HF cache ─────────────────────────────────────────────────────────
|
||||
os.environ.setdefault("HF_HOME", args.hf_cache)
|
||||
os.makedirs(args.hf_cache, exist_ok=True)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# ── Mode: build task list ─────────────────────────────────────────────────
|
||||
if args.build_tasks:
|
||||
build_task_list(args.input_glob, args.task_cache, workers=args.scan_workers)
|
||||
return
|
||||
|
||||
# ── Mode: tokenize ────────────────────────────────────────────────────────
|
||||
|
||||
# SLURM array vars
|
||||
task_id = int(os.environ.get("SLURM_ARRAY_TASK_ID", "0"))
|
||||
num_tasks = int(os.environ.get("SLURM_ARRAY_TASK_COUNT", "1"))
|
||||
|
||||
# Load task list (must exist — run --build-tasks first)
|
||||
if not os.path.exists(args.task_cache):
|
||||
log.error(
|
||||
f"Task list not found: {args.task_cache}\n"
|
||||
f"Run first: python pipeline_pose/snac_finevideo.py --build-tasks"
|
||||
)
|
||||
sys.exit(1)
|
||||
with open(args.task_cache) as f:
|
||||
all_tasks = json.load(f)
|
||||
|
||||
all_vids = sorted(all_tasks.keys())
|
||||
my_vids = all_vids[task_id::num_tasks]
|
||||
log.info(
|
||||
f"Task {task_id}/{num_tasks}: {len(my_vids)}/{len(all_vids)} videos "
|
||||
f"skip_existing={skip_existing}"
|
||||
)
|
||||
|
||||
# ── Load SNAC model ───────────────────────────────────────────────────────
|
||||
from snac import SNAC # imported here to avoid load cost during --build-tasks
|
||||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||||
log.info(f"Loading SNAC model ({SNAC_MODEL}) on {device}...")
|
||||
t_load = time.time()
|
||||
os.environ["HF_HOME"] = args.hf_cache # ensure hub/ subdir is found
|
||||
model = SNAC.from_pretrained(SNAC_MODEL,
|
||||
local_files_only=True).eval().to(device)
|
||||
log.info(f"SNAC loaded ({time.time()-t_load:.1f}s)")
|
||||
|
||||
# ── Process videos ────────────────────────────────────────────────────────
|
||||
cumul = {"ok": 0, "skipped_vid": 0, "failed_audio": 0, "failed_snac": 0, "tokens": 0}
|
||||
t_start = time.time()
|
||||
|
||||
encode_fn = encode_speak if args.format == "speak" else encode_listen
|
||||
|
||||
for idx, vid in enumerate(my_vids, 1):
|
||||
s = process_video(
|
||||
vid, all_tasks[vid], model, device,
|
||||
args.video_dir, args.output_dir, skip_existing,
|
||||
encode_fn=encode_fn,
|
||||
)
|
||||
for k in cumul:
|
||||
cumul[k] += s[k]
|
||||
|
||||
if idx % 100 == 0 or idx == len(my_vids):
|
||||
elapsed = time.time() - t_start
|
||||
rate = idx / elapsed
|
||||
eta = (len(my_vids) - idx) / rate if rate > 0 else 0
|
||||
log.info(
|
||||
f"[{idx:5d}/{len(my_vids)}] vid={vid} "
|
||||
f"ok={s['ok']} skip={s['skipped_vid']} "
|
||||
f"fail_audio={s['failed_audio']} fail_snac={s['failed_snac']} "
|
||||
f"rate={rate:.1f}vid/s ETA={eta/60:.0f}m "
|
||||
f"total_tokens={cumul['tokens']:,}"
|
||||
)
|
||||
|
||||
elapsed = time.time() - t_start
|
||||
log.info(
|
||||
f"DONE task {task_id}: "
|
||||
f"ok={cumul['ok']} skipped={cumul['skipped_vid']} "
|
||||
f"fail_audio={cumul['failed_audio']} fail_snac={cumul['failed_snac']} "
|
||||
f"tokens={cumul['tokens']:,} wall={elapsed:.0f}s"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user